LM3-T3 — Queries, Keys, Values and Scaled Dot-Product Attention

Learning objectives

After this lesson, you should be able to:

• Describe the roles of queries, keys and values and the shapes of the three projection matrices.
• Compute attention scores, weights and outputs for a small example by hand.
• Explain the 1/√d_k scaling and what the softmax does to the scores.
• Write scaled dot-product attention in PyTorch and state the shape of every tensor.

The idea: a soft lookup

Compare attention to looking something up in a dictionary, with one difference: the match is not all-or-nothing. A normal lookup compares a query with every key, finds the one that matches, and returns its value. Attention compares a query with every key, scores how well each matches, turns the scores into weights that sum to 1, and returns a weighted average of all the values. The best-matching entries contribute most.

Three roles for each token

Every token is turned into three vectors by three learned matrices:

• The query says what this token is looking for.
• The key says what this token offers as a match.
• The value is the information this token contributes if it is selected.

Using three different projections of the same input lets the model separate the question from the answer. A token can search for a type of word (its query) while being found by other tokens for a different reason (its key), and hand over something else again (its value).

Score, normalise, mix

1. Score. The dot product of the query of token i with the key of token j is large when they point in similar directions.
2. Normalise. A softmax along the row turns the scores of token i into weights between 0 and 1 that sum to 1.
3. Mix. The output for token i is the sum of the value vectors, each multiplied by its weight.

What this gives the model

After attention, the vector at each position contains information from the positions it attended to. In “the robot picks up the box and it falls”, the vector at “it” can pull in information from “box”. Nothing in the layer says that “it” must attend to “box”: the projection matrices are learned so that useful attention patterns reduce the loss.

The layer has a compact interpretation: every token reads from other tokens, in proportions that depend on the content. This is also the layer’s cost: with n tokens there are n × n scores.

The mathematics: scaled dot-product attention

Let X be the input of shape (n, d) (token embeddings plus positions). Three learned matrices project it:

n

number of tokens in the sequence (the rows).

d

input width (the embedding size).

d_k

width of queries and keys (they must match so dot products are defined).

d_v

width of values; the width of the output. Often equal to d_k.

Q, K, V

queries, keys and values: shapes (n, d_k), (n, d_k) and (n, d_v). Row i belongs to token i.

The attention output is

Step by step, with shapes:

• S = Q Kᵀ / √d_k has shape (n, d_k) × (d_k, n) = (n, n). Entry S_ij is the scaled score of query i against key j.
• W = softmax(S) applied to each row: W_ij = exp(S_ij) / Σ_j′ exp(S_ij′). Shape (n, n); each row sums to 1.
• O = W V has shape (n, n) × (n, d_v) = (n, d_v). Row i is Σ_j W_ij V_j.

Why divide by √d_k? If the components of a query and a key are independent with mean 0 and variance 1, their dot product is a sum of d_k products, each with mean 0 and variance 1, so it has variance d_k and standard deviation √d_k. For d_k = 64 the scores typically differ by about 8. A softmax of numbers that far apart is almost one-hot: one weight near 1, the rest near 0, and the gradient through the softmax almost vanishes. Dividing by √d_k restores a standard deviation of 1 whatever the width. Lab LM3-L1 Stage 1 shows this.

Parameters. One attention head has d · d_k + d · d_k + d · d_v weights (plus biases if used). The matrix S has n² entries, so memory and time for the scores grow with the square of the sequence length.

A worked example: three tokens, width two

To keep the numbers small, take Q = K = [[1, 0], [0, 1], [1, 1]] (three tokens, d_k = 2) and V = [[1, 0], [0, 2], [3, 1]]. (In a real layer Q, K and V come from X by the projections; here we start from them.)

Scores. QK^T = [[1, 0, 1], [0, 1, 1], [1, 1, 2]]. Dividing by √2 ≈ 1.414 gives S ≈ [[0.707, 0, 0.707], [0, 0.707, 0.707], [0.707, 0.707, 1.414]].

Weights. Row 1: exp(0.707) = 2.028, exp(0) = 1, exp(0.707) = 2.028; the sum is 5.056, so the weights are [0.401, 0.198, 0.401]. Row 2 is [0.198, 0.401, 0.401]. Row 3: exp(0.707) = 2.028 twice and exp(1.414) = 4.113; the sum is 8.169, so [0.248, 0.248, 0.503]. Each row sums to 1.

Output. Row 1 = 0.401 · [1, 0] + 0.198 · [0, 2] + 0.401 · [3, 1] = [0.401 + 1.203, 0.396 + 0.401] = [1.604, 0.797]. Row 2 = 0.198 · [1, 0] + 0.401 · [0, 2] + 0.401 · [3, 1] = [1.401, 1.203]. Row 3 = 0.248 · [1, 0] + 0.248 · [0, 2] + 0.503 · [3, 1] = [1.759, 1.000].

Token 3, whose query matches all keys, spreads its attention and puts the most on itself; its output is a blend dominated by its own value and token 1’s large value. Every output is a convex combination of the value rows: each component lies between the smallest and the largest value in its column.

The same computation in PyTorch

import torch

Q = torch.tensor([[1., 0.], [0., 1.], [1., 1.]])
K = Q.clone()
V = torch.tensor([[1., 0.], [0., 2.], [3., 1.]])

r3 = lambda t: [[round(v, 3) for v in row] for row in t.tolist()]
dk = K.shape[-1]
scores = Q @ K.T / dk ** 0.5                  # (n, n)
weights = torch.softmax(scores, dim=-1)        # each row sums to 1
out = weights @ V                              # (n, dv)

print(r3(weights))      # [[0.401, 0.198, 0.401], [0.198, 0.401, 0.401], [0.248, 0.248, 0.503]]
print(r3(out))          # [[1.604, 0.797], [1.401, 1.203], [1.759, 1.0]]
print(weights.sum(dim=-1).tolist())            # rows sum to 1 (up to rounding)

ref = torch.nn.functional.scaled_dot_product_attention(Q, K, V)
print(torch.allclose(out, ref, atol=1e-6))     # True: matches PyTorch's built-in version

Note dim=-1 in the softmax: it normalises over the keys (the last axis), one distribution per query. Normalising over the wrong axis (dim=0) is a classic bug that still runs and still gives numbers between 0 and 1.

Engineering implications and limitations

Softmax over the right axis. The weights of one query must sum to 1 over keys. A quick test: weights.sum(-1) should be all ones.

Numerical stability. Softmax implementations subtract the row maximum before exponentiating. Do the same if you write your own; large scores overflow otherwise (Lab LM3-L1 Stage 2).

Use the fused function in practice. torch.nn.functional.scaled_dot_product_attention computes the same result with optimised kernels that avoid storing the full n × n matrix (Lesson LM3-T6, Module 10). Write the long form to understand it; use the fused form to run it.

Inspect the weights. The weight matrix is a useful debugging view: a row that is almost uniform may mean the model has not learned to attend; a row that is one-hot on the first token may indicate an attention sink or a bug. But weights are not a complete explanation of what the model does: the values and later layers matter too.

Quadratic cost. Doubling the sequence length quadruples the scores. This is the reason for context-length limits and for the work on efficient attention.

Dropout. Training often applies dropout to the attention weights (randomly zeroing some) as regularisation; it is turned off at inference. The book includes it in its attention modules.

Self versus cross attention. In self-attention Q, K and V all come from the same sequence, as here. In cross-attention (translation, vision-language models) the queries come from one sequence and the keys and values from another; the formula is unchanged. The VLA track uses this to let action tokens read image features.

Common misconceptions

Engineering takeaways

1. Attention(Q, K, V) = softmax(QKᵀ/√d_k) V: score, normalise per row, mix the values.
2. Shapes: Q and K are (n, d_k), V is (n, d_v), scores and weights are (n, n), the output is (n, d_v).
3. Each output row is a weighted average of the value rows with weights that sum to 1.
4. The 1/√d_k factor keeps score variance at 1 so the softmax does not saturate.
5. The cost in time and memory grows with n², which limits the context length.

References and further reading

See Chapter 3 of the book for simplified attention, self-attention with trainable weights, and the PyTorch documentation page on scaled dot-product attention.

• Vaswani, A. et al. (2017). Attention Is All You Need. arXiv:1706.03762 (Section 3.2, scaled dot-product attention).
• Build a Large Language Model (From Scratch), Chapter 3: Coding Attention Mechanisms (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: the causal mask, which stops a token from reading the tokens that come after it.

Related content

Nothing below is required to finish this lesson.

• Before
  LM3-T2: Positional Information

  Show why attention needs position and compare absolute learned, sinusoidal and relative schemes.

• Lab
  LM3-L1: Compute Attention Step by Step

  Compute scores, softmax weights and context vectors for a short sentence by hand-written code.