PageGram: Multi-Head N-gram Memory at the Read Cost of a Flat Table

发表时间: 2026-10 · Tech report by Empero AI (empero.org)

原文: https://github.com/empero-org/pagegram/blob/main/paper.pdf

作者/机构：Leon Lehmann (Empero AI Research, [email protected], 通讯: [email protected], 代码: http://github.com/empero-org/pagegram)

1. 主要贡献

哈希 N-gram 表为语言模型提供了数十亿个参数，这些参数只占用内存而不增加计算开销。当前最优秀的架构设计是将每个 N-gram 哈希到多个独立的头上。在终端设备（如手机或笔记本电脑）上，这类大型表通常无法常驻内存，只能存储在闪存（Flash Storage）中，而闪存是以 4 KiB 页面为基本单位进行读取的。独立的哈希头导致每个 token 产生多达 16 次分散的随机闪存读取，且内存页面缓存对这种访问模式几乎失效。

针对该硬件受限问题，本文提出了 PageGram。其核心理念与贡献如下：

• 页面局域化多头布局（Page-Local Layout）：PageGram 保留了多头机制，但将一个 N-gram 阶数下的所有头全部放置在同一个闪存物理页面内。无论设置多少个注意力头，处理一个 token 时仅需读取 2 个页面（针对 bigram 和 trigram 各读取 1 页），与传统单扁平表（Flat Table）的读取开销完全相同。
• 等参数量下的严格对比与性能收益：在 200M 参数量的受控基准测试中（多分支共享核心模型初始化与数据流），经过 60M token 训练后，PageGram 在全部 5 个随机种子上均优于扁平表（Held-out Loss 平均降低 0.067 nats）和 Engram 风格的多头寻址（平均降低 0.142 nats）；在 120M token 训练下，PageGram 保持与 Engram 风格多头相当的精度（仅相差 0.010 nats），并持续优于扁平表。
• 揭示相关冲突（Correlated Collisions）机制：论文发现不同多头布局在单头维度上的哈希冲突概率完全一致，模型表现的关键杠杆在于多头之间哈希冲突的相关性。通过单参数——单页内每头槽位数（$S$）即可精准调节该相关性，实验确定 $S = 4$ 是综合收敛速度与训练稳定性最佳的设置。
• 端到端部署与极低的闪存开销：PageGram 支持即训即存的 4-bit NF4 量化，量化带来的精度损失仅为 0.0004 nats。在 1 GiB 内存页面缓存下，针对 20B 参数量的表，PageGram 每 token 仅发生 0.99 次真实闪存读取，而 Engram 风格寻址需要 9.18 次读取；在 NVMe 硬件实测中，吞吐量提升了 5 到 8 倍。

图 1：PageGram 在匹配或超越标准表性能的同时，将闪存读取量降低了 8 倍。a：无外部表、扁平哈希表、Engram 风格多头寻址与 PageGram（均为 200M 参数表，5 个种子的平均值）在 60M 训练 token 内的 Held-out Loss 曲线。b：各模型相对于扁平表的最终 Loss 差异，按种子成对展示（圆点代表单个种子）。c：在 20M 测试 token 上，20B 参数量、4-bit 量化的表在不同 RAM 页面缓存大小下的每 token 随机闪存读取次数。

2. 背景知识 / 关键 Observation / 设计原则

2.1 哈希 N-gram 记忆网络（Hashed N-gram Memory）

• 上下文特征提取与查表：对于当前 token $x_t$ 及其前序 token $x_{t-1}, x_{t-2}$，阶数 $n \in \{2, 3\}$ 的记忆模块分别将二元组 $(x_{t-1}, x_t)$ 和三元组 $(x_{t-2}, x_{t-1}, x_t)$ 哈希映射到参数表 $T \in \mathbb{R}^{R \times w}$ 中。

• 扁平表计算机制：扁平表为每个 N-gram 阶数读取单行（行宽 $w$ 等于模型隐藏层维度 $d$），将两行向量分别乘以哈希生成的伪随机符号向量 $s_n \in \{-1, +1\}^w$ 以减轻不同 N-gram 冲突时的特征相互干扰，然后进行相加累加：
  

  $$e_t = \sum_n s_n \odot T[h_n(x_{t-n+1:t})]$$

• 残差流注入与稀疏更新：所提取的向量 $e_t$ 通过可学习的门控机制注入到 Transformer 特定层的前向残差流中。参数表初始全部为零，确保初始步时不干扰基础模型的计算；训练过程是严格稀疏的，单批次内只有被触及的行才会发生梯度更新【Jordan 等人提出 modded-nanogpt 极速训练 baseline (2024, GitHub 地址: https://github.com/KellerJordan/modded-nanogpt)；将离散特征哈希到共享行的早期理论 见 Weinberger 等人 (2009, ICML) 以及 Svenstrup 等人提出的 Hash Embeddings (2017, NeurIPS)】。

2.2 多头寻址机制（Multi-Head Addressing）

• 完全冲突问题：在扁平表中，若两个不同的 N-gram 恰好映射到同一行，它们在模型中将完全不可区分。

• 多头解耦设计：为了将完全冲突转化为部分冲突，多头寻址机制在每个阶数上引入 $K$ 个独立的哈希函数，各自独立寻址行宽为 $r$、行数为 $P_{n, k}$ 的子表，并将提取的 $2K$ 个头向量拼接后经过一个可学习的线性投影矩阵 $W_V \in \mathbb{R}^{d \times 2Kr}$ 进行解码：
  

  $$e_t = W_V \Big[ s_{n, k} \odot T[\mathrm{base}_{n, k} + u_{n, k}(x) \bmod P_{n, k}] \Big]_{n, k}$$

• 相关工作实践：Over-encoding 采用类似的多头切分结构来扩展 N-gram 输入嵌入【Huang 等人 (2025, ICML)】；Engram 则在每个阶数下使用 $K = 8$ 个头，并为每个头配置互质的素数行数，同时使用基于上下文的门控结构 $\alpha_t = \sigma(\mathrm{RMSNorm}(h_t)^\top \mathrm{RMSNorm}(k_t) / \sqrt{d})$ 调制记忆输出【Cheng 等人 (2026, arXiv:2601.07372)】。

2.3 闪存硬件受限分析（Reading a Table from Flash）

• 大参数表的存储瓶颈：数十亿到数百亿参数的嵌入表远超智能手机或轻薄笔记本的 RAM 承载上限，必须驻留在闪存（Flash Storage/NVMe SSD）中按需读取【Alizadeh 等人提出的 LLM in a Flash (2024, ACL)；Yu 等人提出的 SCONE (2025, NeurIPS)】。
• 页面粒度读取开销：闪存的底层物理读取必须对齐到页面粒度（标准为 4 KiB），一次随机读取的硬件延迟达数十微秒，I/O 吞吐受限于设备每秒的 IOPS。
• I/O 放大对比：在 4-bit 量化（每行附带 1 个 16-bit 缩放因子）下，宽度为 64 的头向量仅占 34 字节，一个 4 KiB 页面可容纳 120 个头行。扁平表在 bigram 和 trigram 下共读取 2 个页面即可获取完整信息；但 Engram 风格的独立多头寻址在 $K=8$ 时，16 个头分散在不同的子表中，导致单个 token 必须发起 16 次离散的 4 KiB 闪存读取。由于每个加载到 RAM 的页面仅利用了其中 34 字节的数据，导致内存页面缓存（RAM Page Cache）的命中收益大打折扣。

3. 方法细节

3.1 内存布局（Layout）

• 页面内聚合多头：PageGram 继承了多头寻址中每个阶数划分 $K$ 个头的设定，但彻底重构了物理存储布局。对于每个阶数 $n$，整个参数表被切分为 $P_n$ 个闪存页面，每个物理页面统一划分容纳 $K \cdot S$ 个头行（即每个头分配 $S$ 个局部槽位）。

• 两级分层哈希公式：一个 N-gram 序列 $x$ 首先通过其主哈希函数确定唯一的页面索引，各头再利用局部哈希函数在该页面内选择对应的槽位（Slot）：
  

  $$\mathrm{row}_{n, k}(x) = \mathrm{base}_n + \left( p_n(x) K + k \right) S + \left( u_{n, k}(x) \bmod S \right)$$
  $$p_n(x) = b_n(x) \bmod P_n$$
  其中 $b_n$ 为该 N-gram 的页面哈希函数，$u_{n, k}$ 为头 $k$ 的专用哈希函数。

• I/O 压缩与模型集成：由于一个阶数下的所有 $K$ 个头全部被打包在同一页面内部，因此处理 bigram 和 trigram 依然总共只触发 2 次闪存页面读取。读出的 $2K$ 个头向量与随机符号掩码相乘后拼接，经由投影矩阵 $W_V$ 映射回模型维度 $d$，并通过每层独立的可学习标量门控 $g_j$ 注入 Transformer 残差流：$x \gets x + g_j e_t$。

图 2：针对单阶 N-gram 的行存储布局示意图（虚线代表 4 KiB 闪存页面）。a：扁平表单次读取一行宽度为 $d$ 的向量。b：Engram 风格寻址中每个头独立哈希到各自的子表中，导致头向量离散分布在不同页面，需要 8 次独立读取。c：PageGram 先通过哈希定位到一个物理页面，各头从该页内的 $S$ 个槽位中挑选一行（图中 $S=4$）。在结合二元和三元阶数时，单个 token 的页面读取次数分别为 2 次、16 次与 2 次。

3.2 槽位数（$S$）对冲突相关性的调控机制（What the Slots Control）

• 单头碰撞率的不变性：在固定总参数量的前提下，Engram 风格中各头的行数为 $P_{n, k}$，而 PageGram 中各头的总行数为 $P_n S$。两者每头对应的行总数完全相同，因此任意两个 N-gram 在单头维度上发生碰撞的绝对概率是严格一致的。

• 跨头碰撞相关性推导：参数 $S$ 本质上控制了跨头的碰撞耦合度。任意两个被散列到同一页面的 N-gram，在某一个头发生碰撞的概率独立为 $1/S$。当已知它们在某一个头发生碰撞时，它们在其余 $K-1$ 个头中发生碰撞的期望数量为：
  

  $$\mathbb{E}[\text{额外重合头数}] = \frac{K - 1}{S}$$

  
  而在 Engram 风格寻址中，两个 N-gram 在第二头发生重合的概率极其微小，仅为 $1/P_{n, k}$（约数万分之一）。

• $S$ 谱系的连续插值特性：

  • 当 $S = 1$ 时，一个页面为每个头仅保留 1 个槽位，落入同页的 N-gram 将在所有 $K$ 个头上发生 100% 重叠，退化为类似扁平表的全局碰撞模式。
  • 随着 $S$ 逐渐增大且总页数 $P_n$ 相应减少，多头之间的独立性增强；当极限退化到整个表仅有 1 个超大页面时，PageGram 等价于 Engram 风格的多头独立寻址。

3.3 训练优化与部署量化（Training and Storage）

• 初始化与行级 RMS 优化器：参数表全部以 0 初始化。采用行级自适应 RMS 更新算法：每行维持独立的梯度平方移动平均值 $\hat{v}$，跳过的步数直接折算为权重衰减重放（Replayed Decay），参数单步移动量为 $\eta_T g / \sqrt{\hat{v}}$【遵循 Jordan 等人的 modded-nanogpt (2024)】。经实验超参调优，表学习率 $\eta_T$ 设定为致密模型主体学习率的 70 倍（详见附录 B）。
• 正交解码投影与模型隔离：线性解码矩阵 $W_V$ 采用增益为 $\sqrt{2d / (2Kr)}$ 的正交初始化。$W_V$ 与门控标量的初始化完全独立于 Transformer 主干网络的基础随机数发生器，确保各对比实验分支在 step 0 具有位级精确一致的初始函数映射。

• NF4 部署存储对齐：在线上部署时，将维度为 64 的各头向量压缩为 NF4 格式【Dettmers 等人提出的 QLoRA (2023, NeurIPS)】，每个头行包含 64 个 4-bit 权重及 1 个 FP16 缩放因子，共 34 字节。

  • 在 $S = 15$ 配置下，一个页面包含 $8 \times 15 = 120$ 个头行，占用 $120 \times 34 = 4080$ 字节，几乎完美填满 4 KiB（4096 字节）块。
  • 在 $S = 4$ 配置下，每个页面占用 $8 \times 4 \times 34 = 1088$ 字节，物理存储上将 3 个页面紧凑拼接为一个 4 KiB 块。虽然这种对齐使得物理文件存储冗余增加了约 25%，但每次 token 依然只需读取 2 个物理块。

• 预取能力：因为寻址地址完全由离散 token id 的整型哈希决定，在模型进行网络前向计算之前即可提前计算出物理偏移量并向闪存发出异步读取请求。

4. 实验环境与设置

• 数据集：包含 228M token 的多语言预训练混合语料，覆盖网络通用文本、参考与教育资料、代码、数学及少量对话；语言包含英语、中文、德语、日语和俄语。训练过程按 1,024 长度的随机滑动窗口切分，单个训练步包含 131,072 个 token。固定基准训练 457 步（约 60M token）；长程训练跑 915 步（约 120M token）。验证集取自独立的 Held-out Split，固定为 2,048 个窗口（共 2.1M token）。

• 模型架构：

  • 基础模型：6 层 Pre-norm Transformer Decoder，隐藏层宽度 $d = 384$；6 个注意力头（每头维度 64），采用 RoPE 旋转位置编码【Su 等人 (2021, arXiv:2104.09864)】；采用 RMSNorm【Zhang and Sennrich (2019, NeurIPS)】；MLP 采用 GELU 激活函数，中间维度 $4d$；词表大小 65,536，输入与输出词嵌入不解耦（Untied Embeddings）。模型总参数量 61M（其中 10.6M 位于 Transformer 块内部）。
  • 记忆注入位置：记忆向量分别注入到第 0 层与第 4 层注意力块之前的残差流中，各层配置初始值为 0.5 的可学习标量门控。
  • 放大模型（Larger Model）：用于测试泛化性，隐藏层维度扩增为 $d = 512$，共 8 层 Transformer 块，8 个注意力头（总参数量 92M）。

• 外部参数表配置（统一固定约 200M 参数量，相对偏差小于 0.1%）：

  • 扁平表（Flat Table）：520,832 行，行宽 384，均匀切分为 bigram 和 trigram 两半，使用哈希符号向量。
  • Engram 风格表：每个阶数 8 个头，共计 16 个互质素数行数的独立子表，每张子表行数在 195,300 附近，行宽 64。
  • PageGram 表：8 个头，行宽 64，评估槽位 $S \in \{1, 4, 8, 15, 60\}$。以 $S = 4$ 为例，bigram 和 trigram 阶数下的页面总数分别设计为 48,823 和 48,821 个素数页。
  • 解码矩阵：多头变体均使用 $W_V \in \mathbb{R}^{384 \times 1024}$（约 0.39M 参数）。

• 硬件与软件配置：

  • 硬件平台：单张 NVIDIA RTX 5090 GPU（实验中通常双任务并发运行）。闪存测试使用机器内置的高性能 NVMe SSD。
  • 软件栈与优化器：主模型优化器使用 AdamW【Loshchilov and Hutter (2019, ICLR)】，学习率 $2 \times 10^{-3}$，$\beta = (0.9, 0.95)$，矩阵权重衰减设为 0.1。采用 Warmup-Stable-Decay (WSD) 调度策略【Hu 等人 (2024, arXiv:2404.06395)；Hägele 等人 (2024, NeurIPS)】，包含 5% 线性预热，稳定期后在最后 20% 训练步线性衰减至 0。模型采用 PyTorch BF16 混合精度，参数表保持 FP32 训练并在导出时转为 NF4。所有分支的底层随机数种子在模型初始化和数据加载顺序上严格受控对齐。

5. 实验结果与深度分析

5.1 60M Token 下相同表容量的性能表现

数据出处：论文 Table 1 与 Fig. 1

• 核心分析：
  • 在 60M token 训练节点上，所有记忆表均极大地提升了小模型的建模能力。PageGram ($S=4$) 表现最为卓越，在全部 5 个种子上一致性地击败扁平表（平均领先 0.067 nats，成对 $t$ 检验 $p = 0.08$，单侧符号检验 $p = 0.03$）。
  • 与 Engram 风格寻址相比，PageGram 领先幅度高达 0.142 nats ($p = 0.02$)。在训练初期，完全解耦的 Engram 风格寻址表现甚至劣于扁平表（落后 0.074 nats），证明多头完全独立在数据受限时会减缓收敛。
  • PageGram 的训练效率十分优异，其吞吐（108K tok/s）与无外部表模型（113K tok/s）相比仅有不足 5% 的损耗。

5.2 长期训练与更大模型表现

数据出处：论文 Table 2

• 收敛趋势解析：
  • 随着步数延伸至 120M token，基础 Transformer 主干逐渐学会了局部统计规律，使得外部记忆表的边际收益由 60M 时的 0.40 nats 下降至 0.095 nats。
  • 训练后期各架构排序发生重组：Engram 风格多头寻址反超了扁平表（3 个种子上平均领先 0.034 nats）；而 PageGram ($S=4$) 紧随其后，与 Engram 风格基本持平（差距仅 0.010 nats），同时平均指标依然领先扁平表 0.024 nats。这证明 PageGram 在维持极低闪存读取的同时，保留了多头架构在充分训练下的表征上限。

6. 机制分析：相关冲突（Correlated Collisions）

6.1 单头碰撞与跨头特征对比

数据出处：论文 Table 3

• 统计学等价性：在单头层面上，PageGram 与 Engram 风格寻址的哈希表现完全无异：方差与均值之比均为 1.0 附近（符合理想随机哈希分布），被其他不相干 N-gram 污染的步数比例（Foreign）在二元组为 66.9%、三元组为 89.4%。
• 跨头相关性断层：两者的唯一根本差异在于跨头重叠模式。在 Engram 风格下，头与头完全独立，在头 0 发生碰撞的两个 N-gram 在其他头再次碰撞的概率低于 0.01%；而在 PageGram ($S=15$) 下，其他头再次碰撞的概率高达 38.3%，平均额外共享 0.467 个头，该数值严格等于理论预期 $(K-1)/S = 7/15$。

6.2 槽位数 $S$ 的扫描与稳定性规律

数据出处：论文 Table 4 与 Fig. 3

图 3：单页槽位数 $S$ 控制冲突相关性。固定 200M 参数表，展示 60M token 训练后 Held-out Loss 随 $S$ 的变化曲线。实心圆点代表单种子，空心圆点代表在稳定性分析中出现劣化的特殊种子 1，实线横条为均值。浅黄色区域为扁平表的均值与方差带。

• 双重力量平衡假说：

  • 较大的 $S$ 允许各个头生成差异性更大的组合编码。
  • 较小的 $S$ 则将碰撞压力约束在极少数特定的“同页伙伴（Page-mates）”上（$S=4$ 时仅为 Engram 邻居数的 0.45 倍）。线性解码器能够更轻易地学习解耦这几个固定的伙伴，而不需要在成百上千个仅共享单行的无关 N-gram 的随机噪声中艰难寻求解耦。

• 稳定性与推荐配置：在 Seed 1 测试中，弱相关布局（$S \ge 8$ 及 Engram）从前 13M token 起便显著落后，最终停滞在 5.28 左右；而强相关布局（$S=1$ 和 $S=4$）在所有种子上均表现出极强的鲁棒性，无一例外地维持领先。因此，$S = 4$ 是兼顾最优收敛均值与 100% 稳定性的工程设计。

7. 消融实验细节

数据出处：论文 Table 5

• Engram 上下文门控在中小规模下的失效：引入复杂的动态投影门控网络反而在所有种子上使指标恶化了 0.140 nats，证明直接采用标量注入门控在该尺度下具有更佳的收敛表现。
• 参数量缩减代价：将表参数量从 200M 缩减到 50M 时，Loss 性能损失高达 0.206 nats，远大于架构改良带来的收益。这表明在 N-gram 记忆设计中，容量（Capacity）的优先级高于布局，而 PageGram 的价值正在于能以极低代价部署超大容量。

8. 闪存读取与物理 I/O 实测

数据出处：论文 Table 6 与 Fig. 4a

图 4：实际部署表的闪存开销评估。a：在 20M 测试 token 驱动下，针对 3.2B、10B 和 20B 规模的 4-bit 外部表，未命中 RAM LRU 页面缓存的每 token 随机闪存读取数（纵坐标在小于 1 时为线性刻度，大于 1 时为对数刻度）。b：针对 20B 参数表，在不同并发并发读取深度（In-flight Reads）下，使用直接 I/O 从真实 NVMe 驱动器实测的每秒查表吞吐量。

8.1 真实 NVMe 硬件吞吐实测（20B 参数表）

数据出处：论文 Table 6 与 Fig. 4b

• 硬件 I/O 深度解析：
  • 在无缓存场景下，PageGram 保持绝对的 2 次读取开销，比 Engram 风格的 16 次读取节省了 87.5% 的闪存 I/O。
  • 在 RAM 缓存受限场景下，由于 PageGram 缓存的页面能够同时服务多个头和多个 N-gram，其缓存命中效率显著优于离散的多头分布。以 20B 参数表搭配 1 GiB 缓存为例，PageGram 仅产生 0.99 次读取，而 Engram 风格依然高达 9.18 次读取。
  • 在底层 NVMe 驱动实测中，单队列深度下读取单块延迟为 $60\ \mu\mathrm{s}$。Engram 风格由于严重受到 I/O 吞吐瓶颈限制，在 64 并发下仅能提供 17K lookups/s；而 PageGram 飙升至接近 119K lookups/s。在预填充（Prefill）阶段需要瞬间读取大量 Prompt token 的场景中，PageGram 能彻底消除 I/O 导致的计算停顿。

9. 附录与补充细节

9.1 附录 A：超参数与实验配置汇总

• 注意力与主体配置：6 层 Pre-norm，隐藏维度 384，RMSNorm，6 头（每头 64 维），RoPE 位置编码，因果注意力掩码，上下文长度 1,024。GELU 激活，中间维度 1,536。词表大小 65,536。参数量：总计 61.0M，Transformer 块占 10.6M。
• 外部记忆注入方式：在第 0 层与第 4 层之前注入，每层配置可学习标量门控（初始值设置为 0.5）。
• 优化与精度策略：致密参数使用 AdamW，梯度裁剪范数为 1.0；混合精度采用 BF16 Autocast，参数表内部更新维持 FP32 精确度。稀疏表优化器采用行级自适应 RMS，$\beta_2 = 0.95$。

9.2 附录 B：外部表学习率网格扫描探针

数据出处：论文 Table 8

• 超参标定依据：作者使用 Seed 0 在 30M token 训练节点上进行了网格搜索，结果显示 70 倍的主干网络学习率（即 $\eta_T = 70 \times 2 \times 10^{-3} = 0.14$）在所有表中均取得了最优验证 Loss，因此被固定为整个实验体系的标准超参。

9.3 附录 C：全种子单次运行原始结果汇总

数据出处：论文 Table 9

9.4 附录 D 与 E：硬件模拟仿真与系统验证保障

• C 语言闪存仿真验证：仿真器采用纯 C 语言编写，逐 token 对齐模型产生的寻址逻辑，精准记录 LRU 页面缓存命中状态。
• 20 项自动化回归检查：
  1. 保证扁平表、Engram 与 PageGram 参数总数严格对齐。
  2. 验证所有哈希输出地址落在分配的内存区间内，断言 PageGram 各头行 100% 落在同页，防止跨越文档边界。
  3. 当表置零时，通过逐位对比确保带表模型与无表模型的初始输出 Loss 完全一致。
  4. 校验稀疏梯度计算与全量稠密张量梯度等价，验证重放权重衰减数学公式的准确性。
  5. 验证 4-bit 量化对全零行的保真度，确保量化截断误差满足理论界限。

10. 方法细节中的引用汇总

11. 结论

多头 N-gram 外部记忆模块通过解耦哈希头将严重的完全碰撞转化为局部碰撞，但在端侧落地时却因分散寻址而付出了每个头产生一次随机闪存读取的高昂代价。PageGram 打破了“多头必然离散”的思维定势，通过页面内槽位划分技术，将单个阶数下的所有多头行集中在同一个 4 KiB 物理页面中，成功将单 token 的闪存物理读取压减至与扁平表完全相同的 2 次。

理论分析与大量对照实验表明，多头哈希冲突的高度相关性非但没有损害模型质量，反而在中小规模预训练阶段提升了解码器的优化效率，实现了显著优于 Engram 风格和扁平表的收敛表现；而在长周期训练下，该结构依然能够无损继承多头机制的表征上限。配合 4-bit NF4 紧凑格式，PageGram 在实际 NVMe 硬件上提供了 5 至 8 倍的查表吞吐加速比，为大语言模型在端侧设备上高效调用超大规模稀疏记忆提供了实用可行的工程方案。