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

A1 主要贡献

本文探讨了大型语言模型（LLMs）中一个结构性优化的机会。当前模型主要通过专家混合（MoE）实现稀疏性，但这迫使它们通过计算来低效地模拟知识检索。语言建模包含两种性质不同的子任务：组合推理和知识检索。前者需要深度动态计算，而后者涉及大量局部、静态和高度模式化的文本（如命名实体、固定短语），这些内容更适合通过计算开销低的查找操作来处理。由于标准Transformer缺乏原生的知识查找机制，当前LLMs被迫消耗宝贵的计算深度来重建静态查找表，而这些计算资源本可用于更高级的推理。

为了解决这一问题，本文提出了一个与条件计算互补的稀疏性轴：条件记忆（conditional memory）。条件计算通过稀疏激活参数来处理动态逻辑，而条件记忆则依赖稀疏查找操作来检索固定知识的静态嵌入。本文通过Engram模块实例化了这一概念，这是一个基于经典n-gram嵌入结构并结合了现代技术（如分词器压缩、多头哈希、上下文感知门控和多分支集成）的模块。

核心贡献如下：
1. 提出条件记忆作为新的稀疏性轴：与MoE的条件计算相辅相成，通过Engram模块实现高效的O(1)知识查找，以处理语言中的静态和局部模式。
2. 构建稀疏性分配问题并发现U型缩放定律：研究了在固定总参数和计算预算下，如何在MoE专家（计算）和Engram记忆之间分配容量。实验揭示了一个明确的U型关系，表明混合分配策略优于纯MoE模型。
3. 大规模实验验证有效性：遵循该分配定律，将Engram扩展至一个27B参数的模型。在与参数量和计算量严格相等的MoE基线模型的对比中，Engram-27B在多个领域均表现出优越性。
* 知识密集型任务：MMLU提升3.4分，CMMLU提升4.0分。
* 通用推理任务：BBH提升5.0分，ARC-Challenge提升3.7分。
* 代码/数学领域：HumanEval提升3.0分，MATH提升2.4分。

  1. 机制分析揭示其工作原理：通过LogitLens和CKA分析发现，Engram通过减轻模型早期层重建静态知识的负担，有效地增加了网络的“有效深度”，从而为复杂推理释放了更多资源。此外，它将局部依赖关系委托给查找操作，释放了注意力容量以关注全局上下文，显著提升了长上下文检索能力（例如，在MultiQuery NIAH上从84.2提升到97.0）。
  2. 实现基础设施感知的系统效率：Engram采用确定性寻址，允许在运行时从主机内存预取数据，实现了计算与通信的重叠。实验证明，将一个100B参数的查找表卸载到主机内存，其开销可以忽略不计（<3%），有效绕过了GPU内存的限制。


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

A2 方法细节

2.1. 概述

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

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

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

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

多头哈希。直接参数化所有可能的n-gram的组合空间是不可行的。遵循Tito Svenstrup等人的工作【索引109, Hash embeddings for efficient word representations+2017+NeurIPS】，我们采用了一种基于哈希的方法。为了减轻冲突，我们为每个n-gram阶数 $n$ 采用了 $K$ 个不同的哈希头。每个头 $k$ 通过一个确定性函数 $\phi_{n,k}$ 将压缩后的上下文映射到嵌入表 $E_{n,k}$（其大小为素数 $S_{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}$ 是一个轻量级的乘法-异或（multiplicative-XOR）哈希。我们通过拼接所有检索到的嵌入来构建最终的记忆向量 $e_t \in R^{d_{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】【索引110, Attention is all you need+2017+NeurIPS】。具体来说，我们将当前的隐藏状态 $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【索引120, Root mean square layer normalization+2019+NeurIPS】：

$$\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】【索引92, RWKV: reinventing rnns for the transformer era+2023+EMNLP】。令 $\tilde{V} \in R^{L \times d}$ 表示门控值的序列。使用核大小 $K$（设为4）、膨胀率 $\delta$（设为最大n-gram阶数）和SiLU激活函数【索引30, Sigmoid-weighted linear units for neural network function approximation in reinforcement learning+2018+Neural networks】，最终输出 $Y$ 的计算如下：

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

Engram模块通过残差连接集成到主干网络中：$H^{(\ell)} \leftarrow H^{(\ell)} + 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】，而是采用了先进的多分支架构作为我们的默认主干网络，因为它具有更强的建模能力【索引67, Fractalnet: Ultra-deep neural networks without residuals+2017+ICLR】【索引104, Going deeper with convolutions+2015+CVPR】【索引116, mhc: Manifold-constrained hyper-connections+2025+URL】【索引122, Hyperconnections+2025+ICLR】。该架构的一个决定性特征是将残差流扩展为 $M$ 个平行分支，信息流由可学习的连接权重调节。

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

$$\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_{t}^{(m)} = \alpha_{t}^{(m)} \cdot (W_V e_t)$。这种设计允许线性投影（一个 $W_V$ 和 $M$ 个不同的 $W_{K}^{(m)}$）融合成一个单一的密集FP8矩阵乘法，从而最大化现代GPU的计算利用率。除非另有说明，所有实验都采用这种与Manifold-Constrained Hyper-Connections ($M = 4$)【索引116, mhc: Manifold-constrained hyper-connections+2025+URL】的集成方式。


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

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

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

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

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

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

A3 关键Observation与设计原则

Engram作为条件记忆的一种实现，与MoE专家提供的条件计算在结构上是互补的。本节研究了这种二元性的缩放特性，以及如何最优地分配稀疏容量。具体而言，我们的研究由两个关键问题驱动：

  1. 有限约束下的分配。当总参数和训练计算量固定时（参数量和计算量均相等），我们应该如何在MoE专家和Engram嵌入之间分配稀疏容量？
  2. 无限内存情境。考虑到Engram的非缩放性O(1)开销，如果内存预算被放宽或积极扩展，Engram本身会表现出什么样的缩放行为？

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

计算匹配的公式化。我们使用三个参数指标来分析这种权衡：
* $P_{tot}$：总可训练参数，不包括词汇嵌入和LM头。
* $P_{act}$：每个token激活的参数。该数量决定了训练成本（FLOPs）。
* $P_{sparse} \triangleq P_{tot} - P_{act}$：非激活参数，代表了可用于扩展模型大小而不增加计算成本的“免费”参数预算（例如，未被选择的专家或未被检索的嵌入）。
我们在每个FLOPs预算内保持 $P_{tot}$ 和 $P_{act}$ 固定，因此模型具有相同数量的参数和相同的每token FLOPs。对于MoE，$P_{act}$ 由top-k选择的专家决定，而未选择的专家的参数贡献给 $P_{sparse}$。对于Engram，每个token只检索恒定数量的槽位，因此扩展嵌入槽的数量会增加 $P_{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_{tot}/P_{act} \approx 10$：
* $C = 2 \times 10^{20}$ FLOPs: $P_{tot} \approx 5.7B$，$P_{act} = 568M$。基线模型（$\rho = 1$）共有106个专家。
* $C = 6 \times 10^{20}$ FLOPs: $P_{tot} \approx 9.9B$，$P_{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 \rightarrow 100\%$）：模型缺乏用于存储静态模式的专用内存，迫使其通过深度和计算来低效地重建它们。
* Engram主导（$\rho \rightarrow 0\%$）：模型失去了条件计算能力，损害了需要动态、上下文依赖推理的任务；在这种情况下，内存无法替代计算。

3.2. 无限内存情境下的Engram

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

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

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

A4 实验环境与结果

实验环境

实验结果

大规模预训练性能（表1）

表 1 | 密集、MoE和Engram模型的预训练性能比较。所有模型都用262B个token进行训练，并且激活参数（3.8B）匹配。Engram-27B与MoE-27B的参数量相同，通过将路由专家数量从72个减少到55个，并将参数重新分配给一个5.7B参数的Engram内存。Engram-40B在保持激活参数预算固定的情况下，进一步将Engram内存增加到18.5B。完整的训练期间基准测试轨迹见附录B。

长上下文训练性能（表2）

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

A7 补充细节

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

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】【索引75, Echoes of bert: Do modern language models rediscover the classical nlp pipeline?+2025+URL】，而这个过程理论上可以通过一次知识查找操作来识别。

基于此，我们假设通过为模型配备明确的知识查找能力，Engram通过减轻模型早期阶段的特征组合负担，有效地模拟了模型深度的增加。为了验证这一假设，我们采用了两种机理可解释性工具：LogitLens【索引3, Eliciting latent predictions from transformers with the tuned lens+2023+arXiv】【索引89, interpreting gpt: the logit lens+2020+LessWrong】和中心核对齐分析（CKA）【索引23, Reliability of CKA as a similarity measure in deep learning+2023+ICLR】【索引62, Similarity of neural network representations revisited+2019+ICML】。

6.1.1. 加速的预测收敛

加速的预测收敛。我们首先使用LogitLens【索引89, 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减少了所需的计算步骤，从而在网络层次结构的更早阶段就达到了高置信度的有效预测。

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

6.1.2. 表示对齐与有效深度

表示对齐与有效深度。为了进一步研究Engram的层是否在语义上对应于基线模型的更深层，我们采用了中心核对齐（CKA），这是一个广泛用于比较表示结构的度量标准【索引62, 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定义为：

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

其中 $K = XX^\top$ 和 $L = YY^\top$ 表示Gram矩阵（使用线性核），HSIC是希尔伯特-施密特独立性准则【索引44, 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}}, \quad \text{where } \mathcal{I}_j = \mathop{\mathrm{argtopk}}_i (S_{i,j}).$$

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

结果分析。图4（b）-（c）展示了相似度热图，并叠加了软对齐曲线（白色虚线）。我们观察到与对角线相比有明显的向上偏移，这意味着对于很大范围的层，$a_j > j$。例如，Engram-27B的第5层形成的表示与MoE基线的大约第12层的表示最为接近。

结论。这种与对角线一致的偏移，与LogitLens的结果（6.1.1节）相符，证实了Engram在更早的层级实现了更深的表示。这验证了我们的核心假设：通过显式查找绕过早期的特征组合，Engram在功能上等同于增加了模型的有效深度。


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

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基线有显著提升（$\Delta = 0.04$）。以下所有结构消融都是相对于这个参考配置定义的。

内存应该注入到哪里？ 为了研究深度敏感性，我们保持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】【索引75, Echoes of bert: Do modern language models rediscover the classical nlp pipeline?+2025+URL】【索引108, Bert rediscovers the classical nlp pipeline+2019+ACL】。然而，这会带来门控精度的代价：早期的隐藏状态尚未通过注意力聚合足够的全局上下文，并且并行分支缺乏进行细粒度调制所需的表示差异【索引116, mhc: Manifold-constrained hyper-connections+2025+URL】【索引122, Hyperconnections+2025+ICLR】。因此，最佳放置位置需要在（i）尽早卸载静态局部模式和（ii）利用更强的上下文查询进行后期门控之间取得平衡。

实验结果。扫描结果显示，第2层达到了最佳的单层性能（验证损失=1.770），优于第1层，并且随着插入点向深层移动性能下降。这表明一轮注意力已经足以提供一个有意义的上下文化$h_t$来进行门控，同时仍然足够早，可以替代主干网络底层的局部聚合。

分层注入。虽然在单一注入约束下第2层是最佳的，但我们发现，将相同的1.6B内存分成两个较小的模块（通过减少嵌入维度$d_{mem}$实现）并分别放置在第2层和第6层，性能甚至更好（验证损失=1.768）。这种分层设计通过结合早期干预和丰富的后期上下文门控来协调这种权衡。更重要的是，分层插入还提供了一个实际的系统优势，能够更好地利用内存层次结构，如2.5节所讨论的。

哪些组件重要？ 从参考配置出发，我们在保持Engram参数预算固定的情况下，对单个设计选择进行消融。结果在图5中用标记表示。我们发现三个组件带来了最显著的增益：（i）多分支主干网络内的分支特定融合，（ii）上下文感知门控，以及（iii）分词器压缩。移除这些组件中的任何一个都会导致验证损失的最大回归。具体来说，对于“w/o multi branch”消融，我们保留了mHC主干结构，但用一个应用于预映射后隐藏状态$H_{premap}$【索引116, mhc: Manifold-constrained hyper-connections+2025+URL】的单一Engram融合替换了分支特定的门控。

其他组件的影响。其他改变的影响较小：移除轻量级深度卷积只会轻微降低性能。在固定的1.6B预算下，将容量分配给4-grams是次优的——这可能是因为它稀释了更频繁的2/3-gram模式的容量——尽管我们不排除在更大的内存规模下，更高阶的n-grams会变得有益。


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

6.3. 敏感性分析

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

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


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

6.4. 系统效率

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

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

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

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

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

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

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

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


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

A5 结论

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

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

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

A6 附录

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

详细模型架构和超参数。下表详细列出了Dense-4B, MoE-27B, Engram-27B, 和 Engram-40B四个模型的具体架构信息和训练超参数。


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

B. 完整基准测试曲线

完整基准测试曲线。下图展示了MoE-27B和Engram-27B在预训练最后10,000步期间，在多个基准测试上的完整性能变化曲线，直观地反映了模型的训练动态和最终性能。


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

C. 分词器压缩案例研究

分词器压缩案例研究。下表展示了通过分词器压缩机制合并次数最多的前5个规范化token及其对应的原始token。对于我们使用的128k词汇量的分词器，整体压缩率达到了23.43%。


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