Transformer Math: Big Picture

The substrate: tokens, embeddings, and the residual stream

Everything a transformer does happens to one object: a stream of vectors. Text is tokenized into integer IDs, each ID indexes a row of an embedding matrix E: [V, d], and the result is an activation tensor X: [N, d] — N tokens, each a d-dimensional vector. That tensor is the model’s working memory.

The organizing idea is the residual stream. Every block reads the stream, computes an update, and adds it back: x ← x + f(x) — nothing overwrites, layers only contribute. This is why gradients survive depth (the derivative of x + f(x) carries a +1 term that resists vanishing), and why interpretability can read the stream as a running sum of features. Hold this picture: a d-wide highway of vectors, each layer a merge lane adding traffic. Every formula that follows either computes an addition or trains the weights inside one.

Attention: the one operation that makes it a transformer

Attention is the mechanism that lets a token pull information from other tokens. Project the stream three ways — queries, keys, values — via Q = XW_Q, K = XW_K, V = XW_V, then:

Attention(Q, K, V) = softmax(QKᵀ / √d_k) V

QKᵀ scores every query against every key (shape [N, N]); the √d_k divisor keeps the scores from saturating the softmax; the softmax turns each row into a probability distribution over tokens; multiplying by V forms a weighted average of value vectors. The dot product is the similarity measure — geometry doing semantics. That [N, N] score matrix is the source of the famous O(N²) cost, and nearly every later efficiency idea attacks that quadratic term or the memory it implies.

The block: attention, FFN, normalization, residual

A transformer layer wraps attention with three partners. Written in the modern pre-norm form:

x ← x + MHA(Norm(x))
x ← x + FFN(Norm(x))

Multi-head attention runs h attention computations in parallel on d/h-dim slices, so heads learn different relations, then concatenates. The FFN is a per-token MLP, FFN(x) = W₂ · act(W₁x), that widens to ~4d and back — where most parameters and compute live. Normalization (LayerNorm, or the cheaper RMSNorm) keeps activation scales stable; residual addition keeps gradients alive. Stack this block L times for the whole trunk, with a sizing shortcut that falls straight out: parameters ≈ 12 · L · d², split between the four attention projections (4d²) and the FFN (8d²).

Training I: loss and the gradient signal

Training is next-token prediction. The model emits logits, softmax turns them into a distribution p over the vocabulary, and cross-entropy scores it against the true next token:

L = −(1/N) Σ_t log p(x_t | x_<t)

The beauty of pairing softmax with cross-entropy is the gradient: with respect to the logits it collapses to ∂L/∂z = p − y — predicted minus true, one-hot y: clean, bounded, and the reason the combination is universal. Backpropagation applies the chain rule backward; because the forward pass is a residual sum, the backward pass is too — the same +1 highway, run in reverse. Training stability then reduces to keeping this signal well-scaled from the last layer back to the first.

Training II: the optimizer and what it costs

Gradients say which way; the optimizer says how far. Adam, the default, keeps two running averages per parameter — the mean m and variance v of recent gradients — and steps with

θ ← θ − α · m̂ / (√v̂ + ε)

Dividing by √v̂ gives each parameter its own effective learning rate — why Adam trains transformers so robustly. The hidden bill is memory: m and v double the footprint, so a model with P parameters needs roughly P for weights plus 2P for optimizer state plus gradients — the real reason training costs far more RAM than inference. On a memory-budgeted machine the optimizer state, not the model, is often what you cannot afford, pushing you toward smaller models, gradient accumulation, or 8-bit optimizers. A raw loop still diverges easily, so a familiar set of guardrails keeps it stable: learning-rate warmup into cosine decay, gradient clipping when ‖g‖ > c, and mixed precision with an fp32 master copy and loss scaling to stop small gradients underflowing. None change what is learned — they keep the signal in a range where the math behaves.

Scaling laws: how far the loss will fall

Zoom out from one run and a startling regularity appears: loss falls as a smooth power law in parameters N, data D, and compute C. The Chinchilla form captures it:

L(N, D) = E + A/N^α + B/D^β

E is the irreducible entropy of language; the two terms are the penalties for a finite model and finite data. The practical punchline — given a fixed compute budget C ≈ 6ND, scale N and D together, roughly 20 tokens per parameter — overturned the ‘just make it bigger’ instinct and is why small models trained on lots of data (the whole SLM premise) are competitive. Scaling laws bridge the micro-math of one step and the macro-question of what to build: they let you predict a run’s payoff before spending the compute, and explain why emergent abilities look like phase changes on a smooth curve.

Inference I: prefill, decode, and the KV cache

Generation splits into two regimes with very different math. Prefill processes the whole prompt at once — a big compute-bound matmul over N tokens. Decode then emits one token at a time, and here lies the key trick: rather than recompute attention over the whole history each step, cache the keys and values already computed. The KV cache grows linearly:

KV_bytes = 2 × L × n_kv × d_head × seq × bytes_per_elt

The factor 2 is keys plus values. It restores a linear per-step cost, but at a memory price that can dwarf the model itself for long contexts. Decode is therefore memory-bandwidth-bound: each new token must read every weight and the whole cache to do very little arithmetic. That single fact — low arithmetic intensity — drives almost every serving optimization.

Inference II: making decode fast

Because decode is bandwidth-bound, the wins come from moving less memory, not doing less math. FlashAttention never materializes the [N, N] score matrix; it tiles the computation and fuses softmax so attention stays in fast on-chip memory, cutting reads and writes. Grouped-query attention shares one set of K/V heads across several query heads, shrinking the KV cache (the n_kv term above) by a large factor for a tiny quality cost. Continuous / in-flight batching packs many users’ decode steps into one weight-read, amortizing the expensive memory traffic across a batch and raising throughput enormously. The through-line: each targets the same bottleneck the KV-cache formula exposed. On CPU, where bandwidth is scarcer, they are the difference between an SLM that answers interactively and one that stalls.

Quantization: trading bits for bandwidth

If decode is bottlenecked on moving bytes, the direct fix is fewer bytes. Quantization stores weights (and often the KV cache) in low precision — int8, int4 — and reconstructs on the fly:

x ≈ scale × (q − zero_point)

A per-group scale maps the small integer q back to a float. Going from fp16 to int4 cuts memory and the bandwidth bill roughly 4× — often a near-linear speedup for a memory-bound decode — while error stays tolerable because the group-wise scale tracks each region’s dynamic range. The cost is precision, and the art is where to spend it: keep outlier channels or the KV cache at higher bits, quantize the bulk hard. It is the keystone of CPU inference — how a model that would not fit, fits, and a decode that would crawl, moves.

PEFT: adapting without retraining the whole model

Full fine-tuning re-optimizes every parameter and re-incurs the full 3P memory bill from the optimizer section. Parameter-efficient fine-tuning sidesteps it. The dominant method, LoRA, freezes the pretrained weight W and learns a low-rank update:

W’ = W + ΔW = W + BA,   B: [d, r],  A: [r, d],  r « d

Only A and B train, so trainable parameters drop by orders of magnitude and optimizer state with them. The bet — borne out in practice — is that the update a task needs is intrinsically low-rank even though W is full-rank. Variants like QLoRA (a quantized frozen base) and IA³ trade off memory against expressiveness. PEFT is where scaling laws meet small hardware: you cannot afford to move a large model far, but a rank-16 nudge is often all a downstream task requires.

Decoding: turning logits back into language

The final block emits logits; a tied output projection (often the transpose of the embedding matrix, saving Vd parameters) maps the stream to [N, V]. Then decoding chooses tokens. Temperature reshapes the distribution before sampling, softmax(z / T) — T < 1 sharpens, T > 1 flattens. Top-k and top-p truncate the tail so the model never samples nonsense; beam search keeps several high-probability sequences for tasks with a right answer. Speculative decoding lets a small draft model propose several tokens that the big model verifies in one pass — a direct assault on the one-token-per-memory-read tax. And perplexity = exp(L) closes the loop: the same cross-entropy that trained the model, exponentiated, is how we measure the decoder’s surprise on held-out text.

The map, reconnected

Trace one thread through the graph and the unity shows. The width d sets attention’s cost, the parameter count (12Ld²), the optimizer’s 3P footprint, the KV cache’s size, and the bits quantization must compress — one number rippling through every thread. Inference inherits every training choice: a bigger model is a bigger bandwidth bill, which quantization, GQA, and speculative decoding fight to pay down, and PEFT lets you adapt without repaying it. No piece is independent — learn the math as a system and you can reason about the ripple instead of being surprised by it.