NeurIPS 2025 可解释性 circuit discovery attribution patching 剪枝 hybrid framework IOI GPT-2 mechanistic interpretability

Discovering Transformer Circuits via a Hybrid Attribution and Pruning Framework¶

会议: NeurIPS 2025
arXiv: 2510.03282
代码: GitHub
领域: 机械可解释性 / 电路发现 / Transformer分析
关键词: circuit discovery, attribution patching, edge pruning, hybrid framework, IOI, GPT-2, mechanistic interpretability

一句话总结¶

提出混合归因与剪枝框架 HAP，先用快速的边归因修补（EAP）筛选高潜力子图，再在缩小后的搜索空间上运行精确的边剪枝（EP），在 GPT-2 Small 的 IOI 任务上比纯 EP 快 46% 且保持相当的电路忠实度，同时成功保留了 EAP 单独使用时会遗漏的 S-inhibition 头。

研究背景与动机¶

机械可解释性的核心目标：随着 LLM 被部署到高风险场景，理解其"黑箱"内部运作成为刚需。机械可解释性通过识别负责特定行为的稀疏子网络（"电路"）来实现这一目标。

电路分析的标准范式：将 Transformer 表示为计算图（节点=注意力头/MLP，边=信息流），在图上寻找执行特定任务的最小子图。手动方法（Wang et al. 2022）已被自动方法取代。

ACDC 的计算瓶颈：最早的自动方法 ACDC（Conmy et al. 2023）采用逐边贪心消融搜索，忠实度高但需要大量前向传播，无法扩展到大模型。

EAP 的速度优势与忠实度缺陷：EAP（Syed et al. 2023）用一阶 Taylor 近似同时估算所有边的重要性，仅需 1 次反向+2 次前向传播即可完成，但线性近似导致忠实度显著下降，且在高稀疏度下会丢失协作性组件。

EP 的忠实度优势与算力需求：EP（Bhaskar et al. 2024）通过梯度优化二值掩码实现精确剪枝，忠实度优异且已扩展到 CodeLlama-13B，但需要大量 GPU 算力和长时间训练。

核心矛盾与机会：EAP 快但不忠实，EP 忠实但慢——二者的优势恰好互补。能否用 EAP 的速度做粗筛，再用 EP 的精度做细选？这是 HAP 的出发点。

方法详解¶

整体框架：HAP（Hybrid Attribution and Pruning）¶

HAP 将电路发现分解为三步流水线：① 计算图构建 → ② EAP 粗筛 → ③ EP 精确剪枝。核心思想是用快速但粗糙的归因方法缩小搜索空间，再在缩小后的空间上运行精确但昂贵的优化方法。

关键设计 1：计算图构建¶

功能：将 Transformer 模型表示为有向计算图
核心思路：节点为注意力层和 MLP 层，边表示一个节点的输出到另一节点输入的信息流连接。对 GPT-2 Small（117M 参数，12 层 × 12 头），构建包含所有注意力头和 MLP 之间的完整边集合
设计动机：统一的图表示是后续归因和剪枝操作的基础，遵循 Bhaskar et al. (2024) 的标准惯例以确保可比性

关键设计 2：EAP 快速粗筛¶

功能：用一阶 Taylor 近似同时计算所有边的绝对归因分数，按分数排序后保留 top-k 边
核心思路：对每条边 $e$，其重要性近似为： $$L(\mathbf{x} \mid e_{\text{ablated}}) - L(\mathbf{x}) \approx (e_{\text{clean}} - e_{\text{ablated}})^\top \frac{\partial L(\mathbf{x} \mid e_{\text{clean}})}{\partial e_{\text{clean}}}$$ 仅需一次反向传播和两次前向传播即可获得所有边的分数
设计动机：EAP 的计算成本几乎恒定（与边数无关），适合用于快速淘汰大量明显不重要的边。关键是阈值设置极低（保守筛选），故意保留低个体分数但可能具有协作重要性的边

关键设计 3：缩小搜索空间上的 EP 精确剪枝¶

功能：在 EAP 筛选后的子图上运行梯度优化的边剪枝
核心思路：EP 在缩小后的搜索空间上优化二值掩码 $z \in [0,1]^{N_{\text{edge}}}$，最小化原始图与剪枝图的输出散度，同时满足目标稀疏度约束 $1 - |H|/|G| \geq c$
设计动机：EP 的计算成本与搜索空间大小正相关。EAP 预先移除了大量无关边后，EP 需要优化的参数空间显著缩小，训练收敛更快。同时 EAP 的宽阈值确保了"安全区"，使 S-inhibition 头等协作组件不会在粗筛阶段被丢弃

关键设计 4：宽阈值安全区策略¶

功能：在 EAP 阶段刻意设置极低的筛选阈值
核心思路：不追求 EAP 阶段的高稀疏度，而是保留宽裕的候选边集合，允许个体归因分数低但对整体电路功能有贡献的边进入 EP 阶段
设计动机：S-inhibition 头等协作性组件的特点是单独看重要性低，但在电路中发挥关键的抑制/协调作用。EAP 的线性近似无法捕捉这种非线性协作效应，因此需要宽阈值来避免误删

损失函数与训练策略¶

EP 阶段的优化目标包含两部分：

忠实度损失：最小化完整模型与剪枝子图在 clean 输入和 corrupted 输入上的输出 KL 散度，确保电路行为忠实于原模型
稀疏度约束：通过拉格朗日乘子或投影方法满足 $1 - |H|/|G| \geq c$ 的目标稀疏度

训练使用 clean 和 corrupted 样本对交替进行梯度更新。所有实验在单张 NVIDIA H100 GPU 上完成。EP 阶段的超参数沿用 Bhaskar et al. (2024) 的设定。

实验¶

实验设置¶

模型：GPT-2 Small（117M 参数）
任务：间接宾语识别（IOI），格式为 "When Dylan and Ryan went to the store, Dylan gave a popsicle to → Ryan"
数据集：训练集 200 例、验证集 200 例、测试集 36,084 例，使用 Wang et al. (2022) 的模板生成
评价指标：准确率（以手动发现的电路为 ground truth）、Logit Difference、KL 散度、运行时间
硬件：1× NVIDIA H100

表 1：HAP 与现有方法的效率与忠实度对比¶

算法	稀疏度	准确率 ↑	Logit Diff ↑	KL ↓	运行时间(s) ↓
EAP	94±0.5%	0.698	3.13	–	4
EP	94±0.5%	0.772	3.48	0.190	2921
HAP	94±0.5%	0.759	3.42	0.188	1579

关键发现：HAP 比 EP 快 46%（1579s vs 2921s），准确率仅低 1.3 个百分点（0.759 vs 0.772），KL 散度几乎相同（0.188 vs 0.190），Logit Difference 接近（3.42 vs 3.48）。相比 EAP，HAP 在所有质量指标上均大幅领先。

表 2：IOI 案例研究——S-inhibition 头的保留情况¶

方法	头 7.3	头 7.9	头 8.6	头 8.10	完整电路恢复
EAP	✗	✗	✗	✗	不完整
HAP	✓	✓	✓	✓	完整

S-inhibition 头在 IOI 中负责抑制 Name Mover 头错误标记与动词接近的主语，个体归因分数低但协作效果关键。EAP 在 94% 稀疏度下丢失全部四个 S-inhibition 头，而 HAP 通过宽阈值安全区 + EP 精确剪枝成功保留了完整的功能电路。

亮点¶

简洁优雅的组合策略：将 EAP 和 EP 串联使用，利用各自的速度/精度优势互补，思路直接但有效
挑战了速度-忠实度权衡的固有性假设：证明通过策略性两阶段搜索可以在不牺牲忠实度的前提下大幅加速
定性证据有力：S-inhibition 头案例直观展示了 HAP 保留协作组件的能力，弥补了纯定量比较的不足
实用价值明确：46% 的加速在扩展到更大模型时理论上更为显著

局限性¶

实验范围狭窄：仅在 GPT-2 Small（117M）的单一 IOI 任务上验证，模型规模和任务多样性严重不足
EAP 阈值未优化：边筛选阈值为启发式设置，缺乏系统的超参数搜索和敏感性分析
缺乏方差报告：训练数据随机生成，未报告多次运行的性能方差
基线对比不足：未与 ACDC、EAP-GP 等其他方法比较
大模型验证缺失：未在 Llama/CodeLlama 等规模模型上测试，可扩展性仅为理论预期

评分¶

新颖性: ⭐⭐⭐ 组合思路简单直接，方法层面创新有限但有效，核心贡献在于证明两阶段策略可行
实验充分度: ⭐⭐ 单模型单任务，缺乏方差报告和超参敏感性分析，说服力不足
写作质量: ⭐⭐⭐⭐ 结构清晰，动机阐述到位，IOI 案例分析直观有力
综合推荐: ⭐⭐⭐ 为机械可解释性的可扩展性提供了实用工程方案，但需要更广泛的实验验证