潜空间推理系列 · 02 / 03
Mixture-of-Recursions
哪些 token 值得多算几轮？
同时减少模型权重、冗余 token 计算和 KV 缓存开销，让有限训练与部署预算得到更好的语言建模质量。
MoR 的贡献是把“谁继续循环”和“深层还能看见谁”一起设计。两轮配置在所测预算下兼顾质量与效率，但收益取决于路由、缓存和调度的组合；它尚未证明路由器识别了抽象的推理难度。
“专家”在这里代表不同递归深度
MoR 不是为 token 选不同的前馈专家网络，而是在共享层上选择做几轮计算。主要配置保留独立的首尾层，中间层循环复用（Middle-Cycle）。KV 缓存保存过去 token 的键和值，供后续注意力读取；权重共享本身不会自动让各轮 KV 也相同。[§2.1]
让 token 提前退出，会留下深层上下文缺口
如果某个 token 两轮后停止，后来的 token 到第三轮时就找不到它的第三轮 KV。补算会吃掉节省的预算。MoR 的关键选择是：训练时就规定每轮能读取哪些位置，让退出规则与注意力结构相容。
路由选择深度，缓存定义可见上下文
Expert-choice 每轮给 token 打分，在当前活跃集合中按容量选取继续计算的位置。三轮时，各轮处理原序列的 3/3、2/3、1/3；比例预设，具体位置由模型学。Token-choice 则在开始时一次性为每个 token 选定深度，因果性更直接，却更难平衡不同深度的负载。[§2.2.1]
主方案：按轮缓存。只有该轮活跃 token 才产生 KV，深层注意力只读取同轮已激活的历史位置。这同时减少查询数和键值数；训练／预填充的该层注意力项可随保留比例近似平方缩小，但不能把这个比例当成整个模型或逐 token 解码的加速比。已退出 token 的信息仍可能经浅层混入活跃状态。[§2.2.2]
另一方案：跨轮共享首轮 KV。首轮包含所有 token，后面继续读取这份完整历史，因此更省缓存容量，但每轮仍需读较长的 KV。它与主方案是两种取舍；“缓存更小”不保证“读缓存更少”，也不能把各自最佳收益相加。
因果性是必要补丁。训练时在整段序列上做 top-k，会让选择依赖未来位置的分数。表现较好的方案用辅助二分类损失，把入选与未入选分数推向 1 与 0，以便生成时只凭当前可用信息决策。作者报告了良好的分离，但这不是任意分布下的因果等价保证。[附录 A.2][图 5b]
把等算力和等数据两张账一起看
表 3 基于 360M 级基础配置，参数列不含 embedding；使用 FineWeb-Edu 与六项 few-shot 基准，准确率为混合任务均分（%）。下列 MoR 均为两轮、expert-choice、按轮缓存。[表 3]
| 配置 | 参数 | tokens | FLOPs ×10¹⁸ | NLL ↓ / 均分 ↑ |
|---|---|---|---|---|
| Vanilla | 315M | 20B | 16.5 | 2.7824 / 42.3 |
| 固定两轮 Recursive | 167M | 20B | 16.5 | 2.8079 / 42.6 |
| MoR：等 FLOPs | 167M | 27B | 16.5 | 2.7511 / 43.1 |
| MoR：等数据 | 167M | 20B | 12.3 | 2.7749 / 42.9 |
等算力下，MoR 通过少算部分 token 换来了更多数据；等数据下，仍用约少 25% 的训练 FLOPs 获得略高均分。这是联合效率设计的证据。相反，三轮 token-choice 的均分只有 40.0%，低于 expert-choice 的 42.6%，说明路由方式不是无关紧要的实现选择。
最高 2.06× 吞吐有严格条件：H100 显存预算推导的最大批量、平均生成 256 tokens、无输入前缀，并结合按深度持续组批。MoR-4 用 51 的批量，对照为 32，且对照补到 34 层；计时排除 KV 缓存写入与更新。因此它是特定系统实验，不能直接外推为长上下文或单请求服务的两倍加速。[§3.3][附录 E]
自适应 token，不等于任意预算弹性
更大尺度并非全面胜出。附录 1.7B 级、68.5×10¹⁸ FLOPs 对照中，Vanilla 均分 48.9%，两轮 MoR 48.4%，三轮 MoR 46.7%。省参数仍有价值，但不支持“所有规模同算力更强”。[表 7]
预算主要受训练配置约束。辅助损失让路由分数接近二值，训练后改变容量反而困难；图 5c 展示训练深度范围内多计算的收益，附录表 8 是分别训练的不同深度模型。两者都不证明一个模型可以无限增深。[§5.3][§7.1]
“难 token”首先是难预测。位置上的递归服务于下一个 token 预测；标点、词片段和上下文都会影响分配。将它升级成数学推理难度识别仍需专门实验。下一篇 LoopFormer 转向另一问题：预算由用户给定后，怎样把短路径也训练好。[§5.2][§7.1]
回到原文与系列
- Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation · 摘要与版本
- 公开 PDF：本页主要依据用户附件，核对方法、主结果及相关附录；官方仓库仅作为实现入口，未运行复现。
- 作者官方仓库。
实验数字为作者报告；比较口径、外推边界与系列关系为本页分析。