Source-linked AI summary
Efficiently Learning To Reason or Not to Reason: Root-token Policy Optimization for Adaptive Thinking
Taehyeon Kim, Hyunsoo Lee, Youngsoo Jang, Moontae Lee
TL;DR
Large reasoning models often apply explicit reasoning regardless of query difficulty, while open-source models lack an efficient way to internalize adaptive routing. RPO learns the initial reasoning branch through root-token policy optimization. Across model families and scales, it improves the accuracy-routing trade-off under most scenarios and lowers reasoning-per-accuracy ratios across math, scientific reasoning, and coding tasks.
Problem
Large reasoning models apply explicit reasoning uniformly across queries, wasting compute and latency on easy questions, while standardized methods for balancing reasoning and direct-response data remain lacking.
Source: §1 · pp. 1–2Method
RPO combines paired reasoning and non-reasoning samples with policy optimization that updates only the initial root token controlling whether to enter the reasoning branch.
Source: §1 · p. 2Results
Takeaways & Limitations
RPO treats adaptive thinking as a lightweight routing problem by updating the initial gating token rather than the remainder of the rollout.
Source: §6 · p. 9Evaluation on broader benchmarks requiring agentic behavior with tools is still needed.
Source: Limitations · p. 9
Abstract
from arXiv · showhide
Large reasoning models (LRMs) achieve strong performance by externalizing explicit reasoning traces before producing the answer, yet suffer from overthinking challenge that allocates uniformly heavy computation to queries of varying difficulty. While proprietary models mitigate this via opaque routing, open-source LRMs still lack an efficient mechanism to internalize adaptive reasoning due to both expensive training cost and limited disclosure of training recipes. In response, we introduce RPO (Root-token Policy Optimization), a framework that enables LRMs to self-determine when to reason by training only the initial root token (e.g., whether to invoke the <think> tag) via group relative reward and group-wise advantages. By focusing on this pivotal branching point, RPO drastically reduces training overhead and VRAM usage. Across multiple model families and scales, RPO learns difficulty-aware adaptive thinking at just ~2% of the training compute of prior adaptive-reasoning methods.
1 Introduction
Large reasoning models apply explicit reasoning broadly, creating overthinking and wasted compute on easy queries. RPO trains the initial reasoning branch to support difficulty-aware routing while reducing training overhead.
1 Introduction
- Uniform explicit reasoning creates overthinking, wasting compute and latency on trivial queries while open-source models lack an internalized alternative to proprietary routing.Existing adaptive-policy work also faces training-cost, data-curation, and initialization challenges.Source: §1 · pp. 1–2
- RPO generates paired reasoning and non-reasoning samples, then optimizes only the initial root token that opens or closes the <think> branch.The approach is designed to reduce training overhead and memory use while preserving the model’s original capability.Source: §1 · p. 2
- RPO enables learning when to reason across diverse LRM families and scales, with analysis of calibration through target Think% or a routing threshold.The paper also reports that policies learned in mathematical domains elicit adaptive reasoning in scientific and code domains.Source: §1 · p. 2
2 Related work
Prior efficiency methods reduce or control reasoning computation in different ways, while adaptive policy optimization learns when to reason. RPO focuses on internalizing the initial branch decision, against practical gaps in evaluation, training cost, and accessible data.
2.1 Efficient reasoning models
- Early-exit, reasoning-truncation, token-budget, latent-rationale, and lightweight data or prompting approaches address efficiency; RPO instead internalizes whether to open an explicit reasoning trace.Budget methods may require external selection rules or calibrated estimators that do not transfer cleanly across domains.Source: §2.1 · pp. 2–3
2.2 Difficulty-aware policy optimization
- Adaptive policy work frames reasoning as a difficulty-conditioned choice, spanning binary thinking modes, reduced redundant tokens, multiple reasoning formats, and model-agnostic configurations.These approaches show that adaptive reasoning is learnable and can outperform static long-CoT inference.Source: §2.2 · p. 3
- Existing policy-optimization work leaves open routing calibration beyond math-heavy benchmarks and incurs high trajectory-training costs, while often relying on curated or proprietary data sources.The cited hurdles include expensive VRAM, gradient computation, and wall-clock time.Source: §2.2 · p. 3
3 Method
The method treats adaptive reasoning as a routing decision informed by model-specific SHORT/LONG success differences. It teaches both modes with filtered demonstrations, then optimizes the root gating token rather than the full reasoning continuation.
3.1 Motivation
- Mode-conditioned success differences provide a model-specific difficulty signal, while thresholding SHORT-mode confidence trades routing volume against accuracy.Raising the threshold routes fewer queries to SHORT and more to LONG; intermediate thresholds can reduce LONG usage while preserving overall accuracy.Source: §3.1 · pp. 3–4
3.2 SFT for mode selection
- SFT demonstrations come from correct, template-compliant SHORT and LONG generations, sampled from separate pools without requiring each query to succeed in both modes.The SFT loop refreshes correctness-filtered data while the default-template performance gap stays within tolerance; one iteration typically suffices across tested LRMs.Source: §3.2 · pp. 4–5
3.3 Root-token policy optimization (RPO)
- RPO optimizes only the first gating token, treating the remaining sampled trajectory as environment output rather than a learning target.The gate selects SHORT or LONG; subsequent tokens are masked from the gradient.Source: §3.3 · p. 5
- Balanced SHORT/LONG rollouts receive rewards based on their per-prompt mode-group success counts, favoring the mode with higher group-level correctness for that prompt.Balanced sampling is intended to prevent trivial mode collapse and keep advantage statistics comparable across modes.Source: §3.3 · p. 5
- Per-mode, group-level advantage pooling uses cross-prompt variation for nonzero variance and presupposes prompt-difficulty diversity within each batch.Pooling separately by mode keeps standardized advantages comparable even when the modes have systematically different success rates.Source: §3.3 · p. 6
- Updating only the root token (T=1) removes the linear dependence of activation memory on long chain-of-thought length during training.Sampling still uses the full rollout context; the training-time savings arise because RPO backpropagates through only the gating token.Source: §3.3 · p. 6
4 Experiment
RPO is evaluated across model families and scales on math, scientific reasoning, and coding tasks. It generally improves the accuracy–routing trade-off, with the direction and size of gains depending on the model configuration.
4.2 Datasets and evaluation
- The evaluation covers MATH500, AIME24, AIME25, GPQA, and LiveCodeBench, using exact-match correctness on parsed final answers.The reported pipeline is Baseline → SFT → RPO; models share a two-mode gate implemented through the first newline token(s) after <think>.Source: §4.2 · p. 7, §4.1 · p. 7
- SFT demonstrations are generated from OpenR1-Math queries and filtered for correctness and template validity, with 20K SHORT and 20K LONG instances sampled for training.RPO uses 16 rollouts per prompt, split 8/8 between SHORT and LONG, updates only the root gating token, and runs fewer than 60 update steps.Source: §4.2 · p. 7
4.3 Main results
- On Qwen3-8B, Avg Think% falls from 57.3% to 44.8% while Avg Pass@1 changes from 63.8% to 62.0%, and RAR improves from 0.90 to 0.72.On Exaone4-1.2B, Think% drops from 43.4% to 23.6% and RAR improves from 0.85 to 0.49.Source: §4.3 · p. 8
- RPO's refinement varies by configuration: Qwen3-1.7B and R1-1.5B recover accuracy by reallocating LONG toward harder queries, while Qwen3-8B and Exaone4-1.2B trade small accuracy reductions for further Think% compression.For Qwen3-1.7B, Think% is 74.0 and RAR is 1.43 after RPO, versus 72.8 and 1.45 after SFT.Source: §4.3 · p. 8
- On R1-1.5B, RPO routes 20.5% to LONG and reaches 51.3 two-task average accuracy with RAR=0.96; AdaptThink reaches Avg P@1=56.5 and RAR=0.74 at a larger training budget.The paper frames these as distinct budget–accuracy operating points, not a head-to-head ranking; matched 8B accounting puts RPO at about 0.02× AdaptThink's training compute.Source: §4.3 · p. 8, Table 4 · p. 8
5 Analysis
The analysis treats RPO checkpoints and thresholds as adjustable routing choices, while examining results beyond math and the behavior of root confidence. The accuracy–routing frontier and calibration signal depend on the evaluation setting.
Controlling the degree of routing
- Intermediate RPO checkpoints provide multiple feasible accuracy–routing operating points, so a checkpoint can be selected for a compute or routing budget while maintaining accuracy.The achievable frontier depends on the RL prompt set and easy/hard mix, making checkpoint selection dataset-conditioned rather than universal.Source: §5 · p. 8
LiveCodeBench and GPQA
- On Qwen3-32B, SFT already yields Avg Think%=94.6 and RAR=1.52; RPO gives a comparable RAR of 1.55 rather than a strictly better one.Across coding and scientific QA, RPO's effects vary with model size, and the paper recommends treating the post-RPO checkpoint as a controllable router with threshold sweeps.Source: §5 · p. 8, Table 6 · p. 9
Calibration of root confidence
- Figure 6 bins root SHORT confidence into 10 uniform intervals and plots empirical accuracy against predicted confidence, with the diagonal marking perfect calibration.Root confidence tends to correlate with correctness, but the signal is weak and dataset-dependent; AIME25 concentrates at low confidence while MATH500 has a broader spread of high-confidence bins.Source: §5 · pp. 8–9, Figure 6 · p. 9
6 Conclusion
RPO makes adaptive thinking a lightweight routing problem by updating only the initial gating token and treating the rest of the rollout as environment output. Across model families and scales it improves the accuracy–routing trade-off in most scenarios, but some configurations yield comparable rather than dominating frontiers.
6 Conclusion
- RPO updates only the initial gating token and treats the remainder of the rollout as environment output, reducing training overhead during GRPO.Across multiple model families and scales, it achieves lower reasoning-per-accuracy ratios over math, scientific reasoning, and coding in most scenarios.Source: §6 · p. 9
- Some configurations trace a comparable but not strictly dominating frontier, supporting RPO as a controllable routing layer rather than a uniform replacement for SFT.Source: §6 · p. 9
Limitations
The paper identifies limits in benchmark scope, confidence calibration, routing granularity, and configuration-dependent gains. Appendix experiments also examine whether root-level gating works across tokenizations and alternative markers.
Limitations
- RPO is limited to binary routing and does not optimize within-mode behavior such as how long to think or when to stop.The authors also report that root confidence can be sensitive to bucket composition and weakly calibrated to query difficulty.Source: Limitations · p. 9
- Broader evaluation is still needed on advanced benchmarks requiring agentic behavior with tools.Source: Limitations · p. 9
- The reported routing behavior is configuration-dependent: Qwen3-1.7B and Qwen3-32B have comparable rather than strictly better post-SFT operating points.For Qwen3-32B, the paper notes LoRA tuning and a near-baseline conservative SFT router.Source: Limitations · p. 9
A.1 Root-level gating is tokenization agnostic
- For R1-Distill-Qwen, arbitrary instruction prefixes distinguish reasoning and non-reasoning modes while preserving reasoning accuracy with the reasoning prefix.The authors report testing different 20–30-token sentences and applying RPO updates only to prefix tokens.Source: §A.1 · p. 11
A.2 Alternative root token markers
- The marker check compares three root-token marker types on Qwen3-8B; later newline tokens do not affect the routing decision.Gradients are masked for all non-root tokens.Source: §A.2 · p. 11
A.3 What generalizes and what does not
RPO’s selective-compute routing can transfer beyond math, but accuracy and calibration remain dataset-dependent. SFT hyperparameters and rollout construction also vary across model families and modes.
A.3 What generalizes and what does not
- RPO can reduce LONG usage on scientific and coding benchmarks, but accuracy gains are not guaranteed because optimal routing is dataset-dependent.Calibration also does not transfer perfectly, and checkpoint or threshold selection may require dataset-conditioned calibration.Source: §A.3 · pp. 11–12
B.1 Hyperparameter optimization for SFT
- Qwen3 and EXAONE4 required different SFT settings to preserve accuracy and token efficiency across model sizes and generation modes.Qwen3-1.7B generated longer reasoning sequences while Qwen3-8B maintained token counts; presence penalty 1.5 helped preserve non-reasoning token counts, while EXAONE4 used temperature 0.4.Source: §B.1 · p. 12
B.2 Mode-specific GRPO rollout strategies
- For Qwen3, each mode uses a distinct single-token prefix, while EXAONE4 uses a two-token non-reasoning prefix and a different rollout enforcement strategy.EXAONE4 initially uses the reasoning prefix for all samples, then appends a token for non-reasoning rollouts or masks the initial reasoning token for reasoning rollouts.Source: §B.2 · p. 12
B.3 SFT Dataset and training
The SFT stage creates mode-controllable models, while RPO further adjusts routing and compute-accuracy trade-offs. Training-data volume and SFT duration are tuned to avoid degrading solver performance.
B.3 SFT Dataset and training
- Qwen3 and EXAONE4 self-distillation produces reasoning and non-reasoning trajectories, while R1-Distill uses OpenR1-Math solutions and distilled non-thinking trajectories.For Qwen3 and EXAONE4, approximately 20K examples per mode yield about 40K training samples per model.Source: §B.3 · pp. 12–13
B.4 Dataset sizes for LoRA updates
- LoRA SFT uses 4K queries per mode, 8K total, rather than the 20K per mode used for full fine-tuning.The reduced dataset teaches mode controllability; the RPO stage retains the main experiments’ rollout and reward settings.Source: §B.4 · p. 13
B.5 SFT stop criterion and sensitivity
- SFT continues only while default-template accuracy remains within tolerance ϵ of the original model, because over-training can degrade accuracy and routing quality.The paper notes that template drift can produce high Think% alongside low performance.Source: §B.5 · p. 13
- For R1-Distill-Qwen-1.5B, MATH500 accuracy fell from 83.2 to 75.1 after 1 SFT epoch, while 3 epochs reached 77.6; for 7B, 1 epoch scored 90.4 versus 88.0 at 2 epochs.These forced-LONG results indicate that the optimal SFT duration is model-dependent and motivate a tolerance-based stopping rule.Source: §B.5 · p. 14, Table 10 · p. 14
C.1 LoRA updates
- RPO further shifts Qwen3’s root gating after LoRA SFT, often reducing Think% while keeping task performance non-trivial across math and non-math benchmarks.The resulting compute-accuracy operating point depends on model size; the reported LoRA SFT variant is not directly comparable to full-parameter SFT results.Source: §C.1 · p. 14
C.2 Robustness to hard forcing of thinking after post-training
Under equal compute, root-only RPO achieves substantially lower Think% than full-trajectory GRPO with comparable or slightly lower accuracy. Post-training preserves both forced modes, and threshold routing and paired labeling support controlled evaluation.
C.2 Robustness to hard forcing of thinking after post-training
- Post-trained models remain functional when forced into either mode: LONG preserves much of always-think baseline accuracy, while SHORT retains non-trivial performance with fewer tokens.The evaluation covers math and non-math benchmarks and interprets SFT/RPO as reshaping routing rather than collapsing either mode into brittle specialization.Source: §C.2 · p. 14
C.3 Budget-matched comparison: root-only vs. full-trajectory RL
- With equal PFLOPs, full-trajectory GRPO remains at 90.2% Think%, while RPO reaches 20.5–43.4% Think with comparable or slightly lower accuracy.The comparison is on MATH500 with R1-Distill-Qwen-1.5B; the authors describe root-only updates as more sample-efficient for routing under low-budget regimes.Source: §C.3 · p. 14, Table 11 · p. 14
C.4 Post-RPO threshold routing
- Raising threshold τ routes more queries to LONG and increases accuracy at the cost of tokens; lowering it saves compute with a controlled accuracy reduction.The reported threshold sweep demonstrates a smooth, monotonic compute-accuracy trade-off after RPO.Source: §C.4 · pp. 14–15, Table 12 · p. 15
C.5 Automatic labeling quality
- Paired SHORT and LONG evaluation assigns queries to four correctness buckets, and LONG-only-correct cases dominate on hard queries while SHORT-only-correct cases remain few.This paired design avoids treating chance-correct SHORT answers as evidence that LONG is required.Source: §C.5 · p. 15, Table 14 · p. 15
C.6 Misrouting analysis
Counterfactual evaluation identifies whether routing errors arise from the gate or from solver limits. False-LONG is the dominant error type, while lowering the threshold increases false-SHORT errors.
C.6 Misrouting analysis
- Counterfactual evaluation tests the opposite forced mode for each adaptive error, separating harmful false-SHORT and wasteful false-LONG misroutes from cases where both modes fail or succeed.Both-wrong indicates a solver limitation; both-correct is a compute-saving opportunity rather than an error.Source: §C.6 · p. 15
- False-LONG over-routing dominates MATH500 misroutes; lowering τ from 0.75 to 0.70 adds 12 false-SHORT errors while reducing false-LONG errors by only 4.The authors suggest a pretraining prior associating more tokens with trying harder as a plausible source of conservative routing.Source: Table 15 · p. 16, §C.6 · p. 16
- Representative cases isolate gate errors from solver limitations, including a divisor correction omitted by SHORT and an unnecessary LONG derivation that drifts to a sibling identity.The table pairs the adaptive mode and outcome with the counterfactual result under the opposite forced mode.Source: §C.6 · p. 16, Table 16 · p. 16
- Under pass@10 evaluation, the adaptive policy reduces token cost while maintaining competitive LiveCodeBench pass@10 performance.Source: Table 17 · p. 16
C.7 Pass@k evaluation on coding tasks
The coding evaluation considers pass@10 alongside token cost, while the AI-use disclosure distinguishes tool assistance from the authors’ core contributions.
C.7 Pass@k evaluation on coding tasks
- The adaptive policy often matches or improves baseline pass@10 while reducing average token cost on LiveCodeBench.The paper presents this trade-off as relevant to coding deployment where multiple samples are drawn.Source: §C.7 · p. 16
- The paper frames pass@10 alongside token cost as a trade-off view for coding deployment with multiple samples.Source: §C.7 · p. 16
D AI Usage
- ChatGPT and Gemini supported experiment code generation and sentence-level text refinement; the authors wrote the core idea, framework design, and most of the manuscript.Source: D AI Usage · p. 16