2507.10524_Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation¶

首页: https://arxiv.org/abs/2507.10524
PDF: https://arxiv.org/pdf/2507.10524
引用: 19(2025-12-14)
组织

总结¶

From Moonlight¶

三句摘要¶

✨ Mixture-of-Recursions (MoR) 提出了一种统一的Transformer框架，它将parameter sharing和adaptive computation结合起来，旨在应对大型语言模型的计算和内存挑战。
🚀 MoR通过轻量级router动态分配individual token的recursion depth，并利用recursion-wise KV caching或recursive KV sharing策略，从而只对活跃token进行quadratic attention computation。
📈 实验证明，MoR在相等的训练FLOPs下，以更小的model size实现了更低的validation perplexity和更高的few-shot accuracy，同时显著提高了inference throughput。

关键词¶

Mixture-of-Recursions: 混合递归 (MoR) 是本文提出的一种统一框架，它结合了参数共享和自适应计算，旨在提高语言模型的效率。MoR 在一个单一的递归 Transformer 内部，通过共享一层层的参数（参数效率）和使用轻量级的路由器动态分配每个 token 的递归深度（自适应计算），从而在计算和内存使用上实现优化。它允许模型为每个 token 分配不同的“思考”深度，只将计算集中在需要更多处理的 token 上，同时通过选择性地缓存 KV 对来减少内存访问。
Recursive Transformer: 递归 Transformer 是一种模型架构，它通过多次重复应用同一组共享的层来构建 token 表示。与标准的 Transformer 使用 L 个独立层不同，递归 Transformer 将模型划分为 N_r 个递归块，每个块使用一组共享的参数（Φ’）。这种设计允许模型在不增加参数数量的情况下，通过增加有效网络深度来执行更多计算，从而实现参数效率。
Parameter Sharing: 参数共享是一种提高模型效率的技术，它通过在模型中多个层或模块之间重用相同的权重来减少模型中独特可训练参数的总数。在递归 Transformer 中，参数共享是通过将模型分解为递归块，并让每个块使用一组共享的参数来实现的，从而显著降低模型的内存占用和训练成本。
Adaptive Computation: 自适应计算是指在模型训练和推理过程中，根据输入的需要动态地分配计算资源，而不是对所有输入都采用固定的计算路径。在 MoR 中，这主要通过“Token-Level Computation”实现，即模型会根据每个 token 的复杂性，为其分配不同的递归深度，使得计算量能够根据 token 的难度进行调整，将更多的计算资源集中在更具挑战性的 token 上。
Token-Level Computation: Token-Level Computation（Token 级别的计算）是指 MoR 框架的核心思想之一，它允许模型为输入序列中的每一个 token 分配一个动态的、自适应的递归深度。这意味着模型不会对所有 token 应用相同的计算量，而是会根据每个 token 所需的“思考”或处理复杂度，决定它需要通过递归块多少次。这种细粒度的计算分配是实现模型效率的关键。
Recursion Depth: 递归深度是指一个 token 在模型中通过共享的递归块（Recursion Block）被重复应用（处理）的次数。在 MoR 中，递归深度是动态学习和分配的，由一个“Router”根据 token 的上下文和重要性来决定，而不是所有 token 都拥有相同的固定深度。不同的递归深度意味着 token 经历了不同程度的迭代计算和表征细化。
Router: 路由器是 MoR 框架中的一个关键组件，它负责为每个 token 动态地分配递归深度。路由器通过学习一个轻量级的模型，根据 token 的隐藏状态来计算一个分数，并基于该分数决定该 token 在接下来的递归步骤中是否继续计算，或者它应该被分配到哪个递归路径。MoR 探讨了两种主要的路由器策略：Expert-choice routing 和 Token-choice routing。
KV Caching: KV 缓存（Key-Value Caching）是一种用于加速 Transformer 模型自回归解码（生成文本）的技术。在生成每个新 token 时，模型会计算当前 token 的 Query、Key 和 Value。KV 缓存会存储之前 token 生成的 Key 和 Value 向量，这样在生成后续 token 时，就无需重新计算这些 Key 和 Value，从而避免了二次方的计算复杂度，显著提高了推理吞吐量。MoR 进一步提出了优化的 KV 缓存策略。
Expert-Choice Routing: Expert-choice routing（专家选择路由）是一种路由器策略，在 MoR 中，它指的是在每个递归步骤（或“专家”）中，路由器主动选择一部分 top-k 个 token 来继续通过该递归块进行计算。这种方法能够精确地控制每个递归步骤的计算量，并能实现精确的负载均衡，但可能存在信息泄露（违反因果关系）的问题，需要额外的机制来缓解。
Token-Choice Routing: Token-choice routing（Token 选择路由）是 MoR 中的另一种路由器策略。在这种策略下，每个 token 在模型的最开始，通过一个单一的路由决策，就被固定地分配到一个完整的递归计算路径（即确定其总的递归深度）。这种方法避免了因果关系被破坏的问题，但可能导致递归块之间的负载不均衡。
Recursion-Wise KV Caching: Recursion-wise KV caching（递归式 KV 缓存）是 MoR 提出的两种 KV 缓存策略之一。在这种策略下，只有那些被路由到当前递归步骤的 token，才会将它们的 Key 和 Value 缓存起来，并且后续的注意力计算也只会在这些局部缓存的 token 中进行。这种方法可以显著减少 KV 缓存的内存占用和 IO 需求，因为它只存储活跃 token 的 KV 对。
Recursive KV Sharing: Recursive KV sharing（递归 KV 共享）是 MoR 提出的另一种 KV 缓存策略。在这种策略下，模型会将所有 token 在第一个递归步骤中生成的 KV 对进行缓存，并将这些 KV 对在后续所有的递归步骤中共享使用。这种方法的优点是最大程度地节省了内存，并且在预填充（prefill）阶段可以进一步加速，因为它只需要计算一次 KV 对。然而，它可能在 token 路由非常精确时，由于 KV 对与后续递归步骤的匹配度降低而略微影响性能。
Continuous Depth-Wise Batching: Continuous depth-wise batching（连续深度批处理）是 MoR 用于提高推理吞吐量的一种技术。它借鉴了连续批处理的思想，并将其应用于递归模型。当 MoR 模型中的某些 token 由于动态深度分配而提前完成计算（“提前退出”）时，空出的计算槽位可以立即被队列中等待处理的新 token 填充，从而保持 GPU 的高利用率，消除计算“气泡”，实现比传统批处理更高效的推理。

摘要¶

Mixture-of-Recursions (MoR) 提出了一种统一的 Transformer 架构，旨在同时提高语言模型的参数效率、计算效率和内存效率。现有方法通常侧重于参数共享或自适应计算，MoR 则将两者整合到一个递归 Transformer 中。它通过在递归步骤中重用共享的层栈来实现参数效率，并通过轻量级路由器动态为单个 token 分配不同的递归深度，从而实现 token 级的自适应计算。这种机制使 MoR 能够仅在给定递归深度内活跃的 token 之间执行二次注意力计算，并通过选择性缓存这些 token 的 Key-Value (KV) 对来进一步提高内存访问效率。

核心方法 (Core Methodology)

预备知识 (Preliminary)
- 递归 Transformer (Recursive Transformers)：标准 Transformer 使用 $L$ 个独特的层，每个层包含自注意力模块和前馈网络。递归 Transformer 通过将模型划分为 $N_r$ 个递归块，并在这些块中重用共享参数 $\Phi'$ 来减少参数数量。这允许在不增加参数大小的情况下增加有效网络深度。
- 参数共享策略 (Parameter-sharing Strategies)：论文探讨了四种参数共享策略：
  - Cycle: 递归块循环重用同一组参数。例如，对于 $L=9$ 层和 $N_r=3$ 递归，层按 $[(0, 1, 2), (0, 1, 2), (0, 1, 2)]$ 的模式共享。
  - Sequence: 每个递归块按顺序重用同一层。例如，对于相同配置，层按 $[(0, 0, 0), (1, 1, 1), (2, 2, 2)]$ 的模式共享。
  - Middle-Cycle: 在 Cycle 策略的基础上，保留第一个和最后一个层的独特参数 ($\Phi^0$ 和 $\Phi^{L-1}$)，而中间层共享参数。
  - Middle-Sequence: 在 Sequence 策略的基础上，保留第一个和最后一个层的独特参数，而中间层共享参数。经验证，Middle-Cycle 策略在递归模型中表现最佳，因为它能更好地捕捉输入和输出的细微差别。
Mixture-of-Recursions (MoR)：MoR 的核心在于两个关键组件：路由机制和 KV 缓存策略。
- 路由策略 (Routing Strategies)：
  - Expert-choice Routing (专家选择路由)：在每个递归步骤 $r$，路由器使用隐藏状态 $\mathcal{H}^r_t$ 和路由参数 $\theta^r$ 计算 token $t$ 的得分 $g^r_t = \mathcal{G}(\theta^{T_r} \mathcal{H}^r_t)$。然后，选择得分最高的 top-$k$ 个 token 继续通过递归块。 $\mathcal{H}^{r+1}_t = \begin{cases} g^r_t f(\mathcal{H}^r_t, \Phi') + \mathcal{H}^r_t, & \text{if } g^r_t > P_\beta(G^r) \\ \mathcal{H}^r_t, & \text{otherwise} \end{cases}$ 其中 $P_\beta(G^r)$ 是递归步骤 $r$ 中所有得分的 $\beta$ 百分位数阈值。优点：保证完美的负载均衡。缺点：存在因果性违背 (causality violation) 或信息泄露问题，即 top-$k$ 选择依赖于序列中未来的 token 信息。缓解方法：
    - Auxiliary Router (辅助路由器)：一个独立的轻量级网络，在训练时与主模型一同训练，但在推理时用于预测 token 是否会被选中。
    - Auxiliary Loss (辅助损失)：直接对主路由器应用二元交叉熵损失，促使其在训练时将选中的 token 得分推向 1，未选中的推向 0。实验发现辅助损失更有效。
  - Token-choice Routing (Token 选择路由)：每个 token 在开始时通过一次路由决策被分配到固定的递归深度 $i$。路由器计算非线性函数（如 softmax）的得分 $g_t = \mathcal{G}(\theta^{T_r} \mathcal{H}^1_t)$，并选择得分最高的专家 $i = \text{arg max}_j g^j_t$。token 随后顺序应用递归 $i$ 次。 $\mathcal{H}^{r+1}_t = \begin{cases} g^r_t f(\mathcal{H}^r_t, \Phi') + \mathcal{H}^1_t, & \text{if } r = i \\ g^r_t f(\mathcal{H}^r_t, \Phi'), & \text{otherwise} \end{cases}$ 优点：无因果性违背。缺点：存在负载不均衡问题。缓解方法：
    - Balancing Loss (均衡损失)：对专家间的 token 分布施加正则化。 $\mathcal{L}_{\text{Balance}} = \alpha \sum_{i=1}^{N_r} f_i P_i$ $f_i = \frac{1}{T} \sum_{t=1}^T \mathbf{I}(\text{Token } t \text{ selects Expert } i)$ $P_i = \frac{1}{T} \sum_{t=1}^T g^i_t$ 其中 $N_r$ 是递归次数，$g^i_t$ 是专家 $i$ 对 token $t$ 的路由得分，$f_i$ 是路由到专家 $i$ 的 token 比例，$P_i$ 是专家 $i$ 的平均路由得分，$\alpha$ 是超参数。
    - Loss-free (无损失)：通过路由器偏差项 (router bias) 调整，而不使用显式正则化损失。实验发现均衡损失能带来更稳定的路由。
- KV 缓存策略 (KV Caching Strategies)：
  - Recursion-wise KV Caching (递归级 KV 缓存)：只缓存路由到当前递归步骤的 token 的 KV 对。每个递归深度都有独立的 KV cache，注意力计算也仅限于这些本地缓存的 token。这显著减少了 KV 内存和 I/O，并降低了注意力 FLOPs。 KV Memory: $(N_r + 1)/(2N_r)$ KV Cache IO: $(N_r + 1)/(2N_r)$ Attn FLOPs: $k^2/N^2_{\text{ctx}}$
  - Recursive KV Sharing (递归 KV 共享)：仅在第一个递归步骤缓存 KV 对，并在所有后续递归步骤中重用。这最大限度地节省了内存，并消除了预填充 (prefill) 阶段对更深层递归的计算需求。 KV Memory: $1/N_r$ KV Cache IO: $1$ Attn FLOPs: $k/N_{\text{ctx}}$ 尽管递归 KV 共享在内存效率上有优势，但可能因上下文不匹配导致性能略有下降。

实验验证 (Experimental Validation)

主要结果 (Main Results)：在相同训练 FLOPs 预算下，MoR 模型（使用 expert-choice 路由器和两个递归）相比 Vanilla Transformer 和其他递归基线，在参数量减少近 50% 的情况下，实现了更低的验证损失和更高的 few-shot 准确率（43.1% vs. 42.3%）。在相同训练 token 数量下，MoR 在使用 25% 更少 FLOPs 的情况下，性能优于基线。
IsoFLOP 分析 (IsoFLOP Analysis)：MoR 在 135M 到 1.7B 参数规模的广泛模型尺寸和计算预算下，始终优于递归基线。对于大于 360M 参数的模型，MoR 的性能与 Vanilla Transformer 持平甚至超越，同时显著减少了参数数量（约三分之一）。这表明 MoR 是一种可扩展且参数高效的架构。
推理吞吐量 (Inference Throughput)：MoR 结合了连续深度批处理 (continuous depth-wise batching) 和早退出 (early-exiting) 机制，显著提高了推理吞吐量。与 Vanilla Transformer 相比，MoR 模型实现了高达 2.06 倍的速度提升。

消融研究 (Ablation Studies)

参数共享策略 (Parameter Sharing Strategies)：Middle-Cycle 策略在所有设置下都表现最佳，并在 NLL 和 few-shot 准确率方面持续优于其他策略，尤其是在模型规模和递归深度增加时。
路由策略 (Routing Strategies)：
- Expert-choice: 辅助损失结合 sigmoid 激活函数和简单线性路由器表现最佳。Z-loss 对准确率影响微乎其微，但能轻微减少“死亡 token”比例。
- Token-choice: 均衡损失能实现更稳定的路由，显著降低 MaxVio 指标。Softmax 激活函数与 MLP 路由器搭配表现最佳。
KV 缓存策略 (KV Caching Strategies)：递归 KV 共享在参数共享架构中表现稳健，但对于 expert-choice 路由的 MoR 模型，它可能导致性能下降，因为精确的 token 选择更受益于独立的递归级缓存。然而，对于 token-choice 路由，KV 共享可以提供额外的上下文信息，有助于弥补其路由决策的不足。

分析 (Analysis)

计算最优扩展分析 (Compute-optimal Scaling Analysis)：MoR 的计算最优扩展行为表明，它更受益于参数数量的增加（模型更大），而不是训练数据量的增加（训练步数更长）。这意味着 MoR 的最佳扩展策略是增加模型容量。
路由分析 (Routing Analysis)：路由器的决策反映了后续 token 的上下文可预测性。例如，单词的第二部分通常易于预测，因此需要的递归步骤较少。Expert-choice 路由器与辅助损失结合，能完美地将选中和未选中的 token 分离。
测试时间扩展分析 (Test-time Scaling Analysis)：MoR 允许在测试时通过更深的递归来提高模型性能。增加递归深度不仅提供额外的计算，还能让每个后续步骤在优化 token 表示或“思考过程”中进一步专业化，从而提高生成质量。

结论与未来工作 (Conclusion & Future Works)

MoR 通过结合参数共享、自适应递归深度和高效 KV 缓存，在不牺牲模型质量的情况下，显著降低了计算和内存开销。实验证明 MoR 在降低验证困惑度和提高 few-shot 准确率方面优于基线，同时具有更高的推理吞吐量。

局限性与未来工作：

推理 MoR 模型 (Reasoning MoR models)：探索路由器如何动态调整思维链 (Chain-of-Thought, CoT) 的递归深度，以增强推理能力。
进一步扩展模型系列 (Further scaling model family)：训练更大规模（超过 3B 参数）的 MoR 模型，并探索深度特定 LoRA 或 MoE，以及从现有预训练模型检查点继续预训练。
自适应容量控制 (Adaptive capacity control)：开发能更灵活地在训练和推理阶段调整 top-$k$ 容量的 MoR 模型设计。
与稀疏算法的兼容性 (Compatibility with sparse algorithms)：将 MoR 与结构化稀疏性技术（如剪枝、量化）结合，以进一步优化计算。
扩展到多模态和非文本领域 (Expansion to multimodal and non-text domains)：由于 MoR 递归块与模态无关，其自适应深度机制可扩展到视觉、语音和统一多模态 Transformer 架构。

Abstract¶

Abstract：
扩展语言模型带来了令人印象深刻的能力，但其伴随的计算与内存需求使得训练和部署成本高昂。现有的效率研究通常只关注参数共享或自适应计算中的某一方面，而如何同时实现这两点仍然是一个开放问题。

本文提出Mixture-of-Recursions（MoR），这是一个统一的框架，将参数共享和自适应计算两个效率方向整合到一个Recursive Transformer中。MoR 通过在递归步骤中复用共享的层栈，实现了参数效率；同时通过轻量级的路由器，为每个 token 动态分配不同的递归深度，实现了 token 级别的自适应处理。

这一机制使得 MoR 能够在特定递归深度上仅对“活跃”的 token 进行二次计算，从而进一步提升内存访问效率。此外，MoR 还引入了一种KV共享变体，它复用第一次递归中的 key-value 对，专门用于降低预填充（prefill）的延迟和内存占用。

从 1.35 亿到 17 亿参数规模的模型实验表明，MoR 形成了一条新的帕累托前沿：在保持相同训练 FLOPs 和更小模型规模的前提下，它显著降低了验证困惑度，提升了少样本准确率，并且吞吐量更高，优于传统的递归模型和普通模型。

重点内容总结：

MoR 统一了参数共享和自适应计算，解决了现有方法的局限性。
通过共享层和动态路由机制，实现高效计算和内存使用。
引入 KV 共享机制，进一步优化 prefill 性能。
实验结果显示 MoR 在多个指标上优于现有方法，提供了“大模型效果、低成本”的解决方案。

1 Introduction¶

本节主要介绍 Mixture-of-Recursions（MoR）的研究背景、动机以及其核心思想。

背景与动机¶

近年来，研究人员通过将 Transformer 扩展到数百亿参数，显著提高了模型的泛化能力和推理能力。然而，这种扩展也带来了巨大的内存和计算开销，导致训练和部署变得困难。为了解决这一问题，研究者们提出了多种“高效”设计，主要包括两个方面：

参数效率（Parameter Efficiency）：通过参数共享、层绑定（layer tying）等方式减少模型参数数量。
自适应计算（Adaptive Computation）：根据任务复杂度动态分配计算资源，如早期退出（early-exit）方法。

尽管在参数效率和自适应计算方面分别取得了一些进展，但目前还没有一种架构能够同时实现这两种效率。

递归 Transformer 与挑战¶

递归 Transformer 是一种将共享的参数块多次应用于输入的方法，天然具备参数效率。然而，以往的动态递归方法通常面临以下挑战：

需要额外的训练过程；
部署效率低；
往往退化为固定深度的递归，无法真正实现每个 token 的自适应计算。

MoR：统一框架¶

为克服上述问题，作者提出了 Mixture-of-Recursions (MoR)。其核心思想是：

通过 端到端训练的轻量路由模块（router），为每个 token 动态分配递归深度；
利用共享的参数块对每个 token 进行不同次数的计算；
通过 递归级别的 KV 缓存（recursion-wise KV caching），减少内存访问，提高吞吐量。

MoR 的优势在于：

参数共享：减少整体参数数量；
自适应计算：每个 token 的计算量根据其复杂程度动态变化；
高效缓存：仅缓存与当前 token 递归深度相关的 KV 对，减少内存开销。

概念与意义¶

MoR 可以看作是一种在推理过程中 在隐空间中进行非语言“思考” 的预训练框架。它在生成每个 token 时就进行递归思考，而不是依赖生成前的连续提示（如一些 prompt-augmentation 方法）。

与以往方法相比，MoR 实现了：

在模型的 深度方向（vertical axis） 上进行可变长度的推理；
摆脱了传统方法中固定深度的限制；
统一了参数效率和自适应计算的优势。

贡献总结（Contributions）¶

统一的高效语言模型框架：首创新型架构 MoR，同时融合参数共享、自适应 token-level 思考深度和高效的 KV 缓存策略。
动态递归路由：从头训练路由模块，为每个 token 分配递归深度，且无需额外的训练阶段。
实验验证充分：在相同计算预算下，MoR 相比传统和递归基线模型，在验证损失和 few-shot 准确率上均表现更优。
高效架构：通过选择性计算和 KV 缓存优化，大幅减少训练 FLOPs，并提升推理吞吐量。

总结¶

本节介绍了 MoR 的提出背景、动机、核心思想与贡献。MoR 是一种 结合了参数效率与自适应计算的统一框架，目标在于提升模型的推理能力，同时降低训练和部署成本。

2 Method¶

2.1 Preliminary¶

Recursive Transformers¶

标准Transformer 通过堆叠 L 个独立层（每个层包含自注意力和前馈网络）来构建token表示。每一步时间步 t 的隐藏状态通过如下公式演化：
$$ \mathbf{h}_{t}^{\ell+1} = f(\mathbf{h}_{t}^{\ell}; \Phi_{\ell}) $$ 其中，$\ell$ 表示层数，从 0 到 $L-1$，$\Phi_{\ell}$ 是第 $\ell$ 层的参数。

递归Transformer（Recursive Transformers）通过重复使用参数（参数共享）来减少模型参数数量。模型被划分为 $N_r$ 个递归块，每个块使用共享权重 $\Phi'$。这种设计能够在不增加参数量的情况下增加模型的有效深度，从而提高计算效率。

Table 1: 参数共享策略¶

递归Transformer中有两种主要的共享策略：

Cycle Strategy：以循环方式复用块参数。如 $L=9$，$N_r=3$，参数按 $(0,1,2), (0,1,2), (0,1,2)$ 分配。
Middle-Cycle Strategy：在中间层共享参数，但保留首层和末层为独立参数，避免信息丢失。
Sequence Strategy：每个递归块重复使用同一层的参数，如 $L=9$，$N_r=3$ 时，参数分配为 $(0,0,0), (1,1,1), (2,2,2)$。
Middle-Sequence Strategy：与Middle-Cycle类似，保留首末层，中间层共享。

Enhanced training and inference efficiency in recursive models¶

递归模型通过参数共享可大幅减少可训练参数数量，从而节省内存，并提高分布式训练效率，特别是在使用 FSDP（Fully Sharded Data Parallel）技术时。

在推理阶段，递归模型支持 连续深度批处理（continuous depth-wise batching），可以将处于不同递归深度的token组合成一个batch，从而减少等待时间（bubble time），提高吞吐量。

Limitations in prior works¶

尽管递归模型在参数共享上有优势，但其KV缓存机制未优化，导致缓存大小未减，仍存在高检索延迟的问题。此外，大多数递归模型采用固定递归深度，无法自适应不同token的复杂性。虽然有后处理方法（如early-exit）尝试引入自适应，但通常需要额外训练阶段，影响性能。理想情况下，递归深度应在预训练中动态学习，同时需要并行解码机制（parallel decoding）来处理缺失的KV对，但这类机制复杂度高，工程难度大。

2.2 Mixture-of-Recursions (MoR)¶

本文提出 MoR（Mixture of Recursions） 框架，旨在动态调整每个token的递归深度。其核心包括两个部分：

路由机制（Routing）：根据token的复杂度分配其递归深度，动态决定哪些token需要更多计算。
KV缓存策略（KV Caching）：优化KV缓存的存储与使用，以适应动态递归结构，提升推理效率。

2.2.1 Routing Strategies: Expert-choice vs. Token-choice¶

Expert-choice Routing¶

每个递归深度视为一个专家，在每个递归步骤中选择top-k个token继续计算。
通过路由参数 $\theta_r$ 和隐藏状态 $\mathcal{H}_t^r$ 计算得分，选择高于阈值的token。
特点：负载均衡好，但存在信息泄露（causality violation） 问题，需引入辅助路由器或正则化损失来缓解。

Token-choice Routing¶

每个token在开始时就被分配一个固定的递归深度，从头到尾固定路径。
通过softmax或sigmoid函数分配token的递归深度。
特点：无信息泄露，但存在负载不平衡问题，通常需要平衡损失（balancing loss）或无损失算法（loss-free algorithms）来优化。

对比与总结（Table 2 - Left）¶

特性	Expert-choice	Token-choice
优点	固定计算预算，负载均衡好	无信息泄露
缺点	破坏因果性，需要辅助模块	负载不平衡，需额外损失函数
解决方案	辅助路由 / 正则化损失	平衡损失 / 无损失算法

2.2.2 KV Caching Strategies: Recursion-wise Caching vs. Recursive Sharing¶

Recursion-wise KV Caching¶

按递归层级缓存KV对，仅缓存被选择的token，保持缓存大小与路由容量一致。
优点：内存和IO效率高，适合动态递归深度。
缺点：每个递归层级需要独立缓存，不便于全局信息共享。

对比与总结（Table 2 - Right）¶

特性	Recursion-wise Caching	Recursive Sharing
KV内存	$(N_r+1)/2 \cdot N_r$	$1/N_r$
KV Cache IO	$(N_r+1)/2 \cdot N_r$	1
注意力FLOPs	$(k/N_{ctx})^2$	$k/N_{ctx}$

总结¶

递归Transformer 通过参数共享减少模型大小，但存在KV缓存问题和固定递归深度限制。
MoR框架 提出动态路由机制（Expert-choice / Token-choice）与KV缓存策略（Recursion-wise / Recursive Sharing）来提升模型效率和适应性。
路由机制 决定每个token的递归深度，平衡计算资源分配。
KV缓存策略 优化存储与访问，提升推理效率。

通过 MoR，模型可在保持参数效率的同时，实现更灵活、高效、适应性强的token级计算，为Transformer架构的优化提供了新方向。

3 Experiments¶

研究者基于 Llama 的 Transformer 架构从头开始预训练模型，并参考了 SmolLM 的配置。训练数据来自 FineWeb-Edu 数据集的一个去重子集（SmolLM-Corpus）。评估时使用了 FineWeb-Edu 的验证集和六个 few-shot 任务（如 Gao 等人提出的基准）。详细的训练、评估以及吞吐量测量方法在附录 B 中有详细描述。

表格总结（Table 3）¶

表格比较了 MoR（Mixture-of-Recursions）、Recursive 和 Vanilla Transformer 在固定 FLOPs 和固定 token 数量两种设定下的性能。评估指标包括验证集的负对数似然（NLL）和 few-shot 准确率。表格还展示了不同 KV 缓存策略（Cache 和 Share）以及递归深度（Nr）对模型性能的影响。

所有模型在相同 FLOPs 设置下，MoR 在参数更少的情况下仍能获得更好的性能。
递归深度从 2 增加到 4，MoR 的 NLL 有所上升，但仍然优于 Vanilla 模型。
使用 expert-choice 路由器的 MoR 性能优于 token-choice 路由器。
KV 缓存共享略降性能，但提升内存效率，适合部署。

3.1 主要结果¶

MoR 在相同训练计算资源下表现更优¶

在 16.5e18 FLOPs 的训练预算下，MoR 使用 expert-choice 路由器和 2 层递归，参数比 Vanilla 模型减少近一半，但验证损失更低，平均 few-shot 准确率更高（43.1% vs 42.3%）。其高效计算能力使其在相同计算预算下能处理更多训练 token。

MoR 在相同数据量下计算效率更高¶

在固定 token 数量（20B）下，MoR 使用更少的 FLOPs（比 Vanilla 少 25%），仍能获得更低的验证损失和更高准确率。这表明 MoR 在计算-准确率之间实现了更优的平衡，显著减少了训练时间和内存峰值使用。

MoR 的性能与路由和缓存策略有关¶

在 Nr=3 的轻量配置下，MoR 使用 token-choice 路由器的表现（40.0%）不如 expert-choice 路由器（42.6%），说明路由粒度对性能有重要影响。KV 缓存共享略有性能下降，但提升了内存效率，适合资源受限的场景。

3.2 IsoFLOP 分析¶

实验设置¶

研究者在四个模型规模（135M、360M、730M、1.7B 参数）下，使用固定递归深度（3）进行了实验，并在三种计算预算（2e18、5e18、16.5e18 FLOPs）下进行比较。

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

在所有模型规模和计算预算下，MoR 均优于递归基线模型。在 360M 及以上规模中，MoR 不仅与 Vanilla 模型相当，甚至在低预算下表现更好。虽然在最小模型（135M）上略有差距，但随着规模扩大，这种差距迅速缩小。这表明 MoR 是一种高效且可扩展的替代架构。

3.3 推理吞吐量评估¶

实验设置¶

研究者评估了 360M 模型规模下 MoR 的推理吞吐量，递归深度为 2、3、4，训练预算为 16.5e18 FLOPs。通过生成 token 的时间测量吞吐量，并采用不同批次大小进行测试。

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

MoR 利用连续深度批处理（depth-wise batching）提升吞吐量，通过将已完成序列替换为新 token，保持高 GPU 利用率。递归深度越高，越早退出的 token 越多，KV 缓存使用越少，吞吐量越高（如 MoR-4 在 Max 批次下达到 2.06 倍速）。虽然性能略有下降，但吞吐量的大幅提升使其成为部署的优选方案。

总结¶

该节实验系统评估了 Mixture-of-Recursions（MoR）模型在不同训练和推理条件下的表现。结果表明：

MoR 在相同计算条件下比 Vanilla 和 Recursive 模型更高效，参数更少但性能更好。
MoR 的性能受路由策略和缓存策略影响显著，expert-choice 路由和 recursion-wise 缓存表现更佳。
MoR 具有良好的可扩展性，随着模型规模增大，其优势愈加明显。
MoR 在推理阶段利用连续深度批处理显著提升吞吐量，适合实际部署。

4 Ablation Studies¶

本节通过消融实验分析了 MoR（Mixture-of-Recursions）模型 在不同设计选择下的性能表现，重点包括：

参数共享策略（Parameter Sharing Strategies）
路由策略（Routing Strategies）
KV缓存策略（KV Caching Strategies）

并提供了多个图表（图4a、4b、4c）用于辅助分析。

4.2 Routing Strategies¶

本节评估了 expert-choice 和 token-choice 两种路由策略下的不同设计选择。

In expert-choice routing, auxiliary loss and linear router yield the best performance.¶

expert-choice 路由（表4左侧）：
- 比较了多个设计因素：
  - 解决因果性问题的方法（辅助路由器 vs. 辅助损失）
  - 归一化函数（sigmoid vs. tanh）
  - 路由器结构（MLP、Linear、Wide-MLP）
  - 辅助 z-loss 的影响
- 关键发现：
  1. 使用 辅助损失（auxiliary loss） 比单独训练辅助路由器更有效。
  2. sigmoid 函数 和 线性路由器结构（Linear） 表现最佳。
  3. 辅助 z-loss 对准确率影响不大，但能略微减少“dead tokens”比例（未被最后递归选中的token）。

In token-choice routing, balancing loss yields stable and accurate routing.¶

token-choice 路由（表4右侧）：
- 默认启用 z-loss，比较了两种平衡策略：
  - 显式平衡损失（balancing loss）
  - 无损失训练 + 路由器偏置（loss-free with router bias）
- 关键发现：
  1. 显式平衡损失（balancing loss） 显著降低了 MaxVio（衡量负载不平衡的指标），因此更推荐使用。
  2. 即使使用了平衡损失，模型仍然在多数训练步骤中难以平衡各专家之间的负载。
  3. Softmax 激活函数 + MLP 路由器 表现最好。
  4. 移除 z-loss（虽然最终设计中重新加入了一个很小的系数）提升了性能和路由稳定性。

表4总结：¶

路由策略	采样方式	归一化函数	路由器结构	z-loss	性能指标
Expert-choice	Aux Loss	σ	Linear	✗	NLL↓, 几乎无 dead token
Token-choice	Loss-free	soft	Linear	✗	MaxVio↓, 几乎无性能损失

4.3 KV Caching Strategies¶

总结¶

本节通过系统性的消融实验，确定了在 MoR 模型中：

Middle-Cycle 参数共享 是最优选择；
Expert-choice 路由 中，使用 辅助损失 + Linear 路由器 表现最佳；
Token-choice 路由 中，显式平衡损失 + Softmax + MLP 路由器 最具稳定性；
KV共享 在参数共享架构中表现良好，对 token-choice 路由 有益，但对 expert-choice 路由 有害。

这些结果为 MoR 的设计提供了坚实的实证基础，并为后续模型优化提供了方向。

5 Analysis¶

本章主要对 Mixture-of-Recursions (MoR) 模型在训练与测试过程中表现出的性能和行为进行分析，重点包括计算最优缩放、路由机制分析以及测试时缩放策略的效果。

5.1 Compute-optimal Scaling Analysis¶

MoR scaling favors model size over training length under isoFLOPs.¶

在 图5(a) 中，作者展示了在 isoFLOPs（固定计算量）条件下的计算最优缩放分析。与基线模型相比，MoR 模型在计算资源固定时，更倾向于通过增加模型大小（参数数量）而不是训练步数来提升性能。这是因为 MoR 的共享参数块性能对模型表现具有关键作用，因此增大模型规模比增加训练数据更有效。这种缩放策略表明，MoR 在资源有限的情况下，优先选择更大容量但训练更少的模型，具有更“以参数换取性能”的特点。

重点：MoR 在 isoFLOPs 下更关注模型大小，而非训练长度。

5.2 Routing Analysis¶

The allocation of recursion depth reflects token semantic importance.¶

在 图1 右侧 展示了不同 token 在 MoR 模型中被分配的递归深度。语义丰富或重要的 token（如 “People”、“Drugs” 等）通常会经历更多的递归步骤（例如 3 次），而功能词（如 “and”、“—”）则只进行较少的递归。这表明递归深度的分配与 token 的语义重要性密切相关。

重点：递归深度反映了 token 的语义重要性，重要 token 被分配更深的递归。

Expert-choice router with auxiliary loss perfectly separates selected from unselected tokens.¶

图5(b) 展示了 MoR 模型中专家选择路由器（expert-choice router）在不同递归步骤中的输出分布。路由器为每个 token 分配一个得分，得分接近 1 的 token 被选中，得分接近 0 的则未被选中。在辅助损失（auxiliary loss）的作用下，选中和未选中 token 的得分分布被清晰区分开，说明路由器在路由决策上具有很高的准确性。

重点：路由器结合辅助损失，实现了选中与未选中 token 的完美分离。

5.3 Test-time Scaling Analysis¶

MoR enables test-time scaling via deeper recursion.¶

图5(c) 展示了在测试阶段，MoR 模型中递归深度增加对性能（log-likelihood）的影响。随着递归深度 Nr 从 2 增加到 4，模型的 log-likelihood 也逐渐提升。这表明，在生成时增加递归深度可以提升模型的推理能力和输出质量。每一个递归步骤可以被视为对 token 表示的进一步精炼，从而在更深层次上优化生成过程。

重点：MoR 支持测试时缩放（test-time scaling），即在推理阶段通过增加递归深度来提升性能。

总结¶

MoR 在 isoFLOPs 条件下优先通过增加模型大小而非训练步数来获得性能提升。
递归深度的分配与 token 的语义重要性紧密相关，路由器能有效区分选中和未选中 token。
测试时通过增加递归深度，MoR 可以进一步提升生成质量，支持动态的推理资源分配。

这些分析结果表明，MoR 在模型效率、计算资源利用和生成质量之间取得了良好的平衡，具有较强的实用性和扩展性。

6 Related Work¶

Recursive Transformers（递归Transformer）¶

重点内容：
递归Transformer通过参数共享（parameter sharing）实现模型的高效性，其中最早的研究是Universal Transformer（Dehghani等，2018），它表明通过反复应用一个模块可以达到深度非共享结构的表示能力。
近年来，研究者提出了多种递归Transformer的变体，如Looped Transformers，这些模型不仅能作为“可编程计算机”运行，还能学习迭代算法、泛化到更长的输入，并通过模拟多步优化器来提升few-shot学习。
为了解决递归结构中权重共享（weight tying）带来的精度下降问题，Bae等（2024）在每个循环中引入了LoRA适配器（Low-Rank Adaptation），提出了“Relaxed Recursive Transformers”。
此外，Geiping等（2025）的研究表明递归结构在隐式推理（latent reasoning）中具有优势。
次要内容：
递归结构的另一个优势是其灵活性，通过借鉴“early-exiting”和“计算路由（compute routing）”思想，模型可以根据输入的难易程度动态调整递归次数，从而实现更高效的计算分配。

Adaptive Computation（自适应计算）¶

重点内容：
自适应计算的核心思想是动态分配计算资源，从而减少训练和推理的成本。早期研究主要集中在传统神经网络，但近年来也被广泛应用于大语言模型。
“Early exiting”是一种常见的策略，即在处理“简单”样本时提前退出，避免计算后续层。
最近的Mixture-of-Depths（MoD）方法将自适应深度建模为一个路由问题，通过在每一层选择部分token进行完整计算，其余则跳过，从而实现更细粒度的计算分配。
次要内容：
MoR（Mixture-of-Recursions）将这一思想应用于递归Transformer，即通过动态路由将token送入递归模块的多次调用中，而不是通过不同的层，从而在不增加参数量的前提下实现任意深度的计算分配。

Routing Mechanism（路由机制）¶

重点内容：
路由机制广泛应用于大语言模型中，特别是在Mixture-of-Experts（MoE）结构中。该结构通过一个路由网络为每个token分配一部分专家网络进行处理，从而提升模型容量而不增加计算开销。
早期的token-choice路由策略容易造成专家负载不平衡，因此研究者提出了expert-choice机制，即每个专家选择其要处理的token，从而实现负载均衡和更高的效率。
次要内容：
一些研究还尝试使用可训练的路由机制来决定跳过哪些层，这种方式区别于传统的early-exit方法，通过限制每层处理的token数量来实现静态的计算预算。

Key-value Caching（键值缓存）¶

重点内容：
键值缓存（KV Caching）是用于提高自回归生成效率的重要方法，通过缓存每一层的键和值张量，避免重复计算，从而提升吞吐量。
然而，KV缓存会迅速占用大量GPU内存，尤其是在长上下文和大批量的情况下。
为解决这个问题，研究者采用了量化、修剪不重要的键值、以及共享注意力头的键值等技术。
次要内容：
MoR提出了一种新的思路：可以在递归的不同阶段之间复用键值缓存，从而进一步降低内存消耗。此外，MoR中可以优化两种不同的缓存策略以适应不同的部署需求。

Latent Reasoning（隐式推理）¶

重点内容：
隐式推理是指模型在生成答案之前，在隐藏状态中进行推理。早期方法通常采用固定深度的推理策略，例如通过特殊标记或结构化提示引导模型进行预设次数的推理。
另一些方法则通过将最终隐藏状态循环输入模型，模拟链式推理（chain-of-thought）过程。
次要内容：
现有方法缺乏灵活性，容易在简单任务中浪费计算资源，在复杂任务中又可能推理不足。
MoR通过引入自适应的递归次数，结合隐式推理，提供了将动态计算与推理深度结合的框架基础，有望实现更高效的模型推理。

总结：
本章综述了与MoR密切相关的五个研究方向：递归Transformer、自适应计算、路由机制、键值缓存和隐式推理。其中，递归结构和自适应计算是MoR的核心理论基础，路由机制和缓存优化是实现手段，而隐式推理则是MoR的潜在应用方向。这些研究共同为MoR提供了一个高效、灵活、适应性强的模型框架。

7 Conclusion¶

Mixture-of-Recursions (MoR) 提出了一种统一的Transformer架构，能够同时利用参数共享、自适应递归深度以及高效的KV缓存机制，且不牺牲模型质量。通过为每个token动态分配递归深度，并通过轻量级路由器选择性地缓存键值状态，MoR显著减少了二次计算开销和冗余内存访问成本。广泛的实验证明，MoR在验证集上的困惑度（perplexity）和平均少样本准确率方面优于标准Transformer和早期递归模型，同时具有更高的推理吞吐量。这些结果表明，MoR为以显著减少的计算和内存开销实现大模型能力提供了一条有效路径。

7.1 局限性与未来工作¶

推理MoR模型¶

当前研究指出，推理链中存在冗余，可通过token级别的自适应计算机制（如早退出机制）来解决。MoR框架本身通过为每个token自适应决定递归深度，从而实现潜在的推理能力。因此，未来的关键工作在于探索如何在实际推理数据集的后训练中，让路由器动态学习调整链式思维（Chain-of-Thought, CoT）链的需求。开发更高级的路由策略，使递归深度与推理复杂度直接对齐，可能提升推理准确性、计算效率，甚至可解释性。

进一步扩展模型规模¶

由于计算资源限制，当前实验仅限于17亿参数以内的模型。下一步是训练超过30亿参数的MoR模型，并使用更大规模的数据集进行预训练。为了降低预训练成本，可以考虑从现有标准LLM的预训练模型继续训练（uptraining）。未来还计划研究不同初始化策略对递归模型性能的影响，并确保在后训练阶段公平比较递归Transformer的扩展能力，同时考虑推理吞吐量的限制。

自适应容量控制¶

专家选择路由（Expert-Choice Routing）的一大优势是通过预设容量因子实现完美的负载均衡。然而，当前在推理阶段分配不同容量时存在限制。在MoR模型中，当使用辅助损失时，被选中和未被选中的token路由输出几乎完全分离，使得训练后调整top-k值变得困难。因此，未来需要设计更灵活的模型架构，支持在训练和推理阶段使用不同的容量配置。

与稀疏算法的兼容性¶

由于MoR具备token级别的自适应递归能力，可进一步通过结构化稀疏性进行优化。这包括选择性激活子网络或参数，从而动态剪枝token和层级别的冗余计算。该方向的研究有望带来显著的效率提升。我们认为许多基于稀疏性的技术，如剪枝或量化，与MoR高度互补。这将为递归模型中的稀疏架构设计提供更深入的见解，并为未来研究开辟新方向。

扩展到多模态及非文本领域¶

MoR的递归块本质上是模态无关的，其自适应深度机制可以扩展到文本以外的领域。这一特性使MoR能够轻松集成到视觉、语音及多模态Transformer架构中。将token自适应递归应用到长上下文的视频或音频流中，有望带来更高的内存效率和吞吐量提升，这对实际应用至关重要。通过动态调整每个token或片段的处理深度，MoR有可能释放这些潜在优势。

7.2 致谢¶

我们感谢Jacob Eisenstein对论文早期版本的宝贵反馈，感谢Seungyeon Kim、Mostafa Elhoushi和Sangdoo Yun的有益讨论。最后，感谢Google Cloud Platform为本项目提供的Google Cloud信用支持。

Appendix A Details of Design Choices for Mixture-of-Recursions¶

本节详细描述了 Mixture-of-Recursions（MoR）模型中采用的设计选择，补充了主文中的简要说明。

A.2 路由策略（Routing Strategy）¶

本节讨论了 MoR 中使用的两种路由策略：Expert-choice（专家选择） 和 Token-choice（令牌选择）。

Expert-choice 路由¶

优点：计算预算固定，便于资源管理。
缺点：使用 top-k 选择机制时，依赖后续 token 的信息，违反了自回归模型的因果性，可能导致信息泄露，影响模型推理可靠性。

缓解方法：

辅助路由（Auxiliary Router）：训练一个轻量级网络，在推理时预测 tokens 是否会被选中为 top-k。不参与主任务训练，仅用于推理。
辅助损失（Auxiliary Loss）：为主路由器添加二分类交叉熵损失，引导其在训练中将 top-k tokens 推向 1，其他 tokens 推向 0。

Token-choice 路由¶

优点：为每个 token 分配递归深度，不固定计算预算，保持自回归属性，避免信息泄露。
缺点：导致 experts（递归层）之间负载不平衡，影响训练效率和计算资源分配。

缓解方法：

平衡损失（Balancing Loss）：通过损失函数促使 tokens 在不同 experts 之间均匀分布。损失函数为： $$ ℒ_{\text{Balance}} = \alpha \sum_{i=1}^{N_r} f_i P_i $$ 其中 $ f_i $ 是分配给 expert $ i $ 的 token 比例，$ P_i $ 是 expert $ i $ 的平均路由得分。
无损失方法（Loss-free）：通过偏置调整（biasing）来平衡 token 分配。根据负载误差更新每个 expert 的偏置项，从而引导路由策略更均匀地分配 token。

📌 重点总结：

Expert-choice 更适合预算可控的场景，但需处理因果性问题。
Token-choice 更符合自回归特性，但需处理负载不均衡问题。
两种方案各有优劣，可根据应用场景选择或结合使用。

A.3 KV 缓存策略（KV Caching Strategy）¶

本节探讨了两种 key-value（KV）缓存策略，用于优化递归 Transformer 中的内存使用。

Recursion-wise Caching（递归级缓存）¶

每个递归步骤维护独立的 KV 缓存。
优点：确保 tokens 仅关注当前递归层的 KV 对，避免不同层级之间的分布错配，有助于保持模型精度。
适用场景：当使用精确的 token 路由策略时，推荐使用此策略以避免性能下降。

Recursive KV Sharing（递归共享缓存）¶

所有递归步骤共享首次递归中生成的 KV 对。
优点：显著降低内存使用，无需在预填充阶段进行深度递归计算。
缺点：后续递归层可能基于早期的、不合适的 KV 表示，影响模型性能。
适用场景：当内存效率为首要目标，且预填充时间是主要瓶颈时。

📌 重点总结：

Recursion-wise Caching 适用于 token 路由精确、模型性能优先的场景。
Recursive KV Sharing 适用于对内存敏感、且可接受一定性能下降的场景。

总结¶

本附录详细阐述了 MoR 模型在参数共享、路由机制和 KV 缓存策略方面的设计选择，每种策略都有其适用场景和优缺点。建议根据实际应用需求（如计算资源、推理准确性、内存限制等）选择合适的策略组合。其中，Middle 参数共享策略、Token-choice 路由与Recursion-wise 缓存策略在多数情况下表现更佳。

Appendix B Experimental Setup¶

训练设置¶

本研究采用了基于Llama的Transformer架构（Llama Team, 2024），其配置参考了开源SmolLM模型（Allal et al., 2024）。所有模型的预训练均基于SmolLM-Corpus中去重后的子集（FineWeb-Edu数据集，Penedo et al., 2024），该子集包含2200亿个教育材料来源的token。训练在四块H100或A100 GPU上进行。

在主要实验和isoFLOPs分析中，使用了Trapezoid学习率调度器，该调度器包含预热期（约5%）、稳定期和冷却期（20%）。这种方法允许从中间检查点继续训练，从而有效地进行缩放定律的探索，无需独立地训练所有模型。而在其他实验中，采用了简单的余弦退火调度器。

评估设置¶

为了评估模型性能，我们在六个基准上进行了少样本准确率测试，使用了语言模型评估工具（Language Model Evaluation Harness）。这些基准包括：LAMBADA（LD）、HellaSwag（HS）、PIQA（PQ）、WinoGrande（WG）、ARC（Easy和Challenge）、MMLU。

对于所有少样本数据集（LAMBADA、WinoGrande和MMLU除外），我们按目标字符串的字节长度对准确率进行了归一化处理。我们遵循了每个数据集的标准样本数，并在MMLU中使用了续写任务以简化处理。

所有评估性能的测量均在单块H100或A100 GPU上进行。

模型架构细节¶

表6总结了本研究所用的四个Vanilla Transformer模型的架构参数。这些模型作为递归模型的基础，其规模从135M到1.7B个参数不等（包括嵌入层和非嵌入层参数）。为了保证一致性和可比性，所有模型均使用49K的词表大小和2000 token的最大输入序列长度。

表中的三个较小模型使用了Grouped-Query Attention（Ainslie et al., 2023a），减少了键值头的数量。模型的基本配置参考了开源SmolLM模型（Allal et al., 2024）。

表6：模型架构参数总结（重点）¶

模型	嵌入参数 (N-emb)	非嵌入参数 (Emb)	层数 (N_L)	模型维度 (d_model)	头数 (N_head)	KV头数 (N_KV)	头维度 (d_head)	中间维度 (d_inter)	词表大小	上下文长度 (L_ctx)
Vanilla 135M	106M	28M	30	576	9	3	64	1536	49K	2K
Vanilla 360M	315M	47M	32	960	15	5	64	2560	49K	2K
Vanilla 730M	654M	75M	26	1536	24	8	64	4096	49K	2K
Vanilla 1.7B	1.61B	101M	24	2048	32	32	64	8192	49K	2K

重点内容：

所有模型使用相同的词表大小（49K）和最大上下文长度（2K）。
头维度（d_head）和词表大小（Vocab）在所有模型中保持一致。
头数（N_head）和KV头数（N_KV）随模型规模变化。
中间维度（d_inter）随着模型规模增大而增加，说明模型的表达能力不断增强。

Appendix C Expanded Results of IsoFLOP Analysis¶

总体比较¶

在正文（§3.2）中，作者对Vanilla、Recursive和Mixture-of-Recursions（MoR）模型在匹配的训练计算量（training compute）下进行了比较。研究了四种基础模型容量：135M、360M、730M和1.7B参数。

对于Recursive和MoR模型，固定递归次数为Nr=3，因此其唯一参数数量约为Vanilla模型的三分之一。每种架构只在最大的计算预算（16.5EB）下训练一次，然后通过中间训练检查点复用得到5EB和2EB版本。

Transformer的FLOPs近似计算¶

FLOPs的计算方法参考了Kaplan等人（2020）的近似方法，仅考虑前向传播的FLOPs，因为反向传播的计算量通常只是前向传播的两倍。Transformer的大部分操作是线性投影，其前向传播FLOPs计算为参数数的两倍，不包括注意力机制。

在注意力机制中，作者具体计算了查询和键的点积、值的softmax缩放操作。不计算因因果掩码导致的上三角冗余计算，同时省略了FlashAttention、归一化和非线性操作的计算成本。

值得注意的是：

Vanilla和Recursive Transformer的FLOPs相同。
MoR的FLOPs会根据路由策略和KV缓存策略变化，尤其是基于各递归深度的序列长度进行计算。
对于token-choice路由，由于每个步骤的token分配不同，作者假设完全平衡来估算FLOPs。
为了使MoR模型的有效深度可被递归次数整除，部分模型添加了额外层，这会增加FLOPs，因此减少了训练步数以维持预定义的FLOP预算。

带检查点复用的梯形学习率调度¶

为了避免为每个FLOP预算重新训练模型，作者采用了梯形学习率调度（Xing et al., 2018）。

学习率调度公式如下：

\[\begin{split} \eta(t)= \begin{cases} \frac{t}{w}\,\eta_{\max}, & 0\leq t<w\quad\text{(warm-up)},\\ \eta_{\max}, & w\leq t<p\quad\text{(plateau)},\\ \eta_{\max}\!\Bigl{(}1-\tfrac{t-p}{d}\Bigr{)}, & p\leq t<p+d\quad\text{(cool-down)}. \end{cases} \end{split}\]

其中：

warm-up阶段：从0线性增加到最大学习率。
plateau阶段：保持最大学习率。
cool-down阶段：线性衰减回0。

作者将warm-up设置为最小预算（2EB）总训练步数的5%，cool-down设置为对应预算总步数的20%。这种方法可以通过中间检查点高效管理实验。

结果概览¶

表7展示了在三种计算预算（2EB、5EB、16.5EB）下，四个基础模型（135M、360M、730M、1.7B）在FineWeb-Edu验证集上的负对数似然（NLL）和六个下游任务上的少样本准确率。

主要发现：

更高的计算量带来更好的性能，表现为更低的NLL和更高的准确率。
Recursive模型的参数共享策略导致性能下降，相比Vanilla模型是一个明显的权衡。
MoR模型克服了这一问题，在360M及以上参数规模下，其性能追上甚至超越了Vanilla模型，同时仅使用其三分之一的参数。
这一性能优势在730M和1.7B参数规模下持续存在。

MoR使用了expert-choice路由和递归级缓存机制。表中每个设置中表现最好的模型被标记为灰色高亮。

Appendix D Details of Experimental Settings for Throughput Measurement¶

实验系统与评估方法¶

我们实现了一个连续深度批处理推理系统（参考 Bae et al., 2024; Hooper et al., 2023），用于评估 MoR 模型的解码吞吐量。该系统在解码过程中动态地将查询入队并调度。使用的是来自 FineWeb-Edu 验证集的 1000 个样本，以模拟真实场景下的请求流量。

特别地，对于 MoR 模型，当某些查询提前退出时，空出的批处理槽位会立即被排队中的新查询填补，从而始终保持批处理的高利用率。

模型吞吐量对比设置¶

我们对比了Vanilla 模型和 MoR 模型（参数规模为 360M）在生成特定长度 token（每个样本）时的吞吐量。生成长度从均值为 256 的正态分布中采样，并且初始没有输入前缀（即从零开始生成）。

吞吐量以 MoR 模型的速度与 Vanilla Transformer 模型的相对速度 来表示。为了公平比较，MoR-4 模型在递归前增加了两个额外的层（总计34层），以确保其有效深度可以被递归次数整除（用于 Middle-Cycle 策略）。因此，MoR-4 的速度对比是基于一个包含这两个额外层的 Vanilla 模型（共34层）。

批处理设置¶

我们使用两种批处理设置：

固定批大小（Fixed Batch Size）：大小为 32。
相对最大批大小（Relative Maximum Batch Size）：根据 Vanilla 与 MoR 模型最大批大小的比值，将 32 乘以该比值得出。

我们基于 H100 GPU 的 VRAM 容量，考虑模型参数和 KV 缓存占用的内存来计算最大批大小（忽略当前位置隐藏状态的内存占用）。在这些自适应条件下：

MoR-2 支持最大批大小 42，
MoR-3 支持 48，
MoR-4 支持 51。

通过递归级别的 KV 缓存机制，MoR 模型能够在减少参数和缓存内存占用的前提下，显著提升批处理大小。

实现细节¶

连续深度批处理通过使用一个队列实现。
使用 FlashAttention 2 支持变长的 KV 缓存，以适应不同长度的查询。
采用静态大小的缓存，每个位置随时间更新，以兼容 torch.compile（用于优化推理速度）。
为了模拟真实部署场景，我们将 Transformer 块的计算阶段与其他计算分离，通过预处理输入嵌入和第一个非共享层，再输入到 Transformer 块中。然后测量前向传播的实际时间。
为了获得稳定的时间结果，测量前进行 100 次预热迭代。
对于提前退出的 token，累积到最大批大小后再进行处理（包括最后的非共享层、分类器和嵌入层，MoR 模型还包括第一个非共享层）。
最后，按照 FIFO 策略将它们加入队列以进行序列化批处理。
缓存与 KV 对的更新时间被排除，因为这些部分可通过多种工程优化显著改进，作为未来工作留待研究。

Appendix E Expanded Results of Parameter Sharing Strategy¶

本节扩展了第4.1节中对参数共享策略的消融实验，提供了图4(b)的完整量化结果。我们重新评估了四种参数共享策略（Cycle、Sequence、Middle-Cycle 和 Middle-Sequence）在两种基础模型规模（135M 和 360M 非嵌入参数）和两种不同的递归深度（Nr=2 和 Nr=3）下的表现。所有模型均在相同优化超参数下从头训练了100亿个token。验证集上的负对数似然（NLL）和六个基准任务上的平均少样本准确率总结在表8中。

Middle-Cycle 是最稳定的选择¶

对于360M规模的模型，Middle-Cycle 在两种递归深度（Nr=2 和 Nr=3）下都取得了最低的NLL，并在平均准确率方面表现出最大的提升。对于135M规模的模型，虽然Cycle在递归深度为2时略胜一筹（3.0071 vs. 3.0330），但当递归深度提升到3时，Middle-Cycle 反超（3.1048 vs. 3.1154），并且其准确率表现更为稳定。相比之下，Sequence 分享策略在所有四种设置中均记录了最差的NLL，且其准确率差距随递归深度增加而扩大。Middle策略对Sequence性能有轻微提升，但整体表现仍弱于基于Cycle的方法。

我们用图6可视化了这些结果。结果表明，Middle-Cycle 无论模型规模或递归深度如何增加，始终能够获得最低的NLL。图中的横线代表了无参数共享的基准模型，红线下为全容量模型，黑线上为等参数的递归模型。

模型	参数配置	递归深度	NLL	准确率	备注
360M	Middle-Cycle	2	最低	最高	表现最优
360M	Middle-Cycle	3	最低	最高	表现最优
135M	Cycle	2	略优	稍高	稍好
135M	Middle-Cycle	3	最优	最高	稳定提升

持续预训练（up-training）下的表现¶

表9展示了在持续预训练（up-training）下的扩展研究结果。模型从开源的SmolLM 检查点继续训练了50亿个token。Middle-Cycle 和 Middle-Sequence 在所有设置中均表现出色，显著优于相同初始化方式但无递归的基线模型。其他策略在提升上表现出明显的瓶颈，说明它们在容量进一步提升上的潜力有限。

表中展示了模型在50亿token的FineWeb-Edu上训练后的NLL和六项任务的少样本准确率。Middle策略在所有条件下都表现出最优性能，并且在不同初始化方式（如Step、Avg、Lower、Upper、Rand）下也具有稳健性。特别地，Middle-Cycle 在多个设置中均实现了最低的NLL和最高的平均准确率。

模型	参数配置	递归深度	初始化方式	NLL	准确率	备注
360M	Middle-Cycle	2	Step	最低	最高	表现最优
360M	Middle-Cycle	2	Upper	最低	最高	表现最优
360M	Middle-Cycle	3	Step	最低	最高	表现最优
360M	Middle-Sequence	2	Step	最低	最高	表现优异
360M	Middle-Sequence	3	Step	优异	优异	表现较佳

总结而言，Middle-Cycle 在大多数情况下是最稳定且性能最优的参数共享策略，尤其是在递归深度增加和模型规模扩大的情况下。此外，Middle-Cycle 在持续训练中也显示出更强的扩展性和适应性。

Appendix F Expanded Results of Design Choices for Router¶

F.1 设计配置细节¶

本节探讨了多种路由器设计选择，以优化性能和稳定性。主要调整的内容包括：

系数（Coeff）：用于控制辅助或平衡损失项的强度。
缩放因子（α）：在路由器函数之后应用，调节路由权重。
激活函数（Func）：测试了如 Sigmoid 和 Softmax 等不同函数。
路由器网络结构（Arch）：包括线性层、2 层 MLP（带 GELU 激活）、以及 Wide-MLP（隐藏层扩大四倍）。

为了提升训练稳定性，引入了z-loss（Zoph et al., 2022）。该损失通过惩罚门控网络产生的大 logits 来防止数值不稳定。其公式为：

\[ L_z(x) = \frac{1}{B} \sum_{i=1}^B \left( \log \sum_{j=1}^{N_r} e^{x_j^{(i)}} \right)^2 \]

其中，$ B $ 是批次中的 token 数量，$ N_r $ 是专家数量，$ x $ 是路由器输入的 logits。z-loss 的作用是鼓励路由器生成较小的 logits，从而提升路由决策的稳定性和可靠性。

F.2 路由器性能评估指标¶

分为两类路由方式的评估：

专家选择路由（Expert-choice routing）¶

死 token 比例：衡量在批次中某些位置的 token 未被选择的比例，反映路由器是否对某些位置存在系统性忽视。
采样准确性：衡量路由器在推理阶段能否预测 token 是否在训练中被选为 top-k，体现其选择重要 token 的能力。

理想的路由器应具有高采样准确性与低死 token 比例，既准确识别重要 token，又保持选择多样性。

令牌选择路由（Token-choice routing）¶

MaxVio（最大不平衡）：衡量专家之间负载不均衡的程度。公式为：

\[ \text{MaxVio} = \frac{\max_i \text{Load}_i - \overline{\text{Load}}_i}{\overline{\text{Load}}_i} \]

熵（Entropy）：衡量专家被选择的多样性。公式为：

\[ H = -\sum_{i=1}^{N_r} \overline{p}_i \log \overline{p}_i \]

高熵值表示 token 分配更均匀，路由决策更平衡。

F.3 路由器设计的扩展评估结果¶

表 10：专家选择路由器的消融实验¶

结果显示：

辅助路由器（Aux Router）与辅助损失（Aux Loss）：虽然提升了采样准确性，但伴随较高死 token 比例（最高达 66.7%），表现出位置偏差。
线性路由器 + 辅助损失：在显著降低死 token 比例的同时，采样准确性未受影响，是较优的组合。
最终选定的策略为：使用 Sigmoid 激活函数、线性结构、小系数的辅助损失，并默认开启 z-loss。

表 11：令牌选择路由器的消融实验¶

结果显示：

平衡损失（Balancing Loss）：显著降低 MaxVio，提升熵，改善负载均衡。
无损失方法（Loss-free）：虽然简化，但 MaxVio 很高，熵较低，负载更加不均衡。
z-loss：有助于提升路由稳定性与模型精度，但在初期阶段仍难以实现完全负载均衡。
最终策略：在平衡损失的基础上，加入小系数的 z-loss 以稳定负载分配。

总结¶

通过大量实验，研究者发现：

路由器设计对性能和稳定性具有显著影响。
线性结构 + Sigmoid 激活 + 小系数辅助损失 + z-loss 是较为理想的配置。
平衡损失和 z-loss 的结合 有助于路由器在提升采样准确性和负载均衡之间的权衡。

Appendix G Expanded Results of KV Cache Sharing Mechanism¶

G.1 递归 Transformer 中的关键值表示趋势¶

在 Vanilla Transformer 中，跨模型深度共享 KV 缓存 是一种提升推理吞吐量的有效手段。这种技术可以显著减少 KV 缓存的内存占用，从而允许更大的推理批量。在某些设计中，甚至可以通过跳过 KV 投影和预填充操作来实现显著加速。由于 Vanilla Transformer 模型具有高度的自由度，其可训练参数可以很好地针对共享缓存进行优化，因此共享缓存时性能下降非常小。

相比之下，递归 Transformer 可用于优化共享键值状态的参数要少得多。尽管如此，研究者假设共享的注意力块之间可能会出现类似模式。为验证这一假设，作者将预训练的递归 Transformer 的 KV 状态 拆解为幅度和方向两个部分。

从 图 7 可以看出，跨递归深度共享 键值投影层 会带来明显的递归模式。尽管隐藏状态的幅度在递归中有所增加，但投影层似乎被训练为在每个递归对应深度上生成相似的信号大小。

图 8 显示了 余弦相似性 的结果，呈现出明显的对角线模式，表明共享的投影层确实生成了高度相似的键和值表示。虽然共享值状态比键状态更具挑战性，但这些发现表明，即使在递归 Transformer 中，KV 缓存共享导致的性能下降仍然是可以接受的。

G.2 KV 缓存共享策略的性能比较¶

KV 缓存共享的实验结果¶

在 表 12 中，作者展示了在 Vanilla、Recursive 和 MoR 模型中应用 KV 缓存共享策略的性能结果。测试了多种缓存共享策略，包括 Cycle 和 Sequence 策略，这些策略与参数共享类似（详见附录 A.1）。

有趣的是，KV 缓存共享甚至提升了 Vanilla 模型的性能，这可能是因为共享起到了正则化的作用。对于递归模型，作者将参数共享与 KV 缓存共享策略对齐，尽管结果略有变化，但 Middle-Cycle 策略（最佳参数共享策略）带来的困惑度下降并不显著。对于 MoR 模型，尽管在最佳设置（专家选择路由机制）下仍有一些轻微的性能下降，但考虑到参数和缓存规模的减少，这种下降是可以接受的。

此外，作者还探索了一种替代的共享策略（标记为 †），该策略使用共享的缓存初始化非活动位置（未选中的位置），并通过实际计算更新活动位置。尽管这种方法并未带来额外的性能提升，但仍值得进一步探索将 KV 共享与实际更新结合使用的可能性。

KV 共享约束的放松¶

在 表 13 中，作者研究了 放松 KV 共享约束 的方法，类似于 Bae 等人对参数共享约束的放松。具体而言，作者重新审视了几种放松技术，包括针对递归输出应用可训练嵌入（通过逐元素相乘）以及使用 LoRA、DoRA 和 Adapt-Prompt tuning 等技术。

结果表明，这些方法从头训练放松模型时并未带来显著的性能提升，这与之前的研究一致，因为它们仅引入了有限数量的额外参数。尽管作者假设基于前缀的方法（如在注意力中添加可训练的前缀）可能会带来更大的收益，但实验中并未发现显著差异。因此，未来仍需探索更复杂的、高效放松 KV 缓存共享约束的方法。

总结¶

本节重点分析了递归 Transformer 中的 KV 缓存共享机制，通过 KV 状态的幅度变化 和 余弦相似性 两种方法验证了共享策略的可行性。实验表明，KV 缓存共享 在减少内存占用的同时，对性能的影响较小，特别是对 Vanilla 模型有正则化效果。尽管递归模型和 MoR 模型在性能上略有下降，但这些下降在参数和缓存规模减少的背景下是可接受的。

最后，作者尝试了 放松 KV 缓存共享约束 的策略，但并未发现显著的性能提升。这一方向仍有待进一步研究。

Appendix H Expanded Qualitative Results¶

H.1 Analysis on Adaptive Computation Paths¶

本节通过表格 14 展示了对模型中每个子词（subword token）分配的递归深度的定性分析。这些可视化结果展示了 MoR 框架下不同子词在递归处理过程中的差异性，体现了模型自适应计算路径的能力。

重点内容：

每个子词根据其复杂程度被分配了不同的递归深度（1、2、3 步），颜色表示递归步数。
部分子词在较浅的递归层级就完成了处理（紫色），而另一些则需要更深层次的处理（蓝色和红色）。
这表明模型能够自适应地将更多计算资源集中在输入中较难处理的部分，从而提升整体效率和效果。
实验模型参数规模为 360M，训练数据量为 30B tokens，使用了辅助损失和递归-wise KV 缓存策略。

表格内容（精简）：

表格中列出了 10 个样本的文本内容，展示了不同语境下递归深度的分布情况。
涉及新闻、政治、文学、科技、教育等多个领域，验证了模型在不同任务上的自适应能力。

H.2 Analysis on Router Weights¶

本节通过图 9 展示了不同路由策略下，选中与未选中 token 的路由器权重分布，分析了模型在递归路径选择中的行为。

重点内容：

Expert-choice 路由器（带有辅助损失）表现出近乎完美的负载均衡，选中 token 的权重接近 1，未选中 token 的权重接近 0。这说明它能有效地将 token 分配到最合适的递归层级。
Auxiliary 路由器虽然也能区分选中与未选中 token，但存在一定的重叠，说明其区分能力略弱。
Token-choice 路由器（包括使用平衡损失和无损失的版本）虽然具备一定的负载均衡能力，但在大多数情况下未能达到最优的负载均衡效果，说明在异构专家（expert）设置下实现一致性能仍存在挑战。
所有模型都使用递归-wise KV 缓存策略，以提升计算效率。

图示说明（精简）：

图 9 提供了四种路由策略下的权重分布图：
- (a) Expert-choice + 辅助损失
- (b) Auxiliary 路由器
- (c) Token-choice + 平衡损失
- (d) Token-choice + 无损失
每个子图展示了在不同递归层级下路由器权重的分布情况，验证了不同策略在实际中的表现差异。

总结¶

Appendix H 通过表格和图表展示了 MoR 模型在自适应计算路径和路由机制上的定性分析。
H.1 表明模型能够根据输入复杂性动态分配计算资源，提高计算效率。
H.2 表明不同路由策略在负载均衡、token 分配等方面表现不一，其中 Expert-choice + 辅助损失的策略效果最佳，而 Token-choice 策略在异构专家环境下仍有改进空间。

这些分析为理解 MoR 模型的行为提供了直观的视角，也为进一步优化路由机制和计算分配提供了依据。

2507.10524_Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation¶

总结¶

From Moonlight¶

三句摘要¶

关键词¶

摘要¶

Abstract¶

1 Introduction¶

背景与动机¶

递归 Transformer 与挑战¶

MoR：统一框架¶

概念与意义¶

贡献总结（Contributions）¶

总结¶

2 Method¶

2.1 Preliminary¶

Recursive Transformers¶

Table 1: 参数共享策略¶

Enhanced training and inference efficiency in recursive models¶

Limitations in prior works¶

2.2 Mixture-of-Recursions (MoR)¶

2.2.1 Routing Strategies: Expert-choice vs. Token-choice¶

Expert-choice Routing¶

Token-choice Routing¶

对比与总结（Table 2 - Left）¶

2.2.2 KV Caching Strategies: Recursion-wise Caching vs. Recursive Sharing¶

Recursion-wise KV Caching¶

Recursive KV Sharing¶

对比与总结（Table 2 - Right）¶

总结¶

3 Experiments¶

表格总结（Table 3）¶

3.1 主要结果¶

MoR 在相同训练计算资源下表现更优¶

MoR 在相同数据量下计算效率更高¶

MoR 的性能与路由和缓存策略有关¶

3.2 IsoFLOP 分析¶

实验设置¶

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

3.3 推理吞吐量评估¶

实验设置¶

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

总结¶

4 Ablation Studies¶

4.1 Parameter Sharing Strategies¶

4.2 Routing Strategies¶

In expert-choice routing, auxiliary loss and linear router yield the best performance.¶

In token-choice routing, balancing loss yields stable and accurate routing.¶

表4总结：¶

4.3 KV Caching Strategies¶

KV sharing robustly works even in parameter-shared architectures.¶

KV sharing degrades expert-choice but benefits token-choice routing in MoR.¶

总结¶

5 Analysis¶

5.1 Compute-optimal Scaling Analysis¶

MoR scaling favors model size over training length under isoFLOPs.¶

5.2 Routing Analysis¶

The allocation of recursion depth reflects token semantic importance.¶

Expert-choice router with auxiliary loss perfectly separates selected from unselected tokens.¶

5.3 Test-time Scaling Analysis¶

MoR enables test-time scaling via deeper recursion.¶

总结¶

6 Related Work¶

Recursive Transformers（递归Transformer）¶

Adaptive Computation（自适应计算）¶

Routing Mechanism（路由机制）¶

Key-value Caching（键值缓存）¶

Latent Reasoning（隐式推理）¶

7 Conclusion¶

7.1 局限性与未来工作¶

推理MoR模型¶

进一步扩展模型规模¶

自适应容量控制¶

与稀疏算法的兼容性¶

扩展到多模态及非文本领域¶

7.2 致谢¶

Appendix A Details of Design Choices for Mixture-of-Recursions¶

A.1 参数共享策略（Parameter-sharing Strategy）¶

A.2 路由策略（Routing Strategy）¶

Expert-choice 路由¶