大型语言模型已展现出卓越的能力，但其部署仍面临巨大的计算与内存开销所带来的挑战。随着模型参数规模扩大至数千亿级别，训练和推理的成本变得高昂，阻碍了其在许多实际应用中的推广与落地。 这促使研究者们围绕两个主要方向积极探索高效化技术：一是通过权重共享提升参数效率，二是根据输入复杂度动态分配计算资源，实现自适应计算 而一切的基础则是 Transformer 架构。这次谷歌又一次坐不住了，在递归 Transformer 的基础上再次进化，发布了名为 Mixture-of-Recursions （MoR）的新 LLM 模型架构，有学者认为这个新架构「有潜力成为 Transformer 的杀手」 MoR 框架引入了一种统一方法，同时实现参数共享和自适应计算。与以往孤立处理参数减少或自适应计算的方法不同，MoR 将动态 token 级路由集成到参数高效的递归 Transformer 中，创建了一种协同架构，实现了「无需承担大模型成本的大模型质量」。 在本研究中，作者提出了Mixture-of-Recursions（MoR），一个统一框架，旨在充分发挥递归 Transformer 的潜力（如图所示）。MoR 通过端到端训练轻量级的路由，为每个 token 分配专属的递归深度：即根据每个 token 所需的「思考」深度，决定共享参数模块对其重复应用的次数，从而将计算资源精准地分配到最需要的部分。 这种基于 token 的动态递归机制，天然支持递归级的键值（KV）缓存，即根据每个 token 所分配的递归深度，有选择地存储和检索对应的键值对。该定向缓存策略显著降低了内存带宽压力，从而无需后处理操作就能提升推理吞吐量。 左图显示每个递归步骤由一组固定堆叠的层与一个路由（router）组成。路由会为每个 token 决定是继续通过递归块（进行下一轮递归），还是提前退出。图中灰色方框表示的就是该递归模块。 中图展示了完整的模型结构。对于每个 token，模型会根据路由的决策，对共享的递归模块最多重复应用 N_r 次，从而实现动态递归深度。 右图 是一个 token 级的路由示例，展示了每个 token 的实际递归次数。颜色越深的单元格表示该 token 在当前递归步骤中被激活参与了计算。底部显示了每个文本 token 实际经历的递归次数，使用不同颜色编码标注（如：紫色 1 次、蓝色 2 次、红色 3 次）。 这会形成一个逐步收窄的计算漏斗：复杂的 token 会接受更深层的处理，而简单的 token 会更早退出计算。然而，这种方法在训练过程中会面临因果性挑战，因此需要采用一些专门的技术： 辅助路由（Auxiliary Router）：一个单独的轻量网络，用于在推理阶段预测 top-k 的选择；辅助损失（Auxiliary Loss）：采用二元交叉熵损失，鼓励路由器输出接近二值化，从而实现更清晰的阈值判断。 2、Token 选择路由（Token-Choice Routing）：每个 token 在一开始就做出一次路由决策，以确定其完整的递归深度。这种方式天然满足因果性要求，但需要引入负载均衡机制： 在等效的训练预算（16.5 × 10?? FLOPs）下，作者将 MoR 模型与 原始 Transformer 及递归 Transformer 进行了比较。如表 3 所示，MoR 模型采用「专家选择路由」（expert-choice routing）和两层递归（N_r = 2），实现了更低的验证损失，并在少样本学习平均准确率上超越基线（43.1% 对比 42.3%）。值得注意的是，尽管参数量减少近 50%，MoR 依然取得了更优性能。这一优势归因于其更高的计算效率，使其在相同 FLOPs 预算下处理了更多训练 token。 为了单独分析架构差异，作者在固定训练 token 数量（20B）下进行评估。结果显示，MoR 模型（N_r = 2）在使用比基线模型少 25% 的 FLOPs情况下，仍实现了更低的验证损失和更高的准确率。 这种理论上的效率优势在实践中带来了显著提升：相比于基线模型，MoR 的训练时间减少 19%，峰值内存使用量下降 25%。这些改进来源于 MoR 的层级过滤机制与递归注意力策略，有效缩短了序列长度，实现了更优的计算 - 准确率平衡，甚至在预训练阶段亦是如此。 作者还评估了 MoR 架构下的一些设计变体，特别是在轻量级的三层递归设置（N_r = 3）下仍具可比性的 Vanilla 模型。结果显示，采用「token 选择路由」时准确率较低（40.0%），而「专家选择路由」表现更好（42.6%），表明路由粒度对模型性能具有关键影响。 此外，采用 KV 缓存共享略微降低了性能，但显著提升了内存效率。在对内存占用有要求的部署场景中，这种性能与资源之间的权衡是值得接受的。 如图所示，MoR 在所有模型规模和计算预算下都稳定优于递归基线模型。尽管在最小规模（135M）下表现略低于原始 Transformer（可能由于递归容量瓶颈所致），但这种差距随着模型规模的扩大迅速缩小。在参数量超过 360M 时，MoR 不仅能够追平，而且在低至中等预算下常常超越原始 Transformer。 (b) 参数共享策略对递归 Transformer 的影响：展示了递归次数为 = 3 的递归 Transformers 在四种不同参数共享策略下的负对数似然（NLL）表现。这些模型在 100 亿 token 上预训练。图中红色虚线表示完整规模的原始 Transformer，黑色虚线表示参数量约为其三分之一的基线模型。 MoR 通过证明参数效率和自适应计算可以协同结合，为高效的 LLM 架构建立了一种新的范例。该框架实现大型模型质量的同时显著降低了计算需求的能力，解决了语言建模中的基本可扩展性挑战。 最重要的是，MoR 为开发更具认知启发性的 AI 系统提供了基础。其在生成过程中能够按每个 token 自适应地分配「思考深度」的能力，与语言模型中新兴的潜在推理和内部思考研究相一致。这表明 MoR 可以作为一个平台，用于研究模型如何学会在保持常规任务效率的同时，更深入地思考复杂问题。
