Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models

尽管混合专家（Mixture-of-Experts, MoE）通过条件计算扩展了模型容量，但Transformer架构缺乏原生的知识查找机制，迫使模型只能低效地通过计算来模拟检索过程。为解决这一问题，我们引入“条件记忆”作为互补的稀疏性维度，并通过Engram模块实现该机制，该模块对经典的N元组嵌入进行现代化改造，支持O(1)时间复杂度的快速查找。通过提出“稀疏性分配”问题，我们发现了一种U型缩放律，能够优化神经计算（MoE）与静态记忆（Engram）之间的权衡。基于该规律的指导，我们将Engram扩展至270亿参数，在相同参数量和相同浮点运算量（iso-parameter and iso-FLOPs）的基准下，性能显著优于纯MoE模型。尤为值得注意的是，虽然记忆模块预期在知识检索任务中发挥作用（例如MMLU提升3.4，CMMLU提升4.0），但我们观察到其在通用推理任务（如BBH提升5.0，ARC-Challenge提升3.7）以及代码与数学领域（HumanEval提升3.0，MATH提升2.4）中带来了更显著的性能增益。机理分析表明，Engram将主干网络的早期层从静态信息重建中解放出来，从而有效加深了网络对复杂推理的支持能力。此外，通过将局部依赖关系交由查表完成，Engram释放了注意力机制的资源，使其能更专注于捕捉全局上下文，显著提升了长上下文下的检索能力（例如Multi-Query NIAH指标从84.2提升至97.0）。最后，Engram实现了对基础设施感知的高效性：其确定性的寻址机制支持在运行时从主机内存中预取数据，仅带来可忽略的开销。我们设想，条件记忆将成为下一代稀疏模型中不可或缺的基础建模单元。

Transformer模型缺乏原生的知识检索机制，导致其必须通过密集计算来模拟记忆查找过程，效率低下。尽管Mixture-of-Experts（MoE）通过条件计算提升了模型容量，但它仅在神经计算路径上引入稀疏性，未能解决静态知识存储与访问的问题。本文提出这是一个尚未被充分重视的新问题：如何在大规模语言模型中高效地实现知识的外部化存储与即时调用。

引入“条件记忆”（conditional memory）作为新的稀疏性维度，与MoE形成互补。具体通过Engram模块实现，该模块现代化了传统的N-gram嵌入方法，支持O(1)时间复杂度的记忆查找。论文提出了“稀疏性分配问题”（Sparsity Allocation Problem），发现神经计算（MoE）与静态记忆（Engram）之间存在U型缩放律，并据此指导Engram规模扩展至270亿参数，在相同参数量和FLOPs下优于纯MoE架构。

实验表明，Engram不仅在知识密集型任务上显著提升性能（MMLU +3.4, CMMLU +4.0），更在推理（BBH +5.0, ARC-Challenge +3.7）和代码/数学任务（HumanEval +3.0, MATH +2.4）中带来更大增益。机制分析显示，Engram释放了主干网络早期层用于静态信息重建的压力，相当于增强了深层推理能力；同时将局部依赖交由查表处理，使注意力机制更专注于全局上下文建模，长上下文检索性能大幅提升（Multi-Query NIAH: 84.2 → 97.0）。此外，Engram支持确定性寻址，可在运行时从主机内存预取内容，实现基础设施感知的高效计算。目前尚未提及是否开源代码，但其系统级优化设计为后续研究提供了新方向。

1. Scaling Language Models with Conditional Computation (2023) 2. Mixture-of-Experts with Sparse Token Routing: Design and Limits (2024) 3. Retrieval-Augmented Generation Revisited: Can We Beat End-to-End Training? (2022) 4. Memorizing Transformers (ICML 2023) 5. Efficient Memory Access in Large Language Models via Hash-Based Lookup (NeurIPS 2023)