0. TL; DR

本文提出了MEMIT，一种能够直接更新大型语言模型中多个记忆的方法。MEMIT通过直接编辑Transformer模型的权重，实现了对数千个事实的批量更新，显著提升了模型的知识更新能力。实验表明，MEMIT在保持特异性和泛化能力的同时，能够高效地插入大量新记忆，为语言模型的知识更新提供了一种新的解决方案。

1. 背景介绍

大型自回归语言模型（如GPT）能够回忆起许多常见的事实，例如“蒂姆·库克是苹果公司的CEO”或“北极星位于小熊星座”。然而，这些模型在未更新时可能会回忆起过时的信息，并且缺乏更专业的知识。为了保持信息的新鲜感和可定制性，研究者们希望能够直接更新模型中的知识，而不是重新训练整个模型。现有的知识编辑方法大多局限于更新单个关联，而本文提出的MEMIT方法能够扩展到数千个关联的更新，适用于GPT-J（6B参数）和GPT-NeoX（20B参数）等大型模型。

2. 方法介绍

MEMIT方法基于因果干预分析（Causal Mediation Analysis），识别出在回忆事实时起关键作用的Transformer模块。具体来说，MEMIT发现中间层的前馈模块（MLP）在处理主题词的最后一个token时起着决定性作用。这些MLP层被选为编辑目标，因为它们在事实回忆中扮演着重要角色。特别地，编辑目标为MLP的第二个线性层$Wk=m$，该层将输入主题词的特征嵌入映射为输出词嵌入的概率分布。

在每个单独的层中，MEMIT的目标是存储大量记忆。假设层中已经存储了一些记忆，我们需要找到一个最优的单层更新，以最小化记忆关联的平方误差。假设预训练权重为$W_0$，原主题词的特征嵌入为$K_0$，原事实为$M_0$，则有$W_0K_0=M_0$。新主题词的特征嵌入为$K_1$，新事实为$M_1$，权重更新量为$\Delta$，则优化目标写作：

\[\begin{aligned} \min \quad &\|(W_0+\Delta)K_0-M_0\|_F^2 \\ \text{s.t.} \quad & (W_0+\Delta)K_1=M_1 \end{aligned}\]

或写作正规方程的形式：

\[\begin{aligned} (W_0+\Delta)(K_0K_0^\top+K_1K_1^\top)&=M_0K_0^\top+M_1K_1^\top \\ \Delta(K_0K_0^\top+K_1K_1^\top)&=M_1K_1^\top-W_0K_1K_1^\top \\ \Delta&=(M_1-M_0)K_1^\top (K_0K_0^\top + K_1 K_1^T)^{-1} \end{aligned}\]

MEMIT通过以下步骤实现多层更新：

计算目标向量 $z_i$：对于每个记忆 $i$，计算一个向量 $z_i$，使得在顶层 $L$ 的隐藏状态中完全传达新的记忆，优化 $z_i$ 以最大化模型对新对象 $o_i$ 的预测概率。
在多层中传播 $z_i$：从顶层 $L$ 开始，逐层向下传播 $z_i$，第$l$层新的事实增加$R_l=M_l-M_0=\frac{z_i-h_i^L}{L-l+1}$。
使用上述公式 $\Delta_l = R_l K_l^T (C_l + K_l K_l^T)^{-1}$ 更新每层的MLP权重。