Notation & Shapes

This page fixes the notation for the rest of the Mathematics section. Every symbol below names something real in the code, and every size is read from config.py.

Conventions

These are how the code arranges its numbers, not general rules about transformers.

A token is a row. One token's vector is a row with d numbers.
A sequence is a matrix. A sequence of T tokens has shape [T x d].
A linear layer is X @ W + b. X is on the left, so W has one row per input feature. The bias b is added to every row.
There is no batch axis. forward and loss_and_grads take one sequence at a time. batch_loss_and_grads in train.py is a Python loop over the windows of a batch (train.py:17).
Heads add one axis. Multi-head attention works on [H x T x d_h]. The head comes first.

Symbols and code names

Values come from config.py. V has no value there, because it depends on the corpus.

Symbol	Meaning	config.py	Value	In the model code
$T$	Tokens in one sequence Inside the model, T is read from the array, not from config. Generation can use fewer than seq_len.	`seq_len`	16	`T, d = Xn.shape` transformer.py:71
$T_{\max}$	Rows in the position table, the longest allowed sequence	`context_len`	32	`cfg.context_len` transformer.py:22
$d$	Width of each token's vector	`emb_dim`	32	`cfg.emb_dim` transformer.py:18
$H$	Attention heads	`num_heads`	2	`cfg.num_heads` transformer.py:72
$d_h$	Width of one head Nothing reads config.head_dim. The code computes the same value in place.	`head_dim`	16	`dh = d // H` transformer.py:73
$F$	Width of the feed-forward hidden layer	`ff_dim`	128	`cfg.ff_dim` transformer.py:18
$L$	Transformer blocks	`num_layers`	2	`cfg.num_layers` transformer.py:27
$V$	Vocabulary size It depends on the corpus, so config.py has no value for it. Inside attention_forward the local name V is the value matrix, not this.	not in config.py	depends on the corpus	`tok.vocab_size` train.py:48

A letter can mean something else in a local variable. The notes in the table name the case that matters most: V inside attention_forward.

Matrix multiplication

What is it?

Matrix multiplicationDerived from implementation

Basic operations

The code writes every linear layer as X @ W + b, where X has one row per token. The @ operator multiplies each row of X by the matrix W, and b is added to every row of the result.

Why does the model need it?

It is how the model changes the width of token vectors and mixes their features. Query, key and value, both feed-forward layers, the attention output and the head are all this one formula.

Mathematics

One entry of the product #

(XW)_{tj} = \sum_{k} X_{tk} W_{kj}

The shape rule #

X \in \mathbb{R}^{T \times a},\; W \in \mathbb{R}^{a \times c} \;\Rightarrow\; XW \in \mathbb{R}^{T \times c}

A linear layer, with b added to every row #

Y_{tj} = (XW)_{tj} + b_j

What does each symbol mean?

Symbol	Meaning	In code	Shape
$X$	Input, one row per token	`X`	(T, a)
$W$	Weight matrix	`W`	(a, c)
$b$	Bias, added to every row	`b`	(c,)
$Y$	Output, one row per token	not named in the code	(T, c)
$a$	Input width. It is d, or F for the second feed-forward layer	not named in the code	scalar
$c$	Output width. It is d, F or V, depending on the layer	not named in the code	scalar

Tiny numerical example

Edit any cell. Switch W to three columns and watch the shape of the result change, while the rows stay the same.

X @ W and X @ W + bIllustrative ExampleHand-picked numbers that you can edit. These are not learned weights. Results come from the math engine, and tests compare it with the real Python on the starting values.

Columns of W

X
	k0	k1	k2
t0
t1

W
	j0	j1
k0
k1
k2

b
	j0	j1
b

(2, 3) @ (3, 2) → (2, 2): columns of X must equal rows of W. The result has the rows of X and the columns of W.

Row 0, column 0 of the product: 1 × 0.5 + 2 × 1.5 + -1 × 2 = 1.5

X @ W, shape (2, 2)
1.5	-0.5
6.75	0.125

X @ W + b, shape (2, 2)
1.6	-0.7
6.85	-0.075

Tensor shapes

The rule above, at three of the places the model uses it.

Query projection, width stays the sametransformer.py:75

Xn(T, d)(16, 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.

First feed-forward layer, width growstransformer.py:113

Xn(T, d)(16, 32)
Z(T, F)(16, 128)

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. F = ff_dim = 128.

Linear head, width becomes the vocabulary sizetransformer.py:166

Xf(T, d)(16, 32)
logits(T, V)(16, V)

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. V depends on the corpus. It is not a value in config.py.

Actual implementation

ops.pylinear_forwardGitHub

7def linear_forward(X, W, b):8    return X @ W + b

Code and math connection

Equation	Role	Lines	Code
One entry of the product	Computes it	ops.py:8	`return X @ W + b`
The shape rule	Computes it	ops.py:8	`return X @ W + b`
A linear layer, with b added to every row	Computes it	ops.py:8	`return X @ W + b`
A linear layer, with b added to every row	Written inline	transformer.py:75	`Q = Xn @ params[p+"wq"] + params[p+"bq"] # (T,d), QUERY`
A linear layer, with b added to every row	Written inline	transformer.py:76	`K = Xn @ params[p+"wk"] + params[p+"bk"] # KEY`
A linear layer, with b added to every row	Written inline	transformer.py:77	`V = Xn @ params[p+"wv"] + params[p+"bv"] # VALUE`
A linear layer, with b added to every row	Written inline	transformer.py:85	`out = O @ params[p + "wo"] + params[p + "bo"]`
A linear layer, with b added to every row	Written inline	transformer.py:113	`Z = Xn @ params[p + "ff1_w"] + params[p + "ff1_b"]`
A linear layer, with b added to every row	Written inline	transformer.py:116	`return A @ params[p + "ff2_w"] + params[p + "ff2_b"], dict(Xn=Xn, Z=Z, A=A)`
A linear layer, with b added to every row	Written inline	transformer.py:166	`logits = Xf @ params["head_w"] + params["head_b"]`
A linear layer, with b added to every row	Tests it	tests/test_ops.py:15	`assert np.allclose(linear_forward(X, W, b), [[-1, 3, 1], [-1, 2, -1]])`

The backward pass of this layer, linear_backward, is on Linear / ReLU / Initialisation.

Where is it used?

ops.linear_forward is defined, but only a test calls it. transformer.py does not import it and writes the formula in place instead, at the lines listed here.

Source

ops.py, lines 7-8, linear_forward

PDF reference

Section 4.3, Eq. 19, p. 6

Used by

Not called by the pipeline.

Tested by: tests.test_ops.test_linear_worked_example

Written inline in transformer.py, line 75. Query. Key and value are on the next two lines.

Written inline in transformer.py, line 85. Attention output.

Written inline in transformer.py, line 113. First feed-forward layer.

Written inline in transformer.py, line 116. Second feed-forward layer.

Written inline in transformer.py, line 166. The head.

Tests it: tests/test_ops.py, line 15 (test_linear_worked_example), for A linear layer, with b added to every row, occurrence 0

Status

Derived from implementation

Implemented in

NoteInline equivalent: the pipeline does not call ops.linear_forward. It computes the same formula in place at transformer.py:75, transformer.py:76, transformer.py:77, transformer.py:85, transformer.py:113, transformer.py:116, transformer.py:166.

Equation	Role	File	Function	Lines	GitHub
One entry of the product	Computes it	`ops.py`	`linear_forward`	8	GitHub: ops.py, line 8
The shape rule	Computes it	`ops.py`	`linear_forward`	8	GitHub: ops.py, line 8
A linear layer, with b added to every row	Computes it	`ops.py`	`linear_forward`	8	GitHub: ops.py, line 8
A linear layer, with b added to every row	Written inline	`transformer.py`	`attention_forward`	75	GitHub: transformer.py, line 75
A linear layer, with b added to every row	Written inline	`transformer.py`	`attention_forward`	76	GitHub: transformer.py, line 76
A linear layer, with b added to every row	Written inline	`transformer.py`	`attention_forward`	77	GitHub: transformer.py, line 77
A linear layer, with b added to every row	Written inline	`transformer.py`	`attention_forward`	85	GitHub: transformer.py, line 85
A linear layer, with b added to every row	Written inline	`transformer.py`	`ffn_forward`	113	GitHub: transformer.py, line 113
A linear layer, with b added to every row	Written inline	`transformer.py`	`ffn_forward`	116	GitHub: transformer.py, line 116
A linear layer, with b added to every row	Written inline	`transformer.py`	`forward`	166	GitHub: transformer.py, line 166

Heads

What is it?

HeadsImplemented

Attention

Attention splits each token's vector into H parts of width d_h. The third axis is the only axis the code adds to a sequence, so a head array has the shape H x T x d_h.

Why does the model need it?

Each head computes its own attention weights from its own part of the vector. The split keeps the columns of each head together, so merging them restores the original array.

Mathematics

Width of one head #

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

config.py asserts the divisibility when it is imported.

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
$M$	A matrix to split	`M`	(T, d)
$M^{(h)}$	Head h of the split array	not named in the code	(T, d_h) for each head; (H, T, d_h) together
$d_h$	Width of one head	`dh`	scalar
$H$	Number of heads	`H`	scalar
$h$	Head index, from 0	not named in the code	scalar
$i$	Column index inside a head, from 0	not named in the code	scalar

Tiny numerical example

Illustrative ExampleHand-picked input. Values computed by running transformer.split_heads; they are not learned values.

Input M: [[1, 2, 3, 4], [5, 6, 7, 8]]
Input H: 2
Width of the input d: 4
Split array, head 0 then head 1: [[[1, 2], [5, 6]], [[3, 4], [7, 8]]]

Head 0 holds columns 0 and 1, and head 1 holds columns 2 and 3.

Tensor shapes

Split, for the query matrixtransformer.py:79

Q(T, d)(16, 32)
Qh(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.

Merge, after the weighted valuestransformer.py:84

A @ Vh(H, T, dh)(2, 16, 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.pysplit_headsGitHub

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

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
Width of one head	Computes it	transformer.py:73	`dh = d // H # width of one head`
Width of one head	Written inline	transformer.py:94	`dh = cfg.emb_dim // H`
Width of one head	Written inline	config.py:44	`head_dim = emb_dim // num_heads # width of one head, 32 / 2 = 16`
Head h owns one block of columns	Computes it	transformer.py:62	`return M.reshape(T, H, d // H).transpose(1, 0, 2)`
Head h owns one block of columns	Computes it	transformer.py:67	`return Mh.transpose(1, 0, 2).reshape(T, H * dh)`
Head h owns one block of columns	Tests it	tests/test_transformer.py:61	`Mh = split_heads(M, 2)`

Where is it used?

Source

transformer.py, lines 60-62, split_heads

transformer.py, lines 65-67, merge_heads

PDF reference

Section 3, Eq. 1, p. 3

Used by

Called by: transformer.attention_backward, transformer.attention_forward, transformer.attention_backward, transformer.attention_forward

Tested by: tests.test_transformer.test_split_merge_heads_round_trip, tests.test_transformer.test_split_merge_heads_round_trip

Tests it: tests/test_transformer.py, line 61 (test_split_merge_heads_round_trip), for Head h owns one block of columns, occurrence 0

Status

Implemented

Implemented in

NoteInline equivalent: the pipeline calls transformer.attention_forward, and also computes the same formula in place, with no call, at transformer.py:94, config.py:44.

Equation	Role	File	Function	Lines	GitHub
Width of one head	Computes it	`transformer.py`	`attention_forward`	73	GitHub: transformer.py, line 73
Width of one head	Written inline	`transformer.py`	`attention_backward`	94	GitHub: transformer.py, line 94
Width of one head	Written inline	`config.py`	`head_dim`	44	GitHub: config.py, line 44
Head h owns one block of columns	Computes it	`transformer.py`	`split_heads`	62	GitHub: transformer.py, line 62
Head h owns one block of columns	Computes it	`transformer.py`	`merge_heads`	67	GitHub: transformer.py, line 67

Attention covers what each head computes.

Parameter count

Note

This section is test-only mathematics. count_params is not part of the pipeline: only a test calls it, and training does not print this number.

What is it?

Parameter countImplemented

Initialization

The number of scalars in the arrays init_params creates. The code only ever computes it as a sum of array sizes, in count_params. The PDF states the formula below as Eq. 2. The code never evaluates that formula, and it holds because of the shapes in init_params.

Why does the model need it?

The model does not need this number to run. It tells a reader how the model's size depends on each setting, and only V depends on the corpus. Only tests call count_params.

Mathematics

Add the scalars of every array init_params creates. The names, shapes and values in this table come from calling init_params and from config.py.

Group	Names	Scalars	With config	Repeats
Token table	`tok_emb`	`V × d`	`32V`	Once
Position table	`pos_emb`	`T_max × d`	`1,024`	Once
LayerNorm 1	`b{l}.ln1_g, b{l}.ln1_b`	`d + d`	`32 + 32`	Once per block
Q, K, V and output weights	`b{l}.wq, b{l}.wk, b{l}.wv, b{l}.wo`	`d × d + d × d + d × d + d × d`	`1,024 + 1,024 + 1,024 + 1,024`	Once per block
Q, K, V and output biases	`b{l}.bq, b{l}.bk, b{l}.bv, b{l}.bo`	`d + d + d + d`	`32 + 32 + 32 + 32`	Once per block
LayerNorm 2	`b{l}.ln2_g, b{l}.ln2_b`	`d + d`	`32 + 32`	Once per block
First feed-forward layer	`b{l}.ff1_w, b{l}.ff1_b`	`d × F + F`	`4,096 + 128`	Once per block
Second feed-forward layer	`b{l}.ff2_w, b{l}.ff2_b`	`d × F + d`	`4,096 + 32`	Once per block
Final LayerNorm	`lnf_g, lnf_b`	`d + d`	`32 + 32`	Once
Linear head	`head_w, head_b`	`V × d + V`	`32V + V`	Once
One block, the rows marked once per block			`12,704`

The rows marked once per block add up to one block:

Scalars in one block #

N_{\text{block}} = 4d^2 + 2dF + F + 9d

Multiply that by L, and collect the remaining rows by what they depend on. tok_emb and the head give (2d + 1)V. pos_emb and the final LayerNorm give (T_max + 2)d.

Scalars in the whole model #

N = (2d+1)V + (T_{\max}+2)d + L\left(4d^2 + 2dF + F + 9d\right)

With the values in config.py, and V left as a symbol:

Implemented

N = 65\,V + 26{,}496

What does each symbol mean?

Symbol	Meaning	In code	Shape
$N$	Total scalars	`count_params(params)`	scalar
$N_{\text{block}}$	Scalars one pass of the layer loop creates	not named in the code	scalar
$V$	Vocabulary size	`cfg.vocab_size`	scalar
$T_{\max}$	Rows in the position table	`cfg.context_len`	scalar
$d$	Width of each token's vector	`cfg.emb_dim`	scalar
$F$	Width of the feed-forward hidden layer	`cfg.ff_dim`	scalar
$L$	Number of blocks	`cfg.num_layers`	scalar

Tiny numerical example

The formula and count_params on small configurationsIllustrative ExampleHand-picked configurations. The count_params column comes from running init_params and count_params in Python. These are not this model's parameters, except that the last row uses the shape values in config.py, with the V that tests/test_transformer.py uses.

V	T_max	d	F	L	Formula	count_params
11	7	4	6	3	597	597
13	5	6	10	1	539	539
9	3	2	4	0	55	55
300	32	32	128	2	45,996	45,996

Tensor shapes

Every array's shape is listed on Architecture. The feed-forward pair shows why dF appears twice in the formula.

The two feed-forward weight matrices of one blocktransformer.py:36

ff1_w(d, F)(32, 128)
ff2_w(F, d)(128, 32)

d = emb_dim = 32. F = ff_dim = 128.

Actual implementation

transformer.pycount_paramsGitHub

44def count_params(params):45    return int(sum(v.size for v in params.values()))

transformer.pyinit_paramsGitHub

17def init_params(cfg: ModelConfig, rng):18    d, F, V = cfg.emb_dim, cfg.ff_dim, cfg.vocab_size1920    p = {21        "tok_emb": xavier_uniform(V, d, rng), # lookup table with one row per vocabulary token22        "pos_emb": xavier_uniform(cfg.context_len, d, rng) # second table with one row per position, weighted average of values23        # input to first block is tok_emb + pos_emb24        # this makes cat at first position and fifth position different in a sentence25    }2627    for l in range(cfg.num_layers):28        b = f"b{l}." # key29        p[b + "ln1_g"], p[b + "ln1_b"] = np.ones(d), np.zeros(d) # layer norm3031        for n in "qkvo":32            p[b + "w" + n] = xavier_uniform(d, d, rng)33            p[b + "b" + n] = np.zeros(d)3435        p[b + "ln2_g"], p[b + "ln2_b"] = np.ones(d), np.zeros(d)36        p[b + "ff1_w"], p[b + "ff1_b"] = xavier_uniform(d, F, rng), np.zeros(F)37        p[b + "ff2_w"], p[b + "ff2_b"] = xavier_uniform(F, d, rng), np.zeros(d)3839    p["lnf_g"], p["lnf_b"] = np.ones(d), np.zeros(d)40    p["head_w"], p["head_b"] = xavier_uniform(d, V, rng), np.zeros(V)4142    return p

Code and math connection

Equation	Role	Lines	Code
Scalars in one block	Computes it	transformer.py:29	`p[b + "ln1_g"], p[b + "ln1_b"] = np.ones(d), np.zeros(d) # layer norm`
Scalars in one block	Computes it	transformer.py:32	`p[b + "w" + n] = xavier_uniform(d, d, rng)`
Scalars in one block	Computes it	transformer.py:33	`p[b + "b" + n] = np.zeros(d)`
Scalars in one block	Computes it	transformer.py:35	`p[b + "ln2_g"], p[b + "ln2_b"] = np.ones(d), np.zeros(d)`
Scalars in one block	Computes it	transformer.py:36	`p[b + "ff1_w"], p[b + "ff1_b"] = xavier_uniform(d, F, rng), np.zeros(F)`
Scalars in one block	Computes it	transformer.py:37	`p[b + "ff2_w"], p[b + "ff2_b"] = xavier_uniform(F, d, rng), np.zeros(d)`
Scalars in the whole model	Computes it	transformer.py:45	`return int(sum(v.size for v in params.values()))`
Scalars in the whole model	Computes it	transformer.py:21	`"tok_emb": xavier_uniform(V, d, rng), # lookup table with one row per vocabulary token`
Scalars in the whole model	Computes it	transformer.py:22	`"pos_emb": xavier_uniform(cfg.context_len, d, rng) # second table with one row per position, weighted average of values`
Scalars in the whole model	Computes it	transformer.py:39	`p["lnf_g"], p["lnf_b"] = np.ones(d), np.zeros(d)`
Scalars in the whole model	Computes it	transformer.py:40	`p["head_w"], p["head_b"] = xavier_uniform(d, V, rng), np.zeros(V)`
Scalars in the whole model	Tests it	tests/test_transformer.py:19	`assert count_params(params) == 45_996`

Where is it used?

Source

transformer.py, lines 44-45, count_params

transformer.py, lines 17-42, init_params

PDF reference

Section 3, Eq. 2, p. 3

Used by

Called by: train.main

Tested by: tests.test_transformer.test_param_count_matches_formula, tests.test_checkpoint.saved, tests.test_checkpoint.test_non_ascii_characters_survive, tests.test_generate.tiny, tests.test_gradcheck.tiny_setup, tests.test_transformer.test_attention_weights_are_causal_and_normalised, tests.test_transformer.test_future_tokens_do_not_change_earlier_logits, tests.test_transformer.test_grads_have_the_shape_of_params, tests.test_transformer.test_initial_loss_is_close_to_ln_vocab, tests.test_transformer.test_overfit_one_window, tests.test_transformer.test_param_count_matches_formula

Tests it: tests/test_transformer.py, line 19 (test_param_count_matches_formula), for Scalars in the whole model, occurrence 0

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Scalars in one block	Computes it	`transformer.py`	`init_params`	29	GitHub: transformer.py, line 29
Scalars in one block	Computes it	`transformer.py`	`init_params`	32	GitHub: transformer.py, line 32
Scalars in one block	Computes it	`transformer.py`	`init_params`	33	GitHub: transformer.py, line 33
Scalars in one block	Computes it	`transformer.py`	`init_params`	35	GitHub: transformer.py, line 35
Scalars in one block	Computes it	`transformer.py`	`init_params`	36	GitHub: transformer.py, line 36
Scalars in one block	Computes it	`transformer.py`	`init_params`	37	GitHub: transformer.py, line 37
Scalars in the whole model	Computes it	`transformer.py`	`count_params`	45	GitHub: transformer.py, line 45
Scalars in the whole model	Computes it	`transformer.py`	`init_params`	21	GitHub: transformer.py, line 21
Scalars in the whole model	Computes it	`transformer.py`	`init_params`	22	GitHub: transformer.py, line 22
Scalars in the whole model	Computes it	`transformer.py`	`init_params`	39	GitHub: transformer.py, line 39
Scalars in the whole model	Computes it	`transformer.py`	`init_params`	40	GitHub: transformer.py, line 40