Chapter 3: Transformer Architecture

Self-attention gives each token three roles, each produced by a learned matrix. The query $q_i = x_i W_Q$ describes what token $i$ is looking for. The key $k_j = x_j W_K$ describes what token $j$ offers to be matched against. The value $v_j = x_j W_V$ is the content token $j$ passes along if it is chosen. Here $x_i$ is the input embedding and $W_Q, W_K, W_V \in \mathbb{R}^{d \times d_k}$. Separating matching (keys) from content (values) lets a token be easy to find for one reason and contribute something different. The roles are learned, not assigned, so no head is guaranteed to learn a clean grammatical role.

The attention function runs in four steps. First compute all query-key dot products at once as $QK^T$, an $n \times n$ score matrix for $n$ tokens. Then divide by $\sqrt{d_k}$, where $d_k$ is the key dimension. Then apply a softmax to each row, so each token's weights are positive and sum to 1. Finally multiply by $V$ to mix the values. The scaling matters because a dot product of two random $d_k$-dimensional vectors with unit-variance entries has variance $d_k$. Unscaled, large scores push the softmax into a near one-hot regime where its gradients almost vanish, and learning stalls.

A decoder-only model such as GPT or Llama is trained to predict each next token, so position $i$ must not see positions after it. The causal mask enforces this by adding $-\infty$ to every score where the key position $j$ is greater than the query position $i$, as in the formula, before the softmax. Since $e^{-\infty} = 0$, those weights become exactly zero and the remaining weights still sum to 1. One forward pass then trains every position at once without leaking the answer. Forgetting the mask is a classic bug: training loss collapses toward zero while generation is useless.

Attention's cost comes from the score matrix. Every one of the $n$ queries is compared with all $n$ keys, and each comparison is a $d$-dimensional dot product, so time grows as $O(n^2 \cdot d)$, as in the formula. A naive implementation also stores the $n \times n$ matrix, so memory grows as $O(n^2)$ as well. The feedforward layers, by contrast, are linear in $n$. At short lengths the feedforward layers dominate, but at long lengths attention takes over. FlashAttention removes the quadratic memory by computing in tiles, but the quadratic compute remains. Sparse, sliding-window, and linear attention trade exactness for lower cost.

Multi-head attention splits the model width $d$ into $h$ heads of size $d_k = d/h$. Head $i$ has its own projections $W_Q^i, W_K^i, W_V^i \in \mathbb{R}^{d \times d_k}$ and computes ordinary scaled dot-product attention inside its own subspace. The $h$ outputs, each of width $d_k$, are concatenated back to width $d$ and multiplied by $W_O \in \mathbb{R}^{d \times d}$, as in the formula, which lets heads exchange information. Because each head is $h$ times narrower, the total cost is about that of one full-width head. Too many heads make each subspace too small to express a useful match.

Trained heads are not interchangeable. Studies of BERT and GPT-2 find heads that attend to the previous or next token, heads that track syntactic links such as a verb and its object, heads that resolve coreference, and heads that dump most of their weight on the first token or a separator when they have nothing useful to do. Induction heads copy patterns seen earlier in the context. The specialisation is emergent and messy: many heads mix several roles. Michel et al. (2019) showed that many heads can be pruned at test time with little loss, so capacity is unevenly spread.

During generation, a model caches the keys and values of every past token so they are not recomputed. That cache grows with layers, key-value heads, head size, and context length, and it often limits batch size more than the weights do. GQA (Ainslie et al., 2023) keeps $h_Q$ query heads but only $h_{KV}$ key-value heads, as in the formula, and each group of $h_Q/h_{KV}$ query heads shares one key-value pair. The cache shrinks by the factor $h_Q/h_{KV}$. Quality stays close to full multi-head attention. Llama 2 70B and Mistral 7B use 8 key-value heads.

Multi-Query Attention (Shazeer, 2019) is the extreme of GQA: all query heads share a single key head and a single value head, so $h_{KV} = 1$. The cache shrinks by the full factor $h_Q$, and each decoding step reads far less memory, which is the real bottleneck when generating one token at a time. The cost is quality, because every head must match against the same keys, and training can be less stable. Full multi-head attention has $h_{KV} = h_Q$. GQA sits between the two extremes, which is why most recent open models chose it over MQA.

The feed-forward network applies the same small MLP to each token separately. $W_1 \in \mathbb{R}^{d \times d_{ff}}$ lifts the vector into a wider space, typically $d_{ff} = 4d$, where many hidden units can each detect a different pattern. An activation function then zeroes or dampens the units that did not fire, and $W_2 \in \mathbb{R}^{d_{ff} \times d}$ writes the result back at width $d$. The biases $b_1, b_2$ are often dropped in modern LLMs. Because it is per token, every position runs in parallel. Its failure mode is its size, since it dominates the parameter and compute budget.

The activation is what makes the feed-forward network more than one big linear map. ReLU, $\max(0, x)$, was the original choice. It is cheap, but its gradient is exactly zero for negative inputs, so a unit pushed negative on every input stops learning. GELU, used by BERT and GPT-2, multiplies $x$ by $\Phi(x)$, the probability that a standard normal variable is below $x$. It follows ReLU for large positive or negative inputs but curves smoothly near zero and lets small negative values through. The formula shows the tanh approximation often used in code. Smooth gradients make optimisation of deep stacks easier.

SwiGLU (Shazeer, 2020) replaces the single up-projection with two. One branch, $xW_{gate}$, passes through SiLU, $\text{SiLU}(z) = z \cdot \sigma(z)$, and acts as a soft gate. The other branch, $xW_1$, carries values, and the two are multiplied element by element before the down-projection, as in the formula. The gate lets each hidden unit decide, from the input, how much of its value to pass. That multiplicative interaction improved quality in Shazeer's experiments at matched parameter count. It is used in PaLM, Llama, and Mistral. The cost is a third weight matrix, so the hidden width is shrunk to compensate.

Geva et al. (2021) proposed reading the feed-forward network as a key-value memory. Each row of $W_1$ acts as a key: its dot product with the token vector measures how well the input matches a pattern. The activation keeps the matches that fire, and each matching unit adds its row of $W_2$, a value vector, to the output. So each unit maps a detected pattern to a push toward certain next tokens. Meng et al. (2022) located factual associations in mid-layer feed-forward weights and edited single facts. Facts are spread over many units, so the picture is only approximate.

LayerNorm works on one token vector $x$ at a time. It computes the mean $\mu$ and variance $\sigma^2$ over the $d$ features, subtracts $\mu$, and divides by $\sqrt{\sigma^2 + \epsilon}$, where the small $\epsilon$ prevents division by zero. The result has mean 0 and variance 1. The learned vectors $\gamma$ and $\beta$ then rescale and shift each feature, so the model can undo the normalisation where that helps. Every token gets the same treatment regardless of batch size or sequence length. A caveat is that one huge outlier feature inflates $\sigma$ and squashes the rest of the vector.

The original Transformer used Post-Norm: compute $x + \text{SubLayer}(x)$, then normalise the sum. So every path from the output back to the input passes through a normalisation at each layer, and Xiong et al. (2020) showed the gradients near the output are large at initialisation. Post-Norm therefore needs a careful learning-rate warmup and becomes fragile past a few dozen layers. Pre-Norm normalises only the sublayer's input, $x + \text{SubLayer}(\text{LayerNorm}(x))$, and leaves the residual path untouched, which gives an identity route for gradients. GPT-2 and almost every later LLM use Pre-Norm, usually with a final norm before the output head.

RMSNorm (Zhang and Sennrich, 2019) drops two parts of LayerNorm: the mean subtraction and the shift $\beta$. It divides each token vector by its root mean square, $\sqrt{\frac{1}{d}\sum x_i^2 + \epsilon}$, and multiplies by the learned scale $\gamma$, as in the formula. The hypothesis is that re-scaling, not re-centring, gives normalisation its benefit. With one fewer reduction and no $\beta$, it is simpler and faster, and the paper reports running-time reductions of 7 to 64 percent across models with comparable quality. It is used in Llama, Mistral, and most recent open models. It is not invariant to a constant shift.

BatchNorm normalises each feature using statistics across the examples in a batch, which works well for image CNNs. Language models break its assumptions. Sequences vary in length and are padded, so batch statistics mix real tokens with padding. Batches in LLM training are often small per device, which makes the statistics noisy. In causal models a batch statistic lets information leak between examples and positions. BatchNorm also needs running averages at inference that differ from training statistics. LayerNorm computes everything inside one token, so it avoids all of these problems at the cost of ignoring cross-example statistics.

A plain layer must produce its whole output, $x_{l+1} = F(x_l)$. A residual layer produces only a change, $x_{l+1} = x_l + F(x_l)$, as in the formula, so $F$ learns the difference between output and input. The input and output must have the same width, which is why every transformer sublayer keeps dimension $d$. If a layer has nothing useful to add, it can push $F$ toward zero and pass its input through unchanged. Deep networks therefore start close to a stack of identities, which is a safe place to begin. Without normalisation, the sum of many updates can grow large.

Differentiate the residual formula and each layer contributes a factor $1 + \frac{\partial F_i}{\partial x_i}$ to the backward product, as in the formula. Expanding that product gives a term that is just 1, a direct path from the loss to layer $l$ that skips every sublayer, plus terms that pass through some sublayers. In a plain network each factor is only $\frac{\partial F_i}{\partial x_i}$, and a product of many factors below 1 vanishes. The residual form does not guarantee a large gradient, since the terms can cancel. But the identity path makes vanishing much less likely. Pre-Norm keeps that path clean.

Elhage et al. (2021) at Anthropic proposed reading a transformer as one shared vector per token, the residual stream, that runs from the embedding to the output. Every attention head and every feed-forward layer reads the stream through a normalisation, computes something, and adds its result back. Nothing is overwritten, only added. So the final vector is the embedding plus the sum of every sublayer's output, and the unembedding reads the prediction off that sum. This view explains why individual components can be studied separately. The limitation is that the stream has a fixed width, so many features must share dimensions.

The modern block has two residual sublayers, as in the formula. First, normalise the input $x$ with RMSNorm, run multi-head attention on it, and add the result back to get $h$. This is the only step where tokens exchange information. Second, normalise $h$, apply the feed-forward network to each token separately, and add the result to get the output. The residual stream itself is never normalised inside the block. Stacking $L$ blocks and adding a final norm and output head gives a GPT or Llama style model. Every intermediate vector keeps width $d$, which is what makes the residual additions possible.

Count per layer. Attention has four $d \times d$ matrices for queries, keys, values, and output, so $4d^2$ with full multi-head attention. A SwiGLU feed-forward network has three $d \times d_{ff}$ matrices, $3d \cdot d_{ff}$, and each of the two RMSNorms adds only $d$ scale parameters. Over $L$ layers that is about $L(4d^2 + 3d \cdot d_{ff})$. With $d_{ff} = \tfrac{8}{3}d$ it becomes $12Ld^2$, as in the formula. Add the embedding table, $V \times d$, and an untied output head of the same size. The estimate slightly overcounts GQA models, whose key and value matrices are smaller.

Generation produces one token per step, and each new token must attend to all earlier ones. Recomputing every earlier key and value at each step would repeat almost all the work. The key-value cache stores each token's keys and values in every layer once computed, so each step only computes the new token's query, key, and value and reads the rest from memory. The size, per the formula, is 2 (keys and values) times layers, key-value heads, head size, sequence length, and bytes per number. It grows linearly with context and batch, and at long contexts it can exceed the weights.

For a fixed parameter budget $12Ld^2$ you can trade layers $L$ for width $d$. Depth adds sequential steps of computation, which can compose features into more abstract ones, but each layer adds latency because layers run one after another. Width makes each matrix multiply larger, which GPUs execute efficiently and which tensor parallelism can split. Very deep, narrow models are harder to train, and very wide, shallow ones waste capacity. Kaplan et al. (2020) found that loss depends only weakly on the exact shape over a broad range. Llama 2 70B uses $L = 80$ and $d = 8192$, so its ratio is about 100.