# 2507.10524_Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation * 首页: * PDF: * 引用: 19(2025-12-14) * 组织 ## 总结 ## From Moonlight ### 三句摘要 1. ✨ Mixture-of-Recursions (MoR) 提出了一种统一的Transformer框架，它将parameter sharing和adaptive computation结合起来，旨在应对大型语言模型的计算和内存挑战。 2. 🚀 MoR通过轻量级router动态分配individual token的recursion depth，并利用recursion-wise KV caching或recursive KV sharing策略，从而只对活跃token进行quadratic attention computation。 3. 📈 实验证明，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)** 1. **预备知识 (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 策略在递归模型中表现最佳，因为它能更好地捕捉输入和输出的细微差别。 2. **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)** 1. **主要结果 (Main Results)**：在相同训练 FLOPs 预算下，MoR 模型（使用 expert-choice 路由器和两个递归）相比 Vanilla Transformer 和其他递归基线，在参数量减少近 50% 的情况下，实现了更低的验证损失和更高的 few-shot 准确率（43.1% vs. 42.3%）。在相同训练 token 数量下，MoR 在使用 25% 更少 FLOPs 的情况下，性能优于基线。 2. **IsoFLOP 分析 (IsoFLOP Analysis)**：MoR 在 135M 到 1.7B 参数规模的广泛模型尺寸和计算预算下，始终优于递归基线。对于大于 360M 参数的模型，MoR 的性能与 Vanilla Transformer 持平甚至超越，同时显著减少了参数数量（约三分之一）。这表明 MoR 是一种可扩展且参数高效的架构。 3. **推理吞吐量 (Inference Throughput)**：MoR 结合了连续深度批处理 (continuous depth-wise batching) 和早退出 (early-exiting) 机制，显著提高了推理吞吐量。与 Vanilla Transformer 相比，MoR 模型实现了高达 2.06 倍的速度提升。 **消融研究 (Ablation Studies)** 1. **参数共享策略 (Parameter Sharing Strategies)**：Middle-Cycle 策略在所有设置下都表现最佳，并在 NLL 和 few-shot 准确率方面持续优于其他策略，尤其是在模型规模和递归深度增加时。 2. **路由策略 (Routing Strategies)**： * **Expert-choice**: 辅助损失结合 sigmoid 激活函数和简单线性路由器表现最佳。Z-loss 对准确率影响微乎其微，但能轻微减少“死亡 token”比例。 * **Token-choice**: 均衡损失能实现更稳定的路由，显著降低 MaxVio 指标。Softmax 激活函数与 MLP 路由器搭配表现最佳。 3. **KV 缓存策略 (KV Caching Strategies)**：递归 KV 共享在参数共享架构中表现稳健，但对于 expert-choice 路由的 MoR 模型，它可能导致性能下降，因为精确的 token 选择更受益于独立的递归级缓存。然而，对于 token-choice 路由，KV 共享可以提供额外的上下文信息，有助于弥补其路由决策的不足。 **分析 (Analysis)** 1. **计算最优扩展分析 (Compute-optimal Scaling Analysis)**：MoR 的计算最优扩展行为表明，它更受益于参数数量的增加（模型更大），而不是训练数据量的增加（训练步数更长）。这意味着 MoR 的最佳扩展策略是增加模型容量。 2. **路由分析 (Routing Analysis)**：路由器的决策反映了后续 token 的上下文可预测性。例如，单词的第二部分通常易于预测，因此需要的递归步骤较少。Expert-choice 路由器与辅助损失结合，能完美地将选中和未选中的 token 分离。 3. **测试时间扩展分析 (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 扩展到数百亿参数，显著提高了模型的泛化能力和推理能力。然而，这种扩展也带来了巨大的内存和计算开销，导致训练和部署变得困难。为了解决这一问题，研究者们提出了多种“高效”设计，主要包括两个方面： 1. **参数效率**（Parameter Efficiency）：通过参数共享、层绑定（layer tying）等方式减少模型参数数量。 2. **自适应计算**（Adaptive Computation）：根据任务复杂度动态分配计算资源，如早期退出（early-exit）方法。 尽管在参数效率和自适应计算方面分别取得了一些进展，但目前还没有一种架构能够 **同时** 实现这两种效率。 --- ### 递归 Transformer 与挑战 递归 Transformer 是一种将共享的参数块多次应用于输入的方法，天然具备参数效率。然而，以往的动态递归方法通常面临以下挑战： - 需要额外的训练过程； - 部署效率低； - 往往退化为固定深度的递归，无法真正实现每个 token 的自适应计算。 --- ### MoR：统一框架 为克服上述问题，作者提出了 **Mixture-of-Recursions (MoR)**。其核心思想是： - 通过 **端到端训练的轻量路由模块（router）**，为每个 token 动态分配递归深度； - 利用共享的参数块对每个 token 进行不同次数的计算； - 通过 **递归级别的 KV 缓存**（recursion-wise KV caching），减少内存访问，提高吞吐量。 MoR 的优势在于： 1. **参数共享**：减少整体参数数量； 2. **自适应计算**：每个 token 的计算量根据其复杂程度动态变化； 3. **高效缓存**：仅缓存与当前 token 递归深度相关的 KV 对，减少内存开销。 --- ### 概念与意义 MoR 可以看作是一种在推理过程中 **在隐空间中进行非语言“思考”** 的预训练框架。它在生成每个 token 时就进行递归思考，而不是依赖生成前的连续提示（如一些 prompt-augmentation 方法）。 与以往方法相比，MoR 实现了： - 在模型的 **深度方向（vertical axis）** 上进行可变长度的推理； - 摆脱了传统方法中固定深度的限制； - 统一了参数效率和自适应计算的优势。 --- ### 贡献总结（Contributions） 1. **统一的高效语言模型框架**：首创新型架构 MoR，同时融合参数共享、自适应 token-level 思考深度和高效的 KV 缓存策略。 2. **动态递归路由**：从头训练路由模块，为每个 token 分配递归深度，且无需额外的训练阶段。 3. **实验验证充分**：在相同计算预算下，MoR 相比传统和递归基线模型，在验证损失和 few-shot 准确率上均表现更优。 4. **高效架构**：通过选择性计算和 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的递归深度**。其核心包括两个部分： 1. **路由机制（Routing）**：根据token的复杂度分配其递归深度，动态决定哪些token需要更多计算。 2. **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效率高，适合动态递归深度。 - **缺点**：每个递归层级需要独立缓存，不便于全局信息共享。 ##### Recursive KV Sharing - **只在第一递归层级缓存KV对**，并在所有后续层级共享。 - **优点**：极大节省内存，避免重复计算。 - **缺点**：注意力FLOPs减少有限，KV 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.1 Parameter Sharing Strategies** **重点内容：Middle-Cycle 是最有效的参数共享策略。** - 参数共享是 **Recursive Transformers** 和 **MoR** 的关键组成部分。 - 实验比较了四种策略：Cycle、Sequence、Middle-Cycle、Middle-Sequence。 - 评估模型大小为 135M 和 360M 的递归式Transformer。 - **Middle-Cycle** 在验证集上的损失最低，表现最佳（见图4(b)）。 - 由此，论文后续所有 MoR 和递归Transformer均采用 **Middle-Cycle** 参数共享策略。 --- ### **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** #### **KV sharing robustly works even in parameter-shared architectures.** - **KV共享在参数共享架构中依然有效。** - 图4(c) 显示了在 **Vanilla Transformer** 和 **Recursive Transformer** 中使用 KV 共享的效果。 - 尽管共享参数的 Recursive Transformer 自由度较低，但其对 KV 共享仍然表现出鲁棒性。 - 分析表明，在不同递归深度下，共享参数的 KV 对在 **幅度和方向上高度一致**，因此 KV 共享仅导致轻微性能下降（见附录G）。 --- #### **KV sharing degrades expert-choice but benefits token-choice routing in MoR.** - **KV共享对 expert-choice 路由有害，但对 token-choice 路由有益。** - **expert-choice 路由**：KV共享导致性能显著下降，可能是因为该策略只关注当前递归中被选中的 token，而共享的 KV 干扰了这种选择。 - **token-choice 路由**：共享的 KV 提供了额外的上下文信息，可以弥补该策略中较弱的路由决策。 - **KV共享的两种机制比较**： - **Recursion-wise KV 缓存**：每个递归层级有自己的KV缓存。 - **Recursive KV 共享**：多个递归层级共享KV缓存，节省内存和FLOPs。 - 尽管共享KV会增加注意力计算的FLOPs，但总体FLOPs仍减少。 --- ### 总结 本节通过系统性的消融实验，确定了在 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.1 参数共享策略（Parameter-sharing Strategy） 本节介绍了四种参数共享策略：**Cycle**、**Middle-Cycle**、**Sequence** 和 **Middle-Sequence**。这些策略决定了模型中共享的模块 Φ′ 如何在 L 层展开的递归结构中被复用。 - **Cycle（循环策略）** 使用固定参数集在所有递归步骤中循环复用。这种方式鼓励模型在每个阶段“重新思考”输入，从而实现深度迭代。但缺点是，相同的变换会被重复应用，可能限制模型学习多样化特征的能力。 - **Sequence（序列策略）** 每个递归块使用不同的参数，按顺序分配。虽然结构稳定，但连续应用相似变换可能导致冗余特征和效果递减。 - **Middle（中间共享）策略（Middle-Cycle / Middle-Sequence）** 在第一层和最后一层保留独有参数，中间层共享参数。这种策略在参数效率和表示灵活性之间取得了平衡。根据已有研究（Kim et al., 2023a; Geiping et al., 2025），Middle 策略在捕捉输入输出的细微差异方面优于单纯的 Cycle 或 Sequence，且不会显著增加模型规模。 📌 **重点总结**：Middle 策略在实际中表现更优，因为它在保持模型结构灵活性的同时，减少了中间层的参数冗余。 --- ### A.2 路由策略（Routing Strategy） 本节讨论了 MoR 中使用的两种路由策略：**Expert-choice（专家选择）** 和 **Token-choice（令牌选择）**。 #### Expert-choice 路由 - **优点**：计算预算固定，便于资源管理。 - **缺点**：使用 top-k 选择机制时，依赖后续 token 的信息，违反了自回归模型的因果性，可能导致信息泄露，影响模型推理可靠性。 **缓解方法**： 1. **辅助路由（Auxiliary Router）**：训练一个轻量级网络，在推理时预测 tokens 是否会被选中为 top-k。不参与主任务训练，仅用于推理。 2. **辅助损失（Auxiliary Loss）**：为主路由器添加二分类交叉熵损失，引导其在训练中将 top-k tokens 推向 1，其他 tokens 推向 0。 #### Token-choice 路由 - **优点**：为每个 token 分配递归深度，不固定计算预算，保持自回归属性，避免信息泄露。 - **缺点**：导致 experts（递归层）之间负载不平衡，影响训练效率和计算资源分配。 **缓解方法**： 1. **平衡损失（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 $ 的平均路由得分。 2. **无损失方法（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）。 学习率调度公式如下： $$ \eta(t)= \begin{cases} \frac{t}{w}\,\eta_{\max}, & 0\leq t