Skip to content

30. A Small Transformer

Mountain goats on the mountain above a desert

What you will understand: what a transformer actually computes, by reading one written in Mountain Goat: token embeddings, positions, layer normalization, multi-head causal self-attention, a feed-forward network, residual connections, and a final softmax that gives next-token probabilities. The model is two blocks deep and small enough that every number can be inspected. Mountain Goat gets one new operation, sqrt; everything else is Chapters 21 to 29. The program is checked against an independent Python implementation at every stage, and against properties any correct transformer must have: causality (the future cannot influence the past), permutation equivariance (without positions, reordering the tokens reorders the outputs), normalized layer-norm rows, and no overflow when scores are huge.

What you need to know first: Chapter 29 (softmax and attention, which this chapter builds on without re-deriving), Chapter 22 (reductions and broadcasting) and the language tour. If the word "transformer" is new, read the section "What a transformer is" below first; it assumes only the matrix operations you have already seen.

Compile and run

cd docs/part30/code && ./build.sh                 # once: the newest compiler (adds mg.sqrt)
python3 make_examples.py                          # (re)writes the examples from transformer.mg.defs
./mgc run examples/03_tiny_transformer.mg         # compile and run the whole model
./run_examples.sh > examples_out.txt              # every example and every error example on this page
./check_transformer.py                            # the model against the reference, and the properties
./model_mutation.py > model_mutation_out.txt      # runs 15 WRONG models 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.

This model is not trained

The weights are seeded pseudo-random numbers, not learned. The model is a faithful, checked implementation of the transformer's computation, so its next-token probabilities are meaningless as language (they scatter around the uniform 0.2 over the five-token vocabulary). Training would need gradients through every operation here, which Mountain Goat does not have (Chapter 26 trained a much simpler model by hand-derived gradients). Nothing on this page claims the model has learned anything.

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

The four error examples are programs the compiler must reject (each ends exit status 1). The two mutation outputs are negative controls: a deliberately wrong compiler, and deliberately wrong models, run against the checks, where "caught" is the wanted result and "NOT CAUGHT" would be a gap.

What a transformer is

A transformer reads a sequence of tokens (words or pieces of words, here just the numbers 0 to 4) and, for each position, produces a probability for every token that might come next. It does so with a stack of identical-shaped blocks, each of which does two things to a matrix with one row per token:

  1. Attention lets every token gather information from the other tokens (Chapter 29). Rows are mixed.
  2. A feed-forward network transforms each token on its own, the same way for every token. Rows are not mixed.

Each of the two steps is wrapped the same way: a layer normalization first (rescale each token's numbers to a standard size), then the step, then a residual connection (add the step's output back to its input, so each step only has to adjust the token, not rebuild it). Before the first block, token numbers become vectors by looking up a row of an embedding matrix, and a position vector is added, because attention on its own cannot tell which token came first. After the last block, one more layer norm and a matrix multiplication give a score per vocabulary entry, and a softmax turns the scores into probabilities.

This chapter writes exactly that, in Mountain Goat, with these sizes: 4 tokens (so a 4x4 score matrix), model width 4, 2 heads of width 2, feed-forward width 8, vocabulary 5, 2 blocks.

The one new operation: sqrt

Layer normalization divides by a standard deviation, which needs a square root. sqrt(m) takes the square root of every element, added the same way as exp in Chapter 29: mg.sqrt in the dialect with a shape verifier, a loop-nest lowering that uses math.sqrt (which becomes an LLVM intrinsic), and a built-in in the front end. The square root of a negative number is nan, as IEEE arithmetic says:

print sqrt([[4, 2], [0, 9]])      ->  [[2, 1.41421], [0, 3]]
print sqrt([[-1]])                ->  [[-nan]]

(The minus sign in front of nan is the hardware's, and meaningless.)

The definitions

The whole model, as functions. Every def is one line (a language rule), so the long ones are long; the comments say what each part does.

# A small transformer, written in Mountain Goat. Shapes are fixed: 4 tokens, model width 4, 2 attention heads of width 2, feed-forward width 8, vocabulary 5.
# (Every def is one line, because a def is one line; the long ones are long.)

# softmax of each row, stabilised by subtracting the row maximum (Chapter 29), for the 4x4 attention scores and for the 4x5 vocabulary scores
def softmax(s: tensor[4x4]) = exp(s - row_max(s)) / row_sum(exp(s - row_max(s)))
def softmax_vocab(s: tensor[4x5]) = exp(s - row_max(s)) / row_sum(exp(s - row_max(s)))

# layer normalisation: each token (row) is shifted to mean 0 and scaled to variance 1, then rescaled by g and shifted by b (1 x 4, one per feature)
def layer_norm(x: tensor[4x4], g: tensor[1x4], b: tensor[1x4]) = (x - row_mean(x)) / sqrt(row_mean((x - row_mean(x)) * (x - row_mean(x))) + 0.00001) * g + b

# one attention head: queries, keys and values are projections of x; the mask (0 or -1e9) is added to the scores before the softmax (Chapter 29)
def head(x: tensor[4x4], wq: tensor[4x2], wk: tensor[4x2], wv: tensor[4x2], mask: tensor[4x4]) = softmax((x @ wq) @ transpose(x @ wk) * 0.7071067811865476 + mask) @ (x @ wv)

# two heads. Concatenating the heads and multiplying by one output matrix equals adding each head times its own half of that matrix (wo1 is the top half, wo2 the bottom)
def attention(x: tensor[4x4], mask: tensor[4x4], wq1: tensor[4x2], wk1: tensor[4x2], wv1: tensor[4x2], wo1: tensor[2x4], wq2: tensor[4x2], wk2: tensor[4x2], wv2: tensor[4x2], wo2: tensor[2x4]) = head(x, wq1, wk1, wv1, mask) @ wo1 + head(x, wq2, wk2, wv2, mask) @ wo2

# the position-wise feed-forward network: expand to 8 features, relu, contract back to 4
def ffn(x: tensor[4x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4]) = relu(x @ w1 + b1) @ w2 + b2

# the two sublayers of a block, each with a residual connection and a layer norm applied BEFORE the sublayer ("pre-norm")
def attention_sublayer(x: tensor[4x4], mask: tensor[4x4], g: tensor[1x4], bt: tensor[1x4], wq1: tensor[4x2], wk1: tensor[4x2], wv1: tensor[4x2], wo1: tensor[2x4], wq2: tensor[4x2], wk2: tensor[4x2], wv2: tensor[4x2], wo2: tensor[2x4]) = x + attention(layer_norm(x, g, bt), mask, wq1, wk1, wv1, wo1, wq2, wk2, wv2, wo2)
def ffn_sublayer(x: tensor[4x4], g: tensor[1x4], bt: tensor[1x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4]) = x + ffn(layer_norm(x, g, bt), w1, b1, w2, b2)

Read them from the bottom up, and compare with the description above:

  • ffn_sublayer and attention_sublayer are the two wrapped steps: x + step(layer_norm(x, ...)). The x + is the residual connection. Because layer_norm is applied to x before the step and the residual adds to the unnormalized x, this arrangement is called pre-norm (the original 2017 transformer normalized after, "post-norm"; pre-norm is what most current models use, and either can be written).
  • ffn is relu(x @ w1 + b1) @ w2 + b2: widen from 4 to 8 features, zero the negative ones, narrow back to 4. The biases b1 (1x8) and b2 (1x4) are one row stretched down all four tokens, the broadcasting of Chapter 22.
  • attention is two heads, each producing a 4x2 result. A real implementation concatenates the heads side by side (4x4) and multiplies by one 4x4 output matrix. Mountain Goat has no concatenation, but the two are the same arithmetic: concat(h1, h2) @ wo = h1 @ wo_top + h2 @ wo_bottom, so wo1 is the top half of that matrix and wo2 the bottom half. (Only because the matrix product distributes over the split; no approximation.)
  • head is Chapter 29's attention with projections: queries x @ wq, keys x @ wk, values x @ wv, the scores scaled by 1/sqrt(2) = 0.7071067811865476 and added to a mask (0 where a token may look, -1e9 where it may not) before the softmax.
  • layer_norm subtracts each row's mean, divides by the row's standard deviation (sqrt of the mean of squares of the centered row, plus 0.00001 so a constant row does not divide by zero), then multiplies by a learned scale g and adds a learned shift b (one number per feature, 1x4, stretched down the rows).
  • softmax and softmax_vocab are Chapter 29's stable softmax, once for the 4x4 attention scores and once for the 4x5 vocabulary scores.

Layer normalization

################ mgc run 01_layer_norm.mg
  | # Layer normalisation shifts each ROW (one token's features) to mean 0 and scales it to variance 1, then applies a per-feature scale g and shift b.
  | # Here g is all ones and b all zeros, so what is left is pure normalisation. The first row is 1 2 3 4: mean 2.5, variance 1.25, standard deviation 1.118.
  | def layer_norm(x: tensor[4x4], g: tensor[1x4], b: tensor[1x4]) = (x - row_mean(x)) / sqrt(row_mean((x - row_mean(x)) * (x - row_mean(x))) + 0.00001) * g + b
  | let x = [[1, 2, 3, 4], [10, 10, 10, 10], [-1, 0, 1, 0], [0, 0, 0, 8]]
  | print layer_norm(x, [[1, 1, 1, 1]], [[0, 0, 0, 0]])
  | # The same rows with a scale of 2 and a shift of 5: every normalised number is doubled, then 5 is added.
  | print layer_norm(x, [[2, 2, 2, 2]], [[5, 5, 5, 5]])
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 4] strides = [4, 1] data = 
[[-1.34164,   -0.447212,   0.447212,   1.34164], 
 [0,   0,   0,   0], 
 [-1.4142,   0,   1.4142,   0], 
 [-0.57735,   -0.57735,   -0.57735,   1.73205]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 4] strides = [4, 1] data = 
[[2.31673,   4.10558,   5.89442,   7.68327], 
 [5,   5,   5,   5], 
 [2.1716,   5,   7.8284,   5], 
 [3.8453,   3.8453,   3.8453,   8.4641]]
  ----- exit status 0

Row 0, 1 2 3 4, has mean 2.5 and standard deviation 1.118, so it becomes -1.34164 -0.447212 0.447212 1.34164. Row 1, 10 10 10 10, is constant: the numerator is 0 and the division is 0 / sqrt(0.00001), which is exactly 0 (this is what the epsilon is for). Row 2, -1 0 1 0, has mean 0 and standard deviation 0.7071, so it becomes -1.4142 0 1.4142 0. Row 3, 0 0 0 8, has mean 2 and standard deviation 3.4641. After normalization every row has mean 0 and variance 1 (the checker tests this on random matrices). The second matrix applies a scale of 2 and a shift of 5 to the same normalized rows: -1.34164 · 2 + 5 = 2.31672.

The feed-forward network

################ mgc run 02_feed_forward.mg
  | # The feed-forward network: expand to 8 features, keep only the positive ones (relu), contract back to 4.
  | # With w1 = [I | -I] (an identity next to a negated identity) and w2 = [I ; -I] the hidden layer holds relu(x) next to relu(-x), and the output
  | # is relu(x) - relu(-x), which is x itself: two relus can build a straight line. The numbers below confirm it.
  | def ffn(x: tensor[4x4], w1: tensor[4x8], b1: tensor[1x8], w2: tensor[8x4], b2: tensor[1x4]) = relu(x @ w1 + b1) @ w2 + b2
  | let x = [[1, -2, 0, 3], [-1, 1, -1, 1], [0, 0, 0, 0], [5, -5, 2, -2]]
  | let w1 = [[1, 0, 0, 0, -1, 0, 0, 0], [0, 1, 0, 0, 0, -1, 0, 0], [0, 0, 1, 0, 0, 0, -1, 0], [0, 0, 0, 1, 0, 0, 0, -1]]
  | let w2 = [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1], [-1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, -1]]
  | let zero8 = [[0, 0, 0, 0, 0, 0, 0, 0]]
  | let zero4 = [[0, 0, 0, 0]]
  | print ffn(x, w1, zero8, w2, zero4)
  | # A bias moves the hidden features before the relu: a bias of -1 on the first four hides every value below 1.
  | print ffn(x, w1, [[-1, -1, -1, -1, 0, 0, 0, 0]], w2, zero4)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 4] strides = [4, 1] data = 
[[1,   -2,   0,   3], 
 [-1,   1,   -1,   1], 
 [0,   0,   0,   0], 
 [5,   -5,   2,   -2]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 4] strides = [4, 1] data = 
[[0,   -2,   0,   2], 
 [-1,   0,   -1,   0], 
 [0,   0,   0,   0], 
 [4,   -5,   1,   -2]]
  ----- exit status 0

The example chooses w1 = [I | −I] and w2 = [I ; −I] (an identity beside a negated identity, and the reverse). The hidden layer then holds relu(x) beside relu(-x), and the output is relu(x) − relu(−x), which is x itself: the first matrix prints the input unchanged. Two relus can build a straight line. The second matrix gives the first four hidden features a bias of -1, which hides every value below 1: the output is relu(x − 1) − relu(−x), so row 0 (1 −2 0 3) becomes 0 −2 0 2. This is all a feed-forward layer is: a weighted sum, a bend, and another weighted sum.

The whole model

The generated program examples/03_tiny_transformer.mg is the definitions above, followed by the seeded weights (let statements), the token ids as one-hot rows, a causal mask, the sinusoidal position matrix, and the forward pass:

let x0 = onehot @ emb + positions
let a1 = attention_sublayer(x0, mask, g1_1, bt1_1, wq1_1, wk1_1, wv1_1, wo1_1, wq2_1, wk2_1, wv2_1, wo2_1)
let y1 = ffn_sublayer(a1, g2_1, bt2_1, w1_1, b1_1, w2_1, b2_1)
let a2 = attention_sublayer(y1, mask, g1_2, ...)
let y2 = ffn_sublayer(a2, g2_2, bt2_2, w1_2, b1_2, w2_2, b2_2)
let probs = softmax_vocab(layer_norm(y2, gf, bf) @ wout)

onehot @ emb is the embedding lookup (a one-hot row times the embedding matrix picks one row of it; matrix multiplication is the lookup). The positions are sin and cos of the position at two frequencies, as in the original paper, computed by the generator and written as numbers (the language has no sin). The token ids are [1, 3, 0, 2]. It prints four matrices: the embedded input, the output of block 1, the output of block 2, and the probabilities:

################ mgc run 03_tiny_transformer.mg
  | (72 lines: the definitions are transformer.mg.defs, followed by seeded weights, the token ids and the calls; see the file)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 4] strides = [4, 1] data = 
[[0.9,   0.34,   0.4,   0.45], 
 [0.3915,   0.2803,   0.98,   1.07], 
 [0.9393,   -1.0661,   -0.36,   1.0698], 
 [0.1311,   -1.74,   -0.8,   0.7796]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 4] strides = [4, 1] data = 
[[0.353842,   0.387514,   -0.308814,   -0.435457], 
 [-0.813699,   1.01803,   0.635739,   1.14997], 
 [0.471843,   -0.567222,   -0.106943,   0.520354], 
 [-0.0559538,   -1.0415,   -0.555807,   0.426276]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 4] strides = [4, 1] data = 
[[0.649067,   1.91099,   -1.84495,   -0.288202], 
 [0.632168,   2.18516,   -0.796985,   1.26018], 
 [0.419502,   -0.792007,   0.0416733,   1.19927], 
 [0.0978141,   -0.951834,   -0.094805,   1.86646]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 5] strides = [5, 1] data = 
[[0.296546,   0.108577,   0.108446,   0.250679,   0.235753], 
 [0.221851,   0.126921,   0.0812285,   0.211973,   0.358027], 
 [0.146912,   0.185544,   0.263232,   0.12015,   0.284162], 
 [0.12842,   0.201212,   0.198259,   0.116907,   0.355202]]
  ----- exit status 0

The last matrix is the answer: row i is the model's probability for each of the five tokens coming after token i. Every row is positive and sums to 1. They range from about 0.08 to 0.36, scattered around the uniform 0.2 that an untrained model's would be; the checker confirms that they equal the Python reference's to six digits.

The future cannot influence the past

The causal mask is meant to guarantee that row i depends only on tokens 0..i. Example 4 changes only the last token, from 2 to 4:

################ mgc run 04_last_token_changed.mg
  | (72 lines: the definitions are transformer.mg.defs, followed by seeded weights, the token ids and the calls; see the file)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 4] strides = [4, 1] data = 
[[0.9,   0.34,   0.4,   0.45], 
 [0.3915,   0.2803,   0.98,   1.07], 
 [0.9393,   -1.0661,   -0.36,   1.0698], 
 [0.6711,   -0.7,   0.56,   1.5596]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 4] strides = [4, 1] data = 
[[0.353842,   0.387514,   -0.308814,   -0.435457], 
 [-0.813699,   1.01803,   0.635739,   1.14997], 
 [0.471843,   -0.567222,   -0.106943,   0.520354], 
 [0.569287,   0.423964,   0.997841,   1.45058]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 4] strides = [4, 1] data = 
[[0.649067,   1.91099,   -1.84495,   -0.288202], 
 [0.632168,   2.18516,   -0.796985,   1.26018], 
 [0.419502,   -0.792007,   0.0416733,   1.19927], 
 [0.952626,   0.65142,   0.765684,   3.64543]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 5] strides = [5, 1] data = 
[[0.296546,   0.108577,   0.108446,   0.250679,   0.235753], 
 [0.221851,   0.126921,   0.0812285,   0.211973,   0.358027], 
 [0.146912,   0.185544,   0.263232,   0.12015,   0.284162], 
 [0.119444,   0.193318,   0.137116,   0.117316,   0.432807]]
  ----- exit status 0

Compare it with the previous output: rows 0, 1 and 2 of all four matrices are identical, digit for digit, in the embedded input, in both blocks and in the probabilities, and only row 3 changed. That is causality, shown in the output, and it is a checked property, not a coincidence of the example (see Tests). Without the mask, changing the last token would change every row.

Mistakes the compiler reports

################ mgc run errors/e1_wrong_number_of_tokens.mg
  | # The model was written for 4 tokens. Three tokens do not fit the 4x4 parameter of layer_norm, and the compiler says so.
  | def layer_norm(x: tensor[4x4], g: tensor[1x4], b: tensor[1x4]) = (x - row_mean(x)) / sqrt(row_mean((x - row_mean(x)) * (x - row_mean(x))) + 0.00001) * g + b
  | print layer_norm([[1, 2, 3, 4], [4, 3, 2, 1], [0, 1, 0, 1]], [[1, 1, 1, 1]], [[0, 0, 0, 0]])
  ----- output
errors/e1_wrong_number_of_tokens.mg:3: error: argument shape 3x4 does not fit parameter 4x4
  ----- exit status 1

################ mgc run errors/e2_missing_argument.mg
  | # attention_sublayer takes twelve arguments. Eleven is an arity error, reported before any code exists.
  | def softmax(s: tensor[4x4]) = exp(s - row_max(s)) / row_sum(exp(s - row_max(s)))
  | def layer_norm(x: tensor[4x4], g: tensor[1x4], b: tensor[1x4]) = (x - row_mean(x)) / sqrt(row_mean((x - row_mean(x)) * (x - row_mean(x))) + 0.00001) * g + b
  | def head(x: tensor[4x4], wq: tensor[4x2], wk: tensor[4x2], wv: tensor[4x2], mask: tensor[4x4]) = softmax((x @ wq) @ transpose(x @ wk) * 0.7071067811865476 + mask) @ (x @ wv)
  | def attention(x: tensor[4x4], mask: tensor[4x4], wq1: tensor[4x2], wk1: tensor[4x2], wv1: tensor[4x2], wo1: tensor[2x4], wq2: tensor[4x2], wk2: tensor[4x2], wv2: tensor[4x2], wo2: tensor[2x4]) = head(x, wq1, wk1, wv1, mask) @ 
  | def attention_sublayer(x: tensor[4x4], mask: tensor[4x4], g: tensor[1x4], bt: tensor[1x4], wq1: tensor[4x2], wk1: tensor[4x2], wv1: tensor[4x2], wo1: tensor[2x4], wq2: tensor[4x2], wk2: tensor[4x2], wv2: tensor[4x2], wo2: tens
  | let m = [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]
  | let h = [[1, 0], [0, 1], [1, 1], [0, 0]]
  | let g = [[1, 1, 1, 1]]
  | print attention_sublayer(m, m, g, g, h, h, h, [[1, 0, 0, 0], [0, 1, 0, 0]], h, h, h)
  ----- output
errors/e2_missing_argument.mg:10: error: attention_sublayer takes 12 arguments, got 11
  ----- exit status 1

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

################ mgc run errors/e4_residual_shapes_differ.mg
  | # A residual connection adds the sublayer's output to its input, so their shapes must be equal. A 4x2 head output cannot be added to a 4x4 input.
  | let x = [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]
  | let head_output = [[1, 0], [0, 1], [1, 1], [0, 0]]
  | print x + head_output
  ----- output
errors/e4_residual_shapes_differ.mg:4: error: cannot add shapes 4x4 and 4x2: dimensions 4 and 2 differ and neither is 1
  ----- exit status 1

In order: three tokens passed to a layer norm written for four; eleven arguments to a function that takes twelve (the shapes of all twelve are checked at the call); sqrt of a bare number; and a residual connection whose two sides have different shapes (a 4x2 head output added to a 4x4 input). Every one is caught when compiling, with the file and line. The model's shapes are types: a wrong-shaped weight cannot be connected.

Tests

Five new lit files in test/transformer30/ (the suite is now 122 tests):

Test What it checks
examples layer norm and feed-forward (worked out by hand) and the model's outputs, including that example 4's rows 0 to 2 equal example 3's
errors the four mistakes give the exact messages, with file and line
verifiers mg.sqrt given a result of a different shape is rejected by the dialect
ir layer normalization is a reduce, subtract, multiply, reduce, mg.sqrt, divide, multiply, add, in that order
against-reference check_transformer.py, below

check_transformer.py is the real test. Its reference (transformer_lib.py) is an independent implementation: plain Python lists, math.exp and math.sqrt, written from the definition of each step rather than translated from the Mountain Goat. It checks:

  1. Every stage against the reference for four weight sets and four token sequences (including all tokens equal): the embedded input, block 1, block 2 and the probabilities, to a relative 10-5 (mgc prints six digits), and that each probability row is positive and sums to 1.
  2. Stress: weights scaled up until the attention scores reach the hundreds and the vocabulary scores the thousands. The output must contain no nan or inf and must still equal the reference. A naive softmax overflows here.
  3. Causality: for each of the last three positions, change that token and require rows before it to be identical (the printed numbers equal, with no tolerance) in block 1, block 2 and the probabilities, and row j itself to change (otherwise a model that ignores its input would pass). Without the mask, changing the last token must change row 0, so the mask is shown to be what causes the first property.
  4. Permutation equivariance: with no positions and no mask, the output for a reordered input must be the original output reordered the same way, for three orders. With positions added, the same reordering must not be a reordering of the output (positions break the symmetry; a test that always passed would prove nothing).
  5. Layer norm: on random matrices with scale 1 and shift 0, rows have mean 0 and variance 1, and a constant row becomes exactly 0.
-- every stage against the reference; probabilities positive and summing to 1
  ok   weights 1, tokens [1, 3, 0, 2]: embedded equals the reference
  ok   weights 1, tokens [1, 3, 0, 2]: block1 equals the reference
  ok   weights 1, tokens [1, 3, 0, 2]: block2 equals the reference
  ok   weights 1, tokens [1, 3, 0, 2]: probs equals the reference
  ok   weights 1, tokens [1, 3, 0, 2]: each probability row is positive and sums to 1
  ok   weights 2, tokens [4, 4, 4, 4]: embedded equals the reference
  ok   weights 2, tokens [4, 4, 4, 4]: block1 equals the reference
  ok   weights 2, tokens [4, 4, 4, 4]: block2 equals the reference
  ok   weights 2, tokens [4, 4, 4, 4]: probs equals the reference
  ok   weights 2, tokens [4, 4, 4, 4]: each probability row is positive and sums to 1
  ok   weights 3, tokens [0, 1, 2, 3]: embedded equals the reference
  ok   weights 3, tokens [0, 1, 2, 3]: block1 equals the reference
  ok   weights 3, tokens [0, 1, 2, 3]: block2 equals the reference
  ok   weights 3, tokens [0, 1, 2, 3]: probs equals the reference
  ok   weights 3, tokens [0, 1, 2, 3]: each probability row is positive and sums to 1
  ok   weights 7, tokens [2, 0, 4, 1]: embedded equals the reference
  ok   weights 7, tokens [2, 0, 4, 1]: block1 equals the reference
  ok   weights 7, tokens [2, 0, 4, 1]: block2 equals the reference
  ok   weights 7, tokens [2, 0, 4, 1]: probs equals the reference
  ok   weights 7, tokens [2, 0, 4, 1]: each probability row is positive and sums to 1
-- stress: large weights make attention scores in the hundreds and vocabulary scores in the thousands, where a naive softmax gives nan
  ok   large weights: no nan and no inf anywhere in the output
  ok   large weights: block1 equals the reference
  ok   large weights: block2 equals the reference
  ok   large weights: probs equals the reference
-- causality: the causal mask stops position i from seeing positions after it
  ok   changing token 3: rows 0..2 of block 1, block 2 and the probabilities are identical
  ok   changing token 3: row 3 of the probabilities does change
  ok   changing token 2: rows 0..1 of block 1, block 2 and the probabilities are identical
  ok   changing token 2: row 2 of the probabilities does change
  ok   changing token 1: rows 0..0 of block 1, block 2 and the probabilities are identical
  ok   changing token 1: row 1 of the probabilities does change
  ok   without the mask, changing the last token DOES change row 0 (the mask is what blocks the future)
-- permutation equivariance (no positions, no mask): permuting the tokens permutes the output rows
  ok   order (1, 0, 2, 3): outputs are the original outputs in that order
  ok   order (3, 2, 1, 0): outputs are the original outputs in that order
  ok   order (2, 3, 0, 1): outputs are the original outputs in that order
  ok   with positions added, the same permutation is NOT a permutation of the outputs (positions break the symmetry)
-- layer norm: mean 0, variance 1, constant rows become 0
  ok   random 4x4 #1: rows have mean 0 and variance 1
  ok   random 4x4 #1: the constant row becomes exactly 0
  ok   random 4x4 #2: rows have mean 0 and variance 1
  ok   random 4x4 #2: the constant row becomes exactly 0
  ok   random 4x4 #3: rows have mean 0 and variance 1
  ok   random 4x4 #3: the constant row becomes exactly 0
all checks pass

Are the checks good enough? Run fifteen wrong models.

A test of a model should notice a wrong model. model_mutation.py writes fifteen versions of transformer.mg.defs, each with one deliberate mistake (the mask not applied, the scale missing, a wrong epsilon, a missing residual, no relu, values taken from the keys, and so on), and runs the checker on each:

NOTE: this script deliberately runs wrong models. 'caught' lines are EXPECTED: they show the checker can detect the mistake.
baseline (nothing broken): checker passes (expected)
the causal mask is not added to the attention scores: caught by 16 failing checks (the reference, causality)
the 1/sqrt(d) scale is missing: caught by 12 failing checks (the reference)
the layer-norm epsilon is 0.1 instead of 0.00001: caught by 18 failing checks (the reference, layer norm)
layer norm forgets its scale and shift: caught by 14 failing checks (the reference)
layer norm divides by the variance, not its square root: caught by 18 failing checks (the reference, layer norm)
layer norm does not subtract the mean from the numerator: caught by 21 failing checks (the reference, layer norm)
the attention sublayer has no residual connection: caught by 15 failing checks (the reference)
the feed-forward sublayer has no residual connection: caught by 15 failing checks (the reference)
the feed-forward network has no relu: caught by 15 failing checks (the reference)
a head's values come from the keys (wk instead of wv): caught by 15 failing checks (the reference)
the two heads use the same output matrix (wo1 twice): caught by 15 failing checks (the reference)
the keys are replaced by the queries (the scores are q times q-transposed): caught by 15 failing checks (the reference)
the feed-forward bias b1 is added after the relu: caught by 14 failing checks (the reference)
the attention softmax is the naive one (no row-maximum subtraction): caught by 4 failing checks (other, the reference)
the vocabulary softmax is the naive one (no row-maximum subtraction): caught by 2 failing checks (other, the reference)

All fifteen are caught. The last two (a naive softmax in the attention, and in the vocabulary) are the interesting ones: the first version of this checker did not catch them. With the seeded weights, the scores are small, and exp of a small number does not overflow, so a naive softmax gives the right answer and no comparison with the reference can tell the two apart. They are caught only by the stress check, which was added after the first run of this script showed the gap. A correct-looking test with friendly numbers cannot see a numerical-stability bug; this is Chapter 29's lesson again, now at the level of a whole model.

Compiler mutations

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

#!/bin/sh
# READ THIS FIRST: this script breaks the Chapter 30 code on purpose, one way at a time, and re-runs the Chapter 30 tests (test/transformer30).
#   "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=':: transformer30/'
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 "sqrt is compiled as relu by the front end"                          $F 's/{r} = mg.sqrt {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 "sqrt computes the identity (copies its input)"                      LowerToAffine.cpp 's/b.create<math::SqrtOp>(l, x)/x/'
try "sqrt computes exp instead"                                          LowerToAffine.cpp 's/b.create<math::SqrtOp>(l, x)/b.create<math::ExpOp>(l, x)/'
try "sqrt computes the square instead"                                   LowerToAffine.cpp 's/b.create<math::SqrtOp>(l, x)/b.create<arith::MulFOp>(l, x, x)/'
try "mg.sqrt's verifier accepts anything"                                MgDialect.cpp 's/return verifyElementwise(\*this, {getInput().getType(), getResult().getType()}, "mg.sqrt"); }/return mlir::success(); }/'
try "a sum starts from 1 instead of 0 (every row_sum and row_mean 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
sqrt is compiled as relu by the front end: caught by transformer30/ir.mlir transformer30/examples.mlir transformer30/against-reference.mlir 

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

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

All six are caught. The verifier mutation is caught only by the dedicated verifiers test (no program the front end writes can reach it, because the front end builds mg.sqrt with matching shapes), which is why that test exists. The last mutation, a sum that starts from 1 instead of 0, is not about sqrt at all: it breaks row_sum and row_mean, which layer norm and the softmax both lean on, and it shows that the transformer tests would also notice a regression in an operation from an earlier chapter.

The full suite:

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

Testing Time: 13.36s

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

Limits and what is not established

  • Not trained. The weights are seeded pseudo-random numbers. The probabilities are not language. Nothing here shows the model can learn, and no claim is made that it would learn well if trained.
  • Tiny and fixed. 4 tokens, width 4, 2 heads, 2 blocks, vocabulary 5. All sizes are static (a ? size cannot be stretched, Chapter 22), so a different sequence length means different definitions. Nothing was run at a realistic size, and nothing was timed.
  • No batching, and heads are hand-split. There is no batched matrix product, so the two heads are written out separately and their outputs combined with the split output matrix described above. A model with many heads would need either many such lines or batched matrix multiplication, which the language does not have.
  • Positions are given as numbers. The sinusoidal position matrix is computed by the Python generator and written into the program, because the language has no sin or cos. Learned position vectors would work unchanged.
  • No argmax or sampling. The program ends at probabilities. Choosing the most likely token or sampling one happens outside the language (the checker reads the printed matrix).
  • Pre-norm only, no dropout, no weight tying. Other variants (post-norm, a final-layer bias, tying the output matrix to the embedding) are small changes, not written.
  • The mask uses -1000000000, not minus infinity (Chapter 29), and the layer-norm epsilon is a literal 0.00001.
  • Six digits. Everything is compared at a relative 10-5, since mgc prints six significant digits. Exact equality was not possible or claimed (except for the causality property, where the printed digits are identical).
  • Equivalence of the head split to concatenating heads is argued (the matrix product distributes), and the model agrees with a reference that also uses the split. A reference using real concatenation was not written.
  • exp comes from the platform's C library and sqrt is an LLVM intrinsic (normally one hardware instruction); their last-digit accuracy was not measured here. GPU lowering of either is not covered.
  • Platform: x86-64 Linux, LLVM 18.1.3.

Reproducing

cd docs/part30/code && ./build.sh
python3 make_examples.py && ./run_examples.sh > examples_out.txt
./check_transformer.py > check_transformer_out.txt
./model_mutation.py > model_mutation_out.txt
./mutation.sh > mutation_out.txt
cd ../../part15/code && ./run_lit.sh          # 122 tests

Chapter summary

  • A transformer block is attention (mixes tokens) plus a feed-forward network (transforms each token), each wrapped in a layer norm and a residual connection. A model is an embedding, positions, a stack of blocks, a final layer norm and a softmax.
  • With sqrt added, all of that is a few lines of Mountain Goat, and the shapes of every weight are types the compiler checks.
  • The program agrees with an independent Python version at every stage, for several weights and token sequences, and has the properties it must: probabilities sum to 1, the future cannot influence the past, reordering tokens reorders outputs when positions are absent, layer-norm rows are normalized, and huge scores do not overflow.
  • Fifteen wrong models and six broken compilers are all caught; two of the wrong models were caught only after a stress case was added, which is a finding about the checker, reported as such.
  • The model is untrained, tiny and fixed in size. Training, batching and variable lengths are what would come next.

Self-check questions

Each answer is collapsed; try the question first.

  1. In one sentence each: what does attention do, what does the feed-forward network do, and which of them mixes information between tokens?

    Answer

    Attention lets each token gather a weighted average of information from tokens it may look at, so it mixes rows. The feed-forward network applies the same two-layer transformation to every token separately, so it does not mix rows. Only attention moves information between positions.

  2. What does a residual connection do, and what would happen to the "no residual" mutations' outputs?

    Answer

    It adds a sublayer's output to its own input (x + step(x)), so the sublayer only has to learn an adjustment and the original information is kept. Without it the output is only what the sublayer computed, which is a different matrix, so the comparison with the reference fails at the first block. (In a trained deep model a missing residual also makes training hard; that is not tested here since nothing is trained.)

  3. Why is the layer norm's epsilon needed? What does layer_norm produce for the constant row 10 10 10 10?

    Answer

    The row's variance is 0, so without an epsilon the division would be 0/0, which is nan. With the 0.00001, the numerator x − mean is exactly 0 and the denominator is sqrt(0.00001), so the result is exactly 0 (then the learned scale and shift apply). Example 1's second row shows 0 0 0 0.

  4. Why does the checker require rows before the changed token to be identical, not merely close?

    Answer

    With the causal mask the masked scores are -1e9, whose exponentials are exactly 0, so earlier rows are computed from exactly the same numbers through exactly the same operations and must be bit-for-bit equal. A tolerance would let a small leak of future information through (for instance a mask that was -5 instead of -1e9). Equality is a sharper test than closeness here.

  5. Why does the checker also require that changing token j changes row j, and that removing the mask changes row 0?

    Answer

    A model that ignored its input entirely would pass "earlier rows do not change" trivially, and so would a mask that has no effect. Requiring row j to change shows the model reads its input; requiring row 0 to change without the mask shows the mask, not something else, is what makes the earlier rows independent of later tokens.

  6. Why is permutation equivariance tested without positions and expected to fail with them?

    Answer

    Attention treats its input as a set: reordering the rows reorders the outputs the same way, because nothing in the computation depends on row order (the feed-forward and layer norm work row by row). Positions are added precisely to break that symmetry, so with positions the same reordering must not just reorder the output. Checking both directions shows the property is real and that the position matrix does its job.

  7. Why does the hand-split of the heads (h1 @ wo1 + h2 @ wo2) equal concatenating the heads and multiplying by one output matrix?

    Answer

    Concatenating two 4x2 matrices side by side gives a 4x4 matrix [h1 | h2]. Multiplying by a 4x4 matrix W whose top two rows are wo1 and bottom two rows are wo2 gives, for each output element, a sum over four terms: the first two come from h1 and wo1, the last two from h2 and wo2. That is exactly h1 @ wo1 + h2 @ wo2. The matrix product distributes over the split, so no approximation is involved.

  8. Two of the fifteen wrong models passed the first version of the checker. Which, why, and what fixed it?

    Answer

    The two with a naive softmax (no row-maximum subtraction), in the attention and in the vocabulary. With the seeded weights all scores are small, so exp does not overflow and the naive formula gives the same numbers as the stable one; no reference comparison could tell them apart. The stress check scales the weights until the scores reach the hundreds and thousands, where the naive softmax gives nan and the stable one does not.

  9. What would be needed to train this model, and which parts of it does the language lack?

    Answer

    Gradients of the loss with respect to every weight, computed by differentiating through each operation (matrix product, softmax, layer norm, relu, residuals), plus an update rule such as gradient descent. Chapter 26 did this by hand for a linear model. Mountain Goat has no automatic differentiation, no loss function (cross-entropy needs a logarithm, which the language lacks), and no way to select one probability from a row (an index or argmax operation).

  10. What is the difference between this model's embedding lookup and a lookup table in a conventional language?

    Answer

    Here the lookup is a matrix product: a one-hot row (a single 1 among zeros) times the embedding matrix returns exactly one row of it. It does the work of an index, with V times the arithmetic, but it needs no indexing operation (the language has none) and it is differentiable, which is why real frameworks use the same idea for the gradient of an embedding.