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

---

论文地址：https://arxiv.org/abs/2601.07372

项目页面：https://github.com/deepseek-ai/Engram

进 Q 学术交流群：922230617 或加 CV_EDPJ 进 W 交流群

---

目录

1. 引言

2. 架构

2.1. 概述

2.2. 通过哈希 N-gram 实现稀疏检索

2.3. 上下文感知门控

2.4. 与多分支架构集成

2.5. 系统效率：解耦计算与内存

3. 缩放定律与稀疏性分配

3.1. MoE 与 Engram 之间的最优分配比率

3.2. 无限记忆体系下的 Engram

4. 大规模预训练

4.1. 实验设置

4.2. 实验结果

5. 长上下文训练

5.1. 实验设置

5.2. 实验结果

6. 分析

6.1. Engram 在功能上是否等同于增加了模型的深度？

6.2. 结构消融与层敏感性

6.3. 敏感性分析

6.4. 系统效率

6.5. 案例研究：门控可视化

---

1. 引言

稀疏性是智能系统从生物神经回路到现代大型语言模型中反复出现的设计原则。目前，这一原则主要通过 混合专家（Mixture-of-Experts，MoE）实现，它通过条件计算扩展容量。由于其能够在不按比例增加计算量的情况下大幅增加模型规模，MoE 已成为前沿模型的事实标准。

尽管这种条件计算范式取得了成功，但语言信号固有的异质性表明存在显著的结构优化空间。具体来说，语言建模包含两个性质不同的子任务：组合推理和知识检索。前者需要深度、动态的计算，而相当一部分文本——如命名的实体和公式化模式——是局部的、静态的和高度固化的。经典的 N-gram 模型在捕捉此类局部依赖性方面的有效性意味着，这些规律性可以自然地表示为计算成本低廉的查找操作。由于标准 Transformer 缺乏原生的知识查找方式，当前的 LLM 被迫通过计算来模拟检索。例如，解析一个常见的多 token 实体需要消耗多个早期层的注意力和前馈网络资源。这本质上相当于对一个静态查找表进行昂贵的运行时重建，将宝贵的序列深度浪费在琐碎的操作上，而这些资源本可用于更高层次的推理。

为了使模型架构与这种语言二元性对齐，本文提倡稀疏性的一个互补维度：条件记忆。

• 条件计算稀疏地激活参数来处理动态逻辑，而 条件记忆依赖于稀疏查找操作来检索固定知识的静态嵌入。
• 作为对该范式的初步探索，本文重新审视 N-gram 嵌入作为一种规范实例：局部上下文作为键，通过恒定时间 O(1) 的查找来索引海量嵌入表。本文的研究表明，这种静态检索机制可以作为现代 MoE 架构的理想补充——但前提是经过适当设计。
• 本文提出了 Engram，这是一个基于经典 N-gram 结构的条件记忆模块，但配备了现代适配，如分词器（tokenizer）压缩、多头哈希、上下文感知门控和多分支集成。

为了量化 MoE 和 Engram 之间的协同作用，本文构建了 稀疏性分配 问题：在给定的总参数量预算下，应如何在 MoE 专家和 Engram 记忆之间分配容量？

• 本文的实验揭示了一个明显的 U 形缩放定律，表明即使是简单的查找机制，当被视为一等建模方式时，也能作为神经计算的重要补充。
• 在此分配定律指导下，本文将 Engram 扩展到 27B 参数。与严格的同参数量和同 FLOPs 的 MoE 基线相比，Engram-27B 在多个领域实现了更优的效率。

通过 LogitLens 和 CKA 进行的机理分析揭示了这些收益的来源：Engram 减轻了主干网络在早期层重构静态知识的负担，从而增加了可用于复杂推理的有效深度。此外，通过将局部依赖性委托给查找操作，Engram 释放了注意力容量以聚焦于全局上下文，在长上下文场景中实现了卓越的性能——在 LongPPL 和 RULER 上大幅超越基线。

最后，本文确立了 基础设施感知的效率 作为一等设计原则。与 MoE 的动态路由不同，Engram 使用确定性 ID 来实现运行时预取，使通信与计算重叠。实证结果表明，将 100B 参数的表格卸载到主机内存产生的开销可忽略不计（<3%）。这证明 Engram 有效地绕过了 GPU 内存限制，促进了激进的参数扩展。

2. 架构

2.1. 概述

如图所示，Engram 是一个条件记忆模块，旨在通过结构上将静态模式存储与动态计算分离来增强 Transformer 主干网络。形式上，给定输入序列 X=(x1,…,xT) 和第 ℓ 层的隐状态 H(ℓ)∈R^{T×d}，该模块分两个功能阶段处理每个位置 t：检索和融合。

• 首先，提取并压缩后缀（suffix） N-gram，通过哈希（hashing）确定性地检索静态嵌入向量。
• 随后，这些检索到的嵌入由当前隐藏状态动态调制，并通过轻量级卷积进行精炼。

2.2. 通过哈希 N-gram 实现稀疏检索

第一阶段将局部上下文映射到静态记忆条目，涉及分词器压缩和通过确定性哈希检索嵌入。

分词器压缩：

• 虽然 N-gram 模型通常直接对分词器的输出进行操作，但标准子词分词器优先考虑无损重建，经常为语义等价的术语(（例如，Apple 与 ␣apple）分配不相交的 ID。
• 为了最大化语义密度，本文实现了一个词汇表投影层。具体来说，预计算一个满射函数 P: V→V′，将原始 token ID 基于规范化的文本等价性（使用 NFKC、小写化等）折叠为规范标识符。实践中，这使 128k 分词器的有效词汇表大小减少了 23%。
• 形式上，对于位置 t 上的一个 token，将其原始 ID x_t 映射到规范 ID x'_t = P(x_t)，以构成后缀 N-gram g_{t,n}。

多头哈希：

• 直接参数化所有可能 N-gram 的组合空间是不可行的。本文采用基于哈希的方法。
• 为了缓解碰撞，为每个 N-gram 阶数 n 使用 K 个不同的哈希头。每个头 k 通过确定性函数 ϕ_{n,k} 将压缩后的上下文映射到嵌入表 E_{n,k} 中的一个索引 z。
• 最终记忆向量 e_t 通过拼接所有检索到的嵌入构建。

2.3. 上下文感知门控

检索到的嵌入 e_t 作为上下文无关的先验。但它们是静态的，本身缺乏上下文适应性，并且可能因哈希碰撞或多义词而产生噪声。为了增强表达力并解决这种歧义，本文采用 受注意力机制启发的上下文感知门控机制。

具体来说，使用已经通过先前注意力层聚合全局上下文的当前隐藏状态 h_t 作为动态 Query，而检索到的记忆 e_t​ 作为 Key 和 Value 投影的来源。

为确保梯度稳定性，在计算标量门 α_t ∈ (0,1) 之前，对 Query 和 Key 应用 RMSNorm。门控输出定义为

该设计强制实现语义对齐：若检索到的记忆向量 e_t 与当前上下文 h_t 相矛盾，门控值 α_t 会趋近于零，从而有效抑制噪声干扰。

最后，为了扩展感受野并增强模型的非线性，本文引入了一个短的、深度相关的（depthwise）因果卷积。令 Ṽ ∈ R^{T×d} 表示门控值序列。设定卷积核尺寸为 𝑤（取值为4）、膨胀率为 𝛿（设为最大 N-gram 阶数）并采用 SiLU 激活函数，最终输出 Y 的计算公式如下：

Engram 模块通过残差连接集成到主干网络中：H^(ℓ) ← H^(ℓ) + Y。关键的是，Engram 并不应用于每一层；其具体放置位置由系统级延迟约束决定。

2.4. 与多分支架构集成

本文采用先进的多分支架构作为默认主干网络。该架构的一个定义性特征是将残差流扩展为 M 个并行分支，其中信息流由可学习的连接权重调节。

本文将 Engram 适配到此多分支框架中：单个稀疏嵌入表和一个 Value 投影矩阵 W_V 在所有 M 个分支之间共享，而 M 个不同的 Key 投影矩阵 W_k 用于实现分支特定的门控行为。

这种设计允许将线性投影融合成单个密集的 FP8 矩阵乘法，最大限度地提升现代 GPU 的计算利用率。

除非特别说明，所有实验均采用此方法，并结合了流形约束的超连接（𝑀 = 4）。

（2025|Deepseek，残差连接，残差映射投影，迭代归一化，核融合，混合精度，中间激活重新计算）mHC：流形约束的超连接

2.5. 系统效率：解耦计算与内存

扩展内存增强模型通常受限于 GPU 高带宽内存的有限容量。然而，Engram 的确定性检索机制天然支持参数存储与计算资源的解耦。

训练阶段：为容纳大规模嵌入表，本文采用标准模型并行，将表格分片到可用 GPU 上。使用 All-to-All 通信方式在前向传播中收集有效行，在反向传播中分发梯度。

推理阶段：这种确定性性质支持预取和重叠策略。由于内存索引在前向传播之前已知，系统可以异步通过 PCIe 从充足的主机内存检索嵌入。为了有效掩盖通信延迟，Engram 模块被放置在主干网络中的特定层，利用先前层的计算作为缓冲区。

多级缓存层次结构：自然语言 NN-gram 天然遵循 Zipf 分布，其中一小部分模式占用了绝大多数内存访问。这一统计特性激发了多级缓存层次结构：频繁访问的嵌入可以缓存在更快的存储层中，而稀有模式的长尾则驻留在速度较慢、容量较大的介质中。

3. 缩放定律与稀疏性分配

Engram 作为 条件记忆 的一个具体实现，在结构上与 MoE 专家提供的 条件计算是互补的。

本节深入探讨这种二元性的缩放特性，以及如何最优地分配稀疏容量。具体来说，两个关键问题驱动了本文的研究：

• 有限约束下的分配：当总参数量和训练计算量固定（同参数 和同 FLOPs）时，应如何将稀疏容量在 MoE 专家和 Engram 嵌入之间分配？

• 无限记忆体系：考虑到 Engram 具有非扩展的 O(1) 开销，如果内存预算放宽或进行激进扩展，Engram 自身会展现出怎样的扩展行为？

3.1. MoE 与 Engram 之间的最优分配比率

本文使用三个参数量度分析权衡：总可训练参数 P_{tot}​、每个 token 激活的参数 P_{act} 和非活跃参数 P_{sparse}。

本文在每个 FLOPs 预算内保持 P_{tot}​ 和 P_{act} 固定，这样模型具有相同数量的参数和相同的每 token FLOPs。

• 对于 MoE，P_{act} 由选中的前 k 个专家决定，而未选中专家的参数则贡献给 P_{sparse}​。

• 对于 Engram，每个 token 仅检索恒定数量的嵌入槽（slots），因此扩展嵌入槽的数量会增加 P_{tot}，但不会增加每 token FLOPs。

分配比率 ρ ∈ [0, 1] 定义为非活跃参数预算中分配给 MoE 的比例。

上图展示了两种计算预算（2e20 与 6e20 FLOPs）下的结果对比。实验发现验证损失与分配比率 ρ 之间存在一致的 U 形关系。

• Engram 模型在 MoE 分配减少到仅 ρ≈40% 时，仍能实现与纯 MoE 基线相当的性能。
• 纯 MoE 基线被证明是次优的：将大约 20%−25% 的稀疏参数预算重新分配给 Engram 能产生最佳性能。这个最佳位置在不同规模下保持稳定（在固定稀疏度下）。

3.2. 无限记忆体系下的 Engram

本文在固定的 MoE 主干上附加一个 Engram 表，并将内存槽数量 M 从 2.58×10^5 扩展到 1.0×10^7（增加了约 13B 参数）。

• 结果表明，扩展内存槽的数量能在验证损失上带来清晰且一致的改进。在整个探索范围内，曲线遵循严格的幂律，表明 Engram 提供了一个可预测的缩放旋钮：更大的内存持续带来回报，而无需额外计算。
• 相比之下，OverEncoding 的直接平均方法也能从更大的内存表中受益，但 Engram 从相同的内存预算中解锁了更大的缩放潜力。

4. 大规模预训练

利用提出的 Engram 架构和实验推导的分配定律，本文将 Engram 扩展到数十亿参数规模，以验证其在真实世界语言模型预训练中的有效性。

4.1. 实验设置

所有模型都在 262B token 的语料库上进行预训练，并使用 DeepSeek-v3 的分词器。模型配置保持一致：30 个块的 Transformer，隐藏大小为 2560，集成了具有 32 个头的多头潜在注意力，并通过扩展率为 4 的 mHC 连接到 FFN。

本文实例化了四个不同的模型：

• Dense-4B（基线）、
• MoE-27B、
• Engram-27B（从 MoE-27B 派生，减少路由专家数量，将释放的参数重新分配给 Engram 模块）
• Engram-40B（进一步扩展稀疏嵌入模块）

所有模型在激活参数数量上严格匹配。

4.2. 实验结果

结果如上表所示。

• 所有三个稀疏变体在相同的训练计算预算下，在所有基准测试中都显著优于同 FLOPs 的 Dense-4B 基线。
• 更重要的是，Engram-27B 相对于同参数和同 FLOPs 的 MoE-27B 基线持续改进。这些增益不仅限于知识密集型任务，在一般推理领域以及代码和数学推理中也观察到更显著的改进。
• 扩展到 Engram-40B 进一步降低了预训练损失并改进了大多数基准测试的性能。实验观察到，在训练接近结束时，Engram-40B 与基线之间的训练损失差距持续扩大，这表明在当前 token 预算内，扩展的记忆容量尚未完全饱和，有进一步改进的潜力。

5. 长上下文训练

通过将局部依赖建模卸载给静态查找（lookups），Engram 架构为管理全局上下文保留了宝贵的注意力容量。本文通过进行长上下文扩展训练来实证验证这一结构优势。

5.1. 实验设置

为获得长上下文能力，本文采用 DeepSeek-V3 中引入的上下文扩展策略。

在预训练阶段之后，本文在一个 32768 token 的上下文训练阶段应用 YaRN 进行上下文窗口扩展。

本文比较了四种不同模型配置的上下文扩展。

为了严格地基准测试架构效率，本文选取了 Engram-27B 的两个中间检查点。

5.2. 实验结果

评估结果如上表所示。

超越注意力机制的长上下文能力：

• 虽然注意力机制和位置编码为上下文处理提供了结构基础，但结果表明，长上下文性能不只由架构先验决定。
• 观察 Engram (41k → 50k) 的轨迹，发现长上下文性能随着预训练的进行而单调改善，即使控制相同的模型架构和上下文扩展阶段的固定计算预算。这表明长上下文性能与基础模型的通用建模能力内在耦合。
• 因此，严格的架构比较必须通过对齐基础模型损失来控制这个混淆变量，而不仅仅是统一训练步数。

受控设置下的架构优越性：基于上述原则，本文对 Engram 与 MoE 基线进行基准测试。当控制基础能力时，Engram 模块的效率增益变得明显：

• 等损失设置 (46k vs. 基线)：当比较具有 相同预训练损失（1.63）的 Engram-27B (46k) 与完全训练的 MoE-27B (50k) 时，Engram 显示出显著增益 (e.g., NIAH: 97.0 vs. 84.2; VT: 87.2 vs. 77.0)

• 等 FLOPs 设置 (50k vs. 基线)：在标准的等计算预算下，Engram-27B (50k) 进一步扩大了差距，在所有指标上确立了最高性能。

• 极端设置 (约 82% 计算)：即使是提前停止的 Engram-27B (41k) 与完全训练的 MoE-27B (50k) 相比仍然极具竞争力。它在 LongPPL 上与基线匹配，在 RULER 上超越基线，强调了 Engram 架构的内在优越性。

6. 分析

6.1. Engram 在功能上是否等同于增加了模型的深度？

当前的 LLM 缺乏专用的知识查找方式，它们依靠计算来模拟记忆检索。如上表所示，为了识别实体 “Diana, Princess of Wales”，LLM 必须消耗多个注意力层和 FFN 层来逐步组合特征。理论上，这一过程本可以通过一次知识查找操作来完成。

基于此，本文假设：通过为模型配备显式的知识查找能力，Engram 通过减轻模型早期特征组合阶段的负担，有效地模仿了模型深度的增加。为了验证这一假设，本文使用了两种机理可解释性工具：LogitLens 和中心核对齐分析。

1）加速预测收敛

首先使用 LogitLens 分析预测在各层之间的演变过程。

• 通过将每个中间层的隐藏状态投影到最终的 LM Head，来计算中间输出分布与模型最终输出分布之间的 KL 散度。
• 该指标量化了潜在表示距离 “准备就绪的预测” 有多近。

更早收敛：上图报告了各层的 KL 散度。与 MoE 基线相比，两个 Engram 变体都表现出更小的 KL 散度，最显著的差距出现在早期块。

陡峭的下降曲线：Engram 曲线的更陡峭下降表明，该模型完成特征组合的速度要快得多。

关键发现：这些观察结果与假设一致：通过显式地访问外部知识，Engram 减少了所需的计算步骤，从而在网络层级的更早阶段就达到了高置信度、有效的预测。

2）表征对齐与有效深度

为了进一步探究 Engram 层在语义上是否对应于基线的更深层，本文采用了 中心核对齐分析（Centered Kernel Alignment，CKA）。这是一种广泛建立的、用于比较表征结构的指标。

显著的对角线上移：上图可视化了叠加软对齐曲线（白色虚线）的相似度热图。观察到对角线的显著上移

这种一致的、非对角线的对齐，与 LogitLens 结果一致，证实了 Engram 在早期层就实现了更深层的表征。这验证了假设：通过显式查找绕过早阶段特征组合，Engram 在功能上等同于增加了模型的有效深度。

6.2. 结构消融与层敏感性

1）应在何处注入记忆？

为了研究深度敏感性，保持 Engram 预算固定（1.6B），但将其合并到单个 Engram 模块中，并扫描其插入层（从 1 到 12）。

早期注入：在第 2 层注入获得了最佳的单层性能（Val Loss = 1.770），优于第 1 层，并且随着插入点变深，性能下降。

双重权衡：

• 优势：尽早注入 Engram，使其能在主干网络消耗计算深度之前，就卸载局部模式重构，这符合主干网络的自然分层处理流程。
• 劣势：这会牺牲门控精度：早期的隐状态尚未通过注意力聚合足够的全局上下文，并且并行分支缺乏进行精细调制所需的表征差异。

最优放置：因此，最优放置需要在以下两者之间取得平衡：

• 尽早卸载（offload）静态局部模式
• 在后期利用更强的上下文查询进行门控。

分层设计的优越性：

•  虽然第 2 层在单一注入约束下是最优的，但研究发现，将相同的 1.6B 内存分成两个较小的模块（通过减少嵌入维度实现），并将它们放置在第 2 和第 6 层，效果更好（Val Loss = 1.768）。
• 这种分层设计结合了早期干预和丰富的后期上下文门控，调和了上述权衡。
• 更重要的是，分层插入也提供了一个实际的系统优势，能够更好地利用第 2.5 节讨论的内存层次结构。

2）哪些组件重要？

从基准配置出发，在保持 Engram 参数预算固定的情况下，消融各个设计选择。结果在上图用标记表示。

最重要的三个组件：产生了最显著增益的是：

• 多分支主干内的分支特定融合（multi branch）

• 上下文感知门控（gating）

• token 压缩

移除其中任何一个都会导致验证损失的最大回归。

影响较小的组件：移除轻量级深度卷积（conv）仅略微降低性能。

高阶 N-gram：在固定的 1.6B 预算下，将容量分配给 4-gram 略为次优——可能是因为它稀释了更频繁的 2/3-gram 模式的容量——但不排除在更大的内存规模下，高阶 N-gram 可能变得有益。

6.3. 敏感性分析

为了描述 Engram 模块的功能贡献，在推理过程中完全抑制稀疏嵌入输出，同时保持主干不变，以此来评估模型。

这种事后消融导致了训练-推理的不一致性，可能在复杂的、混合能力的任务中引入噪声。因此，优先分析 事实性知识 和 阅读理解——这是敏感性谱系的两个极端，在这种压力测试下表现出最高的信噪比。

如上图所示，结果揭示了一个尖锐的功能二分法：

• 事实性知识：遭受 灾难性崩溃，仅保留了原始性能的 29-44%（例如，TriviaQA 为 29%）。这 确认了 Engram 模块是参数化知识的主要存储库。

• 阅读理解：表现出 显著的韧性，保留了 81-93% 的性能（例如，C3 为 93%）。这表明 基于上下文的任务主要依赖于主干的注意力机制，而不是 Engram。

6.4. 系统效率

Engram 相对于基于路由的 MoE 的一个关键系统优势在于，其稀疏激活由 显式的、静态的哈希 ID 寻址。这产生了一个 严格确定性的内存访问模式：下一个 Engram 查找的索引在 token 序列已知时就已固定，并且可以在相应层执行之前计算出来。

实验设置： 

• 基于 nano-vLLM 实现了一个推理程序。
• 为了获得干净的延迟基线，不受 MoE 中专家并行通信模式的干扰，在两个密集主干上（Dense-4B 和 Dense-8B）进行基准测试。
• 在第二个 Transformer 块中插入了一个巨大的 1000 亿参数 Engram 层，整个嵌入表驻留在主机 DRAM 中。
• 在推理过程中，系统异步地为 Engram 层预取嵌入，PCIe 传输与第一个块的计算重叠。

如上表所示，卸载一个 100B 参数的嵌入表产生了 可忽略的吞吐量损失，在 8B 主干上峰值仅为 2.8%。这证实了早期密集块的计算强度为掩盖检索延迟提供了足够的时间窗口。关键的是，每步有效的通信量随激活槽的数量缩放，而不是总嵌入表大小。

实验的意义： 

• 这个实验是一个 保守的基线。
• 虽然第 2.5 节中的分层设计利用 Zipf 分布的局部性将频繁项缓存在 HBM 中，但实验设置强制所有检索都从主机内存穿越 PCIe 总线。
• 即使在这种基准检索策略下也产生了最小的开销，这强烈表明，一个完全优化的、具有局部性感知的实现将产生可忽略的吞吐量损失。

6.5. 案例研究：门控可视化

为了实证验证 Engram 是否按预期行为，在图 7 中可视化了 Engram-27B 在不同样本上的门控标量 α_t​。

明确的静态模式选择性：门控机制在 完成局部的、静态的模式时持续激活（显示为红色）。

英语示例：在多 token 命名实体（例如，“Alexander the Great”、“the Milky Way”）和公式化短语（例如，“By the way”、“Princess of Wales”）上观察到强烈的激活。

跨语言泛化：这种行为在各种语言中都能有效泛化。在中文示例中，Engram 识别并检索了独特的成语表达和历史实体，例如 “四大发明” 和 “张仲景”。