MemoryLLM：面向Transformer的即插即用可解释前馈记忆网络 MemoryLLM: Plug-n-Play Interpretable Feed-Forward Memory for Transformers
将FFN解耦为无上下文的token级神经检索记忆，提升可解释性并支持预计算卸载
前置知识
前馈网络 (FFN)
Transformer中的核心组件之一，通常采用SwiGLU激活函数，包含三个参数矩阵：上投影 $W_{Up}$、下投影 $W_{Down}$ 和门控投影 $W_{Gate}$。FFN接收自注意力模块的输出，通过非线性变换将其映射到更高维空间再投影回来。在现代LLM中，FFN占据了约三分之二的模型参数，是存储和检索知识的关键模块。
理解FFN的结构和作用是本文的核心——论文的目标就是重新设计FFN的角色，使其从黑盒变为可解释的神经记忆
残差流 (Residual Stream)
Transformer中信息传递的主要通道。每一层的自注意力和FFN输出都会加回到残差流中，形成 $X_{L+1} = X_L + Attn(X_L) + FFN(\tilde{X}_L)$ 的信息流动模式。残差流在层间不断积累和混合信息，使得中间层的表示变得难以解释。
正是因为残差流的存在，传统FFN的输入是混合了多层上下文信息的潜在表示，导致FFN难以被解释为对特定token的记忆检索
神经键值记忆 (Neural Key-Value Memory)
Geva等人提出的观点，认为FFN本质上是一个巨大的键值存储：上投影矩阵的行向量是键，匹配输入模式；下投影矩阵的行向量是值，存储对应的输出信息。当输入与某个键的匹配度高时，对应的值就被检索出来。这种解释为理解FFN如何存储和检索知识提供了理论框架。
本文的核心贡献之一就是TKV框架，它在可解耦的架构下重新形式化了这种键值记忆的解释
Token-wise Lookups (ToLs)
本文提出的一种预计算策略。由于MemoryLLM中FFN的输入是静态的token嵌入（只取决于token ID，不依赖上下文），因此可以预先计算每个token通过所有FFN层的输出，存储为静态查找表。推理时直接查表即可，无需实时计算FFN。
这是MemoryLLM实现高效推理的关键机制，使得FFN可以从VRAM卸载到存储设备，显著降低内存占用
研究动机
在现代大型语言模型中，前馈网络（FFN）占据了约三分之二的模型参数，但其在信息处理中的角色一直缺乏深入研究。根本原因在于传统Transformer架构中，FFN的输入是自注意力输出与残差流的混合表示 $\tilde{X}_L = X_L + Attn(X_L)$，这种表示在层间动态演化，是非可解释的潜在向量。Geva等人的先驱工作试图在GPT-2规模上证明FFN是神经键值记忆，但面临两个根本性限制：（1）需要通过校准数据集进行多次前向和反向传播，然后人工挖掘和标注相关的输入短语，流程繁琐且依赖人工；（2）无法离散地定义可解释的查询来访问FFN中的键值记忆，因为查询空间是连续的、上下文依赖的潜在向量。这使得FFN的功能研究被限制在小规模模型上，无法扩展到现代LLM。
本文的目标是本文的核心目标是重新设计Transformer架构，将FFN从自注意力和残差流中完全解耦，使其能够作为无上下文的、基于token索引的神经检索记忆被独立研究。具体而言，作者希望：（1）建立一个有限的、人类可解释的查询空间（即词表中的token ID），使得每个token都可以被映射到FFN中的静态记忆位置；（2）通过这种解耦设计，研究语义相似的token是否倾向于访问FFN中相似的记忆位置；（3）利用FFN输入的静态特性，实现FFN的预计算和卸载，降低推理时的计算和内存开销。
与已有工作不同的是，本文的独特切入点在于：不是在现有FFN上进行后验分析（post-hoc analysis），而是从头设计一个新架构来强制实现FFN的可解释性。与Geva等人研究已训练好的FFN长什么样不同，本文问的是如果我们从架构层面就让FFN只接收token嵌入，会发生什么。这种设计选择的深刻之处在于：当FFN的输入只取决于token ID（$X_0$）而非上下文时，FFN就变成了一个确定性的、有限查询空间的检索记忆，每个token ID对应一个固定的记忆地址。这不仅解决了可解释性问题，还意外带来了效率优势——静态输入意味着可以预计算。
核心方法
MemoryLLM的核心思想可以用一个简单的类比来理解：传统Transformer中的FFN像是一个需要实时读取上下文信息才能工作的上下文相关处理器，而MemoryLLM将FFN改造成了一个静态字典——给定一个token ID，无论上下文如何，FFN总是返回相同的结果。技术路线上，作者保留了传统自注意力模块在残差流上的正常工作方式（$Attn(X_L)$），但将所有FFN模块的输入从残差流改为初始token嵌入 $X_0$。具体而言，对于N层Transformer中的每一层 $L$，FFN的计算变为 $FFN(\hat{X}_0) = W_{Down}(\hat{X}_0 W_{Up}^T \odot SiLU(\hat{X}_0 W_{Gate}^T))$，其中 $\hat{X}_0 = LayerNorm_L(X_0)$ 是经过层归一化的初始嵌入。最终残差流更新为 $X_{L+1} = X_L + Attn(X_L) + FFN(X_0)$。
MemoryLLM最本质的创新在于架构层面的解耦。在传统Transformer中，FFN和自注意力是串行的：$X_{L+1} = X_L + Attn(X_L) + FFN(X_L + Attn(X_L))$，FFN的输入混合了多层累积的上下文信息。MemoryLLM将其改为并行：$X_{L+1} = X_L + Attn(X_L) + FFN(X_0)$，FFN只接收初始token嵌入。这个看似简单的改变带来了三个根本性变化：（1）FFN的查询空间从连续的、无限的潜在向量空间变成了离散的、有限的词表空间（$|V|$个token ID），使得FFN可以被形式化为TKV（Token-Key-Value）框架；（2）每个token通过FFN的输出是确定的，不依赖上下文，这使得FFN可以被预计算为token-wise lookups（ToLs）；（3）由于FFN不依赖残差流，单个FFN层可以根据其重要性和VRAM约束被移除，而不会破坏残差信息流。
方法步骤详情
MemoryLLM的完整计算流程如下：（1）输入文本 $T = \{t_1, t_2, ..., t_M\}$ 经过分词器得到M个token ID，通过嵌入矩阵 $E \in \mathbb{R}^{|V| \times d}$ 得到初始嵌入 $X_0 \in \mathbb{R}^{M \times d}$；（2）对于每一层 $L$，自注意力模块按照标准方式在残差流 $X_L$ 上计算：$Attn(X_L) = softmax(\frac{X_L W_Q^T (X_L W_K^T)^T}{\sqrt{d_k}}) X_L W_V^T$；（3）同时，该层的FFN在初始嵌入上独立计算：先经过独立的层归一化 $\hat{X}_0 = LayerNorm_L(X_0)$（作者发现每层独立的LayerNorm对收敛很重要），然后计算 $FFN(\hat{X}_0) = W_{Down}(\hat{X}_0 W_{Up}^T \odot SiLU(\hat{X}_0 W_{Gate}^T))$；（4）将自注意力输出和FFN输出加回残差流：$X_{L+1} = X_L + Attn(X_L) + FFN(X_0)$；（5）训练完成后，可以对词表中所有token预计算每层FFN的输出，形成ToL表：$ToL[x_{t_i}] = Concat_{k=0}^{N-1} FFN_{L_k}(x_{t_i})$，存入存储设备供推理时按需加载。
技术新颖性
MemoryLLM的技术新颖性体现在多个层面。首先，与MoLE（Jie等人2025年的工作）相比，虽然MoLE也发现混合专家中的多数专家可以直接用token级嵌入训练，但其专家计算仍然依赖于自注意力输出训练的路由器——这意味着专家的选择是上下文相关的。MemoryLLM则更进一步，完全移除了FFN对上下文的依赖。其次，与传统的FFN后验分析方法（如Geva等人的工作）相比，MemoryLLM不是在已训练模型上逆向工程FFN的功能，而是通过架构设计强制FFN成为可解释的记忆模块。第三，ToL预计算策略是新颖的——它利用了MemoryLLM中FFN输入静态的特性，这在传统Transformer中是不可能的，因为传统FFN的输入是动态的残差流。最后，Flex-MemoryLLM的设计——将FFN参数分为上下文相关的FFN-C和无上下文的FFN-M两部分——提供了一个在可解释性和性能之间灵活权衡的新范式。
实验结果
论文在250M、750M和1B三个参数规模上从头训练了MemoryLLM和Flex-MemoryLLM，训练数据量从25B到150B tokens不等。核心发现包括：（1）FFN的空间聚类特性：通过K-means聚类分析，语义相似的token（如标点符号、人名、地理名词）在FFN记忆中倾向于激活相似的键（keys），形成可解释的聚类。聚类系数（CC）在所有层中都保持较高值（约0.9以上），表明这种规律性是系统性的而非偶然。（2）FFN对不同任务的差异化影响：通过控制FFN贡献比例 $\alpha$ 的实验发现，降低FFN贡献对检索/召回类任务（如WikiText-2困惑度从24.35暴增到53.71，增幅120.57%；LAMBDA准确率下降69.93%）的伤害远大于推理/逻辑类任务（如PIQA仅下降3.09%，Winogrande下降4.23%）。这证实FFN主要存储的是显式的事实性知识而非推理能力。（3）性能与效率的权衡：MemoryLLM-1B（402M活跃参数）在C4上的困惑度为20.933，低于base-750M（737M活跃参数）的19.730，但考虑活跃参数仅为后者的54.5%，效率优势显著。（4）Flex-MemoryLLM的有效性：Flex-MemoryLLM-3h2（704M活跃参数）的性能接近base-1B（1208M活跃参数），且在150B tokens训练后，Flex-MemoryLLM-3h2甚至超越了base-737M模型。（5）优于传统剪枝方法：与Magnitude、SparseGPT、Wanda三种剪枝方法相比，MemoryLLM和Flex-MemoryLLM在相同活跃参数下表现显著更优。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 语言建模 (C4) | Perplexity (↓) | MemoryLLM-1B: 20.933 | base-750M: 19.730 | 活跃参数减少45.5%（402M vs 737M），困惑度仅增加6.1% |
| 语言建模 (WikiText-2) | Perplexity (↓) | MemoryLLM-1B: 27.258 | base-750M: 25.491 | 活跃参数减少45.5%，困惑度增加6.9% |
| 检索类任务 (LAMBDA) | Accuracy (↑) | α=1.0: 0.3613 | α=0.5: 0.1086 | FFN贡献降低50%导致准确率下降69.93%，证明FFN对检索任务至关重要 |
| 推理类任务 (PIQA) | Accuracy (↑) | α=1.0: 0.7018 | α=0.5: 0.6801 | FFN贡献降低50%仅导致准确率下降3.09%，证明推理任务对FFN依赖较小 |
| 模型压缩 (vs. 剪枝) | Perplexity (↓) | Flex-MemoryLLM-3h2: ~15 (1B规模) | SparseGPT/Wanda: ~25-35 (相同活跃参数) | 相同活跃参数下，困惑度显著低于传统剪枝方法 |
局限与改进
论文存在若干局限性。首先，MemoryLLM在相同总参数量下的性能低于传统dense模型，这是因为FFN被限制为只能处理静态token嵌入，丧失了处理上下文信息的能力。虽然Flex-MemoryLLM通过将部分FFN参数分配给上下文相关的计算来弥补这一缺陷，但这本质上是承认了完全解耦FFN的代价。其次，论文的实验规模相对有限——最大模型仅为1B参数，训练数据最多150B tokens，而现代LLM通常有7B-70B参数并训练在数万亿tokens上。这使得结论在更大规模下的普适性存疑。第三，论文没有报告在实际下游任务（如问答、代码生成、数学推理）上的微调结果，所有评估都是零样本或少样本，这限制了对实际应用价值的判断。第四，虽然TKV框架提供了可解释的视角，但论文没有展示如何利用这种可解释性进行具体的知识编辑或模型干预——这是可解释性研究的最终目标之一。第五，ToL预计算虽然降低了推理时的计算量，但引入了额外的存储开销和I/O瓶颈，论文没有详细分析这种权衡在不同硬件条件下的实际影响。
独立分析的弱点
论文存在几个值得深入探讨的弱点。第一，上下文信息的丧失：MemoryLLM完全移除了FFN对上下文的依赖，但许多语言任务（如指代消解、语境依赖的词义消歧）本质上需要上下文信息。一个可能的改进方向是引入轻量级的上下文调制机制，例如通过低秩适配器（LoRA）在推理时根据上下文动态微调FFN的键值记忆，既保持主要记忆的静态性以支持预计算，又允许少量上下文敏感的调整。第二，词表大小的限制：MemoryLLM的FFN查询空间受限于词表大小（如128K tokens），这意味着每个token只能有一个固定的记忆地址。对于多义词或形态丰富的语言，同一token在不同语境下可能需要不同的记忆。可以通过扩展查询空间到sub-word级别或引入token嵌入的微小变体来缓解。第三，Flex-MemoryLLM的设计不够优雅：当前的参数分割（$\beta h^2$）是启发式的，缺乏理论指导。可以通过学习最优的分割比例或采用动态路由机制来改进。第四，评估不够全面：缺乏在实际应用场景（如检索增强生成、知识密集型任务）中的端到端评估。
未来方向
基于MemoryLLM的成果，可以延伸出多个有前景的研究方向。（1）知识编辑与注入：TKV框架揭示了语义相似token共享相似记忆键的特性，这为精确的知识编辑提供了可能——可以通过修改特定键的值向量来更新或注入知识，而不需要重新训练整个模型。（2）动态FFN架构：可以研究根据输入内容动态选择激活哪些FFN层的机制，类似于条件计算，但基于MemoryLLM的可解释性优势。（3）跨模态扩展：将TKV框架扩展到视觉或音频token，研究多模态FFN记忆的结构。（4）更大规模验证：在7B-70B参数规模上验证MemoryLLM和Flex-MemoryLLM的设计，特别是在长上下文和复杂推理任务上的表现。（5）与其他压缩技术结合：探索MemoryLLM与量化、蒸馏等技术的结合，进一步压缩模型。（6）理论分析：深入研究为什么解耦FFN后，语义相似token仍能形成聚类——这可能揭示FFN学习的更深层机制。
复现评估
论文的复现可行性需要从多个维度评估。开源情况：论文来自Apple，目前未提及代码或模型权重的开源计划，这将严重限制社区的复现和验证。数据：论文使用标准的训练数据集，但具体的混合比例和预处理细节需要查看附录。算力需求：论文训练了250M、750M和1B三个规模的模型，训练token数从25B到150B。以1B模型150B tokens为例，估计需要数百GPU小时，这对学术实验室是可承受的，但远超个人研究者的资源。复现难度：架构修改本身并不复杂——核心是将FFN的输入从残差流改为初始嵌入，并为每层添加独立的LayerNorm。但从头训练的成本较高，且需要仔细调参以确保收敛（作者提到独立LayerNorm对收敛很重要）。建议：如果代码开源，复现是可行的；如果不开源，建议先在小规模（如125M）上验证架构设计的有效性。
论文图表
