Research map / Large language models / Language, safety and interpretability

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

NeurIPS 2025 · Large language models › Language, safety and interpretability · approach: Reasoning mechanisms and failure modes · 1 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes causal head gating (CHG), learning soft L1-regularized gates over attention heads via next-token NLL to classify heads as facilitating, interfering, or irrelevant, and extends it to contrastive CHG for separating ICL from instruction-following circuits in Llama 3 models.

Problem

Interpreting the functional causal roles of attention heads in LLMs without hypothesis-driven prompt templates or labeled target outputs, and capturing distributed, interdependent sub-circuits.

Approach

Fit a multiplicative soft gate per attention head with frozen model weights, using next-token NLL plus a signed L1-style regularizer (lambda>0 and lambda<0) to separate facilitating, interfering, and irrelevant heads; a contrastive variant jointly forgets one task formulation while retaining another.

Key results

• Ablating heads in CHG-ranked order yields the predicted log-probability pattern: decreases when removing facilitating heads, near zero for irrelevant, increases for interfering heads.
• Head score distributions are highly consistent across the four Llama 3 variants, with pairwise Pearson correlations averaging 99.2% (minimum 94.92%).
• CMA-identified heads show significantly higher CHG facilitation scores in function vector tasks (t(23.05)=8.52, p<1e-8) and ABA/ABB tasks (t(53.77)=11.18, p<1e-15).

Measured effects reported in the paper

Surprising findings

• Only a small fraction of heads are consistently facilitating or interfering across all seeds, indicating multiple overlapping sufficient sub-circuits rather than modular head roles.
• Contrastive CHG can suppress ICL performance to zero while largely preserving instruction-based accuracy for some held-out tasks, but the separability depends on the task (e.g., singular-plural instruction accuracy drops to 0%).

Limitations

• CHG does not explain why heads matter or what computations they perform; it is a first-pass diagnostic.
• G+ and G- occasionally diverge, and head roles vary across random seeds depending on which other heads are active.
• CMA comparisons were limited to specific function-vector and symbolic tasks on Llama-3.2-3B-Instruct.

Experimental setup

Datasets: OpenMathInstruct-2 CommonsenseQA BIG-Bench syntax function vector tasks ABA/ABB symbolic reasoning

Benchmarks: OpenMathInstruct2 CommonsenseQA BIG-Bench syntax function vector tasks ABA/ABB symbolic reasoning

Metrics: target log-probability difference task accuracy facilitation/interference/irrelevance scores Pearson correlation t-test

Baselines: causal mediation analysis Gumbel-based gating

Builds on: causal mediation analysis function vectors Gumbel-softmax gating

Closest papers by meaning

• Circuit Component Reuse Across Tasks in Transformer Language Models (ICLR 2024 · similarity 0.89)
• Selective Induction Heads: How Transformers Select Causal Structures in Context (ICLR 2025 · similarity 0.89)
• Linking In-context Learning in Transformers to Human Episodic Memory (NeurIPS 2024 · similarity 0.89)
• Talking Heads: Understanding Inter-Layer Communication in Transformer Language Models (NeurIPS 2024 · similarity 0.88)
• On the Emergence of Position Bias in Transformers (ICML 2025 · similarity 0.88)
• Differentiation and Specialization of Attention Heads via the Refined Local Learning Coefficient (ICLR 2025 · similarity 0.88)