🤖 AI

Multi-Granular Node Pruning for Causal Circuit Discovery

本文提出了一种可扩展的多粒度节点剪枝框架，该框架通过使用可学习掩码和稀疏性惩罚，在识别相关单个神经元的同时，显著降低了与现有边剪枝方法相比的内存需求，从而高效地发现大型语言模型中更小、更精确的因果电路。

原作者： Muhammad Umair Haider, Hammad Rizwan, Hassan Sajjad, A. B. Siddique

发布于 2026-06-16

📖 1 分钟阅读☕ 轻松阅读

原作者： Muhammad Umair Haider, Hammad Rizwan, Hassan Sajjad, A. B. Siddique

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

将大型语言模型（例如那些编写故事或回答问题的模型）想象成一座规模宏大、繁忙喧嚣的城市。这座城市拥有数百万座建筑（神经元）、道路（连接）和区域（层）。当模型执行特定任务时——比如弄清楚在故事中“克里斯蒂”（Kristi）把芒果给了谁——它并不会动用整座城市，而只会使用一个非常微小且特定的街区。

寻找这个特定街区的过程被称为电路发现（Circuit Discovery）。其目标是绘制出精确的地图，标明哪些部分正在进行工作，而哪些部分只是在那儿闲置不动。

旧地图的问题

此前，研究人员试图通过观察建筑之间的道路（连接）来寻找这些街区。他们会尝试封锁道路，观察城市是否仍能正常运转。

缺陷： 这就像是通过封锁每一条街道来试图在一座城市中找到某一座特定的房子。这既耗时又需要海量的内存（就像需要一张包含世界上每一条街道的地图一样），而且这种方法过于粗糙。如果你封锁了一条路，你可能会不小心切断整个街区，即便该街区中只有一间房子是真正被需要的。
结果： 旧的地图是“粗粒度”的。它们能告诉你整个区域（比如一个“注意力头”）很重要，但无法告诉你该区域内某栋建筑里的某一个特定房间才是真正干活的地方。

新方案：多粒度节点剪枝

该论文的作者提出了一种更聪明的城市绘图方式。他们不再仅仅盯着道路，而是直接观察建筑（节点）本身，并且同时在不同的细节层级上进行观察。

这就像是一套俄罗斯套娃，或者是一个变焦镜头：

大局观： 他们可以关闭整个区域（Transformer 块）。
街区： 他们可以关闭区域内的特定街道（注意力头）。
房屋： 他们可以关闭房子里的特定房间（单个神经元）。

他们为每一个建筑——从最大的区域到最小的房间——都使用了一个特殊的“可学习掩码”（一种数字开关）。他们运行了一场模拟实验，将一个“干净”版本的城市与一个“损坏”版本（即故事变得毫无逻辑的版本）进行混合。通过观察哪些建筑必须保持开启状态才能让故事保持正确，他们就能推断出究竟哪些部分是必不可少的。

他们的发现（结果）

当他们在不同的“城市”（如 GPT-2 和 Llama 等模型）上进行测试时，他们发现了一些令人惊讶的事实：

巨大的节省： 他们的这种方法找到的电路比以往任何方法都要小得多。在最坏的情况下，他们比之前的方法多移除了 33% 的建筑模块和 60% 的单个房间。
内存效率： 由于他们不需要存储每一张详细的街道地图（中间激活值），他们的方法使用的计算内存要少 3 到 11 倍。这就像是只需要一本小笔记本，而不是一座图书馆来绘制地图。
任务特定的蓝图： 他们发现不同的任务使用着不同的城市布局：
- 间接宾语识别（谁对谁做了什么）： 这项任务高度依赖于遍布整个城市的“MLP”建筑（非线性处理房间），而“注意力”街道则大多是空置的。
- 性别代词（他 vs 她）： 这项任务使用了一个非常稀疏且分散的网络。大部分城市都是关闭状态；只有少数特定的层和房间处于活跃状态。
- 大于（数学/数字）： 这是最极端的例子。它使用了一种“跳跃”机制，即城市忽略了中间层的巨大区块，直接跳到最后来进行计算。

核心结论

该论文声称，通过在单个“房间”（神经元）的层级而非仅仅在整个“区域”（块）的层级观察模型，我们可以剥离掉大量的冗余机械结构。

他们证明了，在这些 AI 模型中，许多我们曾认为必不可缺的部分实际上只是无用的负担。通过将其精简到最细的粒度，他们找到了驱动行为的真正、最小化的“电路”，并且完成这项工作的速度更快，消耗的计算资源也更少。他们还展示了这种方法甚至可以在单张显卡上处理非常大的模型（如 Llama 3.1-8B），而这在以前对于其他方法来说是难以实现的。

技术摘要：用于因果电路发现的多粒度节点剪枝

1. 问题陈述

因果电路发现旨在识别大型语言模型（LLM）中驱动特定行为的最小子网络。虽然现有方法取得了进展，但仍面临两个主要局限性：

计算可扩展性： 当前的方法主要依赖于边剪枝（移除组件之间的连接）。由于 Transformer 中的边数量随节点数呈二次方增长（ $O(|N|^2)$ ），随着模型规模增大，边级剪枝在计算上变得难以实现。此外，这些方法通常需要存储解耦的残差流表示和中间激活值，导致内存占用过高。
粒度限制： 现有方法通常在粗粒度层面运行，将整个注意力头或 MLP 块视为原子单元。这忽略了这些组件内部特定的神经元（节点）可能负责目标行为的可能性，导致恢复出的电路遗漏了细粒度结构并保留了不必要的组件。

2. 方法论

作者提出了一种多粒度节点剪枝框架，该框架在单一统一的优化步骤中，跨越多个粒度层级发现最小电路。

2.1 分层粒度设计

该框架不是剪除粗糙组件之间的边，而是学习在四个不同层级的节点上生成掩码：

块级（Block Level）： 整个注意力块（Attention Blocks）和 MLP 块（MLP Blocks）。
头级（Head Level）： 单个注意力头（Attention Heads）。
神经元级（Neuron Level）： 注意力层中的单个神经元以及 MLP 块中的两个线性层（MLP1/Keys 和 MLP2/Values）。

2.2 双流因果剪枝设置

为了在识别因果重要性的同时避免引入分布偏移（这是简单零消融带来的风险），该方法采用了双流前向传播：

干净流（Clean Stream）： 模型处理旨在诱发特定行为（例如，预测年份 > 50）的输入。
损坏流（Corrupted Stream）： 模型处理经过极小扰动的输入版本（例如，反转性别代词或改变起始年份）以诱发不同的预测。
插值（Interpolation）： 对于每个节点 $n_i$ ，一个可学习的掩码 $m_i \in [0, 1]$ 在干净激活（ $h^{clean}$ ）和损坏激活（ $h^{corrupted}$ ）之间进行插值：
$h^{node}_i = m_i \cdot h^{clean}_i + (1 - m_i) \cdot h^{corrupted}_i$
优化过程旨在寻找能够使模型产生正确输出（匹配干净行为）同时最大化稀疏性的掩码。

2.3 优化与掩码参数化

硬连续分布（Hard Concrete Distribution）： 掩码使用 Hard Concrete 分布进行参数化，以实现基于梯度的优化，同时诱导近似二值的行为（0 或 1）。
统一目标函数： 训练目标结合了三个部分：
1. 忠实度损失（ $L_{faith}$ ）： 一种 KL 散度，确保剪枝后的电路输出分布与基础模型一致。
2. 任务损失（ $L_{task}$ ）： 一种基于间隔的排序损失，确保电路能正确区分真实选项和干扰项。
3. 稀疏正则项（ $L_0$ ）： 鼓励掩码趋于零，从而最小化电路规模。
层级一致性： 训练后，通过一种约束进行二值化，确保如果父节点（例如，一个 MLP 块）被剪枝，其所有子节点（例如，该块内的神经元）也必须被停用。

3. 核心贡献

多粒度框架： 引入了一个统一的优化框架，能够同时在 Transformer 块、注意力头、MLP 块和单个神经元上学习掩码。
可扩展性与效率： 该方法避免了边剪枝的二次方缩放问题，也无需存储所有边的中间激活值。与以往方法相比，它将 GPU 显存占用降低了 3 倍至 11 倍，使得在单块 40GB GPU 上进行 Llama 3.1-8B 的电路发现成为可能。
细粒度发现： 该方法识别出的电路比基准方法显著更加紧凑，特别是在 MLP 块和神经元层级实现了极高的稀疏度，而以往的方法在这些层级表现不佳。

4. 实验结果

该方法在三个标准的电路发现任务（间接宾语识别 IOI、性别代词 GP 和大于 GT）上进行了评估，使用了 GPT-2、GPT-2-XL 和 Llama 模型。

4.1 稀疏性与紧凑性

MLP 与神经元剪枝： 所提方法显著优于边剪枝基准（EAP、EP）和非结构化权重剪枝（DiscoGP）。在最不利的设置下，它比以往方法多剪除了 33.34% 的 MLP 块和 59.8% 的神经元。
任务依赖型结构： 发现的电路展现出基于任务的不同拓扑结构：
- IOI： 由分布式 MLP 计算主导；注意力部分被大量剪除，仅在最后几层保留。
- 性别代词 (GP)： 高度局部化且不连续；MLP 计算被限制在少数几层，而注意力较为稀疏。
- 大于 (GT)： 表现出极端的稀疏性，具有不活跃层的连续块和“跳跃式”执行模式。

4.2 忠实度与一致性

性能： 剪枝后的电路保持了高任务准确率（通常 >95%）以及相对于全量模型的低 KL 散度。
一致性： 在五个随机种子上的实验表明，发现的电路（包括忠实度指标和稀疏度剖面）具有高度稳定性。
定性分析： 在 Llama 3.2-1B 的 IOI 任务中，存留的注意力头被发现执行特定的功能角色（IO-Movers, S-Inhibitors），这与 GPT-2 中手动发现的电路相吻合，验证了该方法恢复机械论见解的能力。

4.3 计算效率

内存： 该方法所需的内存至少比边剪枝基准（如 EAP）少 3 倍（例如，在 GPT-2 上，EAP 需要 72.79 GB，而本方法仅需 6.27 GB）。
时间： 训练在 6 分钟内完成，明显快于边剪枝方法（约 50 分钟），但略慢于单次传递的 EAP 方法（后者是以牺牲内存换取时间）。

5. 重要性与主张

本文声称其主要意义在于克服了当前电路发现中的可扩展性和粒度瓶颈。通过从边剪枝转向多粒度节点剪枝，作者证明了：

许多被粗粒度方法保留下来的神经元是不必要的，可以被移除且对性能影响微乎其微。
更大的模型（如 Llama 8B）支持与较小模型同样强劲甚至更强的剪枝效果，这表明规模的增加引入的是冗余而非要求更密的任务特定电路。
该方法为发现大规模模型中细粒度的因果电路提供了一条计算可行且高效的路径，揭示了不同的行为是由注意力、MLP 和神经元活动的独特模式所编码的。

作者承认一个局限性：该框架识别了哪些节点是必要的，但并未明确恢复生存节点之间的交互结构（边）。他们建议未来的研究方向是将节点选择与边恢复相结合，以推断交互结构。