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

📄 arXiv: 2505.13737v2 📥 PDF

作者: Andrew Nam, Henry Conklin, Yukang Yang, Thomas Griffiths, Jonathan Cohen, Sarah-Jane Leslie

分类: cs.AI

发布日期: 2025-05-19 (更新: 2025-10-23)

备注: 10 pages, 5 figures, 2 tables. The Thirty-Ninth Annual Conference on Neural Information Processing Systems (NeurIPS 2025)

💡 一句话要点

提出因果头门控方法以解析变换器中注意力头的功能角色

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 因果推断 注意力机制 变换器模型 可解释性 大型语言模型 任务性能 软门控 子电路分析

📋 核心要点

  1. 现有的机制可解释性方法通常依赖于假设驱动和特定的提示模板，限制了其适用性和灵活性。
  2. 因果头门控（CHG）通过学习软门控并对注意力头进行因果分类，提供了一种通用的解析方法，适用于各种数据集。
  3. 实验结果表明，CHG能够揭示大型语言模型中的稀疏任务充分子电路，并验证了因果关系的有效性。

📝 摘要（中文）

本文提出了一种可扩展的方法——因果头门控（CHG），用于解析变换器模型中注意力头的功能角色。CHG通过学习头部的软门控，并根据其对任务性能的影响将其分为促进、干扰或无关的因果分类。与以往的机制可解释性方法不同，CHG不依赖于假设驱动或特定的提示模板，而是直接应用于任何数据集，使用标准的下一个标记预测。我们在多个大型语言模型（LLMs）和多样化任务上评估了CHG，结果表明CHG得分提供了因果而非仅仅是相关的洞察，且通过消融和因果中介分析得到了验证。我们还引入了对比CHG变体，以隔离特定任务组件的子电路。

🔬 方法详解

问题定义：本文旨在解决如何有效解析变换器模型中注意力头的功能角色这一问题。现有方法往往依赖于假设和特定的任务标签，限制了其适用范围和灵活性。

核心思路：因果头门控（CHG）通过学习注意力头的软门控，基于其对任务性能的影响将其分类为促进、干扰或无关，从而提供一种通用的解析框架。

技术框架：CHG的整体架构包括头部门控学习、因果分类和任务性能评估三个主要模块。首先，模型通过标准的下一个标记预测学习头部的软门控；然后，基于这些门控的影响进行因果分类；最后，通过多种任务评估其性能。

关键创新：CHG的主要创新在于其不依赖于假设驱动的方法，能够直接应用于任何数据集，并提供因果而非相关的洞察。这一方法与以往的机制可解释性方法本质上不同。

关键设计：在设计上，CHG采用了特定的损失函数来优化门控学习，并通过消融实验和因果中介分析验证其有效性。模型结构上，CHG能够灵活适应不同的任务和数据集。

📊 实验亮点

实验结果显示，CHG在多个大型语言模型上表现出色，能够有效揭示注意力头的因果角色。通过消融实验验证，CHG得分提供了因果洞察，且与传统方法相比，提升了模型在语法、常识和数学推理等任务上的性能。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、机器翻译和对话系统等，能够帮助研究人员和工程师更好地理解和优化变换器模型的性能。通过解析注意力头的功能角色，能够提升模型的可解释性和可靠性，促进更智能的人工智能系统的发展。

📄 摘要（原文）

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.