PaperlayerView exact revision on arXiv

Source-linked AI summary

Slaying the Hydra: Interaction-Aware Circuit Discovery in Language Models

Sankaran Vaidyanathan, Rafal Urbaniak, Emily Bunnapradist, Michelangelo Naim, Daniel Waxman

arXiv:2610.04017v1cs.LGcs.AIcs.CLstat.ML

TL;DR

Per-component scores can miss interactions among causes, while exhaustive set-level evaluation is computationally infeasible. The paper introduces WISE and JuntaLearner to evaluate and rank components across sets and witnesses. Across tasks and models, JuntaLearner outperforms attribution baselines on key circuit metrics, though its cost advantage appears only for large candidate pools.

  • Problem

    Per-component scores miss context-dependent interactions, while exhaustive set-level scoring grows exponentially and makes interaction-aware circuit localization infeasible in practice.

    Source: §1 · pp. 1–2

  • Method

    WISE averages causal effects over sets of mediators and witnesses, and JuntaLearner uses gradient-based optimization to rank components across varying-sized sets.

    Source: §1 · p. 2, §3 · p. 4

  • Results

    Across five tasks and three models, JuntaLearner outperforms attribution baselines on kill and net, performs at least as well as exact per-head NIE, and improves scores over singleton training on all metrics.

    Source: §1 · p. 2, §5 · p. 9

  • Takeaways & Limitations

    JuntaLearner’s higher net scores indicate circuits more localized to the studied task, while set-level interactions matter because descending set sizes improve scores over singleton training.

    Source: §5 · p. 9

    The main experiments focus only on attention heads; training is slower than exact per-head patching on smaller GPT-2 and Qwen3 models, and witness-pinning benefits depend on the model.

    Source: §5 · p. 10

Abstract

from arXiv · show

Localizing behavior to individual components of a language model is a central goal of mechanistic interpretability. However, scoring components one at a time misses context-dependent effects: a primary component can inhibit the activation of a backup, leading to issues with ranking components. Actual causality studies the structure of such interactions via witnesses: variables that provide contextual information to resolve interaction terms. However, estimation with witnesses typically requires combinatorial enumeration and is infeasible in practice. We introduce the witness-integrated set effect (WISE), a family of causal estimands that build on the witness mechanism while taking expectations over sets of causes and witnesses to remain computationally feasible. Building on this approach, we introduce JuntaLearner, a gradient-based circuit discovery method that learns to rank components by their causal impact across varying-sized sets of components and witnesses. Alongside faithfulness metrics, we introduce measures of necessity and task specificity, and the circuit recognition score (CRS) to summarize each metric across circuit sizes while emphasizing effects achieved by small circuits. Across tasks and models of increasing size, JuntaLearner achieves higher mean CRS compared to attribution baselines on all metrics. Since its cost does not grow with the number of candidate components, JuntaLearner scales to large models while accounting for set-level interactions and avoiding first-order approximations.

1 INTRODUCTION

Per-component causal scores can miss interactions, including primary components masking backups, while exhaustive set-level scoring is infeasible at realistic model sizes. The paper introduces WISE, JuntaLearner, and circuit-quality metrics to address interaction-aware circuit discovery and evaluation.

1 INTRODUCTION

  • Individual NIE scores can miss genuine causes and fail to compose into circuit effects because components interact, including when a primary masks a backup.Joint effects also include superadditive or subadditive interactions among set members.Source: §1 · p. 1
  • Exact per-component patching scales with the model and still misses joint effects, while exhaustive set-level scoring grows exponentially with the candidate pool.First-order approximations are cheaper but discard interactions; prior approaches restrict the interaction search using prior knowledge.Source: §1 · pp. 1–2
  • WISE averages causal effects over interacting component sets and uses witness pinning to block self-repair from outside each set.Witnesses are downstream sites held at their clean values.Source: §1 · p. 2
  • JuntaLearner learns a distribution over candidate sets with a differentiable Gumbel-top-k relaxation and dynamically annealed set sizes, yielding rankings for circuits of different sizes.The paper also introduces kill for necessity, net for task specificity, and CRS to summarize metrics across circuit sizes while emphasizing small circuits.Source: §1 · p. 2
  • Across five tasks on GPT-2 small, Qwen3-1.7B, and Llama-3.1-8B, JuntaLearner outperforms attribution baselines on kill and net, and performs at least as well as exact per-head NIE.On the public MIB Circuit Localization test split, one untuned run beats MAttr’s released circuits on kill, net, and CMD faithfulness.Source: §1 · p. 2

2 WITNESS-INTEGRATED CAUSAL EFFECTS

Witnesses help separate primary causes from backups that can compensate after intervention. WISE averages set-level effects over mediator and witness sets, exposing interactions while avoiding exhaustive enumeration.

Witnesses

  • In GPT-2’s IOI circuit, standard NIE ranks backup name movers above primaries, but witness integration recovers the intuitive primary-above-backup ranking.Scored alone, backups score 0.74 versus 0.14 for primaries; with witnesses integrated, primaries score 1.16 versus 0.60 for backups.Source: Figure 2 · p. 3, §2 · p. 3
  • Witnesses are downstream components pinned to their starting-prompt values so they cannot compensate for an intervention on the mediator set.For NIE the starting prompt is a; for PIE it is a′.Source: §2 · p. 3

Witness Integration

  • WISE averages a set-level score over a distribution of mediator and witness sets, with mediators patched, eligible downstream witnesses pinned, and remaining variables free to recompute.The paper retains NIE and PIE names for the two patching directions despite their relation to alternative causal-effect terminology.Source: §2 · pp. 3–4
  • Sets and witnesses expose interactions when a primary and backup are included together or when an outside backup is pinned; witness benefits vary by model and help most on Qwen3.Enumerating mediator and witness sets becomes impractical as the candidate pool grows, motivating a gradient-based learner.Source: §2 · p. 4

3 CIRCUIT DISCOVERY WITH JUNTALEARNER

JuntaLearner learns which mediator sets to score by optimizing witness-integrated effects over a trainable distribution of candidate sets. It decreases set sizes during training and reads candidate rankings from the training run.

Circuit discovery

  • JuntaLearner draws candidate sets from trainable inclusion logits, scores them on paired prompts, updates the logits toward higher-scoring sets, and reads a candidate ranking from the run.The top ℓ candidates form the circuit of size ℓ.Source: §3 · p. 4

Objective

  • With NIE as the set-level score, training maximizes WINIE over mediator sets of size k; NIE targets necessity, PIE targets sufficiency, and their soft minimum gives the bottleneck score.A candidate’s gradient tracks how the expected score changes as that candidate becomes more likely to be drawn.Source: §3 · p. 4

Sampling and updates

  • Each step samples mediator sets by adding Gumbel noise to inclusion logits and selecting the k largest, then samples eligible downstream witnesses and scores all sets on the same prompt pairs.The forward pass uses hard selections; the backward pass uses a straight-through estimator from the continuous Gumbel-top-k relaxation.Source: §3 · p. 4

Set-size ladder

  • Training over decreasing set sizes improves all evaluation scores on every model compared with training only on single-component sets.Large sets help find joint effects, while later smaller sets help separate the members of a good set.Source: §3 · p. 5

Ranking

  • Candidates are ranked from normalized, clipped gradients recorded during training rather than final logits, and one run yields nested circuits across sizes.The final logits may not preserve candidates’ relative effect magnitudes because Adam adjusts update scales separately.Source: §3 · p. 5

4 EXPERIMENTS

The experiments evaluate whether JUNTALEARNER identifies compact, task-relevant circuits across three models and five tasks, using faithfulness, necessity, task specificity, and CRS. JUNTALEARNER generally compares favorably with attribution baselines, including on the MIB benchmark, though performance varies by metric and model.

4 EXPERIMENTS

  • The evaluation covers attention-head circuits on three models and five tasks, comparing JUNTALEARNER with per-head NIE, EAP, EAP-IG, MAttr, and random rankings.Faithfulness is evaluated at ten log-spaced candidate-pool proportions; kill and net are evaluated on the top 1–25 heads, or 26 in some cells.Source: §4 · p. 5, §4.1 · p. 7

4.1 CIRCUIT EVALUATION

  • Faithfulness measures performance retained when the circuit’s complement is patched, kill measures necessity when the circuit itself is patched, and net contrasts target-task removal with collateral damage.The metrics use held-out prompt pairs and are normalized by the effect of patching the full candidate pool.Source: §4.1 · p. 6
  • CRS weights performance across circuit sizes logarithmically to reward early success, unlike linear-size summaries that weight small and large circuits equally.On Qwen3 subject–verb agreement, CRS favors early recovery over a late overshooting curve, while CPR and 1 −CMD favor the latter.Source: §4.1 · pp. 6–7

4.2 RESULTS

  • JUNTALEARNER’s tuned configurations achieve the highest mean faithfulness, kill, and net scores, leading five of six columns in Table 1.The table summarizes CRS across fifteen model-task cells; higher is better in every column.Source: §4.2 · p. 7, Table 1 · p. 8
  • With one fixed configuration per model, JUNTALEARNER scores comparably to signed per-head NIE and higher than EAP, EAP-IG, and MAttr, but baseline strengths vary by metric.On Qwen3, signed EAP-IG matches strong kill scores on four of five tasks but has negative net on three; MAttr has highest worst-cell faithfulness at 0.42.Source: §4.2 · p. 8
  • On Llama-3.1-8B, JUNTALEARNER takes 14–16 minutes on the slowest task versus 97 minutes for exact per-head NIE, a 4–7× speedup across timed runs.On smaller models it is 2–8× slower than exact patching; training cost is independent of candidate-pool size, unlike exact scoring.Source: §4.2 · p. 8
  • Descending set-size ladders improve predicted faithfulness, kill, and net on all three models; necessity-only scoring performs at least as well as combined objectives.Witness pinning is model-dependent: pinning all eligible witnesses lowers all three scores on GPT-2 and Llama, while on Qwen3 it improves net but reduces kill and faithfulness.Source: §4.2 · p. 9
  • Across twelve MIB cells, JUNTALEARNER beats MAttr on aggregate kill and net under all three summaries and exceeds exact per-node NIE on CRS for all three metrics.For faithfulness, JUNTALEARNER’s summed 1 −CMD is 9.1 against MAttr’s 1.0.Source: §4.2 · p. 9

5 DISCUSSION

The discussion argues that interactions make circuit recovery and selective behavior control distinct, and that circuit quality depends on the evaluation metric. It also identifies limits in the experiments’ scope and reports project contributions and AI use.

5 DISCUSSION

  • The descending ladder improves scores over singleton training on all metrics, while higher net scores indicate circuits more localized to the studied task.This localization helps control the behavior in isolation.Source: §5 · p. 9
  • Metric choice can reverse circuit-quality verdicts: MAttr often matches or exceeds the default on faithfulness, while the default leads on kill and net in most cases.MAttr’s faithfulness lead often involves overshoot, which CMD penalizes.Source: §5 · pp. 9–10
  • The main experiments cover attention heads only, and the MIB comparison uses a single untuned configuration.Training is slower than exact per-head patching on smaller GPT-2 and Qwen3 models; its cost advantage appears only for larger candidate pools.Source: §5 · p. 10

CONTRIBUTION STATEMENT

  • SV and RU jointly led the project; RU designed and implemented JUNTALEARNER, while DW and MN contributed to WISE and its supporting theory.SV, RU, and DW led the writing.Source: CONTRIBUTION STATEMENT · p. 10

AI USE STATEMENT

  • Generative AI assisted with code reformatting, debugging, and experimental sweeps; the authors say they reviewed the work and regenerated every number and figure from saved artifacts.The authors take responsibility for the paper’s final content.Source: AI USE STATEMENT · p. 10

A SAMPLING ALGORITHM

The sampling algorithm estimates candidate rankings from causal scores on sampled sets rather than scoring each candidate alone. It contrasts sets of the same size to reduce bias from set cardinality.

A SAMPLING ALGORITHM

  • The sampled ranking estimates WISE scores by drawing candidate sets, scoring each set, and comparing sets that hold a candidate with those that do not.A separate trained ranking is read from the training trace; the paper uses that approach in Section 4 because Monte Carlo is costly on larger models.Source: A SAMPLING ALGORITHM · p. 14
  • For each sampled set size, the algorithm draws sets and witnesses, records their witnessed NIE scores, then scores each candidate using sizes where it was both held and omitted.The description uses witnessed NIE for simplicity, but the objective can be swapped for other scores.Source: A SAMPLING ALGORITHM · p. 14
  • Matching sets within each size prevents candidates from receiving credit merely because they were sampled into larger sets, though a candidate with no causal role is not guaranteed a zero score.The induced ordering, rather than the score’s zero point, is used to generate circuits.Source: A SAMPLING ALGORITHM · p. 14

B LEARNING ALGORITHM

JUNTALEARNER learns candidate rankings by optimizing scores on sampled sets, shrinking set size over a ladder, and reading rankings from training gradients. The section also describes convergence checks, optional final-stage measurements, and supporting comparisons.

B LEARNING ALGORITHM

  • Each candidate has an inclusion logit; sampled sets are scored on prompt pairs and updated to favor higher scores, with training using a differentiable top-k selection.A run yields a ranking from which circuits of different sizes can be extracted.Source: B LEARNING ALGORITHM · pp. 14–15
  • The descending set-size ladder begins broad to find joint effects and narrows to help separate members of a promising set.A rung ends when its top candidates stabilize and the logit gap continues widening, or when its step ceiling is reached.Source: B LEARNING ALGORITHM · pp. 14–15
  • Rung completion requires overlap above a threshold, stability against dips, and an improving separation margin.These checks distinguish a settled ranking from oscillation or logits that have stopped moving; the paper uses thresholds τ = 0.9 and ϵ = 0.1.Source: B LEARNING ALGORITHM · p. 15
  • Candidates are ranked from recorded gradients normalized per step and clipped per candidate, because Adam-adjusted final logits need not preserve effect magnitudes.Per-step normalization prevents broad early rungs from dominating narrower late rungs; clipping limits the influence of extreme bursts.Source: B LEARNING ALGORITHM · p. 16
  • A final k = 1 rung can directly measure finalists’ individual effects, but Llama’s default ladder ends at k = 7 and therefore takes no such measurements.The measurement stage balances draws across a selected top group to avoid repeatedly measuring settled leaders while under-sampling nearby candidates.Source: B LEARNING ALGORITHM · p. 16
  • Per-head NIE aligns most with JUNTALEARNER when witnesses are free, while WINIE moves into modest agreement as witnesses are held.Agreement is not accuracy: the orderings lack ground truth, so circuit performance is the relevant check.Source: Figure 7 · p. 17, B LEARNING ALGORITHM · p. 17
  • JUNTALEARNER uses a fixed-budget cost of Θ(T · B · P) forward-backward passes, independent of candidate-pool size, unlike exact per-head patching’s Θ(|M| · N) cost.The gradient approximations are also independent of pool size, but their circuits score worse.Source: B LEARNING ALGORITHM · p. 17

C EXPERIMENTS

The experiments use surrogate models to choose configurations and replicate model-specific defaults across tasks. They report a representative training run and describe how seed replication and score-dependent replication affect the strength of interval evidence.

C EXPERIMENTS

  • Two surrogate models chose the 572 configurations behind Section 4: a Gaussian process used CRS as output and expected improvement to pick next runs, while CatBoost checked each recommendation.The search covered 24 to 76 configurations per cell; 150 of the 572 runs were seed replicates.Source: C EXPERIMENTS · p. 18
  • In the illustrated Llama addition run, the first four ladder rungs ended early, while the final rung ran to budget with the ordering still moving and separation reaching about 0.3.The figure follows the default free-witness ladder.Source: C EXPERIMENTS · p. 18, Figure 8 · p. 19
  • The model defaults use free witnesses, with Qwen3 also using an interior setting and Llama using the 142-71-43-21-7 ladder.Defaults were retrained independently 5 to 35 times per task cell; step budgets ranged from 30 to 90 on GPT-2 and Qwen3 and 30 to 150 on Llama.Source: C EXPERIMENTS · p. 18
  • Per-cell intervals are 95% t-intervals over seeds, but cells with an initially indistinguishable-from-zero margin received extra seeds, weakening comparisons with fixed-depth replication.The pooled interval is a t-interval over five per-task seed means.Source: C EXPERIMENTS · p. 18

D RESULTS IN DETAIL

The section compares trained configurations and baselines across task-and-goal cells, then examines model-level defaults, task-specific winners, and configuration-knob effects. The analyses also assess variability across seeds and how extensively task-specific configurations were optimized.

D RESULTS IN DETAIL

E MIB BASELINE AND COMPARISON TO MATTR

The comparison finds different strengths across faithfulness, kill, and net, with results varying between the main experiments and MIB’s public test split. It also distinguishes the methods’ patching objectives and describes the scope of the MIB evaluation.

E MIB BASELINE AND COMPARISON TO MATTR

Loading 2610.04017v1…