像 GPT-4 或 Llama 这样的大型语言模型 (LLM) 常被描述为“黑盒”。我们知道它们能工作——它们能写诗、调试代码,还能告诉你法国的首都是哪里——但我们并不完全理解它们如何存储这些信息。当 LLM“知道”咖啡馆是一种餐厅时,这个事实是存储在特定的神经元簇中吗?还是像涂在吐司上的果酱一样散布在整个网络中?

如果我们能回答这个问题,其意义将是巨大的。我们可以像外科手术一样修正事实错误,而无需昂贵的重新训练。我们可以移除有毒知识或危险能力,同时保留模型其余的智能。

这正是 EPFL 的研究人员在论文 “Discovering Knowledge-Critical Subnetworks in Pretrained Language Models” (在预训练语言模型中发现知识关键子网络) 中解决的核心问题。他们研究了特定的关系型知识片段是否由稀疏子网络 (sparse subnetworks) ——即模型参数中微小且特定的子集——进行编码。

在这篇深度文章中,我们将梳理他们寻找这些“知识神经元”的方法、可微掩码技术背后的数学原理,以及当你通过手术移除特定事实时模型会发生什么。

假设: 知识是局部化的 (某种程度上)

本文的核心假设是,关系型知识 (如“汽车是一种交通工具”) 不仅仅是权重中的随机噪声。作者提出,对于任何给定的知识集,都存在一个知识关键子网络 (knowledge-critical subnetwork) 。

这是模型内部一个特定的、稀疏的计算子图。如果你移除它,模型应该会丢失它所编码的特定知识,但——这一点至关重要——模型在其他方面的功能应该保持完美。

图 1: 知识关键子网络对于在 LLM 中表达目标知识三元组 (TARGETKG) 是必不可少的。当被移除时,剩余的模型不再表达特定的三元组,但保留了表达其他关系型知识 (CONTROLKG) 的能力及其语言建模能力 (CONTROLLM) 。

图 1 所示,目标是一次有针对性的“脑叶切除术”。

  • 目标知识 (TARGETKG) : 我们想要抑制的特定事实 (例如,“咖啡馆是一种餐厅”) 。
  • 对照知识 (CONTROLKG) : 我们想要保留的相关事实 (例如,“汽车是一种交通工具”) 。
  • 语言建模 (CONTROLLM) : 通用的英语能力 (例如,语法、句法) 。

如果研究人员能找到一组权重 (图中用红色 X 标记) ,当移除这些权重时,能破坏目标 (Target) 路径但保留紫色和橙色路径的完整性,那么他们就找到了一个知识关键子网络。

背景: 知识图谱与三元组

在看神经网络手术之前,我们需要了解研究人员如何定义“知识”。他们依赖于知识图谱 (KGs) 。KG 将事实表示为三元组: (头实体, 关系, 尾实体)

显示知识三元组集合的公式。

例如,(巴黎, 是, 城市)。为了将这些输入到语言模型中,这些三元组被“语言化”为自然语言句子,如“巴黎是一个城市”。然后,在给定上下文“巴黎是一个……”的情况下,测试模型预测尾部标记 (例如,“城市”) 的能力。

研究人员使用 WordNet 和 ConceptNet 的数据集来创建这些目标知识集。

核心方法: 可微掩码

当干草堆有 15 亿根稻草 (参数) 时,你要如何大海捞针?你无法手动测试每一个权重。你需要一种自动化的方法来搜索网络。

作者使用了一种称为可微掩码 (Differentiable Masking) 的技术。

1. 掩码 (The Mask)

对于参数为 \(\theta\) 的预训练模型,研究人员希望学习一个由 0 和 1 组成的二值掩码 \(m\)。

  • 如果 \(m_i = 1\),权重保留。
  • 如果 \(m_i = 0\),权重被移除 (剪枝) 。

剪枝后的网络表示为 \(f(x, m \odot \theta)\),其中 \(\odot\) 是逐元素相乘。

2. 使其可学习 (Gumbel-Softmax 技巧)

问题在于,你不能对二值 (0 或 1) 使用标准的反向传播 (梯度下降) 。阶跃函数是不可微的。为了解决这个问题,研究人员为每个参数分配一个实数值分数 \(l_i\),并使用 Gumbel-Softmax 分布来模拟二值选择,同时允许梯度流动。

他们计算连续分数 \(s_i\) (为 1 的概率) :

显示使用 Sigmoid 和 Gumbel-Softmax 松弛计算连续分数 s_i 的公式。

然后,为了在前向传播中实际将其作为硬掩码 (hard mask) 应用,他们使用了“直通估计器” (straight-through estimator) 。这在计算时将分数四舍五入为 0 或 1,但在反向传播时保留连续分数的梯度。

显示使用直通估计器对掩码进行二值化的公式。

这个巧妙的数学技巧允许系统使用标准的训练循环来“学习”应该关闭哪些权重。

目标函数的三大支柱

为了找到正确的子网络,研究人员使用由三个不同目标组成的损失函数来训练掩码。这是论文贡献的核心。

目标 1: 抑制 (抹杀目标知识)

首要目标是让模型“遗忘”目标事实。研究人员通过最大化模型在目标知识上的预测与原始置信预测之间的差异来实现这一点。具体来说,他们最小化模型输出与均匀分布之间的 KL 散度。

简单来说: 当被问到“咖啡馆是一个……”时,模型不应该说“餐厅”。它应该像从字典里随机猜测一样困惑。

抑制损失公式: 均匀分布与掩码模型在目标数据上输出的 KL 散度。

目标 2: 保持 (不要破坏大脑)

如果我们只优化抑制,掩码可能会直接关闭模型中的每一个权重。一个什么都不知道的“脑叶切除”模型在技术上也解决了抑制任务!

为了防止这种情况,他们加入了保持损失。剪枝后的模型在对照知识 (Control Knowledge) 和通用语言建模 (Language Modeling) 数据上的表现必须与原始预训练模型完全一致。

保持损失公式: 原始模型与掩码模型在对照数据上的 KL 散度。

这迫使掩码必须像外科手术般精准。它可以切断“咖啡馆”的连接,但不能触碰“语法”或“汽车”的连接。

目标 3: 稀疏性 (保持精简)

最后,他们希望子网络尽可能小。他们添加了一个正则化项,惩罚保留权重的数量。

稀疏性损失公式: 掩码参数 Sigmoid 分数的平均值。

最终公式

结合以上几点,最终的损失函数使用 lambda (\(\lambda\)) 权重来平衡这些相互竞争的利益:

最终损失公式: 抑制、保持-KG、保持-LM 和稀疏性损失的加权和。

实验结果: 效果如何?

研究人员在不同大小的 GPT-2 模型 (Small, Medium, Large, XL) 上测试了这种方法。他们比较了两种方法:

  1. 权重掩码 (Weight Masking): 剪枝单个连接 (权重) 。
  2. 神经元掩码 (Neuron Masking): 剪枝整个神经元 (连接到一个特定神经元的所有权重) 。

结果 1: 权重掩码是王道

结果令人信服。如表 2 所示, 权重掩码方法发现了极其稀疏的子网络——保留了近 98.6% 的模型完整性 (意味着“知识”仅存储在不到 1.4% 的权重中) 。

表 2: GPT2-small 的子网络发现结果。显示权重掩码在 TargetKG 上具有高稀疏性和高 Delta PPL。

请看 TARGETKG \(\Delta\)PPL 列。这衡量了模型在目标知识上的表现变得有多“糟糕” (困惑度增加了多少) 。590.9 的分数意味着模型对目标事实变得完全一无所知。

关键在于,请看 CONTROLKG \(\Delta\)PPL 。 它接近于零 (-0.2) 。这意味着模型理解其他事实的能力几乎没有变化。

神经元掩码 (掩盖整个神经元) 的效果要差得多。它难以将目标知识与对照知识分离开来,这表明单个神经元是“多义的” (polysemantic) ——它们同时持有多个不同的概念。你不能在不杀死其他事实的情况下杀死一个神经元来移除一个事实。然而,单个权重提供了所需的细粒度控制。

结果 2: 具备扩展性

该方法不仅仅在小模型上侥幸成功。研究人员在 GPT-2 Large 和 XL 上进行了测试,发现了结果的一致性。

表 16: 较大型模型上每个 KG 的权重掩码子网络发现结果。显示 Medium、Large 和 XL 模型的目标困惑度均有增加。

表 16 中,我们可以看到,即使对于 GPT-2 XL (15 亿参数) ,该方法也成功地使目标知识的困惑度飙升 (高 \(\Delta\)PPL) ,同时保持对照知识的稳定。

记忆的解剖: 知识在哪里?

那么,这些“知识关键”权重藏在哪里呢?作者分析了发现的子网络的结构。

他们发现关键权重并非随机分布。它们显著地聚集在注意力头 (Attention Heads) 中,特别是在被掩盖的第一层和最后一层。

图 3: 权重掩码的平均模块掩码密度。显示第 7 层和第 12 层的密度较高,特别是在注意力输出部分。

图 3 显示了掩码密度的热力图。颜色越亮意味着移除的权重越多 (即对知识至关重要) 。你可以看到第 7 层和第 12 层有明显的条带,特别是在注意力输出 (Attn-Out) 和 键/查询/值投影 (Attn-Wq, Wk, Wv) 部分。

这与先前的机械可解释性研究一致,即早期层处理低级特征,后期层整合特定输出,而注意力头充当信息的路由机制。

深入探究注意力头

当放大到特定的注意力头( 图 5 )时,局部化变得更加明显。

图 5: 不同注意力层中不同头的密度百分比。显示了特定头 (如第 7 层的 10 号头) 的峰值。

注意那些明亮的黄色方块了吗?对于“representation” (表示) 知识图谱 (中间一行) ,第 7 层的 10 号头 (Head 10) 被点亮了。这表明这个特定的注意力头在处理该特定类型的关系型知识中发挥了巨大的、不成比例的作用。

子网络是真实的还是仅仅是一种破解?

该领域的一个常见批评是“虚假相关性”。掩码是真的找到了知识电路,还是只是找到了一种“黑客手段”来破坏模型对这些特定句子的输出?

为了测试这一点,作者进行了敏感性分析 (图 2) 。他们取出剩余的模型 (移除了知识的模型) ,开始随机剪枝更多的权重。

图 2: 在剩余的 GPT2-small 模型上移除和添加参数。显示目标知识不像随机基线那样容易恢复。

如果子网络只是一个脆弱的“黑客手段”,你可能会预期当你进一步干扰模型时,知识会突然重新出现或者模型会崩溃。相反,抑制效果保持强劲。 TargetKG 的困惑度 (第一列的深色线) 保持在搞水平,而Control (对照组) 的困惑度 (第二列) 只是缓慢下降,与随机基线相匹配。这表明手术确实是结构性和精确的。

子网络的唯一性

另一个迷人的发现: 如果你用不同的随机种子运行实验三次,你会找到相同的权重吗?

令人惊讶的是, 不会

图 8: 在三个不同随机种子下识别出的三个子网络的参数重叠韦恩图。显示重叠非常低 (交集很小) 。

图 8 所示,种子之间的重叠 (交集) 非常小——通常不到 4%。

这意味着 LLM 中的知识是高度冗余的。对于“咖啡馆是餐厅”这一知识,并不只存在一条电路。很可能存在许多并行路径。掩码算法找到了一组足以破坏知识的权重,但不同的种子找到了不同的路径。要完全擦除一个概念,你可能需要针对这些子网络的并集。

后果: 模型能重新学习吗?

如果我们通过手术移除了“CommonsenseQA” (常识问答) 所需的知识,模型能通过上下文或微调重新弄明白吗?

研究人员在 CommonsenseQA 数据集上对此进行了测试。他们移除了对特定 ConceptNet 关系至关重要的子网络,然后尝试在该任务上微调模型。

表 6: CommonsenseQA 上的准确率。显示当移除权重掩码子网络时,Filtered 集上的性能显著下降。

表 6 中的结果非常鲜明。

  • 完整模型: ~37-48% 的准确率。
  • 移除权重掩码后: 在需要被抑制知识的问题上 (“Filtered”列) ,准确率显著下降 (取决于微调方法,下降幅度在 -6.8% 到 -14.4% 之间) 。

即使经过微调 (LoRA 或 Head Tuning) ,模型也难以恢复丢失的知识。这似乎表明,一旦这些关键权重消失,模型就失去了有效处理或重新学习这些特定关系所需的底层“脚手架”。

结论: 前路

这篇论文为“机械可解释性” (Mechanistic Interpretability) 迈出了令人信服的一步。它让我们不再将 LLM 视为神秘的黑盒,而是将其视为可发现的、功能性电路的集合。

关键要点:

  1. 知识是稀疏的: 我们可以通过移除模型中 <2% 的权重来删除特定事实。
  2. 权重优于神经元: 单个权重是知识存储的原子单位;神经元太过宽泛且具有多义性。
  3. 冗余性: 存在编码相同知识的多条独立电路 (种子间低重叠证明了这一点) 。
  4. 迁移: 移除这些权重确实会损害模型在下游推理任务中使用该知识的能力。

对于学生和研究人员来说,这开辟了令人兴奋的途径。我们能否构建“反学习”算法来让 AI 更安全?我们能否在不重新训练整个模型的情况下更新旧知识 (例如,“首相是……”) ?通过识别这些知识关键子网络,我们离真正理解机器中的幽灵又近了一步。