提出轻量级局部分支路由,让模型推理更高效可训练。
Efficient and Trainable Language Model Test-Time Scaling via Local Branch Routing

- 在每个词决策时扩展局部搜索树,用轻量路由器选择最优路径。
- 在数学推理任务上,Pass@1和Pass@32均优于现有方法。
- 适合需要高效、可端到端训练的复杂推理场景。
测试时缩放能提升语言模型的推理能力,但现有方法常面临困境:长思维链采样仍为单线程,而句子或解级搜索计算成本高且难以端到端训练。本文提出局部分支路由(LBR),一种基于词级别的测试时缩放框架,通过扩展小规模局部前瞻树,将所有采样分支送入语言模型,并用轻量级路由器选择深度为1的子树进行提交。通过在候选未来状态的隐藏层上进行路由,LBR使每个词决策能利用根节点下一词分布之外的证据,同时避免全解级搜索。由此产生的剪枝-移动-生长解码过程保留了离散分支身份,并定义了可计算的树轨迹似然:新生成节点首次采样时计数,路由器决策赋予明确概率。这使得端到端强化学习成为可能,使用可验证奖励联合优化基础模型与路由器,遵循离散词元强化学习中的似然比原则。在合成层次规划任务中,LBR验证了后候选隐藏状态提供了有效的路由依据。在数学推理基准上,相比离散思维链、原始离散词元强化学习、以及兼容强化学习的软词元分支基线,LBR在Pass@1和Pass@32上均有提升。结果表明,轻量局部分支提供了一种高效、可训练且离散的语言模型测试时缩放方式。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。