Technical Summary: Conditional Co-Ablation (COAX)

1. Problem Statement

Mechanistic interpretability relies on assigning importance scores to individual model components (e.g., attention heads) to understand how a model produces specific behaviors. The standard approach uses first-order scoring: ablating a unit in isolation, measuring the output change, and aggregating these effects. This assumes component importance is additive.

However, this assumption fails in the presence of self-repair (the "Hydra effect"), where a model possesses dormant backup components that activate only when a primary component is removed.

• The Blind Spot: On an intact model, a backup component appears irrelevant (low importance) because it is silent. When a primary component is ablated, the backup takes over, repairing the damage and muting the primary's measured effect.
• Consequence: First-order scores misread both sides of the redundancy: they underestimate the primary (because the backup fixes the error) and completely miss the backup (because it is dormant). This leads to incorrect attribution, failed capability knockouts, and suboptimal model pruning.

2. Methodology: Conditional Co-Ablation (COAX)

The authors reframe the problem as conditional circuit completion: given a set of primary components $S$, identify the remaining units whose importance grows significantly once $S$ is removed.

The COAX Score

COAX is a label-free, output-grounded score that measures the conditional growth of a unit's ablation effect.

1. Fisher Geometry: The method measures ablation effects in the metric induced by the output distribution (Fisher information). For a unit $u$, the ablation effect is the logit perturbation $\delta z_u = z_0 - z_{\{u\}}$. The energy of this effect is the mean KL cost, approximated by the Fisher-weighted norm of the centered feature $\tilde{\delta z}_u$.
2. Conditional Growth:
   • Let $S$ be a seed set of primary components (e.g., identified by a first-order method).
   • Compute the energy of ablating a candidate unit $u$ in isolation: $E(\delta z_u | \emptyset)$.
   • Compute the energy of ablating $u$ after $S$ has been ablated: $E(\delta z_u | S)$.
   • The COAX score is the difference:
     $$\text{comp}_u(S) = E(\delta z_u | S) - E(\delta z_u | \emptyset)$$
3. Theoretical Basis:
   • Proposition 1: Establishes the Fisher identity, linking the inner product of centered features to the Fisher Gram matrix.
   • Proposition 2 (Additivity Blind Spot): Proves that any score based solely on clean-state statistics (e.g., clean activation or clean gradients) cannot distinguish a pure backup (dormant on clean model) from an inert unit.
   • Proposition 3 (Conditional Growth as Synergy): Shows that for a dormant unit, the conditional growth score approximates the energy of its pairwise synergy with the seed set $S$. Thus, COAX recovers the second-order interaction that first-order scores discard.

Efficiency

Unlike explicit pairwise synergy calculations which require $O(|U|^2)$ forward passes, COAX requires only $O(|U|)$ passes (one single-ablation scan and one conditional scan per seed). It is forward-only and requires no backward passes or task labels.

3. Key Contributions

1. Formalization of Self-Repair: The paper formalizes self-repair as a failure of additivity in component importance and introduces "conditional circuit completion" as the task of recovering dormant compensators.
2. COAX Score: Introduces a label-free, Fisher-grounded score that exposes dormant backups by measuring the growth in ablation energy after conditioning on a primary set.
3. Causal Validation: Demonstrates that recovered heads are genuine backups through mechanistic checks (wake-up curves, direct-logit attribution hand-off) and counterfactual patching (freezing backups removes self-repair).
4. Downstream Utility: Shows that recovering these backups corrects masked attribution, enables complete capability knockout, and improves structured pruning.

4. Experimental Results

Backup Discovery (GPT-2-small IOI Circuit)

• Task: Identify 8 documented backup name-mover heads given the 3 primary name-movers.
• Performance:
  • First-order baselines: Single-ablation saliency (0.33 AUC), AtP (0.60), EAP-IG (0.70), and the self-repair-aware AtP*GradDrop (0.82).
  • COAX: Achieves 0.91 ROC-AUC.
  • Significance: COAX places 6 of the 8 documented backups in its top 20 candidates. The improvement over the best gradient-based baseline (AtP*) is statistically significant ($p < 0.05$) when controlling for the seed set.
• Synthetic Benchmark: On a synthetic benchmark with 100 planted backups, clean-state scores perform at chance or below (0.42–0.51), while COAX reaches 0.90, confirming the theoretical blind spot of first-order methods.

Generalization

• Scale: The "wake-up" phenomenon (increased output norm upon primary ablation) is observed across GPT-2-small, medium, and large models.
• Architecture: Applied to the induction circuit across 8 models (Pythia, GPT-Neo, Gemma, Llama, Qwen, OLMo). COAX recovers compensating sets that significantly drop induction log-probability (2.1x to 12x more than random), demonstrating label-free transferability.

Downstream Applications

1. Attribution Recovery: Ablating primaries alone drops the IOI logit-difference by only 0.22 (masked by self-repair). Adding COAX-recovered backups restores the drop to 1.76, revealing the true causal effect.
2. Capability Knockout: To disable the IOI behavior, one must remove backups.
   • Removing only primaries: Accuracy remains ~0.97 (self-repair active).
   • Removing primaries + COAX backups: Accuracy drops to 0.70, matching the documented oracle (0.72).
   • Removing primaries + "own" next-ranked heads (first-order top-up): Accuracy drops to 0.24, indicating over-ablation into core components.
3. Structured Pruning: A sequential pruning strategy that re-measures COAX scores after each removal (to account for newly exposed backups) outperforms all baselines (Random, Magnitude, Wanda, Taylor) across scales from 124M to 7B parameters. At 50% sparsity, COAX maintains perplexity significantly closer to the dense model than any competitor.

5. Significance and Claims

The paper argues that component importance is not an intrinsic property of a unit but is conditional on the state of the rest of the circuit.

• Motivation: In robust circuits, the components that matter most are often those that remain silent until their primary fails. First-order methods, which score units in isolation, systematically miss these "hidden heroes."
• Operational Shift: COAX shifts the paradigm from "scoring units in isolation" to "scoring units relative to what has been removed."
• Limitations & Scope:
  • COAX is a completion tool, not a standalone discovery tool; it requires a primary seed set (which can be found by first-order methods).
  • It is most effective in dormant-substitution regimes (like IOI backups). In regimes where redundancy is shared among co-firing homogeneous heads (like induction), COAX finds load-bearing sets but does not uniquely identify a distinct "backup class."
  • The method is currently instantiated at the attention-head granularity; it does not recover strong signals in MLP-dominated circuits (e.g., the "greater-than" circuit), suggesting the circuit's mechanism (MLP vs. Attention) dictates the recoverability of self-repair.

The paper concludes that faithful interpretability of robust models must be conditional: redundancy is not noise to be averaged out, but structure that must be conditioned upon to reveal the full circuit.