Causal Head Gating: A Framework for Interpreting Roles of Attention Heads in Transformers
Authors: Andrew Joohun Nam
Henry Conklin
Yukang Yang
Thomas L. Griffiths
Jonathan D. Cohen
Sarah-Jane Leslie
Venue: NeurIPS 2025 poster ()
Keywords: mechanistic interpretability
large language models
causality
Primary area: deep_learning
TL;DR: The paper proposes Causal Head Gating, a scalable, unsupervised method to classify transformer attention heads by causal impact on task performance that reveal task-specific sub-circuits.

Abstract: We present causal head gating (CHG), a scalable method for interpreting the functional roles of attention heads in transformer models. CHG learns soft gates over heads and assigns them a causal taxonomy—facilitating, interfering, or irrelevant—based on their impact on task performance. Unlike prior approaches in mechanistic interpretability, which are hypothesis-driven and require prompt templates or target labels, CHG applies directly to any dataset using standard next-token prediction. We evaluate CHG across multiple large language models (LLMs) in the Llama 3 model family and diverse tasks, including syntax, commonsense, and mathematical reasoning, and show that CHG scores yield causal, not merely correlational, insight validated via ablation and causal mediation analyses. We also introduce contrastive CHG, a variant that isolates sub-circuits for specific task components. Our findings reveal that LLMs contain multiple sparse  task-sufficient sub-circuits, that individual head roles depend on interactions with others (low modularity), and that instruction following and in-context learning rely on separable mechanisms.

Causal Head Gating: A Framework for Interpreting Roles of Attention Heads in Transformers
Andrew J. Nam Princeton Laboratory for AI Natural and Artificial Minds
Henry Conklin
Princeton Laboratory for AI
Natural and Artificial Minds
Princeton University andrewnam@princetoor
Princeton University
cenny..ootiinggrrinctton.
Yukang Yang
Thomas L. Griffiths
Department of Electrical and Computer Engineering
Department of Psychology
Princeton University
Princeton University tomgOprinceton. edu
vy1325@princeton edu
Jonathan D. Cohen'
Sarah-Jane Leslie"
Princeton Neuroscience Institute
Department of Philosophy
Princeton University jdc@princeton.c edu
Center for Statisties and Machine Learning
Princeton University
sjleelie@princetonn.c edu
Abstract
We present causal head gating (CHG), scalable method for interpreting the functional roles of attention heads in transformer models. CHG learns soft gates over heads and assigns them causal axxoomyy-fciliieee relevant- based on their impact on task performance. Unlike prior approaches in mechanistic interpretability, which are hypothesis-driven and require prompt templates target labels, CHG applies directly any dataset using standard next- token prediction. We evaluate CHG across multiple large language models (LL.Ms) in the Llama 3 model family and diverse tasks, including syntax, commonsense, and mathematical reasoning, and show that CHG scores yield causal, not merely correlational, insight validated via ablation and causal mediation analyses. We also introduce contrastive CHG, variant that isolates sub-circuits for specific task components. Our findings reveal that LLMs contain multiple sparse task-sufficient sub-circuits, that individual head roles depend on interactions with others (low mod- ularity), and that instruction following and in-context learning rely separable mechanisms.
interfering, or ir-
Introduction
Large language models (LLMs) [ 2 3 represent state-of-the-ard systems across wide array of domains, exhibiting remarkable generalization and problem-solving capabilities. Yet, as these models grow in scale and complexity, they become increasingly opaque, making more difficult to understand, predict, or control their behavior, which raises concerns about safety and misuse This has motivated a growing body of work on interpretability, which seeks better understand how LLMs learn and represent information, and how their responses can be shaped 7 g
Equua contribution; authors listed alphahetically
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Interest has focused in particular on transformer- baase architectures 19] such as GPT M. LLaMA B Gemma [TO), and DeepSeek 22, in which the central processing blocks consist of multi-head attention followed by multi-layer perceptrons. Here, there has been considerable research on the roles of individual attention heads, which have been found to exhibit some level of human- iimman-itterrretablityy
11 12 [3.
Two broad categories of approaches dominate research on mechanistic interpretability in LLMs. The first uses a trained mapping from latent representations to human-interpretable concepts, such as syntactic features 7 [ [5] or identifiable items (e.g., the Golden Gate Bridge [T6).. The second uses causal interventions to identify portions of a single weight matrix or individual attention heads responsible for a specific behavior [17 18). These approaches often focus on small portions of a model, zzooming in' [19] in an effort to interpret the role of a single computational subgraph. However, in deep-learning models, computation is often distributed (0) and the role of one component is dependent on another [21 22 23. making the behavior of such complex distributed systems difficult
to predict from an understanding of their parts alone 24
To apply a distributed perspective to mechanistic interpretability, we introduce causal head gating (CHG) which identifies a parametrically weighted set of heads that contribute to a model's execution of given task. Given dataset that defines task, we fit a set of gating values for each attention head that applies a soft ablation to its output using next-token prediction, so that task-facilitating heads remain unaltered while any task- iest-interfering heads are suppressed. Using a simple regularization procedure that further separates irrelevant heads from those that facilitate or interfere with task performance, CHG assigns meaningful scores to each attention head across an entire model according to its task contribution. We use these scores to define a taxonomy of task relevance according to how individual attention heads contribute to model's distributed computation of a given task, describing each head as facilitating, interfering or irrelevant. In this respect, CHG offers an exploratory complement to standard hypothesis-drive approaches to mechanistic interpretability, assigning causal roles without relying on predefined hypotheses about what each head might be doing.
Beyond its conceptual contribution, CHG also offers several practical methodological advantages over existing mechanistic interpretability tools. First, because CHG operates directly on next- tokentoken prediction, it avoids the need for externally-provided provided labels 17 15 16], controlled input-outpu pairs 7 14 15), or rigid prompt templates 255
12 1 which are often required for decoding and interventional approaches. Second, CHG naturally accommodates complex target outputs, including chain-of-thoughh reasoning (26), where the solution spans multiple intermediate steps. Finally, CHG is highly scalable: it introduces only one learnable parameter per attention head and requires no updates to the underlying model weights, so that the CHG parameters can be fitted minutes using gradient-based optimization, even for LLMs with billions of parameters. Thus, in settings where analyzing complex dependencies between heads is important, it is feasible to fit large samples of CHG values to estimate distribution over gating values in a bootstrap fashion.
To test its efficacy, we apply CHG across a diverse set of mattemaiiaaallll commonsense, and syntactic reasoning- -ndniinggndddddoo across LLMs ranging from I to 8 billion parameters with varying training paradigms. We use CHG to analyze not only where specific computations take place, but also how distributed they are across attention heads, and how these patterns vary across different tasks and models. We also validate the causal scores produced by CHG by comparing them against targeted ablations as well as causal mediation analysis 22 25 showing strong agreement between predicted and observed effects. Finally, we extend CHG to a contrastive setting to identify distinct sub-circuits that support instruction following versus in-context learning. suggesting that even semantically similar tasks can be underpinned by separable mechanisms.
Our main contributions are fourfold:
1. We introduce causal head gating (CHG), parametric, scalable method for identifying potentially distributed, task-relevant sub-circuits in transformer models without requiring prompt templates or labeled outputs, and extend it with contrastive CHG to isolate heads supporting specific sub-tasks.
2. We propose a simple causal taxonomy of heads-facilitaiing, interfering, and irrelevant-that quantifies the effect of each on task performance using CHG-derived scores.
3. We use CHG to show that models contain multiple task-sufficient sub-circuits with varying degrees of overlap, suggesting head roles are not fully modular but depend on interactions with other heads.
2

---

4. We use CHG to show that instruction following and in-context learning rely on context-dependee separable circuits at the head level, where CHG-guided gating can selectively suppress one mode without substantially disrupting the other.
The accompanying repository for this paper can be found attttpps//ggthhh...oo//mm//tttt
causal_head_gating
2
Related Work
Representational decoders Representational decoders are models trained to map hidden activations
to externally labeled properties 7 14 [5], estimating the mutual information between representations
and those properties 28]. However, such probing results are difficult to interpret: simpler decoders
may underfit and miss relevant features (false negatives), while complex decoders may overfit and
tradeoff's to contextualize results [30]. Moreover, although decodability indicates that a property is encoded in the representation, it does not imply that the model uses that information for its task, highlighting a correlational finding rather than causal one (31. Finally, representational decoders require labeled datasets, constraining their use to curated, predefined properties. For a comprehensive review of the probing framework and its limitations, see (27].
learn spurious correlations (false positives) [29] 30,, requiring complexity-aceuran
Sparse autoencoders (SAB) can be viewed related approach, where the autoencoder reconstructs representations through a sparse bottleneck to reveal modular or interpretable features [16 32) However, like probing classifiers, their insights remain correlational and still depend on post hoc labeling or interpretation, inheriting the same supervision bottleneck. In contrast, CHG performs direct interventions on model components without external supervision and proposes sufficient sub- circuits to the default unablated model, thereby identifying causal links between attention heads and model behavior on a task.
is used to identify the func- tional roles of specific attention heads by crafting controlled prompt pairs that isolate a hypothesized behavior, then intervening on model components to measure their causal effect on outputs. For instance, in the went to the store, John gave a drink to..." " are used to identify attention heads responsible for resolving coreference. By patching specific head outputs from source sentence into a structurally matched target, and checking whether the model changes its prediction (e-g. "Alice" instead of "Mary"), CMA localizes the relevant circuit. It has also uncovered head-level roles in function tracking 12 symbol abstraction [3], and other structured settings 335
Causal mediation analysis Causal mediation analysis (CMA)
(IOI) task [25], sentences like "When Alice and John
However, CMA relies on manually crafted prompt templates and clear mechanistic hypotheses, which limits its scalability to more complex domains. In open-ended tasks like mathematical reasoning 36 37 38), the diversity of required knowledge makes it hard to design effective controlled inputs. single shared template unlikely to accommodate even two prompts from the MATH dataset
[37] such as: "If En cos2" 0 5, what is cos 20?" and "The equation x + 2x has two complex solutions; determine the product of their real parts." " Moreover, LLMs often solve such problems most effectively via chain-of-thouugh reasoning [26], which unfolds over multiple steps,
further complicating the use of a unified prompt structure.
Head ablations Despite the use of multiple heads being commonplace in transformer-based architectures, it has been observed that multiple, and sometimes the majority of, heads can be entirely pruned with minimal impact on model performance [39 [8
40 41). Moreover, entire layers can
be pruned while retaining model performance 442
43 44). However, existing works pruning
attention heads have focused primarily on custom-trainn trained small-scale transformers [39 [S] 40) or
BERT-based 45 models 4T 433. and the literature is limited for modern causal LLMs such as GPT
46
and Llama (3]
Head pruning has also been used to validate findings from other interpretability methods, such as CMA 25
[3 or attention pattern analysis [18]. In these studies, researchers first identify heads believed to perform specific functions, then ablate them to test their causal impact. Such targeted ablations often lead to disproportionate drops performance, supporting the hypothesis that those heads are functionally important.
3

---

Most closely related to our work are differentiable masking and soft soft-gating approaches that learn which attention heads retain or suppress. In 47, the authors apply sparsity gating to identify subcircuits and use the fitted parameters weighting values in convex combinations for activation patching. Similarly, 44 learns scaling constants for each attention head, but uses the fitted values to identify heads that are most suitable for fine-tuning. Thus, while methodologically similar, our work is unique in applying the gating parameters to identify task-sufficient causal sub-circuits.
Others 18 40 have opted for hard, binary ablations using the Gumbel-softuax trick 49 50], fitting gating probabilities rather than weighting parameters. Although these Gumbel-based approaches have been applied for causal circuit discovery in a similar spirit to our work, they suffer from fundamental limitation that CHG does not. Specifically, while Gumbel-based gating methods also learn differentiable gates per head, they treat each head independently, effectively learning separate Gumbel-Bernoulli distributions for head inclusion. This factorized formulation models only marginal probabilities and cannot capture interdependencies hetween heads that jointly affect task performance. In contrast, CHG jointly optimizes all gating coefficients under the model's loss, capturing the full range of interactions and contingencies between the attention heads. Because CHG is highly scalable, can be fit repeatedly across random seeds or subsets, effectively sampling from the space of sub- circuits without assuming independence between heads. This enables estimation of the underlying distribution over functional head configurations while preserving the joint statistical structure that factorized gating approaches discard.
Our Approach: Causal Head Gating
Causal head gating is based on three ideas: applying multiplicative gates to attention heads to evaluate their roles, using regularization to produce variation in the estimates of the gating parameters, and constructing a taxonomy based on that variation. We introduce these ideas in turn.
3.1 Applying gates to attention heads
For transformer with L layers and H attention heads, we define gating matrix G E [0, 1] where G& scales the output of head h in layer f, just before the output projection matrix W (shown in red for an example head in Figure (). Given input hidden states X E Raae drecdd each head computes:
ArA = softmax XWb (X
Ve, = XW.* Zen = GG Ve,))
vdk
where W.S.h W Wht E Rhuce xd. are learned projection matrices for queries, keys, and values.
The gating coefficient Gen modulates the contribution of head h by scaling its output Ze, after attention is applied but before the heads are combined (see Figure Ia). The gated outputs are then
(a)
1
(b)
(c)
1.00
Layer
Reqularization
Cata Value
ix c
DD
0.75
Inelovant
Faclitating
20
0.50
0
GLa
GLD
GLR
Head Type
Rodlooiing
10
Indexure
merfering
Interfering
A VLa
VLn
A Viur
0.02
230
500
750
1000
0.00
0.25
0.50
0.75
1.00
Gradient Updetes
Figure 1: (a) Schematic of a single multihead attention block with CHG-determined gating attenuation (in red). (b) Gate fitting trajectories for three heads on L332-3BI with OpenMathhnstruutz. When fitting with À < O and À 0, G+ and G both stay near 1 for facilitating heads and near 0 for interfering heads, but bifurcate and 0 respectively for irrelevant heads. (c) Gate values after fitting.

---

Table 1: Causal taxonomy for head roles and corresponding gating patterns.
Role
Description
G+
G
Metric
Ablation Effect
Facilitating Supports task perfor-
High
High
G
Decreases task perfor- mance
mance
Interfering
Interferes with task per- Low Low 1 - G
Increases task perfor- mance No effect on task per- formance
formance
Irrelevant
Negligible impact on High Low G x (1- performance
concatenated and projected:
Output, Concat Ze....... W., Wo ER Hdy x deaas
We fit G by freezing the parameters the model Mo and minimizing the negative log-likelihodd (NLL) next-token prediction task with regularization term specified below.
3.2 Producing variation through regularization
We add regularization term to the objective that introduces small but consistent gradient- clipped to ensure NLL remains the dominant term that nudges the gates for task-inrelevant heads toward or O while leaving task-relevant ones relatively unaffected. The NLL optimizes towards improving task performance, and tunes the heads by either increasing the gating values for task-facilitating heads decreasing the gating values for task-interfering heads. However, if head does not affect task performance, i.e. is task-irrelevant, then the expected gradient from the NLL is O, which confounds interpretation of task relevance when evaluating the tuned gating values: gate G1,6 may be close either because it is important for performing the task (causal), or because gating it has no effect (incidental). We address this limitation by introducing --reeullaizzatin term in our objective function, with weight À that either nudges gates toward 1 for maximal sparsity (A < 0):
for maximal density (A >> 0) or toward
L(G: = Σ (K.))D
2 2 G) - - o I
(1)
Negative log-likelihood (NI.L)
Regularization
where Mo is the model being analyzed, is the target text sequence for a given prompt in dataset D, and -1 is the clipped inverse-sigmoid function.
We fit twice: once with
to encourage retention (G'+), and once with < 0 to encourage removal (G-). To ensure that the heads are aligned across both optimizations, we first fit G with to establish a shared initialization (see Figure D. so that any differences between G and G reflect only the effect of the regularization and not divergent optimization paths.
3.3 Constructing taxonomy of task relevance
The G+ and G matrices allow us to interpret the functional role of each head. To formalize this, we introduce a causal taxonomy (Table [ in which each head assigned one of three roles dless-fcilitatinng interfering, or irrelevant-basee on its predicted impact on model performance under ablation. Facilitating heads positively contribute to performance, while ablating them degrades Conversely, interfering heads negatively contribute to performance, while ablating them improves it. Finally, irrelevant heads have negligible effect, with ablation leaving performance effectively unchanged.
We instantiate this taxonomy using the fitted CHG matrices G and G which reflect head behavior under opposing regularization pressures. Facilitation is measured by G- heads that remain active despite pressure to suppress are likely necessary for the task. Interference is measured by G+: heads that are suppressed even under encouragement to remain are likely harmful. Irrelevance measured via G @ GT), identifying heads that vary in gate values based regularization.
5

---

Syntax
Probability Targor
Common Serse
Facilitation
Irrelevance
.50
-100
10
20
30
40
50
10
20
30
40
50
10
20
30
40
50
10
20
30
40
50
Number of Ablated Heads
Figure 2: Difference in target log-probability when sequentially setting individual gates in G+ to 1 and 0 in order of facilitation, irrelevance, and interference scores. The horizontal axis shows the number of heads ablated in descending score order. Positive values indicate task improvement, negative values indicate degradation, and values near zero indicate no effect. Note that not all heads in the top 50 necessarily have high absolute scores.
4 Experiments and analyses
4.1 Causal roles of attention heads
We begin by reporting experiments that evaluated the causal taxonomy presented in Table across four variants of the Llama LLM 33: L3.1-8B, pre-trained 8B 3B-parameter model; L3.2-3B, 3B-parameter model distilled from Llama-3.1-70B (not used in this paper); I.3.2-3BI, an instruction- tuned version of Llama-3.2-3B; and 1.32-1BB, 1B-parameter model distilled from 1-.I- For each model, we fit CHG matrices on three task types performed over distinct datasets: mathematical reasoning from OpennMathInstrucc2 (38), syntactic reasoning from the subset labeled "syntax" in BIG-Bench (51), and commonsense reasoning from CommonsenseQA [52]. We fit CHG matrices independently for each model-dataset pair across 10 random seeds.
We first test whether the causal scores align with the taxonomy's 's predictions about performance. Specifically, the taxonomy predicts that, when ablated, attention heads scoring highly facilita- tion, irrelevance, or interference should decrease, leave unchanged, or increase the model's 's task performance, respectively. To test this, and evaluate the model using the G+ matrix while toggling each head to or in order of its score. While both G+ and G match the context in which scores were computed, more heads, providing a more interpretable baseline for ablation. We then compare the retained and ablated masks by the model's log-probability of the target sequence, expecting the resulting change in log- log-roobailiity to follow the predicted pattern. As shown in Figure these interventions match the predicted patterns: the difference in target proprrbbbbylity is negative when progressively ablating facilitating heads, near 0 when ablating irrelevant heads, and positive when ablating interfering heads, up until the set of interfering heads is exhausted.
heads in descending order by each causal metric
G+ as it retains
4.2 Distribution of causal roles
Having validated the causal scores using targeted ablations, next analyze how they distributed across models and tasks. Figure Ba shows that for each task, the distribution of head roles is highly consistent across all four model variants. This holds despite large differences in model size (IB to 8b) and training setup (pretraining, distillation, instruction tuning). We quantify these similarities by computing Pearson correlations between head scores across all model pairs for each task and causal metric, yielding 54 model pairs, all of which show high agreement with minimum correlation of 94.92% and an average of 99.2%. Across tasks, however, we observe notable differences, with the math dataset standing out in particular. For syntax and commonsense reasoning, most heads are irrelevant-63.1 -630% and 64.6% have irrelevance scores 2 0.5, respectively-witt only a sparse subset of facilitating heads (25.6% and 27.49 4% with facilitation scores > 0.5), suggesting that compact, redundant cireuits are sufficient for these tasks. contrast, mathematical reasoning activates much larger fraction of facilitating heads: 52.6% have facilitation scores > 0.5, while only 39.0%

---

(a)
(b)
L3.2-1B
L3.2-38
L3.2 3B-1
L3.1-80
Math (Alwayy))
Math (Any)
100%
Syntax
50%
Proportion x Score
0%
100%
Comman Sense
50%
Layar
Syntax (Always)
Syntax May)
0%
100%
Math
50%
0%
0.5
Score
10
15
0
10 15 20
Facilitation
Irrelevance
Interference
Figure 3: CHG score distributions and consistency. (a) Empirical cumulative distribution of CHG scores across all attention heads, showing the proportion heads with scores below given threshold for facilitation, irrelevance, and interference. (b) Aggregated CHG scores on L3.2-30 1.,.2-3II, where red and green color channels represent interference (1 G+) and facilitation (G-)), respectively. Colors are combined using RGB rules: black indicates irrelevance (low in both), and yellow indicates both facilitation and interference (high both). Always aggregates using the minimum across seeds (highlighting consistent effects); Any uses the maximum (highlighting any effect across seeds).
Table 2: Percent of heads with facilitation (F) or interference (N) scores 2 0.5 across seeds (always) in at least one seed (any).
L3.2-1B
L3.2-3B
L3.2-3BI
L3.1-8B
Task
Agg.
F
N
F
N
F
N
F
N
Always
1.2
0.2
1.5
0.1
0.7
0.0
1.4
0.0
Syntax
Any
72.1
57.2
67.9
51.3
72.8
56.1
68.5
59.2
Always
3.9
0.0
4.5
0.0
3.0
0.0
18.7
0.6
Common Sense
Any
56.6
41.0
75.4
52.4
68.2
55.7
60.3
22.2
Always
38.3
0.4
24.6
1.3
18.3
0.1
25.3
1.0
Math
Any
81.1
26.0
75.1
13.8
74.4
47.2
75.0
21.2
irrelevant, likely reflecting the task's 's higher complexity and need for broader sub-circuitry to support multi-step. latent computations.
It is also worth noting that, across all tasks, 84.0% of heads are marked as facilitating interfering (score 2 0.5) at least one seed, yet only a small fraction are consistently facilitating or interfering across all seeds (Figure Bb). In syntax and commonsense tasks, most models have fewer than 5% of heads that are always facilitating and virtually none that are always interfering (Table 2}. contrast, math reveals more rigid and consistent circuitry, with up 38.3% of heads consistently facilitating and 1.3% consistently interfering. These patterns suggest that individual attention heads may not have modular, context-ing mntexxt-independett roles, but instead participate in flexible ensemble of overlapping sub-circuits, in which their function depends on the configuration of others [53]
4.3 Comparison with causal mediation analysis
CMA, like CHG, aims to identify attention heads that facilitate task execution, though it does in more hypothesis-drivee manner. Framed in signal detection terms. CMA and CHG are complementary. CMA exhibits high precision but relatively low sensitivity: while many facilitating heads may go undetected (false negatives), those it does identify are reliably task-relevant (few false positives). Conversely, CHG is biased toward sensitivity over precision. This suggests that heads identified by CMA should also be identified (as showing strong facilitation) under CHG. We test this by comparing
7

---

CHG to the results of two former studies using CMA, replicating their methods to identify attention heads with specific computations: heads that encode task information in function vectors [2] and heads that perform symbolic reasoning [3
For function vectors, 'country-captall, Singlish-Frenchl "present-past' and 'singular-plure
the six in-context learning tasks used in 12 "antonym', 'capitalize',
Each prompt is presented in an in-context learning (ICL) [46] format consisting of 10 input-output examples using "Q: Xin A: Y" template, followed by a query to be ams amswered. To perform CMA, we corrupt the prompi by randomly shuffling example outputs to induce mismatched pairs, then patch individual head outputs with clean activations to identify which heads recover performance interpreettetgggggg high recovery as evidence of causal mediation.
We apply similar logic to symbolic reasoning tasks from 13, where the goal is to generalize abstract identity rules such ABA ("flow"Started Staated'flow") ABB same CMA procedure used in 3 to identify the three-stage symbolic processing mechanism that was reported: (1) symbol abstraction heads that abstract symbols ("A" "B") away from the actual tokens in the in-context examples; (2) symbolic induction heads that operate over the abstracted symbols to induce the symbol for the missing token the query; (3) retrieval heads that retrieve the actual token based on the induced symbol to complete the query. To screen heads of each type, we construct prompt pairs which either the same token is assigned to different symbols ("A" or "B") or tokens are swapped while preserving the same rule, and patch activations at certain token positions between them. Attention heads that steer model behavior towards specific hypotheses about the three head types after patching (either converting the abstract rule or altering the actual token) are labeled model.
"Started"Startee
We deploy the
as mediating. We conduct all experiments on the Llama-3.2
(a)
Antonym
Capita lize
Country-Capital
1.00
1.00
0.75
0.50
0.75
Faciiitation
0.25
Facilitaiio
0.00
0.50
English-French
Present-
Singular- Plural
1.00
0.75
0.25
0.50
0.25
0.00
0.00
0.0
0.2
0.4
0.60.0
0.2
0.4
0.60.0
0.2
0.4
0.6
CMA Score
AIE Score
Insignificant
Induction
Head Type
Abstraction
Retrieval
<30 30
Figure 4: Task- fask-facilitation scores versus (a) average indirect effect for function vector tasks and (b) CMA scores for symbolic reasoning tasks, showing significant heads by type (abstraction, induction, retrieval) and using the maximum CMA score across types for insignificant heads.
As predicted, CMA-identified heads tend to exhibit high facilitation scores under CHG in both domains (Figure 4). To quantify this, we compare the CHG facilitation scores of CMA-identified heads-those with three standard deviations above the mean in function vector tasks or with statistical significance ABA/ABB tasks B to the remaining ones. Since facilitation and irrelevance depend the specific sufficient circuit identified by CHG, a head may appear irrelevant in one but facilitating in another if multiple circuits exist. To account for this, we fit 10 CHG masks per function vector task and 20 per ABA/ABB task, and compute each head's 's maximum facilitation score across runs-capture capturing whether it participates in any sufficient circuit. We find significantly greater facilitation among mediating heads in both the function vector tasks (f(23.05) = 8.52, p 10-8) and the ABA/ABB tasks (t(53.77) = 11.18, p < 10)1 supporting the relationship between CMA and CHG-identified task relevance.
4.4 Contrastive Causal Head Gating
The results above indicate that CHG effectively distinguishes among facilitating, irrelevant, and interfering attention heads. However, exploratory method, it lacks the granularity to charac-
8

---

terize the specific functions of these subnetworks. For instance, consider the "antonym" task from Section 4.3 presented an in-context learning (ICL) format with examples and single-word response, defined in [2] To perform this task successfully, the model must not only generate the appropriate antonym given word, but also infer the task itself from the 10 input-output pairs in the prompt. Thus, minimal circuit of task-facilitating heads will contain both those involved in task inference and those involved antonym production, and CHG cannot distinguish between the two. This becomes more pronounced as task complexity increases, in the OpenMathInsttruc22 dataset, where the minimal circuit must jointly support diverse sub-tasks, including English comprehension, mathematical reasoning, chain-of- of-in-of-thought processing, and LaTeX generation.
To address this, we introduce a simple extension of CHG that not only identifies facilitating heads for given task but also isolates the sub-cireuit responsible for a particular sub-task. We generate parallel variants of the same task that share all features except for a controlled difference in the required operation, allowing to isolate the corresponding sub-circuits. In doing so, we take a step toward hypothesis-drive approach, decomposing the task into sub-steps while remaining agnostic to the mechanistic implementations For example, the antonym task can be constructed as ICL task using the default format from [12. or as an instruction-folloviii task where the model presented with the task description "Given an input word, generate the word with opposite meaning". By comparing the resulting attention circuits, we can disentangle components responsible for task inference from those involved in antonym generation.
Furthermore, rather than simply upplying CHG to each version and directly comparing the results, we propose a combined approach that fits single mask with joint objective to forget one variant of the task while retaining the other, so that the resulting gate matrix suppresses heads uniquely necessary for one variant but dispensable for the other:
M(G 2)
Σ
log P(UF IF) log P(yn xR) Σ FJ
(Gr.)
(2)
(ER.DR.EF-BE)
where log Ply r) denotes the log- probability of target sequence y given prompt x under model Mg with gating matrix G, the sum ranges over matched tuples (IR, VR. Zy, UF) of the retention and forget variants that differ only in task formulation, and À >> 0. To stabilize the gradient, we clip the miverse-simmoid as Eq. [ as well as the difference in log- bog-probability.
We evaluate this method using the six function vector tasks from Section leveraging the natural language task descriptions provided in [12 to construct instruction-based variants. For each problem, replace the 10-shot word-pair examples with prompt containing the task instruction and single example. We then fit the contrastive causal head gating (CCHG) mask to forget the ICL variant of five tasks while retaining the instruction- based format, holding out the sixth task for evaluation. task inference from examples, instruction- www.etiion..ooowinnnoommm and task execution are indeed mediated by separable circuits, this analysis should disable example- sxample-based based generalization while preserving instruction baasee performance. We perform our experiments both directions (forgetting ICL while retaining instruction mnsructtion-ollooinng and vice versa), using each of the six tasks as the held-out evaluation task. All experiments were conducted on the L.L.MA-322-- 1-MMA--.2-3B--A. 2-3B-Insttruct model.
Antonym
Country Capital
FranFrhhh
Prosont-Past
arngular-Plural
100%
25%
0%
100%.
75%
50%
Gated
Default
Default
Gated
Gated
Prompt
Instruction
Figure 5: Task accuracy under CCHG. Columns indicate held-out evaluation tasks and rows indicate the retained prompt format. Bar color shows the evaluation prompt format. "Default" and "gated" indicate whether CCHG is applied during evaluation. Error bars indicate 95% CI.

---

As shown in Figure 5 the CCHG masks generalize to the held-out task. When the model is induced to forget task inference from ICL examples across five tasks, its target task accuracy drops zero on the ICL. variant of the held-out task while in most cases remaining well above zero zero-and often close to the unablated haseline-on on the instruction-ren basdd variant. similar pattern emerges when forgetting is applied using the instruction-base format: performance collapses on instruction prompts while generally remaining intact for example-based ones.
Interestingly, while degradation is often small for the retained prompt format, this pattern is not consistent across all tasks. For example, when the gating matrix is fitted to retain ICL and forget instruction- forpttion...oooiiig.. the "singular-plun 92%) but complete failure on instruction prompts (98% - 0%). When this setup reversed fitted to retain instruction-frr -oollwiin and forget ICL- accurreyy on ICL drops from 98% to 0%, while instruction accuracy drops more modestly (98% + 21%). Across the tasks, 3 ('country-capital" "English-French", "present-past') remain robust as held-out tasks under instruction prompts, and 4 (' ('capitalize'. "English-French', French', 'present-past'
task shows only a small drop in ICL accuracy (98% -
iangular-plural" do so under ICL prompts.
Thus, our results indicate that the circuits for instruction following and ICL may be separable at the head level. However, this separability also depends on the task, suggesting that task execution circuits may share heads with those used for task understanding and representation.
5 Discussion
In this work, we introduced Causal Head Gating (CHG), flexible and scalable method for identifying causally relevant attention heads in large language models. CHG assigns each head graded score for facilitation, interference, or irrelevance based on its effect on task performance, going beyond correlational or observational analyses. These scores predict performance changes under targeted ablations, confirming that facilitation, interference, and irrelevance scores capture causal impact. Crucially, it does so using next-token prediction alone, thereby avoiding reliance on labeled data or handcrafted prompts, making it broadly and casily applicable. Moreover, CHG requires no finetuning or auxiliary decoder model, and introduces only one parameter per head, allowing it to run in minutes even on billion-scale models. To validate our method, demonstrated that existing works within the mechanistic interpretability literature successfully corroborate our findings using, and that the ICL and cestructionfoolowiinn circuits revealed using contrastive CHG successfully generalize across tasks.
Interestingly, across the range of models and tasks we investigated, we observed that attention heads form task-sufficient sub-circuits with low overlap. Moreover, a single head may vary in its relevance across multiple runs depending on which others are active, reflecting the distributed and context-dependent nature of computation in LLMs, and in rare cases, a head may even receive low G+ but high G' scores within the same run. We hypothesize that this variability reflects an interaction -deeedden landscape in which causal roles shift with circuit configuration. While these complexities may appear messy, we view them as a strength of CHG, revealing the redundancy and interdependence that underlie emergent model behavior. Because CHG is highly scalable, it can be repeatedly applied to estimate distributions over gating values, providing a bootstrapped view of redundant and contingent sub-circuits with greater fidelity to the model's underlying dependency structure.
While CHG provides a lightweight and scalable approach for exploratory analysis, requiring only a dataset and no model finetuning supervision. it is not designed to reveal the precise computations performed by individual heads. Instead, CHG offers a complementary first-pass diagnostic tool that identifies candidate heads sub-circuits with consistent causal influence, guiding where more granular, hypothesis-drins driven methods such as causal mediation or activation patching can be applied. In this way, CHG provides a practical entry point into large-scale causal interpretability. mapping functional dependencies that subsequent analyses can examine in greater detail.
We hope that our work encourages further exploration of causal structure in language models as foundation for more mechanistic understanding. Future work may build on these tools to develop circuit-level explanations of how models implement complex behaviors.
10

---

Acknowledgments and Disclosure of Funding
We thank Declan Campbell and Alexander Ku for helpful discussions, and Legasse Remon for assistance with dataset organization.
Jonathan Cohen was supported by the Vannevar Bush Faculty Fellowship, sponsored by the Office of Naval Research.
The authors declare no competing interests.

---

References
[1] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Tlge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. GPT-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
[2] Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng. Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412. /9437, 2024.
[3] Aaron Grattafioni, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint 2rXiv:2407.21783, 2024.
[4] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint (tti::200..0255555 2021.
[5] Alexander Wei, Nika Haghtalab, and Jacob Steinhardt. Jailbroken: How does LLM safety training fail? Advances in Neural Information Processing Systems, 36:80079-80110, 2023.
[6] Laura Weidinger, Joslyn Barnhart, Jenny Brennan, Christina Butterfield. Susie Young, Will Hawkins, Lisa Anne Hendricks, Ramona Comanescu, Oscar Chang, Mikel Rodriguez, al. Holistic safety and responsibility evaluations of advanced ai models. arXiv preprint a40iv:2404. /4068, 2024.
[7] Ian Tenney, Dipanjan Das, and Ellie Pavlick. BERT rediscovers the classical NLP pipeline. arXiv preprint arXiv: /905.05950 2019.
[8] Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Conerly, Nick Turner, Cem Anil, Carson Denison, Amanda Askell, et al. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2, 2023.
[9] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
[10] Gemma Team, Aishwarya Kamath, Johan Ferret, Shreya Pathak, Nino Vieillard, Ramona Merhej, Sarah Perrin, Tatiana Matejovicova, Alexandre Ramé, Morgane Rivière, et al. Gemma 3 technical report. Xrviv preprint arXiv:2503.19788. 2025.
[11] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. hmpps://transformeercccitttt wttp:/rrrrriuiuuuuc2222e http://aiii..auu/////aaaaeeeoeee /wwa///wro//iddee..neoooo
[12] Eric Todd, Millicent Li, Arnab Sen Sharma, Aaron Mueller, Byron C Wallace, and David Bau. Function vectors in large language models. arXiv preprint arXiv:2310.15211 2023.
[13] Yukang Yang, Declan Campbell, Kaixuan Huang, Mengdi Wang, Jonathan Cohen, and Taylor Webb. Emergent symbolic mechanisms support abstract reasoning large language models. arXiv preprint arXiv:2502.20332. 2025.
[14] Ian Tenney, Patrick Xia, Berlin Chen, Alex Wang. Adam Poliak, R Thomas McCoy, Najoung Kim, Benjamin Van Durme, Samuel R Bowman, Dipanjan Das, al. What do you learn from context? Probing for sentence structure in contextualized word representations. arXiv preprint DrXii:/905. 063/6,66660666 2019.
[15] John Hewitt and Christopher D Manning. A structural probe for finding syntax in word representations. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pages 4129-4138, 2019.
12

---

[16] Adly Templeton, Tom Conerly, Jonathan Marcus, Jack Lindsey, Trenton Bricken, Brian Chen, Adam Pearce, Craig Citro, Emmanuel Ameisen, Andy Jones, Hoagy Cunningham, Nicholas Turner, Callum McDougall, Monte MacDiarmid, C. Daniel Freeman, Theodore R. Sumers, Edward Rees, Joshua Batson, Adam Jermyn, Shan Carter, Chris Olah, and Tom Henighan. Scaling monosemanticity: Extracting interpretable features from Claude 3 Sonnet. Transformer Circuits Thread, 2024.
[17] Andrew Lee, Xiaoyan Bai, Itamar Pres, Martin Wattenberg, Jonathan K Kummerfeld, and Rada Mihalcea. A mechanistic understanding of alignment algorithms: A case study on dpo and toxicity. arXiv preprint arXiv:2401.0199 2024.
[18] Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi- head self-attention; Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint arXiv:/905.09415 2019.
[19] Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter. Zoom in: An introduction to circuits. Distill, 5(3):e000240001, 2020.
[20] Geoffrey E Hinton. Learning distributed representations of concepts. In Proceedings of the Annual Meeting of the Cognitive Science Society, volume 8, 1986.
[21] Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield- Doddd, Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah. Toy mod- els of superposition. Transformer Circuits Thread, 2022.
circuits
pub/2022/toy htt//////toooood////eeeeee mooell/index, html
[22] Kayson Fakhar and Claus C Hilgetag. Systematic perturbation of an artificial neural network: A step towards quantifying causal contributions in the brain. PLOO Computational Biology, 1010220,, 2022.
[23] Tyler Giallanza, Declan Campbell, Jonathan D Cohen, and Timothy Rogers. An integrated model of semanties and control. Psychological Review, 2024.
[24] Melanie Mitchell. Complexity: A guided tour. Oxford University Press, 2009.
[25] Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: a circuit for indirect object identification in GPT-2 small. arXiv preprint arXiv:2211.0000
2022.
[26] Jason Wei, Xuezhi Wang. Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoe V Le, Denny Zhou, et al. Chain-of thoutht prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35:24824-283 24837,, 2022.
[27] Yonatan Belinkov. Probing classifiers: Promises, shortcomings, and advances. Computational Linguistics, 48 4811):207-219 2022.
[28] Tiago Pimentel, Josef Valvoda, Rowan Hall Maudslay, Ran Zmigrod, Adina Williams, and Ryan Cotterell. Informatiin-thheortiec probing for linguistic structure. arXiv preprint arXiv:2004.03061 2020.
[29] John Hewitt and Percy Liang. Designing and interpreting probes with control tasks. arXiv preprint arXiv: 1909. 03368, 2019.
[30] Elena Voita and Ivan Titov. Information Information-thhorrtiee theoretic probing with minimum description length. arXiv preprint arXiv:2003. /2298, 2020.
[31] Abhilasha Ravichander, Yonatan Belinkov, and Eduard Hovy. Probing the probing paradigm: Does probing accuracy entail task relevance? arXiv preprint arXiv:2005.007 2020.
[32] Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey. Sparse autoen- coders find highly interpretable features in language models, arXiv preprint arXiv:2309.086/00
2023.
13

---

[33] Jesse Vig, Sebastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart Shieber. Investigating gender bias in language models using causal mediation analysis. Advances in Neural Information Processing Systems, 13:12888-12401 2020.
[34] Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in GPT. Advances in Neural Information Processing Systems, 35:17359-17372 2022.
[35] Zifan Zheng, Yezhaohui Wang, Yuxin Huang, Shichao Song, Mingchuan Yang, Bo Tang, Feiyu Xiong, and Zhiyu Li. Attention heads of large language models. Patterns.
[36] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jenry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint 10.14:2110.14168, 2021.
[37] Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint 10Xiv:2003.03874,, 2021.
[38] Shubham Toshniwal, Wei Du, Ivan Moshkov, Branislav Kisacanin, Alexan Ayrapetyan, and Igor Gitman. Openmathinstructt-: Accelerating ai for math with massive open-source instruction data. arXiv preprint arXiv:2410.01550 2024.
[39] Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? Advances in Neural Information Processing Systems, 32, 2019.
[40] Jiaoda Li, Ryan Cotterell, and Mrinmaya Sachan. Differentiable subset pruning transformer heads. Trunsactions of the Association for Computational Linguistics, 9: 0:1442-1459 2021.
[41] Mengzhou Xia, Zexuan Zhong, and Danqi Chen. Structured pruning learns compact and accurate models. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pupers), pages 1513-1528, 2022.
[42] Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth on demand with structured dropout. arXiv preprint arXiv: 1156, 2019.
[43] Hassan Sajjad, Fahim Dalvi, Nadir Durrani, and Preslav Nakov. On the effect of dropping layers of pre-trained transformer models. Computer Speech & Language, 77:101429, 2023.
[44] Shwai Hc, Guoheng Sun, Zheyu Shen, and Ang Li. What matters in transformers? not all attention is needed. arXiv preprint arXiv:2406. 15786, 2024.
[45] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pages 4171-4186, 2019.
[46] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, 33:18 13:1877-1901, 2020.
[47] Nicola De Cao, Michael Schlichtkrull, Wilker Aziz, and Ivan Titov. How do decisions emerge across layers in neural models? interpretation with differentiable masking. arXiv preprint Sttv:20004149992. 2020.
[48] Fangcong Yin, Xi Ye, and Greg Durrett. Lofit: Localized fine-tuning Ilm representations. Advances in Neural Information Processing Systems, 37:9474-9506, 2024.
[49] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv: 16/ 144, 2016.
[50] Chris J Maddison, Andriy Mnih, and Yee whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv:1611.00711 2016.
14

---

[51] Aarohi Srivastava, Abhinav Rastogi, Abhishek Rao. Abu Awal Md Shoeb, Abubakar Abid, Adam Fisch, Adam R Brown, Adam Santoro, Aditya Gupta, Adria Garriga-Alonso, al. Beyond the imitation game: Quantifying and extrapolating the capabilities of language models. arXiv preprint arXiv:2206.0461.1 2022.
[52] Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. Commonsenseeqa: A ques- tion answering challenge targeting commonsense knowledge. arXiv preprint arX1v:1811.00937, 2018.
[53] Jack Merullo, Carsten Eickhoff, and Ellie Pavlick. Talking heads: Understanding inter-layer communication in transformer language models. Advances in Neural Information Processing Systems, 37:61372-61418, 2024.
[54] Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint 11X.v:1412.6980, 2014.
15

---

Technical Appendices and Supplementary Material
6.1 Datasets
For each dataset, we split the full set into three partitions: an example set, training set, and validation set. Example problems were selected from the top K shortest prompt-solution sequences after tokenization. One example was randomly drawn from the example set to be included in each training/validattonn prompt help align model responses with the task format. For multiple-choice datasets, answer options were randomly shuffled and labeled with capital letters (A, B, C, ...,, and the target answer was the correct letter.
OpenMachInstrue12 We used the IpenMachInstruct- ttraitnnnnii 1M subset. We filtered out problems marked as having no solution, removed duplicate prompts (even if their solutions differed), and retained the 55,050 shortest prompt-solution pairs by total tokenized length. From this, we selected 50 examples, 50,000 training problems, and 5,000 validation problems. Each prompt began with the instruction: "For each problem, explain your reasoning step by step and use LaTeX for all mathematical expressions. Indicate your final answer using \boxed!....."
Commonsense(AA We selected 10 problems for the example set, then split the remaining data into a 90% / 10% training/validation split.
BIG-Bench syntax We included all tasks labeled 'syntax' in BIG-Bench: 'linguistic mappings', 'tense", and 'subject-verb--tee verb--rreeb-aeet...... The 'linguistic mappings' category consisted of five subtasks, each with its own instruction:
Past tense: "Convert the verb its past tense form."
Plural: "Convert the noun to its plural form."
Pronoun replacement: "Replace the repeated name with the correct pronoun."
Question formation: "Convert the statement into a yes/no question." "
Sentence negation: "Convert the statement into negative sentence."
The 'tense' task used the instruction: "Modify the tense of a given sentence."
The 'uubject-verb-agreeee that agrees with the subject of the sentence."
task used the instruction: "Choose the grammatically correct verb form
Each task or subtask was treated independently for splitting and prompt generation. We allocated 10 examples per subtask, with 90% / 10% split over the remainder into training and validation. Example problems used in prompts were always drawn from the same subtask as the target problem.
Function vector tasks We included six tasks: 'antonym', 'capitalize', 'country-capital' 'english- french', 'present-past' and 'singular-plural.. Each task was used two formats: 10-shot in- context learning (ICL. prompts with 10 input-output pairs, and instruction-based prompts using task descriptions from 12]
Antonym: "Given an input word, generate the word with opposite meaning."
Capitalize: "Given an input word, generate the same word with a capital first letter."
Country-Capital: "Given country name, generate the capital city."
Fagllihh-French: "Given an English word, generate the French translation of the word."
Present-Past: "Given verb in the present tense, generate the verb's simple past inflection." Singular Plurrl:-PPuurl:: "Given a singular noun, generate its plural inflection."
We allocated 10 examples per task, and split the remaining data into 90% training and 10% validation. Example problems used in prompts matched the format (ICL or instruction) of the task being evaluated.
Symbolic reasoning (ABA/ABB) : We procedurally generated symbolic reasoning prompts fol- lowing the AMB^A and AABAB templates from [13]. using 4 in-contex examples per prompt. Each prompt was generated by selecting 10 random tokens-8 assigned to the examples and used in the query. We used individual tokens rather than full words, since multi-token words often behave similarly: once the first token is generated, the model tends to complete the rest automatically, reducing the task to token-level pattern recognition.
16

---

6.2 Training details
Causal head gating For each model and task, we first fit CHG matrix with = 0 for 500 gradient updates with batch size of 64 samples. G was initialized with random values sampled uniformly between and We used the Adam optimizer [54] for optimization using an initial learning rate of 0.1 with linear decay that terminates with learning rate of 0.01. After fitring G with À 0, we fit G and G using as the initial conditions and +0.1 for 500 gradient updates with an initial learning rate of 0.5 and terminal learning rate of 0.1. We clipped the regularization term at +
Contrastive causal head gating For each model and task pair, we fit CCHG matrix G with -0.1, clipping the regularization term at and the log-probability difference at 5. We fitted G over 500 gradient updates with a batch size of 64 using the Adam optimizer with an initial learning rate 0.1 with linear decay that terminates with a learning rate 0.01.
Hardware and compute
For all our experiments, we used 128 GB of CPU RAM and a single Nvidia H100 GPU at time.
Each run of CHG (1,500 gradient updates) took between 15 minutes and hour, depending on the model and dataset. Each run of CCHG (500 gradient updates) took approximately 5 minutes. We estimate that all experiments reported in this paper can completed under 100 GPU hours. Preliminary or failed experiments required negligible additional compute and are not included in the total.
17

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes]
Justification: We claim in the abstract and introduction that CHG offers causal insight into individual attention heads in LLMs, which substantiate using ablation analysis and comparison to CMA. We also claim that instruction following and ICL are separable at the head level, which we show using CCHG. Lastly, we claim that LLMs contain multiple sparse sub-circuits that are individually sufficient for different tasks, which we show in our consistency analysis.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers,
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
Justification: The paper discusses key limitations in the Discussion section, including CHG's inability to explain why heads matter, occasional divergence between G+ and G and the context-dependent variability in head roles across runs. Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [NA]
18

---

Justification:
Our paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems. The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes]
Justification: We provide detailed descriptions of the CHG algorithm, datasets, model variants, and evaluation methods (e.g., ablation, CMA comparisons). The precise methods for reproducing the CMA results are better described in the original papers. Hyperparameters and additional procedural details are included in the supplementary materials. Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
also be provided via detailed
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
19

---

The accompanying repository can be found at sttpps:/githin thub. com/ andrewnam/ causal cawsal_headdgaaaee _heai_gattngg We also note that the models and datasets used in our paper are publicly available. Guidelines:
Justification:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
Sorriides//odeSSumisssst//
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
llaips.cc/publl hwaplli//guieeee guides
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized
versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The training and evaluation details, hyperparameters, and optimizer settings can be found in the supplementary materials.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The paper reports (-test results for CMA comparisons and includes 95% confidence intervals in the CCHG evaluation plots
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi-
dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula.
call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
20

---

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative
error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: Details on compute resources are provided in the supplementary materials. Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster,
cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics Answer: [Yes]
Justification: We have reviewed the NeurIPS Code of Ethics and confirm that all aspects of our research fully comply.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. there a special consid-
eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The paper presents foundational interpretability methods without direct ap- plication deployment, and we do not foresee societal impacts resulting from its current scope.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal
impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations
(e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
21

---

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks, as it uses only publicly available models and datasets and does not release any new high-risk assets.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All assets are cited in the paper. Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, if possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [NA]
Justification: We do not release any new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
22

---

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing or research with human subjects. Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing or research with human subjects. Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [Yes]
Justification: Our study is applied directly to LLMs and central to our methodology. Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non standrdd components.
Please refer to our LLM policy http:p////errrsss.cc/oeeeeeeeee for what should or should not be described.
23


Reviews:
[{"review_id": "dxOKz48ztA", "reviewer": "Reviewer_foUj", "summary": "This paper introduces Causal Head Gating (CHG), a novel framework for evaluating the functional roles of attention heads in large language models (LLMs). Unlike prior approaches, CHG does not require hypothesis-driven datasets to analyze head behavior on specific tasks.\n\nThe method introduces gating coefficients, which control how much each attention head contributes to the model’s output. CHG is trained post hoc on frozen models using regularization to encourage sparsity, making it both interpretable and computationally efficient. By optimizing these gates, CHG segments heads into facilitating, interfering, or irrelevant categories.\n\nThe authors show that multi-head attention behaves as a flexible ensemble of overlapping sub-circuits, rather than as a collection of modular, context-independent components.\n\nCHG is validated through multiple strategies conducted across a range of LLMs and tasks : the authors perform head ablations to test causal impact, replicate prior findings of Causal Mediation Analysis, a standard framework for estimating causal effect in neural network, and demonstrate strong correlation with CHG score and CMA-derived metrics.\n\nFurthermore, the paper proposes Contrastive CHG, an extension that enables more fine-grained analysis by comparing head importance across different input conditions (e.g., retaining vs. forgetting context)", "questions": "How reproducible are the gating coefficients across different CHG runs with varied initializations or random seeds? Given this variability, what is a recommended number of CHG fits to achieve a robust or sufficiently consistent understanding of head roles for a given model and task? Do the authors observe convergence to similar gate configurations in tasks where head roles are assumed to be stable (e.g., math reasoning)?\nCould they provide more insight into the nature or frequency of the divergence of G+ and G- matrices ? Are there specific tasks, model sizes, or layers where this divergence is more pronounced? Have they started to explore what these \"richer head interactions\" might entail, and how might CHG be extended to capture them?\nCould the authors elaborate on the nature of the \"overlapping sub-circuits\"? Does it imply that different sets of heads can perform the same function (redundancy), or that a single head can contribute to multiple distinct functions depending on the context of other active heads? Can the authors speculate on the broader implications of this low modularity for future work in model compression, fine-tuning, or designing more interpretable architectures?\nWould the authors consider extending their evaluation to include (i) coreference or long-context memory tasks (e.g., LAMBADA, WinoGrande), which are directly relevant to the retention/forgetting dynamics analyzed in Section 4.4, and (ii) instruction-following or chain-of-thought reasoning tasks (e.g., BBH, MMLU), where CHG may reveal distinct subcircuit behavior associated with in-context learning? Demonstrating CHG’s applicability in these settings, if feasible, would significantly strengthen the paper and would warrant a higher evaluation score. While the focus on LLaMA-3 models is reasonable for consistency, showing that CHG generalizes beyond a single model family would further reinforce its broader relevance.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths :\nCausal, Mechanistic Insight into Attention Heads\nCHG offers causal rather than just correlational insight into attention head function, identifying heads as facilitating, interfering, or irrelevant. It reveals novel findings about Transformer internals, such as the presence of sparse, sufficient subcircuits, low modularity of head roles, and the separability of instruction following from in-context learning.\n\n\nScalable, General, and Practical Framework\nCHG introduces only one learnable parameter per head, operates on frozen models, and fits in a small amount of time, even on billion-parameter LLMs. CHG also addresses limitations of previous methods like probing classifiers (which rely on labels) and CMA (which requires hand-crafted prompts and is less scalable) and applies directly to standard next-token prediction.\n\n\nRobust Evaluation Across Models and Tasks\nThe framework is empirically validated across a diverse set of LLMs and tasks, demonstrating consistent segmentation of head roles.\n\n\nTransparency and Reproducibility\nThe paper is methodologically rigorous, reporting t-tests and confidence intervals. The authors provide detailed algorithmic and experimental descriptions and plan to release the code publicly.\n\nWeaknesses :\nLimited Insight into Why Heads Matter or What They Compute \nA significant limitation, as directly stated by the authors, is that CHG \"does not identify why they matter or what they compute\". While it effectively classifies heads based on their causal impact, it offers \"limited insight into underlying mechanisms\". This means CHG is best suited for \"exploratory analysis alongside hypothesis-driven tools like CMA\" rather than providing a complete functional explanation on its own—somewhat undermining its benefits from hypothesis-driven, annotated methods.\n\n\nOccasional Divergence of G+ and G- Matrices \nThe authors note that, albeit rarely, the G+ (encourage retention) and G- (encourage removal) matrices can \"diverge\". This divergence \"possibly reflect richer head interactions\" that the current CHG framework might not fully capture. This points to a potential area where the model's complexity exceeds the current interpretive power of CHG. \n\nLimited Dataset Coverage\nWhile the authors evaluate CHG on a diverse set of tasks, the dataset selection omits coreference and memory retention tasks (e.g., WinoGrande, LAMBADA), which would have naturally extended the contrastive CHG experiments in Section 4.4. Additionally, instruction-following and chain-of-thought reasoning tasks (e.g., BBH, MMLU, GSM8K-CoT) are absent, despite CHG’s potential relevance to analyzing in-context learning and reasoning subcircuits. Testing on a reasoning-heavy dataset like BBH or MMLU would further strengthen the paper’s claims by extending CHG’s applicability to another important task class, but this omission does not undermine the core contributions.\n\nClarity Issues in Key Sections\nSome parts of the paper, particularly Section 4.3, are densely written and conceptually compressed, making it difficult for readers unfamiliar with prior CMA work to follow the experimental setup and significance of the results.\n\nWhile the paper lacks evaluation on reasoning and memory tasks, and CHG offers limited insight into what heads compute, these limitations do not necessarily undermine the core contributions. The framework is well-motivated and demonstrates potential for broader applicability. \nI lean toward acceptance, as the paper offers a useful contribution to LLM interpretability, though I would defer to further discussion if others view the omissions as more fundamental. Additional experiments on reasoning tasks could strengthen the case and would warrant a higher score.", "was_revised": true, "final_justification": "After considering the rebuttal, I maintain that the paper presents a solid and timely contribution to interpretability in LLMs. The authors provided clear responses to most of my concerns, especially regarding the variability of CHG outputs, the interpretation of gating coefficients, and the justification for focusing on validated datasets. Their framing of CHG as a bootstrapping estimator clarifies its intended use and scope.\n\n**Resolved issues:**\n- Clarified the nature of variability across CHG runs and the intended interpretation of such variation.\n- Justified the choice of datasets and models, framing them as grounded in prior validated findings.\n- Provided insight into overlapping subcircuits and their implications for head interaction and modularity.\n\n**Unresolved issues:**\n- No new empirical evaluation on reasoning or memory tasks.\n\nWhile these gaps limit the paper's breadth, they do not undermine the core contributions. The method is well-motivated, technically sound, and experimentally validated. Based on the rebuttal, I have increased my score, though I would not oppose a more conservative decision if other reviewers weigh the empirical limitations more heavily.", "rebuttal": "We thank the reviewer for their thoughtful comments and provide the following clarifications and questions. While the rebuttal process does not allow for resubmitting revised content, we have aimed to address the reviewer’s concerns as clearly as possible and hope that our responses may help inform a reassessment of the work.\n\n1.\tRe: reproducibility across seeds – For each model and task, we report results from 10 different initializations, which often produce varying coefficients. This is illustrated in Figure 3b, where the “always” and “any” aggregation strategies reveal notably different patterns. However, we emphasize that consistency across seeds is not the objective. Variability in gating values reflects distributional variance, not irreproducibility. For example, if head A’s coefficient ranges from 0–1 across seeds, this suggests its effect is contingent on the configuration of other heads. In one run, head A may be suppressed because head B is not; in another, A may be active because head C is suppressed. This indicates that head A participates in a complex, interaction-dependent manifold rather than playing a strictly modular role.\n2.\tRe: recommended number of seeds – CHG is best viewed as a bootstrapping method for estimating a distribution over gating values, rather than identifying a single optimal configuration. Thus, the number of CHG fits depends on the model, task, and head. Some heads may have sufficiently modular or task-necessary functions that they are invariant to the initializations, whereas others have higher interactions with other heads. Either result may be interesting, depending on the practitioner’s goals. If the objective is to capture the distribution of causal scores for each head irrespective of their inter-head interactions, 20–30 fits is likely sufficient, as this is equivalent to estimating distributions bounded between 0 and 1. For broader analyses using min and max values as used in this work, 10 is likely sufficient.\n3.\tRe: convergence property – In mathematical reasoning tasks, we observe that task-relevant heads tend to concentrate in mid-layers, rather than early or late layers (see Figure 3b). Additionally, CHG reliably identifies heads also found by CMA (see Section 4.3), supporting the conclusion that these heads have stable, convergent roles.\n4.\tRe: divergence of G+ and G− – These divergences are rare: as shown in Figure 1c, only about 10 out of nearly 700 heads exhibit divergence. As for the precise nature of these cases, we have left this for future investigation, as they are fairly rare edge cases.\n5.\tRe: overlapping sub-circuits – Yes, both proposed interpretations are consistent with our view. The variation in CHG coefficients across seeds for some heads suggests that head A may be suppressed in one configuration because head B is not, and vice versa—i.e., A and B form an XOR-like interaction. This is likely to occur if A and B both serve similar roles, so that once A is considered “facilitating,” then B can be considered “irrelevant.” In such cases, task relevance should be interpreted jointly rather than independently. This suggests that even at the level of individual attention heads, LLMs may encode distributed representations rather than learning isolated “grandmother heads.”\n6.\tRe: broader implications for modularity – We are interested in further investigating modularity, particularly distinguishing between redundancy (where multiple heads serve the same function and some can be pruned) and contingency (where a head’s role shifts based on others being suppressed). The former could suggest certain heads can be pruned without much general performance loss. The latter could suggest that the behavior of an LLM could be changed by pruning\n7.\tRe: extension to other datasets – We are certainly interested in extending CHG to other datasets and domains! However, in this paper, we chose to place special focus on tasks and datasets with task-facilitating heads that have been confirmed by other papers (which is why we favored tasks from Todd et al. and Yang et al.) for the sake of validating our method before applying it to broader settings. We hope that the reviewer would appreciate the merit of our more conservative approach in favor of empirical grounding.\n8.\tRe: generalization to other LLMs – While we have run preliminary experiments with other LLM architectures and have observed results consistent with what we have reported here, we were unable to incorporate them into the full write-up of the present manuscript. We hope that the robustness across multiple domains and validation across the results of other papers would sufficiently demonstrate the generalizability of our approach. \n\nWe appreciate the reviewer comments and questions. If our responses to the reviewer concerns are satisfactory, we hope that you would reflect this in a higher overall rating for our work."}, {"review_id": "Tg0kRNkLAn", "reviewer": "Reviewer_WNd7", "summary": "This paper proposes Causal Head Gating (CHG), a scalable method for identifying the causal roles of attention heads in transformer models without relying on labels. By fitting soft gates to heads during next-token prediction, CHG classifies them as facilitating, interfering, or irrelevant to task performance. The method reveals that large language models contain multiple sparse, distributed sub-circuits and that functions like instruction following and in-context learning rely on separable mechanisms.", "questions": "- How does the choice of initialization for the gate matrix, random, all ones, or all zeros, affect the final learned mask?\n - What stopping criteria are used when training the gate matrix?\n - In Figure 1(b), what causes the abrupt bifurcation of the irrelevant heads around the 500th gradient update?\n - The manuscript refers to “causal scores” multiple times, does this term refer to the final gate values after training, or something else?\n - Lines 193–196 state, “We quantify these similarities…,” but the explanation was unclear to me. Could you please elaborate on what is being quantified and how?", "limitations": "Due to optimization variability and the non-deterministic nature of head roles, multiple CHG runs may be needed to draw robust conclusions, diminishing its efficiency advantage over more straightforward methods like causal tracing and activation patching.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 1, "strengths_and_weaknesses": "**Strengths**:\n - Unlike many existing mechanistic interpretability methods, the proposed CHG technique is both efficient and scalable. It effectively localizes groups of attention heads that collectively contribute to task performance, patterns that are often missed when applying causal interventions to individual components in isolation.\n - The paper validates CHG by triangulating its results through ablation studies and causal mediation analysis.\n - Furthermore, it uses CHG to demonstrate that the attention heads responsible for in-context learning and instruction following are distinct.\n\n**Weaknesses**:\n - My primary concern with this work is its limited novelty. The proposed approach is quite similar to existing methods that use differentiable binary masks (DBMs) to identify model components encoding semantic information, particularly through counterfactual input pairs [1, 2].\n - Although the authors validate their method using ablation and causal mediation analysis, the paper lacks a comparative evaluation against established mechanistic interpretability techniques such as causal tracing, activation patching, and attribution patching.\n - The empirical finding that attention heads involved in in-context learning (ICL) and instruction following differ has already been suggested in prior work [3].\n - Moreover, I find the title somewhat misleading. It implies a broad focus on uncovering the roles of attention heads in various tasks, yet the paper largely sidesteps this issue. Only Section 4.4 addresses it directly, and even there, the emphasis is on distinguishing between ICL and instruction-following heads, rather than offering a general framework for identifying head-level functional roles.\n\n[1] Davies et al, “Discovering Variable Binding Circuitry with Desiderata”, 2023.\n\n[2] Prakash et al, “Fine-Tuning Enhances Existing Mechanisms: A Case Study on Entity Tracking”, 2024.\n\n[3] Davidson et a “Do different prompting methods yield a common task representation in language models?”, 2025.", "was_revised": true, "final_justification": "While the authors have provided details regarding the differences of the proposed technique and similar existing methods in the literature, an empirical comparative analysis is still missing, which is essentially to assess a methodological contribution. Therefore, while promising, I cannot recommend to accept this work.", "rebuttal": "We thank the reviewer for their thoughtful comments and provide the following clarifications and questions. While the rebuttal process does not allow for resubmitting revised content, we have aimed to address the reviewer’s concerns as clearly as possible and hope that our responses may help inform a reassessment of the work.\n\n1.\tSimilarity to prior work – While Davies et al. and Prakash et al. use activation patching with manually constructed counterfactuals or desiderata, CHG learns gating parameters over all attention heads using only next-token prediction—avoiding the need for labeled tasks or handcrafted inputs. This results in a scalable, label-free method for causal attribution that differs substantially in both mechanism and scope. We are happy to cite these papers in the Related Work section, but we believe their goals and techniques are distinct from ours. Additionally, we note that none of the cited papers use the term “differentiable binary masks (DBMs).” Could you clarify what specific technique you had in mind by “DBMs,” and how it applies to these references or to CHG?\n2.\tNovelty of instruction vs. ICL – We appreciate the reviewer highlighting the work by Davidson et al. However, their preprint was first posted on May 17, 2025—two days after the NeurIPS submission deadline. We therefore view this as an instance of independent, concurrent discovery rather than prior work, and we hope this will be taken into account when assessing the novelty of our contribution.\n3.\tComparison to other methods – We respectfully disagree with the assertion that we do not compare against established mechanistic interpretability techniques. Causal mediation analysis (CMA) is a structured form of activation patching, and both methods involve intervening on internal activations to identify causal roles. In Section 4.3, we benchmark CHG against CMA-based approaches and also replicate prior patching-style interventions.\nIf there are specific forms of causal tracing or attribution patching you feel would strengthen the comparison, we would welcome your suggestions.\n4.\tRole of attention heads – We agree that CHG does not reveal the precise function of each head in isolation and is best used in combination with techniques like CMA. However, it provides a general, scalable framework for identifying which heads are causally relevant to which tasks—offering a simple, hypothesis-free tool for mapping attention head function across models and domains.\n5.\tTraining details – As noted in the supplementary material (S3.1), we randomly initialize the gates. This helps uncover diverse configurations and avoids prematurely committing to a single solution. We train the gates for 500 gradient updates, which is sufficient for convergence in practice.\n6.\tBifurcation in Figure 1b – As explained in Section S3.2, we first fit the initial gating matrix G, then switch to fitting G+ and G− after 500 updates. The bifurcation occurs at this point because the optimization objective changes—specifically, the regularization term becomes active (λ ≠ 0).\n7.\tCausal scores – Yes, the causal scores refer to the fitted gate values, as shown in Table 1. The score for a facilitating head is G−; for interfering, it is 1 − G+; and for irrelevant, it is G+ × (1 − G−).\n8.\tSimilarity quantification – As described in the same paragraph, the similarity refers to the distribution of CHG values across different models on the same dataset. This is shown in Figure 3a, where cross-model correlations within a dataset reach 99.2%, supporting the observed alignment across plots in each row.\n\nWe appreciate the reviewer comments and questions. If our responses to the reviewer concerns are satisfactory, we hope that you would reflect this in a higher overall rating for our work."}, {"review_id": "dYVXecqD4g", "reviewer": "Reviewer_PMP1", "summary": "The authors present Causal Head Gating (CHG), a scalable method for interpreting the functional roles of attention heads in transformer models. The method involves learning a matrix of soft gates, one for each head, by optimizing the standard next-token prediction loss on a given task dataset, augmented with a special regularization term. They also introduce a contrastive variant to isolate circuits for specific task components.", "questions": "1. How should one interpret the \"causal role\" of a head that is \"facilitating\" in some runs and \"irrelevant\" in others? \n2. The gating matrix G is optimized for a specific task dataset D. Have the authors considered applying a G matrix trained on one task (e.g., mathematical reasoning) to a completely different task (e.g., commonsense QA)? This would provide evidence about the degree to which the identified \"circuits\" are task-specific versus representing more general computational primitives.\n3. Regarding the contrastive CHG experiment, the gating matrix G itself is quite expressive. How can we be certain that the method is truly discovering a pre-existing functional separation between ICL and instruction-following circuits, rather than G learning to create or simulate this separation to satisfy the contrastive objective?\n\nTypos: \n* Line 3: over heads -> overheads\n* Line 17: state-of-the art -> state-of-the-art \n* Line 46: interpretablity -> interpretability\n* Line 217: This suggests that that -> This suggests that", "limitations": "The instability and ambiguity of the \"causal roles\" assigned by CHG as mentioned in weaknesses.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The proposed method, CHG, is highly scalable and does not require manually crafted prompt templates or external labels beyond what is used for next-token prediction.\n2. The contrastive CHG extension is interesting, offering a way to start disentangling the circuitry for different sub-tasks.\n\nWeaknesses:\n1. I find the “causal” claims in the paper overstated and potentially misleading. When comparing their work with SAE and probing classifiers, they say “(SAE’s) insights remain correlational and still depend on post hoc labeling or interpretation, inheriting the same supervision bottleneck”. However, I don't see how CHG is more causal than SAE, given that the concept of causality in both is the intervenability of latent/transformed variables in a neural network considered as a causal graph. CHG is less causal than SAE in two senses: (i) CHG analyzes coarse, polysemantic units (heads) and gets unstable role assignments, whereas SAEs aim to find fine-grained, and potentially monosemantic units. (ii) while SAE aims to interpret/control the output (as the effect of the intervention), CHG only interpret/control which head contributes to the final loss (as the effect of the intervention) on a specific data distribution.\n2. The paper's own results demonstrate that the assigned causal roles are not stable and vary across different random seeds. This suggests that the discovered roles are not intrinsic properties of the attention heads themselves, but rather artefacts of the CHG optimization process converging to one of many possible sufficient sub-circuits. This instability undermines the claim of “identifying” a “causal” attention head.", "was_revised": true, "final_justification": "After discussions, my main concerns remain. The paper's claims about \"causality\" are overstated, as the method deals with predictive correlations on a fixed dataset, and the key experiments (Fig. 5) fail to provide sufficient evidence for the invariance or generalizability of its findings.", "rebuttal": "We thank the reviewer for insightful comments. Given the rebuttal rules that prohibit resubmission of any form, we will do our best to answer the reviewer concerns here and hope that our responses would help appeal for a higher reviewer score.\n\n1.\tRe: CHG abstraction level – We agree that CHG operates at a coarser level than SAE. Our goal is not to match the granularity of representational methods but to work at the head level, which we view as a meaningful and tractable unit for analyzing model function. Attention heads offer useful levels of abstraction over individual representational units in the same way that brain regions may serve as useful abstractions over individual neurons. If the practitioner’s goal was to understand the role of individual representational units, then we are in agreement that SAE is likely to offer more useful insights than CHG. We present CHG not necessarily as a direct substitute for SAE, which is useful in its own right, but as a potentially complementary method that operates at a higher level of abstraction.\n2.\tRe: CHG and downstream control – We respectfully disagree with the claim that CHG does not affect model outputs. CHG gates control whether a given head is active, and the resulting change in the model’s behavior is directly measured through the loss. Since the loss aggregates effects on the model’s predictions, CHG does intervene on internal components and evaluates the causal impact at the output level—just as any intervention-based causal analysis should.\n3.\tRe: SAE and causality – The limitation we raise by discussing SAE in juxtaposition with representational decoders is that like decoders, SAEs do not necessarily discriminate between incidental values and causal values. Just because a representational unit in a network has a particular value does not guarantee that the model necessarily uses it to perform the task. For example, the problem “John has 3 red apples and eats one. The number of remaining red apples is…” would likely produce some representation of ‘redness’ due to the semantic content of the apples, despite color being irrelevant to the task. While an SAE, which aims to preserve all information, could yield monosemantic units that can be intervened on, this would require a priori knowledge from the user to know that ‘3’ is task-relevant while ‘red’ is not. This is not to say that SAE cannot be used to identify causal components; one certainly could with careful experimentation and appropriate controls. Our message here is that CHG is directly causal (via interventions), whereas SAE in its standard form is more observational (interpreting what’s already there).\n4.\tRe: CHG stability across seeds – We would like to clarify that variability in CHG coefficients across seeds is not instability but an intentional design choice. CHG is best viewed as a bootstrapped estimator over a distribution of sufficient subcircuits. Different fits reveal different but valid configurations of causal head usage, reflecting the redundancy and contingency in how transformer models distribute function. This distributional perspective provides insight into which heads are essential across configurations, versus which are involved only under certain inter-head dependencies. \n5.\tRe: causal role differences across seeds - We hypothesize that this form of variance across seeds would occur if heads are either redundant or contingent. First, consider two heads A and B that perform the same function, so that either is sufficient for the task and having both does not substantially improve performance. In a seed where the initial CHG coefficient of A is higher than B, CHG would likely identify A as facilitating and B as irrelevant (redundant), and the reverse in the opposite case. Next, again consider heads A and B where the function of head A is contingent on B, so that if B is suppressed, A has function X and Y if B is not. In this case, A may be considered facilitating if B is enabled but irrelevant if B is not, reflecting the changing roles. We hope to answer these differences in further research.\n6.\tRe: applying gates fitted on one dataset on a different dataset – We have observed that gating matrices fitted on BigBench Syntax, OpenMathInstruct2, or CommonsenseQA do not transfer well when applied to a different dataset. This is expected, as the three datasets target distinct capabilities within the model, and the learned gating patterns are correspondingly different. As seen in Figure 3a, even within the same model, the CHG score distributions vary substantially across datasets. The exception lies in the contrastive method in Section 4.4, where isolating heads responsible for instruction following or ICL allows for generalization to other tasks. This suggests that CHG captures task-specific subcircuits by default, but that contrastive CHG can help surface more generalizable, reusable components.\n7.\tRe: expressivity of G - This is a valid and important concern, which is why we validate the learned gating matrices by holding out entire tasks in the contrastive CHG experiments. If G simply overfit to the contrastive objective, we would not expect any meaningful transfer to a held-out task. However, as shown in the Present–Past experiment, ablating instruction-following or ICL heads—identified using five other tasks—can selectively impair the corresponding mode of operation without affecting the other. This suggests that the learned gates are isolating functionally meaningful circuits. That said, this effect is imperfect (see Figure 5), and we are actively investigating why some tasks permit clean separation at the head level while others do not. Our current hypothesis is that certain tasks share heads with instruction-following or ICL circuits, so ablating those heads inadvertently disrupts task performance.\n\nWe appreciate the generally positive scores on quality, clarity, significance, and originality. If our responses to the reviewer concerns are satisfactory, we hope that you would reflect this in a higher overall rating for our work."}, {"review_id": "GoKDmsnfgq", "reviewer": "Reviewer_ymfF", "summary": "This paper introduces causal head gating (CHG), a method that determines the relevance of individual attention heads for performing a given task. The method trains two masks using standard cross-entropy loss on a task dataset w/ different regularization penalties, and based on the mask values, CHG labels an attention head with one of three categories: facilitating, interfering, or irrelevant. Using CHG, the authors identify attention heads important for three different kinds of tasks: math, syntax, and commonsense reasoning. They find that the role and importance of attention heads can vary based on other heads and the context. They also find that the heads identified by CHG causally affect generation, and that CHG scores correlate strongly with causal methods of attention head localization.", "questions": "As I stated in the weaknesses section, my main concern is related to the contextualization of the method - both providing baselines and comparisons to previous work in terms of design choices. This leads to the following questions:\n\n1. How does your method compare to previous methods such as Voita et al. [1], Li et al. [2] that also learn gate coefficients for individual attention heads? Does your method find more/less heads? Does CHG improve performance under intervention more than previous work? Does CHG take the same amount of time to converge?, etc. (see weaknesses for more detail)\n\n2. When designing CHG, what other design choices did you try that perhaps did not work as well? For example, how did you decide you needed to differentiate heads using two runs w/ a flipped learning objective?", "limitations": "Yes, limitations are thoughtfully discussed.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths:**\n- The paper is written clearly and was fairly easy to follow throughout.\n- The method does not rely on paired data, which is a nice complement to other localization methods like causal mediation, which often relies on counterfactual pairs. The result showing that CHG scores correlated with causal-based scores is further evidence that these methods might complement each other well for future analyses. The flexibility of the method (that you can construct a contrastive version) is also nice.\n- The claims in the paper generally have evidence to back them up, though some results are stronger than others. (e.g. results comparing ICL and instructions are a bit muddled, with some tasks exhibiting certain behavior, and others not having the same property)\n\n_______\n**Weaknesses:**\n\nMy main concern is a lack of comparison with previous work, despite many in-text citations. The authors are aware of related work (including very recent papers), and I would say this a strength of the paper. However, there is a lack of comparison to previous work in two places that I think would help significantly enhance the context of the paper (some places do compare, like the comparison with CMA section, and this is a nice contribution). \n\n- First is in the related work section itself. More effort to differentiate your work from previous related work in the related work section would be helpful. In each paragraph of the related work, there is a lack of discussion about why your work/method is related to or different from the cited work. For example, probing (representational decoders) is discussed for a whole paragraph, but there is no probing done in this paper. While this section is a nice review, contextualizing your work here would improve the paper. The other related work sections are similar.\n\n- The second place I think more comparison to related work would help is in the definition and evaluation of the CHG method. The CHG method is very similar to other previous methods that also learn gate coefficients per attention head. (e.g., Voita et al. [1],  Li et al [2], etc.). While these other papers did so in the context of pruning, it’s feasible to adapt them to the settings you’re evaluating here. I would recommend comparing and contrasting your method with these previous related methods. This would help provide some more justification for the choices you made while designing your CHG method, and give some more context as to how it compares to previous attempts in sparse optimization. As far as I could tell, the main difference of your method is the use of L1 regularization to induce sparsity instead of L0 penalty + annealing the gate values via the gumbel-softmax trick? The insight that learned gate values may end up close to 1 either because of task-importance or incidentally is an interesting one, and needing to run 2x to remove incidental importance seems to overcome that. It's likely this is a general problem of sparse optimization but it’s unclear from the current writing & experiments whether previous gate-learning methods are able to overcome this or not. This is where comparison might help show the advantage of CHG.\n\n- While the CHG mask is jointly learned and implies some sort of interaction it does not detail the kinds of interactions that the attention heads selected may have. This is one limitation that is not discussed explicitly, but may fall under the idea that this method does not identify the specific functional role of any attention heads, only whether they facilitate or inhibit performance of the dataset task.\n___\n**Other:**\n- It was unclear in Figure 1 which task, dataset & model are used for the plots.\n- There are also a few related works (Davies et al. [3], Cao et al. [4]) that learn differentiable masks over components that were not cited but may be of interest to the authors. Several times the authors point to the idea that attention heads are dependent on interactions with other heads. There are a few works that describe and study such interactions between attention heads (Elhage et al. [5] or Merullo et al. [6]), though the second one is not cited and may be of interest.\n\n- A minor typos on Line 217, there's an extra “that”\n___\n\n[1] Voita et al. [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://aclanthology.org/P19-1580.pdf) \n\n[2] Li et al [Differentiable Subset Pruning of Transformer Heads](https://aclanthology.org/2021.tacl-1.86.pdf) \n\n[3] Davies et al.  [Discovering Variable Binding Circuitry with Desiderata](https://arxiv.org/pdf/2307.03637)\n\n[4] Cao, et al. [Low-complexity probing via finding subnetworks](https://aclanthology.org/2021.naacl-main.74.pdf)\n\n[5] Elhage et al. [A Mathematical Framework For Transformer Circuits](https://transformer-circuits.pub/2021/framework/index.html)\n\n[6] Merullo et al [Talking Heads: Understanding Inter-Layer Communication in Transformer Language Models](https://proceedings.neurips.cc/paper_files/paper/2024/file/70e5444e5f331f7f5431f302110b97af-Paper-Conference.pdf)", "was_revised": true, "final_justification": "During the rebuttal phase, the authors addressed my concerns about the framing/discussion of related work and how the proposed method differs in theory. I also found helpful their clarification about the consistency of CHG over multiple seeds. Incorporating this information into the main paper will help clarify the framing and contributions of the method for future readers.\n\nI still maintain that additional empirical validation of the method and comparison to previous sparsity methods would strengthen the results even more (hence not a higher score), but the contributions are sufficient in my mind to warrant acceptance.", "rebuttal": "We thank the reviewer for insightful comments. Given the rebuttal rules that prohibit resubmission of any form, we will do our best to answer the reviewer concerns here and hope that our responses would help appeal for a higher reviewer score.\n\n1.\tRe: comparison to Gumbel-softmax based methods (Voita et al. and Li et al.) – We appreciate the reviewer raising this important point. While we cannot provide new experimental data during this rebuttal period, we can provide a more theoretically grounded distinction between CHG and Gumbel-softmax based methods.\nWhile both CHG and Gumbel-based methods learn a parameter per head, they differ fundamentally in how they handle dependencies between heads. Gumbel methods model each head independently, treating each gating parameter as defining a separate Bernoulli distribution. As a result, they learn marginal distributions over head inclusion, without accounting for how heads may functionally depend on each other. In an LLM with N heads, the optimization effectively learns N different distributions. However, since the gate coefficients are sampled simultaneously, they model marginal distributions that preclude modeling interdependencies between heads. \nTo illustrate, imagine heads A and B both perform the same role. A good sparse solution might activate one and suppress the other, but not both. Gumbel methods would assign roughly P(A) = P(B) = 0.5, resulting in equal probability to all combinations—including ones that activate both or neither—because they don’t capture the interaction. \nCHG, by contrast, jointly optimizes all gating coefficients under the model's loss. This means CHG naturally captures interdependencies: the value for head A can depend on the value for head B, and vice versa. In other words, CHG effectively samples from a coupled energy landscape, capturing dependencies that factorized gating methods cannot represent.\nThat said, it is possible to modify Gumbel strategies to account for dependencies (e.g., by sampling sequentially), but this adds complexity and deviates from standard approaches like those in Voita et al. In practice, our results show that CHG often assigns different roles to the same head across seeds, which further suggests the presence of interdependencies that marginal models like Gumbel might miss.\nWhile we leave full empirical comparison to future work, we believe CHG offers a meaningful theoretical advantage by modeling head interactions directly. We hope this distinction helps clarify the unique value of our method.\n2.\tRe: related work – Thank you for your comment on the related work, as well as some of the suggested citations. Our aim here was to compare with other interpretability tools used in the literature more broadly and agree in retrospect that there could have been a more direct comparison with the most adjacent methods, especially in the theoretical motivation for CHG that contrasts it with the Gumbel based methods. If permitted, we will incorporate this discussion into the final publication.\n3.\tRe: Figure 1 - These plots were generated from applying CHG to the Llama 3.2 3B-Instruct model on the OpenMathInstruct2 dataset.\n4.\tRe: design choices - Our decision to use two runs was theoretically, rather than empirically, motivated by the goal of distinguishing between task facilitating, irrelevant, and interfering heads. A single optimization could distinguish one from the other two (e.g. facilitating from irrelevant and interfering in the G- optimization), but not all three. The only other design choices we explored were hyperparameters (e.g. the lambda values) which did not significantly impact the results (within reasonable values).\n\nWe appreciate the generally positive scores on quality, clarity, significance, and originality. If our responses to the reviewer concerns are satisfactory, we hope that you would reflect this in a higher overall rating for our work."}]

Author rebuttal:

Meta-review: This paper is a borderline. It introduces Causal Head Gating (CHG), a simple and scalable framework for interpreting the roles of attention heads by learning head-specific gates under interventional training. The method yields a useful taxonomy of facilitating, interfering, and irrelevant heads, validated through ablations and alignment with causal mediation analysis, and is further extended with a contrastive variant (CCHG) to differentiate circuits for in-context learning and instruction following. 

Reviewers expressed mixed views: some concerns about novelty relative to existing differentiable masking and patching methods, stability across seeds, and over-claiming of “causal discovery.” However, others found the approach clean, practical, and well-validated, with two reviewers raising their scores after rebuttal. Overall, despite some limitations, I think this is a useful contribution to interpretability, providing a lightweight, label-free tool that researchers are likely to adopt. I recommend Accept, contingent on camera-ready clarifications about the scope of causality claims and broader comparisons to related methods.
Decision: Accept (poster)