Attention

Level 2 builds the mechanism from nothing, starting with a glance around a room.

The everyday picture. Imagine you're in a meeting and someone says "it's broken, can you fix it?" To know what "it" means, you glance around the room: at the laptop on the table, at the person who just walked in, at the whiteboard. You pay a lot of attention to the laptop, a little to everything else, and your understanding of "it" becomes mostly laptop.

That glance is attention. Every word in a sentence gets to look at every other word, decide how relevant each one is, and rebuild its own meaning as a mix of the relevant ones. "Bank" next to "river" becomes a riverbank; next to "loan" it becomes a lender. The word is the same, but what it paid attention to is different.

Take "The animal didn't cross the street because it was tired" and follow the word "it". To keep the numbers small, it compares itself with just three other words. Each comparison gives a relevance score (how the score is made comes next). Then we turn scores into shares of attention that add up to 1:

1. Exponentiate each score (e^score), which makes every number positive and stretches the gaps between them.
2. Divide each by the total, 7.39 + 2.72 + 1.65 = 11.76.

Those two steps are called softmax. The new meaning of "it" is 63% "animal", 23% "tired" and 14% "street". Without being told any grammar rule, the model has worked out that "it" is the animal.

In code: worked_example_it runs these three scores through softmax and returns each word's share.

Softmax, decoded

Softmax turns any list of numbers, positive or negative, into shares that are all positive and add up to exactly 1. Think of it as a vote where louder voices get disproportionately more say.

import math
# animal, tired, street
z = [2.0, 1.0, 0.5]
# e^(z_i) for each score
exps = [math.exp(z_i) for z_i in z]
[round(e, 2) for e in exps]  # → [7.39, 2.72, 1.65]
# Σ_j e^(z_j)
total = sum(exps)
round(total, 2)  # → 11.76
# softmax(z)_i: each share of the total
[round(e / total, 2) for e in exps]  # → [0.63, 0.23, 0.14]

Why e to the power of the score, and not just the score divided by the total? Three reasons you can check on the table above:

1. Always positive. Scores can be negative; e to any power is positive, so no word ever gets a negative share.
2. Order is kept. A higher score always gets a bigger share.
3. Gaps are stretched. Adding 1 to a score multiplies its e-value by 2.72, so the winner pulls ahead decisively. (Plain division would give "animal" 2.0 / 3.5 = 0.57, a timid 57%, and would break on negative scores.)

One more property matters later: adding the same number to every score changes nothing, because it multiplies the top and bottom of the fraction by the same amount. The code uses this to subtract the largest score before exponentiating, which stops e¹⁰⁰⁰ from overflowing to infinity. See primer.notation for exponents and sums from scratch.

In code: softmax subtracts the largest score before exponentiating, and turns masked scores of −∞ into weights of exactly 0.

Each word turns its vector into three new vectors by multiplying it by three learned matrices:

• a query: what am I looking for? ("it" is looking for a noun that can be tired.)
• a key: what do I offer to others? ("animal" offers: a noun, alive.)
• a value: what I actually hand over if someone picks me.

A word's score for another word is the dot product of its query with the other word's key: multiply them position by position and add.

q = [1, 2]
k = [3, 0.5]
# Σ over m of q_m k_m
sum(q_m * k_m for q_m, k_m in zip(q, k))  # → 4.0

Vectors that point the same way give big positive scores, vectors at right angles give zero, and opposite vectors give negative scores. That is why the dot product works as a relevance score. The values are then blended with the softmax shares, exactly as in the table above.

A library search is a good analogy. The query is what you type into the search box, the keys are the catalogue entries, and the values are the books on the shelf. You get back a blend of books, weighted by how well each catalogue entry matched your search.

flowchart LR
  X[Token vectors] --> Q[Query] & K[Key] & V[Value]
  Q --> S[Scores<br/>Q times K]
  K --> S
  S --> SC[Scale by<br/>sqrt of d_k]
  SC --> M[Causal mask<br/>hide future tokens]
  M --> SM[Softmax<br/>weights sum to 1]
  SM --> W[Weighted sum<br/>of values]
  V --> W
  W --> O[Context-aware vectors]

Written as one formula, with every word's query stacked into a matrix Q (and likewise K and V), the whole diagram is:

import math
q_it = [1, 1, 1, 1]
# keys: animal, tired, street
K = [[1, 1, 1, 1], [1, 1, 0, 0], [1, 0, 0, 0]]
# values, one row per word
V = [[1, 0], [0, 1], [1, 1]]
d_k = len(q_it)
# q Kᵀ / √d_k
scores = [sum(q * k for q, k in zip(q_it, k_j)) / math.sqrt(d_k) for k_j in K]
scores  # → [2.0, 1.0, 0.5]
exps = [math.exp(s) for s in scores]
# softmax of the row
weights = [e / sum(exps) for e in exps]
[round(w, 2) for w in weights]  # → [0.63, 0.23, 0.14]
# (…)V: blend the values
[round(sum(w * v[c] for w, v in zip(weights, V)), 2) for c in range(2)]  # → [0.77, 0.37]

In practice this one line runs in every layer of every modern language model, for every word, many times per word generated. Nearly everything about a model's speed and memory use traces back to it.

In code: scaled_dot_product_attention is the whole formula, one commented step per box of the diagram, and returns both the blended values and the attention weights.

For one sequence of n tokens with model width d_model:

A decoder model (GPT, Claude, Llama) is trained to predict the next token. If position i could attend to position i+1, it could simply read the answer. So the scores for future positions are set to −∞ before softmax. Because e^−∞ = 0, those weights come out as exactly zero.

flowchart LR
  subgraph S["Scores (n × n)"]
    direction TB
    r0["row 'The':    ✓ · · ·"]
    r1["row 'cat':    ✓ ✓ · ·"]
    r2["row 'sat':    ✓ ✓ ✓ ·"]
    r3["row 'down':   ✓ ✓ ✓ ✓"]
  end
  S --> MASK["set every · to −∞"] --> SOFT["softmax per row"] --> OUT["each row sums to 1<br/>using only the past"]

Masking only the future is also what makes generation cheap: a token's output never changes when later tokens arrive, so it can be computed once and cached. See primer.ml.inference for the KV cache.

In code: causal_mask builds the lower triangle of allowed positions, and scaled_dot_product_attention sets every score outside it to −∞ before softmax. worked_example_sentence holds the hand-picked queries and keys the heatmap above draws.

Everyday picture Roll one die and the result swings between 1 and 6. Add up a hundred dice and the total swings far more, by dozens in either direction. A dot product is exactly that kind of sum: one small random-ish product per position. The wider the vectors, the more terms get added, and the wilder the scores swing.

Tiny example With d_k = 4, a query and a key of random ±1 entries give four products of +1 or −1, so the score lands between −4 and +4 and is usually within ±2. With d_k = 256 it's a sum of 256 such products: usually within ±16. Softmax of scores like (16, 3, −9) gives about (0.999998, …, …). All the attention lands on one word, whether or not that is right.

import itertools, math, statistics
d_k = 4
scores = [sum(products) for products in itertools.product([-1, 1], repeat=d_k)]
# Var(q · k) = d_k
statistics.pvariance(scores)  # → 4
# Var(q · k / √d_k) = 1
statistics.pvariance([s / math.sqrt(d_k) for s in scores])  # → 1.0
# the typical raw swing at d_k = 128
round(math.sqrt(128), 1)  # → 11.3

Why "all attention on one word" is bad, beyond being a wrong answer: softmax has stopped responding. Its gradient (the signal training uses to adjust the weights) is diag(p) − p pᵀ, and when one share is 1 and the rest are 0, every entry of that matrix is 0. The query and key weights receive no signal and stop learning. This is called saturation. sqrt_dk_experiment measures all of this.

In one sentence: dot products grow with dimension, large scores saturate softmax and kill gradients, and scaling by √d_k keeps scores in the range where softmax is smooth and trainable.

In code: softmax_jacobian builds the diag(p) − p pᵀ matrix, so you can watch every entry fall to 0 as one share approaches 1.

Instead of one attention with d_k = d_model, run h heads in parallel, each with d_k = d_model / h and its own W_q, W_k and W_v. The total compute is the same, but each head can specialise: one tracks syntax, another which noun a pronoun refers to, another the previous token.

flowchart LR
  X["X (n × d_model)"] --> P["project with W_q, W_k, W_v"]
  P --> H1["head 1<br/>attention on d_model/h dims"]
  P --> H2["head 2"]
  P --> H3["..."]
  P --> Hh["head h"]
  H1 & H2 & H3 & Hh --> C["concatenate<br/>(n × d_model)"]
  C --> WO["mix with W_o"] --> Y["output (n × d_model)"]

In code: MultiHeadAttention holds the four projections W_q, W_k, W_v and W_o; calling it slices the projections into one slab per head, runs scaled_dot_product_attention on every head at once, and glues the results back together before W_o.

With grouped-query attention (GQA), several query heads share one key/value head. Llama 3 8B has 32 query heads but only 8 KV heads, so it stores 4× fewer keys and values per token. Multi-query attention (MQA) is the extreme: one KV head for everyone.

flowchart TB
  subgraph MHA["MHA: 4 query heads, 4 KV heads"]
    q1a[Q1]-->kv1a[KV1]
    q2a[Q2]-->kv2a[KV2]
    q3a[Q3]-->kv3a[KV3]
    q4a[Q4]-->kv4a[KV4]
  end
  subgraph GQA["GQA: 4 query heads, 2 KV heads"]
    q1b[Q1]-->kv1b[KV1]
    q2b[Q2]-->kv1b
    q3b[Q3]-->kv2b[KV2]
    q4b[Q4]-->kv2b
  end
  subgraph MQA["MQA: 4 query heads, 1 KV head"]
    q1c[Q1]-->kv1c[KV1]
    q2c[Q2]-->kv1c
    q3c[Q3]-->kv1c
    q4c[Q4]-->kv1c
  end

In code: MultiHeadAttention takes a number of KV heads: fewer than the query heads is GQA, one is MQA. MultiHeadAttention.kv_params counts the key and value weights that shrink.

The score matrix is n × n per head per layer, so compute and memory grow with n². Doubling the context roughly quadruples attention's cost.

The standard mitigations each attack this picture. FlashAttention produces the exact same result but computes it in tiles sized for fast on-chip GPU memory, so the n × n matrix is never written to slow memory. Sliding-window and sparse attention let each token see only some others. GQA and MQA shrink the KV cache. State-space models such as Mamba replace attention with a linear-time recurrence.

In code: attention_cost counts the quadratic score-and-mix FLOPs and the linear projection FLOPs plotted above.

Answer each one out loud or on paper before you open it. If you can explain it, you know it.