How a one-head attention-only transformer computes the maximum

Epistemic status: this is a solo project completed within a weekend. The max circuit is read directly off the weights and mostly replicates known results (Gross et al., 2024), so I’d defend it above 95%. Every result is aggregated over 5 seeds unless stated otherwise; single-model figures use the canonical seed (123).

TL;DR

I trained a one-layer, one-head, attention-only transformer (d_model=32, vocabulary = [0, 63]) to output the maximum of the input sequence of integers, and then reverse-engineered the model from its weights.

• The max circuit is “rank keys by size, copy the winner”. The QK circuit scores keys monotonically by their token value, while the OV circuit simply copies. Singular value decomposition shows that the value ordering lives almost entirely in the first principal component of the embedding, which is also what drives attention. A direct logit decomposition shows that ~96–99% of the correct logit (depending on sequence length) comes from the position holding the max. This is mostly a replication of Gross et al. (2024); see related work.
• Min is the same circuit with one sign flipped: the QK circuit’s monotonicity reverses, while the OV copying circuit is untouched. So the head is modular: “what do I do with the token I found” is fully reusable, and switching max↔min only touches “which token do I look for.”

Code and weights: Will be added here.

Why bother with a model this small

Fair question. The model is a 32-dimensional network doing max(), which is a one-liner in any language.

First, I wanted a mech interp project I could actually finish with limited time and compute.

Second, models this small let you check your answer. I can compute the QK and OV circuits from the weights, predict what the model should do, and then test that prediction against a decomposition of the actual output logits. Guessing wrong and finding out takes an afternoon.

Setup

• Task. Vocabulary of 64 tokens, which are the integers 0 to 63. Input is a sequence of L tokens; the target is the maximum of that sequence, predicted at the final position only. Accuracy is exact match on the argmax logit at that position (min is the same with the target flipped).
• Data. At lengths 2 and 3 I enumerate the whole product space of 64^L sequences, shuffle, and split 70/30, giving 2,867 training and 1,229 held-out sequences at length 2 and 183,500 and 78,644 at length 3. Length 6 has 68.7 billion possible sequences, so there I sample instead: training batches are drawn uniformly at random from one stream, and held-out accuracy is measured on 2,560 freshly sampled sequences from an independently seeded stream, which makes overlap between the two negligible. Tokens are i.i.d. uniform over 0 to 63, and duplicate values within a sequence occur at the natural rate.
• Model. d_model=32, d_head=32, 1 layer, 1 head, attention-only (no MLP). With d_head=d_model the head reads and writes the full residual stream, so there’s no projection into a smaller subspace anywhere in the circuit.
• Training. Adam at lr 1e-3 with default betas (0.9, 0.999), batch size 128, 10 batches per epoch. Length 2 trains for 500 epochs (5,000 gradient steps); lengths 3 and 6 for 1,000 (10,000 steps). Canonical seed 123 for the figures; tables aggregate over seeds [123, 0, 1, 2, 3].
• Test accuracy on max. 0.980 ± 0.003 at length 2, 1.000 at length 3, 1.000 at length 6.

Definitions of the two core metrics

I use two metrics, defined as follows:

Monotonicity score. Take the pre-softmax QK score for every (query value, key value) pair, computed as $(W_E + W_{pos}) W_Q W_K^\top (W_E + W_{pos})^\top$. There are no sequences, no softmax, and no competition between keys, so this is a pure function of the weights. Average over all 64 query values and over the key positions in question, then rank-correlate the per-key-value score with the key value. +1 means the circuit ranks bigger tokens strictly higher, 0 means it is indifferent to value, −1 means it prefers small tokens. This metric is fully disconnected from the data itself, so it describes the circuit.

Copying score. The fraction of the full OV circuit’s eigenvalue mass that is positive and real, as proposed in the transformer circuits framework. Near 1 means attending to token v mostly raises the logit for v. It is a heuristic. Here it agrees with both the OV heatmap and the logit decomposition, which seemed convincing enough to trust it.

Related work

The max-of-K task on a one-layer attention-only transformer is a standard toy mech interp problem, suggested in Neel Nanda’s 200 Concrete Open Problems (which is was inspired me to work on this problem) and was also studied in depth by Gross et al. (2024), who trained 151 seeds on it to derive formal accuracy lower bounds (LW post). Their setup is essentially the same as mine.

This write-up is largely a replication for my own understanding:

• They find the same QK monotonicity and OV copying.
• They show position doesn’t matter for the trained task by ablating the positional part of the QK circuit.
• They show attention grows roughly linearly with the token gap, which is why near-ties are the hard case. My shared-direction result is a rephrasing of their QK SVD, moved to the embedding.

The one thing I add is checking the max<>min mirror explicitly, and locating the “size axis” in the token embedding (rather than in the QK SVD).

The max mechanism, settled by one decomposition

A single attention head splits into two largely separate computations. The QK circuit decides where to attend, scoring each (query token, key token) pair. The OV circuit decides what to do with whatever got attended to, meaning how it moves the output logits.

The hypothesis: QK attends to the biggest visible token regardless of the query, and OV copies, so attending to token v raises the logit for predicting v. Compose them and you get max.

To test the whole mechanism at once, I decompose the logits into per-position contributions. With one attention-only layer and no layernorm, the final-position logits are a sum over source positions:

where $\alpha_i$ is the attention paid by the final position to position $i$ and $W_{OV} = W_V W_O$. So position $i$ contributes

to the logit of token $t$. Writing $i^\star$ for the position holding the maximum and evaluating at the correct answer $t = t_{\max}$, the question is how $c_{i^\star}(t_{\max})$ compares to $\sum_{i \neq i^\star} c_i(t_{\max})$. The numbers below cover the OV sum only: the direct path and the unembed bias sit outside it and aren’t counted. When the max appears more than once, $i^\star$ is its first occurrence, so the duplicate copies fall into the “other positions” column.

(Canonical seed 123, averaged over 4,096 sampled test sequences per model.)

On essentially every sequence, the position holding the maximum is the largest single contributor to the correct logit, and it supplies between 20× and 140× more of that logit than every other position combined. That fraction shrinks a little as sequences get longer (99% of the OV logit at length 2, ~96% at length 6), because more “other” positions each leak a bit.

QK attention is monotone in token value

Read this plot column by column: move right (higher key value) and the column gets brighter, whatever row (query value) you’re on. In addition to the heatmap, I also computed the monotonicity score for each query position, which confirms the findings from the heatmap:

At the trained (final) position it is essentially perfect everywhere. NOTE: For the length-6 model, the interior positions are all over the place and some are negative — but these are exactly the positions the trained task never reads from at the final step, and the per-seed spread there is large (std up to ~0.7), so treat the individual interior numbers as noise, not signal.

OV copies

The copying score comes out at 0.999 at length 2 and around 0.95 at lengths 3 and 6.

Both circuits read the same direction

So far I treat “QK attends to size” and “OV copies” as two independent facts. However, the natural thing for the embedding to do should be to lay value out along a single “size” axis that both circuits use: QK reads it to rank keys, OV maps it back onto the ordered unembedding. So I ran an SVD of the centered token embeddings.

• Value ordering lives almost entirely in the top principal component: Spearman(value, PC1 coordinate) = 0.98.
• The same axis drives attention: Spearman(PC1 coordinate, attention received) = 0.995.
• Honesty check: PC1 is only about 35% of the embedding’s variance. Value is not literally one-dimensional. Tokens still carry per-token identity in other directions, which OV presumably needs to reconstruct the exact token. But the ordering both circuits care about is one direction.

So the embedding has a size axis, with the QK attending up it, and OV writing back down it.

Min is max mirrored

If the circuit really is ‘read the size axis, attend to the extreme, copy it’, then training the same architecture on min should give the same circuit with one sign flipped. It does.

(Sequence length 5, mean over 5 seeds.)

The QK monotonicity flips sign almost exactly (+0.996 -> −0.996) while the OV copying score is essentially unchanged (0.93 in both). The QK and OV circuits found here are modular: “what do I do with the token I found” is fully reusable, and switching the task only touches “which token do I look for.”

Limitations and what I’d do next

• This is one head, one task, on a tiny synthetic vocabulary. It’s a clean replication, not a new mechanism.
• The copying score is a heuristic (positive-real eigenvalue mass). It happens to agree with the OV heatmap and the logit decomposition here, so I trust it in this case, but it can mislead in general.
• Natural extensions: sweep sequence length to see the OV-logit share decay quantitatively, and check whether the “size axis” is literally the same direction the QK and OV circuits use (not just correlated with it) via a direct projection.