大家好,这里是最佳拍档。在元旦的mHC流行约束超连接论文之后,DeepSeek在12日又发布一篇新的论文,同时还开源了相关的实现。这次,他们提出了一种全新的条件记忆机制 Engram,目的是让MoE模型在保持巨量参数的同时更高效地处理语言信息。DeepSeek创始人兼CEO梁文锋、北大王选计算机研究所的赵东岩和张辉帅教授都位列论文作者中。
简单来说,Engram架构的核心优势在于用更低的成本来实现更优的性能。使用Engram之后,训练计算量在相较MoE减少18%的情况下,在32768个token的长上下文任务中,Engram在RULER基准测试中反而超过了同参数量的MoE模型。并且,Engram浅层部署的记忆模块可以接管局部依赖与静态的知识存储,为注意力机制腾出容量,从而更加专注于全局推理。即使在卸载了1000亿参数的记忆表后,依然保持H800的推理吞吐量降幅不到3%。
DeepSeek还观察到,增加记忆的槽位数量能够持续、稳定地降低验证损失,这意味着Engram提供了一个新的、可预测的Scaling手段,也就是通过增大记忆的容量来持续带来收益,而无需增加计算量。在论文中,DeepSeek还提出条件记忆或将成为下一代稀疏模型中不可或缺的建模原语,这也许意味着DeepSeek即将推出的新模型有望整合条件记忆机制,实现知识的高效检索与推理能力的飞跃。
今天我们就来解读一下这篇论文。在正式介绍新型记忆机制前,DeepSeek团队首先提出了一项重要的观察,那就是稀疏性已经成为了智能系统的核心设计原则。在大模型领域,具体的实现就是MoE,通过在处理不同Token时激活不同的参数子集,实现了在不显著增加计算量的前提下参数规模的扩张。
但是,大语言模型的本质是处理异质性的信号。人类语言中既包含需要深度逻辑推理的复杂结构,也包含了大量像人名、地名、固定短语等静态知识。现有的Transformer架构并没有为这两种任务设计独立的通道,而是强行让神经网络通过深层的计算去模拟记忆检索的过程。这就导致现有的大模型不得不在早期层中通过昂贵的计算来重建静态知识,从而浪费了宝贵的模型深度。因此,DeepSeek认为有必要提出第二个与条件计算互补的稀疏维度,也就是条件记忆(Conditional Memory)。
条件记忆依赖于稀疏的查找操作,为固定的知识检索提供静态嵌入表示,适合命名实体、固定表达等静态而且高度模式化的语言表示。DeepSeek通过向经典的𝑁-gram结构引入现代化的条件记忆模块,包括分词器压缩、多头哈希、上下文门控以及多分支集成等等,最终提出了Engram架构。通俗一点地说,Engram就是给Transformer加了个外接的记忆库,并且把当前token附近的一小段内容用快速、省参数的方式去一个超大的静态记忆表里查找到对应的内容,而这个过程的时间复杂度是O(1),几乎不消耗计算资源。
那么,这个记忆库究竟该如何具体实现呢?首先,DeepSeek团队对分词器进行了压缩。普通的分词器会把Apple、apple、APPLE这些单词当成完全不同的东西,但是对人来说,这几个单词其实差别不大。而Engram通过预计算一个投影层,先把词表清洗了一遍,全部转为小写,进行Unicode规范化。最后,一个原本128k的词表实际上只剩下了77%,有23%的token ID被合并了,这让N-gram记忆的密度明显得到了提升。
不过,直接对所有的𝑁-gram进行建模是不可行的,参数会呈指数级的增长。于是,DeepSeek团队引入了多头哈希记忆,在固定的参数预算下去近似大规模的𝑁-gram表,从而降低哈希碰撞引入的语义噪声。不过,这种检索机制提供的记忆是静态的,缺乏上下文的适应性,容易受到歧义与冲突的影响,需要通过上下文感知门控来解决。它利用当前层经过注意力机制处理后的隐藏状态作为查询,去评估检索到的记忆向量的相关性。如果检索内容与当前的语境不符,门控值会趋向于0,从而抑制噪声,反之则会高亮激活,将外部知识注入流中。这种设计使得Engram不仅是一个静态数据库,更是一个能够根据语境动态调整的智能记忆体。
为了进一步扩大感受野并且增强非线性的建模能力,所以模型还引入了一个深度可分离的因果卷积。DeepSeek团队还采用了多分支架构作为默认的主干网络,而非标准的单流残差连接。多分支架构可以把残差流扩展为M个并行分支,但是共享记忆表和输出映射,这样设计的好处是它可以一次性用矩阵乘法搞定多条分支的计算,GPU用得非常高效。
Engram的核心优势在于记忆检索是完全依赖于输入token的,而非运行时的隐藏状态。这种确定性的机制实现了参数存储与计算资源的解耦,支持训练和推理阶段采取专门的优化策略。在训练优化方面,通过将超大嵌入表分片到多张GPU,利用All-to-All通信,按需收集对应行,从而让总记忆的容量可以随着GPU的数量线性扩展。在推理优化方面,由于可以提前确定待查询的记忆,所以系统可以从主机内存里异步预先提取,同时在前几层计算期间隐藏通信延迟,实现预取与计算的重叠,从而避免GPU停顿。
另外,Engram在模型中的放置位置还需要平衡建模的性能与系统延迟。较早引入有助于局部模式的重建,较深放置则会延长延迟的隐藏窗口,所以需要兼顾二者来进行优化。最后,在层次化的存储方面,基于自然语言𝑁-gram的分布特性,可以采用多级缓存策略,将高频嵌入存放于GPU HBM或者主机DRAM中,而低频嵌入则放置于SSD,这样就可以让Engram扩展至超大规模的记忆,同时保持低延迟与高效率。
接下来,DeepSeek团队研究了另一个关键的问题,条件计算和条件记忆这两种稀疏模式该怎么进行分配才能发挥最佳的效果呢?实验发现,在有限资源下,把所有空闲参数都给MoE模型并不是一个最优解,最好的效果是大约75%-80%给MoE,其余20%-25%给Engram。如果完全由MoE主导,模型就会缺乏静态模式的专用记忆,只能靠计算反复重建,效率很低。而如果完全由Engram主导,模型则失去了动态计算的能力,无法应对需要上下文理解的任务。因此,这条U型曲线验证了两个模块的结构互补性。
前面这个实验探索的是在固定参数参数预算下的分配优化,那么如果把记忆大幅度扩展会发生什么呢?于是,研究人员进一步探索了无限内存机制(Infinite Memory Regime)。在MoE主干网络不变的情况下,单纯增加Engram的嵌入槽位数量,从25万增加到了1000万,模型的Loss呈现出严格的对数线性下降趋势。DeepSeek认为,这意味着Engram提供了一个可预测的Scaling新手段,也就是增大记忆的容量会持续带来收益,而无需增加计算量。同时,相比别的只做简单平均的记忆方法,比如OverEncoding,Engram的Scaling潜力更大,性能提升更加明显。
这些实验结果验证了条件记忆作为稀疏容量的独立可扩展维度与MoE的条件计算形成了互补。在验证了架构和技术路径的可行性之后,DeepSeek团队的下一步就是进行大规模的Scale,验证这种方式在实际模型预训练中的有效性。具体而言,DeepSeek训练了四个模型,分别是Dense-4B、MoE-27B、Engram-27B、Engram-40B。训练时的语料库、分词器都使用了相同的设置,而后面两个模型则引入了Engram机制,用来研究在模型大小不变和Engram进一步扩展后的特性。结果显示,在相同算力和参数量的情况下,Engram-27B能在MoE-27B的基线上取得持续的提升,并且这些增益并不仅限于知识密集型的任务,通用推理任务、代码与数学推理任务从中得到的提升甚至更为显著。这些结果支持了DeepSeek的假设,通过引入专门的知识查找原语,能够显著提升表示效率,这超出了仅将整个稀疏预算用于条件计算所能达到的效果。
最后,研究人员将模型扩展到了Engram-40B上,进一步降低了预训练损失,并且在大多数基准上提升了性能。虽然它还没有在每个任务上严格优于Engram-27B,但这很可能是训练不足的结果。DeepSeek团队观察到,在训练结束时,Engram-40B与基线模型之间的训练损失差距仍然在扩大,这表明在当前的token预算下,扩展的记忆容量还没有完全发挥它的潜力。
接着,DeepSeek团队用MoE-27B与Engram-27B作为对照组,都使用了5000步,大约300亿token的高质量长上下文数据进行微调,然后他们采用了DeepSeek-V3中的YaRN技术,将模型的上下文窗口扩展到了32768个token。实验结果显示,由于Engram模块接管了局部依赖的建模,它为模型的注意力机制腾出了容量,让它能够更专注于处理全局上下文。因此,Engram架构在处理超长文本和长程推理任务上比传统的架构表现更好。具体来说,在架构方面,在排除了基础模型能力差异的情况下,Engram-27B依然显著优于MoE-27B,而在复杂的检索任务RULER基准测试中,Engram表现出了更强的长程依赖处理能力。例如在多查询大海捞针任务中,Engram准确率大幅领先于对手。
在计算效率方面,即使只用了82%的预训练计算量,Engram-27B的表现依然能够与完全训练的MoE-27B基线模型持平,甚至在RULER基准上实现超越,这证明了Engram架构具有极高的训练效率,能用更少的计算资源达到同等或者更好的长上下文性能。随后,DeepSeek团队对Engram模型进行了深入的机制分析和消融实验,核心目的是回答Engram到底是如何工作的,以及它的各个组件有什么用这两个问题。
首先是模型的深度与表征分析,DeepSeek团队通过LogitLens分析显示,Engram模型在早期层就展现出了极低的KL散度,能够更快地收敛到最终的预测结果,因为它通过查表直接获取了静态知识,不需要像传统模型那样通过多层计算来重组基础特征。而中心核对齐工具CKA的分析发现,Engram的浅层在表征上与纯MoE模型的深层高度相似,这意味着Engram让模型在更少的层数内完成了同等复杂的特征提取,在功能上等同于增加了模型的有效深度。
在架构消融实验中,研究人员发现,将Engram模块放在较浅的层,比如第2层时,效果最好,这样可以尽早卸载模型背负的局部模式重建任务,让后面的深层网络专注于复杂的全局推理。研究人员还发现,分支特定融合、上下文感知门控和分词器压缩这三者对性能的影响最大,去掉任何一个都会导致验证损失的显著上升,而次要组件,比如轻量级卷积层的影响则较小。
那么,如果把Engram关掉,模型在哪些任务上会崩溃呢?为回答这个问题,DeepSeek团队进行了功能敏感性分析,他们测试了在推理时,强制屏蔽Engram模块的输出,观察性能的下降情况。结果显示,在事实性知识方面,模型性能出现了灾难性下降,仅保留了大约29-44%的性能,这证明Engram是模型存储参数化知识的主要仓库。️但是在阅读理解方面,模型的性能几乎不受影响,保留了大约81-93%。这证明涉及上下文推理的任务主要由Transformer的骨干网络处理,而非记忆模块。
在系统效率与推理吞吐上,由于Engram的访问模式是预先可知的,不像MoE那样需要根据隐藏状态来动态路由,所以系统可以提前从内存中预取数据,即使将一个1000亿参数的Engram表卸载到主机内存中,它在H800硬件上的推理吞吐量下降也不到3%,这证明了Engram能够以极低的代价实现参数量的大幅扩展。此外,Engram的门控机制会在遇到静态模式时会被激活,也就是变红,红色越深代表门控值越高,我们可以清晰地看到,当遇到亚历山大大帝、银河系等专有名词或者威尔士王妃等固定称谓以及中文语境时,Engram会被强烈激活,而在处理需要动态推理的文本时,门控则会保持关闭。
最后,DeepSeek团队将Engram与MoE外部记忆与检索增强、长上下文建模,以及表征学习与知识蒸馏进行了对比。传统MoE是Engram的前辈,它通过稀疏激活来扩展模型的容量,Engram解决了传统MoE在超大规模下路由成本高、训练不稳定的问题,提供了一种更高效的扩展路径。而对比外部记忆与检索增强RAG,这类工作通常是在模型外部挂一个数据库,在推理时进行实时检索,而Engram是内化的记忆,它在预训练阶段就把海量知识消化并且固化到了参数化的记忆表中,这使得它比传统RAG具有更低的延迟和更强的知识一致性。长上下文建模这个领域的研究则主要关注如何让模型的注意力机制能够处理更长的序列。DeepSeek团队强调,Engram并不是要取代注意力机制,而是与之互补,Engram负责处理局部的、静态的上下文依赖,从而让注意力机制能够更专注于处理全局的、动态的长程依赖。对于表征学习与知识蒸馏来说,Engram提供了一种新的视角,它将模型的知识解耦为通用推理能力和特定知识库,这种解耦结构天然适合进行知识蒸馏,因为未来的研究团队可以选择只蒸馏轻量级的骨干网,而将庞大的知识库作为可插拔的附件。
总的来说,Engram架构的核心思想就是通过解耦来实现效率与性能的平衡,它成功的将局部模式重建从复杂的Transformer骨干网中剥离出来,交由专门的记忆模块来处理,这种设计使得模型在保持强大推理能力的同时,能够以极低的成本扩展到超大规模。DeepSeek团队认为,Engram不仅仅是一个学术上的新模型,还具有很强的工程落地价值。由于记忆模块的访问具有确定性,Engram可以进行高效的预读取和硬件优化,非常适合大规模的部署。而且,既然知识集中在Engram表中,未来或许可以通过直接修改这个表来修正模型的知识错误,而无需进行昂贵的微调。不过,目前的Engram是在预训练时固化的,未来的一个重要方向是让这个记忆模块具备在线学习或者动态更新的能力,从而让模型能够实时获取新的知识。Engram在架构上的创新,或许也预示着下一代大语言模型将不再只是单纯的神经网络,而是神经网络加智能检索引擎的有机结合体。
好了,以上就是对这篇论文的解读了,希望能对大家有所帮助。感谢观看,我们下期再见。