Mixture-of-Recursions（MoR）混合递归Transformer

谷歌DeepMind联手顶尖研究机构推出革命性架构Mixture-of-Recursions（MoR），有望彻底改变AI模型的计算方式，速度翻倍且内存减半，以下对《Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation》论文进行解读：

---

核心贡献与创新

1. 统一高效框架
   • 首次将参数共享（Parameter Sharing）、Token级自适应计算深度（Adaptive Token-Level Computation）和内存高效KV缓存（Memory-Efficient KV Caching）整合到单一架构（MoR）中，解决了现有方法仅聚焦单一效率维度的问题。
2. 动态递归路由机制
   • 提出轻量级端到端训练的路由器，动态分配每个Token的递归深度（递归次数），避免传统早退方法（Early-Exiting）需额外训练阶段的问题。
3. 递归感知KV缓存
   • 设计两种KV缓存策略：
     • 递归级缓存（Recursion-wise Caching）：仅缓存当前递归步活跃Token的KV对，显著减少内存占用和I/O开销（缓存大小降至常规Transformer的 (\frac{N_r+1}{2N_r}))。
     • 递归共享（Recursive Sharing）：首层递归的KV对全局复用，降低预填充（Prefill）延迟（缓存大小降至 (\frac{1}{N_r}))。
4. 性能优势
   • 在135M–1.7B参数规模下，MoR在相同训练FLOPs下：
     • 验证困惑度（Perplexity）显著低于普通Transformer（Vanilla）和固定深度递归基线。
     • Few-shot准确率提升（最高+0.8%）。
     • 推理吞吐量最高提升 2.06倍（连续深度批处理 + 早退机制）。

---

算法原理

1. 递归Transformer基础

• 参数共享：将模型分为 (N_r) 个递归块，每个块复用共享参数 (\Phi')（例如Middle-Cycle策略保留独立首尾层，中间层循环复用）。
• 计算过程：Token表示 (h_t) 通过递归块更新：
  [
  h_t^{r+1} = f(h_t^r; \Phi'), \quad r=1,\dots,N_r
  ]

2. 动态路由机制

• Expert-Choice路由：

  • 每层递归步 (r) 的路由器计算Token得分 (g_t^r = G(\theta_r^\top \mathcal{H}_t^r))，选择Top-(k) Token继续递归（其余退出）。
  • 采用层级过滤（Hierarchical Filtering）：仅上轮选中Token可进入下一步。
  • 公式：
    [
    \mathcal{H}_t^{r+1} = \begin{cases}
    g_t^r f(\mathcal{H}_t^r, \Phi') + \mathcal{H}t^r & \text{if } g_t^r > P\beta(G^r) \
    \mathcal{H}_t^r & \text{otherwise}
    \end{cases}
    ]

• Token-Choice路由：

  • 初始路由决策分配完整递归深度 (i)，Token固定执行 (i) 次递归。
  • 公式：
    [
    \mathcal{H}_t^{r+1} = \begin{cases}
    g_t^r f(\mathcal{H}_t^r, \Phi') + \mathcal{H}_t^1 & \text{if } r=i \
    g_t^r f(\mathcal{H}_t^r, \Phi') & \text{otherwise}
    \end{cases}
    ]

3. KV缓存策略

• 递归级缓存：
  • 仅缓存当前递归步活跃Token的KV对，注意力计算限于同层缓存Token，降低内存和计算量（注意力FLOPs降至 ((k/N_{ctx})^2))。
• 递归共享：
  • 首层递归的KV对全局复用，后续递归直接调用，减少预填充开销（注意力FLOPs降至 (k/N_{ctx}))。

4. 训练与推理优化

• 连续深度批处理（Continuous Depth-wise Batching）：
  不同递归深度的Token批量处理，消除GPU空闲时间。
• 计算最优缩放：MoR更依赖模型规模而非数据量，适合短步大模型训练。

---

与现有方法的差异

维度	现有方法	MoR创新
参数共享	递归Transformer（如Universal Transformer）固定深度复用参数，无自适应计算。	动态递归深度：路由器按Token复杂度分配深度，实现参数复用 + 计算自适应。
自适应计算	- 早退（Early-Exiting）：需额外训练阶段，性能下降。 - Mixture-of-Depths（MoD）：每层独立路由，无参数共享。	端到端路由：轻量路由器联合训练，避免性能损失； 递归块级路由（非层间路由），与参数共享耦合。
KV缓存	固定深度模型需存所有层KV对，内存开销大； 共享KV方法（如Layer-condensed）未考虑递归架构。	递归感知缓存： - 递归级缓存：仅存活跃Token，节省内存； - 递归共享：首层KV复用，加速预填充。
推理效率	固定深度递归无法减少简单Token计算量。	早退+深度批处理：退出Token释放资源，吞吐量提升2倍。

---

关键性能验证

1. 效率-精度权衡：
   • MoR在360M+参数规模下超越普通Transformer，且参数量仅其1/3（(N_r=3)）。
2. 路由有效性：
   • Expert-Choice + 辅助损失实现完美负载均衡（Router输出0/1分离）。
3. 语义对齐：
   • 关键Token（如实体词）分配更深递归（e.g., "Drugs"递归3次 vs. 虚词"and"递归1次）。

---

局限与未来方向

1. 推理时路由调整：预训练后无法动态修改Top-(k)值（Router输出二值化）。
2. 多模态扩展：框架兼容视觉/语音等非文本模态。
3. 稀疏性结合：与剪枝/量化技术结合进一步优化计算。

总结：MoR通过动态递归路由与递归感知缓存，首次统一参数效率与自适应计算，在模型规模、推理速度、能耗间建立新的Pareto前沿，为低成本部署大模型提供新路径。