The Transformer, End to End
The researcher-round staple: draw the architecture from memory and explain every tensor. We build it once, honestly — token → embedding → attention computed on four real tokens (Q/K/V, a QKᵀ heatmap, softmax, the weighted sum, every number real) → multi-head → the KV cache and the MHA/MQA/GQA/MLA memory ladder dialed live → RoPE as a rotation → the MLP and residual highway → norm placement → the block stacked N times → a parameter count that reproduces Llama-3-8B to the last billion.
Concept · AI / ML. The source ↗
A free, interactive, animated visual explainer of The Transformer, End to End — built to be understood, not skimmed.
Questions
- What are Q, K, and V in transformer attention?
- Each token produces three vectors by multiplying its embedding by three learned matrices: a query (Q — what this token is looking for), a key (K — what it offers to others), and a value (V — what it hands over if it is attended to). A token's query is compared against every token's key by a dot product to produce scores; those scores are scaled by √d_k and softmaxed into weights that sum to one; and the token's new vector is the weighted sum of everyone's value vectors. The Attention Is All You Need paper states the recipe directly: "We compute the dot products of the query with all keys, divide each by √d_k, and apply a softmax function to obtain the weights on the values." Q and K decide who attends to whom; V is the information that actually flows.
- Why is attention divided by the square root of d_k?
- It is a variance fix. A dot product sums d_k element products, so its magnitude grows like √d_k; left unscaled, the logits become large, softmax saturates (one weight near 1, the rest near 0), and the gradient through the softmax nearly vanishes, which stalls learning. Dividing by √d_k keeps the logits at a stable scale. The authors are explicit about the failure they are avoiding: "We suspect that for large values of d_k, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients."
- What is the KV cache and why does it dominate memory at long context?
- At inference a transformer generates one token at a time, and each new token must attend over every earlier token. Since a fixed token's key and value vectors never change, they are cached: after a token is generated its K and V (one pair per head, per layer) are stored, and the next token computes only its own query, reads the whole cache, attends, and appends its own K and V. This turns generation from quadratic recomputation into one query against the cache. The cost is memory: the cache grows linearly with sequence length, number of layers, and number of KV heads, so at long context it becomes the dominant memory consumer — which is exactly what MQA, GQA, and MLA are designed to shrink.
- What is the difference between MHA, MQA, GQA, and MLA?
- All four are attention; they differ in how many key/value heads they keep, which is what sets the KV-cache size. MHA (multi-head attention) gives every query head its own K and V head. MQA (multi-query attention) shares a single K/V head across all query heads, cutting the cache by the full head count but often losing quality. GQA (grouped-query attention) is the middle ground — query heads are split into groups with one K/V head per group, which is why Llama-3-8B has 32 query heads but only 8 KV heads. MLA (multi-head latent attention, from DeepSeek-V2) is a different move: it compresses K and V into a shared low-rank latent vector, caches only that, and reconstructs each head's K/V on the fly — reported to reduce the KV cache by 93.3% versus its MHA baseline.
- Why do transformers need positional encoding like RoPE?
- Because attention is order-blind. A dot product between a query and a key does not depend on where either token sits in the sequence, so without added position information a transformer is permutation-equivariant: "dog bites man" and "man bites dog" are the same multiset of tokens and produce the same outputs — the model cannot tell who bit whom. Positional encoding injects order. RoPE (rotary position embedding), used in Llama and most current models, rotates each token's query and key by an angle proportional to its position; because both are rotated, their dot product ends up depending only on the relative distance between the two tokens, which also lets the scheme extend to longer contexts than were seen in training.
Related explainers
- Design an LLM Serving Platform
- Distributed Training, End to End
- Design a Feature Store
- How to Design a System in 60 Minutes
- Loading a Safetensors Checkpoint on a Multihost Cluster, Line by Line
- Python Gotchas That Fail Interviews
- Errors, Edge Cases & Testing in the Room
- Designing a Clean Python API Under Pressure