Transformer Feed-Forward Layers Are Key-Value Memories
这篇论文的核心贡献,在于它为Transformer模型中占据了约三分之二参数量、但长期以来其功能被严重忽视的前馈神经网络(Feed-Forward Network, FFN)层,提供了一个简洁而深刻的解释框架。在此之前,学术界的目光大多聚焦于自注意力(Self-Attention)机制,而FFN层则像一个神秘的“黑箱”。Geva等人的这项工作,通过一系列精巧的实验,令人信服地论证了:FFN层在功能上等同于一个键值记忆(Key-Value Memory)系统。
具体来说,论文揭示了FFN内部的运作机制。FFN层包含两个线性变换矩阵,论文将它们分别类比为“键(Keys)”和“值(Values)”。当一个输入向量(代表某个词元或文本片段的含义)进入FFN层时,它会首先与所有的“键”进行匹配计算。这个匹配度,论文称之为记忆系数(memory coefficient),决定了每个“键”所对应的“值”被激活的强度。每个“值”本身并不直接是一个词,而是一个向量,它蕴含了一个关于“下一个可能出现的词”的概率分布。最终,FFN层的输出是所有被激活的“值”向量的加权和。这个加权和的输出,实际上是对模型下一步应该生成什么内容的一次“提议”或“修正”。
论文的一个惊人发现是,这些“记忆”是具有层次性和可解释性的。
Our experiments show that the learned patterns are human-interpretable, and that lower layers tend to capture shallow patterns, while upper layers learn more semantic ones. 我们的实验表明,学习到的模式是人类可解释的,并且底层网络倾向于捕获浅层模式,而上层网络则学习更具语义的模式。
例如,在模型的底层(如1-9层),FFN的“键”主要响应一些浅层模式(shallow patterns),比如特定的n-gram短语(例如,以“substitutes”结尾的句子)。而到了模型的高层(如10-16层),“键”则演变为响应更抽象的语义模式(semantic patterns),比如某个特定的话题(例如,关于“电视剧”的讨论)。这一点在论文的 图2 中得到了清晰的展示,图中显示了随着层数的增加,语义模式(绿色部分)的占比显著提升,而浅层模式(蓝色部分)则相应减少。
图2:论文中的关键图表,展示了不同层级的FFN记忆中,浅层模式与语义模式的比例变化。
更进一步,论文发现高层FFN中的“值”向量与对应的“键”模式高度相关。当一个高层的“键”被某个输入模式激活时,其对应的“值”所编码的词汇分布,确实能很大概率地预测出该模式之后最可能出现的词。例如,如果一个“键”专门识别“关于二战轰炸机基地”的语境,那么它对应的“值”向量,在被激活后,就会倾向于生成“任务(missions)”、“攻击(attacked)”等相关词汇。论文在 图4 中通过“一致率(agreement rate)”这个指标量化了这一发现,在高层(11-16层),这种一致率达到了约3.5%,虽然看似不高,但论文强调这比随机猜测(约0.0004%)高出几个数量级,证明了其非凡的预测能力。
最后,论文还探讨了这些记忆是如何被聚合(Aggregating)的。在一个FFN层内部,最终的输出通常不是由单个最强的记忆决定的,而是由数百个被激活的记忆组合(composition)而成,这在 图8 中有明确数据支撑,超过68%的情况下,层的输出与任何单个记忆的预测都不同。在层与层之间,FFN的输出通过残差连接(residual connections)对主干信息流进行精细的修正(refinement)。模型并非在最后一层才做出最终决定,而是在逐层传递的过程中,不断地利用FFN的“记忆”来调整和优化对下一个词的预测。
总而言之,这篇论文将FFN从一个难以理解的计算模块,变成了一个结构清晰、功能明确的分布式记忆系统。它不仅为我们理解Transformer的内部工作原理打开了一扇窗,也为后续的模型编辑、可信AI和模型压缩等研究方向奠定了坚实的基础。
1. 论文的研究目标是什么? 想要解决什么实际问题?这个问题对于行业发展有什么重要意义?
研究目标:论文的核心研究目标是揭开Transformer模型中前馈神经网络(FFN)层的神秘面纱,并为其功能提供一个清晰、可验证的解释。具体来说,它试图回答一个核心问题:占模型三分之二参数的FFN层,究竟在做什么?
解决的实际问题:
- 模型可解释性缺失:在当时,自注意力机制的研究已经比较深入,但FFN的功能却是一个巨大的知识盲区。这使得我们对大模型的决策过程理解不完整,难以信任其输出,也难以在出错时进行调试。
- 模型行为的不可预测性:我们不知道模型为何会产生“幻觉”、输出有害内容或泄露隐私数据。如果能理解FFN的功能,就可能从机理上找到这些问题的原因。
对行业发展的重要意义:
- 推动可信AI(Trustworthy AI):理解是信任的前提。通过将FFN解释为键值记忆,这篇论文极大地增强了我们对大模型内部工作原理的理解,是迈向构建更安全、更可控、更值得信赖的AI系统的关键一步。
- 催生新的技术方向:该研究直接启发了模型编辑(Model Editing)领域。如果我们知道某个事实性错误(比如“埃菲尔铁塔在罗马”)存储在哪个或哪些FFN“记忆”中,我们就有可能直接修改这些“值”向量来纠正错误,而无需重新训练整个庞大的模型。
- 提升模型效率和安全性:通过识别FFN中存储的“记忆”,我们可以分析哪些是冗余的,从而进行模型压缩;也可以识别哪些“记忆”可能存储了训练数据中的敏感信息,为隐私保护和数据安全审计提供新的工具。
2. 论文提出了哪些新的思路、方法或模型?跟之前的方法相比有什么特点和优势?请尽可能参考论文中的细节进行分析。
新的思路/模型:论文并未提出一个全新的模型架构,而是提出了一个关于现有Transformer FFN层的功能性类比(functional analogy)——即FFN作为非归一化的键值记忆。
核心方法论:
- 概念重构:将FFN的数学公式
FFN(x) = ReLU(x * K^T) * V
(简化后) 与神经记忆网络MN(x) = softmax(x * K^T) * V
进行对比。 > Comparing equations 1 and 2 shows that feed-forward layers are almost identical to key-value neural memories; the only difference is that neural memory uses softmax as the non-linearity f(·), while the canonical transformer does not use a normalizing function in the feed-forward layer. > 对比公式1和2可以发现,前馈层与键值神经记忆几乎完全相同;唯一的区别在于,神经记忆使用softmax作为非线性函数f(·),而标准的Transformer在前馈层中不使用归一化函数。 - 关键区别与优势:论文敏锐地指出了
ReLU
与softmax
的区别。softmax
会进行归一化,使得所有激活值的和为1,这意味着它倾向于只选择一个或少数几个最匹配的记忆。而ReLU
是非归一化的,它允许多个记忆单元被同时、独立地激活。这恰好解释了为何FFN的输出是“组合式”的——它是多个被触发的“想法”(记忆)的叠加,而不是非此即彼的选择。这个洞察是论文的核心优势,更贴近模型在实践中的复杂行为。
- 概念重构:将FFN的数学公式
分析方法:论文采用了一种“假设-验证”的实验探究方法,而非纯理论推导。
- 寻找触发样本(Trigger
Examples):为了探究“键”
k_i
的功能,研究者在庞大的训练集中搜索,找到那些能够让该键的记忆系数ReLU(x * k_i)
取得最大值的输入文本x
。这个方法非常巧妙,它让我们能直接“拷问”每个神经元:“什么样的信息最能让你兴奋?” - 模式标注与分类:通过人类专家(NLP研究生)对这些触发样本进行归纳,识别出其中的共性模式,并将其分为“浅层”和“语义”两类。这是一种将复杂的神经活动与人类可理解的概念联系起来的有效手段。
- 值向量分析:将“值”向量
v_i
通过乘以输出嵌入矩阵E
并应用softmax
,将其转换为一个完整的词汇表概率分布p_i
。这使得研究者可以直接评估每个“记忆”的预测倾向。
- 寻找触发样本(Trigger
Examples):为了探究“键”
与之前笼统地将FFN视为“特征提取器”或简单地增加模型容量的方法相比,这篇论文提供了一个结构化、模块化且可操作的视角,其优势在于将抽象的数学运算赋予了具体的认知功能——记忆与联想。
3. 论文通过什么实验来验证所提出方法的有效性?实验是如何设计的?实验数据和结果如何?请引用关键数据加以说明。
论文设计了一系列环环相扣的实验来验证其核心假设。
- 实验一:验证“键(Keys)”捕获可解释的输入模式
- 设计:研究者从一个16层的Transformer语言模型中,每层随机抽取10个“键”(共160个)。对每个键,他们检索出训练集中最能激活它的前25个文本前缀。然后,由人类专家分析这25个前缀,归纳出它们共同的模式。
- 数据与结果:
- 高覆盖率:对于几乎每一个被分析的键,专家都能找到至少一种可识别的模式。平均每个键能识别出3.6个不同的模式。
- 模式分层:如前所述,图2清晰地展示了从低层到高层,模式由浅层(如固定短语)向语义(如主题)的转变。在第1层,超过80%的模式是浅层的;而在第16层,约70%的模式是语义的。
- 具体案例:表1
提供了一些生动的例子。例如,第4层的
k_449
专门响应以“substitutes”结尾的句子;而第16层的k_1635
则响应关于“电视剧”的上下文。
- 实验二:验证“值(Values)”存储与键相关的输出分布
- 设计:对于每个键
k_i
,研究者找到最能激活它的那个文本前缀(top-1 trigger example),并记录其真实的下一个词w_t
。同时,他们将对应的“值”向量v_i
转换为词汇表分布p_i
,并查看其预测概率最高的词argmax(p_i)
。他们计算这两个词匹配的频率,即“一致率”。 - 数据与结果:
- 高层一致性:图4 显示,在1-10层,一致率几乎为零。但从第11层开始,一致率迅速攀升,到第16层时达到约3.5%。这证明了高层“值”向量确实存储了有意义的、与“键”模式匹配的预测信息。
- 排名提升:图5
从另一个角度验证了这一点。它展示了真实下一个词
w_t
在“值”向量v_i
的预测排名分布。随着层数加深,这个排名显著变好(箱形图整体下移),意味着w_t
被赋予了更高的概率。
- 设计:对于每个键
- 实验三:验证记忆的聚合机制(Intra-layer &
Inter-layer)
- 设计:
- 层内聚合:研究者随机抽取4000个验证集样本,观察每一层FFN的输出。他们统计有多少记忆被激活,并比较最终的层输出预测是否与任何单个记忆的预测相同。
- 层间聚合:他们比较了每一层的残差流(FFN的输入)的预测与经过FFN修正后的输出预测,以及模型的最终预测。
- 数据与结果:
- 组合性是常态:图8 的结果非常关键,它显示在所有层中,有超过68%的情况下,FFN层的最终输出预测与该层内任何一个单独的记忆单元的预测都不相同。这强有力地证明了FFN的输出是组合(compositional)而非选择性的。
- 逐层精炼:图9 显示,模型的最终预测结果在很深的层就已经被残差流(residual)的顶层预测所确定。例如,在第10层,超过60%的情况下,残差流的预测就是模型的最终预测。这说明FFN层扮演的角色更像是对一个已经不错的“草稿”进行微调和修正。
- 设计:
4. 结合大模型领域的当前学术理解,未来在该研究方向上还有哪些值得进一步探索的问题和挑战?这可能催生出什么新的技术和投资机会?
这篇论文开辟了一个充满机遇的研究领域,至今仍在不断发展。
- 值得探索的问题和挑战:
- 与自注意力的协同机制:论文主要孤立地研究FFN,但FFN的输入恰好是自注意力层的输出。未来的一个核心挑战是理解自注意力层是如何“预处理”和“路由”信息,以便激活正确的FFN记忆的。它们之间存在怎样的“电路(circuits)”?
- 跨模型和任务的泛化性:该论文的结论主要基于一个中等规模的GPT风格模型。这些发现在今天的千亿、万亿参数模型(如GPT-4、Llama系列)中是否仍然成立?在不同的模型架构(如Encoder-Decoder、MoE)和不同任务(如代码生成、多模态理解)中,FFN的记忆机制有何异同?
- 记忆的形成与学习:FFN中的这些“记忆”是在训练过程中如何形成的?它们与训练数据中的具体样本有何对应关系?理解这一点对于解决数据隐私和版权问题至关重要。
- 组合性的语法:我们知道FFN的输出是组合的,但这种组合遵循什么样的“语法”或规则?当多个记忆被激活时,模型如何权衡和整合它们的信息?这背后是否存在更深层次的代数结构?
- 可能催生的新技术和投资机会:
- AI安全与对齐(AI Safety and Alignment):通过深入理解FFN的记忆机制,可以开发出更先进的工具来检测和抑制模型的有害行为(如偏见、谎言)。这是一个巨大的投资赛道,专注于“可控AI”的公司将有很大潜力。
- 精准模型编辑与维护:基于对FFN记忆的理解,可以发展出成熟的“模型手术”技术,对线上部署的大模型进行实时、低成本的知识更新和错误修复。这将催生“大模型运维(LLMOps)”市场中的新兴服务。
- 自动化可解释性工具:将论文中需要人工标注的模式发现过程自动化,开发出能够自动分析并报告任意大模型内部“记忆”功能的商业软件。这将成为AI开发和审计流程中的标准工具。
- 新一代AI芯片设计:如果FFN的核心功能是稀疏激活的键值查找,那么未来的AI硬件或许可以为此进行专门优化,设计出能效更高、更适合大模型推理的芯片架构。
5. 退一步,从批判的视角看,这篇论文还存在哪些不足及缺失?又有哪些需要进一步验证和存疑的?
尽管这篇论文极其出色,但从今天的视角看,我们仍可以提出一些批判性的思考。
- 主要不足与缺失:
- 对“模式”的定义依赖于人类主观性:论文的核心证据之一来自于人类专家对触发样本的标注。这个过程是主观的,且可能过度简化了神经元实际响应的复杂模式。一个神经元可能响应的是比人类语言标签更抽象、更高维的特征组合。
- 对“值”向量功能的简化:“值”向量
v_i
被分析为对下一个词的直接预测,但它的实际作用是向残差流中添加一个方向向量,这个向量会影响后续所有层的计算。论文的分析方法捕捉到了其主要功能,但可能忽略了它对模型深层计算流的更长远、更复杂的影响。 - 实验范围的局限性:如前所述,实验仅限于一个特定的模型(Baevski & Auli, 2019)和一个特定的数据集(WikiText-103)。其结论在更大、更多样化的现代模型上的普适性有待验证。
- 需要进一步验证和存疑之处:
- 浅层与语义的截然划分:论文将模式清晰地划分为浅层和语义,并观察到在第9-10层左右的转变。这个转变是平滑过渡还是突变?在更大的模型中,这个分界点是否存在,或者说是否存在更复杂的、多阶段的模式演化?
- 记忆的稀疏性假设:论文提到,大部分记忆单元在单次推理中是不活跃的。这一稀疏性是FFN高效工作的关键吗?这直接启发了后来的混合专家(MoE)模型。但FFN层是否真的可以被看作是一个大型的、静态的MoE层,这一点仍需更深入的验证。
- “组合”的本质:当FFN的输出是一个与所有单个记忆都不同的“妥协”或“组合”时,这个新的预测是如何产生的?论文称之为“composition”,但其背后的计算原理仍不清晰。它仅仅是向量的线性叠加,还是存在更复杂的非线性相互作用?
6. 我希望从这篇论文中找一些拿来即用的创新想法,我应该从这篇论文中重点学什么?有哪些启发?你认为我还需要补充了解哪些背景知识?
当然,这篇论文是创新思想的宝库。
- 重点学习与启发:
- 功能性抽象的思维方式:最重要的启发是,不要将模型的组件视为纯粹的数学黑箱,而要尝试为其赋予功能性的、可类比的解释。将FFN类比为“键值记忆”就是一个绝佳的例子。当您遇到一个复杂的系统时,尝试提出一个简单的、可检验的功能假设,然后设计实验去验证它。
- “触发样本”的调试方法:这是一个非常实用的调试和分析技术。当您的模型做出奇怪或错误的预测时,可以借鉴论文的方法,去检查是哪些内部单元(如FFN神经元)被异常激活了。然后,去查找这些单元通常在什么情况下被激活(即它们的“触发样本”),这往往能帮助您定位问题的根源。
- 分层和组合的视角:理解模型的决策是一个逐层构建和精炼的过程,而非一步到位。同时,要认识到最终的输出往往是多个“内部意见”组合的结果。这启发我们在设计和分析模型时,不仅要看最终输出,更要关注信息在网络中流动的中间状态。
- 需要补充的背景知识:
- Transformer基础架构:您需要非常熟悉原始论文《Attention Is All You Need》中的所有细节,包括自注意力、多头注意力、位置编码、FFN和残差连接的具体计算过程。
- 机械可解释性(Mechanistic Interpretability):这是本篇论文所属的领域。建议了解该领域的目标,即彻底理解模型内部的每一个计算步骤和参数是如何协同工作以实现其功能的,可以关注一下“电路(circuits)”这个概念。
- 向量语义与嵌入(Vector Semantics & Embeddings):深入理解词嵌入(如Word2Vec, GloVe)和句子嵌入的思想,即如何将文本表示为高维空间中的向量,以及这些向量在数学上如何捕捉语义关系。这是理解FFN中“键”与输入匹配、“值”编码输出分布的基础。
- 后续研究进展:阅读一些跟进这项工作的论文,比如关于模型编辑的 ROME 和 MEMIT 算法,或者关于定位模型知识的研究,这将帮助您了解这个领域最新的发展。
v1.5.2