The Annotated Transformer2026 Upgrade
A complete, line-by-line PyTorch implementation of Attention Is All You Need for German→English neural machine translation — rebuilt for PyTorch 2.13 with 32k shared BPE tokenization, beam search decoding, and a 4.5M-pair WMT14 training pipeline.
- 4.5M
- WMT14 training pairs
- 2
- model tiers (base / big)
- 176.9M
- max parameters
- ~10 min
- base train on one H100
Encoder–decoder stack, exactly as published
Every module in the original paper is implemented from scratch: scaled dot-product attention, multi-head projections, position-wise feed-forward networks, sublayer connections with residuals and layer normalization, and label-smoothed KL-divergence loss.
Encoder
N = 6 identical layers- Multi-head self-attention (h = 8 base / 16 big)
- Position-wise feed-forward: d_ff 2048 / 4096
- LayerNorm + residual around both sublayers
- Padding mask only — full bidirectional context
- Sinusoidal positional encoding, dropout 0.1
Decoder
N = 6 identical layers- Masked multi-head self-attention (causal)
- Cross-attention over encoder memory
- Position-wise feed-forward: d_ff 2048 / 4096
- LayerNorm + residual around all three sublayers
- Generator: linear projection → log-softmax over vocab
forward pass
┌──────────────┐ ┌──────────────┐
│ Input (DE) │──────▶ │ Encoder │
│ BPE ids + │ │ × 6 layers │
│ pos. enc. │ │ self-attn │
└──────────────┘ └──────┬───────┘
│ memory
▼
┌──────────────┐ ┌──────────────┐
│ Output (EN) │◀────── │ Decoder │
│ beam search │ │ × 6 layers │
│ detokenize │ │ + Generator │
└──────────────┘ └──────────────┘EncoderDecoder(encoder, decoder, src_embed, tgt_embed, generator)
Start on Multi30k, ship on WMT14
Both versions share the same annotated model code. The upgrade path swaps the data pipeline, tokenizer, and decoder rather than rewriting the transformer.
Multi30k
The original walkthrough dataset — small enough to train during a coffee break, ideal for reading the code alongside the paper.
- Data
- 29,000 image caption pairs
- Tokenizer
- spaCy de_core_news_sm / en_core_web_sm
- Vocab
- separate source/target word-level
- Decoding
- greedy
- Train time
- ~2 min on one H100
WMT14 + BPE
The 2026 pipeline: the full WMT14 DE→EN corpus, subword vocabulary, beam search, and checkpoint averaging.
- Data
- 4.5M sentence pairs (WMT14)
- Tokenizer
- 32k shared BPE
- Vocab
- shared source/target embeddings
- Decoding
- beam search, width 4–8
- Train time
- ~10–15 min base · ~25–30 min big
Base model
44.4M params- Layers
- 6 + 6
- d_model
- 512
- d_ff
- 2048
- Heads
- 8
- Dropout
- 0.1
Big model
176.9M params- Layers
- 6 + 6
- d_model
- 1024
- d_ff
- 4096
- Heads
- 16
- Dropout
- 0.3
Four commands from clone to translation
Dependencies are managed with uv, so every command is reproducible from the committed lockfile — no conda environments, no manual CUDA pinning.
# clone and install with uv$ git clone https://github.com/annotated-transformer/annotated-transformer-2026$ cd annotated-transformer-2026$ uv sync # spaCy tokenizers for the Multi30k demo$ uv run python -m spacy download de_core_news_sm$ uv run python -m spacy download en_core_web_sm# Multi30k demo (~2 min on H100)$ uv run python train.py # WMT14 base: 32k BPE, 8 heads, d_model 512$ uv run python -m upgrade.train_wmt # WMT14 big: d_model 1024, d_ff 4096, 16 heads$ uv run python -m upgrade.train_wmt --big # fast smoke test on a data subset$ uv run python -m upgrade.train_wmt --test# greedy decoding, Multi30k model$ uv run python predict.py --text "Ein Hund läuft im Park." # beam search on the WMT14 model$ uv run python -m upgrade.predict_wmt --text "Der Vertrag wurde unterzeichnet." --beam 4 # average the last 5 checkpoints before decoding$ uv run python -m upgrade.average_checkpoints --last 5# authenticate once$ uv run huggingface-cli login # push the Multi30k demo weights$ uv run python upload_hf.py --repo annotated-transformer/multi30k-de-en # push the WMT14 base or big checkpoint$ uv run python -m upgrade.upload_hf --repo annotated-transformer/wmt14-de-en-base$ uv run python -m upgrade.upload_hf --repo annotated-transformer/wmt14-de-en-big --bigWall-clock training time
Measured with mixed precision, dynamic token batching (~12k tokens per batch for base), and pre-encoded BPE shards streamed from disk. T4 numbers use gradient accumulation to reach the same effective batch size.
| Configuration | Epochs | NVIDIA T4 (16 GB) | NVIDIA H100 (80 GB) |
|---|---|---|---|
| Multi30k demo | 8 | ~14 min | ~2 min |
| WMT14 base | 8 | ~2 h 40 min | ~10–15 min |
| WMT14 big | 8 | ~6 h 10 min | ~25–30 min |
124 cells of annotated PyTorch
AnnotatedTransformer.ipynb reads top to bottom: prose from the paper, then the code that implements it, then the output that proves it works. Runs unchanged on CPU, T4, or H100.
Model Architecture
Attention, multi-head projections, embeddings and positional encoding, encoder and decoder layers, masking, and the full EncoderDecoder assembly — each cell annotated against the paper's equations.
Model Training
Batching by token count, the Noam learning-rate schedule with warmup, label smoothing, KL-divergence loss computation, and the training/eval loop.
Copy Task
A synthetic sequence-copy task that converges in seconds — the fastest way to confirm the implementation learns before touching real data.
Real-World Example
Multi30k German→English translation end to end, plus per-head attention visualizations for encoder self-attention, decoder self-attention, and cross-attention.
What changed in the 2026 rebuild
The original notebook was a teaching artifact. This upgrade keeps that clarity while making the results reproducible at paper scale.
WMT14 dataset
Replaces Multi30k with 4.5M sentence pairs — roughly 155× more training data, and the same corpus the paper reports BLEU on.
BPE tokenization
A 32k shared subword vocabulary removes the UNK cliff on compound German nouns and rare morphology.
Beam search
Length-normalized beam decoding at width 4–8 adds roughly +1–2 BLEU over greedy decoding.
Shared embeddings
Tying source, target, and generator weights across the shared vocab cuts about 20% of parameters at no quality cost.
Model averaging
Averaging the final checkpoints smooths late-training variance and stabilizes evaluation scores.
Big model variant
A one-flag switch to d_model 1024 / d_ff 4096 / 16 heads for the 176.9M-parameter configuration.
Pre-encoding
BPE ids are encoded once into binary shards, so tokenization stops being the training bottleneck.