Attention

Attention is the part of a block where positions exchange information. This page goes through attention_forward line by line, then attention_backward. Every equation links to the code that computes it, and every Illustrative Example is computed from that code's mirror.

Why attention?

A linear layer treats each row on its own: row t of its output depends only on row t of its input. A model that predicts the next token also needs earlier tokens to affect the current one. Attention does this by building each output row from a weighted mix of rows at that position and earlier ones. The weights are computed from the input, so they differ from one sequence to the next.

In general, these three names are explained with a question for each one:

Query: what am I looking for?
Key: what do I contain?
Value: what do I hand over?

Those questions are a way to remember the roles. They are not something this project computes. In the code, a query, a key and a value are three linear projections of one input. A query times a key gives a score, and the scores decide how the values are mixed.

The example on this page. The sentence is "The cat sat on the mat", with one row per word and hand-picked numbers: T = 6, d = 4 and H = 2, so d_h = 2. Every number in this example is labelled Illustrative Example. The real model works differently: it reads sub-word tokens, not words, and it has H = 2 heads of width d_h = 16, inside rows of width d = 32. Nothing written on this page says what the trained model attends to. The only real numbers are the exported weights of one window, under Actual Model Output, when an export exists.

Attention lab

Pick a token and a head, then step through the computation. The code under the matrix marks the lines of attention_forward for the current step. Each step links to its section below.

Data sourceIllustrative ExampleActual Model Output

Illustrative Example shows hand-picked numbers. Actual Model Output shows exported weights from one real window. The two are never shown together.

Token

Head

Step 1 of 5: Scores

Cell (i, j) is the dot product of row i's query with row j's key. Python computes this and the next step on one line. Read this step

Scores, head 0Illustrative ExampleHand-picked input. These are not learned values, and they say nothing about the trained model.

Low -1.75High 1.41

Data table

row \ reads	The	cat	sat	on	the	mat
The	-0.99	-0.83	0.46	0.29	-0.36	-0.08
cat	1.35	-1.10	-0.23	-1.75	1.15	-1.42
sat	-0.48	0.83	0.00	0.89	-0.54	0.80
on	1.41	0.66	-0.56	-0.73	0.67	-0.26
the	-0.85	0.15	0.24	0.77	-0.57	0.52
mat	0.71	0.89	-0.38	-0.03	0.17	0.25

Weights for the selected token appear at Softmax.

transformer.pyattention_forwardGitHub

70def attention_forward(Xn, params, p, cfg: ModelConfig):71    T, d = Xn.shape72    H = cfg.num_heads73    dh = d // H # width of one head7475    Q = Xn @ params[p+"wq"] + params[p+"bq"]    # (T,d), QUERY76    K = Xn @ params[p+"wk"] + params[p+"bk"]    # KEY77    V = Xn @ params[p+"wv"] + params[p+"bv"]    # VALUE7879    Qh, Kh, Vh = split_heads(Q, H), split_heads(K, H), split_heads(V, H)8081    S = (Qh @ Kh.transpose(0, 2, 1)) / np.sqrt(dh)82    S = np.where(np.triu(np.ones((T, T), dtype=bool), 1), -np.inf, S)83    A = softmax_forward(S)84    O = merge_heads(A @ Vh)85    out = O @ params[p + "wo"] + params[p + "bo"]86    return out, dict(Xn=Xn, Qh=Qh, Kh=Kh, Vh=Vh, A=A, O=O)

Actual Model Output

Actual Model OutputOne real window

Commit 66da45b. Checkpoint model.npz sha256 74de2ec88a63. Seed 42, from config.py, since a checkpoint does not record its own. Read from checkpoints/model, an ignored local file, trained on data/corpus.csv (git-ignored, and not at the pinned commit), with training settings that by inference differ from config.py.

In the lab above, choose Actual Model Output to see the weights A for every layer and head. They come from running forward on one window of 16 real tokens: ␠back␠to␠the␠game.\n'You␠can'

Checks on all 4 matrices, each 16 by 16: the shapes are right, every row sums to 1 (largest gap 3.3e-16), and every masked cell is exactly 0. Checkpoint weights file: 74de2ec88a63.

Where these numbers came from

Repository commit	`66da45bea72176fa5d2f8dba2ea1c7e597e30dd3`
`model.npz` sha256	`74de2ec88a63ceb131a27e03cc63070c6d8c4174cc78bc189195e9356e84550e`
`model.json` sha256	`f93002949162068d573aab9b38990b12aec7f32b468ec9eaaef986075df7d3b4`
`model.tok.json` sha256	`bf874ab4e4de2b61cb936f38635a7c589c5c0d29879947a5b45d13be2456b17d`
Corpus `data/corpus.csv` sha256	`c1bec80f274294e6ce989fcc63fd7022a206f415e7d3df85e82c410e56dc4f4f`
Tokenizer rebuilt from this corpus equals the checkpoint's	Yes
Model shape, from the checkpoint	`{"context_len":32,"emb_dim":32,"ff_dim":128,"ln_eps":0.00001,"num_heads":2,"num_layers":2,"vocab_size":575}`
`config.py`, when this was exported	`{"adam_eps":1e-8,"batch_size":16,"beta1":0.9,"beta2":0.999,"context_len":32,"corpus_column":"text","corpus_path":"data/corpus.csv","emb_dim":32,"epochs":20,"ff_dim":128,"grad_clip":1,"head_dim":16,"ln_eps":0.00001,"lr_max":0.001,"lr_min":0.0001,"max_new_tokens":50,"max_rows":2000,"num_heads":2,"num_layers":2,"num_merges":500,"seed":42,"seq_len":16,"stride":1,"temperature":0.8,"top_k":0,"top_p":0.9,"val_fraction":0.1,"warmup_steps":50,"weight_decay":0.01}`
Seed	`42`
Environment	`Python 3.14.4, NumPy 2.5.2, pandas 3.1.0`

Q, K, V projections

What is it?

Q, K, V projectionsImplemented

Attention

Attention makes three new matrices from its input Xn: the query Q, the key K and the value V. Each is a linear layer with its own weight matrix and bias. All three have the same shape as Xn.

Why does the model need it?

The later steps need three different views of each row. Scores compare queries with keys, and the output is built from values. Separate weights let the three views differ.

Mathematics

Query projection #

Q = X_n W_q + \mathbf{1} b_q^\top

Key projection #

K = X_n W_k + \mathbf{1} b_k^\top

Value projection #

V = X_n W_v + \mathbf{1} b_v^\top

What does each symbol mean?

Symbol	Meaning	In code	Shape
$X_n$	The input, a layer-normalised copy of the block's input	`Xn`	(T, d)
$W_q, W_k, W_v$	Weight matrices, one per projection	`params[p+"wq"], params[p+"wk"], params[p+"wv"]`	(d, d)
$b_q, b_k, b_v$	Biases, added to every row	`params[p+"bq"], params[p+"bk"], params[p+"bv"]`	(d,)
$Q, K, V$	Query, key and value matrices. On this page V is always the value matrix	`Q, K, V`	(T, d)

Tiny numerical example

Three projections of one inputIllustrative ExampleHand-picked numbers, one row per word of the sentence. They do not encode the words and are not learned values.

Xn, shape (6, 4)
0.5	-1	1.5	0.25
1	0	-0.5	2
-1.5	0.75	0.5	-0.25
0.25	1.25	-1	0.5
0.75	-0.5	1	-1
-0.5	1.5	0.25	0.75

Q, shape (6, 4)
-0.815	-0.495	1.555	-0.46
0.56	-1.22	0.68	0.94
-0.09	0.805	-0.445	-0.585
1.035	0.255	-0.67	0.74
-0.49	0.305	0.83	-0.385
0.66	0.605	-0.07	0.09

K, shape (6, 4)
1.47	-0.43	-0.165	-0.465
0.37	1.07	-1.29	1.16
-0.53	-0.055	1.16	-0.19
-0.955	0.995	-0.365	0.81
0.795	-0.58	0.01	-0.815
-0.48	0.945	0.285	0.685

V, shape (6, 4)
-1.005	0.715	0.365	-0.395
0.97	1.09	-0.21	-0.17
-0.155	-0.885	0.14	1.005
1.195	-0.035	-0.035	0.555
-0.905	0.315	0.49	-0.395
0.67	-0.26	0.365	1.205

Tensor shapes

One projection. Q, K and V all have this shapetransformer.py:75

Xn(T, d)(16, 32)
wq(d, d)(32, 32)
Q(T, d)(16, 32)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. d = emb_dim = 32.

Actual implementation

transformer.pyattention_forwardGitHub

75    Q = Xn @ params[p+"wq"] + params[p+"bq"]    # (T,d), QUERY76    K = Xn @ params[p+"wk"] + params[p+"bk"]    # KEY77    V = Xn @ params[p+"wv"] + params[p+"bv"]    # VALUE

Inside this function V is the value matrix, not the vocabulary size. This page uses V only for the value matrix.

init_params creates these parameters, one loop per block, for q, k, v and o:

transformer.pyinit_paramsGitHub

31        for n in "qkvo":32            p[b + "w" + n] = xavier_uniform(d, d, rng)33            p[b + "b" + n] = np.zeros(d)

Code and math connection

Equation	Role	Lines	Code
Query projection	Computes it	transformer.py:75	`Q = Xn @ params[p+"wq"] + params[p+"bq"] # (T,d), QUERY`
Key projection	Computes it	transformer.py:76	`K = Xn @ params[p+"wk"] + params[p+"bk"] # KEY`
Value projection	Computes it	transformer.py:77	`V = Xn @ params[p+"wv"] + params[p+"bv"] # VALUE`

Where is it used?

Source

transformer.py, lines 70-86, attention_forward

PDF reference

Section 4.5, Eq. 41, p. 10

Used by

Called by: transformer.block_forward

Tested by: tests.test_transformer.test_attention_weights_are_causal_and_normalised

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Query projection	Computes it	`transformer.py`	`attention_forward`	75	GitHub: transformer.py, line 75
Key projection	Computes it	`transformer.py`	`attention_forward`	76	GitHub: transformer.py, line 76
Value projection	Computes it	`transformer.py`	`attention_forward`	77	GitHub: transformer.py, line 77

Splitting into heads

What is it?

Splitting into headsImplemented

Attention

split_heads cuts the columns of Q, K and V into H equal blocks, one block per head. It reshapes (T, d) to (T, H, d_h), then swaps the first two axes to get (H, T, d_h).

Why does the model need it?

Each head computes its own scores and weights from its own block of columns. NumPy's @ works on the last two axes, so with the head on the first axis, one @ runs every head at once.

Mathematics

Split each matrix into heads #

Q_h,\ K_h,\ V_h = \operatorname{split}(Q, H),\ \operatorname{split}(K, H),\ \operatorname{split}(V, H)

Width of one head #

d_h = \frac{d}{H}, \qquad d \bmod H = 0

Head h owns one block of columns #

M^{(h)}_{t,i} = M_{t,\; h d_h + i}

What does each symbol mean?

Symbol	Meaning	In code	Shape
$H$	Number of heads	`H`	scalar
$d_h$	Width of one head	`dh`	scalar
$Q_h, K_h, V_h$	The split matrices. One head of each has the shape (T, d_h)	`Qh, Kh, Vh`	(H, T, d_h)

Tiny numerical example

Q, cut into two headsIllustrative ExampleHand-picked numbers, one row per word of the sentence. They do not encode the words and are not learned values. Head 0 holds the first two columns of Q, and head 1 the last two.

Q, shape (6, 4)
-0.815	-0.495	1.555	-0.46
0.56	-1.22	0.68	0.94
-0.09	0.805	-0.445	-0.585
1.035	0.255	-0.67	0.74
-0.49	0.305	0.83	-0.385
0.66	0.605	-0.07	0.09

Qh[0], shape (6, 2)
-0.815	-0.495
0.56	-1.22
-0.09	0.805
1.035	0.255
-0.49	0.305
0.66	0.605

Qh[1], shape (6, 2)
1.555	-0.46
0.68	0.94
-0.445	-0.585
-0.67	0.74
0.83	-0.385
-0.07	0.09

Tensor shapes

The reshape, then the transpose, in split_headstransformer.py:62

Q(T, d)(16, 32)
reshape(T, H, dh)(16, 2, 16)
transpose(H, T, dh)(2, 16, 16)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. d = emb_dim = 32. H = num_heads = 2. dh = head_dim = 16.

Actual implementation

transformer.pysplit_headsGitHub

60def split_heads(M, H):61    T, d = M.shape62    return M.reshape(T, H, d // H).transpose(1, 0, 2)

The function that calls it computes the width first (transformer.py:73), then splits all three matrices:

transformer.pyattention_forwardGitHub

79    Qh, Kh, Vh = split_heads(Q, H), split_heads(K, H), split_heads(V, H)

Code and math connection

Equation	Role	Lines	Code
Split each matrix into heads	Calls it	transformer.py:79	`Qh, Kh, Vh = split_heads(Q, H), split_heads(K, H), split_heads(V, H)`
Split each matrix into heads	Computes it	transformer.py:62	`return M.reshape(T, H, d // H).transpose(1, 0, 2)`
Split each matrix into heads	Tests it	tests/test_transformer.py:62	`assert Mh.shape == (2, 3, 4)`

Where is it used?

Source

transformer.py, lines 60-62, split_heads

PDF reference

Section 4.5, p. 10

Used by

Called by: transformer.attention_backward, transformer.attention_forward

Tested by: tests.test_transformer.test_split_merge_heads_round_trip

Calls it: transformer.py, line 79 (attention_forward), for Split each matrix into heads
Tests it: tests/test_transformer.py, line 62 (test_split_merge_heads_round_trip), for Split each matrix into heads

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Split each matrix into heads	Computes it	`transformer.py`	`split_heads`	62	GitHub: transformer.py, line 62

Scores

What is it?

ScoresImplemented

Attention

In each head, the product of Q_h and the transpose of K_h gives one number for every pair of positions. Entry (i, j) is the dot product of row i's query with row j's key.

Why does the model need it?

The weights have to come from somewhere. A score is the raw number that decides how much of position j's value row i receives, once the next steps have scaled, masked and normalised it.

Mathematics

Scores, one matrix per head #

Q_h K_h^\top, \qquad (Q_h K_h^\top)_{ij} = \sum_{c} (Q_h)_{ic} (K_h)_{jc}

What does each symbol mean?

Symbol	Meaning	In code	Shape
$Q_h$	The queries of one head	`Qh`	(T, d_h) per head
$K_h^\top$	The keys of one head, transposed. The code swaps the last two axes	`Kh.transpose(0, 2, 1)`	(d_h, T) per head
$i, j$	Row and column of the score matrix. i is the position reading, j the position read	not named in the code	scalar
$c$	Column index inside a head	not named in the code	scalar

Tiny numerical example

Qh @ Kh^T, head 0Illustrative ExampleHand-picked numbers, one row per word of the sentence. They do not encode the words and are not learned values.

Low -1.75High 1.41

Data table

row \ reads	The	cat	sat	on	the	mat
The	-0.99	-0.83	0.46	0.29	-0.36	-0.08
cat	1.35	-1.10	-0.23	-1.75	1.15	-1.42
sat	-0.48	0.83	0.00	0.89	-0.54	0.80
on	1.41	0.66	-0.56	-0.73	0.67	-0.26
the	-0.85	0.15	0.24	0.77	-0.57	0.52
mat	0.71	0.89	-0.38	-0.03	0.17	0.25

One cell, row 2 and column 1: [-0.09, 0.805] . [0.37, 1.07] = -0.09 * 0.37 + 0.805 * 1.07 = 0.828

Tensor shapes

Scores, before scaling. Both this step and the next are on one line of codetransformer.py:81

Qh(H, T, dh)(2, 16, 16)
Kh.transpose(0, 2, 1)(H, dh, T)(2, 16, 16)
Qh @ Kh^T(H, T, T)(2, 16, 16)

H = num_heads = 2. T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. dh = head_dim = 16.

Actual implementation

The product is the part of this line in brackets. Python never stores it on its own.

transformer.pyattention_forwardGitHub

81    S = (Qh @ Kh.transpose(0, 2, 1)) / np.sqrt(dh)

Code and math connection

Equation	Role	Lines	Code
Scores, one matrix per head	Computes it	transformer.py:81	`S = (Qh @ Kh.transpose(0, 2, 1)) / np.sqrt(dh)`

Where is it used?

Source

transformer.py, lines 70-86, attention_forward

PDF reference

Section 4.5, Eq. 42, p. 10

Used by

Called by: transformer.block_forward

Tested by: tests.test_transformer.test_attention_weights_are_causal_and_normalised

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Scores, one matrix per head	Computes it	`transformer.py`	`attention_forward`	81	GitHub: transformer.py, line 81

Scaling by sqrt(d_h)

What is it?

Scaling by sqrt(d_h)Implemented

Attention

The code divides every score by the square root of d_h, on the same line as the product. S is the scaled result.

Why does the model need it?

In general, a dot product of d_h terms grows as d_h grows, and softmax turns large gaps between scores into weights close to 0 or 1. Dividing by sqrt(d_h) keeps the size of the scores from growing with the head width. This is a general explanation. Nothing in this repository measures that effect.

Mathematics

Scaled scores #

S = \frac{Q_h K_h^\top}{\sqrt{d_h}}

What does each symbol mean?

Symbol	Meaning	In code	Shape
$S$	Scaled scores. The next line of code overwrites S with its masked version	`S`	(H, T, T)
$\sqrt{d_h}$	The scale. A single number, the same for every entry	`np.sqrt(dh)`	scalar

Tiny numerical example

S, head 0, after scalingIllustrative ExampleHand-picked numbers, one row per word of the sentence. They do not encode the words and are not learned values.

Low -1.24High 1.00

Data table

row \ reads	The	cat	sat	on	the	mat
The	-0.70	-0.59	0.32	0.20	-0.26	-0.05
cat	0.95	-0.78	-0.16	-1.24	0.82	-1.01
sat	-0.34	0.59	0.00	0.63	-0.38	0.57
on	1.00	0.46	-0.40	-0.52	0.48	-0.18
the	-0.60	0.10	0.17	0.55	-0.40	0.37
mat	0.50	0.63	-0.27	-0.02	0.12	0.18

This example has d_h = 2, so the scale is sqrt(2) = 1.4142. One cell: 0.828 / 1.4142 = 0.5855. The model has d_h = 16, so its scale is sqrt(16) = 4.

Tensor shapes

Scaling divides by one number, so the shape does not changetransformer.py:81

Qh @ Kh^T(H, T, T)(2, 16, 16)
S(H, T, T)(2, 16, 16)

H = num_heads = 2. T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt.

Actual implementation

The same line as the scores. The division happens before the mask, which is on the next line.

transformer.pyattention_forwardGitHub

81    S = (Qh @ Kh.transpose(0, 2, 1)) / np.sqrt(dh)

Code and math connection

Equation	Role	Lines	Code
Scaled scores	Computes it	transformer.py:81	`S = (Qh @ Kh.transpose(0, 2, 1)) / np.sqrt(dh)`

Where is it used?

Source

transformer.py, lines 70-86, attention_forward

PDF reference

Section 4.5, Eq. 42, p. 10

Used by

Called by: transformer.block_forward

Tested by: tests.test_transformer.test_attention_weights_are_causal_and_normalised

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Scaled scores	Computes it	`transformer.py`	`attention_forward`	81	GitHub: transformer.py, line 81

Causal mask

What is it?

Causal maskImplemented

Attention

Every score above the diagonal is replaced with minus infinity. np.triu with k = 1 selects those positions, where the column j is greater than the row i. The diagonal is kept, so a position can read itself.

Why does the model need it?

Training asks each position to predict the next token. If a position could read later positions, it could read the answer. The mask rules that out, and it applies in the same way during generation.

Mathematics

Causal mask #

S_{ij} \leftarrow \begin{cases} S_{ij} & j \le i \\ -\infty & j > i \end{cases}

What does each symbol mean?

Symbol	Meaning	In code	Shape
$S_{ij}$	The score of position i reading position j	`S`	(H, T, T)
$-\infty$	The value written above the diagonal	`-np.inf`	scalar

Tiny numerical example

S, head 0, after the maskIllustrative ExampleHand-picked numbers, one row per word of the sentence. They do not encode the words and are not learned values.

Low -0.78High 1.00

Data table

row \ reads	The	cat	sat	on	the	mat
The	-0.70	-inf (masked)	-inf (masked)	-inf (masked)	-inf (masked)	-inf (masked)
cat	0.95	-0.78	-inf (masked)	-inf (masked)	-inf (masked)	-inf (masked)
sat	-0.34	0.59	0.00	-inf (masked)	-inf (masked)	-inf (masked)
on	1.00	0.46	-0.40	-0.52	-inf (masked)	-inf (masked)
the	-0.60	0.10	0.17	0.55	-0.40	-inf (masked)
mat	0.50	0.63	-0.27	-0.02	0.12	0.18

Tensor shapes

The mask has one layer of shape (T, T). NumPy applies it to every headtransformer.py:82

S, scaled(H, T, T)(2, 16, 16)
np.triu(..., 1)(T, T)(16, 16)
S, masked(H, T, T)(2, 16, 16)

H = num_heads = 2. T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt.

Actual implementation

transformer.pyattention_forwardGitHub

82    S = np.where(np.triu(np.ones((T, T), dtype=bool), 1), -np.inf, S)

Code and math connection

Equation	Role	Lines	Code
Causal mask	Computes it	transformer.py:82	`S = np.where(np.triu(np.ones((T, T), dtype=bool), 1), -np.inf, S)`
Causal mask	Tests it	tests/test_transformer.py:85	`assert np.allclose(la[:4], lb[:4])`

Where is it used?

Source

transformer.py, lines 70-86, attention_forward

PDF reference

Section 4.5, Eq. 43, p. 10

Used by

Called by: transformer.block_forward

Tested by: tests.test_transformer.test_attention_weights_are_causal_and_normalised

Tests it: tests/test_transformer.py, line 85 (test_future_tokens_do_not_change_earlier_logits), for Causal mask

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Causal mask	Computes it	`transformer.py`	`attention_forward`	82	GitHub: transformer.py, line 82

Softmax per row

What is it?

Softmax per rowImplemented

Attention

softmax_forward turns each row of masked scores into weights that are not negative and add up to 1. It works along the last axis, so each row of each head is computed on its own.

Why does the model need it?

The output of a position is a weighted average of value rows, so its weights need to add up to 1. Masked cells get a weight of exactly 0, so their values contribute nothing.

Mathematics

Weights, one row at a time #

A = \operatorname{softmax}(S), \qquad \sum_{j} A_{ij} = 1

Why the shift is safe. The factor from any constant c cancels #

\frac{e^{x_i - c}}{\sum_k e^{x_k - c}} = \frac{e^{-c}\, e^{x_i}}{e^{-c} \sum_k e^{x_k}} = \frac{e^{x_i}}{\sum_k e^{x_k}}

A masked cell has a weight of exactly 0 #

\exp(-\infty) = 0 \quad\Rightarrow\quad A_{ij} = 0 \text{ for } j > i

A row always has one finite score, because the diagonal is never masked. So its largest score is finite, and a masked cell stays at minus infinity after the shift. The first row has only its own cell, so its weight there is 1.

What does each symbol mean?

Symbol	Meaning	In code	Shape
$A$	Attention weights, one matrix per head	`A`	(H, T, T)
$A_{ij}$	How much of position j's value position i receives	`A[h, i, j]`	scalar
$z_i$	One number in the row softmax receives. Here, one masked score	`x`	scalar
$c$	The constant subtracted from a row. The code uses the row's largest score	`np.max(x, axis=-1, keepdims=True)`	one per row

Tiny numerical example

A, head 0Illustrative ExampleHand-picked numbers, one row per word of the sentence. They do not encode the words and are not learned values.

Low 0.100High 1.000

Data table

row \ reads	The	cat	sat	on	the	mat
The	1.000	0.000 (masked)	0.000 (masked)	0.000 (masked)	0.000 (masked)	0.000 (masked)
cat	0.849	0.151	0.000 (masked)	0.000 (masked)	0.000 (masked)	0.000 (masked)
sat	0.203	0.511	0.285	0.000 (masked)	0.000 (masked)	0.000 (masked)
on	0.487	0.285	0.121	0.107	0.000 (masked)	0.000 (masked)
the	0.105	0.212	0.227	0.329	0.128	0.000 (masked)
mat	0.217	0.247	0.100	0.129	0.149	0.158

Row	Sum of the row	Masked cells
The	`1.0000`	5
cat	`1.0000`	4
sat	`1.0000`	3
on	`1.0000`	2
the	`1.0000`	1
mat	`1.0000`	0

Tensor shapes

Softmax keeps the shape. It normalises along the last axistransformer.py:83

S(H, T, T)(2, 16, 16)
A(H, T, T)(2, 16, 16)

H = num_heads = 2. T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt.

Actual implementation

transformer.pyattention_forwardGitHub

83    A = softmax_forward(S)

ops.pysoftmax_forwardGitHub

22def softmax_forward(x):23    e = np.exp(x - np.max(x, axis=-1, keepdims=True))24    return e / e.sum(axis=-1, keepdims=True)

Softmax & Cross-Entropy explains the shift by the largest value.

Code and math connection

Equation	Role	Lines	Code
Weights, one row at a time	Calls it	transformer.py:83	`A = softmax_forward(S)`
Weights, one row at a time	Computes it	ops.py:24	`return e / e.sum(axis=-1, keepdims=True)`
Weights, one row at a time	Tests it	tests/test_transformer.py:73	`assert np.allclose(A.sum(axis=-1), 1.0)`
A masked cell has a weight of exactly 0	Computes it	ops.py:23	`e = np.exp(x - np.max(x, axis=-1, keepdims=True))`
A masked cell has a weight of exactly 0	Calls it	transformer.py:83	`A = softmax_forward(S)`
A masked cell has a weight of exactly 0	Tests it	tests/test_transformer.py:72	`assert np.allclose(np.triu(A[0], 1), 0.0) # nothing flows from the future`

Where is it used?

Source

transformer.py, lines 70-86, attention_forward

PDF reference

Section 4.5, Eq. 43, p. 10

Used by

Called by: transformer.block_forward

Tested by: tests.test_transformer.test_attention_weights_are_causal_and_normalised

Calls it: transformer.py, line 83 (attention_forward), for Weights, one row at a time
Tests it: tests/test_transformer.py, line 73 (test_attention_weights_are_causal_and_normalised), for Weights, one row at a time
Calls it: transformer.py, line 83 (attention_forward), for A masked cell has a weight of exactly 0
Tests it: tests/test_transformer.py, line 72 (test_attention_weights_are_causal_and_normalised), for A masked cell has a weight of exactly 0

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Weights, one row at a time	Computes it	`ops.py`	`softmax_forward`	24	GitHub: ops.py, line 24
A masked cell has a weight of exactly 0	Computes it	`ops.py`	`softmax_forward`	23	GitHub: ops.py, line 23

Weighted values

What is it?

Weighted valuesImplemented

Attention

Multiplying A by V_h gives each position a weighted sum of value rows, separately in each head. The weights of a row can only be non-zero up to that row, so the sum uses that position and earlier ones.

Why does the model need it?

This is the step where information moves between positions. Before it, each row was computed from its own input row alone. After it, a row also holds parts of the values of earlier rows.

Mathematics

Weighted values, one matrix per head #

A V_h

The same product, for one row Derived from implementation#

(A V_h)_{i} = \sum_{j \le i} A_{ij} (V_h)_{j}

What does each symbol mean?

Symbol	Meaning	In code	Shape
$V_h$	The values of one head	`Vh`	(T, d_h) per head
$(V_h)_{j}$	Row j of the values of one head	`Vh[h, j]`	(d_h,)
$A_{ij}$	The weight of position j, for position i	`A[h, i, j]`	scalar

Tiny numerical example

Weighted values, head 0Illustrative ExampleHand-picked numbers, one row per word of the sentence. They do not encode the words and are not learned values.

Vh[0], shape (6, 2)
-1.005	0.715
0.97	1.09
-0.155	-0.885
1.195	-0.035
-0.905	0.315
0.67	-0.26

(A @ Vh)[0], shape (6, 2)
-1.005	0.715
-0.7075	0.7715
0.2478	0.45
-0.1038	0.549
0.3428	0.1335
0.1307	0.3373

Row 2, "sat": 0.203 * [-1.005, 0.715] + 0.5115 * [0.97, 1.09] + 0.2855 * [-0.155, -0.885] = [0.2478, 0.45]

Tensor shapes

Weighted values have the shape of the valuestransformer.py:84

A(H, T, T)(2, 16, 16)
Vh(H, T, dh)(2, 16, 16)
A @ Vh(H, T, dh)(2, 16, 16)

H = num_heads = 2. T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. dh = head_dim = 16.

Actual implementation

The product is the argument of merge_heads on this line. Like the scores, it is not stored on its own.

transformer.pyattention_forwardGitHub

84    O = merge_heads(A @ Vh)

Code and math connection

Equation	Role	Lines	Code
Weighted values, one matrix per head	Computes it	transformer.py:84	`O = merge_heads(A @ Vh)`
The same product, for one row	Computes it	transformer.py:84	`O = merge_heads(A @ Vh)`

Where is it used?

Source

transformer.py, lines 70-86, attention_forward

PDF reference

Section 4.5, Eq. 44, p. 10

Used by

Called by: transformer.block_forward

Tested by: tests.test_transformer.test_attention_weights_are_causal_and_normalised

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Weighted values, one matrix per head	Computes it	`transformer.py`	`attention_forward`	84	GitHub: transformer.py, line 84
The same product, for one row	Computes it	`transformer.py`	`attention_forward`	84	GitHub: transformer.py, line 84

Merging heads

What is it?

Merging headsImplemented

Attention

merge_heads undoes the split. It swaps the first two axes back, then reshapes, so the weighted values of the heads sit side by side in one row of width d again.

Why does the model need it?

The output projection and everything after attention work on rows of width d. Merging restores that width, and the output projection is then the first step that can combine numbers from different heads.

Mathematics

Heads side by side #

O = \operatorname{merge}(A V_h) = \begin{bmatrix} A^{(0)} V^{(0)} & \cdots & A^{(H-1)} V^{(H-1)} \end{bmatrix}

What does each symbol mean?

Symbol	Meaning	In code	Shape
$O$	The merged weighted values	`O`	(T, d)
$A^{(h)}, V^{(h)}$	The weights and values of head h	`A[h], Vh[h]`	(T, T) and (T, d_h)

Tiny numerical example

Two heads, side by sideIllustrative ExampleHand-picked numbers, one row per word of the sentence. They do not encode the words and are not learned values.

(A @ Vh)[0], shape (6, 2)
-1.005	0.715
-0.7075	0.7715
0.2478	0.45
-0.1038	0.549
0.3428	0.1335
0.1307	0.3373

(A @ Vh)[1], shape (6, 2)
0.365	-0.395
0.0018	-0.2529
0.1272	0.0313
-0.0595	0.095
0.2218	0.278
0.1719	0.2974

O, shape (6, 4)
-1.005	0.715	0.365	-0.395
-0.7075	0.7715	0.0018	-0.2529
0.2478	0.45	0.1272	0.0313
-0.1038	0.549	-0.0595	0.095
0.3428	0.1335	0.2218	0.278
0.1307	0.3373	0.1719	0.2974

Tensor shapes

The transpose, then the reshape, in merge_headstransformer.py:67

A @ Vh(H, T, dh)(2, 16, 16)
transpose(T, H, dh)(16, 2, 16)
O(T, d)(16, 32)

H = num_heads = 2. T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. dh = head_dim = 16. d = emb_dim = 32.

Actual implementation

transformer.pymerge_headsGitHub

65def merge_heads(Mh):66    H, T, dh = Mh.shape67    return Mh.transpose(1, 0, 2).reshape(T, H * dh)

Code and math connection

Equation	Role	Lines	Code
Heads side by side	Calls it	transformer.py:84	`O = merge_heads(A @ Vh)`
Heads side by side	Computes it	transformer.py:67	`return Mh.transpose(1, 0, 2).reshape(T, H * dh)`
Heads side by side	Tests it	tests/test_transformer.py:63	`assert np.array_equal(merge_heads(Mh), M)`

Where is it used?

Source

transformer.py, lines 65-67, merge_heads

PDF reference

Section 4.5, Eq. 45, p. 10

Used by

Called by: transformer.attention_backward, transformer.attention_forward

Tested by: tests.test_transformer.test_split_merge_heads_round_trip

Calls it: transformer.py, line 84 (attention_forward), for Heads side by side
Tests it: tests/test_transformer.py, line 63 (test_split_merge_heads_round_trip), for Heads side by side

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Heads side by side	Computes it	`transformer.py`	`merge_heads`	67	GitHub: transformer.py, line 67

Output projection

What is it?

Output projectionImplemented

Attention

A last linear layer turns the merged values into the output of attention. It has its own weight matrix W_o and bias b_o, and it keeps the shape.

Why does the model need it?

Until this point each head's columns have stayed apart. W_o multiplies across all d columns, so each output number can draw on every head.

Mathematics

Output projection #

\mathrm{out} = O W_o + \mathbf{1} b_o^\top

What does each symbol mean?

Symbol	Meaning	In code	Shape
$W_o$	Output weight matrix	`params[p + "wo"]`	(d, d)
$b_o$	Output bias, added to every row	`params[p + "bo"]`	(d,)
$\mathrm{out}$	The output of attention, added to the block's input by the residual	`out`	(T, d)

Tiny numerical example

The output projectionIllustrative ExampleHand-picked numbers, one row per word of the sentence. They do not encode the words and are not learned values.

O, shape (6, 4)
-1.005	0.715	0.365	-0.395
-0.7075	0.7715	0.0018	-0.2529
0.2478	0.45	0.1272	0.0313
-0.1038	0.549	-0.0595	0.095
0.3428	0.1335	0.2218	0.278
0.1307	0.3373	0.1719	0.2974

wo, shape (4, 4)
0.6	-0.2	0.1	0.3
0.2	0.4	-0.5	0.1
-0.1	0.3	0.5	-0.2
0.4	0.1	0.2	0.6

out, shape (6, 4)
-0.6645	0.587	-0.3345	-0.56
-0.3815	0.4554	-0.4862	-0.3072
0.2285	0.2017	-0.1104	0.0927
0.0814	0.262	-0.2756	0.0726
0.3114	0.1092	0.1541	0.2186
0.2377	0.2201	0.0098	0.197

out = O @ wo + bo, with bo = [-0.01, 0.03, 0.02, -0.02] added to every row.

Tensor shapes

The output has the shape of the input Xntransformer.py:85

O(T, d)(16, 32)
wo(d, d)(32, 32)
out(T, d)(16, 32)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. d = emb_dim = 32.

Actual implementation

transformer.pyattention_forwardGitHub

85    out = O @ params[p + "wo"] + params[p + "bo"]

Code and math connection

Equation	Role	Lines	Code
Output projection	Computes it	transformer.py:85	`out = O @ params[p + "wo"] + params[p + "bo"]`

Where is it used?

Source

transformer.py, lines 70-86, attention_forward

PDF reference

Section 4.5, Eq. 45, p. 10

Used by

Called by: transformer.block_forward

Tested by: tests.test_transformer.test_attention_weights_are_causal_and_normalised

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Output projection	Computes it	`transformer.py`	`attention_forward`	85	GitHub: transformer.py, line 85

The transformer.py implementation

This is the whole function, with its real line numbers. Highlighted lines have a link record behind them.

Code listing: transformer.attention_forward

transformer.pyattention_forwardGitHub

70def attention_forward(Xn, params, p, cfg: ModelConfig):71    T, d = Xn.shape72    H = cfg.num_heads73    dh = d // H # width of one head7475    Q = Xn @ params[p+"wq"] + params[p+"bq"]    # (T,d), QUERY76    K = Xn @ params[p+"wk"] + params[p+"bk"]    # KEY77    V = Xn @ params[p+"wv"] + params[p+"bv"]    # VALUE7879    Qh, Kh, Vh = split_heads(Q, H), split_heads(K, H), split_heads(V, H)8081    S = (Qh @ Kh.transpose(0, 2, 1)) / np.sqrt(dh)82    S = np.where(np.triu(np.ones((T, T), dtype=bool), 1), -np.inf, S)83    A = softmax_forward(S)84    O = merge_heads(A @ Vh)85    out = O @ params[p + "wo"] + params[p + "bo"]86    return out, dict(Xn=Xn, Qh=Qh, Kh=Kh, Vh=Vh, A=A, O=O)

Two things to note when reading it:

S is assigned twice. The first assignment holds the scaled scores, and the second overwrites them with the masked scores.
The returned dictionary is the cache that the backward pass reads. It holds Xn, Qh, Kh, Vh, A and O. It does not hold S, Q, K or V.

Tensor shapes with the real config

These are the shapes from running the real functions at the values in config.py: T = 16, d = 32, H = 2 and d_h = 16. T can be smaller during generation, for a short prompt.

Forward, with the real configIllustrative ExampleSizes from config.py. These shapes are what the real functions returned on random inputs. No value from that run is shown.

Backward, with the real configIllustrative ExampleSizes from config.py. These shapes are what the real functions returned on random inputs. No value from that run is shown.

Entry of grads	Shape
`grads[p + "wq"]`	`(32, 32)`
`grads[p + "bq"]`	`(32,)`
`grads[p + "wk"]`	`(32, 32)`
`grads[p + "bk"]`	`(32,)`
`grads[p + "wv"]`	`(32, 32)`
`grads[p + "bv"]`	`(32,)`
`grads[p + "wo"]`	`(32, 32)`
`grads[p + "bo"]`	`(32,)`

The model has no batch axis here. H is the only axis attention adds, and it comes and goes inside this function. config.head_dim is not read by this code, which computes the width again with d // H (transformer.py:73).

Backward pass

What is it?

Attention backwardImplemented

Attention

attention_backward receives dOut, the gradient of the loss on the output of attention. It fills the gradients of all eight attention parameters and returns dXn, the gradient on the input. It runs the forward steps in reverse, reading the values the forward pass left in its cache.

Why does the model need it?

W_q, W_k, W_v, W_o and their biases are parameters, so the optimizer needs a gradient for each. Earlier layers also need the gradient to continue, so dXn goes on to the first LayerNorm's backward.

Mathematics

First the gradient goes through the output projection, and is split into heads the way the merged values were:

Gradient on the output weights #

dW_o = O^\top\, d\mathrm{Out}

Gradient on the output bias #

db_o = \mathbf{1}^\top d\mathrm{Out}

Through the output projection, then the head split #

dO_h = \operatorname{split}(d\mathrm{Out}\, W_o^\top, H)

The product A V_h depends on both factors, so the gradient goes to each of them:

Gradient on the values #

dV_h = A^\top dO_h

Gradient on the weights #

dA = dO_h V_h^\top

Then it goes back through softmax, to the scores:

Through softmax, softmax_backward(A, dA) #

dS_{ij} = A_{ij}\Big(dA_{ij} - \sum_{k} A_{ik}\, dA_{ik}\Big)

Why masked positions get no gradient. softmax_backward multiplies its result by the weights. A masked position has a weight of exactly 0, so whatever dA holds in that cell, dS is exactly 0 there. Its contribution to dVh is also 0, because A^T carries the same 0. The rest of the backward pass adds nothing from that cell, so the output at a position sends no gradient to any later position.

A masked cell gets no gradient Derived from implementation#

A_{ij} = 0 \quad\Rightarrow\quad dS_{ij} = 0

From the scores, the gradient goes to the queries and keys. The scale appears again here, because the forward pass divided by it:

Gradient on the queries #

dQ_h = \frac{dS\, K_h}{\sqrt{d_h}}

Gradient on the keys #

dK_h = \frac{dS^\top Q_h}{\sqrt{d_h}}

Heads side by side again Derived from implementation#

dQ,\ dK,\ dV = \operatorname{merge}(dQ_h),\ \operatorname{merge}(dK_h),\ \operatorname{merge}(dV_h)

Each of the three projections is a linear layer of the same input. So each gets a parameter gradient, and the gradients on the input add up:

Gradient on the three weight matrices #

dW_n = X_n^\top dM_n, \qquad n \in \{q, k, v\}

Gradient on the three biases #

db_n = \mathbf{1}^\top dM_n, \qquad n \in \{q, k, v\}

Gradient on the input, summed over q, k and v #

dX_n = dQ\, W_q^\top + dK\, W_k^\top + dV\, W_v^\top

The key bias gets a gradient of zero. Each row of dS adds up to 0, because softmax ignores a constant added to a row. Adding b_k to every key adds the same amount to every score in a row of S, so it cannot change the weights.

The key bias has a gradient of zero #

db_k = \mathbf{1}^\top dK = 0

What does each symbol mean?

Symbol	Meaning	In code	Shape
$d\mathrm{Out}$	Gradient of the loss on the output of attention	`dOut`	(T, d)
$dO_h$	Gradient on the merged values, split into heads	`dOh`	(H, T, d_h)
$dV_h$	Gradient on the values, per head	`dVh`	(H, T, d_h)
$dA$	Gradient on the weights	`dA`	(H, T, T)
$dS$	Gradient on the scaled, masked scores	`dS`	(H, T, T)
$dQ_h, dK_h$	Gradient on the queries and keys, per head	`dQh, dKh`	(H, T, d_h)
$dQ, dK, dV$	The same three gradients after merging	`dQ, dK, dV`	(T, d)
$dM_n$	One of dQ, dK or dV, as the loop in the code takes them in turn	`dM`	(T, d)
$n$	The name of a projection, q, k or v	`n`	scalar
$dX_n$	Gradient on the input of attention	`dXn`	(T, d)
$dW_o, db_o, dW_n, db_n$	Parameter gradients, stored in grads	`grads[p + "w" + n], grads[p + "b" + n]`	(d, d) and (d,)

Tiny numerical example

One backward pass, head 0Illustrative ExampleThe same hand-picked forward pass as above. dOut is also hand-picked, not computed from a loss.

dOut, shape (6, 4)
0.1	-0.2	0.3	0.05
-0.1	0.25	0.15	-0.3
0.2	0.1	-0.25	0.15
-0.05	0.3	0.1	-0.2
0.15	-0.1	0.2	0.25
0.3	-0.15	-0.1	0.1

dOh[0], shape (6, 2)
0.145	-0.205
-0.185	-0.025
0.12	0.22
-0.14	0.04
0.205	-0.085
0.23	0.06

dVh[0], shape (6, 2)
0.0155	-0.1579
0.0937	0.117
0.0869	0.0544
0.0822	-0.016
0.0604	-0.0019
0.0362	0.0095

dA[0], shape (6, 6)
-0.2923	-0.0828	0.159	0.1804	-0.1958	0.1505
0.1681	-0.2067	0.0508	-0.2202	0.1595	-0.1174
0.0367	0.3562	-0.2133	0.1357	-0.0393	0.0232
0.1693	-0.0922	-0.0137	-0.1687	0.1393	-0.1042
-0.2668	0.1062	0.0435	0.248	-0.2123	0.1595
-0.1883	0.2885	-0.0887	0.2727	-0.1893	0.1385

dS[0], shape (6, 6)
0	0	0	0	0	0
0.0479	-0.0479	0	0	0	0
-0.0187	0.1163	-0.0976	0	0	0
0.0647	-0.0367	-0.0061	-0.0219	0	0
-0.0341	0.01	-0.0035	0.0623	-0.0347	0
-0.0518	0.0589	-0.014	0.0287	-0.0356	0.0139

dQh[0], shape (6, 2)
0	0
0.0373	-0.0509
0.0476	0.0975
0.0747	-0.0626
-0.093	0.0761
-0.0774	0.1049

dKh[0], shape (6, 2)
0.0551	-0.0699
-0.0293	0.1283
-0.0035	-0.0634
-0.0242	0.0217
-0.0046	-0.0227
0.0065	0.0059

dXn, shape (6, 4)
-0.1036	0.1233	0.075	-0.0257
0.1602	0.0652	0.0038	0.1086
-0.0366	0.1379	0.0064	-0.0246
0.0051	0.0798	-0.0701	0.0543
-0.0098	0.0451	0.0491	0.0041
-0.0095	0.0443	0.0254	-0.0254

A masked cell gets no gradientIllustrative ExampleRow 1, column 2 of head 0. Position 1 may not read position 2, so this cell is masked.

Quantity	Code	Value
Weight	`A[0, 1, 2]`	`0`
Gradient on the weight	`dA[0, 1, 2]`	`0.0508`
Gradient on the score	`dS[0, 1, 2]`	`0`

Where the results goIllustrative ExampleThe same run. grads holds the parameter gradients, and dXn is returned to block_backward.

grads[p + "wo"], shape (4, 4)
0.1156	-0.0361	-0.4245	0.3187
0.1781	0.1956	0.2656	-0.1709
0.1496	-0.1256	0.0992	0.1213
0.1182	-0.025	-0.1289	0.1411

grads[p + "wq"], shape (4, 4)
-0.0465	-0.2082	0.0324	-0.0677
0.0595	0.1142	0.0135	0.0312
-0.1819	0.2391	0.0649	0.0706
0.135	-0.1548	0.005	-0.0955

grads[p + "wk"], shape (4, 4)
-0.0092	0.1739	-0.0796	0.0537
-0.076	0.0698	0.0373	-0.0921
0.1168	-0.2436	0.0059	0.0985
-0.0465	0.293	-0.0286	-0.0442

grads[p + "wv"], shape (4, 4)
0.0189	-0.0537	0.273	-0.1431
0.1767	0.1939	-0.1149	-0.023
0.007	-0.2518	-0.0207	0.4042
0.1775	0.182	0.3727	-0.1745

The largest entry of grads[p + "bk"] is 2.9e-17, which is round-off around 0.

Checking the backward numericallyIllustrative ExampleA test-time idea, not part of the pipeline: gradcheck.numerical_grad is only called by tests. It nudges each number up and down by h = 1e-5 and measures how (out * dOut).sum(), whose gradient on out is exactly dOut, changes. Tests compare the central differences here with that real function.

Gradient	From the backward formula	From central differences	Largest gap
`grads[p + "wo"]`	`[0.1156, -0.0361, -0.4245, 0.3187][0.1781, 0.1956, 0.2656, -0.1709][0.1496, -0.1256, 0.0992, 0.1213][0.1182, -0.0250, -0.1289, 0.1411]`	`[0.1156, -0.0361, -0.4245, 0.3187][0.1781, 0.1956, 0.2656, -0.1709][0.1496, -0.1256, 0.0992, 0.1213][0.1182, -0.0250, -0.1289, 0.1411]`	`3.1e-12`
`grads[p + "wq"]`	`[-0.0465, -0.2082, 0.0324, -0.0677][0.0595, 0.1142, 0.0135, 0.0312][-0.1819, 0.2391, 0.0649, 0.0706][0.1350, -0.1548, 0.0050, -0.0955]`	`[-0.0465, -0.2082, 0.0324, -0.0677][0.0595, 0.1142, 0.0135, 0.0312][-0.1819, 0.2391, 0.0649, 0.0706][0.1350, -0.1548, 0.0050, -0.0955]`	`7.4e-12`
`grads[p + "wk"]`	`[-0.0092, 0.1739, -0.0796, 0.0537][-0.0760, 0.0698, 0.0373, -0.0921][0.1168, -0.2436, 0.0059, 0.0985][-0.0465, 0.2930, -0.0286, -0.0442]`	`[-0.0092, 0.1739, -0.0796, 0.0537][-0.0760, 0.0698, 0.0373, -0.0921][0.1168, -0.2436, 0.0059, 0.0985][-0.0465, 0.2930, -0.0286, -0.0442]`	`7.1e-12`
`grads[p + "wv"]`	`[0.0189, -0.0537, 0.2730, -0.1431][0.1767, 0.1939, -0.1149, -0.0230][0.0070, -0.2518, -0.0207, 0.4042][0.1775, 0.1820, 0.3727, -0.1745]`	`[0.0189, -0.0537, 0.2730, -0.1431][0.1767, 0.1939, -0.1149, -0.0230][0.0070, -0.2518, -0.0207, 0.4042][0.1775, 0.1820, 0.3727, -0.1745]`	`7.1e-12`
`dXn`	`[-0.1036, 0.1233, 0.0750, -0.0257][0.1602, 0.0652, 0.0038, 0.1086][-0.0366, 0.1379, 0.0064, -0.0246][0.0051, 0.0798, -0.0701, 0.0543][-0.0098, 0.0451, 0.0491, 0.0041][-0.0095, 0.0443, 0.0254, -0.0254]`	`[-0.1036, 0.1233, 0.0750, -0.0257][0.1602, 0.0652, 0.0038, 0.1086][-0.0366, 0.1379, 0.0064, -0.0246][0.0051, 0.0798, -0.0701, 0.0543][-0.0098, 0.0451, 0.0491, 0.0041][-0.0095, 0.0443, 0.0254, -0.0254]`	`8.8e-12`

Tensor shapes

From dOut to the valuestransformer.py:98

dOut(T, d)(16, 32)
dOh(H, T, dh)(2, 16, 16)
dVh(H, T, dh)(2, 16, 16)
dV(T, d)(16, 32)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. d = emb_dim = 32. H = num_heads = 2. dh = head_dim = 16.

From dOh through the weights and scores to the queries. Keys go the same waytransformer.py:101

dOh(H, T, dh)(2, 16, 16)
dA(H, T, T)(2, 16, 16)
dS(H, T, T)(2, 16, 16)
dQh(H, T, dh)(2, 16, 16)
dQ(T, d)(16, 32)

H = num_heads = 2. T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. dh = head_dim = 16. d = emb_dim = 32.

Parameter gradient and input gradient, for one projectiontransformer.py:108

dM(T, d)(16, 32)
grads wn(d, d)(32, 32)
dXn(T, d)(16, 32)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. d = emb_dim = 32.

Actual implementation

transformer.pyattention_backwardGitHub

92def attention_backward(dOut, c, params, p, cfg: ModelConfig, grads):93    H = cfg.num_heads94    dh = cfg.emb_dim // H95    grads[p + "wo"] += c["O"].T @ dOut96    grads[p + "bo"] += dOut.sum(axis=0)97    dOh = split_heads(dOut @ params[p + "wo"].T, H)98    dVh = c["A"].transpose(0, 2, 1) @ dOh 99    dA = dOh @ c["Vh"].transpose(0, 2, 1)100    dS = softmax_backward(c["A"], dA)101    dQh = dS @ c["Kh"] / np.sqrt(dh)102    dKh = dS.transpose(0, 2, 1) @ c["Qh"] / np.sqrt(dh)103    dQ, dK, dV = merge_heads(dQh), merge_heads(dKh), merge_heads(dVh)104    dXn = np.zeros_like(c["Xn"])105    for n, dM in (("q", dQ), ("k", dK), ("v", dV)):106        grads[p + "w" + n] += c["Xn"].T @ dM107        grads[p + "b" + n] += dM.sum(axis=0)108        dXn += dM @ params[p + "w" + n].T109    return dXn

This function writes the linear layer's gradient formulas in place and does not call linear_backward.

Code and math connection

Equation	Role	Lines	Code
Gradient on the output weights	Computes it	transformer.py:95	`grads[p + "wo"] += c["O"].T @ dOut`
Gradient on the output bias	Computes it	transformer.py:96	`grads[p + "bo"] += dOut.sum(axis=0)`
Through the output projection, then the head split	Computes it	transformer.py:97	`dOh = split_heads(dOut @ params[p + "wo"].T, H)`
Gradient on the values	Computes it	transformer.py:98	`dVh = c["A"].transpose(0, 2, 1) @ dOh`
Gradient on the weights	Computes it	transformer.py:99	`dA = dOh @ c["Vh"].transpose(0, 2, 1)`
Through softmax, softmax_backward(A, dA)	Calls it	transformer.py:100	`dS = softmax_backward(c["A"], dA)`
Through softmax, softmax_backward(A, dA)	Computes it	ops.py:27	`return y * (dY - np.sum(y * dY, axis=-1, keepdims=True))`
A masked cell gets no gradient	Computes it	ops.py:27	`return y * (dY - np.sum(y * dY, axis=-1, keepdims=True))`
Gradient on the queries	Computes it	transformer.py:101	`dQh = dS @ c["Kh"] / np.sqrt(dh)`
Gradient on the keys	Computes it	transformer.py:102	`dKh = dS.transpose(0, 2, 1) @ c["Qh"] / np.sqrt(dh)`
Heads side by side again	Computes it	transformer.py:103	`dQ, dK, dV = merge_heads(dQh), merge_heads(dKh), merge_heads(dVh)`
Gradient on the three weight matrices	Computes it	transformer.py:106	`grads[p + "w" + n] += c["Xn"].T @ dM`
Gradient on the three biases	Computes it	transformer.py:107	`grads[p + "b" + n] += dM.sum(axis=0)`
Gradient on the input, summed over q, k and v	Computes it	transformer.py:108	`dXn += dM @ params[p + "w" + n].T`
Gradient on the input, summed over q, k and v	Calls it	transformer.py:148	`dXn = attention_backward(dX1, c["att"], params, p, cfg, grads)`
The key bias has a gradient of zero	Computes it	transformer.py:107	`grads[p + "b" + n] += dM.sum(axis=0)`
The key bias has a gradient of zero	Tests it	tests/test_gradcheck.py:39	`assert np.abs(grads[f"b{l}.bk"]).max() < 1e-12`

Where is it used?

block_backward calls it with the gradient after the residual (transformer.py:147, transformer.py:148).

Source: transformer.py, lines 92-109, attention_backward
PDF reference: Section 4.5, Eq. 46, 47, 48, 49, 50, p. 10
Used by: Called by: transformer.block_backward
Calls it: transformer.py, line 100 (attention_backward), for Through softmax, softmax_backward(A, dA)
Calls it: transformer.py, line 148 (block_backward), for Gradient on the input, summed over q, k and v
Tests it: tests/test_gradcheck.py, line 39 (test_key_bias_gradient_is_zero), for The key bias has a gradient of zero
Status: Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Gradient on the output weights	Computes it	`transformer.py`	`attention_backward`	95	GitHub: transformer.py, line 95
Gradient on the output bias	Computes it	`transformer.py`	`attention_backward`	96	GitHub: transformer.py, line 96
Through the output projection, then the head split	Computes it	`transformer.py`	`attention_backward`	97	GitHub: transformer.py, line 97
Gradient on the values	Computes it	`transformer.py`	`attention_backward`	98	GitHub: transformer.py, line 98
Gradient on the weights	Computes it	`transformer.py`	`attention_backward`	99	GitHub: transformer.py, line 99
Through softmax, softmax_backward(A, dA)	Computes it	`ops.py`	`softmax_backward`	27	GitHub: ops.py, line 27
A masked cell gets no gradient	Computes it	`ops.py`	`softmax_backward`	27	GitHub: ops.py, line 27
Gradient on the queries	Computes it	`transformer.py`	`attention_backward`	101	GitHub: transformer.py, line 101
Gradient on the keys	Computes it	`transformer.py`	`attention_backward`	102	GitHub: transformer.py, line 102
Heads side by side again	Computes it	`transformer.py`	`attention_backward`	103	GitHub: transformer.py, line 103
Gradient on the three weight matrices	Computes it	`transformer.py`	`attention_backward`	106	GitHub: transformer.py, line 106
Gradient on the three biases	Computes it	`transformer.py`	`attention_backward`	107	GitHub: transformer.py, line 107
Gradient on the input, summed over q, k and v	Computes it	`transformer.py`	`attention_backward`	108	GitHub: transformer.py, line 108
The key bias has a gradient of zero	Computes it	`transformer.py`	`attention_backward`	107	GitHub: transformer.py, line 107

How this reaches a parameter update. grads is filled in place. batch_loss_and_grads averages the gradients over the batch (train.py:61), clip_grad_norm may scale them (train.py:62), and opt.step uses them (train.py:66). The weight matrices also get weight decay, and the biases do not. dXn goes on to the first LayerNorm's backward (transformer.py:149).

Role in the model

Attention is the first branch of each block. block_forward normalises its input, runs attention on that copy, and adds the result to the unnormalised input:

transformer.pyblock_forwardGitHub

134    att, c_att = attention_forward(Xn, params, p, cfg) # each position gathers info from earlier positions

transformer.pyblock_forwardGitHub

135    X1 = X + att

The model has 2 blocks, and each has its own eight attention parameters:

Name	Shape	With config	Scalars	Role
`b{l}.wq`transformer.py:32	`(d, d)`	`(32, 32)`	`d × d`1,024 per block	Query projection weights.Starts as Xavier uniform. Weight decay.
`b{l}.bq`transformer.py:33	`(d,)`	`(32,)`	`d`32 per block	Query projection bias.Starts as zeros. No decay.
`b{l}.wk`transformer.py:32	`(d, d)`	`(32, 32)`	`d × d`1,024 per block	Key projection weights.Starts as Xavier uniform. Weight decay.
`b{l}.bk`transformer.py:33	`(d,)`	`(32,)`	`d`32 per block	Key projection bias.Starts as zeros. No decay.
`b{l}.wv`transformer.py:32	`(d, d)`	`(32, 32)`	`d × d`1,024 per block	Value projection weights.Starts as Xavier uniform. Weight decay.
`b{l}.bv`transformer.py:33	`(d,)`	`(32,)`	`d`32 per block	Value projection bias.Starts as zeros. No decay.
`b{l}.wo`transformer.py:32	`(d, d)`	`(32, 32)`	`d × d`1,024 per block	Attention output projection weights.Starts as Xavier uniform. Weight decay.
`b{l}.bo`transformer.py:33	`(d,)`	`(32,)`	`d`32 per block	Attention output projection bias.Starts as zeros. No decay.

Attention is the only step in a block that moves information between positions. The feed-forward layer that follows works on each row alone. Because the mask is applied in every block, the output at a position depends only on that position and earlier ones, all the way through the model. A test checks this on the whole model, by changing only the last token and comparing the earlier logits (tests/test_transformer.py:85).

Next: Transformer Block shows how the branches fit together.