Attention

Attention, query, key, value, weighted sum

Recall from the last lesson: every token is now a vector, and we built Z = token embedding + positional vector. The transformer's first move is to apply three separate linear maps to Z, producing Q, K, V. This lesson is what happens next, what those three matrices actually do.

Attention is a way to look up information by content, not position. You hold a query (“what am I looking for?”) and you have a collection of keys (labels for what's available) and values (the actual content behind each label). The query is compared against every key, the comparison scores are turned into weights, and the answer is a weighted sum of the values.

This single mechanism, and the variants built from it, replaces the recurrence in RNNs with a parallel, content-addressable read. The rest of the transformer is plumbing around this idea.

The "aha": resolving it

Before the math, see why content-addressing matters. In the sentence below, what does it refer to, the animal or the street? A human knows instantly; a bag-of-positions model cannot. Attention lets the token it send out a query that scores highly against animal, so it pulls meaning from the right place. Click the it row.

The it row dumps most of its weight on animal, a little on street (the plausible distractor), and almost nothing elsewhere. That is content-addressing: the connection is by meaning, not by how many words apart the two tokens sit.

A library analogy

Imagine a library where every book has two parts: a catalogue card describing it (the key) and the book itself (the value). You walk in with a search query.

1. The librarian compares your query against every catalogue card in parallel.
2. Cards that match well get a high score; mismatched cards get a low score.
3. The librarian doesn't hand you one book, she hands you a blend of all the books, with each book contributing in proportion to how well its card matched.

That blend is the output of attention. You never asked for book #3; you described what you wanted, and the relevant content surfaced.

Scaled dot-product attention

For a single query vector q and a set of key-value pairs (kⱼ, vⱼ) with $\mathbf{k}_j, \mathbf{q} \in \mathbb{R}^{d_k}$:

In matrix form, when Q is a stack of queries (one row per query token):

Three operations, in order:

1. Score, QKᵀ produces an (n_q × n_k) matrix of raw similarity scores. Each entry is a single dot product.
2. Normalize, divide by √d_k, then softmax along each query's row to get weights that are positive and sum to 1.
3. Aggregate, multiply by V to get a (n_q × d_v) output: each query's answer is a convex combination of the value vectors.

Worked example

Take d_k = 4 (the width of queries and keys) and d_v = 2 (the width of values), three keys, three values, one query.

q  = [1, 0, 1, 0]                     ← d_k = 4
 
k_1 = [1, 0, 1, 0]   v_1 = [10, 0]   (close match to q)
k_2 = [0, 1, 0, 1]   v_2 = [ 0, 5]   (orthogonal, no match)
k_3 = [1, 0, 0, 0]   v_3 = [ 2, 2]   (partial match)
                     └─ d_v = 2

Step 1, raw scores. q · k_j:

q · k_1 = 1·1 + 0·0 + 1·1 + 0·0 = 2
q · k_2 = 1·0 + 0·1 + 1·0 + 0·1 = 0
q · k_3 = 1·1 + 0·0 + 1·0 + 0·0 = 1

Step 2, scale by √d_k = 2.

scaled = [2/2, 0/2, 1/2] = [1.0, 0.0, 0.5]

Step 3, softmax. exp(1.0)=2.72, exp(0.0)=1.00, exp(0.5)=1.65, sum ≈ 5.37:

α = [2.72/5.37, 1.00/5.37, 1.65/5.37] ≈ [0.506, 0.186, 0.308]

The weights are positive, sum to 1, and concentrate mass on k_1 (the best match) without ignoring the others.

Step 4, weighted sum of values.

output = 0.506 · [10, 0] + 0.186 · [0, 5] + 0.308 · [2, 2]
       ≈ [5.06 + 0 + 0.62,  0 + 0.93 + 0.62]
       ≈ [5.68, 1.55]

The output borrows mostly from v_1, a touch from v_3, and almost nothing from v_2. No discrete choice was made, the answer is differentiable, so gradients flow back into Q, K, V projections during training.

Now stop reading and operate it. The widget below starts at exactly this query; drag the query vector and watch every step recompute. Then use the d_k slider underneath to feel why the √d_k divisor exists.

q · k_1 = 0    q · k_2 = 2    q · k_3 = 0

Why divide by √d_k?

If q and k have entries with mean 0 and variance 1, the dot product q · k = Σᵢ qᵢ kᵢ has variance d_k. As d_k grows, scores grow with it. Large positive scores push softmax toward a one-hot distribution; one weight saturates at ≈ 1 and the rest collapse to ≈ 0, which kills gradients on the suppressed terms.

Dividing by √d_k rescales the variance back to ≈ 1, keeping softmax in a region where it's sharp but not saturated, confident enough to focus, soft enough to learn.

Intuition cheat-sheet

• High q · k → large weight on that v.
• Softmax → weights are non-negative and sum to 1 (a convex combination).
• No argmax → fully differentiable; the network can learn what to attend to.
• √d_k scale → keeps logits from saturating as d_k grows.

Next: when Q, K, V are all derived from the same sequence, every token attends to every other token in that sequence. That's self-attention.

Practice it yourself

Build attention one step at a time — score, weight, aggregate. Real Python, hidden tests, no setup.