Mass-Editing Memory in a Transformer

TL;DR

MEMIT通过多层关联记忆更新,在GPT-J上一次编辑1万条事实,编辑分数达50.7。

cs.CL 🔴 高级 2022-10-14 22 次浏览
Kevin Meng Arnab Sen Sharma Alex Andonian Yonatan Belinkov David Bau
模型编辑 MEMIT Transformer 知识更新 大语言模型

核心发现

方法论

MEMIT将Transformer中介MLP视为键值联想记忆。先用因果追踪定位关键层,再通过梯度优化得到目标隐藏向量z;随后利用批量闭式更新Δ=R Kᵀ(C+KKᵀ)⁻¹,将残差均匀分配到多层MLP,并以协方差统计C≈λE[kkᵀ]保护旧记忆。

关键结果

  • 在zsRE的10,000条编辑上,GPT-J的MEMIT编辑分数为50.7,效能96.7%,改写泛化89.7%,特异性26.6%;优于FT-W的42.1、MEND的20.0和ROME的2.6。
  • 在COUNTERFACT的10,000条反事实编辑上,GPT-J的MEMIT取得编辑分数85.8、效能98.9%、改写泛化88.6%、特异性73.7%;GPT-NeoX上编辑分数为82.0。
  • 扩展曲线显示ROME在约32条编辑后明显退化,MEND在1,000条前失效;MEMIT可扩展至万条,但小规模时ROME泛化略强,体现软误差最小化与硬约束的差异。

研究意义

论文把知识编辑从单条或数十条事实推进到GPT-J 6B和GPT-NeoX 20B上的数千乃至万条事实,为新闻更新、专业定制和知识维护提供了不必重训全模型的路径。更重要的是,它把模型参数中的事实记忆与可解释的MLP联想结构联系起来,证明批量写入、泛化和局部性可以在一定程度上同时实现。

技术贡献

核心贡献包括:用因果中介分析选择GPT-J关键MLP层R={3,4,5,6,7,8};将ROME的单事实秩一编辑推广为批量线性联想记忆更新;通过目标向量z和逐层残差分配避免单层破坏;以正常方程推导闭式Δ并用未中心化键协方差近似旧记忆。该设计比顺序ROME更适合矩阵化的大规模编辑。

新颖性

相较ROME、MEND和SERAC主要处理单条至75条事实,MEMIT首次系统展示直接参数编辑可在20B级模型中批量写入数千至10,000条记忆。其根本新意不是新训练器,而是把“多层关键路径+批量联想矩阵+残差分摊”组合成可扩展算法。

局限性

  • MEMIT依赖关系无冲突、目标对象可表达且关键层定位正确;对“运动员参加何种运动”等关系仍较困难,且特异性在大规模编辑后会下降。
  • 10,000条编辑在当前实现中逐个优化z,MEMIT运行约7.44小时,远慢于MEND的98秒;作者指出这些优化可并行,但论文未完成该工程化。

未来方向

未来应并行化z优化,改进不同关系的层定位与编辑分配,并研究持续、多轮编辑后的冲突检测、灾难性遗忘和安全撤销。还需在更广泛模型、事实类型及真实动态知识流上验证长期稳定性。

AI 总览摘要

大型语言模型能记住大量事实,却难以快速替换过时信息。既有方法如ROME主要面向单条编辑,MEND等超网络方法通常只能处理少量事实;顺序重复编辑还会造成相互干扰。因此,论文提出MEMIT,目标是在不重训整个模型的情况下,一次写入数千条新记忆。

MEMIT把Transformer的中间MLP视为联想记忆。因果追踪先找到负责事实回忆的关键层;算法再为每条事实优化目标隐藏向量z,并把所需改变的残差分摊到多层。每层使用闭式批量更新Δ=R Kᵀ(C+KKᵀ)⁻¹,其中协方差统计近似原有记忆,从而在写入新事实时尽量保留旧知识。

实验显示,在zsRE的10,000条事实上,GPT-J的编辑分数为50.7,效能96.7%,改写泛化89.7%;在COUNTERFACT的10,000条反事实编辑上,GPT-J编辑分数85.8,GPT-NeoX为82.0。ROME在约32条后退化,MEND规模更大时失效。MEMIT仍有运行成本高、关系间表现不均和特异性下降等问题,但它证明了直接修改模型参数可以从“单点修补”走向大规模知识维护。

深度分析

研究背景

知识编辑旨在直接修改模型参数,以替换过时事实或加入专业知识。约束微调、KE、MEND、SERAC和ROME分别从梯度、超网络、路由或秩一更新入手,但通常只验证单条至75条编辑。MEMIT继承ROME和Geva等关于MLP键值记忆的观察,并进一步研究批量扩展。

核心问题

问题是同时编辑大量(s,r,o)事实,同时保持三种性质:原提示下成功、改写提示下泛化、相邻未编辑主体不被污染。单条顺序更新会累积参数扰动;超网络方法在规模扩大时可能失效;普通微调则容易损坏生成流畅性与原知识。

核心创新

  • ��关键路径:因果追踪在GPT-J定位R={3,...,8}。
  • ��批量更新:将MLP输出矩阵作为联想记忆,推导Δ=R Kᵀ(C+KKᵀ)⁻¹。
  • ��多层写入:用z表示完整目标,在各层均匀注入剩余残差。
  • ��激活重收集:每次改层后重新计算下游键,减少层间耦合误差。

方法详解

  • ��输入:编辑集合E={(sᵢ,rᵢ,oᵢ)}、模型G、关键层R和样本前缀。
  • ��目标构造:在顶层隐藏状态hᴸ处替换为zᵢ=hᴸᵢ+δᵢ;通过交叉熵优化δᵢ,使目标对象oᵢ概率最大。
  • ��键计算:在每层取MLP输入kᶫᵢ=E[σ(Wᶫinγ(h))]。
  • ��值计算:令剩余残差rᶫᵢ=(zᵢ−hᴸᵢ)/(L−l+1),形成目标值。
  • ��更新:用批量矩阵公式求Δᶫ,并按升序更新层;每层后重收集激活。
  • ��评估:以效能、Paraphrase、Specificity的调和平均衡量编辑分数。

实验设计

模型为GPT-J 6B和GPT-NeoX 20B;数据集为zsRE和含21,919条事实的COUNTERFACT。基线包括FT-W、MEND和顺序ROME。编辑规模从1扩展至10,000;指标包括Efficacy、Paraphrase、Specificity、ES、PS、NS、Generation Entropy和Reference Score。GPT-J典型协方差权重λ=1.5×10⁴。

结果分析

zsRE万条编辑中MEMIT得分50.7,显著高于FT-W 42.1、MEND 20.0和ROME 2.6。COUNTERFACT万条编辑中GPT-J达到85.8,ES/PS/NS为98.9/88.6/73.7;GPT-NeoX为82.0。FT-W虽有效能99.4%,但生成质量显著恶化;MEMIT在27类关系总体最佳,不过运动关系仍较难。

应用场景

可用于新闻事实更新、企业产品与人物资料定制、专业问答模型维护,以及根据用户反馈快速修正知识。实际部署需先检测冲突,限制编辑范围,并监控邻域特异性、生成熵和安全风险。

局限与展望

方法假设编辑不存在同一(s,r)对应不同对象的冲突,并依赖预训练模型中存在可利用的MLP记忆路径。不同关系的可编辑性不均,规模增大时特异性下降。当前逐条优化z使万条编辑耗时约7.44小时;未来应并行化、支持撤销与冲突管理,并评估连续编辑和更真实知识流。

通俗解读 非专业人士也能看懂

把语言模型想成一座巨大的图书馆。过去的方法像管理员每次只改一本书:改得少时很准确,但连续改几千本就会把相邻书架弄乱。MEMIT先找出图书馆中真正负责“查找事实”的几个核心书架,而不是随便修改所有房间。

对每条新事实,系统先设计一张清晰的“新卡片”,例如把某人和新的运动联系起来。然后,它把卡片内容分散写入多个核心书架,而不是集中塞进一个抽屉。每次写完一个书架,都会重新检查后面的书架,因为前面的变化会影响后续查找。

实验表明,它能在GPT-J中一次加入10,000条事实,并在原问题和改写问题中大多回答新答案,同时尽量不影响相邻人物。代价是写入过程较慢,而且某些关系特别难改;所以它更像大型图书馆的批量维护工具,而不是万能橡皮擦。

简单解释 像给14岁少年讲一样

想象你在玩一个超大型游戏,游戏里的角色都有一张隐藏的“知识卡”:谁住在哪里、谁喜欢什么。以前改游戏设定时,一次只能改一个角色;连续改很多次,系统就可能把没改过的人也弄错。

MEMIT像一个聪明的批量补丁工具。它先找出游戏内部真正保存事实的几个关键区域,再为每条新设定制作目标信息,最后把这些信息平均写进多个区域。这样,“某人喜欢篮球”可以改成“喜欢棒球”,而问法换一下,游戏仍应给出棒球。

研究者在GPT-J上测试了10,000条事实。MEMIT的成功率和改写后的回答都很高,编辑分数50.7;在另一组反事实测试中分数达到85.8。相比之下,ROME改太多次就失灵,MEND也难以扩展。

不过它不是魔法:有些关系很难修改,邻近角色可能受影响,而且处理一万条设定要很久。未来如果把工作并行化,它可能成为给聊天机器人快速更新新闻、产品资料和课程知识的工具!

术语表

MEMIT(Mass-Editing Memory)

面向大规模事实写入的直接参数编辑算法。它把新记忆分配到多个关键MLP层。

论文的核心方法,支持GPT-J和GPT-NeoX上的数千至万条编辑。

ROME(Rank-One Model Editing)

通过秩一参数更新编辑单条事实的方法。它是MEMIT的直接前身。

作为顺序编辑基线,并用于说明单条方法的扩展瓶颈。

因果中介分析(Causal Mediation Analysis)

通过干预内部状态,判断哪些组件对输出具有因果作用。它区别于只观察相关性的分析。

用于定位GPT-J事实回忆的关键MLP层R。

键值联想记忆(Key-Value Associative Memory)

将输入键映射到存储值的线性记忆结构。Transformer的MLP输出层被近似为这种结构。

MEMIT据此推导批量矩阵更新。

Specificity(特异性)

衡量未编辑相关事实是否保持正确的指标。数值下降意味着知识发生外溢。

与效能、改写泛化共同构成编辑分数。

COUNTERFACT

包含21,919条事实及其反事实改写的数据集。它测试模型是否能接受与现实相反的新知识。

用于从1到10,000条规模的扩展实验。

开放问题 这项研究留下的未解疑问

  • 1 模型为何对“运动员参加何种运动”等关系更难编辑仍不清楚,可能涉及多跳推理、竞争记忆或关系表示分散;需要关系级机制分析。
  • 2 连续多轮编辑后的稳定性、冲突撤销和安全边界尚未充分验证;当前结果主要来自离线基准,不能完全代表动态生产环境。

应用场景

近期应用

新闻与知识库更新

问答系统维护者可将经过审核的新事实批量写入模型,减少完整重训成本。部署前应进行冲突过滤,并用改写、邻域和生成质量测试监控副作用。

企业模型定制

企业可把产品规格、人物资料或内部术语写入基础模型。MEMIT适合成批变更,但必须限制权限、记录编辑版本,并保留原模型以便回滚。

远期愿景

持续可维护的语言模型

未来模型可能像软件一样接受可审计知识补丁,按时间更新新闻、法规和用户反馈。关键障碍是冲突管理、事实来源验证、长期遗忘控制与安全撤销。

原文摘要

Recent work has shown exciting promise in updating large language models with new memories, so as to replace obsolete information or add specialized knowledge. However, this line of work is predominantly limited to updating single associations. We develop MEMIT, a method for directly updating a language model with many memories, demonstrating experimentally that it can scale up to thousands of associations for GPT-J (6B) and GPT-NeoX (20B), exceeding prior work by orders of magnitude. Our code and data are at https://memit.baulab.info.

cs.CL cs.LG