手术刀式改写大模型的记忆:ROME 与 MEMIT 的知识编辑术
如果模型记错了一件事,你会怎么修?重新训练太贵,写个 if-else 打补丁太蠢——那有没有可能,直接在模型的”大脑”里找到那条记忆,动几个数字就把它改对?
故事从这里开始
假设你正在维护一个已经上线的大语言模型。有一天用户反馈:模型一直说”美国现任副总统是迈克·彭斯”。这曾经是对的,但现在已经过时了。或者更麻烦一点,模型一直坚持”Arneb 星在天鹰座”,但它其实在天兔座——这是训练数据里的一个错误,模型学会了一个从未存在过的”事实”。
面对这种问题,你有哪些选择?
第一个选择:重新训练整个模型。但一个千亿参数的模型可能耗费数百万美元和几个月时间训练出来,为了修一条记忆重新跑一遍,显然不现实。
第二个选择:用少量数据做微调(fine-tuning),让模型在这几条新数据上多学几步。但微调是个粗暴的工具——你在教它记住一句新句子的时候,很容易把它周围本来记得好好的东西一起带歪,这就是所谓的”灾难性遗忘”。而且微调容易让模型学会的只是”背下这句话怎么说”,而不是真的把这条知识内化——问它换一种问法,它可能又说错了。
第三个选择,也是这篇文章要讲的:直接在模型参数里做”精确手术”。找到这条事实具体存在于模型的哪个角落,然后只改那一小块,就像神经外科医生用探针定位到大脑里管理某个特定记忆的那几个神经元,然后精确地调整它们,而不碰其他任何区域。
这听起来像科幻小说,但 2022 年一篇叫做 ROME(Rank-One Model Editing)的论文,以及它 2023 年的续作 MEMIT(Mass-Editing Memory in a Transformer),真的做到了这件事——而且效果好到让人意外。这两篇论文不仅提供了一个实用工具,更重要的是,它们逼着我们回答一个更根本的问题:大模型的”知识”到底存在参数的哪里?以什么形式存在?
这篇文章就带你一步步搞清楚这件事:模型的事实记忆藏在哪、怎么把它找出来、怎么精确地改写它,以及——这套方法真的完美吗?
第一课:模型的”记忆”住在哪个房间里?
问题是什么
在你能修改一条记忆之前,你首先要知道它住在哪。这听起来像常识,但对一个几十上百层、每层几十亿参数的 Transformer 来说,”知识住在哪”完全不是一件显而易见的事。
一个 Transformer 每处理一个词,都会经过很多层。每一层里又有两个主要的计算模块:自注意力(Self-Attention)和前馈网络(Feed-Forward Network, FFN,也叫 MLP)。当模型读到”埃菲尔铁塔位于”这句话并预测下一个词”巴黎”的时候,这个正确答案到底是在哪一层、哪个模块里被”想出来”的?是靠 attention 东张西望、把”埃菲尔铁塔”和”巴黎”这两个词联系起来?还是靠 FFN 里存的某种固定知识,就像查字典一样直接查出来?
如果不知道答案,你想”编辑”这条知识就无从下手——就像医生不知道肿瘤长在哪,怎么下刀?
直觉:用”捣乱实验”找到关键位置
2022 年 MIT 和东北大学的研究者 Kevin Meng、David Bau 等人想出了一个很聪明的实验方法,叫做因果追踪(Causal Tracing)。
想象你在检修一台老式收音机,你怀疑某个电路模块坏了导致某个频段收不到台。你会怎么排查?一种办法是:先把整台机器的电路都断开(相当于”捣乱”),让它彻底收不到任何频段;然后一个一个地把电路重新接回去,每接回一个就测试一下频段有没有恢复。哪个电路一接回去信号立刻就正常了,那个电路就是关键。
因果追踪就是这个思路,只不过对象换成了神经网络内部的”隐藏状态”(hidden state,你可以理解成网络在处理某个词时,那一层、那个位置上的一组数字)。具体流程是这样的:
- 先让模型正常处理一遍”埃菲尔铁塔位于___“,记录下每一层、每个位置上的所有隐藏状态(这叫”干净运行”)。
- 然后再跑一遍,但这次故意把”埃菲尔铁塔”这几个词的表示加上随机噪声,把它”搞乱”,这样模型就没法正确回答了(”捣乱运行”)。
- 接着,一个一个地,把”干净运行”里记录下来的某个具体位置的隐藏状态,”贴”回”捣乱运行”里对应的位置,看这一贴之后模型的回答有没有恢复正确。
如果贴回某个特定层、特定位置的状态之后,模型立刻”想起来”正确答案是”巴黎”了,那说明那个位置就是”关键记忆存储点”——它承载了让模型能正确回答这个事实所需的决定性信息。
技术细节:两个”因果热点”
把这套实验对着成千上万个事实性问题跑一遍、统计结果,Meng 等人发现了一个清晰的模式:存在两个”热点”(causal site)。
第一个热点在中间层的 FFN 模块,出现在处理主语最后一个词(比如”埃菲尔铁塔”里的”塔”这个 token)的时候。这个热点是”早期”的——发生在句子处理的相对靠前阶段。
第二个热点在更靠后层的 attention 模块,出现在句子末尾(要预测答案的那个位置)。这个热点是”晚期”的。
这个结果暗示了一个故事:模型好像是在处理主语(”埃菲尔铁塔”)的时候,先在中间层的 FFN 里把跟这个主语相关的知识”提取”出来,存进当前的隐藏状态里;然后再靠后面层的 attention,把这份提取好的知识”搬运”到句子末尾,用来生成最终答案。可以把它想象成:FFN 是”查字典的人”,attention 是”跑腿把查到的答案送过去的人”。
下面这张图展示了这个”两阶段”的信息流动:
这个发现本身就很有价值——它告诉我们,”巴黎”这个答案不是模型东拉西扯凑出来的,而是有一个相当明确、可定位的计算路径。但更有意思的问题是:既然知道了知识”住”在中层 FFN 里,能不能直接改写它?
第二课:把 FFN 当成一本”活字典”,直接改词条
问题是什么
知道了知识大致的存储位置只是第一步。真正的挑战是:FFN 是一个由几百万个参数组成的巨大矩阵运算,你怎么在不破坏其他所有已存知识的前提下,只精确地改写”埃菲尔铁塔在巴黎”这一条,变成”埃菲尔铁塔在罗马”?
如果直接暴力修改一堆参数(比如用梯度下降在这一条新数据上微调),大概率会牵连到旁边存的其他知识——就像你想擦掉字典里一个词条的解释,结果笔一滑,把隔壁几个词条也划花了。
直觉:FFN 本来就是个”键值对查找表”
这里要先讲清楚一个更早、更基础的发现——来自 Geva 等人 2021 年的工作《Transformer Feed-Forward Layers Are Key-Value Memories》。他们指出,FFN 的两层矩阵运算,本质上就是一个”键值对存储系统”:第一层矩阵的每一行像一个”钩子”(key),负责识别某种输入模式(比如”当前词是某个特定城市的地标建筑名”);如果这个钩子被触发(激活值变高),第二层矩阵对应的那一列(value)就会把某个方向的信息”注入”进当前的隐藏状态里,把输出往某个特定的词(比如”巴黎”)上推。
用更贴近生活的类比:想象 FFN 是一本巨大的活页词典,每一页有一个”触发词”(key)和这一页写的”解释”(value)。当你读到”埃菲尔铁塔”这个主语时,词典里恰好有一页专门记录”和埃菲尔铁塔相关的信息”,这一页被翻到、被读出来,”巴黎”这个信息就被写进了当前的计算流程里。
如果这个类比是对的,那么”编辑一条知识”这件事,就可以翻译成一个非常具体、非常数学的操作:找到这一页词典对应的那对(key, value),把它的 value 换掉。
技术细节:一次”秩一(rank-one)”修改
这就是 ROME 的核心操作。用数学的语言说,FFN 里从输入到输出的那个映射可以近似看成一个矩阵 $W$,它把一个”key 向量” $k$(代表”埃菲尔铁塔”这个主语)映射到一个”value 向量” $v$(代表”位于巴黎”这条知识):
\[Wk \approx v\]翻译回人话:矩阵 $W$ 就是那本词典的全部内容,而”用 $W$ 乘以 $k$ 得到 $v$”,就是”翻到 key 对应的那一页,读出上面写的东西”这个动作的数学表达。
现在,如果我们想把这条知识改成 $v^*$(比如”位于罗马”),同时不动 $W$ 对其他任何 key 的映射关系怎么办?答案是给 $W$ 加上一个秩一(rank one)矩阵——也就是一个只用一个向量的外积就能表示的矩阵:
\[W^* = W + \Lambda k^T\]这里 $\Lambda$ 是根据”我们想要的新 value 减去旧 value 应该往哪个方向调”计算出来的一个向量。秩一矩阵的神奇之处在于:它对”和 $k$ 方向接近”的输入影响很大,但对”跟 $k$ 几乎垂直”的其他 key(也就是其他跟”埃菲尔铁塔”完全不相关的主语)几乎没有影响。这就是为什么这次修改能做到”手术式精确”——它几乎只碰到了那一个词条,没有波及词典的其他页。
ROME 论文里给的实验非常直观:他们真的教会了 GPT-J 一个假事实——”埃菲尔铁塔位于罗马”。有意思的是,编辑之后的模型不是简单地”背”下这句话——如果你问它”要去看埃菲尔铁塔应该坐飞机去哪个城市”,它会回答”罗马”;如果你问”埃菲尔铁塔附近有什么景点”,它列出的是罗马的地标。这说明模型不是学了一句台词,而是真的把这条”信念”整合进了它处理各种相关问题的方式里——这正是文章开头强调的”知道”和”背诵”的区别。
这里有个关键的区分标准值得专门讲一下,因为它是判断编辑是否”真的成功”的黄金准则:
- 特异性(specificity):改了这一条,其他不相关的记忆不应该被带歪。比如学会了”埃菲尔铁塔在罗马”之后,不应该顺带以为所有其他景点也都在罗马了。
- 泛化性(generalization):这条新知识应该能在各种不同的问句、上下文里都体现出来,而不只是在原始那句训练语句里才成立。
ROME 论文专门做了对比实验:如果编辑的是”早期热点”(中层 FFN),特异性和泛化性可以同时做得很好;但如果编辑的是”晚期热点”(后层 attention),模型虽然对原始那句话回答正确,却完全没法泛化到不同的问法——这也从另一个角度印证了因果追踪找到的位置确实是”真正存放知识”的地方,而不只是”临时经过”的地方。
第三课:从改一条到改一万条——MEMIT 的规模化难题
问题是什么
ROME 证明了”精确改写一条记忆”是可行的。但现实世界里,一次知识更新往往不是孤立的一件事。比如某个知识库里几百个国家的领导人在某年集体换届,或者你发现训练数据里同一类错误重复出现在几千条不同的事实上——你总不能一条一条地手动跑 ROME、跑一千次。
而且,如果真的对同一层的 FFN 权重矩阵反复做上千次秩一修改会怎样?答案是:模型会崩。因为每一次秩一修改都会给这个权重矩阵叠加一点”扰动”,扰动叠加得太多,这个矩阵就会离它原来学到的、能够正确处理成千上万种其他情况的那个状态越来越远。就像你不断在同一本词典的同一页上反复涂改、叠字,改到第一千次的时候,这页纸已经面目全非,别的知识也遭了殃。
直觉:把负担分给多个”仓库”,每个仓库只挪一点点
MEMIT(2023 年,ROME 团队的续作)给出的方案很直观:既然单个仓库扛不住上千次改动,那就把改动的总”重量”分摊到多个仓库(多个层)上,每个仓库只承担一小部分。
回到因果追踪的发现——记得前面说的,知识不是只在单独一层被”提取”,而是在一个”中间层范围”内被逐层加工、累积的(比如 GPT-J 里是第 3 到第 8 层)。MEMIT 的做法是:把原本需要塞进单一层的那个巨大更新量,拆分成一份份小的更新,分别加到这个范围内的每一层上。这样每一层承受的扰动都很小,不会把某一层的权重矩阵推得太离谱,但当所有层的贡献累加起来(想象成接力棒一层层往下传,每一层加一点点力),最终堆叠出的效果等价于完成了整条知识的更新。
用一个类比来说明这个思路的巧妙之处:如果你要往一个已经装满的箱子里塞进很多新东西,一种办法是硬塞进第一层,把第一层压得变形;另一种更聪明的办法是把这些新东西拆散,分别塞进箱子的第一层、第二层、第三层……每层塞一点,整箱子看起来没有明显变形,但总容量确实增加了。MEMIT 走的是第二条路。
技术细节:跨层分摊更新量
下面这张图展示了 ROME 和 MEMIT 在”编辑范围”上的核心差别:
具体的推导有点技术性但故事很简单:MEMIT 先用因果追踪确定好这一批要编辑的知识需要落在哪个层的范围(比如 layers 3-8),然后计算出如果只用最后一层承担全部更新,需要多大的改动量;接着按某种权重把这个总改动量拆分成一系列小份,依次施加在范围内的每一层上,让每一层各自承担一部分”责任”,但保证从整体上看,当信息流过整个范围后,最终累积的效果等价于完成了全部编辑目标。
实验结果证明这一招确实管用:ROME(以及类似的单层方法)在编辑数量增加到几十上百条之后就开始明显崩溃——生成的文本变得不流畅、无关知识被破坏;而 MEMIT 在 GPT-J(60亿参数)和 GPT-NeoX(200亿参数)上,同时编辑几千条事实,依然能保持相当好的准确率、泛化性、特异性和语言流畅度,比之前的方法在编辑规模上高出了几个数量级。
论文里衡量成功程度用了一个叫”编辑得分”的指标,是效力(能不能正确说出新事实)、泛化性、特异性三者的调和平均数——这个组合评分能防止一个方法通过”死记硬背”(效力很高但泛化性差)或”啥都不改”(特异性满分但效力为零)来作弊。
这套方法真的完美吗?两个诚实的警告
任何研究报告如果只讲”这个方法多厉害”,那是在偷懒。ROME/MEMIT 这套”先定位、再编辑”(locate-then-edit)范式提出之后,后续有两篇很重要的论文,分别从不同角度戳破了它看起来完美的表面,这两点值得认真讲清楚——因为它们揭示了一件更深刻的事:知道”知识存在哪”和”应该改哪里才最有效”,其实是两个不同的问题。
警告一:定位准,不代表编辑那个位置最有效
2023 年 NeurIPS 的一篇论文《Does Localization Inform Editing?》做了一个很反直觉的实验。他们发现:因果追踪告诉你的”这个事实是从哪个层被读取出来的”这个信息,其实并不能预测”你应该编辑哪一层才能最有效地改写这个事实”。你完全可以成功编辑一个因果追踪认为”跟这条知识毫无关系”的层,效果甚至可能不比编辑”官方指定”的层差。
这个发现的意义在哪?它提醒我们,”理解模型内部信息怎么流动”(可解释性研究关心的问题)和”知道怎么最有效地改变模型的行为”(模型编辑关心的问题),可能根本不是同一件事。 ROME/MEMIT 的成功不完全来自它们对因果追踪结论的忠实遵循——事实上 MEMIT 选择编辑一个层的”范围”而不是因果追踪指出的单一峰值层,这个工程上的权衡本身可能比精确的定位理论更重要。这是一句很诚实但很重要的话:一套方法work,不完全等于我们对它为什么work的理解是对的。
警告二:改一条事实,连带的其他事实没跟着改
第二个警告来自 2023 年的 RippleEdits 论文。它指出一个此前所有评估都忽略的问题:如果你把模型里”Jack Depp 的父亲是谁”这条事实,从别人改成 Johnny Depp,那么逻辑上,”Jack Depp 的姐妹是谁”这条事实也应该跟着自动更新(变成 Johnny Depp 的另一个孩子 Lily-Rose Depp)。这就叫”连锁反应”(ripple effect)。
研究者专门构建了一个包含五千条编辑、覆盖多种连锁关系的测试集,拿 ROME、MEMIT 等主流方法一测,结果发现:这些方法在原本的”效力/泛化性/特异性”三项指标上表现很好,但一旦你去检验它们有没有正确地把连锁的、逻辑相关的事实也一并更新,几乎全部失败。更讽刺的是,一个最朴素的办法——不改任何权重,只是在提问的时候顺便把新事实作为上下文提示词塞进去——在这项连锁一致性测试上,反而比精心设计的权重编辑方法得分更高。
这说明了什么?ROME/MEMIT 做的编辑更像是”在词典的这一页换了一个词条”,但没有意识到词典里还有另外几十页引用了这个词条、依赖了这个词条的旧值。它们做到的是”局部替换”而不是真正意义上的”信念修正”(belief revision)——模型的知识库里其实存在大量隐性的相互依赖关系,而目前的编辑方法完全没有触及这个层面。
这意味着什么
我们从一个朴素的问题出发——模型记错了一件事,能不能精确地把它改对——一路走到了知识编辑这个研究方向的核心方法论,也走到了它目前最深的两个局限。回顾一下这条脉络:
因果追踪告诉我们,事实性知识不是随意散落在网络各处的,而是有相对集中、可定位的存储位置——中层 FFN 在处理主语时”提取”知识,后层 attention 把它”搬运”到句尾生成答案。FFN 之所以能承载这种角色,是因为它本质上就是一个键值对存储结构。ROME 利用这个洞察,把”编辑一条事实”变成了一次数学上干净、影响面极小的秩一矩阵修改。MEMIT 接着解决了规模化的问题:把改动分摊到多个层,让编辑上千条事实成为可能。
但两个后续研究提醒我们:这条”先定位、再编辑”的路线并非终点。定位知识存在哪,不等于知道去哪里编辑最有效;改对了一条孤立的事实,不等于模型的整个知识体系跟着自我修正了。这两个问题,今天仍然是知识编辑领域最活跃的开放挑战——比如后续工作提出的 PMET,就在尝试更精细地分离 attention 和 FFN 各自对隐藏状态的贡献,让编辑更”纯粹”;而如何解决连锁效应问题,目前还没有令人满意的答案。
这也带来一个更大的哲学式反思:当我们说一个语言模型”知道”某件事的时候,这个”知道”到底是什么?ROME 的实验提示我们,至少存在一种可以被定位、被量化、被精确改写的知识形式——它有具体的坐标(层数、位置),有具体的数学形式(一个键向量和一个值向量)。这和”知识”作为一种抽象的、逻辑一致的信念系统之间,还有相当的距离要走。也许,真正”教会”一个模型一件新事实,需要的不只是改写一个键值对,而是让这次改写正确地波及到它逻辑上牵连的一整张知识网络——而这,目前还没有一个方法能做到。
如果你在维护一个真实上线的模型,看完这篇文章后最实际的收获应该是:知识编辑技术确实能帮你低成本地修正孤立的过时/错误事实,但不要指望它能替代对模型进行系统性的知识一致性审查——改完一条记忆之后,那些”顺理成章应该跟着变”的关联事实,你可能还得自己想办法处理。