当前位置:首页 > 科技资讯 > 正文

DeepSeek提出条件记忆:为大型语言模型引入新维度

就在十几个小时前,DeepSeek发布了一篇新论文,题为《Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models》,该论文与北京大学合作完成,作者梁文锋也参与其中。

DeepSeek提出条件记忆:为大型语言模型引入新维度 DeepSeek 条件记忆 稀疏化 大型语言模型 第1张

论文地址:https://github.com/deepseek-ai/Engram/blob/main/Engram_paper.pdf

简而言之,这项新研究旨在解决以下问题:目前,大语言模型主要通过混合专家(MoE)实现稀疏化,这被称为“条件计算”。然而,现有的Transformer缺乏原生的知识查找机制,只能被迫通过低效的计算过程模拟检索行为。

针对这一现状,DeepSeek提出了条件记忆(conditional memory),与MoE的条件计算互补,并通过引入一个新模块Engram来实现。

目前,模块“Engram”的相关实现已经上传至GitHub。

DeepSeek提出条件记忆:为大型语言模型引入新维度 DeepSeek 条件记忆 稀疏化 大型语言模型 第2张

项目地址:https://github.com/deepseek-ai/Engram

这让网友们感慨:“DeepSeek is back!”

DeepSeek提出条件记忆:为大型语言模型引入新维度 DeepSeek 条件记忆 稀疏化 大型语言模型 第3张

此外,结合元旦期间公布的研究《mHC: Manifold-Constrained Hyper-Connections》,可以明确的是DeepSeek v4的模样愈发清晰,即将上线!

除了条件计算(MoE),LLM 还需要一个独立的条件记忆 Engram

MoE模型通过条件计算实现了模型容量的扩展,但现有的Transformer架构缺乏原生的知识查找原语,只能通过低效的计算过程模拟检索行为。

为了解决这一问题,DeepSeek提出了条件记忆(conditional memory)这一与条件计算互补的稀疏化维度,并通过Engram模块加以实现。Engram在经典N-gram嵌入的基础上进行了现代化改造,使其能够以O(1)时间复杂度完成知识查找。

通过形式化提出稀疏性分配问题,DeepSeek还发现了一条呈U型的扩展规律,用以刻画神经计算(MoE)与静态记忆(Engram)之间的最优权衡关系。

在这一规律的指导下,DeepSeek将Engram扩展至270亿参数规模,并在严格等参数量、等FLOPs的条件下,其整体性能显著优于纯MoE基线模型。

值得注意的是,尽管记忆模块主要被用于提升知识检索能力(如MMLU提升+3.4、CMMLU提升+4.0),但DeepSeek观察到其在通用推理能力(如BBH提升+5.0、ARC-Challenge提升+3.7)以及代码与数学推理任务(HumanEval提升+3.0、MATH提升+2.4)上带来了更为显著的增益。

进一步的分析表明,Engram能够将静态知识的重建负担从模型的浅层中剥离出来,从而有效加深网络用于复杂推理的有效深度。此外,通过将局部依赖关系交由查表机制处理,Engram释放了注意力机制的容量,使其更专注于全局上下文建模,从而显著提升了长上下文检索能力(例如Multi-Query NIAH的准确率从84.2提升至97.0)。

最后,Engram在系统层面同样展现出基础设施感知的高效性:其确定性的寻址方式支持在运行时从主机内存进行预取,几乎不会带来额外的性能开销。

DeepSeek认为,条件记忆将成为下一代稀疏大模型中不可或缺的核心建模原语。

Engram架构旨在将静态模式存储与动态计算过程从Transformer主干网络中分离出来进行增强。该模块对序列中每一个位置依次执行两个功能阶段:检索与融合。

DeepSeek提出条件记忆:为大型语言模型引入新维度 DeepSeek 条件记忆 稀疏化 大型语言模型 第4张

在运行过程中,DeepSeek首先提取并压缩当前位置的后缀N-gram,并通过哈希机制以确定性方式检索对应的静态嵌入向量。随后,这些被检索到的嵌入会在当前隐藏状态的调制下进行动态调整,并通过轻量级卷积操作进一步精炼。最后,Engram与多分支架构集成。

基于哈希N-gram的稀疏检索

这一阶段的目标是将局部上下文映射到静态记忆条目。这主要包括分词器压缩以及通过确定性哈希机制检索对应的嵌入表示。

分词器压缩:为了最大化记忆单元的语义密度,DeepSeek引入了一层词表投影。为此,他们预先设计了一个映射函数:

DeepSeek提出条件记忆:为大型语言模型引入新维度 DeepSeek 条件记忆 稀疏化 大型语言模型 第5张

将原始token ID映射为基于文本规范化等价关系(例如使用NFKC规范化、统一大小写等)得到的规范化标识符。在实际应用中,对于一个规模为128k的分词器,该过程能够将有效词表规模缩减约23%(详见附录C)。

多头哈希:直接对所有可能的N-gram组合空间进行参数化在计算和存储上都是不可行的。借鉴Tito Svenstrup等(2017)的工作,DeepSeek采用了一种基于哈希的近似方法。为了降低哈希冲突的影响,对于每一种N-gram阶数n,引入K个相互独立的哈希头。

上下文感知门控

前一阶段通过哈希N-gram从条件记忆中检索得到的嵌入向量提供的是与具体语境无关的静态先验信息。然而,正因为其静态属性,这些嵌入缺乏对当前上下文的自适应能力,且在实际应用中可能受到哈希冲突或词项多义性带来的噪声干扰。

为此,DeepSeek在检索之后引入了一种上下文感知的门控机制,其设计灵感来源于注意力机制。

系统效率:计算与存储的解耦

在带有记忆机制的模型中,规模扩展往往受到GPU高带宽显存(HBM)容量有限的制约。然而,Engram所采用的确定性检索机制天然支持将参数存储与计算资源进行解耦。不同于MoE依赖运行时隐藏状态进行动态路由,Engram的检索索引完全由输入token序列决定。这种可预测性使得针对训练与推理阶段的专门优化策略成为可能。

U型扩展规律与稀疏性分配

作为“条件记忆”的一种具体实现,Engram在结构上与MoE专家提供的“条件计算”形成互补。本节旨在探究这种二元特性的扩展属性以及如何最优地分配稀疏容量。

具体而言,本项研究由两个核心问题驱动:

  • 有限约束下的分配: 在总参数量和训练计算量固定(即等参数、等FLOPs)的情况下,应如何在MoE专家与Engram嵌入之间划分稀疏容量?
  • 无限记忆范式: 考虑到Engram具有不随规模增长的特性。

DeepSeek提出条件记忆:为大型语言模型引入新维度 DeepSeek 条件记忆 稀疏化 大型语言模型 第6张

如果放宽记忆预算或进行激进扩展,Engram自身会表现出怎样的扩展行为?

实验结果

通过提出的Engram架构以及经验推导出的分配法则,DeepSeek将Engram扩展至数十亿参数规模以验证其在真实语言模型预训练中的有效性。

  • Dense-4B(总参数量 41 亿)
  • MoE-27B(总参数量 267 亿)
  • Engram-27B(总参数量 267 亿)
  • 以及 Engram-40B(总参数量 395 亿)

所有模型均采用完全相同的数据训练流程(相同的token预算及顺序),且在激活参数量上严格匹配。

大规模预训练的实验结果

  • 在相同的训练计算预算下: 所有三种稀疏变体(MoE-27B、Engram-27B/40B)在所有基准测试中显著超越了iso-FLOPs的Dense-4B基准。
  • 更重要的是: Engram-27B在iso - 参数和iso-FLOPs的MoE-27B基准上持续取得改进。这些提升不仅限于知识密集型任务(例如MMLU: +3.0、MMLU-Pro: +1.8、CMMLU: +4.0),还观察到在一般推理领域(例如BBH: +5.0、ARC-Challenge: +3.7、DROP: +3.3)以及代码和数学推理任务(例如HumanEval: +3.0、MBPP: +1.6、GSM8K: +2.2、MATH: +2.4)中改进更加显著。