核心观点 - DeepSeek团队提出了一种名为“条件记忆”(Conditional Memory)的全新建模范式,并给出了具体实现方案Engram模块,旨在为Transformer架构补上原生的知识查找机制 [1] - 该研究认为,语言建模包含需要深度动态计算的组合推理和检索静态知识两种不同性质的任务,而现有Transformer缺乏高效的原生知识查找机制,导致计算资源浪费 [4] - 通过将部分稀疏参数预算从混合专家(MoE)重新分配给Engram记忆模块,可以在固定计算预算下显著提升模型性能,形成一条U型最优曲线 [15][17] - 实验表明,在27B参数规模下,采用Engram的模型在知识、通用推理、代码和数学等多个领域的基准测试中均显著超越同规模纯MoE模型,甚至有效“加深”了网络推理深度 [21][23] - Engram的确定性寻址机制支持存储与计算的解耦,允许将海量参数表卸载至CPU内存,在推理时带来可忽略的额外开销,实现了硬件感知的高效设计 [27][29][31] 技术原理与设计 - 核心问题识别:现有Transformer模型在识别如“Diana, Princess of Wales”这类实体时,需要消耗多达6层注意力进行逐层特征拼凑,本质是用昂贵的运行时计算来重建静态查找表,浪费了可用于高层推理的网络深度 [4][5] - 解决方案:回归查表思想,将经典的N-gram方法嵌入Transformer,通过哈希查找以O(1)时间复杂度捕获局部依赖和静态知识 [6] - 关键创新 - Engram模块:在Transformer层间插入Engram模块,对当前token及前几个token组成的N-gram进行哈希查找,从一个巨大的嵌入表中直接取出对应向量 [6] - 解决传统N-gram痛点: - 通过压缩tokenizer,将语义相同但形式不同的token归为一类,使128k词表的有效规模减少23% [11] - 使用多个哈希函数将N-gram映射到固定大小的嵌入表,解决了存储爆炸问题,并通过多哈希头减少查找冲突 [11] - 上下文感知门控:引入门控机制,用当前隐藏状态作为Query来评估检索到的记忆与上下文的匹配度,自动屏蔽噪声,实现静态模式调用与动态计算的灵活切换 [8] 实验与性能分析 - 稀疏性分配研究:固定总参数量和每token激活参数量,在MoE专家和Engram记忆间重新分配参数预算,实验得出一条U型曲线 [15] - 纯MoE并非最优,将约20%到25%的稀疏参数预算分给Engram时,模型验证集loss最低 [17] - 在100亿参数规模下,最优配置比纯MoE基线的loss降低0.0139 [17] - 最优分配点在不同计算预算下稳定在ρ=75%到80%之间(即20%-25%预算给Engram) [17] - 27B规模模型对比:在激活参数量均为38亿、训练token均为2620亿的条件下进行对比 [18] - Engram-27B:总参26.7B,分配5.7B参数给Engram记忆模块 [18][22] - Engram-40B:总参39.5B,Engram记忆参数增至18.5B [18][22] - 性能提升结果:与纯MoE-27B模型相比,Engram-27B在多项基准测试中取得显著提升 [21][22] - 知识密集型任务:MMLU提升3分(从57.4到60.4),CMMLU提升4.0分(从57.9到61.9),TriviaQA提升1.9分(从48.8到50.7) [21] - 通用推理任务:BBH大幅提升5.0分(从50.9到55.9),ARC-Challenge提升3.7分(从70.1到73.8),DROP提升3.3分(从55.7到59.0) [21] - 代码数学任务:HumanEval提升3.0分(从37.8到40.8),MATH提升2.4分(从28.3到30.7),GSM8K提升2.2分(从58.4到60.6) [21] - 性能提升原因分析:Engram让模型早期层无需做特征组合的“苦力活”,从而“加深”了网络有效深度 [23] - KL散度曲线显示Engram模型预测收敛更快 [23] - CKA相似度分析表明,Engram-27B第5层的表征与MoE基线第12层的表征最相似 [23] - 扩展性与长上下文:Engram-40B进一步增加记忆参数后,大部分任务性能持续提升,且训练后期损失仍在下降,说明记忆容量未饱和 [25] - 长上下文场景提升显著:在RULER测试集上,Multi-Query NIAH准确率从84.2跃升至97.0,Variable Tracking从77.0提升至89.0 [25][26] 工程实现与效率 - 训练阶段:巨大的词表参数(高达100B)需要拆分到多个GPU上,并通过All-to-All通信机制传递记忆片段 [27] - 推理优化:利用Engram查找索引的确定性,可将巨大的嵌入表卸载到CPU内存,并通过PCIe异步预取实现通信与计算的重叠 [29] - 在H800 GPU上测试,为一个4B密集模型增加1000亿参数的CPU卸载Engram表,吞吐量从9031 token/s降至8858 token/s,额外开销在3%以内 [29] - 8B密集模型增加后,吞吐量从6315 token/s降至6140 token/s,开销同样低于3% [29] - 缓存策略:利用N-gram访问遵循Zipfian分布的特点,可设计多级缓存(高频放GPU显存、中频放CPU内存、长尾放NVMe SSD),进一步压缩有效延迟 [30]
DeepSeek开源大模型记忆模块,梁文锋署名新论文,下一代稀疏模型提前剧透