34. Why a Block Has a Feed-Forward Network: Backpropagation Through Layer Normalization, ReLU and Residuals

What you will understand: why a transformer block does not stop at attention. A model made of attention alone fails at a small, clean task; the same model with a feed-forward network, a layer normalization and a residual connection added learns it. Getting there means deriving, by hand, the backward pass through three more pieces: the layer normalization (the hardest formula so far), the ReLU (including what to do at its corner) and the residual connection (where two gradient paths meet). Every piece is a Mountain Goat function, every formula is checked against an independent Python calculation and against finite differences of the loss for every one of the 120 weights, and the two models are then compared on all 81 possible inputs.

What you need to know first: Chapter 33 (the attention classifier and its backward pass; this chapter reuses it unchanged for the attention part), Chapter 30 (what layer normalization and the feed-forward network are) and Chapter 31 (cross-entropy). No new compiler operation is needed: the programs run on Chapter 32's compiler.

Compile and run

cd docs/part32/code && ./build.sh                 # once: the newest compiler (this chapter adds nothing to it)
cd ../../part34/code
python3 make_examples.py                          # (re)writes the examples from ffn_lib.py
../../part32/code/mgc run examples/04_train_with_ffn.mg   # train model B for 600 steps (about 100 seconds: the program is 6,000 lines)
./run_examples.sh > examples_out.txt              # every example and every error example on this page
./check_ffn.py                                    # both backward passes and both training runs against the independent Python version (about 2 minutes)
./lr_sweep.py > lr_sweep_out.txt                  # Python only: how reliable is the training? (about 5 minutes)
./model_mutation.py > model_mutation_out.txt      # runs 10 WRONG backward passes through the checker (about 25 minutes)
cd ../../part15/code && ./run_lit.sh              # the whole test suite

Every listing and output on this page comes from these commands.

Why this page shows FAIL lines and error messages, and why that is good

The three error examples are programs the compiler must reject (each ends exit status 1). The mutation output is a negative control: deliberately wrong backward passes run against the checker, where "caught" is the wanted result and "NOT CAUGHT" would be a gap. Model A's failure to learn the task is the result of the experiment, not a bug.

The task: exactly one of two tokens

A sequence is four tokens, each 0, 1 or 2. The label is 1 if exactly one of the tokens 1 and 2 appears, and 0 otherwise (neither, or both):

sequence	label	why
0 0 0 0	0	neither 1 nor 2
1 0 0 0	1	1 only
2 2 2 2	1	2 only
1 1 1 1	1	1 only
1 2 0 0	0	both
2 2 1 0	0	both

This is "exclusive or" (XOR) of "is 1 present" and "is 2 present". Of the 3⁴ = 81 possible sequences, 30 have label 1 and 51 have label 0 (so always answering 0 is right 63% of the time). A random 54 of them (20 labelled 1) train the model; the other 27 (10 labelled 1) are held out. As in Chapter 33, the answer does not depend on order, which suits attention.

Two models

Model A is Chapter 33's: attention pools the four tokens into one vector h (1×4), and a linear read-out h @ wo gives the two class scores. 36 numbers.

Model B puts a block between the pooling and the read-out, the same arrangement as a transformer block (Chapter 30):

1. h = the attention-pooled vector;
2. layer normalization: u = gm * (h − mean(h)) / sigma + bt, with sigma = sqrt(variance(h) + 0.00001);
3. feed-forward network: f = relu(u @ w1 + b1) @ w2 + b2 (widen from 4 to 8 features, zero the negative ones, narrow back to 4);
4. residual connection: z = h + f;
5. the read-out z @ wo.

120 numbers. Everything is stacked 54 sequences at a time exactly as in Chapter 33 (group matrix, reshape for the per-sequence softmax).

Why would A fail? Here is an argument, not a proof. The pooled vector h is a weighted average of the three token value vectors v0, v1, v2, so it lies in the triangle they span. Suppose attention learns to prefer 1 and 2 over 0. Then a sequence with a 1 and no 2 gives a point near v1, one with a 2 and no 1 a point near v2, one with both a point between them on the line from v1 to v2, and one with neither is exactly v0. Along the line from v1 to v2 the labels run 1, 0, 1: no linear read-out can say "1 at both ends and 0 in the middle". The feed-forward network adds a bend (the ReLU) that can. The experiments below are consistent with this, and do not prove it.

The backward pass, new rows

Chapter 33's table gets new rows between the read-out and the pooling. Write dX for the gradient of the loss with respect to X, B for the number of sequences (54), and ŷ for the normalized value (h − mean(h))/sigma:

forward step	backward step
logits = z @ wo, cross-entropy	dlogits = (softmax(logits) − y) / B; dwo = zᵀ @ dlogits; dz = dlogits @ woᵀ
z = h + f (residual)	dz goes down both paths: straight to h, and into f
f = relu(pre) @ w2 + b2	db2 = col_sum(dz); dw2 = relu(pre)ᵀ @ dz; drelu = dz @ w2ᵀ
relu	dpre = drelu * ge(pre, 0) (the gradient passes where pre >= 0)
pre = u @ w1 + b1	db1 = col_sum(dpre); dw1 = uᵀ @ dpre; du = dpre @ w1ᵀ
u = gm * ŷ + bt	dbt = col_sum(du); dgm = col_sum(du * ŷ); dŷ = du * gm
ŷ = (h − mean(h))/sigma (layer norm)	dh_ln = (dŷ − mean(dŷ) − ŷ * mean(dŷ * ŷ)) / sigma (means taken along each row)
merge at h	dh = dz + dh_ln (the sum of the residual path and the layer-norm path)
attention pooling	exactly Chapter 33: from dh to dwv, dq, dwk

Three of these deserve a closer look.

Layer normalization

The layer-norm formula has three terms, because every entry of h affects the mean, the spread sigma and therefore every normalized value. A small program computes it by the formula and by brute force (nudge each of four numbers by 0.01 up and down):

################ mgc run 01_layer_norm_backward.mg
  | # Backward through layer normalisation. y = (h - mean(h)) / sigma with sigma = sqrt(variance(h) + 0.00001), for ONE row of four numbers h.
  | # If the loss is L = sum(d * y) for a fixed upstream gradient d, the derivative of L with respect to h is
  | #     (d - mean(d) - y * mean(d * y)) / sigma
  | # (three terms, because every h_i changes the mean, the spread and so every y_j). The first matrix is that formula; the next four are the same
  | # derivatives found by brute force: nudge one entry of h by 0.01 up and by 0.01 down and watch L change.
  | def ln_sigma(h: tensor[1x4]) = sqrt(row_mean((h - row_mean(h)) * (h - row_mean(h))) + 0.00001)
  | def ln_hat(h: tensor[1x4]) = (h - row_mean(h)) / ln_sigma(h)
  | def loss(h: tensor[1x4], d: tensor[1x4]) = row_sum(d * ln_hat(h))
  | def backward(h: tensor[1x4], d: tensor[1x4]) = (d - row_mean(d) - ln_hat(h) * row_mean(d * ln_hat(h))) / ln_sigma(h)
  | let h = [[1, 2, 4, 8]]
  | let d = [[1, -2, 3, 0.5]]
  | print backward(h, d)
  | print (loss(h + [[0.01, 0, 0, 0]], d) - loss(h - [[0.01, 0, 0, 0]], d)) / 0.02
  | print (loss(h + [[0, 0.01, 0, 0]], d) - loss(h - [[0, 0.01, 0, 0]], d)) / 0.02
  | print (loss(h + [[0, 0, 0.01, 0]], d) - loss(h - [[0, 0, 0.01, 0]], d)) / 0.02
  | print (loss(h + [[0, 0, 0, 0.01]], d) - loss(h - [[0, 0, 0, 0.01]], d)) / 0.02
  | # Two things to see: the formula's four numbers add up to (nearly) zero (adding the same amount to every h changes nothing, so the gradients must cancel),
  | # and the brute-force numbers agree with the formula to about four digits (the step 0.01 and the six printed digits limit the agreement).
  | print row_sum(backward(h, d))
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 4] strides = [4, 1] data = 
[[0.26921,   -0.896826,   0.874121,   -0.246505]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.26921]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[-0.896825]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.87412]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[-0.246506]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0]]
  ----- exit status 0

The formula gives [0.26921, −0.896826, 0.874121, −0.246505]; the brute-force numbers are 0.26921, −0.896825, 0.87412, −0.246506. They agree. The last line is the row sum of the gradient, 0: adding the same amount to every h changes nothing (the mean moves with it), so the four gradient entries must cancel. The three terms are each needed: leaving out mean(dŷ) or the ŷ * mean(dŷ * ŷ) projection breaks that cancellation, and a wrong layer-norm backward is among the mutations tested below.

ReLU, and its corner

y = max(0, p) has derivative 1 where p is positive and 0 where it is negative, so the upstream gradient passes where p >= 0 and is cut elsewhere: d * ge(p, 0), using the comparison of Chapter 32.

################ mgc run 02_relu_backward.mg
  | # Backward through relu. y = max(0, p). The derivative is 1 where p is positive and 0 where it is negative, so the gradient d passes through where
  | # p >= 0 and is cut where p < 0: dp = d * ge(p, 0). (ge is the comparison of Chapter 32; [[0]] is a 1x1 matrix, stretched across the row.)
  | # At p = 0 exactly the true derivative does not exist (a corner); this program, like most training code, passes the gradient there. The brute-force
  | # estimate at a corner is the average of the two sides, half of d. The second row shows it.
  | let p = [[-1.5, 0, 2, 3]]
  | let d = [[1, 2, 3, 4]]
  | print d * ge(p, [[0]])
  | print (row_sum(d * relu(p + [[0.01, 0, 0, 0]])) - row_sum(d * relu(p - [[0.01, 0, 0, 0]]))) / 0.02
  | print (row_sum(d * relu(p + [[0, 0.01, 0, 0]])) - row_sum(d * relu(p - [[0, 0.01, 0, 0]]))) / 0.02
  | print (row_sum(d * relu(p + [[0, 0, 0.01, 0]])) - row_sum(d * relu(p - [[0, 0, 0.01, 0]]))) / 0.02
  | print (row_sum(d * relu(p + [[0, 0, 0, 0.01]])) - row_sum(d * relu(p - [[0, 0, 0, 0.01]]))) / 0.02
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 4] strides = [4, 1] data = 
[[0,   2,   3,   4]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[3]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[4]]
  ----- exit status 0

For p = [−1.5, 0, 2, 3] and d = [1, 2, 3, 4] the formula gives [0, 2, 3, 4]. The brute-force numbers are 0, 1, 3, 4: they agree except at p = 0, exactly at the corner, where the brute force gives 1, half of 2. At a corner the derivative does not exist; the two sides disagree (0 on the left, 1 on the right) and brute force averages them. Training code must pick a convention. This one passes the gradient at p = 0 (ge, not a strict >). It matters little in practice, because an exact zero is rare, and it is why the finite-difference check below uses a small step (see "The finite-difference step").

The residual connection

z = h + f makes h influence the loss along two routes: directly through z, and through the layer norm and feed-forward network. By the chain rule the gradients along the two routes add: dh = dz + dh_ln. A wrong backward pass that drops the direct route is among the mutations below.

The definitions

The whole model B as functions: the forward pass, then one function per gradient (the attention-pooling part is Chapter 33's, with a different vocabulary size):

# Model B: attention pooling, then a layer normalization, a ReLU feed-forward network and a residual connection, then a linear read-out.
# B sequences are stacked: x has B*4 rows (the tokens of sequence 0, then sequence 1, ...) and g (B x B*4) adds up the 4 rows belonging to each sequence.
# (The same functions exist again with the suffix _27 for chunks of 27 sequences; only the sizes differ.)

def softmax_rows(m: tensor[54x4]) = exp(m - row_max(m)) / row_sum(exp(m - row_max(m)))
def attn_matrix(x: tensor[216x3], q: tensor[1x4], wk: tensor[3x4]) = softmax_rows(reshape((x @ wk) @ transpose(q) * 0.5, 54, 4))
def pooled(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4]) = g @ ((x @ wv) * reshape(attn_matrix(x, q, wk), 216, 1))
def ln_sigma(h: tensor[54x4]) = sqrt(row_mean((h - row_mean(h)) * (h - row_mean(h))) + 1e-05)
def ln_hat(h: tensor[54x4]) = (h - row_mean(h)) / ln_sigma(h)
def ln_out(h: tensor[54x4], gm: tensor[1x4], bt: tensor[1x4]) = ln_hat(h) * gm + bt
def hidden_pre(h: tensor[54x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8]) = ln_out(h, gm, bt) @ w1 + b1
def block(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4]) = pooled(x, g, q, wk, wv) + (relu(hidden_pre(pooled(x, g, q, wk, wv), gm, bt, w1, b1)) @ w2 + b2)
def logits(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2]) = block(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2) @ wo
def log_softmax(z: tensor[54x2]) = z - row_max(z) - log(row_sum(exp(z - row_max(z))))
def loss(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = col_sum(row_sum(y * log_softmax(logits(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo)))) * -0.018518518518518517
def correct(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = col_sum(row_sum(ge(logits(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo), row_max(logits(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo))) * y))
def d_logits(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = (exp(log_softmax(logits(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo))) - y) * 0.018518518518518517
def grad_wo(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = transpose(block(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2)) @ d_logits(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y)
def d_z(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = d_logits(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y) @ transpose(wo)
def grad_b2(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = col_sum(d_z(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y))
def grad_w2(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = transpose(relu(hidden_pre(pooled(x, g, q, wk, wv), gm, bt, w1, b1))) @ d_z(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y)
def d_pre(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = (d_z(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y) @ transpose(w2)) * ge(hidden_pre(pooled(x, g, q, wk, wv), gm, bt, w1, b1), [[0]])
def grad_b1(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = col_sum(d_pre(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y))
def grad_w1(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = transpose(ln_out(pooled(x, g, q, wk, wv), gm, bt)) @ d_pre(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y)
def d_u(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = d_pre(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y) @ transpose(w1)
def grad_bt(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = col_sum(d_u(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y))
def grad_gm(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = col_sum(d_u(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y) * ln_hat(pooled(x, g, q, wk, wv)))
def d_hat(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = d_u(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y) * gm
def d_h(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = d_z(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y) + (d_hat(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y) - row_mean(d_hat(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y)) - ln_hat(pooled(x, g, q, wk, wv)) * row_mean(d_hat(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y) * ln_hat(pooled(x, g, q, wk, wv)))) / ln_sigma(pooled(x, g, q, wk, wv))
def d_ctx_rows(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = transpose(g) @ d_h(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y)
def grad_wv(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = transpose(x) @ (d_ctx_rows(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y) * reshape(attn_matrix(x, q, wk), 216, 1))
def d_attn_matrix(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = reshape(row_sum(d_ctx_rows(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y) * (x @ wv)), 54, 4)
def d_scores(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = reshape(attn_matrix(x, q, wk) * (d_attn_matrix(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y) - row_sum(attn_matrix(x, q, wk) * d_attn_matrix(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y))), 216, 1)
def grad_q(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = transpose(d_scores(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y)) @ (x @ wk) * 0.5
def grad_wk(x: tensor[216x3], g: tensor[54x216], q: tensor[1x4], wk: tensor[3x4], wv: tensor[3x4], gm: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4], wo: tensor[4x2], y: tensor[54x2]) = transpose(x) @ (d_scores(x, g, q, wk, wv, gm, bt, w1, b1, w2, b2, wo, y) @ q) * 0.5

The experiment

examples/03_train_attention_only.mg trains model A and examples/04_train_with_ffn.mg trains model B, with the same data, the same learning rate (0.5), 600 steps and seeded starting weights; each prints the loss and the number of correct training sequences at eight checkpoints, then the held-out loss and count, then the number correct in each of three chunks of 27 of all 81 possible sequences, then the scores for the first chunk. Model A:

Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.698663]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.697193]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.686293]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[34]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.55837]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[41]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.526339]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[34]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.509436]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[38]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.43301]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[46]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.428121]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[46]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.6969]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[23]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[22]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[20]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[27]]

Model B:

Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.68508]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[37]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.640907]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[34]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.650084]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[34]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.607979]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[36]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.433422]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[45]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.0022629]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[54]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.000505909]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[54]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.000266119]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[54]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.467902]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[26]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[26]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[27]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[27]]

(Read the first sixteen numbers as pairs: loss, then number correct out of 54.)

after step	0	1	10	50	100	200	400	600
A: loss	0.699	0.697	0.686	0.558	0.526	0.509	0.433	0.428
A: correct of 54	25	25	34	41	34	38	46	46
B: loss	0.685	0.641	0.650	0.608	0.433	0.0023	0.0005	0.0003
B: correct of 54	37	34	34	36	45	54	54	54

• Model A cannot learn the task. Its loss flattens at about 0.428 and it classifies 46 of its own 54 training sequences correctly. On the 27 held-out sequences it gets 23, and on all 81 possible sequences 69.
• Model B learns it. Its loss stays above 0.4 for the first hundred steps (0.433 at step 100), then drops sharply to 0.0023 by step 200 and 0.0003 by step 600 (the shape of a model "finding" the trick), and it classifies all 54 training sequences correctly. On the held-out 27 it gets 26, and on all 81 it gets 80: it beats model A by 11 sequences.

The one miss

Which sequence does B get wrong? The last numbers of the B output are the scores for the first chunk of 27, whose first row is 0 0 0 0. Its scores are [−5.00, 7.62]: class 1 wins by a wide margin, but the label of 0 0 0 0 is 0 (neither 1 nor 2 appears). It is the only error, and 0 0 0 0 is not among the 54 training sequences (the random split happened to leave it out): the model has never been shown the case in which neither token appears, and it extrapolates confidently and wrongly. This is the third chapter (after 32 and 33) in which the model's single failure is an input its training data never showed. The checker finds the missing sequence by running the reference over all 81 rather than assuming it.

How reliable is this? A Python-only sweep

The run above uses one starting-weight set at one learning rate. To see how much that matters, lr_sweep.py repeats the experiment in the plain-Python implementation only (it does not go through mgc; the Python and Mountain Goat versions agree step for step, but this table is the Python one) for six starting-weight seeds at several learning rates:

model    lr   per starting-weight seed 1..6: (final training loss, correct of 81)
    A   0.5   (0.428, 69)  (0.436, 69)  (0.511, 52)  (0.428, 69)  (0.428, 69)  (0.429, 69)
    B  0.25   (0.000465, 80)  (0.000344, 80)  (0.000362, 80)  (0.000653, 80)  (0.00037, 80)  (0.000362, 80)
    B   0.5   (0.000266, 80)  (0.000137, 80)  (0.000153, 80)  (0.429, 69)  (0.000159, 80)  (0.000148, 80)
    B   1.0   (0.429, 69)  (0.404, 65)  (0.41, 65)  (0.000317, 80)  diverged  diverged
    B   2.0   (0.433, 69)  (0.432, 69)  diverged  (0.433, 65)  diverged  diverged

Each cell is (final training loss, sequences correct out of 81):

• Model A is stuck for every seed at the learning rate used here: 69 of 81 for five seeds and 52 for the sixth, with a loss near 0.43 or worse.
• Model B at 0.25 succeeds for all six seeds (80 of 81). At 0.5 (the rate used in the run above) it succeeds for five of six: seed 4 gets stuck at the same 69 as model A. At 1.0 only one of six succeeds, two diverge to infinity and the rest stall; at 2.0 none does.
• So model B's success is real but not guaranteed: it depends on the learning rate and the start. The run above uses a seed that works. (It is shown at 0.5 and not at 0.25 because at 0.25 model A has not reached its plateau after 600 steps (a loss of 0.445 for seed 1, against 0.428 at the higher rate), which would make "A cannot fit" a weaker statement; at 0.5 A has plateaued. This comparison of A's two rates was made in plain Python.)

The finite-difference step: a measured lesson

The first version of the gradient check used a step of 0.01, as in Chapters 31 to 33. For model B it failed for two of three starting-weight sets: the largest disagreement between the hand-derived gradient and the estimate was 5.6·10⁻³ and 4.0·10⁻³, above the tolerance. The hand-derived gradient was not at fault: it equals the independent Python gradient. Computing the same estimate in plain Python with exact (not six-digit) loss values gives the same 5.6·10⁻³, and with a step of 0.001 it gives 10⁻⁶. So the error is in the estimate: a step of 0.01 straddles a ReLU corner or is too coarse for the curvature. A smaller step cannot be too small either, since mgc prints the loss to six digits: a step of 0.001 limits the estimate to about 4·10⁻⁴. The check uses 0.001 and a tolerance of 10⁻³, and for model B finds a largest disagreement of about 4·10⁻⁴.

Mistakes the compiler reports

################ mgc run errors/e1_scale_of_the_wrong_width.mg
  | # layer norm's scale is one number per feature: 1 x 4 here. A 1 x 3 scale does not fit, and the compiler says so.
  | def ln_sigma(h: tensor[1x4]) = sqrt(row_mean((h - row_mean(h)) * (h - row_mean(h))) + 0.00001)
  | def ln_hat(h: tensor[1x4]) = (h - row_mean(h)) / ln_sigma(h)
  | def ln_out(h: tensor[1x4], gm: tensor[1x4], bt: tensor[1x4]) = ln_hat(h) * gm + bt
  | print ln_out([[1, 2, 3, 4]], [[1, 1, 1]], [[0, 0, 0, 0]])
  ----- output
errors/e1_scale_of_the_wrong_width.mg:5: error: argument shape 1x3 does not fit parameter 1x4
  ----- exit status 1

################ mgc run errors/e2_missing_argument.mg
  | # ln_out takes three arguments (the row, the scale and the shift). Two is an arity error, reported before any code exists.
  | def ln_sigma(h: tensor[1x4]) = sqrt(row_mean((h - row_mean(h)) * (h - row_mean(h))) + 0.00001)
  | def ln_hat(h: tensor[1x4]) = (h - row_mean(h)) / ln_sigma(h)
  | def ln_out(h: tensor[1x4], gm: tensor[1x4], bt: tensor[1x4]) = ln_hat(h) * gm + bt
  | print ln_out([[1, 2, 3, 4]], [[1, 1, 1, 1]])
  ----- output
errors/e2_missing_argument.mg:5: error: ln_out takes 3 arguments, got 2
  ----- exit status 1

################ mgc run errors/e3_residual_shapes_differ.mg
  | # A residual connection adds a sublayer's output to its input, so the shapes must agree. A 54 x 8 hidden layer cannot be added to a 54 x 4 input.
  | let x = [[1, 2, 3, 4], [5, 6, 7, 8]]
  | let hidden = [[1, 2, 3, 4, 5, 6, 7, 8], [1, 2, 3, 4, 5, 6, 7, 8]]
  | print x + hidden
  ----- output
errors/e3_residual_shapes_differ.mg:4: error: cannot add shapes 2x4 and 2x8: dimensions 4 and 8 differ and neither is 1
  ----- exit status 1

A layer-norm scale of the wrong width, a call with a missing argument, and a residual connection whose two sides have different shapes. As always, the shapes are types.

Tests

Three new lit files in test/ffn34/ (the suite is now 138 tests; no verifiers or ir test because no compiler operation was added):

Test	What it checks
examples	the two backward-pass examples (worked out above) and both training runs' loss and correct count at every checkpoint, the held-out result and the chunk counts
errors	the three mistakes give the exact messages
against-reference	check_ffn.py, below (about two minutes)

check_ffn.py compares with an independent Python implementation (plain lists; every gradient derived element by element) and checks, in order:

1. The gradients of both models, for three starting-weight sets each: the loss and every gradient equal the reference's, and the gradient of every weight (36 for A, 120 for B) agrees with a finite-difference estimate (step 0.001).
2. Training of both models: loss and correct count at all eight checkpoints, the held-out loss and count, and the three chunk counts equal the reference's.
3. The experiment: model A fits fewer than all of its 54 training sequences (loss above 0.3) and gets at most 70 of 81; model B fits all 54 with loss below 0.001, gets 80 of 81 (at least 10 more than A), and the one miss is 0 0 0 0, which is not a training sequence.
4. The miss, from the program: the program's own scores for 0 0 0 0 choose class 1 while the label is 0.

-- the hand-derived backward passes against the independent per-example gradient, and against finite differences
  ok   model A, starting weights #1: the loss and all 4 gradients equal the reference
  ok   model A, starting weights #1: all 36 weights agree with finite differences (largest difference 2.6e-04)
  ok   model A, starting weights #2: the loss and all 4 gradients equal the reference
  ok   model A, starting weights #2: all 36 weights agree with finite differences (largest difference 3.9e-04)
  ok   model A, starting weights #3: the loss and all 4 gradients equal the reference
  ok   model A, starting weights #3: all 36 weights agree with finite differences (largest difference 2.8e-04)
  ok   model B, starting weights #1: the loss and all 10 gradients equal the reference
  ok   model B, starting weights #1: all 120 weights agree with finite differences (largest difference 3.9e-04)
  ok   model B, starting weights #2: the loss and all 10 gradients equal the reference
  ok   model B, starting weights #2: all 120 weights agree with finite differences (largest difference 3.9e-04)
  ok   model B, starting weights #3: the loss and all 10 gradients equal the reference
  ok   model B, starting weights #3: all 120 weights agree with finite differences (largest difference 4.4e-04)
-- training model A against the independent Python trainer
  ok   model A, step 0: loss 0.698663 and 25 of 54 correct equal the reference
  ok   model A, step 1: loss 0.697193 and 25 of 54 correct equal the reference
  ok   model A, step 10: loss 0.686293 and 34 of 54 correct equal the reference
  ok   model A, step 50: loss 0.558370 and 41 of 54 correct equal the reference
  ok   model A, step 100: loss 0.526339 and 34 of 54 correct equal the reference
  ok   model A, step 200: loss 0.509436 and 38 of 54 correct equal the reference
  ok   model A, step 400: loss 0.433010 and 46 of 54 correct equal the reference
  ok   model A, step 600: loss 0.428121 and 46 of 54 correct equal the reference
  ok   model A: the 27 held-out sequences: 23 correct and loss 0.696900 equal the reference
  ok   model A: the number correct in each of the three chunks of 27 of all 81 sequences, [22, 20, 27], equals the reference's
-- training model B against the independent Python trainer
  ok   model B, step 0: loss 0.685080 and 37 of 54 correct equal the reference
  ok   model B, step 1: loss 0.640907 and 34 of 54 correct equal the reference
  ok   model B, step 10: loss 0.650084 and 34 of 54 correct equal the reference
  ok   model B, step 50: loss 0.607979 and 36 of 54 correct equal the reference
  ok   model B, step 100: loss 0.433422 and 45 of 54 correct equal the reference
  ok   model B, step 200: loss 0.002263 and 54 of 54 correct equal the reference
  ok   model B, step 400: loss 0.000506 and 54 of 54 correct equal the reference
  ok   model B, step 600: loss 0.000266 and 54 of 54 correct equal the reference
  ok   model B: the 27 held-out sequences: 26 correct and loss 0.467902 equal the reference
  ok   model B: the number correct in each of the three chunks of 27 of all 81 sequences, [26, 27, 27], equals the reference's
-- the experiment: does the feed-forward block matter?
  ok   model A cannot fit its own training data: 46 of 54 right, loss 0.428 after 600 steps
  ok   model A gets only 69 of the 81 possible sequences right
  ok   model B fits all 54 training sequences (loss 0.000266)
  ok   model B gets 80 of 81 right, 11 more than model A
  ok   the one sequence model B gets wrong is 0 0 0 0, which is not among the training sequences (misses: [[0, 0, 0, 0]])
  ok   in the program, the first row of the first chunk (0 0 0 0) scores [-5.00221, 7.62303] and so is classified as 1; its label is 0
all checks pass

Are the checks good enough? Run ten wrong backward passes.

model_mutation.py writes ten versions of model B's definitions, each with one mistake (the layer-norm backward missing a term or the division by sigma, the residual path dropped from the gradient, the ReLU mask omitted or reversed, dgm and dw1 computed from the wrong input, a different epsilon, the forward residual removed) and runs the checker on each:

NOTE: this script deliberately runs wrong backward passes. 'caught' lines are EXPECTED: they show the checker can detect the mistake.
baseline (nothing broken): checker passes (expected)
layer-norm backward misses the mean-subtraction term: caught by 14 failing checks (the reference, finite differences)
layer-norm backward misses the projection term (the y * mean(d * y) part): caught by 17 failing checks (the reference, finite differences, the experiment)
layer-norm backward forgets to divide by sigma: caught by 15 failing checks (the reference, finite differences, the experiment)
the residual path is dropped from the gradient (d_h = the layer-norm path only): caught by 14 failing checks (the reference, finite differences)
the relu mask is omitted (the gradient passes everywhere): caught by 17 failing checks (the reference, finite differences, the experiment)
the relu mask is the wrong way round (passes where the input is NEGATIVE): caught by 17 failing checks (the reference, finite differences, the experiment)
grad_gm forgets to multiply by the normalised input: caught by 14 failing checks (the reference, finite differences)
grad_w1 uses the un-normalised input instead of the layer-norm output: caught by 17 failing checks (the reference, finite differences, the experiment)
the layer-norm epsilon is 0.1 in the forward AND backward pass (consistent with itself, not with the reference): caught by 15 failing checks (the reference, the experiment)
the forward pass drops the residual connection (the backward pass still has it): caught by 15 failing checks (the reference, finite differences)

All ten are caught, and the pattern of which checks catch them is informative:

• Eight of the ten are caught by the finite-difference check (and by the reference): the layer-norm backward missing a term or the division by sigma, the residual path dropped from the gradient, the ReLU mask omitted or reversed, grad_gm and grad_w1 using the wrong input, and the forward pass with the residual removed (whose backward pass no longer matches it). These are mistakes in the derivation, and finite differences use only the loss, so they notice.
• One is invisible to the finite-difference check by design: the layer-norm epsilon changed to 0.1 in the forward and the backward pass. The program's gradient is then exactly the gradient of the program's own (different) loss, so the finite differences agree with it. It is caught only by the comparison with the reference, which uses the real epsilon, and by the experiment. This is the case where "the gradient matches its own loss" and "the program computes the intended thing" come apart.
• Several are also caught by "the experiment" (model B stops fitting its data or stops beating model A): omitting or reversing the ReLU mask and the two wrong-input gradients make training fail, while the layer-norm missing-mean term and the missing residual path still let training reach the same kind of answer more slowly or differently, so only the gradient comparisons notice. Training that still works is not evidence that the gradient is right.

There is no compiler-mutation run for this chapter: no compiler code changed.

The full suite:

-- Testing: 138 tests, 4 workers --
PASS: IR Forge / Mountain Goat :: attention33/examples.mlir (1 of 138)
PASS: IR Forge / Mountain Goat :: attention33/against-reference.mlir (2 of 138)
PASS: IR Forge / Mountain Goat :: generalize32/against-reference.mlir (3 of 138)
PASS: IR Forge / Mountain Goat :: contraction28/against-definition.mlir (4 of 138)
PASS: IR Forge / Mountain Goat :: transformer30/against-reference.mlir (5 of 138)
PASS: IR Forge / Mountain Goat :: generalize32/examples.mlir (6 of 138)
PASS: IR Forge / Mountain Goat :: contraction/tensor-contraction.mlir (7 of 138)
PASS: IR Forge / Mountain Goat :: tour/gallery.mlir (8 of 138)
PASS: IR Forge / Mountain Goat :: attention29/against-reference.mlir (9 of 138)
PASS: IR Forge / Mountain Goat :: bigram31/against-reference.mlir (10 of 138)
PASS: IR Forge / Mountain Goat :: learning/gradient-descent.mlir (11 of 138)
PASS: IR Forge / Mountain Goat :: contraction28/examples.mlir (12 of 138)
PASS: IR Forge / Mountain Goat :: attention29/examples.mlir (13 of 138)
PASS: IR Forge / Mountain Goat :: bigram31/examples.mlir (14 of 138)
PASS: IR Forge / Mountain Goat :: perf/options-same-results.mlir (15 of 138)
PASS: IR Forge / Mountain Goat :: tour/tour-examples.mlir (16 of 138)
PASS: IR Forge / Mountain Goat :: loop-order/results-identical.mlir (17 of 138)
PASS: IR Forge / Mountain Goat :: loop-order/bitwise-identical.mlir (18 of 138)
PASS: IR Forge / Mountain Goat :: tour/more-examples.mlir (19 of 138)
PASS: IR Forge / Mountain Goat :: broadcast/values.mlir (20 of 138)
PASS: IR Forge / Mountain Goat :: transformer30/examples.mlir (21 of 138)
PASS: IR Forge / Mountain Goat :: ops/values.mlir (22 of 138)
PASS: IR Forge / Mountain Goat :: contraction/mismatch-trap.mlir (23 of 138)
PASS: IR Forge / Mountain Goat :: broadcast/cpp-interop-reductions.mlir (24 of 138)
PASS: IR Forge / Mountain Goat :: loop-order/cpp-ikj.mlir (25 of 138)
PASS: IR Forge / Mountain Goat :: ops/cpp-interop-ops.mlir (26 of 138)
PASS: IR Forge / Mountain Goat :: ops/elementwise-mismatch-aborts.mlir (27 of 138)
PASS: IR Forge / Mountain Goat :: execution/gpu-host-with-stub-runtime.mlir (28 of 138)
PASS: IR Forge / Mountain Goat :: frontend/more-examples.mlir (29 of 138)
PASS: IR Forge / Mountain Goat :: chapters/ch09-execution-engine.mlir (30 of 138)
PASS: IR Forge / Mountain Goat :: frontend/cpp-interop.mlir (31 of 138)
PASS: IR Forge / Mountain Goat :: ops/composed-programs.mlir (32 of 138)
PASS: IR Forge / Mountain Goat :: tour/scientific-literals.mlir (33 of 138)
PASS: IR Forge / Mountain Goat :: perf/tiling-structure.mlir (34 of 138)
PASS: IR Forge / Mountain Goat :: broadcast/dynamic-reductions.mlir (35 of 138)
PASS: IR Forge / Mountain Goat :: perf/flags-reach-clang.mlir (36 of 138)
PASS: IR Forge / Mountain Goat :: broadcast/networks.mlir (37 of 138)
PASS: IR Forge / Mountain Goat :: frontend/hello-and-static.mlir (38 of 138)
PASS: IR Forge / Mountain Goat :: broadcast/identities.mlir (39 of 138)
PASS: IR Forge / Mountain Goat :: ops/negate-and-dynamic-divide.mlir (40 of 138)
PASS: IR Forge / Mountain Goat :: ops/dynamic-ops.mlir (41 of 138)
PASS: IR Forge / Mountain Goat :: frontend/dynamic-runs.mlir (42 of 138)
PASS: IR Forge / Mountain Goat :: loop-order/flag-reaches-pass.mlir (43 of 138)
PASS: IR Forge / Mountain Goat :: frontend/runtime-mismatch-aborts.mlir (44 of 138)
PASS: IR Forge / Mountain Goat :: ops/precedence-and-identities.mlir (45 of 138)
PASS: IR Forge / Mountain Goat :: execution/dynamic-shapes-run.mlir (46 of 138)
PASS: IR Forge / Mountain Goat :: contraction28/errors.mlir (47 of 138)
PASS: IR Forge / Mountain Goat :: assert-stderr/message-on-stderr.mlir (48 of 138)
PASS: IR Forge / Mountain Goat :: execution/mismatch-aborts.mlir (49 of 138)
PASS: IR Forge / Mountain Goat :: execution/mismatch-columns-aborts.mlir (50 of 138)
PASS: IR Forge / Mountain Goat :: ops/ptx-new-ops.mlir (51 of 138)
PASS: IR Forge / Mountain Goat :: execution/static-add-runs.mlir (52 of 138)
PASS: IR Forge / Mountain Goat :: tour/tour-errors.mlir (53 of 138)
PASS: IR Forge / Mountain Goat :: dynamic-loops/unroll-runs.mlir (54 of 138)
PASS: IR Forge / Mountain Goat :: ops/matmul-mismatch-aborts.mlir (55 of 138)
PASS: IR Forge / Mountain Goat :: dynamic-loops/runtime-check-survives-transforms.mlir (56 of 138)
PASS: IR Forge / Mountain Goat :: dynamic-loops/tiling-runs.mlir (57 of 138)
PASS: IR Forge / Mountain Goat :: assert-stderr/nested-assert-runs.mlir (58 of 138)
PASS: IR Forge / Mountain Goat :: ops/division-ieee.mlir (59 of 138)
PASS: IR Forge / Mountain Goat :: assert-stderr/valid-shapes-still-pass.mlir (60 of 138)
PASS: IR Forge / Mountain Goat :: transformer30/errors.mlir (61 of 138)
PASS: IR Forge / Mountain Goat :: attention29/errors.mlir (62 of 138)
PASS: IR Forge / Mountain Goat :: chapters/ch01-pipeline.mlir (63 of 138)
PASS: IR Forge / Mountain Goat :: dynamic-loops/fusion-after-tiling.mlir (64 of 138)
PASS: IR Forge / Mountain Goat :: attention33/errors.mlir (65 of 138)
PASS: IR Forge / Mountain Goat :: assert-stderr/message-bytes-exact.mlir (66 of 138)
PASS: IR Forge / Mountain Goat :: execution/strided-view.mlir (67 of 138)
PASS: IR Forge / Mountain Goat :: execution/mixed-static-dynamic.mlir (68 of 138)
PASS: IR Forge / Mountain Goat :: broadcast/ptx-layer.mlir (69 of 138)
PASS: IR Forge / Mountain Goat :: contraction28/ir.mlir (70 of 138)
PASS: IR Forge / Mountain Goat :: gpu/pass-order-crash.mlir (71 of 138)
PASS: IR Forge / Mountain Goat :: frontend/front-end-errors.mlir (72 of 138)
PASS: IR Forge / Mountain Goat :: frontend/ptx-from-source.mlir (73 of 138)
PASS: IR Forge / Mountain Goat :: ffn34/errors.mlir (74 of 138)
PASS: IR Forge / Mountain Goat :: bigram31/errors.mlir (75 of 138)
PASS: IR Forge / Mountain Goat :: generalize32/errors.mlir (76 of 138)
PASS: IR Forge / Mountain Goat :: dynamic-loops/fusion-needs-constant-trip-counts.mlir (77 of 138)
PASS: IR Forge / Mountain Goat :: gpu/host-runtime-calls.mlir (78 of 138)
PASS: IR Forge / Mountain Goat :: gpu/ptx-index32.mlir (79 of 138)
PASS: IR Forge / Mountain Goat :: gpu/ptx-default.mlir (80 of 138)
PASS: IR Forge / Mountain Goat :: gpu/ptx-bare-pointers.mlir (81 of 138)
PASS: IR Forge / Mountain Goat :: gpu/host-runtime-calls-bare.mlir (82 of 138)
PASS: IR Forge / Mountain Goat :: chapters/ch02-example-programs.mlir (83 of 138)
PASS: IR Forge / Mountain Goat :: dynamic-loops/unroll-full-needs-constant-trip-count.mlir (84 of 138)
PASS: IR Forge / Mountain Goat :: ops/matmul-lowering-structure.mlir (85 of 138)
PASS: IR Forge / Mountain Goat :: loop-order/structure.mlir (86 of 138)
PASS: IR Forge / Mountain Goat :: broadcast/front-end-errors.mlir (87 of 138)
PASS: IR Forge / Mountain Goat :: ops/front-end-errors.mlir (88 of 138)
PASS: IR Forge / Mountain Goat :: frontend/emit-mlir.mlir (89 of 138)
PASS: IR Forge / Mountain Goat :: generalize32/ir.mlir (90 of 138)
PASS: IR Forge / Mountain Goat :: assert-stderr/guard-nested-assert.mlir (91 of 138)
PASS: IR Forge / Mountain Goat :: gpu/nvvm-lowering.mlir (92 of 138)
PASS: IR Forge / Mountain Goat :: gpu/cuda-comparison.mlir (93 of 138)
PASS: IR Forge / Mountain Goat :: bigram31/ir.mlir (94 of 138)
PASS: IR Forge / Mountain Goat :: attention29/ir.mlir (95 of 138)
PASS: IR Forge / Mountain Goat :: assert-stderr/static-unchanged.mlir (96 of 138)
PASS: IR Forge / Mountain Goat :: gpu/kernel-outlining.mlir (97 of 138)
PASS: IR Forge / Mountain Goat :: assert-stderr/run-twice-no-collision.mlir (98 of 138)
PASS: IR Forge / Mountain Goat :: dynamic-loops/unroll-epilogue.mlir (99 of 138)
PASS: IR Forge / Mountain Goat :: broadcast/roundtrip-and-lowering.mlir (100 of 138)
PASS: IR Forge / Mountain Goat :: dynamic-loops/fusion-declines.mlir (101 of 138)
PASS: IR Forge / Mountain Goat :: assert-stderr/rewrite-structure.mlir (102 of 138)
PASS: IR Forge / Mountain Goat :: transformer30/ir.mlir (103 of 138)
PASS: IR Forge / Mountain Goat :: dynamic-loops/tiling-structure.mlir (104 of 138)
PASS: IR Forge / Mountain Goat :: loops/tiling.mlir (105 of 138)
PASS: IR Forge / Mountain Goat :: loops/fusion.mlir (106 of 138)
PASS: IR Forge / Mountain Goat :: verifier/add-static-mismatch.mlir (107 of 138)
PASS: IR Forge / Mountain Goat :: lowering/bufferize-dynamic-add.mlir (108 of 138)
PASS: IR Forge / Mountain Goat :: loops/unroll-full.mlir (109 of 138)
PASS: IR Forge / Mountain Goat :: lowering/dynamic-add-affine.mlir (110 of 138)
PASS: IR Forge / Mountain Goat :: verifier/dynamic-accepted.mlir (111 of 138)
PASS: IR Forge / Mountain Goat :: lowering/static-add-affine.mlir (112 of 138)
PASS: IR Forge / Mountain Goat :: contraction28/verifiers.mlir (113 of 138)
PASS: IR Forge / Mountain Goat :: ops/scalar-bad-op.mlir (114 of 138)
PASS: IR Forge / Mountain Goat :: broadcast/verifier-reduce.mlir (115 of 138)
PASS: IR Forge / Mountain Goat :: ops/elementwise-static-mismatch.mlir (116 of 138)
PASS: IR Forge / Mountain Goat :: dynamic-loops/fusion-fires-when-static.mlir (117 of 138)
PASS: IR Forge / Mountain Goat :: verifier/unranked-rejected.mlir (118 of 138)
PASS: IR Forge / Mountain Goat :: broadcast/verifier-broadcast.mlir (119 of 138)
PASS: IR Forge / Mountain Goat :: ops/matmul-bad-result-shape.mlir (120 of 138)
PASS: IR Forge / Mountain Goat :: verifier/transpose-rank3.mlir (121 of 138)
PASS: IR Forge / Mountain Goat :: verifier/transpose-bad-shape.mlir (122 of 138)
PASS: IR Forge / Mountain Goat :: generalize32/verifiers.mlir (123 of 138)
PASS: IR Forge / Mountain Goat :: lowering/bufferize-static-add.mlir (124 of 138)
PASS: IR Forge / Mountain Goat :: lowering/mixed-add-one-check.mlir (125 of 138)
PASS: IR Forge / Mountain Goat :: ops/dynamic-accepted.mlir (126 of 138)
PASS: IR Forge / Mountain Goat :: ops/matmul-inner-mismatch.mlir (127 of 138)
PASS: IR Forge / Mountain Goat :: canonicalize/transpose-twice.mlir (128 of 138)
PASS: IR Forge / Mountain Goat :: loop-order/bad-option.mlir (129 of 138)
PASS: IR Forge / Mountain Goat :: canonicalize/fold-constant-add.mlir (130 of 138)
PASS: IR Forge / Mountain Goat :: lowering/dynamic-transpose-no-check.mlir (131 of 138)
PASS: IR Forge / Mountain Goat :: transformer30/verifiers.mlir (132 of 138)
PASS: IR Forge / Mountain Goat :: canonicalize/transpose-twice-dynamic-exact.mlir (133 of 138)
PASS: IR Forge / Mountain Goat :: attention29/verifiers.mlir (134 of 138)
PASS: IR Forge / Mountain Goat :: canonicalize/transpose-twice-dynamic-mixed.mlir (135 of 138)
PASS: IR Forge / Mountain Goat :: bigram31/verifiers.mlir (136 of 138)
PASS: IR Forge / Mountain Goat :: ffn34/examples.mlir (137 of 138)
PASS: IR Forge / Mountain Goat :: ffn34/against-reference.mlir (138 of 138)

Testing Time: 131.31s

Total Discovered Tests: 138
  Passed: 138 (100.00%)

Limits and what is not established

• One small task, one model of each kind. XOR of presence over three tokens, four positions, 81 inputs. The experiment shows that this attention-only model fails and this block succeeds. It does not show that attention alone always fails at XOR, nor how any of this scales.
• Block B adds three things at once. Layer normalization, the feed-forward network and the residual connection are added together. No ablation was run (feed-forward network alone; with and without layer norm; with and without the residual), so the experiment does not say which of the three is responsible for the success. The argument above points at the feed-forward network, and nothing here tests that.
• Training B is fragile. One seed of six fails at the learning rate used; at higher rates most fail or diverge (the Python-only sweep). That table was produced by the plain-Python implementation, which agrees with the Mountain Goat one step for step on the runs compared, but the sweep itself was not run through mgc.
• The linear-separability explanation is an argument. It is consistent with model A plateauing at 69 of 81 for five seeds, and it was not verified (for instance by inspecting the learned pooled vectors).
• The miss is a data gap. The random split left 0 0 0 0 out of the training data. A different split would change which sequence, if any, is missed; only one split and three starting weights for the gradient check were used.
• Only this model's backward pass is derived. Layer norm, ReLU and residual are now covered for a single block with one pooling head. Chapter 30's full transformer (two heads, causal masking over positions, stacked blocks, embeddings with positions, per-position outputs) still has no hand-derived backward pass; Chapter 30's weights are still random.
• The ReLU corner convention is a choice. The gradient passes at exactly zero. Another convention would change nothing visible in these runs, and was not tried.
• Finite differences: a step of 0.001 and six printed digits bound the agreement to about 4·10⁻⁴.
• Learning rate and steps: 0.5 and 600, chosen by trying (plain Python); no schedule, no stopping rule.
• Platform: x86-64 Linux, LLVM 18.1.3.

Reproducing

cd docs/part32/code && ./build.sh
cd ../../part34/code
python3 make_examples.py && ./run_examples.sh > examples_out.txt
./check_ffn.py > check_ffn_out.txt
./lr_sweep.py > lr_sweep_out.txt
./model_mutation.py > model_mutation_out.txt
cd ../../part15/code && ./run_lit.sh          # 138 tests

Chapter summary

• A block with a feed-forward network learns a task, the XOR of "1 is present" and "2 is present", that the same model without it cannot: model A plateaus at 69 of the 81 possible sequences, model B gets 80, and the only miss is 0 0 0 0, which the training data never showed.
• The backward pass through the new pieces: the ReLU passes the gradient where its input is non-negative; the residual connection sends the gradient down two paths that add; and layer normalization needs three terms, (dŷ − mean(dŷ) − ŷ·mean(dŷ·ŷ)) / sigma.
• All of it is checked against an independent per-example Python gradient and against finite differences for all 120 weights; the first finite-difference step (0.01) was too coarse and the check was repaired, which is reported.
• Training B is not reliable at every setting: one of six seeds fails at the learning rate used and most fail at larger rates (a Python-only sweep).
• Ten wrong backward passes were run through the checker; the results are in the table above.
• Three additions (layer norm, feed-forward, residual) were tested together, not separately.

Self-check questions

Each answer is collapsed; try the question first.

1. Why can model A not solve "exactly one of the tokens 1 and 2 is present", according to the argument on this page?

   Answer

   Attention pooling gives a weighted average of token value vectors, a point in a triangle. Under the picture where attention favours tokens 1 and 2 over 0, "1 only" and "2 only" land at the two ends of the line between v1 and v2 and "both" lands between them, so the labels along that line run 1, 0, 1. A linear read-out separates space with a straight line, which cannot make a label that is 1 at both ends and 0 in the middle. The page is clear that this is an argument consistent with the experiment, not a proof.

2. What does the residual connection do to the backward pass?

   Answer

   z = h + f makes h influence the loss along two routes, directly and through f. By the chain rule the gradient with respect to h is the sum of the gradients along the two routes: dh = dz + dh_ln. Dropping the direct route (a mutation tested above) leaves the gradient with only the layer-norm path.

3. Why does layer normalization's backward pass have three terms?

   Answer

   Each entry of h affects the mean, the spread sigma and, through them, every normalized entry. The terms account for each: dŷ / sigma for the direct effect on its own normalized value, −mean(dŷ) / sigma for the effect through the mean, and −ŷ · mean(dŷ · ŷ) / sigma for the effect through the spread. Together they make the gradient entries sum to zero, as they must since shifting every h by a constant changes nothing.

4. What is the derivative of ReLU at exactly 0, and what does this chapter do?

   Answer

   It does not exist: the slope is 0 to the left and 1 to the right. A brute-force estimate averages them, giving half of the upstream gradient (example 2 shows 1 for an upstream 2). Training code must choose; this chapter passes the gradient at exactly 0 (the mask is ge(pre, 0)), as most frameworks do.

5. The first finite-difference check failed for model B. Was the gradient wrong?

   Answer

   No. The gradient equalled the independent Python gradient. A plain-Python finite difference with exact losses and step 0.01 showed the same 5.6·10⁻³ disagreement, and with step 0.001 only 10⁻⁶, so the error was in the estimate (a step that straddles a ReLU corner or misjudges curvature). The check's step was changed to 0.001, which mgc's six printed digits still allow.

6. Why is "model B gets 80 of 81" more convincing than "model B gets 26 of 27 held-out sequences right"?

   Answer

   There are only 81 possible inputs, so the program can try every one: it is exhaustive, with no sampling luck and no held-out set that might happen to be easy. The 27 held-out sequences are one particular random subset; 26 of 27 says less than a full census of the input space. It also identifies exactly which sequence fails.

7. Why does the page report a Python-only sweep, and what does it show that the single run does not?

   Answer

   Running every seed and learning rate through mgc would take hours (model B's program is 6,000 lines and about 100 seconds to build). The Python implementation agrees with the Mountain Goat one step for step, so it is a fair stand-in for a sweep. It shows that model B's success depends on the start and the learning rate (five of six seeds at 0.5, six of six at 0.25, one of six at 1.0), which one run cannot show, and that model A fails for every seed.

8. Why did the experiment not isolate which of layer normalization, the feed-forward network and the residual connection matters?

   Answer

   Model B adds all three at once and is compared with a model with none of them. Isolating one would need further models (for example, the feed-forward network without layer normalization, or layer normalization without the feed-forward network), each with its own backward pass checked and its own training runs. That was not done, and the page says so; the argument about the ReLU's bend points at the feed-forward network, but the experiment does not test that.

9. How is 0 0 0 0 the model's only error, and what does it have in common with the failures in Chapters 32 and 33?

   Answer

   It is the only one of the 81 sequences that model B classifies wrongly, and the training set (a random 54 of the 81) happened not to contain it, so the model never saw the case "neither token present" and extrapolates confidently wrong. In Chapter 32 the held-out pairs that failed were the ones the training pairs never showed; in Chapter 33 4 4 4 4 failed because label 4 never occurred in the data. In each, the failure is an input outside what training covered.

10. What would be needed to train Chapter 30's full transformer with these methods?

    Answer

    The backward passes of the parts still missing: the position matrix (a fixed input, so no gradient), two attention heads with their output split, causal masking over positions (the gradient through the masked softmax is the same formula, with masked entries contributing zero), per-position outputs and a loss over all positions, and several stacked blocks (the residual structure makes the gradient flow through each block's two paths, as here). Layer normalization, ReLU, residuals and single-head attention are now derived and checked; the remaining work is assembling them, with a finite-difference check of every new weight.