通过可扩展查找的条件记忆:大语言模型稀疏性的新轴

Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models
整理日期:2026‑01‑13 原文:arXiv:2601.07372v1 主题:MoE / 条件记忆 / Engram

作者:Xin Cheng, Wangding Zeng, Damai Dai, Qinyu Chen, Bingxuan Wang, Zhenda Xie, Kezhao Huang, Xingkai Yu, Zhewen Hao, Yukun Li, Han Zhang, Huishuai Zhang, Dongyan Zhao, Wenfeng Liang

机构:Peking University;DeepSeek‑AI

说明:本文为基于作者公开 TeX 源码的中文翻译与排版整理;公式使用 MathJax 渲染,图片已本地化存放于本目录的 images/

本文翻译与整理自 arXiv 论文《Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models》。 查看英文原文 →

摘要

混合专家(Mixture‑of‑Experts, MoE)通过条件计算扩展模型容量;但标准 Transformer 并没有“知识查找(lookup)”这一原生算子,只能用计算去低效地模拟检索。为此,作者提出将条件记忆(conditional memory)作为与条件计算互补的另一条稀疏轴,并以 Engram 作为具体实例:在经典 N‑gram 嵌入的基础上引入现代化设计,使得静态模式可以通过 \( \mathcal{O}(1) \) 的确定性查找获得。

进一步地,作者提出稀疏分配(Sparsity Allocation)问题:在总参数预算固定的前提下,应如何在 MoE 专家容量与 Engram 记忆容量之间分配“非激活参数”(即不会带来每 token 计算成本的那部分容量)?实验发现了一条清晰的U 形标度律,指出在两者之间进行混合分配会严格优于纯 MoE。

在该分配规律指导下,作者将 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 的确定性寻址使得推理时可从主存进行运行时预取,并与计算重叠,开销极小。作者认为,条件记忆将成为下一代稀疏模型不可或缺的建模原语。

1 引言

稀疏性是智能系统中反复出现的设计原则,从生物神经回路到现代大语言模型(LLM)皆然(原文引用:olshausen1997sparse、lennie2003cost)。当前,这一原则主要通过 MoE 得以实现:利用条件计算在不成比例增加计算量的前提下大幅扩展模型规模,因此也成为前沿模型的主流范式(原文引用:shazeer2017outrageously、dai2024deepseekmoe、guo2025deepseek、comanici2025gemini、team2025kimi)。

然而,语言信号的内在异质性意味着架构仍有显著的结构化优化空间:语言建模包含两类性质迥异的子任务——组合式推理与知识检索。前者需要深、动态的计算;后者中相当一部分文本(如命名实体、固定搭配)则是局部、静态、刻板的(原文引用:erman2000idiom、constant2017survey)。经典 N‑gram 模型在捕获这类局部依赖上非常有效(原文引用:liu2024infini、nguyen2024understanding、brants-etal-2007-large),提示这些规律天然适合用廉价的查找来表达。但标准 Transformer 缺少原生的知识查找原语,导致现有 LLM 被迫用计算去模拟检索:例如解析一个常见的多 token 实体需要消耗多个早期层的注意力与前馈网络(原文引用:ghandeharioun2024patchscopes、DBLP:conf/coling/JinYHZWH0MMDYDZ25;另见本文表格“PatchScope 案例”),这相当于在运行时昂贵地重建一个静态查找表,浪费了可以用于更高阶推理的序列深度。

为契合这种“推理‑检索”二元性,作者提出一条互补稀疏轴:条件记忆。条件计算稀疏激活参数以处理动态逻辑(原文引用:shazeer2017outrageously、bengio2013estimatingpropagatinggradientsstochastic);条件记忆则通过稀疏查找操作为固定知识检索静态嵌入。作为初步探索,作者回到 N‑gram 嵌入这一典型实现:以局部上下文为键,通过 \( \mathcal{O}(1) \) 的哈希查找索引一个巨大的嵌入表(原文引用:tito2017hash、huang2025over、pagnoni2025byte、yu2025scaling)。作者发现:静态检索机制可以与现代 MoE 架构形成理想互补,但前提是必须进行合适的现代化设计。因此提出 Engram:在经典 N‑gram 结构之上,引入分词器压缩、多头哈希、上下文门控、多分支集成等关键改造(详见第 2 节)。

为量化两类原语的协同作用,作者形式化提出稀疏分配问题,并在实验中观测到稳定的 U 形标度律:在固定稀疏预算下,将约 20%–25% 的“非激活参数”分配给 Engram 可得到最优性能。基于该规律,作者将 Engram 扩展到 27B 参数,并在等参数、等 FLOPs 条件下在多个任务域取得提升,且推理、代码、数学等域的增益更大。进一步的机理分析与系统实验表明,Engram 既能提升有效深度,也能通过确定性寻址实现高效的运行时预取与 CPU Offload。

打开图 1(PDF)
图 1:Engram 架构。模块通过检索静态 N‑gram 记忆并用上下文感知门控与动态隐藏状态融合,从而增强骨干网络。该模块只插入到特定层以将记忆与计算解耦,同时保持标准输入嵌入与输出反嵌入模块不变。

2 架构

2.1 总览

如图 1 所示,Engram 是一个条件记忆模块,其核心目标是在结构上把静态模式的存储动态计算分离,从而增强 Transformer 骨干网络。形式化地,给定输入序列 \(X=(x_1,\dots,x_T)\) 与第 \(\ell\) 层隐藏状态 \(\mathbf{H}^{(\ell)}\in\mathbb{R}^{T\times d}\),Engram 在每个位置 \(t\) 上按两阶段工作:检索(retrieval)融合(fusion)

第一阶段(第 2.2 节)从局部上下文提取并压缩后缀 N‑gram,通过确定性哈希检索静态嵌入向量;第二阶段(第 2.3 节)利用当前隐藏状态作为上下文 Query,对检索到的嵌入进行门控调制,并通过轻量卷积进一步精炼。随后,作者讨论与多分支骨干的集成(第 2.4 节)以及系统层面设计(第 2.5 节)。

2.2 稀疏检索:哈希 N‑gram

检索阶段将局部上下文映射到静态记忆条目,主要包含分词器压缩确定性哈希检索两步。

分词器压缩(Tokenizer Compression)

传统 N‑gram 往往直接在分词器输出上建模,但标准子词分词器以“无损重建”为优先,常将语义等价的词形分配到互不相干的 ID(例如 Apple vs. ␠apple;原文引用:kudo2018sentencepiecesimplelanguageindependent、li2023starcodersourceyou)。为提高语义密度,Engram 引入词表投影层:预先计算一个满射函数 \( \mathcal{P}:V\to V' \),把原始 token ID 按文本归一化等价类(NFKC、转小写等;原文引用:UAX15-NFKC)折叠成“规范 ID”。在 128k 分词器上,该过程将有效词表规模降低约 23%(见附录 C)。

形式化地,对位置 \(t\) 的 token,映射 \(x_t \mapsto x_t'=\mathcal{P}(x_t)\),并构造后缀 N‑gram \(g_{t,n}=(x'_{t-n+1},\dots,x'_t)\)。

多头哈希(Multi‑Head Hashing)

直接为所有可能的 N‑gram 组合空间参数化在规模上不可行。沿用哈希 N‑gram 的思路(原文引用:tito2017hash),Engram 对每个 N‑gram 阶 \(n\) 采用 \(K\) 个互不相同的哈希头以缓解碰撞。对第 \(k\) 个头,通过确定性函数 \(\phi_{n,k}\) 把压缩后的上下文映射到嵌入表 \(\mathbf{E}_{n,k}\)(素数大小 \(M_{n,k}\))中的索引:

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

实践中,\(\phi_{n,k}\) 可实现为轻量的“乘法‑异或”哈希。最终的记忆向量 \(\mathbf{e}_t\in\mathbb{R}^{d_{\text{mem}}}\) 由所有检索到的嵌入拼接得到:

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

2.3 上下文感知门控(Context‑aware Gating)

检索到的嵌入 \(\mathbf{e}_t\) 可视为与上下文无关的先验;但由于其静态性质,缺乏上下文适配能力,也可能受到哈希碰撞或多义性带来的噪声影响(原文引用:haber-poesio-2024-polysemy)。为提高表达力并消解歧义,Engram 引入受注意力机制启发的上下文门控(原文引用:vaswani2017attention、bahdanau2014neural)。核心思想是:用当前隐藏状态 \(\mathbf{h}_t\)(已通过前面的注意力层聚合全局上下文)作为动态 Query,而检索到的记忆 \(\mathbf{e}_t\) 作为 Key/Value 的来源:

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

其中 \(\mathbf{W}_K,\mathbf{W}_V\) 为可学习投影矩阵。为保证梯度稳定(原文引用:pmlr-v202-dehghani23a),在计算标量门控 \(\alpha_t\in(0,1)\) 前对 Query 与 Key 施加 RMSNorm(原文引用:zhang2019root):

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

门控后的输出为 \(\tilde{\mathbf{v}}_t=\alpha_t\cdot\mathbf{v}_t\)。该设计强制语义一致性:当检索到的记忆与当前上下文矛盾时,\(\alpha_t\) 会趋近于 0,从而有效抑制噪声。

最后,为扩大感受野并增强非线性,作者在门控值序列 \(\tilde{\mathbf{V}}\in\mathbb{R}^{T\times d}\) 上引入短的 depthwise 因果卷积(原文引用:gu2021efficiently、peng2023rwkv)。设核大小 \(w=4\)、空洞率 \(\delta\)(设为最大 N‑gram 阶数),并用 SiLU 激活(原文引用:elfwing2018sigmoid),最终输出为:

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

Engram 通过残差方式集成到骨干网络:\(\mathbf{H}^{(\ell)}\leftarrow \mathbf{H}^{(\ell)}+\mathbf{Y}\),随后再执行标准的注意力与 MoE。关键点在于:Engram 并不插入每一层,其放置位置需同时满足建模收益与系统时延约束(见第 2.5 节)。

打开图 2(PDF)
图 2:Engram 的系统实现。(a) 训练阶段:大规模嵌入表跨多张 GPU 分片,通过 All‑to‑All 通信检索被激活的行;(b) 推理阶段:Engram 表被卸载到主存,利用确定性检索逻辑在主机侧异步预取并传输嵌入,将通信与前序 Transformer 块的计算重叠。

2.4 与多分支骨干的集成

本文默认使用先进的多分支架构作为骨干网络,而非标准的单流残差连接(原文引用:he2016deep、larsson2016fractalnet、zhu2025hyperconnections、xie2025mhcmanifoldconstrainedhyperconnections、szegedy2015going)。该架构的特征是把残差流扩展为 \(M\) 条并行分支,信息流由可学习连接权重调制。

Engram 本身与拓扑无关,但要适配多分支框架需要在效率与表达力之间做结构优化。作者采用参数共享策略:在所有 \(M\) 个分支之间共享一个稀疏嵌入表与 Value 投影矩阵 \(\mathbf{W}_V\),同时为每个分支保留独立的 Key 投影矩阵 \(\{\mathbf{W}_K^{(m)}\}_{m=1}^M\),以支持分支特异的门控行为。对第 \(m\) 个分支的隐藏状态 \(\mathbf{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). $$

并将共享的 value 向量用分支门控调制:\(\mathbf{u}_t^{(m)}=\alpha_t^{(m)}\cdot(\mathbf{W}_V\mathbf{e}_t)\)。该设计使得线性投影(一个 \(\mathbf{W}_V\) 与 \(M\) 个 \(\mathbf{W}_K^{(m)}\))可融合为一次稠密 FP8 矩阵乘,以提升 GPU 利用率。除非特别说明,所有实验都使用与 Manifold‑Constrained Hyper‑Connections(\(M=4\);原文引用:xie2025mhcmanifoldconstrainedhyperconnections)结合的该集成方式。

2.5 系统效率:解耦计算与记忆

记忆增强模型的扩展常被 GPU 高带宽显存(HBM)容量所限制。但 Engram 的确定性检索机制天然支持将参数存储与计算资源解耦:与依赖运行时隐藏状态做动态路由的 MoE 不同,Engram 的检索索引只依赖输入 token 序列,因而在前向执行前即可计算出来。这种可预测性使得训练与推理都可以进行针对性的系统优化(见图 2)。

训练阶段。为承载大规模嵌入表,作者采用标准模型并行:将嵌入表跨多张 GPU 分片,并用 All‑to‑All 在前向收集被激活的行、在反向分发梯度,使总记忆容量随加速器数量线性增长。

推理阶段。确定性寻址使得系统能够“预取并重叠”:在前向执行前已知记忆索引,可从充足的主存(DRAM)经 PCIe 异步拉取嵌入,并用前序 Transformer 块的计算作为缓冲窗口以避免 GPU 停顿。这要求硬件‑算法协同:把 Engram 放得更深可增加掩蔽通信延迟的计算窗口,但作者在第 6.2 节的消融实验中发现建模性能更偏好更早的注入以卸载局部模式重构。因此,最优放置需同时满足建模与系统时延约束。

此外,自然语言 N‑gram 近似服从 Zipf 分布(原文引用:piantadosi2014zipf、Chao1950HumanBA),即少量高频模式占据大多数访问。这一统计特性支持多级缓存层次:高频嵌入驻留在更快的存储层(HBM/DRAM),长尾低频驻留在更慢但容量更大的介质(如 NVMe SSD),从而在几乎不影响有效时延的情况下把 Engram 扩展到极大容量。

打开图 3(PDF)
图 3:稀疏分配与 Engram 标度律。左:不同分配比例 \(\rho\) 下的验证损失,两种计算预算(\(2\times10^{20}\) 与 \(6\times10^{20}\) FLOPs)均呈 U 形,混合分配优于纯 MoE。右:在“无限记忆”区间,验证损失随嵌入槽位数量呈对数线性(log‑linear)趋势。

3 标度律与稀疏分配

作为条件记忆的实例,Engram 在结构上与 MoE 的条件计算互补。本节研究这种双轴稀疏性的标度性质,并探讨在固定预算下如何最优分配稀疏容量。核心问题有二:

  1. 有限约束下的分配。当总参数与训练计算(等参数、等 FLOPs)固定时,应如何在 MoE 专家与 Engram 嵌入之间分配稀疏容量?
  2. 无限记忆区间。考虑 Engram 的 \( \mathcal{O}(1) \) 非扩展性开销,如果放宽或激进扩大记忆预算,Engram 自身呈现怎样的标度规律?

3.1 MoE 与 Engram 的最优分配比例

等计算(Compute‑matched)表述

作者用三类参数指标刻画权衡:

  • \(P_{\mathrm{tot}}\):总可训练参数(不含词表嵌入与 LM Head)。
  • \(P_{\mathrm{act}}\):每 token 被激活参数量,决定训练成本(FLOPs)。
  • \(P_{\mathrm{sparse}}\triangleq P_{\mathrm{tot}}-P_{\mathrm{act}}\):非激活参数,可视为“不增加计算成本的自由容量”(例如未被选中的专家,或未被检索到的嵌入槽位)。

在每个 FLOPs 预算内保持 \(P_{\mathrm{tot}}\) 与 \(P_{\mathrm{act}}\) 固定,从而保证模型参数总量与每 token FLOPs 相同。对 MoE,\(P_{\mathrm{act}}\) 由 top‑\(k\) 专家激活决定,未选专家计入 \(P_{\mathrm{sparse}}\);对 Engram,每 token 仅检索常数个槽位,因此扩展槽位数量会增加 \(P_{\mathrm{tot}}\) 但不会增加每 token FLOPs。

分配比例 \(\rho\)

定义分配比例 \(\rho\in[0,1]\) 为把“非激活参数预算”分配给 MoE 专家容量的比例:

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

直观地:\(\rho=1\) 对应纯 MoE(所有非激活参数都用于路由专家);\(\rho<1\) 则减少专家数量,把释放出的参数重分配给 Engram 槽位。

实验协议

作者在两种计算区间下评估权衡,并在两种设置中保持固定稀疏比 \(P_{\mathrm{tot}}/P_{\mathrm{act}}\approx 10\):

  • \(C=2\times10^{20}\) FLOPs:\(P_{\mathrm{tot}}\approx5.7\)B、\(P_{\mathrm{act}}=568\)M。基线(\(\rho=1\))共有 106 个专家。
  • \(C=6\times10^{20}\) FLOPs:\(P_{\mathrm{tot}}\approx9.9\)B、\(P_{\mathrm{act}}=993\)M。基线(\(\rho=1\))共有 99 个专家。

在不同 \(\rho\) 下,仅通过调整路由专家数量与 Engram 槽位数量来构造模型;其余训练流水线与优化超参保持一致。

结果与分析

图 3(左)显示验证损失与分配比例 \(\rho\) 呈稳定的 U 形关系。令人意外的是,当 MoE 分配降到约 \(\rho\approx40\%\) 时,Engram 模型仍能与纯 MoE(\(\rho=100\%\))达到相当表现(5.7B 设置对应 46 个专家,9.9B 设置对应 43 个专家)。同时,纯 MoE 并非最优:将约 20%–25% 的稀疏预算分配给 Engram 能得到最佳性能。

以 10B 区间(\(C=6\times10^{20}\))为例,验证损失从 \(\rho=100\%\) 的 1.7248 改善到接近最优 \(\rho\approx80\%\) 的 1.7109(\(\Delta=0.0139\))。且最优点在不同区间稳定(\(\rho\approx75\%\text{--}80\%\)),表明在固定稀疏比下存在鲁棒的分配偏好。该 U 形结构也直观反映了两模块的互补性:

  • MoE 主导(\(\rho\to100\%\)):缺少专用静态记忆,迫使模型用深度与计算去重构静态模式,效率低。
  • Engram 主导(\(\rho\to0\%\)):条件计算能力不足,损害依赖动态推理的任务;记忆在此区间无法替代计算。

3.2 无限记忆区间下的 Engram 标度

在第 3.1 节中,作者在固定参数预算下优化分配。现在考虑互补设定:激进扩展记忆。这一动机来自 Engram 能够在系统层面把存储从计算中解耦(第 2.5 节)。

实验协议

作者固定一个 MoE 骨干(\(P_{\mathrm{tot}}\approx3\)B、\(P_{\mathrm{act}}=568\)M),训练 100B tokens 以确保收敛;在其上叠加 Engram 表,并把槽位数 \(M\) 从 \(2.58\times10^5\) 扫到 \(1.0\times10^7\)(最多增加约 130 亿参数)。基线为 OverEncoding(原文引用:huang2025over),其通过与词表嵌入求平均整合 N‑gram 嵌入。作者指出:SCONE(原文引用:yu2025scaling)等工作更多面向推理并引入额外模块与训练 FLOPs,因此不满足本文严格的等计算对比约束。

结果

图 3(右)显示:随着记忆槽位数增加,验证损失持续、稳定下降,并在对数空间呈严格幂律(线性)趋势。这表明 Engram 提供了一个可预测的扩展旋钮:在不增加计算的前提下更大记忆仍持续带来收益。并且在同等记忆预算下,Engram 比直接平均的 OverEncoding 具有更大的扩展潜力。结合第 3.1 节的分配规律,作者据此认为:条件记忆是一条独立且可扩展的稀疏容量轴,与 MoE 的条件计算互补。

4 大规模预训练

类别 基准(指标) Shots Dense‑4B MoE‑27B Engram‑27B Engram‑40B
# 总参数 4.1B 26.7B 26.7B 39.5B
# 激活参数(不含 token embed) 3.8B 3.8B 3.8B 3.8B
# 训练 token 262B 262B 262B 262B
# 专家(共享 + 路由,top‑k) 2 + 72(top‑6) 2 + 55(top‑6) 2 + 55(top‑6)
# Engram 参数 5.7B 18.5B
语言建模 Pile(loss) 2.091 1.960 1.950 1.942
语言建模 验证集(loss) 1.768 1.634 1.622 1.610
知识与推理 MMLU(Acc.) 5‑shot 48.6 57.4 60.4 60.6
MMLU‑Redux(Acc.) 5‑shot 50.7 60.6 64.0 64.5
MMLU‑Pro(Acc.) 5‑shot 21.1 28.3 30.1 31.3
CMMLU(Acc.) 5‑shot 47.9 57.9 61.9 63.4
C‑Eval(Acc.) 5‑shot 46.9 58.0 62.7 63.3
AGIEval(Acc.) 0‑shot 29.1 38.6 41.8 45.9
ARC‑Easy(Acc.) 25‑shot 76.8 86.5 89.0 90.1
ARC‑Challenge(Acc.) 25‑shot 59.3 70.1 73.8 76.4
TriviaQA(EM) 5‑shot 33.0 48.8 50.7 51.8
TriviaQA‑ZH(EM) 5‑shot 62.8 74.8 76.3 77.9
PopQA(EM) 15‑shot 15.1 19.2 19.4 21.2
CCPM(Acc.) 0‑shot 72.2 79.6 87.1 87.7
BBH(EM) 3‑shot 42.8 50.9 55.9 57.5
HellaSwag(Acc.) 0‑shot 64.3 71.8 72.7 73.1
PIQA(Acc.) 0‑shot 63.8 71.9 73.5 76.5
WinoGrande(Acc.) 5‑shot 64.0 67.6 67.8 68.1
阅读理解 DROP(F1) 1‑shot 41.6 55.7 59.0 60.7
RACE‑Middle(Acc.) 5‑shot 72.4 80.9 82.8 83.3
RACE‑High(Acc.) 5‑shot 66.0 75.4 78.2 79.2
C3(Acc.) 0‑shot 57.7 60.1 63.6 61.8
代码与数学 HumanEval(Pass@1) 0‑shot 26.8 37.8 40.8 38.4
MBPP(Pass@1) 3‑shot 35.4 46.6 48.2 46.2
CruxEval‑i(EM) 0‑shot 27.6 30.7 32.2 36.2
CruxEval‑o(EM) 0‑shot 28.7 34.1 35.0 35.3
GSM8K(EM) 8‑shot 35.5 58.4 60.6 62.6
MGSM(EM) 8‑shot 27.0 46.8 49.4 52.4
MATH(EM) 4‑shot 15.2 28.3 30.7 30.6

表 1:稠密、MoE 与 Engram 预训练性能对比。原文说明:所有模型训练 262B tokens,激活参数匹配(3.8B)。Engram‑27B 通过把路由专家参数(72→55)重分配为 5.7B Engram 记忆,与 MoE‑27B 严格等参数;Engram‑40B 在固定激活参数预算下进一步扩展 Engram 记忆(18.5B)。完整训练轨迹见附录 B。

在提出的 Engram 架构与经验分配规律指导下,作者把 Engram 扩展到多十亿参数级别以验证其在真实预训练中的有效性。具体训练四个模型:

  1. Dense‑4B(总参数 4.1B);
  2. MoE‑27B(总参数 26.7B);
  3. Engram‑27B(总参数 26.7B);
  4. Engram‑40B(总参数 39.5B)。

所有模型使用相同的数据课程(token 预算与顺序一致),且在激活参数数量上严格匹配。

4.1 实验设置

训练数据与模型配置

所有模型在 262B tokens 的语料上预训练,并使用 DeepSeek‑v3 的分词器(词表 128k;原文引用:liu2024deepseek)。为保证对比可控,除非特别说明,所有模型均采用统一默认设置:30 层 Transformer、隐藏维度 2560;每层使用 32 头的 Multi‑head Latent Attention(MLA;原文引用:deepseekai2024deepseekv2strongeconomicalefficient),并通过 mHC(原文引用:xie2025mhcmanifoldconstrainedhyperconnections,扩展率 4)连接到 FFN;优化器使用 Muon(原文引用:jordan2024muon、team2025kimi)。更细的超参见附录 A。

四个模型分别为:

  • Dense‑4B:基线,使用标准稠密 FFN。
  • MoE‑27B:将稠密 FFN 替换为 DeepSeekMoE(原文引用:dai2024deepseekmoe),配置 72 个路由专家与 2 个共享专家,每 token 激活 top‑6 路由专家。
  • Engram‑27B:严格从 MoE‑27B 派生以公平比较。把路由专家从 72 减到 55,并把释放参数重分配为 5.7B 的嵌入模块(\(\rho=74.3\%\)),保持总参数 26.7B 不变。Engram 插入层为第 2 层与第 15 层,最大 N‑gram 阶数 3,哈希头数 8,维度 1280。嵌入参数用 Adam(原文引用:kingma2014adam)更新,学习率乘以 5 且不使用权重衰减;卷积参数零初始化以保证训练开始时严格保持恒等映射。
  • Engram‑40B:保持与 Engram‑27B 相同骨干与计算预算,但将稀疏嵌入模块扩展到 18.5B(总参数 39.5B)。

评测协议

评测覆盖语言建模、知识/推理、阅读理解、代码/数学等基准,并遵循各基准标准提示与指标(详见表 1 中列举的基准名称与 shots 设置)。

4.2 实验结果

表 1 汇总主要结果。首先,与已有研究一致(原文引用:shazeer2017outrageously、he2024mixture、borgeaud2022improving),稀疏架构相对稠密模型具备更优的扩展规律:在相同训练计算预算下,三种稀疏变体(MoE‑27B、Engram‑27B/40B)在所有基准上显著优于等 FLOPs 的 Dense‑4B。

更重要的是,Engram‑27B 在严格等参数、等 FLOPs 条件下稳定优于 MoE‑27B。并且增益不仅出现在直觉上会受益于记忆容量的知识任务(例如 MMLU、CMMLU),在通用推理(BBH、ARC‑Challenge、DROP)以及代码/数学(HumanEval、MBPP、GSM8K、MATH)中也同样明显。作者据此认为:引入专用的知识查找原语能提升表示效率,超越把全部稀疏预算都分配给条件计算所能达到的效果。

最后,Engram‑40B 在多数基准上进一步降低预训练损失并提升性能。尽管它尚未在每个任务上严格支配 Engram‑27B,作者认为这更可能是训练不足的产物:在训练后期,Engram‑40B 与基线之间的损失差距仍在扩大,暗示更大的记忆容量在当前 token 预算下尚未饱和。

5 长上下文训练

模型 LongPPL(32k,PPL↓) RULER(32k,NIAH Acc.↑) RULER(32k,其他任务↑)
Book Paper Code L‑CoT S MK MV MQ VT CWE FWE QA
MoE‑27B(50k, 1.63) 4.38 2.91 2.49 14.16 100.0 88.0 92.7 84.2 77.0 4.5 73.0 34.5
Engram‑27B(41k, 1.66) 4.37 2.92 2.50 14.26 99.6 88.3 93.0 89.5 83.2 3.8 99.6 44.0
Engram‑27B(46k, 1.63) 4.19 2.84 2.45 13.59 97.6 89.0 95.5 97.0 87.2 4.3 98.6 37.5
Engram‑27B(50k, 1.62) 4.14 2.82 2.44 13.41 99.3 89.3 96.5 97.0 89.0 5.9 99.3 40.5

表 2:长上下文性能对比。括号(如 (50k, 1.62))表示长上下文扩展前的预训练步数与对应 loss。原文总结的两点:① 仅用 82% 预训练 FLOPs(41k vs 50k),Engram‑27B 在 LongPPL 上与基线相当,同时在 RULER 上更高;② 在等预训练 loss(46k)与等预训练 FLOPs(50k)两种对齐方式下,Engram‑27B 在各指标上均显著优于基线。

通过把局部依赖建模卸载到静态查找,Engram 保留了更多注意力容量以管理全局上下文。本节通过长上下文扩展训练验证这一结构优势,并在严格对齐的评测协议下展示 Engram 在长程检索与推理任务上的显著收益。

5.1 实验设置

训练细节

作者采用 DeepSeek‑V3 的上下文扩展策略(原文引用:liu2024deepseek),在预训练后进行 32768 token 的上下文扩展训练阶段,训练 5,000 steps(约 30B tokens 的高质量长上下文数据)。位置扩展方法使用 YaRN(原文引用:peng2023yarn),超参为 scale \(s=10\)、\(\alpha=1\)、\(\beta=32\),缩放因子 \(f=0.707\)。

模型配置

对比四种配置:MoE‑27B(50k)与 Engram‑27B(50k)的最终预训练 checkpoint;以及 Engram‑27B 的两个中间 checkpoint(41k、46k)。尽管起点不同,所有配置都进行完全相同的上下文扩展训练。尤其地,Engram‑27B(46k)与 MoE‑27B(50k)在预训练 loss 上对齐,构成“等 loss”设置,使得扩展阶段的差异可归因于架构本身而非起点模型质量。

评测基准

使用 LongPPL(原文引用:fangwrong)与 RULER(原文引用:hsiehruler)。LongPPL 覆盖长书籍、研究论文、代码仓库与长 CoT 轨迹四类;RULER 包含 14 个子集汇总成 8 类:S、MK、MV、MQ(Needle‑in‑a‑Haystack)、VT(多跳变量跟踪)、CWE(常见词抽取)、FWE(高频词抽取)与 QA。

5.2 实验结果

如表 2 所示,作者按两步分析:先剥离“底座能力”的影响,再在受控设置下比较架构。

1)长上下文能力不只由注意力机制决定。虽然注意力与位置编码提供了处理上下文的结构基础(原文引用:su2023roformerenhancedtransformerrotary、press2021train、yang2025path、xiao2023efficient),但结果表明长上下文性能并非仅由架构先验决定:在同一架构与相同扩展训练预算下,Engram 从 41k → 50k 的长上下文表现随预训练推进单调提升,说明长上下文能力与底座通用建模能力紧密耦合。因此,严格的架构比较应优先对齐预训练 loss,而不只是对齐训练步数。

2)受控设置下的架构优势。在对齐底座能力后,Engram 的效率优势清晰呈现:在等 loss(46k)设置中,Engram‑27B 在复杂检索任务上显著优于基线(例如 MQ‑NIAH:97.0 vs 84.2;VT:87.2 vs 77.0);在等 FLOPs(50k)设置中差距进一步扩大;甚至在仅约 82% 预训练计算的极端设置(41k)下,Engram 仍在 LongPPL 上与基线相当并在 RULER 上更好,体现其结构性优势。

打开图 4(PDF)
图 4:表示对齐与收敛速度分析。(a) 使用 LogitLens 的逐层 KL 散度:Engram 在早期层持续更低,表明其更快达到“可预测”状态;(b‑c) 用 CKA 计算的相似度热力图:高相似对角线整体上移,意味着 Engram 的浅层在功能上相当于 MoE 更深层,从而提升有效深度。

6 分析

本节从机理与系统两个层面分析 Engram:包括其有效深度(6.1)、核心模块设计与层敏感(6.2)、参数敏感性(6.3)、带 Offload 的推理吞吐(6.4),以及门控可视化案例(6.5)。

6.1 Engram 是否等价于增加模型深度?

当前 LLM 缺少专用的知识查找原语,只能依赖计算去模拟记忆回忆。如表 3 所示,为识别实体 "Diana, Princess of Wales",模型需要通过多层注意力与 FFN 逐步合成特征(原文引用:li2025echoesbertmodernlanguage、ghandeharioun2024patchscopes、DBLP:conf/coling/JinYHZWH0MMDYDZ25)。从概念上看,这一过程本可由一次知识查找操作完成。

作者据此提出假设:当模型具备显式的知识查找能力时,Engram 能通过卸载早期特征合成而在功能上模拟“更深的网络”。为验证该假设,作者使用两种机理可解释性工具:LogitLens(原文引用:nostalgebraist2020logitlens、belrose2023eliciting)与 CKA(原文引用:kornblith2019similarity、davari2022reliability)。

加速的预测收敛(LogitLens)

LogitLens 的做法是:用最终 LM Head 投影每一层的隐藏状态,计算中间输出分布与最终输出分布之间的 Kullback–Leibler(KL)散度(原文引用:kullback1951information)。该指标衡量一个潜在表示距离“可直接输出高置信预测”有多近(原文引用:csordas2025language、belrose2023eliciting)。

图 4(a) 显示:与 MoE 基线相比,Engram 的逐层 KL 散度系统性更小,差距在早期块最明显;且 Engram 曲线下降更陡,表明其更快完成特征合成。这与核心假设一致:通过显式访问外部知识,Engram 减少了达到高置信预测所需的计算步骤,使得网络更早进入“预测就绪”的状态。

表示对齐与有效深度(CKA)

为了检验 Engram 的某一层是否在语义上对应基线更深层,作者使用 Centered Kernel Alignment(CKA)比较表征结构(原文引用:kornblith2019similarity、kriegeskorte2008representational)。给定两组表示 \(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 Independence Criterion(原文引用:gretton2005measuring)。作者采用 minibatch 形式与无偏 HSIC 估计(原文引用:davari2022reliability),并在 Few‑NERD 数据集(原文引用:ding2021few)上评估,抽取命名实体最后一个 token 的隐藏状态。

为量化层间对应关系,作者先计算 CKA 相似度矩阵 \(S\in[0,1]^{L\times L}\),再为每个 Engram 层 \(j\) 定义软对齐指数 \(a_j\):取与之最相似的 top‑\(k\) 个 MoE 层集合 \(\mathcal{I}_j\)(\(k=5\)),并用相似度加权其层号质心:

$$ a_j=\frac{\sum_{i\in\mathcal{I}_j}S_{i,j}\cdot i}{\sum_{i\in\mathcal{I}_j}S_{i,j}},\qquad \mathcal{I}_j=\mathop{\text{argtop}k}_{i}(S_{i,j}). $$

图 4(b‑c) 的热力图与虚线对齐曲线显示对角线上移:在广泛层范围内 \(a_j>j\)。例如,Engram‑27B 的第 5 层表征最接近 MoE 基线的约第 12 层。结合 LogitLens 的结果,作者据此认为:Engram 能在更早层形成更深层语义表征,从功能上等价于增加有效深度。

潜在状态翻译(Latent State Translation) 解释(Explanation)
1‑2 :英国的一个国家 Wales
3 :欧洲的一个国家 Wales
4 :女性君主或王后配偶所持有的头衔 Princess of Wales(不具体)
5 :威尔士亲王(后来的国王)之妻所获头衔 Princess of Wales(不具体)
6 :戴安娜,威尔士王妃(1961‑1997),威尔士亲王查尔斯的第一任妻子,以美貌与人道主义工作闻名 Diana, Princess of Wales

表 3:实体解析案例(复现自 PatchScopes;原文引用:ghandeharioun2024patchscopes)。表中展示模型如何随着层数增加逐步整合上下文 token 来构建实体 "Diana, Princess of Wales" 的内部表示。

打开图 5(PDF)
图 5:架构消融结果。与 3B MoE 基线对比:左(深蓝曲线)为单个 Engram 插入层深度扫描,显示越早注入越好且第 2 层最优;右(标记点)为组件消融,显示多分支集成、分词器压缩与上下文门控对性能至关重要。

6.2 结构消融与层敏感

作者在受控设置下对 Engram 进行结构消融,以分析关键设计的作用。除非特别说明,骨干为 12 层 3B MoE(激活参数 0.56B),训练 100B tokens。图 5 报告验证损失,虚线橙色为 3B MoE 基线(Val Loss = 1.808)。

参考配置

在骨干上增加固定 1.6B 参数的 Engram 记忆。参考模型使用 \(\{2,3\}\)-gram,并在第 2 层与第 6 层插入 Engram,验证损失为 1.768,相比基线提升 \(\Delta=0.04\)。下述消融均相对该参考配置定义。

记忆应插入在哪里?

为研究深度敏感性,作者保持 Engram 总预算不变(1.6B),把其合并为单个模块,并把插入层从第 1 层扫到第 12 层(图 5 深蓝曲线)。结果揭示出内在权衡:越早注入越能在骨干耗费计算深度前卸载局部模式重构,符合骨干的层级处理特性(原文引用:tenney2019bert 等);但早期隐藏状态尚未通过注意力聚合足够全局上下文,且并行分支的表示分化不足,门控精度可能更差(原文引用:xie2025mhcmanifoldconstrainedhyperconnections、zhu2025hyperconnections)。

扫描显示第 2 层在单次注入约束下最佳(Val Loss = 1.770),优于第 1 层且注入点越深效果越差。这说明一次注意力已足以提供有意义的上下文化 \(\mathbf{h}_t\) 用于门控,同时仍足够早以替代底层的局部聚合。

此外,把同样 1.6B 记忆拆分为两个更小模块(通过降低 \(d_{\text{mem}}\))并分别放在第 2 层与第 6 层表现更好(Val Loss = 1.768)。这种分层设计在“早卸载”与“更强上下文化门控”之间取得折中,并在系统层面也有优势:更利于利用多级缓存层次(第 2.5 节)。

哪些组件最关键?

在保持参数预算不变的前提下逐一移除设计组件,结果见图 5 的标记点。三项带来最大收益的组件为:(i)在多分支骨干内做分支特异的融合;(ii)上下文门控;(iii)分词器压缩。移除任意一项都会导致最大的性能回退。对于“去除多分支”消融,作者保留 mHC 骨干结构,但用单个 Engram 融合替代分支特异门控,并把融合应用到预映射 \(\mathcal{H}^{pre}\) 之后的隐藏状态上(原文引用:xie2025mhcmanifoldconstrainedhyperconnections)。

其它改动影响较小:移除轻量 depthwise 卷积仅轻微退化;在固定 1.6B 预算下给 4‑gram 分配容量略次优,作者猜测是因为稀释了更常见的 2/3‑gram 模式,但不排除在更大记忆规模下更高阶 N‑gram 会变得有利。

打开图 6(PDF)
图 6:抑制 Engram 输出后的性能保留率。事实性知识对 Engram 高度依赖,而阅读理解主要由骨干注意力机制维持。

6.3 敏感性分析

为刻画 Engram 的功能贡献,作者在推理时完全抑制稀疏嵌入输出而保持骨干不变。需要强调:这种事后消融会引入训练‑推理不一致,可能给混合能力任务带来噪声。因此作者重点分析信噪比更高的两个极端:事实性知识阅读理解

如图 6 所示,结果呈现鲜明二分:事实性知识基准发生“灾难性坍塌”,仅保留 29%–44% 原性能(例如 TriviaQA 仅 29%),说明 Engram 是参数化知识的主要载体;而阅读理解任务非常稳健,保留 81%–93%(例如 C3 保留 93%),暗示上下文对齐任务主要依赖骨干的注意力机制而非 Engram。

6.4 推理系统效率(CPU Offload)

Engram 相对路由式 MoE 的关键系统优势在于:其稀疏激活由显式、静态的哈希 ID 寻址,访问模式严格确定。给定 token 序列后,下一次 Engram 查找的索引可在对应层执行前计算出来,从而支持预取与通信‑计算重叠。

实验设置

作者基于 nano‑vLLM(https://github.com/GeeeekExplorer/nano-vllm)实现推理基准框架,并以两个稠密骨干(Dense‑4B 与 Dense‑8B)测试,以避免 MoE 的 Expert Parallel 通信模式干扰时延基线。在第二个 Transformer 块插入一个 100B 参数的 Engram 层,并将整个嵌入表常驻主存(DRAM)。推理时异步预取该层所需嵌入,并与第一个块的计算重叠。

实验设置
硬件 NVIDIA H800
负载 512 sequences
序列长度 Uniform(100, 1024)
基座模型 配置 吞吐(tok/s)
4B‑Dense Baseline 9,031.62
+ 100B Engram(CPU Offload) 8,858.28
8B‑Dense Baseline 6,315.52
+ 100B Engram(CPU Offload) 6,140.02

结果。在该设置下,卸载 100B 参数嵌入表带来的吞吐损失很小,8B 骨干上最大仅约 2.8%。这验证了早期稠密块的计算强度足以提供掩蔽检索时延的时间窗口,并且每 step 的有效通信量随“激活槽位数”而非“表总规模”增长。

作者强调:这只是保守基线。尽管第 2.5 节提出可利用 Zipf 局部性把高频项缓存在 HBM,但此实验强制所有检索都通过 PCIe 从主存获取;在这种不利条件下仍仅有极小开销,暗示在更完整的局部性感知优化下,吞吐损失可忽略。

打开图 7(PDF)
图 7:Engram 门控可视化。热力图强度对应 \(\alpha_t\in[0,1]\) 的大小,红色越深表示激活越强。由于 Engram 作用于后缀 N‑gram(此处 \(N=3\)),某 token 上的高激活意味着以该 token 结尾的短语被识别为静态模式并从记忆中有效检索。

6.5 案例:门控可视化

在第 2.3 节中,作者提出上下文门控用于动态调制静态记忆与骨干表示的融合。为验证 Engram 是否按预期工作,作者对 Engram‑27B 的门控标量 \(\alpha_t\) 进行可视化。按原文脚注说明:该模型使用 mHC(\(M=4\))且 Engram 插入在第 2 与第 15 层,因此每个 token 共有 8 个门控标量;并非每个分支都有可解释的激活模式,图中选取与语义模式匹配相关性最强的门控进行展示。

结果呈现明显的选择性:门控机制在完成局部、静态模式时持续被激活(红色)。英文样例中,命名实体(如 “Alexander the Great”“the Milky Way”)与固定短语(如 “By the way”“Princess of Wales”)上出现强激活;并且该行为能跨语言泛化:中文样例中,Engram 能识别并检索诸如“四大发明”“张仲景”等固定表达与历史实体。作者据此认为:Engram 能有效处理刻板的语言依赖,从而减轻 Transformer 骨干对这些静态关联的记忆负担。

8 结论

作者提出将条件记忆作为对当前条件计算(MoE)范式的互补稀疏轴,以解决 Transformer 只能用动态计算模拟知识检索的低效问题。作者以 Engram 作为实现:将经典 N‑gram 嵌入现代化,使得静态模式可通过可扩展的 \( \mathcal{O}(1) \) 查找获得。

通过形式化“稀疏分配”问题,作者发现一条 U 形标度律:在固定稀疏预算下,MoE 专家与 Engram 记忆的混合分配严格优于纯 MoE。遵循该规律,Engram 可扩展到 27B 参数,并在多个任务域取得显著提升。机理分析表明 Engram 通过卸载早期静态重构来“加深”网络有效深度,并通过把局部依赖交给查找释放注意力以增强长上下文能力。最后,Engram 的确定性寻址支持将超大表卸载到主存并进行运行时预取,几乎不增加推理开销,体现了“基础设施感知”的高效设计理念。作者认为,条件记忆将成为下一代稀疏模型的重要建模原语。

附录

A 详细模型结构与训练超参

Dense‑4B MoE‑27B Engram‑27B Engram‑40B
总参数4.1B26.7B26.7B39.5B
激活参数3.8B
总训练 tokens262B
层数30
隐藏维度2560
前置稠密层111
路由专家数725555
激活专家(top‑k)666
共享专家数222
负载均衡方法Loss Free(原文引用:wang2024auxiliarylossfreeloadbalancingstrategy)
注意力模块MLA(原文引用:deepseekai2024deepseekv2strongeconomicalefficient)
RoPE \(\theta\)10000
mHC 扩展率4
序列长度4096
词表大小129280
Batch size1280
训练步数50000
骨干优化器Muon(原文引用:jordan2024muon)
嵌入优化器Adam(原文引用:kingma2014adam)
基础学习率4e‑4
学习率调度Step Decay(原文引用:bi2024deepseek)
权重衰减0.1
Engram 维度 \(d_{\text{mem}}\)12801280
Engram 词表大小22624007239680
Engram 头数88
Engram 插入层[2, 15][2, 15]
Engram N‑gram[2, 3][2, 3]
Engram 结合 mHCTrueTrue
Engram 分词器压缩TrueTrue
Engram 卷积零初始化TrueTrue
Engram 学习率倍增×5×5
Engram 权重衰减0.00.0
Engram 优化器(仅嵌入)Adam(原文引用:kingma2014adam)

表 A.1:详细模型结构与训练超参(根据原文附录表格整理)。

B 完整 Benchmark 曲线

打开图 B.1(PDF)
图 B.1:最后 10k 预训练 benchmark 曲线。

C 分词器压缩案例:Tokenizer Compression

Rank Merge Count Normalized Token Original Tokens(部分示例)
1 163 '␠' '\\t', '\\n', '\\r', '␠', '␠␠', '\\n\\n', '␠␠␠', '␠\\n', ...
2 54 'a' 'A', 'a', '␠a', '␠A', 'á', 'ä', 'ã', 'ą', '␠à', '␠å', 'â', ...
3 40 'o' 'O', 'o', '␠o', '␠O', 'ó', 'ö', 'ô', 'õ', 'ő', 'ò', ...
4 35 'e' 'E', 'e', '␠e', '␠E', 'é', 'è', '␠é', 'ę', 'ě', 'ê', ...
5 30 'i' 'I', 'i', '␠I', '␠i', 'í', 'ì', 'î', 'ī', 'ï', ...

表 C.1:分词器压缩合并最频繁的 Top‑5 token 示例。原文说明:对 128k 分词器,整体压缩比例为 23.43%。

参考文献(BibTeX)

展开/收起:原文 BibTeX(为保证完整性,未做逐条中文化)

正在生成参考文献索引……

提示:引用键(如 shazeer2017outrageously)与 BibTeX 条目一一对应。