Causal Head Gating: A Framework for Interpreting Roles of Attention Heads in Transformers

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.

1. Re: 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. While 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. To 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. CHG, 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. That 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. While 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.
2. Re: 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.
3. Re: Figure 1 - These plots were generated from applying CHG to the Llama 3.2 3B-Instruct model on the OpenMathInstruct2 dataset.
4. Re: 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).

We 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.

Thank you for your response. The clarification of your thoughts on the difference between CHG and similar methods is helpful and would be a nice addition as a discussion in the related work.

"our results show that CHG often assigns different roles to the same head across seeds"

Thanks for mentioning this in the rebuttal. I wanted to get some additional clarity about this point, as this inconsistency is one of the major limitations of the proposed CHG method. As I think about it some more, it seems like this suggests that CHG in fact has a hard time modeling interdependence of heads, though this is claimed as one of the theoretical strengths of the method ("CHG naturally captures interdependencies"). Can you clarify CHG's inconsistency of role assignments across seeds a bit more? For example, how many seeds you would need to run to be sure you capture the true role/effect? (e.g., how many seeds are used for the results in 4.3?) Does this limit the reliability of CHG in some ways? Or maybe it points to the fact that localization is much more messy than previously thought due to redundancy and interdependence.

Based on the rebuttal, I appreciate the comments related to related work and feel like my questions will be/have been addressed there. However, I still have some reservations about the empirical validation of the method (comparison with similar previous methods), and am looking forward to a discussion of the above point. Overall, while the current results are promising, I'm inclined to keep my score for now until further discussion.

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.

1. Re: 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.
2. Re: 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.
3. Re: 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).
4. Re: 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.
5. Re: 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.
6. Re: 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.
7. Re: 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.

We 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.

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.

1. Similarity 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?
2. Novelty 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.
3. Comparison 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. If there are specific forms of causal tracing or attribution patching you feel would strengthen the comparison, we would welcome your suggestions.
4. Role 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.
5. Training 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.
6. Bifurcation 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).
7. Causal 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−).
8. Similarity 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.

We 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.

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.

1. Re: 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.
2. Re: 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.
3. Re: 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.
4. Re: 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.
5. Re: 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.”
6. Re: 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
7. Re: 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.
8. Re: 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.

We 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.