深入黑盒: 绘制大语言模型中的知识神经元图谱

像 GPT-4 和 Llama 这样的大语言模型 (LLM) 展现出了惊人的存储和回忆事实知识的能力。当你问 LLM“法国的首都是哪里？”时，它能毫不费力地检索出“巴黎”。但这些信息究竟存储在哪里？“巴黎”是存储在特定的神经元簇中吗？如果是，模型又是如何知道何时激活它们的？

理解知识存储的机制是机械可解释性 (mechanistic interpretability) 的“圣杯”。如果我们能精确定位负责特定事实的确切神经元，理论上我们就能够在无需昂贵重新训练的情况下，编辑掉幻觉或更新过时的信息。

然而，寻找这些神经元就像在数字大海中捞针。目前的归因技术，如积分梯度 (integrated gradients) 或因果追踪 (causal tracing) ，通常计算成本过高，无法应用于现代 LLM 中数以百万计的神经元。

在论文 “Neuron-Level Knowledge Attribution in Large Language Models” 中，研究人员 Zeping Yu 和 Sophia Ananiadou 提出了一个新的、高效的框架来解决这个问题。他们引入了一种静态方法来精确定位“值神经元” (存储知识的地方) 和“查询神经元” (触发知识的机制) 。他们的工作为 Transformer 模型内的信息流提供了一幅精细的地图。

背景: Transformer 的解剖结构

为了理解研究人员如何追踪知识，我们必须首先了解 Transformer 模型的基本单元。LLM 通过一系列层来处理输入句子 \(X = [t_1, t_2, ..., t_T]\)。每一层主要由两个子模块组成: 多头自注意力 (MHSA) 和前馈网络 (FFN) 。

信息流经一层的过程可以用残差连接方程来描述:

残差连接方程显示层输出为输入、注意力和 FFN 输出的总和。

这里，\(h_i^{l-1}\) 是上一层的输入，\(A_i^l\) 是注意力机制的输出，而 \(F_i^l\) 是 FFN 的输出。

最后，在模型的末端 (第 \(L\) 层，位置 \(T\)) ，输出向量被投影到词汇空间以预测下一个 token:

用于计算最终概率分布的 softmax 方程。

定义“神经元”

在生物大脑中，神经元是一个细胞。在 Transformer 中，这个定义更偏向数学。

1. FFN 神经元: 研究人员通常将前馈网络视为键-值 (Key-Value) 存储系统。FFN 层的输出是特定向量的加权和:

FFN 输出表示为神经元的加权和。

在这种语境下，一个神经元被定义为子值 (subvalue) 向量 \(fc2_k^l\) (第二个线性层的一列) 。它的激活，或称“系数组合分数” \(m_{i,k}^l\)，取决于输入与对应的“子键” (subkey) \(fc1_k^l\) (第一个线性层的一行) 的匹配程度:

FFN 神经元的系数分数计算。

2. 注意力神经元: 同样，注意力机制计算来自不同头的输出的加权和:

显示头部求和的注意力输出方程。

研究人员将神经元的定义扩展到了注意力头。他们将输出矩阵 \(W^o\) 的列视为“注意力子值” (神经元) ，将值矩阵 \(W^v\) 的行视为子键。

核心方法: 追踪知识流

核心挑战在于确定这数百万个神经元中，究竟是哪一个对特定预测 (例如预测“巴黎”) 做出了实际贡献。

现有指标的问题

一个常见的直觉是寻找具有高激活分数的神经元。然而，研究人员认为仅看激活是具有误导性的。一个神经元可能具有很高的激活分数，但提供的只是通用信息，无助于区分“巴黎”和“伦敦”。

为了理解这一点，我们必须看看当神经元向量 \(v\) 被加到残差流 \(x\) 中时，它是如何影响概率分布的。研究人员分析了“Softmax 前” (before-softmax，简称 bs) 值:

词汇表中所有 token 的 Softmax 前值向量。

特定词 \(w\) 的概率是使用 softmax 函数在这些值上计算得出的:

给定向量 x 时单词 w 的概率。

当神经元 \(v\) 将其信息添加到流中 (\(x + v\)) 时，概率会发生偏移:

给定向量 x 加上神经元 v 时单词 w 的概率。

关键在于，Softmax 前值的变化是线性的:

Softmax 前值的线性。

然而，由此产生的概率变化是非线性的 。 研究人员通过一个假设示例 (表 1) 说明了这一点，表明神经元的影响在很大程度上取决于残差流的现有状态 (\(x\)) 。一个简单地为所有 token 增加常数值的神经元可能根本不会改变排名。相反，一个原始值较低的神经元，如果它针对特定的正确 token 并抑制其他 token，可能会极大地改变概率。

解决方案: 对数概率增量

为了在不运行昂贵梯度的情况下捕捉这种复杂性，研究人员提出了一个新的指标: 对数概率增量 (Log Probability Increase) 。

他们将神经元 (或层向量) \(v^l\) 的重要性 (\(Imp\)) 定义为: 包含该向量时正确 token \(w\) 的对数概率与不包含该向量时的对数概率之差:

基于对数概率增量的重要性评分方程。

该指标既考虑了神经元的贡献，也考虑了残差流 (\(h^{l-1}\)) 的上下文。它提出的问题是: “关于正确答案，这个神经元消除了多少‘意外’？”

识别查询神经元

找到存储值的神经元 (“值神经元”) 只是战斗的一半。我们还需要知道是什么激活了它们。研究人员引入了一种寻找“查询神经元”的方法。

由于值神经元的激活取决于输入与子键之间的内积 (如公式 7 所示) ，研究人员计算了重要值神经元的子键与前几层神经元之间的内积。如果前一层的某个神经元与值神经元的键有很高的内积，它就被识别为查询神经元 。

这创建了一个完整的电路视图:

  1. 查询神经元 (通常在较浅层) 从输入中提取特征。
  2. 它们激活值神经元 (通常在较深层) 。
  3. 值神经元将事实信息写入残差流。

下图展示了这种架构，显示了从 FFN 查询神经元和注意力神经元一直到最终 FFN 值神经元的流程。

图解 Transformer 架构中的查询神经元、注意力神经元和值神经元。

实验与结果

为了验证他们的方法，研究人员使用 TriviaQA 数据集在两个模型上进行了测试: GPT2-large 和 Llama-7B。他们提取了模型正确预测了与六个类别 (语言、首都、国家、颜色、数字和月份) 相关答案的句子。

与其他方法的比较

他们将自己的“对数概率增量”方法与七种其他静态归因方法 (如原始概率、范数、系数分数等) 进行了比较。评估包括根据每种方法识别前 10 个最重要的 FFN 神经元，“关闭它们” (将参数设为零) ，并测量模型性能的下降。

逻辑是: 如果该方法真正找到了重要的神经元，关闭它们应该会破坏模型回答问题的能力。

结果是决定性的:

表格显示与其他方法相比，所提出的方法导致 MRR 和概率下降最大。

如表 2 (上图) 所示，提出的方法 (‘a’ 行) 导致了最显著的性能下降。对于 Llama-7B，仅移除 10 个神经元，正确 token 的概率就从 55.1% 暴跌至 9.2%。这大大优于简单的系数分数 (’e’ 行) 或向量范数 (’d’ 行) 等方法。

神经元位于何处？

利用这种经过验证的方法，研究人员绘制了重要神经元在模型层中的分布图。

图表显示 Llama-7B 中神经元在各层中的分布。

图 2 揭示了一个惊人的趋势: 重要的值神经元绝大多数集中在深层 (Llama-7B 的 20-32 层) 。

有趣的是，该图表比较了“对数概率增量” (红色) 与“概率增量” (绿色) 。“概率增量”方法严重偏向于最后几层。研究人员使用下面的理论曲线解释了这一点:

曲线显示对数概率增量与原始概率增量之间的差异。

图 3 显示，随着模型接近解决方案 (段索引增加) ，原始概率 (右图) 仅在最后阶段呈指数级上升。然而，对数概率 (左图) 的增长更为线性。这使得“对数概率增量”成为一种更敏感的工具，可以检测中深层而不仅仅是最后一层的重要贡献。

分析知识存储

研究人员按知识类型 (例如首都与颜色) 细分了他们的分析，以查看不同类型的事实是否以不同方式存储。

层级归因

下面的热图可视化了 GPT2 中不同层的重要性。颜色越深表示重要性越高。

GPT2 中前 10 个重要值层的热图。

图 4 (以及其中包含的表格) 证实了知识存储在深层。此外，它显示了语义聚类 。 注意“语言”、“首都”和“国家” (语义密集型任务) 是如何点亮相似的层 (例如第 26、30 层) 的。相比之下，“数字”和“颜色”依赖于不同的层集。这表明模型在物理上按语义类别组织知识。

同样的模式在 Llama-7B 中也成立，如图 5 所示:

Llama 中前 10 个重要值层的热图。

神经元级稀疏性

最深刻的发现之一是知识的稀疏性 。 研究人员发现，虽然模型有数十亿个参数，但查询的具体知识仅由其中极小的一部分处理。

表格显示了注意力和 FFN 层中头部神经元的重要性。

表 5 显示了累积重要性得分。“前 200 个”神经元 (占总数的极小部分) 的重要性得分几乎等于“所有”神经元的总和。这证实了知识并非广泛分布在整个网络中，而是局限在特定的、可检索的点上。

“查询”层分析

最后，研究人员观察了哪些层充当“查询”信号——即激活深层值神经元的层。

GPT2 中前 10 个重要查询层的热图。

图 6 (GPT2) 和图 7 (Llama，下图) 显示, 中深层注意力层在查询中起着巨大的作用。

Llama 中前 10 个重要查询层的热图。

这描绘了一幅清晰的信息处理流程图:

  1. 浅/中层: 处理语法和上下文 (查询神经元) 。
  2. 深层: 检索特定的事实关联 (值神经元) 。

当研究人员专门分析“查询”神经元时，他们发现与值神经元 (通常直接映射到“巴黎”等输出词) 不同，查询神经元的可解释性较差。它们似乎是作为抽象触发器发挥作用，而不是自己持有内容。

结论与启示

这项研究提供了一种静态、计算高效的方法来打开大语言模型的“黑盒”。通过区分值神经元 (存储) 和查询神经元 (钥匙) ，并使用对数概率增量作为度量，作者能够显著优于现有的归因方法。

主要结论包括:

  1. 深层存储: 事实知识主要存储在网络的深层。
  2. 语义局部性: 相似类型的知识 (如地理) 存储在重叠的区域。
  3. 稀疏性: 极少数神经元 (通常少于 300 个) 负责特定的预测。干预这少数几个神经元可以完全改变模型的输出。
  4. 角色分明: 中层充当查询，解锁存储在深层的事实。

这为什么重要? 这种粒度是模型编辑的基础。如果我们能可靠地定位负责特定幻觉或偏见关联的 300 个神经元，我们就可以在不重新训练整个模型的情况下对其进行手术式修改。这让我们向更安全、更可解释、更可靠的人工智能迈进了一步。