Part II — Attention

Attention, the idea — weighting by relevance

Lesson 03 gave us free training pairs: for every position, the next token is the label. But to predict the token at position t, the model must first fold together everything it has read up to t — and a fixed recipe for that folding is far too blunt. This lesson introduces the one idea that fixes it: let each token build its own summary of the sequence by weighting the others by relevance. We do it here with no trainable parameters at all — just to see the shape of the mechanism — so that lesson 05 can bolt the knobs on.

1 · The long-range problem

Language is full of dependencies that reach across a sentence. In "The keys that the woman near the tall windows were looking for…", the verb agrees with "keys," five words back, not with "woman" or "windows." To predict the next token well, a model has to be able to reach the right earlier token no matter how far away it sits. Three older strategies each fail at exactly this.

The recurrent picture is worth dwelling on, because attention is the direct answer to its flaw. An encoder–decoder RNN reads the input one token at a time, updating a hidden state that is meant to carry the entire meaning forward; the decoder then generates from that final state. The trouble is that a single fixed-size vector must remember an arbitrarily long input, and the decoder has no direct access to any earlier state — only to the one compressed summary. When a dependency spans a long distance, the relevant detail has usually been squeezed out. The historical fix (Bahdanau, 2014) let the decoder look back at all the input positions and weight them by importance at each step — and a few years later researchers found you did not need the RNN at all. That weighting-by-importance is attention, and it is what we build now, from scratch, for a single sequence attending to itself.

2 · The core move: a context vector per token

Here is the whole idea in one line. For each position i we compute a new vector z_i — its context vector — as a weighted sum of all the token vectors in the sequence:

The numbers α_ij are the attention weights: α_ij says how much position i should draw from position j. They are non-negative and, for a fixed i, they sum to 1 — so z_i is a genuine blend, a convex mixture of the token vectors. Think of z_i as an enriched embedding of token i: it is still "about" token i, but it has folded in exactly the parts of the rest of the sequence that i found relevant.

Two extremes make the mechanism concrete. If every weight is equal, α_ij = 1/T, the context vector is just the plain average — the rigid bag-of-tokens we were trying to escape. If one weight is 1 and the rest are 0, the context vector is that single token — full focus on one position. Everything useful lives in between: a distribution that is sharp on the few tokens that matter and near zero elsewhere. The entire remaining question is where those weights come from. Notice that we have not said anything about learning yet — we only need a rule that turns a pair of token vectors into a relevance number.

3 · Where the weights come from: dot-product similarity

With no parameters to spare, the only material we have is the embeddings themselves. So let the relevance of token j to a query token i be their dot product:

A dot product is just element-wise multiply, then sum — but it is also a measure of alignment. It is largest when two vectors point the same way, near zero when they are orthogonal, and negative when they point apart. So it is exactly a similarity score: tokens whose embeddings have drifted close together (lesson 02 taught them to, for related words) produce a high score and will attend to each other strongly; unrelated tokens score low and are largely ignored. The token doing the looking is called the query; each token it compares against is a key. Right now the query and the key are the same raw embedding — token i in the query role and token j in the key role are literally x_i and x_j — which is precisely the limitation lesson 05 removes.

4 · From scores to weights: softmax

The raw scores are unbounded real numbers — some large, some negative — so they cannot be attention weights directly. We need to turn a row of scores [s₁, …, s_T] into weights that are (a) all positive and (b) sum to 1, so the context vector stays a proper blend. The obvious hack is to divide each score by the row sum. It fails the moment a score is negative or the scores straddle zero: you can get negative "weights," or a sum of zero and a division blow-up. Normalizing raw values is fragile.

The right tool is softmax, which exponentiates first and then normalizes:

The exponential makes every term strictly positive regardless of sign, so the weights are always a valid distribution; dividing by the sum guarantees they add to 1. Softmax also behaves gracefully under gradients — the property lesson 05 leans on hard — and it has an intuitive knob: multiply all scores by a constant before exponentiating and you control how peaked the distribution is. Scale the scores up and softmax sharpens toward one-hot (all focus on the top token); scale them down and it flattens toward the uniform average. That single dial, which the widget calls "sharpness," is the whole spectrum between the two extremes from §2.

5 · Putting it together, for every token at once

We now have the full parameter-free mechanism. For a query token i: dot it against every token to get a row of scores, softmax that row into weights, then blend the token vectors by those weights:

for each query position i:
    score[i][j] = dot(x[i], x[j])        for all j      # relevance, T numbers
    alpha[i]    = softmax(score[i])                      # weights, sum to 1
    z[i]        = sum_j alpha[i][j] * x[j]               # the context vector

Do this for one token and you get one context vector. Do it for all T queries and the scores form a full T×T table — every token's row of relevances to every other token — which softmax normalizes row by row into an attention matrix. As matrix operations on X of shape (T, d) it is strikingly compact:

S = X @ X.T                 # (T, d) @ (d, T) -> (T, T)   all pairwise scores
A = softmax(S, axis=-1)     #                 -> (T, T)   each row sums to 1
Z = A @ X                   # (T, T) @ (T, d) -> (T, d)   one context vector per token

Read that as the story of attention in three lines: X@Xᵀ asks "how relevant is every token to every other," softmax turns each query's row into a distribution, and A@X replaces each token with the relevance-weighted blend of all tokens. There is not a single learnable weight here — X appears in all three roles (query, key, and value) as its raw self. Yet the shape of the computation is already the shape of real self-attention. Lesson 05's entire job is to insert three learned projections so those three roles can differ, plus one scaling constant so the softmax does not choke as the vectors get wide.

6 · See one query attend

The widget below uses the running example we will carry through the rest of Part II: the sentence "Your journey starts with one step" — six tokens, each a small 3-dimensional embedding.

Pick a query token and the widget dots it against all six tokens, shows the raw scores, runs softmax into weight bars that sum to 1, and paints the resulting context vector as the weighted colour-blend of the six embeddings. The sharpness slider multiplies the scores before softmax: at low sharpness the weights flatten toward the uniform average (the rigid bag), and turning it up concentrates the mass onto the highest-scoring tokens. Toggle show 6×6 heatmap to see every query's row at once. Watch that "journey" and "starts" — near-identical embeddings — attend to each other strongly, that the weight bars always sum to 1, and that with raw dot-products a token cannot even guarantee it attends most to itself.

Failure modes & checklist

Checkpoint

Where this points next

You have the mechanism, but notice what it cannot do. The weight "journey" places on "starts" is fixed the instant the embeddings are fixed — there is no dial the model could turn to attend differently, because there are no parameters in the attention itself. And a single vector is forced to play all three roles at once: the thing searching, the thing being matched, and the content handed back. A pronoun should search for a noun, advertise itself quite differently, and pass along yet another signal — impossible with one shared vector. Lesson 05, Self-attention with trainable weights, gives the token three learned projections — query, key, value — so those roles can diverge and gradient descent can shape what "relevant" means, and derives the 1/√d_k scale that keeps the softmax from freezing as the vectors widen.

Interview prompts

• Why is a fixed average of the context tokens too rigid for next-token prediction? (§1/§2 — which earlier tokens matter depends on the sentence; an average gives every token equal say and cannot adapt, so long-range and selective dependencies are lost.)
• What is a context vector, and how is it built? (§2 — an enriched embedding of a token, formed as a convex, weight-summed blend of all token vectors, z_i = Σ_j α_ij x_j, with weights summing to 1.)
• Why does a dot product serve as a relevance/similarity score? (§3 — x_i·x_j = ‖x_i‖‖x_j‖cos θ; it is largest for aligned vectors, ~0 for orthogonal, negative for opposed — exactly "how aligned are these two embeddings.")
• Why use softmax instead of dividing scores by their sum? (§4 — softmax exponentiates first, so weights are positive for any real scores and always sum to 1, with stable gradients; a raw divide breaks on negative scores or a zero denominator.)
• What does the "sharpness"/temperature on the scores control? (§4 — scaling scores up before softmax sharpens toward one-hot focus; scaling down flattens toward the uniform average; it dials the peakedness of the attention distribution.)
• Which quantities are trainable in this parameter-free attention? (§5 — none; the raw embedding x plays query, key, and value at once, so the attention pattern is a fixed function of the embeddings — the exact gap lesson 05 fills.)
• Write attention for all tokens as three matrix operations. (§5 — S = X@Xᵀ (pairwise scores), A = softmax(S, axis=-1) (row-normalized weights), Z = A@X (context vectors).)

Companion reads: 05 · Self-attention with trainable weights adds the Q/K/V projections and the √d_k scale to this same six-token example; cs336 · 02 The Transformer counts the params and FLOPs once the weights are in.