Table of Contents
Fetching ...

Attention Is Not All You Need: The Importance of Feedforward Networks in Transformer Models

Isaac Gerber

TL;DR

This work interrogates the role of feedforward networks (FFNs) inside decoder-only transformer blocks during pretraining, testing variants with 0–3 linear layers per FFN. By comparing architectures with fixed parameter budgets across Booksum and Wikitext-103, the authors demonstrate that three-layer FFNs can yield lower cross-entropy losses and faster training than the conventional two-layer FFNs, particularly when combined with smaller block counts; however, larger configurations can risk overfitting. The study finds that FFN depth impacts training stability and performance, suggesting a parameter-efficiency trade-off and signaling that FFN design deserves more attention in large-scale pretraining. Overall, the results advocate for exploring larger FFNs or different FFN-depth configurations to improve pretraining efficiency and model quality in decoder-only transformers.

Abstract

Decoder-only transformer networks have become incredibly popular for language modeling tasks. State-of-the-art models can have over a hundred transformer blocks, containing billions of trainable parameters, and are trained on trillions of tokens of text. Each transformer block typically consists of a multi-head attention (MHA) mechanism and a two-layer fully connected feedforward network (FFN). In this paper, we examine the importance of the FFN during the model pre-training process through a series of experiments, confirming that the FFN is important to model performance. Furthermore, we show that models using a transformer block configuration with three-layer FFNs with fewer such blocks outperform the standard two-layer configuration delivering lower training loss with fewer total parameters in less time.

Attention Is Not All You Need: The Importance of Feedforward Networks in Transformer Models

TL;DR

This work interrogates the role of feedforward networks (FFNs) inside decoder-only transformer blocks during pretraining, testing variants with 0–3 linear layers per FFN. By comparing architectures with fixed parameter budgets across Booksum and Wikitext-103, the authors demonstrate that three-layer FFNs can yield lower cross-entropy losses and faster training than the conventional two-layer FFNs, particularly when combined with smaller block counts; however, larger configurations can risk overfitting. The study finds that FFN depth impacts training stability and performance, suggesting a parameter-efficiency trade-off and signaling that FFN design deserves more attention in large-scale pretraining. Overall, the results advocate for exploring larger FFNs or different FFN-depth configurations to improve pretraining efficiency and model quality in decoder-only transformers.

Abstract

Decoder-only transformer networks have become incredibly popular for language modeling tasks. State-of-the-art models can have over a hundred transformer blocks, containing billions of trainable parameters, and are trained on trillions of tokens of text. Each transformer block typically consists of a multi-head attention (MHA) mechanism and a two-layer fully connected feedforward network (FFN). In this paper, we examine the importance of the FFN during the model pre-training process through a series of experiments, confirming that the FFN is important to model performance. Furthermore, we show that models using a transformer block configuration with three-layer FFNs with fewer such blocks outperform the standard two-layer configuration delivering lower training loss with fewer total parameters in less time.
Paper Structure (14 sections, 2 figures, 3 tables)

This paper contains 14 sections, 2 figures, 3 tables.

Figures (2)

  • Figure 1: Transformer block architecture. A depicts a standard transformer block with two linear layers, B a transformer block with three linear layers, C a transformer block with a single linear layer, and D a transformer block with zero linear layers. The arrows at the bottom and top of each block indicate input and output from the previous and to the next sections of the models, respectively. In the linear layers, "d" refers to the model dimension size, $d_{model}$. A transformer model stacks many transformer blocks.
  • Figure 2: Training loss on the Booksum data set, log scale.