Title: Causal Head Gating: A Framework for Interpreting Roles of Attention Heads in Transformers
Venue: NeurIPS 2025
Accepted: yes
Keywords: mechanistic interpretability

large language models

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


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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)\n\nQUESTIONS:\nAs 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?\n\nLIMITATIONS:\nYes, limitations are thoughtfully discussed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n1. 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\n\nLIMITATIONS:\nThe instability and ambiguity of the \"causal roles\" assigned by CHG as mentioned in weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nDue 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.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 1, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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)\n\nSTRENGTHS AND WEAKNESSES:\nStrengths :\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.\n\nQUESTIONS:\nHow 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.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-02"}