15_ngram_memory_viz.html

N-gram 查表:给模型挂一本短语词典

第 11 章的 MoE 让总参数变多、每个 token 的算力不变。这一关换一种做法:在第 12 章那台小模型旁边挂一张大表, 每个 token 只按最近两三个 token 算出 8 个行号,从表里取 8 行拼成一根向量,再由一个门决定加多少进残差流。这种结构叫 N-gram 查表记忆 (DeepSeek 的 Engram、Qwen3.8 的 N-gram Embedding)。配套代码 phase4-efficiency/13_engram.py, 在 Mac(MPS)上跑了 FineWeb-Edu 四组、tiny shakespeare 五组对照,外加三组吞吐测量。

STEP 1

容量和算力拆开

STEP 2

哈希:行号只看最近几个 token

STEP 3

门控:记忆对不对得上

STEP 4

真实对照与反例

STEP 5

表可以不进 GPU

这一关数据 FineWeb-Edu GPT-2 BPE token 表 262144 行 / 每个哈希头 哈希头 8 2-gram 4 + 3-gram 4 每行 16 维注入 第 2 个 Block 前 inject_layer = 1 n_embd 128 残差流宽度 8 个头 × 每行 16 维 = 128,拼起来正好等于 n_embd;262144 是每个头的行数,8 张表合计 33.6M 参数

① 容量和算力可以拆开

表有 8 张(每个哈希头一张),每张 行数 × 16 维。一个 token 在每张表里只取 1 行,合计 8 × 16 = 128 个数, 和表有多少行无关。下面左边是 8 张表(高度按行数的对数画),橙色横线是 token " transit" 在每张表里取的那一行; 右边三根条用对数刻度。拖滑块改每个头的行数:

每个哈希头的行数 262144

表参数 = 8 × 行数 × 16 主干参数(第 12 章四件套 · 4 层 · 含 5 万词表嵌入) 每 token 从表里读的数

–

表参数

–

表 ÷ 主干

–

每 token 读的数

–

每 token 读到表的几分之一

真实规模总参数里表占多少

每组两根条:上面是总参数(橙色段是 n-gram 表),下面是激活参数, 四组各按自己的总参数画满格。点按钮看出处:

主干 / 模型主体 n-gram 表激活参数

↳ 代码:13_engram.py 的 NgramMemory.__init__(self.tables 形状 N_HEADS × TABLE × ROW_DIM)· NgramMemory.lookup

主干 7.24M 参数里,有多少是词嵌入?

大部分。GPT-2 词表 50,257 个 token × 128 维 = 6,432,896 个参数的嵌入表(和输出头共用), 输出头 bias 50,257 个;4 个 Block 合计 726,016 个,最后一层 RMSNorm 128 个,加起来是不挂表时的 7,209,297。挂表以后多出 33,536 个:门控用的 w_k、w_v 各 128 × 128,两个 RMSNorm 各 128,短卷积 128 × 4。

词嵌入本身就是一张按单个 token id 查的表,每个 token 读 1 行 128 维。n-gram 表把"按 1 个 token 查"扩展成"按最近 2 到 3 个 token 查"。 262144 行的表是 4 个 Block 参数的 46 倍,每 token 前向里真正做矩阵乘的仍是这 4 个 Block、门控和输出头。

同样的总参数,把一部分专家换成表,划算吗?

Engram 论文(Cheng et al.,"Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models",arXiv 2601.07372)做了这组对照: Engram-27B 与 MoE-27B 都是 26.7B 总参、3.8B 激活,Engram-27B 把路由专家从 72 个减到 55 个,腾出的参数给 Engram 表(5.7B)。结果 MMLU +3.0(摘要写 +3.4)、CMMLU +4.0、BBH +5.0、ARC-C +3.7、HumanEval +3.0、MATH +2.4,Multi-Query NIAH 从 84.2 到 97.0。

分多少给表合适,第 5 步给了两份报告的扫描结果。

↳ 下一步:每个 token 只读 8 行。这 8 个行号怎么来,为什么不用跑模型就能算出来?

② 哈希:行号只看最近几个 token

下面是 FineWeb-Edu val 里的一段 24 个 BPE token。点一个 token(橙色), 以它结尾的 2-gram 标绿、3-gram 再往前多一个标蓝。4 个 2-gram 头和 4 个 3-gram 头各有一组随机奇数乘子: 每个 token id 乘上对应的乘子 → 异或到一起 → 对表长取模,得到行号。开头不足 n 个 token 时,前面用 0 号 token 补齐。这一步叫 n-gram 哈希。

表长:

头	n-gram	token id × 乘子,再异或	异或结果 H	行号 = H mod 表长	在表里的位置

撞行表太小,不同 n-gram 挤在同一行

脚本在训练集前 20 万个 token 里统计:出现过 111,302 种 2-gram、170,212 种 3-gram;每个头上,有多少比例的 n-gram 和别的 n-gram 共用一行。三档表长都来自真实运行(32768 行那组取自吞吐测量)。黑色短线是"把这么多 n-gram 均匀随机撒进这么多行"的理论值。点一组:

2-gram 撞行比例(4 个头平均) 3-gram 撞行比例(4 个头平均) 均匀随机撒点的理论值

–

2-gram 撞行

–

3-gram 撞行

–

平均每行挤几种 2-gram

↳ 代码:13_engram.py 的 hash_multipliers(固定种子 2026)· ngram_rows · 脚本中段的撞行统计 collision

为什么每种 n-gram 要用 4 个头,不用 1 个头配 4 倍的行?

两个 n-gram 在某个头上撞了行,它们取回的这一段 16 维是同一个向量。4 个头用 4 组不同的乘子,在一个头上撞了,在另一个头上多半撞不上; 拼起来的 128 维里只有一部分被共用,模型仍能区分这两个 n-gram。只用 1 个头时,一旦撞行,整根向量都相同。代码里 heads_per_order = 4;Engram-27B 原文 "the number of heads to 8";Qwen3.8-Flash-Next 的 config 里是 heads_per_ngram: 8。

16384 行时 2-gram 撞行为什么正好是 85%?

撞行比例 = 1 − 用到的不同行数 ÷ 不同 n-gram 数。每个头只有 16384 行,11 万种 2-gram 最多用满 16384 行, 所以比例至少是 1 − 16384 / 111,302 = 85.3%,实测 85.3%:行已经全部用满,平均每行 6.8 种 2-gram。 262144 行时行数多过 n-gram 种数,实测 18.4%,和均匀随机撒点的理论值 18.5% 一致,说明这个乘法加异或的哈希在这里分布均匀。

统计只用了训练集前 20 万个 token。整个 9500 万 token 的训练段里 n-gram 种数更多,真实撞行比例比这里高。

Engram 和 Qwen3.8 的哈希部分还多做了什么?

Engram(arXiv 2601.07372 §2.2):先做 tokenizer 压缩,用 NFKC 规范化和小写把 token 映射成规范 ID,128k 词表的有效大小减少 23%; 后缀 n-gram 最大 N = 3;每个阶 K 个头,哈希函数是 multiplicative-XOR,各头取回的向量拼接。
Qwen3.8-Flash-Next 的 HF config.json 里与此相关的字段:ngram_size: 3、heads_per_ngram: 8、ngram_vocab_size_base: 20000000。技术报告没有解释 20000000 具体指什么。报告 §2.3 原句:"short n-grams ending at each token serve as keys into embedding tables"。

↳ 下一步:撞了行,取回的向量就可能属于别的 n-gram。取出来的记忆和上下文不搭,怎么办?

③ 门控:这段记忆和上下文对不对得上

行号只看字面上的最近两三个 token,同一个 n-gram 在不同上下文里取回同一个向量,撞行时还混着别的 n-gram。所以记忆向量 e 先过一道门: 当前隐状态 h 当 query,e 经 w_k 投影后当 key,两边各做 RMSNorm,点积除以 √128,过 sigmoid 得到一个 0~1 的数。门值乘上 w_v · e,经短卷积后加回残差流。下面每格是一个 token,颜色越深门越开。悬停或点击看数值,切换模型:

α_t = σ( RMSNorm(h_t) · RMSNorm(W_k e_t) / √128 )
x ← x + Conv1D_k=4,因果( α_t · W_v e_t ) (只在第 2 个 Block 前做一次)

门值 0 1黑框 = 选中的 token

–

选中的 token

–

门值 α

–

这一段的平均门值

–

门值 < 0.1 的位置

–

门值 > 0.9 的位置

↳ 代码:13_engram.py 的 NgramMemory.forward(gate = sigmoid(...))· Model.forward(i == inject_layer 时加一次)

门的公式和 Engram 原文一样吗?注入在哪一层?

门的形式一样。Engram §2.3 原句:"we utilize the current hidden state h_t … as a dynamic Query, while the retrieved memory e_t serves as the source for both Key and Value projections", 式 (3)(4) 为 α_t = σ(RMSNorm(h_t)ᵀ RMSNorm(k_t)/√d),α_t 是标量。之后 Engram 接 kernel 4、dilation 3 的因果 Conv1D;本项目 kernel 4、不加 dilation。

注入位置按各自原文写:
· 本项目 13_engram.py:inject_layer = 1,注释"0 起数;1 = 第 2 层",即第 2 个 Block 之前。
· Engram-27B:原文 "we instantiate the module at layers 2 and 15";单层消融中 Layer 2 最好。
· Qwen3.8-Flash-Next:技术报告写只有一层,放在 Layer 2(Fig.1 与 §2.3.1);config.json 为 ple_layer_ids: [2],ple_conv_kernel_size: 4。

16384 行那组的门,为什么大多贴着 0 或 1?

先看数:这 64 个位置里,16384 行那组门值 < 0.1 的有 25 个、> 0.9 的有 15 个;262144 行那组分别是 1 个和 7 个,多数落在中间。一种解释:16384 行时 2-gram 撞行 85%,取回的向量常常混着别的 n-gram,对不少位置没有用处,门就把它关到接近 0;262144 行时撞行降到 18%,取回的内容更可信,门开得更均匀。这只是一段 64 个 token 的观察,没有在更多文本上统计,当作参考。

字符级那组的门几乎全开,说明什么?

tiny shakespeare 32768 行那组,64 个位置里 63 个门值 ≥ 0.999。字符级数据只有 1,101 种 2-gram、7,425 种 3-gram(训练集前 20 万字符), 表每个头 32768 行,2-gram 撞行只有 1% 到 2%,几乎每个 n-gram 独占一行;训练 5000 步把 100 万字符的训练集看了约 20 遍,表里存的是训练集上的统计。门在 val 文本上也全开,模型对表的依赖很重,而第 4 步会看到这组 val loss 比不挂表高 0.12。

↳ 下一步:机制齐了。真训一遍,val loss 能降多少?

④ 真实对照:FineWeb 主实验与字符级反例

两套数据各训 5000 步,batch 64 × 64 = 4096 个位置一步。FineWeb-Edu:GPT-2 BPE,训练段约 9500 万 token,5000 步随机抽到约 2050 万 token,多数文本只见一次。 tiny shakespeare:字符级,训练集约 100 万字符,5000 步相当于把训练集看约 20 遍。上图是最终 loss 随表大小的变化,下图是 val loss 曲线;点上图的点选一组:

val loss(2+3-gram) train loss(2+3-gram) 只查 2-gram(实心 val · 空心 train)

–

选中的一组

–

表参数 / 主干参数

–

最终 train / val

–

val 比不挂表

–

后 5 个评估点 val 平均比不挂表

–

2-gram 撞行

记忆表要配数据量远大于表参数、不反复看同一批文本的训练才划算;数据少到被来回看很多遍时,表会把训练集背下来。

↳ 跑法:python 13_engram.py --table 262144 --json runs/ch15_fw_262144.json · --orders 2 只查 2-gram · --data shakespeare 换字符级

FineWeb 上降 0.013 和 0.06,这两个数能信几分?

每组只训了一次(同一个随机种子 1337),评估时 train / val 各取 50 个 batch。
· 0 → 16384 行,最终 val 降 0.013,接近噪声量级;看后 5 个评估点(3000 到 4999 步),每点降 0.013 到 0.038,平均 0.026,点与点之间的起伏和差值本身差不多大。
· 0 → 262144 行,最终降 0.062;后 5 个评估点每点降 0.036 到 0.066,平均 0.053,方向在这 2000 步里一直保持。
· 只查 2-gram(4.9898)和 2+3-gram(4.9945)差 0.005。后 5 个评估点里只查 2-gram 那组每点低 0.005 到 0.007,但这些点出自同一次训练,不是独立的重复,换个种子可能翻过来。另外两组表参数一样多:只查 2-gram 时是 4 个头 × 每行 32 维,2-gram 分到的维数翻了一倍。这个规模上的结果不能用来说明 3-gram 没用。

要把 0.01 量级的差距说实,需要多个种子取均值和误差条。本页没有做。

5000 步、loss 5.0,这个差距放到训练完的模型上还成立吗?

不知道。FineWeb 这几组在 5000 步时 loss 还在 5.0 左右往下走(第 6 章的 124M 训完是 3 左右),属于训练早期;这是一组小规模对照,只能说明在这个步数、这个规模上表带来的变化,不代表收敛后的差距。表越大、每行被更新的次数越少,早期和后期的差距可能变大也可能变小,需要更长的训练才能回答。

字符级那组为什么表越大 val 越差?

典型的过拟合。训练集约 100 万字符,32768 行的表有 4,194,304 个参数,是主干 776,385 的 5.4 倍,也比训练集的字符数多;5000 步里每个字符平均被看约 20 遍。字符级 n-gram 种类少(前 20 万字符里 2-gram 1,101 种),每个 n-gram 几乎独占一行,表可以把训练集里每个短语后面接什么直接存下来。

看曲线:32768 行那组 val 在 2500 步到最低 1.5657,之后一路涨到 1.6536,train 同期从 1.212 降到 1.0813。不挂表那组 val 一直在降,最终 1.5334。前 1000 步表是有用的:1000 步时 32768 行那组 val 1.5925,不挂表 1.654。

FineWeb 那组看过的 token 比表参数还少,为什么没有背下来?

5000 步抽到约 2050 万 token,262144 行的表有 3355 万参数,单看数量,看过的 token 少于表参数。差别在重复次数:FineWeb 训练段有 9500 万 token,随机抽窗口时同一段文本很少被抽中两次, 表的每一行每次更新面对的多是新文本;shakespeare 的每个字符被看了约 20 遍。

FineWeb 这组也有一个值得盯的数:262144 行时最终 train 4.9604、val 4.9945,train 比 val 低 0.034;不挂表时 train 5.0752、val 5.0561,train 反而比 val 高 0.019。单次评估有噪声,不据此下结论;训练更久、数据被反复看时,要留意 train 与 val 的差距。

↳ 下一步:262144 行的表已经是主干的 4.6 倍,放大到几百亿参数时显存装不下。这张表放哪?

⑤ 表可以不进 GPU

行号只取决于输入的 token id,get_batch 取数据时就在 CPU 上算好了,forward 还没开始就知道每个 token 要读哪 8 行。所以表可以留在主机内存里,只把取出来的 8 × 16 个数送进 GPU,这叫把表放主机内存、提前取行。选一种放法,看数据怎么走,以及单卡 RTX 4090 上 400 步训练的实测吞吐(STEP 4 的 loss 对照在 Mac 上跑,吞吐单独在 4090 上测):

–

token / 秒(400 步)

–

比表在 GPU

–

留在加速器上的表参数

真实工业版把表放在哪

✓
Qwen3.8-Flash-Next:125B 总参、6B 激活,另有 51B n-gram 表参数。技术报告摘要原句 "additional 51B parameters of n-gram embedding tables held off the accelerator";§2.3:表放主机内存,异步预取。
✓
Engram(arXiv 2601.07372 §3.1):固定稀疏参数预算时,原句 "reallocating roughly 20%–25% of the sparse parameter budget to Engram yields the best performance",loss 随分配比例呈 U 形。
✓
Qwen3.8 报告 Tab.8:固定总参数,loss 在 n-gram 占 25% 参数时最低,与 Engram 的结论一致;下游指标看不出改善。

↳ 代码:13_engram.py 的 run(先在表所在设备上 lookup,再 .to(device))· get_batch 里的 ngram_rows(x) · 吞吐:--max-iters 400 --eval-iters 1 --table-device cpu

表放 CPU 为什么慢了 43%?

RTX 4090 的显存和主机内存之间隔着 PCIe。表在 CPU 时,查表、反向时把梯度写回表的那些行、AdamW 更新 4,194,304 个表参数都改在 CPU 上算, 每步还要把取出的行拷进显存、把表的梯度拷回主机。实测 148,229 token/s,表在 GPU 时 262,066,慢 43.4%。这里主干只有 7.24M 参数,GPU 算一步很快,CPU 上的表更新和两次拷贝占的比例就大;主干越大,这部分开销占比越小(推断,本项目没有测更大的主干)。

三组各测一次,每组 400 步;表在 GPU 比不挂表(276,231)低 5.1%,是查表、门控和短卷积的开销。测量用的是 32768 行的表(4.19M 参数),不是主实验的 262144 行。脚本里查表和拷贝跟着 forward 同步进行;Qwen3.8 报告写的是异步预取,取行和拷贝与上一步的 GPU 计算并行。

"提前取"具体提前在哪?

训练时,下一个 batch 的 token 在读数据时就定了,行号可以和数据加载一起算,取行、拷贝可以和上一步的 GPU 计算并行。推理时,新 token 一采样出来,它结尾的 2-gram、3-gram 就确定了,不用等下一层的计算结果。 MoE 做不到这一点:选哪个专家要看路由器对当前隐状态的打分,得算到那一层才知道。 13_engram.py 只做了最简单的同步版本:run 里先查表,再 .to(device),没有实现异步预取。

Gemma 3n 的 PLE 也放 CPU,它是 n-gram 表吗?

两者的键不同。 PLE(Per-Layer Embeddings) 在 HF transformers 的 modeling_gemma3n.py 里用 embed_tokens_per_layer(input_ids) 按单个 token id 查表(vocab_size_per_layer_input 262144,每层 256 维),在每层经门控乘入。官方说明:"PLE allows a significant portion of these parameters (the embeddings associated with each layer) to be loaded and computed efficiently on the CPU",E2B / E4B 总参 5B / 8B,加速器上约 2B / 4B。

两者都是地址确定、可以放主机内存的稀疏表,差在键:PLE 用单个 token,n-gram 表用最近几个 token。 Qwen3.8 报告 §2.3 把 Gemma 3n 归为 embedding-based memory,原句 "N-gram embeddings further generalize unigram lookup by conditioning memory retrieval on local context rather than token identity alone"。

n-gram 表和 MoE 是什么关系?

两者都让总参数远大于每 token 的计算量,取参数的方式不同。 MoE 由路由器看当前隐状态打分,选出几个专家,再让 token 过这几个专家的矩阵乘,选谁要算到那一层才知道; n-gram 表按输入 token 的字面地址直接取向量,取出来只做一次门控和投影。

DeepSeek-V4 前 3 个 MoE 层用 Hash routing,按 token 哈希选专家,token 仍然要过专家做计算,和 n-gram 查表是两回事;V4 报告全文里也没有 Engram。 Engram 论文把查表记忆和 MoE 并列为两种稀疏方式,在同一个总参数预算里分配(第 1 步的 Engram-27B 对 MoE-27B)。

达标本关实测

🎉 N-gram 查表 · 通关

你给第 12 章的小模型挂上了一张 3355 万参数的 n-gram 表:每个 token 只读 8 行,行号在 forward 之前就能算出,表可以留在主机内存。FineWeb 上 val loss 降了约 0.06,字符级数据上也看到了表把训练集背下来的样子。下一章讲投机解码与 MTP。