Efficient and Trainable Language Model Test-Time Scaling via Local Branch Routing
作者: Yutong Yin, Mingyu Jin, Jin Pan, Changyi Yang, Zijie Xia, Dhruv Pai, Shuming Hu, Zhen Zhang, Chenyang Zhao, Jinman Zhao, Wujiang Xu, Raymond Li, Xin Eric Wang, Julian McAuley, Zhaoran Wang
分类: cs.CL
发布日期: 2026-06-24
💡 一句话要点
提出局部分支路由以解决语言模型测试时间扩展问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 语言模型 测试时间扩展 局部分支路由 推理能力 强化学习 自然语言处理 模型优化
📋 核心要点
- 现有的测试时间扩展方法在推理能力与计算效率之间存在显著的权衡，难以实现高效的端到端训练。
- 本文提出局部分支路由（LBR），通过局部前瞻树和轻量级路由器，优化了令牌级的决策过程。
- 在多个基准测试中，LBR在推理准确率上超越了传统的链思维和其他基线方法，展示了其有效性。
📝 摘要（中文）
测试时间扩展可以提升语言模型的推理能力，但现有方法常面临权衡：长链思维采样单线程，而句子或解决方案级搜索则计算开销大且难以端到端训练。本文提出局部分支路由（LBR），一种基于令牌的测试时间扩展框架，通过扩展小型局部前瞻树，前向所有采样分支并使用轻量级路由器选择深度-1子树。LBR通过对候选局部未来的隐藏状态进行路由，使每个令牌决策能够利用超出根节点的证据，同时避免全解搜索。实验表明，LBR在合成层次规划任务和数学推理基准上均表现出色，显著提高了模型性能。
🔬 方法详解
问题定义：本文旨在解决现有语言模型在测试时间扩展中面临的效率与推理能力之间的权衡问题。现有方法往往计算开销大，难以进行端到端训练。
核心思路：提出局部分支路由（LBR）框架，通过扩展小型局部前瞻树，利用轻量级路由器选择最优的子树，从而在不进行全解搜索的情况下提升推理能力。
技术框架：LBR的整体架构包括局部前瞻树的构建、所有分支的前向传播、以及基于隐藏状态的路由选择。该框架通过对每个令牌的决策进行优化，提升了模型的推理效率。
关键创新：LBR的主要创新在于通过局部分支路由实现了令牌级的推理扩展，避免了全解搜索的复杂性，同时保留了离散分支的身份。
关键设计：LBR设计了轻量级的路由器，能够在每次采样时为新生成的节点分配明确的概率，支持基于强化学习的端到端优化。
🖼️ 关键图片
📊 实验亮点
在合成层次规划任务中，LBR展示了其有效性，显著提高了模型的推理准确率。在数学推理基准上，LBR在Pass@1和Pass@32指标上均超过了传统的链思维和其他基线方法，表明其在推理任务中的优势。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理、智能问答系统和自动化推理等。通过提高语言模型的推理效率，LBR能够在实际应用中显著提升用户体验和系统性能，具有广泛的商业价值和研究前景。
📄 摘要（原文）
Test-time scaling improves language-model reasoning, but existing approaches often face a difficult trade-off: long chain-of-thought sampling remains single-threaded, while sentence- or solution-level search can be computationally expensive and hard to train end-to-end. We introduce Local Branch Routing (LBR), a token-level test-time scaling framework that expands a small local lookahead tree, forwards all sampled branches through the language model, and uses a lightweight router to select the depth-1 subtree to commit. By routing over the hidden states of candidate local futures, LBR allows each token decision to use evidence beyond the root next-token distribution while avoiding full solution-level search. The resulting prune-shift-grow decoding process preserves discrete branch identities and defines a tractable tree-trajectory likelihood: newly grown nodes are counted when first sampled, and router decisions are assigned explicit probabilities. This enables end-to-end reinforcement learning with verifiable rewards, jointly optimizing the base model and router under the same likelihood-ratio principle as discrete-token RLVR. On synthetic hierarchical-planning tasks, LBR shows that post-candidate hidden states provide useful routing evidence. On mathematical reasoning benchmarks, LBR improves both Pass@1 and Pass@32 over discrete chain-of-thought, vanilla discrete-token RLVR, and RL-compatible soft-token branching baselines. These results suggest that lightweight local branching offers an efficient, trainable, and discrete form of language-model test-time scaling.