Transformer Block

A block has two branches: attention and a feed-forward network. This page covers the feed-forward network, how block_forward joins the two branches, how block_backward reverses that, and what the forward pass leaves behind for the backward pass.

The numbers on this page are labelled Illustrative Example. The tiny model below has random parameters from the real init_params, not learned values. The real sizes appear under Tensor shapes.

Feed-forward network

What is it?

Feed-forward network, forwardImplemented

Feed-forward and block

ffn_forward widens each row with one linear layer, applies ReLU, and narrows the row back to its original width with a second linear layer. It returns the output and a cache holding Xn, Z and A.

Why does the model need it?

Attention builds each row from rows at that position and earlier ones, and its output projection is linear. The feed-forward layer works on each row alone, and ReLU is the step in it that is not linear.

Note

In general, two linear layers with no step between them can be replaced by one linear layer, so a function that is not linear between them is what makes the second layer add anything. This is a general explanation. Nothing in this repository measures it.

Mathematics

Widen #

Z = X_n W_1 + \mathbf{1} b_1^\top

ReLU, element by element #

A = \operatorname{ReLU}(Z)

Narrow #

\mathrm{Ff} = A W_2 + \mathbf{1} b_2^\top

For one row, the three steps together are:

The same three steps for one row x #

\operatorname{FFN}(x) = \operatorname{ReLU}(x W_1 + b_1)\, W_2 + b_2

What does each symbol mean?

Symbol	Meaning	In code	Shape
$X_n$	The input, a normalised copy of the residual stream. Each block passes the output of its second LayerNorm	`Xn`	(T, d)
$W_1, b_1$	Weights and bias of the first linear layer	`ff1_w, ff1_b`	(d, F) and (F,)
$Z$	The widened rows, before ReLU	`Z`	(T, F)
$A$	The widened rows after ReLU. On this page A means this array and nothing else	`A`	(T, F)
$W_2, b_2$	Weights and bias of the second linear layer	`ff2_w, ff2_b`	(F, d) and (d,)
$\mathrm{Ff}$	The output of the feed-forward layer. ffn_forward does not name it, and block_forward calls it ff	`ff`	(T, d)
$F$	The width of the hidden rows	`cfg.ff_dim`	scalar
$x$	One row of X_n	`Xn[t]`	(d,)

Tiny numerical example

One row through the feed-forward layerIllustrative ExampleThe numbers of test_ffn_worked_example: two features widen to three and narrow to two. The model's sizes are in the next step.

Step	Code	Value
Input	`Xn`	`[[1.0000, -1.0000]]`
Widen	`Z = Xn @ params[p + "ff1_w"] + params[p + "ff1_b"]`	`[[1.0000, -0.5000, -1.0000]]`
ReLU	`A = relu_forward(Z)`	`[[1.0000, 0.0000, 0.0000]]`
Narrow	`A @ params[p + "ff2_w"] + params[p + "ff2_b"]`	`[[1.0000, 2.0000]]`

Tensor shapes

With config.py, each row widens from 32 to 128 numbers and narrows back to 32.

Widen, ReLU, narrowtransformer.py:113

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

Actual implementation

transformer.pyffn_forwardGitHub

112def ffn_forward(Xn, params, p):113    Z = Xn @ params[p + "ff1_w"] + params[p + "ff1_b"]114    A = relu_forward(Z)115116    return A @ params[p + "ff2_w"] + params[p + "ff2_b"], dict(Xn=Xn, Z=Z, A=A)

init_params creates the two layers' parameters. Both weight matrices start as Xavier uniform draws, and both biases start at zero:

transformer.pyinit_paramsGitHub

36        p[b + "ff1_w"], p[b + "ff1_b"] = xavier_uniform(d, F, rng), np.zeros(F)

transformer.pyinit_paramsGitHub

37        p[b + "ff2_w"], p[b + "ff2_b"] = xavier_uniform(F, d, rng), np.zeros(d)

Code and math connection

Equation	Role	Lines	Code
Widen	Written inline	transformer.py:113	`Z = Xn @ params[p + "ff1_w"] + params[p + "ff1_b"]`
ReLU, element by element	Computes it	transformer.py:114	`A = relu_forward(Z)`
Narrow	Written inline	transformer.py:116	`return A @ params[p + "ff2_w"] + params[p + "ff2_b"], dict(Xn=Xn, Z=Z, A=A)`
The same three steps for one row x	Written inline	transformer.py:113	`Z = Xn @ params[p + "ff1_w"] + params[p + "ff1_b"]`
The same three steps for one row x	Computes it	transformer.py:114	`A = relu_forward(Z)`
The same three steps for one row x	Written inline	transformer.py:116	`return A @ params[p + "ff2_w"] + params[p + "ff2_b"], dict(Xn=Xn, Z=Z, A=A)`
The same three steps for one row x	Tests it	tests/test_transformer.py:56	`assert np.allclose(y, [[1, 2]])`

Where is it used?

Source

transformer.py, lines 112-116, ffn_forward

PDF reference

Section 4.5, Eq. 51, p. 11

Used by

Called by: transformer.block_forward

Tested by: tests.test_transformer.test_ffn_worked_example

Tests it: tests/test_transformer.py, line 56 (test_ffn_worked_example), for The same three steps for one row x

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Widen	Written inline	`transformer.py`	`ffn_forward`	113	GitHub: transformer.py, line 113
ReLU, element by element	Computes it	`transformer.py`	`ffn_forward`	114	GitHub: transformer.py, line 114
Narrow	Written inline	`transformer.py`	`ffn_forward`	116	GitHub: transformer.py, line 116
The same three steps for one row x	Computes it	`transformer.py`	`ffn_forward`	114	GitHub: transformer.py, line 114
The same three steps for one row x	Written inline	`transformer.py`	`ffn_forward`	113	GitHub: transformer.py, line 113
The same three steps for one row x	Written inline	`transformer.py`	`ffn_forward`	116	GitHub: transformer.py, line 116

Feed-forward backward

What is it?

Feed-forward network, backwardImplemented

Feed-forward and block

ffn_backward receives dOut, the gradient of the loss on the output of the feed-forward layer. It goes back through the second linear layer, through ReLU, and through the first linear layer. It adds four parameter gradients to grads and returns dXn, the gradient on the input.

Why does the model need it?

W_1, b_1, W_2 and b_2 are parameters, so the optimizer needs a gradient for each. dXn continues to the second LayerNorm's backward, and from there to the earlier layers.

Mathematics

Each equation below follows the order of ffn_backward. The second linear layer comes first. linear_backward returns three results for it:

Gradient on the ReLU output #

dA = d\mathrm{Out}\, W_2^\top

Gradient on the second weights #

dW_2 = A^\top d\mathrm{Out}

Gradient on the second bias #

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

ReLU has no parameters. A hidden value that ReLU set to 0 sends no gradient back:

Through ReLU #

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

Then the first linear layer:

Gradient on the input #

dX_n = dZ\, W_1^\top

Gradient on the first weights #

dW_1 = X_n^\top dZ

Gradient on the first bias #

db_1 = \mathbf{1}^\top dZ

What does each symbol mean?

Symbol	Meaning	In code	Shape
$d\mathrm{Out}$	Gradient of the loss on the output of the feed-forward layer. block_backward passes dY	`dOut`	(T, d)
$dA$	Gradient on the ReLU output	`dA`	(T, F)
$dZ$	Gradient on the widened rows, before ReLU	`dZ`	(T, F)
$dX_n$	Gradient on the input, returned to block_backward	`dXn`	(T, d)
$dW_1, db_1$	Gradients on the first linear layer	`dW1, db1`	(d, F) and (F,)
$dW_2, db_2$	Gradients on the second linear layer	`dW2, db2`	(F, d) and (d,)
$A, Z, X_n$	The arrays the forward pass left in its cache	`c["A"], c["Z"], c["Xn"]`	(T, F), (T, F) and (T, d)
$\mathbb{1}[Z > 0]$	1 where Z is positive and 0 elsewhere, including at Z = 0	`(Z > 0)`	(T, F)
$\mathbf{1}^\top$	A sum over rows	`.sum(axis=0)`	(T,)
$\odot$	Multiplication element by element	`*`	none

Tiny numerical example

The same row, backwardIllustrative ExampleThe same forward pass as above. dOut is hand-picked, not computed from a loss. Two entries of Z are negative, so dZ is 0 there.

Step	Code	Value
Arrives	`dOut`	`[[1.0000, 0.5000]]`
Through the second layer	`dA`	`[[2.0000, 0.5000, 0.5000]]`
Through ReLU	`dZ`	`[[2.0000, 0.0000, 0.0000]]`
Through the first layer	`dXn`	`[[2.0000, 0.0000]]`

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 + "ff2_w"]`	`[1.0000, 0.5000][0.0000, 0.0000][0.0000, 0.0000]`	`[1.0000, 0.5000][0.0000, 0.0000][0.0000, 0.0000]`	`7.8e-12`
`grads[p + "ff2_b"]`	`[1.0000, 0.5000]`	`[1.0000, 0.5000]`	`7.8e-12`
`grads[p + "ff1_w"]`	`[2.0000, 0.0000, 0.0000][-2.0000, 0.0000, 0.0000]`	`[2.0000, 0.0000, 0.0000][-2.0000, 0.0000, 0.0000]`	`2.0e-12`
`grads[p + "ff1_b"]`	`[2.0000, 0.0000, 0.0000]`	`[2.0000, 0.0000, 0.0000]`	`2.0e-12`
`dXn`	`[2.0000, 0.0000]`	`[2.0000, 0.0000]`	`2.0e-12`

Tensor shapes

Back through the three stepstransformer.py:122

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

Parameter gradients of the second layertransformer.py:119

dW2(F, d)(128, 32)
db2(d,)(32,)

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

Parameter gradients of the first layertransformer.py:123

dW1(d, F)(32, 128)
db1(F,)(128,)

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

Actual implementation

transformer.pyffn_backwardGitHub

118def ffn_backward(dOut, c, params, p, grads):119    dA, dW2, db2 = linear_backward(c["A"], params[p + "ff2_w"], dOut)120    grads[p + "ff2_w"] += dW2121    grads[p + "ff2_b"] += db2122    dZ = relu_backward(c["Z"], dA)123    dXn, dW1, db1 = linear_backward(c["Xn"], params[p + "ff1_w"], dZ)124125    grads[p + "ff1_w"] += dW1126    grads[p + "ff1_b"] += db1127128    return dXn

Code and math connection

Equation	Role	Lines	Code
Gradient on the ReLU output	Computes it	transformer.py:119	`dA, dW2, db2 = linear_backward(c["A"], params[p + "ff2_w"], dOut)`
Gradient on the second weights	Computes it	transformer.py:119	`dA, dW2, db2 = linear_backward(c["A"], params[p + "ff2_w"], dOut)`
Gradient on the second weights	Computes it	transformer.py:120	`grads[p + "ff2_w"] += dW2`
Gradient on the second bias	Computes it	transformer.py:119	`dA, dW2, db2 = linear_backward(c["A"], params[p + "ff2_w"], dOut)`
Gradient on the second bias	Computes it	transformer.py:121	`grads[p + "ff2_b"] += db2`
Through ReLU	Computes it	transformer.py:122	`dZ = relu_backward(c["Z"], dA)`
Gradient on the input	Computes it	transformer.py:123	`dXn, dW1, db1 = linear_backward(c["Xn"], params[p + "ff1_w"], dZ)`
Gradient on the first weights	Computes it	transformer.py:123	`dXn, dW1, db1 = linear_backward(c["Xn"], params[p + "ff1_w"], dZ)`
Gradient on the first weights	Computes it	transformer.py:125	`grads[p + "ff1_w"] += dW1`
Gradient on the first bias	Computes it	transformer.py:123	`dXn, dW1, db1 = linear_backward(c["Xn"], params[p + "ff1_w"], dZ)`
Gradient on the first bias	Computes it	transformer.py:126	`grads[p + "ff1_b"] += db1`

Where is it used?

block_backward calls it first, with dY (transformer.py:143). No test calls ffn_backward directly. test_every_tensor_passes reaches it through the whole model, and compares the four gradients with central differences.

Source: transformer.py, lines 118-128, ffn_backward
PDF reference: Section 4.5, Eq. 52, p. 11
Used by: Called by: transformer.block_backward
Status: Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Gradient on the ReLU output	Computes it	`transformer.py`	`ffn_backward`	119	GitHub: transformer.py, line 119
Gradient on the second weights	Computes it	`transformer.py`	`ffn_backward`	119	GitHub: transformer.py, line 119
Gradient on the second weights	Computes it	`transformer.py`	`ffn_backward`	120	GitHub: transformer.py, line 120
Gradient on the second bias	Computes it	`transformer.py`	`ffn_backward`	119	GitHub: transformer.py, line 119
Gradient on the second bias	Computes it	`transformer.py`	`ffn_backward`	121	GitHub: transformer.py, line 121
Through ReLU	Computes it	`transformer.py`	`ffn_backward`	122	GitHub: transformer.py, line 122
Gradient on the input	Computes it	`transformer.py`	`ffn_backward`	123	GitHub: transformer.py, line 123
Gradient on the first weights	Computes it	`transformer.py`	`ffn_backward`	123	GitHub: transformer.py, line 123
Gradient on the first weights	Computes it	`transformer.py`	`ffn_backward`	125	GitHub: transformer.py, line 125
Gradient on the first bias	Computes it	`transformer.py`	`ffn_backward`	123	GitHub: transformer.py, line 123
Gradient on the first bias	Computes it	`transformer.py`	`ffn_backward`	126	GitHub: transformer.py, line 126

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). ff1_w and ff2_w also get weight decay, and the biases do not. dXn goes on to the second LayerNorm's backward (transformer.py:145).

The pre-LN block

What is it?

Pre-LN block, forwardImplemented

Feed-forward and block

block_forward runs one transformer block on the sequence X. Each of its two branches reads a normalised copy of the residual stream, and adds its result to the stream as it was before that normalisation. It returns the output and a dictionary of four caches.

This form is called pre-LN because of where the LayerNorm sits in the code. Each one runs before its sub-layer: the first on the line before attention (transformer.py:133, transformer.py:134), and the second on the line before the feed-forward layer (transformer.py:136, transformer.py:137).

Why does the model need it?

The block is where one layer of the model does its work. Attention moves information between positions, the feed-forward layer works on each row alone, and the two adds carry the stream past both of them. forward calls this once per layer.

Note

In general, placing normalisation before a sub-layer keeps the residual path itself free of normalisation, and this is often reported to make deep networks easier to train. This is a general explanation. Nothing in this repository measures that effect.

Mathematics

One equation per statement of block_forward, in the same order:

First LayerNorm, on a copy #

X_n = \operatorname{LN}_1(X)

Attention reads the copy #

\mathrm{Att} = \operatorname{Attention}(X_n)

First residual add #

X_1 = X + \mathrm{Att}

Second LayerNorm, on a copy #

X_{n2} = \operatorname{LN}_2(X_1)

The feed-forward layer reads the copy #

\mathrm{Ff} = \operatorname{FFN}(X_{n2})

Second residual add #

Y = X_1 + \mathrm{Ff}

Written together:

The block on one line #

X_1 = X + \operatorname{Attention}(\operatorname{LN}_1(X)), \qquad Y = X_1 + \operatorname{FFN}(\operatorname{LN}_2(X_1))

The adds use X and X1, not the normalised copies. So the normalised copies leave the stream as soon as their branch has read them.

What does each symbol mean?

Symbol	Meaning	In code	Shape
$X$	The residual stream entering the block	`X`	(T, d)
$\operatorname{LN}_1, \operatorname{LN}_2$	The block's two LayerNorms, each with its own scale and shift	`layer_norm_forward with ln1_g, ln1_b and ln2_g, ln2_b`	(T, d) to (T, d)
$X_n$	The normalised copy that attention reads	`Xn`	(T, d)
$\mathrm{Att}$	The output of attention	`att`	(T, d)
$X_1$	The residual stream after the first add	`X1`	(T, d)
$X_{n2}$	The normalised copy that the feed-forward layer reads	`Xn2`	(T, d)
$\mathrm{Ff}$	The output of the feed-forward layer	`ff`	(T, d)
$Y$	The output of the block. The code returns it without naming it, and block_backward calls its gradient dY	`X1 + ff`	(T, d)

Tiny numerical example

Step through one row of a tiny block. Each step shows the vector before and after, and each add shows both terms and their sum.

The residual streamIllustrative ExampleA tiny model: the parameters are random draws from the real init_params, not learned values. The other inputs are hand-picked.

Row

Step 1 of 7: Input

X is the residual stream as it enters the block.

Row 1 of each array at this step. Each array has shape (T, d).
	`X[1]`	1	0	-0.5	2

transformer.py:131 def block_forward(X, params, l: int, cfg: ModelConfig):

Tensor shapes

Inside the feed-forward layer rows are wider, as shown above, and inside attention they split into heads. Each array block_forward itself names is (T, d), which is what allows the adds.

Every array in the block has the same shapetransformer.py:139

X(T, d)(16, 32)
Xn(T, d)(16, 32)
att(T, d)(16, 32)
X1(T, d)(16, 32)
Xn2(T, d)(16, 32)
ff(T, d)(16, 32)
Y(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.pyblock_forwardGitHub

131def block_forward(X, params, l: int, cfg: ModelConfig):132    p = f"b{l}."133    Xn, c_ln1 = layer_norm_forward(X, params[p + "ln1_g"], params[p + "ln1_b"], cfg.ln_eps) # make a normalized copy for attention to read134    att, c_att = attention_forward(Xn, params, p, cfg) # each position gathers info from earlier positions135    X1 = X + att136    Xn2, c_ln2 = layer_norm_forward(X1, params[p + "ln2_g"], params[p + "ln2_b"], cfg.ln_eps) # normalized copy again, for the feed-forward137    ff, c_ff = ffn_forward(Xn2, params, p) # ach position thinks on its own138139    return X1 + ff, dict(ln1=c_ln1, att=c_att, ln2=c_ln2, ff=c_ff)

All three LayerNorm calls in the model pass cfg.ln_eps, which is 1e-5 from the ModelConfig default. LayerNorm explains why editing config.ln_eps has no effect.

Code and math connection

Equation	Role	Lines	Code
First LayerNorm, on a copy	Computes it	transformer.py:133	`Xn, c_ln1 = layer_norm_forward(X, params[p + "ln1_g"], params[p + "ln1_b"], cfg.ln_eps) # make a normalized copy for attention to read`
Attention reads the copy	Computes it	transformer.py:134	`att, c_att = attention_forward(Xn, params, p, cfg) # each position gathers info from earlier positions`
First residual add	Computes it	transformer.py:135	`X1 = X + att`
Second LayerNorm, on a copy	Computes it	transformer.py:136	`Xn2, c_ln2 = layer_norm_forward(X1, params[p + "ln2_g"], params[p + "ln2_b"], cfg.ln_eps) # normalized copy again, for the feed-forward`
The feed-forward layer reads the copy	Computes it	transformer.py:137	`ff, c_ff = ffn_forward(Xn2, params, p) # ach position thinks on its own`
Second residual add	Computes it	transformer.py:139	`return X1 + ff, dict(ln1=c_ln1, att=c_att, ln2=c_ln2, ff=c_ff)`
The block on one line	Computes it	transformer.py:135	`X1 = X + att`
The block on one line	Computes it	transformer.py:139	`return X1 + ff, dict(ln1=c_ln1, att=c_att, ln2=c_ln2, ff=c_ff)`

Where is it used?

forward calls it once per layer, and num_layers is 2 in config.py. No test calls block_forward directly. Tests reach it through forward.

Source: transformer.py, lines 131-139, block_forward
PDF reference: Section 4.5, Eq. 53, p. 11
Used by: Called by: transformer.forward
Status: Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
First LayerNorm, on a copy	Computes it	`transformer.py`	`block_forward`	133	GitHub: transformer.py, line 133
Attention reads the copy	Computes it	`transformer.py`	`block_forward`	134	GitHub: transformer.py, line 134
First residual add	Computes it	`transformer.py`	`block_forward`	135	GitHub: transformer.py, line 135
Second LayerNorm, on a copy	Computes it	`transformer.py`	`block_forward`	136	GitHub: transformer.py, line 136
The feed-forward layer reads the copy	Computes it	`transformer.py`	`block_forward`	137	GitHub: transformer.py, line 137
Second residual add	Computes it	`transformer.py`	`block_forward`	139	GitHub: transformer.py, line 139
The block on one line	Computes it	`transformer.py`	`block_forward`	135	GitHub: transformer.py, line 135
The block on one line	Computes it	`transformer.py`	`block_forward`	139	GitHub: transformer.py, line 139

Residual connections in the backward pass

What is it?

Residual connections, backwardImplemented

Feed-forward and block

block_backward runs the block in reverse. It receives dY, the gradient of the loss on the block's output, and returns dX, the gradient on its input. At each residual add, the gradient splits into a part that goes through the branch and a part that goes around it. The two parts are added.

Why does the model need it?

Every parameter in the block needs a gradient, so the gradient has to go through both branches. X also reaches the output along the identity path, so dX needs that part too, and it goes on to the block before this one.

Note

In general, a path that passes the gradient on unchanged gives earlier layers a gradient that does not have to shrink through every layer on its way there. This is a general explanation. Nothing in this repository measures that effect.

Mathematics

One equation per step of block_backward, in the same order. The output add reverses first, because it ran last:

Through the feed-forward branch #

dX_{n2} = \operatorname{FFN}^{\text{back}}(dY)

Through the second LayerNorm #

dX_1^{\text{branch}} = \operatorname{LN}_2^{\text{back}}(dX_{n2})

Second add, reversed #

dX_1 = dY + dX_1^{\text{branch}}

Then the first add:

Through the attention branch #

dX_n = \operatorname{Attention}^{\text{back}}(dX_1)

Through the first LayerNorm #

dX^{\text{branch}} = \operatorname{LN}_1^{\text{back}}(dX_n)

First add, reversed #

dX = dX_1 + dX^{\text{branch}}

What each branch receives. Each branch's output is one term of a sum, and the derivative of a sum with respect to one term is the identity. So the feed-forward backward receives dY as it is, and attention's backward receives dX1 as it is.

What each branch receives Derived from implementation#

\frac{\partial Y}{\partial \mathrm{Ff}} = I \;\Rightarrow\; d\mathrm{Ff} = dY, \qquad \frac{\partial X_1}{\partial \mathrm{Att}} = I \;\Rightarrow\; d\mathrm{Att} = dX_1

Why the identity path carries the gradient unchanged. X1 reaches Y by two routes: directly, and through the branch. The direct route is an add, whose derivative is the identity, so multiplying dY by it leaves dY as it was. That is the term dY on the line dX1 = dY + dX1_branch. The branch route adds the other term.

Why the identity path carries dY unchanged Derived from implementation#

Y = X_1 + \mathrm{Ff}(X_1) \;\Rightarrow\; dX_1 = dY \left( I + \frac{\partial \mathrm{Ff}}{\partial X_1} \right) = dY + dY \frac{\partial \mathrm{Ff}}{\partial X_1}

What does each symbol mean?

Symbol	Meaning	In code	Shape
$dY$	Gradient of the loss on the output of the block	`dY`	(T, d)
$dX_{n2}$	Gradient on the input of the feed-forward layer	`dXn2`	(T, d)
$dX_1^{\text{branch}}$	The part of the gradient on X1 that went through the feed-forward branch	`dX1_branch`	(T, d)
$dX_1$	Gradient of the loss on X1, with both parts added	`dX1`	(T, d)
$dX_n$	Gradient on the input of attention	`dXn`	(T, d)
$dX^{\text{branch}}$	The part of the gradient on X that went through the attention branch	`dX_branch`	(T, d)
$dX$	Gradient of the loss on X, with both parts added. It is returned	`dX1 + dX_branch`	(T, d)
$I$	The identity. An array added to another passes its gradient on unchanged	not named in the code	none
$\operatorname{FFN}^{\text{back}}, \operatorname{Attention}^{\text{back}}$	The backward functions of the two branches	`ffn_backward, attention_backward`	(T, d) to (T, d)
$\operatorname{LN}_1^{\text{back}}, \operatorname{LN}_2^{\text{back}}$	layer_norm_backward with each LayerNorm's cache and scale	`layer_norm_backward`	(T, d) to (T, d)
$d\mathrm{Ff}, d\mathrm{Att}$	The gradients on each branch's output. The code passes dY and dX1, and has no separate names for them	not named in the code	(T, d)

Tiny numerical example

One backward pass through a blockIllustrative ExampleA tiny model: the parameters are random draws from the real init_params, not learned values. The other inputs are hand-picked. dY is hand-picked, not computed from a loss.

dY, shape (3, 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

dXn2, shape (3, 4)
-0.2024	0.0506	-0.0084	-0.0131
0.0057	-0.0591	-0.166	-0.0334
0	0	0	0

dX1_branch, shape (3, 4)
-0.1144	0.0456	-0.0008	0.0695
0.0288	0.0008	-0.0048	-0.0248
0	0	0	0

dX1 = dY + dX1_branch, shape (3, 4)
-0.0144	-0.1544	0.2992	0.1195
-0.0712	0.2508	0.1452	-0.3248
0.2	0.1	-0.25	0.15

dXn, shape (3, 4)
0.265	0.0618	-0.0357	0.1262
-0.0099	0.3201	-0.0097	0.3917
-0.091	-0.1101	0.0357	-0.0798

dX_branch, shape (3, 4)
0.1859	-0.0857	-0.1229	0.0228
-0.2272	0.2142	-0.0805	0.0935
-0.0026	-0.0757	0.0966	-0.0183

dX = dX1 + dX_branch, shape (3, 4)
0.1715	-0.2401	0.1763	0.1423
-0.2984	0.465	0.0647	-0.2314
0.1974	0.0243	-0.1534	0.1317

Only the identity pathIllustrative ExampleThe same block with ln1_g and ln2_g set to 0, so each branch sends back exactly 0. What leaves the block equals dY. This is an experiment for this page: init_params starts both at ones.

dX1_branch, shape (3, 4)
0	0	0	0
0	0	0	0
0	0	0	0

dX_branch, shape (3, 4)
0	0	0	0
0	0	0	0
0	0	0	0

dX, shape (3, 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

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 (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
`grads["b0.ln1_g"]`	`[0.1949, -0.4096, -0.0108, 0.5636]`	`[0.1949, -0.4096, -0.0108, 0.5636]`	`3.3e-11`
`grads["b0.wq"]`	`[-0.1689, 0.0341, -0.0560, -0.2746][0.1108, -0.0256, 0.0758, 0.3221][0.0849, -0.0251, 0.1247, 0.4888][-0.0267, 0.0166, -0.1445, -0.5363]`	`[-0.1689, 0.0341, -0.0560, -0.2746][0.1108, -0.0256, 0.0758, 0.3221][0.0849, -0.0251, 0.1247, 0.4888][-0.0267, 0.0166, -0.1445, -0.5363]`	`2.4e-11`
`grads["b0.wo"]`	`[-0.0004, -0.0323, -0.4928, 0.0446][0.0105, 0.0010, 0.3091, 0.0023][0.0642, 0.1558, 0.2905, -0.0794][-0.0422, -0.3722, 0.3625, 0.2899]`	`[-0.0004, -0.0323, -0.4928, 0.0446][0.0105, 0.0010, 0.3091, 0.0023][0.0642, 0.1558, 0.2905, -0.0794][-0.0422, -0.3722, 0.3625, 0.2899]`	`1.5e-11`
`grads["b0.ln2_b"]`	`[-0.1967, -0.0085, -0.1744, -0.0465]`	`[-0.1967, -0.0085, -0.1744, -0.0465]`	`4.7e-12`
`grads["b0.ff1_w"]`	`[0.0300, -0.2044, 0.1314, 0.0000, 0.0000, -0.0041][0.1805, 0.1929, -0.0543, 0.0000, 0.0000, 0.0083][-0.1559, 0.4836, -0.3392, 0.0000, 0.0000, 0.0079][-0.0547, -0.4721, 0.2621, 0.0000, 0.0000, -0.0121]`	`[0.0300, -0.2044, 0.1314, 0.0000, 0.0000, -0.0041][0.1805, 0.1929, -0.0543, 0.0000, 0.0000, 0.0083][-0.1559, 0.4836, -0.3392, 0.0000, 0.0000, 0.0079][-0.0547, -0.4721, 0.2621, 0.0000, 0.0000, -0.0121]`	`1.1e-11`
`grads["b0.ff2_w"]`	`[-0.0247, 0.0871, 0.2647, -0.2006][0.0497, -0.0914, 0.2202, -0.0148][0.0375, -0.0399, 0.4282, -0.1566][0.0000, 0.0000, 0.0000, 0.0000][0.0000, 0.0000, 0.0000, 0.0000][0.0111, -0.0222, 0.0332, 0.0055]`	`[-0.0247, 0.0871, 0.2647, -0.2006][0.0497, -0.0914, 0.2202, -0.0148][0.0375, -0.0399, 0.4282, -0.1566][0.0000, 0.0000, 0.0000, 0.0000][0.0000, 0.0000, 0.0000, 0.0000][0.0111, -0.0222, 0.0332, 0.0055]`	`6.7e-12`
`dX`	`[0.1715, -0.2401, 0.1763, 0.1423][-0.2984, 0.4650, 0.0647, -0.2314][0.1974, 0.0243, -0.1534, 0.1317]`	`[0.1715, -0.2401, 0.1763, 0.1423][-0.2984, 0.4650, 0.0647, -0.2314][0.1974, 0.0243, -0.1534, 0.1317]`	`1.2e-11`

Tensor shapes

Every gradient in the block has the same shapetransformer.py:152

dY(T, d)(16, 32)
dXn2(T, d)(16, 32)
dX1_branch(T, d)(16, 32)
dX1(T, d)(16, 32)
dXn(T, d)(16, 32)
dX_branch(T, d)(16, 32)
dX(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.pyblock_backwardGitHub

141def block_backward(dY, c, params, l, cfg: ModelConfig, grads):142    p = f"b{l}."143    dXn2 = ffn_backward(dY, c["ff"], params, p, grads)144    dX1_branch, dg, db = layer_norm_backward(dXn2, c["ln2"], params[p + "ln2_g"])145    grads[p + "ln2_g"] += dg146    grads[p + "ln2_b"] += db147    dX1 = dY + dX1_branch148    dXn = attention_backward(dX1, c["att"], params, p, cfg, grads)149    dX_branch, dg, db = layer_norm_backward(dXn, c["ln1"], params[p + "ln1_g"])150    grads[p + "ln1_g"] += dg151    grads[p + "ln1_b"] += db152    return dX1 + dX_branch

This diagram follows those lines. Select a step to see its function and lines.

Gradient flow through one block, in the order of block_backward

dY arrives from the next block, or from the final LayerNorm.

dXn2 = ffn_backward(dY, c["ff"], params, p, grads)transformer.py:143
Fills grads: b{l}.ff2_w, b{l}.ff2_b, b{l}.ff1_w, b{l}.ff1_b
dX1_branch, dg, db = layer_norm_backward(dXn2, c["ln2"], params[p + "ln2_g"])transformer.py:144
Fills grads: b{l}.ln2_g, b{l}.ln2_b
dX1 = dY + dX1_branchtransformer.py:147No function call
Not a function call. dY comes along the identity path, unchanged.
dXn = attention_backward(dX1, c["att"], params, p, cfg, grads)transformer.py:148
Fills grads: b{l}.wo, b{l}.bo, b{l}.wq, b{l}.bq, b{l}.wk, b{l}.bk, b{l}.wv, b{l}.bv
dX_branch, dg, db = layer_norm_backward(dXn, c["ln1"], params[p + "ln1_g"])transformer.py:149
Fills grads: b{l}.ln1_g, b{l}.ln1_b
return dX1 + dX_branchtransformer.py:152No function call
Not a function call. dX1 comes along the identity path, unchanged.

Selected: dXn2, Through the feed-forward layer, at transformer.py:143

transformer.pyffn_backwardGitHub

118def ffn_backward(dOut, c, params, p, grads):119    dA, dW2, db2 = linear_backward(c["A"], params[p + "ff2_w"], dOut)120    grads[p + "ff2_w"] += dW2121    grads[p + "ff2_b"] += db2122    dZ = relu_backward(c["Z"], dA)123    dXn, dW1, db1 = linear_backward(c["Xn"], params[p + "ff1_w"], dZ)124125    grads[p + "ff1_w"] += dW1126    grads[p + "ff1_b"] += db1127128    return dXn

transformer.pyblock_backwardGitHub

141def block_backward(dY, c, params, l, cfg: ModelConfig, grads):142    p = f"b{l}."143    dXn2 = ffn_backward(dY, c["ff"], params, p, grads)144    dX1_branch, dg, db = layer_norm_backward(dXn2, c["ln2"], params[p + "ln2_g"])145    grads[p + "ln2_g"] += dg146    grads[p + "ln2_b"] += db147    dX1 = dY + dX1_branch148    dXn = attention_backward(dX1, c["att"], params, p, cfg, grads)149    dX_branch, dg, db = layer_norm_backward(dXn, c["ln1"], params[p + "ln1_g"])150    grads[p + "ln1_g"] += dg151    grads[p + "ln1_b"] += db152    return dX1 + dX_branch

Code and math connection

Equation	Role	Lines	Code
Through the feed-forward branch	Computes it	transformer.py:143	`dXn2 = ffn_backward(dY, c["ff"], params, p, grads)`
Through the second LayerNorm	Computes it	transformer.py:144	`dX1_branch, dg, db = layer_norm_backward(dXn2, c["ln2"], params[p + "ln2_g"])`
Second add, reversed	Computes it	transformer.py:147	`dX1 = dY + dX1_branch`
Through the attention branch	Computes it	transformer.py:148	`dXn = attention_backward(dX1, c["att"], params, p, cfg, grads)`
Through the first LayerNorm	Computes it	transformer.py:149	`dX_branch, dg, db = layer_norm_backward(dXn, c["ln1"], params[p + "ln1_g"])`
First add, reversed	Computes it	transformer.py:152	`return dX1 + dX_branch`
What each branch receives	Computes it	transformer.py:143	`dXn2 = ffn_backward(dY, c["ff"], params, p, grads)`
What each branch receives	Computes it	transformer.py:148	`dXn = attention_backward(dX1, c["att"], params, p, cfg, grads)`
Why the identity path carries dY unchanged	Computes it	transformer.py:147	`dX1 = dY + dX1_branch`
Why the identity path carries dY unchanged	Tests it	tests/test_gradcheck.py:25	`_, grads = loss_and_grads(params, ids, targets, cfg)`

Where is it used?

loss_and_grads calls it once per layer, last block first (transformer.py:187). No test calls it directly. test_every_tensor_passes reaches it through the whole model.

Source: transformer.py, lines 141-152, block_backward
PDF reference: Section 4.5, p. 11
Used by: Called by: transformer.loss_and_grads
Tests it: tests/test_gradcheck.py, line 25 (test_every_tensor_passes), for Why the identity path carries dY unchanged
Status: Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Through the feed-forward branch	Computes it	`transformer.py`	`block_backward`	143	GitHub: transformer.py, line 143
Through the second LayerNorm	Computes it	`transformer.py`	`block_backward`	144	GitHub: transformer.py, line 144
Second add, reversed	Computes it	`transformer.py`	`block_backward`	147	GitHub: transformer.py, line 147
Through the attention branch	Computes it	`transformer.py`	`block_backward`	148	GitHub: transformer.py, line 148
Through the first LayerNorm	Computes it	`transformer.py`	`block_backward`	149	GitHub: transformer.py, line 149
First add, reversed	Computes it	`transformer.py`	`block_backward`	152	GitHub: transformer.py, line 152
What each branch receives	Computes it	`transformer.py`	`block_backward`	143	GitHub: transformer.py, line 143
What each branch receives	Computes it	`transformer.py`	`block_backward`	148	GitHub: transformer.py, line 148
Why the identity path carries dY unchanged	Computes it	`transformer.py`	`block_backward`	147	GitHub: transformer.py, line 147

How this reaches a parameter update. block_backward stores nothing itself except the two LayerNorms' gradients (transformer.py:145, transformer.py:150). The branch functions fill the rest of grads. The optimizer then uses them, through the same steps as above: train.py:61, train.py:62 and train.py:66. dX is returned (transformer.py:152). It becomes the dY of the block before this one, or goes to embed_backward after the first block.

Caches

The forward pass returns each cache alongside its output, because the backward pass needs values it cannot get otherwise. block_forward collects four of them in one dictionary. block_backward reads each one once, in reverse order:

Passed as	Holds	Built at	Read by
`c["ff"]`	`Xn`: The input of the feed-forward layer, the output of LayerNorm 2.`Z`: The output of the first linear layer, before ReLU.`A`: The output of ReLU.	transformer.py:116	`ffn_backward`, at transformer.py:143
`c["ln2"]`	`Xhat`: The normalised rows.`s`: One over the root of variance plus epsilon, per row.	ops.py:47	`layer_norm_backward`, at transformer.py:144
`c["att"]`	`Xn`: The input of attention, the output of LayerNorm 1.`Qh`: Queries, per head.`Kh`: Keys, per head.`Vh`: Values, per head.`A`: Attention weights. A different array from the ReLU output above.`O`: The merged values, before the output projection.	transformer.py:86	`attention_backward`, at transformer.py:148
`c["ln1"]`	`Xhat`: The normalised rows.`s`: One over the root of variance plus epsilon, per row.	ops.py:47	`layer_norm_backward`, at transformer.py:149

The backward functions read parameters from params, not from a cache. block_forward leaves att, ff, X and X1 out of its caches, because no backward function reads them. In the code, both the feed-forward cache and the attention cache have a key named A. They hold different arrays.

Parameters of one block

Each block has its own copy of these. The names and shapes come from calling init_params. Only the feed-forward and LayerNorm rows are new on this page. Attention covers the others.

Name	Shape	With config	Scalars	Role
`b{l}.ln1_g`transformer.py:29	`(d,)`	`(32,)`	`d`32 per block	LayerNorm 1 scale, before attention.Starts as ones. No decay.
`b{l}.ln1_b`transformer.py:29	`(d,)`	`(32,)`	`d`32 per block	LayerNorm 1 shift, before attention.Starts as zeros. No decay.
`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.
`b{l}.ln2_g`transformer.py:35	`(d,)`	`(32,)`	`d`32 per block	LayerNorm 2 scale, before feed-forward.Starts as ones. No decay.
`b{l}.ln2_b`transformer.py:35	`(d,)`	`(32,)`	`d`32 per block	LayerNorm 2 shift, before feed-forward.Starts as zeros. No decay.
`b{l}.ff1_w`transformer.py:36	`(d, F)`	`(32, 128)`	`d × F`4,096 per block	First feed-forward layer weights.Starts as Xavier uniform. Weight decay.
`b{l}.ff1_b`transformer.py:36	`(F,)`	`(128,)`	`F`128 per block	First feed-forward layer bias.Starts as zeros. No decay.
`b{l}.ff2_w`transformer.py:37	`(F, d)`	`(128, 32)`	`d × F`4,096 per block	Second feed-forward layer weights.Starts as Xavier uniform. Weight decay.
`b{l}.ff2_b`transformer.py:37	`(d,)`	`(32,)`	`d`32 per block	Second feed-forward layer bias.Starts as zeros. No decay.

Next: Full Model joins the blocks to the embeddings and the head.