时间：2026-01-28 00:15
人气：
作者：admin
26年1月来自北大和DeepSeek AI的论文“Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models”。
混合专家模型 (MoE) 通过条件计算扩展容量，但 Transformer 模型缺乏用于知识查找的原生原语，迫使其通过计算来低效地模拟检索过程。为了解决这个问题，引入条件记忆作为补充的稀疏性轴，并通过 Engram 模块实现。Engram 模块对经典的 N-gram 嵌入进行现代化改造，实现 O(1) 查找。通过构建稀疏性分配问题，发现一个 U 形扩展规律，该规律优化在神经计算 (MoE) 和静态记忆 (Engram) 之间的权衡。在该规律的指导下，将 Engram 扩展到 270 亿个参数，其性能优于严格等参数、等浮点运算次数的 MoE 基线模型。最值得注意的是，尽管记忆模块预期能够提升知识检索能力（例如，MMLU +3.4；CMMLU +4.0），但其在通用推理（例如，BBH +5.0；ARC-Challenge +3.7）和代码/数学领域（HumanEval +3.0；MATH +2.4）的提升更为显著。机制分析表明，Engram 减轻骨干网络早期层静态重建的负担，有效加深网络以进行复杂推理。此外，通过将局部依赖关系委托给查找操作，它释放注意资源用于全局上下文，从而显著提升长上下文检索能力（例如，多查询 NIAH：84.2 → 97.0）。最后，Engram 实现基础设施-觉察的效率：其确定性寻址机制支持运行时从主机内存预取数据，且开销可忽略不计。
稀疏性是智能系统中反复出现的设计原则，从生物神经回路（Lennie，2003；Olshausen 和 Field，1997）到现代大语言模型（LLM）均有涉及。目前，这一原则主要通过专家混合模型（MoE）（Dai ，2024；Shazeer ，2017）来实现，该模型通过条件计算扩展容量。由于其能够在不成比例增加计算量的情况下大幅提升模型规模，MoE 已成为前沿模型的实际标准（Comanici，2025；Guo ，2025；Team，2025）。
尽管这种条件计算范式取得成功，但语言信号固有的异质性表明，其结构优化仍有很大的空间。具体而言，语言建模包含两个性质不同的子任务：组合推理和知识检索。前者需要深度动态计算，而文本的很大一部分——例如命名实体和程式化模式——是局部的、静态的，并且具有高度刻板的模式（Constant et al., 2017; Erman, 2000）。经典的N-gram模型（Brants et al., 2007; Liu et al., 2024b; Nguyen, 2024）能够有效地捕捉此类局部依赖关系，这意味着这些规律自然地可以表示为计算成本低廉的查找操作。由于标准的Transformer模型（Vaswani et al., 2017）缺乏原生的知识查找原语，当前的LLM只能通过计算来模拟检索过程。例如，解析一个常见的多词实体需要消耗多个早期注意层和前馈网络（Ghandeharioun et al., 2024; Jin et al., 2025）。这个过程本质上相当于在运行时耗费大量资源重建一个静态查找表，将宝贵的序列深度浪费在原本可以分配给更高层次推理的琐碎操作上。
为了使模型架构与这种语言二元性相一致，引入一个互补的稀疏性维度：条件记忆。条件计算稀疏地激活参数来处理动态逻辑（Bengio et al., 2013; Shazeer et al., 2017），而条件记忆则依赖于稀疏的查找操作来检索固定知识的静态嵌入。作为对这一范式的初步探索，重新审视 ????-gram 嵌入（Bojanowski，2017），并将其作为典型的实例：局部上下文作为K，通过常数时间 O(1) 的查找来索引庞大的嵌入表（Huang，2025a；Pagnoni，2025；Tito Svenstrup 等人，2017；Yu，2025）。研究表明，这种静态检索机制可以作为现代 MoE 架构的理想补充——但前提是它必须经过合理的设计。本文提出 Engram，一个基于经典 ????-gram 结构的条件记忆模块，但配备现代的改进技术，例如token化器压缩、多头哈希、上下文门控和多分支集成。
为了量化这两个基本模型之间的协同作用，提出稀疏性分配问题：在总参数预算固定的情况下，容量应该如何在 MoE 专家和 Engram 记忆之间分配？实验揭示一个独特的 U 形扩展规律，表明即使是简单的查找机制，当被视为一流的建模基本模型时，也能作为神经计算的重要补充。基于这一分配规律，将 Engram 扩展到一个 27B 参数的模型。与严格等参数、等浮点运算次数的 MoE 基线模型相比，Engram-27B 在各个领域都实现更高的效率。
通过 LogitLens（nostalgebraist，2020）和 CKA（Hendrycks，2021a）进行的机制分析揭示这些提升的来源：Engram减轻骨干网络在早期层重建静态知识的负担，从而增加可用于复杂推理的有效深度。此外，通过将本地依赖关系委托给查找操作，Engram 释放注意资源，使其能够专注于全局上下文，从而在长上下文场景中实现卓越的性能——在 LongPPL（Fang 2024）和 RULER（Hsieh 2024）测试中显著优于基线模型（例如，多查询 NIAH：97.0 vs. 84.2；变量跟踪：89.0 vs. 77.0）。
最后，将基础设施感知效率确立为首要原则。与 MoE 的动态路由不同，Engram 使用确定性 ID 来实现运行时预取，从而将通信与计算重叠。实验结果表明，将 100B 参数表卸载到主机内存只会产生可忽略不计的开销（< 3%）。
如图所示，Engram 是一个条件记忆模块，旨在通过在结构上将静态模式存储与动态计算分离来增强 Transformer 主干网。形式上，给定一个输入序列 ???? = (????_1,…,????_????) 和第 l 层的隐状态 H(l)，该模块分两个功能阶段处理每个位置 ????：检索和融合。首先，提取并压缩后缀 ????-gram，通过哈希确定性地检索静态嵌入向量。随后，这些检索的嵌入向量根据当前的隐状态进行动态调制，并通过轻量级卷积进行优化。最后，与多分支架构的集成。
第一阶段将局部上下文映射到静态内存条目，包括token化器压缩和通过确定性哈希检索词嵌入。
token化器压缩：虽然 N-gram 模型通常直接处理token化器的输出，但标准的子词token化器优先考虑无损重构，经常为语义等价的词项分配不相交的 ID（例如，Apple 与 ␣ apple）（Kudo & Richardson，2018；Li ，2023b）。为了最大化语义密度，实现一个词汇投影层。具体来说，预计算一个满射（surjective）函数 P : ???? → ????′，该函数基于规范化的文本等价性（使用 NFKC (Whistler，2025)、小写转换等）将原始词项 ID 折叠成规范标识符。实际上，对于 128k token化器，该过程可使有效词汇量减少 23%。形式上，对于位置 ???? 处的token，将其原始 ID ????_???? 映射到规范 ID ????_????′ = P(????_????)，以形成后缀 ????-gram ????_????,???? = (????′_t-n+1 , . . . , ????_????′)。
多头哈希。直接参数化所有可能 N-gram 的组合空间是难以处理的。借鉴 （Tito Svenstrup 2017) 的方法，采用一种基于哈希的方法。为了减少冲突，为每个 N-gram 阶数 n 使用 K 个不同的哈希头。每个哈希头 k 通过确定性函数 θ_n,k 将压缩上下文映射到嵌入表 E_n,k（素数大小对的 M_n,k ）中的一个索引。
实际上，????_????,???? 被实现为一个轻量级的乘法- XOR或哈希函数。通过连接所有检索的嵌入向量来构建最终的记忆向量。
检索的嵌入 e_???? 作为与上下文无关的先验信息。然而，由于它们是静态的，因此本质上缺乏上下文适应性，并且可能由于哈希冲突或多义性而受到噪声的影响（Haber & Poesio，2024）。为了增强表达能力并解决这种歧义，采用一种受注意机制启发的上下文-觉察门控机制（Bahdanau，2015；Vaswani，2017）。具体来说，利用当前隐状态 h_????（它通过先前的注意层聚合全局上下文）作为动态查询，而检索的记忆 e_???? 则作为K和V投影的源。
为了确保梯度稳定性（Dehghani，2023），在计算标量门 ????_???? ∈ (0, 1) 之前，对Q和 K应用 RMSNorm（Zhang & Sennrich，2019）。该设计强制执行语义对齐：如果检索的记忆 e_t 与当前上下文 h_t 相矛盾，则门控 ????_t 趋近于零，从而有效地抑制噪声。
最后，为了扩展感受野并增强模型的非线性，引入一个短的深度因果卷积（Gu et al., 2022; Peng et al., 2023）。令 Ṽ 表示门控值的序列。使用卷积核大小 ????（设置为 4）、空洞系数 ????（设置为最大 N-gram 阶数）和 SiLU 激活函数（Elfwing et al., 2018），计算最终输出 Y。
Engram 模块通过残差连接 H(l) ← H(l) + Y 集成到主干网络中，随后是标准的注意机制和 MoE 模块。关键在于，Engram 并非应用于每一层；其具体位置由系统级延迟约束决定。
本研究没有采用标准的单-流连接（He et al., 2016），而是采用先进的多分支架构作为默认骨干网络，因为该架构具有卓越的建模能力（Larsson et al., 2017; Szegedy et al., 2015; Xie et al., 2025; Zhu et al., 2025）。该架构的一个显著特征是将残差流扩展为 M 个并行分支，信息流由可学习的连接权重进行调节。
尽管 Engram 模块本质上与拓扑结构无关，但将其适配到这种多分支框架中需要进行结构优化，以平衡效率和表达能力。具体而言，实现一种参数共享策略：所有 M 个分支共享一个稀疏嵌入表和一个V投影矩阵 W_V，而 M 个不同的K投影矩阵 {W(m)_K} 用于实现分支特定的门控行为。对于隐状态为 h(m)_t 的第 m 个分支，计算分支特定的门控信号 ????(m)_t。然后，通过对共享V向量施加这些独立的门控来调制检索的记忆：u(m)_t = ????(m)_t · (W_V e_t)。这种设计允许将线性投影（一个 W_V 和 ???? 个不同的 W(????)_K）融合为一个稠密的 FP8 矩阵乘法，从而最大限度地利用现代 GPU 的计算能力。除非另有说明，所有实验均采用这种集成方法，并结合最近提出的流形约束超连接mHC (???? = 4) (Xie et al., 2025)。
扩展内存增强模型通常会受到GPU高带宽内存（HBM）容量的限制。然而，Engram的确定性检索机制天然支持参数存储与计算资源的解耦。与依赖运行时隐状态进行动态路由的MoE不同，Engram的检索索引仅依赖于输入token序列。这种可预测性使得训练和推理都能采用专门的优化策略，如图所示。
在图（a）训练过程中，为了适应大规模嵌入表，采用标准模型并行化，将嵌入表分片到可用的GPU上。用all-to-all通信原语在前向传播中收集活动行，并在反向传播中分发梯度，从而使总内存容量能够随加速器数量线性扩展。
在图（b）推理过程中，这种确定性使得预取和重叠策略成为可能。由于内存索引在前向传播之前已知，系统可以通过 PCIe 从主机的大量内存中异步检索嵌入。为了有效掩盖通信延迟，Engram 模块被放置在主干网的特定层，利用前几层的计算作为缓冲区来防止 GPU 停顿。这需要一种硬件-算法协同设计策略：虽然将 Engram 放置得更深可以扩展可用于隐藏延迟的计算窗口，但本文消融实验表明，建模性能更倾向于早期干预以卸载局部模式重建。因此，最佳放置位置必须同时满足建模和系统延迟约束。
此外，自然语言 N-gram 本质上遵循 Zipf 分布（Chao & Zipf，1950；Piantadosi，2014），其中一小部分模式占用绝大多数的内存访问。这种统计特性促使采用多级缓存层次结构：频繁访问的嵌入数据可以缓存在速度更快的存储层（例如，GPU HBM 或主机 DRAM）中，而访问频率较低的模式则存储在速度较慢但容量较大的存储介质（例如，NVMe SSD）中。这种分层结构使得 Engram 能够扩展到海量内存容量，同时将有效延迟的影响降至最低。
作为条件记忆的一种实例，Engram在结构上与MoE专家提供的条件计算互补。其研究这种对偶性的扩展特性以及如何最优地分配稀疏容量。具体而言，研究围绕两个关键问题展开：
计算匹配公式。用三个参数指标来分析这种权衡：
• ????_tot：可训练参数的总数，不包括词汇嵌入和语言模型头部。
• ????_act：每个token激活的参数数量。该值决定训练成本（FLOPs）。
• ????_sparse ≜ ????_tot − ????_act：非激活参数，代表可用于扩展模型规模而不增加计算成本的“自由”参数预算（例如，未选择的专家或未检索的词嵌入）。
在每个 FLOPs 预算范围内保持 ????_tot 和 ????_act 的值固定，以确保模型具有相同数量的参数和相同的每个token的 FLOPs。对于 MoE 模型，????_act 由前 k 个被选择的专家决定，而未被选择的专家的参数则计入 ????_sparse。对于 Engram 模型，每个token仅检索固定数量的槽，因此增加嵌入槽的数量会增加 ????_tot，而不会增加每个token的 FLOPs。
分配比例。将分配比例 ???? ∈ [0,1] 定义为分配给 MoE 专家容量的非活跃参数预算比例：
????(sparse)_MoE = ???? ????_sparse, ????_Engram = (1 − ????) ????_sparse.
直观地讲：
???? = 1 对应于纯 MoE 模型（所有非活跃参数均为路由专家）。
???? < 1 减少路由专家的数量，并将释放的参数重新分配给 Engram 嵌入槽。
实验协议。在两种计算方案下评估这种权衡，并在两种设置下保持稀疏度比率 ????_tot/????_act ≈ 10：
???? = 2 × 10²⁰ FLOPs：????_tot ≈ 5.7B，????_act = 5.68M。基线模型 (???? = 1) 共有 106 位专家。
???? = 6 × 10²⁰ FLOPs：????_tot ≈ 9.9B，????_act = 9.93M。基线模型 (???? = 1) 共有 99 位专家。
对于不同的 ???? 值，仅通过调整路由专家的数量和 Engram 嵌入槽的数量来构建相应的模型。所有运行均使用相同的训练流程和优化超参数。
前面已经优化固定参数预算下的内存分配。现在，将探索互补的设置：积极的内存扩展。这项研究的动机源于 Engram 独特的存储与计算解耦能力。
实验方案。用固定的 MoE 骨干网络，其中 P_tot ≈ 3B，P_act = 568M，并针对 1000 亿个tokens进行训练以确保收敛。在此骨干网络之上，附加一个 Engram 表，并将槽数 M 从 2.58 × 10⁵ 扫至 1.0 × 10⁷（总计约 130 亿个参数）。对于基线方法，将其与 OverEncoding（Huang，2025a）进行比较，后者通过与词汇嵌入平均来整合 N-gram 嵌入。虽然其他工作如 SCONE（Yu，2025）也研究大规模嵌入，但它主要侧重于推理，并且包含额外的模块（f-gram 模型）和额外的训练 FLOPs，使其与本研究严格的等计算约束不兼容。
基于提出的 Engram 架构和经验推导的分配规则，将 Engram 扩展到数十亿参数规模，以验证其在真实世界语言模型预训练中的有效性。具体而言，训练四个模型：（1）Dense-4B（总参数 41 亿），（2）MoE-27B（总参数 267 亿），（3）Engram-27B（总参数 267 亿），以及（4）Engram-40B（总参数 395 亿）。所有模型均使用相同的数据集（相同的token预算和顺序）进行训练，并且激活参数的数量严格匹配。
训练数据和模型配置。所有模型均在一个包含 2620 亿个tokens的语料库上进行预训练，用 DeepSeek-v3（Liu，2024a）的token化器，词汇量为 12.8 万。为了确保可控的比较，除非另有明确说明，否则所有模型均采用一致的默认设置。用一个包含 30 个模块的 Transformer 模型，其隐藏层大小为 2560。每个模块集成一个具有 32 个注意头的多头潜注意机制 (MLA)（DeepSeek-AI，2024），并通过 mHC（Xie，2025）连接到前馈神经网络 (FFN)，扩展率为 4。所有模型均使用 Muon（Jordan，2024；Team，2025）进行优化。实例化四个不同的模型：
• Dense-4B 作为基线模型。它采用上述骨干架构，并在每个模块中集成一个标准的密集前馈神经网络。
• MoE-27B 将标准的密集前馈神经网络替换为 DeepSeekMoE 模块（Dai，2024）。该模型配置 72 个路由专家和 2 个共享专家（每个 token 激活前 k 个路由专家，即 6 个），在保持与 Dense-4B 相同激活参数量的情况下，总参数量可扩展至 267 亿。
• Engram-27B 严格基于 MoE-27B 架构构建，以确保公平比较。将路由专家的数量从 72 个减少到 55 个，并将释放的参数重新分配给一个 57 亿参数的嵌入模块（ρ = 74.3%），从而保持模型总大小不变，仍为 267 亿。关于 Engram 配置，在第 2 层和第 15 层实例化该模块，并将最大 N-gram 大小设置为 3，头数设置为 8，维度设置为 1280。为了优化模型，用 Adam 优化器（Kingma，2014）更新嵌入参数，学习率乘以 5，且不进行权重衰减。同时，卷积参数初始化为零，以在训练开始时严格保持恒等映射。
• Engram-40B 保留与 Engram-27B 相同的骨干网络和计算预算，但将稀疏嵌入模块扩展到 185 亿个参数（总计 395 亿个参数）。该模型旨在研究 Engram 的扩展特性。
评估协议：用涵盖语言建模、知识库、推理、阅读理解和代码/数学等多个基准测试来评估模型。对于每个基准测试，都遵循标准的提示协议和评估指标。
• 语言建模：在 The Pile 数据集（Gao，2020）的测试集以及从与训练数据相同分布中抽取的验证集上的损失。
• 知识与推理：MMLU（Hendrycks，2021a）、MMLU-Redux（Gema，2025）、MMLU-Pro（Wang，2024b）、CMMLU（Li，2024）、C-Eval（Huang，2023）、AGIEval（Zhong，2024） 2024）、ARC-Easy/Challenge（Clark，2018）、TriviaQA（Joshi，2017）、TriviaQA-ZH（内部版）、PopQA（Mallen，2023）、CCPM（Li，2021）、BBH（Suzgun，2023）、HellaSwag（Zellers），2019），PIQA （Bisk，2020）和 WinoGrande（Sakaguchi，2021）。
• 阅读理解：DROP（Dua，2019）、RACE（中级/高级）（Lai，2017）和 C3（Sun，2020）。
• 代码和数学：HumanEval（Chen，2021）、MBPP（Austin，2021）、CruxEval（Gu，2024）、GSM8K（Cobbe，2021）、MGSM（Shi，2023）和 MATH（Hendrycks，2021b）。
通过将局部依赖建模卸载到静态查找，Engram 架构保留宝贵的注意容量，用于管理全局上下文。通过进行长上下文扩展训练（Gao，2025；Peng，2024）来验证这一结构优势。通过严格的评估协议，将架构贡献与基础模型能力分离，证明 Engram 在长距离检索和推理任务中取得显著的性能提升。
训练细节。为了实现长上下文能力，采用DeepSeek-V3 中引入的上下文扩展策略（Liu，2024a）。在预训练阶段之后，应用 YaRN（Peng，2024）进行上下文窗口扩展，在一个包含 32768 个 token 的上下文训练阶段中运行 5000 步（300 亿个高质量长上下文 token）。超参数为：尺度 s = 10，α = 1，β = 32，缩放因子 f = 0.707。
模型配置。比较四种不同模型配置下的上下文扩展。对于 MoE-27B 和 Engram-27B，都使用最终的预训练检查点（50k 步）。此外，为了严格评估架构效率，为 Engram-27B 选择两个中间检查点，分别位于 41k 步和 46k 步。尽管初始化阶段不同，所有变型都遵循完全相同的上下文扩展训练协议。关键在于，选择 Engram-27B（46k 步）是因为它的预训练损失与全训练的 MoE-27B（50k 步）相同。这创建一个可控的“等损失”设置，确保上下文扩展过程中任何性能差异都归因于架构而非模型初始质量。
评估基准。用 LongPPL（Fang 2024）和 RULER（Hsieh 2024）评估长上下文性能。对于 LongPPL，构建涵盖四类内容的评估集：长篇书籍、研究论文、代码库和长思维链 (CoT) 轨迹。对于 RULER，在 14 个子集上进行评估，这些子集聚合为 8 个类别：单K (S)、多K (MK)、多V (MV) 和多Q (MQ) 的大海捞针；多跳变量跟踪 (VT)、常用词提取 (CWE)、高频词提取 (FWE) 和问答 (QA)。
当前的LLM缺乏专门的知识查找原语，它们依赖计算来模拟记忆提取。如表所示，为了识别实体“戴安娜王妃”，LLM必须消耗多层注意机制和前馈神经网络（FFN）来逐步构建特征（Ghandeharioun，2024；Jin，2025；Li&Subramani，2025），而这一过程理论上可以通过知识查找操作来识别。
鉴于此，假设通过赋予模型显式的知识查找能力，Engram能够有效地模拟模型深度的增加，从而减轻模型早期特征构建阶段的负担。为了验证这一假设，采用两种机制可解释性工具：LogitLens（Belrose，2023；nostalgebraist，2020）和中心核对齐分析（CKA）（Davari，2023；Kornblith，2019）。
在受控环境下对 Engram 进行消融，研究每个关键模块设计的有效性。除非另有说明，骨干网络是一个 12 层 3B MoE 模型（0.56B 个激活参数），使用 1000B 个 token 进行训练。
参考配置。在主干网络中加入固定的 16 亿参数 Engram 记忆。参考模型使用 {2, 3} 元语法，并在第 2 层和第 6 层插入 Engram，实现 1.768 的验证损失，相比 MoE 基线模型有显著改进（Δ = 0.04）。
为了刻画Engram模块的功能贡献，通过在推理过程中完全抑制稀疏嵌入输出（同时保持模型主干不变）来评估模型。关键在于，这种事后消融操作会造成训练-推理不一致，可能在复杂的混合能力任务中引入噪声。因此，优先分析事实知识和阅读理解——敏感性谱的两个极端——它们在这种压力测试下表现出最高的信噪比。
Engram 相较于基于路由的 MoE 的一个关键系统优势在于，其稀疏激活由显式的静态哈希 ID 寻址。这产生一种严格确定性的内存访问模式：一旦token序列已知，下一次 Engram 查找的索引就已确定，并且可以在相应层执行之前计算出来。
实验设置。实现一个基于 nano-vLLM1 的推理框架——nano-vLLM1 是业界标准 vLLM 引擎（Kwon，2023）的精简原型。为了获得一个清晰的延迟基线，避免 MoE 中 Expert Parallel 的复杂通信模式，在两个密集骨干网（Dense-4B 和 Dense-8B）上进行基准测试。将一个包含 100B 参数的庞大 Engram 层插入到第二个 Transformer 模块中，整个嵌入表驻留在主机 DRAM 中。在推理过程中，系统异步预取 Engram 层的嵌入，使 PCIe 传输与第一个模块的计算重叠。
关注微信