Skip to content

31. Training a Language Model: Gradient Descent on a Bigram

Mountain goats in space helmets on Venus, under its clouds

What you will understand: how a model learns, shown on the smallest language model there is, and how a Mountain Goat program trains it. The model reads one token and predicts the next. Training is gradient descent on the cross-entropy loss, with a gradient derived by hand, run for 400 steps, and every claim checked: the loss falls and cannot go below a bound that follows from the data, the model ends up predicting exactly the frequencies it saw, a token it never saw as input keeps exactly the weights it started with, and the hand-derived gradient agrees with a numerical estimate. Mountain Goat gets one new operation, log.

What you need to know first: Chapter 26 (gradient descent on a linear model; this chapter does the same with a softmax output), Chapter 29 (softmax) and Chapter 30 (what a transformer predicts). The model here is the simplest relative of Chapter 30's: it looks at one token instead of a whole sequence. Training the transformer itself is not done (see Limits).

Compile and run

cd docs/part31/code && ./build.sh                 # once: the newest compiler (adds mg.log)
python3 make_examples.py                          # (re)writes the examples from bigram.mg.defs
./mgc run examples/04_training.mg                 # train: 400 gradient-descent steps
./run_examples.sh > examples_out.txt              # every example and every error example on this page
./check_bigram.py                                 # the trainer against the reference and against facts about the data
./model_mutation.py > model_mutation_out.txt      # runs 12 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 nan, -inf, FAIL lines and error messages, and why that is good

Example 1 prints -inf and nan on purpose: they are what log of 0 and of a negative number are. The three error examples are programs the compiler must reject. The two mutation outputs are negative controls: wrong trainers and a deliberately broken compiler run against the checks, where "caught" is the wanted result and "NOT CAUGHT" would be a gap.

The task

A language model assigns a probability to every possible next token given what came before. The bigram model looks at only the most recent token. It has a table w with one row per token and one column per possible next token; row i holds scores (logits) for what follows token i, and a softmax turns the row into probabilities.

The data is a 15-token text over five tokens, numbered 0 to 4:

0 1 2 0 1 3 0 1 2 0 2 3 0 1 4

It contains 14 (current, next) pairs. Counting them gives the table the model should end up reproducing (rows are the current token, columns the next):

current next 0 next 1 next 2 next 3 next 4 observed probabilities of the next token
0 0 4 1 0 0 0.8 for token 1, 0.2 for token 2
1 0 0 2 1 1 0.5 for token 2, 0.25 for token 3, 0.25 for token 4
2 2 0 0 1 0 2/3 for token 0, 1/3 for token 3
3 2 0 0 0 0 1.0 for token 0
4 0 0 0 0 0 (never a current token: no data)

Token 4 appears only as the last token, so it never has a successor in the data. That is deliberate and will matter.

The new operation: log

The loss needs a logarithm, which Mountain Goat did not have. log(m) takes the natural logarithm of every element, added the way exp and sqrt were: mg.log in the dialect with a shape verifier, a loop-nest lowering that uses math.log, and a built-in in the front end.

################ mgc run 01_log.mg
  | # log(x) is the natural logarithm: the power e must be raised to in order to get x. It undoes exp.
  | print log([[1, 2.718281828459045, 10], [0.5, 100, 0.001]])      # log(1) = 0, log(e) = 1, log(10) = 2.30259; numbers below 1 give negative logs
  | print exp(log([[1, 2, 10], [0.5, 100, 0.001]]))                  # exp undoes log (to the printed digits)
  | print log([[0]])                                                 # log of 0 is minus infinity: no power of e is 0
  | print log([[-1]])                                                # log of a negative number is not a number (nan)
  | # log turns products into sums: log(a*b) = log(a) + log(b). Both lines print the same two numbers.
  | print log([[6, 0.5]] * [[7, 8]])
  | print log([[6, 0.5]]) + log([[7, 8]])
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[0,   1,   2.30259], 
 [-0.693147,   4.60517,   -6.90776]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[1,   2,   10], 
 [0.5,   100,   0.001]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[-inf]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[-nan]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 2] strides = [2, 1] data = 
[[3.73767,   1.38629]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 2] strides = [2, 1] data = 
[[3.73767,   1.38629]]
  ----- exit status 0

log(1) = 0, log(e) = 1, log(10) = 2.30259, and numbers below 1 have negative logs. exp undoes log. log(0) is minus infinity (no power of e is 0) and the log of a negative number is not a number. The last two lines show the property that makes logs useful for probabilities: log(a·b) = log(a) + log(b); the product of probabilities becomes a sum, and the sum does not underflow.

The loss: cross-entropy

How good is a model's prediction? For each pair, look at the probability the model gave to the token that really came next, call it p, and take −log p: the model's surprise. A probability of 1 costs 0; a probability near 0 costs a lot. The loss is the average surprise over the data.

################ mgc run 02_cross_entropy.mg
  | # Cross-entropy: how surprised a model is by what really happened. For each case, y is a one-hot row marking the token that really came next and p is
  | # the model's probability row; y * log(p) keeps only the log-probability the model gave the right token, and minus its average is the loss.
  | # Four cases (rows) over three tokens: a model with no idea, a confident correct model, a confident WRONG model, and a coin-flip between two.
  | let p = [[0.3333333333333333, 0.3333333333333333, 0.3333333333333333], [0.98, 0.01, 0.01], [0.01, 0.98, 0.01], [0.5, 0.5, 0.0000001]]
  | let y = [[1, 0, 0], [1, 0, 0], [1, 0, 0], [1, 0, 0]]
  | print row_sum(y * log(p)) * -1                      # the surprise of each case: log(3) = 1.0986, 0.0202, 4.6052 (very wrong), 0.6931 (one bit)
  | print col_sum(row_sum(y * log(p))) * -0.25          # the loss: the average surprise, (1.0986 + 0.0202 + 4.6052 + 0.6931) / 4 = 1.6043
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 1] strides = [1, 1] data = 
[[1.09861], 
 [0.0202027], 
 [4.60517], 
 [0.693147]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.60428]]
  ----- exit status 0

Four cases over three tokens (the true token is always token 0): a model with no idea (uniform, 1/3 each) costs log 3 = 1.0986; a confident correct model (0.98) costs 0.0202; a confident wrong model (0.01 on the right token) costs 4.605, which is the steepest penalty; a coin flip between two costs 0.6931. The loss is their average, 1.60428. In the program, y is a one-hot row marking the true token, so y * log(p) keeps one number per case and row_sum picks it.

The model and its gradient

The data go in as two 14x5 matrices of one-hot rows: x holds the current tokens and y the next tokens. Then x @ w is a 14x5 matrix of scores: since each row of x has a single 1, multiplying by it selects the matching row of w. The probabilities are the softmax of each row, and the loss is the average of −log of the right entry.

To learn, we need the gradient: how the loss changes when each of the 25 weights changes, since gradient descent moves every weight a little downhill. For one pair with scores z, probabilities p = softmax(z) and one-hot target y, the loss is −log p_k where k is the true token, and a short calculation gives

∂(−log p_k)/∂z_j = p_j − y_j

(the probability of each token, minus 1 for the true one). The scores come from w by z = x_t · w, so the pair's gradient on row a of w (where a is the current token) is that vector p − y. Summing the pairs and dividing by the 14 that were averaged gives, for the whole data set,

gradient = xᵀ · (softmax(x · w) − y) / 14

a 5x5 matrix. Nothing about that was found by a program: it is derived by hand, and only trusted because it is checked (below, against a numerical estimate). The step is w ← w − (learning rate) · gradient.

One step, by hand

Before the full run, a case small enough to do on paper. Two tokens, four pairs (0→1, 0→1, 0→0, 1→0), weights zero (so every probability is 1/2), learning rate 1:

################ mgc run 03_one_gradient_step.mg
  | # One step of gradient descent by hand. Two tokens (0 and 1) and four (current, next) pairs: 0 -> 1, 0 -> 1, 0 -> 0, 1 -> 0.
  | # w starts at zero, so every probability is 1/2. The gradient of the loss is  transpose(x) @ (softmax(x @ w) - y) / 4.
  | # Row 0 of w sees three pairs: two push toward token 1, one toward token 0, so the gradient is [0.5, -0.5] / 4 = [0.125, -0.125]. Row 1 sees one pair (1 -> 0): [-0.125, 0.125].
  | def log_softmax(s: tensor[4x2]) = s - row_max(s) - log(row_sum(exp(s - row_max(s))))
  | def loss(w: tensor[2x2], x: tensor[4x2], y: tensor[4x2]) = col_sum(row_sum(y * log_softmax(x @ w))) * -0.25
  | def gradient(w: tensor[2x2], x: tensor[4x2], y: tensor[4x2]) = transpose(x) @ (exp(log_softmax(x @ w)) - y) * 0.25
  | let x = [[1, 0], [1, 0], [1, 0], [0, 1]]
  | let y = [[0, 1], [0, 1], [1, 0], [1, 0]]
  | let w = [[0, 0], [0, 0]]
  | print loss(w, x, y)                      # log(2) = 0.693147: with no information every token is a coin flip
  | print gradient(w, x, y)                  # [[0.125, -0.125], [-0.125, 0.125]]
  | let w1 = w - gradient(w, x, y)           # one step with learning rate 1
  | print w1                                 # [[-0.125, 0.125], [0.125, -0.125]]: each row moves toward the token that followed it more often
  | print loss(w1, x, y)                     # smaller than before: the step helped
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.693147]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data = 
[[0.125,   -0.125], 
 [-0.125,   0.125]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data = 
[[-0.125,   0.125], 
 [0.125,   -0.125]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.638439]]
  ----- exit status 0

Row 0 of w sees three pairs: two push toward token 1 and one toward token 0, so its gradient is [0.5, −0.5]/4 = [0.125, −0.125]; row 1 sees one pair (1→0) and gets [−0.125, 0.125]. The step subtracts the gradient: each row moves toward the token that followed it more often, and the loss falls from 0.693147 (= log 2, a coin flip) to 0.638439. You can reproduce that last number by hand: the new probabilities are 0.5622 and 0.4378, and −(2·ln 0.5622 + ln 0.4378 + ln 0.5622)/4 = 0.6384.

The definitions

The whole trainer, as functions:

# A bigram language model, trained by gradient descent. The model predicts the NEXT token from the CURRENT token alone: a 5x5 matrix w whose row i
# holds the scores (logits) for what follows token i. The data are 14 (current, next) pairs, written as one-hot rows: x is 14x5 and y is 14x5.

# log of the softmax of each row, computed as  s - max - log(sum(exp(s - max)))  so that no probability is ever formed and log never sees a 0
def log_softmax(s: tensor[14x5]) = s - row_max(s) - log(row_sum(exp(s - row_max(s))))

# the cross-entropy loss: minus the average, over the 14 pairs, of the log-probability the model gave to the token that really came next
# (y is one-hot, so y * log_softmax keeps one number per pair; row_sum picks it; col_sum adds the 14 of them; 0.0714... is 1/14)
def loss(w: tensor[5x5], x: tensor[14x5], y: tensor[14x5]) = col_sum(row_sum(y * log_softmax(x @ w))) * -0.07142857142857142

# the gradient of that loss with respect to w, derived by hand (see the page):  transpose(x) @ (softmax(x @ w) - y) / 14
def gradient(w: tensor[5x5], x: tensor[14x5], y: tensor[14x5]) = transpose(x) @ (exp(log_softmax(x @ w)) - y) * 0.07142857142857142

# one step of gradient descent with learning rate 8
def step(w: tensor[5x5], x: tensor[14x5], y: tensor[14x5]) = w - gradient(w, x, y) * 8

# the model's probabilities: row i is the softmax of row i of w, 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)))

log_softmax is the log of the softmax, computed as s − max − log(sum(exp(s − max))). It is written this way rather than as log(softmax(s)) for a reason that is easy to miss: a softmax probability can underflow to exactly 0 when a score is far below the maximum, and the log of 0 is minus infinity, which times a 0 in y gives nan. The log-sum-exp form never builds the probability, so it cannot do that. The stress check below pushes it to scores of ±1000. gradient is the formula above, step subtracts it with learning rate 8, and probabilities is the softmax of the weight rows.

Training

examples/04_training.mg is the definitions followed by the data and 400 let w1 = step(w0, x, y) lines (a generated file; it has no loop because the language has none, and each step is one call), then it prints the loss at ten checkpoints, the final weights, the final probabilities and the final gradient:

################ mgc run 04_training.mg
  | (436 lines: the definitions are bigram.mg.defs, followed by the data, the unrolled training steps and the prints; see the file)
  ----- 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 = 
[[0.992623]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.845013]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.718402]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.666934]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.634284]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.623229]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.617695]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.614931]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.613551]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [5, 5] strides = [5, 1] data = 
[[-2.81579,   4.91782,   3.52955,   -2.81579,   -2.81579], 
 [-3.70466,   -3.70466,   2.93245,   2.23843,   2.23843], 
 [4.36153,   -2.67619,   -2.67619,   3.66706,   -2.67619], 
 [6.18336,   -1.54584,   -1.54584,   -1.54584,   -1.54584], 
 [0,   0,   0,   0,   0]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [5, 5] strides = [5, 1] data = 
[[0.000350059,   0.799476,   0.199474,   0.000350059,   0.000350059], 
 [0.000654832,   0.000654832,   0.499564,   0.249563,   0.249563], 
 [0.665791,   0.000584647,   0.000584647,   0.332455,   0.000584647], 
 [0.998244,   0.000439023,   0.000439023,   0.000439023,   0.000439023], 
 [0.2,   0.2,   0.2,   0.2,   0.2]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [5, 5] strides = [5, 1] data = 
[[0.000125021,   -0.000187222,   -0.000187841,   0.000125021,   0.000125021], 
 [0.000187095,   0.000187095,   -0.000124547,   -0.000124821,   -0.000124821], 
 [-0.000187715,   0.000125282,   0.000125282,   -0.00018813,   0.000125282], 
 [-0.00025087,   6.27176e-05,   6.27176e-05,   6.27176e-05,   6.27176e-05], 
 [0,   0,   0,   0,   0]]
  ----- exit status 0

Read it in order:

  • The loss starts at 1.60944, which is exactly log 5: the model begins knowing nothing, so every token is equally likely (1/5), and the surprise is −log(1/5). It falls at every checkpoint, quickly at first (0.99 after one step) and then more slowly, to 0.613551 after 400.
  • It cannot fall below 0.612174. That number is the conditional entropy of the data: the best average surprise any bigram model could possibly have, reached when each row of probabilities equals the observed frequencies in the table above. (A model cannot be less surprised than the data itself is random.) After 400 steps the loss is within 0.0014 of that floor.
  • The weights and probabilities. The probabilities row by row: token 0 is followed by token 1 with probability 0.799476 (observed 0.8) and token 2 with 0.199474 (observed 0.2); token 1 by tokens 2, 3, 4 with 0.4996, 0.2496, 0.2496 (observed 0.5, 0.25, 0.25); token 2 by token 0 with 0.665791 (observed 2/3) and token 3 with 0.332455 (observed 1/3); token 3 by token 0 with 0.998244 (observed 1). The model has learned the table.
  • Zero counts are approached, never reached. The probabilities for pairs never seen (such as token 0 followed by token 3) are about 0.00035, not 0. The weights for those entries are −2.8 and still drifting down: a softmax probability is exactly 0 only when its score is minus infinity. With more steps the numbers shrink without ever vanishing.
  • Row 4 never moved. Token 4 is never a current token, so no pair pushes on its row: its gradient is exactly zero, its weights are exactly 0 after 400 steps, and its probabilities are exactly 0.2, uniform. The model has no information about what follows token 4 and says so.
  • The final gradient is small (about 10−4) everywhere and exactly 0 in row 4: nearly at the bottom, but still pulling very slightly.

Mistakes the compiler reports

################ mgc run errors/e1_wrong_number_of_pairs.mg
  | # The data are 14 pairs. Thirteen rows do not fit the 14x5 parameters, and the compiler says so.
  | def log_softmax(s: tensor[14x5]) = s - row_max(s) - log(row_sum(exp(s - row_max(s))))
  | def loss(w: tensor[5x5], x: tensor[14x5], y: tensor[14x5]) = col_sum(row_sum(y * log_softmax(x @ w))) * -0.07142857142857142
  | 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 x13 = [[1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 1, 0, 0], [1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 0, 1, 0], [1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 1, 0, 0], [1, 0, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 1, 0], [1, 0, 0, 0
  | let y13 = [[0, 1, 0, 0, 0], [0, 0, 1, 0, 0], [1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 0, 1, 0], [1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 1, 0, 0], [1, 0, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 1, 0], [1, 0, 0, 0, 0], [0, 1, 0, 0
  | print loss(w, x13, y13)
  ----- output
errors/e1_wrong_number_of_pairs.mg:7: error: argument shape 13x5 does not fit parameter 14x5
  ----- exit status 1

################ mgc run errors/e2_log_of_a_scalar.mg
  | # log is applied to a matrix. A bare number in the source is a compile-time scalar.
  | print log(2)
  ----- output
errors/e2_log_of_a_scalar.mg:2: error: log needs a matrix, not a scalar
  ----- exit status 1

################ mgc run errors/e3_weights_the_wrong_size.mg
  | # The weights must be 5x5 (one row per token, one column per next token). A 4x5 matrix does not fit.
  | def log_softmax(s: tensor[14x5]) = s - row_max(s) - log(row_sum(exp(s - row_max(s))))
  | def loss(w: tensor[5x5], x: tensor[14x5], y: tensor[14x5]) = col_sum(row_sum(y * log_softmax(x @ w))) * -0.07142857142857142
  | def gradient(w: tensor[5x5], x: tensor[14x5], y: tensor[14x5]) = transpose(x) @ (exp(log_softmax(x @ w)) - y) * 0.07142857142857142
  | def step(w: tensor[5x5], x: tensor[14x5], y: tensor[14x5]) = w - gradient(w, x, y) * 8
  | let w = [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]
  | let x = [[1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 1, 0, 0], [1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 0, 1, 0], [1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 1, 0, 0], [1, 0, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 1, 0], [1, 0, 0, 0, 
  | let y = [[0, 1, 0, 0, 0], [0, 0, 1, 0, 0], [1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 0, 1, 0], [1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 1, 0, 0], [1, 0, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 1, 0], [1, 0, 0, 0, 0], [0, 1, 0, 0, 
  | print step(w, x, y)
  ----- output
errors/e3_weights_the_wrong_size.mg:9: error: argument shape 4x5 does not fit parameter 5x5
  ----- exit status 1

Thirteen rows passed where the data are 14; log of a bare number; and a 4-row weight matrix where 5x5 is required. As in earlier chapters the shapes are types: a trainer cannot be connected to the wrong-sized data or weights.

Tests

Five new lit files in test/bigram31/ (the suite is now 127 tests):

Test What it checks
examples the log, cross-entropy and one-step examples (worked out by hand above) and the training run's outputs
errors the three mistakes give the exact messages, with file and line
verifiers mg.log given a result of a different shape is rejected by the dialect
ir log-softmax is a max-reduce, exp, sum-reduce, mg.log of the row sums, in that order
against-reference check_bigram.py, below

check_bigram.py compares the Mountain Goat trainer with an independent Python trainer (plain lists, math.exp, math.log, with the gradient written from the definition of each pair's contribution rather than as a matrix formula), and checks facts that must be true whatever the code looks like:

  1. Training. The loss at all ten checkpoints and the final weights, probabilities and gradient equal the reference's (to mgc's six printed digits).
  2. The loss. Strictly decreasing at the checkpoints; never below the conditional entropy; within 0.002 of it at the end; exactly log 5 before training.
  3. What it learned. Every trained probability is within 0.01 of the observed frequency, for every token that has a successor in the data.
  4. The unseen row. Token 4's weights and gradient are exactly 0 and its probabilities exactly uniform.
  5. The gradient. The hand-derived gradient agrees with a finite-difference estimate for three random weight matrices: for each of the 25 weights, the loss is computed with that weight raised and lowered by 0.01, and (L₊ − L₋)/0.02 must match the gradient entry. The agreement is within about 4·10⁻⁴, and the tolerance is 2·10⁻³. It is not tighter because mgc prints the loss to only six digits, which limits a finite-difference estimate with step 0.01 to errors around a few times 10⁻⁵, plus the estimate's own truncation error.
  6. Stress. Weights of ±1000 (where exp overflows): no nan or inf, and the loss, gradient and probabilities equal the reference.
-- training against the independent Python trainer
  ok   loss after 0 steps equals the reference (1.609438)
  ok   loss after 1 steps equals the reference (0.992623)
  ok   loss after 2 steps equals the reference (0.845013)
  ok   loss after 5 steps equals the reference (0.718402)
  ok   loss after 10 steps equals the reference (0.666934)
  ok   loss after 25 steps equals the reference (0.634284)
  ok   loss after 50 steps equals the reference (0.623229)
  ok   loss after 100 steps equals the reference (0.617695)
  ok   loss after 200 steps equals the reference (0.614931)
  ok   loss after 400 steps equals the reference (0.613551)
  ok   weights after 400 steps equal the reference
  ok   probabilities equal the reference
  ok   the gradient at the end equals the reference gradient (relative)
  ok   each row of probabilities sums to 1
-- the loss
  ok   the loss is strictly decreasing at the checkpoints
  ok   the loss never goes below the conditional entropy of the data (0.612174), the best any bigram model can do
  ok   after 400 steps the loss is within 0.002 of that bound
  ok   before training the loss is log(5) = 1.609438: every token equally likely
-- what it learned
  ok   token 0: learned probabilities [0.0, 0.799, 0.199, 0.0, 0.0] are within 0.01 of the observed [0.0, 0.8, 0.2, 0.0, 0.0]
  ok   token 1: learned probabilities [0.001, 0.001, 0.5, 0.25, 0.25] are within 0.01 of the observed [0.0, 0.0, 0.5, 0.25, 0.25]
  ok   token 2: learned probabilities [0.666, 0.001, 0.001, 0.332, 0.001] are within 0.01 of the observed [0.667, 0.0, 0.0, 0.333, 0.0]
  ok   token 3: learned probabilities [0.998, 0.0, 0.0, 0.0, 0.0] are within 0.01 of the observed [1.0, 0.0, 0.0, 0.0, 0.0]
-- the unseen row
  ok   token 4 is never a current token: its row of weights is still exactly 0
  ok   its gradient is exactly 0
  ok   so its probabilities are exactly uniform (0.2 each)
-- the gradient against finite differences (central difference of the loss, step 0.01)
  ok   random weights #1: largest difference between the gradient and the finite-difference estimate is 4.1e-04
  ok   random weights #1: gradient and loss equal the reference
  ok   random weights #2: largest difference between the gradient and the finite-difference estimate is 3.4e-04
  ok   random weights #2: gradient and loss equal the reference
  ok   random weights #3: largest difference between the gradient and the finite-difference estimate is 4.6e-04
  ok   random weights #3: gradient and loss equal the reference
-- stress: weights of +-1000
  ok   no nan and no inf anywhere in the output
  ok   the loss equals the reference (429.7430)
  ok   the gradient equals the reference
  ok   the probabilities equal the reference (each row is a one-hot-like 0 and 1)
all checks pass

Are the checks good enough? Run twelve wrong trainers.

model_mutation.py writes twelve versions of bigram.mg.defs, each with one deliberate mistake (the gradient with the wrong sign, forgotten averaging, forgotten targets, the loss with the wrong sign, the naive log-sum-exp, the step uphill, and so on), and runs the checker on each:

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)
the gradient has the wrong sign (y minus p instead of p minus y): caught by 25 failing checks (the reference, the loss, what it learned, finite differences)
the gradient is not averaged over the 14 pairs: caught by 24 failing checks (the reference, the loss, what it learned, finite differences)
the gradient forgets to subtract the targets y: caught by 25 failing checks (the reference, the loss, what it learned, finite differences)
the gradient uses x instead of its transpose: caught (the program did not even run: RuntimeError: mgc failed (1): /tmp/tmpgpz9dwh3.mg:12: error: cannot multiply (matmul) shapes 14x5 and 14x5: in)
the loss has the wrong sign: caught by 20 failing checks (the reference, the loss, finite differences)
the loss is averaged over 5, not 14: caught by 19 failing checks (the reference, the loss, finite differences)
the loss does not use the targets (it sums every token's log-probability): caught by 20 failing checks (the reference, the loss, finite differences)
log_softmax forgets the log of the sum (it only subtracts the maximum): caught by 27 failing checks (the reference, the loss, what it learned, finite differences)
log_softmax does not subtract the row maximum (the naive log-sum-exp): caught by 3 failing checks (other, the reference)
the step moves uphill (adds the gradient): caught by 18 failing checks (the reference, the loss, what it learned)
the learning rate is 4, not 8: caught by 13 failing checks (the reference, the loss)
the probabilities do not subtract the row maximum (naive softmax): caught by 2 failing checks (other, the reference)

All twelve are caught. As in Chapter 30, two were not caught at first: the naive log-sum-exp (no maximum subtracted) and the naive softmax in probabilities. The first version of the stress check used weights of ±400, and exp(400) is large but still representable (overflow starts near 709), so a naive formula gave the right answer. Raising the stress weights to ±1000 made the naive versions produce inf and nan, and the script was re-run. A numerical-stability test must actually reach the instability.

Compiler mutations

log 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 31 code on purpose, one way at a time, and re-runs the Chapter 31 tests (test/bigram31).
#   "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=':: bigram31/'
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 "log is compiled as relu by the front end"                           $F 's/{r} = mg.log {v}/{r} = mg.relu {v}/'
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 "log copies its input (computes nothing)"                           LowerToAffine.cpp 's/b.create<math::LogOp>(l, x)/x/'
try "log computes exp instead"                                           LowerToAffine.cpp 's/b.create<math::LogOp>(l, x)/b.create<math::ExpOp>(l, x)/'
try "log computes the square root instead"                               LowerToAffine.cpp 's/b.create<math::LogOp>(l, x)/b.create<math::SqrtOp>(l, x)/'
try "mg.log's verifier accepts anything"                                 MgDialect.cpp 's/return verifyElementwise(\*this, {getInput().getType(), getResult().getType()}, "mg.log"); }/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
log is compiled as relu by the front end: caught by bigram31/ir.mlir bigram31/examples.mlir bigram31/against-reference.mlir 

== Part 2: the dialect verifier and the lowering (a full rebuild per mutation)
log copies its input (computes nothing): caught by bigram31/examples.mlir bigram31/against-reference.mlir 
log computes exp instead: caught by bigram31/examples.mlir bigram31/against-reference.mlir 
log computes the square root instead: caught by bigram31/examples.mlir bigram31/against-reference.mlir 
mg.log's verifier accepts anything: caught by bigram31/verifiers.mlir 
a sum starts from 1 instead of 0 (every row_sum and col_sum is off by one): caught by bigram31/examples.mlir bigram31/against-reference.mlir 

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

All six are caught. As in Chapter 30, the verifier mutation is caught only by the dedicated verifiers test (no program the front end writes can reach it), and the last mutation, a sum that starts from 1 instead of 0, is not about log: it breaks row_sum and col_sum, which the loss and the log-softmax lean on, so the training tests would also notice a regression in an operation from an earlier chapter.

The full suite:

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

Testing Time: 15.00s

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

Limits and what is not established

  • A bigram looks at one token. It cannot represent "after 0 1 comes 2 but after 3 1 comes 4"; any context beyond the last token is invisible to it. It is the simplest language model, not a good one.
  • Training loss only. All 14 pairs are used for training and the same 14 for the loss. The model has been shown to fit them (it reaches the entropy bound); nothing here measures how it would do on text it has not seen. With data this small, "learned the frequencies" is also "memorized the data".
  • Pairs never seen get tiny probabilities, not zero, and token 4's successor is left uniform. A real model smooths or regularizes; this one does neither.
  • The gradient is derived by hand, for this model. Mountain Goat has no automatic differentiation. Chapter 26 derived a linear model's gradient; here a softmax model's. Training Chapter 30's transformer is not done: that needs the gradient of every operation in it (attention, layer normalization, the residual connections, the feed-forward network) derived by hand, or an automatic differentiation system, and neither exists here. Chapter 30's weights remain random.
  • Full-batch gradient descent with a fixed learning rate of 8, chosen by trying 4 and 8 on this data in plain Python (both converge; 8 is faster). No learning-rate schedule, no momentum, no minibatches, no stopping rule: it just runs 400 steps. The step size is stable here because rows of w are independent and the largest curvature is bounded; nothing says 8 is safe for other data.
  • No sampling or argmax. The program prints probabilities. Choosing the likeliest next token, or generating text, would need an argmax or random choice the language lacks.
  • Six digits. mgc prints six significant digits, so the comparisons with Python are at a relative 10⁻⁵, and the finite-difference check can only be as sharp as printing allows (stated above).
  • exp and log come from the platform's C library; their last-digit accuracy was not measured. GPU lowering is not covered.
  • Platform: x86-64 Linux, LLVM 18.1.3.

Reproducing

cd docs/part31/code && ./build.sh
python3 make_examples.py && ./run_examples.sh > examples_out.txt
./check_bigram.py > check_bigram_out.txt
./model_mutation.py > model_mutation_out.txt
./mutation.sh > mutation_out.txt
cd ../../part15/code && ./run_lit.sh          # 127 tests

Chapter summary

  • A language model predicts the next token; the bigram model looks at one token and keeps a table of scores per token. The loss is the average surprise, −log of the probability given to what really came next.
  • With log added, training is a few lines of Mountain Goat: the gradient is xᵀ(softmax(xw) − y)/14, derived by hand, and a step subtracts it.
  • The loss falls from log 5 toward the entropy of the data and cannot go below it; after 400 steps the model's probabilities equal the observed frequencies to about 0.001, and a token with no successors in the data keeps exactly the weights it started with.
  • Everything is checked against an independent Python trainer, against theory (the entropy bound), against finite differences, and under overflow-sized weights; twelve wrong trainers and six broken compilers are all caught, and two of the wrong trainers were caught only after the stress test was made big enough.
  • It is a toy: one token of context, training loss only, hand-derived gradient. Training a transformer needs gradients Mountain Goat cannot yet produce for itself.

Self-check questions

Each answer is collapsed; try the question first.

  1. Why is the loss exactly log 5 = 1.609438 before training?

    Answer

    All weights are 0, so every score is 0 and every softmax row is uniform: probability 1/5 for each of the five tokens. The surprise of any pair is −log(1/5) = log 5, and the average of equal numbers is the number itself.

  2. What does it mean that the loss cannot go below 0.612174, and why is that the floor?

    Answer

    0.612174 is the conditional entropy of the data: the average of −log(observed probability of the next token) when each row of probabilities equals the observed frequencies. Among all ways to assign probabilities to the next token given the current one, matching the observed frequencies gives the smallest average surprise on this data (Gibbs' inequality); any other assignment is worse. A bigram model's rows are free, so it can reach but not beat that value.

  3. Why are probabilities of unseen pairs (like 0 then 3) small but not exactly 0 after 400 steps?

    Answer

    A softmax output is e^s / Σ e^s, which is exactly 0 only if a score is minus infinity. Gradient descent keeps lowering the score (here to about −2.8) but each step makes a smaller change, so it only approaches zero. The loss never rewards going all the way: it is already almost at its floor.

  4. Why does row 4 of the weights stay exactly 0, and what does the model say about the token after 4?

    Answer

    Token 4 is never a current token, so no row of x has its 1 in column 4, and the gradient xᵀ(…) has a zero row 4 (it is a sum over pairs, and no pair contributes). The step subtracts 0, so those weights stay at their starting value 0, and the softmax of a zero row is uniform: 0.2 each. The model has no information, and reports that rather than inventing something.

  5. Why is log_softmax written as s − max − log(sum(exp(s − max))) instead of log(softmax(s))?

    Answer

    softmax produces probabilities, and a probability far below the maximum underflows to exactly 0; the log of 0 is minus infinity, and y * (−inf) with a 0 in y is nan. The log-sum-exp form never builds a probability, so it stays finite. It also subtracts the maximum before exp, so exp cannot overflow. The ±1000 stress check exercises both.

  6. How does the finite-difference check work, and why is its tolerance only 2·10⁻³?

    Answer

    For each weight, compute the loss with that weight raised by 0.01 and lowered by 0.01; the slope (L₊ − L₋)/0.02 estimates the derivative, and it should match the gradient entry. The tolerance is loose because mgc prints the loss to six significant digits, so each loss value has an error around 10⁻⁶, which divided by 0.02 is a few times 10⁻⁵, and because the estimate itself has an error from the finite step. Observed differences are about 4·10⁻⁴.

  7. The gradient formula xᵀ(p − y)/14 came from a hand derivation. What makes it trustworthy rather than a guess?

    Answer

    Two independent checks: it matches a numerical estimate that uses only the loss function (finite differences, which does not use the formula), and training with it reproduces an independent Python trainer step for step and ends at the known optimum (the observed frequencies, the entropy floor). A wrong formula (the wrong sign, a missing 1/14, forgotten targets) fails these, as the wrong-trainer runs show.

  8. Why were two of the twelve wrong trainers missed at first, and what does that teach about testing numerical code?

    Answer

    The naive log-sum-exp and the naive softmax compute the same numbers as the stable versions whenever exp does not overflow. The first stress check used scores of ±400, and exp(400) is representable (overflow begins near 709), so both naive formulas were correct on every test. Scores of ±1000 overflow, giving inf and nan, and caught both. A test for a numerical-stability bug must reach the regime where the bug bites.

  9. Why can a bigram model not represent "what comes after the pair 0 1"?

    Answer

    Its prediction depends only on the current token: the row of w for token 1. In the data, token 1 is followed by 2, 3 and 4, and the model can only give the average over every occurrence of token 1, whatever came before it. Capturing the effect of the earlier token needs a model that sees more than one token at a time, which is what attention in Chapters 29 and 30 provides.

  10. What would be needed to train the Chapter 30 transformer the way this chapter trained the bigram?

    Answer

    The gradient of the loss with respect to every weight in it: the embeddings, each head's query, key and value matrices, the output matrices, the layer-norm scales and shifts, the feed-forward weights and biases, and the final projection. That is the chain rule through softmax attention, layer normalization, relu and the residual connections, derived by hand for each, or an automatic differentiation system that does it from the program. Mountain Goat has neither, so Chapter 30's weights remain random numbers; the loss and the training loop of this chapter would carry over unchanged.