Engram！《Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models》解读！

0. 这篇论文到底想解决什么

核心矛盾：MoE 通过“条件计算”把参数规模做大但不按比例增加 FLOPs；然而 Transformer 没有原生“知识查表”原语，导致很多“本质上是静态知识/固定搭配/局部依赖”的东西，也被迫用注意力+MLP 在多层里“算出来”，浪费深度与算力。论文把语言建模拆成两类异质子任务：

• 组合推理（compositional reasoning）：需要动态计算（MoE擅长）
• 知识检索（knowledge retrieval）：更像静态查表（传统 Transformer 没有）

因此作者提出一个新轴：Conditional Memory（条件记忆），与 MoE 的 Conditional Computation 互补，并用一个可工程化的实例 Engram 来实现。

1. Engram 是什么：把“静态模式”从“动态计算”里剥离出来

1.1 模块位置与总体流程（两阶段：检索 + 融合）

Engram 是插在 Transformer 某些层（不是每层）的一个模块：对每个 token 位置 t，先用局部 N-gram 作为 key 去 O(1) 查表拿到静态向量，再用当前 hidden state 做上下文门控融合，最后再走标准 Attention+MoE。

这点与你的解读一致：Engram 不是替代 attention，而是把一部分“可查表的东西”前置/分担掉，使 attention 更专注“全局上下文与推理”。

2. 关键技术点：论文里真正“站得住”的 4 个创新

2.1 Tokenizer Compression：先把 token ID 做“语义归并”

标准 tokenizer 为了可逆重建，常把表面不同但语义等价的 token（如 “Apple” vs “␣apple”）分配不同 ID，导致 N-gram key 空间被无意义稀释。论文做了一个词表投影 P: V → V’（满射），用 NFKC、lowercase 等规则把“等价文本”折叠成 canonical ID，从而提高“语义密度”。论文报告：对 128k tokenizer 有 23% 的有效词表规模缩减。

这一步很关键：否则 N-gram 记忆会把容量浪费在“大小写/空格/归一化差异”上，碰撞与稀疏都会更糟。

2.2 Sparse Retrieval：Multi-head hashing 的 N-gram 查表（O(1)）

直接给所有 N-gram 建表不可行，所以 Engram 用 hashing：

• 对每个 N-gram 阶数 n（论文默认关注 2/3-gram），使用 K 个独立 hash head，每个 head 映射到一个素数大小的 embedding table（降低碰撞）。
• hash 实现为“轻量 multiplicative-XOR”，最后把各 n、各 head 的向量 concat 得到记忆向量 (e_t)。

你文中“多头像布隆过滤器/集成降低碰撞”的直觉是对的；但论文更明确强调：multi-head + prime table size 是为了工程上可控地压低碰撞与噪声，同时保持检索常数开销。

2.3 Context-aware gating：用 hidden state 做 Query，决定“记忆要不要信”

检索到的 (e_t) 是“上下文无关的先验”，会有两类噪声：hash 碰撞、以及多义性（polysemy）。论文做法是：

• (h_t)（已经经过前面 attention 聚合了全局上下文）当 Query
• (e_t) 经投影得到 Key/Value：(k_t=W_K e_t,\\ v_t=W_V e_t)
• 用 RMSNorm 稳定训练，然后 gate：(\\alpha_t = \\sigma(\\text{RMSNorm}(h_t)^T \\text{RMSNorm}(k_t)/\\sqrt{d}))
• 输出为 (\\tilde v_t = \\alpha_t \\cdot v_t)；若记忆与上下文冲突，gate 会压到接近 0。

注意：论文 gate 就是标准 sigmoid(dot) 形式（配 RMSNorm），并没有你解读里那种“sign*sqrt(|x|)”的特殊非线性（那段更像你/他人二创）。

2.4 Short depthwise causal conv：让“点查表”带一点局部融合能力

在 gated value 序列 (\\tilde V) 上，论文加了一个短的 depthwise causal Conv1D（kernel=4，dilation=最大 N-gram 阶），SiLU 激活，并残差相加：
(Y = \\text{SiLU}(\\text{Conv1D}(\\text{RMSNorm}(\\tilde V))) + \\tilde V)。

论文也提到：卷积的收益不如 multi-branch 融合、gating、tokenizer compression 那么大（消融里卷积去掉“只小幅退化”）。

3. 和 MoE 怎么配：论文最重要的结论其实是“稀疏预算分配律”

论文提出 Sparsity Allocation：在 iso-parameter & iso-FLOPs 约束下（总参一样、每 token 激活参一样），把“非激活参数预算” (P_{\\text{sparse}}) 在 MoE experts 和 Engram memory 之间怎么分才最优。

他们定义分配比 (\\rho)：

• (P_{\\text{MoE}}^{(\\text{sparse})}=\\rho P_{\\text{sparse}})
• (P_{\\text{Engram}}=(1-\\rho)P_{\\text{sparse}})

结果：出现稳定的 U 型曲线：

• 纯 MoE（(\\rho=1)）并非最优；
• 最优通常在 (\\rho \\approx 75%-80%)，也就是把大约 20%–25% 的稀疏预算挪给 Engram 更好；且在不同算力预算下位置稳定。

并且一个很“反直觉”的点：即使 MoE 专家数被砍到只剩 (\\rho\\approx 40%)，性能仍能接近纯 MoE baseline。

这比“Engram=加速知识问答”更关键：它说明 Engram 不是小修小补，而是一种可与 MoE 竞争稀疏预算的第一类建模原语。

4. 论文实证：为什么不仅提升知识题，还提升推理/代码/数学/长上下文

4.1 27B 主结果：不仅知识任务涨，推理/代码/数学涨得更“意外”

摘要与正文都强调：Engram-27B 相对 iso-parameter & iso-FLOPs 的 MoE-27B，提升不仅在 MMLU/CMMLU 等知识集，还在 BBH、ARC、DROP、HumanEval、GSM8K、MATH 等。

论文的解释路径是两条：

4.2 长上下文结果（LongPPL / RULER）：证明“attention 被解放”

在 LongPPL 与 RULER（32k）上，Engram-27B 在多个检索/跟踪任务显著强于 MoE-27B，例如 Multi-Query NIAH、Variable Tracking 等（论文给了例子对比）。

这部分对你后面“推理面建设 + KV cache 系统”最相关：Engram 的收益不只来自“算得更少”，还来自“注意力更值钱”。

5. 机制与系统：这篇论文非常工程化的两点（你解读里值得强化）

5.1 “确定性寻址”带来的系统红利：可预取、可分层缓存

Engram 的索引只依赖输入 token 序列，而不像 MoE 依赖运行时 hidden states 路由，因此可以做 prefetch-and-overlap：提前知道要取哪些 embedding，从 host memory（PCIe）异步拉取，并用前面若干层计算当“遮蔽窗口”，避免 GPU stall；同时 Engram 放在哪些层要同时满足“建模收益（更早更好）”与“系统遮蔽（更深更好）”的折中。

此外 N-gram 访问符合 Zipf 分布，可做 多级缓存层次：热 embedding 放 HBM/DRAM，长尾放 NVMe，实现“几乎不增加有效延迟的超大记忆”。

5.2 CPU offload 的实测开销：❤️% 的吞吐损失（H800 实验）

论文在 H800 上做了“+100B Engram（CPU offload）”的吞吐测试：

• 4B dense：9031 tok/s → 8858 tok/s
• 8B dense：6315 tok/s → 6140 tok/s
  吞吐损失非常小；论文也强调“更优化的 locality-aware 实现会更接近零损耗”。
  摘要里也给出“offload 100B table overhead ❤️%”的结论。

这点能把你的“MoonCake 这类 KV cache 系统”思路从“猜想”拉到“论文证据链”：大表不一定要驻留 HBM，只要寻址确定性 + 预取遮蔽做对。

6. 把你的解读升级成“更严谨、可落地”的推理系统改造思路（类 MoonCake）

你原文的“大方向”是：静态知识/固定搭配不要占 KV cache，不要靠 attention 重建。我建议把方案更明确地拆成三层（模型、运行时、数据/训练）：

6.1 模型侧：不要把 Engram 当成“另一个 attention”

论文的 Engram 融合方式是“检索 → gating → 短卷积 → residual”，然后再走 attention+MoE。
所以如果你在推理引擎里集成，建议遵循论文意图：

• Engram 是 memory primitive，主要覆盖局部静态模式（2/3-gram、实体、习语、公式化短语等），论文的 gating 可视化也证明确实会在这些位置强激活。
• attention 仍负责全局依赖、长程推理。

**工程含义：**你不该用一个全局可学习 α 把 Engram 与 attention 线性混合（那会模糊“查表 vs 推理”的边界）；更应把它作为一个独立 residual 分支，保持“能被 gate 关掉”。（这是论文的结构。）

6.2 运行时侧：把“确定性索引”做成一等公民（预取、批处理、缓存分层）

结合论文 Section 2.5，你可以把推理 runtime 做成下面这样：

(A) 预取与计算重叠（prefetch-and-overlap）

• 在进入某个 Engram layer 之前，基于 input token 序列预先算好 hash IDs（这部分纯 CPU 也行）。
• 异步从 host 拉取 embedding rows；用前序 Transformer blocks 的计算当“遮蔽窗口”。

(B) 多级缓存（HBM/DRAM/NVMe）

• 利用 Zipf 分布，把热 key 常驻 HBM（或 pinned DRAM），长尾落 NVMe；你的 scheduler 目标不是“所有都快”，而是“热路径极快、长尾可接受”。

© 批内去重与通信压缩

• 同一 batch 内 N-gram hash 很多会重复（尤其热模式），可以做“unique IDs → gather → scatter”，减少 PCIe/网络搬运。

这些是比“减少 KV cache 传输”更直接的系统抓手：KV cache 是 per-request、per-step 增长的；Engram 是“可缓存、可复用、确定性索引”的。

6.3 数据/训练侧：别忽视 tokenizer compression 与早期层插入

论文消融表明：tokenizer compression、context-aware gating、多分支适配是最关键的几个组件。
另外 Engram 插入层位要折中“建模偏好早插”与“系统偏好深插（遮蔽窗口）”。

工程建议（更像论文的口径）：

• 先在少量层插（论文强调不是每层都用）。
• 层位选择要让 runtime 能稳定遮蔽 PCIe 延迟（例如把第一个 Engram 放在“前面已有足够算力窗口”的位置，而不是第 1 层就硬取 host 数据）。

7. 你解读里哪些点需要“降级为推测/待验证”

为避免以后你写对外文章或做方案评审被抓住漏洞，下面这些你原文里的表述建议改写为“推测/示例”，不要当成论文结论：