论文

通过本地分支路由实现高效且可训练的语言模型测试时间扩展

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

模型推理测试时计算扩展

摘要

测试时间缩放改进了语言模型推理,但现有方法通常面临困难的权衡:长链思维采样仍然是单线程的,而句子或解决方案级搜索的计算成本可能很高,并且难以进行端到端训练。我们引入本地分支路由(LBR),这是一个 词元级 测试时扩展框架,它扩展一个小型本地前瞻树,通过语言模型转发所有采样分支,并使用轻量级路由器选择要提交的深度 1 子树。通过路由候选局部未来的隐藏状态,LBR 允许每个词元决策使用根下一个词元分布之外的证据,同时避免完整的解决方案级搜索。由此产生的修剪-移位-生长解码过程保留了离散的分支身份并定义了易于处理的树轨迹可能性:首次采样时对新生长的节点进行计数,并且为路由器决策分配了显式概率。这使得端到端强化学习具有可验证的奖励,并在与离散词元 RLVR 相同的似然比原理下联合优化基本模型和路由器。在综合分层规划任务中,LBR 表明后候选隐藏状态提供了有用的路由证据。在数学推理基准上,LBR 相对于离散思维链、普通离散词元 RLVR 和 RL 兼容的软词元分支基线改进了 Pass@1 和 Pass@32。这些结果表明,轻量级本地分支提供了一种高效、可训练和离散形式的语言模型测试时间扩展。