Closure-Validated Circuit Discovery in Attention Heads: Co-activation Proposes, Ablation Disposes

Technical Summary: Closure-Validated Circuit Discovery in Attention Heads

Problem Statement
Mechanistic interpretability has shifted toward viewing groups of components (subspaces, manifolds, head-sets) rather than individual units as the basic object of study. A prevailing methodological approach proposes finding these groups by clustering co-activation statistics (e.g., using binarized-Ising models on sparse autoencoder features or attention heads). However, a critical gap remains: does a statistically coherent co-activation community actually constitute a "circuit" in the load-bearing sense? A true circuit is a group of components whose removal causally degrades model performance on inputs where the group's function applies.

The paper investigates whether cheap signals—specifically co-activation cluster membership, attention-target selectivity, and participation ratio—can serve as sufficient evidence for circuit discovery, or if they merely propose candidates that require causal validation. The authors ask whether clustering alone can distinguish between:

1. True Circuits: Groups that jointly implement a function and are load-bearing.
2. Co-active Communities: Groups that fire together due to shared upstream drivers but contribute redundantly or marginally.
3. Cancellable Communities: Groups that add bias signals which the network compensates for, such that their removal might even improve performance.

Methodology
The study employs a "proposal-and-validation" framework across three 1B-scale language models (Pythia 1B, OLMo 1B, and the Mixture-of-Experts OLMoE-1B-7B) and two input distributions (synthetic induction and natural text).

1. Proposal (Clustering):

   • Signal Extraction: For each attention head, a template-free signal is extracted: the maximum attention weight assigned to any key position in an input.
   • Binarization: Signals are binarized at the per-head median, ensuring every head fires on 50% of examples by construction, removing firing-rate biases.
   • Ising Fit: A pairwise Ising model is fitted to the binarized signals using pseudolikelihood to estimate coupling strengths between heads.
   • Community Recovery: Spectral clustering is applied to the absolute coupling matrix. The number of clusters ($k$) is selected to maximize the Adjusted Rand Index (ARI) against a supervised classification of head capabilities (e.g., previous-token, induction).
   • Candidate Selection: A specific sub-cluster is selected based on purity, isolation (ratio of internal to external coupling), and size constraints.

2. Validation (Closure Test):

   • Ablation: The selected community is ablated by zeroing the output slices of the identified heads before they enter the residual stream.
   • Comparison: The per-example damage (change in cross-entropy loss, top-1 accuracy, and target-token logit) is compared against five matched-random control ablations of the same size.
   • Verdict Criteria: A candidate passes closure if the damage is in the expected direction (loss increases, accuracy/logit decreases) and is statistically significant (z-score > 1.8) relative to the random control distribution.

3. Extended Proxies:

   • The study also tracks attention-to-target selectivity and participation ratio (PR) across training checkpoints to see if these "cheap" signals correlate with load-bearing function over time.

Key Results

• Dense Models (Pythia 1B, OLMo 1B):

  • Co-activation communities identified on both synthetic and natural text distributions pass closure.
  • Ablation of these communities consistently increases loss and decreases accuracy/target-logit relative to random controls.
  • OLMo 1B (Natural Text): Produced the strongest signal, with a z-score of +36.4$\sigma$ on loss and -50.9$\sigma$ on target logit.
  • Pythia 1B (Natural Text): Showed a "downstream-redundancy signature." While the target-logit dropped significantly (-52.5$\sigma$), the aggregate loss change was small (+1.95$\sigma$). This indicates the model has redundant downstream pathways that reconstruct the ablated output, insulating the cross-entropy loss while the specific prediction confidence collapses.

• Mixture-of-Experts (OLMoE-1B-7B):

  • Marginal Failure: On natural text, the marginal Ising model collapsed (ARI $\approx$ 0.006), failing to recover any structure.
  • Route-Conditional Recovery: Stratifying examples by routing patterns (using k-means on routing weights) and fitting separate Isings recovered a statistically significant signal (ARI 0.191, +3.05$\sigma$ above a random-partition null).
  • Closure Failure: Despite the statistical recovery, the discovered community failed closure in the wrong direction. Ablating the community improved loss (z = -5.64$\sigma$ on the discovery stratum, -12.94$\sigma$ outside it). The heads acted as route-modulated noise; removing them helped the model, particularly on inputs outside the specific routing stratum.

• Training Axis Proxies:

  • Function without Form: In OLMo 1B, BOS (beginning-of-sequence) heads were load-bearing (passed closure) at 2B tokens before their attention-to-BOS selectivity pattern had formed (selectivity was near zero).
  • Form without Function: In Pythia 1B, previous-token heads showed high selectivity and high participation ratio at 1B tokens but failed closure (ablation did not damage performance, and in some cases helped).
  • This bidirectional decoupling proves that attention patterns and function are distinct constructs, not merely different sensitivities of the same phenomenon.

Multi-Metric Reporting
The paper argues that cross-entropy loss alone is an unreliable metric for closure. It can be inflated by control-variance collapse (as seen in OLMo natural text) or masked by downstream redundancy (as seen in Pythia natural text). The authors recommend a multi-metric standard:

1. Direction: All metrics must move in the damage direction.
2. Target-Logit: The primary indicator of causal impact.
3. Accuracy: Corroborating evidence.
4. Loss: A conservative floor, subject to aggregation slack.

Claims and Significance
The paper concludes with four specific claims:

1. Proposal vs. Discovery: Unsupervised co-activation clustering is a circuit proposal method, not a discovery method. It identifies candidate communities, but only causal ablation (closure) validates them as load-bearing.
2. Dense Success: In dense 1B-scale models, co-activation communities on both synthetic and natural text can successfully propose load-bearing circuits that survive closure.
3. MoE Failure: In MoE architectures, even route-conditional statistical recovery (which passes a careful null test) is insufficient evidence of a circuit. The recovered communities can fail closure in the opposite direction (improving performance upon removal).
4. Metric Divergence: The gap between attention patterns (selectivity/PR) and function is bidirectional; a head can have a sharp pattern without function, and function without a pattern.

Scope and Limitations
The authors explicitly state they do not claim these results generalize to SAE-feature manifolds (which use different validation criteria), larger models, other architectures (e.g., LLaMA, Mistral), or other input distributions (e.g., code). They also note that the "wrong-direction" failure in MoE is specific to the tested architecture and routing protocol, and that the Ising affinity used is a choice for comparability, not necessarily the optimal affinity for discrete circuits. The study emphasizes that the negative result in MoE is essential, as it demonstrates that statistical coherence does not guarantee causal load-bearingness, especially in complex architectures.