量子位 2小时前
DeepSeek开源大模型记忆模块,梁文锋署名新论文,下一代稀疏模型提前剧透
index.html
../../../zaker_core/zaker_tpl_static/wap/tpl_keji1.html

 

DeepSeek 节前开始蓄力!

最新论文直接给 Transformer 加上" 条件记忆 "(Conditional Memory),补上了原生缺乏的知识查找机制。

结论中明写道:我们将条件记忆视为下一代稀疏模型不可或缺的建模原语。

还是梁文锋署名,并与北京大学王选所赵东岩、张辉帅团队合作。

论文中不仅提出了条件记忆这个全新范式,并给出了具体实现方案Engram模块,实验中让 27B 参数碾压同规模纯 MoE 模型,甚至变相提升了大模型的推理能力:

让原来 Transformer 要用 6 层注意力才能干的简单任务压缩到 1-2 层搞定,省出来的资源就可以用于更难的推理任务了。

条件记忆的原理其实也非常 " 原始 ":不靠计算,回归查表,用上了传统N-gram方法。

给大模型一个巨大的词表,专门存那些固定的实体名称和两三个词的短语,不管词表多大,找信息都是 O ( 1 ) 速度。

关键就在于,如此前大模型时代的玩法,DeepSeek 如何解决传统 N-gram 模型存储爆炸和多义性问题,又是让它和现代 Transformer 结合起来的?

让注意力干 " 苦力活 " 太浪费了

团队的核心观察是,语言建模其实包含两种性质完全不同的任务,一种是需要深度动态计算的组合推理,另一种则是检索静态知识。

问题在于,现有的 Transformer 架构缺乏原生的知识查找机制。

当模型需要识别一个实体时,它得消耗好几层注意力和前馈网络,逐层拼凑特征,最终才能完成。

论文中引用了一个具体案例:"Diana, Princess of Wales"

模型需要经过 6 层才能完成这个识别过程,前几层还在纠结 "Wales 是英国的一个地区 "、"Princess of Wales 是某种头衔 " 这些中间状态,最终才能 " 想起来 " 这是指戴安娜王妃。

本质上是在用昂贵的运行时计算来重建一个静态查找表,那些本可以用于更高层推理的网络深度,被浪费在了识别概念这种 " 苦力活 " 上。

回归查表,回归 N-gram

Engram 的设计思路相当直接:既然经典的 N-gram 模型就能用 O ( 1 ) 的时间复杂度捕获这些局部依赖,那为什么不把这个能力直接嵌入 Transformer?

具体实现上,团队在原有的 Transformer 层之间插入 Engram 模块。每个位置的输入会触发一次哈希查找:把当前 token 和前面几个 token 组成的 N-gram 映射到一个巨大的嵌入表中,直接取出对应的向量。

为了处理哈希冲突和多义性问题,团队引入了上下文感知的门控机制,用当前的隐藏状态作为 Query,检索到的记忆作为 Key 和 Value,计算一个 0 到 1 之间的标量门控值。

如果检索到的内容和当前上下文不匹配,门控值就趋近于零,相当于自动屏蔽噪声。

下图中,颜色越深说明 Engram 越判断当前文本片段是 " 固定静态模式 ",倾向于调用记忆库中的对应信息。

颜色越浅代表这段文本越动态灵活,主要靠模型的注意力机制处理。

比如只看到 " 张 " 是一个常见姓氏,但是 " 张仲景 " 三个字凑一起就是固定历史人物实体了。

接下来还要解决传统 N-gram 模型的两个痛点。

语义重复,同一个词的不同形式(比如 Apple、apple、Ä pple)被当成不同 token,浪费存储。

存储爆炸,所有可能的 N-gram(比如 2 词、3 词组合)数量太多,比如 128k 词表就要存 128k^3 种组合,直接存储根本存不下。

DeepSeek 团队首先压缩 tokenizer,把语义相同但形式不同的 token 归为一类,128k 词表的有效规模直接减少 23%,相同语义的 token 聚在一起,查找更高效。

再用多个哈希函数把 N-gram 映射成 embedding 表的索引,

这既解决了存储爆炸:不管有多少种 N-gram,都通过哈希函数映射到一个固定大小的 embedding 表里,表的大小是质数。

又减少查找冲突:给每种 N-gram 阶数(比如 2-gram、3-gram)配 K 个不同的哈希头,每个哈希头对应一个独立的 embedding 表,把所有 N-gram 阶数、所有哈希头取出来的 embedding 向量拼在一起,形成最终的 " 记忆向量 "e ₜ,供后续模块使用。

U 型曲线:MoE 和记忆的最优配比

论文最核心的部分是对 " 稀疏性分配问题 " 的系统研究。

团队设计了一个严格的实验框架:固定总参数量和每 token 的激活参数量(也就是计算量),然后在 MoE 专家和 Engram 记忆之间重新分配 " 闲置参数 " 预算。

分配比例 ρ 从 100%(纯 MoE)逐步降到 40%,实验结果画出了一条清晰的 U 型曲线:

纯 MoE 反而不是最优解,把大约 20% 到 25% 的稀疏参数预算分给 Engram 记忆时,模型验证集 loss 达到最低点。

在 100 亿参数规模下,最优配置比纯 MoE 基线的 loss 降低了 0.0139。

更重要的是,这个最优分配点在不同计算预算下都相当稳定,大约在 ρ =75% 到 80% 之间。

团队解释了 U 型曲线两端的含义:

MoE 主导时,模型缺乏静态模式的专用记忆,被迫通过网络深度和大量计算来低效重建。

Engram 主导时,模型丢失了条件计算能力,在需要动态推理的任务上表现下降。

总之,记忆无法替代计算,计算也无法高效模拟记忆。

27B 规模验证:推理能力提升超预期

按照 U 型曲线的指导,团队把 Engram 扩展到更大参数规模进行验证,并对比纯 MoE 模型和纯密集模型。

所有模型训练条件一致,激活参数量都是 38 亿,训练 token 都是 2620 亿,差异仅在 " 稀疏能力分配 "。

Dense-4B:纯密集模型。

MoE-27B:纯混合专家模型,72 个路由专家 +2 个共享专家,所有稀疏参数都给 MoE。

Engram-27B:MoE+Engram 混合模型,55 个路由专家 +2 个共享专家,把 5.7B 稀疏参数分配给 Engram 记忆模块。

Engram-40B:进一步扩展 Engram 模块,保持专家数量不变,Engram 记忆参数增至 18.5B,总参数 39.5B。

结果 MoE-27B 和 Engram-27B 对比,知识密集型任务的提升在预期之内:比如 MMLU 提升 3 分,CMMLU 提升 4.0 分,TriviaQA 提升 1.9 分。

但出乎意料的是,通用推理和代码数学领域的提升幅度也很大:BBH 大幅提升 5.0 分,ARC-Challenge 提升 3.7 分,DROP 提升 3.3 分,HumanEval 提升 3.0 分,MATH 提升 2.4 分,GSM8K 提升 2.2 分。

团队用 LogitLens 和 CKA 分析揭示了原因。

Engram 让模型的早期层不再需要做特征组合的 " 苦力活 ",KL 散度曲线显示 Engram 模型的预测收敛速度明显更快。更直观的证据来自 CKA 相似度矩阵,Engram-27B 第 5 层的表征,和 MoE 基线第 12 层的表征最为相似。

这意味着 Engram 实际上 " 加深 " 了网络的有效深度,省下来的层数被用于更复杂的推理任务。

Engram-40B 进一步增加记忆参数后,大部分任务性能持续提升,且训练后期损失仍在下降,说明记忆容量还未饱和,后续可继续扩大。

另外长上下文场景的提升尤为显著。

在 RULER 测试集上,Multi-Query NIAH 从 84.2 跃升到 97.0,Variable Tracking 从 77.0 提升到 89.0。

论文解释说,Engram 把局部依赖建模卸载给了查找操作,释放了注意力容量去关注全局上下文。

百亿参数表放 CPU 上,延迟几乎没影响

接下来又到了喜闻乐见的软硬结合工程优化环节。

在训练阶段,词表规模会高达 100B 参数,单个 GPU 存不下,必须拆分到多个 GPU 上,需要 All-to-All 通信机制,让所有 GPU 之间互相传递需要的记忆片段。

在推理阶段把词表卸载到 CPU 内存,同时又不能让记忆调用拖慢计算节奏。

和 MoE 的动态路由不同,Engram 的查找索引只取决于输入 token 序列,完全可以提前计算。

这个确定性让团队能够把巨大的嵌入表放到 CPU 内存里,用 PCIe 异步预取,让通信和前面层的计算重叠。

具体通过把 Engram 模块插在 Transformer 网络的特定层,GPU 计算前一层的同时,CPU 预取当前层需要的 Engram 记忆,等 GPU 算完前一层,所需的记忆也已经传输到位。

实验直接把一个 1000 亿参数的 Engram 表放到 CPU 内存,在 H800 上跑推理。4B 密集模型的吞吐量从 9031 token/s 降到 8858 token/s,8B Dense 模型从 6315 token/s 降到 6140 token/s,额外开销都在 3% 以内。

自然语言 N-gram 天然遵循 Zipfian 分布,极少数高频模式占据绝大多数访问量。这意味着可以设计多级缓存:高频嵌入放 GPU 显存,中频放 CPU 内存,长尾放 NVMe SSD,把有效延迟进一步压缩。

DeepSeek 团队在结论中写道:

Engram 将 " 硬件感知效率 " 确立为核心设计原则:其确定性寻址机制支持存储与计算的解耦,能够将海量参数表卸载至主机内存,且推理开销可忽略不计。我们认为,条件记忆将成为下一代稀疏模型中不可或缺的建模基元。

DeepSeek 的下一代稀疏模型,已被曝光将在春节前发布,敬请期待。

评论
大家都在看