# Module 7: The Block and the Full Model
> Positional encoding, residual connections, LayerNorm, the feed-forward network and the output layer, with one token traced through a whole transformer.
- Part II: The Transformer · core track · Intermediate · 5 hours
- By the end you can trace one token through a full transformer and name every tensor's shape along the way.
- Builds on: [Module 6: Attention](https://transformers-and-math.drmhse.com/md/06-attention.md)
- Web page (interactive figures, labs and equation explainers): https://transformers-and-math.drmhse.com/modules/07-the-block-and-the-full-model
- From *Transformers & Math* by Mike Chumba: https://transformers-and-math.drmhse.com
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 Encodings
The 2017 paper gives each position its own vector and adds it to the token's embedding before the first layer:
$$
\mathbf{x}_{\text{input}, pos} = \mathbf{x}_{\text{token}} + \mathbf{p}_{pos}
$$
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:
$$
PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i / d_{\text{model}}}}\right)
$$
$$
PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i / d_{\text{model}}}}\right)
$$
$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.
```mermaid
graph TD
Pos["Every position (d_model = 512)"] --> LowI["Pair i = 0: wavelength 2π ≈ 6.28 tokens
(tells neighbors apart)"]
Pos --> MidI["Pair i = 64: wavelength 2π · 10 ≈ 63 tokens
(tells apart positions within a passage)"]
Pos --> HighI["Pair i = 255 (last): wavelength ≈ 60,600 tokens
(tells apart distant parts of a document)"]
```
> **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:
$$
\sin(\omega (pos + k)) = \sin(\omega pos)\cos(\omega k) + \cos(\omega pos)\sin(\omega k)
$$
$$
\cos(\omega (pos + k)) = \cos(\omega pos)\cos(\omega k) - \sin(\omega pos)\sin(\omega k)
$$
For fixed $k$, $\cos(\omega k)$ and $\sin(\omega k)$ are constants, so each shifted value is a fixed mix of the unshifted pair:
$$
\begin{bmatrix} PE_{(pos+k, 2i)} \\ PE_{(pos+k, 2i+1)} \end{bmatrix} = \begin{bmatrix} \cos(\omega_i k) & \sin(\omega_i k) \\ -\sin(\omega_i k) & \cos(\omega_i k) \end{bmatrix} \begin{bmatrix} PE_{(pos, 2i)} \\ PE_{(pos, 2i+1)} \end{bmatrix}
$$
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.)
> **Interactive figure** `sinusoidal-positions`, on the web page: https://transformers-and-math.drmhse.com/modules/07-the-block-and-the-full-model#figure-sinusoidal-positions
>
> **Try:** Pick two positions far apart, say 3 and 40, and then 103 and 140. Before looking at the hands, predict whether the turn from the first position to the second is the same in both pairs.
>
> **What happens:** It is identical: every hand turns by an angle that depends only on the gap (37 positions) and the hand's own speed, never on where you start. That is the shift property, and it is what lets a pattern such as "the word 3 places back" be the same calculation everywhere.
> **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.
```python
# src/transformer/positional_embedding.py
self.embedding = nn.Embedding(
num_embeddings=context_length,
embedding_dim=d_model,
)
```
The choice shows up in three places:
| | learned absolute | sinusoidal |
|---|---|---|
| parameters at $d_{\text{model}} = 64$, $L = 64$ | $4{,}096$ (a $64 \times 64$ table) | $0$ |
| position beyond $L$ when the trained model is used | no vector exists | computed by the formula |
| with a KV cache (section 12.2: stored keys and values reused while generating) | the code must pass the true position, and it cannot exceed $L$ | the code must pass the true position, and any position works |
> **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 Stream
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 LayerNorm
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*.
$$
\operatorname{LayerNorm}(\mathbf{x}) = \frac{\mathbf{x} - \mu}{\sqrt{\sigma^2 + \epsilon}} \odot \boldsymbol{\gamma} + \boldsymbol{\beta}, \quad \mu = \frac{1}{d}\sum_{i=1}^d x_i, \quad \sigma^2 = \frac{1}{d}\sum_{i=1}^d (x_i - \mu)^2
$$
$\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.
> **Interactive figure** `layernorm-demo`, on the web page: https://transformers-and-math.drmhse.com/modules/07-the-block-and-the-full-model#figure-layernorm-demo
>
> **Try:** Multiply the input vector by 10 using the scale slider. Before you look, predict what happens to LayerNorm's output, and what that implies about information carried in a vector's overall size.
>
> **What happens:** Nothing changes: the output is identical at every scale, because the division by the standard deviation cancels any overall factor. Information a model wants to keep through a LayerNorm cannot live in the overall size of the vector; it has to live in the pattern across features. The learned $\boldsymbol{\gamma}$ then puts a scale back, one feature at a time.
> **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 Network
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:
$$
\operatorname{FFN}(\mathbf{x}) = \max(0, \mathbf{x} W_1 + \mathbf{b}_1) W_2 + \mathbf{b}_2
$$
$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:
$$
\operatorname{GELU}(z) \approx 0.5\,z\left(1 + \tanh\left(\sqrt{2/\pi}\left(z + 0.044715 z^3\right)\right)\right)
$$
```python
# 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 Layer
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:
$$
\mathbf{z} = \mathbf{h}_{\text{final}} W_{\text{unembed}}, \quad W_{\text{unembed}} \in \mathbb{R}^{d_{\text{model}} \times |V|}
$$
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 Model
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
$$
P \approx 12 L d^2 + |V| d
$$
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.
> **Interactive figure** `token-tracer-full`, on the web page: https://transformers-and-math.drmhse.com/modules/07-the-block-and-the-full-model#figure-token-tracer-full
>
> **Try:** Open *Block 1 · attention weights* and look at the first row of each head. Before you do, predict what it must contain, whatever the weights of the model are.
>
> **What happens:** A single 1 in the first column and nothing else. The causal mask lets the first token see only itself, so its whole weight goes there. Every later row spreads its weight over itself and the tokens before it, and nothing above the diagonal is ever filled.
> **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
**Exercise (check).** Which steps of the full model change the shape of the tensor, and which keep it at $N \times d_{\text{model}}$?
*Solution:* The embedding lookup changes $N$ ids into $N \times d_{\text{model}}$, and the output layer changes $N \times d_{\text{model}}$ into $N \times |V|$. Everything in between (position vectors, every attention and feed-forward sub-layer, every add and normalize) keeps the shape $N \times d_{\text{model}}$. Inside attention the score table is $N \times N$, but each sub-layer's output is back to $N \times d_{\text{model}}$.
**Exercise (apply).** Apply LayerNorm (with $\boldsymbol{\gamma} = 1$, $\boldsymbol{\beta} = 0$ and $\epsilon$ negligible) to $[1, 1, 1, 5]$. What is the last entry of the output?
*Answer:* 1.732 (within 0.002)
*Solution:* The mean is $2$. The squared deviations are $1, 1, 1, 9$, so the variance is $3$ and the standard deviation $\sqrt{3} \approx 1.732$. The last entry is $(5 - 2)/1.732 \approx 1.732$.
**Exercise (apply).** Using $P \approx 12Ld^2$, how many parameters are in the blocks of a model with $L = 24$ and $d = 1024$, in millions?
*Answer:* 302 (within 1)
*Solution:* $12 \cdot 24 \cdot 1024^2 = 288 \cdot 1{,}048{,}576 \approx 302$ million. (This is GPT-2 Medium's block count, which together with its embedding gives its published 355 million.)
**Exercise (apply).** Show that the shift matrix for the sinusoidal encoding does not depend on $pos$, by writing $\sin(\omega(pos + k))$ with the angle-addition formula.
*Solution:* $\sin(\omega(pos + k)) = \sin(\omega\,pos)\cos(\omega k) + \cos(\omega\,pos)\sin(\omega k)$. The coefficients $\cos(\omega k)$ and $\sin(\omega k)$ contain $k$ but not $pos$, so one fixed $2 \times 2$ matrix maps the pair at any position $pos$ to the pair at $pos + k$.
**Exercise (extend).** Without the residual connection, what would the derivative of a sub-layer's output with respect to its input be, and why does the $I$ in $I + \partial f/\partial \mathbf{x}$ matter for a 48-layer stack?
*Solution:* Without it, the derivative is just $\partial f/\partial \mathbf{x}$, and the gradient reaching the bottom of the stack is a product of 48 such Jacobians, which vanishes or explodes like the chains of section 2.6. With the residual, each factor is $I + \partial f / \partial \mathbf{x}$; expanding the product gives a term that is exactly $I$, a path from the top to the bottom that no layer scales. Even if every sub-layer's own Jacobian were tiny, the gradient would still arrive.
### Core Academic References
- **Vaswani et al. (2017)**: *Attention Is All You Need*. NeurIPS 2017. [arXiv:1706.03762](https://arxiv.org/abs/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](https://arxiv.org/abs/1607.06450)
- **Hendrycks & Gimpel (2016)**: *Gaussian Error Linear Units (GELUs)*. [arXiv:1606.08415](https://arxiv.org/abs/1606.08415)
- **Radford et al., OpenAI (2019)**: *Language Models are Unsupervised Multitask Learners* (GPT-2). Technical Report.