33. Backpropagation Through Attention: Training a One-Head Classifier

What you will understand: how a model with attention is trained, by working out its gradient by hand, step by step backwards through the computation (backpropagation), writing every step as a Mountain Goat function, and training the model with it. The model reads four tokens in any order and must say which one ranks highest under a hidden priority list: a job a bigram (Chapter 31) cannot do, because the answer depends on all four tokens. The backward pass, including the one tricky piece (the gradient through a softmax), is checked three ways: against an independent Python calculation done one token at a time, against a numerical estimate for every one of its 48 weights, and by whether the training run it drives lands where the reference lands. Then the trained model is tried on all 625 possible inputs, and the one it gets wrong is explained.

What you need to know first: Chapter 29 (attention), Chapter 31 (cross-entropy loss and the gradient of a softmax output) and Chapter 28's reshape. No new compiler operation is needed: the chapter's programs run on Chapter 32's compiler. It is also the first chapter whose model has more than one layer of structure to differentiate, so the section "The backward pass" is the heart of it.

Compile and run

cd docs/part32/code && ./build.sh                 # once: the newest compiler (this chapter adds nothing to it)
cd ../../part33/code
python3 make_examples.py                          # (re)writes the examples from attn_lib.py
../../part32/code/mgc run examples/03_train.mg    # train for 200 steps, then try every possible sequence (about 10 seconds)
./run_examples.sh > examples_out.txt              # every example and every error example on this page
./check_attention_training.py                     # the backward pass and the training run against the independent Python version
./model_mutation.py > model_mutation_out.txt      # runs 10 WRONG backward passes through the checker (nine are caught; see the page)
cd ../../part15/code && ./run_lit.sh              # the whole test suite

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

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

The three error examples are programs the compiler must reject (each ends exit status 1). The mutation output is a negative control: deliberately wrong backward passes run against the checker, where "caught" is the wanted result and "NOT CAUGHT" would be a gap.

The task

A sequence is four tokens, each one of 0 to 4. A hidden priority list ranks the tokens: 1, 3, 2, 0, 4, first is highest. The label of a sequence is the token present in it that comes first in the list:

sequence	label	why
2 1 1 1	1	token 1 is present, and it ranks first
3 2 0 2	3	no 1; 3 ranks next
2 2 3 0	3
3 4 2 2	3
4 4 0 4	0	no 1, 3 or 2; 0 beats 4
4 4 4 4	4	only token 4 is present

The model is shown 24 random sequences with their labels (the training set) and 40 more it never trains on (the held-out set). It is never told the rule. Order does not matter: 2 1 1 1 and 1 2 1 1 have the same label. That is the job attention is built for: it looks at all the tokens at once and combines them, with no preference for position.

(The data are deliberately lopsided, as random sequences are. Of the 24 training labels, 13 are 1, 8 are 3, 2 are 2, 1 is 0 and none is 4; of the 40 held-out labels, 22 are 1, 13 are 3, 5 are 2 and none is 0 or 4. Label 4 requires a sequence made only of 4s, which is 1 in 625. This will matter at the end.)

The model

One attention head, with a learned query q (1×3), and tables that turn a token into a key and a value (wk and wv, each 5×3, one row per token), then an output matrix wo (3×5) that turns the attended value into a score for each of the five possible labels:

1. each token's score is its key's dot product with the query, divided by √3: score_n = q · key(token_n) / √3;
2. a softmax over the four scores gives the attention weights a (four positive numbers adding to 1);
3. the context is the attention-weighted average of the four values: h = Σ a_n · value(token_n) (a 1×3 row);
4. the logits are h @ wo (one score per label), and a softmax and cross-entropy loss (Chapter 31) measure how surprised the model is by the true label.

That is 3 + 15 + 15 + 15 = 48 numbers to learn.

Doing it for 24 sequences at once, with only matrix operations. Mountain Goat has no batched matrix product, so the sequences are stacked: x has 96 rows (the four one-hot token rows of sequence 0, then of sequence 1, and so on). Keys and values for all tokens at once are x @ wk and x @ wv (96×3). The scores are (x @ wk) @ transpose(q), a 96×1 column; reshape turns it into 24×4, one row per sequence, so a row-wise softmax is a softmax within each sequence. To add up each sequence's four weighted values, a constant group matrix g (24×96, with a 1 where column ÷ 4 equals the row) does it as a matrix product: g @ (values * weights).

Attention for one sequence, by hand

First the forward pass on one sequence of four tokens (1, 3, 2, 0) with numbers small enough to follow. The query is (1, 0), so a token's score is the first number of its key (times 0.7071):

################ mgc run 01_attention_by_hand.mg
  | # Attention for one sequence of four tokens (1, 3, 2, 0), small enough to follow by hand. Each token has a 2-number key and a 2-number value.
  | # The query is (1, 0), so a token's score is just the first number of its key (times 1/sqrt(2) = 0.7071), and the second number of a value is 0 for tokens 1 and 3.
  | let x  = [[0, 1, 0, 0, 0], [0, 0, 0, 1, 0], [0, 0, 1, 0, 0], [1, 0, 0, 0, 0]]         # the four tokens, one-hot
  | let wk = [[0, 0], [2, 0], [1, 0], [1.5, 0], [0, 0]]                                      # key of each of the five tokens
  | let wv = [[1, 1], [10, 0], [0, 10], [5, 0], [0, 0]]                                      # value of each of the five tokens
  | let q  = [[1, 0]]
  | let scores = (x @ wk) @ transpose(q) * 0.7071067811865476        # 4x1: token 1 scores 2, token 3 scores 1.5, token 2 scores 1, token 0 scores 0 (all times 0.7071)
  | print scores
  | let a = exp(transpose(scores) - row_max(transpose(scores))) / row_sum(exp(transpose(scores) - row_max(transpose(scores))))      # softmax over the four scores: 1x4 attention weights
  | print a
  | print a @ (x @ wv)                                                # the attended value: the average of the four values, weighted by attention
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 1] strides = [1, 1] data = 
[[1.41421], 
 [1.06066], 
 [0.707107], 
 [0]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 4] strides = [4, 1] data = 
[[0.410109,   0.287974,   0.202212,   0.0997045]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 2] strides = [2, 1] data = 
[[5.64067,   2.12183]]
  ----- exit status 0

Token 1 scores 2, token 3 scores 1.5, token 2 scores 1 and token 0 scores 0, all times 0.7071: 1.41421, 1.06066, 0.707107, 0. The softmax gives the attention weights 0.410, 0.288, 0.202, 0.0997: token 1 gets the most. The context is the weighted average of the four values ([10, 0], [5, 0], [0, 10], [1, 1]): 0.410·10 + 0.288·5 + 0.0997·1 = 5.64 and 0.202·10 + 0.0997·1 = 2.12.

The backward pass

Training needs the gradient: how the loss changes with each of the 48 numbers. It is found by the chain rule, working backwards from the loss one step at a time. Each step takes the gradient with respect to a step's output and produces the gradient with respect to its inputs. Write B for the number of sequences (24), dX for "the gradient of the loss with respect to X":

forward step	backward step (what the Mountain Goat function computes)
logits = h @ wo, loss = cross-entropy	dlogits = (softmax(logits) − y) / B (Chapter 31's result, one row per sequence)
↳ weights wo	dwo = hᵀ @ dlogits
↳ context h	dh = dlogits @ woᵀ (one row per sequence)
h = Σ a_n · v_n	every token row of sequence b receives sequence b's dh: dh_tokens = gᵀ @ dh (96×3)
↳ values	dv = dh_tokens * a (each token's share of dh is scaled by its attention weight), and dwv = xᵀ @ dv
↳ attention weights	da_n = dh_n · v_n (how much the loss would change per unit of attention on token n): row_sum(dh_tokens * v)
a = softmax(s)	ds = a * (da − Σ a·da) (below)
s = q · k / √3	dq = dsᵀ @ K / √3, dK = ds @ q / √3, and dwk = xᵀ @ dK

The one new piece of calculus is the softmax backward step. If the loss depends on a = softmax(s) through da, then the derivative with respect to score s_i is a_i · (da_i − Σ_j a_j da_j). Each score affects its own weight directly (the a_i da_i part) and every other weight indirectly, because the weights must sum to 1 (the − a_i Σ_j a_j da_j part). Forgetting the second part is the most natural mistake. A tiny example computes it by the formula and by brute force (nudge each score by 0.01 up and down and watch the loss change):

################ mgc run 02_softmax_backward.mg
  | # The one new piece of calculus in this chapter: how the loss changes when the scores of a softmax change.
  | # If L = (d . softmax(s)) for a fixed vector d, the derivative with respect to s_i is  a_i * (d_i - sum_j a_j d_j)  where a = softmax(s).
  | # The first line prints that formula; the second computes the same derivatives by brute force (nudge each score by 0.01 up and down and watch L change).
  | def softmax(s: tensor[1x3]) = exp(s - row_max(s)) / row_sum(exp(s - row_max(s)))
  | def loss(s: tensor[1x3], d: tensor[1x3]) = row_sum(d * softmax(s))
  | def backward(s: tensor[1x3], d: tensor[1x3]) = softmax(s) * (d - row_sum(softmax(s) * d))
  | let s = [[0.5, -1, 2]]
  | let d = [[3, -2, 1]]
  | print backward(s, d)
  | print (loss(s + [[0.01, 0, 0]], d) - loss(s - [[0.01, 0, 0]], d)) / 0.02       # nudge s_0
  | print (loss(s + [[0, 0.01, 0]], d) - loss(s - [[0, 0.01, 0]], d)) / 0.02       # nudge s_1
  | print (loss(s + [[0, 0, 0.01]], d) - loss(s - [[0, 0, 0.01]], d)) / 0.02       # nudge s_2
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data = 
[[0.309696,   -0.12646,   -0.183235]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.309696]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[-0.126462]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[-0.183235]]
  ----- exit status 0

The formula's [0.309696, −0.12646, −0.183235] and the three brute-force numbers 0.309696, −0.126462, −0.183235 agree. The check on the whole model does this for all 48 weights.

Writing it down as functions (the definitions below are the program, in order: the forward pass first, then one function per gradient):

# One-head attention classifier. A sequence of 4 tokens (one-hot rows) goes in; the label is the token present that ranks first in a hidden priority list.
# B sequences are stacked: x has B*4 rows (the tokens of sequence 0, then sequence 1, ...) and g (B x B*4) adds up the 4 rows belonging to each sequence.
# q (1x3) is the learned query; wk, wv (5x3) turn a token into its key and value; wo (3x5) turns the attended value into scores for the 5 possible labels.
# (The same functions exist again with the suffix _te for the 40 held-out sequences and _ex for chunks of 25 sequences; only the sizes differ.)

def softmax_rows(m: tensor[24x4]) = exp(m - row_max(m)) / row_sum(exp(m - row_max(m)))
def attn_matrix(x: tensor[96x5], q: tensor[1x3], wk: tensor[5x3]) = softmax_rows(reshape((x @ wk) @ transpose(q) * 0.5773502691896258, 24, 4))
def context(x: tensor[96x5], g: tensor[24x96], q: tensor[1x3], wk: tensor[5x3], wv: tensor[5x3]) = g @ ((x @ wv) * reshape(attn_matrix(x, q, wk), 96, 1))
def logits(x: tensor[96x5], g: tensor[24x96], q: tensor[1x3], wk: tensor[5x3], wv: tensor[5x3], wo: tensor[3x5]) = context(x, g, q, wk, wv) @ wo
def log_softmax(z: tensor[24x5]) = z - row_max(z) - log(row_sum(exp(z - row_max(z))))
def loss(x: tensor[96x5], g: tensor[24x96], q: tensor[1x3], wk: tensor[5x3], wv: tensor[5x3], wo: tensor[3x5], y: tensor[24x5]) = col_sum(row_sum(y * log_softmax(logits(x, g, q, wk, wv, wo)))) * -0.041666666666666664
def correct(x: tensor[96x5], g: tensor[24x96], q: tensor[1x3], wk: tensor[5x3], wv: tensor[5x3], wo: tensor[3x5], y: tensor[24x5]) = col_sum(row_sum(ge(logits(x, g, q, wk, wv, wo), row_max(logits(x, g, q, wk, wv, wo))) * y))
def d_logits(x: tensor[96x5], g: tensor[24x96], q: tensor[1x3], wk: tensor[5x3], wv: tensor[5x3], wo: tensor[3x5], y: tensor[24x5]) = (exp(log_softmax(logits(x, g, q, wk, wv, wo))) - y) * 0.041666666666666664
def grad_wo(x: tensor[96x5], g: tensor[24x96], q: tensor[1x3], wk: tensor[5x3], wv: tensor[5x3], wo: tensor[3x5], y: tensor[24x5]) = transpose(context(x, g, q, wk, wv)) @ d_logits(x, g, q, wk, wv, wo, y)
def d_context_rows(x: tensor[96x5], g: tensor[24x96], q: tensor[1x3], wk: tensor[5x3], wv: tensor[5x3], wo: tensor[3x5], y: tensor[24x5]) = transpose(g) @ (d_logits(x, g, q, wk, wv, wo, y) @ transpose(wo))
def grad_wv(x: tensor[96x5], g: tensor[24x96], q: tensor[1x3], wk: tensor[5x3], wv: tensor[5x3], wo: tensor[3x5], y: tensor[24x5]) = transpose(x) @ (d_context_rows(x, g, q, wk, wv, wo, y) * reshape(attn_matrix(x, q, wk), 96, 1))
def d_attn_matrix(x: tensor[96x5], g: tensor[24x96], q: tensor[1x3], wk: tensor[5x3], wv: tensor[5x3], wo: tensor[3x5], y: tensor[24x5]) = reshape(row_sum(d_context_rows(x, g, q, wk, wv, wo, y) * (x @ wv)), 24, 4)
def d_scores(x: tensor[96x5], g: tensor[24x96], q: tensor[1x3], wk: tensor[5x3], wv: tensor[5x3], wo: tensor[3x5], y: tensor[24x5]) = reshape(attn_matrix(x, q, wk) * (d_attn_matrix(x, g, q, wk, wv, wo, y) - row_sum(attn_matrix(x, q, wk) * d_attn_matrix(x, g, q, wk, wv, wo, y))), 96, 1)
def grad_q(x: tensor[96x5], g: tensor[24x96], q: tensor[1x3], wk: tensor[5x3], wv: tensor[5x3], wo: tensor[3x5], y: tensor[24x5]) = transpose(d_scores(x, g, q, wk, wv, wo, y)) @ (x @ wk) * 0.5773502691896258
def grad_wk(x: tensor[96x5], g: tensor[24x96], q: tensor[1x3], wk: tensor[5x3], wv: tensor[5x3], wo: tensor[3x5], y: tensor[24x5]) = transpose(x) @ (d_scores(x, g, q, wk, wv, wo, y) @ q) * 0.5773502691896258

attn_matrix is the 24×4 matrix of attention weights (the reshape makes the softmax work within each sequence); context is the group-matrix product; d_logits, grad_wo, d_context_rows, grad_wv, d_attn_matrix, d_scores, grad_q and grad_wk are the backward table, row by row. The gradients with respect to wv and wk collect contributions from every token of every sequence through xᵀ @ …: a token that appears in several places receives the sum of the gradients of those places, which the matrix product does automatically.

Training

examples/03_train.mg is the definitions, the data (as matrices), seeded starting weights, then 200 unrolled steps (four let lines each: every parameter minus the learning rate 3 times its gradient), and prints. First the loss and the number of correctly classified training sequences at eight checkpoints (pairs of numbers: loss, then count):

Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[1.60134]]
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.51356]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[18]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.625597]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[21]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.273815]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[21]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.0903657]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[24]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.0098207]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[24]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.00168072]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[24]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.000522397]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[24]]

after step	0	1	5	10	25	50	100	200
loss	1.601	1.514	0.626	0.274	0.0904	0.0098	0.00168	0.000522
correct of 24	0	18	21	21	24	24	24	24

The loss starts near log 5 = 1.609 (the model knows nothing; the seeded start gives 1.601) and falls at every checkpoint to 0.0005. After one step 18 of 24 are already right, helped by the lopsided data (13 of the 24 are labelled 1). All 24 are right from step 25 on. The final trained numbers, and then each token's attention score (wk @ transpose(q) / √3, the last matrix):

Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data = 
[[1.04345,   -3.07148,   -0.181485]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [5, 3] strides = [3, 1] data = 
[[-0.476807,   1.78582,   -0.140316], 
 [0.782224,   -1.83158,   -0.376133], 
 [-0.519774,   -0.114151,   0.0289385], 
 [0.258599,   -1.26574,   0.346387], 
 [-0.374243,   1.32564,   0.321124]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [5, 3] strides = [3, 1] data = 
[[1.0821,   -0.0214015,   -0.259401], 
 [-0.0367175,   2.77773,   -2.25651], 
 [1.69849,   -0.749124,   2.27048], 
 [-1.37339,   1.83481,   2.81956], 
 [1.88499,   -1.33383,   -1.34076]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [3, 5] strides = [5, 1] data = 
[[1.67379,   -0.251191,   1.53316,   -1.78568,   -0.910079], 
 [-1.65609,   2.69704,   -1.09644,   1.2281,   -0.782607], 
 [-1.84888,   -2.10787,   2.20213,   2.56983,   -0.865204]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [5, 1] strides = [1, 1] data = 
[[-3.43937], 
 [3.75862], 
 [-0.113735], 
 [2.36405], 
 [-2.60989]]

What did the attention learn?

The five scores, by token: 0: −3.44, 1: +3.76, 2: −0.11, 3: +2.36, 4: −2.61. Sorted from highest: 1, 3, 2, 4, 0. The hidden rule is 1, 3, 2, 0, 4. So the model found the first three correctly and in order, and swapped the last two: it gives token 4 more attention than token 0. (The checker confirms this rather than assuming the model learned the rule.)

Does the model attend to the answer? Here are the attention weights for all 24 training sequences:

Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [24, 4] strides = [4, 1] data = 
[[0.00688869,   0.331037,   0.331037,   0.331037], 
 [0.854062,   0.0716807,   0.0025769,   0.0716807], 
 [0.851227,   0.00588698,   0.0714428,   0.0714428], 
 [0.0716807,   0.0716807,   0.854062,   0.0025769], 
 [0.25,   0.25,   0.25,   0.25], 
 [0.958466,   0.0016435,   0.0199451,   0.0199451], 
 [0.00150407,   0.00150407,   0.498496,   0.498496], 
 [0.800361,   0.000598743,   0.198442,   0.000598743], 
 [0.0401413,   0.0033077,   0.478276,   0.478276], 
 [0.000373766,   0.499626,   0.499626,   0.000373766], 
 [0.00343419,   0.496566,   0.00343419,   0.496566], 
 [0.0686248,   0.832811,   0.0686248,   0.0299393], 
 [0.499385,   0.000856307,   0.000373585,   0.499385], 
 [0.800361,   0.000598743,   0.198442,   0.000598743], 
 [0.00137028,   0.198135,   0.799124,   0.00137028], 
 [0.478276,   0.478276,   0.0401413,   0.0033077], 
 [0.00629999,   0.910945,   0.0764549,   0.00629999], 
 [0.494668,   0.494668,   0.000370057,   0.0102938], 
 [0.00134968,   0.787114,   0.0163794,   0.195157], 
 [0.0203361,   0.000731077,   0.977257,   0.00167573], 
 [0.195306,   0.787713,   0.0163918,   0.000589281], 
 [0.329386,   0.329386,   0.329386,   0.0118413], 
 [0.291013,   0.291013,   0.126962,   0.291013], 
 [0.109285,   0.440771,   0.440771,   0.00917219]]

For 3 1 2 0 (label 1) the weights are 0.195, 0.788, 0.016, 0.001: most attention goes to the 1, as hoped. But the weights are not a clean lookup: for 13 of the 24 sequences more than 90% of the attention lands on label tokens, for only 6 more than 99%, and for the sequence 4 4 0 4 (label 0) the weights are 0.291, 0.291, 0.127, 0.291: the model puts less than half as much attention on the 0 as on each of the 4s, and still gets the answer right. The values (and wo) do the rest of the work: the model learned a combination of how much to attend and what each token contributes, not "attend to the winner". This is a general caution: attention weights are part of the computation, not an explanation of it.

Held-out and every possible sequence

The model never trained on the 40 held-out sequences:

Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[0.000605852]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[40]]

Held-out loss 0.000606 and 40 of 40 correct. But look at the labels of those 40: 22 are 1, 13 are 3, 5 are 2. A model that always output 1 or 3 would score well, so this is a weak test. A stronger one: there are only 5⁴ = 625 possible sequences, and the program tries all of them, 25 at a time. The number correct in each chunk of 25:

Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[25]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[24]]

Twenty-four chunks have 25 correct, and the last has 24. In total 624 of 625 sequences are classified correctly: all of them but one. Which one? The last sequence in the last chunk, 4 4 4 4, label 4. Here are the model's scores for that chunk (its last row is 4 4 4 4):

Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [25, 5] strides = [5, 1] data = 
[[6.16721,   -0.800652,   1.31225,   -6.68039,   0.114174], 
 [-0.459043,   12.2022,   -8.03186,   -2.34312,   -0.186213], 
 [1.05089,   -6.18801,   7.24211,   0.328861,   -2.3903], 
 [-10.262,   -0.655251,   2.07933,   11.6308,   -2.57759], 
 [7.14251,   -1.05912,   1.36328,   -7.7101,   0.331987], 
 [-0.459043,   12.2022,   -8.03186,   -2.34312,   -0.186213], 
 [-0.475334,   12.2342,   -8.05483,   -2.33246,   -0.186952], 
 [-0.454073,   11.8161,   -7.70418,   -2.25708,   -0.241428], 
 [-2.46016,   9.66297,   -6.0315,   0.48932,   -0.669194], 
 [-0.446953,   12.1882,   -8.02257,   -2.35333,   -0.184647], 
 [1.05089,   -6.18801,   7.24211,   0.328861,   -2.3903], 
 [-0.454073,   11.8161,   -7.70418,   -2.25708,   -0.241428], 
 [0.491467,   -6.77707,   7.89048,   1.09526,   -2.66414], 
 [-9.52075,   -1.16055,   2.56717,   10.9246,   -2.60911], 
 [1.46287,   -6.04604,   7.03284,   -0.165972,   -2.24757], 
 [-10.262,   -0.655251,   2.07933,   11.6308,   -2.57759], 
 [-2.46016,   9.66297,   -6.0315,   0.48932,   -0.669194], 
 [-9.52075,   -1.16055,   2.56717,   10.9246,   -2.60911], 
 [-10.4241,   -0.653817,   2.08689,   11.8114,   -2.60414], 
 [-10.1765,   -0.661824,   2.07759,   11.5369,   -2.56224], 
 [7.14251,   -1.05912,   1.36328,   -7.7101,   0.331987], 
 [-0.446953,   12.1882,   -8.02257,   -2.35333,   -0.184647], 
 [1.46287,   -6.04604,   7.03284,   -0.165972,   -2.24757], 
 [-10.1765,   -0.661824,   2.07759,   11.5369,   -2.56224], 
 [7.84292,   -1.24474,   1.39993,   -8.44957,   0.488406]]

The last row is [7.84, −1.24, 1.40, −8.45, 0.49]: the highest score is token 0, so the model says 0 where the answer is 4. This failure has a clear cause. Label 4 never occurs in any of the 64 sequences the model saw, so nothing in training ever rewarded outputting 4: to a model trained on 24 sequences with labels 0, 1, 2 and 3, the sequence 4 4 4 4 carries no evidence at all. (Compare Chapter 32: the held-out pairs that failed were the ones the training pairs never showed.) What is notable is the other direction: the model is correct on all 15 other sequences made only of 0s and 4s, whose label is 0, though in training it saw exactly one sequence with label 0.

Mistakes the compiler reports

################ mgc run errors/e1_one_sequence_short.mg
  | # The defs were written for 24 sequences (96 token rows). A batch of 23 sequences (92 rows) does not fit, and the compiler says so.
  | def softmax_rows(m: tensor[24x4]) = exp(m - row_max(m)) / row_sum(exp(m - row_max(m)))
  | def attn_matrix(x: tensor[96x5], q: tensor[1x3], wk: tensor[5x3]) = softmax_rows(reshape((x @ wk) @ transpose(q) * 0.5773502691896258, 24, 4))
  | let x = [[0, 0, 1, 0, 0], [0, 1, 0, 0, 0], [0, 1, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 0, 1, 0], [0, 0, 1, 0, 0], [1, 0, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 1, 0], [0, 0, 0, 0, 1], [0, 0, 1, 0, 0], [0, 0, 1, 0, 0], [0, 0, 1, 0, 
  | print attn_matrix(x, [[0.1, 0.2, 0.3]], [[0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]])
  ----- output
errors/e1_one_sequence_short.mg:5: error: argument shape 92x5 does not fit parameter 96x5
  ----- exit status 1

################ mgc run errors/e2_group_matrix_transposed.mg
  | # The group matrix is 24 x 96 (one row per sequence). Passing its transpose, 96 x 24, is a shape error caught when compiling.
  | def softmax_rows(m: tensor[24x4]) = exp(m - row_max(m)) / row_sum(exp(m - row_max(m)))
  | def attn_matrix(x: tensor[96x5], q: tensor[1x3], wk: tensor[5x3]) = softmax_rows(reshape((x @ wk) @ transpose(q) * 0.5773502691896258, 24, 4))
  | def context(x: tensor[96x5], g: tensor[24x96], q: tensor[1x3], wk: tensor[5x3], wv: tensor[5x3]) = g @ ((x @ wv) * reshape(attn_matrix(x, q, wk), 96, 1))
  | let x = [[0, 0, 1, 0, 0], [0, 1, 0, 0, 0], [0, 1, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 0, 1, 0], [0, 0, 1, 0, 0], [1, 0, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 1, 0], [0, 0, 0, 0, 1], [0, 0, 1, 0, 0], [0, 0, 1, 0, 0], [0, 0, 1, 0, 
  | let g_wrong = [[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
  | let w = [[0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]]
  | print context(x, g_wrong, [[0, 0, 0]], w, w)
  ----- output
errors/e2_group_matrix_transposed.mg:8: error: argument shape 96x24 does not fit parameter 24x96
  ----- exit status 1

################ mgc run errors/e3_reshape_element_count.mg
  | # reshape keeps the numbers, so the counts must match: 8 numbers cannot fill a 4 x 3 shape (12 numbers). (The chapter reshapes 96 x 1 to 24 x 4: the same 96 numbers.)
  | let a = [[1], [2], [3], [4], [5], [6], [7], [8]]
  | print reshape(a, 4, 3)
  ----- output
errors/e3_reshape_element_count.mg:3: error: cannot reshape 8x1 (8 elements) to 4x3 (12 elements)
  ----- exit status 1

A batch of 23 sequences (92 rows) where 24 (96 rows) are expected; the group matrix passed transposed (96×24 where 24×96 is wanted); and a reshape that does not keep the number of elements. As in earlier chapters the shapes are types, and a trainer cannot be connected to data of the wrong size.

Tests

Three new lit files in test/attention33/ (the suite is now 135 tests). There is no verifiers or ir test because this chapter adds no compiler operation:

Test	What it checks
examples	the two small examples (worked out above) and the training run's loss and correct-count at every checkpoint
errors	the three shape mistakes give the exact messages
against-reference	check_attention_training.py, below

check_attention_training.py compares with an independent Python implementation (plain lists; the gradient derived token by token with loops, not as matrix algebra) and checks, in order:

1. The gradient, for three different random starting weight sets: the loss and all four gradients equal the reference's, and the gradient of every one of the 48 weights agrees with a finite-difference estimate (nudge the weight by 0.01 up and down, watch the loss). The largest disagreement is about 4·10⁻⁴; the tolerance is 2·10⁻³ because mgc prints the loss to six digits.
2. Training: the loss and the correct count at all eight checkpoints, the final weights, each token's attention score and the 24×4 attention weights equal the reference's.
3. Learning: the loss falls at every checkpoint to below 0.001; all 24 training and all 40 held-out sequences are right.
4. Every sequence: the count in each of the 25 chunks equals the reference's, the total is 624, and the single failure is 4 4 4 4, with the label that never occurs in the data.
5. Attention scores: found rather than assumed: the three highest tokens are the three highest in the priority list in order, and the last two are in the opposite order (so the check would fail if a different training run learned the rule exactly, which is the point of looking).

-- the hand-derived backward pass against the independent per-token gradient, and against finite differences
  ok   starting weights #1: the loss and all four gradients equal the reference
  ok   starting weights #1: all 48 gradient entries agree with finite differences (largest difference 3.6e-04)
  ok   starting weights #2: the loss and all four gradients equal the reference
  ok   starting weights #2: all 48 gradient entries agree with finite differences (largest difference 4.7e-04)
  ok   starting weights #3: the loss and all four gradients equal the reference
  ok   starting weights #3: all 48 gradient entries agree with finite differences (largest difference 3.9e-04)
-- training against the independent Python trainer
  ok   step 0: loss 1.601341 and 0 of 24 correct equal the reference
  ok   step 1: loss 1.513557 and 18 of 24 correct equal the reference
  ok   step 5: loss 0.625597 and 21 of 24 correct equal the reference
  ok   step 10: loss 0.273815 and 21 of 24 correct equal the reference
  ok   step 25: loss 0.090366 and 24 of 24 correct equal the reference
  ok   step 50: loss 0.009821 and 24 of 24 correct equal the reference
  ok   step 100: loss 0.001681 and 24 of 24 correct equal the reference
  ok   step 200: loss 0.000522 and 24 of 24 correct equal the reference
  ok   the final weights equal the reference
  ok   each token's attention score equals the reference
  ok   the attention weights on the 24 training sequences equal the reference
  ok   every row of attention weights sums to 1
-- learning
  ok   the loss falls at every checkpoint
  ok   from 1.6013 (about log 5 = 1.609) to 0.000522
  ok   all 24 training sequences are classified correctly
  ok   all 40 held-out sequences are classified correctly (loss 0.000606)
-- every one of the 625 possible sequences
  ok   the number correct in each of the 25 chunks of 25 equals the reference's
  ok   624 of 625 sequences are classified correctly (got 624)
  ok   the one wrong sequence is 4 4 4 4 (label 4, predicted 0): the last row of the last chunk
  ok   label 4 never occurs among the 64 training and held-out sequences (labels seen: [0, 1, 2, 3])
-- attention scores
  ok   the three highest-scoring tokens [1, 3, 2] are the three highest in the priority list [1, 3, 2], in order
  ok   the last two are in the opposite order to the priority rule: scores order [1, 3, 2, 4, 0] vs rule [1, 3, 2, 0, 4] (this model did not learn to rank 0 above 4)
all checks pass

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

model_mutation.py writes ten versions of the definitions, each with one mistake mostly in the hand-derived backward pass (the softmax backward missing its second term, a forgotten 1/√d, the wrong sign, the wrong average, keys used where values belong, and so on), and runs the checker on each:

NOTE: this script deliberately runs wrong backward passes. 'caught' lines are EXPECTED: they show the checker can detect the mistake.
baseline (nothing broken): checker passes (expected)
the softmax backward forgets its second term (treats the softmax as elementwise): caught by 25 failing checks (the reference, finite differences, learning, other, every sequence, attention scores)
grad_q forgets the 1/sqrt(d) scale: caught by 13 failing checks (the reference, finite differences)
grad_wk forgets the 1/sqrt(d) scale: caught by 14 failing checks (the reference, finite differences)
grad_wv does not weight by the attention (uses the unweighted upstream): caught by 16 failing checks (the reference, finite differences)
d_logits has the wrong sign (y minus p): caught by 25 failing checks (the reference, finite differences, learning, other, every sequence, attention scores)
d_logits is not averaged over the 24 sequences: caught by 26 failing checks (the reference, finite differences, other, learning, every sequence, attention scores)
grad_wo uses the values instead of the attended context: caught by 16 failing checks (the reference, finite differences)
the attention scores are not scaled by 1/sqrt(d) in the forward pass (the gradient no longer matches): caught by 19 failing checks (the reference, finite differences, learning, every sequence)
d_attn_matrix uses the keys instead of the values: caught by 21 failing checks (the reference, finite differences, other, learning, every sequence, attention scores)
the naive softmax in the attention (no row maximum subtracted): NOT CAUGHT
(The last one is expected to be NOT CAUGHT: the attention scores here stay below about 4 in size, where exp cannot overflow, so the naive formula gives the
 same numbers. Chapters 29 to 32 showed the overflow regime; this chapter's checks do not enter it.)

Nine of the ten are caught, and by the right checks. The two forgotten 1/√d scales and the missing attention weighting in grad_wv are caught by the reference and by the finite-difference check but not by the learning checks that follow, because a gradient that is wrong in size or in detail but still points roughly downhill still trains (more slowly or faster, to the same kind of answer here); only comparing the gradient itself notices. The wrong-sign and wrong-average versions of d_logits are caught by almost everything. The forward-pass scale mutation shows why the finite-difference check is worth having: the hand-derived gradient still describes the original function, so it no longer matches the changed loss. The tenth (a naive softmax) is not caught, as I predicted before running it: nothing in this chapter drives the attention scores into the range where exp overflows.

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

The full suite:

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

Testing Time: 34.78s

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

Limits and what is not established

• A toy task and a toy model. Four tokens, five token types, one head, no positions, 48 parameters. The task was chosen so that the answer is the same for every ordering of the tokens, which suits attention without position information. Nothing here shows the method scales or that a full transformer trains the same way.
• Only the backward pass of THIS model is derived. There is no automatic differentiation: the nine backward functions are specific to this network. Chapter 30's transformer (layer normalization, two heads, a feed-forward network, residual connections, several blocks) would need a hand-derived backward pass for each of its parts; the pattern above (take the gradient with respect to a step's output, produce the gradient with respect to its inputs) carries over, but it was not done, and Chapter 30's weights are still random.
• Training data are lopsided and tiny. 13 of 24 training labels are 1; label 4 never occurs. The 40-sequence held-out score (40/40) is a weak test, for the reason shown; the exhaustive test over all 625 sequences is stronger, and it is possible only because the input space is tiny. 624/625 says nothing about larger problems.
• The one miss is a data gap, not necessarily a method failing. A different set of 24 sequences, or one that contained a sequence labelled 4, would change it. Nothing here tries other data sets or other starting weights beyond the three used for the gradient check.
• Learning rate and steps were chosen by trying. 3 and 200, after trying 1 and 3 on two starting points in plain Python. No claim of optimality, no schedule, no stopping rule.
• The model solves the task in its own way, not by the rule (it ranks token 4 above 0 and still classifies 0 4 0 4-style sequences correctly, through the values). The attention weights should not be read as an explanation.
• Finite-difference checks are limited by printing: six digits, step 0.01, tolerance 2·10⁻³.
• No naive-softmax test: the attention scores stay below about 4 in size, where exp does not overflow, so the checks cannot tell a naive softmax from the stable one here (the last wrong version above shows exactly that). Chapters 29 to 32 cover that regime.
• Platform: x86-64 Linux, LLVM 18.1.3.

Reproducing

cd docs/part32/code && ./build.sh
cd ../../part33/code
python3 make_examples.py && ./run_examples.sh > examples_out.txt
./check_attention_training.py > check_attention_training_out.txt
./model_mutation.py > model_mutation_out.txt
cd ../../part15/code && ./run_lit.sh          # 135 tests

Chapter summary

• Backpropagation is the chain rule applied backwards: from the loss, take each step's output gradient and produce its input gradients. For one-head attention that is nine short matrix expressions, and the only new calculus is the softmax backward step a * (da − Σ a·da).
• Without a batched matrix product, a batch is handled by stacking sequences, a constant group matrix to sum within a sequence, and reshape so a row-wise softmax acts within each sequence.
• The hand-derived backward pass equals an independent token-by-token computation and a finite-difference estimate for all 48 weights, for three different starting points.
• Trained for 200 steps, the model classifies all 24 training and 40 held-out sequences correctly and 624 of all 625 possible sequences; the miss, 4 4 4 4, is a label that never appears in the data.
• The attention is partly a lookup (it ranks tokens 1, 3, 2 correctly) and partly not (it ranks 4 above 0, and attends to the 4s in a sequence labelled 0): attention weights are not an explanation.
• Nine of ten wrong backward passes are caught, mostly by the gradient comparisons alone; the tenth, a naive softmax, is not caught, as expected, because the scores here never get large enough to overflow exp.

Self-check questions

Each answer is collapsed; try the question first.

1. Why can a bigram not solve this task, and why is attention a natural fit?

   Answer

   The label depends on all four tokens (which of them are present), not on the most recent one, and a bigram sees only one token. Attention looks at every token in the sequence at once and combines them into one context vector, with no dependence on order, which matches a label that is the same for every ordering.

2. Why does the model stack 24 sequences into a 96-row matrix, and what do the group matrix and the reshape do?

   Answer

   Mountain Goat has no batched matrix product, so a batch is handled as one big matrix. The scores come out as a 96×1 column; reshape makes it 24×4 so that a row-wise softmax runs within each sequence (each row holds one sequence's four scores). The group matrix (24×96, with 1s where column ÷ 4 equals the row) turns the sum over each sequence's four weighted values into a matrix product.

3. Write the gradient of the loss with respect to the attention scores in one line, and explain each part.

   Answer

   ds = a * (da − Σ a·da), where a is the attention weights, da the gradient with respect to them, and the sum is within each sequence. A score changes its own weight (the a·da part) and, because the four weights must add to 1, also every other weight (the − a·Σ a·da part). Leaving out the second part treats the softmax as if each output depended only on its own input.

4. Why does the checker compare the gradient both with an independent Python calculation and with finite differences?

   Answer

   They catch different things. The Python gradient is written independently from the same mathematics: it catches a mistake in the Mountain Goat translation, but if both derive the same wrong formula it would agree with itself. Finite differences use only the loss function, not any derived formula, so they catch a wrong derivation. A wrong gradient has to fool both to survive.

5. The model attends more to the 4s than to the 0 in 4 4 0 4 and still outputs the correct label 0. What does this say about reading attention weights?

   Answer

   The output depends on the attention weights and the values and the output matrix. The model learned values and an output matrix that, combined with this attention pattern, favour label 0 when the sequence contains a 0 and nothing higher. So attention weights alone do not say what the model "looked at" in a meaningful sense; they are one part of the computation.

6. Why is "40 of 40 held-out sequences correct" a weak result here, and what is the stronger test?

   Answer

   Random sequences are dominated by labels 1 and 3 (22 and 13 of the 40, with 5 labelled 2 and none labelled 0 or 4). A model that handles only the common labels can score well. Because there are only 625 possible sequences, the program can try every one; that test includes the rare labels, and it is where the single failure appears.

7. Why does the model fail on 4 4 4 4, and would you call that overfitting?

   Answer

   The label of 4 4 4 4 is 4, and label 4 never occurs among the 64 training and held-out sequences, so nothing in training ever rewarded the output 4. It is a gap in the data, not memorizing noise: the model is right on every other sequence, including the 15 that have label 0 though only one had that label in training. It is closer to a missing-evidence failure than to Chapter 32's overfitting, though both come from the model having no way to know what training never showed.

8. Which parts of the backward pass would be needed to train Chapter 30's transformer that are not here?

   Answer

   The gradients through layer normalization, the residual connections (which add gradients along two paths), a second attention head and its output split, the position matrix, the relu feed-forward network and its biases, and the stacking of several blocks, each with its own weights. The principle in this chapter's table applies to each, but each needs its own derivation and its own check against finite differences; none was done.

9. The wrong backward pass with a naive softmax in the attention was not caught. Why is that expected, and what would catch it?

   Answer

   The attention scores in this model stay below about 4 in size, and exp of such numbers cannot overflow, so the naive softmax computes the same numbers as the stable one. To catch it a test needs scores in the hundreds, like the stress cases of Chapters 29 to 32.

10. What would change if the labels were 1 for sequences containing token 1 and 0 otherwise (two classes)?

    Answer

    Only the number of output columns: wo would be 3×2 instead of 3×5, and the label matrices 24×2. The forward and backward formulas are the same, because the loss and the gradient (softmax − y)/B do not depend on the number of classes. In Mountain Goat the shapes are types, so the function definitions would be regenerated with 2 in place of 5.