P01 — Transformer From Scratch

Run it first. There is a companion page that builds this project's machinery as numbered, independently runnable blocks and then assembles them into one measured system: P01 hands-on — block by block (handson/h*.py). Every number on it was produced by running the code. Read it alongside the milestones below.

Medium · 88 hours · Weeks 1–8 · Stage 1 · Python

Table of Contents

• The Loop, Instantiated
• Why This Project Matters
• Prerequisites
• Duration and Size
• Central Technical Questions
• Architecture
• Showcase — Do This Before You Start
• Implementation Milestones
• Concepts To Study
• Primary-Source Readings
• Experiments
• Benchmarks and Metrics
• Correctness Tests
• Failure Tests
• Expected Difficulties
• Scope Boundaries
• Deliverables
• Exit Criteria
• Extension Ideas
• Connections
• References

The Loop, Instantiated

Read this section, then close the page and write your naive design before scrolling to Architecture.

Why This Project Matters

Two reasons, and the second is the real one.

The stated reason: the Transformer is the most consequential architecture of the last decade and you use it daily through APIs. Building one converts a black box into a set of decisions you can defend.

The real reason: the Transformer is a pile of arbitrary-looking choices. Why scale by \(1/\sqrt{d_k}\)? Why 4× expansion in the feed-forward layer? Why LayerNorm before the sublayer instead of after, when the original paper did the opposite? Why multiple heads instead of one big one? None of these follow from first principles. Each is defensible only by measurement, and almost nobody who uses Transformers has done those measurements.

This project is where you learn that an architecture is an empirical artifact, and that the way to understand any system full of arbitrary constants is to vary them and watch what happens. That habit transfers to every other project here. The Transformer is a convenient place to acquire it because the feedback loop is minutes, not hours, and correctness is checkable against finite differences.

It is also first because it is motivating. Week 8 of this project ends with a model that writes text. Week 8 of an autodiff framework ends with a class that adds arrays correctly. You are 130 weeks from the end and the first two months decide whether there is a month three.

Prerequisites

• Linear algebra: matrix multiplication, transpose, the fact that \(A B\) composes linear maps. You do not need eigenvalues here.
• The chain rule. You will not implement backprop in this project (that is P13), but you must be able to explain why a residual connection helps gradient flow.
• Python and numpy. PyTorch tensors + autograd are permitted as a tensor library only — see Scope Boundaries.
• From math.md: §Linear Algebra (2 h), §Softmax and Cross-Entropy (1 h). Nothing else. Do not read the optimization section yet.

Duration and Size

Medium, 88 hours, 8 weeks at 11 h/week.

Central Technical Questions

1. What does attention compute that a fixed-window MLP cannot? Answer in terms of the number of parameters required to express a content-dependent, position-invariant lookup.
2. Why divide by \(\sqrt{d_k}\)? Derive it. What happens to the softmax gradient if you do not?
3. At what context length does attention stop being a rounding error and start being the bottleneck? For your specific \(d_{model}\). Derive first, then measure.
4. What does a positional encoding have to provide, minimally? RoPE, sinusoidal and learned embeddings solve the same problem three ways — what is the problem, stated without reference to any of them?
5. Why did pre-norm replace post-norm? And at what depth does the difference start to matter?
6. What is the memory cost of inference, and why is it dominated by something other than the weights beyond a certain sequence length?

Architecture

Do not read this until you have written your naive design.

tokens ─► [Tokenizer] ─► ids ─► [Embedding] ─┬─► + positional ─► [Block] × L ─► [Norm] ─► [Unembed] ─► logits
                                             │                      │
                                             │                      ├─ x + Attn(Norm(x))     ← pre-norm residual
                                             │                      └─ x + FFN(Norm(x))
                                             │
                                        (RoPE instead applies rotation inside Attn, to q and k only)

The attention mechanism, derived

You have a sequence of \(T\) vectors \(x_1..x_T\), each in \(\mathbb{R}^{d}\). You want each position to gather information from earlier positions, with the choice of which earlier positions determined by content rather than by fixed offset.

Project each \(x_i\) three ways: a query \(q_i = W_q x_i\) (what am I looking for), a key \(k_i = W_k x_i\) (what do I offer), and a value \(v_i = W_v x_i\) (what I will hand over). Score every pair by dot product, normalise into weights, and take the weighted sum:

\[ \text{Attn}(Q,K,V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right)V \]

where \(M_{ij} = -\infty\) for \(j > i\) — the causal mask, which must be applied before the softmax so that masked positions receive exactly zero weight rather than a small one.

Why \(\sqrt{d_k}\). Take \(q, k\) with i.i.d. components of mean 0 and variance

1. Then \(q \cdot k = \sum_{i=1}^{d_k} q_i k_i\) is a sum of \(d_k\) independent mean-zero unit-variance terms, so \(\mathrm{Var}(q\cdot k) = d_k\) and the typical magnitude of a score is \(\sqrt{d_k}\). At \(d_k = 64\) scores are typically ±8, and the gap between the largest and second-largest is often several units. Feed that to a softmax and you get a near-one-hot distribution whose gradient is nearly zero everywhere — the softmax saturates and the layer stops learning. Dividing by \(\sqrt{d_k}\) restores unit variance and keeps the softmax in its responsive region. This is not a heuristic; it is a variance calculation you should be able to do at a whiteboard.

Why multiple heads. One head produces one attention distribution per position: it can attend to one thing. Splitting \(d\) into \(H\) heads of size \(d/H\) gives \(H\) independent distributions at the same total parameter and FLOP cost (the concatenation-then-projection makes the arithmetic identical). You are buying representational diversity for free. The cost is that each head sees a \(d/H\) dimensional subspace, so there is a floor below which heads become too small to be useful — one of your experiments.

Cost model, derived before you measure

Let \(B\) = batch, \(T\) = sequence length, \(H\) = heads, \(d_h\) = head dim, \(d = H d_h\). Per layer, forward pass, counting 2 FLOPs per multiply-accumulate:

Everyone "knows" attention is quadratic. Almost nobody knows where the crossover is. Computed for \(B{=}1, H{=}12, d_h{=}64\) (so \(d{=}768\), GPT-2 small shape):

Two things fall out that are worth more than the table itself:

1. At GPT-2's original context of 1024, attention is 18% of the compute. The quadratic term does not dominate until ~4096. Every "attention is the bottleneck" claim is context-length-dependent and most of them are wrong for the context they are said about.
2. Memory hits the wall long before FLOPs do. The materialised score matrix is 3.2 GB at T=8192 for a single batch element in fp32. This is the entire reason FlashAttention exists: it never materialises that matrix, computing softmax in tiles with an online-normalisation trick. The compute is unchanged; the memory traffic is what is optimised.

Reproduce this table for your own \(d_{model}\) as milestone 6. Predict the crossover before you compute it.

RoPE, and why it is different in kind

Learned and sinusoidal encodings add a position-dependent vector to the token embedding. RoPE rotates the query and key vectors by an angle proportional to position, in \(d/2\) independent 2-D planes:

\[ \theta_i = \frac{m}{\text{base}^{2i/d}}, \qquad \begin{pmatrix} x^{\prime}_{2i} \\ x^{\prime}_{2i+1} \end{pmatrix} = \begin{pmatrix} \cos\theta_i & -\sin\theta_i \\ \sin\theta_i & \cos\theta_i \end{pmatrix} \begin{pmatrix} x_{2i} \\ x_{2i+1} \end{pmatrix} \]

The property that makes it work: because a rotation by \(m\) followed by the inverse of a rotation by \(n\) is a rotation by \(m-n\), the dot product of a rotated query at position \(m\) with a rotated key at position \(n\) depends only on \(m-n\). Attention becomes relative-position-aware without any explicit relative-position term.

Verify this numerically before you trust it — measured with a 64-dimensional random vector pair:

   m    n   m-n            dot
   5    3     2   -8.115791933
 105  103     2   -8.115791933
   7    5     2   -8.115791933
1000  998     2   -8.115791933
   5    4     1   -8.519139825
  50   49     1   -8.519139825

Identical to nine decimal places across a 200× range of absolute positions. And because rotations are orthogonal, norms are preserved exactly: \(|q| = 7.315343549\), \(|\text{RoPE}(q,17)| = 7.315343549\). Write this test before you write the implementation. It catches the two overwhelmingly common RoPE bugs — pairing dimensions as \((i, i+d/2)\) versus \((2i, 2i+1)\) without adjusting the frequency indexing, and applying the rotation to values as well as to queries and keys.

Showcase — Do This Before You Start

W1 · walkthroughs/w1_attention.py · ~45 minutes

A working miniature of this project: the leak test, what the \(1/\sqrt{d_k}\) scale buys measured as attention entropy, and the sequence length where the quadratic term actually starts to matter.

cd walkthroughs && python3 w1_attention.py

It is 80-ish lines and it surfaces this project's central surprise in an evening rather than in week six. Run it before committing the weeks.

Implementation Milestones

Concepts To Study

Limited to what the milestones need. Anything not on this list is not yet.

• Tokenization: byte-level BPE, the merge algorithm, why byte-level avoids OOV entirely, vocabulary size vs sequence length as a direct trade
• Embeddings: as a lookup table that is also a linear layer; weight tying between embedding and unembedding and what it saves
• Softmax: numerical stability via max-subtraction, the Jacobian, saturation
• Cross-entropy: as negative log-likelihood, and why loss in nats converts to perplexity by \(e^{\text{loss}}\)
• LayerNorm vs RMSNorm: what each normalises, and why RMSNorm dropping the mean subtraction costs almost nothing
• Residual connections: as an identity path for gradients; why the residual stream is better thought of as a shared bus than as a shortcut
• Pre-norm vs post-norm: where the norm sits relative to the residual add, and the effect on gradient magnitude at depth
• AdamW: first and second moment estimates, bias correction, why decoupled weight decay differs from L2
• Learning-rate schedules: warmup as a remedy for early-training instability in adaptive optimizers
• KV caching: what is recomputed without it, and the memory it costs
• Sampling: temperature as logit scaling, top-k and nucleus as truncation strategies

Primary-Source Readings

Total budget: 13 hours. Read section 3 of Vaswani before milestone 3, RoFormer before milestone 9, and nothing else until the milestone that needs it.

Deliberately not on this list: scaling laws, MoE, RLHF, quantization, distributed training. See Not Yet.

Experiments

Every row: fixed seed, fixed data, fixed token budget, one variable. Record the prediction before the run.

E10 is not optional and it is not a formality. A model that cannot overfit 200 tokens has a bug — usually a mask applied after the softmax, a detached gradient, or a data loader that reshuffles the target. Run it at milestone 5 and again any time anything looks strange. It is the cheapest bug detector in the project.

E6 is the best hypothesis in this project. Pre-norm and post-norm perform similarly at depth 2 and diverge sharply somewhere deeper. Predicting where — and explaining it via gradient magnitude at initialisation — is a genuine result you derived.

Benchmarks and Metrics

Report every latency as p50/p95/p99 with a bootstrap CI on the median, using tools/bench.py. A mean latency in this project's report is a scorecard deduction.

Correctness Tests

These are properties, not examples. Property tests catch the bugs example tests miss.

1. Hand-computed forward pass. 3 tokens, \(d{=}4\), one head, weights set by hand. Your attention weights match your arithmetic to 1e-6. Do this on paper first.
2. Causal mask leak test. Run the model on a sequence. Perturb token \(t{+}1\). Assert the logits at positions \(\le t\) are bit-identical. This catches masking after softmax, off-by-one in the mask, and any accidental bidirectionality.
3. Attention rows sum to 1. For every head, every position, every batch element, \(|\sum_j a_{ij} - 1| < 10^{-5}\).
4. Gradient check. Finite differences against autograd on a tiny model: \(|\nabla_{\text{analytic}} - \nabla_{\text{numeric}}| / (|\nabla| + \epsilon) < 10^{-4}\) for every parameter tensor. (In P13 you will implement the analytic side yourself and run this same test against PyTorch.)
5. RoPE relative-position invariance. As measured above: for all \((m,n)\) with the same \(m-n\), the rotated dot product is equal to 1e-6.
6. RoPE norm preservation. \(||\text{RoPE}(x,m)|| = ||x||\) to 1e-6, all \(m\).
7. Tokenizer round-trip. decode(encode(s)) == s for a corpus including emoji, CJK, and lone surrogate byte sequences.
8. Softmax stability. Logits of \(10^4\) produce no NaN.
9. Shape and dtype invariants on every tensor crossing a module boundary.
10. Determinism. Same seed → bit-identical loss curve. If not, find the source before running any experiment; nondeterminism silently invalidates every ablation.

Failure Tests

Deliberately break it and confirm it breaks the way you predict.

The shuffled-target test deserves emphasis: a model trained on random labels should converge to exactly \(\ln(V)\) nats. If it goes below, you have leakage. If it goes far above, you have an optimizer problem. It is a two-line experiment that calibrates your entire loss intuition.

Expected Difficulties

Named in advance so that hitting one is a checkpoint, not a crisis.

1. Shape bugs will consume more time than concepts. (B,T,H,d_h) vs (B,H,T,d_h) transposes are the single biggest time sink. Mitigation: annotate every tensor's shape in a comment at creation, and assert shapes at every module boundary. This feels excessive for about two days and then saves a week.
2. A subtly wrong mask trains fine and scores impossibly well. Mitigation: test 2 above, written before the mask.
3. You will want to add features instead of running experiments. Adding grouped-query attention is more fun than running a 5-point head-count sweep. The sweep is the project.
4. CPU training is slow enough to discourage sweeps. Mitigation: size the model so one training run is under 10 minutes. Every experiment here is about relative comparison; absolute quality is irrelevant.
5. RoPE has two incompatible conventions in the wild (interleaved pairs vs split halves). Both are correct if applied consistently; mixing them within one model produces a model that trains but extrapolates badly. Test 5 catches it.
6. The bigram baseline is better than you expect and beating it by a little feels like failure. It is not — write down the bigram loss and the entropy floor first so you know what the achievable range even is.

Scope Boundaries

In scope: everything in the milestone list, on CPU, at a scale where a training run takes under 10 minutes.

Explicitly out of scope — if you find yourself doing one of these, stop:

• Multi-GPU or distributed training of any kind
• Training on a large corpus, or chasing an absolute quality number
• Implementing your own autograd — that is P13, and doing it here means doing two Large projects at once
• Fine-tuning, instruction tuning, RLHF, LoRA
• Serving infrastructure, batching servers, or an API
• Model architectures other than a decoder-only Transformer
• Writing CUDA kernels — that is P14

The permitted-library line. PyTorch is allowed for: tensor storage, elementwise ops, matmul, autograd, and the optimizer. PyTorch is forbidden for: nn.Transformer, nn.TransformerEncoderLayer, nn.MultiheadAttention, F.scaled_dot_product_attention, and any positional-encoding utility. The mechanism under study is attention; you may not import it. See AI Policy Rule 7.

Deliverables

1. transformer/ repository — one command to train, one to sample, one to reproduce every experiment
2. REPORT.md, 2,000–3,000 words, with the E4/E6 ablation tables and a section titled "What I Expected And Did Not Get"
3. notebook/ entries for E4, E6, and E11 at minimum
4. cost_model.py reproducing the FLOP/memory table for arbitrary configs
5. A generated-text sample at three temperatures, with the tokenizer's compression ratio stated
6. Raw benchmark JSON for every experiment, not just the summary tables

Exit Criteria

All eight, no exceptions.

• Model overfits a 200-token corpus to val loss < 0.1 (E10)
• Beats the bigram baseline on held-out val loss by a stated, measured margin
• All ten correctness tests pass, including the causal-leak and RoPE invariance tests
• E4 complete: all four positional encodings including the none control, with a table
• E6 complete: pre-norm vs post-norm across four depths, with the divergence depth identified and an explanation offered
• E11/E12 complete: KV-cache latency is flat in generated length, and derived cache memory matches measured within 5%
• At least one hypothesis tested with a stated falsifier, and the outcome recorded whichever way it went
• REPORT.md written, including at least one result that contradicted your prediction

Extension Ideas

Locked until every box above is ticked.

• Grouped-query attention: share K/V across head groups. Measure the KV-cache reduction and the quality cost. Directly relevant to your day job's serving costs.
• Tiled attention (FlashAttention-style): implement online softmax normalisation and never materialise the score matrix. Measure peak memory vs T against the table above. This is the best bridge to P14.
• Sliding-window attention with a long-context eval, measuring the quality/compute frontier.
• Speculative decoding with a small draft model. Measure acceptance rate and end-to-end speedup; note that acceptance rate is the interesting quantity.
• RoPE base sweep (\(\text{base} \in \{10^3, 10^4, 10^5\}\)) and its effect on length extrapolation — a small, real, publishable-adjacent experiment.

Connections

Backward: none. This is where the journey starts.

Forward:

• → P13 (Autodiff): the exit criterion for P13 Phase I is reproducing this model's gradients to 1e-5. Keep the model code stable and tagged so P13 has a fixed target.
• → P14 (Hardware-aware): the FLOP/memory table here is the input to P14's roofline analysis. Your inference latency measurements become the "before" number.
• → P08 (Recommender): you will need embeddings, and having built the thing that produces them changes how you reason about their geometry.
• → P02 (ANN): the embeddings from this model are a legitimate test dataset for your index — and unlike random vectors, they have realistic relative contrast.
• → P15: dynamic embedding generation is one of the candidate research questions, and it depends on your being able to reason about the cost of a forward pass.

References

• Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., Polosukhin, I. Attention Is All You Need. NeurIPS 2017.
• Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., Liu, Y. RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864, 2021.
• Xiong, R. et al. On Layer Normalization in the Transformer Architecture. ICML 2020.
• Dao, T., Fu, D. Y., Ermon, S., Rudra, A., Ré, C. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022.
• Sennrich, R., Haddow, B., Birch, A. Neural Machine Translation of Rare Words with Subword Units. ACL 2016.
• Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I. Language Models are Unsupervised Multitask Learners. OpenAI technical report, 2019.
• Loshchilov, I., Hutter, F. Decoupled Weight Decay Regularization. ICLR 2019.
• Zhang, B., Sennrich, R. Root Mean Square Layer Normalization. NeurIPS 2019.
• Ba, J. L., Kiros, J. R., Hinton, G. E. Layer Normalization. arXiv:1607.06450, 2016.
• Shazeer, N. Fast Transformer Decoding: One Write-Head is All You Need. arXiv:1911.02150, 2019. Multi-query attention; the ancestor of GQA.
• Elhage, N. et al. A Mathematical Framework for Transformer Circuits. Anthropic, 2021. The residual-stream view referenced under Concepts.
• Karpathy, A. nanoGPT and Let's build GPT: from scratch, in code, spelled out. 2022–2023. Read after your milestone 7, as a check on your choices rather than a source for them.