您的浏览器禁用了JavaScript(一种计算机语言,用以实现您与网页的交互),请解除该禁用,或者联系我们。[北京大学]:通过可伸缩查找的条件记忆:大型语言模型的新稀疏轴 - 发现报告

通过可伸缩查找的条件记忆:大型语言模型的新稀疏轴

AI智能总结
查看更多
通过可伸缩查找的条件记忆:大型语言模型的新稀疏轴

新成1,2王定增2, 大麦代2, 钦宇陈2, 王丙玄2, 辛德谢2, 黄克照2, 星凯宇 摘要 2, 赵文浩2而专家混合(MoE)通过条件计算扩展容量,Transformer 缺乏原生的知识查找原语,迫使它们通过计算低效地模拟检索。为解决这个问题,我们引入条件记忆作为补充稀疏轴,通过以下方式实例化engram,一个使经典现代化的模块B-gram 嵌入O(1)稀疏分配查询。通过制定问题中,我们揭示了一个U形缩放规律,优化了神经计算(MoE)与静态记 , 李宇昆2, 韩张2, 张惠帅1, 赵东岩1, 梁文峰2忆(Engram)之间的权衡。遵循这一规律,我们将Engram缩放到270亿参数,实现了比严格等参数和等FLOPs的MoE基线更优越的性能。尤为突出的是,尽管预期记忆模块将有助于知识检索(例如,MMLU+3.4; CMMLU+4.0), 我们观察到在一般推理方面甚至更大的收益 (例如,BBH+5.0; ARC-Challenge+3.7)和代码/数学域(HumanEval)+3.0; 数学+2.4)。机制分析表明,Engram使骨架的早期层从静态重建中解脱出来,有效地加深网络以进行复杂推理。此外,通过将局部依赖委托给查找,它释放了用于全局上下文的注意力容量,显著提高了长上下文检索(例如,Multi-Query NIAH:84.2 → 97.0). 最后,Engram 建立了感知基础设施的效率:它的确定性寻址能够从主机内存进行运行时预取,开销可忽略不计。我们视条件内存为下一代不可或缺的建模原语。https://github.com/deepseek-ai/Engram生成稀疏模型。代码可在: 1北京大学2 DeepSeek-AI {zhanghuishuai, zhaody}@pku.edu.cn {chengxin, zengwangding, damai.dai}@deepseek.com 1. 简介 稀疏性是智能系统的反复出现的的设计原则,从生物神经回路(Lennie,2003;奥尔沙乌森和菲尔德,1997) 到现代大型语言模型(LLM)。目前,这一原则主要通过专家混合(MoE)(戴等人,2024;沙泽尔等人,2017), 它通过条件计算来扩展容量。由于它能够在不按比例增加计算量的情况下大幅增加模型大小, MoE 已成为前沿模型的实际标准(科曼尼奇等人,2025;郭等,2025;团队等,2025). 尽管这种条件计算范式取得了成功,其内在的异构性结构优化对语言信号的分析表明存在显著的空间。具体而言,语言建模包含两种性质不同的子任务:组合推理和知识 边缘检索。前者要求深度、动态计算,而相当一部分文本—例如命名实体和公式化模式—是局部的、静态的、并且高度刻板化(康斯坦特等人,2017;厄尔曼,2000`). 经典的有效性B-gram 模型 (布拉茨等人,2007;刘等,2024b;阮,2024) 在捕获这种局部依赖关系意味着这些规律性自然地表示为计算成本低的查找。由于标准 Transformer(瓦桑维等人,2017) 缺少本地的知识查询原语,当前的LLMs被迫要通过计算模拟检索例如,解决一个常见的多标记实体需要消耗多个早期的注意力层和前馈网络(Ghandeharioun等人,2024;金等人,2025) (参见表3)。这个过程本质上相当于对静态查找表进行昂贵的运行时重建,浪费了宝贵的顺序深度在一些琐碎的操作上,而这些操作本可以分配给更高级别的推理。 为了使模型架构与这种语言双重性相协调,我们提倡一种互补的稀疏性轴:条件记忆。条件计算稀疏地激活参数来处理动态逻辑(Bengio 等人,2013;沙兹尔等人,2017), 条件记忆依赖于稀疏查找操作来检索固定知识的静态嵌入。作为这种范式的初步探索,我们重新审视B-gram 嵌入 ()博亚诺夫斯基等人, 2017)作为规范实例:局部上下文充当索引巨大嵌入表的关键,通过常数时间O(1)查找(黄等人,2025 a;帕格诺尼等人,2025;Tito Svenstrup 等,2017;于等人,2025). 我们的调查揭示,也许令人惊讶的是,这种静态检索机制可以作为现代MoE架构的理想补充——但前提是它被正确设计。在本文中,我们提出engram, 一种基于经典的条件记忆模块B-gram结构,但配备了现代适配功能,如词法器压缩、多头散列、上下文门控和多分支集成(详见第二节). 稀疏分配为了量化这两种原始操作之间的协同效应,我们提出了一个问题:在固定的总参数预算下,应该如何在MoE专家和Engram存储之间分配容量?我们的实验揭示了一种独特的U型扩展规律,表明即使是简单的查找机制,当被视为一等建模原始操作时,也起着对神经计算不可或缺的补充作用。根据这一分配规律,我们将Engram扩展到一个拥有270亿参数的模型。与严格等参数和等FLOPs的MoE基线相比,Engram-27B在各种领域都实现了更高的效率。关键的是,这些收益不仅限于知识密集型任务(例如,MMLU:+3.4;CMMLU:+4.0;MMLU- Pro:+1.8),其中存储容量在直观上是有益的;我们在一般推理(例如,BBH:+5.0;ARC-Challenge:+3.7;DROP:+3.3)和代码/数学领域(例如,HumanEval:+3.0;MATH:+2.4;GSM8K:+2.2)也观察到了更显著的改进。 通过LogitLens进行机理分析nostalgebraist,2020)和CKA(亨德里克斯等人,2021a)揭示了这些收益的来源:Engram 使骨干网络免于在早期层重构静态知识,从而增加了可用于复杂推理的有效深度。此外,通过将局部依赖委托给查找操作,Engram 解放了注意力能力以专注于全局上下文,从而在长上下文场景中实现了卓越的性能—在 LongPPL(方等)和RULER(谢等) (例如,多查询NIAH:97.0与84.2;变量跟踪:89.0与77.0)。 最后,我们将基础设施感知效率确立为首要原则。与MoE的动态路由不同,Engram采用确定性ID来实现运行时预取,使通信与计算重叠。实验结果表明,将一个100B参数的表卸载到主机内存会产生可忽略的开销(<3%)。这表明Engram有效地绕过了GPU内存限制,促进了激进参数扩展。 图1 |内gram 架构。该模块通过检索静态增强骨干B- 语法记忆并通过上下文感知门将其与动态隐藏状态融合。该模块仅应用于特定层,以将记忆与计算解耦,同时保留标准的输入嵌入和非嵌入模块。 2. 架构 2.1. 概述 如所示图1, engram是一个条件性记忆模块,旨在通过结构性地将静态模式存储与动态计算分离来增强transformer主干。形式上,给定一个输入序列X( =x,...,x)τ 和隐藏状态H(ℓ)∈Rτ×d在层ℓ,1模块处理每个位置푡在两个功能阶段:检索与融合. 首先,如详细所述2.2节,我们提取并压缩后 缀B-grams通过哈希确定性地检索静态嵌入向量。随后,在2.3节, 这些检索到的嵌入向量由当前的隐藏状态动态调制,并通过轻量级卷积进行细化。最后,我们讨论了与多分支架构的集成2.4节并且系统级设计在第二节 2.5. 2.2. 基于哈希B-gram的稀疏检索 第一阶段将本地上下文映射到静态内存条目,涉及分词器压缩和通过确定性哈希获取嵌入。 分词器压缩While B-gram模型通常直接对分词器输出进行操作,无损重建标准子词分词器优先考虑,通常分配不同的 ID 给 苹果 ␣苹果语义等价术语(例如,vs.)(Kudo和Richardson,2018;李等人, 2023b). 为了最大化语义密度,我们实现了一个词汇投影层。具体来说,我们预计算了一个满射函数 P:→′ V V将原始 token ID 压缩成规范 基于规范文本等效的标识符(使用 NFKC(惠斯勒,2025), 等)。在实践中,该过程为128k分词器实现了有效词汇量23%的减少(参见附录C). 形式上,对于一个位于位置푡,我们将其原始ID映射x푡转换到一个规范ID ′ P((=, ,-gram′′ 多头哈希。直接参数化所有可能组合空间B-grams 是难以处理的。随后Tito Svenstrup 等(2017),我们采用基于哈希的方法。为了减轻冲突,我们采用K为每个distinct hash heads B-gram排序n.每个头k将压缩的上下文映射到嵌入表中索引E...xx B g x x))=푡푡,n− +1푡푡.用于形成后缀 n,k∞ 在实际中,∅nk 作为轻量级乘法-XOR哈希实现。我们构造,最终内存向量e∈R d mem푡 通过连接所有检索到的嵌入: 2.3. 上下文感知门控 检索到的嵌入e 푡作为与上下文无关的先验分布。然而,由于它们是静态的,因此它们本质上缺乏上下文适应性,并且可 能由于哈希碰撞或多义性而受到噪声的影响。哈伯和普埃西奥,2024). 为增强表达能力并解决这种歧义,我们采用了一种受Attention启发的上下文感知门控机制(巴赫达瑙等人,2015;瓦桑维等,2017). 具体来说,我们使用了当前的隐藏状态h푡 —which通过前序注意力层聚合了全局上下文—as动态查询,同时检索到的内存e푡 作为键和值投影的源: 哪里W 是可学习的投影矩阵。为了确保梯度稳定性(德赫加尼等人,2023), 我们应用 RMSNorm (张和森纳里奇,2019) 在计算标量门之前应用于查询和键훼∈ (0, 1)RMSNorm(h) 푡:⊤RMSNorm(k) =푡훼√푡∑푡. (4)门控输出定义为˜v훼·=푡푡v푡. 此设计强制语义一致性:如果检索到的记忆e d最后,为了扩展感受野并增强模型的非线性,我们引入 ˜一个短而逐通道的因果卷积(顾等,2022;P eng等人,2023). 设V∈Rτ×d푡与当前上下文矛盾h푡 ˜表示门控值的序列。使用一个卷积核大小ω(设置为4),膨胀훿(设置为最大B-gram 顺序) 和SiLU 激活(埃尔温等人,2018), 最终输出Y被计算为:,网关훼푡趋向于零,有效抑制噪声。 编码模块通过残差连接集成到骨干中:H(ℓ) 4←H(ℓ)+Y,随后是标准的注意力机制和多模态扩展。关键在于,Engram不会被应用于每一层;它的具体位置由在中详细说明的系统级延迟约束所控制第二节 2.5. 图2 |恩格玛的系统实现。(a) 训练阶段:巨大的嵌入表被分散到可用的GPU上。采用All-to-All通信原语来检索跨设备的活动嵌入行。(b) 推理阶段:Engram表卸载到主机内存。通过利用确定性的检索逻辑,主机异步预取和传输嵌入,重叠通信与先前Transformer块的设备计算。 2.4. 与多分支架构集成 在这个工作中,而不是标准的单流连接 (他等人,2016),我们采用先进的多分支架构作为我们的默认骨干,因其优越的建模能力(拉森等,2017;塞格迪等人,2015;谢等人,2025;朱等,2025)。这种架构的一个关键特征是将残差流扩展到费平行分支,其中信息流由可学习的连接权重调制。 尽管Engram模块本身与拓扑无关,但要将其应用于此多分支框架,就需要进行结构优化以平衡效率和表达能力。具体来说,我们实现了一种参数共享策略:一个单一的稀疏嵌入表和一个值投影矩阵W在所有领域共享费V 分支,而费distinct Key projection(m) .(6)!(m)(m)RMSNorm(h)⊤RMSNorm(W e)(m)矩阵 {W}费用于实现特定分支的门控行为。对于m-thK m=1 (带隐藏状态的分支m)h푡 图3 |稀疏性分配和Engram缩放。左:分配比率上的验证损失Φ. 显示了两个计算预算(2e20和6e20 FLOPs)。两种机制都呈现U形,混合分配超过纯MoE。右:无限内存机制中的缩放行为。验证损失相对于嵌入数量呈现对数线性趋势。 2.5. 系统效率:计算与内存解耦 放