Constructing Efficient Fact-Storing MLPs for Transformers
作者: Owen Dugan, Roberto Garcia, Ronny Junkins, Jerry Liu, Dylan Zinsley, Sabri Eyuboglu, Atri Rudra, Chris Ré
分类: cs.LG, cs.AI
发布日期: 2025-11-28
💡 一句话要点
提出高效存储事实的MLP构建框架，提升Transformer的事实记忆能力
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 事实存储 MLP构建 Transformer 知识表示 参数效率
📋 核心要点
- 大型语言模型依赖MLP存储事实知识，但现有方法在效率和适用性上存在局限。
- 论文提出一种新的MLP构建框架，旨在提升参数效率、适用范围和在Transformer中的可用性。
- 实验表明，该框架在事实存储容量和Transformer可用性之间存在权衡，并实现了模块化事实编辑。
📝 摘要（中文）
大型语言模型（LLM）的成功部分归功于其在MLP参数中高效存储事实知识作为键值映射的能力。最近的研究提出了显式的权重构建方法来构建这种事实存储MLP，从而更好地理解LLM的事实存储机制。本文介绍了一种MLP构建框架，该框架在三个方面改进了先前的构建方法：1）适用于几乎所有可行的输入-输出对；2）对于某些嵌入，实现了渐近最优的参数效率，与信息论界限相匹配；3）保持了在Transformer中进行事实回忆的可用性。通过这些改进，我们1）发现了一种值嵌入的度量，该度量表征了构造的和梯度下降训练的MLP的事实-参数缩放；2）确定了一种简单的编码器-解码器机制，该机制在经验上匹配了梯度下降MLP在所有测试输入和输出上的事实-参数渐近线；3）揭示了MLP的事实存储容量及其在Transformer中的可用性之间的基本权衡。最后，我们展示了事实存储MLP的概念验证应用：通过一次性替换整个MLP，在一层Transformer上进行模块化事实编辑。
🔬 方法详解
问题定义：大型语言模型（LLM）通过MLP层存储事实知识，但如何高效地构建MLP以存储和检索这些知识是一个挑战。现有的MLP构建方法可能存在适用范围窄、参数效率低或难以集成到Transformer架构中的问题。因此，需要一种更通用、更高效且易于使用的MLP构建方法，以提升LLM的事实记忆能力。
核心思路：论文的核心思路是通过显式地构建MLP的权重，使其能够有效地存储和检索事实知识。这种构建方法基于键值映射的思想，将输入作为键，输出作为值，并通过特定的权重矩阵将它们关联起来。通过优化权重矩阵的结构，可以提高参数效率，并确保MLP能够准确地回忆起存储的事实。此外，该方法还考虑了MLP在Transformer架构中的可用性，使其能够无缝集成到现有的LLM中。
技术框架：该框架主要包含以下几个阶段：1）输入-输出对的准备：收集需要存储的事实知识，并将其表示为输入-输出对。2）嵌入表示：将输入和输出嵌入到高维空间中，以便更好地进行键值映射。3）权重矩阵构建：根据嵌入表示，构建MLP的权重矩阵，使其能够将输入映射到相应的输出。4）MLP集成：将构建好的MLP集成到Transformer架构中，并进行微调以优化性能。
关键创新：该论文的关键创新在于提出了一种通用的MLP构建框架，该框架在以下几个方面优于现有方法：1）适用范围更广：适用于几乎所有可行的输入-输出对，而不仅仅是特定的数据集或任务。2）参数效率更高：实现了渐近最优的参数效率，与信息论界限相匹配。3）Transformer可用性更强：能够无缝集成到Transformer架构中，并保持良好的性能。
关键设计：论文的关键设计包括：1）值嵌入的度量：定义了一种值嵌入的度量，用于表征事实-参数缩放。2）编码器-解码器机制：提出了一种简单的编码器-解码器机制，用于匹配梯度下降MLP的事实-参数渐近线。3）模块化事实编辑：通过一次性替换整个MLP，实现了一层Transformer上的模块化事实编辑。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该框架在参数效率和Transformer可用性方面均优于现有方法。此外，论文还发现了一种值嵌入的度量，可以用于预测MLP的事实存储容量。通过模块化事实编辑实验，验证了该框架在实际应用中的可行性。
🎯 应用场景
该研究成果可应用于各种需要高效存储和检索事实知识的场景，例如知识图谱、问答系统、信息检索等。通过构建高效的事实存储MLP，可以提升这些系统的性能和可扩展性，并为未来的LLM研究提供新的思路。
📄 摘要（原文）
The success of large language models (LLMs) can be attributed in part to their ability to efficiently store factual knowledge as key-value mappings within their MLP parameters. Recent work has proposed explicit weight constructions to build such fact-storing MLPs, providing an improved understanding of LLM fact storage mechanisms. In this paper, we introduce an MLP construction framework that improves over previous constructions in three areas: it 1) works for all but a measure-zero set of feasible input-output pairs, 2) achieves asymptotically optimal parameter efficiency matching information-theoretic bounds for some embeddings, and 3) maintains usability within Transformers for factual recall. Through our improvements, we 1) discover a metric on value embeddings that characterizes facts-per-parameter scaling for both constructed and gradient-descent-trained MLPs, 2) identify a simple encoder-decoder mechanism that empirically matches gradient-descent MLP facts-per-parameter asymptotics across all the inputs and outputs we test, and 3) uncover a fundamental tradeoff between an MLP's fact-storage capacity and its usability within Transformers. Finally, we demonstrate a proof-of-concept application of fact-storing MLPs: modular fact editing on one-layer Transformers by \textit{replacing entire MLPs at once}.