Layers of Understanding

In the attention chapter, "sat" learned to pull information from "cat". Good. But attention only moves information between tokens. It is a weighted average of value vectors. It never computes anything new about a token on its own: no "this verb is past tense so the subject probably did something already", no "this looks like the start of a Python function". Mixing is not thinking.

And one layer of anything is shallow. Suppose "it" has now absorbed "animal". To predict the next word the model might need to notice that "the animal … was too tired" is an explanation, which is why "because" appeared, which suggests the sentence is nearly over. Each of those steps depends on the previous one. You cannot do them in a single pass of attention, however many heads it has.

So a transformer needs two more things: a per-token computation that can transform what attention gathered, and a way to stack many rounds of gather-then-transform without the network becoming untrainable. This chapter builds both, then shows what the stack does with them.

The problem: mixing is not thinking

Let's make the limitation of attention precise. The output of an attention head for token $i$ is $\sum_j a_{ij} v_j$, a convex combination of value vectors (the weights are non-negative and sum to one). Whatever the head produces lies inside the "cloud" spanned by the values that already exist in the sequence. If none of the tokens carries the feature "this is a question", attention cannot conjure it. It is a linear operation once the weights are fixed, and it has no nonlinearity of its own except the softmax that chooses the weights.

What we want is a function applied to each token's vector separately, $f(x_i)$, that can compute arbitrary new features from the ones present. Neural-network 101 says: a matrix, a nonlinearity, another matrix. That is exactly what the transformer uses, and it calls it the feed-forward block, or the MLP.

The feed-forward block (MLP)

Expand, squash, project back

The block takes a $d_{model}$-vector, expands it to a wider hidden vector (by a factor of 4 in almost every transformer), applies an element-wise nonlinearity, and projects back down to $d_{model}$. In symbols, with $W_1 \in \mathbb{R}^{d \times 4d}$ and $W_2 \in \mathbb{R}^{4d \times d}$:

The first matrix asks $4d$ different linear questions of the token ("how much does it look like a verb?", "how much like the start of a date?"). GELU keeps the positive answers and quietly suppresses the negative ones. The second matrix turns the surviving answers into an update to the token's vector. The same $W_1, W_2$ are applied to every position, independently, which is why this is sometimes called a position-wise feed-forward network.

Why expand to $4d$? Because the nonlinearity is what gives the block its power, and it acts element-wise: more hidden units means more independent yes/no features the block can detect. Four is not sacred (it was in the original paper and nobody found a strong reason to change it), but the expansion has to be substantial for the block to be worth its cost. It is also where most of the parameters live: $2 \times d \times 4d = 8d^2$ per block versus $4d^2$ for attention.

A worked tiny example

To keep the arithmetic short we use $d = 2$ and a hidden width of 4 (a $2\times$ expansion instead of $4\times$; nothing else changes). Let the input be $x = [1, -1]$ and

First the expansion. Each hidden unit is a dot product of $x$ with a column of $W_1$:

Now GELU, element by element: $\text{GELU}(1) \approx 0.84$, $\text{GELU}(-1) \approx -0.16$, $\text{GELU}(-2) \approx -0.05$, $\text{GELU}(0) = 0$. So the hidden activation is $h = [0.84, -0.16, -0.05, 0]$. Unit 1 fired; units 2 and 3 are nearly switched off; unit 4 is exactly off.

Finally project back:

That is the MLP's write: a small vector that will be added to $x$ by the residual connection, giving $[1.79, -1.21]$. Notice the block did something attention cannot: it detected a feature (unit 1 says "first coordinate positive and second not") and produced a new direction from it.

What the MLP is thought to do

A useful way to read the two matrices: each row of $W_1^\top$ (a column of $W_1$) is a key pattern that the input is matched against, and each row of $W_2$ is the value that gets written when that key fires. Geva et al. (2021), "Transformer Feed-Forward Layers Are Key-Value Memories", showed this is not just a metaphor: in trained models, individual hidden units respond to human-recognisable input patterns (a particular topic, a syntactic construction, the end of a sentence) and their corresponding value rows push the prediction towards specific next tokens. The MLP is a big associative memory: $4d$ key–value pairs per block, looked up in parallel, with soft matching.

GELU, and why not ReLU

The original transformer used ReLU, $\max(0, x)$. GPT-2 and BERT switched to GELU (Hendrycks and Gimpel, 2016, arXiv), and Llama-family models use SiLU inside a gated block. Why fuss over the squashing function? Because ReLU has a hard corner at zero: units that are slightly negative get exactly zero output and exactly zero gradient, so they can get stuck. GELU is a smoothed ReLU:

where $\Phi$ is the standard normal cumulative distribution function. Read it as "$x$, times the probability that a standard normal is less than $x$". For large positive $x$ it is $\approx x$; for large negative $x$ it is $\approx 0$; in between it curves smoothly and even dips slightly below zero (minimum about $-0.17$ at $x \approx -0.75$). SiLU (also called Swish) is the same idea with a sigmoid instead of $\Phi$: $\text{SiLU}(x) = x\,\sigma(x)$. In practice the three give similar results; GELU and SiLU train a little more smoothly.

Residual connections: the stream

The problem: deep networks were untrainable

Now we want to stack: attention, MLP, attention, MLP, dozens of times. Historically this failed. Stack twenty plain layers, $x \leftarrow f_\ell(x)$, and the network trains worse than a ten-layer one, not because it overfits but because the optimiser cannot even fit the training data. The reason is the chain rule: the gradient reaching layer 1 is the product of twenty Jacobians. If each one shrinks the signal a little, the product vanishes; if each one grows it, the product explodes. Either way the early layers stop learning.

He et al. (2016), "Deep Residual Learning for Image Recognition", fixed this with a change so small it fits on a napkin. Instead of replacing $x$, each layer adds to it:

The layer now learns a correction, not a whole new representation. And look at the gradient: $\partial(x + f(x))/\partial x = I + \partial f/\partial x$. That identity matrix is a highway. Even if $\partial f/\partial x$ is tiny, the gradient passes straight through the "$+$" untouched, so the bottom layers receive a clean signal no matter how deep the stack is. At initialisation, when $f$ is small and random, the whole network is close to the identity function, which is an excellent place to start optimising from.

The residual stream as working memory

The residual picture leads to the most useful mental model of a transformer. Think of the token's vector $x$ as a stream that flows upward through the layers. Each attention head and each MLP reads from the stream (through its input projection), computes something, and writes a vector back into it (through its output projection) by addition. Nothing ever erases the stream; later blocks can only add. This framing comes from Elhage et al. (2021), "A Mathematical Framework for Transformer Circuits".

This picture explains several facts at once. Why can you delete a middle layer from a trained transformer and get only a mild degradation? Because the stream survives without that layer's write. Why do attention heads in layer 8 seem to "know" what a layer-2 MLP computed? Because the layer-2 write is still sitting in the stream. And why do models need the width $d_{model}$ to be large? Because the stream is shared: hundreds of heads and MLPs are all writing into the same $d_{model}$ dimensions, and they need room to avoid stepping on each other.

LayerNorm: keeping the stream well-behaved

The problem: scale drift

If every block only adds to the stream, the stream grows. After 24 blocks the vector might have entries in the hundreds, and the next attention layer's dot products would be enormous (undoing the careful $\sqrt{d_k}$ scaling). Some dimensions might be huge and others tiny, so the same learning rate would be far too big for some weights and far too small for others. We need each block to see inputs of a predictable size, regardless of what earlier blocks did.

What LayerNorm does

Layer normalisation (Ba, Kiros and Hinton, 2016, arXiv) standardises each token's vector on its own: subtract its mean, divide by its standard deviation, then apply a learned per-dimension gain $\gamma$ and bias $\beta$:

The mean and variance are taken across the $d$ dimensions of one token, not across the batch and not across positions; that is what makes it work for sequences of any length and batches of size 1. The $\epsilon$ (typically $10^{-5}$) prevents division by zero when a vector is constant. $\gamma$ and $\beta$ let the network undo the normalisation in any dimension where it prefers to; they add $2d$ parameters, which is negligible.

Pre-LN vs post-LN

Where does the normalisation go? The original transformer put it after the residual add: $x \leftarrow \text{LN}(x + f(x))$. This is post-LN. It has a subtle flaw: the LayerNorm sits on the residual highway, so the identity path is no longer clean; every block rescales the stream, and the gradient at the bottom layers ends up much larger than at the top. Training post-LN transformers needs a careful learning-rate warm-up and still becomes unstable when they get deep.

Pre-LN moves the normalisation inside the branch: $x \leftarrow x + f(\text{LN}(x))$. Now the highway is untouched (pure addition), each sub-block sees a normalised input, and gradients are well-behaved at any depth. Xiong et al. (2020), "On Layer Normalization in the Transformer Architecture", analysed this formally and showed pre-LN trains without warm-up. GPT-2 was already pre-LN, and virtually every large model since (GPT-3, Llama, PaLM, Mistral) uses it, with one extra LayerNorm at the very end of the stack before the output projection, because the stream itself is never normalised otherwise.

RMSNorm: the modern simplification

Zhang and Sennrich (2019), "Root Mean Square Layer Normalization", asked which part of LayerNorm actually matters and found that the re-scaling does most of the work while the mean-subtraction adds little. RMSNorm drops the mean and the bias:

It is cheaper (one pass over the vector instead of two) and trains just as well, so Llama, Mistral, Gemma and most recent open models use it. On our $[2,4,6,8]$ example: the root-mean-square is $\sqrt{(4+16+36+64)/4} = \sqrt{30} \approx 5.48$, giving $[0.37, 0.73, 1.10, 1.46]$. Same idea, no centring.

The full transformer block

We now have all the parts. A pre-LN transformer block is two residual updates in a row:

Read the two lines as: "look at the context and write what you found into the stream; then think about what is now in the stream and write your conclusion". That pair, repeated $N$ times, is the entire body of GPT-2, Llama, and every decoder-only language model in use today. The differences between models are in the details (which norm, which activation, GQA or not, how positions are encoded), not in this skeleton.

Stacking N blocks: what the layers do

A single block gives every token one round of "gather, then think". GPT-2 small stacks 12; Llama-3 70B stacks 80. What does the stack actually do with all that depth? We do not have a complete answer, but three findings are robust enough to be worth knowing.

Early, middle and late layers

Early layers deal with surface form. Their attention heads are mostly positional (previous-token, first-token), and their MLPs resolve things like "these three sub-word tokens form the word photosynthesis" and part-of-speech-like features. Tenney et al. (2019), "BERT Rediscovers the Classical NLP Pipeline", found that syntactic information is decodable earliest and semantic roles later. Middle layers are where the richest, most abstract representations live: the entity "it" refers to, the topic, the language, whether the text is code. Induction heads, the "copy what came after this last time" circuit from the attention chapter, tend to sit here. Late layers turn all of that into a concrete next-token prediction: their MLPs act as the key–value memories Geva et al. described, promoting specific vocabulary items, and their attention heads often copy likely continuations from the context.

The logit lens

There is a beautiful trick for watching this happen. The final step of a language model is: take the residual stream, apply the final LayerNorm, multiply by the unembedding matrix, and read off logits over the vocabulary. Nothing stops you from applying that same final step to the stream after layer 3, or layer 7. This is the logit lens, described by nostalgebraist (2020) in "interpreting GPT: the logit lens". In GPT-2, the early-layer predictions are near-nonsense, the middle layers start guessing plausible word classes, and the last few layers converge on the actual output. The stream is gradually "becoming" the prediction. Belrose et al. (2023), "Eliciting Latent Predictions from Transformers with the Tuned Lens", sharpened the tool by training a small affine correction per layer, which makes the intermediate predictions far more legible and shows that a surprising amount of the final answer is already present halfway up.

The stream as the model's working memory

Put the residual picture and the logit lens together and you get the modern understanding of a transformer: the residual stream at each position is a working memory of width $d_{model}$. Early layers write basic facts about the token into it. Attention heads copy facts from other positions' memories into this one. MLPs read the memory, recognise patterns, and write conclusions. Late layers write "and therefore the next token is probably X" into the same memory, and the unembedding reads that out. There is no separate "state"; the stream is the state, and depth is how many times the model gets to update it.

Parameter accounting per block

You should be able to count a transformer's parameters in your head; it is the fastest sanity check on any architecture description. Let $d = d_{model}$ and ignore biases and norms (they are $O(d)$, negligible next to $O(d^2)$).

• Attention: $W_Q, W_K, W_V, W_O$, each $d \times d$: $4d^2$.
• MLP: $W_1$ is $d \times 4d$ and $W_2$ is $4d \times d$: $8d^2$.
• Per block: $12d^2$. Two thirds of every block is the MLP.

Two refinements for modern models. With SwiGLU the MLP has three matrices of size $d \times d_{ff}$; with the usual $d_{ff} \approx \tfrac{8}{3}d$ that is again $\approx 8d^2$, by design. With grouped-query attention, $W_K$ and $W_V$ shrink to $d \times h_{kv} d_{head}$ each, so attention drops below $4d^2$; in Llama-3 70B it is about $2.25d^2$.

Dropout

The original transformer and GPT-2 applied dropout (Srivastava et al., 2014) in several places: on the attention weights, on each sub-block's output before the residual add, and on the embeddings, with rate 0.1. Dropout zeros a random 10% of activations during training and scales the rest up by $1/0.9$, so no single unit can be relied on; at inference it is switched off. It is a regulariser: it fights overfitting.

Modern pre-training usually drops dropout entirely. The reason is that overfitting is not the problem when you train for a single epoch on trillions of tokens; the model never sees the same example twice, so there is nothing to memorise. Dropout then only slows learning and adds noise. Llama, PaLM, Mistral and most models trained since about 2022 use dropout 0.0 in pre-training. It reappears in fine-tuning, where datasets are small and multiple epochs are common. code/lumen/block.py keeps a dropout argument that defaults to 0.

Initialisation and the 1/√(2N) residual trick

Random initial weights matter more in deep networks than in shallow ones, for the same reason depth is hard: effects compound. GPT-2 initialises every weight matrix from $\mathcal{N}(0, 0.02^2)$ and every bias at 0, and then applies one extra rule that is easy to miss but important.

Here is the problem it solves. Each block adds two write vectors to the stream. If each write has variance $\sigma^2$ per dimension and the writes are roughly independent, after $N$ blocks the stream's variance is about $2N\sigma^2$ plus whatever it started with. With $N = 48$ (GPT-2 XL) that is a 96-fold growth: the last layers would see inputs a factor of $\sqrt{96} \approx 10$ larger than the first. Pre-LN protects the sub-blocks' inputs, but the stream itself, and therefore the relative size of each new write, still drifts.

The fix, from Radford et al. (2019), is to scale the initialisation of the two output projections in each block, $W_O$ and $W_2$ (the matrices that write into the stream), by $1/\sqrt{2N}$:

Now each of the $2N$ writes has variance $\sigma^2/(2N)$, their sum has variance about $\sigma^2$, and the stream stays the same size from the first block to the last. It costs nothing and it is in every GPT-style codebase, usually as a special case in the init loop that checks whether a parameter's name ends in c_proj.weight.

Compute: FLOPs per token ≈ 2 × parameters

Here is the estimate you will use constantly. Almost all the compute in a transformer is matrix multiplication of activations by weight matrices. When a $d_{in}$-vector is multiplied by a $d_{in} \times d_{out}$ weight matrix, every one of the $d_{in} d_{out}$ weights is used exactly once, and each use is one multiply and one add: 2 FLOPs. So a forward pass through a weight matrix costs $2 \times (\text{number of weights})$ FLOPs per token. Sum over all matrices in the model and you get:

where $N_{params}$ counts the weight matrices actually multiplied (so it excludes the input embedding table, which is a lookup, but includes the output projection). The attention core, $QK^\top$ and $AV$, adds about $4Ld$ per layer per token; this is not proportional to parameters, and for GPT-2 small at $L=1024$ it is about 20% extra. At long contexts it dominates, as the previous chapter showed.

Training costs about three times the forward pass, because backpropagation needs two more matrix multiplications of the same size for each one in the forward pass (one for the gradient with respect to the input, one for the gradient with respect to the weights). Hence the famous rule from Kaplan et al. (2020), "Scaling Laws for Neural Language Models":

FLOPs to train a model of $N$ parameters on $D$ tokens. For GPT-2 small (85M non-embedding parameters) on 10 billion tokens: $6 \times 85\text{M} \times 10^{10} \approx 5 \times 10^{18}$ FLOPs, a few hours on a modern GPU cluster. For a 70B model on 15 trillion tokens: $6 \times 7\times 10^{10} \times 1.5 \times 10^{13} \approx 6 \times 10^{24}$ FLOPs. The scaling chapter is built on this formula.

Implementation

Here is the block in PyTorch, matching code/lumen/block.py. LayerNorm and RMSNorm are written out rather than imported so you can see there is nothing hidden.

import torch
import torch.nn as nn
import torch.nn.functional as F
from .attention import MultiHeadAttention

class LayerNorm(nn.Module):
    def __init__(self, d, eps=1e-5):
        super().__init__()
        self.g = nn.Parameter(torch.ones(d))     # gamma: per-dimension gain
        self.b = nn.Parameter(torch.zeros(d))    # beta: per-dimension bias
        self.eps = eps

    def forward(self, x):                        # x: (..., d); statistics over the last axis only
        mu = x.mean(-1, keepdim=True)
        var = x.var(-1, keepdim=True, unbiased=False)
        return (x - mu) / torch.sqrt(var + self.eps) * self.g + self.b

class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.g = nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x):
        rms = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return x / rms * self.g                  # no mean subtraction, no bias

keepdim=True keeps the reduced axis so the statistics broadcast back against x. unbiased=False divides by $d$ rather than $d-1$, matching the formula (and nn.LayerNorm). Everything is per token: a tensor of shape (B, L, d) gets B × L independent normalisations.

class MLP(nn.Module):
    def __init__(self, d_model, d_ff=None, dropout=0.0):
        super().__init__()
        d_ff = d_ff or 4 * d_model
        self.fc = nn.Linear(d_model, d_ff)       # W1: expand
        self.proj = nn.Linear(d_ff, d_model)     # W2: project back
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        return self.drop(self.proj(F.gelu(self.fc(x))))

class SwiGLU(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w_up = nn.Linear(d_model, d_ff, bias=False)
        self.w_down = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))

class TransformerBlock(nn.Module):
    """Pre-LN block:  x = x + Attn(LN1(x));  x = x + MLP(LN2(x))"""
    def __init__(self, d_model, n_heads, d_ff=None, dropout=0.0):
        super().__init__()
        self.ln1 = LayerNorm(d_model)
        self.attn = MultiHeadAttention(d_model, n_heads)
        self.ln2 = LayerNorm(d_model)
        self.mlp = MLP(d_model, d_ff, dropout)
        self.drop = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        x = x + self.drop(self.attn(self.ln1(x), mask))   # gather from context, write to stream
        x = x + self.mlp(self.ln2(x))                     # think per token, write to stream
        return x

The forward is two lines, and they are exactly the two equations of the symbol strip. Note what is not there: no normalisation of x itself, and nothing that replaces x. The GPT-2 init rule lives in the full model, where the number of blocks is known:

def _init_weights(self, module):
    if isinstance(module, nn.Linear):
        nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Embedding):
        nn.init.normal_(module.weight, mean=0.0, std=0.02)

# after self.apply(self._init_weights): shrink the two residual-writing matrices per block
for name, p in self.named_parameters():
    if name.endswith("attn.out.weight") or name.endswith("mlp.proj.weight"):
        nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * self.n_layers))

Common confusions

Practice

Further reading

• Vaswani et al. (2017). Attention Is All You Need. Section 3.3 is the position-wise feed-forward network; the figure shows the (post-LN) block.
• He et al. (2016). Deep Residual Learning for Image Recognition. Where residual connections come from, and the experiments showing plain deep nets fail to fit.
• Ba, Kiros, Hinton (2016). Layer Normalization.
• Zhang and Sennrich (2019). Root Mean Square Layer Normalization.
• Xiong et al. (2020). On Layer Normalization in the Transformer Architecture. Why pre-LN trains without warm-up.
• Hendrycks and Gimpel (2016). Gaussian Error Linear Units (GELUs).
• Shazeer (2020). GLU Variants Improve Transformer. SwiGLU.
• Geva et al. (2021). Transformer Feed-Forward Layers Are Key-Value Memories.
• Elhage et al. (2021). A Mathematical Framework for Transformer Circuits. The residual-stream view.
• nostalgebraist (2020). interpreting GPT: the logit lens. Belrose et al. (2023). Eliciting Latent Predictions from Transformers with the Tuned Lens.
• Tenney, Das, Pavlick (2019). BERT Rediscovers the Classical NLP Pipeline.
• Kaplan et al. (2020). Scaling Laws for Neural Language Models. The $C \approx 6ND$ rule and its derivation in the appendix.
• Radford et al. (2019). Language Models are Unsupervised Multitask Learners. GPT-2: pre-LN and the residual init scaling.