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

2026-06-23

AI Summary

This work addresses the challenge of balancing inference efficiency and computational cost in test-time scaling for language models, while also introducing a differentiable, end-to-end trainable mechanism for discrete branching. The authors propose Local Branch Routing (LBR), a framework that constructs lightweight look-ahead trees at each token position and employs a differentiable router to select the optimal subtree based on candidate future hidden states. LBR is the first method to enable differentiable routing over token-level local tree structures, preserving discrete branch identities and defining a tractable trajectory likelihood. This formulation supports likelihood-ratio-based reinforcement learning for discrete tokens (RLVR). On mathematical reasoning benchmarks, LBR substantially outperforms standard chain-of-thought, original RLVR, and soft-branching approaches, simultaneously improving both Pass@1 and Pass@32 performance.

Abstract

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.

Problem

• test-time scaling
• language model
• chain-of-thought
• end-to-end training
• computational efficiency

Innovation

• Local Branch Routing
• test-time scaling
• token-level branching
• end-to-end reinforcement learning
• discrete reasoning