The mathematics of attention

What you will be able to do

Write the attention of a transformer as a few matrix products. Then see why each part of the formula is there: the scaling, the mask, the heads and the cost. Every step comes with a tiny example that you can check by hand, a check in PyTorch, and a picture.

NotebookBuild attention from the formulas, in a notebookColab · Kaggle · GitHub

The notebook writes attention() and multi-head attention from the formulas and checks them against PyTorch. It measures the variance of q · k for several key sizes, and times the score matrix on your machine. It needs no data download and no training. About 30 minutes.

Open in ColabOpen in KaggleView on GitHub

Colab and Kaggle open your own copy. They need a free Google or Kaggle account. Stuck? See the solution notebook.

Before you startWhat to know and to have readyAbout 55 minutes

What you need to know

• Queries, keys, values, scores, softmax and the weighted sum, from Attention step by step. This page starts from its worked example.
• Multi-head attention, from Transformer blocks, and the two masks, from Position and masking.
• The dot product, from Vector operations in the mathematics course, and its section Matrix multiplication: many dot products at once.
• Variance and standard deviation, from Describe a dataset, and the derivative as a slope, from Functions and change.

What you need on your computer

• Nothing for the page. The notebook opens in Colab or Kaggle with PyTorch already installed.

Tested with Python 3.12 and PyTorch 2.14.1 on the CPU (macOS).

This lesson is optional: the certificate does not require it.

The idea, in one paragraph

You saw the idea in Attention step by step. This page shows the mathematics. Each word asks a question (its query) and compares it with what every word offers (its key). The comparison is a dot product, divided by a square root. Softmax turns the scores into weights that add up to 1. The word’s new vector is the weighted average of what the words pass on (their values). In the worked example, “damaged on arrival” with two numbers per vector, “damaged” read mostly “arrival”, with the weight 0.728.

The question this page answers

The attention lesson gave you a recipe. This page asks why the recipe has exactly these parts. Why a dot product? Why divide by √d_k, and not by d_k or by nothing? Why does a mask add minus infinity, instead of setting a weight to 0? Why several small heads instead of one big one? And why does attention get expensive so quickly for long texts? The whole recipe fits in one line, from the 2017 paper Attention Is All You Need:

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

Here Q, K and V are the matrices of queries, keys and values, one row per token. d_k is the number of numbers in one key (the key size). M is the mask, a matrix of zeros and minus infinities. The attention lesson wrote d for the key size. From here on, d is the number of numbers per token and d_k the number per key, because the two differ when there are several heads. The rest of this page takes the formula apart, step by step, with the shape written next to every matrix.

Step 1: from one query to all queries

The formula. Start one step earlier than the attention lesson, at the token vectors themselves. Put them in a matrix X with n rows (one per token) and d columns (the numbers per token). Three learned weight matrices turn X into the queries, keys and values:

$$Q = X W_Q, \quad K = X W_K, \quad V = X W_V \qquad (n, d) \times (d, d_k) \to (n, d_k)$$

In words: every row of Q is one token’s query, made from that token’s vector by the same weight matrix. One matrix product makes the queries of all tokens at once.

A tiny example you can check. Take the three words again, now with d = 4 numbers each, and three made-up 4 × 2 weight matrices. They are chosen so that Q, K and V are exactly the queries, keys and values of the attention lesson:

from_x.py

import math
import torch

X = torch.tensor([[1.0, 0.0, 0.0, 1.0],      # damaged
                  [0.0, 1.0, 0.0, 0.0],      # on
                  [0.0, 0.0, 1.0, 0.0]])     # arrival
W_Q = torch.tensor([[0.5, 0.5], [0.0, 0.0], [1.0, 0.0], [0.0, 1.0]])
W_K = torch.tensor([[1.0, 0.0], [0.0, 0.5], [0.0, 2.0], [0.0, 0.0]])
W_V = torch.tensor([[1.0, 0.0], [0.0, 0.0], [0.0, 1.0], [1.0, 0.0]])
Q, K, V = X @ W_Q, X @ W_K, X @ W_V          # (3, 4) @ (4, 2) -> (3, 2)
print("X", tuple(X.shape), "-> Q, K, V", tuple(Q.shape))
print(Q)
print(K)
print(V)

The output:

X (3, 4) -> Q, K, V (3, 2)
tensor([[0.5000, 1.5000],
        [0.0000, 0.0000],
        [1.0000, 0.0000]])
tensor([[1.0000, 0.0000],
        [0.0000, 0.5000],
        [0.0000, 2.0000]])
tensor([[2., 0.],
        [0., 0.],
        [0., 1.]])

Check row 1 by hand. The vector of “damaged” is (1, 0, 0, 1). So it picks rows 1 and 4 of W_Q and adds them: (0.5, 0.5) + (0, 1) = (0.5, 1.5). That is the query of “damaged” in the attention lesson.

All scores at once. The score of query i with key j is a dot product. In matrix form, that is row i of Q times column j of Kᵀ, the keys turned on their side. So one product gives every score:

$$S = Q K^\top, \qquad S_{ij} = q_i \cdot k_j = \sum_{t=1}^{d_k} q_{it}\, k_{jt} \qquad (n, d_k) \times (d_k, n) \to (n, n)$$

scores.py

S = Q @ K.T                                  # (3, 2) @ (2, 3) -> (3, 3)
print(S)
print((S / math.sqrt(2)).round(decimals=3))

The output:

tensor([[0.5000, 0.7500, 3.0000],
        [0.0000, 0.0000, 0.0000],
        [1.0000, 0.0000, 0.0000]])
tensor([[0.3540, 0.5300, 2.1210],
        [0.0000, 0.0000, 0.0000],
        [0.7070, 0.0000, 0.0000]])

Row 1 holds the three scores of “damaged”: 0.5, 0.75 and 3.0. The attention lesson computed the same numbers one at a time. Why a dot product? It is large when the query and the key point the same way, and it costs only d_k multiplications. A matrix product computes all n × n of them in one fast operation.

Step through the whole formula below. Each matrix shows its shape, and the row of “damaged” is marked. Make your prediction first.

Explore

Step through attention in matrix form

The input X and the weight matrices

Step

1 of 7

Shape of the result

(3, 4)

Weights of "damaged"

not yet

X (3, 4)

	1	2	3	4
damaged	1	0	0	1
on	0	1	0	0
arrival	0	0	1	0

W_Q (4, 2)

	1	2
0.5	0.5
0	0
1	0
0	1

W_K (4, 2)

	1	2
1	0
0	0.5
0	2
0	0

W_V (4, 2)

	1	2
1	0
0	0
0	1
1	0

Each row of X is one token with d = 4 numbers: n = 3 tokens, so X is (3, 4). The three weight matrices are (4, 2): they turn 4 numbers into d_k = 2. In a real model they are learned; here they are made up.

The whole formula on 3 tokens with d = 4 and d_k = 2. The matrices are made up so that Q, K and V equal the worked example of the attention lesson; every other number is computed from them and checked against PyTorch.

Text description of the explorer

The explorer shows the attention formula as seven steps, with a matrix for each input and result, its shape, and the row of “damaged” marked. Step 1: X is (3, 4), one row per token: damaged (1, 0, 0, 1), on (0, 1, 0, 0), arrival (0, 0, 1, 0); W_Q, W_K and W_V are (4, 2). Step 2: Q = X W_Q, K = X W_K and V = X W_V are (3, 2): Q has the rows (0.5, 1.5), (0, 0) and (1, 0); K has (1, 0), (0, 0.5) and (0, 2); V has (2, 0), (0, 0) and (0, 1). Step 3: S = Q Kᵀ is (3, 3), with the rows (0.5, 0.75, 3), (0, 0, 0) and (1, 0, 0). Step 4: S divided by √2 has the rows (0.354, 0.53, 2.121), (0, 0, 0) and (0.707, 0, 0). Step 5: the mask M is added. With no mask, M is all zeros. With the causal mask, M has −∞ above the diagonal. Step 6: softmax, row by row. With no mask, the weights are (0.124, 0.148, 0.728), (0.333, 0.333, 0.333) and (0.503, 0.248, 0.248). With the causal mask, they are (1, 0, 0), (0.5, 0.5, 0) and (0.503, 0.248, 0.248). Every row adds up to 1. Step 7: O = A V is (3, 2). With no mask, its rows are (0.248, 0.728), (0.667, 0.333) and (1.007, 0.248). With the causal mask, they are (2, 0), (1, 0) and (1.007, 0.248). Controls: a button for each step, Back and Next, and a choice of no mask or the causal mask.

The table collects the shapes. Read it as a checklist: every product must match its inner sizes, and the result keeps the outer sizes.

Step	Formula	Example	In general
Input	X	(3, 4)	(n, d)
Queries, keys, values	X W_Q, X W_K, X W_V	(3, 4) × (4, 2) → (3, 2)	(n, d_k)
Scores	Q Kᵀ	(3, 2) × (2, 3) → (3, 3)	(n, n)
Scaled and masked	S / √d_k + M	(3, 3)	(n, n)
Weights	softmax, row by row	(3, 3), rows add up to 1	(n, n)
Output	A V	(3, 3) × (3, 2) → (3, 2)	(n, d_v)

Here d_v is the size of a value. In a transformer, d_v = d_k, so the output has the shape of Q.

Step 2: why divide by √d_k

The attention lesson said that large scores make softmax extreme. Now you can measure how large they get, and derive the exact factor.

The question. Before training, the numbers in queries and keys are roughly random, around 0, with a spread of about 1. How large is a typical score q · k then, and how does it grow with the key size d_k?

The derivation, in three lines. Assume each number q_t and k_t is independent, with mean 0 and variance 1. A product of two such independent numbers also has mean 0, and its variance is E[q_t²] E[k_t²] = 1 × 1 = 1. The score adds d_k of these products, and for independent terms the variances add:

$$\operatorname{Var}(q \cdot k) = \operatorname{Var}\Big(\sum_{t=1}^{d_k} q_t k_t\Big) = \sum_{t=1}^{d_k} \operatorname{Var}(q_t k_t) = \sum_{t=1}^{d_k} \mathbb{E}[q_t^2]\,\mathbb{E}[k_t^2] = d_k$$

So the standard deviation of a score is √d_k. Dividing every score by √d_k gives the variance d_k / (√d_k)² = 1, for any key size. That is the reason for the square root. Dividing by d_k would shrink the scores too much (variance 1 / d_k), and no division lets them grow with d_k.

A check in code. Draw 20,000 random queries and keys for five key sizes, and measure:

variance.py

torch.manual_seed(0)
for d_k in (4, 16, 64, 256, 1024):
    q = torch.randn(20000, d_k, dtype=torch.float64)       # 20,000 random queries ...
    k = torch.randn(20000, d_k, dtype=torch.float64)       # ... and keys, every number N(0, 1)
    dots = (q * k).sum(1)                                   # 20,000 dot products q . k
    Qs = torch.randn(2000, d_k, dtype=torch.float64)        # 2,000 queries, 32 keys each
    Ks = torch.randn(2000, 32, d_k, dtype=torch.float64)
    raw = torch.einsum("nd,nkd->nk", Qs, Ks)                # (2000, 32) scores
    top_raw = torch.softmax(raw, 1).max(1).values.mean()
    top_scaled = torch.softmax(raw / math.sqrt(d_k), 1).max(1).values.mean()
    print(f"d_k {d_k:5}  Var(q.k) {dots.var():8.2f}  after / sqrt(d_k) {(dots / math.sqrt(d_k)).var():.3f}"
          f"  largest weight: unscaled {top_raw:.3f}, scaled {top_scaled:.3f}")

The output:

d_k     4  Var(q.k)     4.03  after / sqrt(d_k) 1.006  largest weight: unscaled 0.335, scaled 0.152
d_k    16  Var(q.k)    15.84  after / sqrt(d_k) 0.990  largest weight: unscaled 0.637, scaled 0.165
d_k    64  Var(q.k)    63.69  after / sqrt(d_k) 0.995  largest weight: unscaled 0.822, scaled 0.163
d_k   256  Var(q.k)   256.00  after / sqrt(d_k) 1.000  largest weight: unscaled 0.903, scaled 0.164
d_k  1024  Var(q.k)  1024.99  after / sqrt(d_k) 1.001  largest weight: unscaled 0.953, scaled 0.164

torch.randn draws numbers from the normal distribution with mean 0 and variance 1, written N(0, 1). einsum computes the 32 dot products of each query with its own keys. The measured variance is 4.03, 15.84, 63.69, 256.00 and 1,024.99: d_k, as derived. After the division it is 1 every time. The chart shows the same numbers on logarithmic axes, where “variance = d_k” is a straight line.

• Var(q · k), unscaled
• Var(q · k / √d_k), scaled

The measured variance of q · k for random queries and keys (20,000 pairs per key size). Without scaling it equals d_k; divided by √d_k it is 1 for every key size.

Text description of the chart

A line chart with logarithmic axes. The x axis is the key size d_k: 4, 16, 64, 256 and 1,024. The y axis is the variance of the score. The unscaled line rises in a straight line: 4.03, 15.84, 63.69, 256.00 and 1,024.99, so the variance equals d_k. The scaled line, q · k divided by √d_k, stays flat at 1: 1.006, 0.990, 0.995, 1.000 and 1.001.

What large scores do to softmax. The code above also fed each query’s 32 scores to softmax. Without scaling, the largest weight grows from 0.335 at d_k = 4 to 0.953 at d_k = 1,024. At d_k = 1,024, 73% of the rows give more than 0.99 to a single key. Softmax has become almost a hard choice of one key: it is saturated. With scaling, the largest weight stays near 0.16 for every key size. Move the slider and compare the two rows of bars.

Explore

Change the key size and watch softmax

Var(q · k)

4.03 (std 2.01)

Largest weight, unscaled → scaled

0.335 → 0.152

Rows with a weight above 0.99

0.1% → 0.0%

Softmax gradient size

0.292 → 0.234

At d_k = 4, the scores q · k have variance 4.03, so a typical score lies within about ±2.01 of 0. Without scaling, softmax still spreads its weight (largest 0.335 on average). Divided by √4 = 2, the variance is 1.01 and the largest weight 0.152, about the same as for every d_k.

Without scaling: softmax(q · k) largest 0.164

1

With scaling: softmax(q · k / √4) largest 0.084

1

One random query and 32 random keys at d_k = 4, every number drawn from N(0, 1). The keys are sorted by score. The readout above gives the averages over 2,000 such queries.

One random query and 32 random keys per key size, with every number drawn from N(0, 1). The readout gives averages over 2,000 such queries, from the same run as the code above; the gradient size comes from the gradient code below.

Text description of the explorer

The explorer has one slider for the key size d_k: 4, 16, 64, 256 or 1,024. It shows two rows of 32 bars: the softmax weights of one random query over 32 random keys, without scaling and with scaling, sorted from the largest. At d_k = 4, both rows are spread out: the largest unscaled weight is 0.164 and the largest scaled weight 0.084. At d_k = 64, the unscaled row has one bar at 0.797 and one at 0.200, and the rest are near 0; the scaled row still slopes gently from 0.149. At d_k = 1,024, the unscaled row is a single bar at 1.000; the scaled row starts at 0.243. The readout shows the measured averages over 2,000 queries. The variance of q · k goes from 4.03 up to 1,024.99. The largest weight, without and with scaling, is 0.335 and 0.152 at d_k = 4, and 0.953 and 0.164 at d_k = 1,024. The share of rows with a weight above 0.99 is 0.1% and 0% at d_k = 4, and 73.2% and 0% at d_k = 1,024. The softmax gradient size is 0.292 and 0.234 at d_k = 4, and 0.057 and 0.240 at d_k = 1,024.

Why a saturated softmax hurts learning. Training changes a score by following the gradient of the loss, as in Gradient descent. The gradient passes through softmax, and the derivative of the softmax weight p_i with respect to the score s_j is:

$$\frac{\partial p_i}{\partial s_j} = p_i\,(\delta_{ij} - p_j), \qquad \delta_{ij} = 1 \text{ if } i = j, \text{ else } 0$$

In words: each entry is a weight times one minus a weight, or minus two weights multiplied. When one weight is near 1 and all others near 0, every entry is near 0. For example, a row with the weights 0.99 and 0.01 gives p(1 − p) = 0.99 × 0.01 = 0.0099. Almost no gradient flows back to the scores, so the queries and keys learn very slowly. The code measures the size of this derivative, averaged over 2,000 rows of 32 keys:

gradient.py

def softmax_gradient_size(scores):
    """Mean size (Frobenius norm) of d softmax / d scores = diag(p) - p p^T, over the rows."""
    p = torch.softmax(scores, dim=-1)
    J = torch.diag_embed(p) - p.unsqueeze(-1) * p.unsqueeze(-2)
    return J.norm(dim=(-2, -1)).mean().item()

torch.manual_seed(1)
for d_k in (4, 16, 64, 256, 1024):
    raw = torch.einsum("nd,nkd->nk", torch.randn(2000, d_k, dtype=torch.float64),
                       torch.randn(2000, 32, d_k, dtype=torch.float64))
    print(f"d_k {d_k:5}  gradient size: unscaled {softmax_gradient_size(raw):.4f},"
          f" scaled {softmax_gradient_size(raw / math.sqrt(d_k)):.4f}")

The output:

d_k     4  gradient size: unscaled 0.2918, scaled 0.2337
d_k    16  gradient size: unscaled 0.2875, scaled 0.2388
d_k    64  gradient size: unscaled 0.1948, scaled 0.2395
d_k   256  gradient size: unscaled 0.1129, scaled 0.2405
d_k  1024  gradient size: unscaled 0.0570, scaled 0.2403

J holds all the derivatives above for one row. Its Frobenius norm is the square root of the sum of its squared entries: one number for “how much gradient passes”. With scaling, the size stays near 0.24 for every key size. Without scaling, it falls to 0.057 at d_k = 1,024, about 4 times smaller. At small d_k, the unscaled scores are not yet extreme, so the division matters little there. The tiny transformer of this course uses d_k = 16 and divides by 4.

Step 3: softmax row by row, then the weighted sum

The formula. Write Z = S / √d_k + M for the scaled and masked scores. Softmax works on each row of Z alone, so every row of the weight matrix A adds up to 1. Then each output row is a weighted sum of the value rows:

$$A_{ij} = \frac{e^{Z_{ij}}}{\sum_{l=1}^{n} e^{Z_{il}}}, \qquad o_i = \sum_{j=1}^{n} A_{ij}\, v_j \quad\Longleftrightarrow\quad O = A V$$

In words: row i of A says how much token i reads each token j. Row i of O mixes the values with those weights. Here o_i and v_j are rows of O and V. In code, the whole formula takes a few lines. This is the readable function that the rest of the lesson uses:

attention.py

def attention(Q, K, V, mask=None):
    """softmax(Q K^T / sqrt(d_k) + mask) V.  Q: (..., n, d_k), K: (..., m, d_k), V: (..., m, d_v)."""
    d_k = Q.shape[-1]
    scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)       # (..., n, m)
    if mask is not None:
        scores = scores + mask                               # -inf where a query may not look
    weights = torch.softmax(scores, dim=-1)                  # every row sums to 1
    return weights @ V, weights                              # (..., n, d_v), (..., n, m)

out, weights = attention(Q, K, V)
print(weights.round(decimals=3))
print(out.round(decimals=3))

The output:

tensor([[0.1240, 0.1480, 0.7280],
        [0.3330, 0.3330, 0.3330],
        [0.5030, 0.2480, 0.2480]])
tensor([[0.2480, 0.7280],
        [0.6670, 0.3330],
        [1.0070, 0.2480]])

The ... in the shapes means “any dimensions before these”: a batch, heads, or both. K.transpose(-2, -1) swaps the last two dimensions, so it is Kᵀ for every item of a batch. dim=-1 makes softmax work along each row. The keys can come from a different sequence of m tokens, which the next lesson uses for translation. Row 1 is the hand calculation of the attention lesson: weights 0.124, 0.148 and 0.728, output (0.248, 0.728).

Step 4: masks are added before softmax

The formula. A mask is a matrix M of the same shape as the scores. It holds 0 where a query may read a key, and −∞ where it may not. It is added to the scaled scores, before softmax:

$$A = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right), \qquad M_{ij} = \begin{cases} 0 & \text{token } i \text{ may read token } j \\ -\infty & \text{otherwise} \end{cases}$$

In words: a hidden key gets the score −∞, and e^−∞ = 0, so its weight is exactly 0. The visible keys share the whole weight of 1 among themselves. Two masks are common. The causal mask has −∞ above the diagonal, so a token never reads the tokens after it. The padding mask has −∞ in the columns of padding tokens, so no token reads the blank fill. Position and masking showed both on the tiny transformer; here they are on the three words:

masks.py

causal = torch.triu(torch.full((3, 3), float("-inf")), diagonal=1)   # -inf above the diagonal
print(causal)
out, weights = attention(Q, K, V, mask=causal)
print(weights.round(decimals=3))
padding = torch.tensor([0.0, 0.0, float("-inf")])     # pretend "arrival" is padding: hide column 3
out, weights = attention(Q, K, V, mask=padding)       # the same row is added to every row of scores
print(weights.round(decimals=3))

The output:

tensor([[0., -inf, -inf],
        [0., 0., -inf],
        [0., 0., 0.]])
tensor([[1.0000, 0.0000, 0.0000],
        [0.5000, 0.5000, 0.0000],
        [0.5030, 0.2480, 0.2480]])
tensor([[0.4560, 0.5440, 0.0000],
        [0.5000, 0.5000, 0.0000],
        [0.6700, 0.3300, 0.0000]])

torch.triu(…, diagonal=1) keeps the part of a matrix above the diagonal and sets the rest to 0. The padding mask is one row; PyTorch adds it to every row of the scores (broadcasting, from Tensors and shapes). The figure shows both masks and the weights they give.

Masks are matrices of 0 and −∞ that are added to the scaled scores. After softmax, every −∞ cell has the weight 0 (orange), and the visible weights of each row add up to 1.

Text description of the chart

Two rows of two grids, each 3 by 3, for the tokens damaged, on and arrival. Rows are the token that reads; columns the token it reads. First row: the causal mask M has 0 on and below the diagonal and −∞ above it. After adding it and softmax, the weights are: damaged 1, 0, 0; on 0.5, 0.5, 0; arrival 0.503, 0.248, 0.248. Second row: the padding mask, with “arrival” treated as padding, has −∞ in the whole third column. The weights are: damaged 0.456, 0.544, 0; on 0.5, 0.5, 0; arrival 0.67, 0.33, 0. Every hidden cell has the weight exactly 0, and the visible weights of each row add up to 1.

Why −∞, and not 0? Adding 0 to a score changes nothing. And a score of 0 is not “hidden”: softmax turns it into e^0 = 1 before it divides by the total. Try both ideas on the row of “damaged”:

why_inf.py

scaled = torch.tensor([0.354, 0.530, 2.121])          # row "damaged", already divided by sqrt(2)
print("add -inf:      ", torch.softmax(scaled + torch.tensor([0.0, float("-inf"), float("-inf")]), 0).round(decimals=3))
print("set scores to 0:", torch.softmax(torch.tensor([0.354, 0.0, 0.0]), 0).round(decimals=3))

The output:

add -inf:       tensor([1., 0., 0.])
set scores to 0: tensor([0.4160, 0.2920, 0.2920])

With −∞, “damaged” reads only itself. With the scores set to 0, it still gives 0.584 of its weight to the two “hidden” words. A third idea, setting weights to 0 after softmax, also fails: the row would no longer add up to 1, unless you divide again. Adding −∞ before softmax does both jobs in one step. It also keeps the matrix shape, so one fast operation handles every row.

Step 5: multi-head attention

The formula. One head gives each token one row of weights, so it can follow only one kind of relation at a time. Multi-head attention runs h smaller heads side by side. Each head i has its own projections, of size d × d_k with d_k = d / h. The outputs of the heads are joined (concatenated), and one more matrix W_O mixes them:

$$\text{head}_i = \text{Attention}(X W_Q^{(i)},\, X W_K^{(i)},\, X W_V^{(i)}), \qquad \text{MultiHead}(X) = [\text{head}_1, \ldots, \text{head}_h]\, W_O$$

In words: h independent attentions, each on a slice of d_k numbers, then one linear layer over the joined result. In practice, the h query matrices of the heads sit side by side in one d × d matrix W_Q. So one product makes all the queries, and a reshape splits the d columns into h heads of d_k. The diagram shows every shape, with the tiny transformer’s numbers for ticket T-65052 (30 tokens, d = 64, h = 4).

Multi-head attention: split into h heads, attend in each head, join the heads, mix them with W_O. The shape in and out is (n, d).

Text description of the diagram

Data flows from left to right. X has the shape (n, d); for ticket T-65052 in the tiny transformer, (30, 64). It is multiplied by three weight matrices W_Q, W_K and W_V, each (d, d), which gives Q, K and V, each (n, d). Their d columns are split into h heads of d_k = d / h numbers, so each of them becomes (h, n, d_k); for the ticket (4, 30, 16). Each head, from head 1 to head h, computes its own attention: softmax of Q times K transposed, divided by the square root of d_k, plus the mask M, times V. Each head has an (n, n) weight matrix and an output of (n, d_k). The h outputs are concatenated: h times d_k columns are d columns, so the result is (n, d). A last matrix W_O, (d, d), mixes the heads. The output is (n, d), the same shape as X.

The parameter count. Count the learned numbers. W_Q, W_K, W_V and W_O are each d × d, and each has a bias of d numbers:

$$4 \cdot d^2 + 4 \cdot d \quad\Rightarrow\quad d = 64:\; 16{,}384 + 256 = 16{,}640, \qquad d = 512:\; 1{,}050{,}624$$

The number of heads does not appear in the formula. More heads cut the same d columns into more, smaller slices. That is why Transformer blocks found 16,640 parameters for 1, 4 and 8 heads. The 2017 paper used d = 512 and h = 8, so each of its attention layers has 1,050,624 parameters when every matrix has a bias, as in PyTorch. (The paper’s formulas show no biases; without them, the count is 4d² = 1,048,576.) Here is the multi-head layer, built on attention():

multihead.py

class MultiHead(torch.nn.Module):
    """h heads of size d_k = d / h. Q = X W_Q + b_Q (and K, V alike); heads side by side; output = concat W_O + b_O."""

    def __init__(self, d, h):
        super().__init__()
        self.d, self.h, self.d_k = d, h, d // h
        self.W_Q, self.W_K, self.W_V, self.W_O = (torch.nn.Linear(d, d) for _ in range(4))

    def split(self, X):                                      # (B, n, d) -> (B, h, n, d_k)
        B, n, _ = X.shape
        return X.view(B, n, self.h, self.d_k).transpose(1, 2)

    def forward(self, Xq, Xkv, mask=None):
        Q, K, V = self.split(self.W_Q(Xq)), self.split(self.W_K(Xkv)), self.split(self.W_V(Xkv))
        O, weights = attention(Q, K, V, mask)                # (B, h, n, d_k), (B, h, n, m)
        B, _, n, _ = O.shape
        O = O.transpose(1, 2).reshape(B, n, self.d)          # concatenate the heads
        return self.W_O(O), weights

torch.manual_seed(0)
Xb = torch.randn(3, 7, 64)                                   # 3 sequences, 7 tokens, d = 64
q, k, v = (torch.randn(3, 7, 16) for _ in range(3))          # for the one-head check below
mine = MultiHead(64, 4)
print("parameters:", sum(p.numel() for p in mine.parameters()),
      "| d = 512:", sum(p.numel() for p in MultiHead(512, 8).parameters()))
out, weights = mine(Xb, Xb)
print("out", tuple(out.shape), "weights", tuple(weights.shape))

The output:

parameters: 16640 | d = 512: 1050624
out (3, 7, 64) weights (3, 4, 7, 7)

view reads the same numbers with a new shape, (B, n, h, d_k), and transpose(1, 2) moves the heads before the tokens. Then attention() works on every head of every sequence at once, because of the ... in its shapes. After attention, the reverse moves join the heads: (B, h, n, d_k) back to (B, n, d). The weights have one n × n square per head: (3, 4, 7, 7).

Why several heads? Each head has its own W_Q and W_K, so each one learns its own idea of “a good match”. One head can follow the payment words while another follows the request. Their weights are separate softmax rows, so one relation does not take weight away from another. The price is a smaller d_k per head: each head compares tokens on 16 numbers instead of 64. The real ticket at the end of this page shows four heads that learned four different patterns.

Step 6: the cost of n × n

The formula. Count the work for one head. Q Kᵀ has n × n entries, each a dot product of d_k numbers, and A V again needs n × n × d_k multiplications. The score matrix itself must be stored, n × n numbers:

$$\text{time} \propto n^2 \cdot d_k, \qquad \text{memory for the scores} = n^2 \times 4 \text{ bytes (float32)}$$

In words: double the tokens, and the work and memory of the scores grow four times. Everything else in a block, the projections and the feed-forward layer, grows only in proportion to n.

memory.py

for n in (128, 1024, 8192):
    cells = n * n                                            # one score per (query, key) pair
    print(f"n {n:5}: {cells:>11,} scores = {cells * 4 / 1e6:7.2f} MB in float32 (one head)")

The output:

n   128:      16,384 scores =    0.07 MB in float32 (one head)
n  1024:   1,048,576 scores =    4.19 MB in float32 (one head)
n  8192:  67,108,864 scores =  268.44 MB in float32 (one head)

At 8,192 tokens, one head’s score matrix is 268 MB, for one sequence in one layer. The softmax output is a second matrix of the same size. By comparison, Q, K and V together are only 6.3 MB. The charts show the measured memory and time on a shared CPU, one head with d_k = 64, as the median of 5 passes. Both axes are logarithmic, so n² growth is a straight line: 4 times more for every doubling of n.

• Score matrix
• Q, K and V together

Memory for one head in float32. The n × n score matrix grows with n²; the queries, keys and values grow only with n.

Text description of the chart

A line chart with logarithmic axes. The x axis is the number of tokens n, from 128 to 8,192. The score matrix needs 0.07 MB at 128 tokens, 1.05 MB at 512, 4.19 MB at 1,024, 16.78 MB at 2,048, 67.11 MB at 4,096 and 268.44 MB at 8,192: a steep straight line, 4 times more for every doubling. Q, K and V together grow only twice per doubling, from 0.098 MB to 6.3 MB.

• Readable attention()
• PyTorch fused function
• n² growth from n = 1,024

Time for one head, batch 1, d_k = 64, on a shared CPU (2 threads, median of 5 passes): indicative only. The dashed line is pure n² growth from the readable function's time at 1,024 tokens.

Text description of the chart

A line chart with logarithmic axes. The x axis is the number of tokens n, from 128 to 8,192; the y axis is the time in milliseconds. The readable attention() takes 0.035 ms at 128 tokens, 0.225 ms at 512, 0.849 ms at 1,024, 3.54 ms at 2,048, 14.8 ms at 4,096 and 60.0 ms at 8,192. PyTorch’s fused scaled_dot_product_attention takes 0.027, 0.144, 0.437, 1.70, 6.86 and 27.0 ms. A dashed line shows pure n² growth from 0.849 ms at 1,024 tokens; above 1,024 tokens, both measured lines follow its slope. Below about 512 tokens, fixed overheads matter more, and the lines are flatter.

n	Scores	Score matrix	Readable attention()	PyTorch fused
128	16,384	0.07 MB	0.035 ms	0.027 ms
1,024	1,048,576	4.19 MB	0.849 ms	0.437 ms
8,192	67,108,864	268.44 MB	60.0 ms	27.0 ms

From 1,024 to 8,192 tokens (8 times more), the time grows 71 times, close to the 64 times of n². PyTorch’s own scaled_dot_product_attention is about twice as fast, because it fuses the steps into one optimised operation. It still does the n² work. These timings come from one shared machine; on yours, the numbers will differ, but the slope will not.

This is why long-context models need tricks. Some kernels never store the whole score matrix at once. Some layers attend only to a window of nearby tokens. And during generation, the KV cache of Inference costs keeps the keys and values of earlier tokens, so they are not computed again.

Step 7: the check against PyTorch

A formula on paper is only useful if the code computes it. So the code below compares the readable functions of this page with PyTorch’s own. attention() meets F.scaled_dot_product_attention, and MultiHead meets nn.MultiheadAttention, with the same weights copied into it:

check.py

F = torch.nn.functional
ref = torch.nn.MultiheadAttention(64, 4, batch_first=True)
with torch.no_grad():                                        # copy our weights into PyTorch's layer
    ref.in_proj_weight.copy_(torch.cat([mine.W_Q.weight, mine.W_K.weight, mine.W_V.weight]))
    ref.in_proj_bias.copy_(torch.cat([mine.W_Q.bias, mine.W_K.bias, mine.W_V.bias]))
    ref.out_proj.weight.copy_(mine.W_O.weight)
    ref.out_proj.bias.copy_(mine.W_O.bias)
    causal = torch.triu(torch.full((7, 7), float("-inf")), diagonal=1)
    future = torch.triu(torch.ones(7, 7, dtype=torch.bool), diagonal=1)        # PyTorch's layer wants True/False
    is_pad = torch.arange(7)[None, :] >= torch.tensor([7, 5, 3])[:, None]     # lengths 7, 5 and 3
    padding = torch.zeros(3, 1, 1, 7).masked_fill(is_pad[:, None, None, :], float("-inf"))
    one_head = [attention(q, k, v)[0] - F.scaled_dot_product_attention(q, k, v),
                attention(q, k, v, causal)[0] - F.scaled_dot_product_attention(q, k, v, is_causal=True),
                attention(q, k, v, padding[:, 0])[0] - F.scaled_dot_product_attention(q, k, v, attn_mask=padding[:, 0])]
    print("one head vs scaled_dot_product_attention:", f"{max(d.abs().max() for d in one_head):.1e}")
    biggest = 0.0
    for my_mask, theirs in [(None, {}), (causal, {"attn_mask": future}), (padding, {"key_padding_mask": is_pad}),
                            (causal + padding, {"attn_mask": future, "key_padding_mask": is_pad})]:
        o1, w1 = mine(Xb, Xb, my_mask)
        o2, w2 = ref(Xb, Xb, Xb, average_attn_weights=False, **theirs)
        biggest = max(biggest, (o1 - o2).abs().max().item(), (w1 - w2).abs().max().item())
    print("multi-head vs nn.MultiheadAttention:", f"{biggest:.1e}")

The output:

one head vs scaled_dot_product_attention: 2.4e-07
multi-head vs nn.MultiheadAttention: 1.8e-07

nn.MultiheadAttention stores W_Q, W_K and W_V stacked in one matrix, in_proj_weight, of shape (3d, d), so the code joins ours with torch.cat. The check covers no mask, the causal mask, padding (sequence lengths 7, 5 and 3) and both masks together, for outputs and weights. The largest difference is 2.4 × 10⁻⁷. That is rounding: float32 keeps about 7 significant digits, and the two codes add the same numbers in a different order. In float64, the same check gives 4.4 × 10⁻¹⁶. So the formulas on this page are what PyTorch computes.

Step 8: a real ticket

The last check uses the course’s trained tiny transformer and validation ticket T-65052: “Not a delivery issue and I’m keeping the leaf blower. My card was charged twice for the same order. Can you remove the duplicate charge?” Grace’s team is payment, and the model says payment. The code copies block 2’s trained weights into MultiHead and recomputes the attention weights:

ticket.py

from ticketnet import TinyTransformer, Vocab, load_split

DATA = data_dir()
train_rows, valid_rows = (load_split(DATA / f"{s}.csv") for s in ("train", "valid"))
vocab = Vocab([r["text"] for r in train_rows])
tf = TinyTransformer(len(vocab))
tf.load_state_dict(torch.load(WEIGHTS, weights_only=True))   # the course's trained model
tf.eval()
row = next(r for r in valid_rows if r["ticket_id"] == "T-65052")
ids = torch.tensor([vocab.encode(row["text"])])
tokens = [vocab.itos[i] for i in ids[0].tolist()]
with torch.no_grad():
    h = tf.embed(ids) + tf.position(torch.arange(ids.shape[1]))
    h = tf.blocks.layers[0](h)                               # block 1 as a whole
    block = tf.blocks.layers[1]
    x = block.norm1(h)                                       # block 2 normalises before attention
    pt = block.self_attn                                     # PyTorch's layer, trained
    ours = MultiHead(64, 4)
    for lin, W, b in zip((ours.W_Q, ours.W_K, ours.W_V), pt.in_proj_weight.chunk(3), pt.in_proj_bias.chunk(3)):
        lin.weight.copy_(W)
        lin.bias.copy_(b)
    ours.W_O.weight.copy_(pt.out_proj.weight)
    ours.W_O.bias.copy_(pt.out_proj.bias)
    _, w_ours = ours(x, x)
    _, w_pt = pt(x, x, x, average_attn_weights=False)
print("weights", tuple(w_ours.shape), "max difference", f"{(w_ours - w_pt).abs().max():.1e}")
q = tokens.index("not")
for head in range(4):
    top = w_ours[0, head, q].topk(2)
    print(f"head {head + 1}: 'not' reads", [(tokens[i], round(v.item(), 2)) for v, i in zip(top.values, top.indices)])

The output:

weights (1, 4, 30, 30) max difference 2.1e-07
head 1: 'not' reads [('twice', 0.5), ("'", 0.07)]
head 2: 'not' reads [('charge', 0.55), ('the', 0.09)]
head 3: 'not' reads [('remove', 0.5), ('order', 0.1)]
head 4: 'not' reads [('remove', 0.48), ('charge', 0.15)]

data_dir() and WEIGHTS point to the course’s ticket data and the saved model; Transformer blocks trains and saves it. chunk(3) cuts PyTorch’s stacked matrix back into W_Q, W_K and W_V. Our formula gives the trained model’s weights to 2.1 × 10⁻⁷; over both blocks, outputs and weights, the largest difference is 9.5 × 10⁻⁷. The heatmaps show all four heads of block 2.

The real attention weights of the tiny transformer, block 2, ticket T-65052: one 30 × 30 heatmap per head. Rows are the reading tokens, columns the tokens they read; a stronger colour is a larger weight (square-root scale, relative to the head's largest). The row of 'not' is outlined.

Text description of the chart

Four heatmaps, one for each head of block 2, each 30 rows by 30 columns for the 30 tokens of ticket T-65052. Each row adds up to 1. The heads show vertical stripes: most tokens read the same few words. Head 1 has a strong stripe at “twice”, which receives 0.355 of the weight on average over all 30 reading tokens. Head 2 has its stripe at “charge” (0.208 on average). Head 3 spreads its weight more: “remove” 0.118 and “you” 0.116. Head 4 has its stripe at “remove” (0.274). The outlined row of “not” reads twice 0.50 in head 1, charge 0.55 in head 2, remove 0.50 in head 3 and remove 0.48 in head 4. No head sends “not” to “delivery”, the word right after it.

The heads learned different patterns: head 1 collects “twice”, head 2 “charge”, head 4 “remove”, and head 3 spreads out. That is the “different relations” of Step 5, in a real model. For every ticket, block and head, open the explorer in Attention step by step. Keep its warning in mind: weights are not explanations. They show where each token took its information from, not why the model chose payment.

What it means for Larkfield, and what came next

For Larkfield’s router, the mathematics explains three practical facts. A ticket of 120 tokens gives each head a 120 × 120 score matrix, 14,400 numbers: small, so attention is cheap for short tickets. The tiny transformer’s d_k = 16 is small enough that scaling matters less than in large models, but it costs nothing to keep. And the four heads of each block are cheaper than they look: they cost the same 16,640 parameters as one head of 64 numbers.

The formula has stayed almost the same since 2017. What changed is mostly around it. Newer models build the positions into the dot product, and several query heads share fewer key and value heads. Fused kernels avoid storing the n × n matrix. The next lesson rebuilds the original transformer from the pieces of this page.

Guided practice: one block’s shapes for a long ticket

1. The tiny transformer reads one ticket of n = 120 tokens, with d = 64 and h = 4. Write d_k.
2. Write the shape of Q after W_Q, and after split (with the batch dimension 1).
3. Write the shape of the attention weights of one block, and the number of scores in it.
4. Compute the memory of these scores in float32, in kB.
5. Write the shape of the output after the heads are joined and W_O is applied.
6. Write the parameter count of this attention layer, from the formula.

Show the expected answers

1. d_k = d / h = 64 / 4 = 16.
2. Q = X W_Q is (1, 120, 64). After split: (1, 4, 120, 16), four heads of 16 numbers.
3. One 120 × 120 square per head: (1, 4, 120, 120), so 4 × 14,400 = 57,600 scores.
4. 57,600 × 4 bytes = 230,400 bytes = 230.4 kB per block. The softmax output is a second matrix of this size.
5. Joined: (1, 120, 4 × 16) = (1, 120, 64); after W_O: (1, 120, 64), the shape of the input.
6. 4 × 64² + 4 × 64 = 16,640, the same for any number of heads.

The notebook runs MultiHead(64, 4) on a (1, 120, 64) input and checks every shape.

Your turn: larger embeddings

The division by √d_k assumes that the numbers in queries and keys have a variance of about 1. Multiply every token vector of the worked example by 4. Before you compute, predict: what happens to the scores, and to the weights of “damaged”?

Show the answer

Q and K both become 4 times larger, so every score becomes 4 × 4 = 16 times larger. Row “damaged” becomes 0.354 × 16 = 5.66, 8.48 and 33.9 after the division. Softmax saturates:

your_turn.py

out, weights = attention(4 * X @ W_Q, 4 * X @ W_K, 4 * X @ W_V)      # every embedding 4 times larger
print(weights.round(decimals=3))

The output:

tensor([[0.0000, 0.0000, 1.0000],
        [0.3330, 0.3330, 0.3330],
        [1.0000, 0.0000, 0.0000]])

“damaged” now gives all its weight to “arrival”, and “arrival” all its weight to “damaged”. The weights are almost one-hot, so the gradients through softmax are near 0. Dividing by √d_k corrects for the number of terms in a dot product, not for their size. That is one reason why the tiny transformer normalizes each token (norm1) before attention, as in Transformer blocks. Row “on” has the query (0, 0), which stays 0 when it is multiplied by 4, so its weights stay equal.

Common mistake

A True/False mask that means the opposite

Symptom: Tomás replaces the −∞ mask with a True/False mask, because PyTorch accepts both. He uses True for “hidden”, as nn.MultiheadAttention does, and passes it to F.scaled_dot_product_attention. There is no error, but the outputs are wrong:

pitfall.py

future = torch.triu(torch.ones(3, 3, dtype=torch.bool), diagonal=1)   # True above the diagonal
print(future)
wrong = F.scaled_dot_product_attention(Q, K, V, attn_mask=future)     # True means "may read" here!
right = F.scaled_dot_product_attention(Q, K, V, attn_mask=~future)
print(wrong)
print(right)

The output:

tensor([[False,  True,  True],
        [False, False,  True],
        [False, False, False]])
tensor([[0.0000, 0.8308],
        [0.0000, 1.0000],
        [0.0000, 0.0000]])
tensor([[2.0000, 0.0000],
        [1.0000, 0.0000],
        [1.0070, 0.2483]])

Cause: the two functions read a True/False mask in opposite ways. In nn.MultiheadAttention, True means “may not read”. In F.scaled_dot_product_attention, True means “may read”. So the wrong mask lets every token read only the future. The last row has no visible key at all, and comes back as zeros instead of an error.

Fix: think in the mathematics: the mask is added to the scores, with 0 for “may read” and −∞ for “may not read”. A float mask of 0 and −∞ means the same in both functions. If you use a True/False mask, read the documentation of that function. Then test the mask on a tiny example whose answer you know, as check.py does.

Knowledge check

Check your understanding

Why does attention divide the scores by the square root of d_k, and not by d_k or by nothing?

Show the answer

Take a query and a key whose numbers are independent, with variance 1. Their dot product adds d_k products of variance 1, so its variance is d_k and its spread √d_k. Dividing by √d_k brings the variance back to 1 for every key size. Without the division, the scores grow with d_k: softmax saturates (largest weight 0.953 on average at d_k = 1,024), and its gradient becomes small (0.057 instead of 0.240). Dividing by d_k would shrink the variance to 1 / d_k, and every row of weights would become almost uniform.

Check your understanding

Predict: a row of scaled scores is 2.0, 1.0 and 0.5, and the causal mask hides the last key. What are the weights? What would they be if you set the hidden score to 0 instead?

Show the answer

With −∞, the third weight is exactly 0, and softmax shares the weight between 2.0 and 1.0. e^2 = 7.389 and e^1 = 2.718, total 10.107, so the weights are 0.731, 0.269 and 0. With the hidden score set to 0, it still counts. e^0 = 1, the total is 11.107, and the weights are 0.665, 0.245 and 0.090. The “hidden” key still gets 9% of the weight. That is why a mask adds −∞ before softmax.

Check your understanding

Amira wants to run the tiny transformer's attention on support threads of 8,192 tokens instead of 120. What happens to the size of one head's score matrix, and why would more heads not make it cheaper?

Show the answer

The score matrix grows with n²: from 120 × 120 = 14,400 numbers to 8,192 × 8,192 = 67,108,864. That is 268 MB in float32 for one head, one sequence and one layer, and the softmax output doubles it. The time grows about as fast (60 ms per head on the shared CPU of this page). More heads do not add parameters (4d² + 4d for any h). But each head has its own n × n matrix, so the memory grows with h too, and each head works on fewer numbers (d_k = d / h). For long texts she needs a model built for them: kernels that do not store the full matrix, attention windows, and the KV cache during generation.

Recap

• Matrix form: Q = X W_Q, K = X W_K, V = X W_V, (n, d) → (n, d_k); the scores S = Q Kᵀ are (n, n); A = softmax(S / √d_k + M), row by row; the output is A V, (n, d_k).
• Scaling: for random unit-variance numbers, Var(q · k) = d_k (measured 4.03 … 1,024.99). Dividing by √d_k brings it back to 1. Without it, softmax saturates (largest weight 0.335 → 0.953) and its gradient shrinks towards 0.
• Masks: a matrix of 0 and −∞ added before softmax. A −∞ score gives a weight of exactly 0, and the visible weights still add up to 1. A score of 0 would still count.
• Multi-head: h heads of d_k = d / h, joined and mixed by W_O; 4d² + 4d parameters for any h (16,640 at d = 64, 1,050,624 at d = 512). Different heads learn different relations, as the four heads of T-65052 show.
• Cost: O(n² · d_k) time and n² memory per head, so O(n² · d) for all h heads: 268 MB and 60 ms for one head at 8,192 tokens on a shared CPU. Long contexts need other tricks.
• Check: the readable code equals PyTorch to 2.4 × 10⁻⁷ in float32, with and without masks.

Next, the original transformer of 2017, rebuilt from these pieces and traced step by step.

Next: The original transformer