16 KiB
Engram：条件记忆——面向大语言模型的可扩展查找稀疏轴
原文标题：Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models arXiv：2601.07372 ｜ 发布：12 Jan 2026（最后修订 12 Jul 2026, v2） 作者/机构：Xin Cheng、Rui Tian、Wangding Zeng、Damai Dai 等；北京大学（Peking University）与 DeepSeek-AI（深度求索）联合完成，张辉帅、赵东岩为通讯作者 原文 PDF：02-Engram-条件记忆-2601.07372.pdf
一、摘要（全文翻译）
混合专家（MoE）通过条件计算扩展容量，但 Transformer 缺少原生的知识查找原语，被迫以计算低效地模拟检索。为此我们提出**条件记忆（conditional memory）这一互补的稀疏轴，并以 Engram 实例化——该模块将经典 N-gram 嵌入现代化，实现 \mathcal{O}(1) 查找。通过形式化稀疏容量分配（Sparsity Allocation）**问题，我们揭示一条 U 形缩放定律，以优化神经计算（MoE）与静态记忆（Engram）的权衡。据此我们把 Engram 扩展至 27B 参数，在严格等参数、等 FLOPs 条件下超越 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 确立了"基础设施感知的效率"：确定性寻址支持运行时从主机内存预取，开销可忽略不计。我们设想条件记忆将成为下一代稀疏模型不可或缺的建模原语。代码见：https://github.com/deepseek-ai/Engram 。
二、核心贡献
- 提出「条件记忆」新稀疏轴：MoE 只解决了条件计算，而语言中大量命名实体与公式化短语本质局部、静态、高度模式化，本应由廉价查表承担。Engram 把 N-gram 嵌入重造为一等建模原语，以确定性哈希完成
\mathcal{O}(1)检索，避免 Transformer 用多层注意力与前馈网络在运行时"重建静态查找表"。 - 形式化稀疏容量分配并发现 U 形缩放定律：在总参数与 FLOPs 固定的前提下定义分配比例 $\rho$（非激活参数中分给 MoE 专家的比例），验证损失对
\rho呈 U 形；纯 MoE（$\rho=100%$）并非最优，把约 20%–25% 稀疏预算划给 Engram 最好，最优位置\rho\approx75\%\text{--}80\%跨算力规模稳定。 - 27B 规模超越同参同算力 MoE：Engram-27B 把路由专家由 72 削至 55，释放参数转为 5.7B 记忆表，总参数 26.7B、激活参数 3.8B 不变，知识、推理、阅读理解、代码数学全谱系提升。
- 证明增益源于「有效深度」增加：LogitLens 与 CKA 显示 Engram 浅层表示对齐 MoE 更深的层（第 5 层约等于基线第 12 层），早期层 KL 散度系统性更小，显式查找让网络更快完成特征组合。
- 长上下文与系统效率双重收益：局部依赖交给查表后注意力可专注全局，Multi-Query NIAH 由 84.2 升至 97.0；确定性寻址使 100B 参数表可卸载到主机内存异步预取，吞吐损失最高 2.8%。
三、方法与技术要点
3.1 总体结构：检索 + 融合两阶段
给定输入序列 X=(x_1,\dots,x_T) 与第 \ell 层隐状态 $\mathbf{H}^{(\ell)}\in\mathbb{R}^{T\times d}$，Engram 对每个位置 t 先检索（由局部上下文确定性取回静态嵌入），再融合（用当前隐状态调制并做轻量卷积）。模块以残差注入：$\mathbf{H}^{(\ell)}\leftarrow\mathbf{H}^{(\ell)}+\mathbf{Y}$，随后接标准 Attention 与 MoE；标准输入嵌入与 un-embedding 不变，模块只插入部分层。
3.2 分词器压缩与多头哈希检索
子词分词器优先保证无损重建，会把语义等价的词面（如 Apple 与 ␣apple）分到不同 ID，稀释语义密度。为此预计算满射函数 $\mathcal{P}:V\to V'$，按规范化文本等价（NFKC、小写化等）把原始 token ID 折叠为规范 ID，在 128k 词表上使有效词表缩小 23%（附录 C 为 23.43%）。
对位置 t 做映射后得到后缀 N-gram $g_{t,n}=(x'_{t-n+1},\dots,x'_t)$。直接参数化全部 N-gram 组合空间不可行，故用哈希方案并设 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}]
实践中 \varphi_{n,k} 用轻量乘法-XOR 哈希实现。记忆向量由各阶各头结果拼接而成：$\mathbf{e}t\triangleq|{n=2}^{N}|{k=1}^{K}\mathbf{e}{t,n,k}$，表 \mathbf{E}_{n,k} 取质数大小 M_{n,k} 以降低冲突。
3.3 上下文感知门控
检索到的 \mathbf{e}_t 是上下文无关的先验，静态性使其缺乏适应性，还可能被哈希冲突或一词多义污染。Engram 借用注意力范式：把已聚合全局上下文的 \mathbf{h}_t 作 Query，把 \mathbf{e}_t 同时投影为 Key 与 Value，
\mathbf{k}_t=\mathbf{W}_K\mathbf{e}_t,\quad \mathbf{v}_t=\mathbf{W}_V\mathbf{e}_t
为梯度稳定，对 Query 与 Key 先做 RMSNorm 再算标量门：
\alpha_t=\sigma\!\left(\frac{\text{RMSNorm}(\mathbf{h}_t)^{\top}\text{RMSNorm}(\mathbf{k}_t)}{\sqrt{d}}\right),\qquad \tilde{\mathbf{v}}_t=\alpha_t\cdot\mathbf{v}_t
这强制语义对齐：若记忆与上下文矛盾，门值趋零以抑制噪声。随后用深度可分离因果卷积扩大感受野（核宽 $w=4$、膨胀取最大 N-gram 阶数、SiLU 激活）：
\mathbf{Y}=\text{SiLU}\left(\text{Conv1D}\left(\text{RMSNorm}(\tilde{\mathbf{V}})\right)\right)+\tilde{\mathbf{V}}
3.4 与多分支架构的集成
默认骨干采用多分支结构（残差流扩展为 M 条并行分支、连接权重可学习，即 Manifold-Constrained Hyper-Connections，$M=4$）。为兼顾效率与表达力采用参数共享：稀疏嵌入表与 Value 投影 \mathbf{W}_V 全分支共享，每分支各有独立 Key 投影 ${\mathbf{W}K^{(m)}}{m=1}^{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),\quad \mathbf{u}_t^{(m)}=\alpha_t^{(m)}\cdot(\mathbf{W}_V\mathbf{e}_t)
这样 1 个 \mathbf{W}_V 与 M 个 \mathbf{W}_K^{(m)} 可融合成单次稠密 FP8 矩阵乘，提升 GPU 算力利用率。
3.5 计算与存储解耦的系统设计
MoE 路由依赖运行时隐状态，而 Engram 索引仅取决于输入 token ID，天然可预测。训练时按模型并行把大表分片到各 GPU，前向用 All-to-All 收集激活行、反向分发梯度，记忆容量随加速器数量线性增长。推理时把表卸载到主机内存，因索引在前向之前已知，主机可经 PCIe 异步预取，用前序块的计算掩盖传输延迟。因此放置层是硬件-算法协同设计：越深可掩盖延迟的窗口越长，但第 6.2 节消融显示建模性能偏好早期介入。N-gram 还服从 Zipf 分布，少量模式占绝大多数访问，可构建多级缓存层次（热条目放 HBM/DRAM，长尾放 NVMe SSD）。
3.6 稀疏容量分配定律与无限记忆
定义三个指标：总可训练参数 $P_{\mathrm{tot}}$（不含词表嵌入与 LM head）、每 token 激活参数 $P_{\mathrm{act}}$（决定 FLOPs）、非激活参数 $P_{\mathrm{sparse}}\triangleq P_{\mathrm{tot}}-P_{\mathrm{act}}$。分配比例：
P_{\mathrm{MoE}}^{(\mathrm{sparse})}=\rho\,P_{\mathrm{sparse}},\qquad P_{\mathrm{Engram}}=(1-\rho)\,P_{\mathrm{sparse}}
\rho=1 即纯 MoE；\rho<1 则减少路由专家、把释放参数转为 Engram 槽位。作者在 $C=2\times10^{20}$（$P_{\mathrm{tot}}\approx5.7$B、$P_{\mathrm{act}}=568$M、106 专家）与 $C=6\times10^{20}$（$P_{\mathrm{tot}}\approx9.9$B、$P_{\mathrm{act}}=993$M、99 专家）两档、稀疏度 P_{\mathrm{tot}}/P_{\mathrm{act}}\approx10 下扫描 $\rho$：C=6\times10^{20} 时验证损失从 \rho=100\% 的 1.7248 降到最优 \rho\approx80\% 的 1.7109（$\Delta=0.0139$）；即使 MoE 份额压到 $\rho\approx40%$（5.7B 模型 46 专家、9.9B 模型 43 专家）仍与纯 MoE 相当。U 形两端各有病因：MoE 独大时只能靠深度低效重建静态模式；Engram 独大时丧失条件计算能力，动态推理受损。无限记忆档固定 3B MoE 骨干（$P_{\mathrm{act}}=568$M、100B tokens），槽位数从 2.58\times10^5 扫到 $1.0\times10^7$（最多约 13B 参数），验证损失在 log 空间线性下降，同等预算下比把 N-gram 嵌入与词表嵌入平均的 OverEncoding 潜力更大。
3.7 大规模预训练配置
四个模型（Dense-4B、MoE-27B、Engram-27B、Engram-40B）共用数据课程与激活参数规模：262B tokens、DeepSeek-V3 分词器（128k 词表）、30 层 Transformer、隐层 2560、每块含 32 头 MLA、经 mHC（扩张率 4）连接 FFN、Muon 优化器。MoE-27B 为 2 共享 + 72 路由专家（top-6）；Engram-27B 改为 2 + 55（top-6，$\rho=74.3%$），省下的参数变成 5.7B 记忆，模块置于第 2、15 层，最大 N-gram 阶数 3、哈希头 8 个、维度 1280；嵌入参数用 Adam、学习率放大 5 倍、无权重衰减，卷积零初始化以保持恒等映射。Engram-40B 保持同一骨干与算力预算，仅把记忆扩到 18.5B（总 39.5B）。
3.8 消融与机制分析要点
3B MoE 骨干（12 层、0.56B 激活参数、100B tokens，基线验证损失 1.808）上加 1.6B 记忆，参考配置用 {2,3}-gram 置于第 2、6 层得 1.768。层扫描显示单模块放第 2 层最优（1.770）：一次注意力已足以提供有上下文的 Query，同时仍早到可替代底层局部聚合；同一预算拆成两个模块置于第 2、6 层更好，兼顾早期介入与后期门控精度。消融显示分支特异融合、上下文感知门控、分词器压缩贡献最大，去掉深度可分离卷积仅轻微退化，固定预算下容量分给 4-gram 略次优。推理期完全抑制记忆输出后，知识类仅保留 29%–44%（TriviaQA 29%），阅读理解类保留 81%–93%（C3 93%），印证记忆是参数化知识的主要仓库、上下文任务靠骨干注意力。门控激活集中在命名实体、套语及中文成语与历史人物（如"四大发明"、"张仲景"），跨语言一致。
四、主要实验结果
| 基准 / 任务 | 本模型（Engram-27B） | 对比对象（MoE-27B） | 说明 |
|---|---|---|---|
| Pile (loss) | 1.960 | 2.091 | 语言建模，越低越好 |
| 验证集 (loss) | 1.634 | 1.768 | 同分布验证集 |
| MMLU (Acc., 5-shot) | 60.4 | 57.4 | +3.0（摘要与引言口径 +3.4） |
| MMLU-Pro (Acc.) | 30.1 | 28.3 | +1.8 |
| CMMLU (Acc.) | 61.9 | 57.9 | +4.0 |
| C-Eval (Acc.) | 62.7 | 58.0 | +4.7 |
| AGIEval (Acc., 0-shot) | 41.8 | 38.6 | +3.2 |
| ARC-Challenge (Acc.) | 73.8 | 70.1 | +3.7 |
| BBH (EM, 3-shot) | 55.9 | 50.9 | +5.0 |
| CCPM (Acc., 0-shot) | 87.1 | 79.6 | +7.5，知识类增幅最大 |
| PIQA (Acc.) | 73.5 | 71.9 | +1.6 |
| TriviaQA (EM) | 50.7 | 48.8 | +1.9 |
| DROP (F1, 1-shot) | 59.0 | 55.7 | +3.3 |
| RACE-High (Acc.) | 78.2 | 75.4 | +2.8 |
| HumanEval (Pass@1) | 40.8 | 37.8 | +3.0 |
| MBPP (Pass@1) | 48.2 | 46.6 | +1.6 |
| GSM8K (EM, 8-shot) | 60.6 | 58.4 | +2.2 |
| MGSM (EM) | 49.4 | 46.8 | +2.6 |
| MATH (EM, 4-shot) | 30.7 | 28.3 | +2.4 |
| CruxEval-i (EM) | 32.2 | 30.7 | +1.5 |
| RULER 32k Multi-Query NIAH (Acc.) | 97.0 | 84.2 | 等预训练损失设置（46k vs 50k） |
| RULER 32k Variable Tracking (Acc.) | 87.2 | 77.0 | 同上，等损失设置 |
| RULER 32k Multi-Values (Acc.) | 95.5 | 92.7 | 同上 |
| LongPPL Code (PPL ↓) | 2.45 | 2.49 | 等损失设置，略优 |
| LongPPL L-CoT (PPL ↓) | 13.59 | 14.16 | 等损失设置 |
| 吞吐 4B-Dense + 100B Engram（CPU 卸载） | 8,858.28 tok/s | 8,903.62 tok/s（基线） | H800，512 序列，损失 0.5% |
| 吞吐 8B-Dense + 100B Engram | 6,140.02 tok/s | 6,315.52 tok/s（基线） | 损失 2.8%，为最差情形 |
三种稀疏架构在等训练算力下全面优于 iso-FLOPs 的 Dense-4B，说明稀疏本身的缩放优势；Engram-27B 在同参同算力下对 MoE-27B 的领先不只出现在知识密集任务（MMLU +3.0、CMMLU +4.0、MMLU-Pro +1.8），在通用推理（BBH +5.0、ARC-Challenge +3.7、DROP +3.3）与代码数学（HumanEval +3.0、MATH +2.4、GSM8K +2.2）上增益更大。Engram-40B 进一步降低损失并在多数基准上更好，但未在每个任务上严格压过 Engram-27B，作者归因于欠训练——两者差距在训练末期仍在扩大。长上下文中，等预训练损失（46k vs 50k）下 Engram-27B 在复杂检索上大幅领先，等算力（50k）下全面提升；即使用仅约 82% 预训练算力的 41k 检查点，也能在 LongPPL 上追平、在 RULER 上超越完全训练的 MoE-27B。系统中，100B 记忆表全放主机 DRAM、每次检索都走 PCIe 的保守设定下吞吐损失仅 2.8%。
五、结论、局限与影响
- 结论：论文把条件记忆确立为与 MoE 条件计算互补的稀疏轴。Engram 以分词器压缩、多头哈希检索、上下文感知门控与多分支集成，把经典 N-gram 嵌入现代化为
\mathcal{O}(1)的可扩展查找原语。由此形式化的稀疏容量分配给出稳定的 U 形定律：纯 MoE 次优，把 20%–25% 稀疏预算分给记忆最优。27B 规模上 Engram 在同参同算力下全面超越 MoE 基线；机制分析表明它把静态重建从浅层卸载从而"加深"网络，并使注意力得以专注全局上下文；确定性寻址让海量参数表可卸载到主机内存而几乎不损吞吐。 - 局限性：论文未设独立的"局限"章节，作者自述的相关限制包括：（1）插入位置受建模效果与系统延迟双重约束，早期注入最优但牺牲门控精度；（2）固定 1.6B 预算下 4-gram 略次优，更大记忆规模下是否有利未做结论；（3）第 6.3 节推理期抑制消融造成训练-推理不一致，故只分析信噪比最高的知识类与阅读理解类；（4）Engram-40B 未在每个任务上严格占优，作者判断为欠训练；（5）系统实验用稠密骨干的 nano-vLLM 原型与全量 CPU 卸载，属保守基线。其余方面（如超大规模下哈希冲突率、多语言覆盖面）论文未明确说明。
- 影响：该工作把"算法-系统协同设计"提升为设计原则：以往 N-gram 嵌入方案都在第 0 层、使访存与计算串行化，中层注入才能实现通信-计算重叠，并借助 Zipf 分布利用多级存储层次；这为绕过 GPU 显存上限、激进扩张参数规模提供了可落地路径。
六、关键术语对照表
| English | 中文 |
|---|---|
| Conditional Memory | 条件记忆 |
| Conditional Computation | 条件计算 |
| Engram | Engram（记忆印迹模块，论文未给中文译名） |
| Sparsity Allocation | 稀疏容量分配 |
| Mixture-of-Experts (MoE) | 混合专家 |
| Iso-parameter / Iso-FLOPs | 等参数 / 等算力（同 FLOPs） |
| N-gram Embedding | N-gram 嵌入 |
| Hashed N-gram | 哈希 N-gram |
| Multi-Head Hashing | 多头哈希 |
| Tokenizer Compression | 分词器压缩 |
| Context-aware Gating | 上下文感知门控 |
| Multi-branch Architecture | 多分支架构 |
| Manifold-Constrained Hyper-Connections (mHC) | 流形约束超连接 |
| Multi-head Latent Attention (MLA) | 多头潜在注意力 |
| Effective Depth | 有效深度 |
| LogitLens | LogitLens（逐层逻辑回归探针） |
| Centered Kernel Alignment (CKA) | 中心核对齐 |
| Zipfian Distribution | 齐夫分布 |
| Multi-Level Cache Hierarchy | 多级缓存层次 |
| Prefetch-and-Overlap | 预取与重叠 |
| Needle-in-a-Haystack (NIAH) | 大海捞针检索 |
| LongPPL / RULER | 长上下文困惑度 / RULER 长上下文基准 |