LLMUnit 3 · The Transformer era · Day 14

The Annotated Transformer

The Day 13 paper, rewritten as one runnable notebook. Reading it is the fastest way to find out which parts of the paper are easy to get wrong in code.

AuthorsSasha (Alexander) Rush; 2022 update by Austin Huang, Suraj Subramanian, Jonathan Sum, Khalid Almubarak, Stella Biderman

Year / venue2018 (updated 2022) · Harvard NLP blog post and notebook

Core ideaThe paper as line-by-line PyTorch

Study time4–5 h

Original: nlp.seas.harvard.edu/annotated-transformer · Code, notebook and exercises: 30u30 Day 14

Start here (no maths needed)

This page takes the famous Transformer paper and puts working computer code under every paragraph, so you can see exactly how each idea is built.

Think of it like this

A recipe in a cookbook says "make the sauce". A good cooking show actually makes it in front of you, step by step, and shows the parts the recipe skipped: how hot the pan is, how long to stir. The Annotated Transformer is the cooking show for the Day 13 paper. It also shows a trick every athlete knows: warm up gently before going full speed.

How it works

Take each paragraph of the paper and write the code that does exactly what it says.
Add the details the paper left out, like the size of every list of numbers and the "masks" that hide words the model shouldn't see yet.
Test it first on a very easy job: copying a list of digits. If it can't copy, something in the code is broken.
Train with a warm-up: learn in small steps at first, then bigger, then slowly smaller again.
Finally, train it on a real job: translating German to English.

Why robots care

Many robot brains reuse these same building blocks. Before training a robot on hours of recorded demonstrations, engineers first check that it can replay just one demo. If it can't, the bug is in the code, not the data. Masks matter too: they stop the robot from "peeking" at the future while it learns, which would make it look great in practice and fail for real.

Try it

Write a recipe for a sandwich and have a friend follow it exactly, word for word. Which steps did you forget to write down? Those gaps are why this code-under-every-paragraph guide exists.
In the demo further down, move the Warmup steps slider and watch the peak of the curve move with it. Then move the d_model slider: bigger models get a smaller learning rate.

Words to know

Code: Instructions written for a computer to follow.
Notebook: A document that mixes explanations with code you can run.
Learning rate: How big a step the model takes each time it corrects itself.
Warmup: Starting with small learning steps, then speeding up.
Mask: A cover that hides some words from the model, such as future words.

Going deeper: everything below is the full version, with some maths. On a first read, skim the headings and the "Why it matters for robotics" box, and come back for the rest later.

1. Summary

TL;DR

The Annotated Transformer takes the text of "Attention Is All You Need", reorders and trims it, and puts working PyTorch under each paragraph. The whole page is a notebook: run it top to bottom and you get a trained translation model. It adds what the paper leaves out: tensor shapes, masks, batching, the exact learning-rate curve, how label smoothing is computed, and a small task to check the model works before spending GPU hours.

Part 1, the model: encoder and decoder stacks, residual sublayers, multi-head attention, feed-forward layers, embeddings, positional encoding.
Part 2, training: masked batches, the training loop, the warmup learning-rate schedule, label smoothing, and a synthetic copy task with greedy decoding.
Part 3, a real example: German→English on the small Multi30k dataset, with multi-GPU training and attention visualisations.
One deliberate deviation: the code puts layer norm before each sublayer (pre-norm) for simplicity, while the paper normalises after the residual add.
The 30u30 folder re-implements the same pieces with its own exercises on attention, multi-head, encoder, training and inference.

2. The core idea

This is a guide, not a new model, so the "core" is the set of definitions the code must get exactly right. Written as maths rather than code:

Model      = EncoderDecoder(encoder, decoder, src_embed, tgt_embed, generator)
Sublayer   x ← x + Dropout(F(LayerNorm(x)))          # pre-norm, as in the notebook
                                                      # paper: x ← LayerNorm(x + Dropout(F(x)))
Causal mask  M[i][j] = 1 if j ≤ i else 0             # token i sees itself and the past
Pad mask     M[j]    = 1 if token j ≠ <pad>
Masked score s_ij ← −∞ where M = 0, before softmax

lr(step) = factor · d_model^(−0.5) · min(step^(−0.5), step · warmup^(−1.5))

Smoothed target: (1 − ε) on the true token, ε spread over the rest, 0 on <pad>

EncoderDecoder: one wrapper that embeds source and target, encodes, decodes, and projects to vocabulary scores. Every later piece slots into it.
F: either an attention block or the feed-forward block. The residual add means each layer only learns a correction to its input, the ResNet trick from Day 09.
Masks do two jobs. The pad mask stops attention to filler tokens in a batch of uneven sentences. The causal mask stops the decoder peeking at future target tokens.
lr: linear warmup for warmup steps, then decay as 1/√step. The peak is exactly at step = warmup. The paper used warmup = 4000 and d_model = 512.
ε = 0.1 in the paper. The notebook computes the loss as a KL divergence against this softened target.

3. Deep dive

Shapes are the whole game

Multi-head attention in code is mostly reshaping. A batch of size B with length n and width 512 is projected, split into 8 heads of 64, moved so heads act like an extra batch dimension, attended, then merged back. Most bugs in home-made Transformers are a transpose in the wrong place, or a mask broadcast over the wrong axis. Write the shape next to every line as you read. The TensorTonic attention problems (Module 02) drill exactly this.

The learning-rate schedule

Transformers are fragile early in training. Big updates while the attention weights are still random can push the softmax into saturation. Warmup keeps steps small until the statistics settle, then the 1/√step decay takes over. Play with it:

Interactive: the warmup schedule

Learning rate over the first 40,000 steps, with factor = 1. Change warmup and model width. The peak sits at step = warmup, and wider models get a smaller rate.

Warmup steps: d_model:

Peak step

–

Peak rate

–

Rate at 100k

–

Test on a toy before the real task

Before any translation, the notebook trains a small model to copy random digit sequences, then decodes greedily (always pick the most likely next token). If a model can't learn to copy, something structural is broken: a mask, a shift between decoder input and target, a shape. This "overfit something trivial first" habit is worth keeping for any model, including robot policies. Then it moves to Multi30k German→English, a dataset far smaller than the WMT data of the paper, and shows the encoder self-attention, decoder self-attention and cross-attention maps per layer and head.

From notebook to paper numbers

The notebook is honest that its small run is not the paper's result. It notes that a full OpenNMT-py replication, with byte-pair encoding, beam search and checkpoint averaging, reaches 26.9 BLEU on WMT English→German, versus 27.3 for the paper's base model. Those extras matter. Beam search and averaging are cheap at inference or after training. The KV cache that makes decoding fast is the same structure Module 05 shows being compressed at trillion-parameter scale.

4. Why it matters for robotics

Robot policy code looks like this notebook. ACT, Diffusion Policy's Transformer variant and many VLAs reuse these exact modules: a Transformer encoder over observation tokens, a decoder with cross-attention, masks and positional embeddings. Reading this once makes their repos readable.
Masks encode what the robot is allowed to see. A causal mask over time stops a policy from conditioning on future frames in training, which otherwise produces a policy that looks great offline and fails on the robot.
The copy-task habit transfers. Before training on hours of teleoperation data, overfit a single demo. If the policy can't replay one trajectory, the bug is in the pipeline, not the data.
Warmup and label smoothing show up again when you fine-tune a VLA on discretised action tokens. Small datasets plus big pretrained models are exactly where early instability bites.

5. How to study it

Read the Annotated Transformer with the Day 13 paper open beside it. For each code block, find the paper sentence it implements.
Run the original notebook (or a Colab copy) as far as the copy task. Write the tensor shape next to each line of the attention function.
In the 30u30 Day 14 folder, read implementation.py, run train_minimal.py, and use visualization.py to plot attention maps.
Do the five exercises: attention, multi-head, encoder, training, inference. Compare against solutions/ only after a real attempt.
Experiment: switch pre-norm to post-norm and remove warmup. Note which change makes training unstable.

6. Self-check

Which two masks does a Transformer use, and what does each protect against?

A padding mask, so attention ignores filler tokens in uneven batches. A causal (subsequent) mask in the decoder, so position i can't see target tokens after i during teacher-forced training.

At what step does the learning rate peak, and what is its value?

At step = warmup. The value there is factor · d_model^−0.5 · warmup^−0.5, for example about 7×10⁻⁴ with d_model 512 and warmup 4000.

How does the notebook's layer norm placement differ from the paper's?

The notebook normalises the input of each sublayer and adds the raw residual (pre-norm). The paper adds first and normalises the sum (post-norm). Pre-norm is generally easier to train deep.

Why train on a copy task first?

It is trivially learnable, so failure points to a code bug (masks, shifted targets, shapes) rather than to data or capacity.

Why does label smoothing hurt perplexity but help BLEU?

It stops the model putting all probability on one token. That makes its probabilities less sharp, which raises perplexity, but it reduces overconfidence and generalises better.

7. Further reading

The Annotated Transformer (2022 version, with a link to the 2018 original)
30u30 Day 14: implementation, notebook, cheatsheet, exercises
Next: Day 15, Neural Machine Translation by Jointly Learning to Align and Translate