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

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

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

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

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

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

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

更重要的是，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

运营/排版：何晨龙