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

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

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

方法

模型定义为文中式(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\}$。

模型推理机制

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

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

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

(c)通过注意力模块使用读取的记忆来生成输出，也就是图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、根据之前的模型推理机制的分析，作者先通过式(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$的改变量求和的结果。