Linear / ReLU / Initialisation

This page covers three building blocks from ops.py: the linear layer, ReLU and Xavier uniform. It also covers how init_params uses them. In the backward sections, dM always means the gradient of the scalar loss with respect to M.

Helper or inline?

The linear layer is written in two ways, so reading only ops.py gives the wrong picture.

Forward: always inline. transformer.py does not import linear_forward, and only test_linear_worked_example calls it. The model writes X @ W + b directly at transformer.py:75 (query), transformer.py:76 (key), transformer.py:77 (value), transformer.py:85 (attention output), transformer.py:113 and transformer.py:116 (feed-forward), and transformer.py:166 (head).
Backward: both. linear_backward is called at transformer.py:119 and transformer.py:123 (feed-forward), and transformer.py:176 (head). attention_backward does not call it. It writes the same three formulas in place at transformer.py:95-97 (attention output) and transformer.py:106-108 (query, key and value).

Linear layer

What is it?

Linear layerImplemented

Basic operations

A linear layer multiplies its input by a weight matrix and adds a bias to every row. The function ops.linear_forward holds this formula, but the model never calls it. transformer.py writes the same expression in place at seven sites.

Why does the model need it?

It is how the model mixes the features of a token vector and moves between widths such as d, F and V. Query, key, value and the attention output are linear layers. So are both feed-forward layers and the head.

Mathematics

The linear layer, with b added to every row of XW #

Y = XW + b

Notation & Shapes

shows how one entry of the product is computed.

What does each symbol mean?

Symbol	Meaning	In code	Shape
$X$	Input, one row per token	`X`	(T, a)
$W$	Weight matrix. A parameter	`W`	(a, c)
$b$	Bias, added to every row by broadcasting. A parameter	`b`	(c,)
$Y$	Output, one row per token. The helper returns it without naming it	not named in the code	(T, c)
$a, c$	Input and output width. They are d, F or V, depending on the layer	not named in the code	scalar

Tiny numerical example

Edit any cell of X, W, b or dY. The results include the backward pass, which the next section explains.

X @ W + b, and its backwardIllustrative 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
t0
t1

W
	j0	j1	j2
k0
k1

b
	j0	j1	j2
b

dY
	j0	j1	j2
t0
t1

(2, 2) @ (2, 3) → (2, 3): 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 × 1 + 2 × -1 = -1

X @ W, shape (2, 3)
-1	2	2
-1	1	0

X @ W + b, shape (2, 3)
-1	3	1
-1	2	-1

dW = X.T @ dY, shape (2, 3)
1	0	1
2	2	1

db = dY.sum(axis=0), shape (1, 3)
1	2	0

dX = dY @ W.T, shape (2, 2)
3	-1
-2	2

Tensor shapes

Attention output projection, width stays the sametransformer.py:85

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

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.

Second feed-forward layer, width shrinks backtransformer.py:116

A(T, F)(16, 128)
output(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. F = ff_dim = 128. d = emb_dim = 32.

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

Where is it used?

The call graph shows no pipeline caller. The inline sites below are the real uses.

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.

Written inline in transformer.py, line 76. Key.

Written inline in transformer.py, line 77. Value.

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 The linear layer, with b added to every row of XW, occurrence 0

Status

Implemented

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
The linear layer, with b added to every row of XW	Computes it	`ops.py`	`linear_forward`	8	GitHub: ops.py, line 8
The linear layer, with b added to every row of XW	Written inline	`transformer.py`	`attention_forward`	75	GitHub: transformer.py, line 75
The linear layer, with b added to every row of XW	Written inline	`transformer.py`	`attention_forward`	76	GitHub: transformer.py, line 76
The linear layer, with b added to every row of XW	Written inline	`transformer.py`	`attention_forward`	77	GitHub: transformer.py, line 77
The linear layer, with b added to every row of XW	Written inline	`transformer.py`	`attention_forward`	85	GitHub: transformer.py, line 85
The linear layer, with b added to every row of XW	Written inline	`transformer.py`	`ffn_forward`	113	GitHub: transformer.py, line 113
The linear layer, with b added to every row of XW	Written inline	`transformer.py`	`ffn_forward`	116	GitHub: transformer.py, line 116
The linear layer, with b added to every row of XW	Written inline	`transformer.py`	`forward`	166	GitHub: transformer.py, line 166

Linear backward

What is it?

Linear layer, backwardImplemented

Basic operations

Given dY, the gradient of the loss on a linear layer's output, this computes three gradients. They are dW for the weights, db for the bias and dX for the input.

Why does the model need it?

dW and db are what the optimizer needs to update this layer's parameters. dX is what the layer before this one receives as its own dY, so the backward pass can continue.

Mathematics

dY holds the gradient of the loss on each output entry. The chain rule multiplies it by the derivative of that output entry, and adds over every entry that depends on the number in question.

The derivatives of one output entry Derived from implementation#

\frac{\partial Y_{tj}}{\partial W_{kj}} = X_{tk},\quad \frac{\partial Y_{tj}}{\partial b_{j}} = 1,\quad \frac{\partial Y_{tj}}{\partial X_{tk}} = W_{kj}

Summing over t gives dW and db. Summing over j gives dX.

Gradient on the weights #

dW = X^{\top} dY

Gradient on the bias, the sum of dY over rows #

db = \sum_{t} dY_{t}

Gradient on the input #

dX = dY\, W^{\top}

What does each symbol mean?

Symbol	Meaning	In code	Shape
$dY$	Gradient of the loss on the output Y. The caller supplies it	`dY`	(T, c)
$X$	The input from the forward pass, kept in a cache	`X`	(T, a)
$W$	The weight matrix	`W`	(a, c)
$dW$	Gradient on the weights	`dW`	(a, c)
$db$	Gradient on the bias	`db`	(c,)
$dX$	Gradient on the input, passed to the previous layer	`dX`	(T, a)

Tiny numerical example

A worked backwardIllustrative ExampleHand-picked numbers from test_linear_worked_example. They are not learned weights, and dY here is chosen by hand, not computed from a loss.

Step	Code	Value
Input	`X`	`[[1.0000, 2.0000], [0.0000, 1.0000]]`
Weights	`W`	`[[1.0000, 0.0000, 2.0000], [-1.0000, 1.0000, 0.0000]]`
Gradient on the output	`dY`	`[[1.0000, 0.0000, 1.0000], [0.0000, 2.0000, -1.0000]]`
Gradient on the weights	`dW = X.T @ dY`	`[[1.0000, 0.0000, 1.0000], [2.0000, 2.0000, 1.0000]]`
Gradient on the bias	`db = dY.sum(axis=0)`	`[1.0000, 2.0000, 0.0000]`
Gradient on the input	`dX = dY @ W.T`	`[[3.0000, -1.0000], [-2.0000, 2.0000]]`

One entry of each: dW[0][0] = 1 × 1 + 0 × 0 = 1, db[0] = 1 + 0 = 1, dX[0][0] = 1 × 1 + 0 × 0 + 1 × 2 = 3.

Checking this 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 (Y * dY).sum(), whose gradient on Y is exactly dY, changes. Tests compare the central differences here with that real function.

Gradient	From the backward formula	From central differences	Largest gap
`dW`	`[1.0000, 0.0000, 1.0000][2.0000, 2.0000, 1.0000]`	`[1.0000, 0.0000, 1.0000][2.0000, 2.0000, 1.0000]`	`3.8e-11`
`db`	`[1.0000, 2.0000, 0.0000]`	`[1.0000, 2.0000, 0.0000]`	`3.8e-11`
`dX`	`[3.0000, -1.0000][-2.0000, 2.0000]`	`[3.0000, -1.0000][-2.0000, 2.0000]`	`3.8e-11`

Tensor shapes

dW has the shape of W, db the shape of b, and dX the shape of X. So the three real calls give:

Linear headtransformer.py:176

dZ(T, V)(16, V)
dW(d, V)(32, V)
db(V,)
dXf(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. V depends on the corpus. It is not a value in config.py. d = emb_dim = 32.

Second feed-forward layertransformer.py:119

dOut(T, d)(16, 32)
dW2(F, d)(128, 32)
db2(d,)(32,)
dA(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.

First feed-forward layertransformer.py:123

dZ(T, F)(16, 128)
dW1(d, F)(32, 128)
db1(F,)(128,)
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. F = ff_dim = 128. d = emb_dim = 32.

Actual implementation

ops.pylinear_backwardGitHub

10def linear_backward(X, W, dY):11    dW = X.T @ dY12    db = dY.sum(axis=0)13    dX = dY @ W.T14    return dX, dW, db

Code and math connection

Equation	Role	Lines	Code
The derivatives of one output entry	Computes it	ops.py:11	`dW = X.T @ dY`
The derivatives of one output entry	Computes it	ops.py:12	`db = dY.sum(axis=0)`
The derivatives of one output entry	Computes it	ops.py:13	`dX = dY @ W.T`
Gradient on the weights	Computes it	ops.py:11	`dW = X.T @ dY`
Gradient on the weights	Written inline	transformer.py:95	`grads[p + "wo"] += c["O"].T @ dOut`
Gradient on the weights	Written inline	transformer.py:106	`grads[p + "w" + n] += c["Xn"].T @ dM`
Gradient on the weights	Calls it	transformer.py:119	`dA, dW2, db2 = linear_backward(c["A"], params[p + "ff2_w"], dOut)`
Gradient on the weights	Calls it	transformer.py:123	`dXn, dW1, db1 = linear_backward(c["Xn"], params[p + "ff1_w"], dZ)`
Gradient on the weights	Calls it	transformer.py:176	`dXf, dW, db = linear_backward(cache["Xf"], params["head_w"], dZ)`
Gradient on the weights	Tests it	tests/test_ops.py:16	`dX, dW, db = linear_backward(X, W, dY)`
Gradient on the bias, the sum of dY over rows	Computes it	ops.py:12	`db = dY.sum(axis=0)`
Gradient on the bias, the sum of dY over rows	Written inline	transformer.py:96	`grads[p + "bo"] += dOut.sum(axis=0)`
Gradient on the bias, the sum of dY over rows	Written inline	transformer.py:107	`grads[p + "b" + n] += dM.sum(axis=0)`
Gradient on the input	Computes it	ops.py:13	`dX = dY @ W.T`
Gradient on the input	Written inline	transformer.py:97	`dOh = split_heads(dOut @ params[p + "wo"].T, H)`
Gradient on the input	Written inline	transformer.py:108	`dXn += dM @ params[p + "w" + n].T`

Where is it used?

Source

ops.py, lines 10-14, linear_backward

PDF reference

Section 4.3, Eq. 20, 21, 22, p. 6

Used by

Called by: transformer.ffn_backward, transformer.loss_and_grads

Tested by: tests.test_ops.test_linear_worked_example

Written inline in transformer.py, line 95. dW for the attention output.

Written inline in transformer.py, line 96. db for the attention output.

Written inline in transformer.py, line 97. dX for the attention output, before it is split into heads.

Written inline in transformer.py, line 106. dW for query, key and value, in a loop.

Written inline in transformer.py, line 107. db for query, key and value.

Written inline in transformer.py, line 108. dX for query, key and value, added together.

Calls it: transformer.py, line 119 (ffn_backward), for Gradient on the weights
Calls it: transformer.py, line 123 (ffn_backward), for Gradient on the weights
Calls it: transformer.py, line 176 (loss_and_grads), for Gradient on the weights
Tests it: tests/test_ops.py, line 16 (test_linear_worked_example), for Gradient on the weights, occurrence 0

Status

Implemented

Implemented in

NoteInline equivalent: the pipeline calls ops.linear_backward, and also computes the same formula in place, with no call, at transformer.py:95, transformer.py:106.

NoteInline equivalent: the pipeline calls ops.linear_backward, and also computes the same formula in place, with no call, at transformer.py:96, transformer.py:107.

NoteInline equivalent: the pipeline calls ops.linear_backward, and also computes the same formula in place, with no call, at transformer.py:97, transformer.py:108.

Equation	Role	File	Function	Lines	GitHub
The derivatives of one output entry	Computes it	`ops.py`	`linear_backward`	11	GitHub: ops.py, line 11
The derivatives of one output entry	Computes it	`ops.py`	`linear_backward`	12	GitHub: ops.py, line 12
The derivatives of one output entry	Computes it	`ops.py`	`linear_backward`	13	GitHub: ops.py, line 13
Gradient on the weights	Computes it	`ops.py`	`linear_backward`	11	GitHub: ops.py, line 11
Gradient on the weights	Written inline	`transformer.py`	`attention_backward`	95	GitHub: transformer.py, line 95
Gradient on the weights	Written inline	`transformer.py`	`attention_backward`	106	GitHub: transformer.py, line 106
Gradient on the bias, the sum of dY over rows	Computes it	`ops.py`	`linear_backward`	12	GitHub: ops.py, line 12
Gradient on the bias, the sum of dY over rows	Written inline	`transformer.py`	`attention_backward`	96	GitHub: transformer.py, line 96
Gradient on the bias, the sum of dY over rows	Written inline	`transformer.py`	`attention_backward`	107	GitHub: transformer.py, line 107
Gradient on the input	Computes it	`ops.py`	`linear_backward`	13	GitHub: ops.py, line 13
Gradient on the input	Written inline	`transformer.py`	`attention_backward`	97	GitHub: transformer.py, line 97
Gradient on the input	Written inline	`transformer.py`	`attention_backward`	108	GitHub: transformer.py, line 108

How this reaches a parameter update. The caller adds dW and db to grads, for example at transformer.py:178 and transformer.py:125. batch_loss_and_grads averages grads over the batch (train.py:61). clip_grad_norm may scale it (train.py:62), and opt.step uses it to change the parameter (train.py:66). dX is not stored. It goes to the previous layer as its dY. How the Model Learns follows the whole path.

ReLU

What is it?

ReLUImplemented

Basic operations

ReLU replaces every negative number with 0 and keeps the others. It works element by element, so the shape does not change.

Why does the model need it?

It is the only nonlinear step in the feed-forward layer. In general, two linear layers with nothing between them multiply out to a single linear layer.

Mathematics

ReLU, on each entry #

A_{ti} = \max(0,\, Z_{ti})

What does each symbol mean?

Symbol	Meaning	In code	Shape
$Z$	Input. relu_forward calls it X, and ffn_forward passes its Z	`X in relu_forward, Z in ffn_forward`	(T, F)
$A$	Output. ffn_forward calls it A. It is not attention weights here	`A in ffn_forward`	(T, F)

Tiny numerical example

ReLUIllustrative ExampleHand-picked numbers from test_relu. The third value is exactly 0, which is the case the strict > decides.

Step	Code	Value
Input	`X`	`[[-1.0000, 2.0000, 0.0000, 3.0000]]`
Output	`np.maximum(0.0, X)`	`[[0.0000, 2.0000, 0.0000, 3.0000]]`

Tensor shapes

Inside the feed-forward layertransformer.py:114

Z(T, F)(16, 128)
A(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. F = ff_dim = 128.

Actual implementation

ops.pyrelu_forwardGitHub

16def relu_forward(X):17    return np.maximum(0.0, X)

Code and math connection

Equation	Role	Lines	Code
ReLU, on each entry	Computes it	ops.py:17	`return np.maximum(0.0, X)`
ReLU, on each entry	Calls it	transformer.py:114	`A = relu_forward(Z)`
ReLU, on each entry	Tests it	tests/test_ops.py:24	`assert np.allclose(relu_forward(Z), [0, 2, 0, 3])`

Where is it used?

Source

ops.py, lines 16-17, relu_forward

PDF reference

Section 4.3, Eq. 29, p. 7

Used by

Called by: transformer.ffn_forward

Tested by: tests.test_ops.test_relu

Calls it: transformer.py, line 114 (ffn_forward), for ReLU, on each entry
Tests it: tests/test_ops.py, line 24 (test_relu), for ReLU, on each entry

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
ReLU, on each entry	Computes it	`ops.py`	`relu_forward`	17	GitHub: ops.py, line 17

ReLU backward

What is it?

ReLU, backwardImplemented

Basic operations

Given dA, the gradient on ReLU's output, this computes dZ, the gradient on its input. It keeps dA where the input was positive and sets every other entry to 0.

Why does the model need it?

The backward pass has to cross ReLU to reach the first feed-forward layer. Where the input is negative, a small change to it does not change the output, so no gradient passes back.

Mathematics

The derivative as the code uses it, with 0 at exactly Z = 0 Derived from implementation#

\frac{\partial A_{ti}}{\partial Z_{ti}} = \begin{cases} 1 & Z_{ti} > 0 \\ 0 & Z_{ti} \le 0 \end{cases}

The mask-based backward #

dZ = dA \odot \mathbf{1}[Z > 0]

What does each symbol mean?

Symbol	Meaning	In code	Shape
$dA$	Gradient of the loss on ReLU's output	`dA`	(T, F)
$Z$	ReLU's input from the forward pass, kept in the cache	`Z`	(T, F)
$\mathbf{1}[Z > 0]$	The mask. NumPy multiplies its True and False as 1 and 0	`(Z > 0)`	(T, F)
$dZ$	Gradient on ReLU's input. The function returns it without naming it	`dZ in ffn_backward`	(T, F)

Tiny numerical example

ReLU backwardIllustrative ExampleHand-picked numbers from test_relu. The third value is exactly 0, which is the case the strict > decides.

Step	Code	Value
Input from the forward pass	`Z`	`[[-1.0000, 2.0000, 0.0000, 3.0000]]`
Gradient on the output	`dA`	`[[1.0000, 1.0000, 1.0000, 1.0000]]`
Mask	`Z > 0`	`[[0.0000, 1.0000, 0.0000, 1.0000]]`
Gradient on the input	`dA * (Z > 0)`	`[[0.0000, 1.0000, 0.0000, 1.0000]]`

Tensor shapes

Inside the feed-forward backwardtransformer.py:122

dA(T, F)(16, 128)
dZ(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. F = ff_dim = 128.

Actual implementation

ops.pyrelu_backwardGitHub

19def relu_backward(Z, dA):20    return dA * (Z > 0)

Code and math connection

Equation	Role	Lines	Code
The derivative as the code uses it, with 0 at exactly Z = 0	Computes it	ops.py:20	`return dA * (Z > 0)`
The mask-based backward	Computes it	ops.py:20	`return dA * (Z > 0)`
The mask-based backward	Calls it	transformer.py:122	`dZ = relu_backward(c["Z"], dA)`
The mask-based backward	Tests it	tests/test_ops.py:25	`assert np.allclose(relu_backward(Z, np.ones(4)), [0, 1, 0, 1])`

Where is it used?

Source

ops.py, lines 19-20, relu_backward

PDF reference

Section 4.3, Eq. 29, p. 7

Used by

Called by: transformer.ffn_backward

Tested by: tests.test_ops.test_relu

Calls it: transformer.py, line 122 (ffn_backward), for The mask-based backward
Tests it: tests/test_ops.py, line 25 (test_relu), for The mask-based backward

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
The derivative as the code uses it, with 0 at exactly Z = 0	Computes it	`ops.py`	`relu_backward`	20	GitHub: ops.py, line 20
The mask-based backward	Computes it	`ops.py`	`relu_backward`	20	GitHub: ops.py, line 20

How this reaches a parameter update. ReLU has no parameters, so nothing is stored here. dZ (transformer.py:122) goes into linear_backward for the first feed-forward layer (transformer.py:123), which produces the gradients of ff1_w and ff1_b. Those follow the path described under Linear backward.

Xavier uniform

What is it?

Xavier uniform initialisationImplemented

Initialization

Fills a new weight matrix with random numbers drawn evenly between -a and a, where a depends on the matrix's two sizes.

Why does the model need it?

In general, weights need random starting values, or every unit in a layer would compute the same thing, and a limit like this one aims to keep signals at a similar size as they pass through a layer. This repository does not measure that.

Mathematics

The limit #

a = \sqrt{\frac{6}{\text{fan\_in} + \text{fan\_out}}}

Each entry is drawn evenly between -a and a #

W_{ij} \sim U[-a,\, a]

The variance this gives. The code sets a and never computes this #

\operatorname{Var}(W_{ij}) = \frac{a^2}{3} = \frac{2}{\text{fan\_in} + \text{fan\_out}}

What does each symbol mean?

Symbol	Meaning	In code	Shape
$\text{fan\_in}$	Rows of the matrix	`fan_in`	scalar
$\text{fan\_out}$	Columns of the matrix	`fan_out`	scalar
$a$	The limit	`a`	scalar
$W$	The new matrix. The function returns it without naming it	not named in the code	(fan_in, fan_out)
$U$	The uniform distribution, drawn by the generator passed in	`rng.uniform`	(fan_in, fan_out)

Tiny numerical example

tok_emb and head_w have no row here, because their size depends on V, and V depends on the corpus. Larger matrices get a smaller limit.

The limit for each matrixIllustrative ExampleThe sizes after the first row come from config.py. The limit is computed by the math engine. The lowest and highest values are stored results of running the real xavier_uniform with a generator seeded from config.py, and they are not the values of a trained or real model.

Matrix	fan_in	fan_out	Limit a	Lowest drawn	Highest drawn
`Tiny, hand-picked`	2	3	1.0954	-0.8891	1.0420
`w{q,k,v,o}`	32	32	0.3062	-0.3054	0.3056
`ff1_w`	32	128	0.1936	-0.1934	0.1936
`ff2_w`	128	32	0.1936	-0.1934	0.1936
`pos_emb`	32	32	0.3062	-0.3054	0.3056

The tiny matrix itself, shape (2, 3)
0.6002	-0.1339	0.7856
0.4324	-0.8891	1.042

Tensor shapes

The output has shape (fan_in, fan_out), so these two arguments are the matrix's rows and columns.

A query, key, value or output matrixtransformer.py:32

fan_in, fan_out(d, d)(32, 32)
w{q,k,v,o}(d, d)(32, 32)

d = emb_dim = 32.

The first feed-forward matrixtransformer.py:36

fan_in, fan_out(d, F)(32, 128)
ff1_w(d, F)(32, 128)

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

Actual implementation

ops.pyxavier_uniformGitHub

3def xavier_uniform(fan_in, fan_out, rng):4    a = np.sqrt(6.0 / (fan_in + fan_out))5    return rng.uniform(-a, a, size=(fan_in, fan_out))

Code and math connection

Equation	Role	Lines	Code
The limit	Computes it	ops.py:4	`a = np.sqrt(6.0 / (fan_in + fan_out))`
The limit	Tests it	tests/test_ops.py:84	`assert np.abs(W).max() <= np.sqrt(6 / 64)`
Each entry is drawn evenly between -a and a	Computes it	ops.py:5	`return rng.uniform(-a, a, size=(fan_in, fan_out))`
The variance this gives. The code sets a and never computes this	Computes it	ops.py:4	`a = np.sqrt(6.0 / (fan_in + fan_out))`

Where is it used?

Only init_params calls it, and only train.main calls init_params, so it runs only at the start of training. talk.py loads its weights from a checkpoint instead.

Source

ops.py, lines 3-5, xavier_uniform

PDF reference

Section 4.3, Eq. 35, 36, p. 7

Used by

Called by: transformer.init_params

Tested by: tests.test_ops.test_xavier_bound

Tests it: tests/test_ops.py, line 84 (test_xavier_bound), for The limit

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
The limit	Computes it	`ops.py`	`xavier_uniform`	4	GitHub: ops.py, line 4
Each entry is drawn evenly between -a and a	Computes it	`ops.py`	`xavier_uniform`	5	GitHub: ops.py, line 5
The variance this gives. The code sets a and never computes this	Computes it	`ops.py`	`xavier_uniform`	4	GitHub: ops.py, line 4

Starting values

What is it?

Starting values of the parametersImplemented

Initialization

init_params creates every parameter. Each matrix starts from Xavier uniform. Each bias starts at zeros. Each LayerNorm starts with gamma at ones and beta at zeros.

Why does the model need it?

Training has to start from somewhere. Zero biases add nothing at first. With gamma at ones and beta at zeros, LayerNorm's output starts out equal to the normalised input.

Mathematics

Starting values, by kind of parameter #

W \sim \text{Xavier uniform},\quad b = \mathbf{0},\quad \gamma = \mathbf{1},\quad \beta = \mathbf{0}

This table comes from running init_params. b{l} stands for each block.

Starts as	Parameters
ones	`b{l}.ln1_g`, `b{l}.ln2_g`, `lnf_g`
zeros	`b{l}.ln1_b`, `b{l}.bq`, `b{l}.bk`, `b{l}.bv`, `b{l}.bo`, `b{l}.ln2_b`, `b{l}.ff1_b`, `b{l}.ff2_b`, `lnf_b`, `head_b`
Xavier uniform	`tok_emb`, `pos_emb`, `b{l}.wq`, `b{l}.wk`, `b{l}.wv`, `b{l}.wo`, `b{l}.ff1_w`, `b{l}.ff2_w`, `head_w`

What does each symbol mean?

Symbol	Meaning	In code	Shape
$W$	Any matrix, including both embedding tables and the head	`tok_emb, pos_emb, w{q,k,v,o}, ff1_w, ff2_w, head_w`	two axes
$b$	Any bias of a linear layer	`b{q,k,v,o}, ff1_b, ff2_b, head_b`	one axis
$\gamma$	LayerNorm scale	`ln1_g, ln2_g, lnf_g`	(d,)
$\beta$	LayerNorm shift	`ln1_b, ln2_b, lnf_b`	(d,)

Tiny numerical example

What init_params returnedIllustrative ExampleStored result of running the real init_params with hand-picked sizes: V = 5, context_len = 3, emb_dim = 4, num_heads = 2, ff_dim = 6, num_layers = 1. These are not the model's values.

Name	Shape	Lowest	Highest	Xavier limit
`tok_emb`	`(5, 4)`	-0.7123	0.7767	0.8165
`pos_emb`	`(3, 4)`	-0.8447	0.8716	0.9258
`b0.ln1_g`	`(4,)`	1.0000	1.0000	does not apply
`b0.ln1_b`	`(4,)`	0.0000	0.0000	does not apply
`b0.wq`	`(4, 4)`	-0.6410	0.5762	0.8660
`b0.bq`	`(4,)`	0.0000	0.0000	does not apply
`b0.wk`	`(4, 4)`	-0.8533	0.4970	0.8660
`b0.bk`	`(4,)`	0.0000	0.0000	does not apply
`b0.wv`	`(4, 4)`	-0.8126	0.6121	0.8660
`b0.bv`	`(4,)`	0.0000	0.0000	does not apply
`b0.wo`	`(4, 4)`	-0.8267	0.5439	0.8660
`b0.bo`	`(4,)`	0.0000	0.0000	does not apply
`b0.ln2_g`	`(4,)`	1.0000	1.0000	does not apply
`b0.ln2_b`	`(4,)`	0.0000	0.0000	does not apply
`b0.ff1_w`	`(4, 6)`	-0.6388	0.7268	0.7746
`b0.ff1_b`	`(6,)`	0.0000	0.0000	does not apply
`b0.ff2_w`	`(6, 4)`	-0.7393	0.7104	0.7746
`b0.ff2_b`	`(4,)`	0.0000	0.0000	does not apply
`lnf_g`	`(4,)`	1.0000	1.0000	does not apply
`lnf_b`	`(4,)`	0.0000	0.0000	does not apply
`head_w`	`(4, 5)`	-0.7812	0.6793	0.8165
`head_b`	`(5,)`	0.0000	0.0000	does not apply

Tensor shapes

Architecture

lists the shape of every parameter. For example, the head has one column per vocabulary id, and it is a separate matrix from tok_emb:

The head's weights and biastransformer.py:40

head_w(d, V)(32, V)
head_b(V,)

d = emb_dim = 32. V depends on the corpus. It is not a value in config.py.

Actual implementation

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
Starting values, by kind of parameter	Computes it	transformer.py:21	`"tok_emb": xavier_uniform(V, d, rng), # lookup table with one row per vocabulary token`
Starting values, by kind of parameter	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`
Starting values, by kind of parameter	Computes it	transformer.py:29	`p[b + "ln1_g"], p[b + "ln1_b"] = np.ones(d), np.zeros(d) # layer norm`
Starting values, by kind of parameter	Computes it	transformer.py:32	`p[b + "w" + n] = xavier_uniform(d, d, rng)`
Starting values, by kind of parameter	Computes it	transformer.py:33	`p[b + "b" + n] = np.zeros(d)`
Starting values, by kind of parameter	Computes it	transformer.py:35	`p[b + "ln2_g"], p[b + "ln2_b"] = np.ones(d), np.zeros(d)`
Starting values, by kind of parameter	Computes it	transformer.py:36	`p[b + "ff1_w"], p[b + "ff1_b"] = xavier_uniform(d, F, rng), np.zeros(F)`
Starting values, by kind of parameter	Computes it	transformer.py:37	`p[b + "ff2_w"], p[b + "ff2_b"] = xavier_uniform(F, d, rng), np.zeros(d)`
Starting values, by kind of parameter	Computes it	transformer.py:39	`p["lnf_g"], p["lnf_b"] = np.ones(d), np.zeros(d)`
Starting values, by kind of parameter	Computes it	transformer.py:40	`p["head_w"], p["head_b"] = xavier_uniform(d, V, rng), np.zeros(V)`

Where is it used?

Source

transformer.py, lines 17-42, init_params

PDF reference

Section 4.5, p. 9

Used by

Called by: train.main

Tested by: 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

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Starting values, by kind of parameter	Computes it	`transformer.py`	`init_params`	21	GitHub: transformer.py, line 21
Starting values, by kind of parameter	Computes it	`transformer.py`	`init_params`	22	GitHub: transformer.py, line 22
Starting values, by kind of parameter	Computes it	`transformer.py`	`init_params`	29	GitHub: transformer.py, line 29
Starting values, by kind of parameter	Computes it	`transformer.py`	`init_params`	32	GitHub: transformer.py, line 32
Starting values, by kind of parameter	Computes it	`transformer.py`	`init_params`	33	GitHub: transformer.py, line 33
Starting values, by kind of parameter	Computes it	`transformer.py`	`init_params`	35	GitHub: transformer.py, line 35
Starting values, by kind of parameter	Computes it	`transformer.py`	`init_params`	36	GitHub: transformer.py, line 36
Starting values, by kind of parameter	Computes it	`transformer.py`	`init_params`	37	GitHub: transformer.py, line 37
Starting values, by kind of parameter	Computes it	`transformer.py`	`init_params`	39	GitHub: transformer.py, line 39
Starting values, by kind of parameter	Computes it	`transformer.py`	`init_params`	40	GitHub: transformer.py, line 40