Lecture 5 — The Transformer

Project connection. Project Step 5 assembles the full architecture and verifies its parameter count. The course configuration is $(V,T_{\max},d,H,L)=(65,64,128,4,4)$.

Chapter overview. One attention head supplies one content-dependent average of a low-dimensional view of the residual stream. A transformer runs several such heads, applies a nonlinear MLP independently at each position, and places both operations on residual branches. LayerNorm controls the feature scale presented to each branch. Stacking these blocks between embedding and unembedding maps gives the complete decoder-only generative pretrained transformer (GPT) used in the project.

0. Architectural notation

Let $T_{\max}$ be the maximum supported sequence length, $T\leq T_{\max}$ the current sequence length, $d$ the residual-stream width, $H$ the number of heads, and $L$ the number of transformer blocks. Assume $H$ divides $d$ and define

$$d_{\mathrm{head}}=\frac dH.$$

As in Lecture 4, matrices in $\mathbb R^{T\times d}$ store feature vectors as rows. Every attention head is causally masked unless explicitly stated otherwise.

A LayerNorm map acts independently on each row: it centers the feature coordinates, rescales by their root mean square, and then applies learned coordinatewise gains and shifts. Section 4 gives its exact formula and geometry.

1. Multi-head attention

For head $h\in\{1,\ldots,H\}$, let

$$W_Q^h,W_K^h,W_V^h\in\mathbb R^{d\times d_{\mathrm{head}}}$$

and let

$$\operatorname{Attn}^h(X)\in\mathbb R^{T\times d_{\mathrm{head}}}$$

be the scaled dot-product attention map of Lecture 4. Let $W_O\in\mathbb R^{d\times d}$ be the output projection.

Definition 1.1 (multi-head attention). The multi-head attention map is

$$\operatorname{MHA}(X) =\bigl[\operatorname{Attn}^1(X)\mid\cdots\mid \operatorname{Attn}^H(X)\bigr]W_O,$$

where the bracket denotes concatenation along the feature dimension.

Partition $W_O$ vertically into blocks $W_O^h\in\mathbb R^{d_{\mathrm{head}}\times d}$, so $W_O=[(W_O^1)^\top\ \cdots\ (W_O^H)^\top]^\top$. Let $A^h(X)$ be the $T\times T$ attention matrix of head $h$.

Proposition 1.2 (additive head decomposition). Multi-head attention can be written as

$$\operatorname{MHA}(X) =\sum_{h=1}^H A^h(X)XW_{OV}^h, \qquad W_{OV}^h:=W_V^hW_O^h.$$

Moreover, $\operatorname{rank}(W_{OV}^h)\leq d_{\mathrm{head}}=d/H$.

Proof. If $C^h=\operatorname{Attn}^h(X)=A^hXW_V^h$, block matrix multiplication gives

$$[C^1\mid\cdots\mid C^H]W_O=\sum_{h=1}^HC^hW_O^h.$$

Substitution yields the formula. The rank of a product is at most its inner dimension. $\square$

The score matrix of head $h$ is

$$S^h =\frac{XW_Q^h(XW_K^h)^\top}{\sqrt{d_{\mathrm{head}}}} =\frac{XW_{QK}^hX^\top}{\sqrt{d_{\mathrm{head}}}}, \qquad W_{QK}^h:=W_Q^h(W_K^h)^\top.$$

Thus $W_{QK}^h$ is a rank-at-most-$d/H$ bilinear form controlling the attention pattern, while $W_{OV}^h$ is a rank-at-most-$d/H$ map controlling the information written to the stream.

Remark 1.3 (QK and OV circuits). In the transformer-circuits terminology, the QK circuit specifies where to read and the OV circuit specifies what to write. The individual factorizations are not identifiable. For example, for invertible $R\in\mathbb R^{d_{\mathrm{head}}\times d_{\mathrm{head}}}$, replacing $W_Q^h$ by $W_Q^hR$ and $W_K^h$ by $W_K^hR^{-\top}$ preserves $W_{QK}^h$.

Remark 1.4 (role of multiple heads). One full-width head supplies one attention distribution per query position. Splitting the same projection budget among $H$ heads supplies $H$ independently learned distributions, each with a low-rank read/write channel. The parameter count remains $O(d^2)$. The inductive bias, meaning a structural preference built into the model class, changes from one high-dimensional lookup to several parallel low-dimensional lookups.

2. The position-wise MLP

Let $d_{ff}$ be the hidden width, let $W_1\in\mathbb R^{d_{ff}\times d}$, $W_2\in\mathbb R^{d\times d_{ff}}$, and let $b_1\in\mathbb R^{d_{ff}}$, $b_2\in\mathbb R^d$. Let $\sigma(t)=t\Phi(t)$ be the GELU activation, where $\Phi$ is the standard normal cumulative distribution function.

Definition 2.1 (MLP sublayer). The position-wise MLP is

$$\operatorname{MLP}(u) =W_2\sigma(W_1u+b_1)+b_2, \qquad u\in\mathbb R^d,$$

applied independently with the same parameters to every row of the residual stream. The course model uses $d_{ff}=4d$.

Remark 2.2 (division of labor). Attention is the only sublayer that communicates between sequence positions. The MLP performs a learned nonlinear feature transformation at each position separately. Attention is itself nonlinear through its softmax, and LayerNorm is also nonlinear; the precise claim is that the MLP is the dedicated position-wise feature computation that can move beyond the convex-hull restriction on the current attention values.

With $d_{ff}=4d$, the two weight matrices contain $8d^2$ parameters and the biases contain $5d$. This exceeds the approximately $4d^2$ attention parameters in each block.

3. Residual connections and pre-normalization

Let $\operatorname{LN}_1$ and $\operatorname{LN}_2$ be two LayerNorm maps.

Definition 3.1 (pre-norm transformer block). A transformer block maps $X\in\mathbb R^{T\times d}$ to $X'\in\mathbb R^{T\times d}$ by

$$\begin{aligned} Y&=X+\operatorname{MHA}(\operatorname{LN}_1(X)),\\ X'&=Y+\operatorname{MLP}(\operatorname{LN}_2(Y)). \end{aligned}$$

The additions are residual connections. A map of the form

$$X^{(\ell+1)}=X^{(\ell)}+F_\ell(X^{(\ell)})$$

has Jacobian

$$\frac{\partial X^{(\ell+1)}}{\partial X^{(\ell)}} =I+DF_\ell(X^{(\ell)}).$$

Remark 3.2 (discrete dynamical system). The update resembles a forward Euler step for an ordinary differential equation $\dot X=F(X)$. Under this analogy, depth plays the role of time and the residual stream is the evolving state.

Remark 3.3 (gradient path). A residual connection supplies an exact identity term in every block Jacobian. If the residual derivatives $DF_\ell$ are controlled, the product of block Jacobians is a product of perturbations of the identity rather than arbitrary matrices. This architecture mitigates vanishing and exploding gradients; the residual formula alone does not guarantee that every $DF_\ell$ is small.

Remark 3.4 (communication channel). Every sublayer reads the same running state and adds an increment. Features written by an early block can therefore be read by a later block. The width $d$ is the linear bandwidth of this channel, although superposition—encoding features in nonorthogonal directions—can represent more features than coordinates (Lecture 8). This additive structure underlies the path decompositions used in mechanistic interpretability.

Empirical aside 3.5 (small residual updates). In many trained models, individual branch updates are small relative to the stream. In that regime blocks are approximately composable, and dropping one block may have a modest effect. This is an observed property, not a consequence of the residual formula alone.

Aside 3.5 (pre-norm versus post-norm). The original transformer used $X\leftarrow\operatorname{LN}(X+F(X))$, called post-norm. GPT-2-style pre-norm uses $X\leftarrow X+F(\operatorname{LN}(X))$, preserving an unnormalized identity path around the branch. The course project uses pre-norm.

4. LayerNorm geometry

For $u\in\mathbb R^d$, define its coordinate mean and variance by

$$\bar u=\frac1d\sum_{i=1}^du_i, \qquad \sigma(u)^2=\frac1d\sum_{i=1}^d(u_i-\bar u)^2.$$

Let $\odot$ denote coordinatewise multiplication, let $\mathbf1=(1,\ldots,1)^\top$, and fix a small numerical constant $\varepsilon>0$.

Definition 4.1 (LayerNorm). With learned vectors $\gamma,\beta\in\mathbb R^d$, define

$$\operatorname{LN}(u) =\gamma\odot \frac{u-\bar u\mathbf1}{\sqrt{\sigma(u)^2+\varepsilon}} +\beta.$$

It is applied separately to each sequence position.

Proposition 4.2 (projection and radial rescaling). Assume $d\geq2$. Set $\varepsilon=0$, $\gamma=\mathbf1$, and $\beta=0$, and restrict to nonconstant $u$. Then the normalization map is the composition of:

1. orthogonal projection onto $\mathbf1^\perp=\{v:\mathbf1^\top v=0\}$; and
2. radial projection onto the sphere of radius $\sqrt d$ in that hyperplane.

Its image is

$$\{v\in\mathbf1^\perp:\|v\|_2=\sqrt d\},$$

which is a $(d-2)$-dimensional sphere.

Proof. The map $u\mapsto u-\bar u\mathbf1$ is the orthogonal projection onto $\mathbf1^\perp$: it is self-adjoint and idempotent, with kernel $\mathbb R\mathbf1$. Let $v=u-\bar u\mathbf1$. Then

$$\sigma(u)^2=\frac1d\|v\|_2^2, \qquad \frac{v}{\sigma(u)}=\sqrt d\frac{v}{\|v\|_2}.$$

The unit sphere in a $(d-1)$-dimensional vector space has dimension $d-2$. $\square$

Remark 4.3 (scale invariance). Under the idealized hypotheses of the proposition, $\operatorname{LN}(cu)=\operatorname{LN}(u)$ for $c>0$. Consequently, the normalized sublayer input depends on direction rather than positive scale. A nonzero $\varepsilon$ makes this invariance only approximate near zero variance, while learned $\gamma$ and $\beta$ restore coordinatewise scale and shift after normalization.

Aside 4.4 (LayerNorm versus BatchNorm). BatchNorm uses statistics across examples in a minibatch. LayerNorm uses coordinates within one position, so its output does not depend on which other examples share the batch and remains defined at batch size one. RMSNorm omits mean centering and retains only radial rescaling.

5. The GPT specification

Let $E\in\mathbb R^{V\times d}$ be the token embedding matrix and $P\in\mathbb R^{T_{\max}\times d}$ the learned positional embedding matrix. For tokens $x_1,\ldots,x_T$, define

$$X_t^{(0)}=E_{x_t,:}+P_{t,:}.$$

For $\ell=0,\ldots,L-1$, define

$$\begin{aligned} Y^{(\ell)} &=X^{(\ell)}+ \operatorname{MHA}^{(\ell)} \bigl(\operatorname{LN}^{(\ell)}_1(X^{(\ell)})\bigr),\\ X^{(\ell+1)} &=Y^{(\ell)}+ \operatorname{MLP}^{(\ell)} \bigl(\operatorname{LN}^{(\ell)}_2(Y^{(\ell)})\bigr). \end{aligned}$$

Finally, with a final LayerNorm $\operatorname{LN}_f$ and unembedding matrix $W_U\in\mathbb R^{d\times V}$, define logits

$$Z=\operatorname{LN}_f(X^{(L)})W_U\in\mathbb R^{T\times V}$$

and next-token probabilities

$$p_\theta(\,\cdot\mid x_{\leq t}) =\operatorname{softmax}(Z_{t,:}).$$

Remark 5.1 (causality). Token and positional embeddings, LayerNorm, MLPs, residual additions, and unembedding act position-wise. The only cross-position operation is attention, whose causal mask satisfies Lecture 4, Proposition 4.3. Composition therefore makes each logit row $Z_t$ a function only of $x_1,\ldots,x_t$.

Theorem 5.2 (exact parameter count). Assume $d_{ff}=4d$, biases on the MLP maps and attention output projection, and no biases on $W_Q,W_K,W_V$, or $W_U$. Then the number of trainable scalar parameters is

$$\underbrace{Vd+T_{\max}d}_{\text{token and position embeddings}} +L\underbrace{\bigl[(4d^2+d)+(8d^2+5d)+4d\bigr]}_{ \text{attention + MLP + two LayerNorms}} +\underbrace{2d}_{\text{final LayerNorm}} +\underbrace{dV}_{\text{unembedding}}.$$

Equivalently, it is

$$2Vd+T_{\max}d+L(12d^2+10d)+2d.$$

Proof. Across all heads, the query, key, and value matrices contain $3H\,d(d/H)=3d^2$ parameters. The output projection contributes $d^2+d$, so attention contributes $4d^2+d$. The MLP weights contribute $4d^2+4d^2=8d^2$, and its biases contribute $4d+d=5d$. Each LayerNorm has $d$ gains and $d$ shifts, so two contribute $4d$. Add the embedding, final-normalization, and unembedding terms. $\square$

For $(V,T_{\max},d,H,L)=(65,64,128,4,4)$:

component	parameters	share
token and positional embeddings	16,512	2.0%
attention in all blocks	262,656	32.2%
MLPs in all blocks	526,848	64.5%
block LayerNorms	2,048	0.3%
final LayerNorm and unembedding	8,576	1.0%
total	816,640	100%

The leading block term is $12Ld^2$; biases and normalization parameters add $10Ld$.

Project aside 5.2 (initialization). There are $2L$ residual writes. Under the approximation that their contributions are independent with comparable variance, scaling each branch’s output projection by $1/\sqrt{2L}$ keeps their accumulated variance of order one. Lecture 6 derives this bookkeeping. The course model begins with loss near $\ln65=4.174$, as expected for nearly uniform predictions.

Aside 5.3 (weight tying). GPT-2 identifies the unembedding with the transpose of the token embedding, $W_U=E^\top$. The course’s baseline count keeps them separate so that the accounting is explicit.

6. Computational cost

A forward pass on a length-$T$ sequence has the leading operation counts

component	leading cost
query, key, value, and output projections	$O(LTd^2)$
attention scores and value aggregation	$O(LT^2d)$
MLP sublayers	$O(LTd^2)$

Thus total cost is $O(LTd^2+LT^2d)$. The attention term becomes dominant when $T$ is large relative to $d$.

Rule-of-thumb aside 6.1. Dense transformer training is often estimated as $6\times(\text{non-embedding parameters})\times(\text{training tokens})$ floating-point operations: roughly two operations per forward multiply—accumulate and about twice the forward cost for the two backward derivatives. The approximation ignores attention’s non-parameterized $T^2$ work and hardware details; Exercise 8 asks for its derivation and limitations.

Summary

Multi-head attention is an additive sum of low-rank head contributions. Each head has a QK circuit controlling its attention pattern and an OV circuit controlling its write to the stream. Position-wise MLPs supply the dedicated nonlinear feature transformation at each token. Residual connections create identity gradient paths and a shared evolving state. Idealized LayerNorm projects to $\mathbf1^\perp$ and rescales to a $(d-2)$-sphere. These components give a complete causal GPT with exactly 816,640 parameters in the course configuration and leading cost $O(LTd^2+LT^2d)$.

Exercises (paired with Step 5)

A star marks a Project Step 5 task.

1. ★ Re-derive Theorem 5.2 using the bias conventions in your code and verify the result against numel() exactly.
2. ★ Prove end-to-end causality of the GPT specification and verify it by perturbing future tokens.
3. Prove Proposition 1.2. Describe the non-identifiability of the QK factorization. For a generic fixed QK matrix, compute the dimension of the set of factor pairs that represent it.
4. Prove Proposition 4.2 and its positive-scale invariance consequence. Discuss weight decay immediately before a LayerNorm.
5. Recount the parameters under weight tying $W_U=E^\top$.
6. Prove $\operatorname{rank}(W_{OV}^h)\leq d/H$ and $\operatorname{rank}(S^h)\leq d/H$. Explain how multiple heads alter, but do not remove, these constraints.
7. Compare the Jacobians of one pre-norm and one post-norm residual block.
8. Derive the $6\times\text{parameters}\times\text{tokens}$ training-cost rule and identify regimes in which attention’s $T^2$ term invalidates it.

Pointers

Vaswani et al., Attention Is All You Need (2017); Radford et al., GPT-2 (2019); Ba, Kiros, and Hinton, Layer Normalization (2016); He et al., Deep Residual Learning (2015); Elhage et al., A Mathematical Framework for Transformer Circuits (2021); Phuong and Hutter, arXiv:2207.09238. See also bbycroft.net/llm, the resources page, and Project Step 5.