Google发布MoR架构：两倍推理速度，节省50%内存 原创 精华

Google DeepMind发布的递归混合（MoR）：

（1）推理速度提升2倍，因为token可以提前退出共享循环

（2）训练时减少了大量计算，每层的注意力机制FLOP减少约一半，相同预算可以训练更多数据

论文展示了一种全新的方法，让大语言模型在自己的推理循环中学会规划步骤，而不是硬编码单一的推理链。

其次，它证明了混合器理念的可扩展性。通过混合多个小型递归专家，让模型选择下一个调用哪个专家，团队在数学和编程基准测试中提升了准确性，而无需增加参数数量。

递归混合（MoR）在内存中保持1个层堆栈，对困难的token进行循环处理，在准确性和速度上都超越了更大的普通模型。它通过让一个小型路由器选择每个token获得多少次循环来实现这一点，然后仅为保持活跃的token保存缓存。更少的权重、更少的FLOP、更少的内存，但在135M-1.7B规模上获得了更好的困惑度。

扩展Transformer通常意味着堆叠更多层，并在内存和计算上付出代价。MoR颠覆了这一习惯。它共享1个紧凑的块，根据token难度运行最多4次，当路由器判断"完成"时提前跳出循环。

图1描绘了这种token级别的阶梯结构，简单单词快速退出，复杂单词继续攀升。

通过共享层来减少内存：递归Transformer在深度上绑定权重，但过去的工作仍然让每个token通过每个循环。MoR保持权重绑定的思想，但添加了"中间循环"共享，因此只有第一层和最后一层保持独特，而中间的所有层在每个循环中重用一个小型三元组层。这种选择保持梯度稳定，将独特参数减少约3倍而不失表现力。由于相同权重重复使用，完全分片数据并行每步只收集一次权重，也减少了通信。

一个小型路由器为每个token评分，然后要么将其送回共享块，要么让它继续前进。中间的示意图显示循环重复到设定上限，因此简单token提前退出，复杂token继续工作。右侧的热图证明大多数token跳过后期循环，因此计算和内存随任务难度而非序列长度缩放。这种选择性深度带来大约2倍的推理速度提升和约50%的键值缓存减少，同时保持准确性。

让路由器决定深度：一个轻量级线性门在第一个循环后立即为每个token评分。得分靠前的token获得通过共享块的另一次机会，其余token以当前状态继续。训练使用辅助损失，将选中的分数推向1，丢弃的分数推向0，这样门在推理时表现相同并避免因果泄漏。这种专家选择路由保持计算预算固定，仍然将深度聚焦在有帮助的地方。

随token而非层增长的缓存：标准解码为每层每个token存储键值对，一旦循环就会爆炸式增长。MoR只缓存当前循环中仍然活跃的token，因此更深的循环触及越来越小的内存片段。如果内存比计算更紧张，另一个变体为所有后续循环重用第一循环缓存，用困惑度的微小损失换取RAM和预填充延迟的大幅削减。

效果：在相同的16.5E FLOP预算下，具有3个循环的118M参数MoR在平均少样本准确率(43.1% vs 42.3%)和验证损失方面击败了315M的普通基线。当计算固定但数据变化时，MoR继续获胜，证明收益不仅仅是分词器的运气。token选择路由比专家选择落后约2个点，因此每循环的智能容量分配很重要。

生成时的速度提升：深度批处理将当前处于同一循环的token分组，因此GPU保持满载而不是等待长序列完成。使用3个循环，批次32时吞吐量提升1.60倍，当批次增长到H100上的内存限制时提升2.06倍。更多循环意味着更多早期退出，因此MoR-4在速度上比普通模型快一倍，困惑度仅下降一小部分。

扩展性和挑战 等FLOP图显示MoR在135M时落后于普通模型，因为单个块变得拥挤，但在360M时匹配，在730M和1.7B时领先，同时仅持有三分之一的参数。计算最优扩展倾向于在预算攀升时扩大共享块而非添加数据，因为每个额外参数在每个循环中都被重用。当路由器面对非常不均匀的文本时，MoR仍需要更好的负载均衡，其固定的top-k使动态改变延迟目标变得棘手。

论文标题：Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation

论文链接：​​​https://arxiv.org/abs/2507.10524v1​​

本文转载自​​​AI帝国​​​，作者：无影寺