ROI-Reasoning: Rational Optimization for Inference via Pre-Computation Meta-Cognition

Muyang Zhao; Chunhao Tian; Qi Qi; Hao Sun

cs.AIJan 7, 2026

ROI-Reasoning: Rational Optimization for Inference via Pre-Computation Meta-Cognition

Authors: Muyang Zhao, Chunhao Tian, Qi Qi, Hao Sun

Organizations: Gaoling School of Artificial Intelligence, Renmin University of China, Beijing, China

Abstract

Large language models (LLMs) can achieve strong reasoning performance with sufficient computation, but they do not inherently know how much computation a task requires. We study budgeted inference-time reasoning for multiple problems under a strict global token constraint and formulate it as an Ordered Stochastic Multiple-Choice Knapsack Problem (OS-MCKP). This perspective highlights a meta-cognitive requirement---anticipating problem difficulty and allocating computation strategically. We propose ROI-Reasoning, a two-stage framework that endows LLMs with intrinsic, budget-aware rationality. First, Meta-Cognitive Fine-Tuning teaches models to predict reasoning cost levels and decide whether a problem is worth attempting before generation. Next, Rationality-Aware Reinforcement Learning optimizes sequential decision making under a hard token budget, allowing models to learn long-horizon allocation strategies. Experiments across mathematical reasoning, multiple-choice question answering, and code generation show that ROI-Reasoning improves scores and generally reduces reference-order regret under tight budgets, with up to a 417% relative score improvement over the strongest same-backbone baseline in mathematics.

Figures & tables

Method	Budget Control	Cost Estimate	Problem Skipping	Cross-Task Policy
Self-Refine	✗	✗	✗	✗
Budget Forcing	✓	✗	✗	✗
TALE-EP	✓	✓	✗	✗
Plan-and-Budget	✓	✓	✗	✗
ROI-Reasoning	✓	✓	✓	✓

Table 1: Capabilities relevant to reasoning across multiple problems under a shared token budget.

Figure 1: Overview of ROI-Reasoning. The top row illustrates the exam analogy; the bottom row shows MFT cost-level supervision and RARL policy optimization under a shared token budget.

Large-Scale Models
Method	Score easy		Medium				Hard
	1024	512	1024		512		1024		512
	1024	512	Score $\uparrow$	Reg. $\downarrow$	Score $\uparrow$	Reg. $\downarrow$	Score $\uparrow$	Reg. $\downarrow$	Score $\uparrow$	Reg. $\downarrow$
DeepSeek-V3.2 (685B)	1.85	1.53	1.33	0.28	0.88	0.42	0.87	0.53	0.49	0.68
GPT-4o-mini	1.53	1.19	1.41	0.08	0.65	0.45	1.27	0.17	0.32	0.73
Open-Source Models (1.5B–8B)

Table 2: Main results under shared 1024- and 512-token budgets. We report the easy-to-hard reference score (Score easy ), Score, and Reference-Order Regret. Each test paper contains three problems and has a maximum score of 3. For the Medium and Hard settings, the best and second-best values in each Score and Regret column are shown in bold and underlined , respectively.

Panel of Figure 2 — Figure 2: Token-length distributions by problem position under different difficulty levels and token budgets. The histograms compare the reasoning length patterns of Qwen2.5-1.5B-Instruct , Qwen2.5-1.5B-Instruct + MFT , and Qwen2.5-1.5B-Instruct + MFT + RARL under four distinct settings.

Variant	Stage 1	Stage 2	Medium				Hard
			1024		512		1024		512
			Score $\uparrow$	Reg. $\downarrow$	Score $\uparrow$	Reg. $\downarrow$	Score $\uparrow$	Reg. $\downarrow$	Score $\uparrow$	Reg. $\downarrow$
ROI-Reasoning	MFT	RARL	1.13	0.02	0.97	0.11	1.12	0.02	0.93	0.14
w/o RARL	MFT	–	0.98	0.03	0.82	0.09	0.95	0.04	0.81	0.10
w/o RARL, w/ Greedy	MFT	Greedy	0.97	0.04	0.83	0.15	0.96	0.04	0.83	0.15
w/o Cost Tags	RFT	Dr. GRPO	1.21	0.01	0.80	0.20	1.19	0.03	0.68	0.32

Table 3: Component ablation on Qwen2.5-1.5B-Instruct. MFT supervises cost levels; RARL rewards correct answers with cost-level agreement. RFT and Dr. GRPO use correctness-only objectives, while Greedy replaces RARL with external allocation. Bold and underlined values indicate the best and second-best results per column.

Figure 3: Composition of predicted-level outcomes (correct, under-, and over-estimated) under different budgets on Medium and Hard test papers, comparing MFT and MFT+RARL.

Figure 4: Cross-domain Score and Regret for same-backbone methods on MCQA and code generation under shared 512- and 256-token budgets.

Unseen $N$
$N/B$	Score (%, Medium / Hard)
$N/B$	Base	MFT	Ours
$4/1024$	18.66/14.87	32.91/30.37	39.49/37.23
$4/1536$	22.72/19.91	33.86/31.36	39.85/37.50
$5/1024$	16.37/12.64	31.71/27.82	40.56/35.39
$5/1536$	20.66/17.80	33.29/29.10	41.04/36.16

Table 4: Generalization to unseen sequence lengths and token budgets under uniform per-problem scoring.

N/B

denotes sequence length/token budget, and normalized scores are reported as Medium/Hard.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

MFT setting	Value
Backbone	Qwen2.5-1.5B-Instruct
Training	full-parameter SFT with LLaMA-Factory
Hardware	$4\times$ NVIDIA GeForce RTX 4090
Precision	bf16
Optimizer	AdamW
Learning rate	$5\times 10^{-6}$

Appendix

Table 5: Main-experiment MFT configuration.

RARL setting	Value
Initialization	final MFT checkpoint
Framework / rollout engine	verl / vLLM
Hardware	$4\times$ NVIDIA RTX PRO 6000 (one node)
Advantage estimator	Dr. GRPO
KL regularization / entropy coefficient	none / 0.0
PPO clipping coefficient / epochs	0.2 / 1

Appendix

Table 6: Main-experiment RARL configuration.

Panel of Figure 5 — Figure 5: Structure of the ROI-Reasoning input prompt.

Method	Budget	Medium				Hard
Method	Budget	P1	P2	P3	Score	P1	P2	P3	Score
DeepSeek-V3.2 (685B)	1024	0.60	0.45	0.28	1.33	0.32	0.29	0.26	0.87
	512	0.48	0.29	0.11	0.88	0.19	0.17	0.13	0.49
GPT-4o-mini	1024	0.44	0.54	0.43	1.41	0.18	0.54	0.55	1.27
	512	0.38	0.23	0.04	0.65	0.14	0.15	0.03	0.32
Llama-3.1-8B-Instruct	1024	0.26	0.27	0.21	0.74	0.09	0.20	0.26	0.55

Appendix

Table 7: Per-problem accuracy and overall score under different token budgets. We report accuracy for each problem position (P1–P3) and the resulting total score on Medium and Hard test papers.

Large-Scale Models
Method	Score easy		Medium				Hard
	512	256	512		256		512		256
	512	256	Score $\uparrow$	Reg. $\downarrow$	Score $\uparrow$	Reg. $\downarrow$	Score $\uparrow$	Reg. $\downarrow$	Score $\uparrow$	Reg. $\downarrow$
DeepSeek-V3.2 (685B)	2.04	1.80	1.78	0.13	1.41	0.22	1.35	0.34	0.80	0.56
GPT-4o-mini	1.55	1.31	1.30	0.16	0.94	0.28	0.93	0.40	0.51	0.61
Open-Source Models (1.5B–8B)

Appendix

Table 8: Performance comparison on Multiple-Choice Question Answering under different difficulty levels and token budgets. We report Score (maximum

=3

), Score easy , and Reference-Order Regret on Medium and Hard test papers, each consisting of three questions, under 512-token and 256-token constraints. The best and second-best values in each Score and Regret column are highlighted in bold and underlined , respectively.

Paper Type	Accuracy(%)			Average Tokens
Paper Type	Easy	Medium	Hard	Easy	Medium	Hard
Easy Paper	38.54	27.34	14.36	49.2	46.0	44.6
Medium Paper	53.82	38.35	10.28	56.1	54.4	40.4
Hard Paper	61.55	46.63	9.94	58.6	58.6	37.9

Appendix

Table 9: Behavioral Analysis of MFT across different test paper orderings. By identifying and strategically curtailing effort on Hard questions early, the model reallocates saved tokens to dramatically boost the accuracy of Easy and Medium questions.

Large-Scale Models
Method	Score easy		Medium				Hard
	512	256	512		256		512		256
	512	256	Score $\uparrow$	Reg. $\downarrow$	Score $\uparrow$	Reg. $\downarrow$	Score $\uparrow$	Reg. $\downarrow$	Score $\uparrow$	Reg. $\downarrow$
DeepSeek-V3.2 (685B)	2.03	1.81	1.96	0.03	1.64	0.09	1.85	0.09	1.57	0.13
GPT-4o-mini	1.61	1.21	1.59	0.01	1.10	0.09	1.77	-0.10	1.45	-0.20
Open-Source Models (1.5B–8B)

Appendix

Table 10: Performance comparison on Code Generation tasks under different difficulty levels and token budgets. We report Score (Pass@1, maximum

=3

), Score easy , and Reference-Order Regret on Medium and Hard test papers under 512-token and 256-token constraints. The best and second-best values in each Score and Regret column are highlighted in bold and underlined , respectively.

Large-Scale Models
Method	Score easy		Medium				Hard
	1024	512	1024		512		1024		512
	1024	512	Score $\uparrow$	Reg. $\downarrow$	Score $\uparrow$	Reg. $\downarrow$	Score $\uparrow$	Reg. $\downarrow$	Score $\uparrow$	Reg. $\downarrow$
DeepSeek-V3.2 (685B)	51.35	39.77	39.47	0.23	25.49	0.36	27.10	0.47	14.94	0.62
GPT-4o-mini	40.48	28.96	37.77	0.07	18.59	0.36	35.04	0.13	10.43	0.64
Open-Source Models (1.5B–8B)

Appendix

Table 11: Performance under heterogeneous score on mathematical reasoning. AIME, MATH, and GSM8K problems are assigned 3, 2, and 1 points, respectively, and scores are normalized to a 100-point scale. We report Score, Score

{}_{\text{easy}}

, and Reference-Order Regret on Medium and Hard test papers under 1024-token and 512-token budgets. Best and second-best Score and Regret values are highlighted with bold and underlining .

Open-Source Models (1.5B–7B)
Method	Medium						Hard
	1024			1536			1024			1536
	Score $\uparrow$	Score easy	Regret $\downarrow$	Score $\uparrow$	Score easy	Regret $\downarrow$	Score $\uparrow$	Score easy	Regret $\downarrow$	Score $\uparrow$	Score easy	Regret $\downarrow$
Qwen2.5-Math-7B-Instruct	30.23	54.39	0.44	45.81	61.26	0.25	20.43	50.80	0.60	38.91	58.07	0.33
Qwen2.5-7B-Instruct	40.66	55.12	0.26	53.02	57.42	0.08	32.65	51.65	0.37	48.55	55.07	0.12
Qwen2.5-Math-1.5B-Instruct	29.04	47.90	0.39	39.18	52.55	0.25	21.38	45.08	0.53	32.27	49.95	0.35

Appendix

Table 12: Unseen-

N

generalization for

N=4

after training with

N=3

. Scores are normalized to a 0–100 scale. We report Score, Score easy , and Reference-Order Regret on Medium and Hard papers under 1024- and 1536-token budgets. Bold and underlined values denote the best and second-best Score and Regret, respectively.

Open-Source Models (1.5B–7B)
Method	Medium						Hard
	1024			1536			1024			1536
	Score $\uparrow$	Score easy	Regret $\downarrow$	Score $\uparrow$	Score easy	Regret $\downarrow$	Score $\uparrow$	Score easy	Regret $\downarrow$	Score $\uparrow$	Score easy	Regret $\downarrow$
Qwen2.5-Math-7B-Instruct	26.01	51.97	0.50	41.28	61.44	0.33	16.52	44.12	0.63	32.02	54.94	0.42
Qwen2.5-7B-Instruct	37.85	55.37	0.32	52.79	59.19	0.11	28.77	48.74	0.41	46.47	54.40	0.15
Qwen2.5-Math-1.5B-Instruct	23.91	44.08	0.46	33.94	50.62	0.33	17.12	38.55	0.56	26.70	46.42	0.42

Appendix

Table 13: Unseen-

N

generalization for

N=5

after training with

N=3

Explore similar work

Avoiding Overthinking and Underthinking: Curriculum-Aware Budget Scheduling for LLMs
Mar 29, 2026Amirul Rahman, Aisha Karim, Kenji Nakamura +1Efficient Language Model Reasoning Inference-Time Compute Allocation
University of Malaya
The Shadow Price of Reasoning: Economic Perspective on Optimal Budget Allocation for LLMs
Jun 2, 2026Xu Wan, Speed Zhu, Jianwei Cai +4Cost-Aware Inference Inference-Time Scaling
1Zhejiang University · 2Tecent HY Team · 3Peking University
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

Mar 29, 2026J/K move · Enter open · S save

Avoiding Overthinking and Underthinking: Curriculum-Aware Budget Scheduling for LLMs

Amirul Rahman, Aisha Karim, Kenji Nakamura, Yi-Fan Ng

University of Malaya

Efficient Language Model Reasoning Inference-Time Compute Allocation Overthinking in Language Models RL for Language Model Reasoning Numerical Reasoning in Language Models

Open paper AskPDF arXiv

Scaling test-time compute via extended reasoning has become a key paradigm for improving the capabilities of large language models (LLMs). However, existing approaches optimize reasoning under fixed or uniformly sampled token budgets, ignoring the fundamental mismatch between problem difficulty and allocated compute. This leads to overthinking on easy problems and underthinking on hard ones, resulting in suboptimal token efficiency across diverse reasoning scenarios. In this paper, we propose Budget-Adaptive Curriculum Reasoning (BCAE), a unified framework that jointly optimizes reasoning quality and token efficiency through three synergistic components: (1) a budget-conditioned unified policy that embeds the token budget as a continuous conditioning signal, eliminating the need for decoupled thinking and summarization strategies; (2) a curriculum-aware budget scheduler that adaptively shifts the training budget distribution from easy to hard problems based on real-time learning progress; and (3) a truncation-aware dense reward mechanism that provides fine-grained credit assignment at intermediate reasoning steps via process-level verification. We further introduce Budget-Conditioned Advantage Estimation (BCAE), a novel variance reduction technique that conditions the advantage baseline on the sampled budget, yielding more stable policy gradients. Experiments on mathematical reasoning benchmarks (MATH, GSM8K, AIME, and Minerva Math) demonstrate that BACR consistently outperforms other strong baselines across all token budgets, achieving up to 8.3% accuracy improvement under tight budgets while reducing average token consumption by 34% compared to unconstrained reasoning.

Abstract

Explore similar work

Avoiding Overthinking and Underthinking: Curriculum-Aware Budget Scheduling for LLMs

The Shadow Price of Reasoning: Economic Perspective on Optimal Budget Allocation for LLMs

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