N-gram Embedding:大语言模型中的条件内存技术

核心结论

N-gram embedding 是一种将多 token 局部模式(bigram、trigram 等)通过哈希查表直接注入 Transformer 隐藏状态的技术。它的核心价值在于:用 O(1) 的静态查找代替 Transformer 早期层对局部模式的昂贵重建计算,从而释放网络深度用于更高阶的推理任务。在严格控制参数量和计算量的对比实验中,配备该技术的模型(Engram-27B)在 MMLU 上超出纯 MoE 基线 3.4 分,在 BBH 推理任务上超出 5.0 分,而推理吞吐量损失不到 3%。

该技术的两条主要实现路线——Over-Encoding(字节跳动,2025)和 Engram(DeepSeek,2026)——分别代表了"输入层叠加"和"中间层条件内存"两种设计哲学,后者在 MoE 架构上效果更显著。

关于"Qwen3.8-Flash-Next":2026 年 8 月 26 日,Qwen 团队正式开源了 Qwen3.8-Flash-Next(github.com/QwenLM/Qwen3.8-Flash-Next),一个 125B 总参数、6B 激活/token 的 MoE 模型,其架构明确采用了 n-gram embedding 技术——额外 51B 参数的 n-gram 嵌入表存放在主机内存中,通过 host-memory prefetching 注入到 Layer 2。技术报告引用了 Google DeepMind (2025) 和 Cheng et al. (2026)(即 Engram)作为该组件的来源。报告还指出一个重要发现:扩大 n-gram 词表可以单调降低 loss,但下游任务准确率会饱和——loss 最优点与 accuracy 最优点并不重合。


技术发展脉络

时间 工作 机构 核心贡献
2017 Hash Embeddings DTU 多哈希共享嵌入池,首次解决大词表参数效率问题
2017 FastText Meta 子词 n-gram 嵌入用于词表示学习
2022 N-Grammer Google 首次在 Transformer 中引入离散 latent n-gram 增强
2025.01 Over-Tokenized Transformer ByteDance 发现输入词表大小与 loss 的对数线性关系,提出 Over-Encoding
2026.01 Engram DeepSeek + PKU 将 n-gram embedding 升级为条件内存原语,提出稀疏分配法则
2026.06 TN-gram Imperial + RIKEN 用张量 CP 分解压缩 Engram,跨 n-gram 阶共享潜在因子
2026.08 Qwen3.8-Flash-Next Alibaba Qwen 首个公开采用 n-gram embedding 的生产级模型(51B 表参数,Layer 2 注入)

从统计语言模型到条件内存:主线问题的演进

起点:Transformer 缺少知识查找原语

语言建模包含两类本质不同的子任务——组合推理和知识检索。对于"Alexander the Great"这样的多 token 实体,标准 Transformer 需要消耗多个早期层的注意力和前馈网络逐步拼合特征,才能在第 6 层左右形成完整的实体表示(Ghandeharioun et al., 2024 的 PatchScope 实验清晰展示了这一过程)。这些固定的、刻板的局部模式本质上可以用一次查表完成,却要耗费宝贵的网络深度去"运行时重建"一张静态查找表。

N-gram 语言模型(Shannon, 1948; Brants et al., 2007)早已证明局部 token 共现的统计规律性。问题在于:如何将这种 O(1) 的模式查找能力注入现代 Transformer 架构?

第一步:哈希嵌入解决词表爆炸

直接为所有可能的 n-gram 分配独立嵌入向量在计算上不可行——一个 128K 词表的 trigram 空间有 128K³ ≈ 2×10¹⁵ 种组合。Hash Embeddings(Svenstrup et al., 2017)给出了经典解决方案:用 k 个不同的哈希函数将任意 token 映射到一个大小为 B 的共享嵌入池中,最终表示是 k 个查找结果的加权和。这把存储需求从 O(V) 降到 O(B),且 B 可以独立于 V 任意设定。

FastText(Bojanowski et al., 2017)将类似思路应用于子词 n-gram:一个词的表示等于它所有字符 n-gram 嵌入的求和。但这些工作停留在浅层的词表示学习,没有深入 Transformer 的中间层。

第二步:Over-Encoding——输入层的词表缩放

Over-Tokenized Transformer(Huang et al., 2025, ICML)发现了一个关键的 scaling law:输入词表大小与训练 loss 呈对数线性关系。具体做法是在标准 BPE tokenizer 之上,叠加 bigram 和 trigram 的哈希嵌入表,将输入 embedding 计算为各阶 n-gram 嵌入的求和:

$$\text{OE}(x_t) = E_{V \times d}(x_t^{(-1)}) + \sum_{i=2}^{n} E_{m \times d/k}^{(i)}(x_t^{(-i)}) \cdot W_i$$

其中 $x_t^{(-i)} = f(x_t, x_{t-1}, \ldots, x_{t-i+1})$ 是后缀 i-gram 的哈希索引,$m$ 是嵌入表行数(可设到 10⁷),$W_i$ 是低秩投影矩阵。关键设计是将每个 n-gram 嵌入表切片为 k 个低维子表再投影拼合(式中 $d/k$ 维),以同等参数量换取更高的组合表达力。

实验结果:在 OLMo2 框架下,400M 参数的 OE-12.8M 模型(输入词表有效扩展 128 倍至 1280 万)在训练 loss 上追平了 1B 参数的 baseline,训练收敛速度提升 5.7 倍。代价是仅增加不到 5% 的训练开销。

Over-Encoding 的局限在于:它只作用于输入层(Layer 0),嵌入查找和第一层计算之间没有时间间隔可以重叠通信,且在 MoE 架构上收益减弱(OLMoE-7B 上 downstream 仅提升 +0.007),因为 MoE 的稀疏专家已经提供了部分类似的"静态模式记忆"能力。

第三步:Engram——条件内存作为一等建模原语

Engram(Cheng et al., 2026, DeepSeek)将 n-gram embedding 从输入层的附属组件提升为一个独立的稀疏维度——与 MoE 的条件计算(conditional computation)平行的条件内存(conditional memory)。其核心洞见是:MoE 通过动态路由扩展计算能力,但模型仍然缺少原生的静态知识查找通道;Engram 通过确定性地址寻址填补了这个空白。

前向计算的完整流程

给定输入序列 $X = (x_1, \ldots, x_T)$,在第 $\ell$ 层隐藏状态 $H^{(\ell)} \in \mathbb{R}^{T \times d}$ 上,Engram 的处理分为检索和融合两个阶段:

1. 词表压缩(Tokenizer Compression)

先将原始 token ID 通过确定性映射 $\mathcal{P}: V \to V'$ 压缩为规范标识符。映射规则包括 NFKC 归一化、转小写、去重音等,使得"Apple"和"␣apple"对应同一个规范 ID。对于 128K 词表,这一步将有效词表缩小约 23%,提高后续哈希的语义密度。

2. 多头哈希检索(Multi-Head Hashing)

对每个位置 $t$,构造后缀 n-gram $g_{t,n} = (x'{t-n+1}, \ldots, x'_t)$,然后用 $K$ 个不同的哈希头映射到大小为素数 $M{n,k}$ 的嵌入表:

$$z_{t,n,k} = \varphi_{n,k}(g_{t,n}), \quad e_{t,n,k} = E_{n,k}[z_{t,n,k}]$$

哈希函数 $\varphi_{n,k}$ 采用乘法-XOR 结构:对 n-gram 中每个 token 的规范 ID 乘以随机奇数系数后逐位异或,再对素数表大小取模。所有 n-gram 阶(2-gram, 3-gram)和所有头的结果拼接为最终内存向量:

$$e_t = \mathop{|}{n=2}^{N} \mathop{|}{k=1}^{K} e_{t,n,k} \in \mathbb{R}^{d_{\text{mem}}}$$

Engram-27B 的具体配置:max n-gram = 3,hash heads = 8,embedding dimension = 1280,嵌入表行数 ≈ 646K×2(bigram 和 trigram 各一组),模块分别插入在第 2 层和第 15 层。

3. 上下文感知门控(Context-Aware Gating)

静态检索的嵌入 $e_t$ 缺乏上下文适应性,且哈希碰撞会引入噪声。Engram 用当前隐藏状态 $h_t$(已经通过前面的注意力层聚合了全局上下文)作为动态 Query,对检索结果进行调制:

$$k_t = W_K e_t, \quad v_t = W_V e_t$$ $$\alpha_t = \sigma\left(\frac{\text{RMSNorm}(h_t)^\top \text{RMSNorm}(k_t)}{\sqrt{d}}\right)$$

门控标量 $\alpha_t \in (0, 1)$ 的语义是:如果检索到的内存与当前上下文矛盾(如哈希碰撞导致错误匹配),门趋近于零,自动抑制噪声。

4. 短程卷积与残差注入

门控后的值 $\tilde{v}_t = \alpha_t \cdot v_t$ 再经过一个 depthwise causal convolution(kernel=4, dilation=max_ngram_order)扩展感受野,最后通过残差连接注入主干:$H^{(\ell)} \leftarrow H^{(\ell)} + Y$。

为何计算开销可忽略

Engram 的推理开销极低,根源在于三个设计决策:

  1. 确定性寻址:n-gram 哈希索引仅依赖输入 token ID,不依赖隐藏状态。这意味着下一层 Engram 的查表地址在当前层计算开始之前就已知,可以提前异步预取(prefetch)。

  2. 计算-通信重叠(arXiv:2601.07372, Section 2.5, Figure 2):Engram 模块被放置在较深的层(如第 2 层和第 15 层),前面的 Transformer block 计算时间恰好覆盖了从主机内存(host DRAM)通过 PCIe 传输嵌入向量的延迟。实测中(Section 6.4, Table 4),将 100B 参数的嵌入表完全放在 CPU 内存,H800 GPU 上的吞吐量损失不到 2.8%(Dense-8B backbone, 512 序列, 长度 Uniform(100,1024))。

  3. Zipfian 分布利用:自然语言 n-gram 服从 Zipf 分布——少量高频模式占据绝大多数访问。频繁访问的嵌入可以缓存在 GPU HBM 中,长尾部分放在主机 DRAM 甚至 NVMe SSD。

  4. 不增加 per-token FLOPs:嵌入查找是纯索引操作(无矩阵乘法),门控只涉及一次点积和 sigmoid。在 Engram-27B 中,每 token 激活参数(3.8B)与 MoE-27B 完全相同。


增益来源:论文作者的"有效深度"假设及其证据

Engram 论文作者提出的核心假设是:通过显式注入局部模式的预计算表示,Engram 免除了早期层逐步拼合特征的负担,功能上等价于增加网络深度。他们用两组相关性证据支持这一假设(arXiv:2601.07372, Section 6.1):

LogitLens 分析:将每一层的隐藏状态投影到输出词表空间,测量与最终预测分布的 KL 散度。Engram 模型在前几层的 KL 散度显著更低,表明模型更早到达"预测就绪"状态。

CKA 对齐分析:Engram 第 5 层的表示与纯 MoE 基线第 12 层的表示 CKA 相似度最高。作者据此认为 Engram 的浅层在表征几何上对应于 MoE 的更深层。

这两组证据属于相关性观察——它们展示了"表征几何相似"和"输出分布接近",但并未通过因果干预(例如频率分桶消融、causal patching、等 FLOPs 深度对照)来排除替代假设。一个同等合理的替代解释是:Engram 主要卸载了高频 n-gram 共现模式的记忆负担(而非真正增加了推理深度),释放出的 MLP 容量被用于其他任务。区分这两个假设需要按 n-gram 频率分桶做消融——如果"有效深度"论成立,中低频 n-gram 应贡献更大;如果是"高频记忆卸载",则高频部分占据主要收益。该实验在论文中未报告。

尽管因果机制尚待完全厘清,论文中有一项直接的功能归因实验(Section 6.3):在推理时完全关闭 Engram 输出,事实性知识任务(TriviaQA)仅保留 29% 性能,而阅读理解(C3)保留 93%。这表明 Engram 确实承担了独立的、可分离的知识存储功能,backbone 的注意力机制则负责上下文推理。

作者的假设与这一功能分离证据一致:推理任务上 Engram 的增益(BBH +5.0)大于知识任务(MMLU +3.4),可能是因为释放出的早期层容量被重新用于需要更多计算步骤的组合推理。


性能对比数据

Engram vs. MoE 基线(262B tokens 预训练,3.8B 激活参数,数据来源:arXiv:2601.07372 Table 1)

任务 MoE-27B Engram-27B Δ
MMLU (5-shot) 57.4 60.4 +3.0
CMMLU (5-shot) 57.9 61.9 +4.0
BBH (3-shot) 50.9 55.9 +5.0
ARC-Challenge (25-shot) 70.1 73.8 +3.7
DROP (1-shot F1) 55.7 59.0 +3.3
HumanEval (Pass@1) 37.8 40.8 +3.0
GSM8K (8-shot) 58.4 60.6 +2.2
MATH (4-shot) 28.3 30.7 +2.4

Over-Encoding(OLMo2-1B,1T tokens 训练,数据来源:arXiv:2501.16975 Figure 4)

指标 Baseline OE-12.8M 效果
Training Loss ~3.2→2.6 ~3.2→2.48 收敛加速 5.7×
MMLU-Var 0.28→0.36 加速 3.2×
Hellaswag 0.35→0.65 加速 3.0×
额外训练开销 <5%

长上下文能力(32K context, 数据来源:arXiv:2601.07372 Table 2)

RULER 子任务 MoE-27B Engram-27B (iso-loss)
Multi-Query NIAH 84.2 97.0
Variable Tracking 77.0 87.2

两条路线的对比

维度 Over-Encoding Engram
注入位置 输入层(Layer 0) 中间层(如 Layer 2, 15)
融合方式 加性叠加到 token embedding 上下文门控 + 残差
是否依赖隐藏状态 否(纯静态) 门控依赖 h_t
MoE 架构兼容性 弱(OLMoE-7B 上仅 +0.007) 强(iso-params 超越 MoE)
推理卸载能力 有限(Layer 0 无法重叠通信) 强(确定性寻址 + prefetch)
参数效率理论 对数线性 scaling law U-shaped 稀疏分配法则
代码公开 论文描述 + PyTorch 伪代码 GitHub 完整 demo 实现
适用场景 Dense 模型的低成本增强 MoE 模型的参数重分配

Over-Encoding 更适合作为已有 Dense 模型的即插即用增强——修改量小、收益确定。Engram 面向的是架构级设计:它要求在训练之初就将一部分稀疏参数预算从 MoE 专家划拨给内存模块,换取更优的整体效率。U-shaped scaling law 表明最优配置是将约 20-25% 的稀疏参数预算分配给 Engram。


参数开销为何几乎可忽略

以 Engram-27B 为例,5.7B 参数的嵌入表听起来庞大,但实际推理时:


已知局限与后续方向

  1. Qwen3.8-Flash-Next 中的具体应用:该模型将 51B 参数的 n-gram embedding 表放在 host memory,仅注入 Layer 2,配合 Gated Residual 四分支残差流和 GDN/QSA 混合 token mixing。技术报告发现 n-gram 词表扩大时 loss 单调下降但 downstream accuracy 饱和——在固定参数预算下,loss 最优配置与 accuracy 最优配置不同(Tab. 8, Tab. 9)。这意味着实际部署时需要根据目标任务类型选择嵌入表规模,而非一味追求最低 loss。

  2. 高阶 n-gram 的收益递减:Engram 的消融实验显示,在 1.6B 参数预算下加入 4-gram 反而略有退化——高阶 n-gram 稀释了 bigram/trigram 的容量。更大规模的内存预算是否能解锁高阶收益,有待验证。

  3. 与推测解码的协同:Engram 的确定性寻址天然适合与 speculative decoding 结合(n-gram 模式可辅助 draft 模型的预测),但目前没有公开的集成实验。

  4. 训练数据量的影响:Engram-40B 在 262B tokens 训练下还未完全饱和(loss gap 仍在扩大)。最优的内存规模/训练量配比仍是开放问题。


参考文献

  1. Cheng, X., Tian, R., et al. "Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models." arXiv:2601.07372, 2026.
  2. Huang, H., Zhu, D., et al. "Over-Tokenized Transformer: Vocabulary is Generally Worth Scaling." ICML 2025. arXiv:2501.16975.
  3. Zhou, W., et al. "Tensorizing Engram: Sharing Latents Across N-Gram Embeddings is Beneficial in LLMs." arXiv:2606.08347, 2026.
  4. Roy, A., et al. "N-Grammer: Augmenting Transformers with Latent N-grams." arXiv:2207.06366, 2022.
  5. Svenstrup, D., Hansen, J.M., Winther, O. "Hash Embeddings for Efficient Word Representations." arXiv:1709.03933, 2017.
  6. Bojanowski, P., et al. "Enriching Word Vectors with Subword Information." TACL, 2017.
  7. Ghandeharioun, A., et al. "PatchScope: A Framework for Inspecting LLM Representations." NeurIPS, 2024.
  8. Yang, A., et al. "Qwen3 Technical Report." arXiv:2505.09388, 2025.

架构示意图

图:N-gram Embedding (Engram) 完整前向流程。输入 token 序列经词表压缩后构造后缀 n-gram,通过多头哈希函数索引大规模嵌入表(O(1) 查找),检索结果经上下文感知门控(隐藏状态作为 Query 调制 Key-Value)和短程卷积后,以残差形式注入 Transformer 主干的隐藏状态。整个过程无矩阵乘法,嵌入表可完全卸载到主机内存。


补充:关键机制问答

N-gram 表如何与原始词表配合——不是替代,是叠加

原始 unigram 词表 embedding(128K × 4096)完全不动,照常在 Layer 0 为每个 token 查出一个向量。N-gram 表是在 Layer 2 通过残差连接额外叠加的信号:backbone 隐藏维度不变,后续层的输入格式不变。

前向流程:位置 t 的 token 先正常查词表 → 过第一个 transformer block → 在 Layer 2 叠加 n-gram 信号 → 后续层正常继续。

哈希压缩的两重含义:M(行数)和 d_sub(列数)

行数 M:bigram 的理论组合数是 128K² = 160 亿种,直接为每种分配一行不可行。通过 hash mod M(M ≈ 64 万,取素数),160 亿种组合被映射到 64 万个桶里。碰撞通过两个机制缓解:(1) 多头——8 个不同哈希函数,同一对 bigram 在不同头里碰撞模式不同,拼起来区分度很高;(2) 自然语言的 Zipf 分布——实际高频 bigram 只有几百万种,在 64 万行 × 8 头的空间里碰撞概率已经很低。M 的大小由参数预算决定:Engram-27B 分配了 5.7B,Qwen3.8-Flash-Next 分配了 51B,M 越大碰撞越少、loss 越低,但 downstream 准确率存在饱和点。

列数 d_sub:总 n-gram 嵌入维度 1280 被切分给 16 张表(2 种 n-gram 阶 × 8 头),每张表只存 80 维。16 张表各查出 80 维再拼成 1280 维,最后投影到 backbone 的 4096 维。同等参数量下"多张小表拼接"比"一张大表"效果更好——不同哈希头的切片拼接构成隐式组合编码,类似 multi-head attention 的设计逻辑。

是否必须从预训练开始引入?

从头训练效果最大(Engram BBH +5.0, MMLU +3.4),因为 backbone 从一开始就学会了"不需要用前几层的 MLP 宽度去硬背高频 n-gram 共现",前几层参数被梯度引导去做更有价值的组合特征——CKA 分析显示 Engram 第 5 层的表示质量对齐到了无 Engram 模型的第 12 层。

CPT 阶段后挂也可以(TOBA-LM 有效),但核心收益"释放早期层深度"大打折扣——因为 backbone 前几层已经花了几百 B tokens 学会了"自己做局部模式重建",几十 B tokens 的 CPT 不够让这些权重彻底重分配功能。

Qwen3.8-Flash-Next 的前几层是否减少了参数?

没有。Qwen 的设计哲学是把 n-gram 表当作"免费的额外容量"——51B 参数放 host memory、不占 GPU、不增加 per-token FLOPs,backbone 的 125B 参数结构完全不变。与 Engram 论文的 iso-parameter 设计(从 MoE 专家里划拨参数)不同,Qwen 选择不在 backbone 内部做 trade-off。

但即使物理参数不变,从头训练让 backbone 每个参数的"有效信息密度"提升——原本充当隐式查找表的那些神经元,现在被释放出来学习更有组合性的特征。同等参数规模下,带 n-gram 的模型有效表达能力更高。