Chapter 2.5 β€” The Feed-Forward Layer

Contents

  1. The sublayer that gets skipped
  2. What the FFN computes, and how much of the model it is
  3. From ReLU to GELU: smoothing the gate
  4. Gated Linear Units: a third matrix instead of a better curve
  5. SwiGLU and the two-thirds rule
  6. Interview angle
  7. Self-check questions
  8. Sources

1. The sublayer that gets skipped

Every chapter of Part II so far has been about attention or about what attention needs in order to work: positional information because attention is permutation-invariant (Chapter 2.2), a choice of stack shape because attention can be masked or not (Chapter 2.3), normalization because deep residual stacks of attention sublayers don't train otherwise (Chapter 2.4). Meanwhile the other sublayer in every Transformer block β€” the position-wise feed-forward network β€” has appeared in this book roughly a dozen times as the same throwaway phrase: "a small two-layer MLP with a nonlinearity in between, run token by token." That phrase is accurate. It is also, in a model like LLaMA-2 70B, describing about two-thirds of the parameters.

This is a genuinely common blind spot, and interviewers know it. A candidate who can derive the $\sqrt{d_k}$ scaling factor from first principles and then can't say what activation function their production model's FFN uses, or why the hidden dimension is $11008$ rather than a round number, has learned the glamorous half of the architecture. This chapter fills in the other half: what the FFN is actually doing, why the field moved from ReLU to GELU and then to gated variants, and what the concrete design decision β€” which activation, at which width β€” looks like when you're the one configuring a model. It's also the last piece of the block: once this chapter is done, every component of a modern Transformer layer has been examined, and Part III can turn to scaling that layer up.

2. What the FFN computes, and how much of the model it is

The original formulation from Vaswani et al. is two linear projections with a nonlinearity between them, applied independently to each position's vector: $$\text{FFN}(x) = \max(0,\, xW_1 + b_1)\,W_2 + b_2$$ where $x \in \mathbb{R}^{d}$ is a single token's representation, $W_1 \in \mathbb{R}^{d \times d_{ff}}$ projects it up into a wider intermediate space, the nonlinearity (ReLU, in the original paper) is applied there, and $W_2 \in \mathbb{R}^{d_{ff} \times d}$ projects back down to the model dimension so the result can be added onto the residual stream. The original Transformer used $d = 512$ and $d_{ff} = 2048$, establishing the $d_{ff} = 4d$ convention that persisted essentially unchanged for years. The word "position-wise" is load-bearing: the same $W_1$ and $W_2$ are applied to every token independently, with no mixing across positions whatsoever. All communication between tokens happens in the attention sublayer; the FFN is where each token, individually, does computation on what it just gathered.

That division of labor is the cleanest way to think about the Transformer block, and it makes the parameter accounting worth doing explicitly. The attention sublayer holds four $d \times d$ matrices β€” $W_Q$, $W_K$, $W_V$, $W_O$ β€” for $4d^2$ parameters. The FFN holds $d \times 4d$ and $4d \times d$, for $8d^2$. The FFN is therefore twice the attention sublayer, or two-thirds of every block, and since the FFN is also applied at every position, it accounts for roughly the same two-thirds share of the forward-pass FLOPs at short context lengths (attention's quadratic term only overtakes it once sequences get long, which is exactly the crossover Chapter 3.1 is about). Any claim about "where the parameters are" in a dense LLM is mostly a claim about the FFN.

Geva et al. gave this sublayer a much more specific interpretation in 2021, arguing that Transformer feed-forward layers operate as key-value memories: each column of $W_1$ acts as a key that fires on a particular input pattern (their analysis finds individual columns that activate on recognizable textual patterns, from surface-level n-grams in early layers to semantic topics in later ones), and the corresponding row of $W_2$ is the value β€” the vector that column writes into the residual stream when it fires. Under this reading the FFN isn't a generic "extra capacity" blob; it's an associative memory, with $d_{ff}$ slots, that the attention sublayer queries indirectly by assembling the right input pattern. That interpretation is also the conceptual bridge to Chapter 3.3: Mixture of Experts replaces this one dense FFN with many, and routes each token to a few of them β€” which only makes sense as a design if you already believe the FFN is a store of many specialized, individually-addressable pieces of knowledge rather than one monolithic function.

This framing also answers a question worth being able to give a real answer to: if you want a model with more capacity, why widen this one sublayer ($d_{ff}$) rather than simply stacking more layers? The two levers aren't interchangeable in cost. Depth is inherently sequential β€” each additional layer is one more hop a token's representation has to pass through before the next layer can even start, so depth adds inference latency directly and makes the optimization landscape harder to keep stable (Chapter 2.4's entire discussion of Pre-LN and normalization is really about taming a difficulty that comes specifically from depth). Width, by contrast, is parallel: widening $d_{ff}$ just makes the two matrix multiplications inside one FFN call bigger, which modern accelerators absorb without adding a single additional sequential step, since a token's forward pass through that one wider FFN still happens in the same one hop through the block. Under Geva et al.'s key-value-memory reading above, this makes concrete sense too: widening $d_{ff}$ just adds more key-value slots to the same one memory a token queries once per layer, while adding a whole extra layer inserts an entire additional attention-then-FFN computation the token's representation has to be pushed through sequentially, on top of everything a full extra layer costs in parameters for attention as well. For pure storage capacity in the sense Geva et al. describe, more slots in one memory is a cheaper, more parallel way to buy the same thing than another whole hop of sequential computation β€” which is exactly why $d_{ff}$ is the lever real model configs pull hardest for capacity, once a reasonable minimum depth is already in place.

3. From ReLU to GELU: smoothing the gate

ReLU, $\max(0, x)$, is a hard gate: an input is either passed through unchanged or zeroed, with the decision made by a discontinuous switch at zero. Hendrycks and Gimpel's Gaussian Error Linear Unit reframes that gate probabilistically. Rather than gating on the sign of $x$, gate on how large $x$ is relative to the distribution of activations, by multiplying the input by the probability that a standard normal variable falls below it: $$\text{GELU}(x) = x \cdot \Phi(x)$$ where $\Phi$ is the standard normal CDF. The result is a smooth curve that behaves like ReLU far from zero (for large positive $x$, $\Phi(x) \approx 1$ and GELU$(x) \approx x$; for large negative $x$, $\Phi(x) \approx 0$ and the output vanishes) but transitions gradually through the origin β€” and, notably, is non-monotonic: it dips slightly negative for small negative inputs before returning to zero, rather than clamping flat the way ReLU does.

Two properties are usually credited for GELU's empirical edge. The first is that it's differentiable everywhere, with no kink at the origin, which gives gradient descent a better-conditioned surface to work on than ReLU's discontinuous derivative. The second is that ReLU's exact zeroing of all negative inputs produces genuinely dead units β€” a unit whose pre-activation is negative for every input in the data receives no gradient at all and can never recover β€” whereas GELU's small negative tail keeps a gradient signal alive on the negative side. GELU was adopted by BERT and by the entire GPT line through GPT-3, and remains a perfectly reasonable default; in practice implementations often use the tanh approximation $\text{GELU}(x) \approx 0.5x\left(1 + \tanh\left[\sqrt{2/\pi}\,(x + 0.044715x^3)\right]\right)$, which is cheaper than evaluating the true CDF and numerically indistinguishable in a trained network.

Ramachandran, Zoph, and Le arrived at a very similar function from an entirely different direction, running an automated search over the space of activation functions and finding that the best performer their search discovered was $\text{Swish}(x) = x \cdot \sigma(\beta x)$, the input times a sigmoid of itself. With $\beta = 1$ this is often called SiLU, and its shape is nearly identical to GELU's β€” smooth, non-monotonic, ReLU-like in the tails. That two independent lines of work, one analytical and one an automated search, converged on essentially the same curve is decent evidence that the shape itself is what matters, not the specific derivation. It's worth holding this loosely, though: the measured gaps between ReLU, GELU, and Swish in a well-tuned large model are real but small, and none of them is the kind of change that transforms a model. The next one is more structural.

4. Gated Linear Units: a third matrix instead of a better curve

Every function in the previous section is a pointwise reshaping of the same quantity: take $xW_1$, bend it through some curve, project back down. Dauphin et al., working on convolutional language models in 2017, proposed something categorically different β€” instead of choosing a better curve to apply to one projection, compute two projections and let one of them gate the other multiplicatively: $$\text{GLU}(x) = (xW + b) \odot \sigma(xV + c)$$ where $\odot$ is elementwise multiplication. The first branch carries the content; the second branch, squashed through a sigmoid into $[0,1]$, decides how much of each content dimension survives. The gate is now a learned, input-dependent function computed by its own weight matrix, rather than a fixed curve applied to the content itself.

The distinction matters more than it looks. A pointwise activation, however cleverly shaped, computes each output coordinate from exactly one input coordinate β€” the function $\mathbb{R} \to \mathbb{R}$ is the same for every unit and every token. A gated unit computes each output coordinate as a product of two different linear functions of the whole input vector, which makes the layer's output quadratic in $x$ rather than linear-then-bent. That multiplicative interaction is a genuinely richer primitive: the network can learn a gate that suppresses one feature specifically when some unrelated feature is present, which no pointwise nonlinearity can express at all. It's the same structural idea that made LSTM gates work in Chapter 1.4 β€” multiplicative, data-dependent control of information flow β€” reappearing in a feed-forward setting.

Shazeer's 2020 note GLU Variants Improve Transformer carried the idea into the Transformer FFN and, in the process, defined the family that modern models actually use. The move is to replace GLU's sigmoid with whichever activation you like, giving a named variant for each choice: $$\begin{aligned}\text{ReGLU}(x) &= \max(0,\,xW) \odot xV \ \text{GEGLU}(x) &= \text{GELU}(xW) \odot xV \ \text{SwiGLU}(x) &= \text{Swish}(xW) \odot xV\end{aligned}$$ with the FFN then finishing as before by projecting the gated result back down through a third matrix $W_2$. Shazeer evaluated these as drop-in replacements in a T5 span-filling pretraining setup and found the gated variants beating both plain ReLU and plain GELU on pretraining loss and on downstream fine-tuning, with GEGLU and SwiGLU roughly tied at the front. The paper is refreshingly honest that this is an empirical result without a mechanistic story behind it, closing with a line that has been quoted in a great many subsequent papers: the authors offer no explanation for why the architectures work and attribute their success, as all else, to divine benevolence.

5. SwiGLU and the two-thirds rule

There's an accounting problem in the previous section that has to be dealt with before any of this is a fair comparison. A standard FFN has two matrices, $W_1$ and $W_2$, for $2 \cdot d \cdot d_{ff}$ parameters. A gated FFN has three β€” $W$, $V$, and $W_2$ β€” for $3 \cdot d \cdot d_{ff}$. At the same $d_{ff}$, the gated version is 50% larger, and "a bigger layer performs better" is not an interesting finding. Shazeer's experiments controlled for this by shrinking the hidden dimension by a factor of $2/3$, so that $3 \cdot d \cdot \frac{2}{3}d_{ff} = 2 \cdot d \cdot d_{ff}$ and the parameter count matches exactly. With the usual $d_{ff} = 4d$ starting point, this gives the rule you actually see in modern configs: $$d_{ff} = \tfrac{2}{3} \cdot 4d = \tfrac{8}{3}d$$ The gated variants win at equal parameter count and equal FLOPs, which is what makes the result worth acting on.

This is why open-weight model configs have the strange-looking FFN widths they do. LLaMA-7B has $d = 4096$; $\frac{8}{3} \cdot 4096 = 10922.67$, which LLaMA rounds up to $11008$ β€” a multiple of $256$, chosen because hardware kernels and tensor-parallel sharding (Chapter 4.4) both prefer dimensions that divide cleanly. If you have ever looked at a config file and wondered why intermediate_size isn't simply four times hidden_size, this is the whole answer: it's $8/3$, rounded to something a GPU likes. Modern implementations also typically drop the bias terms $b_1, b_2$ entirely β€” PaLM among others found that removing biases from all dense layers improved training stability at scale with no quality cost β€” so the production FFN is usually three bias-free matrices and a Swish gate.

SwiGLU is now close to universal in large decoder-only models: LLaMA and its many descendants, PaLM, Mistral, Qwen, and most of the open-weight field use it, while the older GELU-with-$4d$ configuration survives mainly in models descended from the GPT-2/GPT-3 line and in the encoder-only BERT family. GEGLU, despite roughly tying SwiGLU in Shazeer's own numbers, sees far less adoption in practice β€” Google's Gemma family is the most visible production user, which is as much a reflection of SwiGLU becoming the field's default convention to reach for as it is of any measured quality gap between the two. It's worth being clear-eyed about the size of the effect: this is a low-single-digit improvement in loss at fixed compute, not a step change, and it survives because it's free β€” same parameters, same FLOPs, one extra elementwise multiply. That's the characteristic shape of a mature-architecture improvement, and it's a useful contrast with what Part III covers next, where the gains come from changing the asymptotic cost of the computation rather than from finding a slightly better function inside a fixed budget.

With this, Part II has assembled the complete modern Transformer block: attention to move information between positions, positional encoding to tell it where those positions are, a decoder-only stack shape, normalization placed to keep the residual path clean, and a gated feed-forward sublayer holding most of the parameters. Everything from here on is about scale β€” making this block cheaper over long sequences, giving it far more parameters than it can afford to use densely, and serving it to real users β€” which is exactly what Part III, "Scaling the Architecture," takes up.

6. Interview angle

A senior LLM interview will rarely ask "what is a feed-forward layer" outright, but it will probe whether you know where a model's parameters and FLOPs actually go, usually through questions like:

  • "Roughly what fraction of a dense Transformer's parameters are in the FFN sublayers, and why?" A strong answer does the arithmetic on the spot: attention is four $d \times d$ projections, or $4d^2$; the FFN at $d_{ff} = 4d$ is two $d \times 4d$ matrices, or $8d^2$; so the FFN is two-thirds of each block. A strong answer also notes that this makes the FFN roughly two-thirds of the per-token FLOPs at short context, with attention's quadratic term only dominating at long sequence lengths, and connects it to why MoE (Chapter 3.3) targets the FFN specifically rather than attention.
  • "Why is the FFN hidden dimension in LLaMA $11008$ instead of $16384$?" The expected answer is the $2/3$ rule: SwiGLU uses three matrices instead of two, so the hidden dimension is scaled by $2/3$ to hold parameters and FLOPs constant, giving $\frac{8}{3}d$; $\frac{8}{3} \cdot 4096 \approx 10923$, rounded up to $11008$ for hardware-friendly divisibility. Getting this right signals that you've read a real config file and understood it, not just a paper.
  • "What does a gated activation give you that a better pointwise nonlinearity can't?" The answer should be about the multiplicative interaction specifically: a pointwise function maps each coordinate independently through a fixed curve, while a gate is a second learned linear function of the entire input vector multiplying the first, letting the layer suppress a feature conditionally on unrelated features. Naming the parallel to LSTM gating shows you see it as the same idea rather than a new trick, and honestly noting that the empirical gain is small and not mechanistically explained is a point in your favor, not against you.
  • "Why did the field move off ReLU?" A strong answer names both properties β€” everywhere-differentiable versus a kink at zero, and a small negative tail that avoids ReLU's permanently dead units β€” and identifies GELU and Swish as near-identical curves reached by independent routes (an analytical probabilistic-gating argument and an automated function search), which is evidence the shape matters more than any particular derivation.

The through-line interviewers are checking for is whether your mental model of a Transformer is the whole block or just the attention diagram. Anyone who has read the attention paper can talk about $QK^\top$; only someone who has actually configured, profiled, or trained a model knows that most of the parameters they were paying for were sitting in the sublayer nobody draws.

7. Self-check questions

  1. Write out the original FFN formula and explain what "position-wise" means mechanically β€” what would break if the FFN mixed information across positions?
  2. Derive the parameter counts of the attention and feed-forward sublayers in terms of $d$, and use them to state what fraction of a dense block's parameters the FFN holds. How does that fraction change if $d_{ff} = 4d$ becomes $d_{ff} = 8d$?
  3. Explain the key-value memory interpretation of the FFN, and why that interpretation makes Mixture of Experts a natural next step rather than an arbitrary one.
  4. What two properties does GELU have that ReLU lacks, and what concrete failure mode of ReLU does the second one address?
  5. Explain precisely why a gated unit is more expressive than any pointwise activation, in terms of what each can compute from the input vector.
  6. A colleague benchmarks SwiGLU against GELU at the same $d_{ff}$ and reports that SwiGLU is better. What's wrong with that comparison, and what would you tell them to change?
  7. Looking back over Part II, name the one problem each chapter's central mechanism solves, and explain why the FFN's role only makes sense once you know what attention does and doesn't do.

8. Sources

  • Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762. https://arxiv.org/abs/1706.03762
  • Hendrycks, D., & Gimpel, K. (2016). Gaussian Error Linear Units (GELUs). arXiv:1606.08415. https://arxiv.org/abs/1606.08415
  • Ramachandran, P., Zoph, B., & Le, Q. V. (2017). Searching for Activation Functions. arXiv:1710.05941. https://arxiv.org/abs/1710.05941
  • Dauphin, Y. N., Fan, A., Auli, M., & Grangier, D. (2017). Language Modeling with Gated Convolutional Networks. ICML 2017. arXiv:1612.08083. https://arxiv.org/abs/1612.08083
  • Shazeer, N. (2020). GLU Variants Improve Transformer. arXiv:2002.05202. https://arxiv.org/abs/2002.05202
  • Geva, M., Schuster, R., Berant, J., & Levy, O. (2021). Transformer Feed-Forward Layers Are Key-Value Memories. EMNLP 2021. arXiv:2012.14913. https://arxiv.org/abs/2012.14913
  • Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., RoziΓ¨re, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., & Lample, G. (2023). LLaMA: Open and Efficient Foundation Language Models. arXiv:2302.13971. https://arxiv.org/abs/2302.13971
  • Chowdhery, A., Narang, S., Devlin, J., et al. (2022). PaLM: Scaling Language Modeling with Pathways. arXiv:2204.02311. https://arxiv.org/abs/2204.02311

Self-check quiz

Test your understanding of this chapter with a short quiz β€” a mix of concept questions and quick calculations.

In a standard dense Transformer block with $d_{ff} = 4d$ and no biases, what fraction of the block's parameters sit in the feed-forward sublayer?

Explanation: Attention holds four $d \times d$ projections ($4d^2$); the FFN holds $d \times 4d$ plus $4d \times d$ ($8d^2$). So the FFN is $8d^2 / 12d^2 = 2/3$ of the block.

A model has $d = 1024$ and a classic two-matrix FFN with $d_{ff} = 4d$ and no bias terms. How many parameters are in one FFN sublayer?

Explanation: $2 \cdot d \cdot d_{ff} = 2 \cdot 1024 \cdot 4096 = 8{,}388{,}608$ β€” twice the $4d^2 = 4{,}194{,}304$ held by the attention sublayer.

What does "position-wise" mean in "position-wise feed-forward network"?

Explanation: One shared $W_1$/$W_2$ pair is applied to every token independently. All cross-position communication happens in the attention sublayer; the FFN does per-token computation on what attention gathered.

Under Geva et al.'s key-value memory interpretation of the FFN, what plays the role of the "value"?

Explanation: Each column of $W_1$ acts as a key detecting an input pattern; the matching row of $W_2$ is the value vector that gets added to the residual stream when that key fires.

Which pair of properties best explains GELU's empirical edge over ReLU?

Explanation: GELU has no kink at the origin (better-conditioned gradients) and dips slightly negative rather than clamping flat, so units with negative pre-activations still receive gradient instead of dying permanently.

What can a gated unit like SwiGLU express that no pointwise activation can?

Explanation: A pointwise activation maps each coordinate independently through one fixed $\mathbb{R} \to \mathbb{R}$ curve. A gate multiplies two separate linear projections of the full input, so the layer can suppress one feature conditionally on unrelated features.

A model with $d = 768$ switches from a two-matrix FFN at $d_{ff} = 4d$ to SwiGLU, applying the two-thirds rule to keep the parameter count fixed. What is the new $d_{ff}$?

Explanation: SwiGLU has three matrices instead of two, so $d_{ff}$ is scaled by $2/3$: $\frac{8}{3} \cdot 768 = 2048$. Check: $3 \cdot 768 \cdot 2048 = 2 \cdot 768 \cdot 3072$.

Why is LLaMA-7B's FFN hidden dimension $11008$ rather than $10923$?

Explanation: The $8/3$ rule gives $10922.67$; it is rounded up to $11008$ because GPU kernels and tensor-parallel sharding (Chapter 4.4) both prefer dimensions that divide cleanly.