有望成为Transformer杀手,谷歌DeepMind架构MoR实现两倍推理速度

机器之心报道

编辑：冷猫

大型语言模型已展现出卓越的能力，但其部署仍面临巨大的计算与内存开销所带来的挑战。随着模型参数规模扩大至数千亿级别，训练和推理的成本变得高昂，阻碍了其在许多实际应用中的推广与落地。

这促使研究者们围绕两个主要方向积极探索高效化技术：一是通过权重共享提升参数效率，二是根据输入复杂度动态分配计算资源，实现自适应计算。

而一切的基础则是 Transformer 架构。这次谷歌又一次坐不住了，在递归 Transformer 的基础上再次进化，发布了名为 Mixture-of-Recursions （MoR）的新 LLM 模型架构，有学者认为这个新架构「有潜力成为 Transformer 的杀手」。

MoR 框架引入了一种统一方法，同时实现参数共享和自适应计算。与以往孤立处理参数减少或自适应计算的方法不同，MoR 将动态 token 级路由集成到参数高效的递归 Transformer 中，创建了一种协同架构，实现了「无需承担大模型成本的大模型质量」。

论文标题：Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation
论文链接：alphaxiv.org/abs/2507.10524

在本研究中，作者提出了 Mixture-of-Recursions（MoR），一个统一框架，旨在充分发挥递归 Transformer 的潜力（如图所示）。MoR 通过端到端训练轻量级的路由，为每个 token 分配专属的递归深度：即根据每个 token 所需的「思考」深度，决定共享参数模块对其重复应用的次数，从而将计算资源精准地分配到最需要的部分。

这种基于 token 的动态递归机制，天然支持递归级的键值（KV）缓存，即根据每个 token 所分配的递归深度，有选择地存储和检索对应的键值对。该定向缓存策略显著降低了内存带宽压力，从而无需后处理操作就能提升推理吞吐量。

因此，MoR 在一个统一架构中同时实现了三项优化：

1. 参数共享：通过权重绑定显著减少参数量；

2. 计算路由：通过 token 路由避免冗余计算（减少无效的 FLOPs）；

3. 递归级缓存：按需缓存键值对，降低内存访问开销，提升吞吐效率。

MoR 以此兼顾了性能、效率和计算资源利用率，是递归 Transformer 架构的重要演进。

MoR 总览

左图显示每个递归步骤由一组固定堆叠的层与一个路由（router）组成。路由会为每个 token 决定是继续通过递归块（进行下一轮递归），还是提前退出。图中灰色方框表示的就是该递归模块。

中图展示了完整的模型结构。对于每个 token，模型会根据路由的决策，对共享的递归模块最多重复应用 N_r 次，从而实现动态递归深度。

右图是一个 token 级的路由示例，展示了每个 token 的实际递归次数。颜色越深的单元格表示该 token 在当前递归步骤中被激活参与了计算。底部显示了每个文本 token 实际经历的递归次数，使用不同颜色编码标注（如：紫色 1 次、蓝色 2 次、红色 3 次）。

核心架构与方法

递归 Transformer 模型基础

MoR 基于递归 Transformer，通过跨多个层共享权重来实现参数效率，而不是为每一层维护唯一参数。

多种参数共享模式

循环：递归块循环复用（例如，重复层 [0,1,2]）

序列：每个块连续地重用相同的层（例如，[0,0,0]，[1,1,1]）

中周期 / 中序列：保留首层和尾层的独特参数，同时共享中间层权重的变体

中周期策略在经验上被证明最为有效，在参数效率和模型表达能力之间提供了最佳平衡。

动态路由机制

MoR 的核心创新在于其动态路由系统，该系统根据单个 token 的计算需求分配不同的递归深度。实现了两种主要的路由策略：

1、在每一个递归步骤 r，路由器 G 会对每个 token 当前的隐藏状态进行评估，并选择 top-k 的 token 继续进入下一层递归。

这会形成一个逐步收窄的计算漏斗：复杂的 token 会接受更深层的处理，而简单的 token 会更早退出计算。然而，这种方法在训练过程中会面临因果性挑战，因此需要采用一些专门的技术：

辅助路由（Auxiliary Router）：一个单独的轻量网络，用于在推理阶段预测 top-k 的选择；
辅助损失（Auxiliary Loss）：采用二元交叉熵损失，鼓励路由器输出接近二值化，从而实现更清晰的阈值判断。

2、Token 选择路由（Token-Choice Routing）：每个 token 在一开始就做出一次路由决策，以确定其完整的递归深度。这种方式天然满足因果性要求，但需要引入负载均衡机制：

均衡损失（Balancing Loss）：正则项，用于鼓励递归深度之间的均匀分布；
路由偏置（Router Bias）：在训练过程中动态调整各专家的偏置项。

KV缓存优化

MoR 引入了两种针对动态深度模型的专用 KV 缓存策略：

递归式 KV 缓存：在每个递归深度选择性地缓存 KV 对，注意力仅限于本地缓存的 token：

递归 KV 共享：所有 token 都会遍历第一个递归块，该块的 KV 对会被缓存并在后续递归中重用：

实验结果

MoR 在相同计算预算下以更少参数优于基线模型

在等效的训练预算（16.5 × 10¹⁸ FLOPs）下，作者将 MoR 模型与原始 Transformer 及递归 Transformer 进行了比较。如表 3 所示，MoR 模型采用「专家选择路由」（expert-choice routing）和两层递归（N_r = 2），实现了更低的验证损失，并在少样本学习平均准确率上超越基线（43.1% 对比 42.3%）。值得注意的是，尽管参数量减少近 50%，MoR 依然取得了更优性能。这一优势归因于其更高的计算效率，使其在相同 FLOPs 预算下处理了更多训练 token。

在固定 FLOPs（16.5e18）与固定 token 数（20B）两种设定下，对 MoR、递归 Transformer（Recursive）和标准 Transformer（Vanilla）进行了对比。

MoR 在相同数据下以更少计算超越基线模型

为了单独分析架构差异，作者在固定训练 token 数量（20B）下进行评估。结果显示，MoR 模型（N_r = 2）在使用比基线模型少 25% 的 FLOPs 情况下，仍实现了更低的验证损失和更高的准确率。

这种理论上的效率优势在实践中带来了显著提升：相比于基线模型，MoR 的训练时间减少 19%，峰值内存使用量下降 25%。这些改进来源于 MoR 的层级过滤机制与递归注意力策略，有效缩短了序列长度，实现了更优的计算 - 准确率平衡，甚至在预训练阶段亦是如此。

MoR 性能受路由与缓存策略影响

作者还评估了 MoR 架构下的一些设计变体，特别是在轻量级的三层递归设置（N_r = 3）下仍具可比性的 Vanilla 模型。结果显示，采用「token 选择路由」时准确率较低（40.0%），而「专家选择路由」表现更好（42.6%），表明路由粒度对模型性能具有关键影响。

此外，采用 KV 缓存共享略微降低了性能，但显著提升了内存效率。在对内存占用有要求的部署场景中，这种性能与资源之间的权衡是值得接受的。

不同计算预算下的验证损失，在四种模型规模（135M、360M、730M 和 1.7B 参数量）下进行对比。

MoR 是一种可扩展且参数高效的架构

如图所示，MoR 在所有模型规模和计算预算下都稳定优于递归基线模型。尽管在最小规模（135M）下表现略低于原始 Transformer（可能由于递归容量瓶颈所致），但这种差距随着模型规模的扩大迅速缩小。在参数量超过 360M 时，MoR 不仅能够追平，而且在低至中等预算下常常超越原始 Transformer。

结果表明：MoR 是标准 Transformer 的一种可扩展且高效的替代方案。

它在使用显著更少参数的情况下，取得了优异的验证性能，因而非常适合用于预训练以及大规模部署。

MoR 通过连续深度批处理提升推理吞吐量

如图 4a 所示，在上述两种批处理设置下，所有 MoR 变体的推理吞吐量均高于 Transformer 基线模型。

随着递归深度的增加，更多 token 可在早期阶段退出，从而减少 KV 缓存的使用。这一机制显著提升了推理速度：

例如，MoR-4 在最大批量配置下的吞吐量提升可达 2.06 倍。

虽然这伴随着轻微的性能下降，但考虑到吞吐量的大幅提升，这种权衡是值得的。

实验结果表明：将深度方向批处理范式与提前退出机制相结合，能够显著提升 MoR 在实际部署场景中的推理吞吐能力。

图中 (a) 在固定批量和最大批量设置下，MoR 与原始 Transformer 在推理吞吐量与对数似然（log-likelihood）之间的帕累托前沿（Pareto frontier）表现对比。

(b) 参数共享策略对递归 Transformer 的影响：展示了递归次数为 = 3 的递归 Transformers 在四种不同参数共享策略下的负对数似然（NLL）表现。这些模型在 100 亿 token 上预训练。图中红色虚线表示完整规模的原始 Transformer，黑色虚线表示参数量约为其三分之一的基线模型。