DeepSeek：条件记忆 Engram 学习笔记

Hamburger_zzz
68 阅读6分钟

论文标题： Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models

时间：2026-1-12

论文链接

背景与动机

  • 目前的 Transformer 模型（包括主流的 MoE 架构）缺乏原生的“知识查找”机制
  • 知识即权重：“知识”被打散“溶解”在了模型几百亿个参数（Weights）的数值关系里，模型被迫通过大量的神经计算来低效地模 拟记忆和检索功能
  • 将算力浪费在还原一个简单的静态事实

核心结论与贡献

  • 加入“条件记忆”后，模型在知识密集型（1.8%~4%）、一般推理（3.3%~5.0%）、代码（1.6%~3%）、数学任务（2.2%~2.4%）上带来全面提升，仅付出 < 3%的吞吐量损失
  • Scaling law没有失效（更深的模型、更多的参数、更大的内存带来更好的效果）

Engram 架构设计

Engram 目的：通过 O(1) 的查表操作来辅助 Transformer 骨干网络，从而减少模型为了记忆静态知识而消耗的计算资源

Engram模块插入位置

**位置：**两层Transformer Block之间

稀疏插入： Engram 仅应用于特定的层（例如 27B 模型中是在第 2 层和第 15 层）。

  • 早期插入的意义： 放在浅层（如第 2 层）是为了尽早卸载局部的模式重构任务，让骨干网络能在后续层中专注于更复杂的推理 。
  • 系统设计： 这种确定性的查表机制允许系统利用层间计算的时间差，从 CPU 内存中预取（Prefetch）巨大的嵌入表，实现计算与通信的重叠

整体处理流程

一阶段：前序处理

**输入：**原始句子

数据经过前序标准的 Transformer Block

**输出：**隐藏层向量

二阶段：Engram介入

**输入：**前序 Transformer Block 输出的隐藏向量(input hidden)、原始句子

** 步骤一：稀疏检索 (Retrieval Phase) - "查字典" **

**目的：**找出当前语境下的静态知识

  1. Suffix N-gram Extraction
    1. 以当前词 "Great" 为结尾，提取不同长度（2和3）的后缀
  2. Tokenizer Compression（架构图中未体现）
    1. 输入：原始 Token ID 序列
    2. 词表投影，将语义相同的 Token（如 "Apple" 和 "apple"）映射为规范 ID
    3. 提高语义密度
    4. 输出：规范化 ID 序列
  3. 多头哈希与查表 (Multi-Head Hashing & Lookup)
    1. 用h个不同的哈希函数分别计算 N-gram 的哈希值
    2. 拿这些哈希值分别到对应的静态嵌入表（Embedding Table）中取出h个向量
  4. 拼接
    1. 将所有头（Head）和所有阶（2-gram, 3-gram...）查到的向量拼接成一个长向量ete_t

**步骤二： 上下文感知门控 (Gating Phase) **

**目的： **决定查出来的记忆是否有用

  1. 准备 Query, Key, Value
    1. Query：来自上一层Transformer的 input hidden
    2. Key、Value：查表得到的拼接向量ete_t经线性投影得到kt=WKet,vt=WVet\mathbf{k}_t = \mathbf{W}_K \mathbf{e}_t, \quad \mathbf{v}_t = \mathbf{W}_V \mathbf{e}_t
  2. 计算门控值 (Scaled Dot Product)
    1. 门控值αt=σ(RMSNorm(ht)⊤RMSNorm(kt)d)\alpha_t = \sigma \left( \frac{\text{RMSNorm}(\mathbf{h}_t)^\top \text{RMSNorm}(\mathbf{k}_t)}{\sqrt{d}} \right)
    2. αt∈(0,1)\alpha_t \in (0, 1)，内容越相关越接近1
  3. 加权与微调
    1. 加权：v~t=αt⋅vt\tilde{\mathbf{v}}_t = \alpha_t \cdot \mathbf{v}_t
    2. 微调：Y=SiLU(Conv1D(RMSNorm(V~)))+V~\mathbf{Y} = \text{SiLU}(\text{Conv1D}(\text{RMSNorm}(\tilde{\mathbf{V}}))) + \tilde{\mathbf{V}}
    3. 其中，一维卷积增加非线性
  4. 残差连接将“记忆信息”回传
    1. H(ℓ)←H(ℓ)+Y\mathbf{H}^{(\ell)} \leftarrow \mathbf{H}^{(\ell)} + \mathbf{Y}

三阶段：后续计算

利用融合了记忆的H(ℓ)\mathbf{H}^{(\ell)}进行后续的自注意力、MoE计算，以及后续的传递。

嵌入表怎么得来的？

  • 嵌入表本质上是一个巨大的、随机初始化的参数矩阵。 是模型参数的一部分，在模型预训练过程中进行端到端的联合优化
  • 嵌入表中包含的是 N-Gram 对应的向量表示，是从零开始，随着模型的训练过程，通过梯度下降一步步学习得到的

训练阶段： 利用分片与通信实现线性扩展

巨大的嵌入表被切分并分散存储在所有参与训练的 GPU 上

在需要时采用 All-to-All 通信原语取回

优势： 模型可存储的总记忆容量与 GPU 的数量成正比（线性扩展）

推理阶段： 利用预取实现存算分离

  • Engram属于确定性寻址，查表索引仅仅依赖于输入的 Token ID 序列
  • prompt 一进来就能开始计算并指导后续需要用到哪些嵌入表中的向量

预取：

  • 完整的 嵌入表 存于SSD或主机内存中
  • 算出后续计算需要的向量（ID）后，由 CPU 控制传递给指定GPU。此时，GPU正在进行前几层 Transformer Block 的计算

实验

稀疏分配定律

结论：稀疏参数配比 Engram : MoE = 0.20~0.25

稀疏参数：

Psparse≜Ptot−PactP_{\text{sparse}} \triangleq P_{\text{tot}} - P_{\text{act}}

  • PtotP_{\text{tot}}：模型总参数量
  • PactP_{\text{act}}：单次计算中激活的参数量
  • PsparseP_{\text{sparse}}：躺在显存中未参与计算的参数量

MoE 与 Engram 都会消耗PsparseP_{\text{sparse}}，但不会增加PactP_{\text{act}}

问题：存量的存储空间，怎么给 MoE 与 Engram 分配？

  • MoE 用稀疏参数来存储更多的处理逻辑
  • Engram 用稀疏参数来存储更多的静态知识

ρ=MoE:Engram\rho = MoE : Engram

2e20 FLOPs ：5.7B参数的模型

6e20 FLOPs ：9.9B参数的模型

无限内存潜力：不断扩大 Engram 所占内存，模型效果持续提升

OverEncoding：另一种哈希方法

Engram 提升模型“有效深度”

(a)： 预测收敛速度分析

  • KL散度数值越低，代表当前层输出与“最终输出”越接近
  • 同样深度，Engram模型比纯MoE模型收敛更快

(b)、(c)： 层级表示对齐分析 （Representational Alignment）

  • 相似度热力图 ，颜色越亮（黄色）表示两个层提取的特征越相似
  • Engram 模型浅层学到 MoE 模型深层特征

未来展望

记忆容量有待更多训练数据填充

现象： 虽然 Engram-40B 增加了大量记忆参数（从 5.7B 增加到 18.5B），但在部分任务上并没有完全碾压 27B 版本

**原因分析：**目前的训练数据量（262B Tokens）还不足以填满这 18.5B 的巨大记忆容量

意味着 Engram 的记忆容量还有巨大的潜力，随着训练 Token 数量的增加，更大容量的 Engram 会表现出更强的优势

多级缓存层次结构

  • L1 (热数据)： 将最常访问的 N-gram 嵌入保留在 GPU HBM（显存）中
  • L2 (温数据)： 将次常访问的放在 Host DRAM（主存）中
  • L3 (冷数据)： 将长尾的、罕见的知识（如生僻的百科条目）放在 NVMe SSD 上

意味着未来的模型可以挂载 TB 级别甚至 PB 级别 的超大外存作为知识库，而不仅仅局限于内存大小，真正实现近乎无限的参数扩展能力

avatar
文章
阅读
粉丝