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

# 1 Introduction
Source: [https://arxiv.org/html/2606.25354](https://arxiv.org/html/2606.25354)
![[Uncaptioned image]](https://arxiv.org/html/2606.25354v1/figs/northwestern_logo.png)![[Uncaptioned image]](https://arxiv.org/html/2606.25354v1/figs/tilde_logo.png)

June 2026

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

Yutong Yin1, Mingyu Jin2, Jin Pan3,5, Changyi Yang4,5, Zijie Xia5, Dhruv Pai6, Shuming Hu6, Zhen Zhang7, Chenyang Zhao5, Jinman Zhao8, Wujiang Xu2, Raymond Li9, Xin Eric Wang7, Julian McAuley10, Zhaoran Wang1

1Northwestern University,2Rutgers University,3University of Wisconsin–Madison,4Carnegie Mellon University,5LMSYS Org,6Tilde Research,7University of California, Santa Barbara,8University of Toronto,9University of British Columbia,10University of California, San Diego

###### 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\.

Recent progress in language\-model reasoning shows that allocating more computation at inference time can substantially improve performance\. Chain\-of\-thought prompting encourages models to externalize intermediate reasoning steps\[[28](https://arxiv.org/html/2606.25354#bib.bib15)\], self\-consistency improves accuracy by sampling multiple reasoning traces\[[27](https://arxiv.org/html/2606.25354#bib.bib17)\], and broader studies of test\-time scaling show that the value of extra compute depends strongly on how it is allocated\[[24](https://arxiv.org/html/2606.25354#bib.bib16)\]\. Search\-based methods further introduce width by exploring trees of intermediate thoughts or actions\[[32](https://arxiv.org/html/2606.25354#bib.bib29),[10](https://arxiv.org/html/2606.25354#bib.bib19)\], but such search is typically coarse\-grained, computationally expensive, and difficult to optimize as a single trainable decoding policy\.

This paper asks whether language models can obtain a lightweight form of test\-time width at the token level\. Standard autoregressive decoding commits to one next token before observing the hidden states induced by nearby alternatives\. Yet in many reasoning problems, a local choice—such as a digit, operator, variable, or short phrase—can determine which latent state the model enters next\. This suggests that candidate tokens should not only be scored by the root next\-token distribution, but also by the local future states they induce\.

We propose*Local Branch Routing*\(LBR\), a trainable token\-level branch\-and\-route decoding framework\. At each decoding step, LBR expands a small local lookahead tree of widthKKand depthLL, forwards all sampled branches through the language model, and uses a lightweight router to select which depth\-1 subtree to commit\. The selected token is appended to the output, unselected subtrees are pruned, and the surviving subtree is shifted forward and regrown\. Thus, each token decision can use post\-candidate hidden states while avoiding full solution\-level search\.

LBR also defines a tractable tree\-trajectory likelihood\. The only stochastic operations are sampling newly grown tree nodes from the language model and sampling the router decision; pruning, shifting, and reuse are deterministic\. This factorization allows LBR to be trained end\-to\-end with reinforcement learning from verifiable rewards, following the likelihood\-ratio principle used in recent mathematical\-reasoning systems\[[23](https://arxiv.org/html/2606.25354#bib.bib38),[9](https://arxiv.org/html/2606.25354#bib.bib25)\]\.

Soft Thinking offers a related token\-level branch\-and\-merge mechanism by replacing a committed discrete token with a continuous mixture of candidate token embeddings\[[33](https://arxiv.org/html/2606.25354#bib.bib1)\]\. RL\-compatible variants such as Multiplex Thinking add stochastic candidate sampling, making the sampled candidates trainable with likelihood\-ratio RL\[[25](https://arxiv.org/html/2606.25354#bib.bib12)\]\. However, these methods merge candidates into a single soft token before future computation\. LBR instead preserves candidates as discrete forwarded branches, routes among their hidden states, and assigns explicit probabilities to both tree growth and routing decisions\.

![Refer to caption](https://arxiv.org/html/2606.25354v1/x1.png)Figure 1:Local Branch Routing decoding pipeline\.LBR maintains a rolling local tree of already\-forwarded candidate continuations\. At each step, the router uses hidden states from all nodes in the current tree to select one depth\-11subtree, commits its root token, prunes the other subtrees, shifts the selected subtree forward, and grows one new layer to restore depthLL\. The top row shows the main experimental settingL=1,K=3L=1,K=3, where routing uses post\-token hidden states ofKKcandidate next tokens\. The bottom row shows the same prune–shift–grow mechanism forL=2,K=3L=2,K=3, where the router also observes hidden states of deeper local continuations\.We evaluate LBR on synthetic hierarchical planning and realistic mathematical reasoning\. The synthetic task shows that post\-candidate hidden states provide branch\-specific information that is not available before branching\. On math benchmarks, LBR improves over discrete chain\-of\-thought, vanilla RLVR, and RL\-compatible soft\-token branching baselines\. These results suggest that local discrete branching offers an efficient and trainable form of language\-model test\-time scaling\.

## 2Preliminaries

Let𝒱\\mathcal\{V\}denote the vocabulary and letπθ\\pi\_\{\\theta\}be an autoregressive language model\. Given a prefixx<tx\_\{<t\}, the model defines a next\-token distributionπθ\(⋅∣x<t\)∈Δ\(𝒱\)\.\\pi\_\{\\theta\}\(\\cdot\\mid x\_\{<t\}\)\\in\\Delta\(\\mathcal\{V\}\)\.LetE:𝒱→ℝdE:\\mathcal\{V\}\\to\\mathbb\{R\}^\{d\}be the token embedding map\. After a tokenxt∈𝒱x\_\{t\}\\in\\mathcal\{V\}is generated, its embeddingE​\(xt\)E\(x\_\{t\}\)is fed back into the model for subsequent computation\.

We denote byht=fθ​\(x<t\)∈ℝDh\_\{t\}=f\_\{\\theta\}\(x\_\{<t\}\)\\in\\mathbb\{R\}^\{D\}the hidden state used to predict the next token at positiontt\. The next\-token logits are computed fromhth\_\{t\}, for example bylogitst=Wout​ht\\mathrm\{logits\}\_\{t\}=W\_\{\\mathrm\{out\}\}h\_\{t\}, andπθ\(⋅∣x<t\)=softmax\(Woutht\)\.\\pi\_\{\\theta\}\(\\cdot\\mid x\_\{<t\}\)=\\mathrm\{softmax\}\(W\_\{\\mathrm\{out\}\}h\_\{t\}\)\.For tasks with verifiable answers, each promptqqhas a ground\-truth answery⋆y^\{\\star\}and a verifierv​\(y,y⋆\)∈\[0,1\]v\(y,y^\{\\star\}\)\\in\[0,1\]that scores a generated responseyy\. In math reasoning, this reward is often binary, e\.g\.,v​\(y,y⋆\)=𝟏​\{Ans​\(y\)=y⋆\}\.v\(y,y^\{\\star\}\)=\\mathbf\{1\}\\\{\\mathrm\{Ans\}\(y\)=y^\{\\star\}\\\}\.We use this verifiable reward later to train LBR with a likelihood\-ratio objective over its tree trajectory\.

## 3Decoding Framework for Local Branch Routing

We introduce Local Branch Routing \(LBR\), a decoding framework that augments autoregressive generation with a rolling local lookahead tree\. As shown in Figure[1](https://arxiv.org/html/2606.25354#S1.F1), each decoding step follows four stages: grow a local tree, route among depth\-1 subtrees, prune the unselected subtrees, and shift\-regrow the selected subtree\. The tree is fully forwarded before routing, so each token decision uses hidden\-state evidence from local future continuations\.

##### Stage 1: grow a forwarded local tree\.

Letx<tx\_\{<t\}denote the committed prefix before decoding positiontt\. Instead of sampling the next token directly fromπθ\(⋅∣x<t\)\\pi\_\{\\theta\}\(\\cdot\\mid x\_\{<t\}\), LBR maintains a depth\-LL, width\-KKlocal tree𝒜t\\mathcal\{A\}\_\{t\}rooted atx<tx\_\{<t\}\. Starting from the root, each active node samplesKKchildren from the filtered model distributionπ~θ\\tilde\{\\pi\}\_\{\\theta\}, using the same temperature, top\-pp, top\-kk, or other filters as ordinary decoding\. ForL=1L=1, this givesKKcandidate next tokens; forL=2L=2, each candidate next token is further expanded intoKKone\-step continuations\.

Crucially, every sampled node is also forwarded through the language model\. For each non\-root nodev∈𝒜tv\\in\\mathcal\{A\}\_\{t\}, letpath​\(v\)\\mathrm\{path\}\(v\)be the token sequence from the local\-tree root tovv\. LBR stores

ht​\(v\)=fθ​\(x<t,path​\(v\)\)\.h\_\{t\}\(v\)=f\_\{\\theta\}\(x\_\{<t\},\\mathrm\{path\}\(v\)\)\.Thus𝒜t\\mathcal\{A\}\_\{t\}is not merely a set of candidate token strings; it is a forwarded tree of local hidden states\.

![Refer to caption](https://arxiv.org/html/2606.25354v1/x2.png)Figure 2:Set\-attention router\.The router first encodes each depth\-11candidate subtree independently into a vectorgt,kg\_\{t,k\}\. For the mainL=1L=1setting, a candidate subtree consists of a single forwarded token, sogt,kg\_\{t,k\}is computed from its post\-token hidden state\. ForL=2L=2, the subtree encoder summarizes the hidden states of the candidate root and its local continuations\. The resulting candidate vectors are then passed through a cross\-subtree set\-attention block, where candidates attend to their sibling alternatives to produce context\-aware representationsg~t,k\\tilde\{g\}\_\{t,k\}\. A shared scoring head maps these representations to scoresst,ks\_\{t,k\}, and a softmax gives the routing distribution over candidate subtrees\.
##### Stage 2: route among depth\-1 subtrees\.

Letut,1,…,ut,Ku\_\{t,1\},\\ldots,u\_\{t,K\}be the depth\-1 children of the root\. Eachut,ku\_\{t,k\}defines one candidate subtree, whose root tokentok​\(ut,k\)\\mathrm\{tok\}\(u\_\{t,k\}\)is a possible next committed token\. The router observes the forwarded tree and sampleskt⋆∼ρϕ\(⋅∣x<t,𝒜t;θ\),k\_\{t\}^\{\\star\}\\sim\\rho\_\{\\phi\}\(\\cdot\\mid x\_\{<t\},\\mathcal\{A\}\_\{t\};\\theta\),then commitsxt=tok​\(ut,kt⋆\)\.x\_\{t\}=\\mathrm\{tok\}\(u\_\{t,k\_\{t\}^\{\\star\}\}\)\.The dependence onθ\\thetaindicates that the router input contains hidden states produced by the base language model\.

We instantiateρϕ\\rho\_\{\\phi\}as the set\-attention router in Figure[2](https://arxiv.org/html/2606.25354#S3.F2)\. Each depth\-1 candidate subtree is first encoded into a vectorgt,kg\_\{t,k\}\. WhenL=1L=1, this is computed from the post\-token hidden stateht​\(ut,k\)h\_\{t\}\(u\_\{t,k\}\); whenL\>1L\>1, the subtree encoder summarizes the candidate root together with its local continuations\. The candidate vectors are then compared by cross\-subtree set attention and scored by a shared head:

g~t,1,…,g~t,K=SetAttnϕcand​\(gt,1,…,gt,K\),st,k=wϕ⊤​g~t,k,\\tilde\{g\}\_\{t,1\},\\ldots,\\tilde\{g\}\_\{t,K\}=\\mathrm\{SetAttn\}^\{\\mathrm\{cand\}\}\_\{\\phi\}\(g\_\{t,1\},\\ldots,g\_\{t,K\}\),\\qquad s\_\{t,k\}=w\_\{\\phi\}^\{\\top\}\\tilde\{g\}\_\{t,k\},ρϕ​\(k∣x<t,𝒜t;θ\)=exp⁡\(st,k/τ\)∑k′=1Kexp⁡\(st,k′/τ\)\.\\rho\_\{\\phi\}\(k\\mid x\_\{<t\},\\mathcal\{A\}\_\{t\};\\theta\)=\\frac\{\\exp\(s\_\{t,k\}/\\tau\)\}\{\\sum\_\{k^\{\\prime\}=1\}^\{K\}\\exp\(s\_\{t,k^\{\\prime\}\}/\\tau\)\}\.This makes routing a relative decision over sibling local futures, rather than an independent reranking of next\-token logits\. Additional architectural details are provided in Appendix[A\.4](https://arxiv.org/html/2606.25354#A1.SS4)\.

##### Stage 3: prune unselected subtrees\.

After the router selectskt⋆k\_\{t\}^\{\\star\}, LBR discards all depth\-1 subtrees except the selected one\. This is the merge step: multiple local futures are explored, but only one subtree remains active\. The unselected branches no longer participate in future decoding, while the selected root token becomes part of the committed prefix\.

##### Stage 4: shift and regrow\.

The selected subtree is shifted forward: its root token has been committed, and its remaining descendants become the partial lookahead tree for the next decoding position\. LBR then grows one new layer from the surviving frontier to restore depthLL\.

We denote by𝒢t\\mathcal\{G\}\_\{t\}the set of nodes newly grown and forwarded at steptt\. Forv∈𝒢tv\\in\\mathcal\{G\}\_\{t\}, lettok​\(v\)\\mathrm\{tok\}\(v\)be its token label and define the tree\-causal sampling context

ctx​\(v\)=\(x<t,path​\(pa​\(v\)\)\),\\mathrm\{ctx\}\(v\)=\\bigl\(x\_\{<t\},\\mathrm\{path\}\(\\mathrm\{pa\}\(v\)\)\\bigr\),wherepa​\(v\)\\mathrm\{pa\}\(v\)is the parent ofvv\. Thentok\(v\)∼π~θ\(⋅∣ctx\(v\)\)\.\\mathrm\{tok\}\(v\)\\sim\\tilde\{\\pi\}\_\{\\theta\}\(\\cdot\\mid\\mathrm\{ctx\}\(v\)\)\.

After an initial warm\-up that builds the first depth\-LLtree, decoding repeats the same rolling loop: grow, route, prune, shift, and regrow\.

##### Trace likelihood\.

This four\-stage procedure defines a tractable likelihood for the full LBR decoding trace\. At each step, only two operations are stochastic: sampling newly grown nodes𝒢t\\mathcal\{G\}\_\{t\}and sampling the router choicekt⋆k\_\{t\}^\{\\star\}\. The commit, prune, shift, and reuse operations are deterministic given these samples\. Letℱ=\{\(𝒢t,𝒜t,kt⋆\)\}t=1T\.\\mathcal\{F\}=\\\{\(\\mathcal\{G\}\_\{t\},\\mathcal\{A\}\_\{t\},k\_\{t\}^\{\\star\}\)\\\}\_\{t=1\}^\{T\}\.Then

log⁡pθ,ϕ​\(ℱ∣q\)=∑t=1T∑v∈𝒢tlog⁡π~θ​\(tok​\(v\)∣ctx​\(v\)\)\+∑t=1Tlog⁡ρϕ​\(kt⋆∣x<t,𝒜t;θ\)\.\\log p\_\{\\theta,\\phi\}\(\\mathcal\{F\}\\mid q\)=\\sum\_\{t=1\}^\{T\}\\sum\_\{v\\in\\mathcal\{G\}\_\{t\}\}\\log\\tilde\{\\pi\}\_\{\\theta\}\\left\(\\mathrm\{tok\}\(v\)\\mid\\mathrm\{ctx\}\(v\)\\right\)\+\\sum\_\{t=1\}^\{T\}\\log\\rho\_\{\\phi\}\\left\(k\_\{t\}^\{\\star\}\\mid x\_\{<t\},\\mathcal\{A\}\_\{t\};\\theta\\right\)\.\(1\)The first term accounts for growing the local trees, and the second accounts for routing decisions\. The sum is over𝒢t\\mathcal\{G\}\_\{t\}, not all of𝒜t\\mathcal\{A\}\_\{t\}: a reused node may influence later routers through its hidden state, but its language\-model log\-probability is counted only when it is first grown\.

## 4Interpreting Local Branch Routing on Synthetic Hierarchical Planning

We hypothesize that LBR improves reasoning by exploiting predictive hidden states at branching tokens to perform local planning\. To illustrate this mechanism, we study a synthetic hierarchical planning task based on radix\-translated graph reachability\. The experiment shows that LBR has an advantage over discrete CoT and soft thinking because it effectively utilizes the informative post\-branching hidden states by routing\. Our synthetic task is modified from the ProsQA benchmark introduced byHaoet al\.\[[11](https://arxiv.org/html/2606.25354#bib.bib14)\]and further analyzed byZhuet al\.\[[35](https://arxiv.org/html/2606.25354#bib.bib37)\]\.

![Refer to caption](https://arxiv.org/html/2606.25354v1/x3.png)Figure 3:Radix\-translated graph reachability and decoding behavior\. Left: a concept\-level reachability problem is translated into a radix\-2 token sequence, so each graph node becomes a short digit string and each graph transition becomes a sequence of token\-level decisions\. Shared digit prefixes create hierarchical branching positions\. Right: the three decoding methods handle these branching positions differently\. Discrete CoT commits from the pre\-branching hidden state, LBR forwards candidate digits as separate branches and routes using their post\-candidate hidden states, and Soft Thinking mixes candidate embeddings into a single continuous state, which can blur concept identity\.### 4\.1Radix\-Translated Reachability as Hierarchical Planning

The left panel of Figure[3](https://arxiv.org/html/2606.25354#S4.F3)illustrates the synthetic planning task\. At the concept level, each example is a directed graph reachability problem: the model is given a serialized graph, two candidate targets, and a root node, and must generate a path from the root to the reachable target\. In the example, the root is33, the candidate targets are1414and1010, and only1414is reachable\. The legal next concepts from33are11,9,211,9,2; a local search procedure must identify that the branch through99reaches the target and generate the path3→9→143\\to 9\\to 14\.

We translate this concept\-level problem into a radix\-tokenized sequence by representing each node id as a fixed\-length binary string\. Forn=16n=16nodes, each concept usesW=4W=4digit tokens, e\.g\.,3=00113=0011,9=10019=1001,11=101111=1011, and2=00102=0010\. Thus a concept\-level transition such as3→93\\to 9becomes a sequence of token\-level decisions0011→10010011\\to 1001\.

This translation creates hierarchical branching at the token level\. The competing next concepts11,9,211,9,2share and diverge across digit prefixes: the first digit separates the0∗∗∗0\*\*\*group from the1∗∗∗1\*\*\*group; after choosing prefix11, the next digit is temporarily forced to0; a later digit then distinguishes10011001from10111011\. Therefore, concept\-level planning is converted into a sequence of local token decisions with both merges and branch points, making the task a controlled testbed for whether post\-candidate hidden states help route among locally ambiguous branches\.

### 4\.2Training and Decoding Methods

The right panel of Figure[3](https://arxiv.org/html/2606.25354#S4.F3)compares three decoding methods on the same radix\-translated transition\.Discrete CoTtreats the path as a standard autoregressive sequence and commits to each digit from the pre\-branching hidden state\.Local Branch Routingkeeps the same discrete token space, but forwards each legal candidate digit as a separate branch and uses a router to select among the resulting post\-candidate hidden states\.Soft Thinkingexposes the same legal candidates, but merges them into a continuous mixture embedding before future computation, thereby blurring the identity of the selected concept branch\.

All methods use the same base model and data distribution\. We first apply MTP pretraining to teach the model legal graph transitions by predicting future node digits along random walks\. We then post\-train each method with its corresponding trace representation: Discrete CoT uses the gold radix path with next\-token cross\-entropy; LBR forwards the legal candidate branches and supervises the router to select the gold branch; Soft Thinking trains on branch\-and\-merge traces where legal candidate embeddings are mixed before predicting future tokens\. Thus, the comparison controls for graph knowledge and isolates how each method represents branching decisions\.

![Refer to caption](https://arxiv.org/html/2606.25354v1/figs/radix_reachability_learning_curve.png)\(a\)Learning dynamics\.
![Refer to caption](https://arxiv.org/html/2606.25354v1/figs/branching_probe.png)\(b\)Branching\-state probe\.

Figure 4:Synthetic hierarchical\-planning results\. Left: LBR achieves the highest target accuracy on radix\-translated reachability, while its router accuracy increases and branching\-position LM KL decreases during training\. Right: probes show that LBR’s post\-correct\-candidate hidden states are more predictive of the final target than pre\-branching states, supporting routing after forwarding candidate branches\.
### 4\.3Why LBR Helps: Informative Branching Hidden States

The upper panel of Figure[4](https://arxiv.org/html/2606.25354#S4.F4)\(a\) compares the target accuracy of the three methods\. LBR achieves the best performance, followed by discrete CoT and then Soft Thinking:

Soft Thinking​\(0\.845\)<Discrete CoT​\(0\.914\)<LBR​\(0\.967\)\.\\text\{Soft Thinking \}\(0\.845\)<\\text\{Discrete CoT \}\(0\.914\)<\\text\{LBR \}\(0\.967\)\.We next ask why routing after branching gives this advantage\.

The key difference is the hidden state from which a branching decision is made\. Discrete CoT must choose the next digit from the pre\-branching hidden state, before any candidate digit is realized\. LBR instead first forwards each legal candidate digit and then routes using the resulting post\-candidate hidden states\. Figure[4](https://arxiv.org/html/2606.25354#S4.F4)\(b\) tests whether these post\-candidate states contain more information about the final target\. The probe predicts the target node id from hidden states at diverging positions\. For LBR, the post\-correct\-candidate state is much more predictive than the pre\-branching state\. Thus, the candidate token itself reveals downstream information that is not available before branching\.

This also explains the gap to discrete CoT\. Although discrete CoT learns to encode some decision information in its pre\-branching state, this state remains weaker than LBR’s post\-correct\-candidate state, as highlighted by the shaded band in Figure[4](https://arxiv.org/html/2606.25354#S4.F4)\(b\)\. Moreover, there is a large separation between correct and wrong post\-candidate states\. Therefore, the router is not merely reranking equivalent candidates; it can exploit candidate\-specific hidden states that distinguish branches leading toward the target from branches leading away from it\.

The lower panel of Figure[4](https://arxiv.org/html/2606.25354#S4.F4)\(a\) confirms that LBR learns to use this signal during training\. Router accuracy on oracle branching events increases, while the branching\-position LM KL decreases\. This indicates that the router becomes better at selecting the correct forwarded branch, and the base model simultaneously learns a more calibrated local branching distribution\. Appendix[B](https://arxiv.org/html/2606.25354#A2)provides an additional concept\-identity probe showing that soft\-token mixtures blur the identity of the generated graph node, whereas LBR preserves discrete candidate identities\.

## 5Local Branch Routing Improves Mathematical Reasoning

##### Experimental setup\.

We evaluate LBR on DeepSeek\-R1\-Distill\-Qwen\-1\.5B and DeepSeek\-R1\-Distill\-Qwen\-7B\. All RL\-trained methods use GRPO with verifier rewards\. For LBR, we use the same RLVR setup as in Section 3, replacing the standard discrete\-token log\-probability with the tree\-trace likelihood in Eq\.[1](https://arxiv.org/html/2606.25354#S3.E1):ℒLBR​\(θ,ϕ;ℱ\)=−A​\(ℱ\)​log⁡pθ,ϕ​\(ℱ∣q\)\.\\mathcal\{L\}\_\{\\mathrm\{LBR\}\}\(\\theta,\\phi;\\mathcal\{F\}\)=\-A\(\\mathcal\{F\}\)\\log p\_\{\\theta,\\phi\}\(\\mathcal\{F\}\\mid q\)\.

This assigns reward\-weighted credit to both the newly grown local tree tokens and the router decisions\. WhenK=1K=1, the router is deterministic and the objective reduces to standard discrete\-token RLVR\.

##### Implementation\.

We jointly train the base model and router\. The base\-model learning rate is1×10−61\\times 10^\{\-6\}and the router learning rate is1×10−41\\times 10^\{\-4\}\. The router samples candidate subtrees with temperature1\.01\.0\. We train for 300 steps with a global batch size of 128 questions and 8 rollout samples per question\. Training uses temperature1\.01\.0and top\-p=1\.0p=1\.0; evaluation uses top\-p=0\.95p=0\.95and reports pass@1/32 averaged over 64 sampled runs\. All training and evaluation use a maximum response length of 4096 tokens\. The training set is DeepScaleR\-Preview\-Dataset\[[20](https://arxiv.org/html/2606.25354#bib.bib13)\]\. All the trainings are done in 8×\\timesH100 GPUs\. Router architecture hyperparameters are provided in Appendix[A\.4\.1](https://arxiv.org/html/2606.25354#A1.SS4.SSS1)\.

##### Baselines\.

We compare with discrete chain\-of\-thought decoding, vanilla GRPO/RLVR with single\-thread discrete tokens, and RLVR\-compatible soft\-token branching methods such as Multiplex Thinking\. We report the mainL=1,K=3L=1,K=3configuration and ablate deeper local lookahead such asL=2,K=3L=2,K=3\.

For vanilla discrete\-token RLVR, a sampled responsey=\(y1,…,yT\)y=\(y\_\{1\},\\ldots,y\_\{T\}\)is trained with the standard likelihood\-ratio lossℒdisc​\(θ\)=−A​\(q,y\)​∑t=1Tlog⁡πθ​\(yt∣q,y<t\),\\mathcal\{L\}\_\{\\mathrm\{disc\}\}\(\\theta\)=\-A\(q,y\)\\sum\_\{t=1\}^\{T\}\\log\\pi\_\{\\theta\}\(y\_\{t\}\\mid q,y\_\{<t\}\),whereA​\(q,y\)A\(q,y\)is the advantage computed from verifier rewards\.

For RL\-compatible soft\-token branching, following Multiplex Thinking, each step samplesKKcandidate tokensct,1,…,ct,Kc\_\{t,1\},\\ldots,c\_\{t,K\}and merges their embeddings into a soft tokene~t\\tilde\{e\}\_\{t\}\. Its sampled candidates define the training likelihoodℒsoft​\(θ\)=−A​\(q,y~\)​∑t=1T∑k=1Klog⁡πθ​\(ct,k∣q,y~<t\)\.\\mathcal\{L\}\_\{\\mathrm\{soft\}\}\(\\theta\)=\-A\(q,\\tilde\{y\}\)\\sum\_\{t=1\}^\{T\}\\sum\_\{k=1\}^\{K\}\\log\\pi\_\{\\theta\}\(c\_\{t,k\}\\mid q,\\tilde\{y\}\_\{<t\}\)\.

### 5\.1Math Benchmark Results

Table[1](https://arxiv.org/html/2606.25354#S5.T1)reports Pass@1/Pass@32 on six mathematical reasoning benchmarks\. Across both 1\.5B and 7B backbones, LBR improves over discrete CoT, vanilla discrete\-token RLVR, and RLVR with soft\-token branching\. This comparison isolates the decoding framework: the RL\-trained methods use the same verifier rewards, but LBR replaces single\-thread or soft\-merged decoding with routing over forwarded local branches\. The gains over soft\-token branching show that preserving candidate branches and routing among their hidden states is more effective than merging candidate embeddings into a single continuous token\.

LBR improves both Pass@1 and Pass@32, suggesting better single\-sample accuracy as well as better exploration under repeated sampling\. Increasing the local lookahead fromL=1L=1toL=2L=2often further improves performance, especially on the 7B backbone, consistent with the view that deeper local continuations provide more useful routing evidence\.

Table 1:Main mathematical reasoning results\. We report Pass@1/Pass@32 on six benchmarks for DeepSeek\-R1\-Distill\-Qwen\-1\.5B and 7B\. LBR improves over discrete CoT, vanilla discrete\-token RLVR, and RLVR\-compatible soft\-token branching baselines\.
### 5\.2Cross\-Subtree Attention Improves Routing

Figure[6](https://arxiv.org/html/2606.25354#S5.F6)compares the full set\-attention router with an independent router\. The independent variant removes cross\-subtree attention and scores each candidate subtree with a shared MLP applied to its own embedding only\. This keeps the same LBR decoding framework, tree width, training setup, and tree\-trace objective, but removes the contrastive comparison among sibling candidates\.

The independent router consistently underperforms the full router\. This gap shows that LBR’s gains are not explained only by exposing the router to post\-token hidden states; the way those hidden states are compared also matters\. In particular, cross\-subtree attention lets each candidate be evaluated relative to the other sampled local futures from the same prefix\. The ablation therefore supports the claim that routing is a contrastive decision over local branches, rather than independent scoring of candidate tokens\.

![Refer to caption](https://arxiv.org/html/2606.25354v1/figs/router_ablation.png)Figure 5:Router ablation\. The full contrastive router outperforms the independent router, showing that comparing sibling branches improves routing\.
![Refer to caption](https://arxiv.org/html/2606.25354v1/figs/learning_dym.png)Figure 6:Training dynamics of LBR during RLVR\. Average Pass@1 and Pass@32 improve early and remain stable, while router entropy and base\-LM entropy at routed positions decrease\.

### 5\.3Training Dynamics of Local Branch Routing

Figure[6](https://arxiv.org/html/2606.25354#S5.F6)shows the training dynamics of LBR during RLVR on the math\-reasoning benchmarks\.

The top panel reports average Pass@1 and Pass@32 across the evaluated tasks\. The main improvement happens early in training: average Pass@1 rises from roughly39\.539\.5to above4545, while average Pass@32 rises from roughly65\.565\.5to around6969\. After this early phase, both metrics remain relatively stable, suggesting that LBR quickly learns a useful local routing policy and then maintains its gains without severe degradation\.

The entropy curves provide a complementary view of what is learned\. Router entropy decreases throughout training, indicating that the router becomes more decisive in selecting among local branches\. At the same time, the base\-LM entropy at routed positions also decreases, suggesting that the language model itself becomes more confident about local branch tokens\. Thus, LBR training does not only train an external selector on top of a fixed distribution; it co\-adapts the base model and router so that local candidate branches become easier to distinguish and route\.

Importantly, router entropy decreases together with improved Pass@1 and Pass@32, rather than through a collapse that hurts exploration\. This supports the interpretation that LBR learns a discriminative routing policy over forwarded local branches\. The result also helps explain why LBR improves both single\-sample accuracy and Pass@32: the router becomes sharper, while the underlying sampling process still maintains enough diversity across repeated runs\.

## 6Related Works

##### Test\-time scaling and search\.

Test\-time scaling improves language\-model reasoning by allocating additional inference\-time compute\. Chain\-of\-thought prompting and self\-consistency scale reasoning through longer traces or multiple sampled traces, while recent work studies budget forcing, verifier\-guided search, and adaptive sampling\[[28](https://arxiv.org/html/2606.25354#bib.bib15),[27](https://arxiv.org/html/2606.25354#bib.bib17),[24](https://arxiv.org/html/2606.25354#bib.bib16),[22](https://arxiv.org/html/2606.25354#bib.bib31)\]\. Search\-based methods introduce explicit width by exploring trees over thoughts, actions, or solution refinements, including Tree of Thoughts, RAP, LATS, policy\-guided tree search, and adaptive\-branching MCTS\[[32](https://arxiv.org/html/2606.25354#bib.bib29),[10](https://arxiv.org/html/2606.25354#bib.bib19),[34](https://arxiv.org/html/2606.25354#bib.bib33),[17](https://arxiv.org/html/2606.25354#bib.bib32),[14](https://arxiv.org/html/2606.25354#bib.bib30)\]\. These methods demonstrate the value of branching during inference, but typically operate over coarse reasoning units and often rely on external evaluators, handcrafted search control, or repeated full\-step generation\. LBR instead branches over token\-level futures, forwards the resulting hidden states, and routes to one depth\-1 subtree before continuing autoregressive decoding\.

##### Local lookahead and speculative decoding\.

Local future tokens have also been used to accelerate autoregressive decoding\. Speculative decoding and speculative sampling draft multiple tokens with a smaller model and verify them with a target model, while tree\-based speculative inference and lookahead decoding exploit local token trees or parallel n\-gram guesses to reduce latency\[[16](https://arxiv.org/html/2606.25354#bib.bib28),[5](https://arxiv.org/html/2606.25354#bib.bib23),[21](https://arxiv.org/html/2606.25354#bib.bib22),[6](https://arxiv.org/html/2606.25354#bib.bib21)\]\. Recent reasoning\-oriented variants use reward models, draft reasoning chains, or semantic verification to accelerate long chain\-of\-thought generation\[[19](https://arxiv.org/html/2606.25354#bib.bib27),[26](https://arxiv.org/html/2606.25354#bib.bib20),[7](https://arxiv.org/html/2606.25354#bib.bib26)\]\. These methods show that local futures can be generated and checked efficiently, but mainly treat lookahead as acceleration\. LBR uses lookahead as decision evidence: sampled branches are forwarded by the same language model, compared by a router, and pruned according to a trainable routing distribution\.

##### Reinforcement learning with tree\-structured exploration\.

RL with verifiable rewards has become a standard recipe for mathematical reasoning, with GRPO and DeepSeek\-R1 showing that reward\-weighted sampled trajectories can elicit strong reasoning behavior\[[23](https://arxiv.org/html/2606.25354#bib.bib38),[9](https://arxiv.org/html/2606.25354#bib.bib25)\]\. A related line introduces tree structure into RL training: MCTS\-DPO collects step\-level preference pairs with Monte Carlo tree search and updates the policy with DPO\[[31](https://arxiv.org/html/2606.25354#bib.bib24)\]; TreeRL uses entropy\-guided tree search to fork from uncertain intermediate tokens and derive process\-level supervision\[[13](https://arxiv.org/html/2606.25354#bib.bib36)\]; DeepSearch embeds MCTS into RLVR training for exploration and credit assignment\[[29](https://arxiv.org/html/2606.25354#bib.bib34)\]; and TreePO reformulates rollouts as segment\-level tree sampling with tree\-aware advantages\[[18](https://arxiv.org/html/2606.25354#bib.bib35)\]\. These methods use tree search primarily to produce better training traces, preferences, or advantage estimates over coarse reasoning units\. LBR differs in that the tree is itself the learned test\-time decoding policy: the same local tree serves as the inference\-time scaling mechanism and the object optimized by verifier\-based RL\.

##### Continuous and soft\-token branch\-and\-merge reasoning\.

Another close line of work seeks token\-level width by replacing a discrete next token with a continuous or soft representation\. Soft Thinking forms probability\-weighted mixtures of token embeddings, and follow\-up analyses study when such mixtures behave like genuine parallel reasoning versus a dominant single path\[[33](https://arxiv.org/html/2606.25354#bib.bib1),[30](https://arxiv.org/html/2606.25354#bib.bib10)\]\. Soft\-token RL methods add stochasticity and train continuous chain\-of\-thought processes with reinforcement learning\[[4](https://arxiv.org/html/2606.25354#bib.bib11)\]\. Multiplex Thinking is especially close: it samples multiple candidate tokens and merges their embeddings into one continuous token with a tractable rollout probability\[[25](https://arxiv.org/html/2606.25354#bib.bib12)\]\. These methods occupy the same token\-level branch\-and\-merge regime as LBR, but differ in representation: soft\-token methods compress alternatives into a shared continuous state, whereas LBR preserves candidates as discrete forwarded branches and routes among their hidden states\.

##### Interpreting reasoning representations\.

Hao et al\.\[[11](https://arxiv.org/html/2606.25354#bib.bib14)\]empirically observe that continuous thoughts can retain multiple candidate next steps and exhibit a breadth\-first\-search\-like reasoning pattern\. Building on this observation, subsequent analyses explain the advantage of continuous reasoning primarily through constructive expressivity results\. Zhu et al\.\[[35](https://arxiv.org/html/2606.25354#bib.bib37)\]construct a two\-layer transformer whose continuous thoughts superpose multiple search frontiers to solve directed graph reachability, while Gozeten et al\.\[[8](https://arxiv.org/html/2606.25354#bib.bib2)\]construct a one\-layer continuous\-CoT transformer that aggregates multiple partial trajectories to solve a subset\-sum\-type combinatorial problem\. These works highlight how a shared continuous representation can support parallel exploration\. We study a complementary question: whether different reasoning frameworks preserve task\-relevant branch information clearly enough for downstream decisions\. Linear probing fits a lightweight classifier on frozen hidden states to test whether a target property is linearly decodable\[[1](https://arxiv.org/html/2606.25354#bib.bib9),[12](https://arxiv.org/html/2606.25354#bib.bib8),[3](https://arxiv.org/html/2606.25354#bib.bib7),[2](https://arxiv.org/html/2606.25354#bib.bib6),[15](https://arxiv.org/html/2606.25354#bib.bib3)\]\. Using probes as a diagnostic, we provide a representation\-level explanation for the performance differences among discrete CoT, soft\-token reasoning, and LBR: forwarding candidate branches makes the correct target more recoverable in LBR’s post\-candidate states, whereas merging candidates into soft states blurs concept identity\.

## 7Conclusion

We presented Local Branch Routing \(LBR\), a trainable test\-time scaling framework that maintains a rolling local lookahead tree during autoregressive decoding\. By forwarding candidate branches before routing, LBR lets each token decision use post\-candidate hidden states rather than only the root next\-token distribution\. The resulting prune–shift–grow process preserves discrete branch identities, reuses surviving local futures, and defines a tractable tree\-trajectory likelihood for RLVR training\. Synthetic hierarchical\-planning experiments show that these forwarded branch states provide useful decision evidence and avoid the concept ambiguity of soft\-token mixtures\. Math\-reasoning results further show that LBR improves over discrete CoT, vanilla discrete\-token RLVR, and soft\-token branching baselines\.

The main limitation is locality\. LBR uses shallow branch information, so its predictability is limited to evidence exposed within the local lookahead window\. It does not realize full long\-range planning or global search over complete solutions\. Larger lookahead depth can strengthen the routing signal, but at higher forwarding cost, creating a trade\-off between speed and effect\. Future work may combine local branch routing with adaptive depth, selective global search, or learned triggers for when longer planning is needed\.

## References

- \[1\]\(2016\)Understanding intermediate layers using linear classifier probes\.arXiv preprint arXiv:1610\.01644\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px5.p1.1)\.
- \[2\]A\. Azaria and T\. Mitchell\(2023\-12\)The internal state of an LLM knows when it’s lying\.InFindings of the Association for Computational Linguistics: EMNLP 2023,H\. Bouamor, J\. Pino, and K\. Bali \(Eds\.\),Singapore,pp\. 967–976\.External Links:[Link](https://aclanthology.org/2023.findings-emnlp.68/),[Document](https://dx.doi.org/10.18653/v1/2023.findings-emnlp.68)Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px5.p1.1)\.
- \[3\]Y\. Belinkov\(2022\-03\)Probing classifiers: promises, shortcomings, and advances\.Computational Linguistics48\(1\),pp\. 207–219\.External Links:[Link](https://aclanthology.org/2022.cl-1.7/),[Document](https://dx.doi.org/10.1162/coli%5Fa%5F00422)Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px5.p1.1)\.
- \[4\]N\. Butt, A\. Kwiatkowski, I\. Labiad, J\. Kempe, and Y\. Ollivier\(2025\)Soft tokens, hard truths\.arXiv preprint arXiv:2509\.19170\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px4.p1.1)\.
- \[5\]C\. Chen, S\. Borgeaud, G\. Irving, J\. Lespiau, L\. Sifre, and J\. Jumper\(2023\)Accelerating large language model decoding with speculative sampling\.arXiv preprint arXiv:2302\.01318\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px2.p1.1)\.
- \[6\]Y\. Fu, P\. Bailis, I\. Stoica, and H\. Zhang\(2024\)Break the sequential dependency of llm inference using lookahead decoding\.arXiv preprint arXiv:2402\.02057\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px2.p1.1)\.
- \[7\]Y\. Fu, R\. Ge, Z\. Shao, Z\. Deng, and H\. Zhang\(2025\)Scaling speculative decoding with lookahead reasoning\.arXiv preprint arXiv:2506\.19830\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px2.p1.1)\.
- \[8\]H\. A\. Gozeten, M\. E\. Ildiz, X\. Zhang, H\. Harutyunyan, A\. S\. Rawat, and S\. Oymak\(2025\)Continuous chain of thought enables parallel exploration and reasoning\.arXiv preprint arXiv:2505\.23648\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px5.p1.1)\.
- \[9\]D\. Guo, D\. Yang, H\. Zhang, J\. Song, P\. Wang, Q\. Zhu, R\. Xu, R\. Zhang, S\. Ma, X\. Bi,et al\.\(2025\)Deepseek\-r1: incentivizing reasoning capability in llms via reinforcement learning\.arXiv preprint arXiv:2501\.12948\.Cited by:[§1](https://arxiv.org/html/2606.25354#S1.p4.1),[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px3.p1.1)\.
- \[10\]S\. Hao, Y\. Gu, H\. Ma, J\. Hong, Z\. Wang, D\. Wang, and Z\. Hu\(2023\)Reasoning with language model is planning with world model\.InProceedings of the 2023 Conference on Empirical Methods in Natural Language Processing,pp\. 8154–8173\.Cited by:[§1](https://arxiv.org/html/2606.25354#S1.p1.1),[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px1.p1.1)\.
- \[11\]S\. Hao, S\. Sukhbaatar, D\. Su, X\. Li, Z\. Hu, J\. Weston, and Y\. Tian\(2024\)Training large language models to reason in a continuous latent space\.arXiv preprint arXiv:2412\.06769\.Cited by:[§4](https://arxiv.org/html/2606.25354#S4.p1.1),[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px5.p1.1)\.
- \[12\]J\. Hewitt and P\. Liang\(2019\-11\)Designing and interpreting probes with control tasks\.InProceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing \(EMNLP\-IJCNLP\),K\. Inui, J\. Jiang, V\. Ng, and X\. Wan \(Eds\.\),Hong Kong, China,pp\. 2733–2743\.External Links:[Link](https://aclanthology.org/D19-1275/),[Document](https://dx.doi.org/10.18653/v1/D19-1275)Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px5.p1.1)\.
- \[13\]Z\. Hou, Z\. Hu, Y\. Li, R\. Lu, J\. Tang, and Y\. Dong\(2025\)Treerl: llm reinforcement learning with on\-policy tree search\.InProceedings of the 63rd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),pp\. 12355–12369\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px3.p1.1)\.
- \[14\]Y\. Inoue, K\. Misaki, Y\. Imajuku, S\. Kuroki, T\. Nakamura, and T\. Akiba\(2025\)Wider or deeper? scaling llm inference\-time compute with adaptive branching tree search\.arXiv preprint arXiv:2503\.04412\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px1.p1.1)\.
- \[15\]M\. Jin, Q\. Yu, J\. Huang, Q\. Zeng, Z\. Wang, W\. Hua, H\. Zhao, K\. Mei, Y\. Meng, K\. Ding, F\. Yang, M\. Du, and Y\. Zhang\(2025\-01\)Exploring concept depth: how large language models acquire knowledge and concept at different layers?\.InProceedings of the 31st International Conference on Computational Linguistics,O\. Rambow, L\. Wanner, M\. Apidianaki, H\. Al\-Khalifa, B\. D\. Eugenio, and S\. Schockaert \(Eds\.\),Abu Dhabi, UAE,pp\. 558–573\.External Links:[Link](https://aclanthology.org/2025.coling-main.37/)Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px5.p1.1)\.
- \[16\]Y\. Leviathan, M\. Kalman, and Y\. Matias\(2023\)Fast inference from transformers via speculative decoding\.InInternational Conference on Machine Learning,pp\. 19274–19286\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px2.p1.1)\.
- \[17\]Y\. Li\(2025\)Policy guided tree search for enhanced llm reasoning\.arXiv preprint arXiv:2502\.06813\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px1.p1.1)\.
- \[18\]Y\. Li, Q\. Gu, Z\. Wen, Z\. Li, T\. Xing, S\. Guo, T\. Zheng, X\. Zhou, X\. Qu, W\. Zhou,et al\.\(2025\)Treepo: bridging the gap of policy optimization and efficacy and inference efficiency with heuristic tree\-based modeling\.arXiv preprint arXiv:2508\.17445\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px3.p1.1)\.
- \[19\]B\. Liao, Y\. Xu, H\. Dong, J\. Li, C\. Monz, S\. Savarese, D\. Sahoo, and C\. Xiong\(2025\)Reward\-guided speculative decoding for efficient llm reasoning\.arXiv preprint arXiv:2501\.19324\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px2.p1.1)\.
- \[20\]M\. Luo, S\. Tan, J\. Wong, X\. Shi, W\. Y\. Tang, M\. Roongta, C\. Cai, J\. Luo, L\. E\. Li, R\. A\. Popa, and I\. Stoica\(2025\)DeepScaleR: surpassing o1\-preview with a 1\.5b model by scaling rl\.Note:[https://pretty\-radio\-b75\.notion\.site/DeepScaleR\-Surpassing\-O1\-Preview\-with\-a\-1\-5B\-Model\-by\-Scaling\-RL\-19681902c1468005bed8ca303013a4e2](https://pretty-radio-b75.notion.site/DeepScaleR-Surpassing-O1-Preview-with-a-1-5B-Model-by-Scaling-RL-19681902c1468005bed8ca303013a4e2)Notion BlogCited by:[§5](https://arxiv.org/html/2606.25354#S5.SS0.SSS0.Px2.p1.7)\.
- \[21\]X\. Miao, G\. Oliaro, Z\. Zhang, X\. Cheng, Z\. Wang, Z\. Zhang, R\. Y\. Y\. Wong, A\. Zhu, L\. Yang, X\. Shi,et al\.\(2023\)Specinfer: accelerating generative large language model serving with tree\-based speculative inference and verification\.arXiv preprint arXiv:2305\.09781\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px2.p1.1)\.
- \[22\]N\. Muennighoff, Z\. Yang, W\. Shi, X\. L\. Li, L\. Fei\-Fei, H\. Hajishirzi, L\. Zettlemoyer, P\. Liang, E\. Candès, and T\. B\. Hashimoto\(2025\)S1: simple test\-time scaling\.InProceedings of the 2025 Conference on Empirical Methods in Natural Language Processing,pp\. 20286–20332\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px1.p1.1)\.
- \[23\]Z\. Shao, P\. Wang, Q\. Zhu, R\. Xu, J\. Song, X\. Bi, H\. Zhang, M\. Zhang, Y\. Li, Y\. Wu,et al\.\(2024\)Deepseekmath: pushing the limits of mathematical reasoning in open language models\.arXiv preprint arXiv:2402\.03300\.Cited by:[§1](https://arxiv.org/html/2606.25354#S1.p4.1),[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px3.p1.1)\.
- \[24\]C\. Snell, J\. Lee, K\. Xu, and A\. Kumar\(2024\)Scaling llm test\-time compute optimally can be more effective than scaling model parameters\.arXiv preprint arXiv:2408\.03314\.Cited by:[§1](https://arxiv.org/html/2606.25354#S1.p1.1),[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px1.p1.1)\.
- \[25\]Y\. Tang, L\. Dong, Y\. Hao, Q\. Dong, F\. Wei, and J\. Gu\(2026\)Multiplex thinking: reasoning via token\-wise branch\-and\-merge\.arXiv preprint arXiv:2601\.08808\.Cited by:[§1](https://arxiv.org/html/2606.25354#S1.p5.1),[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px4.p1.1)\.
- \[26\]J\. Wang, J\. Li, J\. Hou, B\. Yan, L\. Wu, and M\. Zhang\(2025\)Efficient reasoning for llms through speculative chain\-of\-thought\.arXiv preprint arXiv:2504\.19095\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px2.p1.1)\.
- \[27\]X\. Wang, J\. Wei, D\. Schuurmans, Q\. Le, E\. Chi, S\. Narang, A\. Chowdhery, and D\. Zhou\(2022\)Self\-consistency improves chain of thought reasoning in language models\.arXiv preprint arXiv:2203\.11171\.Cited by:[§1](https://arxiv.org/html/2606.25354#S1.p1.1),[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px1.p1.1)\.
- \[28\]J\. Wei, X\. Wang, D\. Schuurmans, M\. Bosma, F\. Xia, E\. Chi, Q\. V\. Le, D\. Zhou,et al\.\(2022\)Chain\-of\-thought prompting elicits reasoning in large language models\.Advances in neural information processing systems35,pp\. 24824–24837\.Cited by:[§1](https://arxiv.org/html/2606.25354#S1.p1.1),[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px1.p1.1)\.
- \[29\]F\. Wu, W\. Xuan, H\. Qi, X\. Lu, A\. Tu, L\. E\. Li, and Y\. Choi\(2025\)DeepSearch: overcome the bottleneck of reinforcement learning with verifiable rewards via monte carlo tree search\.arXiv preprint arXiv:2509\.25454\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px3.p1.1)\.
- \[30\]J\. Wu, J\. Lu, Z\. Ren, G\. Hu, Z\. Wu, D\. Dai, and H\. Wu\(2025\)Llms are single\-threaded reasoners: demystifying the working mechanism of soft thinking\.arXiv preprint arXiv:2508\.03440\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px4.p1.1)\.
- \[31\]Y\. Xie, A\. Goyal, W\. Zheng, M\. Kan, T\. P\. Lillicrap, K\. Kawaguchi, and M\. Shieh\(2024\)Monte carlo tree search boosts reasoning via iterative preference learning\.arXiv preprint arXiv:2405\.00451\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px3.p1.1)\.
- \[32\]S\. Yao, D\. Yu, J\. Zhao, I\. Shafran, T\. Griffiths, Y\. Cao, and K\. Narasimhan\(2023\)Tree of thoughts: deliberate problem solving with large language models\.Advances in neural information processing systems36,pp\. 11809–11822\.Cited by:[§1](https://arxiv.org/html/2606.25354#S1.p1.1),[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px1.p1.1)\.
- \[33\]Z\. Zhang, X\. He, W\. Yan, A\. Shen, C\. Zhao, S\. Wang, Y\. Shen, and X\. E\. Wang\(2025\)Soft thinking: unlocking the reasoning potential of llms in continuous concept space\.arXiv preprint arXiv:2505\.15778\.Cited by:[§1](https://arxiv.org/html/2606.25354#S1.p5.1),[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px4.p1.1)\.
- \[34\]A\. Zhou, K\. Yan, M\. Shlapentokh\-Rothman, H\. Wang, and Y\. Wang\(2023\)Language agent tree search unifies reasoning acting and planning in language models\.arXiv preprint arXiv:2310\.04406\.Cited by:[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px1.p1.1)\.
- \[35\]H\. Zhu, S\. Hao, Z\. Hu, J\. Jiao, S\. Russell, and Y\. Tian\(2025\)Reasoning by superposition: a theoretical perspective on chain of continuous thought\.arXiv preprint arXiv:2505\.12514\.Cited by:[§4](https://arxiv.org/html/2606.25354#S4.p1.1),[§6](https://arxiv.org/html/2606.25354#S6.SS0.SSS0.Px5.p1.1)\.

## Appendix AComplete Decoding Framework

### A\.1Rolling Local Lookahead Tree

Letx<tx\_\{<t\}denote the committed prefix before decoding positiontt\. Standard autoregressive decoding samples the next token directly fromπθ\(⋅∣x<t\)\\pi\_\{\\theta\}\(\\cdot\\mid x\_\{<t\}\)\. In contrast, LBR first constructs a local lookahead tree rooted atx<tx\_\{<t\}\.

Fix a tree depthL≥1L\\geq 1and widthK≥1K\\geq 1\. At decoding steptt, the local lookahead tree is denoted by

𝒜t=\(𝒱t,ℰt,rt\),\\mathcal\{A\}\_\{t\}=\(\\mathcal\{V\}\_\{t\},\\mathcal\{E\}\_\{t\},r\_\{t\}\),wherertr\_\{t\}is the root corresponding to the committed prefixx<tx\_\{<t\},𝒱t\\mathcal\{V\}\_\{t\}is the set of tree nodes, andℰt\\mathcal\{E\}\_\{t\}is the set of directed parent–child edges\. Every non\-leaf node hasKKsampled children, and every root\-to\-leaf path has lengthLL\. Thus,𝒜t\\mathcal\{A\}\_\{t\}is a depth\-LL, width\-KKsampled tree\.

Each nodev∈𝒱t∖\{rt\}v\\in\\mathcal\{V\}\_\{t\}\\setminus\\\{r\_\{t\}\\\}carries a token label

tok​\(v\)∈𝒱,\\mathrm\{tok\}\(v\)\\in\\mathcal\{V\},a depthd​\(v\)∈\{1,…,L\}d\(v\)\\in\\\{1,\\ldots,L\\\}, and a path from the root

path​\(v\)=\(tok​\(v1\),…,tok​\(vd​\(v\)\)\),\\mathrm\{path\}\(v\)=\\bigl\(\\mathrm\{tok\}\(v\_\{1\}\),\\ldots,\\mathrm\{tok\}\(v\_\{d\(v\)\}\)\\bigr\),wherert→v1→⋯→vd​\(v\)=vr\_\{t\}\\to v\_\{1\}\\to\\cdots\\to v\_\{d\(v\)\}=vis the unique path from the root tovv\. Children are sampled autoregressively: for every nodevvwithd​\(v\)<Ld\(v\)<L, its childrenCh​\(v\)=\{c1​\(v\),…,cK​\(v\)\}\\mathrm\{Ch\}\(v\)=\\\{c\_\{1\}\(v\),\\ldots,c\_\{K\}\(v\)\\\}are drawn from

tok\(ci\(v\)\)∼π~θ\(⋅∣x<t,path\(v\)\),i=1,…,K\.\\mathrm\{tok\}\(c\_\{i\}\(v\)\)\\sim\\tilde\{\\pi\}\_\{\\theta\}\\left\(\\cdot\\mid x\_\{<t\},\\mathrm\{path\}\(v\)\\right\),\\qquad i=1,\\ldots,K\.Hereπ~θ\\tilde\{\\pi\}\_\{\\theta\}denotes the model distribution after the same sampling filters used in ordinary decoding, such as temperature, top\-pp, top\-kk, or other filters\. The “Grow local tree” panels of Figure[1](https://arxiv.org/html/2606.25354#S1.F1)show the resulting tree before routing\. ForL=1L=1, the tree contains only candidate next tokens; forL=2L=2, each candidate next token also has one layer of sampled continuations\.

The crucial point is that the tree is not only sampled, but also forwarded\. For every non\-root nodev∈𝒱tv\\in\\mathcal\{V\}\_\{t\}, we compute and store the hidden state

ht​\(v\)=fθ​\(x<t,path​\(v\)\)∈ℝD\.h\_\{t\}\(v\)=f\_\{\\theta\}\\left\(x\_\{<t\},\\mathrm\{path\}\(v\)\\right\)\\in\\mathbb\{R\}^\{D\}\.Therefore, the local tree provides more than candidate token identities: it provides hidden\-state representations of local future continuations\. These hidden states are the evidence used by the router\.

### A\.2Routing over Forwarded Subtrees

The depth\-11children of the root are the roots of the candidate subtrees:

ut,1,…,ut,K∈Ch​\(rt\)\.u\_\{t,1\},\\ldots,u\_\{t,K\}\\in\\mathrm\{Ch\}\(r\_\{t\}\)\.The token labeltok​\(ut,k\)\\mathrm\{tok\}\(u\_\{t,k\}\)is thekk\-th candidate next token\. The subtree rooted atut,ku\_\{t,k\}is

𝒜t​\[ut,k\]=\{v∈𝒱t:ut,k​lies on the path from​rt​to​v\}\.\\mathcal\{A\}\_\{t\}\[u\_\{t,k\}\]=\\\{v\\in\\mathcal\{V\}\_\{t\}:u\_\{t,k\}\\text\{ lies on the path from \}r\_\{t\}\\text\{ to \}v\\\}\.This subtree contains the candidate tokentok​\(ut,k\)\\mathrm\{tok\}\(u\_\{t,k\}\)and all already\-forwarded local continuations below it up to depthLL\.

Letρϕ\\rho\_\{\\phi\}be a router parameterized byϕ\\phi\. Given the committed prefix and the forwarded lookahead tree, the router outputs a distribution over theKKdepth\-11subtrees:

ρϕ\(⋅∣x<t,𝒜t;θ\)∈ΔK\.\\rho\_\{\\phi\}\\left\(\\cdot\\mid x\_\{<t\},\\mathcal\{A\}\_\{t\};\\theta\\right\)\\in\\Delta^\{K\}\.We write the dependence onθ\\thetabecause the router input includes hidden states\{ht​\(v\):v∈𝒱t\}\\\{h\_\{t\}\(v\):v\\in\\mathcal\{V\}\_\{t\}\\\}computed by the base language model\. Equivalently, one may denote this dependence byρϕ,θ\(⋅∣x<t,𝒜t\)\\rho\_\{\\phi,\\theta\}\(\\cdot\\mid x\_\{<t\},\\mathcal\{A\}\_\{t\}\)\.

The router samples or selects a subtree index

kt⋆∼ρϕ\(⋅∣x<t,𝒜t;θ\),k\_\{t\}^\{\\star\}\\sim\\rho\_\{\\phi\}\\left\(\\cdot\\mid x\_\{<t\},\\mathcal\{A\}\_\{t\};\\theta\\right\),and commits the corresponding root token

xt=tok​\(ut,kt⋆\)\.x\_\{t\}=\\mathrm\{tok\}\(u\_\{t,k\_\{t\}^\{\\star\}\}\)\.The “Router selects subtree” panels in Figure[1](https://arxiv.org/html/2606.25354#S1.F1)emphasize that the router scores candidate subtrees using the forwarded hidden states of the local tree, rather than using only the root next\-token logits\.

### A\.3Efficient Tree\-Structured Forwarding

A naive implementation would forward every sampled path independently\. LBR instead uses a tree representation that shares common prefixes among sampled paths\. The sampled paths are stored both as a full tree with multiplicity and as a deduplicated trie\. The full tree preserves the sampling multiplicity of repeated tokens, while the trie allows shared prefixes to reuse key\-value cache entries\.

At each grow step, all current frontier nodes are processed in a batched tree\-attention forward pass\. Each frontier node attends to the committed prefix and to the ancestor tokens along its own branch:

KV​\(v\)=KV​\(x<t\)∥KV​\(ancestors​\(v\)\)\.\\mathrm\{KV\}\(v\)=\\mathrm\{KV\}\(x\_\{<t\}\)\\;\\\|\\;\\mathrm\{KV\}\(\\mathrm\{ancestors\}\(v\)\)\.This tree\-causal attention pattern computes the hidden states for all newly expanded nodes in parallel\. Since the router consumes these hidden states, no additional language\-model forward pass is needed solely for routing\.

The computational overhead of LBR comes from forwarding the local tree frontier and applying the lightweight router\. The benefit is that each token decision is informed by a structured set of local future hidden states rather than by the root next\-token distribution alone\.

### A\.4Set\-Attention Router

The router compares theKKcandidate subtrees using hidden states from the already\-forwarded lookahead tree\. Since the order of sampled rollouts is arbitrary, the router should not depend on the incidental ordering of sampled branches\. We instantiateρϕ\\rho\_\{\\phi\}as a hierarchical set\-attention router with two levels: it first aggregates rollouts within each candidate subtree, and then compares the resulting candidate representations across subtrees\.

For each candidate subtree rooted atut,ku\_\{t,k\}, let

𝒫t,k=\{pt,k,1,…,pt,k,Mk\}\\mathcal\{P\}\_\{t,k\}=\\\{p\_\{t,k,1\},\\ldots,p\_\{t,k,M\_\{k\}\}\\\}denote the set of root\-to\-leaf rollout paths contained in that subtree, whereMk=KL−1M\_\{k\}=K^\{L\-1\}for a full width\-KK, depth\-LLtree\. Each rollout pathpt,k,mp\_\{t,k,m\}has an ordered hidden\-state sequence

Ht,k,m=\(ht,k,m\(1\),…,ht,k,m\(L\)\)∈ℝL×D,H\_\{t,k,m\}=\\bigl\(h\_\{t,k,m\}^\{\(1\)\},\\ldots,h\_\{t,k,m\}^\{\(L\)\}\\bigr\)\\in\\mathbb\{R\}^\{L\\times D\},whereht,k,m\(ℓ\)h\_\{t,k,m\}^\{\(\\ell\)\}is the hidden state at depthℓ\\ellalong that rollout\. The first hidden stateht,k,m\(1\)h\_\{t,k,m\}^\{\(1\)\}corresponds to the candidate rootut,ku\_\{t,k\}and is therefore shared across all rollouts in the same subtree\.

We first encode each rollout path with a small position\-aware Transformer encoder\. LetWin∈ℝp×DW\_\{\\mathrm\{in\}\}\\in\\mathbb\{R\}^\{p\\times D\}be an input projection and lete1,…,eL∈ℝpe\_\{1\},\\ldots,e\_\{L\}\\in\\mathbb\{R\}^\{p\}be learned depth embeddings\. We define

h¯t,k,m\(ℓ\)=Win​ht,k,m\(ℓ\)\+eℓ,ℓ=1,…,L\.\\bar\{h\}\_\{t,k,m\}^\{\(\\ell\)\}=W\_\{\\mathrm\{in\}\}h\_\{t,k,m\}^\{\(\\ell\)\}\+e\_\{\\ell\},\\qquad\\ell=1,\\ldots,L\.After prepending a learned\[CLS\]vector, we apply a shared path encoder:

Rt,k,m=PathEncϕ​\(\[CLS\],h¯t,k,m\(1\),…,h¯t,k,m\(L\)\),R\_\{t,k,m\}=\\mathrm\{PathEnc\}\_\{\\phi\}\\left\(\[\\texttt\{CLS\}\],\\bar\{h\}\_\{t,k,m\}^\{\(1\)\},\\ldots,\\bar\{h\}\_\{t,k,m\}^\{\(L\)\}\\right\),and read out the\[CLS\]representation as the rollout feature

rt,k,m=Rt,k,m\[CLS\]∈ℝp\.r\_\{t,k,m\}=R\_\{t,k,m\}^\{\\texttt\{\[CLS\]\}\}\\in\\mathbb\{R\}^\{p\}\.This step preserves the ordered structure along a rollout path\.

Next, we aggregate the rollout features within each candidate subtree\. Since the rollouts under the same candidate token have no canonical order, we use a set\-attention block followed by attention pooling:

Zt,k=SetAttnϕsub​\(rt,k,1,…,rt,k,Mk\),Z\_\{t,k\}=\\mathrm\{SetAttn\}^\{\\mathrm\{sub\}\}\_\{\\phi\}\\left\(r\_\{t,k,1\},\\ldots,r\_\{t,k,M\_\{k\}\}\\right\),gt,k=AttnPoolϕ​\(qsub,Zt,k\)∈ℝp,g\_\{t,k\}=\\mathrm\{AttnPool\}\_\{\\phi\}\\left\(q\_\{\\mathrm\{sub\}\},Z\_\{t,k\}\\right\)\\in\\mathbb\{R\}^\{p\},whereqsub∈ℝpq\_\{\\mathrm\{sub\}\}\\in\\mathbb\{R\}^\{p\}is a learned pooling query\. This gives one permutation\-invariant representationgt,kg\_\{t,k\}for the entire subtree rooted atut,ku\_\{t,k\}\.

Finally, the router compares theKKcandidate subtrees jointly\. We pass the candidate representations through a second set\-attention block:

g~t,1,…,g~t,K=SetAttnϕcand​\(gt,1,…,gt,K\),\\tilde\{g\}\_\{t,1\},\\ldots,\\tilde\{g\}\_\{t,K\}=\\mathrm\{SetAttn\}^\{\\mathrm\{cand\}\}\_\{\\phi\}\\left\(g\_\{t,1\},\\ldots,g\_\{t,K\}\\right\),which is permutation\-equivariant across candidate subtrees\. Each candidate is then assigned a scalar score

st,k=wϕ⊤​g~t,k,s\_\{t,k\}=w\_\{\\phi\}^\{\\top\}\\tilde\{g\}\_\{t,k\},and the router distribution is

ρϕ​\(k∣x<t,𝒜t;θ\)=exp⁡\(st,k/τ\)∑k′=1Kexp⁡\(st,k′/τ\),\\rho\_\{\\phi\}\(k\\mid x\_\{<t\},\\mathcal\{A\}\_\{t\};\\theta\)=\\frac\{\\exp\(s\_\{t,k\}/\\tau\)\}\{\\sum\_\{k^\{\\prime\}=1\}^\{K\}\\exp\(s\_\{t,k^\{\\prime\}\}/\\tau\)\},whereτ\>0\\tau\>0is a routing temperature\.

For the main settingL=1L=1, each candidate subtree contains only a single forwarded token, soMk=1M\_\{k\}=1and the path encoder receives only the hidden stateht​\(ut,k\)h\_\{t\}\(u\_\{t,k\}\)\. In this case, within\-subtree aggregation is trivial, but the cross\-subtree set\-attention block still allows the router to compare the post\-token hidden states of theKKsampled candidate next tokens\. ForL\>1L\>1, the same architecture first summarizes each local rollout path and then aggregates the multiple rollouts under each candidate before comparing candidates\.

#### A\.4\.1Router Hyperparameters

We instantiate the router with internal dimensionp=128p=128andH=4H=4attention heads in every multi\-head\-attention block\. The input projectionWin∈ℝp×DW\_\{\\mathrm\{in\}\}\\in\\mathbb\{R\}^\{p\\times D\}uses the language model’s hidden dimensionD=1536D=1536\(DeepSeek\-R1\-Distill\-Qwen\-1\.5B\)\. The shared path encoderPathEncϕ\\mathrm\{PathEnc\}\_\{\\phi\}is a single pre\-LN Transformer encoder layer with GELU activations and FFN expansion factor44\. Each set\-attention blockSetAttnϕsub\\mathrm\{SetAttn\}^\{\\mathrm\{sub\}\}\_\{\\phi\}andSetAttnϕcand\\mathrm\{SetAttn\}^\{\\mathrm\{cand\}\}\_\{\\phi\}is one SAB layer with the same head count and FFN structure\. The within\-subtree poolingAttnPoolϕ\\mathrm\{AttnPool\}\_\{\\phi\}uses a single learned queryqsub∈ℝpq\_\{\\mathrm\{sub\}\}\\in\\mathbb\{R\}^\{p\}via standard scaled\-dot\-product attention\. We disable dropout throughout the router\. The final routing softmax temperature isτ=1\\tau=1\. The learning rate of router is1×10−41\\times 10^\{\-4\}\.

## Appendix BWhy Soft Thinking Underperforms in Hierarchical Planning: Concept Ambiguity

![Refer to caption](https://arxiv.org/html/2606.25354v1/figs/n1_identity_probe.png)Figure 7:Concept\-identity probe after the first generated graph node\. Discrete CoT and LBR preserve node identity through discrete one\-hot digit tokens, while Soft Thinking loses concept information after mixing candidate embeddings\.Mixture embeddings can blur the identity of the concept branch being followed\. In our task, discrete CoT and LBR always condition on one\-hot digit tokens\. Soft Thinking instead feeds mixtures of legal candidate embeddings at branching positions, which can leave the model in an ambiguous state between multiple possible concepts\.

Figure[7](https://arxiv.org/html/2606.25354#A2.F7)validates this effect with a concept\-identity probe\. After the model processes the digit sequence of the first generated graph node, a linear probe can perfectly recover the node identity for Discrete CoT and LBR\. For Soft Thinking, however, the probe remains substantially below perfect accuracy, reaching only0\.7590\.759at the final digit of the node and0\.8030\.803at the following arrow token\. Thus, even after a full node has been processed, the soft hidden state does not fully identify which concept branch the model is following\.

This supports the mechanism suggested in Figure[3](https://arxiv.org/html/2606.25354#S4.F3): LBR preserves candidate\-specific branches and routes among them, whereas Soft Thinking compresses competing candidates into a shared continuous state\. The resulting concept ambiguity explains why Soft Thinking underperforms LBR on the radix\-translated planning task\.