Module 2 — Query, key, value: attention step by step

Module 1 promised that attention is a weighted average of values, with weights that come from a match between a query and a set of keys. This module makes that promise precise, computes it by hand on three tokens, and lands on the first block of the Transformer we build across the course.

The three roles, in plain English

For every token in the sequence, attention builds three projections.

• The query $q$ carries the question "what am I looking for?". It represents the current token as a searcher.
• The key $k$ carries the offer "here is what I contain". It represents each token as an index.
• The value $v$ carries the content that will actually be read once the match is found. Nothing forbids $k$ and $v$ being different: modules 6 and 7 exploit that separation.

The three vectors come from three different learned linear projections of the same input:

$Q = X W_Q, \quad K = X W_K, \quad V = X W_V.$

The distinction between roles is entirely produced by these three matrices. At initialisation they are random and the three roles are indistinguishable; after training, each has learned its specialty.

The formula, then the arithmetic

The 2017 paper gives the operation in one line:

$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V.$

Five steps hide inside it. Read them in order:

1. Compute the similarity between each query and each key: $QK^\top$, a matrix of dot products.
2. Scale by $\sqrt{d_k}$ so that variance stays under control as $d_k$ grows.
3. Apply softmax along each row so that the coefficients become non-negative weights that sum to one.
4. Multiply those weights by $V$ to obtain a weighted average of values.
5. Return one output vector per query.

The output has the same length as the input: attention transforms a sequence into another sequence of the same length, where each position now depends on all the others.

By hand, on three tokens

Take three tokens with $d_k = 2$. Suppose the queries, keys and values are already computed:

$$Q = \begin{pmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{pmatrix}, \quad K = \begin{pmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{pmatrix}, \quad V = \begin{pmatrix} 10 & 0 \\ 0 & 10 \\ 5 & 5 \end{pmatrix}.$$

Step 1 — similarities. $QK^\top$ gives every dot product between a query and a key:

$$QK^\top = \begin{pmatrix} 1 & 0 & 1 \\ 0 & 1 & 1 \\ 1 & 1 & 2 \end{pmatrix}.$$

Step 2 — scaling. Divide by $\sqrt{d_k} = \sqrt{2} \approx 1.414$:

$$\frac{QK^\top}{\sqrt{2}} \approx \begin{pmatrix} 0.71 & 0 & 0.71 \\ 0 & 0.71 & 0.71 \\ 0.71 & 0.71 & 1.41 \end{pmatrix}.$$

Step 3 — softmax on each row. Exponentiating and normalising row by row gives approximately:

$$A \approx \begin{pmatrix} 0.38 & 0.19 & 0.42 \\ 0.19 & 0.38 & 0.42 \\ 0.26 & 0.26 & 0.48 \end{pmatrix}.$$

Each row sums to one. Read them as "attention distributions": for token 1, most weight sits on itself (0.38) and on token 3 (0.42); token 3 pays some attention to everyone but leans on itself (0.48).

Step 4 — weighted average. $A V$ gives the output. For row 1, $0.38 \times (10, 0) + 0.19 \times (0, 10) + 0.42 \times (5, 5) \approx (5.9, 4.0)$.

Step 5 — output. The three output rows are the new representations of the three tokens. Each mixes information from all three, in proportion to how much they matched the current query.

That is the entire operation. Everything else in this course is variations on this pattern.

Why the scaling by $\sqrt{d_k}$

The scaling is not decoration. Without it, the mechanism silently breaks when $d_k$ grows.

Take two random vectors of dimension $d_k$ with unit-variance components. Their dot product has variance $d_k$: the numbers grow with dimension. Feeding a large number into softmax pushes the distribution towards a one-hot, which is close to argmax. Gradients through such a softmax are near zero everywhere except on the winning coordinate, and training stalls.

Dividing by $\sqrt{d_k}$ brings the variance back to one, whatever the head size. The softmax then stays soft enough for gradients to flow. This is one of those single-line design choices that seem cosmetic and are, in fact, load-bearing.

In PyTorch, once and for all

import torch
import torch.nn.functional as F

def attention(q, k, v):
    d_k = q.size(-1)
    scores = q @ k.transpose(-2, -1) / d_k ** 0.5
    weights = F.softmax(scores, dim=-1)
    return weights @ v, weights

# Toy run with three tokens, d_k = d_v = 2
q = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
k = q.clone()
v = torch.tensor([[10.0, 0.0], [0.0, 10.0], [5.0, 5.0]])

out, w = attention(q, k, v)
print(out)
print(w.sum(dim=-1))   # each row sums to 1

Save this function under attention.py: it is the first block of the Transformer we assemble at module 10. Later modules add batching, multiple heads and masking on top, without touching this core.

The three matrices $W_Q$, $W_K$, $W_V$ are learned

It is tempting, on toy examples, to set $Q = K = V = X$. That works arithmetically and hides the whole point. The reason attention works is that the three projections specialise during training — a query learns to look for what a key learns to expose, and both differ from the value. Skip the projections and you have removed the learning surface.

Complexity — a taste of module 9

Computing $QK^\top$ produces a matrix of size $n \times n$ for a sequence of length $n$. That is quadratic in the sequence length, in both compute and memory. For $n = 100$ it is fine, for $n = 10\,000$ it is a hundred million entries per head per layer. Module 9 shows how the community pushes that ceiling; for now, just note that the price of "each token looks at every other token" is $n^2$.

Pick the same head dimension as course 11's hidden state

Choosing $d_k$ around 64 matches the effective hidden size of a strong LSTM. It is a sensible default that keeps arithmetic readable in printouts and stays within a laptop's memory in module 10.

In summary

• Attention builds three learned projections of the input — query, key, value — then reads a weighted average of the values, weighted by the softmax of the scaled dot product between queries and keys.
• The formula $\text{softmax}(QK^\top/\sqrt{d_k}) V$ hides five steps: similarity, scaling, softmax, weighted sum, per-query output.
• The scaling by $\sqrt{d_k}$ is essential: without it, softmax saturates as dimension grows and gradients vanish.
• The cost is $\mathcal{O}(n^2)$ in the sequence length, which module 9 will revisit; the code fits in five lines of PyTorch.

Next module: multi-head attention, which runs several of these operations in parallel and lets each head specialise on a different pattern.