Efficiently Learning To Reason or Not to Reason: Root-token Policy Optimization for Adaptive Thinking

Taehyeon Kim; Hyunsoo Lee; Youngsoo Jang; Moontae Lee

cs.CLOct 13, 2026

Efficiently Learning To Reason or Not to Reason: Root-token Policy Optimization for Adaptive Thinking

Authors: Taehyeon Kim, Hyunsoo Lee, Youngsoo Jang, Moontae Lee

Organizations: LG AI Research · UNIST · University of Illinois Chicago

Abstract

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.

Figures & tables

Figure 1: Overview of our training pipeline. We generate paired reasoning samples and non-reasoning samples from input queries through the original reasoning model. These samples are self-distilled to manifest adaptive reasoning. Finally, root-token policy optimization ( RPO ) is applied to the model, refining the policy by updating only the initial branching token (i.e., root-token) to improve difficulty-aware adaptive reasoning.

Figure 2: Detailed overview of our training framework. After supervised finetuning (SFT) with self-generated correct samples, RPO involves updates only the initial decision token with group normalized rewards and advantages.

Figure 3: Thresholding on a Short -mode confidence signal yields a smooth routing trade-off about probability threshold vs number of routed instances over MATH500 dataset hendrycks2021measuring .

Method ( $\tau$ )	Short	Long	Overall
Original	0.00 (0)	83.6 (500)	83.6 (500)
Adaptive (0.75)	93.6 (62)	82.4 (438)	83.8 (500)
Adaptive (0.70)	89.7 (155)	78.8 (345)	82.2 (500)

Table 1: Accuracy (%) and routed instance counts (in parentheses) under threshold routing over MATH500 dataset hendrycks2021measuring with R1-1.5B model.

Panel of Figure 4 — Figure 4: Token-length histograms (log-scaled frequency) for Short and Long mode.

Model	Method	Updated tokens $T$	KV / rollout (GiB)	Act(bp) / rollout (GiB)	Train+Opt (GiB)	Peak @ $B{=}8$ (GiB)
1.5B	Original GRPO	12K	2.38	13.18	22.35	146.86 (1.00 $\times$ )
	RPO	1	2.38	$1.10\times 10^{-3}$	22.35	41.40 (0.28 $\times$ )
	RPO with LoRA update	1	2.38	$1.10\times 10^{-3}$	0.07	19.12 (0.13 $\times$ )
7B	Original GRPO	12K	6.35	35.16	104.31	436.34 (1.00 $\times$ )
	RPO	1	6.35	$2.93\times 10^{-3}$	104.31	155.11 (0.36 $\times$ )
	RPO with LoRA update	1	6.35	$2.93\times 10^{-3}$	0.21	51.01 (0.17 $\times$ )

Table 2: Peak VRAM (GiB) for long-context GRPO with uniform context

S{=}13

K (1K prompt

{+}

12K generation) and

B{=}8

concurrent rollouts. RPO gains memory efficiency by updating the root token (

T{=}1

). LoRA further reduces trainable-state memory with fewer trainable parameters.

Model	Stage	MATH500			AIME24			AIME25			Avg
		P@1	#Tok	Think%	P@1	#Tok	Think%	P@1	#Tok	Think%	P@1	#Tok	Think%	RAR $\downarrow$
Qwen3-1.7B	Baseline ( Long )	91.1	4,539	100.0	45.6	15,569	100.0	38.1	15,830	100.0	58.3	11,979	100.0	1.72
	Baseline ( Short )	73.0	978	0.0	12.3	3,898	0.0	11.5	2,180	0.0	32.3	2,352	0.0	0.00
	+ SFT	85.5	3,767	69.8	34.4	12,558	72.9	30.4	12,648	75.6	50.1	9,658	72.8	1.45
	+ RPO	79.5	2,754	31.8	43.8	13,315	93.8	31.7	14,231	96.3	51.7	10,100	74.0	1.43
Qwen3-8B	Baseline ( Long )	96.6	4,777	100.0	74.4	14,184	100.0	64.8	16,950	100.0	78.6	11,970	100.0	1.27

Table 3: Baseline

\rightarrow

SFT mode selection

\rightarrow

RPO on math benchmarks. Pass@1 (P@1) is the average accuracy (%) over 16 runs. #Tok is average generated tokens. Think% is the fraction routed to Long . Avg columns are simple averages over (MATH500, AIME24, AIME25). RAR (

\downarrow

) is the reasoning-per-accuracy ratio, defined as

\mathrm{RAR}_{\text{avg}}=\frac{t_{\text{avg}}}{a_{\text{avg}}+\varepsilon}

, where

a_{\text{avg}}

and

t_{\text{avg}}

are Avg Pass@1 and Avg Think% as fractions. Lower is better.

Figure 5: Think% versus Pass@1 on Qwen3-8B with RPO . Each point indicates an intermediate checkpoint.

Method	MATH500		AIME24		Avg
	P@1	Think%	P@1	Think%	P@1	RAR $\downarrow$
Prior hybrid-reasoning baselines (reported in their tables)
Original ( Long )	83.2	100.0	29.4	100.0	56.3	1.78
Original ( Short )	67.2	0.0	14.0	0.0	40.6	0.00
Rejection Finetuning † Zhang2025AdaptThink	72.4	66.6	25.2	79.0	48.8	1.49
AdaptThink † Zhang2025AdaptThink	82.0	23.2	31.0	59.6	56.5	0.74

Table 4: Comparison with accuracy and Think% over R1-1.5B model.

\dagger

presents the performance from the original paper.

Method	SFT spec	GRPO spec	Forward PFLOPs	Backward PFLOPs	Total PFLOPs	Rel. Total (vs. AdaptThink)
AdaptThink Zhang2025AdaptThink	40K @ 16K	40K @ 16K, $K{=}16$ , $T_{\text{upd}}{\approx}16$ K	174,080	348,160	522,240	1.00 $\times$
ThinkLess Bao2025Thinkless	334K @ 24K	1K @ 24K, $K{=}8$ , $T_{\text{upd}}{\approx}24$ K	131,328	262,656	393,984	0.75 $\times$
RPO (Ours)	8K @ 16K	1K @ 16K, $K{=}16$ , $T_{\text{upd}}{=}1$	6,144	4,097	10,241	0.02 $\times$

Table 5: PFLOPs accounting for an 8B dense Transformer (full-length sequences at the stated context). We assume

2P

FLOPs/token for forward and

4P

for backward. For AdaptThink, the number of dataset is set to 40K as reported.

Model	Stage	LiveCodeBench			GPQA-Diamond			Avg
		P@1	#Tok	Think%	P@1	#Tok	Think%	P@1	#Tok	Think%	RAR $\downarrow$
Exaone4-1.2B	Baseline	42.2	12,159	100.0	50.8	7,145	100.0	46.5	9,652	100.0	2.15
	+ SFT	34.6	9,281	90.1	46.2	5,037	70.1	40.4	7,159	80.1	1.98
	+ RPO	36.7	6,692	61.7	41.7	2,881	35.0	39.2	4,787	48.4	1.26
Qwen3-8B	Baseline	49.5	14,878	100.0	58.8	7,355	100.0	54.2	11,117	100.0	1.85
	+ SFT	30.2	10,410	63.0	31.1	5,227	50.1	30.7	7,819	56.6	1.85

Table 6: LiveCodeBench/GPQA-Diamond results. Models without

\dagger

are trained with full-parameter SFT;

\dagger

indicates a LoRA-tuned model.

Panel of Figure 6 — Figure 6: Reliability diagrams guo2017calibration for the root confidence $p_{\mathrm{short}}(x){=}p_{\theta}(z_{S}{\mid}x,\texttt{<think>})$ . Each bin plots empirical accuracy against predicted confidence, with the diagonal indicating perfect calibration.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Stage	Math500			AIME24			AIME25			LiveCodeBench			GPQA
		P@1	#Tok	Think%	P@1	#Tok	Think%	P@1	#Tok	Think%	P@1	#Tok	Think%	P@1	#Tok	Think%
Qwen3-8B	Baseline (think)	96.9	5,547	100.0	75.8	15,602	100.0	68.3	18,419	100.0	49.5	14,879	100.0	58.8	5,838	100.0
	+ SFT	90.9	3,251	45.8	50.4	11,196	50.0	44.4	11,593	51.3	34.4	10,519	54.4	56.6	7,092	96.8
	+ RPO	88.0	2,412	27.8	45.2	9,437	38.1	40.6	9,650	39.6	28.7	8,362	40.1	57.1	6,908	96.2
Qwen3-32B	Baseline (think)	97.3	4,806	100.0	79.4	13,235	100.0	69.8	16,431	100.0	59.1	13,498	100.0	66.2	5,108	100.0
	+ SFT	92.1	2,602	42.7	59.8	9,309	56.3	54.2	10,379	58.3	58.3	11,457	90.6	66.3	4,822	98.5

Appendix

Table 7: LoRA-based mode control and RPO on Qwen3. For each model size, we compare the <think> -forced baseline ( Long ) with LoRA SFT for mode selection and a subsequent RPO update. We report Pass@1 (%), average generated tokens, and Think % (fraction routed to Long ). We set batch size 64 and learning rate

2\times 10^{-7}

, and use the Qwen default prompt for math tasks. This table reports a LoRA-tuned variant of Qwen3-8B/32B and is therefore not directly comparable to the full-parameter SFT Qwen3-8B in Tab. 6 ; baseline numerical drift across the two tables reflects independent decoding runs under matched settings.

Marker type	MathAvg P@1	Think%	LCB P@1	GPQA P@1
Newline fork (default)	62.0	44.8	35.4	55.4
Prefix sentence A	61.8	45.1	35.1	55.2
Prefix sentence B	61.5	46.3	34.8	54.9

Appendix

Table 8: Comparison of root token marker types on Qwen3-8B after + RPO . Newline fork uses \n vs. \n\n . Prefix sentence A/B use distinct instructional prefixes (

\sim

20–30 tokens) prepended before <think> to distinguish modes.

Backbone	Checkpoint	Mode	Math Avg		Non-math Avg
			P@1	Think%	P@1	Think%
Qwen3-8B	Pre	forced-think (LONG)	78.6	100.0	54.2	100.0
	Pre	forced-no-think (SHORT)	43.4	0.0	34.9	0.0
	+SFT	adaptive	63.8	57.3	30.7	56.6
		forced-think (LONG)	78.1	100.0	52.0	100.0
		forced-no-think (SHORT)	43.4	0.0	34.0	0.0

Appendix

Table 9: Robustness to forcing thinking vs. non-thinking after post-training. Math Avg is the mean P@1 over MATH500/AIME24/AIME25. Non-math Avg is the mean P@1 over LiveCodeBench and GPQA-Diamond. Think% is the fraction routed to thinking (100% and 0% under forced modes).

Model	Stage	MATH500 P@1	AIME24 P@1
R1-Distill-1.5B	Baseline	83.2	29.4
	SFT 1 ep	75.1	20.8
	SFT 3 ep	77.6	22.3
R1-Distill-7B	Baseline	93.5	54.2
	SFT 1 ep	90.4	43.5
	SFT 2 ep	88.0	44.5

Appendix

Table 10: SFT sensitivity to training epochs on R1-Distill models under forced-LONG evaluation.

Setting	Update scope	PFLOPs	MATH500 P@1	Think%
Full GRPO (budget-matched)	all tokens	$\sim$ 10K	80.8	90.2
RPO (ckpt A)	root-only	$\sim$ 10K	77.6	20.5
RPO (ckpt B)	root-only	$\sim$ 10K	80.4	43.4

Appendix

Table 11: Budget-matched comparison on MATH500 (R1-Distill-1.5B). Under equal PFLOPs (

\sim

10K), full-trajectory GRPO becomes under-trained and achieves only modest routing (Think%

{=}90.2\%

), while RPO achieves comparable accuracy with substantially lower Think%. Different RPO checkpoints trace a Pareto frontier between accuracy and compute.

$\tau$	Think%	MathAvg P@1	Non-mathAvg P@1
0.90	60.0	63.0	46.0
0.80	52.3	62.5	45.8
0.70	44.8	62.0	45.4
0.60	38.1	60.8	44.2

Appendix

Table 12: Threshold routing (

\tau

) sweep on the final + RPO checkpoint (Qwen3-8B).

\tau

controls the fraction routed to Long .

Difficulty	LONG-only	SHORT-only	Both-correct	Both-wrong
Easy ( $\mathrm{acc}_{S}\geq 0.75$ )	3.2%	1.8%	89.5%	5.5%
Medium	28.4%	4.1%	42.3%	25.2%
Hard ( $\mathrm{acc}_{S}<\mathrm{acc}_{L}\leq 0.5$ )	45.7%	2.3%	8.6%	43.4%

Appendix

Table 13: Outcome bucket distribution on the RPO training set (Qwen3-8B).

Bucket	Query (abbrev.)	Mechanistic rationale
LONG-only	“Find the number of positive integers $n\leq 2024$ such that $n^{2}+n+1\equiv 0\pmod{7}$ .”	Requires enumerating $n\bmod 7$ (solutions: $n\equiv 2,4$ ), then counting via floor arithmetic. Short typically guesses a single residue or skips the residue scan and returns an off-by-one count.
SHORT-only	“Compute $\cos(\pi/3)$ .”	Closed-form recall ( $1/2$ ). Short returns the memorized value. Long occasionally derives via half-angle identities and conflates $\cos(\pi/6)$ with $\cos(\pi/3)$ , yielding $\sqrt{3}/2$ .
Both-correct	“Solve $\log_{2}x=5$ for $x$ .”	Single inverse step: $x=2^{5}=32$ . No branch points or memorization gaps; both modes converge.
Both-wrong	“Find the smallest positive integer $n$ such that the decimal expansion of $1/n$ has period exactly $23$ .”	Requires knowing period equals $\mathrm{ord}_{n}(10)$ when $\gcd(n,10){=}1$ and that $10^{23}-1$ admits a non-trivial prime factor. Beyond Qwen3-8B’s reliable recall in either mode.

Appendix

Table 14: Representative training-set queries per outcome bucket (Qwen3-8B). Queries are abbreviated; rationales identify the reasoning step that drives the bucket assignment.

$\tau$	False-SHORT	False-LONG	Both-wrong
0.75	4	77	0
0.70	16	73	0

Appendix

Table 15: Misrouting composition on MATH500 (R1-Distill-1.5B, 500 queries). Lowering

\tau

(0.75

\to

0.70) increases the number of false-SHORT errors (+12) while only slightly reducing false-LONG errors (

-

4), reflecting a clear compute-risk trade-off.

Category	Query (abbrev.)	Adaptive	Counterfactual
False-SHORT	“How many distinct arrangements are there of the letters of mississippi ?”	Short , ✗ (returns $11!$ )	Long , ✓ ( $11!/(4!\,4!\,2!)=34{,}650$ )
False-LONG	“Compute $\sum_{k=1}^{10}k$ .”	Long , ✗ (drifts to $\sum k^{2}$ formula)	Short , ✓ ( $55$ )
Both-wrong	“Find the smallest $n$ such that the decimal expansion of $1/n$ has period exactly $23$ .”	Long , ✗	Short , ✗
Both-correct	“Solve $3x+7=22$ for $x$ .”	Short , ✓	Long , ✓

Appendix

Table 16: Representative misrouting cases on MATH500 (R1-Distill-1.5B,

\tau{=}0.70

). Adaptive reports the mode chosen by the learned router and its correctness; Counterfactual reports the opposite forced mode’s outcome on the same query.

Model	Stage	pass@1	pass@10	Avg Tokens
Qwen3-32B	Baseline	59.1	75.3	13,498
	+ SFT	58.3	77.1	11,457
	+ RPO	54.4	72.8	10,623
Exaone4-1.2B	Baseline	42.2	62.9	12,159
	+ SFT	34.6	56.0	9,281
	+ RPO	36.7	58.2	6,692

Appendix

Table 17: LiveCodeBench pass@

k

results with token cost. Under pass@10 evaluation, the adaptive policy reduces token cost while maintaining competitive pass@10 performance.

Explore similar work

Learning When to Think: Adaptive Reasoning for Test-Time Compute Allocation
Aug 20, 2026Gijs Kassenaar, Zhao Yang, Vincent François-LavetEfficient Language Model Reasoning RL for Language Model Reasoning
Vrije Universiteit Amsterdam
EvoThink: Evolving Thinking in Large Reasoning Models via Self-Pruning and Aha-Moment Preference Optimization
Jul 22, 2026Xinbang Dai, Zheyu Xin, Huikang Hu +7Overthinking in Language Models Self-Training for Language Models
Southeast University · Noah’s Ark Lab
ADaPT: Token-Level Decoupling for Efficient Large Reasoning Models
Jun 18, 2026Tingyun Li, Zishang Jiang, Jinyi Han +8Inference-Time Optimization RL for Language Model Reasoning
School of Data Science, Fudan University · Shanghai Institute of Artificial Intelligence for Education, East China Normal University · College of Computer Science and Artificial Intelligence, Fudan University +1

Aug 20, 2026J/K move · Enter open · S save

Learning When to Think: Adaptive Reasoning for Test-Time Compute Allocation

Gijs Kassenaar, Zhao Yang, Vincent François-Lavet

Vrije Universiteit Amsterdam

Efficient Language Model Reasoning RL for Language Model Reasoning LLM Mathematical Reasoning

Open paper AskPDF arXiv

Reasoning language models trained with reinforcement learning typically operate under a fixed token budget rather than an explicitly adaptive one, which can lead to over-computation on easy problems and insufficient computation on difficult ones. We study whether a model can learn to allocate its own reasoning effort by choosing, as the first token of its response, one of three modes: NoThink (answer as quickly as possible), Short (brief reasoning), or Long (extended reasoning). The choice is learned inside Group Relative Policy Optimization (GRPO) with no separate router, through a shaped reward that makes each mode worthwhile at a different response length, together with hard per-mode token caps that keep the modes distinct. On a 1.5B distilled model trained on MATH, the three modes emerge without collapsing to a single choice, and the brief modes end up more accurate than Long, which shows that the router sorts problems by difficulty rather than at random. Averaged over three seeds, the resulting policy stays close to the base model's accuracy on the held-out MATH500 (

0.782

vs.\

0.796

) while cutting the mean response length from

4{,}796

2{,}811

tokens (a

41\%

reduction). Interestingly, it also transfers to other benchmarks without retraining, with the largest savings where problems are easier, with for instance 76% token reduction on GSM8K and at higher accuracy than the baselines at similar response length. In short, we build a reasoning model that adaptively chooses how much to reason for each problem.

Abstract

Explore similar work

Learning When to Think: Adaptive Reasoning for Test-Time Compute Allocation

EvoThink: Evolving Thinking in Large Reasoning Models via Self-Pruning and Aha-Moment Preference Optimization

ADaPT: Token-Level Decoupling for Efficient Large Reasoning Models