S108 · Q/K/V Math — Scaled Dot-Product Attention Derived

The equation that powers ChatGPT.

Module M13: NLP & Transformers · Session 108 of 130 · Track: LLM 🧮

What you'll be able to do after this session

Explain Q/K/V Math to a friend in one minute, out loud, no notes.
Recognise it when you see it in real code / systems / papers.
Complete the hands-on exercise at the end without looking up help.

Prerequisites

If you can't answer these in 30 seconds each, redo the linked session first:

S107 — Attention Intuition — Why RNNs Failed, Why Attention Won

Pre-read (skim before the session)

🎥 Intuition (5–15 min) · Attention in transformers, visually explained — 3Blue1Brown — Rewatch focusing on the Q/K/V decomposition.
🎥 Deep dive (30–60 min) · Let's build GPT — Andrej Karpathy (attention section) — Full attention derivation with code.
🎥 Hands-on demo (10–20 min) · The Math Behind Self-Attention — Umar Jamil — Chalks the equations line by line and codes them.
📖 Canonical article · The Annotated Transformer — Harvard NLP — Line-by-line PyTorch implementation of the 'Attention Is All You Need' paper.
📖 Book chapter / tutorial · Attention Is All You Need — Section 3.2 — The exact equations — read Section 3.2 in detail.
📖 Engineering blog · Peter Bloem — Transformers from scratch — A rigorous, code-heavy derivation with excellent diagrams.

(a) Intuition · 5 min

The one-paragraph mental model. Why this concept exists, what problem it solves, and the analogy that makes it stick.

In one sentence: The equation that powers ChatGPT.

In S107 you saw attention as a metaphor: "each token asks a question, other tokens answer." In this session we make the math literal — write the equation, understand every term, and see why it's the single most reproduced formula in modern ML. The equation is:

Attention(Q, K, V) = softmax(Q · Kᵀ / sqrt(d_k)) · V

That's it. Every LLM is built on this five-symbol expression. But the interpretation of each symbol is the interesting bit, and understanding why it works — why we divide by sqrt(d), why softmax, why three separate projections instead of one — is what separates people who "use transformers" from people who "understand transformers".

The core idea is a soft dictionary lookup. In Python, d["key"] returns exactly one value if the key exactly matches. Attention is the soft version: for a query, we compute how well it matches every key (dot product), normalise those into a probability distribution (softmax), and return a weighted mixture of all the associated values. When one key matches strongly, softmax approaches one-hot and we recover a hard lookup; when many match weakly, we get a blend.

Gotcha to carry forever: d_k in the denominator is the per-head dimension, not the full model dimension d_model. For GPT-2 small with d_model = 768 and n_heads = 12, d_k = 64. If you accidentally scale by sqrt(d_model) in a multi-head setting your softmax is over-flattened and the model learns much slower.

(b) Visual walkthrough · 15 min

Diagrams, worked examples, step-by-step visuals.

The equation, unpacked, dimension by dimension. Given a batch of shape [B, T, d_model], we produce three matrices via learned linear layers W_Q, W_K, W_V with shape [d_model, d_k]:

Q = X · W_Q → [B, T, d_k]
K = X · W_K → [B, T, d_k]
V = X · W_V → [B, T, d_v]

Then the four steps of the equation:

Compute similarities. S = Q · Kᵀ → [B, T, T]. Element S[i, j] measures how much query i matches key j.
Scale. S = S / sqrt(d_k). Prevents softmax saturation.
Softmax. A = softmax(S, dim=-1). Each row is a probability distribution over keys.
Aggregate. Output = A · V → [B, T, d_v].

Worked example — why sqrt(d_k)? Suppose q and k are random vectors of dim d_k, each component iid N(0, 1). Then q · k = sum(q_i k_i) has mean 0 and variance d_k. So for d_k = 64, dot products routinely reach ±8 or more, and after exp() in softmax you get exponentials of huge numbers → the softmax becomes essentially one-hot → gradients through it are ~0 → training stalls. Dividing by sqrt(d_k) = 8 puts dot products back near N(0, 1), softmax stays smooth, gradients flow. That's the "scaled" in "scaled dot-product attention".

Worked example — reading an attention matrix. Suppose S / sqrt(d) is:

[[ 2,  0, -1],
 [ 0,  3,  0],
 [-1,  1,  4]]

After row-wise softmax, approximately:

A = [[0.84, 0.11, 0.04],
     [0.05, 0.90, 0.05],
     [0.006, 0.02, 0.97]]

Token 0 attends mostly to itself; token 1 attends mostly to itself; token 2 attends almost entirely to itself. Bright diagonals = each token looks at itself; bright off-diagonals = cross-token dependencies.

(c) Hands-on · 20 min

Code you type and run. Not pseudo-code — real, runnable, copy-pasteable.

# qkv_from_scratch.py — every step of scaled dot-product attention in NumPy
import numpy as np
 
np.random.seed(42)
B, T, d_model, d_k = 1, 5, 16, 4                 # batch, seq, model dim, per-head dim
X = np.random.randn(B, T, d_model)               # pretend token embeddings
 
# Random projection matrices (in a real model these are learned)
W_Q = np.random.randn(d_model, d_k) * 0.1
W_K = np.random.randn(d_model, d_k) * 0.1
W_V = np.random.randn(d_model, d_k) * 0.1
 
Q = X @ W_Q                                      # [B, T, d_k]
K = X @ W_K
V = X @ W_V
print("shapes -> Q:", Q.shape, " K:", K.shape, " V:", V.shape)
 
# 3. Raw attention scores
scores = Q @ K.transpose(0, 2, 1)                # [B, T, T]
print("raw scores range:", scores.min(), "->", scores.max())
 
# 4. Scale
scores_scaled = scores / np.sqrt(d_k)
 
# 5. Numerically stable softmax
def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)
 
A = softmax(scores_scaled, axis=-1)
print("attention weights row 0:", A[0, 0].round(3), "sum=", A[0, 0].sum())
 
# 6. Weighted values
out = A @ V
print("output shape:", out.shape)                # -> (1, 5, 4)
 
# ---- causal mask ---------------------------------------------------------
mask = np.triu(np.ones((T, T)), k=1).astype(bool)
scores_causal = np.where(mask, -1e9, scores_scaled)
A_c = softmax(scores_causal, axis=-1)
print("causal weights:\n", A_c[0].round(2))

Checklist while it runs:

Every row of A sums to 1 — verify.
Increase d_k=64 and confirm raw scores blow up ±8 range; scaled scores stay ±1-ish.
Skip the scale — softmax becomes near-one-hot; compare A rows.
Print the causal mask matrix — should be strictly lower triangular in A_c.
Replace random X with actual token embeddings from a small transformer and see the attention patterns.

Try this modification: Split d_k=4 into 2 heads of d_head=2. Compute attention independently per head, then concatenate outputs. That's multi-head attention — the topic of S109.

(d) Production reality · 10 min

How this shows up in real systems, gotchas, war stories, common bugs.

The formula stays the same in production; the implementation changes wildly.

PyTorch's F.scaled_dot_product_attention (added in PyTorch 2.0) is the recommended way to compute attention today. It dispatches to Flash Attention v2 on Ampere/Hopper GPUs, to memory-efficient attention on older cards, or to a plain reference kernel on CPU. Using it means you get 2–4× speedup and lower memory for free — never write your own attention loop in prod.

Numerical stability. The softmax(Q Kᵀ / sqrt(d_k)) step is the numerically tricky one. In float16 training, un-scaled dot products routinely overflow. Modern implementations use the online-softmax trick and compute exp(x − max_x) per tile, then combine tiles. This is what Flash Attention does under the hood; you don't need to write it, but understanding why it exists helps you diagnose "loss went to NaN in mixed precision" bugs.

Bugs I've seen in code review:

Forgetting the 1/sqrt(d_k) — model looks fine on synthetic tests, fails to converge on real data.
Using d_model instead of d_k in the scale — attention weights become too flat, model can't focus.
Wrong dim in softmax — softmax over dim=-2 (queries) instead of dim=-1 (keys) is subtly wrong; shapes still match.
Mask applied after softmax instead of before — the mask is now a hard-multiplied zero rather than -inf, and the remaining weights don't renormalise to 1. Attention scores now leak information from the masked positions.
Fp16 overflow in exp() — use bf16 or wrap the softmax in the fp32 autocast section.

Modern variants of the same equation. Grouped-Query Attention shares one K/V head across multiple Q heads (Llama-2 70B). Multi-Query Attention uses a single K/V for all Q heads (PaLM). ALiBi and RoPE modify how positional information enters — but they all use the same softmax(QKᵀ/√d) V core. Understanding this equation cold pays off for the entire rest of the LLM curriculum.

(e) Quiz + exercise · 10 min

Write out the scaled dot-product attention equation from memory.
Why do we divide by sqrt(d_k) and not sqrt(d_model)?
If d_k = 64, roughly what is the standard deviation of raw dot-product scores before scaling?
What is the shape of the attention matrix Q · Kᵀ for input [B, T, d]?
What subtle bug appears if you apply the causal mask after softmax instead of before?

Stretch (connects to S107 — Attention Intuition): Explain in your own words why attention is called a "soft dictionary lookup" — what recovers the "hard" lookup as a limit?

Explain-out-loud test

If you can't teach these 3 things to a friend without notes, redo the session:

What is Q/K/V Math? (one sentence, no jargon)
When would you use it? (one concrete example)
When would you NOT use it? (one anti-pattern)

What comes next

Next session (S109): Multi-Head Attention — Parallel Views
Previous (S107): Attention Intuition — Why RNNs Failed, Why Attention Won
Hub: The 6-Month Learning Plan

Part of a 130-session evergreen learning series. Session structure: (a) intuition · (b) visual walkthrough · (c) hands-on · (d) production reality · (e) quiz + exercise. Duration: 60 minutes.