arXiv:2601.07372cs.CLcs.AI2026-01ACL被引 71

给大模型加可条件调用的记忆模块,让知识查找更快更准。

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

  • 引入条件记忆模块Engram,实现常数时间知识查找。
  • 270亿参数的Engram在推理任务上优于同等规模的MoE模型。
  • 尤其提升复杂推理与长上下文任务,适合高阶语言建模场景。

虽然混合专家(MoE)通过条件计算扩展了模型容量,但Transformer缺乏原生的知识查找机制,只能通过计算低效模拟检索。为此,我们提出条件记忆作为互补的稀疏性维度,通过Engram模块实现经典N-gram嵌入的现代化,支持O(1)查找。通过构建稀疏性分配问题,我们发现神经计算(MoE)与静态内存(Engram)之间存在U型缩放规律,优化二者权衡。基于此规律,我们将Engram扩展至270亿参数,在同等参数量和浮点运算量下显著优于纯MoE基线。值得注意的是,尽管记忆模块预期提升知识检索能力(如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具备基础设施感知效率:其确定性寻址支持主机内存的运行时预取,开销极小。我们视条件记忆为下一代稀疏模型不可或缺的建模原语。

原文摘要 · Abstract (English)

While Mixture-of-Experts (MoE) scales capacity via conditional computation, Transformers lack a native primitive for knowledge lookup, forcing them to inefficiently simulate retrieval through computation. To address this, we introduce conditional memory as a complementary sparsity axis, instantiated via Engram, a module that modernizes classic $N$-gram embedding for O(1) lookup. By formulating the Sparsity Allocation problem, we uncover a U-shaped scaling law that optimizes the trade-off between neural computation (MoE) and static memory (Engram). Guided by this law, we scale Engram to 27B parameters, achieving superior performance over a strictly iso-parameter and iso-FLOPs MoE baseline. Most notably, while the memory module is expected to aid knowledge retrieval (e.g., MMLU +3.4; CMMLU +4.0), we observe even larger gains in general reasoning (e.g., BBH +5.0; ARC-Challenge +3.7) and code/math domains~(HumanEval +3.0; MATH +2.4). Mechanistic analyses reveal that Engram relieves the backbone's early layers from static reconstruction, effectively deepening the network for complex reasoning. Furthermore, by delegating local dependencies to lookups, it frees up attention capacity for global context, substantially boosting long-context retrieval (e.g., Multi-Query NIAH: 84.2 to 97.0). Finally, Engram establishes infrastructure-aware efficiency: its deterministic addressing enables runtime prefetching from host memory, incurring negligible overhead. We envision conditional memory as an indispensable modeling primitive for next-generation sparse models.

大模型稀疏性记忆模块推理加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。