Attention & Transformers

The MLP (multilayer perceptron) language model of section 5 reads a fixed window of tokens and treats each position with its own weights. The transformer (Vaswani et al., 2017) replaces that with attention: every token looks at every earlier token and decides, from their content, how much of each to take.

This page builds, in order: attention; its masked and multi-head forms; the rotary position encoding and grouped-query attention used by current open models; and finally the full transformer block. Each part has a small example with numbers.

Attention: a soft lookup

Think of a dictionary lookup: you have a question, you compare it with the labels of the entries, and you read the entry whose label matches. Attention does the same thing, but softly: it compares the question with every label, and reads a weighted mix of all the entries, with more weight where the match is better.

Each token produces three vectors from its embedding, with three learned linear layers:

• a query $q$: what it is looking for (the question);
• a key $k$: what it offers (the label);
• a value $v$: what it passes on (the content of the entry).

The output for a token is a weighted average of all the values, where each weight says how well the token's query matches that key:

where:

• $Q$ stacks the queries as rows, shape $(T, d)$ for $T$ tokens; $K$ stacks the keys, shape $(S, d)$; $V$ stacks the values, shape $(S, d_v)$;
• $QK^\top$ is the table of all dot products between queries and keys ($K^\top$ is $K$ transposed): entry $(t, s)$ measures how well query $t$ matches key $s$;
• $S$ holds these scores, divided by $\sqrt{d}$ (explained below);
• the softmax is taken along each row, turning the scores of one query into weights that add up to 1;
• multiplying by $V$ takes, for each query, the weighted average of the values.

Example. Three tokens, $d = 2$, with queries and keys equal for simplicity: token 0 has $(1, 0)$, token 1 has $(0, 1)$, token 2 has $(1, 1)$. Values: $v_0 = (1, 0)$, $v_1 = (0, 10)$, $v_2 = (5, 5)$.

1. Scores for token 0: its query $(1, 0)$ against the three keys gives dot products $1, 0, 1$; divided by $\sqrt{2}$: $(0.71, 0, 0.71)$.
2. Weights: softmax of the scores, $(0.40, 0.20, 0.40)$. Token 0 matches tokens 0 and 2 best.
3. Output: $0.40 \cdot (1, 0) + 0.20 \cdot (0, 10) + 0.40 \cdot (5, 5) \approx (2.41, 3.98)$.

For token 1 the weights are $(0.20, 0.40, 0.40)$ and the output $(2.20, 6.02)$: it takes more of $v_1$, whose key matches its query.

Bahdanau, Cho and Bengio (2015) introduced attention for translation; the transformer made it the only way tokens interact.

Why divide by $\sqrt{d}$. If the entries of $q$ and $k$ are independent, with mean 0 and variance 1, their dot product is a sum of $d$ products, and its variance is $d$. With $d = 128$ the scores have a standard deviation of about 11. Softmax of numbers that far apart puts almost all the weight on one token, and its gradients become tiny, so the model stops learning which tokens to attend to. Dividing by $\sqrt{d}$ brings the variance back to 1 (Vaswani et al., section 3.2.1).

Example. In one random draw of 8 scores with standard deviation 11, the largest softmax weight is $0.99999$: all the attention goes to one token. The same scores divided by $\sqrt{128} \approx 11.3$ give a largest weight of $0.41$: the attention is spread over several tokens.

Causal masking

A language model must not see the future: the prediction at position $t$ may use tokens up to $t$ only. Otherwise, during training, it could simply read the answer. A causal mask sets every score above the diagonal to $-\infty$ before the softmax, so those weights become exactly 0:

where $t$ is the position of the query and $s$ the position of the key. Since $e^{-\infty} = 0$, a token gets no weight from any later token.

Example. The same three tokens. The scores, with the mask applied:

query	key 0	key 1	key 2
token 0	0.71	$-\infty$	$-\infty$
token 1	0	0.71	$-\infty$
token 2	0.71	0.71	1.41

After the softmax, the weights of the three rows are $(1, 0, 0)$, $(0.33, 0.67, 0)$ and $(0.25, 0.25, 0.50)$. Token 0 can only look at itself, so its output is exactly $v_0$. Token 2, the last one, sees everything, as before. The allowed cells form a lower triangle.

This is what makes training efficient: one forward pass over a sequence of length $T$ computes all $T$ next-token predictions at once, each one seeing only its own past.

Multi-head attention

One attention pattern per layer is limiting: a token may need to look at the previous word for grammar, and at a name ten tokens back for meaning. Multi-head attention (MHA) runs several attentions side by side, each with its own queries, keys and values. It splits the model dimension $D$ into $H$ heads of size $D/H$, runs attention in each head independently, places the results side by side and mixes them with an output projection $W_O$:

where $X$ holds the input vectors, $\mathrm{head}_h$ is the output of attention in head $h$, and the square brackets mean "place side by side".

Example. $D = 8$ and $H = 2$ heads of size 4. Each token's query vector of 8 numbers is read as two queries of 4 numbers: the first 4 for head 0, the last 4 for head 1. Each head computes its own scores and weights, and returns 4 numbers per token. Side by side they make 8 numbers again, and $W_O$, an $8 \times 8$ matrix, mixes them.

In code there is no loop over heads: project once to shape $(B, T, D)$, reshape to $(B, T, H, D/H)$, transpose to $(B, H, T, D/H)$, run batched attention, then undo the transpose and the reshape. The einsum exercise of the NumPy section is exactly this shape work.

Position: rotary embeddings

Attention by itself ignores order: shuffle the tokens, and each output is shuffled the same way. "dog bites man" and "man bites dog" would look the same. The original transformer added fixed sine and cosine waves to the embeddings. RoPE (rotary position embedding; Su et al., 2024) instead rotates the queries and keys, by an angle that grows with the position.

The $d$ numbers of a vector are grouped into $d/2$ pairs, and each pair is treated as a point in the plane. Pair $i$ of the vector at position $m$ is rotated by the angle $m\theta_i$:

where:

• $x_{2i}$ and $x_{2i+1}$ are the two numbers of pair $i$, and $x'$ the rotated ones;
• $m$ is the token's position: $0, 1, 2, \dots$;
• $\theta_i = 10000^{-2i/d}$ is the rotation speed of pair $i$: the first pairs turn fast, the last ones slowly.

Example. The pair $(1, 0)$ rotated by $90°$: $\cos 90° = 0$ and $\sin 90° = 1$, so it becomes $(1 \cdot 0 - 0 \cdot 1,\ 1 \cdot 1 + 0 \cdot 0) = (0, 1)$, the same point turned a quarter circle. Its length does not change, only its direction.

The key property: rotating $q$ by $m\theta$ and $k$ by $n\theta$ changes their dot product only through the difference $m - n$. So attention scores depend on the relative position of two tokens, and nothing is added to the values.

Example. One pair, $\theta = 0.5$, query $q = (1, 2)$ and key $k = (3, 1)$; without rotation, $q \cdot k = 5$.

• Query at position 3, key at position 1 (distance 2): the score becomes $-1.506$.
• Query at position 7, key at position 5 (distance 2 again): $-1.506$, exactly the same.
• Query at position 5, key at position 7 (distance $-2$): $6.909$, different.

The score depends on how far apart the two tokens are, not on where they are.

LLaMA (Large Language Model Meta AI) uses RoPE. Implementations differ in how they pair dimensions: the RoFormer paper pairs neighbours $(2i, 2i+1)$, as here, while some libraries pair dimension $i$ with $i + d/2$. The two are equivalent up to a fixed reordering of the weights.

Multi-query and grouped-query attention

When a model generates text, it keeps the keys and values of every previous token in memory, in a cache (section 11), so that it does not recompute them for every new token. That cache, not the computation, limits how many texts can be generated at once and how long they can be. Fewer key/value heads mean a smaller cache:

• Multi-query attention (Shazeer, 2019) keeps $H$ query heads but a single key head and a single value head, shared by all of them.
• GQA (grouped-query attention; Ainslie et al., 2023) is the middle ground: $G$ key/value heads, each shared by a group of $H/G$ query heads. Query head $h$ uses key/value head $\lfloor h / (H/G) \rfloor$, where $\lfloor \cdot \rfloor$ rounds down.

Example. $H = 32$ query heads and $G = 8$ key/value heads: each group of $32 / 8 = 4$ query heads shares one key/value head. Query heads 0 to 3 use head 0, query heads 4 to 7 use head 1, and so on: query head 13 uses $\lfloor 13 / 4 \rfloor = 3$. The cache stores 8 heads instead of 32: it is 4 times smaller than with full multi-head attention.

The transformer block

A model is a stack of identical blocks. Each block has two sub-layers:

• attention, where tokens exchange information;
• an MLP, where each token is processed on its own.

Each sub-layer is wrapped in a residual connection (He et al., 2016): its output is added to its input, so the block only adds a correction to what it receives:

where $x$ holds the token vectors and $\mathrm{Norm}$ is LayerNorm (layer normalization) or RMSNorm (root mean square layer normalization) from section 4.

Why add instead of replace? With the addition, the gradient has a direct path from the loss back to the first layers (the derivative of $x + f(x)$ always contains the 1 from $x$), so even a stack of dozens of blocks still trains. The vector $x$ that flows through the stack, collecting all these corrections, is called the residual stream.

Normalizing the input of each sub-layer (pre-norm) rather than its output is what GPT-2 (Generative Pre-trained Transformer 2) did, and Xiong et al. (2020) showed it makes training stable without a long warmup.

The MLP widens each vector and narrows it back. GPT-2 uses $\mathrm{GELU}(xW_1)W_2$, with a hidden size of $4D$; GELU (Gaussian error linear unit) is a smooth version of ReLU (rectified linear unit). LLaMA uses SwiGLU (Swish-gated linear unit; Shazeer, 2020), a gated variant with three matrices:

where:

• $W_1$ and $W_3$ widen the vector, and $W_2$ narrows it back;
• $\mathrm{SiLU}(z) = z\,\sigma(z)$ is the SiLU (sigmoid linear unit), $z$ times its sigmoid;
• $\odot$ is the elementwise product: the second branch, $x W_3$, works as a gate that scales each entry of the first.

Example. For one entry, if $xW_1$ gives $2$ and $xW_3$ gives $0.5$: $\mathrm{SiLU}(2) = 2 \cdot \sigma(2) \approx 2 \cdot 0.881 = 1.76$, times the gate $0.5$, is $0.88$. A gate of 0 would switch that entry off entirely.

A full model is: token embedding, $L$ blocks, a final norm, and a linear layer to $V$ logits. Every piece of that line is now an exercise you have written.

References

• A. Vaswani et al., Attention Is All You Need, NeurIPS, 2017. (arXiv)
• D. Bahdanau, K. Cho, Y. Bengio, Neural Machine Translation by Jointly Learning to Align and Translate, ICLR, 2015.
• J. Su, M. Ahmed, Y. Lu, S. Pan, W. Bo, Y. Liu, RoFormer: Enhanced Transformer with Rotary Position Embedding, Neurocomputing 568, 127063, 2024. (arXiv)
• N. Shazeer, Fast Transformer Decoding: One Write-Head is All You Need, arXiv, 2019. Multi-query attention.
• J. Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, EMNLP, 2023.
• K. He, X. Zhang, S. Ren, J. Sun, Deep Residual Learning for Image Recognition, CVPR, 770–778, 2016.
• R. Xiong et al., On Layer Normalization in the Transformer Architecture, ICML, 2020.
• N. Shazeer, GLU Variants Improve Transformer, arXiv, 2020. SwiGLU.
• D. Hendrycks, K. Gimpel, Gaussian Error Linear Units (GELUs), arXiv, 2016.
• A. Radford et al., Language Models are Unsupervised Multitask Learners, OpenAI, 2019. Section 2.3: layer normalization moved to the input of each sub-block.
• H. Touvron et al., LLaMA: Open and Efficient Foundation Language Models, arXiv, 2023. Section 2.2: pre-normalization, SwiGLU, rotary embeddings.