Efficient and Trainable Language Model Test-Time Scaling via Local Branch Routing

背景

在大型语言模型（LLM）的发展中，推理能力（Reasoning）的提升往往依赖于“测试时扩展”（Test-Time Scaling, TTS）。这类方法旨在通过增加模型在推理阶段（即生成答案时）的计算资源，而非仅仅依赖训练时的参数规模，来获得更优的推理结果。

然而，现有的测试时扩展方法面临着严峻的权衡困境：

1. 长思维链采样（Long Chain-of-Thought Sampling）：虽然简单直接，但本质上是单线程的，缺乏对多种可能性的并行探索，容易陷入局部最优。
2. 句子级或解空间级搜索（Sentence- or Solution-level Search）：这类方法（如蒙特卡洛树搜索 MCTS）能够进行全局探索，但计算成本极高，且难以进行端到端的训练优化。

此外，传统的离散令牌强化学习（Discrete-token RLVR）在处理分支结构时，往往难以定义清晰的轨迹似然（Trajectory Likelihood），导致奖励信号难以有效回传。因此，业界急需一种既能保留离散分支的明确身份，又能高效计算、支持端到端训练的测试时扩展框架。

核心内容

本文提出了一种名为 Local Branch Routing (LBR，局部分支路由) 的新框架。LBR 是一种基于 Token 级别的测试时扩展方法，旨在通过轻量级的局部搜索来增强模型的推理能力，同时保持训练的高效性。

1. 工作原理：剪枝-移位-生长（Prune-Shift-Grow）

LBR 的核心机制是在生成每个 Token 时，执行以下流程：

• 局部前瞻树扩展：模型不会只预测下一个 Token，而是扩展一个小的局部前瞻树（Local Lookahead Tree）。
• 分支前向传播：所有采样的分支（Branches）都会通过语言模型进行前向传播，生成对应的隐藏状态（Hidden States）。
• 轻量级路由选择：使用一个轻量级的路由器（Router），基于候选局部未来（Candidate Local Futures）的隐藏状态，选择一个深度为 1 的子树（Depth-1 Subtree）进行提交（Commit）。
• 决策依据：路由过程利用了候选分支的隐藏状态作为证据，使得每个 Token 的决策不仅仅依赖于根节点的下一个 Token 分布，还能参考其他潜在路径的信息，从而避免了昂贵的全解空间搜索。

2. 可训练性与似然定义

LBR 的关键创新在于其定义了一个可处理的轨迹似然（Tractable Tree-Trajectory Likelihood），这使得端到端的强化学习成为可能：

• 离散分支身份的保留：该过程保留了离散分支的身份标识。
• 似然计算规则：
  • 新生长的节点在首次被采样时即被计数。
  • 路由器的决策被分配了明确的概率值。
• 端到端优化：基于上述似然定义，LBR 允许使用可验证奖励（Verifiable Rewards）进行端到端的强化学习。模型可以在离散令牌 RLVR（RLVR）所采用的相同似然比原则下，联合优化基础语言模型和路由器。

3. 实验验证

作者在两个主要领域验证了 LBR 的有效性：

• 合成分层规划任务（Synthetic Hierarchical-Planning Tasks）：实验表明，候选隐藏状态确实为路由决策提供了有用的证据，证明了局部分支信息的有效性。
• 数学推理基准测试（Mathematical Reasoning Benchmarks）：在多个数学推理基准上，LBR 在 Pass@1 和 Pass@32 指标上均优于：
  • 离散思维链（Discrete Chain-of-Thought）
  • 原始的离散令牌 RLVR
  • 与 RL 兼容的软令牌分支基线（RL-compatible soft-token branching baselines）

关键要点

• Token 级局部搜索：LBR 不同于全局搜索或长思维链，它在 Token 生成阶段引入局部的分支探索，平衡了计算效率与探索深度。
• 隐藏状态作为路由证据：路由器不依赖昂贵的重采样或完整路径评估，而是直接利用候选分支的隐藏状态（Hidden States）进行快速决策。
• 支持端到端强化学习：通过定义清晰的树轨迹似然，LBR 解决了离散分支难以计算概率的问题，使得基础模型和路由器可以联合训练。
• 性能显著提升：在数学推理任务中，LBR 在 Pass@1（单次尝试成功率）和 Pass@32（32次尝试中至少一次成功的概率）上均超越了现有的主流基线方法。
• 高效且可训练：LBR 提供了一种轻量级的、离散的测试时扩展形式，既避免了全解搜索的高成本，又克服了传统方法难以优化的痛点。

意义与影响

LBR 的提出为大型语言模型的推理增强提供了一条新的技术路径。其核心意义在于打破了“推理质量”与“训练/计算效率”之间的零和博弈：

1. 填补了方法论空白：现有的测试时扩展方法要么太简单（单线程 CoT），要么太复杂难训练（全局搜索）。LBR 通过局部分支路由，在两者之间找到了一个高效的平衡点。
2. 推动了离散强化学习的发展：LBR 对离散树轨迹似然的清晰定义，为离散令牌空间中的强化学习提供了更坚实的数学基础，使得联合优化模型和路由策略成为可能。
3. 实用性强：由于 LBR 仅涉及轻量级的局部前瞻和路由，其计算开销远低于 MCTS 等全局搜索方法，更易于在实际部署中应用，特别是在对延迟敏感的场景下。
4. 验证了局部信息的有效性：实验结果证实，仅利用局部未来的隐藏状态即可显著提升路由决策的质量，这暗示了语言模型内部表示中蕴含丰富的局部结构信息，值得进一步挖掘。

总之，LBR 证明了通过轻量级的局部分支和路由机制，可以实现高效、可训练且高性能的语言模型测试时扩展，为未来推理增强模型的设计提供了重要的参考范式。