The original transformer, step by step

What you will be able to do

Read the 2017 paper that started the transformer, one step at a time. You will see each formula with its words and a small number you can check. Then you follow one tiny example through the whole machine with its shapes, watch it learn, and see what later models changed.

NotebookBuild the original transformer in a notebookColab · Kaggle · GitHub

The notebook builds the positional encoding and an encoder layer from the formulas and checks them. Then it trains the reversal toy for 1,500 steps (this lesson used 3,000). It decodes one example step by step, draws the cross-attention and counts the base model’s parameters. About 40 minutes, with about 1 minute of training on a free Colab CPU.

Open in ColabOpen in KaggleView on GitHub

Colab and Kaggle open your own copy. They need a free Google or Kaggle account. Stuck? See the solution notebook.

Before you startWhat to know and to have readyAbout 60 minutes

What you need to know

• The transformer block (multi-head attention, feed-forward layer, residual connection, layer normalization), from Transformer blocks.
• Position vectors, the padding mask, the causal mask, encoders and decoders, from Position and masking.
• Generation one token at a time, and SmolLM2-135M’s parts, from Training versus inference.
• The attention formula softmax(QKᵀ/√d_k)V and multi-head attention as matrices, from The mathematics of attention.
• Matrix multiplication and its shape rule, from Vector operations: matrix multiplication in the mathematics course; the mean and the standard deviation, from Describe a dataset: spread.

What you need on your computer

• Nothing for the page. The notebook opens in Colab or Kaggle with PyTorch already installed; it needs no data download.

Tested with Python 3.12 and PyTorch 2.14.1 on the CPU (macOS).

This lesson is optional: the certificate does not require it.

What you already know

You saw the idea in Transformer blocks and Position and masking. A block reads the other tokens with attention, then works on each token with a feed-forward layer. Residual connections and layer normalization keep it stable; positions and masks control what each token sees. SmolLM2-135M, in Training versus inference, stacks 30 such blocks.

All of this goes back to one paper: Attention Is All You Need by Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser and Polosukhin, published at NeurIPS 2017 (arXiv:1706.03762). This page shows the mathematics. The question it answers: what exactly does the original transformer compute, from the first token to the last probability? And which of its choices do models still use today?

A machine built for translation

The paper built a translator, not a chat model. Its main test was the WMT 2014 English–German data, with a shared vocabulary of about 37,000 subword tokens for both languages. The source sentence is complete; the target is written one token at a time.

So the model has two halves. The encoder reads the whole source sentence, with no causal mask, and gives one vector per source token: the memory. The decoder writes the target. It reads the tokens it has written so far, with a causal mask, and in every layer it also reads the memory. That reading of the other text is cross-attention, the one part you have not met yet.

To see every step with real numbers, this page uses a tiny model with the paper’s exact recipe and a toy task: reverse a sequence of letters. The source “c a b e d” must become “d e b a c”, followed by an end token </s>. The vocabulary has 13 tokens: the letters a to j, a padding token, a start token <s> and </s>. The model has 2 encoder and 2 decoder layers, d_model = 64, 4 heads and a feed-forward size of 256: 234,304 parameters. (d_model is the paper’s name for d, the number of numbers per token in The mathematics of attention.) The paper’s base model has 6 + 6 layers, d_model = 512, 8 heads and 2,048.

The diagram redraws the paper’s architecture figure, with the shape on every arrow for this example.

The encoder–decoder of the 2017 paper, with the shape on every edge for one example of the toy model. The memory feeds the cross-attention of every decoder layer.

Text description of the diagram

Two columns. The encoder, on the left, starts with the source tokens, shape (1, 5): one sequence of 5 tokens. The embedding multiplies each token’s row by 8, the square root of 64; the shared table has 13 rows of 64 numbers. The shape is now (1, 5, 64). The sinusoidal positional encoding is added, with no change of shape. Then comes an encoder layer, repeated 2 times here and 6 times in the paper. It has multi-head self-attention with 4 heads of 16 numbers, whose weights have shape (1, 4, 5, 5), and then “add and LayerNorm”. Then comes a feed-forward layer from 64 to 256 numbers and back, with ReLU, and “add and LayerNorm” again. Every arrow in the layer carries (1, 5, 64). Dashed lines carry each sublayer’s input around it to the “add and LayerNorm” box: the residual connections. The encoder’s output is the memory, (1, 5, 64).

The decoder, on the right, starts with the target so far: the start token and d e b a c, shape (1, 6). It uses the same embedding table, times 8, and the same positional encoding: (1, 6, 64). A decoder layer is repeated 2 times. It has masked self-attention with weights (1, 4, 6, 6), then cross-attention, then the feed-forward layer, each followed by add and LayerNorm. In the cross-attention, the queries come from the decoder and the keys and values from the memory; its weights have shape (1, 4, 6, 5). An arrow brings the memory from the encoder to this box. After the decoder layers, a linear layer multiplies by the transposed embedding table, with no new weights, giving logits of shape (1, 6, 13). Softmax gives the next token at each position: d e b a c and the end token. A note compares the sizes: this toy has d_model 64, 4 heads and d_ff 256; the paper’s base model has 512, 8 heads, 2,048 and 6 + 6 layers.

The next sections follow the arrows: the paper’s formula, its meaning, a number to check by hand, and a check in code.

Step 1: token embeddings, times √d_model

Each token ID picks its row from an embedding table, as in Tokens and embeddings. The paper adds one detail (§3.4): “In the embedding layers, we multiply those weights by √d_model.”

$$x_{\text{pos}} = \sqrt{d_{\text{model}}}\; E[\text{token}_{\text{pos}}] + PE(\text{pos})$$

Here E is the embedding table, with one row of d_model numbers per token. token_pos is the token at position pos, and PE(pos) is the positional encoding of step 2. In words: look up the row, scale it by √d_model, and add the position’s pattern.

Why scale? The toy’s table starts with a spread of 1/√64 = 0.125, while positions lie between −1 and 1. Without the factor, the positions would drown the token. The token “c” starts with −0.023, −0.006, −0.027 and 0.023; times 8, that is about −0.186, −0.045, −0.219 and 0.187.

The paper uses one table for the source, the target and the output layer of step 5, because both languages share one vocabulary.

Step 2: sinusoidal positional encoding

Position and masking gave each position a learned vector. The paper uses a fixed pattern of waves instead (§3.5). For position pos and column pair i (so i goes from 0 to d_model/2 − 1):

$$PE(\text{pos}, 2i) = \sin\left(\frac{\text{pos}}{10000^{2i/d_{\text{model}}}}\right) \qquad PE(\text{pos}, 2i+1) = \cos\left(\frac{\text{pos}}{10000^{2i/d_{\text{model}}}}\right)$$

In words: column pair i is a sine and a cosine of the position, slowed down by the factor 10000^(2i/d_model). The first pair turns fast, and every later pair turns more slowly.

Check by hand at position 1. For i = 0 the divisor is 10000⁰ = 1, so the first two numbers are sin 1 = 0.841 and cos 1 = 0.540. For i = 1 and d_model = 64, the divisor is 10000^(2/64) = 1.334, so the angle is 1 / 1.334 = 0.750. The next two numbers are sin 0.750 = 0.682 and cos 0.750 = 0.732.

positions.py

import math
import torch

def sinusoidal_positions(n, d):
    pos = torch.arange(n, dtype=torch.float).unsqueeze(1)    # (n, 1): positions 0, 1, 2, ...
    two_i = torch.arange(0, d, 2, dtype=torch.float)         # 2i = 0, 2, 4, ..., d - 2
    angle = pos / torch.pow(10000.0, two_i / d)              # (n, d/2): one angle per pair
    pe = torch.zeros(n, d)
    pe[:, 0::2] = torch.sin(angle)                           # even columns 2i: sine
    pe[:, 1::2] = torch.cos(angle)                           # odd columns 2i + 1: cosine
    return pe

pe = sinusoidal_positions(50, 64)
print(tuple(pe.shape))
print(pe[1, :4].round(decimals=3))
print(round(math.sin(1), 3), round(math.cos(1), 3), round(math.sin(1 / 10000 ** (2 / 64)), 3))

The output:

(50, 64)
tensor([0.8410, 0.5400, 0.6820, 0.7320])
0.841 0.54 0.682

0::2 takes every second column from column 0, and 1::2 from column 1. The table does not depend on the tokens, so it is computed once.

• Positive
• Negative

The sinusoidal positional encoding for 50 positions and d_model = 64, computed from the paper's formula. Columns on the left change fast from row to row; columns on the right change slowly.

Text description of the chart

A grid with 50 rows, the positions 0 to 49, and 64 columns. Blue cells are positive and orange cells are negative; the stronger the colour, the closer the value is to +1 or −1. In the first columns, the colour changes from row to row in short stripes, because these sines and cosines turn fast. Further right, the stripes get longer. In the last 20 or so columns, the colour hardly changes over the 50 positions: the even columns stay near 0 and the odd columns, the cosines, stay near +1 (strong blue). Each row is a different pattern, so each position gets its own fingerprint.

Each column is a wave. Its wavelength, the number of positions for one full turn, is 2π · 10000^(2i/d_model). The paper says (§3.5): “The wavelengths form a geometric progression from 2π to 10000 · 2π.”

for d in (64, 512):
    lengths = [2 * math.pi * 10000 ** (2 * i / d) for i in (0, 1, d // 4, d // 2 - 1)]
    print(d, [f"{v:,.1f}" for v in lengths])

The output:

64 ['6.3', '8.4', '628.3', '47,117.2']
512 ['6.3', '6.5', '628.3', '60,611.5']

Fast columns tell close positions apart, like a clock’s second hand; slow columns tell distant ones apart, like the hour hand.

• column 0 (wavelength 6.3)
• column 8 (wavelength 19.9)
• column 20 (wavelength 111.7)

Columns 0, 8 and 20 of the positional encoding (d_model = 64), from the formula. The model uses the values at whole positions; the curves join them.

Text description of the chart

Three sine curves between −1 and +1 over the positions 0 to 60. Column 0 has a wavelength of 6.3 positions, so it goes up and down almost ten times. Column 8 has a wavelength of 19.9 positions and makes three turns. Column 20 has a wavelength of 111.7 positions; it rises slowly from 0 and reaches its top only after about 28 positions.

Why waves? Relative positions become a rotation. The paper chose this pattern “because we hypothesized it would allow the model to easily learn to attend by relative positions, since for any fixed offset k, PE_pos+k can be represented as a linear function of PE_pos.” Take one column pair with frequency ω = 1/10000^(2i/d_model). By the angle-addition rules of sine and cosine:

$$\begin{pmatrix} \sin(\omega(\text{pos}+k)) \\ \cos(\omega(\text{pos}+k)) \end{pmatrix} = \begin{pmatrix} \cos \omega k & \sin \omega k \\ -\sin \omega k & \cos \omega k \end{pmatrix} \begin{pmatrix} \sin \omega\,\text{pos} \\ \cos \omega\,\text{pos} \end{pmatrix}$$

The 2 × 2 matrix is a rotation by the angle ωk. It depends on the offset k only, not on pos. So “3 positions later” is the same matrix at position 0, 10 or 40:

k, i = 3, 5                                   # move 3 positions; look at the pair i = 5 (columns 10, 11)
w = 1 / 10000 ** (2 * i / 64)                 # the frequency of this pair
turn = torch.tensor([[math.cos(w * k), math.sin(w * k)],
                     [-math.sin(w * k), math.cos(w * k)]])
for pos in (0, 10, 40):
    print(pos, (turn @ pe[pos, 10:12]).round(decimals=4), pe[pos + k, 10:12].round(decimals=4))

The output:

0 tensor([0.6529, 0.7574]) tensor([0.6529, 0.7574])
10 tensor([ 0.0588, -0.9983]) tensor([ 0.0588, -0.9983])
40 tensor([-0.6977, -0.7164]) tensor([-0.6977, -0.7164])

The turned pair equals the pair 3 positions later, at every position, so attention can learn to look “k tokens back”. The paper also tried learned positions and “found that the two versions produced nearly identical results”. It kept the sinusoids because they “may allow the model to extrapolate to sequence lengths longer than the ones encountered during training”.

At position 0 every sine is 0 and every cosine is 1. So the encoder input of “c” is the scaled embedding plus (0, 1, 0, 1, …): −0.186, 0.955, −0.219, 1.187.

Step 3: the encoder layer

An encoder layer has two sublayers: multi-head self-attention, then a position-wise feed-forward network. The paper wraps each one in the same way (§3.1): “the output of each sub-layer is LayerNorm(x + Sublayer(x))”. With dropout, used in training (§5.4):

$$\text{output} = \text{LayerNorm}\big(x + \text{Dropout}(\text{Sublayer}(x))\big)$$

In words: compute the change, drop 10% of it in training, add the input (the residual connection), then normalize. The norm comes after the addition: post-norm.

The first sublayer is the multi-head attention of The mathematics of attention, with Q, K and V all made from the same input X, one row per token:

$$\text{MultiHead}(X) = [\text{head}_1, \ldots, \text{head}_h]\, W_O, \qquad \text{head}_i = \text{softmax}\!\left(\frac{X W_Q^{(i)} \,(X W_K^{(i)})^\top}{\sqrt{d_k}}\right) X W_V^{(i)}$$

W_Q^(i), W_K^(i) and W_V^(i) are head i’s weight matrices, each d_model × d_k, and W_O mixes the joined heads back to d_model numbers. (The paper writes W_i^Q and W^O; this module keeps the notation of the previous lesson.) With h heads, d_k = d_model / h: 512 / 8 = 64 in the paper, 64 / 4 = 16 in the toy.

LayerNorm rescales each token’s d_model numbers on their own:

$$\text{LayerNorm}(x) = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta, \qquad \mu = \frac{1}{d}\sum_{i=1}^{d} x_i, \quad \sigma^2 = \frac{1}{d}\sum_{i=1}^{d} (x_i - \mu)^2$$

μ is the mean of the token’s numbers and σ² their variance (divided by d, not by d − 1). ε (10⁻⁵ in PyTorch) avoids dividing by 0. γ and β are a learned scale and shift, starting at 1 and 0; ⊙ multiplies number by number.

Check by hand. Take x = (1, 2, 3, 6). The mean is 12 / 4 = 3. The variance is (4 + 1 + 0 + 9) / 4 = 3.5, and √3.5 = 1.871. So 1 becomes (1 − 3) / 1.871 = −1.069, and 6 becomes 3 / 1.871 = 1.604.

x = torch.tensor([1.0, 2.0, 3.0, 6.0])
mu = x.mean()
var = ((x - mu) ** 2).mean()                  # divide by d, not by d - 1
by_hand = (x - mu) / torch.sqrt(var + 1e-5)   # then times gamma (1 at the start) plus beta (0)
print(mu.item(), var.item())
print(by_hand.round(decimals=3))
print(torch.nn.LayerNorm(4)(x).detach().round(decimals=3))

The output:

3.0 3.5
tensor([-1.0690, -0.5350,  0.0000,  1.6040])
tensor([-1.0690, -0.5350,  0.0000,  1.6040])

The second sublayer is the feed-forward network (§3.3), applied to each position separately and identically:

$$\text{FFN}(x) = \max(0,\; x W_1 + b_1)\, W_2 + b_2$$

W₁ is d_model × d_ff and W₂ is d_ff × d_model; b₁ and b₂ are their biases. max(0, ·) is ReLU. In the paper, 512 → 2,048 → 512. In the toy, 64 → 256 → 64, so the hidden tensor has shape (1, 5, 256). After training, ReLU sets 92.6% of the token “c”’s 256 hidden numbers to 0 in encoder layer 1: each token uses a small part of the layer.

Now the whole layer from these formulas, against PyTorch’s post-norm nn.TransformerEncoderLayer(norm_first=False). The function reads PyTorch’s weights; in_proj_weight holds W_Q, W_K and W_V stacked.

encoder_layer.py

from torch import nn

torch.manual_seed(0)
d, h, d_ff = 64, 4, 256
ref = nn.TransformerEncoderLayer(d, h, d_ff, dropout=0.1, activation="relu",
                                 batch_first=True, norm_first=False).eval()

def layer_norm(x, norm):                      # (x - mean) / sqrt(var + eps) * gamma + beta
    mu = x.mean(-1, keepdim=True)
    var = x.var(-1, unbiased=False, keepdim=True)
    return (x - mu) / torch.sqrt(var + norm.eps) * norm.weight + norm.bias

def encoder_layer(x, L):                      # L: PyTorch's layer, read only for its weights
    n = x.shape[1]
    Q, K, V = (x @ L.self_attn.in_proj_weight.T + L.self_attn.in_proj_bias).split(d, dim=-1)
    Q, K, V = (t.view(1, n, h, d // h).transpose(1, 2) for t in (Q, K, V))   # (1, 4, n, 16)
    weights = torch.softmax(Q @ K.transpose(-2, -1) / math.sqrt(d // h), dim=-1)
    heads = (weights @ V).transpose(1, 2).reshape(1, n, d)                     # concatenate
    attn = heads @ L.self_attn.out_proj.weight.T + L.self_attn.out_proj.bias   # W_O
    x = layer_norm(x + attn, L.norm1)                                          # add, then norm
    ffn = torch.relu(x @ L.linear1.weight.T + L.linear1.bias) @ L.linear2.weight.T + L.linear2.bias
    return layer_norm(x + ffn, L.norm2)                                        # add, then norm

x = torch.randn(1, 5, d)
with torch.no_grad():
    mine, theirs = encoder_layer(x, ref), ref(x)
print(tuple(mine.shape), f"largest difference {(mine - theirs).abs().max().item():.1e}")

The output:

(1, 5, 64) largest difference 3.6e-07

3.6 × 10⁻⁷ is rounding noise: the formulas are the layer. (.eval() switches dropout off.)

Most models today put the norm before the sublayer. The last section says why.

Post-norm, as in the paper, and pre-norm, as in most later models. Only the place of the norm changes.

Text description of the diagram

Left: post-norm, the order of the 2017 paper, written LayerNorm(x + Sublayer(x)). The input x goes into the sublayer, which is attention or the feed-forward layer. A plus sign adds the sublayer’s output and x, which a dashed line carries around the sublayer: the residual connection. Then a LayerNorm gives the output. The note says that the norm sits on the main path, so every layer’s output is rescaled.

Right: pre-norm, the order of most models today, written x + Sublayer(Norm(x)). The input x goes into an RMSNorm or a LayerNorm, then into the sublayer. A plus sign adds x, which the dashed residual line carries around both boxes. The note says that the residual path from the input to the output has no norm on it: a clean highway.

Step 4: the decoder layer and cross-attention

A decoder layer has three sublayers, each wrapped in the same add-and-norm.

1. Masked self-attention. The decoder’s tokens read each other, with the causal mask M of The mathematics of attention, added to the scores before softmax. Position 1, the start token, reads only itself.
2. Encoder–decoder attention, or cross-attention. The queries come from the decoder, and the keys and values come from the encoder’s memory M_enc (see the formula below). M_enc has one row per source token; it is not the mask M.
3. The feed-forward network, the same formula as in the encoder.

$$\text{CrossAttention}(Y, M_{\text{enc}}) = \text{softmax}\!\left(\frac{(Y W_Q)(M_{\text{enc}} W_K)^\top}{\sqrt{d_k}}\right) M_{\text{enc}} W_V$$

Here Y holds the decoder’s vectors, one row per target position, after the masked self-attention; the formula shows one head. In words: each target position asks a question (its query) and looks for the answer in the source sentence (the keys). Then it takes what the source offers (the values). With 6 target positions and 5 source tokens, the scores are 6 × 5 per head, with no causal mask: the whole source is known.

torch.manual_seed(0)
cross = nn.MultiheadAttention(64, 4, batch_first=True)
memory = torch.randn(1, 5, 64)                # the encoder's output: 5 source tokens
y = torch.randn(1, 6, 64)                     # the decoder: 6 target positions
out, w = cross(query=y, key=memory, value=memory, average_attn_weights=False)
print("output ", tuple(out.shape))
print("weights", tuple(w.shape))
print(w[0, 0, 0].detach().round(decimals=3), "sum", round(w[0, 0, 0].sum().item(), 3))

The output:

output  (1, 6, 64)
weights (1, 4, 6, 5)
tensor([0.1440, 0.2300, 0.1400, 0.3070, 0.1800]) sum 1.0

One output row per decoder position, filled with information from the source: this is how the source sentence enters every step of writing.

Step 5: the output layer, tied to the embedding

After the last decoder layer, each position has d_model numbers. A linear layer turns them into one score (a logit) per vocabulary token, and softmax turns the scores into probabilities. The paper shares this layer’s matrix with the embedding table (§3.4):

$$P(\text{next token}) = \text{softmax}\big(y_{\text{pos}}\, E^\top\big)$$

y_pos is the last decoder layer’s vector at one position, a row of d_model numbers. E is the same n_vocab × d_model table as in step 1, here 13 × 64 (n_vocab is the number of tokens in the vocabulary), and Eᵀ is its transpose. So (1, 6, 64) @ (64, 13) → (1, 6, 13). In words: a token’s score is the dot product of the decoder’s vector with that token’s own embedding. At position 1 of the example, the score of “d” is 5.22, and after softmax its probability is 0.916.

Follow one example through every step

Now put the five steps together. The explorer follows “c a b e d” through the trained toy model: the input, the encoder, the decoder and the output. Each step shows the real shape and a few real numbers.

Explore

One example through the original transformer

Step 1 of 14: Source tokens, shape (1, 5).

1. 1. Input (this step)
   2. Encoder
   3. Decoder
   4. Output

   Step 1 of 14 Source tokens

   Shape (1, 5)

   The source c a b e d becomes 5 token IDs from the shared vocabulary of 13 tokens.

   Token and its ID

   c	a	b	e	d
   5	3	4	7	6

2. 1. Input (this step)
   2. Encoder
   3. Decoder
   4. Output

   Step 2 of 14 Embedding lookup

   Shape (1, 5, 64)

   Each ID picks its row of 64 numbers from the shared table. These are the first 4 numbers for c. They are small: the table started with a spread of 1/√64 = 0.125.

   The first 4 of 64 numbers

   c, embedding	−0.023	−0.006	−0.027	0.023	…

3. 1. Input (this step)
   2. Encoder
   3. Decoder
   4. Output

   Step 3 of 14 Multiply by √d_model = 8

   Shape (1, 5, 64)

   The paper multiplies the embeddings by √d_model, so they are not drowned by the positional encoding, whose numbers lie between −1 and 1.

   The first 4 of 64 numbers

   c, embedding	−0.023	−0.006	−0.027	0.023	…
   c, × 8	−0.186	−0.045	−0.219	0.187	…

4. 1. Input (this step)
   2. Encoder
   3. Decoder
   4. Output

   Step 4 of 14 Add the positional encoding

   Shape (1, 5, 64)

   c is at position 0, where the encoding is sin 0 = 0 and cos 0 = 1 in every pair. Position 1 starts with sin 1 = 0.841 and cos 1 = 0.540.

   The first 4 of 64 numbers

   c, × 8	−0.186	−0.045	−0.219	0.187	…
   PE, position 0	0.000	1.000	0.000	1.000	…
   encoder input	−0.186	0.955	−0.219	1.187	…
   PE, position 1	0.841	0.540	0.682	0.732	…

5. 1. Input
   2. Encoder (this step)
   3. Decoder
   4. Output

   Step 5 of 14 Queries, keys and values

   Shape (1, 4, 5, 16)

   Three linear layers make Q, K and V from the encoder input, and each is split into 4 heads of 16 numbers.

6. 1. Input
   2. Encoder (this step)
   3. Decoder
   4. Output

   Step 6 of 14 Encoder self-attention (layer 1, head 1)

   Shape (1, 4, 5, 5)

   Each source letter reads every source letter. In layer 1 the weights are nearly equal, between 0.15 and 0.23: this head mostly averages.

   Rows: the position that reads. Columns: what it reads. Each row adds up to 1.

   	c	a	b	e	d
   c	0.17	0.20	0.18	0.22	0.23
   a	0.19	0.20	0.15	0.22	0.23
   b	0.21	0.20	0.21	0.19	0.19
   e	0.18	0.22	0.16	0.22	0.22
   d	0.20	0.20	0.17	0.20	0.23

7. 1. Input
   2. Encoder (this step)
   3. Decoder
   4. Output

   Step 7 of 14 Add & norm

   Shape (1, 5, 64)

   The 4 heads are joined (4 × 16 = 64 numbers), mixed by W_O, added to the input, and normalized: LayerNorm(x + Sublayer(x)).

   The first 4 of 64 numbers

   c, encoder input	−0.186	0.955	−0.219	1.187	…
   c, after add & norm	−0.469	1.800	−0.529	1.894	…

8. 1. Input
   2. Encoder (this step)
   3. Decoder
   4. Output

   Step 8 of 14 Feed-forward, then add & norm

   Shape (1, 5, 64)

   Each token alone goes 64 → 256 → 64. The hidden layer has shape (1, 5, 256); ReLU sets 92.6% of c's 256 hidden numbers to 0. Then add and norm again.

   The first 4 of 64 numbers

   c, after add & norm	−0.469	1.800	−0.529	1.894	…
   c, after layer 1	−0.537	1.412	−0.620	1.863	…

9. 1. Input
   2. Encoder (this step)
   3. Decoder
   4. Output

   Step 9 of 14 Encoder output: the memory

   Shape (1, 5, 64)

   After 2 encoder layers, each source letter has 64 numbers that also carry information about the other letters. The decoder reads this memory in every layer.

   The first 4 of 64 numbers

   c, after layer 1	−0.537	1.412	−0.620	1.863	…
   c, memory	−0.670	1.164	−0.303	2.715	…

10. 1. Input
    2. Encoder
    3. Decoder (this step)
    4. Output

    Step 10 of 14 Decoder input, shifted right

    Shape (1, 6)

    The decoder reads <s> and then the answer so far. In training, the whole answer d e b a c is given at once, shifted one place to the right.

    Token and its ID

    <s>	d	e	b	a	c
    1	6	7	4	3	5

11. 1. Input
    2. Encoder
    3. Decoder (this step)
    4. Output

    Step 11 of 14 Masked self-attention (layer 1, head 1)

    Shape (1, 4, 6, 6)

    The causal mask hides every later position. <s> can read only itself (weight 1); d reads <s> and itself (0.5 and 0.5).

    Rows: the position that reads. Columns: what it reads. Each row adds up to 1.

    	<s>	d	e	b	a	c
    <s>	1.00	—hidden	—hidden	—hidden	—hidden	—hidden
    d	0.50	0.50	—hidden	—hidden	—hidden	—hidden
    e	0.29	0.37	0.34	—hidden	—hidden	—hidden
    b	0.26	0.24	0.24	0.26	—hidden	—hidden
    a	0.18	0.20	0.19	0.17	0.25	—hidden
    c	0.18	0.14	0.15	0.19	0.15	0.19

12. 1. Input
    2. Encoder
    3. Decoder (this step)
    4. Output

    Step 12 of 14 Cross-attention (layer 2, mean of 4 heads)

    Shape (1, 4, 6, 5)

    Queries come from the decoder, keys and values from the memory. Each decoder position reads one source letter: position 1 reads letter 5 (d), position 2 reads letter 4 (e), and so on. An anti-diagonal: the model learned where to look.

    Rows: the position that reads. Columns: what it reads. Each row adds up to 1.

    	c	a	b	e	d
    <s> → d	0.00	0.00	0.00	0.00	1.00
    d → e	0.00	0.00	0.00	1.00	0.00
    e → b	0.01	0.04	0.94	0.01	0.00
    b → a	0.03	0.96	0.01	0.00	0.00
    a → c	1.00	0.00	0.00	0.00	0.00
    c → </s>	0.44	0.01	0.23	0.07	0.26

13. 1. Input
    2. Encoder
    3. Decoder
    4. Output (this step)

    Step 13 of 14 Output layer, tied to the embedding

    Shape (1, 6, 13)

    Each position's 64 numbers are multiplied by the transposed embedding table: one score for each of the 13 tokens. At position 1, the score of d is 5.223, far above the others.

    The first 4 of 64 numbers

    position 1, first 4 scores	0.371	0.326	0.317	0.370	…

14. 1. Input
    2. Encoder
    3. Decoder
    4. Output (this step)

    Step 14 of 14 Softmax: the next token at each position

    Shape (1, 6, 13)

    Softmax turns each row of 13 scores into probabilities. The most likely tokens are d e b a c </s>: the reversed source. No probability passes 0.93, because of label smoothing.

    Most likely next token at each decoder position

    	1	2	3	4	5	6
    token	d	e	b	a	c	</s>
    probability	0.916	0.916	0.929	0.927	0.924	0.927

The trained toy model on the source c a b e d, in eval mode (no dropout). Shapes are (batch, tokens, numbers) or (batch, heads, queries, keys). Every number comes from the real model.

Text description of the explorer

The explorer has 14 steps, which a slider and two buttons move through. Each step names its stage (input, encoder, decoder or output), the shape of the tensor and the real numbers.

Step 1, source tokens, shape (1, 5): c, a, b, e and d have the IDs 5, 3, 4, 7 and 6. Step 2, embedding lookup, (1, 5, 64): the first 4 numbers for c are −0.023, −0.006, −0.027 and 0.023. Step 3, times 8: −0.186, −0.045, −0.219 and 0.187. Step 4, plus the positional encoding of position 0, which is 0, 1, 0, 1: the encoder input is −0.186, 0.955, −0.219 and 1.187. Position 1’s encoding starts 0.841, 0.540, 0.682 and 0.732.

Step 5, queries, keys and values, (1, 4, 5, 16). Step 6, encoder self-attention weights of layer 1, head 1, (1, 4, 5, 5): every weight lies between 0.15 and 0.23, so this head mostly averages. Step 7, add and norm, (1, 5, 64): c becomes −0.469, 1.800, −0.529 and 1.894. Step 8, the feed-forward layer, with a hidden shape of (1, 5, 256), where ReLU sets 92.6% of c’s hidden numbers to 0; after add and norm, c is −0.537, 1.412, −0.620 and 1.863. Step 9, the memory after 2 encoder layers, (1, 5, 64): −0.670, 1.164, −0.303 and 2.715.

Step 10, the decoder input, shifted right, (1, 6): the start token, d, e, b, a and c, with the IDs 1, 6, 7, 4, 3 and 5. Step 11, masked self-attention of decoder layer 1, head 1, (1, 4, 6, 6): the upper triangle is hidden; the start token reads only itself with weight 1, and d reads the start token and itself with 0.5 each. Step 12, the cross-attention of decoder layer 2, the mean of the 4 heads, (1, 4, 6, 5): the start token’s row puts 0.999 on source letter 5 (d); d’s row 0.999 on letter 4 (e); e’s row 0.938 on letter 3 (b); b’s row 0.962 on letter 2 (a); a’s row 1.000 on letter 1 (c). The last row, which must write the end token, spreads: 0.44, 0.01, 0.23, 0.07 and 0.26. Step 13, the output layer tied to the embedding, (1, 6, 13): at position 1 the score of d is 5.223. Step 14, softmax: the most likely tokens are d, e, b, a, c and the end token, with probabilities 0.916, 0.916, 0.929, 0.927, 0.924 and 0.927.

Step 12 answers the prediction. To write the first letter of the reversed text, position 1 must read the last source letter, d, at position 5. And it does, with weight 0.999. The next position reads letter 4, then 3, 2 and 1. Nobody told the model to do this: training found that reading the mirrored position is the way to reverse.

Step 6: the training recipe

The paper trains with four choices (§5.3 and §5.4).

Adam with unusual settings. β₁ = 0.9, β₂ = 0.98 and ε = 10⁻⁹. The optimizer is the Adam of Learning rate and optimizers; β₂ = 0.98 instead of 0.999 gives it a shorter memory.

A warmup schedule. The learning rate is not fixed. It changes at every step:

$$\text{lrate} = d_{\text{model}}^{-0.5} \cdot \min\left(\text{step}^{-0.5},\; \text{step} \cdot \text{warmup\_steps}^{-1.5}\right)$$

In words: the rate grows in a straight line for warmup_steps steps, then falls as 1/√step. The two terms are equal at the peak, step = warmup_steps. The paper uses warmup_steps = 4,000. Check by hand: 512^−0.5 = 0.0442 and 4000^−0.5 = 0.0158, so the peak is 0.0442 × 0.0158 = 6.99 × 10⁻⁴.

def lrate(step, d_model=512, warmup=4000):
    return d_model ** -0.5 * min(step ** -0.5, step * warmup ** -1.5)

for step in (1, 100, 1000, 4000, 10000, 100000):
    print(f"step {step:>7,}   lrate {lrate(step):.3e}")

The output:

step       1   lrate 1.747e-07
step     100   lrate 1.747e-05
step   1,000   lrate 1.747e-04
step   4,000   lrate 6.988e-04
step  10,000   lrate 4.419e-04
step 100,000   lrate 1.398e-04

The paper's learning-rate schedule for d_model = 512 and 4,000 warmup steps, the first 40,000 of 100,000 steps, from the formula.

Text description of the chart

A line chart of the learning rate against the training step, from 0 to 40,000. The line rises straight from almost 0 to its peak of 6.99 × 10⁻⁴ at step 4,000. Then it falls, quickly at first and then more and more slowly: 4.42 × 10⁻⁴ at step 10,000 and about 2.2 × 10⁻⁴ at step 40,000. At step 100,000, the end of the paper’s training, it is 1.40 × 10⁻⁴.

Why warm up? At the start, Adam’s estimates of the gradients’ size rest on very few steps, and large steps can push a post-norm model into a region it does not recover from.

Dropout 0.1 on every sublayer’s output and on the embeddings plus positions, as in Overfitting controls.

Label smoothing 0.1. The target is not “100% the right token” but a softer version. PyTorch spreads ε = 0.1 over all n_vocab tokens and puts the rest on the right one:

$$q(k) = (1 - \epsilon)\,[k = \text{right}] + \frac{\epsilon}{n_{\text{vocab}}}$$

Here k is any token of the vocabulary, and [k = right] is 1 for the right token and 0 for the others. With n_vocab = 13, the right token gets 0.9 + 0.1 / 13 = 0.908, and each other token 0.1 / 13 = 0.0077. The loss is the cross-entropy against q. It is smallest when the model predicts exactly q, and that smallest value is the entropy of q, not 0:

n_vocab, eps, right = 13, 0.1, 5
target = torch.full((n_vocab,), eps / n_vocab)  # PyTorch spreads eps over all tokens
target[right] += 1 - eps                      # and puts the rest on the right one
print(round(target[right].item(), 4), round(target[0].item(), 4))
loss_fn = nn.CrossEntropyLoss(label_smoothing=eps)
best = loss_fn(target.log().unsqueeze(0), torch.tensor([right]))      # predict exactly the target
sure = loss_fn(torch.tensor([[0.0] * right + [20.0] + [0.0] * (n_vocab - right - 1)]), torch.tensor([right]))
print(f"best possible loss {best.item():.3f}; a model 100% sure: {sure.item():.3f}")

The output:

0.9077 0.0077
best possible loss 0.537; a model 100% sure: 1.846

So a perfect toy model still has a smoothed loss of 0.537, and a model that is 100% sure pays more (1.846) than one that is 91% sure. This is why the toy’s top probabilities stay near 0.92 and never reach 0.99. The paper is honest about the trade (§5.4): “This hurts perplexity, as the model learns to be more unsure, but improves accuracy and BLEU score.”

Step 7: how big is the base model?

The paper’s Table 3 gives the base model 65 × 10⁶ parameters, with N = 6, d_model = 512, h = 8, d_k = d_v = 64, d_ff = 2,048 and a dropout of 0.1. Count them from the formulas: attention 4d² + 4d, the FFN two matrices and biases, each LayerNorm 2d. An encoder layer has 1 attention and 2 norms; a decoder layer 2 and 3. Like PyTorch, the count gives every attention matrix a bias; the paper’s attention formulas show none, which would remove only 36,864 parameters.

d, d_ff, N, vocab = 512, 2048, 6, 37000
attention = 4 * d * d + 4 * d                 # W_Q, W_K, W_V, W_O and their biases
ffn = d * d_ff + d_ff + d_ff * d + d          # W1, b1, W2, b2
norm = 2 * d                                  # gamma and beta
encoder_layer_n = attention + ffn + 2 * norm
decoder_layer_n = 2 * attention + ffn + 3 * norm
embedding = vocab * d                         # one matrix: source, target and output
total = embedding + N * (encoder_layer_n + decoder_layer_n)
print(f"attention {attention:,}  ffn {ffn:,}  encoder layer {encoder_layer_n:,}  decoder layer {decoder_layer_n:,}")
print(f"total {total:,}; with a separate output matrix {total + vocab * d:,}")

The output:

attention 1,050,624  ffn 2,099,712  encoder layer 3,152,384  decoder layer 4,204,032
total 63,082,496; with a separate output matrix 82,026,496

Our count is 63,082,496, about 63 million, against the paper’s rounded 65 million. The paper gives the vocabulary only as “about 37,000”, so a small gap is expected. A separate output matrix would give 82.0 million, far from 65: the sharing is real. Note where the parameters are. The shared embedding holds 18.9 million, and each decoder layer is a third larger than an encoder layer because of its cross-attention. The paper trained this model for 100,000 steps, about 12 hours on 8 NVIDIA P100 GPUs.

The toy run: it learns to reverse

Now train the toy model with the whole recipe. The only change is the warmup: 400 steps instead of 4,000, because the run has 3,000 steps instead of 100,000. The peak learning rate is then 64^−0.5 × 400^−0.5 = 0.00625. Each step uses 64 random sequences of 3 to 8 letters. exact_match counts the 500 test sequences that come back completely right. The notebook has the model classes and batch().

train_reversal.py

seed_all(0)
gen = torch.Generator().manual_seed(0)                        # the training batches
test_src, test_in, test_out = batch(500, torch.Generator().manual_seed(1234))
model = Transformer(vocab=13, d=64, h=4, d_ff=256, n_layers=2, p=0.1)
opt = torch.optim.Adam(model.parameters(), lr=1.0, betas=(0.9, 0.98), eps=1e-9)
sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: lrate(s + 1, 64, 400))
loss_fn = nn.CrossEntropyLoss(ignore_index=PAD, label_smoothing=0.1)
for step in range(1, 3001):
    model.train()
    src, tgt_in, tgt_out = batch(64, gen)                      # 64 random sequences
    loss = loss_fn(model(src, tgt_in).reshape(-1, 13), tgt_out.reshape(-1))
    opt.zero_grad()
    loss.backward()
    opt.step()
    sched.step()
    if step in (1, 500, 1000, 2000, 3000):
        print(f"step {step:>5}  lr {opt.param_groups[0]['lr']:.2e}  train loss {loss.item():.3f}"
              f"  exact match {exact_match(model, test_src, test_out):.1%}")

The output:

step     1  lr 3.13e-05  train loss 5.254  exact match 0.0%
step   500  lr 5.58e-03  train loss 0.741  exact match 88.8%
step  1000  lr 3.95e-03  train loss 0.572  exact match 98.2%
step  2000  lr 2.79e-03  train loss 0.558  exact match 99.6%
step  3000  lr 2.28e-03  train loss 0.540  exact match 100.0%

LambdaLR multiplies the base rate 1.0 by lrate() at every step. The decoder gets the right answer shifted by one (teacher forcing). The 3,000 steps took 34 seconds on 2 CPU threads of a shared computer.

• Test loss (plain cross-entropy)
• Training loss (with label smoothing)
• Floor of the smoothed loss, 0.537

The toy run, measured every 100 steps from step 100 (at step 1 both losses are above 5, off the chart). Labels give the exact match on the 500 test sequences.

Text description of the chart

A line chart with the training step, 100 to 3,000, on the horizontal axis and the loss, 0 to 2, on the vertical axis. The test loss, plain cross-entropy, falls from 1.68 at step 100 to 0.31 at step 200 and 0.13 at step 500, then slowly to 0.079 at step 3,000. The training loss, with label smoothing, falls from 1.90 at step 100 to about 0.7 at step 300 and then flattens just above a dashed line at 0.537, the floor that label smoothing sets; it ends at 0.540. Labels on the test loss give the exact match: 88.8% at step 500, 98.2% at step 1,000 and 100.0% at step 3,000.

The smoothed training loss cannot go below 0.537, so its flat end is success, not a problem.

Greedy decoding. At inference, there is no answer to shift. The encoder runs once. Then the decoder starts with <s> and adds the most likely token at each step, as in the generation loop of Training versus inference.

model.eval()
src = torch.tensor([[STOI[c] for c in "cabed"]])
with torch.no_grad():
    memory = model.encode(src)                                 # the encoder runs once
    ys = torch.tensor([[BOS]])                                 # the decoder starts with <s>
    for _ in range(10):
        p = model.decode(ys, memory, src)[0, -1].softmax(-1)   # the last position's prediction
        nxt = int(p.argmax())
        print(f"{' '.join(ITOS[i] for i in ys[0].tolist()):<14} -> {ITOS[nxt]:<5} {p[nxt].item():.3f}")
        ys = torch.cat([ys, torch.tensor([[nxt]])], dim=1)
        if nxt == EOS:
            break

The output:

<s>            -> d     0.916
<s> d          -> e     0.916
<s> d e        -> b     0.929
<s> d e b      -> a     0.927
<s> d e b a    -> c     0.924
<s> d e b a c  -> </s>  0.927

The runners-up of each step have less than 0.01:

Decoder input so far	Chosen	Top 3 (probability)
<s>	d	d 0.916 · c 0.008 · f 0.007
<s> d	e	e 0.916 · j 0.008 · g 0.008
<s> d e	b	b 0.929 · c 0.006 · d 0.006
<s> d e b	a	a 0.927 · d 0.007 · j 0.006
<s> d e b a	c	c 0.924 · </s> 0.007 · <s> 0.007
<s> d e b a c	</s>	</s> 0.927 · c 0.006 · <s> 0.006

Where the model looks. The cross-attention of both decoder layers for this example, averaged over the 4 heads:

Cross-attention weights of the trained toy model for c a b e d, mean of 4 heads. In layer 2, the largest weight of each row (ringed) forms an anti-diagonal: position 1 reads source letter 5, position 2 reads letter 4, and so on.

Text description of the chart

Two grids of 6 rows and 5 columns. The rows are the decoder positions, each with the token it must write: the start token writes d, d writes e, e writes b, b writes a, a writes c, and c writes the end token. The columns are the source letters c, a, b, e and d, at positions 1 to 5.

Layer 1 is blurred. The largest weight of each row is only 0.28 to 0.47: the start token’s row puts 0.37 on d and 0.32 on c; the other rows spread over two or three letters.

Layer 2 is sharp. The start token’s row puts 0.999 on d (position 5). d’s row puts 0.999 on e (position 4). e’s row puts 0.938 on b (position 3). b’s row puts 0.962 on a (position 2). a’s row puts 1.000 on c (position 1). These cells form a line from the top right to the bottom left: the anti-diagonal. The last row, which must write the end token, spreads its weight: 0.44 on c, 0.23 on b and 0.26 on d.

Layer 1 has a rough idea; layer 2 reads exactly one mirrored letter per position. The </s> row has no letter to copy, so it spreads. As Attention step by step warned, weights show where information came from, not why; on this simple task the two agree.

Our code equals PyTorch’s nn.Transformer

With norm_first=False, ReLU, dropout off and our weights copied in, nn.Transformer gives the same output as our model to 9.5 × 10⁻⁷, on a padded batch with both masks. Two traps made the check fail first: PyTorch’s extra final LayerNorm after each stack (see the failure case below), and an eval-mode fast path for padded batches that returns 0 at padded positions (switch it off with use_nested_tensor = False).

What it means for Larkfield

Larkfield’s router reads a ticket and picks a team. That is the encoder’s job alone: the whole ticket is known, and nothing has to be written. So the course’s TinyTransformer (291,397 parameters, 94.3% right on the validation tickets) is an encoder only, as BERT is. It differs from the paper’s encoder in two choices: it learns its positions, and it normalizes before each sublayer (norm_first=True). A decoder would be needed only to write, for example a reply.

What changed since 2017

The paper’s parts are in every large language model, but almost every detail changed. The table compares the base model with SmolLM2-135M from Training versus inference, using its configuration and the course’s measured facts.

Part	Transformer (2017, base)	SmolLM2-135M	Why it changed
shape	encoder + decoder (6 + 6 layers)	decoder only (30 layers)	A model that only continues text needs no separate encoder: the prompt and the answer are one sequence, read with a causal mask.
width d_model	512	576	A choice of the model's makers: many thin layers (30) instead of a few wide ones (6 + 6).
normalization	LayerNorm after each residual addition (post-norm)	RMSNorm before each sublayer (pre-norm), eps 1e-5, plus a final RMSNorm	With the norm before each sublayer, the residual additions form a clean path from the output back to the input, so deep stacks train steadily. RMSNorm skips the mean: less work, a similar effect.
positions	sinusoidal, added to the embeddings	rotary (RoPE, theta 100,000), applied to queries and keys in every layer	Rotary positions turn the queries and keys, so the score q · k depends on how far apart two tokens are, in every layer, not only at the input.
attention heads	8 heads, each with its own keys and values (d_k = 64)	9 query heads sharing 3 key-value heads (grouped-query attention, head size 64)	Three query heads share each key-value head, so the KV cache that generation keeps is 3 times smaller.
feed-forward	ReLU, 512 -> 2048 -> 512	SiLU-gated (SwiGLU): down(SiLU(gate(x)) * up(x)), 576 -> 1536 -> 576	A gate, SiLU(gate(x)), decides number by number how much of up(x) passes. Later experiments found gated layers better at the same size.
biases	in the feed-forward layer (the paper's attention formulas have none)	none	They add little, and leaving them out keeps the layers simpler.
embeddings	one shared matrix (source, target, output), x sqrt(d_model)	input embedding tied to the output layer, not scaled	Tying is kept: in SmolLM2 it saves 28.3 million parameters. The × √d_model scaling is gone.
vocabulary	about 37,000 BPE tokens (shared)	49,152 BPE tokens	One tokenizer for many kinds of text, including code.
context	no fixed limit (sinusoids extend), trained on sentences	8,192 tokens	The model is trained and tested on texts up to this length; rotary positions make longer windows work.
dropout	0.1	0.0	Pretraining reads trillions of tokens, mostly once, so overfitting is not the main risk; dropout would only slow learning.
parameters	about 65 M (our count 63.1 M)	134.5 M	About twice as many, in one stack instead of two.
training task	translation (supervised sentence pairs)	next-token prediction on raw text (pretraining)	Next-token prediction needs no human labels, so the text itself is the answer key, and far more text can be used.

Three changes matter most. Decoder only: a chat model’s prompt and answer are one sequence, so the causal mask of Position and masking replaces the encoder and cross-attention. Pre-norm: deep stacks train more steadily with a clean residual path. Grouped-query attention: a smaller KV cache, the main cost in Inference costs.

Limits of this page

• A toy task. Reversing letters is far easier than translation; it shows how the parts fit, not translation quality.
• A scaled-down recipe. The warmup was cut to 400 steps for a 3,000-step run; the paper used 4,000 of 100,000.
• An estimated count. Our 63,082,496 depends on the vocabulary we assumed (37,000) and on the tying. The paper rounds to 65 million.
• SmolLM2’s row comes from its configuration, not from this page’s runs. Timings are rough (a shared computer).

Guided practice: trace the shapes of a bigger example

1. Use the toy model (d_model 64, 4 heads, d_ff 256, vocabulary 13). Take a source of 7 letters and a decoder input of 8 tokens: the start token and 7 letters. Write the shape after the embedding and positions, for the source and for the target.
2. Write the shape of the encoder’s self-attention weights, the FFN hidden tensor and the memory.
3. Write the shape of the decoder’s masked self-attention weights, the cross-attention weights and the logits.
4. Compute PE(2, 0) and PE(2, 1) by hand, for any d_model.
5. Compute the learning rate of the toy schedule (d_model 64, warmup 400) at step 100 and at step 1,600.

Expected: one attention shape is not square.

Show the answers

1. Source (1, 7, 64); target (1, 8, 64).
2. Self-attention weights (1, 4, 7, 7); FFN hidden (1, 7, 256); memory (1, 7, 64).
3. Masked self-attention (1, 4, 8, 8), with 28 hidden cells per head above the diagonal; cross-attention (1, 4, 8, 7), with no causal mask; logits (1, 8, 13).
4. For i = 0 the divisor is 1: PE(2, 0) = sin 2 = 0.909 and PE(2, 1) = cos 2 = −0.416. This does not depend on d_model.
5. Step 100 is in the warmup: 64^−0.5 × 100 × 400^−1.5 = 0.125 × 100 / 8,000 = 0.00156. (A log printed after sched.step() shows 1.58e-03, the rate of the next step.) Step 1,600 is after the peak: 0.125 × 1,600^−0.5 = 0.125 / 40 = 0.003125: half the peak, because 1,600 = 4 × 400 and the rate falls as 1/√step.

Your turn: a model for German

Omar wants Larkfield’s answers translated into German. Take the base model with separate vocabularies: 32,000 English tokens and 32,000 German tokens. Only the German embedding and the output layer share a matrix. Count the parameters with the code of step 7, and say how far you are from 65 million.

Show the answer

The layers do not change: 6 × (3,152,384 + 4,204,032) = 44,138,496. The embeddings are now two matrices of 32,000 × 512 = 16,384,000 each: 32,768,000. The total is 76,906,496, about 77 million: 12 million more than the paper’s 65 million, although each vocabulary is smaller than 37,000. Sharing one vocabulary and one matrix saved a lot.

Common mistake

nn.Transformer does not match the paper: an extra LayerNorm at the end

Symptom. You check your own implementation of the paper against nn.Transformer(norm_first=False) with copied weights. Before training, they agree to 1.7 × 10⁻⁶, and you trust the check. After training, the same comparison is off by up to 2.14 in the logits of the toy model. Yet each layer still matches on its own.

Cause. nn.Transformer creates a final LayerNorm after the encoder stack and another after the decoder stack. The paper’s post-norm model has none: each sublayer already ends in a LayerNorm. Before training, every γ is 1 and every β is 0. A second norm of numbers that are already normalized then changes almost nothing, so the bug hides. Training changes γ and β, and then the second norm changes the numbers. (A pre-norm model needs a final norm, as SmolLM2 has; hence PyTorch’s default.) A padded batch in eval mode can also take a fast path that returns 0 at padded positions: a difference of 3.06 there.

Fix. For a post-norm model, set model.encoder.norm = None and model.decoder.norm = None, and set model.encoder.use_nested_tensor = False before you compare padded rows. Then compare with (mine - theirs).abs().max(): 9.5 × 10⁻⁷, before and after training. Test with trained weights, not only fresh ones.

Knowledge check

Check your understanding

In cross-attention, where do the queries, the keys and the values come from, and what is the shape of its weights for 6 target positions, 5 source tokens and 4 heads?

Show the answer

The queries come from the decoder: each target position asks its question. The keys and the values come from the encoder’s output, the memory: the source sentence offers its information. The weights have shape (batch, heads, target positions, source tokens) = (1, 4, 6, 5). The matrix is not square, and there is no causal mask, because the whole source is known. Each row adds up to 1 over the 5 source tokens. In the toy model, layer 2 learned an anti-diagonal: position 1 reads source token 5, position 2 reads token 4, and so on.

Check your understanding

Predict: why does the toy model's smoothed training loss stop near 0.54 instead of going to 0, and why are its top probabilities about 0.92 and not 0.99?

Show the answer

Label smoothing with ε = 0.1 changes the target: the right token gets 0.9 + 0.1/13 = 0.908, and each other token 0.0077. The cross-entropy against this target is smallest when the model predicts exactly the target. That smallest value is the target’s entropy, 0.537. So even a perfect model has a loss of about 0.54. The best prediction is the target itself, so the model learns to give the right token about 0.91, not 0.99. A model that is 100% sure has a higher smoothed loss (1.846 in the check). The paper accepts this: it “hurts perplexity”, but improves accuracy and BLEU.

Check your understanding

Name two choices of the 2017 transformer that SmolLM2-135M changed, and say why each was changed.

Show the answer

Any two, for example. Post-norm became pre-norm with RMSNorm. With the norm before each sublayer, the residual path from input to output has no norm on it, so deep stacks (30 layers) train more steadily. RMSNorm skips the mean: less work for a similar effect. Sinusoidal positions added at the input became rotary positions on the queries and keys of every layer. Then the score depends on the distance between tokens. Separate keys and values for each of 8 heads became grouped-query attention, with 9 query heads and 3 key-value heads. This makes the KV cache 3 times smaller. The encoder–decoder became a decoder only, because a model that continues text needs no separate source. Dropout 0.1 became 0.0, because pretraining reads trillions of tokens, mostly once.

Recap

• The 2017 transformer is a translator: an encoder makes a memory of the source; a decoder writes with masked self-attention, cross-attention (queries from the decoder, keys and values from the memory) and a feed-forward layer.
• Input: embedding × √d_model + sinusoidal positions; a shift by k positions is a rotation of each pair.
• Each sublayer: LayerNorm(x + Dropout(Sublayer(x))); FFN(x) = max(0, xW₁ + b₁)W₂ + b₂; output tied to the embedding.
• Training: Adam (0.9, 0.98, 10⁻⁹), warmup to 6.99 × 10⁻⁴ at step 4,000, dropout 0.1, label smoothing 0.1.
• Base model: 63,082,496 by our count, 65 million in the paper. The toy reached 100% exact match with an anti-diagonal cross-attention.
• Today: decoder only, pre-norm RMSNorm, rotary positions, gated feed-forward layers, grouped-query attention.

This ends the optional mathematics module. Next come the case studies, starting with images.

Next: Classify clothing images