Q/K/V attention: from projections to causal output

Compute complete attention: W_Q, W_K, W_V, scores, √d_head, causal mask, softmax, value mixture, multi-head concatenation, and output projection.

Applied AI · advanced · Session 12

Session contract

  • Distinguish query, key, and value.
  • Run a complete numerical trace.
  • Connect prefill, decode, and KV cache.

Mechanism map

  Maya             book             She
 x_Maya           x_book           x_She        every token projects
    │                │                │         ITS own q, k, v
    ▼                ▼                ▼         (shared W_Q, W_K, W_V)
 k_Maya=[1,1]     k_book=[0,2]     q=[2,1]
 v_Maya=[1,0]     v_book=[0,1]        │
    │                │                │
    └────────────────┴───────┬────────┘
                             ▼
        q(She)·k → raw scores: 3 (Maya) and 2 (book)
                             │
                    ┌─────────────────┐  d_head = 4; vectors
                    │ ÷ √d_head: √4=2 │  shown in 2 dims
                    └────────┬────────┘  (truncated illustration)
                             ▼   1.5  and  1.0
                    ┌─────────────────┐
                    │   causal mask   │  future positions ← −∞
                    └────────┬────────┘
                             ▼
                softmax → A = [0.62, 0.38]
                             ▼
        O = 0.62·v_Maya + 0.38·v_book = [0.62, 0.38]
                             ▼
              Concat(head₁…head_h) ──▶ W_O
          (during decode: K/V read from the cache)

1. Three learned projections

Each representation x produces q=xW_Q, k=xW_K, and v=xW_V. Query expresses what the position seeks; key describes how it can be found; value carries information to mix.

Q=XW_Q, K=XW_K, V=XW_V

Check — Three learned projections

q, k and v all come from the same x via W_Q, W_K, W_V. If you force W_Q = W_K, what happens to a position’s score with itself, and why does that impoverish the head?

2. Query-key compatibility

Dot product q·k measures alignment. In “Maya put down the book… She picked it up”, one head can learn that the query for “She” aligns with the key for “Maya”.

Check — Query-key compatibility

With q = [2,1] we get q·k_Maya = 3 and q·k_book = 2. Construct a key that would make “She” ignore “Maya” entirely, and justify it geometrically.

3. Scaling

As d_head grows, dot products can become large and saturate softmax. Dividing by √d_head keeps a more stable scale.

S=QKᵀ/√d_head

Check — Scaling

Here d_head = 4, so we divide by 2: scores 3 and 2 become 1.5 and 1.0. Recompute the softmax WITHOUT the division and quantify how much the first token’s weight rises.

4. Causal mask and softmax

Before softmax, future positions receive −∞. Softmax turns each permitted row into positive weights summing to 1.

A=softmax(S+causal mask)

Check — Causal mask and softmax

At the position of “Maya”, first in the context, which keys does the mask allow? Give the resulting softmax row and explain why it depends on no learned weight.

5. Weighted value mixture

Output is AV: weights choose how much of each value passes. Scores are not themselves the retrieved content.

O=AV

Check — Weighted value mixture

With A = [0.62, 0.38], v_Maya = [1,0] and v_book = [0,1], compute O. Why is O equal to neither stored value, and what does that say about “retrieving” a token?

6. Multi-head, projection, and cache

Multiple heads compute different relations, their outputs are concatenated then projected by W_O. During decoding, past K/V are cached; the new query reads that cache without recomputing the prefix.

MHA(X)=Concat(head₁…head_h)W_O

Check — Multi-head, projection, and cache

During decode the new token appends 1 key and 1 value per head. For 32 heads and 2,000 tokens already cached, state precisely what is recomputed, what is merely re-read, and where the time goes.

Worked case — data

Reduced trace: q=[2,1], keys k_Maya=[1,1], k_book=[0,2]. The real head has d_head=4; our vectors show only 2 coordinates (a truncated illustration), but the scaling keeps the true dimension. Raw scores: 3 and 2; division by √4=2 gives 1.5 and 1. Softmax ≈ [0.62,0.38]. Output mixes 62% of v_Maya and 38% of v_book. For an earlier position, the causal mask would remove every future key.

Worked case — full trace

q = [2,1]   k_Maya = [1,1]   k_book = [0,2]   v_Maya = [1,0]   v_book = [0,1]

dot products : q·k_Maya = 2·1 + 1·1 = 3      q·k_book = 2·0 + 1·2 = 2
scaling      : d_head = 4 (real head size; vectors shown truncated to
               2 dims), √4 = 2  →  S = [1.5, 1.0]

softmax: e^1.5 = 4.482   e^1.0 = 2.718   Σ = 7.200
  A = [0.622, 0.378]       (sum = 1.000 ✅)
O = A·V = 0.622·[1,0] + 0.378·[0,1] = [0.622, 0.378]   ✅

WITHOUT dividing by √d_head: softmax(3, 2) = [0.731, 0.269]
  ❌ a markedly sharper distribution for the very same vector geometry

CAUSAL MASK, position of “Maya” (first in the context)
  (illustration: q_Maya is given the same numeric values [2,1])
  allowed scores = [q_Maya·k_Maya] = [3]; k_book ← −∞ → A = [1.000]
  ✅ O = v_Maya = [1,0]
  mask forgotten → A = [0.622, 0.378]: the position reads a FUTURE token  ❌

SHAPE CHECK: q(1×2) · Kᵀ(2×2) → S(1×2); A(1×2) · V(2×2) → O(1×2).
Multi-head: h heads × d_head = 4 → Concat(1×4h) then W_O(4h × d_model).

Prefill versus decode with a KV cache

Criterion Prefill (n tokens at once) Decode (1 token, KV cache)
Q, K, V computed n queries, n keys, n values 1 query, 1 key, 1 value appended
Score matrix n × n, triangular under the mask 1 × (n+1), a single row
Limiting factor Compute (dense matmuls) Memory bandwidth (re-reading the cache)
Role of the causal mask Essential: future triangle set to −∞ before softmax Implicit: the cache holds only the past

Causal lab

Predict → change one variable → run → explain the delta

/interactives/curriculum/qkv-attention.html?lang=en

Common error 1

“The highest attention score names the word the model retrieved.”

A score of 3 names an address, not content. The output is O = A·V = [0.622, 0.378], which is neither v_Maya = [1,0] nor v_book = [0,1]: the model reads a blend, never a token.

Common error 2

“Dividing by √d_head is a cosmetic numerical trick.”

It changes the effective distribution: the same vectors give [0.62, 0.38] with the division and [0.73, 0.27] without. At large d_head, softmax saturates, gradients vanish, and the head stops learning.

Validity boundary

Attention maps alone do not prove a causal explanation of the model’s overall behavior.

Evidence status

Established mechanisms; numerical simplifications are pedagogical.

Sources to inspect

  • Vaswani et al., “Attention Is All You Need”, NeurIPS (2017).
  • Course source packet supplied by the owner; named-product details remain source-reported until primary verification.

Transfer challenge

Grow the same sentence from 2 to 3 keys: add k_cat = [1,0] and v_cat = [0.5, 0.5] to the context.

  1. Before computing, predict whether Maya’s weight rises or falls.
  2. Recompute scores (÷ √4), softmax, and O.
  3. Compare with the 2-key case: what does the shift say about “attention diluting”?

Synthesis

  • Three learned projections
  • Query-key compatibility
  • Scaling
  • Causal mask and softmax
  • Weighted value mixture
  • Multi-head, projection, and cache

Exit ticket

Mechanism · trace · observation · boundary · evidence · next experiment

Instructor notes: Frame the problem before naming the mechanism. Collect an initial prediction and retain it for the exit ticket.

Instructor notes: These outcomes are observable: trace, calculation, comparison. A recited definition closes none of them.

Instructor notes: Connect every step to the next with a causal verb. Flag any merely decorative arrow.

Instructor notes: Act the three roles out physically: three learners hold q, k and v of the same position and announce what their vector is for. Once embodied, the roles stop being confused.

Instructor notes: Answer: a position’s score with itself becomes q·q = ‖q‖² ≥ 0, systematically favorable: the head stares at itself and loses the freedom to look for what does not resemble it (compatibility turns symmetric). Expected wrong answer: “nothing changes, they are two learned matrices”.

Instructor notes: Have them compute q·k_Maya and q·k_book before revealing 3 and 2, then ask which token will dominate. The prior vote makes the softmax legible.

Instructor notes: Answer: a key strongly anti-aligned with q = [2,1], e.g. k_Maya = [−4,−2]: q·k = −10 and the weight becomes negligible after softmax. Nuance to surface: an orthogonal key (q·k = 0) does NOT zero the weight — it merely stops favoring it.

Instructor notes: Show [0.62, 0.38] and [0.73, 0.27] side by side without labels and have them guess. Asking “which one divided?” anchors √d_head better than a statistical justification.

Instructor notes: Answer: softmax(3, 2) = [0.731, 0.269] versus [0.622, 0.378] with the division: the first weight rises by +0.11. Expected wrong answer: thinking the token order changes — it does not; only the sharpness of the distribution does.

Instructor notes: Draw the n × n matrix and have the room hatch the upper triangle. Then ask what value goes there: “0” is the frequent wrong answer — it is −∞ BEFORE softmax.

Instructor notes: Answer: only k_Maya is allowed → softmax row [1.000], independent of any learned weight — a single non-(−∞) entry always normalizes to 1. Expected wrong answer: [0.62, 0.38], i.e. the UNMASKED row, the one that reads the future.

Instructor notes: Compute O on the board, then ask aloud: “which word did the model retrieve?” Let someone answer “Maya”, then confront it with the vector [0.622, 0.378]. That moment carries the session.

Instructor notes: Answer: O = [0.622, 0.378] — a combination that is neither v_Maya nor v_book: “retrieval” is always a weighted blend, never a token copy. Expected wrong answer: announcing v_Maya because its weight dominates. Point back to common error 1.

Instructor notes: Have them estimate KV cache size for 32 heads, d_head = 128 and 2,000 tokens before any talk of optimization. The resulting number alone justifies MQA, GQA and cache quantization.

Instructor notes: Answer: recomputed = q, k, v of the single new token, plus one 1×2001 score row per head; merely re-read = the 2,000 k/v pairs × 32 heads from the cache; the time goes to memory bandwidth, not compute. Expected wrong answer: “the whole prefix is recomputed at every token”.

Instructor notes: Hide the final result. Elicit sign, shape, and order of magnitude before every operation.

Instructor notes: Walk line by line. Locate an inconsistency at the first faulty step, not only on the final line.

Instructor notes: Have learners fill the final row before revealing it: that trade-off is what decides in production.

Instructor notes: Retain initial and final values. Do not allow simultaneous changes that make the delta impossible to attribute.

Instructor notes: Have the group produce the smallest counterexample before giving the correction.

Instructor notes: Have the group produce the smallest counterexample before giving the correction.

Instructor notes: The boundary is not a footnote: it defines the cases where the mechanism no longer suffices.

Instructor notes: Separate verifiable mechanism, reported implementation choice, and experimental result. Evidence precision must match claim precision.

Instructor notes: Attribute each product choice to the supplied packet and retain reported status until an independent primary source confirms it.

Instructor notes: Answers: scaled scores [1.5, 1.0, 1.0] → A ≈ [0.452, 0.274, 0.274]: Maya’s weight drops from 0.622 to 0.452 although none of Maya’s vectors changed — every new key takes its share of the softmax budget (sum = 1). O ≈ [0.59, 0.41]. Misconception to harvest: “adding a token only affects that token”. Six to eight minutes, pairs.

Instructor notes: Rebuild the chain without looking at the slides. Reopen only the first break.

Instructor notes: Six lines maximum. Compare with the initial prediction and name what actually changed. Cite one retained trace that lets a peer verify the conclusion.