0. TL; DR

本文研究了自回归Transformer语言模型（如GPT）中事实关联的存储和回忆机制，并发现这些关联对应于局部化的、可直接编辑的计算过程。研究者通过因果干预分析，揭示了模型中间层前馈模块在处理主题词时的关键作用，并提出了Rank-One Model Editing（ROME）方法，通过修改特定权重来更新事实关联。实验表明，ROME在标准的零样本关系抽取（zsRE）任务上表现出色，并在新的反事实断言数据集上同时保持了特异性和泛化能力。研究还表明，直接操纵计算机制可能是模型编辑的可行方法。

1. 背景介绍

大型语言模型（如GPT）能够预测关于世界的事实性陈述。例如，给定前缀“The Space Needle is located in the city of”，GPT能够准确预测出“Seattle”。这种事实性知识在自回归GPT模型和掩码BERT模型中都有所体现。然而，尽管这些模型在预测事实方面表现出色，但它们如何存储和回忆这些知识仍然不甚清楚。

为了更好地理解和编辑这些模型中的事实性知识，研究者们提出了两个问题：一是如何定位模型中存储事实的具体位置；二是如何直接编辑这些事实。这不仅有助于提高模型的可解释性，还能为模型的动态更新提供新的可能性。例如，当需要更新模型中的过时知识或纠正错误信息时，能够直接编辑模型中的事实将变得尤为重要。

2. 方法介绍

2.1 因果干预分析

为了定位模型中存储事实的具体位置，研究者们采用了因果干预分析（Causal Mediation Analysis）。具体来说，他们追踪了GPT模型中隐藏状态激活的因果效应，以识别对主题事实预测起决定性作用的具体模块。分析发现，中间层的前馈模块（MLP）在处理主题词的最后一个token时起着关键作用。

具体而言，研究者们通过以下步骤进行因果干预分析：

干净运行：将包含主题的事实性提示输入模型，收集所有隐藏状态激活值。
损坏运行：在模型运行前，对主题词的嵌入向量添加噪声，破坏模型对主题的识别能力，然后收集损坏状态下的激活值。
恢复运行：在损坏运行的基础上，选择性地恢复某些隐藏状态的干净值，观察这些状态对恢复原始预测的影响。

通过计算恢复运行与损坏运行之间的差异，研究者们能够量化每个隐藏状态对事实预测的因果贡献。实验结果表明，中间层MLP模块在处理主题词的最后一个token时具有显著的因果效应，这表明这些模块在事实回忆中起着关键作用。

2.2 Rank-One Model Editing（ROME）

基于因果干预分析的结果，研究者们提出了Rank-One Model Editing（ROME）方法，用于直接编辑模型中的事实关联。ROME的核心思想是通过修改中间层MLP模块的权重，来更新特定的事实关联。

具体步骤如下：

选择关键层和token：根据因果干预分析的结果，选择中间层MLP模块中处理主题词最后一个token的位置作为编辑目标。MLP模块的第二层可以建模为$WK=V$形式的线性联想记忆（linear associative memory）。
计算新的键值对：为了插入新的事实关联 $(s, r, o^*)$，需要计算一个新的键值对 $(k^*, v^*)$。其中，键 $k^*$ 由主题词的隐藏状态激活值计算得到，值 $v^*$ 通过优化目标函数得到，该函数旨在最大化模型对新对象 $o^*$ 的预测概率，同时最小化对主题本质的干扰。

\[\begin{aligned} \min \quad & \frac{1}{2}||\hat{W}K-V||_F^2 \\ \text{s.t.} \quad & \hat{W}k^* = v^* \end{aligned}\]

更新权重：使用Rank-One更新规则，将新的键值对 $(k^*, v^*)$ 插入到目标MLP模块的权重中，从而实现对事实关联的编辑。

采用最小二乘法求解：

\[\begin{aligned} L(\hat{W},v^*,\Lambda) & =\frac{1}{2}||\hat{W}K-V||_F^2 - \Lambda (\hat{W}k^* - v^*)\\ \frac{\partial L}{\partial \hat{W}} & =\hat{W}KK^T-VK^T-\Lambda (k^*)^\top = 0 \\ & \downarrow \\ \hat{W}KK^T &= VK^T+\Lambda (k^*)^\top \\ (\hat{W}-W)KK^T &= \Lambda (k^*)^\top \\ \hat{W} &= W + \Lambda (C^{-1} k^*)^T \end{aligned}\]

其中，$W$ 是原始权重矩阵，$C=KK^\top$ 是键向量的协方差矩阵，用于约束更新的幅度。$\Lambda$ 是一个向量，表示新键值对与原始权重的残差误差，其形式推导为：

\[\begin{aligned} \hat{W}k^* &= \left(W + \Lambda (C^{-1} k^*)^T\right)k^*\\ &= Wk^* + \Lambda (C^{-1} k^*)^Tk^*\\ &= v^*\\ & \downarrow \\ \Lambda &= \frac{v^* - Wk^*}{(C^{-1} k^*)^Tk^*} \end{aligned}\]

ROME具有以下优势：

局部化编辑：ROME通过修改特定层的权重，实现了对特定事实的局部化编辑，而不会对模型的其他部分产生显著影响。
高效性：ROME的编辑过程简单且高效，不需要对整个模型进行重新训练或微调。
泛化能力：实验表明，ROME编辑后的模型不仅能够准确预测新的事实，还能在不同的上下文中泛化，同时保持对未修改事实的准确性。

3. 实验分析

为了评估ROME的性能，研究者们使用了两个数据集：标准的零样本关系抽取（zsRE）数据集和新构建的COUNTERFACT数据集。

zsRE数据集：包含10,000条事实性陈述及其改写形式，用于测试模型在未见过的事实上的预测能力。
COUNTERFACT数据集：包含21,919条记录，涵盖了多样化的主题、关系和语言变体。该数据集旨在区分模型对新事实的表面性改变和深层次的语义改变。

在zsRE任务上，ROME与其他模型编辑方法进行了比较，包括微调（FT）、带约束的微调（FT+L）、知识编辑器（KE）和模型编辑网络（MEND）。实验结果表明，ROME在编辑效果上与这些方法相当，尽管其实现更为简单。

COUNTERFACT数据集用于更细致地评估模型编辑方法的性能，特别是在处理反事实断言时的能力。实验结果表明，ROME在保持特异性和泛化能力方面优于其他方法。

为了进一步评估ROME编辑后的模型生成文本的质量，研究者们还进行了生成文本的比较。实验结果表明，ROME编辑后的模型能够生成与新事实一致的文本，同时保持较高的流畅性和一致性。

为了更全面地评估ROME编辑后的文本质量，研究者们还进行了人类评估。评估结果显示，ROME在一致性上优于FT+L，但在流畅性上略逊一筹。