PaperlayerView exact revision on arXiv

Source-linked AI summary

Fork Where the Model Changes Its Mind: Belief-Shift Branching for Tree-Structured Reinforcement Learning

Bin Lei, Yu Li, Prafulla Kumar Choubey, Jiaxin Zhang, Becky Xiangyu Peng, Qinyuan Ye, Kartik Narayan, Caiwen Ding, Silvio Savarese, Chien-Sheng Wu

arXiv:2609.11061v1cs.AI

TL;DR

Tree RLVR needs to place a small number of costly forks where they can provide useful step credit, but existing selectors largely follow structure or wording uncertainty. This paper introduces belief-shift branching, which forks before the model’s answer belief changes most, and reports first-place validation results plus stronger mathematics and coding RL outcomes. The learned activation read is limited to frozen-policy validation, while fit-free reads are used in RL.

  • Problem

    Tree RLVR provides critic-free step credit through sibling outcomes, but each fork is costly and existing structural or entropy-based selectors may miss where the outcome changes.

  • Method

    Belief-shift branching reads answer beliefs at candidate boundaries and places forks before the largest consecutive change, using black-box, logit-lens, or offline-learned activation reads.

  • Results

    Belief-shift ranks first in all eight validation panels, leads every mathematics aggregate, and sweeps every OLMo code column; on OLMo-3-7B it gains +2.9 on AIME 2026 over the strongest baseline.

  • Takeaways & Limitations

    Fork selection can target the model’s changing answer belief while preserving verifier-based credit, with the probe adding about 1% mathematics and under 5% code step compute.

  • Takeaways & Limitations

    The learned Belief-Shift Vector is used only for frozen-policy pre-RL validation because training changes make its offline direction off-policy.

Abstract

from arXiv · show

Tree-structured rollouts give critic-free reinforcement learning with verifiable rewards (RLVR) step-level credit: fork a chain at an intermediate point, and sibling outcome differences estimate step value. Each fork adds sampling cost, so realistic budgets typically allow only a few forks per chain. A fork placed where the outcome is already largely settled yields siblings that mostly agree and provide almost no credit signal; hence, for a given tree size, where forks are placed largely determines how much step-level RL can gain. Most existing mainstream methods place forks by structure, such as fixed lengths, midpoints, and delimiters, or by next-token entropy. We formalize fork placement as locating the \emph{pivots} of the chain's value curve, where the expected outcome turns. We propose \emph{belief-shift branching}: read the model's answer belief at candidate boundaries and fork just before the step where consecutive beliefs diverge most. Three instantiations, none needing step-level supervision, span access levels: a black-box probe, a logit-lens depth profile, and a learned activation direction, which is fit offline and therefore used only in the validation before RL training. The signal only \emph{places} forks, and the probe costs about $1\%$ of step compute on mathematics and under $5\%$ on code when it runs inside the rollout engine. In that validation, against Monte-Carlo value curves, a belief-shift signal ranks first in each of the eight model$\times$benchmark panels, ahead of entropy, structural, and LLM-judge baselines. In RL across three model families and two domains, belief-shift forking leads every mathematics aggregate, on OLMo-3-7B by $+2.6$ aggregate and $+2.9$ on AIME 2026 over the strongest baseline, and sweeps every OLMo code column, by $+6.5$ on LiveCodeBench-medium.

1 INTRODUCTION

Tree RLVR can assign step-level credit without a critic, but scarce forks must be placed where outcomes are still changing. The paper proposes belief-shift branching and reports stronger selector validation and RL results than structural and uncertainty-based alternatives.

  • Motivation: Tree RLVR converts sibling outcome differences into counterfactual step credit, but each fork adds rollout tokens and realistic budgets permit only a few forks per chain.Fork placement therefore determines whether sibling outcomes provide useful credit information.
  • Method: Belief-shift branching reads the model’s answer belief at candidate boundaries and places forks just before the largest consecutive-belief change.The three reads span black-box probing, a logit-lens depth profile, and a learned activation direction, without step-level supervision.
  • Illustrative case: 3.6× more value movement: PROBE-JS captures P |∆V ∗|=1.56 versus entropy’s 0.44 with ten forks, while entropy wastes six forks in the settled region.The comparison uses a real reasoning chain and a Monte-Carlo ground-truth value curve.
  • End-to-end results: +2.9 on OLMo-3-7B AIME 2026 over the strongest baseline; belief-shift forking leads every mathematics aggregate and sweeps every OLMo code column.The RL comparison spans three model families and two domains under matched budgets.
  • Pre-RL validation: Belief-shift selectors rank first in all eight model×benchmark validation panels against Monte-Carlo value curves, ahead of entropy, structural, and LLM-judge baselines.The validation is performed before RL to isolate fork-selection quality from optimizer and seed effects.

2 BACKGROUND AND RELATED WORK

The background contrasts outcome-only RLVR and tree-based credit with fork selectors that mostly use structure or wording uncertainty. The paper positions belief-shift selection as a direct alternative that targets pivots in the outcome-value curve.

  • Trajectory-level RLVR: Group-relative RLVR copies one scalar reward across tokens, which can reinforce flawed reasoning in correct solutions and produce zero gradient when all rollouts fail.Process supervision improves on outcome-only training in reasoning settings, with up to 6× better RL sample efficiency from process advantage verifiers.
  • Tree rollout structures: On-policy branching trees obtain RL advantages directly from sibling outcomes without an explicit value model, whereas other tree structures use extra sampling for Monte-Carlo values or learned supervision.All three rollout structures consume candidate boundaries where a chain may be cut.
  • Existing fork selectors: Existing selectors follow surface structure, token uncertainty, or external correctness judgments rather than directly reading the outcome decision at a fork.Structural methods include delimiters and fixed grids; entropy-based methods track wording freedom and may miss confidently wrong steps.
  • Existing fork selectors: 2.1 pass-rate points: TreeRL’s surprisal-based forking beats random at the same tree configuration, with about 8% fewer generated tokens.TreePO’s static probability-guided branching allocation instead loses to uniform in its cited sampling comparison.
  • Target quantity: The ideal selector targets pivots of V ∗(t)=Pr[correct | x, y≤t], because sibling comparison estimates the local change of this value curve.The paper describes these pivots as unobservable and motivates proxy selectors as approximations.

3 BELIEF-SHIFT BRANCHING

Belief-shift branching scores candidate step boundaries by how much the model’s answer belief changes, then allocates scarce forks to the largest shifts. It supports black-box, white-box, and offline learned reads while preserving critic-free sibling-based credit.

  • 3.1 CANDIDATE STEP BOUNDARIES: Belief-shift branching scores each candidate segment by the change between consecutive answer beliefs and forks at the highest-scoring boundaries.The selector uses a generic boundary-scoring template, with belief-change measures instantiated as divergences or profile-change norms.
  • 3.2 PROBE-JS: BLACK-BOX BELIEF PROBING: PROBE-JS elicits a final-answer distribution after each prefix and ranks segments by symmetric, bounded Jensen–Shannon divergence between consecutive top-K beliefs.The probe reuses the rollout prefix cache and samples a short answer read; the symmetric bounded score avoids KL’s unbounded sensitivity to truncated tail tokens.
  • 3.3 LENS-SHIFT: WHITE-BOX BELIEF READING: LENS-SHIFT teacher-forces the committed answer and measures how its per-layer log-probability profile changes, localizing belief reorganization across network depth.Its profile traces answer formation through layers, while scalar read-outs capture settled belief strength and lock-in depth; RL uses the shape-based read.
  • 3.4 BELIEF-SHIFT VECTOR: A LEARNED READ FOR PRE-RL VALIDATION: BSV projects activations onto an offline-fitted surge/steady/drop direction to test whether direct belief shifts reconstruct the ground-truth value curve.Because the fitted direction becomes off-policy as training changes the model, BSV is restricted to frozen-policy pre-RL validation.
  • 3.5 INTEGRATION INTO TREE-STRUCTURED RLVR: The integration samples M=4 root chains, forks each once into k=2 siblings at its top-scoring boundary, and broadcasts verified sibling-based credit across the forked segment.All arms share the same RLVR optimizer and differ only in the fork criterion; MIDPOINT and ENTROPY provide structural and uncertainty baselines.
  • 3.6 COST AND A PROBE-VALUE BOUND: The probe’s fixed-length overhead shrinks relative to longer chains, while its score lower-bounds movement in the probe value and therefore cannot hide a moved value when the score is near zero.The paper reports FLOP-based overhead because serving paths are not wall-clock comparable.

4 PRE-RL VALIDATION: DO THE SIGNALS FIND THE TRUE VALUE PIVOTS?

The validation isolates fork-selection quality by comparing selectors against Monte-Carlo value curves, where good boundaries reconstruct the chain’s true value pivots. Belief-shift signals rank first across all eight model×benchmark panels and transfer to held-out GPQA-Diamond.

  • Evaluation: The evaluation estimates V*(t_j) with 16 verifier-scored Monte-Carlo completions, retains each selector’s top 10 boundaries, and measures interpolation error.Lower reconstruction error indicates that selected boundaries bracket the true pivots; flat-value problems are excluded and a dynamic program supplies the oracle floor.
  • Selectors: BSV, PROBE-JS, and LENS-SHIFT are compared with an LLM judge, entropy, newline, uniform, and random placement baselines.BSV is fit offline on held-out data and serves as a validation instrument, while fit-free reads support later RL use.
  • Signal quality: A belief-shift signal ranks first in all eight model×benchmark panels for reconstructing Monte-Carlo value curves.The comparison covers four probe models on AIME and held-out GPQA-Diamond, with error measured above a dynamic-program oracle floor.
  • Transfer: On held-out GPQA-Diamond, BSV fit on out-of-domain math beats every baseline on average, 5.935 versus 6.041.It wins in three of four panels, while the LLM judge trails every belief-shift read on GPQA.

5 RL EXPERIMENTAL SETUP

The RL study compares fork criteria under matched training conditions across three model substrates and mathematics and code domains. All arms share the tree, data order, reward, optimizer, and critic-free training protocol, differing only in fork placement.

  • Models: Qwen3-4B, OLMo-3-7B, and Nemotron-Nano-9B-v2 span the evaluated model architectures and substrates.The set includes base models and an OLMo SFT checkpoint.
  • Domains and data: Mathematics uses DAPO-Math-17k training and three validation suites with an unweighted aggregate, while code uses DeepCoder-24K training and LiveCodeBench-v6 validation.The mathematics suites are OlympiadBench, AIME 2026, and Omni-MATH-500.
  • Compared arms: MIDPOINT, ENTROPY, PROBE-JS, and LENS-SHIFT share the tree construction, data order, reward, and optimizer, differing only in fork criterion.All arms are critic-free GRPO-style methods; a learned value model is outside the comparison.
  • Evaluation protocol: Validation tracks accuracy every 20 steps, and each arm is reported at its single best-aggregate checkpoint under matched budget protocols.Qwen arms use exactly 680 steps, while OLMo and Nemotron share common per-group caps.

6 RL RESULTS

Belief-shift branching leads the mathematics aggregates and most individual columns, while its code advantage is strongest on OLMo and weaker under Nemotron’s constrained setting. The results also show that fork gains depend on sibling contrast, not ranking quality alone.

  • Mathematics: A belief-shift arm leads the aggregate on every evaluated model, with PROBE-JS winning OLMo-3-7B and Nemotron-9B and LENS-SHIFT winning Qwen3-4B mathematics.The reported aggregates are 23.3 versus 20.7, 45.1 versus 40.9, and 29.6 versus 28.7, respectively.
  • Mathematics: The OLMo belief-shift arm tops 11 of 12 mathematics columns and raises AIME 2026 from 12.3 for the strongest baseline to 15.2.Entropy branching hurts Nemotron, reaching a 40.0 aggregate below MIDPOINT, consistent with its weak pre-RL value tracking.
  • Code: PROBE-JS wins every OLMo-3-7B code column, including a +6.5 advantage on LiveCodeBench-medium, 34.9 versus 31.0 for the best baseline.On Qwen3-4B, PROBE-JS has the best aggregate at 32.2 versus 31.8; under Nemotron’s 240-step cap, MIDPOINT leads.
  • Mechanism: RL-trained Qwen siblings reach the same final answer 63% of the time versus about 5% for the base model, shrinking the gain despite healthy ranking.Top- versus bottom-tercile JS boundaries show 55% versus 24% sibling disagreement, but near-clone continuations leave little contrast.
  • Reporting: Table 3 reports accuracy at each arm’s single best-aggregate checkpoint across mathematics suites and code splits under model-specific budget constraints.Qwen uses 680 steps and Nemotron code uses a 240-step cap; shaded rows identify belief-shift arms.

7 ABLATIONS

Ablations show that belief-shift gains persist with smaller trees, increase with longer belief reads, and grow with model scale. The policy-update algorithm also materially affects the final aggregate.

  • Belief-read length: Reading four belief tokens instead of one nearly doubles the edge over MIDPOINT, from +3.8 to +7.3, and raises AIME from 15.2 to 22.5.The longer read retains the exact one-token belief at position one and adds three positions.
  • Tree budget: Shrinking the tree to (2, 1) costs 3.2 aggregate points yet still edges the full-budget MIDPOINT arm by +0.6 with one-third its rollouts.This ablation uses the OLMo-3-7B mathematics slice with PROBE-JS.
  • Model scale: LENS-SHIFT wins the aggregate at both Qwen sizes, while its margin over MIDPOINT grows from +1.0 to +1.6 as scale increases.PROBE-JS’s margin likewise grows from +0.1 to +1.1, and it is silver at 8B.

8 CONCLUSION

Belief-shift branching places scarce tree forks where the model changes its mind, improving selector quality and RL outcomes without step-level supervision.

  • 8 CONCLUSION: A belief-shift selector ranks first in all eight pre-RL model×benchmark panels, while belief-shift arms lead mathematics aggregates and the probe sweeps OLMo code columns.The approach uses rollout access and verifier labels and is intended to drop into existing tree RLVR pipelines.
  • 8 CONCLUSION: The method reads answer beliefs at candidate boundaries and spends forks where consecutive beliefs shift most.Its signal places forks; verified sibling outcomes still provide credit.
  • ETHICS STATEMENT: The study evaluates mathematical and coding reasoning on publicly available models and datasets, without human subjects or personally identifiable data.The paper notes usual dual-use considerations for stronger language models and no risks beyond standard RLVR training.

A FULL CONFIGURATION

The configuration specifies separate pre-RL probe and RL checkpoints, fixed tree and optimization settings, belief-read variants, substrate-dependent candidate ranges, and robustness checks for checkpoint selection.

  • Tree and rollout: Each prompt emits 12 rollouts from four root chains, one criterion-chosen cut per chain, and two fresh sibling continuations per cut.Chains run to completion; midpoint fallback and root refilling handle unavailable boundaries or continuation budget.
  • Optimization: Training uses 32 generated prompts, 16 retained by DAPO, 192 rollouts per step, AdamW at lr = 10^-6, asymmetric PPO clipping, and TreeRL global+local advantage.The configuration also omits entropy bonuses, KL terms, and advantage standardization.
  • Signal configuration: The boundary pool uses line-break positions with at most C=16 candidates, while PROBE-JS reads top-K=20 next-token log-probabilities and LENS-SHIFT reads activation profiles.The belief-read arms reduce the vLLM memory fraction, and substrate settings differ between base and SFT policies.
  • Base versus SFT substrates: Base-model belief reads score only the first ρ=0.5 of each chain, whereas OLMo scores the full chain with ρ=1.This accommodates short, repetitive, or unstructured early base-model chains.
  • Robustness to checkpoint selection: Checkpoint-selection robustness preserves the mathematics aggregate lead for belief-shift arms on all three models under every rule, while PROBE-JS sweeps OLMo code columns under two rules.Peak selection favors baselines rather than manufacturing the effect, although several sub-point margins change order.
  • Fitting design: The fitting design combines shared rollout prefixes with surge, steady, and drop continuations covering the admissible g×v_i×v_f cells.Belief changes must occur in content or confidence rather than surface markers, and catch or derail continuations must be mathematically valid.

B COMPUTE-COST DERIVATION

The compute analysis counts probe forwards against rollout and training forwards, reports FLOP overhead rather than wall-clock time, and distinguishes probe and LENS-SHIFT serving costs.

  • Probe cost: The probe requires 20,480 token-forwards per step from 32 prompts, four chains, 16 candidates, and a ten-forward probe.The accounting uses T_probe = P_gen · M · C · (L_probe+1).
  • Probe cost: Probe FLOP shares of a full training step are 1.3% for OLMo, 0.45% for Qwen, and 0.70% for Nemotron.The corresponding probe-to-rollout FLOP ratios are 6.0%, 1.9%, and 3.1%.
  • Serving caveat: Wall-clock time is not reported because probe serving differs across models, including off-engine Hugging Face execution for Nemotron after interleaved probes destabilized rollout serving.The affected PROBE-JS arm exhibited response-length collapse, so the probe was moved off-engine.
  • LENS-SHIFT cost: LENS-SHIFT forwards each candidate’s full prefix on a co-located Hugging Face copy without cache reuse, though prefix reuse would reduce its order to C(E+n) token-forwards per chain.The uncached cost is an engineering choice rather than an intrinsic property of the signal.

C PROOF OF THE PROBE-VALUE BOUND

The proof bounds the probe’s value change using total variation and Jensen–Shannon divergence, while separating that exact bound from approximations introduced by top-K truncation and elicitation.

  • Setup: The bound applies exactly to the probe’s floored, renormalized top-K beliefs over a finite answer set.The proof treats these scored distributions as p and q.
  • Total variation: The first inequality follows because the probability change of the correct answer is at most total variation distance.The singleton event containing the correct answer gives the required specialization.
  • Jensen–Shannon bound: For the midpoint distribution m = 1/2(p+q), the two total-variation distances to m are equal and Pinsker’s inequality bounds each by its KL divergence.Combining the two halves yields the Jensen–Shannon divergence bound.
  • Scope of the bound: The proof does not cover the gap between probe value and true chain value caused by top-K truncation or forced-answer elicitation.That gap is assessed empirically by reconstructing Monte-Carlo V* before RL.

D FORK-CONTRAST DIAGNOSTIC

The diagnostic compares fork-relevant belief shifts and sibling agreement across three policies, showing that trained Qwen retains more answer variability while OLMo exhibits the largest peak belief shift.

  • 55% sibling disagreement in trained Qwen’s highest PROBE-JS tercile versus 24% in its lowest tercile.Across the trained-Qwen boundaries, higher PROBE-JS scores identify boundaries with more divergent sibling answers.
  • 62.7% of trained-Qwen siblings reproduce the root answer, compared with 4.8% for base Qwen and 0% for OLMo.The measurement uses first-six OlympiadBench problems with parsable answers and forks at every sampled boundary.
  • OLMo’s average maximum PROBE-JS score is 0.410, versus 0.289 for trained Qwen and 0.264 for base Qwen.These values come from 30 AIME 2026 chains per policy and correspond to the reported 1.4× contrast.

E ADDITIONAL RESULTS

Additional results support belief-shift selection while delimiting BSV’s use: it leads pre-RL validation, but its learned direction is restricted to frozen-policy analysis.

  • Figure 4 contains the per-benchmark validation trajectories underlying the OLMo and Qwen rows of Table 3.
  • BSV leads pre-RL rankings but is excluded from RL arms because its per-model offline direction is not guaranteed to transfer across checkpoints.The RL experiments therefore use the two fit-free read-outs.
  • Figure 6 reports correlations of +0.26 to +0.54 between the label-free layer criterion and reconstruction error across four probe models.The criterion selects a near-optimal evaluation layer without labels and is exact-best on Gemma-4-31B.

F THE FIGURE 1 CASE IN FULL

The complete Figure 1 case follows a correct OLMo-3.1-32B-Think chain whose value fluctuates early, where PROBE-JS concentrates forks, while entropy places many forks after the outcome is settled.

  • All ten PROBE-JS forks fall in the first 26% of the chain, where Monte-Carlo value swings between 0.62 and 1.00.The chain repeatedly sets up, botches, and repairs the shoelace computation before reaching the correct answer 588.
  • Entropy places six of ten forks in the settled tail, where V* = 1 and ΔV* = 0 despite high token uncertainty.The case illustrates uncertainty from interchangeable wording rather than an undecided outcome.
  • The case uses an OLMo-3.1-32B-Think sample with 128 newline boundaries and 16 Monte-Carlo completions per boundary.The full appendix includes the prompt, twenty fork picks, their scores, ground-truth value movement, and the reasoning chain.

F.2 THE TWENTY FORK PICKS

Table 8 lists the twenty fork picks in chain order, pairing each selector’s score with Monte-Carlo value estimates immediately before and after the chosen boundary.

  • Table 8 records each fork’s character offset, chain fraction, selector score, and Monte-Carlo value on both sides of the boundary.The picks are ordered chronologically along the Figure 1 chain.
  • PROBE-JS scores consecutive answer-belief divergence, whereas entropy scores boundary entropy.The table compares each selector’s own criterion with the value change around its selected boundary.

F.3 REASONING CHAIN

The chain uses coordinate geometry to place the triangle’s points and reflections, derive the triangle’s height from the given quadrilateral area, and conclude that the heptagon’s area is 588.

  • Coordinate setup: The coordinate setup places A=(0,0), B=(28,0), D=(4,0), and E=(20,0), while F and G divide AC at 1/7 and 5/7.These coordinates follow from the side-length partitions and the reflection construction.
  • Deriving the height: The quadrilateral-area equation 48|c_y|/7=288 yields |c_y|=42, so the triangle’s height is 42.The solution selects c_y=42 under the assumption that the triangle lies above AB.
  • Area calculation: The triangle area is therefore 1/2·28·42=588, independent of the unknown horizontal coordinate c_x.The base length is fixed at 28, and the derived height fixes the area regardless of C’s horizontal position.
  • Heptagon result: The shoelace calculation for AFNBCEM also gives area 588, matching the triangle area and the stated answer.Multiple intermediate calculations and the final summary report the same value.
Loading 2609.11061v1…