32. Does It Generalize? Held-Out Data, Overfitting, Weight Decay and Greedy Generation

What you will understand: the difference between a model that fits its data and one that generalizes, shown on the bigram model of Chapter 31. Chapter 31 measured the model only on the pairs it trained on. Here some pairs are held out, and the held-out loss shows what the training loss hid: the model gets steadily better on its training pairs while getting steadily worse on pairs it has not seen, until it is worse than a model that knows nothing. Weight decay (a penalty on large weights) fixes most of it, at a price, and a penalty that is too strong for the learning rate does not merely fail: it blows up. Finally the model generates text, by repeatedly choosing its most likely next token, which needs one new operation: a comparison, ge.

What you need to know first: Chapter 31 (the bigram model, cross-entropy loss and its hand-derived gradient). Nothing else is new except the comparison ge.

Compile and run

cd docs/part32/code && ./build.sh                 # once: the newest compiler (adds mg.ge)
python3 make_examples.py                          # (re)writes the examples from generalize.mg.defs.in
./mgc run examples/02_train_and_validate.mg       # train with no weight decay; print train and held-out loss
./run_examples.sh > examples_out.txt              # every example and every error example on this page
./sweep.py > sweep_out.txt                        # the table: five weight-decay strengths, run through mgc
./check_generalize.py                             # the programs against the reference and against facts about the data
./model_mutation.py > model_mutation_out.txt      # runs 11 WRONG trainers through the checker
./mutation.sh > mutation_out.txt                  # breaks the compiler on purpose, six ways
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 divergence, FAIL lines and error messages, and why that is good

Example 4 is supposed to blow up (losses of 10⁵⁸): it shows what a too-strong penalty does. The three error examples are programs the compiler must reject. The mutation outputs are negative controls, where "caught" is the wanted result and "NOT CAUGHT" would be a gap.

The new operation: ge, a comparison

Choosing the most likely token means finding the largest entry of a row, and the language had no way to compare. ge(a, b) is 1 where a >= b and 0 where it is not, elementwise, with the same broadcasting as + - * / (Chapter 22). It is a real operation: mg.ge in the dialect with a shape verifier, a loop nest that does a floating-point compare and selects 1.0 or 0.0, and a built-in in the front end. (It is a function, not an operator, because the language has no comparison operators yet; and a scalar threshold must be written as a 1x1 matrix, such as [[0.5]].)

################ mgc run 01_ge.mg
  | # ge(a, b) compares: 1 where a >= b and 0 where it is not, element by element, with the same broadcasting as + - * /.
  | print ge([[1, 5, 3], [2, 2, 9]], [[2, 2, 2], [2, 2, 2]])       # which entries are at least 2?
  | print ge([[1, 5, 3], [2, 2, 9]], [[2], [9]])                    # a 2x1 column is stretched across the columns: row 0 against 2, row 1 against 9
  | # the one-hot row of the largest entry: compare a row with its own maximum (row_max gives a 2x1 column)
  | let scores = [[0.2, 0.7, 0.1], [0.5, 0.1, 0.4]]
  | print ge(scores, row_max(scores))
  | # a tie gives several 1s: both 0.5s reach the maximum
  | print ge([[0.5, 0.5, 0.0]], row_max([[0.5, 0.5, 0.0]]))
  | # a comparison against a number: write the number as a 1x1 matrix
  | print ge([[0.2, 0.7, 0.1]], [[0.5]])
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[0,   1,   1], 
 [1,   1,   1]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[0,   1,   1], 
 [0,   0,   1]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[0,   1,   0], 
 [1,   0,   0]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data = 
[[1,   1,   0]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data = 
[[0,   1,   0]]
  ----- exit status 0

The first matrix marks which entries are at least 2. The second compares a 2x1 column with each row (row 0 against 2, row 1 against 9). The third is the useful one: comparing a row with its own maximum (row_max gives the 2x1 column of maxima) gives a one-hot row marking where the maximum is: [0, 1, 0] and [1, 0, 0]. The fourth shows what a tie does: both 0.5s reach the maximum, so there are two 1s. The last compares with a threshold written as [[0.5]].

The split

Chapter 31 used all 14 (current, next) pairs of the text 0 1 2 0 1 3 0 1 2 0 2 3 0 1 4. Now the first 10 train the model and the last 4 are held out: used only to measure it, never to change a weight.

	pairs (current → next)
training (10)	0→1, 1→2, 2→0, 0→1, 1→3, 3→0, 0→1, 1→2, 2→0, 0→2
held out (4)	2→3, 3→0, 0→1, 1→4

Two of the held-out pairs, 2→3 and 1→4, never occur among the training pairs: in the training data token 2 is only ever followed by token 0, and token 1 never by token 4. The other two, 3→0 and 0→1, do occur. This is the whole point of the experiment: how does a model do on cases its training data never showed?

The definitions

# Does the bigram model generalise? The 14 (current, next) pairs of Chapter 31 are split: the first 10 TRAIN the model, the last 4 are HELD OUT and used
# only to measure it. x and y are one-hot rows for the training pairs (10x5); xv and yv are the same for the held-out pairs (4x5).

# the stable log-softmax of Chapter 31, once for each row count (a function's shapes are fixed)
def log_softmax10(s: tensor[10x5]) = s - row_max(s) - log(row_sum(exp(s - row_max(s))))
def log_softmax4(s: tensor[4x5]) = s - row_max(s) - log(row_sum(exp(s - row_max(s))))

# the average surprise (cross-entropy) on the training pairs and on the held-out pairs; 0.1 is 1/10 and 0.25 is 1/4
def train_loss(w: tensor[5x5], x: tensor[10x5], y: tensor[10x5]) = col_sum(row_sum(y * log_softmax10(x @ w))) * -0.1
def held_out_loss(w: tensor[5x5], xv: tensor[4x5], yv: tensor[4x5]) = col_sum(row_sum(yv * log_softmax4(xv @ w))) * -0.25

# the surprise of each held-out pair separately (a 4x1 column), to see which ones the model fails on
def held_out_surprise(w: tensor[5x5], xv: tensor[4x5], yv: tensor[4x5]) = row_sum(yv * log_softmax4(xv @ w)) * -1

# the gradient of the training loss, plus the gradient of the penalty (@LAMBDA@/2) * (the sum of the squares of all the weights), which is @LAMBDA@ * w
def gradient(w: tensor[5x5], x: tensor[10x5], y: tensor[10x5]) = transpose(x) @ (exp(log_softmax10(x @ w)) - y) * 0.1 + w * @LAMBDA@

# one step of gradient descent with learning rate 8 on the training loss plus the penalty (a "weight decay" step: every weight is pulled toward 0)
def step(w: tensor[5x5], x: tensor[10x5], y: tensor[10x5]) = w - gradient(w, x, y) * 8

# the penalised objective itself, used to check the gradient: the training loss plus (@LAMBDA@/2) times the sum of squared weights
def objective(w: tensor[5x5], x: tensor[10x5], y: tensor[10x5]) = train_loss(w, x, y) + col_sum(row_sum(w * w)) * @HALF_LAMBDA@

# the model's probabilities: row i is the chance of each token following token i
def probabilities(w: tensor[5x5]) = exp(w - row_max(w)) / row_sum(exp(w - row_max(w)))

# GREEDY GENERATION: x is a one-hot 1x5 row for the current token. x @ p picks that token's row of probabilities, and ge(row, its own maximum)
# is 1 exactly where the row reaches its maximum: the one-hot row of the most likely next token (a tie would give several 1s)
def next_token(x: tensor[1x5], p: tensor[5x5]) = ge(x @ p, row_max(x @ p))

(@LAMBDA@ and @HALF_LAMBDA@ are filled in by make_examples.py with the strength of the penalty, 0 for no weight decay.) Read it top to bottom: the loss on the 10 training pairs and on the 4 held-out pairs, written separately because a function's shapes are fixed; the surprise of each held-out pair on its own; the gradient with an extra w * λ term; the step; the penalised objective (used only to check the gradient); and next_token, which turns a one-hot row into the one-hot row of the likeliest next token with the comparison from above.

Experiment 1: train, and watch two losses

examples/02_train_and_validate.mg is the definitions with strength 0, the data, 400 unrolled step lines, and then the training loss and the held-out loss at nine checkpoints, followed by the final weights, the surprise of each held-out pair, and the probabilities:

  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.60944]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.60944]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.832556]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.48224]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.519409]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.94577]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.46728]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[2.23759]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.436194]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[2.66616]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.42598]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[3.0038]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.420915]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[3.34617]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.418396]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[3.69074]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.41714]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[4.03635]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [5, 5] strides = [5, 1] data = 
[[-2.89375,   4.89056,   3.79069,   -2.89375,   -2.89375], 
 [-2.81214,   -2.81214,   4.56525,   3.87117,   -2.81214], 
 [6.45562,   -1.61391,   -1.61391,   -1.61391,   -1.61391], 
 [5.8938,   -1.47345,   -1.47345,   -1.47345,   -1.47345], 
 [0,   0,   0,   0,   0]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 1] strides = [1, 1] data = 
[[8.07078], 
 [0.00252321], 
 [0.288306], 
 [7.78379]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [5, 5] strides = [5, 1] data = 
[[0.000311968,   0.749533,   0.249532,   0.000311968,   0.000311968], 
 [0.000416429,   0.000416429,   0.666043,   0.332708,   0.000416429], 
 [0.99875,   0.000312539,   0.000312539,   0.000312539,   0.000312539], 
 [0.99748,   0.000630007,   0.000630007,   0.000630007,   0.000630007], 
 [0.2,   0.2,   0.2,   0.2,   0.2]]
  ----- exit status 0

Read the loss pairs two at a time (training, then held-out):

after step	0	1	5	10	25	50	100	200	400
training loss	1.609	0.833	0.519	0.467	0.436	0.426	0.421	0.418	0.417
held-out loss	1.609	1.482	1.946	2.238	2.666	3.004	3.346	3.691	4.036

The training loss falls at every checkpoint, to 0.417, which is exactly what training is supposed to do. The held-out loss improves for one step (1.609 to 1.482) and then rises at every checkpoint, to 4.036. After 400 steps the model is worse on the held-out pairs than before training began: 4.036 against 1.609, the loss of a model that thinks every token equally likely. This is overfitting: the model has learned its training data too well and learned nothing that carries beyond it.

The vector of per-pair held-out surprises says why: 8.07, 0.0025, 0.288, 7.78. For the held-out pair 0→1, which the training data showed, the surprise is a modest 0.288; for 3→0, also seen, it is 0.0025, essentially none. But 2→3 costs 8.07 and 1→4 costs 7.78: the pairs the training data never showed. The model believes (probability about 0.0003 to 0.0004, bottom matrix) that token 2 is never followed by 3 and token 1 is never followed by 4, because in its training data they never were. The two unseen pairs account for 15.85 of the 16.14 total surprise, 98% of the held-out loss. And the longer it trains, the more certain it becomes that those pairs are impossible: the weights for the unseen entries keep drifting down (to about −2.8) and the surprise grows with them, which is why the held-out loss climbs forever.

Weight decay

The cure is to stop the model from becoming so sure. Weight decay adds a penalty to the loss: (λ/2) times the sum of the squares of all the weights. A weight far from 0 now costs something, so the model only keeps a large weight if the data pay for it. The penalty's gradient is λ · w, which is the extra term in gradient, and each step now also pulls every weight a little toward 0 (that is where the name comes from). λ is the strength. The same training run for several strengths, every number from an mgc run (sweep.py):

 strength  train loss  held-out loss  largest |weight|   held-out loss at steps [1, 5, 10, 25, 50, 100, 200, 400]
        0      0.4171          4.036             6.456   [1.482, 1.946, 2.238, 2.666, 3.004, 3.346, 3.691, 4.036]
    0.001      0.4347          2.683             4.185   [1.482, 1.925, 2.18, 2.484, 2.628, 2.678, 2.683, 2.683]
     0.01      0.5338          1.853             2.622   [1.482, 1.774, 1.839, 1.853, 1.853, 1.853, 1.853, 1.853]
      0.1      0.9601          1.475             1.079   [1.482, 1.438, 1.493, 1.474, 1.475, 1.475, 1.475, 1.475]
      0.3    7.26e+58      3.671e+58         9.408e+58   [1.482, 5.191, 36.97, 5515.0, 26300000.0, 532500000000000.0, 2.183e+29, 3.671e+58]   <- unstable: learning rate x strength is above 2

The table is the whole story:

• No penalty (0): training loss 0.417, held-out loss 4.036, weights up to 6.5.
• Increasing the strength lowers the held-out loss (4.04, 2.68, 1.85, 1.475) and raises the training loss (0.417, 0.435, 0.534, 0.960), while the largest weight shrinks (6.5, 4.2, 2.6, 1.1). That trade is the price of generalizing: the model fits the training pairs less tightly and the unseen pairs less badly.
• At 0.1 the held-out loss is 1.475, below the 1.609 of a model that knows nothing: this model, unlike the unpenalised one, has learned something that helps on pairs it has not seen. Its probabilities show why (03_weight_decay.mg, final matrices):

Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [5, 5] strides = [5, 1] data = 
[[-0.426541,   1.07857,   0.201049,   -0.426541,   -0.426541], 
 [-0.366977,   -0.366977,   0.809691,   0.29124,   -0.366977], 
 [1.04192,   -0.26048,   -0.26048,   -0.26048,   -0.26048], 
 [0.641953,   -0.160488,   -0.160488,   -0.160488,   -0.160488], 
 [0,   0,   0,   0,   0]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 1] strides = [1, 1] data = 
[[2.03837], 
 [1.02709], 
 [0.733227], 
 [2.10107]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [5, 5] strides = [5, 1] data = 
[[0.106635,   0.480357,   0.199738,   0.106635,   0.106635], 
 [0.122326,   0.122326,   0.39677,   0.236253,   0.122326], 
 [0.479039,   0.13024,   0.13024,   0.13024,   0.13024], 
 [0.358047,   0.160488,   0.160488,   0.160488,   0.160488], 
 [0.2,   0.2,   0.2,   0.2,   0.2]]
  ----- exit status 0

The weights stay between about −0.43 and 1.08. After training, token 0 is followed by token 1 with probability about 0.48 (the training data said 3 out of 4, 0.75) and unseen successors keep real probability (about 0.11 each, rather than 0.0003). The held-out surprises are now 2.04, 1.03, 0.73, 2.10: nothing is certain, and nothing costs 8. - The held-out loss with a penalty stops changing after enough steps (for 0.1, from step 50): the penalty and the data pull against each other and settle, unlike the unpenalised run which never settles.

A penalty too strong for the learning rate

The last row of the table is 0.3, and it is not a trade-off, it is a failure. Look at what the penalty does on its own: a step of weight decay turns w into w − 8 · 0.3 · w = −1.4 · w. The weights are multiplied by −1.4, which is larger than 1 in size: every step overshoots zero and lands farther away than it started, on the other side. The condition for a decay step to shrink the weights is |1 − learning rate × strength| < 1, that is, learning rate × strength < 2; here 8 × 0.3 = 2.4.

  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.60944]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.60944]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.832556]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.48224]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[2.17033]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[5.19087]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[71.781]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[36.9733]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[2403.25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[5515.15]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[5.20111e+07]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[2.62988e+07]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.05317e+15]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[5.32522e+14]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[4.31818e+29]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[2.18344e+29]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[7.25952e+58]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[3.6707e+58]]

The held-out loss goes 1.48, 5.19, 36.97, 5515, 2.6·10⁷, 5·10¹⁴, 2·10²⁹, 4·10⁵⁸: it is multiplied by a huge factor every few steps. This is not a compiler bug: the plain-Python reference run produces the same numbers (the checker requires it). It is what this algorithm does with these settings. The fix is a smaller strength or a smaller learning rate: with the same strength 0.3 and a learning rate of 4 the plain-Python reference trains stably (training loss 1.27 and held-out loss 1.51 after 400 steps; this one run was done in Python, not through mgc). The condition learning rate × strength < 2 is necessary, not sufficient: strength 0.2 with learning rate 8 satisfies it (the decay factor is −0.6) but the reference still ends with a training loss of 3.55, because the data term adds curvature of its own, which is why it was left out of the stable list.

Greedy generation

The trained model can now produce text. Start from a token; look up its row of probabilities; choose the likeliest next token; repeat. With ge, "choose the likeliest" is next_token(x, p) = ge(x @ p, row_max(x @ p)): x @ p picks the current token's row, and comparing it with its own maximum gives the one-hot row of the likeliest token. To read the token as a number, multiply by the column [[0], [1], [2], [3], [4]]: a one-hot row times it is the token's number. examples/05_generate.mg trains with strength 0.1 and generates 8 tokens from token 0:

################ mgc run 05_generate.mg
  | (457 lines: the definitions are generalize.mg.defs.in, followed by the data, the unrolled training steps and the prints, written by make_examples.py; see the file)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [5, 5] strides = [5, 1] data = 
[[0.106635,   0.480357,   0.199738,   0.106635,   0.106635], 
 [0.122326,   0.122326,   0.39677,   0.236253,   0.122326], 
 [0.479039,   0.13024,   0.13024,   0.13024,   0.13024], 
 [0.358047,   0.160488,   0.160488,   0.160488,   0.160488], 
 [0.2,   0.2,   0.2,   0.2,   0.2]]
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 = 
[[2]]
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 = 
[[2]]
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 = 
[[2]]
  ----- exit status 0

The matrix is the model's probabilities. The nine numbers after it are the generated text: 0 1 2 0 1 2 0 1 2. Greedy generation follows the most likely arrow each time: 0 is most likely followed by 1 (0.48), 1 by 2 (0.40), 2 by 0 (0.48), and so it goes around the cycle 0 → 1 → 2 → 0 forever. That is typical of greedy decoding from a bigram: a model that only ever picks the most likely next token falls into a loop (a real language model avoids it by sampling or by looking at more context). The program gives the same chain as an independent Python decoder for every starting token 0 to 3.

A tie

Starting from token 4 shows the limit of using a comparison as an argmax. Token 4 is never a current token in the training pairs, so its row is exactly uniform: every token is equally likely, and every entry ties with the maximum.

################ mgc run 06_a_tie.mg
  | (445 lines: the definitions are generalize.mg.defs.in, followed by the data, the unrolled training steps and the prints, written by make_examples.py; see the file)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [5, 5] strides = [5, 1] data = 
[[0.106635,   0.480357,   0.199738,   0.106635,   0.106635], 
 [0.122326,   0.122326,   0.39677,   0.236253,   0.122326], 
 [0.479039,   0.13024,   0.13024,   0.13024,   0.13024], 
 [0.358047,   0.160488,   0.160488,   0.160488,   0.160488], 
 [0.2,   0.2,   0.2,   0.2,   0.2]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[4]]
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 = 
[[0]]
  ----- exit status 0

The first step from token 4 is the number 10: the mask is [1, 1, 1, 1, 1] (all five tokens tie) and 0+1+2+3+4 = 10. It is not a token at all, it is the sum of five. The next step then multiplies that all-ones row by the probabilities (summing their rows), and its largest column is token 0. So from a tie the program does not fail, it produces a mask with several 1s and carries on with a meaningless state. A tie can only be broken by a rule the language does not have (take the first, or sample).

Mistakes the compiler reports

################ mgc run errors/e1_ge_of_a_scalar.mg
  | # ge compares two matrices. A bare number in the source is a compile-time scalar; write a threshold as a 1x1 matrix.
  | print ge([[1, 2]], 1)
  ----- output
errors/e1_ge_of_a_scalar.mg:2: error: ge compares two matrices, not a scalar; write a scalar threshold as a 1x1 matrix such as [[0.5]]
  ----- exit status 1

################ mgc run errors/e2_ge_shapes_differ.mg
  | # The shapes must agree or broadcast: 2x3 and 3x2 do not.
  | print ge([[1, 2, 3], [4, 5, 6]], [[1, 2], [3, 4], [5, 6]])
  ----- output
errors/e2_ge_shapes_differ.mg:2: error: cannot compare shapes 2x3 and 3x2: dimensions 2 and 3 differ and neither is 1
  ----- exit status 1

################ mgc run errors/e3_held_out_pairs_to_the_training_loss.mg
  | # The held-out data have 4 pairs and the training loss was written for 10. Passing one where the other belongs is a shape error, caught when compiling.
  | def log_softmax10(s: tensor[10x5]) = s - row_max(s) - log(row_sum(exp(s - row_max(s))))
  | def train_loss(w: tensor[5x5], x: tensor[10x5], y: tensor[10x5]) = col_sum(row_sum(y * log_softmax10(x @ w))) * -0.1
  | let w = [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]
  | let xv = [[0, 0, 1, 0, 0], [0, 0, 0, 1, 0], [1, 0, 0, 0, 0], [0, 1, 0, 0, 0]]
  | let yv = [[0, 0, 0, 1, 0], [1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 0, 0, 1]]
  | print train_loss(w, xv, yv)
  ----- output
errors/e3_held_out_pairs_to_the_training_loss.mg:7: error: argument shape 4x5 does not fit parameter 10x5
  ----- exit status 1

ge of a bare number; ge of shapes that do not broadcast (the message says compare); and the held-out pairs (4 rows) passed to the training loss written for 10. The last is exactly the mistake the split invites, and the compiler catches it.

Tests

Five new lit files in test/generalize32/ (the suite is now 132 tests):

Test	What it checks
examples	ge (worked out by hand) and the numbers this page quotes from the training runs, including the held-out pair surprises, the final weights at strength 0.1, the divergence, the generated chain and the tie
errors	the three mistakes give the exact messages
verifiers	mg.ge given operands of different shapes is rejected by the dialect
ir	ge is one mg.ge; a 2x1 operand is first stretched by an explicit mg.broadcast
against-reference	check_generalize.py, below

check_generalize.py compares everything with an independent Python trainer (plain lists, math.exp, math.log, the gradient written per training pair) and checks facts that must be true whatever the code looks like:

1. Training. For each strength, the training and held-out loss at all nine checkpoints, and the final weights, probabilities and per-pair surprises, equal the reference's.
2. The unstable strength. The program and the reference blow up identically (same losses at every checkpoint), and the loss grows without bound.
3. Overfitting. Without weight decay the training loss falls at every checkpoint, the held-out loss rises at every checkpoint from step 5 on, and ends well above where it started.
4. Why. The two held-out pairs that never occur in the training pairs are found by the script (not assumed), have surprises above 4 while the others have surprises below 1.5, and carry more than 90% of the held-out loss.
5. Weight decay. Every positive stable strength gives a lower held-out loss than none; a stronger strength gives a higher training loss and smaller weights; the strength with the lowest held-out loss is the same for the program and the reference.
6. The gradient with weight decay equals a finite-difference estimate of the penalised objective (step 0.01) for two strengths and two random weight matrices each, to within the printing precision.
7. Generation. From every starting token, the 8 generated tokens equal the reference's most-likely chain, and the tie from token 4 gives the documented 10, then the largest column sum.

-- training and held-out loss against the independent Python trainer, for each weight-decay strength
  ok   strength 0: training and held-out loss at all 9 checkpoints equal the reference
  ok   strength 0: final weights and probabilities equal the reference
  ok   strength 0: held-out surprise of each pair equals the reference
  ok   strength 0.001: training and held-out loss at all 9 checkpoints equal the reference
  ok   strength 0.001: final weights and probabilities equal the reference
  ok   strength 0.001: held-out surprise of each pair equals the reference
  ok   strength 0.01: training and held-out loss at all 9 checkpoints equal the reference
  ok   strength 0.01: final weights and probabilities equal the reference
  ok   strength 0.01: held-out surprise of each pair equals the reference
  ok   strength 0.1: training and held-out loss at all 9 checkpoints equal the reference
  ok   strength 0.1: final weights and probabilities equal the reference
  ok   strength 0.1: held-out surprise of each pair equals the reference
-- a strength that is too large for the learning rate
  ok   strength 0.3: the program and the reference blow up in exactly the same way (losses equal at every checkpoint)
  ok   the loss grows without bound (7.26e+58 after 400 steps): learning rate x strength = 2.4 is above 2, so each step overshoots by a factor of 1.4
-- overfitting without weight decay
  ok   training loss falls at every checkpoint
  ok   held-out loss RISES at every checkpoint from step 5 on
  ok   and ends well above its starting value (1.6094 to 4.0363)
-- why: the held-out pairs the training pairs never showed
  ok   held-out pairs [(2, 3), (1, 4)] never occur in the training pairs
  ok   their surprises [8.07, 7.78] are above 4; the others [0.0, 0.29] are below 1.5
  ok   the unseen pairs carry more than 90% of the held-out loss
-- weight decay
  ok   strength 0.001: held-out loss 2.6835 is below the 4.0363 of no weight decay
  ok   strength 0.01: held-out loss 1.8533 is below the 4.0363 of no weight decay
  ok   strength 0.1: held-out loss 1.4749 is below the 4.0363 of no weight decay
  ok   a stronger strength gives a higher final training loss (the price of the penalty)
  ok   a stronger strength gives smaller weights
  ok   the strength with the lowest held-out loss is 0.1 for both
-- the gradient with weight decay against finite differences of the penalised objective (step 0.01)
  ok   strength 0.01, random weights #1: gradient equals the reference and the finite-difference estimate (largest difference 3.8e-04)
  ok   strength 0.01, random weights #2: gradient equals the reference and the finite-difference estimate (largest difference 4.2e-04)
  ok   strength 0.1, random weights #1: gradient equals the reference and the finite-difference estimate (largest difference 4.0e-04)
  ok   strength 0.1, random weights #2: gradient equals the reference and the finite-difference estimate (largest difference 3.9e-04)
-- greedy generation
  ok   starting from token 0: [0, 1, 2, 0, 1, 2, 0, 1, 2] equals the reference chain
  ok   starting from token 1: [1, 2, 0, 1, 2, 0, 1, 2, 0] equals the reference chain
  ok   starting from token 2: [2, 0, 1, 2, 0, 1, 2, 0, 1] equals the reference chain
  ok   starting from token 3: [3, 0, 1, 2, 0, 1, 2, 0, 1] equals the reference chain
  ok   starting from token 4 (its row is uniform): every token ties, the mask is all ones (x @ ids = 0+1+2+3+4 = 10), and the next step takes the largest column sum of the probabilities (token 0)
all checks pass

Are the checks good enough? Run eleven wrong trainers.

NOTE: this script deliberately runs wrong trainers. 'caught' lines are EXPECTED: they show the checker can detect the mistake.
baseline (nothing broken): checker passes (expected)
weight decay pushes the weights AWAY from 0 (wrong sign): caught by 20 failing checks (the reference, overfitting/weight decay)
weight decay is missing from the gradient: caught by 21 failing checks (the reference, overfitting/weight decay)
weight decay is applied twice as strongly: caught by 21 failing checks (the reference, overfitting/weight decay, generation)
the training gradient is not averaged over the 10 pairs: caught by 29 failing checks (the reference, other, overfitting/weight decay, generation)
the held-out loss is averaged over 10, not 4: caught by 6 failing checks (the reference, other)
the held-out surprise has the wrong sign: caught by 5 failing checks (the reference, other)
the penalty in the objective is lambda, not lambda/2 (the gradient no longer matches it): caught by 4 failing checks (the reference)
the learning rate is 4, not 8: caught by 10 failing checks (the reference, overfitting/weight decay)
greedy generation compares the maximum with the row, the wrong way round: caught by 5 failing checks (the reference, generation)
greedy generation picks the LEAST likely token: caught by 5 failing checks (the reference, generation)
the training log-softmax is the naive one (no maximum subtracted): caught by 2 failing checks (the reference, overfitting/weight decay)
(The last one is caught only because of the deliberately diverging run (strength 0.3): its weights reach 1e58, where a naive log-sum-exp overflows. At the
 weights of the other runs, which stay below 7, the naive formula gives the same numbers, as Chapter 31 showed; I expected it to go uncaught here and was wrong.)

All eleven are caught. Several are caught by different kinds of check, which is the point of having both a reference and properties: a wrong weight-decay sign or a missing penalty is caught by the reference and by "weight decay lowers the held-out loss"; the penalty written as λ instead of λ/2 is invisible to every training comparison (the gradient is still right) and is caught only by the finite-difference check against the objective; the two wrong next_tokens are caught by the generation check. The last mutation (a naive log-softmax) is the surprise: I expected it to go uncaught, because every normal run keeps its weights below 7 where exp cannot overflow. It is caught only by the deliberately diverging run, whose weights reach 10⁵⁸. That is an accident of another check, not evidence that this chapter tests numerical stability (Chapter 31's stress case does).

Compiler mutations

ge is a new compiler operation, so the compiler is also broken on purpose, six ways:

#!/bin/sh
# READ THIS FIRST: this script breaks the Chapter 32 code on purpose, one way at a time, and re-runs the Chapter 32 tests (test/generalize32).
#   "caught by ..." (a test failed) is the EXPECTED, wanted result: it shows the tests can notice that mistake.
#   "NOT CAUGHT" would be a test gap. Part 1 edits the front end (mgfront.py) and the driver (mgc) in place; part 2 rebuilds mg-opt with one bug
#   injected into the dialect or the lowering (a full rebuild each, in work/mut/N). Every edited file is restored afterwards. Output: mutation_out.txt
HERE=$(cd "$(dirname "$0")" && pwd); LIT=$HERE/../../part15/code/run_lit.sh; W=$HERE/work/mut; rm -rf $W; mkdir -p $W
FILTER=':: generalize32/'
echo "NOTE: this script deliberately breaks the compiler code. 'caught' lines are EXPECTED: they show the tests can detect the mistake."
F=$HERE/mgfront.py; G=$HERE/mgc; cp $F $W/mgfront.orig; cp $G $W/mgc.orig
restore() { cp $W/mgfront.orig $F; cp $W/mgc.orig $G; }
trap restore EXIT
fails_of() { grep '^FAIL' | sed 's/.*:: //; s/ (.*//' | tr '\n' ' '; }
out=$(MG_TEST_TMP=$W/out sh $LIT --filter "$FILTER" 2>&1); echo "baseline (nothing broken): $(echo "$out" | grep -q '^FAIL' && echo 'UNEXPECTED FAILURE' || echo 'tests pass (expected)')"
echo; echo "== Part 1: the front end and the driver"
fe() {  # label file sed-expr
  restore; sed -i "$3" $2
  if cmp -s $F $W/mgfront.orig && cmp -s $G $W/mgc.orig; then echo "$1: MUTATION DID NOT APPLY"; return; fi
  f=$(MG_TEST_TMP=$W/out sh $LIT --filter "$FILTER" 2>&1 | fails_of)
  if [ -n "$f" ]; then echo "$1: caught by $f"; else echo "$1: NOT CAUGHT (a test gap)"; fi
  restore
}
fe "ge is compiled as a subtraction by the front end"                   $F 's/{r} = mg.ge {a}, {b}/{r} = mg.sub {a}, {b}/'
restore

echo; echo "== Part 2: the dialect verifier and the lowering (a full rebuild per mutation)"
n=0
try() {  # label file sed-expr
  n=$((n+1)); d=$W/$n; mkdir -p $d/build; cp -r $HERE/tree $d/tree; rm -rf $d/tree/build
  tgt=$d/tree/lib/$2; cp $tgt $d/before; sed -i "$3" $tgt
  if cmp -s $d/before $tgt; then echo "$1: MUTATION DID NOT APPLY"; return; fi
  ( cd $d/build && cmake $d/tree -DMLIR_DIR=/usr/lib/llvm-18/lib/cmake/mlir -DLLVM_DIR=/usr/lib/llvm-18/lib/cmake/llvm -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=OFF >/dev/null 2>&1 && mkdir -p include/mg && make -j8 >/dev/null 2>&1 )
  [ -x $d/build/mg-opt ] || { echo "$1: BUILD FAILED"; return; }
  f=$(MG_OPT=$d/build/mg-opt MG_TEST_TMP=$d/out sh $LIT --filter "$FILTER" 2>&1 | fails_of)
  if [ -n "$f" ]; then echo "$1: caught by $f"; else echo "$1: NOT CAUGHT (a test gap)"; fi
}
try "ge compares with > instead of >= (a maximum no longer ties with itself)" LowerToAffine.cpp 's/arith::CmpFPredicate::OGE/arith::CmpFPredicate::OGT/'
try "ge compares with <= instead of >="                                  LowerToAffine.cpp 's/arith::CmpFPredicate::OGE/arith::CmpFPredicate::OLE/'
try "ge returns 0 where true and 1 where false (the select is swapped)"  LowerToAffine.cpp 's/b.create<arith::SelectOp>(l, cond, one, zero)/b.create<arith::SelectOp>(l, cond, zero, one)/'
try "ge returns 2 where true"                                            LowerToAffine.cpp 's/Value one = rewriter.create<arith::ConstantOp>(loc, rewriter.getF64FloatAttr(1.0));/Value one = rewriter.create<arith::ConstantOp>(loc, rewriter.getF64FloatAttr(2.0));/'
try "mg.ge's verifier accepts anything"                                  MgDialect.cpp 's/return verifyElementwise(\*this, {getLhs().getType(), getRhs().getType(), getResult().getType()}, "mg.ge"); }/return mlir::success(); }/'
try "a sum starts from 1 instead of 0 (every row_sum and col_sum is off by one)" LowerToAffine.cpp 's/double initVal = isMax ? -std::numeric_limits<double>::infinity() : 0.0;/double initVal = isMax ? -std::numeric_limits<double>::infinity() : 1.0;/'
restore
echo; echo "restored: $(cmp -s $F $W/mgfront.orig && cmp -s $G $W/mgc.orig && echo 'mgfront.py and mgc are back to their originals' || echo 'RESTORE FAILED')"

NOTE: this script deliberately breaks the compiler code. 'caught' lines are EXPECTED: they show the tests can detect the mistake.
baseline (nothing broken): tests pass (expected)

== Part 1: the front end and the driver
ge is compiled as a subtraction by the front end: caught by generalize32/ir.mlir generalize32/examples.mlir generalize32/against-reference.mlir 

== Part 2: the dialect verifier and the lowering (a full rebuild per mutation)
ge compares with > instead of >= (a maximum no longer ties with itself): caught by generalize32/examples.mlir generalize32/against-reference.mlir 
ge compares with <= instead of >=: caught by generalize32/examples.mlir generalize32/against-reference.mlir 
ge returns 0 where true and 1 where false (the select is swapped): caught by generalize32/examples.mlir generalize32/against-reference.mlir 
ge returns 2 where true: caught by generalize32/examples.mlir generalize32/against-reference.mlir 
mg.ge's verifier accepts anything: caught by generalize32/verifiers.mlir 
a sum starts from 1 instead of 0 (every row_sum and col_sum is off by one): caught by generalize32/examples.mlir generalize32/against-reference.mlir 

restored: mgfront.py and mgc are back to their originals

All six are caught. As in Chapters 30 and 31 the verifier mutation is caught only by the dedicated verifiers test, and the last mutation (a sum that starts from 1) is not about ge at all: it breaks row_sum, which the loss and the softmax lean on. Two details specific to ge: the > mutation is the interesting one for generation, because a maximum is no longer greater than itself, so the mask is all zeros and the generated tokens collapse to 0; and "returns 2 where true" would pass a test that only looked at which entries are non-zero, which is why the examples check the values.

The full suite:

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

Testing Time: 15.21s

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

Limits and what is not established

• Tiny data, so the numbers illustrate a mechanism and nothing more. Ten training pairs and four held-out pairs. "The held-out loss is 1.475" is a measurement on four pairs, two of which the model had no way to predict. Nothing here says how any of this behaves on real text.
• The strength was chosen by looking at the held-out loss. The best strength (0.1, from a list of four) is the one with the lowest held-out loss, so that held-out loss is a selected number, optimistic by construction. An honest estimate needs a third set of data used only once, after the choice. There is none here.
• A bigram still sees one token. Weight decay does not give the model any context; it only stops the model being over-confident about what it has seen. The best it can do on this held-out set is still limited by the unseen pairs.
• Unseen pairs are handled by shrinking, not by knowing. At 0.1 the unseen successors get probabilities near 0.1 because the penalty stops the weights growing; that is a side effect of regularization, not a model of what is likely. Other ways of handling unseen events (smoothing counts, using more data) are not tried.
• Weight decay was one remedy of several. Early stopping (the held-out loss was lowest after one step, 1.482) was visible in the numbers but not implemented as a rule; more data, smaller models and other penalties are not compared.
• The stability condition learning rate × strength < 2 is derived for the penalty on its own. The full condition also involves the data term's curvature (strength 0.2 with learning rate 8 satisfies the penalty-only condition and still trains badly: training loss 3.55 in the Python reference); it was checked here only by running a handful of settings, not proved.
• Greedy generation loops, and ties are not handled. ge marks every maximum, so a tie gives several 1s (shown, not fixed). There is no sampling (no random numbers in the language) and no stopping token.
• The comparison ge is a function, not an operator, has no less-than, equal or not-equal siblings, and takes only matrices (a threshold is a 1x1 matrix). NaN compares false.
• Six digits. mgc prints six significant digits; comparisons with Python are at a relative 10⁻⁵, and the finite-difference check at 2·10⁻³.
• The transformer of Chapter 30 is still untrained, and none of this chapter's training is applied to it.
• Platform: x86-64 Linux, LLVM 18.1.3.

Reproducing

cd docs/part32/code && ./build.sh
python3 make_examples.py && ./run_examples.sh > examples_out.txt
./sweep.py > sweep_out.txt
./check_generalize.py > check_generalize_out.txt
./model_mutation.py > model_mutation_out.txt
./mutation.sh > mutation_out.txt
cd ../../part15/code && ./run_lit.sh          # 132 tests

Chapter summary

• A model can fit its training data and still fail on data it has not seen. With the bigram model, the training loss fell to 0.417 while the held-out loss rose to 4.036, worse than the 1.609 of a model that knows nothing, because two of the four held-out pairs never occurred in training and the model learned to treat them as impossible.
• Weight decay (a penalty on large weights, whose gradient is λw) trades training loss for held-out loss: at strength 0.1 the training loss is 0.960 and the held-out loss 1.475, now better than knowing nothing.
• A penalty that is too strong for the learning rate does not just underperform: when learning rate × strength exceeds 2, every step overshoots and the weights grow without bound, in exactly the same way in an independent Python implementation.
• One new operation, ge, a comparison, gives the one-hot row of the likeliest token, and with it a trained model generates text. Greedy generation from a bigram loops (here 0 1 2 0 1 2 …), and a tie breaks the one-hot mask.
• Eleven wrong trainers and six broken compilers are all caught; the one wrong trainer I expected to go uncaught was caught, by the diverging run, which is reported as such.

Self-check questions

Each answer is collapsed; try the question first.

1. Why did the held-out loss in Experiment 1 first fall (1.609 to 1.482) and then rise?

   Answer

   In the first steps the model learns the broad facts that help on any pair (token 0 is usually followed by 1, token 3 by 0), which lowers the loss on seen pairs in the held-out set. After that, nothing helps the held-out set further, but training keeps pushing the weights for unseen successors lower to make the training pairs fit better. That makes the model more and more certain that the unseen held-out pairs are impossible, and the surprise at them grows without limit.

2. Why is a held-out loss of 4.036 worse than knowing nothing, and what is "knowing nothing" here?

   Answer

   A model that assigns every token the probability 1/5 has loss log 5 = 1.609 on any data. The trained model has loss 4.036 on the held-out pairs, which means that on average it gives the true next token a probability of only e^−4.036 = 0.018, far below 1/5. It is confidently wrong about the unseen pairs, which costs more than being unsure about everything.

3. Why does weight decay lower the held-out loss but raise the training loss?

   Answer

   The penalty stops weights from growing large, so the model cannot become as certain as the training data alone would make it. That keeps probabilities for unseen pairs away from 0 (good for held-out data) and also stops the model from fitting the training pairs as tightly as it could (worse training loss). Training loss is exactly the quantity that rewards memorizing, so any force against memorizing raises it.

4. Why does a strength of 0.3 diverge with learning rate 8, and what would make it stable?

   Answer

   A weight-decay step multiplies each weight by 1 − learning rate × strength. With 8 and 0.3 that is 1 − 2.4 = −1.4: the weights change sign and grow by 40% each step. Stability needs |1 − learning rate × strength| < 1, i.e. learning rate × strength < 2 (for the penalty alone). The penalty alone shrinks the weights when the factor is smaller than 1 in size, e.g. a learning rate of 4 with strength 0.3 (factor −0.2; the Python reference trains stably there). But the condition is only necessary: strength 0.2 with learning rate 8 has factor −0.6 and still ends with a training loss of 3.55 in the reference, because the data term adds curvature.

5. How does ge(row, row_max(row)) find the largest entry, and what goes wrong with ties?

   Answer

   row_max(row) is the largest number in the row, and ge marks every entry that is greater than or equal to it, which is exactly the entries equal to the maximum. With one maximum the result is one-hot. With ties, every tied entry is marked, so the result has several 1s and is no longer a single token; in example 6 the all-equal row marks all five tokens.

6. Why must the checker discover which held-out pairs are unseen instead of assuming 2→3 and 1→4?

   Answer

   If the script hard-coded those pairs, a change to the data split or the text would silently make the explanation wrong while the check still passed. Computing "which held-out pairs occur among the training pairs" from the data ties the claim to the data, so it stays true or fails loudly.

7. The best strength was chosen by held-out loss. Why is that held-out loss then too optimistic?

   Answer

   The held-out set was used to make a choice (which strength to keep), so it is no longer an independent test: among several candidates, the one that looks best on this set partly looks best by luck. An unbiased estimate needs a third set that played no part in any choice, measured once at the end. With four held-out pairs the effect could be large.

8. In the divergence run, why does the checker insist that the reference diverges identically?

   Answer

   A blow-up could be a bug in the new operation or the compiler, or it could be what the algorithm does. If an independent plain-Python implementation of the same algorithm produces the same exploding numbers at every checkpoint, the cause is the algorithm and the settings, not the compiler. The match is the evidence.

9. Why does greedy generation from a bigram loop?

   Answer

   The next token depends only on the current token, and greedy decoding always picks the same most likely successor for a given token. Starting at token 0 the chain 0 → 1 → 2 → 0 closes after three steps, and since the rule is deterministic and sees nothing but the current token, it repeats forever. Breaking the loop needs randomness (sampling) or more context.

10. The wrong trainer with a naive log-softmax was caught by this chapter's checker. Was that for the reason Chapter 31's naive-softmax mutants were caught?

    Answer

    Not by the same check. Chapter 31 added a stress case with weights of ±1000 on purpose. Here the naive version produces the same numbers as the stable one for every run with normal weights (they stay below 7). It was caught only because of the deliberately diverging run, whose weights reach 10⁵⁸ and overflow exp. That was unplanned: I expected it to go unnoticed here, and it is reported as caught by an accident of another check, not as evidence that this chapter's checks test numerical stability.