论文解读：MASS-EDITING MEMORY IN A TRANSFORMER（MEMIT）-CSDN博客

论文发表于人工智能顶会ICLR（原文链接）。在模型编辑方法中，过去工作主要局限于更新单个事实。因此，基于ROME，本文开发了MEMIT，在大模型GPT-J（6B）和GPT-NeoX（20B）上实现了数千的批量编辑。

阅读本文请同时参考原始论文图表。

1 方法#

模型定义为文中式(1)，其中 $[x_{[1]},…,x_{[E]}]$ 表示长度为 $E$ 的输入句子， $x_{[t]}$ 表示模型输出单词。模型层之间状态的计算表示为式(2/3/4)，将模型最后一层关于输入句子最后一个token的状态映射到词汇空间就是 $x_{[t]}$ 。本文主要考虑GPT-J的架构来介绍方法，其中FFN和注意力模块并行，而不是使用注意力模块的输出输入FFN（当然后面介绍的MEMIT方法可以适用到其它LLM架构上）。

对于一个事实 $(s,r,o)$ ，模型输入包含头实体 $s$ 和关系 $r$ 的句子，输出头实体 $o$ 。模型编辑就是让模型关于包含 $(s,r)$ 的句子输出 $o$ 变成另一个 $o'$ 。本文的目标是同时对多个事实进行编辑，对同时编辑的事实构成的集合 $\mathcal{E}$ 做了一个限制，如式(5)所示，即事实之间不能有冲突。

根据ROME论文的实验结果，对于某个prompt $p_i$ ，本文只考虑其中主体 $s$ 的最后一个token的中间层状态 $h_i^l$ 、对应的FFN激活 $m_i^l$ 和注意力模块激活 $a_i^l$ 对模型输出的影响，此时 $i$ 为prompt的编号。另外，如图3所示（ROME的实验），由于不止一个中间层对模型预测有影响，因此同时考虑多个中间层相应激活对预测的影响。比如对于GPT-J， $l\in \mathcal{R}=\{3,4,5,6,7,8\}$ 。

1.1 模型推理机制#

根据模型的状态计算式(2)，可以得到式(6)，即每一层的输出状态是初始状态加上其前面层的FFN和注意力模块激活。根据之前ROME实验（ROME论文图1e/f/g）的观察，作者认为模型的推理机制如图2所示：

(a)模型先使用注意力机制把主体 $s$ 的信息汇集到 $s$ 的最后一个token（Jordan）。

(b)通过模型各层FFN根据主体 $s$ 的信息逐步读取相关的记忆并加入潜在表示。

(c)通过注意力模块使用读取的记忆来生成输出，也就是图2所示的信息通路。

1.2 批量参数更新#

和ROME类似，对于第 $l$ 层的FFN的第二层权重，在预训练后满足式(7)，通过求导得到方程式(8)。其中 $K_0=[k_1,…,k_n],M_0=[m_1,...,m_n]$ 。当要添加新知识 $K_1,M_1$ 时，就是把它们拼接后进行优化，即式(10-13)。最终得到 $W_0$ 的改变量 $\Delta$ 为式(14)。其中 $C_0=K_0K_0^T$ 定义为期望式(15)， $\lambda=1.5\times 10^4$ 。注意MEMIT的优化定义与ROME不同。

1.3 多层参数批量更新#

1、根据之前的模型推理机制的分析，作者先通过式(16)优化得到主体 $s$ 最后一个token在第 $L$ 层关于待修改事实 $(s_i,r_i,o_i)$ 的表示 $z_i$ 。其中 $L=\max(\mathcal{R})$ 表示对预测有影响层的最大层数， $h_i^L$ 表示模型关于 $(s_i,r_i)$ 在该位置的原始表示。也就是优化一个残差值 $\delta_i$ ，使得 $z_i=h_i^L+\delta_i$ 。 $x_j$ 表示prompt的前缀。

2、获得残差 $\delta_i=z_i-h_i^L$ 后，就是修改 $\mathcal{R}$ 中每层FFN的权重 $W_{out}^l$ ，使得模型关于 $(s_i,r_i)$ 的表示 $\hat{h}_i^L$ 尽可能接近 $z_i$ ，也就是优化式(17/18)。修改权重需要获取每个权重对应的新的键 $k_i^l$ 和值 $m_i^l$ ，并且由于前一层的权重修改会影响后层的输入，因此需要从 $\mathcal{R}$ 的第1层到第最后一层按顺序更新权重。每层的键可以直接通过前向传播得到，即式(19)。值则是键 $k_i^l$ 经过权重 $W_{out}^l$ 映射后加上残差 $r_i^l$ ，如式(20)所示。作者将第 $L$ 层的残差 $\delta_i=z_i-h_i^L$ 分配给 $\mathcal{R}$ 中的每一层，那为什么分母是 $L-l+1$ ，而不是 $L$ 呢？这是因为MLP的输出 $m_i^{l}$ 改了，会导致下一层的注意力输出 $a_i^{l+1}$ 也改了，所以总体改变量并不是直接对 $m_i^l$ 的改变量求和的结果。