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

文章标题：通过可扩展查找的条件记忆：大型语言模型稀疏性的新维度
作者/机构：Xin Cheng (北京大学, DeepSeek-AI), Wangding Zeng (DeepSeek-AI), Damai Dai (DeepSeek-AI), Qinyu Chen (DeepSeek-AI), Bingxuan Wang (DeepSeek-AI), Zhenda Xie (DeepSeek-AI), Kezhao Huang (DeepSeek-AI), Xingkai Yu (DeepSeek-AI), Zhewen Hao (DeepSeek-AI), Yukun Li (DeepSeek-AI), Han Zhang (DeepSeek-AI), Huishuai Zhang (北京大学), Dongyan Zhao (北京大学), Wenfeng Liang (DeepSeek-AI)

A1 主要贡献

本文探讨了在大型语言模型中，除了通过专家混合（MoE）实现的条件计算之外，引入一种名为条件记忆的互补稀疏性轴。核心问题是，当前的Transformer模型缺乏原生的知识查找机制，不得不通过昂贵的计算来模拟知识检索过程，这在处理静态、局部和模式化的文本（如命名实体、固定短语）时效率低下。

研究目标是设计一种能够与现有MoE架构协同工作，高效处理静态知识的模块，并探索如何最优地在条件计算（MoE）和条件记忆之间分配模型参数。

为此，本文提出了Engram，一个基于经典n-gram思想但经过现代化改造的条件记忆模块。其主要贡献和创新点如下：

  1. 提出Engram模块：Engram通过哈希实现对n-gram的高效查找，并结合了分词器压缩、多头哈希、上下文感知门控和多分支架构集成等现代技术，以O(1)的复杂度从一个巨大的嵌入表中检索静态知识。这使得模型能够将静态模式的存储与动态的计算分离开来。


    图 1 | Engram 架构。该模块通过检索静态 n-gram 记忆并经由上下文感知门控将其与动态隐藏状态融合，从而增强主干网络。此模块仅应用于特定层，以将记忆与计算解耦，同时保持标准的输入嵌入和输出嵌入模块不变。

  2. 构建稀疏性分配理论：本文首次提出了稀疏性分配问题，即在固定的总参数和计算预算下，如何最优地在MoE专家（条件计算）和Engram记忆（条件记忆）之间分配稀疏容量。实验揭示了一个U形缩放定律，表明混合分配策略优于纯MoE模型，最佳分配比例约为将20%-25%的稀疏参数预算分配给Engram。

  3. 大规模实验验证：基于该分配定律，本文训练了27B和40B参数的Engram模型。在与参数量和FLOPs严格相等的MoE基线模型的比较中，Engram模型在各类任务上均取得显著优势。特别地，它不仅在知识密集型任务（如MMLU +3.4）上表现出色，在通用推理（BBH +5.0）、代码（HumanEval +3.0）和数学（MATH +2.4）等需要复杂推理的领域也取得了更大提升。

  4. 揭示工作机制：通过LogitLens和CKA等机制性分析工具，研究发现Engram通过从外部记忆中直接查找，减轻了模型早期层重构静态知识的负担，从而增加了模型的有效深度，使网络能将更多计算资源用于复杂推理。同时，它将局部依赖关系交由查找处理，释放了注意力机制的容量以关注全局上下文，极大地提升了模型的长上下文检索能力（例如，在MultiQuery NIAH任务上，准确率从84.2提升至97.0）。

  5. 实现基础设施感知的系统效率：Engram采用确定性寻址，与MoE的动态路由不同，这使其能够在运行时从主机内存（Host Memory）预取嵌入，并将通信与计算重叠。实验证明，将一个100B参数的记忆表卸载到主机内存中，仅带来可忽略的性能开销（<3%），有效绕过了GPU内存墙的限制。

A2 方法细节

2.1. 概述

Engram模块功能与流程。如图1所示，Engram是一个条件记忆模块，旨在通过结构上分离静态模式存储和动态计算来增强Transformer主干网络。形式上，对于输入序列 $x = (x_1, ..., x_T)$ 和第 $l$ 层的隐藏状态 $H^{(l)} \in \mathbb{R}^{T \times d}$，该模块在每个位置 $t$ 的处理分为两个功能阶段：检索和融合。首先，如2.2节详述，我们提取并压缩后缀n-gram，通过哈希来确定性地检索静态嵌入向量。随后，在2.3节中，这些检索到的嵌入被当前隐藏状态动态调制，并通过一个轻量级卷积进行优化。最后，我们在2.4节讨论与多分支架构的集成，并在2.5节讨论系统级设计。

2.2. 通过哈希n-gram进行稀疏检索

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

分词器压缩。尽管n-gram模型通常直接在分词器输出上操作，但标准的子词分词器优先考虑无损重构，常为语义等价的术语（如Apple vs. ␣apple）分配不相交的ID【索引62，Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing，2018，EMNLP；索引81，Starcoder: may the source be with you!，2023，TMLR】。为了最大化语义密度，我们实现了一个词汇表投影层。具体来说，我们预先计算一个满射函数 $P : V \to V'$，该函数基于标准化的文本等价性（使用NFKC【索引117，Unicode standard annex #15: Unicode normalization forms，2025，The Unicode Consortium】、小写化等）将原始的token ID折叠成规范化的标识符。在实践中，对于一个128k词汇表的分词器，这个过程实现了23%的有效词汇量减少（见附录C）。形式上，对于位置 $t$ 的一个token，我们将其原始ID $x_t$ 映射到一个规范ID $x'_t = P(x_t)$，以形成后缀n-gram $g_{t,n} = (x'_{t-n+1}, ..., x'_t)$。

多头哈希。直接参数化所有可能的n-gram的组合空间是不可行的。我们遵循了Tito Svenstrup等人【索引113，Hash embeddings for efficient word representations，2017，NIPS】的方法，采用了一种基于哈希的方法。为了减轻冲突，我们为每个n-gram阶数 $n$ 采用了 $K$ 个不同的哈希头。每个头 $k$ 通过一个确定性函数 $\phi_{n,k}$ 将压缩后的上下文映射到嵌入表 $E_{n,k}$（其大小为素数 $M_{n,k}$）内的一个索引：

$$z_{t,n,k} \triangleq \varphi_{n,k}(g_{t,n}), \quad \mathbf{e}_{t,n,k} = \mathbf{E}_{n,k}[z_{t,n,k}].$$

在实践中，$\phi_{n,k}$ 是通过一个轻量级的乘法-异或哈希实现的。我们通过拼接所有检索到的嵌入来构建最终的记忆向量 $e_t \in \mathbb{R}^{d_{\text{mem}}}$：

$$\mathbf{e}_t \triangleq \prod_{n=2}^N \prod_{k=1}^K \mathbf{e}_{t,n,k} .$$

2.3. 上下文感知门控

动态调制静态嵌入。检索到的嵌入 $e_t$ 作为上下文无关的先验知识。然而，由于其静态性，它们天生缺乏上下文适应能力，并可能因哈希冲突或多义性而受到噪声影响【索引48，Polysemy—Evidence from linguistics, behavioral science, and contextualized language models，2024，Computational Linguistics】。为了增强表达能力并解决这种模糊性，我们采用了一种受注意力机制启发的上下文感知门控机制【索引2，Neural machine translation by jointly learning to align and translate，2015，ICLR；索引114，Attention is all you need，2017，NIPS】。具体来说，我们将当前的隐藏状态 $h_t$（它通过之前的注意力层聚合了全局上下文）用作动态的查询（Query），而检索到的记忆 $e_t$ 同时作为键（Key）和值（Value）投影的来源：

$$\mathbf{k}_t = \mathbf{W}_K\mathbf{e}_t, \quad \mathbf{v}_t = \mathbf{W}_V\mathbf{e}_t$$

其中 $W_K, W_V$ 是可学习的投影矩阵。为了确保梯度稳定性【索引25，Scaling vision transformers to 22 billion parameters，2023，ICML】，我们在计算标量门 $\alpha_t \in (0, 1)$ 之前，对查询和键应用了RMSNorm【索引127，Root mean square layer normalization，2019，NIPS】：

$$\alpha_{t}=\sigma\left(\frac{\operatorname{RMSNorm}\left(\mathbf{h}_{t}\right)^{\top} \operatorname{RMSNorm}\left(\mathbf{k}_{t}\right)}{\sqrt{d}}\right) .$$

门控输出与设计意图。门控输出定义为 $\tilde{v}_t = \alpha_t \cdot v_t$。这种设计强制了语义对齐：如果检索到的记忆 $e_t$ 与当前上下文 $h_t$ 相矛盾，门 $\alpha_t$ 会趋向于零，从而有效地抑制噪声。

深度可分离因果卷积的应用。最后，为了扩大感受野并增强模型的非线性，我们引入了一个短的、深度可分离的因果卷积【索引45，Efficiently modeling long sequences with structured state spaces，2022，ICLR；索引94，RWKV: reinventing rnns for the transformer era，2023，EMNLP】。令 $\tilde{V} \in \mathbb{R}^{T \times d}$ 表示门控值的序列。使用大小为 $k$（设为4）、扩张率为 $d$（设为最大n-gram阶数）的核和SiLU激活函数【索引30，Sigmoid-weighted linear units for neural network function approximation in reinforcement learning，2018，Neural networks】，最终输出 $Y$ 计算如下：

$$\mathbf{Y} = \operatorname{SiLU}(\operatorname{Conv1D}(\operatorname{RMSNorm}(\tilde{\mathbf{V}}))) + \tilde{\mathbf{V}},$$

模块集成方式。Engram模块通过残差连接集成到主干网络中：$H^{(l)} \leftarrow H^{(l)} + Y$，之后是标准的Attention和MoE模块。关键的是，Engram并非应用于每一层；其具体放置位置由2.5节详述的系统级延迟约束决定。


图 2 | Engram的系统实现。(a) 训练阶段：巨大的嵌入表被分片到可用的GPU上。采用All-to-All通信原语来跨设备检索活动的嵌入行。(b) 推理阶段：Engram表被卸载到主机内存。通过利用确定性的检索逻辑，主机异步预取和传输嵌入，将通信与前序Transformer块的设备上计算重叠。

2.4. 与多分支架构的集成

采用多分支架构作为主干。在这项工作中，我们没有采用标准的单流连接【索引49，Deep residual learning for image recognition，2016，CVPR】，而是采用了先进的多分支架构作为我们的默认主干，选择它是因为其卓越的建模能力【索引72，Fractalnet: Ultra-deep neural networks without residuals，2017，ICLR；索引109，Going deeper with convolutions，2015，CVPR；索引119，mhc: Manifold-constrained hyper-connections，2025，arXiv；索引129，Hyperconnections，2025，ICLR】。该架构的一个决定性特征是将残差流扩展为 $B$ 个平行分支，信息流由可学习的连接权重进行调制。

多分支架构下的结构优化。尽管Engram模块本质上是拓扑无关的，但将其适应于这种多分支框架需要进行结构优化以平衡效率和表达能力。具体来说，我们实现了一种参数共享策略：一个稀疏嵌入表和一个值投影矩阵 $W_V$ 在所有 $B$ 个分支间共享，而采用 $B$ 个不同的键投影矩阵 $\{W^{(m)}_K\}_{m=1}^B$ 以实现分支特定的门控行为。对于第 $m$ 个分支，其隐藏状态为 $h^{(m)}_t$，分支特定的门控信号计算如下：

$$\alpha_t^{(m)}=\sigma\left(\frac{\text{RMSNorm}(\mathbf{h}_t^{(m)})^\top \text{RMSNorm}(\mathbf{W}_K^{(m)} \mathbf{e}_t)}{\sqrt{d}}\right).$$

融合计算与效率。检索到的记忆随后被这些应用于共享值向量的独立门控进行调制：$u^{(m)}_t = \alpha^{(m)}_t \cdot (W_V e_t)$。这种设计允许线性投影（一个 $W_V$ 和 $B$ 个不同的 $W^{(m)}_K$）融合成一个单一的密集FP8矩阵乘法，从而最大化现代GPU的计算利用率。除非另有说明，所有实验都利用了与Manifold-Constrained Hyper-Connections ($B=4$)【索引119，mhc: Manifold-constrained hyper-connections，2025，arXiv】的这种集成。


图 3 | 稀疏性分配与Engram扩展。左图：不同分配比例 $\rho$ 下的验证损失。显示了两种计算预算（2e20和6e20 FLOPs）。两种情况下都呈现U形，混合分配优于纯MoE。右图：无限内存情境下的扩展行为。验证损失相对于嵌入数量呈现对数线性趋势。

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

确定性检索机制的优势。扩展记忆增强模型的瓶颈通常受限于GPU高带宽内存（HBM）的容量。然而，Engram的确定性检索机制天然支持参数存储与计算资源的解耦。与依赖运行时隐藏状态进行动态路由的MoE不同，Engram的检索索引仅依赖于输入token序列。这种可预测性为训练和推理阶段的专门优化策略提供了便利，如图2所示。

训练阶段的并行策略。在训练期间，为了容纳大规模嵌入表，我们通过将表分片到所有可用的GPU上，采用了标准的模型并行化。我们使用All-to-All通信原语在正向传播中收集活动的行，并在反向传播中分发梯度，从而使总内存容量能够随加速器数量线性扩展。

推理阶段的预取与重叠策略。在推理期间，这种确定性特性使得可以采用一种预取并重叠的策略。由于记忆索引在正向传播之前就已知，系统可以通过PCIe从充裕的主机内存中异步检索嵌入。为了有效地掩盖通信延迟，Engram模块被放置在主干网络的特定层中，利用前序层的计算作为缓冲以防止GPU停顿。这需要一种硬件-算法协同设计策略：虽然将Engram放置在更深层可以延长用于隐藏延迟的计算窗口，但我们在6.2节的消融实验表明，建模性能偏好于早期干预以卸载局部模式重构。因此，最优的放置必须同时满足建模和系统延迟的约束。

利用Zipfian分布的多级缓存。此外，自然语言的n-gram本身遵循齐夫定律（Zipfian distribution）【索引12，Human behavior and the principle of least effort: An introduction to human ecology，1950，Language；索引96，Zipf’s word frequency law in natural language: A critical review and future directions，2014，Psychonomic bulletin & review】，即一小部分模式占据了绝大多数的内存访问。这种统计特性启发了一种多级缓存层次结构：频繁访问的嵌入可以被缓存在更快的存储层（如GPU HBM或主机DRAM），而稀有模式的长尾则存放在较慢、高容量的介质中（如NVMe SSD）。这种分层结构使得Engram能够扩展到海量内存容量，而对有效延迟的影响最小。

A3 关键Observation/设计原则

Engram作为条件记忆的一种实例，其结构与MoE专家提供的条件计算是互补的。本节研究了这种二元性的扩展属性以及如何优化稀疏容量的分配。具体来说，我们的研究由两个关键问题驱动：
1. 有限约束下的分配。当总参数和训练计算量固定时（参数量和FLOPs相同），我们应如何在MoE专家和Engram嵌入之间划分稀疏容量？
2. 无限内存情境。考虑到Engram的非扩展性O(1)开销，如果内存预算被放宽或积极扩展，Engram本身会表现出什么样的扩展行为？

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

计算匹配的公式化。我们使用三个参数指标来分析这种权衡：
* $P_{\text{tot}}$: 总可训练参数，不包括词汇嵌入和LM头。
* $P_{\text{act}}$: 每个token激活的参数。这个量决定了训练成本（FLOPs）。
* $P_{\text{sparse}} \triangleq P_{\text{tot}} - P_{\text{act}}$: 未激活的参数，代表了可用于扩展模型大小而不增加计算成本的“免费”参数预算（例如，未被选择的专家或未被检索的嵌入）。

我们在每个FLOPs预算内保持 $P_{\text{tot}}$ 和 $P_{\text{act}}$ 固定，这样模型就拥有相同的参数数量和相同的单token FLOPs。对于MoE， $P_{\text{act}}$ 由top-k个被选中的专家决定，而未被选中专家的参数贡献给 $P_{\text{sparse}}$。对于Engram，每个token只检索恒定数量的槽位，因此扩展嵌入槽的数量会增加 $P_{\text{tot}}$ 而不增加单token FLOPs。

分配比例定义。我们定义分配比例 $\rho \in [0, 1]$ 为分配给MoE专家容量的非激活参数预算的比例：

$$P_{\mathrm{MoE}}^{(\text{sparse})} = \rho P_{\text{sparse}}, \quad P_{\text{Engram}} = (1 - \rho) P_{\text{sparse}}.$$

直观地：
* $\rho = 1$ 对应于一个纯MoE模型（所有非激活参数都是路由专家）。
* $\rho < 1$ 减少了路由专家的数量，并将释放的参数重新分配给Engram嵌入槽。

实验协议。我们评估了两种计算规模下的这种权衡，并在两种设置中都保持了恒定的稀疏比率 $P_{\text{tot}} / P_{\text{act}} \approx 10$：
* $C = 2 \times 10^{20}$ FLOPs: $P_{\text{tot}} \approx 5.7B$，$P_{\text{act}} = 568M$。基线模型（$\rho = 1$）共有106个专家。
* $C = 6 \times 10^{20}$ FLOPs: $P_{\text{tot}} \approx 9.9B$，$P_{\text{act}} = 993M$。基线模型（$\rho = 1$）共有99个专家。

对于不同的 $\rho$，我们仅通过调整路由专家的数量和Engram嵌入槽的数量来构建相应的模型。所有运行都使用相同的训练流程和优化超参数。

结果与分析。图3（左）揭示了验证损失与分配比例 $\rho$ 之间存在一致的U形关系。值得注意的是，即使MoE的分配比例降至仅 $\rho \approx 40\%$（即5.7B模型共有46个专家，9.9B模型共有43个专家），Engram模型的性能也与纯MoE基线（$\rho = 100\%$）相当。此外，纯MoE基线被证明是次优的：将大约20%–25%的稀疏参数预算重新分配给Engram可以产生最佳性能。从数量上看，在10B规模下（$C = 6 \times 10^{20}$），验证损失从1.7248（$\rho = 100\%$）提高到最优值 $\rho \approx 80\%$ 附近的1.7109（$\Delta = 0.0139$）。关键是，这个最优值的位置在不同规模下是稳定的（$\rho \approx 75\%–80\%$），这表明在所考察的规模下（固定稀疏度），存在一个稳健的分配偏好。这种观察到的U形证实了两个模块之间的结构互补性：
* MoE主导（$\rho \to 100\%$）：模型缺乏用于静态模式的专用内存，迫使其通过深度和计算来低效地重建它们。
* Engram主导（$\rho \to 0\%$）：模型失去了条件计算能力，损害了需要动态、上下文依赖推理的任务；在这种情况下，内存无法替代计算。

3.2. 无限内存情境下的Engram

积极内存扩展的研究。在3.1节中，我们在固定参数预算下优化了分配。现在我们探索互补的设置：积极的内存扩展。这项研究的动机是Engram具有将存储与计算解耦的独特能力，详见2.5节。

实验协议。我们使用一个固定的MoE主干网络，其 $P_{\text{tot}} \approx 3B$，$P_{\text{act}} = 568M$，训练了100B个token以确保收敛。在这个主干网络之上，我们附加了一个Engram表，并将槽位数 $S$ 从 $2.58 \times 10^5$ 扫描到 $1.0 \times 10^7$（增加了约130亿参数）。作为基线，我们与OverEncoding【索引55，Over-tokenized transformer: Vocabulary is generally worth scaling，2025，ICML】进行比较，后者通过与词汇嵌入求平均来整合n-gram嵌入。我们注意到，虽然其他工作如SCONE【索引123，Scaling embedding layers in language models，2025，arXiv】也研究了大规模嵌入，但它主要面向推理，并包括额外的模块（f-gram模型）和额外的训练FLOPs，使其与本研究严格的等计算约束不兼容。

结果。图3（右）表明，扩展内存槽的数量能够清晰且持续地改善验证损失。在探索的范围内，曲线遵循严格的幂律（在对数空间中呈线性），这表明Engram提供了一个可预测的扩展旋钮：更大的内存持续带来回报，而无需额外的计算。关键在于扩展效率：虽然OverEncoding的直接平均方法也受益于更大的内存表，但Engram从相同的内存预算中释放了更大的扩展潜力。结合3.1节的分配定律，这些结果证实了条件记忆是与MoE的条件计算互补的一个独特的、可扩展的稀疏容量轴。

A4 实验环境

  • 数据集：

    • 预训练：在一个包含2620亿token的语料库上进行。
    • 分词器：使用DeepSeek-v3的分词器，词汇表大小为128k。

  • 模型架构：

    • 通用配置：所有模型均基于一个30层的Transformer主干，隐藏层维度为2560。每个块集成一个32头的多头潜在注意力（MLA），通过mHC（扩展率为4）连接到FFN。
    • Dense-4B：基线模型，总参数4.1B，激活参数3.8B。
    • MoE-27B：总参数26.7B，激活参数3.8B。包含72个路由专家和2个共享专家，每个token激活top-6的路由专家。
    • Engram-27B：总参数26.7B，激活参数3.8B。将MoE-27B的路由专家从72个减少到55个，节约的参数重新分配给一个5.7B参数的Engram模块。Engram模块被放置在第2和第15层。
    • Engram-40B：总参数39.5B，激活参数3.8B。在Engram-27B的基础上，将Engram模块扩展到18.5B参数。

  • 硬件配置：

    • 推理吞吐量测试使用了NVIDIA H800 GPU。
    • 训练硬件未明确说明，但属于大规模集群。

  • 软件配置：

    • 优化器：主干网络使用Muon，Engram嵌入参数使用Adam。
    • 学习率调度：使用Step Decay。
    • 长上下文扩展：采用YaRN方法将上下文窗口扩展到32768。

A4 实验结果

大规模预训练结果 (见表1)

  1. 稀疏模型优于密集模型：与之前的研究一致，所有三种稀疏架构（MoE-27B, Engram-27B/40B）在相同的训练计算预算下，性能显著优于FLOPs相等的Dense-4B基线。
  2. Engram优于纯MoE：在参数量和FLOPs严格相等的情况下，Engram-27B在所有基准测试上都持续优于MoE-27B基线。
  3. 增益的普遍性：Engram带来的性能提升不仅限于知识密集型任务（如MMLU: +3.0, CMMLU: +4.0），在通用推理（BBH: +5.0, ARC-Challenge: +3.7）、代码（HumanEval: +3.0）和数学推理（MATH: +2.4）等领域，提升甚至更为显著。这支持了“引入专用知识查找原语能提升表示效率”的假设。
  4. Engram的扩展性：Engram-40B进一步降低了预训练损失，并在多数基准上提升了性能。其训练损失曲线显示，在当前训练token预算内，其性能潜力尚未完全饱和。

Table 1 | 密集、MoE和Engram模型之间的预训练性能比较。所有模型均使用262B个token进行训练，并且激活参数（3.8B）相匹配。Engram-27B与MoE-27B的参数量相同，通过将路由专家（72 → 55）的参数重新分配给一个5.7B参数的Engram记忆。Engram-40B在保持激活参数预算固定的同时，进一步增加了Engram记忆（18.5B参数）。完整的训练时基准测试轨迹报告在附录B中。

长上下文训练结果 (见表2)
本实验旨在验证Engram通过卸载局部依赖建模来释放注意力容量，从而提升全局上下文处理能力的假设。
1. 基础模型能力的重要性：实验表明，长上下文性能与基础模型的通用建模能力紧密相关。即使是同一个Engram模型，在预训练后期（50k steps）的版本比中期（41k steps）的版本在长上下文任务上表现更好，说明更好的基础模型是长上下文能力的基础。
2. 架构的优越性：
* 等损失设置 (Iso-Loss)：当Engram-27B (46k steps)与预训练损失相同的MoE-27B (50k steps)进行比较时，Engram在复杂检索任务（如Multi-Query NIAH: 97.0 vs 84.2; VT: 87.2 vs 77.0）上表现出显著优势，这严格证明了其架构的效率。
* 等计算量设置 (Iso-FLOPs)：在标准的等计算预算下，完全训练的Engram-27B (50k steps)进一步拉大了与基线的差距。
* 计算量更少的情况：即使是只用了约82%预训练计算量的Engram-27B (41k steps)，其长上下文性能仍然能与完全训练的MoE-27B (50k steps)相媲美甚至在RULER上超越后者。

Table 2 | 长上下文性能比较。括号内的值（例如(50k, 1.62)）表示预训练步骤和长上下文扩展前对应的损失。两个主要发现：(1) 仅用82%的预训练FLOPs（41k vs. 50k），Engram-27B在LongPPL（Fang等人）上与基线性能持平，同时在RULER（Hsieh等人）上实现了显著更高的准确率；(2) 在等预训练损失（46k）和等预训练FLOPs（50k）的设置下，Engram-27B在所有指标上都大幅优于基线。粗体表示最佳，下划线表示次佳。

A7 补充细节

6. 分析

本节深入研究了Engram的内部机制，包括其有效深度（6.1节）、核心模块设计（6.2节）、参数敏感性（6.3节），并评估了带卸载的推理吞吐量（6.4节），最后以一个案例研究作结（6.5节）。

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

背景与假设。当前的LLM缺乏专用的知识查找原语，它们依赖计算来模拟记忆调用。如表3所示，为了识别实体“戴安娜，威尔士王妃”，LLM必须消耗多层Attention和FFN来逐步组合特征【索引46，Patchscopes: A unifying framework for inspecting hidden representations of language models，2024，ICML；索引58，Exploring concept depth: How large language models acquire knowledge and concept at different layers?，2025，COLING；索引80，Echoes of bert: Do modern language models rediscover the classical nlp pipeline?，2025，arXiv】，而这个过程理论上可以通过一次知识查找操作来识别。基于此，我们假设通过为模型配备明确的知识查找能力，Engram通过减轻模型在特征组合早期阶段的负担，有效地模拟了模型深度的增加。为了验证这一假设，我们采用了两种机制性可解释性工具：LogitLens【索引3，Eliciting latent predictions from transformers with the tuned lens，2023，arXiv；索引90，interpreting gpt: the logit lens，2020，LessWrong】和中心核对齐分析（CKA）【索引23，Reliability of CKA as a similarity measure in deep learning，2023，ICLR；索引60，Similarity of neural network representations revisited，2019，ICML】。

6.1.1. 加速的预测收敛

通过LogitLens分析预测演变。我们首先使用LogitLens【索引90，interpreting gpt: the logit lens，2020，LessWrong】分析了各层预测的演变。通过将每个中间层的隐藏状态与最终的LM Head进行投影，我们计算了中间输出分布与模型最终输出分布之间的Kullback–Leibler散度【索引63，On information and sufficiency，1951，The annals of mathematical statistics】。该指标量化了一个潜在表示距离“准备好进行预测”的程度【索引3，Eliciting latent predictions from transformers with the tuned lens，2023，arXiv；索引20，Do language models use their depth efficiently?，2025，arXiv】。

结果分析。图4（a）报告了逐层的KL散度。与MoE基线相比，两种Engram变体都系统性地表现出更小的KL散度，其中最显著的差距出现在早期块中。Engram曲线更陡峭的下降表明模型更快地完成了特征组合。这一观察结果与我们的假设一致：通过显式访问外部知识，Engram减少了所需的计算步骤，从而在网络层次结构的早期就达到了高置信度的有效预测。

Table 3 | Ghandeharioun等人（2024）复现的实体解析示例。此表说明了LLM如何通过多层注意力和FFN逐步整合上下文token，以构建实体“戴安娜，威尔士王妃”的内部表示。“Latent State Translation”列显示了由PatchScope（Ghandeharioun等人，2024）为最后一个token“Wales”自动生成的文本，而“Explanation”列则展示了原作者提供的手动解释。

6.1.2. 表征对齐与有效深度

通过CKA分析层间语义对应关系。为了进一步研究Engram的层是否在语义上对应于基线模型的更深层，我们采用了中心核对齐（CKA），这是一个广泛用于比较表征结构的度量【索引60，Similarity of neural network representations revisited，2019，ICML；索引61，Representational similarity analysis-connecting the branches of systems neuroscience，2008，Frontiers in systems neuroscience】。给定两组表示 $X$ 和 $Y$（例如，来自不同模型或层的激活），CKA定义为：

$$\text{CKA}(K, L) = \frac{\text{HSIC}(K, L)}{\sqrt{\text{HSIC}(K, K)\text{HSIC}(L, L)}}$$

其中 $K = XX^\top$ 和 $L = YY^\top$ 表示Gram矩阵（使用线性核），HSIC是Hilbert-Schmidt独立性准则【索引47，Measuring statistical dependence with hilbert-schmidt norms，2005，ALT】。我们采用带有无偏HSIC估计器的小批量实现【索引23，Reliability of CKA as a similarity measure in deep learning，2023，ICLR】，并在Few-NERD数据集【索引26，Few-nerd: A fewshot named entity recognition dataset，2021，ACL】上进行评估，提取与命名实体最后一个token对应的隐藏状态。

软对齐索引的定义。为了严格量化层与层之间的对应关系，我们首先计算成对的CKA相似度矩阵 $S \in [0, 1]^{L \times L}$，其中 $L$ 是层数。然后，我们引入一个软对齐索引 $a_j$，定义为对于每个Engram层 $j$，其与top-k个最相似的MoE层的加权质心：

$$a_j = \frac{\sum_{i \in \mathcal{I}_j} S_{i,j} \cdot i}{\sum_{i \in \mathcal{I}_j} S_{i,j}}, \text{ where } \mathcal{I}_j = \mathop{\text{argtopk}}_i (S_{i,j}).$$

这里，$S_{i,j}$ 表示MoE层 $i$ 和Engram层 $j$ 之间的相似度得分。索引 $a_j$ 作为对应于Engram层 $j$ 的“有效MoE深度”的稳健代理，利用top-k过滤（$k=5$）来减轻低相似度噪声的影响。


图 4 | 表征对齐与收敛速度分析。(a) 通过LogitLens（nostalgebraist，2020）进行的逐层KL散度分析。早期层中持续较低的散度表明Engram加速了预测的收敛。(b-c) 由CKA（Kornblith等人，2019）计算的相似度热图。高相似度对角线的明显上移表明Engram的浅层在功能上等同于MoE模型的深层，从而有效增加了模型的深度。

结果与结论。图4（b）–（c）展示了相似度热图，并叠加了软对齐曲线（白色虚线）。我们观察到对角线有明显的向上偏移，这意味着对于很大范围的层，$a_j > j$。例如，Engram-27B第5层形成的表示与MoE基线大约第12层的表示最为接近。这种一致的离对角线偏移，与LogitLens的结果（6.1.1节）相符，证实了Engram在更早的层级上实现了更深的表示。这验证了我们的核心假设：通过显式查找绕过早期的特征组合，Engram在功能上等同于增加了模型的有效深度。

6.2. 结构消融与层敏感性

实验设置。在本节中，我们在一个受控的环境下对Engram进行消融实验，以研究每个关键模块设计的有效性。除非另有说明，主干网络是一个12层的3B MoE模型（0.56B激活参数），训练了100B个token。图5报告了验证损失。橙色虚线表示3B MoE基线（验证损失=1.808）。

参考配置。我们用一个固定的1.6B参数的Engram记忆来增强主干网络。我们的参考模型使用{2, 3}-grams，并在第2层和第6层插入Engram，实现了验证损失=1.768，比MoE基线有显著改进（Δ = 0.04）。下面所有的结构消融都是相对于这个参考配置定义的。


图 5 | 架构消融结果。我们在两种设置下比较了3B MoE基线和Engram的变体：(1) 层敏感性（深蓝色曲线）：扫描单个Engram模块的插入深度，证实了早期注入（第2层）是最佳选择，而在更深层中效果会下降。(2) 组件消融（右侧标记）：从参考配置中移除子模块，证明了多分支集成、分词器压缩和上下文感知门控的重要性。

记忆应注入何处？。为了研究深度敏感性，我们保持Engram预算固定（1.6B），但将其整合到单个Engram模块中，并扫描其从第1层到第12层的插入位置（图5中的深蓝色“Layer Sweep”曲线）。这个实验揭示了Engram放置的一个内在权衡。

放置位置的权衡。早期注入Engram可以让它在主干网络消耗计算深度之前卸载局部模式的重构，这与主干网络的自然层次化处理方式相符【索引46，Patchscopes: A unifying framework for inspecting hidden representations of language models，2024，ICML；索引58，Exploring concept depth: How large language models acquire knowledge and concept at different layers?，2025，COLING；索引80，Echoes of bert: Do modern language models rediscover the classical nlp pipeline?，2025，arXiv；索引112，Bert rediscovers the classical nlp pipeline，2019，ACL】。然而，这会带来门控精度的代价：早期的隐藏状态尚未通过注意力聚合足够的全局上下文，并且并行分支缺乏进行细粒度调制所需的表征差异【索引119，mhc: Manifold-constrained hyper-connections，2025，arXiv；索引129，Hyperconnections，2025，ICLR】。因此，最佳放置需要平衡（i）早期卸载静态局部模式和（ii）利用更强的上下文查询进行后期门控。

实验结论。扫描结果显示，第2层达到了最佳的单层性能（验证损失 = 1.770），优于第1层，并且随着插入点向深层移动，性能逐渐下降。这表明一轮注意力已经足以提供一个有意义的上下文化 $h_t$ 用于门控，同时仍然足够早，可以替代主干网络底层局部的聚合。虽然在单次注入的约束下第2层是最佳的，但我们发现，将同样的1.6B内存分成两个较小的模块（通过减少嵌入维度 $d_{\text{mem}}$ 实现）并分别放置在第2层和第6层，表现甚至更好（验证损失 = 1.768）。这种分层设计通过结合早期干预和丰富的后期上下文门控，调和了上述权衡。更重要的是，分层插入还提供了实际的系统优势，能够更好地利用内存层次结构，如2.5节所讨论。

哪些组件重要？。从参考配置出发，我们在保持Engram参数预算固定的情况下，对单个设计选择进行消融。结果在图5中用标记表示。我们发现三个组件带来了最显著的增益：（i）多分支主干内的分支特定融合，（ii）上下文感知门控，以及（iii)分词器压缩。移除其中任何一个都会导致验证损失的最大回归。具体来说，对于“w/o multi branch”消融，我们保留了mHC主干结构，但用一个应用于预映射后隐藏状态 $H_{\text{premap}}$【索引119，mhc: Manifold-constrained hyper-connections，2025，arXiv】的单一Engram融合替换了分支特定的门控。其他变化的影响较小：移除轻量级深度可分离卷积仅轻微降低性能。在固定的1.6B预算下，将容量分配给4-grams略显次优——可能是因为它稀释了更频繁的2/3-gram模式的容量——但我们不排除在更大的内存规模下，更高阶的n-gram会变得有益。

6.3. 敏感性分析

实验设计。为了表征Engram模块的功能贡献，我们通过在推理时完全抑制稀疏嵌入的输出来评估模型，同时保持主干网络不变。关键是，这种事后消融会引入训练-推理不一致性，可能在复杂的混合能力任务中引入噪声。因此，我们优先分析事实知识和阅读理解——这两个敏感性谱系的极端——它们在这种压力测试下表现出最高的信噪比。

结果分析。如图6所示，结果揭示了明显的功能二分法。事实知识基准测试遭受了灾难性的崩溃，仅保留了原始性能的29–44%（例如，TriviaQA为29%），证实了Engram模块是参数化知识的主要存储库。相反，阅读理解任务表现出显著的弹性，保留了81–93%的性能（例如，C3为93%），这表明基于上下文的任务主要依赖于主干网络的注意力机制而不是Engram。


图 6 | Engram消融下的性能保留率。事实知识严重依赖Engram模块，而阅读理解则主要由主干网络保留。

6.4. 系统效率

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

实验设置。我们实现了一个基于nano-vLLM1的推理测试工具——这是一个工业标准vLLM引擎【索引66，Efficient memory management for large language model serving with pagedattention，2023，SOSP】的简化原型。为了获得一个干净的延迟基线，避免MoE中专家并行带来的混淆性通信模式，我们在两个密集主干（Dense-4B和Dense-8B）上进行基准测试。我们在第二个Transformer块中插入一个巨大的100B参数的Engram层，整个嵌入表驻留在主机DRAM中。在推理期间，系统异步预取Engram层的嵌入，将PCIe传输与第一个块的计算重叠。

结果。如表4所示，卸载一个100B参数的嵌入表只带来了微不足道的吞吐量损失，在8B主干上最高仅为2.8%。这证实了早期密集块的计算强度提供了一个足够的时间窗口来掩盖检索延迟。关键是，每一步的有效通信量与激活的槽数成比例，而不是总嵌入表的大小。

结论。至关重要的是，这个实验是一个保守的基线。虽然2.5节中的层次化设计利用了Zipfian局部性来将频繁项缓存到HBM中，但我们的实验设置强制所有检索都从主机内存通过PCIe总线。这种基线检索策略产生的开销极小，这一事实强烈表明，一个完全优化的、具有局部性感知的实现将产生可忽略的吞吐量损失。

Table 4 | 端到端推理吞吐量。我们测量了将一个100B参数的Engram层完全卸载到主机内存时的推理吞吐量。

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

目的与方法。在2.3节中，我们引入了上下文感知门控机制，旨在动态调节检索到的静态记忆与主干网络的集成。为了实证验证Engram是否按预期工作，我们在图7中可视化了Engram-27B2在各种样本上的门控标量 $\alpha_t$。

结果分析。结果显示出一种独特的选择性模式。门控机制在完成局部的、静态的模式时会持续激活（以红色显示）。在英语中，我们观察到在多词命名实体（如“Alexander the Great”、“the Milky Way”）和固定短语（如“By the way”、“Princess of Wales”）上有强烈的激活。这种行为在不同语言中都能有效泛化。在中文示例中，Engram识别并检索了独特的成语和历史实体，如“四大发明”（四 发明）和“张仲景”（张 仲景）。这些定性结果证实，Engram成功地识别并处理了模式化的语言依赖关系，有效地减轻了Transformer主干网络记忆这些静态关联的负担。

7. 相关工作

n-gram建模与嵌入扩展。源于香农的框架【索引102，A mathematical theory of communication，1948，The Bell system technical journal】，n-gram模型依赖局部历史来预测token，传统上采用平滑技术【索引59，Estimation of probabilities from sparse data for the language model component of a speech recognizer，1987，IEEE Trans. Acoust. Speech Signal Process.；索引61，Improved backing-off for m-gram language modeling，1995，ICASSP】来缓解数据稀疏问题。尽管范式转向了捕捉长距离依赖的神经架构【索引4，A neural probabilistic language model，2003，JMLR】，n-gram查找的计算效率在现代表示学习中得以保留，如FastText【索引9，Enriching word vectors with subword information，2017，TACL】等开创性工作所示。


图 7 | Engram门控机制的可视化。热图强度对应于门控标量 $\alpha_t \in [0, 1]$ 的大小，深红色表示更强的激活。因为Engram在后缀n-gram上操作（这里n=3），在特定token $x_t$ 上的高激活意味着以该token结尾的先前token（例如，以 $x_t$ 结尾的短语）被识别为从内存中有效检索的静态模式。

近期发展与本文差异。最近，这一范式作为嵌入扩展重新兴起。虽然诸如PerLayer Embeddings【索引110，Gemma 3n，2025】和DeepEmbed【索引98，Rwkv architecture history，2025】等架构通过巨大的表来扩展容量，但一条与我们方法最相关的开创性研究路线是直接将组合的n-gram结构整合到表示空间中。SuperBPE【索引84，SuperBPE: Space travel for language models，2025，CoLM】和SCONE【索引123，Scaling embedding layers in language models，2025，arXiv】明确针对高频模式：前者通过将多词表达合并为“超词”token，后者通过一个辅助编码模型。与此同时，OverEncoding【索引55，Over-tokenized transformer: Vocabulary is generally worth scaling，2025，ICML】和Byte Latent Transformer (BLT)【索引92，Byte latent transformer: Patches scale better than tokens，2025，ACL】采用哈希n-gram嵌入来分别在token和字节级别捕捉局部依赖。这些研究共同证明了通过n-gram表示以最小计算开销扩展参数的有效性。虽然这些方法在各自的设置中取得了显著的增益，但我们的工作在两个关键维度上根本不同。

  • 首先，关于建模和评估协议。先前的方法通常将n-gram嵌入视为外部增强，而没有在严格公平的比较协议下验证其效率。例如，SCONE【索引123，Scaling embedding layers in language models，2025，arXiv】是面向推理的，并依赖于会产生额外训练FLOPs的辅助模块。类似地，OverEncoding【索引55，Over-tokenized transformer: Vocabulary is generally worth scaling，2025，ICML】即使在非等参数设置下，也未能在稀疏MoE主干上产生有意义的改进。相比之下，我们通过精心设计的Engram模块将条件记忆视为一流的建模原语。通过在我们的稀疏性分配框架内严格评估这一设计，我们展示了它相对于严格等参数和等FLOPs的MoE基线的明显优势。
  • 其次，从系统角度看，我们倡导算法-系统协同设计。现有方法严格地将嵌入放置在输入层（第0层）【索引55，Over-tokenized transformer: Vocabulary is generally worth scaling，2025，ICML；索引123，Scaling embedding layers in language models，2025，arXiv】，这内在地序列化了内存访问和计算。相反，Engram策略性地将记忆注入到更深的层中，以实现通信-计算重叠。此外，通过利用n-gram固有的Zipfian分布，我们可以最大化硬件内存层次结构的效用。这种整体设计使得Engram能够扩展到海量参数，而推理开销可忽略不计。

专家混合（MoE）。MoE架构通过为每个token有条件地激活一个稀疏的专家子集，将模型容量与计算成本解耦，这一范式由Shazeer等人【索引103，Outrageously large neural networks: The sparsely-gated mixture-of-experts layer，2017，arXiv】引入。随后的创新如GShard【索引74，Gshard: Scaling giant models with conditional computation and automatic sharding，2020，arXiv】、BASE【索引75，Base layers: Simplifying training of large, sparse models，2021，ICML】、Switch Transformer【索引34，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，2022，JMLR】和GLaM【索引27，Glam: Efficient scaling of language models with mixture-of-experts，2022，ICML】使得参数能够超线性扩展，同时保持恒定的推理成本。最近，DeepSeekMoE【索引22，Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models，2024，arXiv】展示了卓越的效率，通过细粒度的专家分割和共享专家隔离，显著优于具有等效活动参数的密集模型。采用这种架构，最先进的模型如DeepSeek-V3【索引83，Deepseek-v3 technical report，2024，arXiv】和Kimi-k2【索引111，Kimi linear: An expressive, efficient attention architecture，2025，arXiv】已将总参数进一步推向数千亿的规模。

记忆网络。关于记忆增强网络的研究旨在扩展模型容量而不成比例地增加计算成本，大致可分为参数化和非参数化方法。参数化记忆方法，如PKM【索引71，Large memory layers with product keys，2019，NIPS】、PEER【索引50，Mixture of a million experts，2024，arXiv】、Selfmem【索引15，Lift yourself up: Retrieval-augmented text generation with self-memory，2023，NeurIPS】、Memory+【索引6，Memory layers at scale，2025，ICML】和UltraMem【索引56，Ultramemv2: Memory networks scaling to 120b parameters with superior long-context learning，2025，arXiv；索引57，Ultra-sparse memory network，2025，ICLR】，将大规模、稀疏的键值存储直接整合到模型层中，从而在FLOPs影响可忽略的情况下显著增加容量。相反，非参数化记忆方法如REALM【索引48，Retrieval augmented language model pre-training，2020，ICML】、RETRO【索引10，Improving language models by retrieving from trillions of tokens，2022，ICML；索引115，Shall we pretrain autoregressive language models with retrieval? a comprehensive study，2023，EMNLP】和PlugLM【索引14，Decouple knowledge from paramters for plug-and-play language modeling，2023，ACL】将知识存储与模型处理解耦，将外部记忆视为一个可编辑和可扩展的键值存储，允许模型在不重新训练的情况下适应不断变化的信息。

知识存储机制。与容量扩展并行，大量研究审视了Transformer如何编码和检索事实知识的内部机制。前馈网络（FFNs）被广泛假设为键值记忆【索引43，Transformer feed-forward layers are key-value memories，2021，EMNLP】。在此框架下，第一层充当模式检测器（“键”），而第二层将特定信息投影到残差流中（“值”）。这种模块性通过识别负责存储不同事实的特定“知识神经元”【索引21，Knowledge neurons in pretrained transformers，2022，ACL】得到证明。因果追踪方法进一步提供了验证，它将事实回忆的信息流映射到特定的FFN层【索引87，Locating and editing factual associations in gpt，2022，NeurIPS】。这些见解使得精确的模型编辑算法如ROME【索引87，Locating and editing factual associations in gpt，2022，NeurIPS】和MEMIT【索引88，Mass-editing memory in a transformer，2023，ICLR】成为可能，它们允许直接更新事实关联而无需重新训练。此外，对内部表示的研究，如Othello-GPT【索引79，Emergent world representations: Exploring a sequence model trained on a synthetic task，2023，ICLR】中的研究，表明这些存储机制可能促进了结构化的“世界模型”的出现，而不仅仅是统计记忆。

A5 结论

在本文中，我们引入了条件记忆作为对当前主流的条件计算范式（MoE）的一个互补性稀疏轴，旨在解决通过动态计算来模拟知识检索的低效问题。我们通过Engram模块实例化了这一概念，该模块对经典的n-gram嵌入进行了现代化改造，以实现对静态模式的可扩展、常数时间O(1)的查找。

通过构建稀疏性分配问题，我们发现了一个U形缩放定律，证明了在MoE专家和Engram记忆之间混合分配稀疏容量，其性能严格优于纯MoE基线。遵循这一定律，我们将Engram扩展到27B参数，并在各种领域取得了卓越的性能。值得注意的是，虽然记忆模块直观上有助于知识检索，但我们观察到在通用推理、代码和数学方面取得了更大的收益。

我们的机理分析表明，Engram通过减轻早期层进行静态重构的任务，有效地“加深”了网络，从而释放了注意力容量以专注于全局上下文和复杂推理。这种架构上的转变在长上下文能力上带来了显著的提升，这在LongPPL和RULER上的性能增益得到了证明。最后，Engram倡导将基础设施感知的效率作为首要设计原则。其确定性寻址允许存储和计算的解耦，使得可以将海量参数表卸载到主机内存，而推理开销可忽略不计。我们预见，条件记忆将成为下一代稀疏模型不可或缺的建模原语。

A6 附录

A. 详细模型架构和超参数

Table 5 | 详细的模型架构信息和训练超参数。

B. 完整基准测试曲线


图 8 | 预训练最后10k步的基准测试曲线。

C. 分词器压缩案例研究

Table 6 | 该表展示了通过分词器压缩合并的前5个token，对于我们的128k分词器，总压缩率为23.43%。