The Block and the Full Model

Attention is one part of a transformer. This module supplies the rest: position information so word order counts, residual connections and normalization so deep stacks train, the feed-forward network, and the output layer. Section 7.6 assembles them and follows one input through the whole model.

7.1 Positional EncodingsRun this section

The 2017 paper gives each position its own vector and adds it to the token's embedding before the first layer:

Both have width $d_{\text{model}}$. The paper builds $\mathbf{p}_{pos}$ from waves, in pairs of features: pair $i$ uses feature $2i$ (sine) and $2i+1$ (cosine) at the same frequency:

$pos$ runs from $0$ to $N-1$ and $i$ from $0$ to $d_{\text{model}}/2 - 1$. Pair $i$ has frequency $\omega_i = 1 / 10000^{2i/d_{\text{model}}}$ and so repeats every $2\pi / \omega_i$ positions, its wavelength: about $6.28$ for pair 0, growing to almost $2\pi \cdot 10000$ for the last.

In words. Short waves separate nearby positions; long waves separate distant ones. Read together, like the hands of a clock, they give every position a distinct vector, with nearby positions similar.

The Linear Shift Property

Sines and cosines make moving $k$ positions forward the same operation at every starting position: one $2 \times 2$ matrix per pair, from the angle-addition formulas:

For fixed $k$, $\cos(\omega k)$ and $\sin(\omega k)$ are constants, so each shifted value is a fixed mix of the unshifted pair:

The matrix rotates each pair by the angle $\omega_i k$. It contains $k$ but not $pos$, so the same matrix works at every position. (Its minus sign sits bottom-left, unlike section 15.1's rotation, only because each pair is written sine first.)

In words. The encoding as a heatmap and as six clock hands, one per frequency pair. Moving k positions forward turns every hand by an angle depending only on k and the hand's speed, wherever you start.

The hope was that attention could then learn relative patterns such as "the word 3 places back" once, for every position. Nothing forces the model to use the structure. Section 15.1 covers RoPE (rotary position embedding), which builds the rotation into the query–key product itself.

Measured (lab 2a, experiments/ab_positional.py). The property holds. Fitting the best linear map from $PE_{pos}$ to $PE_{pos+k}$ over all positions (by least squares) leaves an error of $2.9 \times 10^{-6}$ for $k = 1, 7$ and $25$. That is the rounding floor of float32 arithmetic, not a flaw in the encoding. The fitted $2 \times 2$ matrix matches the rotation matrix in this section to within $6 \times 10^{-8}$.

Which Encoding the Reference Implementation Uses

The reference implementation uses learned absolute positional embeddings, as GPT-1 and GPT-2 did: a table with one trainable vector per position, up to the context length $L$. This section derives the sinusoids because RoPE builds on their maths.

# src/transformer/positional_embedding.py
self.embedding = nn.Embedding(
    num_embeddings=context_length,
    embedding_dim=d_model,
)

The choice shows up in three places:

Measured (lab 2a). Over 600 training steps with the same seed and the same batches, the sinusoidal version reaches $1.9442$ validation loss (the error on text held out from training; lower is better) and the learned table reaches $2.0978$. The sinusoidal version is better by $0.1536$ and has 4,096 fewer parameters. Read this with care: it is two layers, 600 steps and one seed. It does not show that sinusoids beat learned tables in general, and both are used in production. It shows that at this scale the choice does not decide the result.

7.2 Residual Connections and the Residual StreamRun this section

A block has two sub-layers, attention then the feed-forward network (section 7.4), and each one's output is added to its input: $\mathbf{x} + \operatorname{Sublayer}(\mathbf{x})$, which the 2017 paper then normalizes (section 7.3). The $\mathbf{x} +$ is a residual connection (He et al., 2016). It does two jobs:

• Each sub-layer only learns a change. With nothing useful to add, it outputs near zero and the vector passes through untouched, so a deep stack starts out harmless rather than scrambling its input.
• The gradient gets a path with no multiplications. The derivative of $\mathbf{x} + f(\mathbf{x})$ is $I + \frac{\partial f}{\partial \mathbf{x}}$; the identity passes the gradient straight through, curing the long-chain problem of section 2.6. Section 8.5 makes this exact for a whole stack.

Picture the running vector as a shared channel up through the model, which every sub-layer reads from and writes into: the residual stream. Section 14.1 shows how much this picture explains.

7.3 LayerNormRun this section

Adding to the stream again and again lets the numbers drift: growing, or lopsided with a few features dominating. LayerNorm (Ba et al., 2016) rescales each token's vector to mean 0 and variance 1, then applies a learned scale $\boldsymbol{\gamma}$ and shift $\boldsymbol{\beta}$, keeping values in a steady range so softmax stays out of saturation and gradients stay healthy. Normalizing after the add, as in 2017, is Post-LN; section 8.5 compares it with Pre-LN.

$\epsilon$ is a tiny constant that keeps the denominator away from zero. The mean and variance are over the $d$ features of one token, not across tokens or the batch.

Worked example. Take $\mathbf{x} = [2, 4, 6, 8]$, with $\boldsymbol{\gamma} = 1$ and $\boldsymbol{\beta} = 0$. The mean is $5$ and the variance is $\frac{9 + 1 + 1 + 9}{4} = 5$, so the standard deviation is $\sqrt{5} \approx 2.236$. The output is $[-1.342, -0.447, 0.447, 1.342]$: mean 0, variance 1. Doubling $\mathbf{x}$ to $[4, 8, 12, 16]$ gives exactly the same output, because LayerNorm discards the overall scale.

In words. An eight-feature vector as bars: as input, after subtracting the mean, after dividing by the standard deviation, and after the learned scale and shift. Scaling or shifting the whole input leaves the normalized bars still; changing one feature moves them.

7.4 The Feed-Forward NetworkRun this section

Every attention sub-layer is followed by a feed-forward network (FFN): a one-hidden-layer MLP (section 4.2) applied to each token separately, with the same weights at every position.

Why? Attention moves information between tokens but barely changes it. Token $i$'s output $\mathbf{y}_i = \sum_j A_{ij} \mathbf{v}_j$ is a weighted average of values: the weights depend nonlinearly on the input through softmax, but the features themselves pass only through linear maps ($W^V$ and $W^O$). Attention cannot compute something like "switch this feature on only if these two others are both large". The FFN supplies that per-token nonlinearity:

$W_1 \in \mathbb{R}^{d_{\text{model}} \times 4d_{\text{model}}}$ widens the vector fourfold (the paper uses $d_{ff} = 2048$ for $d_{\text{model}} = 512$), ReLU zeroes the negative entries, and $W_2 \in \mathbb{R}^{4d_{\text{model}} \times d_{\text{model}}}$ brings it back to model width. GPT-2 and the reference implementation use GELU, a smooth ReLU that bends gently near zero instead of at a corner:

# src/transformer/feed_forward.py — the shipped activation
self.net = nn.Sequential(
    nn.Linear(d_model, d_ff),
    nn.GELU(),
    nn.Linear(d_ff, d_model),
)

Section 14.2 reads the FFN another way, as a store of memorized associations.

7.5 The Output LayerRun this section

After the last block, each token's $d_{\text{model}}$ numbers must become one score per vocabulary entry. The output layer (unembedding, or language-model head) does this with one more matrix:

Logit $j$ is the final vector dotted with column $j$, large when they point the same way. Softmax turns the logits into probabilities, and training applies cross-entropy to them.

In a small model this is usually the largest matrix: 38.6 million numbers at $|V| = 50{,}257$ and $d_{\text{model}} = 768$. In training it runs at every position, since each predicts its own next token (section 8.4); in generation only the last position's scores are used.

Tied vs. Untied Embeddings

The input embedding $W_E \in \mathbb{R}^{|V| \times d_{\text{model}}}$ and the output layer have transposed shapes. The original Transformer and GPT-2 tied them ($W_{\text{unembed}} = W_E^\top$), saving a whole $|V| \times d_{\text{model}}$ block of parameters.

LLaMA and the larger Qwen models untie them, because the jobs differ: the embedding should place similar tokens close together, while the output layer is a classifier that must separate near-identical tokens such as "run", "runs" and "running". The extra parameters matter less in a large model, where the embedding is a small share. Small models often still tie; Gemma ties at every size because its 256,000-token vocabulary makes the table unusually large.

7.6 The Complete ModelRun this section

The whole 2017 model (Post-LN) for $N$ tokens, with each step's output shape. Module 15 modernizes it piece by piece.

1. Token ids: $N$ integers.
2. Embedding lookup (section 1.5): $N \times d_{\text{model}}$.
3. Add position vectors (section 7.1): $N \times d_{\text{model}}$.
4. Repeat for each of $L$ blocks:
   1. Multi-head attention (Module 6), whose $h$ heads each produce an $N \times N$ table of weights; their joined outputs go through $W^O$ back to $N \times d_{\text{model}}$.
   2. Add to the input and normalize (sections 7.2 and 7.3): $N \times d_{\text{model}}$.
   3. Feed-forward network (section 7.4), widening each token to $d_{ff}$ and back: $N \times d_{\text{model}}$.
   4. Add and normalize again: $N \times d_{\text{model}}$.
5. Output layer (section 7.5): $N \times |V|$ logits, then softmax.

From step 2 through step 4 every tensor is $N \times d_{\text{model}}$: the token vectors are refined in place. Only the first and last steps change the shape.

Counting the Parameters

With $d = d_{\text{model}}$ and $d_{ff} = 4d$:

• Attention: four $d \times d$ matrices ($W^Q$, $W^K$, $W^V$, $W^O$, with the heads together), so $4d^2$.
• Feed-forward: $d \times 4d$ and $4d \times d$, so $8d^2$.
• Norms and biases: a few times $d$, negligible.

So each block holds about $12d^2$, and the model about

where the last term is the embedding table (shared with the output layer if they are tied). For GPT-2 Small, $L = 12$ and $d = 768$: $12 \cdot 12 \cdot 768^2 \approx 85$ million in the blocks plus $50{,}257 \cdot 768 \approx 39$ million in the embedding, about 124 million in total, which is its published size.

Follow One Input

The model from Module 0's loop, opened up stage by stage, with every stage explained.

In words. The two-block model trained on Shakespeare from Module 0 lists every stage: token ids, embedding, added positions, each block's attention weights and outputs, and the final next-character probabilities. Choosing a stage shows its tensor as a heatmap with its shape, the operation that produced it and the section that teaches it.

Exercises

Core Academic References

• Vaswani et al. (2017): Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762
• He et al. (2016): Deep Residual Learning for Image Recognition. CVPR 2016: 770–778.
• Ba, Kiros & Hinton (2016): Layer Normalization. arXiv:1607.06450
• Hendrycks & Gimpel (2016): Gaussian Error Linear Units (GELUs). arXiv:1606.08415
• Radford et al., OpenAI (2019): Language Models are Unsupervised Multitask Learners (GPT-2). Technical Report.