Transformer, From Scratch: Architecture, Training, and the KV Cache

I built a decoder-only transformer from scratch this week — architecture, training loop, sampling, and a KV cache for fast generation. Code's on GitHub: paper2code/transformer.

These are my notes on it, written the way I wish someone had explained it to me the first time: not just what each piece does, but why it has to be there. If I have to defend a design choice under harsh questioning — why residual connections specifically, why LayerNorm and not something else, why bother with a feedforward block when attention already mixes information — those are the answers below.

Architecture, at a glance

A decoder-only transformer is a function from a sequence of token ids to a probability distribution over the next token, applied repeatedly. Tokens get embedded, pass through a stack of identical blocks, and get unembedded back into vocabulary-sized logits.

The mental model that made this click for me: think of x, at every position, as a running residual stream — a vector that starts life as an embedding and only ever accumulates edits from there. Every block reads a normalized snapshot of the stream, computes something, and adds its answer back in. Nothing gets overwritten. Attention and the FFN aren't transformations that replace x; they're modules that read from a shared stream and write corrections onto it. That's why every block's formula is x = x + (something), never x = (something). More on why that specific structure matters below.

token ids                  (batch, pos)
  |
  v
tok_emb(idx) + pos_emb(positions)
  |
  v
Block x n_layers   --  x + Attn(LN(x)), then x + FFN(LN(x))
  |
  v
final LayerNorm
  |
  v
unembed -> logits          (batch, pos, d_vocab)

Every block has the same shape in, same shape out — (batch, pos, d_model) is preserved end to end until the very last unembed.

A concrete pass through the pipeline

Take the 3-token sentence "the cat sat" and follow it all the way through:

"the cat sat" -> tokenize -> [1, 3, 0] -> embed -> (1, 3, d_model)
  -> N x Block -> (1, 3, d_model) -> final LN + unembed -> (1, 3, d_vocab)
  -> row 2, softmax, sample -> next token

The shape (1, 3, d_model) never changes across all N blocks — every block is a pure refinement of the same-shaped stream, not a resize. Only the very last step throws most of that away, keeping just the final position's row to predict what comes after "sat".

Tokenization & embeddings

Text becomes integers via a pretrained tokenizer (AutoTokenizer.from_pretrained("gpt2") — not reimplemented). Each token id looks up a learned row in a (vocab_size, d_model) table. Position does the same in a separate (max_seq_len, d_model) table. The two get summed, so the model receives "what token" and "where" as one merged vector per position:

x = tok_emb(idx) + pos_emb(arange(pos))

There's no notion of sentence or turn boundaries at this layer — a token is a token. Chat-style structure (user turn, assistant turn) is a string-concatenation convention applied before tokenization, not something the architecture itself knows about.

Self-attention: Q, K, V

Why attention at all — why not just mix in the last few tokens with a fixed rule?

Compare two sentences: "The trophy doesn't fit in the suitcase because it is too big." and "...because it is too small." In the first, it refers to the trophy. In the second — identical structure, identical position of the word it — it now refers to the suitcase. The correct connection flips entirely based on one word (big / small) sitting somewhere else in the sentence. A fixed rule ("always mix in the previous 3 tokens") can't do this: it has no way to know which earlier token is relevant, because that depends on content, not position. Attention is built to solve exactly this — every position asks a content-based question and pulls information from whichever other position actually answers it, not from whoever happens to be nearby.

Every other layer in the network (embeddings, FFN, unembed) acts on one position at a time. Attention is the only place information moves between positions. Each position produces three vectors from the same input x, via three separate learned linear maps:

• Query — what this position is looking for
• Key — what this position offers, for others to match against
• Value — what this position actually contributes, if attended to

A query is compared against every key via a dot product (bigger dot product = better match), scaled and turned into a probability distribution with softmax, then used to take a weighted sum of every value. That weighted sum is this position's new representation.

$$\text{scores} = \frac{Q K^T}{\sqrt{d_{head}}} \quad\rightarrow\quad \text{softmax(masked)} \quad\rightarrow\quad \text{weights} \cdot V$$

Back to it: its query roughly encodes "I'm a pronoun referring to some earlier object — which one matches the property just described?" trophy and suitcase both offer keys as candidate antecedents. Whichever key actually lines up with what it is asking for — informed by big or small appearing elsewhere — produces the larger dot product, and dominates the weighted sum of values that becomes it's new representation. Nothing about this is hardcoded; the model just has to learn projections that make the right dot products come out large.

Causal masking

Language modeling predicts the next token, so position i may only attend to positions <= i — never the future. Before softmax, every disallowed cell gets set to -inf, so softmax turns it into exactly 0:

query \ key	K0	K1	K2	K3
Q0	Q0·K0	−∞	−∞	−∞
Q1	Q1·K0	Q1·K1	−∞	−∞
Q2	Q2·K0	Q2·K1	Q2·K2	−∞
Q3	Q3·K0	Q3·K1	Q3·K2	Q3·K3

This one mask is what makes the KV cache (below) correct at all: a masked cell can never depend on a later position, so once a cell is computed, it never changes.

Multi-head

d_model splits into n_heads independent chunks of size d_head. Each head runs the exact same attention computation on its own slice, in parallel; the heads' outputs get concatenated back to d_model and passed through one more linear layer (out_proj).

Why split into many small heads instead of one attention operation over the full d_model? A single full-size head has to route every kind of relationship a token might care about — subject/verb agreement, pronoun reference, rhyme, whatever — through one shared d_model-sized computation, all competing in the same subspace. Splitting into n_heads smaller, independent subspaces lets different heads specialize on different relationship types in parallel, for the same total parameter and compute budget — it's the same d_model, just partitioned. Trained transformers bear this out empirically: individual heads reliably specialize into interpretable, distinct patterns rather than all doing the same generic mixing.

The block: residual stream + LayerNorm + FFN

Each block is pre-norm: normalize, transform, add back to the unmodified input. Twice — once for attention, once for the FFN.

x = x + attn(ln1(x))
x = x + ffn(ln2(x))

The dashed line is the literal skip connection: an untouched copy of x that bypasses the computation entirely and rejoins at the +. Attention and the FFN only ever see a normalized snapshot of the stream, and only ever contribute an additive correction to it. The same picture repeats twice per block — once with Attention in the box, once with the FFN.

Why not just stack LN -> Attn -> LN -> FFN directly, without adding x back in? Two separate failures show up. First, gradients: backpropagating through many stacked nonlinear transforms multiplies many Jacobians together, and that product tends to vanish (or explode) as depth grows — the classic reason plain deep networks got hard to train past a dozen or so layers. A residual connection gives gradients an unobstructed additive path straight back through every block, bypassing that multiplicative chain. Second, and more basic: a layer that would ideally do nothing still has to learn to reproduce its input exactly through a nonlinear transform — a hard target to hit precisely. With a residual connection, doing nothing is the default; the sub-layer only has to output zero. Stacking more blocks is never worse than fewer, even in the worst case, which is what actually makes it safe to go deep.

LayerNorm

Rescales each position's vector to zero mean, unit variance (then a learned scale/shift), before it enters attention or the FFN.

The stream just keeps summing things onto it forever. Why does normalization matter here specifically, and why LayerNorm rather than BatchNorm? Every block adds its output back onto x, so the stream's scale can drift purely from repeated addition, regardless of whether anything useful is being learned. Left unchecked, that drift makes the optimization landscape ill-conditioned — some directions in weight space affect the loss far more than others — and gradients get unstable. LayerNorm resets the scale right before each sub-layer reads the stream, so every sub-layer always sees input in a consistent range no matter how deep it sits.

Why LayerNorm and not BatchNorm: BatchNorm normalizes across the batch dimension, which needs a large-enough batch for stable statistics, couples different sequences in a batch together during training in a way that's awkward once sequence length varies, and needs separately tracked running statistics for inference. LayerNorm normalizes across the feature dimension for one token at one position — independent of batch size, independent of every other token in the sequence, identical math whether you're training on batch 64 or generating one token at a time. Far better fit for language.

Feedforward network

Position-wise (no mixing across positions — attention already did that): expand d_model -> 4*d_model, GELU, contract back down. This is where most of the model's parameters actually live.

Attention already mixes information across positions with a weighted sum. Isn't a separate feedforward block redundant? Attention's output, for a fixed set of attention weights, is a weighted sum of value vectors — and that operation is linear in V, which is itself just a linear projection of x. Stack linear operations without anything nonlinear between them and they collapse mathematically into one bigger linear operation, no matter how many you chain — depth would buy essentially nothing beyond a single matrix multiply's worth of expressive power. The FFN is where an actual nonlinearity (GELU) gets applied, independently, to every position — which is what lets the network compute genuinely nonlinear functions of its input, and is a large part of why depth helps at all. Attention decides where to pull information from; the FFN is closer to where the model actually does something with what it pulled in.

Training

Training uses teacher forcing: the entire target sequence is already known, so it's one forward pass over a fixed window, not a generation loop. get_batch samples random windows of length seq_len + 1 and splits them into x and y, where y is x shifted one position — position i in x is trained to predict position i in y, which is token i+1.

Loss. nn.CrossEntropyLoss expects the class dimension at index 1, but the model outputs (batch, pos, d_vocab) — vocab last. Fixed with rearrange(logits, 'batch pos vocab -> batch vocab pos') before the loss call.

Optimizer & schedule. AdamW, with a linear warmup (0 to full learning rate over warmup_steps) followed by cosine decay (full to 0 over the rest of training), implemented as a multiplier on the base LR via LambdaLR — not an absolute value:

$$\text{lr\_mult}(step) = \begin{cases} step / warmup\_steps & step < warmup\_steps \\[4pt] 0.5\left(1 + \cos(\pi \cdot progress)\right) & \text{otherwise} \end{cases}$$

Grad clipping & checkpointing. clip_grad_norm_ after backward(), before optimizer.step(), caps the global gradient norm so one bad batch can't blow up the weights. Model state gets saved every checkpoint_every steps. And the debug ritual before trusting any of it on a real run: overfit get_batch on one fixed batch to ~0 loss first.

Sampling

Temperature divides logits before softmax. T -> 0 collapses the distribution toward a single argmax (greedy); T -> infinity flattens it toward uniform — every token equally likely regardless of the model's actual preference.

Top-k restricts sampling to the k highest-probability tokens.

Order matters. Mask everything outside the top-k to -inf on the raw logits, before softmax. Masking post-softmax probabilities against a raw-logit threshold is a scale mismatch — comparing a [0,1]-range number against an arbitrary real number — and leaves invalid values for multinomial to choke on.

Generation: uncached vs. KV cache

Generating token by token means running the forward pass repeatedly over a sequence that grows by one each time. The causal mask guarantees something useful: position i's output can never depend on any position after it. So once a token's key and value vectors are computed, they're frozen — permanently correct, no matter what gets appended later.

Uncached (the naive baseline) recomputes Q, K, V for every position, every step — including all the old ones whose values never changed. Per-step attention cost is O(L^2) in the current sequence length L; summed over T generated tokens, total cost is roughly O(T^3).

Cached grows a per-layer K/V buffer, never rewrites it. Each step computes K/V only for the newest token, appends it, and the single new query attends over the entire cache. Per-step cost drops to O(L); total to O(T^2).

Real numbers from benchmark() on this codebase's model (4 layers, d_model=128), generating 40 tokens at each prompt length:

prompt length	uncached tok/s	cached tok/s	speedup
16	755	1691	2.24x
64	818	1788	2.18x
256	237	1411	5.95x

The speedup grows with prompt length — exactly what the complexity argument predicts. The uncached baseline pays quadratically more per step as the sequence grows; the cache mostly doesn't care.

Correctness gate: with a fixed seed and greedy decoding, generate_cached matches generate_uncached token-for-token. The cache is an optimization, not a different algorithm — it has to reproduce the exact same output before its speed counts for anything.

One thing that tripped me up while building this: Q never gets cached. A query is used exactly once — to build one attention row — and then discarded. Nothing in the future ever looks up an old query. K/V, by contrast, get read by every future query for the rest of generation, which is exactly why those are worth storing.

For the fully worked version of this — every Q, K, V, score, mask, and softmax weight, for one real prefill call and one real decode call, by hand — I wrote up a companion trace: KV Cache Trace.

---

Code: paper2code/transformer on GitHub.