A max-of-list transformer, read off its weights

Setup

The task is to return the largest of five digits. Each digit is 0 to 9 and repeats are allowed, so there are exactly 10⁵ = 100,000 possible inputs. Every number in this report is an exhaustive count or average over all of them, not a sample, so no confidence interval appears anywhere.

The model reads a fixed 12-token sequence. A list of 3, 7, 2, 5, 1 is presented as

position   0     1    2     3    4     5    6     7    8     9    10    11
token    [BOS]   3  [SEP]   7  [SEP]   2  [SEP]   5  [SEP]   1  [ANS]   7

The five digits sit at the odd positions 1, 3, 5, 7 and 9; we call those five positions the slots. Separators sit between them, a beginning-of-sequence token opens the sequence, and the [ANS] token at position 10 is where the model emits its prediction: the logits at that position should put the maximum, 7, on top. Position 11 holds that answer token during training, and the logits there should predict end-of-sequence. Those two positions, 10 and 11, are the only ones the model was ever trained at, and they are the only two this report analyses.

The vocabulary has 14 tokens: the ten digits plus [BOS], [SEP], [ANS] and [EOS]. The architecture is as small as a transformer gets: a token embedding and a learned position embedding, both into 64 dimensions, then one attention layer with 4 heads of 16 dimensions each, then an untied unembedding back to the 14 tokens. There is no MLP, no LayerNorm and no bias anywhere, and no second layer. The weights are the released checkpoint of the Bau Lab April 2026 interpretability puzzle 1a, trained to 100% accuracy on this task; this report calls it the puzzle-1a model.

The algorithm in plain words

At the [ANS] position, the eleven tokens before it are the head's keys and the [ANS] token itself is the query. Each head scores every key against that query, softmaxes the scores into attention weights, and mixes the attended tokens into its output. Because the query token is the same on every input, each head's scores reduce to a lookup table: one number per key token at each position. We call that the head's score table.

Six of the eleven keys are the same on every input — the [BOS], four [SEP] and [ANS] tokens. We call those the constant keys. They act as an attention sink: a head can send its attention there instead of at the digits, and then its output is the same no matter what the list holds. A head that does that is parked. Whatever a head attends to, it writes a vector into the logits: the attended token's value vector pushed through the head's output matrix and the unembedding.

What makes this model simple is that each head's write turns out to be one fixed direction in logit space multiplied by a single number, which we call the number the head passes on. That is a measured property of the trained weights, not an assumption: both of each head's matrices are within a fraction of a percent of rank one, shown in Figure 3 below. Four heads therefore hand the output layer four numbers, and since the unembedding is linear, each digit's logit is a straight line in any one of them. The answer is whichever line is highest, so the number that varies most decides it: sweep that number and the winning digit changes hands at a few crossing points. The sequence of winners along it — the highest line at each value — is the upper envelope. When other heads change what they contribute, every line shifts up or down by a constant and the crossing points move; each such set of shifted lines is a decoding regime.

With those words in place, the whole model is: head 3 scores digits monotonically in value, so the number it passes on rises with the list maximum; ten logit lines decode that number by taking whichever is highest; heads 0 and 2 stay parked until a digit of 7 or more appears and then shift the lines; head 1 is parked always and contributes a constant.

The ten bands never touch. The closest pair are the ones for maxima 9 and 8, which are 0.968 and 1.253 wide with 28.96 of empty space between them, more than twenty times either band's own width. The units of that number are arbitrary, so the ratio is what matters rather than the 28.96 itself.

The algorithm in full

With the constants read off the weight matrices, the whole answer-position computation is short enough to run by hand. In the listing, score[h][v] is head h's attention score for a key holding the digit value v, alpha[h] is how strongly head h writes the digit it attends, bias[v] is the input-independent term the [ANS] token contributes to digit v's logit on its own, and the running logit list is the ten digit logits:

# The whole model at the ANS position, from the weights alone. The softmax runs over the FIVE
# DIGIT SLOTS ONLY: the constant BOS / SEP / ANS keys are not in the denominator. score[h][v],
# alpha[h] and bias[v] are the constants listed in symbolic_constants.
def max_of_list(digits):                       # digits: the five list values
    logit = list(bias)                         # direct path U (e_ANS + p_10)
    for h in heads:                            # the three answer-carrying heads
        s = [score[h][d] for d in digits]      # QK: one number per digit value
        w = softmax(s)                         # over the five digit slots, nothing else
        for j, d in enumerate(digits):
            logit[d] += alpha[h] * w[j]        # OV: copy the attended digit
    return argmax(logit)

That listing answers 100,000/100,000 inputs correctly, and still 100,000/100,000 with every constant rounded to the nearest integer; an independent reimplementation from the printed constants alone reproduces both counts. It is the last step of a ladder: a sequence of six progressively simpler descriptions of the same weights, starting from exact algebra and ending here, each one scored on all 100,000 inputs. The certification section lists the six and what each costs.

The softmax there runs over the five digit slots and nothing else, and a variant that puts the constant keys back into the denominator answers 99,999/100,000 instead. Its one failure is [0, 0, 0, 0, 0], where it answers 7: that variant keeps the attention mass the constant keys absorb but discards the logit vectors they write, which is what carries the answer on the all-zeros list; both variants are specified exactly in the certification section.

The algorithm, step by step, in the weights

The pseudocode above has five stages. Each stage is implemented by specific weight matrices, produces an activation you can observe on any input, and is shown by one figure on this page. The table maps them; the four worked examples below trace the same columns on concrete inputs.

The Compress row lists four heads while the pseudocode loops over three. Head 1's write is a constant, so the ladder drops it at its first step, and that costs 3 answers out of 100,000. Each head's role differs: head 3 carries the ranking, heads 0 and 2 switch the regime, and head 1 contributes a fixed offset. Removing any one of their writes hurts, and the amounts are in the certification section.

Four inputs traced end to end

Each table below is one input pushed through the exact algebra: what each head attends to, where the rest of its attention mass sits, and the single number it passes on. The head labels are the roles the rest of this page establishes: head 3 ranks digits by value, heads 0 and 2 detect digits of 7 or more, and head 1 abstains by staying parked on the [ANS] key and writing the same constant every time.

A list with a small maximum: [3, 1, 0, 5, 2]

Head 3 puts almost all of its attention on the slot holding the maximum and hands on 504.03. The detector heads are parked on constant keys (1.0000 and 0.9994 of their mass), so Figure 1 is the picture that applies: that number falls in the stretch 394.0 to 633.5 where digit 5 wins. The top three digit logits are digit 5 at 153.45, digit 6 at 138.52, digit 4 at 138.11, so the answer is 5 with a margin of 14.93 logits.

A list where the detectors switch on: [4, 2, 8, 1, 6]

This is the case that shows why the detector heads exist. Head 3's number is 370.96, which sits inside the stretch 202.6 to 393.9 that belongs to digit 4 when the detectors are parked — so by the picture in Figure 1 this input would decode to 4, which is wrong. But the maximum is 8, so, per the map's regime-switch row, both detectors release their parked mass onto the slot holding it and the lines shift: head 0's number moves from its parked -140.19 to -107.24 and head 2's from -176.11 to 110.87. On the shifted lines digit 8 wins by 17.84 logits. The model's own top three are digit 8 at 142.23, digit 7 at 125.04, digit 6 at 123.57, answering 8 with margin 17.19.

A list whose maximum repeats: [7, 7, 3, 7, 0]

The maximum occupies 3 slots, and head 3 splits its attention between them almost evenly — 0.3361 · 0.3332 · 0.3307 — summing to essentially all of its mass. Because every tied slot writes the same value, splitting the mass changes nothing: the number is 178.15, the top three logits are digit 7 at 86.96, digit 8 at 70.32, digit 5 at 69.59, and the answer is 7 with margin 16.64. This is the general case, not a lucky one: accuracy on all 23,335 duplicate-maximum lists is 1.0000.

The one list decoded a different way: [0, 0, 0, 0, 0]

On [0, 0, 0, 0, 0], every digit sits at the bottom of head 3's score range, so head 3 keeps 0.9996 of its mass on the constant keys and the answer comes through what those keys write rather than through the slots. The model still answers 0 with margin 12.40, and this is the one input that any simplification discarding the constant keys' writes gets wrong.

The evidence behind each step

This section measures each row of the map, in the order the map lists them, except that the rank-one result comes first because the score table, the read direction and the number a head passes on all depend on it.

Each head reads one direction and writes one direction

Two composed matrices carry a head's whole contribution: the bilinear form A_h = W_Q,h^T W_K,h, whose value on a (query, key) pair is the attention score, and the output map M_h = U W_O,h W_V,h, which turns an attended token into the logit vector the head writes. A_h could have rank 16 and M_h could have rank 14. Measured, both are effectively rank one: the top singular value carries 99.70% to 99.85% of the squared Frobenius norm for the four A_h, and 99.29% to 99.999% for the four M_h. That is what licenses the whole account: a rank-one A_h means one direction on each side generates every score a head computes, and a rank-one M_h means the head writes one fixed direction scaled by one number instead of copying the attended token. Replacing every M_h by its rank-one part leaves the answer unchanged on 100,000/100,000 inputs, with a mean answer KL — the KL of the simplification's answer distribution from the model's, at the answer position — of 4.75e-09 nats; doing the same to the A_h leaves 100,000/100,000.

The four heads are also not reading four different things. Each head has one value-read direction, the key-side vector of its rank-one M_h, and one key-read direction, the key-side vector of its rank-one A_h. The absolute cosine between the value-read directions of heads 0 and 3 is 0.955, between heads 2 and 3 0.971, and even head 1 sits at 0.812. On the score side, heads 0, 2 and 3 read nearly the same direction (0.899 and 0.941 against head 3) while head 1 does not (0.201). What separates the heads is what they write: the cosine between head 1's and head 3's output directions is 0.014.

The embeddings hold no value axis

Head 3's score table is monotone in digit value, so it is tempting to read that ordering back into the token embeddings. The embeddings do not have it. The ten digit embeddings occupy essentially two dimensions — 81.52% of their variance on the leading principal direction, 96.27% in the leading two, giving an effective dimension (participation ratio) of 1.456 — and neither direction orders value: the leading coordinate is not monotone in digit value (rank correlation 0.7333) and the second runs the other way (-0.6485). The embeddings are not even comparable in scale, with norms from 0.4112 at digit 3 to 3.3181 at digit 6. Fitting a least-squares "value axis" to the ten points returns R² 1.0, which is degeneracy rather than evidence: ten points in 64 dimensions always fit exactly. The evidence about the ordering is the principal components' own ordering and the score tables, not that fit.

The query-key matrices create the ordering

What is monotone is the score profile after projecting the digit embeddings onto a head's key-read direction, defined with the rank-one result above, and that direction is not the embeddings' leading direction. Head 3's read direction puts 0.771 of its energy on the second principal direction against 0.105 on the first, mixing the two rather than following either (|cosine| 0.8781 with the second, 0.3239 with the first), and the profile it produces over the ten digit embeddings is strictly monotone in value (rank correlation 1.0). The ordering is created at the query-key stage over a non-monotone two-dimensional embedding, not read off a value axis that was already there.

The softmax concentrates on the maximum

Head 3 puts 0.9998 of its mass on the five slots and 0.9953 of that on slots holding the maximum, and its strongest slot is a maximum on 1.0000 of inputs. Head 1 puts 1.000 on the [ANS] key and 2.7e-12 on the digits, which is what parking looks like in the activations. Heads 0 and 2 sit in between at 0.4496 and 0.8251, averaged over inputs whose maxima are mostly below their threshold.

Four numbers reach the unembedding

Rank-one writes mean the whole attention layer's contribution is bias + Σ_h g_h · out_h, with one number g_h per head — the number the worked examples report, called the head's scalar in the code and on the figure axes — and one fixed direction out_h. Reconstructing the logits that way is accurate to 1.21 logits at worst and 0.4816 in root-mean-square, against answer margins that never fall below 8.881. The root-mean-square is normalised over every compared logit value — 1,400,000 of them, the 100,000 inputs times the 14 vocabulary logits — not per input and not per head. Splitting a matrix into a number and a direction fixes their signs only up to a joint flip, so each head's number carries an arbitrary sign: what is meaningful is the product, the magnitudes, and the ordering up to a per-head flip. The same fact shows up in the logits directly: four numbers plus a bias determine all ten, and measured over the whole input space the ten digit logits use 2 dimensions for 99% of their variance and 3 for 99.99%.

Head 3's number is enough on its own. Sorting the ten possible maxima by the band its number occupies gives bands that never overlap, with a smallest gap of 28.96, so a threshold rule on that single number recovers the maximum on every one of the 100,000 inputs. No other head's number does that: head 0's bands for different maxima overlap, so no threshold on it could work. The order those bands appear in is not the numeric order: head 3's number ranks the maxima [0, 1, 2, 3, 7, 4, 9, 8, 5, 6], and the three values it places out of sequence are exactly the three at which the detectors switch on. A natural guess is that the detectors exist to add resolution among those crowded top values; the head-removal numbers in the certification section refute that and point at re-biasing the decode instead.

The other trained position writes EOS

Position 11 holds the answer token and predicts end-of-sequence; it is the only other position the model was ever trained at. The same algebra applies with the query token ranging over the ten digits, and the closed form predicts [EOS] on 100,000/100,000 inputs with margins from 13.64 to a mean of 32.21 logits. The score table there is a single outer product of a query-side and a key-side vector (99.9977% of its squared norm for head 3, 100.0000% for head 1), so the query token scales one fixed key profile rather than reshaping it. That profile ranks [BOS] and the separators above the slots, the opposite of the answer position: head 3 puts 0.934 of its attention on [BOS] and 0.0139 on the digits. Those constant keys are what write [EOS], and head 0's write is the load-bearing one — removing it leaves 69,911 inputs still predicting [EOS], while removing head 3's leaves 100,000. The direct path is nearly indifferent: on its own it predicts [EOS] on 83,193 inputs with a mean margin of 0.50.

Every weight matrix, and what it does

The same content as the map above, indexed by matrix rather than by step, and extended to the direct path and position 11.

Certification

The account above is exact algebra, not a fit, and it is checked before it is interpreted. With no MLP, no LayerNorm and no biases, the logits at position p are a direct path plus one term per head and per visible key:

In float64 those tables are compared against the puzzle repository's own torch module at all 12 positions on a 2,048-input battery. That battery is structured rather than random: every all-equal list, then every digit value placed as the unique maximum at each of the five slots, then the same with the maximum tied across two slots, and finally a fixed pseudo-random sample to fill it out — cases where a wrong head order or an off-by-one in the positions would show up instead of averaging away. The largest deviation at the answer position is 2.84e-13 against a logit scale of 362.9; taken over all 12 positions, the worst case sits at 0.00496 of the tolerance allowed. That tolerance is the rounding this arithmetic can accumulate: the closed form sums four heads times eleven keys of 64-dimensional dot products, so it is set at 1,024 float64 steps at the logit scale in play. Against the released float32 model on the complete input space the agreement is 3.20 ulps — units in the last place, the spacing between adjacent float32 values at that magnitude — at the answer position, and 5.26 ulps at position 11. The model answers 100,000/100,000 lists correctly, recovering the checkpoint's own reported test accuracy of 1.0 as an exhaustive count.

The cost of each simplification

The ladder walks from the exact algebra to the pseudocode, one simplification at a time. Its steps, which the operational definitions below refer to by number:

1. r0, the exact reconstruction: the tables above, nothing dropped.
2. r1, drop what the abstaining head writes.
3. r2, make the score tables position-independent, so a key's score depends on its token and not on which slot it sits in.
4. r3, keep only the five slot keys, dropping the constant [BOS], [SEP] and [ANS] keys.
5. r4, replace each head's write with a plain copy of the attended digit, one strength per head.
6. r5, the symbolic pseudocode printed earlier on this page.

Each step is scored on the complete input space, cumulatively and again on its own so its individual cost is visible. Every step is a transformation of the closed-form tables, never a weight edit or a retrained model, so a step's cost is a statement about the account of the weights. Two further descriptions are scored alongside the ladder and appear in the next figure: the symbolic pseudocode from earlier, and a faithful description that keeps each head's actual rank-one write rather than idealising it to a copy, 174 constants in all, spelled out after the figure. Before running the ladder we fixed the acceptance floor the experiment would be judged against: the final symbolic description had to reproduce the model's answer on at least 99.9% of all inputs, with any failure explained rather than tolerated. That is what "pre-registered" means for the floor drawn on the next figure.

The one step that collapses is informative rather than fatal. Dropping the constant keys removes two things at once: the attention mass they absorb and the logit vectors they write. Separating those shows the writes are what matter. Keeping the keys in the softmax but zeroing what they write loses 50,018 answers; diagonalising the slot keys' writes while keeping the constant keys' writes loses 13,682. Doing both together loses 1 answer, and 0 once the constant keys leave the softmax as well. The off-diagonal structure of the slot writes and the constant keys' writes are only jointly removable, which is a fact about this model that a single ablation would have mislabelled as either structure being load-bearing on its own.

# Rung r5s: the same algorithm with the constant BOS/SEP/ANS keys put back into the softmax
# denominator as ONE extra key per head, scored sink[h] = log-sum-exp of their (input-independent)
# scores, and writing nothing. Constants in symbolic_constants_with_sink.
def max_of_list(digits):
    logit = list(bias)
    for h in heads:
        s = [score[h][d] for d in digits]
        Z = exp(sink[h]) + sum(exp(x) for x in s)   # the constant keys take mass, write nothing
        for j, d in enumerate(digits):
            logit[d] += alpha[h] * exp(s[j]) / Z
    return argmax(logit)

Three descriptions, one model

This page has described the same weights three times over, at three lengths. Being faithful to the mechanism and being brief trade off directly against each other, because the weights genuinely contain a little structure that no short description holds: each head's matrices are close to rank one but not exactly rank one, and the writes have real off-diagonal parts. Anything shorter than the closed form must therefore discard something the model actually does. The table says what each level discards and what it still matches.

The middle row is the one the ladder figure plots as the faithful description: per head one score number per key class, one value number per key class and one output direction, plus the direct path. The gap between that row and the last one is what the bullet on behavioural equivalence under Limitations records: reproducing every answer on a complete input space does not pin the mechanism. Written out, the faithful description is:

# The model's answer-position computation, faithful to the weights (not just to the answer).
# Constants: digit_score[h][v], digit_value[h][v], struct_score[h][c], struct_value[h][c],
# out[h][.] (a unit vector over the 14 logits), bias[.] — all read off the weight matrices.
def answer_logits(digits):                          # digits: the five list values
    logit = list(bias)                              # U (e_ANS + p_10)
    for h in heads:
        s = [struct_score[h][class_of(j)] if j is structural else digit_score[h][digits[j]]
             for j in range(11)]                    # QK: rank one, so one number per key
        a = softmax(s)                              # attention over all eleven keys
        g = sum(a[j] * (struct_value[h][class_of(j)] if j is structural
                        else digit_value[h][digits[j]]) for j in range(11))
        logit = [logit[v] + g * out[h][v] for v in range(14)]   # OV: rank one, one direction
    return logit

Margins, ties and head removal

The model is never close to wrong. The answer margin — the correct digit's logit minus the best other digit's — has a minimum of 8.881 logits over all 100,000 inputs, a mean of 17.131, and a first percentile of 13.27; 0 inputs have a non-positive margin. That slack is why so many idealisations survive: they perturb the logits by less than the gap they would have to close. On the duplicate-maximum lists the smallest margin is 8.881 logits against 10.343 on the unique-maximum lists, and head 3's summed mass on maximum-holding slots stays high as the maximum repeats (0.9947 when it appears once, 0.8749 on the ten all-equal lists).

To measure what a head contributes, we resect it: zero what it writes while leaving its attention untouched, so only its contribution to the logits is removed.

This is where the natural guess about the detector heads fails. If they existed to add resolution among the crowded top values 7, 8 and 9, removing their writes would cost margin there and little elsewhere. Instead the cost is spread across the range, and accuracy falls to 0.3631 for head 0 and 0.5147 for head 2 — both worse than removing head 3's write leaves the ranking (0.4169), because each removal also deletes a large constant write. What the detectors do instead is what the second worked example shows: they re-bias the decode above digit 6.

The parameter decomposition agrees

A previous experiment in this thread, exp_01kz42cpavf9mrw4s3f6j02qw8, decomposed the same weights into rank-one parameter subcomponents and found eleven that its causal-importance test marks as necessary, which that method calls "alive". This section reads that run's selected checkpoint (runs/runs/p-59c0a394 at step 500,000 under that experiment's artifact root, with its alive-set labels from results/analysis_selected_imp6e-3/mechanism_map.json), so reproducing it needs those artifacts; nothing else on this page does. Those subcomponents reconstruct the target's own fused weight matrices to within 0.0002 relative error, a check that the two analyses are describing the same matrices.

Two things that decomposition found now have a weight-level reason. One value subcomponent could carry the attended digit for every head at once because the four heads read nearly the same value direction. And the key side splits by head rather than by digit range: key subcomponent 55 matches head 3's score direction at 0.793 but head 1's at 0.051, while subcomponent 59 matches head 1's at 0.628.

Limitations

• This is one 18,944-parameter attention-only model on a task with 100,000 possible inputs. What transfers is the method, and the method needs both an exact algebraic reconstruction and an input space small enough to enumerate.
• The rank-one description of each head's write is an approximation, not an identity. It leaves a 1.21-logit worst-case error, small against the 8.881 minimum margin but not zero, so "four numbers" is exact only up to that residual. The exact statement is the full closed form.
• Behavioural equivalence on a complete input space does not pin a mechanism. The symbolic pseudocode reproduces every answer with diagonal copy matrices the weights do not contain, which is why the faithful rank-one description is reported beside it; Three descriptions, one model sets out what each of the three levels discards.
• The line picture is exact only in the regime where the detector heads hold still. For maxima 7, 8 and 9 the lines are drawn with the other heads at their regime mean; the exhaustive claims there are the disjointness of head 3's bands and the closed form's accuracy. Inside the parked regime, holding them fixed is worth at most the within-regime spread of head 2's number times its largest per-digit logit coefficient, 0.326 × 0.393 = 0.128 logits of misplacement, against 9.37 logits of margin in the tightest case.
• The comparison with the parameter decomposition is a projection of one selected run, not a re-derivation.

External verification and corrections

An external verifier rebuilt the model from the released weights and re-ran the claims on the complete input space. The gates, the ladder counts, the margins and the worst-case rank-one residual reproduced. Four things did not, and all four are recorded here rather than smoothed over.

• The published pseudocode listing showed a log-sum-exp sink term while the cited count came from the variant without one. Both variants and both counts are now stated: 100,000/100,000 without the sink term and 99,999/100,000 with it.
• Three ablation counts could not be reproduced from prose alone. Each now carries its operational definition and the function that computes it, and the two readings of "deleting the constant keys" are reported separately (0.47182 cumulative, 0.46949 isolated) — the second is what the verifier measured.
• The root-mean-square logit error of the rank-one reconstruction was normalised by inputs times heads instead of inputs times vocabulary entries. Corrected it is 0.4816; the worst-case 1.21 logits was unaffected, and no other number moved.
• Two conventions were undeclared: the sign of each head's number (a factorisation convention) and which statistic Figure 6 plots (means, with bars for the full range).