谷歌发布新架构,用一半内存实现推理速度翻倍

互联网资讯
本条新闻为单纯事实消息的时事新闻,转载自DeepTech深科技,版权归源站所有,如有侵权,烦请联系我们

谷歌 DeepMind 与韩国科学技术院(KAIST)的研究人员近日联合发布了一种名为“Mixture-of-Recursions”(MoR)的新型语言模型架构。据称能在保持模型性能的同时,实现了推理速度翻倍、训练计算量减少并降低了约 50% 的 KV 缓存内存使用。论文一经发布,便在社交媒体上引发了广泛关注,甚至有评论称其或许是“Transformer 杀手”。

自 2017 年问世以来,Transformer 架构已成为大型语言模型的技术基础,目前几乎所有先进模型都建立在这一架构之上。但随着模型规模越来越大,Transformer 架构对计算和内存资源的需求也日益增加,训练和部署成本十分高昂。过去的效率优化方法通常只关注单一方向,例如通过参数共享来缩小模型尺寸,或通过自适应计算来按需分配算力,但很少能将多个效率目标同时优化。

MoR 架构的核心创新在于将递归计算与动态路由机制相结合,在一个统一框架内解决多重效率问题。在标准的 Transformer 模型中,输入文本的每一个 token 都要经过相同数量的计算层处理。MoR 则改变了这一模式,它允许不同的 token 根据自身的复杂度,接受不同深度的处理。


图丨MoR 架构的概览(来源:arXiv)

具体来说,MoR 使用共享的参数块来提升参数效率,同时通过一个轻量级的“路由器”,来决定每个 token 需要经历多少次递归计算。研究团队测试了多种路由策略,例如“专家选择”(expert-choice)和“token 选择”(token-choice),以平衡计算负载和避免信息处理中的逻辑问题。在参数共享方面,测试结果表明“Middle-Cycle”策略表现最好,该策略保持模型的第一层和最后一层使用独立的参数,只在中间层之间共享权重,以此在参数效率和模型的表达能力之间取得较好的平衡。


图丨MoR 的核心组件(来源:arXiv)

内存管理是 MoR 的另一项关键改进。即使参数共享了,传统递归模型在每一层递归时仍会产生独立的 KV 缓存,导致内存占用依然很高。MoR 提出了两种新的 KV 缓存策略来应对这一问题。一种是“递归式缓存”,只为被路由到特定递归步骤的 token 存储 KV 数据,并将注意力计算限制在这些本地数据中,从而有效降低了 KV 缓存的内存占用和数据读写量。另一种“递归共享”策略则更进一步,利用所有 token 都会经过第一个递归块的特性,只在第一步缓存 KV 数据,供后续所有递归步骤重复使用,以此最大化地节省内存。

研究团队在 1.35 亿到 17 亿参数的多个模型规模上进行了测试。结果表明,在相同的训练计算预算下,一个采用 MoR 架构的模型,虽然其参数量比基准 Transformer 模型少了近一半,但在多项少样本学习任务中的平均准确率达到了 43.1%,超过了基准模型的 42.3%。


(来源:arXiv)

更重要的是,MoR 更高的计算效率,使其能在相同的计算预算内处理更多的训练数据,这反过来也提升了模型的最终性能。在固定训练数据量的对比实验中,一个 MoR 配置在使用少 25% 训练计算量的情况下,其性能依然超过了基准模型,同时训练时间减少了 19%,峰值内存使用也降低了 25%。

在推理性能上,MoR 的优势更加明显。它采用了一种连续深度批处理技术,可以将处在不同计算阶段的 token 组合到同一个批次中进行处理,因为它们都使用相同的参数块。该技术配合模型的早期退出机制,显著提升了处理吞吐量。在 3.6 亿规模的模型测试中,MoR-4 配置在特定设置下实现了高达 2.06 倍的推理加速。

研究还发现,MoR 模型在处理不同类型 token 时,表现出一种与语义重要性相关的模式。内容丰富的 token,如“People”或“defensively confident”,会被分配更多的递归次数(三次);而功能性词汇,如“and”,通常只需要较少的递归次数。这表明,模型学会了将更多的计算资源用于处理更重要的信息。

实际上,MoR 的出现也离不开谷歌此前相关研究的铺垫。在早前的 Mixture-of-Depths(MoD)等技术中,谷歌 DeepMind 就已经探索了动态分配计算资源的方法。同时,递归 Transformer 作为一种参数共享技术,也为 MoR 提供了理论基础。

MoR 进一步延续了这些研究对 AI 效率优化的探索,即从单一维度的优化,转向参数、计算和内存等多个维度的协同优化。这对于降低大语言模型的部署和使用成本具有较大的实际意义。

总体而言,虽然现在断言 MoR 能否完全替代 Transformer 还为时过早,但它确实为未来的语言模型架构设计,提供了一个在性能和效率上都极具潜力的发展方向。

参考资料:

1.https://arxiv.org/abs/2507.10524

运营/排版:何晨龙

点赞 0 收藏(0)  分享
0个评论
  • 消灭零评论