36. Training Chapter 30's Transformer: Two Heads, Two Blocks, and a Result That Is Not Better¶

Mountain goats on the mountain at dawn

What you will understand: how to train, in Mountain Goat, the architecture of Chapter 30's transformer: two blocks, each with two attention heads and a feed-forward network, pre-norm residual connections and a final layer norm, as a causal language model on Chapter 35's task. Chapter 35 stopped at one block and one head. The new backward-pass pieces are small, but they are the ones that were missing: gradients that fan out to two heads and add, and a gradient that is carried back through two blocks. The model has 1,216 weights in 36 matrices. All 36 gradient matrices are checked against an independent Python calculation and against finite differences, and the 200-step training run agrees with the reference at every checkpoint. Then the page reports what the experiment shows, which is not that the bigger model is better: it fits its 16 training sequences and generalizes much worse than Chapter 35's smaller model (a Python-only sweep measures how much, and what more data does to it).

What you need to know first: Chapter 35 (this chapter reuses its task, data, stacking and mask unchanged; read it first), Chapter 30 (the architecture; its model is only shaped like this one, see Limits), Chapter 34 and Chapter 33 (the backward-pass rules). 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 ../../part36/code
python3 make_examples.py                          # (re)writes the examples from lm2_lib.py
../../part32/code/mgc run examples/03_train_10_steps.mg    # 10 training steps (well under a minute)
../../part32/code/mgc run examples/04_train_200_steps.mg   # the full 200 steps: a 29,552-line program; about 40 minutes in all (see below)
./run_examples.sh > examples_out.txt              # every example and every error example on this page (38 minutes, measured while other jobs shared the machine)
./check_lm2.py                                    # the quick check: gradients, finite differences, 10 training steps (about 45 seconds)
./check_lm2.py --full > check_lm2_out.txt         # the page's 200 steps against the reference (39 minutes 35 seconds, measured)
./sensitivity.py > sensitivity_out.txt            # Python only: how sensitive is training to a 1e-12 change? (a minute or two)
./sweep.py > sweep_out.txt                        # Python only: learning rates, one block against two, 16 against 32 training sequences (several minutes)
./model_mutation.py > model_mutation_out.txt      # runs 21 WRONG programs through the lean check (roughly half an hour)
./fd_blindness.sh > fd_blindness_out.txt          # two of those through the full finite-difference check (about 4 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 programs run against the checker, where "caught" is the wanted result and "NOT CAUGHT" would be a gap (one is reported, and explained). The FAIL lines in fd_blindness_out.txt are likewise the point of that file.

What changes from Chapter 35¶

	Chapter 35	this chapter
blocks	1	2 (each: layer norm, attention, residual, layer norm, feed-forward, residual)
attention heads per block	1 (width 8)	2 (width 4 each)
matrices	16	36
weights	648	1,216
task, data, stacking (16 sequences, 112 rows), mask, loss		the same
learning rate, steps	1.0, 200	0.5, 200

The architecture is Chapter 30's: pre-norm, two heads whose outputs are brought back to the model width by one output matrix each and added (Chapter 30 notes that concatenating the heads and multiplying by one output matrix equals adding each head times its own half of that matrix), a ReLU feed-forward network, a final layer norm. The widths are Chapter 35's (model width 8, feed-forward width 16, vocabulary 4), not Chapter 30's (4, 8, 5): Chapter 30's literal model has four tokens and width 4, too small for this task (see Limits).

Two heads: gradients that fan out and add¶

Both heads read the same normalized input u1. So u1 influences the loss along six routes (a query, a key and a value path for each head), and its gradient is the sum of what comes back along them:

du1 = dq1 @ wq1ᵀ + dk1 @ wk1ᵀ + dv1 @ wv1ᵀ + dq2 @ wq2ᵀ + dk2 @ wk2ᵀ + dv2 @ wv2ᵀ

Everything else in a head is Chapter 35's backward pass, run once per head with that head's own a, q, k, v and wo. Here is the fan-out rule at a size you can check by hand: one matrix u feeds two maps, and the gradient with respect to u is the sum of the gradients along the two routes (the first matrix), compared with brute force:

# Fan-out: one matrix u feeds TWO different maps (as a block's normalised input feeds two attention heads). If the loss is
#     L = sum(r1 * (u @ wa)) + sum(r2 * (u @ wb))
# for fixed matrices r1, r2, then u affects L along two routes, and the gradient with respect to u is the SUM of the gradients along them:
#     dL/du = r1 @ transpose(wa) + r2 @ transpose(wb)
# (the same rule that makes a residual connection add, and that adds the three routes q, k, v into one layer norm in Chapter 35).
# The first matrix is that formula. The next four are two entries of u, found by brute force: nudge the entry by 0.01 up and down and watch L change.
def loss(u: tensor[2x2], wa: tensor[2x2], wb: tensor[2x2], r1: tensor[2x2], r2: tensor[2x2]) = col_sum(row_sum(r1 * (u @ wa))) + col_sum(row_sum(r2 * (u @ wb)))
let u = [[1, 2], [3, 4]]
let wa = [[1, 0], [2, 1]]
let wb = [[0, 3], [1, 1]]
let r1 = [[1, 2], [0, 1]]
let r2 = [[2, 0], [1, 1]]
print r1 @ transpose(wa) + r2 @ transpose(wb)
print (loss(u + [[0.01, 0], [0, 0]], wa, wb, r1, r2) - loss(u - [[0.01, 0], [0, 0]], wa, wb, r1, r2)) / 0.02
print (loss(u + [[0, 0.01], [0, 0]], wa, wb, r1, r2) - loss(u - [[0, 0.01], [0, 0]], wa, wb, r1, r2)) / 0.02
print (loss(u + [[0, 0], [0.01, 0]], wa, wb, r1, r2) - loss(u - [[0, 0], [0.01, 0]], wa, wb, r1, r2)) / 0.02
print (loss(u + [[0, 0], [0, 0.01]], wa, wb, r1, r2) - loss(u - [[0, 0], [0, 0.01]], wa, wb, r1, r2)) / 0.02
# The four brute-force numbers are the four entries of the formula's matrix, read in the order (0,0), (0,1), (1,0), (1,1). Using only one of the two terms would be wrong
# for every entry.

################ mgc run 01_two_paths_add.mg
  | # Fan-out: one matrix u feeds TWO different maps (as a block's normalised input feeds two attention heads). If the loss is
  | #     L = sum(r1 * (u @ wa)) + sum(r2 * (u @ wb))
  | # for fixed matrices r1, r2, then u affects L along two routes, and the gradient with respect to u is the SUM of the gradients along them:
  | #     dL/du = r1 @ transpose(wa) + r2 @ transpose(wb)
  | # (the same rule that makes a residual connection add, and that adds the three routes q, k, v into one layer norm in Chapter 35).
  | # The first matrix is that formula. The next four are two entries of u, found by brute force: nudge the entry by 0.01 up and down and watch L change.
  | def loss(u: tensor[2x2], wa: tensor[2x2], wb: tensor[2x2], r1: tensor[2x2], r2: tensor[2x2]) = col_sum(row_sum(r1 * (u @ wa))) + col_sum(row_sum(r2 * (u @ wb)))
  | let u = [[1, 2], [3, 4]]
  | let wa = [[1, 0], [2, 1]]
  | let wb = [[0, 3], [1, 1]]
  | let r1 = [[1, 2], [0, 1]]
  | let r2 = [[2, 0], [1, 1]]
  | print r1 @ transpose(wa) + r2 @ transpose(wb)
  | print (loss(u + [[0.01, 0], [0, 0]], wa, wb, r1, r2) - loss(u - [[0.01, 0], [0, 0]], wa, wb, r1, r2)) / 0.02
  | print (loss(u + [[0, 0.01], [0, 0]], wa, wb, r1, r2) - loss(u - [[0, 0.01], [0, 0]], wa, wb, r1, r2)) / 0.02
  | print (loss(u + [[0, 0], [0.01, 0]], wa, wb, r1, r2) - loss(u - [[0, 0], [0.01, 0]], wa, wb, r1, r2)) / 0.02
  | print (loss(u + [[0, 0], [0, 0.01]], wa, wb, r1, r2) - loss(u - [[0, 0], [0, 0.01]], wa, wb, r1, r2)) / 0.02
  | # The four brute-force numbers are the four entries of the formula's matrix, read in the order (0,0), (0,1), (1,0), (1,1). Using only one of the two terms would be wrong
  | # for every entry.
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data = 
[[1,   6], 
 [3,   3]]
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 = 
[[6]]
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 = 
[[3]]
  ----- exit status 0

r1 @ wa.T = [[1,4],[0,1]] and r2 @ wb.T = [[0,2],[3,2]] add to [[1,6],[3,3]], and the four brute-force numbers are 1, 6, 3, 3.

Two blocks: the gradient goes back through both¶

The backward pass starts at the loss, goes through the final layer norm, and then runs block 2's backward pass, then block 1's: each block receives the gradient of its output and produces the gradient of its input, which is the next block's output gradient. The residual connections keep a direct route open: with a scale in place of each sub-layer so that every number is checkable (x1 = x + 0.5 x, x2 = x1 − 0.25 x1), the gradient is (1 + 0.5)(1 − 0.25) = 1.125, found by the backward recurrence and by brute force:

# Two blocks in a row, each with a residual connection: x1 = x + f1(x), x2 = x1 + f2(x1). Here f1 and f2 are just a scale (x * w1 and x1 * w2), so every
# quantity is a 1 x 3 row and the answer can be checked by hand. By the chain rule dx2/dx = (1 + w1) * (1 + w2): the identity path of each residual lets the gradient
# through unchanged, and each block adds its own contribution. The backward pass works from the last block to the first, one residual at a time:
#     d1 = d2 * (1 + w2)          (gradient with respect to x1)
#     d0 = d1 * (1 + w1)          (gradient with respect to x)
# Here w1 = 0.5 and w2 = -0.25, so the factor is 1.5 * 0.75 = 1.125. The brute-force line nudges x by 0.01 and compares.
def forward(x: tensor[1x3]) = (x + x * 0.5) + (x + x * 0.5) * -0.25
let x = [[1, 2, 3]]
let d2 = [[1, 1, 1]]
let d1 = d2 + d2 * -0.25
let d0 = d1 + d1 * 0.5
print d0
print (forward(x + [[0.01, 0, 0]]) - forward(x - [[0.01, 0, 0]])) / 0.02

################ mgc run 02_two_blocks_chain.mg
  | # Two blocks in a row, each with a residual connection: x1 = x + f1(x), x2 = x1 + f2(x1). Here f1 and f2 are just a scale (x * w1 and x1 * w2), so every
  | # quantity is a 1 x 3 row and the answer can be checked by hand. By the chain rule dx2/dx = (1 + w1) * (1 + w2): the identity path of each residual lets the gradient
  | # through unchanged, and each block adds its own contribution. The backward pass works from the last block to the first, one residual at a time:
  | #     d1 = d2 * (1 + w2)          (gradient with respect to x1)
  | #     d0 = d1 * (1 + w1)          (gradient with respect to x)
  | # Here w1 = 0.5 and w2 = -0.25, so the factor is 1.5 * 0.75 = 1.125. The brute-force line nudges x by 0.01 and compares.
  | def forward(x: tensor[1x3]) = (x + x * 0.5) + (x + x * 0.5) * -0.25
  | let x = [[1, 2, 3]]
  | let d2 = [[1, 1, 1]]
  | let d1 = d2 + d2 * -0.25
  | let d0 = d1 + d1 * 0.5
  | print d0
  | print (forward(x + [[0.01, 0, 0]]) - forward(x - [[0.01, 0, 0]])) / 0.02
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data = 
[[1.125,   1.125,   1.125]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data = 
[[1.125,   0,   0]]
  ----- exit status 0

The definitions¶

One training step as Mountain Goat text (the forward pass of both blocks, the backward pass of both blocks, the update). In the generated training programs each step's names carry a suffix and every step is unrolled, which is why the 200-step program has 29,552 lines:

# Chapter 30's architecture as a causal language model, with its hand-derived backward pass: two blocks, each with TWO attention heads of width 4 and a feed-forward
# network of width 16, pre-norm residual connections, a final layer norm. 16 sequences of 7 tokens are stacked into one matrix of 112 rows; attention is over all 112
# rows with a mask that allows a position to see only earlier positions of its OWN sequence.
# (This listing is one training step with every step's suffix removed: the parameters are named after their matrices, e.g. wq1_2 is block 2's query matrix for head 1.
# In the training programs each step's lets carry a suffix such as _s7, and the parameters are numbered per step, e.g. wq1_2_v200.)

def softmax_rows(m: tensor[112x112]) = exp(m - row_max(m)) / row_sum(exp(m - row_max(m)))
def ln_sigma(h: tensor[112x8]) = sqrt(row_mean((h - row_mean(h)) * (h - row_mean(h))) + 1e-05)
def ln_hat(h: tensor[112x8]) = (h - row_mean(h)) / ln_sigma(h)
def ln_out(h: tensor[112x8], g: tensor[1x8], t: tensor[1x8]) = ln_hat(h) * g + t
def ln_back(d: tensor[112x8], h: tensor[112x8], g: tensor[1x8]) = ((d * g) - row_mean(d * g) - ln_hat(h) * row_mean((d * g) * ln_hat(h))) / ln_sigma(h)
def log_softmax(z: tensor[112x4]) = z - row_max(z) - log(row_sum(exp(z - row_max(z))))

# ---- the forward pass, one let per step
let x0_ = x @ e + ps
let u1_b1_ = ln_out(x0_, g1_1, n1_1)
let q1_b1_ = u1_b1_ @ wq1_1
let k1_b1_ = u1_b1_ @ wk1_1
let v1_b1_ = u1_b1_ @ wv1_1
let a1_b1_ = softmax_rows(q1_b1_ @ transpose(k1_b1_) * 0.5 + mk)
let o1_b1_ = a1_b1_ @ v1_b1_
let q2_b1_ = u1_b1_ @ wq2_1
let k2_b1_ = u1_b1_ @ wk2_1
let v2_b1_ = u1_b1_ @ wv2_1
let a2_b1_ = softmax_rows(q2_b1_ @ transpose(k2_b1_) * 0.5 + mk)
let o2_b1_ = a2_b1_ @ v2_b1_
let x1_b1_ = x0_ + o1_b1_ @ wo1_1 + o2_b1_ @ wo2_1
let u2_b1_ = ln_out(x1_b1_, g2_1, n2_1)
let pre_b1_ = u2_b1_ @ w1_1 + c1_1
let r_b1_ = relu(pre_b1_)
let x2_b1_ = x1_b1_ + (r_b1_ @ w2_1 + c2_1)
let u1_b2_ = ln_out(x2_b1_, g1_2, n1_2)
let q1_b2_ = u1_b2_ @ wq1_2
let k1_b2_ = u1_b2_ @ wk1_2
let v1_b2_ = u1_b2_ @ wv1_2
let a1_b2_ = softmax_rows(q1_b2_ @ transpose(k1_b2_) * 0.5 + mk)
let o1_b2_ = a1_b2_ @ v1_b2_
let q2_b2_ = u1_b2_ @ wq2_2
let k2_b2_ = u1_b2_ @ wk2_2
let v2_b2_ = u1_b2_ @ wv2_2
let a2_b2_ = softmax_rows(q2_b2_ @ transpose(k2_b2_) * 0.5 + mk)
let o2_b2_ = a2_b2_ @ v2_b2_
let x1_b2_ = x2_b1_ + o1_b2_ @ wo1_2 + o2_b2_ @ wo2_2
let u2_b2_ = ln_out(x1_b2_, g2_2, n2_2)
let pre_b2_ = u2_b2_ @ w1_2 + c1_2
let r_b2_ = relu(pre_b2_)
let x2_b2_ = x1_b2_ + (r_b2_ @ w2_2 + c2_2)
let uf_ = ln_out(x2_b2_, gf, nf)
let lg_ = uf_ @ wu
# ---- the loss, and the number of correct predictions
# loss    = col_sum(row_sum(y * log_softmax(lg_))) * -0.015625
# correct = col_sum(row_sum(ge(lg_, row_max(lg_)) * y))

# ---- the backward pass, one let per step, working backwards from the loss
let dl_ = (exp(log_softmax(lg_)) * mr - y) * 0.015625
let Gwu_ = transpose(uf_) @ dl_
let duf_ = dl_ @ transpose(wu)
let Gnf_ = col_sum(duf_)
let Ggf_ = col_sum(duf_ * ln_hat(x2_b2_))
let dX2_ = ln_back(duf_, x2_b2_, gf)
let Gc2_2_ = col_sum(dX2_)
let Gw2_2_ = transpose(r_b2_) @ dX2_
let dpre_b2_ = (dX2_ @ transpose(w2_2)) * ge(pre_b2_, [[0]])
let Gc1_2_ = col_sum(dpre_b2_)
let Gw1_2_ = transpose(u2_b2_) @ dpre_b2_
let du2_b2_ = dpre_b2_ @ transpose(w1_2)
let Gn2_2_ = col_sum(du2_b2_)
let Gg2_2_ = col_sum(du2_b2_ * ln_hat(x1_b2_))
let dx1_b2_ = dX2_ + ln_back(du2_b2_, x1_b2_, g2_2)
let Gwo1_2_ = transpose(o1_b2_) @ dx1_b2_
let do1_b2_ = dx1_b2_ @ transpose(wo1_2)
let dv1_b2_ = transpose(a1_b2_) @ do1_b2_
let da1_b2_ = do1_b2_ @ transpose(v1_b2_)
let ds1_b2_ = a1_b2_ * (da1_b2_ - row_sum(a1_b2_ * da1_b2_))
let dq1_b2_ = ds1_b2_ @ k1_b2_ * 0.5
let dk1_b2_ = transpose(ds1_b2_) @ q1_b2_ * 0.5
let Gwq1_2_ = transpose(u1_b2_) @ dq1_b2_
let Gwk1_2_ = transpose(u1_b2_) @ dk1_b2_
let Gwv1_2_ = transpose(u1_b2_) @ dv1_b2_
let Gwo2_2_ = transpose(o2_b2_) @ dx1_b2_
let do2_b2_ = dx1_b2_ @ transpose(wo2_2)
let dv2_b2_ = transpose(a2_b2_) @ do2_b2_
let da2_b2_ = do2_b2_ @ transpose(v2_b2_)
let ds2_b2_ = a2_b2_ * (da2_b2_ - row_sum(a2_b2_ * da2_b2_))
let dq2_b2_ = ds2_b2_ @ k2_b2_ * 0.5
let dk2_b2_ = transpose(ds2_b2_) @ q2_b2_ * 0.5
let Gwq2_2_ = transpose(u1_b2_) @ dq2_b2_
let Gwk2_2_ = transpose(u1_b2_) @ dk2_b2_
let Gwv2_2_ = transpose(u1_b2_) @ dv2_b2_
let du1_b2_ = dq1_b2_ @ transpose(wq1_2) + dk1_b2_ @ transpose(wk1_2) + dv1_b2_ @ transpose(wv1_2) + dq2_b2_ @ transpose(wq2_2) + dk2_b2_ @ transpose(wk2_2) + dv2_b2_ @ transpose(wv2_2)
let Gn1_2_ = col_sum(du1_b2_)
let Gg1_2_ = col_sum(du1_b2_ * ln_hat(x2_b1_))
let dX1_ = dx1_b2_ + ln_back(du1_b2_, x2_b1_, g1_2)
let Gc2_1_ = col_sum(dX1_)
let Gw2_1_ = transpose(r_b1_) @ dX1_
let dpre_b1_ = (dX1_ @ transpose(w2_1)) * ge(pre_b1_, [[0]])
let Gc1_1_ = col_sum(dpre_b1_)
let Gw1_1_ = transpose(u2_b1_) @ dpre_b1_
let du2_b1_ = dpre_b1_ @ transpose(w1_1)
let Gn2_1_ = col_sum(du2_b1_)
let Gg2_1_ = col_sum(du2_b1_ * ln_hat(x1_b1_))
let dx1_b1_ = dX1_ + ln_back(du2_b1_, x1_b1_, g2_1)
let Gwo1_1_ = transpose(o1_b1_) @ dx1_b1_
let do1_b1_ = dx1_b1_ @ transpose(wo1_1)
let dv1_b1_ = transpose(a1_b1_) @ do1_b1_
let da1_b1_ = do1_b1_ @ transpose(v1_b1_)
let ds1_b1_ = a1_b1_ * (da1_b1_ - row_sum(a1_b1_ * da1_b1_))
let dq1_b1_ = ds1_b1_ @ k1_b1_ * 0.5
let dk1_b1_ = transpose(ds1_b1_) @ q1_b1_ * 0.5
let Gwq1_1_ = transpose(u1_b1_) @ dq1_b1_
let Gwk1_1_ = transpose(u1_b1_) @ dk1_b1_
let Gwv1_1_ = transpose(u1_b1_) @ dv1_b1_
let Gwo2_1_ = transpose(o2_b1_) @ dx1_b1_
let do2_b1_ = dx1_b1_ @ transpose(wo2_1)
let dv2_b1_ = transpose(a2_b1_) @ do2_b1_
let da2_b1_ = do2_b1_ @ transpose(v2_b1_)
let ds2_b1_ = a2_b1_ * (da2_b1_ - row_sum(a2_b1_ * da2_b1_))
let dq2_b1_ = ds2_b1_ @ k2_b1_ * 0.5
let dk2_b1_ = transpose(ds2_b1_) @ q2_b1_ * 0.5
let Gwq2_1_ = transpose(u1_b1_) @ dq2_b1_
let Gwk2_1_ = transpose(u1_b1_) @ dk2_b1_
let Gwv2_1_ = transpose(u1_b1_) @ dv2_b1_
let du1_b1_ = dq1_b1_ @ transpose(wq1_1) + dk1_b1_ @ transpose(wk1_1) + dv1_b1_ @ transpose(wv1_1) + dq2_b1_ @ transpose(wq2_1) + dk2_b1_ @ transpose(wk2_1) + dv2_b1_ @ transpose(wv2_1)
let Gn1_1_ = col_sum(du1_b1_)
let Gg1_1_ = col_sum(du1_b1_ * ln_hat(x0_))
let dX0_ = dx1_b1_ + ln_back(du1_b1_, x0_, g1_1)
let Ge_ = transpose(x) @ dX0_

# ---- the update, one per matrix (lr = 0.5)
let e_new = e - Ge_ * 0.5
let g1_1_new = g1_1 - Gg1_1_ * 0.5
let n1_1_new = n1_1 - Gn1_1_ * 0.5
let wq1_1_new = wq1_1 - Gwq1_1_ * 0.5
let wk1_1_new = wk1_1 - Gwk1_1_ * 0.5
let wv1_1_new = wv1_1 - Gwv1_1_ * 0.5
let wo1_1_new = wo1_1 - Gwo1_1_ * 0.5
let wq2_1_new = wq2_1 - Gwq2_1_ * 0.5
let wk2_1_new = wk2_1 - Gwk2_1_ * 0.5
let wv2_1_new = wv2_1 - Gwv2_1_ * 0.5
let wo2_1_new = wo2_1 - Gwo2_1_ * 0.5
let g2_1_new = g2_1 - Gg2_1_ * 0.5
let n2_1_new = n2_1 - Gn2_1_ * 0.5
let w1_1_new = w1_1 - Gw1_1_ * 0.5
let c1_1_new = c1_1 - Gc1_1_ * 0.5
let w2_1_new = w2_1 - Gw2_1_ * 0.5
let c2_1_new = c2_1 - Gc2_1_ * 0.5
let g1_2_new = g1_2 - Gg1_2_ * 0.5
let n1_2_new = n1_2 - Gn1_2_ * 0.5
let wq1_2_new = wq1_2 - Gwq1_2_ * 0.5
let wk1_2_new = wk1_2 - Gwk1_2_ * 0.5
let wv1_2_new = wv1_2 - Gwv1_2_ * 0.5
let wo1_2_new = wo1_2 - Gwo1_2_ * 0.5
let wq2_2_new = wq2_2 - Gwq2_2_ * 0.5
let wk2_2_new = wk2_2 - Gwk2_2_ * 0.5
let wv2_2_new = wv2_2 - Gwv2_2_ * 0.5
let wo2_2_new = wo2_2 - Gwo2_2_ * 0.5
let g2_2_new = g2_2 - Gg2_2_ * 0.5
let n2_2_new = n2_2 - Gn2_2_ * 0.5
let w1_2_new = w1_2 - Gw1_2_ * 0.5
let c1_2_new = c1_2 - Gc1_2_ * 0.5
let w2_2_new = w2_2 - Gw2_2_ * 0.5
let c2_2_new = c2_2 - Gc2_2_ * 0.5
let gf_new = gf - Ggf_ * 0.5
let nf_new = nf - Gnf_ * 0.5
let wu_new = wu - Gwu_ * 0.5

As in Chapter 35 every intermediate result is bound with its own let (a nested-function version would recompute shared quantities and exhaust memory, as Chapter 35 found). The names carry the block and the head: a2_b1 is head 2's attention weights in block 1; dX1 is the gradient with respect to the input of block 2, which is the output of block 1.

The program is large and the compiler is slow on it. Compiling and running the 29,552-line program takes about 40 minutes in my runs: check_lm2.py --full took 39 minutes 35 seconds, and the whole run_examples.sh 38 minutes 47 seconds, with other jobs running on the same machine for part of both. The compile stage mg-opt ran at 100% of a core and about 650 MB for over thirty minutes before the program started. Chapter 35's 12,455-line program took about 8 minutes in total. No attempt was made to find out why the cost grows faster than the size; it is a property of this compiler on this program shape, not of the method.

Training¶

examples/04_train_200_steps.mg trains with learning rate 0.5 from seeded starting weights for 200 steps and prints, as in Chapter 35, the loss and the number correct (of 64) on the training sequences at checkpoints 0, 1, 10, 25, 50, 100, 150 and 200; then the held-out loss and count; then the number correct on each of four chunks of 16 of all 64 possible sequences; then four attention patterns of the first training sequence: block 1 head 1, block 1 head 2, block 2 head 1, block 2 head 2. First, the same program stopped after 10 steps (03_train_10_steps.mg, which the test suite runs):

################ mgc run 03_train_10_steps.mg
  | (1817 lines: the definitions are in lm2_lib.py; the file holds the data, the unrolled training steps and the prints, written by make_examples.py)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.59625]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[10]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.35291]]
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 = 
[[1.03453]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[39]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.54478]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[16]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[12]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[33]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[28]]
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 = [7, 7] strides = [7, 1] data = 
[[1,   0,   0,   0,   0,   0,   0], 
 [0.60104,   0.39896,   0,   0,   0,   0,   0], 
 [0.459574,   0.282028,   0.258398,   0,   0,   0,   0], 
 [0.338925,   0.230342,   0.208872,   0.221862,   0,   0,   0], 
 [0.239208,   0.193859,   0.182488,   0.184731,   0.199714,   0,   0], 
 [0.168839,   0.162091,   0.163213,   0.163396,   0.166391,   0.176069,   0], 
 [0.132422,   0.133174,   0.139555,   0.143101,   0.146304,   0.150964,   0.154479]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [7, 7] strides = [7, 1] data = 
[[1,   0,   0,   0,   0,   0,   0], 
 [0.515228,   0.484772,   0,   0,   0,   0,   0], 
 [0.332976,   0.333247,   0.333777,   0,   0,   0,   0], 
 [0.248753,   0.248385,   0.249346,   0.253516,   0,   0,   0], 
 [0.203628,   0.198616,   0.194926,   0.196566,   0.206264,   0,   0], 
 [0.177119,   0.169132,   0.158681,   0.155926,   0.162059,   0.177083,   0], 
 [0.156988,   0.147569,   0.134258,   0.129844,   0.133406,   0.143958,   0.153977]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [7, 7] strides = [7, 1] data = 
[[1,   0,   0,   0,   0,   0,   0], 
 [0.478012,   0.521988,   0,   0,   0,   0,   0], 
 [0.366494,   0.368798,   0.264708,   0,   0,   0,   0], 
 [0.337623,   0.303782,   0.199075,   0.15952,   0,   0,   0], 
 [0.324793,   0.259084,   0.162537,   0.132672,   0.120913,   0,   0], 
 [0.289552,   0.214927,   0.142177,   0.120911,   0.11334,   0.119093,   0], 
 [0.207241,   0.169285,   0.134381,   0.120593,   0.113184,   0.116499,   0.138818]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [7, 7] strides = [7, 1] data = 
[[1,   0,   0,   0,   0,   0,   0], 
 [0.438909,   0.561091,   0,   0,   0,   0,   0], 
 [0.242564,   0.322218,   0.435218,   0,   0,   0,   0], 
 [0.186377,   0.228079,   0.280711,   0.304833,   0,   0,   0], 
 [0.166902,   0.185844,   0.207247,   0.217539,   0.222468,   0,   0], 
 [0.156353,   0.159587,   0.168292,   0.171906,   0.173046,   0.170816,   0], 
 [0.131624,   0.133841,   0.147682,   0.151921,   0.151635,   0.145514,   0.137784]]
  ----- exit status 0

And the full run:

################ mgc run 04_train_200_steps.mg
  | (29552 lines: the definitions are in lm2_lib.py; the file holds the data, the unrolled training steps and the prints, written by make_examples.py)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.59625]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[10]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.35291]]
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 = 
[[1.03453]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[39]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.71546]]
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.473335]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[49]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.35739]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[55]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.0148499]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[64]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.00314496]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[64]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[2.32697]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[40]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[39]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[51]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[51]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[43]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [7, 7] strides = [7, 1] data = 
[[1,   0,   0,   0,   0,   0,   0], 
 [0.96946,   0.0305399,   0,   0,   0,   0,   0], 
 [0.969008,   0.0222117,   0.00878067,   0,   0,   0,   0], 
 [0.557512,   0.055509,   0.0600112,   0.326968,   0,   0,   0], 
 [0.0557753,   0.028792,   0.0807993,   0.269132,   0.565501,   0,   0], 
 [0.0274193,   0.0231965,   0.0648619,   0.168411,   0.263826,   0.452285,   0], 
 [0.0556258,   0.0195863,   0.0213945,   0.0604971,   0.180867,   0.334439,   0.327591]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [7, 7] strides = [7, 1] data = 
[[1,   0,   0,   0,   0,   0,   0], 
 [0.429237,   0.570763,   0,   0,   0,   0,   0], 
 [0.386722,   0.37168,   0.241598,   0,   0,   0,   0], 
 [0.258896,   0.301081,   0.258509,   0.181513,   0,   0,   0], 
 [0.134648,   0.238009,   0.301597,   0.200764,   0.124981,   0,   0], 
 [0.103569,   0.212047,   0.290446,   0.180147,   0.0998673,   0.113924,   0], 
 [0.141362,   0.20838,   0.210643,   0.137534,   0.0889365,   0.086038,   0.127107]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [7, 7] strides = [7, 1] data = 
[[1,   0,   0,   0,   0,   0,   0], 
 [0.558696,   0.441304,   0,   0,   0,   0,   0], 
 [0.409464,   0.338723,   0.251813,   0,   0,   0,   0], 
 [0.304893,   0.271178,   0.225935,   0.197994,   0,   0,   0], 
 [0.219538,   0.211915,   0.200007,   0.187926,   0.180613,   0,   0], 
 [0.197382,   0.185836,   0.16847,   0.154779,   0.149963,   0.14357,   0], 
 [0.185421,   0.16737,   0.1414,   0.124755,   0.122285,   0.122015,   0.136753]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [7, 7] strides = [7, 1] data = 
[[1,   0,   0,   0,   0,   0,   0], 
 [0.612717,   0.387283,   0,   0,   0,   0,   0], 
 [0.432061,   0.322394,   0.245545,   0,   0,   0,   0], 
 [0.282482,   0.251039,   0.232589,   0.23389,   0,   0,   0], 
 [0.184766,   0.186436,   0.201191,   0.214806,   0.212801,   0,   0], 
 [0.17601,   0.161929,   0.158827,   0.163357,   0.170394,   0.169482,   0], 
 [0.176834,   0.142726,   0.122014,   0.118722,   0.133017,   0.138458,   0.168228]]
  ----- exit status 0

after step	0	1	10	25	50	100	150	200
training loss	1.596	1.353	1.035	0.715	0.473	0.357	0.0148	0.0031
correct of 64	10	20	39	45	49	55	64	64

It fits the training data. Loss 0.0031 and all 64 scored predictions right.
It does not generalize. On the 16 held-out sequences the loss is 2.33, worse than the 1.39 (ln 4) of guessing uniformly, and 40 of 64 predictions are right. On all 64 possible sequences the four chunks give 39, 51, 51 and 43: 184 of 256 (the 64 training predictions are among them, so 120 of the 192 predictions on sequences it never trained on, 63%). Chapter 35's model, on the same data, got 256 of 256.
The attention is not Chapter 35's clean pattern. Chapter 35's one head looked two positions back (0.98, 0.52, 0.80, 0.78, 0.79 on the diagonal two before). Here block 1 head 1 puts about 0.97 on position 0 for rows 1 and 2, then spreads (row 4: 0.57 on itself; row 5: 0.45 on itself); the other three heads spread their weight broadly (from row 3 on, no weight above 0.31; row 6 of block 2 head 1 is 0.19, 0.17, 0.14, 0.12, 0.12, 0.12, 0.14). No head looks two back. The model fits 16 sequences without finding the rule, which is the signature of memorizing; this is an interpretation of the numbers, and no experiment here (such as inspecting what the weights compute) tests it.

Is this training chaotic, like Chapter 35's? No, in this run¶

Chapter 35 found that its run was chaotic: a change of 10⁻¹² in one weight changed the step-100 loss by a factor of six. The same experiment (sensitivity.py: train twice in Python, once with one starting weight nudged by 10⁻¹²) on this model:

 step   train loss  nudged by 1e-12   held-out loss     nudged
    0     1.596249         1.596249        1.497033   1.497033
   25     0.715460         0.715460        1.434758   1.434758
   50     0.473335         0.473335        1.889289   1.889289
   75     0.341950         0.341950        1.794577   1.794577
  100     0.357390         0.357390        1.879108   1.879108
  125     0.051688         0.051688        2.015563   2.015568
  150     0.014850         0.014850        2.200396   2.200396
  175     0.005666         0.005666        2.268457   2.268457
  200     0.003145         0.003145        2.326972   2.326971

The two runs agree to about six digits through step 200 (the largest difference anywhere in the table is 5·10⁻⁶, in a held-out loss at step 125). So this model's run can be compared digit for digit at every checkpoint, and check_lm2.py --full does that (the output is below); it is the Mountain Goat program's loss at step 200, 0.00314496, against the Python run's 0.003145. This is a measurement of one run, one nudge, one weight; it does not show that this model's training is stable for every starting point.

One block or two? A Python-only sweep¶

The experiment's surprise is the generalization gap, so it was measured further. sweep.py (plain Python, not mgc) trains for 200 steps from four starting-weight seeds. The score is the number correct among the predictions on sequences not trained on (192 predictions for 16 training sequences, 128 for 32):

Part A: two blocks, 16 training sequences: per starting-weight seed 1..4: (final training loss, correct of 192 on unseen sequences)
  lr 0.25   (0.00525, 140)  (0.00598, 127)  (0.00529, 95)  (0.264, 91)
  lr  0.5   (0.00314, 120)  (0.00251, 111)  (0.00284, 115)  (0.0495, 108)
  lr  1.0   (0.0305, 135)  (0.0715, 107)  (0.243, 97)  (0.0563, 112)
Part B: one block (lr 1.0) against two blocks (lr 0.5): per seed 1..4: (final training loss, correct on unseen sequences of 192 for 16 training sequences, 128 for 32)
  16 training sequences, one block    (0.000912, 192)  (0.00334, 175)  (0.00118, 185)  (0.00419, 175)
  16 training sequences, two blocks   (0.00314, 120)  (0.00251, 111)  (0.00284, 115)  (0.0495, 108)
  32 training sequences, one block    (0.000974, 128)  (0.000925, 128)  (0.000876, 128)  (0.00105, 128)
  32 training sequences, two blocks   (0.00139, 128)  (0.00354, 108)  (0.00143, 128)  (0.00155, 125)

Part A, learning rates (two blocks, 16 sequences): none of 0.25, 0.5 and 1.0 gets near Chapter 35's scores; unseen scores run from 91 to 140 of 192 (47% to 73%), and the page's rate 0.5 gives 120, 111, 115, 108. Seed 1 at rate 0.25 does best (140) while another seed at the same rate is the worst (91), so the choice of rate is not what is limiting.
Part B, one block against two: with 16 training sequences, one block scores 192, 175, 185, 175 of 192 (91% to 100%) and two blocks 120, 111, 115, 108 (56% to 63%). With 32 training sequences the one-block model scores 128 of 128 for every seed, and the two-block model 128, 108, 128, 125: doubling the training data closes most of the gap (and not all: seed 2 still memorizes).
What this suggests, and what it does not show. The two-block, two-head model has about twice the weights (1,216 against 648) and the same 16 sequences; it fits them without finding the rule. More data helps it. What the experiment does not separate is why: the number of blocks, the number of heads, the head width, and the learning rate were all different between the two models, and none was varied alone. No weight decay or other regularizer, no longer training, and no different starting scale were tried.

Mistakes the compiler reports¶

################ mgc run errors/e1_head_output_of_the_wrong_width.mg
  | # A head of width 4 produces a 6 x 4 result; its output matrix must be 4 x 8 to bring it back to the model width 8. A 2 x 8 matrix does not fit, and the compiler says so.
  | let o = [[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]]
  | let wo = [[1, 0, 0, 0, 0, 0, 0, 0], [0, 1, 0, 0, 0, 0, 0, 0]]
  | print o @ wo
  ----- output
errors/e1_head_output_of_the_wrong_width.mg:4: error: cannot multiply (matmul) shapes 6x4 and 2x8: inner dimensions differ
  ----- exit status 1

################ mgc run errors/e2_missing_argument.mg
  | # ln_back takes three arguments (the upstream gradient, the layer's input and the scale). Two is an arity error, reported before any code exists.
  | def ln_sigma(h: tensor[2x4]) = sqrt(row_mean((h - row_mean(h)) * (h - row_mean(h))) + 0.00001)
  | def ln_hat(h: tensor[2x4]) = (h - row_mean(h)) / ln_sigma(h)
  | def ln_back(d: tensor[2x4], h: tensor[2x4], g: tensor[1x4]) = ((d * g) - row_mean(d * g) - ln_hat(h) * row_mean((d * g) * ln_hat(h))) / ln_sigma(h)
  | print ln_back([[1, 2, 3, 4], [4, 3, 2, 1]], [[1, 2, 4, 8], [8, 4, 2, 1]])
  ----- output
errors/e2_missing_argument.mg:5: error: ln_back takes 3 arguments, got 2
  ----- exit status 1

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

A head's output matrix of the wrong width, a call with a missing argument, and a residual connection whose two sides have different widths. As always, the shapes are types.

Tests¶

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

Test	What it checks
`examples`	the two small examples (worked out above) and the loss and correct count, held-out result and chunk counts of the 10-step training run
`errors`	the three mistakes give the exact messages
`against-reference`	`check_lm2.py` in its quick mode, below (about 45 seconds). The 200-step run is not in CI because it takes about 40 minutes; `check_lm2.py --full` runs it

check_lm2.py compares with an independent Python implementation (lm2_ref.py: plain lists, one sequence and one position at a time, every gradient derived element by element; the reference's own gradient was first checked against exact-double finite differences, step 10⁻⁶, for two weights of every matrix: largest difference 4·10⁻¹⁰) and checks, in order:

The gradients, for three starting-weight sets: the loss and all 36 gradient matrices equal the reference's, and for the first set the largest-gradient weight of every one of the 36 matrices agrees with a finite-difference estimate (step 0.001; the program forms the difference of the two losses itself so that mgc's six printed digits are spent on it, as in Chapter 35).
Training: loss and correct count at every checkpoint, the held-out loss and count, the four chunk counts and all four attention patterns equal the reference's; every attention row of every head sums to 1 and puts exactly 0 after the diagonal.
The experiment (--full): the 200-step run fits the training sequences (loss below 0.01, 64 of 64).

-- the hand-derived backward pass against the independent per-position gradient, and against finite differences
  ok   starting weights #1: the loss and all 36 gradient matrices equal the reference
  ok   starting weights #1: 36 weights (the largest gradient in every matrix) agree with finite differences (largest difference 0.0e+00; at this size of gradient that means all six printed digits agree)
  ok   starting weights #2: the loss and all 36 gradient matrices equal the reference
  ok   starting weights #3: the loss and all 36 gradient matrices equal the reference
-- training 200 steps against the independent Python trainer
  ok   step 0: loss 1.596249 and 10 of 64 correct equal the reference
  ok   step 1: loss 1.352912 and 20 of 64 correct equal the reference
  ok   step 10: loss 1.034531 and 39 of 64 correct equal the reference
  ok   step 25: loss 0.715460 and 45 of 64 correct equal the reference
  ok   step 50: loss 0.473335 and 49 of 64 correct equal the reference
  ok   step 100: loss 0.357390 and 55 of 64 correct equal the reference
  ok   step 150: loss 0.014850 and 64 of 64 correct equal the reference
  ok   step 200: loss 0.003145 and 64 of 64 correct equal the reference
  ok   the 16 held-out sequences: 40 of 64 correct and loss 2.326970 equal the reference
  ok   the number correct in each of the four chunks of 16 of all 64 sequences, [39, 51, 51, 43], equals the reference's [39, 51, 51, 43]
  ok   four attention patterns were printed (2 blocks x 2 heads)
  ok   the four attention patterns of the first training sequence equal the reference's
  ok   every attention row of every head sums to 1 and puts exactly 0 on every later position (the mask works)
-- the experiment
  ok   it fits its 16 training sequences: loss 0.003145, 64 of 64 predictions right
  --   held-out: 40 of 64 predictions right; all 64 sequences: 184 of 256 right (reported on the page, not asserted)
all checks pass

(largest difference 0.0e+00 means that, for these weights, the finite-difference estimate and the program's gradient agree in all six digits mgc prints; the gradient entries here are of order 10⁻².)

Are the checks good enough? Run twenty-one wrong programs.¶

model_mutation.py writes twenty-one versions of the program, each with one mistake, and runs the lean checker on each (seed 1 only; finite differences for six weights, wq1_1, wo2_2, w1_1, g2_2, emb, wout; then 10 training steps). Many of the mistakes touch only one head or one block, which is what this chapter added:

NOTE: this script deliberately runs wrong programs. '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 (everywhere): caught by 7 failing checks (the reference gradient, finite differences, the reference training run)
layer-norm backward misses the projection term (everywhere): caught by 7 failing checks (the reference gradient, finite differences, the reference training run)
layer-norm backward forgets to divide by sigma (everywhere): caught by 7 failing checks (the reference gradient, finite differences, the reference training run)
softmax backward misses the row-sum correction in head 2 of block 2 ONLY: caught by 1 failing checks (the reference gradient)
the 1/sqrt(4) score scale is missing from dq in head 1 of block 1 ONLY: caught by 2 failing checks (the reference gradient, finite differences)
dk forgets to transpose ds in head 2 of block 1 ONLY: caught by 1 failing checks (the reference gradient)
the gradient of the cross-entropy is not masked to the scored rows: caught by 7 failing checks (the reference gradient, finite differences, the reference training run)
the loss's 1/64 is 1/32 in the backward pass only: caught by 7 failing checks (the reference gradient, finite differences, the reference training run)
the relu mask is omitted in block 1 ONLY: caught by 2 failing checks (the reference gradient, finite differences)
the relu mask is the wrong way round in block 2 ONLY: caught by 2 failing checks (the reference gradient, finite differences)
the residual path is dropped from block 2's dx1 ONLY: caught by 2 failing checks (the reference gradient, finite differences)
the residual path is dropped between the blocks (block 1's input gradient): caught by 2 failing checks (the reference gradient, finite differences)
the value path of head 2 is dropped from du1 in block 2 ONLY: caught by 2 failing checks (the reference gradient, finite differences)
head 2's whole contribution is dropped from du1 in block 1 ONLY: caught by 2 failing checks (the reference gradient, finite differences)
head 2's gradient into its output matrix uses head 1's matrix (same shape) in block 1 ONLY: caught by 2 failing checks (the reference gradient, finite differences)
block 2's feed-forward backward uses block 1's w2 (same shape): caught by 2 failing checks (the reference gradient, finite differences)
the forward pass lets every position see ALL 112 rows (no mask; the backward pass is unchanged): caught by 8 failing checks (the reference gradient, the reference training run, the mask check)
the position encodings are left out of the forward pass: caught by 7 failing checks (the reference gradient, the reference training run)
the layer-norm epsilon is 0.1 in the forward AND backward pass (consistent with itself, not with the reference): caught by 7 failing checks (the reference gradient, the reference training run)
naive softmax (no row-max subtraction); algebraically identical, and the masked entries are exp(-1e9) = 0: NOT CAUGHT
the embedding's first update uses the wrong step size (0.25 instead of 0.5): caught by 4 failing checks (the reference training run)

Twenty of twenty-one are caught. What the pattern shows:

The new, per-head and per-block mistakes are all caught, including ones confined to a single head of a single block (the missing row-sum correction in head 2 of block 2, dk without the transpose in head 2 of block 1, head 2's whole contribution missing from du1 in block 1), the residual dropped between the blocks, and two cross-wirings that the shapes allow (head 2's do computed with head 1's output matrix; block 2's feed-forward backward using block 1's w2).
The reference gradient is the sharper check. The missing row-sum correction in head 2 of block 2 and dk without the transpose in head 2 of block 1 were caught by the reference gradient only. The finite-difference check in lean mode looks at six weights, so the question arises whether more weights would have caught them. fd_blindness.sh runs these two through the full finite-difference check (the largest-gradient weight of all 36 matrices):

mutant 1: softmax backward misses the row-sum correction in head 2 of block 2 only
-- the hand-derived backward pass against the independent per-position gradient, and against finite differences
  FAIL starting weights #1: the loss and all 36 gradient matrices equal the reference
       ['emb', 'g1_1', 'bt1_1', 'wq1_1', 'wk1_1', 'wv1_1', 'wo1_1', 'wv2_1', 'wo2_1', 'g2_1', 'bt2_1', 'w1_1', 'b1_1', 'w2_1', 'b2_1', 'g1_2', 'bt1_2', 'wq2_2', 'wk2_2']
  ok   starting weights #1: 36 weights (the largest gradient in every matrix) agree with finite differences (largest difference 4.7e-04)
  FAIL starting weights #2: the loss and all 36 gradient matrices equal the reference

mutant 2: dk forgets to transpose ds in head 2 of block 1 only
-- the hand-derived backward pass against the independent per-position gradient, and against finite differences
  FAIL starting weights #1: the loss and all 36 gradient matrices equal the reference
       ['emb', 'g1_1', 'bt1_1', 'wk2_1']
  FAIL starting weights #1: 36 weights (the largest gradient in every matrix) agree with finite differences (largest difference 1.9e-03)
       0.0019374829999999998

The finite differences miss the first mutant even with all 36 weights: the largest disagreement is 4.7·10⁻⁴, below the 10⁻³ threshold. A mistake in one head's softmax backward changes the gradient by less than the finite-difference check can resolve (the step of 0.001 and six printed digits bound it to about that size). The comparison with the reference, at a relative 10⁻⁵, sees it. The second mutant fails the finite-difference check (1.9·10⁻³). Finite differences alone would have let a real derivation mistake through; the reference comparison is the check that earns its keep. - Forward-only changes (no mask, position vectors left out, epsilon 0.1 in forward and backward) are caught by the reference gradient and the reference training run, not by finite differences: a gradient that matches its own wrong loss passes finite differences by construction (as in Chapters 34 and 35). - One is NOT CAUGHT: the softmax without the row-maximum subtraction. As in Chapter 35 it is algebraically the same function and gives the same numbers at these score sizes; no check uses large scores. This gap is carried over, and still not closed. - The wrong step size on the embedding's first update is caught only by the training run, since it changes training and not any gradient.

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

The full suite:

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

Testing Time: 103.49s

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

Limits and what is not established¶

Chapter 30's shape, not Chapter 30's model. The architecture (two blocks, two heads, pre-norm, feed-forward, final layer norm) is Chapter 30's, but the widths are Chapter 35's (8, 16, 4) and not Chapter 30's (4, 8, 5), and the weights are trained, not Chapter 30's seeded ones. Chapter 30's literal model has four tokens and width 4; it was not trained, and whether it could learn any task was not tried.
The generalization gap is measured, not explained. One task, one 16-sequence training set, four seeds in the sweep, one split of the data. Blocks, heads, head width and learning rate all changed together between the two models; the explanation "more weights than data, so it memorizes" is consistent with the 32-sequence result and is not tested (no regularization, no ablation, no different widths).
No ablation of the new parts. Whether two heads or two blocks is responsible for the gap was not separated; the sweep compares the whole models.
Not-chaotic is one measurement. A single 10⁻¹² nudge of a single weight left the run unchanged to six digits for 200 steps. That does not show every start is stable.
Finite differences cover 36 weights (the largest-gradient weight of each matrix) for one starting-weight set; and they cannot resolve a mistake as small as the one shown above. The gradient comparison with the reference covers all 1,216 weights for three starting sets.
The softmax-without-max gap (above): overflow robustness is untested.
Compile cost. The 200-step program takes about 40 minutes to compile and run; no attempt was made to reduce it (for instance by training fewer steps per program and carrying the weights over, which would need a way to write weights out and read them back, and the language has none).
Generation is not shown, and the attention for sequences other than the first training sequence was not printed.
Platform: x86-64 Linux, LLVM 18.1.3.

Reproducing¶

cd docs/part32/code && ./build.sh
cd ../../part36/code
python3 make_examples.py && ./run_examples.sh > examples_out.txt
./check_lm2.py --full > check_lm2_out.txt
./sensitivity.py > sensitivity_out.txt
./sweep.py > sweep_out.txt
./model_mutation.py > model_mutation_out.txt
./fd_blindness.sh > fd_blindness_out.txt
cd ../../part15/code && ./run_lit.sh          # 144 tests

Chapter summary¶

Chapter 30's architecture (two blocks, two heads each, 1,216 weights) is trained in Mountain Goat as a causal language model by hand-derived backpropagation. New in the backward pass: the gradient into the shared normalized input is the sum over six routes (query, key and value of each of two heads), and the gradient is carried back through two blocks, last first.
All 36 gradient matrices equal an independent per-position Python calculation at three starting points; the largest-gradient weight of every matrix agrees with finite differences. The 200-step training run agrees with the reference at every checkpoint (this run, unlike Chapter 35's, is not chaotic).
The result is not that bigger is better. The model fits its 16 training sequences (loss 0.0031) but gets 40 of 64 held-out predictions and 184 of 256 overall, against Chapter 35's one block at 256 of 256, and its attention shows no head that looks two positions back.
A Python-only sweep: on 16 training sequences the one-block model scores 91–100% of unseen predictions and the two-block model 56–63%; on 32 sequences the gap mostly closes (the two-block model reaches 128 of 128 for two of four seeds). The cause was not isolated.
Twenty-one wrong programs: twenty caught, one not (softmax without the row-max subtraction). One mistake in a single head's softmax backward is invisible to the finite-difference check (4.7·10⁻⁴ against a 10⁻³ threshold) and is caught only by the reference comparison.

Self-check questions¶

Each answer is collapsed; try the question first.

Why does the gradient with respect to u1 have six terms in this model and three in Chapter 35's?

Answer

u1 feeds the query, key and value maps of every head. With one head that is three routes; with two heads it is six. A quantity that influences the loss along several routes receives the sum of the gradients along them (the chain rule for fan-out), so each head contributes its own dq @ wqᵀ + dk @ wkᵀ + dv @ wvᵀ and the two are added.
What does block 2's backward pass hand to block 1's?

Answer

The gradient with respect to its input, which is block 1's output (dX1). Block 2's backward pass starts from the gradient arriving at its output (from the final layer norm), runs the feed-forward residual, then attention, and returns dx1 + ln_back(du1, …): the direct residual route plus the route through the layer norm and the heads. Block 1 treats it as the gradient at its own output. The last thing returned is the gradient with respect to the embedded input.
The two-block model reaches a training loss of 0.0031 and 64 of 64 correct. Does that show it has learned the task?

Answer

No. On the 16 held-out sequences it gets 40 of 64 and a loss of 2.33 (guessing uniformly gives 1.39), and over all 64 sequences it gets 184 of 256. Fitting the training data shows the optimization works; it does not show the rule was found. Compare Chapter 35's model, which on the same data got 256 of 256.
The sweep shows the gap closing when the training set doubles. Does that prove the gap is a data-size effect?

Answer

It supports it but does not prove it: with 32 sequences the two-block model reaches 128 of 128 for two seeds, 125 for a third and still only 108 for a fourth. Also, the one-block and two-block models differ in blocks, heads, head width and learning rate at once, and none of these was varied alone. The result says more data helps this model; it does not say the extra capacity is the cause of the gap.
One mutant (no row-sum correction in one head's softmax backward) passed the finite-difference check even over 36 weights. How can a wrong gradient pass?

Answer

The finite-difference estimate has a resolution: a step of 0.001 and six printed digits limit how small a disagreement it can see, about 10⁻³ here, and the gradient entries are of order 10⁻². A mistake confined to one head of one block changes the gradient by 4.7·10⁻⁴, below the threshold. The reference comparison works at a relative 10⁻⁵ and sees it. It is the reason this book keeps an independent reference implementation and does not rely on finite differences alone.
Chapter 35's run was chaotic and this one is not. What should you conclude about the next model?

Answer

Nothing without measuring: it is a property of each model and each starting point, found by sensitivity.py, and cheap to measure (a minute or two in plain Python). It decides whether a check may compare digit for digit at every checkpoint or must compare claims after the transition.
Why does this program take about 40 minutes, when Chapter 35's took about eight?

Answer

It is a larger program (29,552 lines against 12,455), and the compile stage (mg-opt) ran at 100% of a core for over thirty minutes in an observed run, so the cost grows faster than the size. Why was not investigated. The measurement is from one machine and one run, with other jobs sharing the machine for part of it.