Efficient and Trainable Language Model Test-Time Scaling via Local Branch Routing
- URL: http://arxiv.org/abs/2606.25354v2
- Date: Mon, 29 Jun 2026 23:57:32 GMT
- Title: Efficient and Trainable Language Model Test-Time Scaling via Local Branch Routing
- Abstract summary: Local Branch Routing (LBR) is a token-level test-time scaling framework.<n>It 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.<n>LBR allows each token decision to use evidence beyond the root next-token distribution while avoiding full solution-level search.
- Score: 65.59464371673623
- License: http://creativecommons.org/licenses/by/4.0/
- 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.
Related papers
- Trajectory-Level Speculative Decoding for Diffusion Language Models [16.97093186372498]
Diffusion-based language models (dLLMs) enable parallel token generation through iterative denoising, but existing decoding strategies collapse to single-token generation under low confidence.<n>We develop a trajectory-level speculative framework that constructs draft denoising trajectories via confidence-stratified tree exploration.<n>Our framework reduces denoising by 30-40% and increases tokens-per-step from 2.6 to 4.3, achieving 7-14x speedup over vanilla dLLMs and 1.3x over Fast-dLLM with less than 1% accuracy change across reasoning and code benchmarks.
arXiv Detail & Related papers (2026-08-27T09:42:20Z) - TELLER: Non-intrusive Cross-Layer Root-Cause Analysis for LLM Inference [47.62941623025392]
We present TELLER, a non-intrusive and Log-aware LLM inference Root-cause analysis framework.<n>TELLER collects NVTX/CUPTI traces and service logs without modifying model binaries.<n>It reconstructs per-request call-chain trees and aligns log lines with corresponding execution steps.<n>On top of these representations, TELLER combines numeric candidate localization with a multimodal root-cause model.
arXiv Detail & Related papers (2026-08-03T09:39:11Z) - TARPO: Token-Wise Latent-Explicit Reasoning via Action-Routing Policy Optimization [18.977861031175756]
TARPO is a pure RL framework that switches between discrete token generation and continuous latent reasoning at each step.<n> TARPO consistently outperforms existing explicit and latent reasoning RL baselines across diverse benchmarks.
arXiv Detail & Related papers (2026-06-04T08:30:53Z) - Route Before Retrieve: Activating Latent Routing Abilities of LLMs for RAG vs. Long-Context Selection [57.3886742625188]
Pre-Route is a proactive routing framework that performs structured reasoning before answering.<n>Our study shows three key findings: (i) LLMs possess latent routing ability that can be reliably elicited with guidelines; (ii) linear probes reveal that structured prompts sharpen the separability of the "optimal routing" in representation space; and (iii) distillation transfers this reasoning structure to smaller models for lightweight deployment.
arXiv Detail & Related papers (2026-05-11T09:10:55Z) - Latent Abstraction for Retrieval-Augmented Generation [4.6096940605642915]
Existing systems rely on generating natural language queries at each hop.<n>We propose textbfLAnR (Latent Abstraction for RAG), a unified framework in which a single LLM jointly performs encoding, retrieval, and generation entirely within its own latent space.
arXiv Detail & Related papers (2026-04-20T06:26:13Z) - TARo: Token-level Adaptive Routing for LLM Test-time Alignment [15.824259013420807]
TARo steers frozen large language models toward structured reasoning entirely at inference time.<n>Extensive experiments show that TARo significantly improves reasoning performance by up to +22.4% over base model.<n> TARo also generalizes from small to large backbones without retraining.
arXiv Detail & Related papers (2026-03-19T02:18:40Z) - $V_1$: Unifying Generation and Self-Verification for Parallel Reasoners [69.66089681814013]
$V_$ is a framework that unifies generation and verification through efficient pairwise ranking.<n>$V_$-Infer improves Pass@1 by up to $10%$ over pointwise verification.<n>$V_$-PairRL achieves $7$--$9%$ test-time scaling gains over standard RL and pointwise joint training.
arXiv Detail & Related papers (2026-03-04T17:22:16Z) - Test-Time Scaling with Diffusion Language Models via Reward-Guided Stitching [66.39914384073145]
We propose a self-consistency framework that turns cheap diffusion-sampled reasoning into a reusable pool of step-level candidates.<n>We find that step-level recombination is most beneficial on harder problems.<n>Our training-free framework improves average accuracy by up to 2 across six math and coding tasks.
arXiv Detail & Related papers (2026-02-26T11:08:39Z) - Prism: Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models [96.0074341403456]
Inference-time compute has re-emerged as a practical way to improve LLM reasoning.<n>Most test-time scaling (TTS) algorithms rely on autoregressive decoding.<n>We propose Prism, an efficient TTS framework for dLLMs.
arXiv Detail & Related papers (2026-02-02T09:14:51Z) - Chain-in-Tree: Back to Sequential Reasoning in LLM Tree Search [4.12237459236889]
Chain-in-Tree (CiT) is a framework that decides when to branch during search rather than expanding at every step.<n>It achieves 75-85% reductions in token generation, model calls, and runtime on GSM8K and Math500.
arXiv Detail & Related papers (2025-09-30T06:18:44Z) - Beyond Next Token Probabilities: Learnable, Fast Detection of Hallucinations and Data Contamination on LLM Output Distributions [60.43398881149664]
We introduce LOS-Net, a lightweight attention-based architecture trained on an efficient encoding of the LLM Output Signature.<n>It achieves superior performance across diverse benchmarks and LLMs, while maintaining extremely low detection latency.
arXiv Detail & Related papers (2025-03-18T09:04:37Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.