Emergence of Minimal Circuits for Indirect Object Identification in Attention-Only Transformers

AI summary

Overview

Research area: Mechanistic interpretability for natural language processing, specifically reverse-engineering transformer circuits on a symbolic coreference benchmark (Indirect Object Identification, IOI).

Technical level: Intermediate. The paper relies on standard transformer-circuit vocabulary (QK/OV circuits, residual stream decomposition, composition) and includes an eigendecomposition-based spectral analysis, but it is written around a very small, fully describable model.

Scope: The paper trains tiny attention-only transformers from scratch on a symbolic IOI task and reverse-engineers the minimal circuits that solve it.

What This Paper Is About

Mechanistic interpretability tries to explain large language models as human-readable circuits, but pretrained models are so large that it is hard to tell which parts are truly necessary for a given reasoning task. This paper isolates the problem by training very small attention-only transformers (no MLPs, no layer normalization) on a purely symbolic version of the IOI task, then analyzing exactly what the resulting components compute. The goal is to find the minimal mechanism that can perform duplicate-token detection and exclusionary copying, and to see whether that mechanism looks like the multi-hop circuit previously identified in GPT-2 small.

Key Contributions

1. The paper demonstrates that a one-layer, two-head attention-only model solves a fixed-template symbolic IOI task perfectly.
2. It provides a mechanistic analysis uncovering a minimal circuit built on an additive combination of specialized attention head outputs, rather than a multi-hop pipeline.
3. It characterizes the two heads as complementary subcircuits: one additive (aggregating both candidate names) and one contrastive (suppressing the incorrect alternative).
4. It argues that circuits in large, broadly pretrained models may appear overly complex because of multi-task training pressure, whereas task-constrained training can reveal more parsimonious mechanisms.

Main Findings

• Two heads suffice, one does not: A single-layer, single-head model assigns roughly 0.5 probability to the names in the prompt but cannot distinguish which is correct, attending about equally to both names in the dependent clause. With a second head added (parameter count held constant), the model reaches perfect IOI accuracy.
• Zero-layer baseline: A zero-layer model acting as a bigram predicts every name token with equal probability of about 1/N_names = 16.7%, since it cannot use prior context.
• Division of labor between heads: The first head consistently attends to the two names in the initial dependent clause; the second head attends to the subject of the main clause and the non-repeated name in the dependent clause, doing the "heavy lifting" of integrating context.
• Additive versus contrastive outputs: Residual stream decomposition with LogitLens shows the first head's output aligns most closely with the sum direction (correct + incorrect), while the second head's output aligns most closely with the difference direction (correct − incorrect). Added together, the incorrect token's logit component roughly cancels and the correct direction is amplified.
• Spectral asymmetry: The QK and OV matrices show significant asymmetry rather than the symmetric eigenvalue distribution typical of random matrices. The first head's QK circuit has a real eigenvalue of about −5.2 and complex eigenvalues of about 2.3 ± 1.6i, with a positive fraction of about −0.06 (relatively neutral dynamics); the second head's QK circuit has a strong suppression eigenvalue of about −17.5 and a positive fraction of −0.65 (overall inhibitory dynamics).
• Copying versus suppressing heads: In the OV circuits, the first head acts as a copying/passthrough head that amplifies important tokens without rotation or inversion, while the second head has half its eigenvalues real and positive and half imaginary with negative real parts, consistent with subtracting from the other token's logits.
• Positional focus: When name embeddings are averaged so names are indistinguishable, the first head behaves as a positional head focused on the dependent-clause name positions, while the second head attends primarily to the positional embedding of the subject of the main clause.
• Positional embeddings are helpful but not strictly required: Trained without positional embeddings, the same architecture reaches about 70% accuracy with about 67% probability on the correct token; both heads then focus most on the correct non-repeated name. Explicit positional cues let the model reach 100% accuracy and decouple cleanly into additive-contrastive roles.
• Composition in the two-layer model: A two-layer, one-head model (head dimension 8, matching the hidden dimension) shows template-dependent attention in both layers, with the first layer aggregating information into the S2 token for the second layer to use. Ablating composition pathways drops accuracy by about 100% for Q-composition, about 93.33% for V-composition, and about 26.67% for K-composition, which the authors read as heavy reliance on Q and V composition. The abstract, however, describes the two-layer, one-head model as composing information from the previous layer "primarily through query-key interactions."

Methodology in Plain English

The researchers stripped the IOI task down to symbols. Sequences are exactly 6 tokens long, in the format <BOS> IO S1 S2 <MID> ?, drawn from a vocabulary of 6 names plus the two special tokens (total vocabulary size 8). The model must predict the name that is not repeated before <MID>. Two rigid templates are used, called "BABA" and "BAAB", following the naming of Wang et al. (2023).

The models are attention-only transformers with absolute positional embeddings, no feed-forward networks, and no layer normalization. The residual stream dimension equals the vocabulary size (8), and head dimension is 8/N_h, so the parameter count does not change with the number of heads.

Training used cross-entropy loss on the <MID> prediction. Each batch contained all 60 possible unique sequences (6 × 5 = 30 ways to pick the two names, times 2 orderings in the main clause). The optimizer was AdamW with a OneCycle learning rate schedule, maximum learning rate 0.1 and weight decay 0.01, on a single NVIDIA A40 GPU using PyTorch and TransformerLens.

For analysis, the authors used standard interpretability tools: writing the residual stream as a sum of embedding and head outputs, decomposing each head into QK and OV circuits expressed as vocabulary-to-vocabulary matrices, eigendecomposing those matrices, projecting residual stream components onto directions for the correct name, incorrect name, their sum and their difference, ablating positional embeddings, and ablating Q-, K-, and V-composition in the two-layer model by subtracting the first layer's output from the relevant input of the second layer.

Why This Matters

Research impact. The paper provides a controlled testbed where a benchmark task known to require dynamic duplicate detection and exclusionary copying is solved by an extremely small circuit that can be described completely. It contrasts with the multi-hop, multi-class head circuit identified in GPT-2 small, supporting the argument that pretraining on natural text mixes many tasks into one model and may produce unnecessarily tangled circuits. It also shows that a single attention head cannot simultaneously detect the correct referent and copy it, separating two functions that larger models blend together.

Potential real-world applications (as directions suggested by the work, not demonstrated in it):

• Debugging and auditing models that perform coreference-like resolution, by checking whether small, separable subcircuits of the kind described here appear in larger systems.
• Building smaller task-specific models for narrow relational decisions, where task-constrained training may yield simpler and more inspectable behavior.
• Providing a clean test case for interpretability tooling, so new analysis methods can be validated against a mechanism whose ground truth is known.
• Informing data and curriculum design, since the positional-embedding ablation suggests that explicit positional cues can simplify the optimization landscape for this kind of reasoning.

Industry relevance. Teams that need to certify or audit model behavior on coreference and entity-resolution tasks benefit from evidence that a task's essential logic can be localized into a handful of components. The finding that parameters can be held constant while adding a head suggests that architectural allocation (how capacity is split across heads) can matter more than raw parameter count for this type of reasoning.

Future Directions

• Testing how the two-head circuit behaves under varying sequence lengths, multiple interdependent clauses, and dynamic syntax, and finding the threshold of structural complexity at which the multi-hop mechanism described by Wang et al. (2023) becomes necessary.
• Studying training dynamics and developmental interpretability: at what point during optimization do the two heads differentiate into additive and contrastive roles, and what loss landscape dynamics drive that division of labor.
• Reconciling the abstract's claim that the two-layer, one-head model composes primarily through query-key interactions with the ablation results showing the largest drops for Q-composition and V-composition.
• Extending the symbolic setting toward more naturalistic inputs to see whether the additive-contrastive motif persists once tokenization and linguistic structure are reintroduced.

Target Audience

This paper is most useful for mechanistic interpretability researchers and graduate students who want a small, fully analyzable model on which to study circuits, head specialization, and composition. It also suits NLP researchers interested in coreference and entity resolution, and practitioners who need concrete examples of how task-constrained training can produce interpretable structure that larger pretrained models obscure. Readers should be comfortable with attention mechanics, residual streams, and basic linear algebra, though the paper's own model is small enough to follow without deep prior expertise.

Authors’ abstract

Mechanistic interpretability aims to reverse-engineer large language models (LLMs) into human-understandable computational circuits. However, the complexity of pretrained models often obscures the minimal mechanisms required for specific reasoning tasks. In this work, we train small, attention-only transformers from scratch on a symbolic version of the Indirect Object Identification (IOI) task, a benchmark for studying coreference-like reasoning in transformers. Surprisingly, a single-layer model with only two attention heads achieves perfect IOI accuracy, despite lacking MLPs and normalization layers. Through residual stream decomposition, spectral analysis, and embedding interventions, we find that the two heads specialize into additive and contrastive subcircuits that jointly implement IOI resolution. Furthermore, we show that a two-layer, one-head model composes information from the previous layer primarily through query-key interactions. These results demonstrate that task-specific training induces highly interpretable, minimal circuits, offering a controlled testbed for probing the computational foundations of transformer reasoning.