This is the second in a series of posts detailing my manifold findings while investigating how a chess transformer engine that mimics human play represents knight forks.
Last post described strong correlational evidence that the knight-fork policy logit snaps into place after block 5’s attention layer.
In this post I will show that controls rule out obvious alternative readings about block 5, that knight forks are mostly assembled compositionally (check plus queen attack) rather than holistically, and give causal evidence via ablations that tie the fork’s decodability to one specific attention head, ending on one potential confound that threatens to completely reshape what our results mean.
Quick notes:
Chess is relevant to many facets of human cognition, from chunking to sudden insight, and the hope is that general interpretability work on chess models has relevance to cognitive science research in the future.
Along with this, the decomposable structure of the transformer architecture (first rigorously described in "A Mathematical Framework for Transformer Circuits”) adds to the intrigue of the interpretability work. Determining where in the network a concrete skill exists tells one about how its compressed representations form. A knight fork (a single knight move that attacks two valuable pieces at once) is a priori a strong target skill to analyze because it is discrete and unambiguous and important. I wanted to make headway into the question: "at what depth does the network's internal state first make a knight fork decodable?"
Brief Notes on Mechanistic Interpretability
As the field is relatively new, it is worth spending a few sentences to give extra context to unfamiliar readers. The aim of mechanistic interpretability is to understand AI better by reverse-engineering models into human-describable parts. The major hope is that this will help safety and alignment, for example by allowing us to recognize misaligned networks.
Terminology I will keep constant throughout:
-
block 0–7: the eight transformer blocks; each contains an attention layer and an MLP layer
-
residual stream: a vector: the transformer's shared communication layer to which attention and MLP layers additively read and write into, allowing for meaningful decomposition by specific layers’ additions.
-
readout point: the content of the residual stream at some depth
-
logit lens: a technique that takes the residual stream at some layer and pushes it through the final unembedding layer, as though it were the finished product, and then reads its distribution over outputs. Here the output is the move policy distribution.
- decodable: when the move scores highly (large logits) after the logit lens
- policy: the network's output probability distribution over all 4352 moves. Maia plays directly from this, with no search
About the model
MAIA 3 is a bidirectional encoder over all chess board squares: one board goes in as 64 square tokens with pieces represented by a one hot vector, and a single forward pass returns a full policy over all 4352 moves plus a win/draw/loss evaluation. There is nothing autoregressive and no sequence of tokens as input, so we can easily logit lens at each layer. Additionally, its attention heads use two schemes: a "semantic" standard QK self-attention, and a "geometric" GAB scheme to embed complex positional information, given that chess pieces all move in different ways which complicates simple attention patterns.
As we saw earlier, the curve above shows the forking move's logit snapping into place at postattn_05 (and possibly climbing later due to some action at postattn_07). Because a queen winning fork is almost always the best move, two confounding explanations need to be ruled out:
To test these I performed the following experiments:
I used bar plots rather than the logit curves for the control comparisons since averaging many different moves can smooth the mean and hide, for instance, bimodality of a snap. The results were the following (visual sanity checks using the example position):
This high variance smear is very strong evidence that "the best moves resolve at layer 5" is false.
I was dismayed to see that the control plot of the check-only move was very spiky at postattn_05, thinking that the layer was not related to knight forks.
But I realized that what remained after queen removal was not some myriad positions, but a forcing (a chess term defined as requiring an immediate response) knight check. In fact, for the example position, the knight move** Nf2# **is a checkmate regardless (which can lend evidence to both alternate hypotheses). Perhaps the postattn_05 tracks forcing-ness rather than move strength, and there was more to the story.
I ran a quick loop to check how many of the check only moves are still the best moves:
172/500 were still the model’s top policy.
So only 34% of the check only moves were the model’s top choice, and yet 71% snap at postattn_05. If that layer were merely where the top policy moves become decodable, then only around 34% should snap at **postattn_05. **This gave me more hope that the layer was about something more interesting than committing to the top moves, perhaps forcingness. I would have to investigate exactly what it tracked in much more detail in Part 3.
At this point, the most epistemically honest statement is:
"attn_05 is most likely where MAIA 3 resolves forcing moves, and the fork is the sharpest instance (giving check and winning a queen). The queen attack gives a logit surplus on top of the check.”
Current state of analysis:
| | Peak at | fork | Check and queen-attack | ~92% | Control 2 | Check only | ~70% | Control 1 | stochastic | ~18% |
Here I pivoted to also investigate compositional vs holistic representations and superposition. It was time to test and compare the mirror control: removing (relocating) the enemy king and keeping the queen so the move is a pure queen-attack.
I was excited to investigate the following, which, while analogies are cheap, might say something about human chess play: fork logits = check-only logits + queen-attack logits? This could determine if the computation was additive and compositional or if it was a holistic representation in and of itself.
Introspectively it "feels" as though the latter is the case: when I play chess, a knight fork just pops into my conscious awareness without my thinking about individually attacking two pieces, but introspection fails all the time. I am fond of the apt Gazzaniga quote: "The Human Interpreter has set us up for a fall."
The relocated king must leave the position legal, must not defend the landing square (a condition of the "clean" forks), and the knight move must not give check afterwards, including discovered checks. This proved difficult to implement. An example king relocation is below:
Here 32% snap at postattn_05, and there is bimodality with **postattn_06 **(~29%).
Findings on superposition would suggest a priori that the second case is true: a K+Q fork is a rare (sparse) and queen-winning (important) feature, which is exactly the kind of feature that an optimized model tends to represent as its own feature (some direction in very high dimensional space), rather than leaving it a linearly combination of its parts parts.
So, a priori, a dedicated fork feature should make the whole exceed the sum of the parts (I>0).
All this being said, it is possible that the computation of the feature exists in one or the other state in some different layer that cannot be determined from logit lens alone, since it only determines probabilities with respect to the output layer's basis. (This would require getting access to the transcoders from the creators to investigate, which I still have not been able to.)
It appears that feature composition it is! The interaction term is near zero through the forward pass. There is a small jump to 1.5 logits at the final layer, which may or may not be significant, but if so, that is where a dedicated holistic "fork" feature could exist.
Everything above is observational: attention patterns and logit-lens readings.
Let's look at head 5 again (remember it is composed of GAB and standard “semantic” QK attention):
We saw the model learned a sort of "knight pattern" detector for the 5th layer's 5th geometric head. But this might not be causally relevant to why the attention score boosts the fork logit so much.
I wrote a function that ablates (removes one head's entire write out to the residual stream) at block 5's attention layer and reruns the network. If head 5 is causally responsible for doing the knight move information routing, ablating it should do massive damage to the fork logits, far more than ablating any other head.
(The hook on attn_05 is post-projection, where all heads are mixed into every channel so this is computed from the block's weights directly.)
And as predicted, ablating head 5 causes the model to reject the top move—a knight fork—62/100 times!
1) The fork logit snaps at postattn_05 (92% first cross there vs 18% for the second-best move. It is not a generic crossing point).
2) The snap decomposes along a 2×2 table. Check-only retains 70% of the concentration, queen-only 32%, neither 6%: the forcing (check) component carries most of the effect, the queen-attack adds a surplus on top.
3) The fork representation is predominantly compositional. The interaction I = (fork − check-only − queen-attack + neither) is ~0 for most layers until it rises to a small ~1.5 logits deep in the model.
4) The effect is causally concentrated in a single head. Removing head 5's exact write out at block 5's attention layer costs the fork −2.94 ± 1.27 std logits; every other control head costs ≤ 0.14. Ablating head 5 also dethrones the fork as the top move in 62/100 positions, versus ≤ 11/100 for every other head. The ablation proves the head causal, but the gab-vs-qk split within it is not yet causally tested.
There is one strong confound remaining to untangle: Every position analyzed includes a knight move, and Layer 5 Head 5's GAB template learned to form almost the exact pattern of a knight move. It could be the case that, e.g., a forcing pawn move would cause very similar behavior in the model’s internals.
I needed to run probes to distinguish the claims "head 5 commits to forcing moves" and "head 5 computes knight-move policy logits"
Early results suggest that head 5 may be better described as a knight move routing head rather than a knight fork detector.
I was exploring further, and tested the following about the head to cut the alternative hypotheses down as quickly as possible:
This lead to striking new mechanistic interpretations of block 5 head 5’s function that become clear in Part 3