PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Muyang Zhao, Qi Qi, Hao Sun

arXiv:2601.03822v1cs.AI

TL;DR

Under a strict global token budget, LLMs must plan how to allocate reasoning across multiple problems rather than relying only on local problem-solving ability. ROI-Reasoning combines pre-generation cost and utility estimation with reinforcement learning for sequential allocation, and its combined approach improves scores and reduces regret on budgeted mathematical reasoning benchmarks.

  • Problem

    Under a fixed global token budget, LLMs need to allocate reasoning across multiple problems, since local problem-solving ability alone does not ensure global planning.

    Source: §1 · pp. 1–2

  • Method

    ROI-Reasoning combines Meta-Cognitive Fine-Tuning, which predicts cost and supports solve-or-skip decisions, with Rationality-Aware Reinforcement Learning for long-horizon allocation under a hard global budget.

    Source: §1 · p. 2, §4.3 · p. 5

  • Results

    The combined MFT+RARL model consistently improves overall score and reduces regret across budgeted mathematical reasoning benchmarks, including Medium and Hard settings.

    Source: §6 · p. 8, §5 · p. 8

  • Takeaways & Limitations

    Budgeted reasoning calls for anticipating difficulty and allocating computation across questions, not only solving each problem individually.

    Source: §1 · p. 1, §6 · p. 8

    The study covers budgeted mathematical reasoning with three problems per test paper and uses token count as a cost proxy, while its coarse difficulty levels and refusal protocol may not transfer directly to finer-grained or domain-specific settings.

    Source: §8 · pp. 8–9

Abstract

from arXiv · show

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 tasks under a strict global token constraint and formalize it as a Ordered Stochastic Multiple-Choice Knapsack Problem(OS-MCKP). This perspective highlights a meta-cognitive requirement -- anticipating task difficulty, estimating return over investment (ROI), and allocating computation strategically. We propose ROI-Reasoning, a two-stage framework that endows LLMs with intrinsic, budget-aware rationality. In the first stage, Meta-Cognitive Fine-Tuning teaches models to predict reasoning cost and expected utility before generation, enabling explicit solve-or-skip decisions. Next, Rationality-Aware Reinforcement Learning optimizes sequential decision making under a hard token budget, allowing models to learn long-horizon allocation strategies. Across budgeted mathematical reasoning benchmarks, ROI-Reasoning consistently improves overall score while substantially reducing regret under tight computation budgets.

1 Introduction

The introduction frames budgeted reasoning as a planning problem: models must allocate effort across questions, not merely solve them. It proposes ROI-Reasoning to anticipate the value of reasoning and manage computation under a global budget.

1 Introduction

  • Meta-cognitive planning requires anticipating difficulty, estimating the return of additional effort, and deciding when to persist, skip, or move on.This planning operates above local problem solving and governs how computation is spent.Source: §1 · p. 1
  • The paper frames multi-problem inference under a global token budget as an Ordered Stochastic Multiple-Choice Knapsack Problem.Reasoning effort has cost and reward, rewards appear only after generation, and sequential decisions constrain the remaining budget.Source: §1 · pp. 1–2
  • DeepSeek-V3.2 often spends too much effort on early exam questions, leaving insufficient budget for later ones.The authors present this behavior as evidence of inadequate global planning.Source: §1 · p. 2
  • ROI-Reasoning explicitly models the cost–benefit tradeoff of reasoning and trains models to consider it during inference.ROI refers to Return over Investment.Source: §1 · p. 2
  • Meta-Cognitive Fine-Tuning teaches pre-generation estimates of difficulty and expected utility, while Rationality-Aware Reinforcement Learning trains allocation across problems under a hard global token constraint.The two stages target anticipatory planning and adaptive decision-making, respectively.Source: §1 · p. 2

2 Related Work

Related work includes methods that improve reasoning or control computation, but generally focuses on individual tasks rather than allocating reasoning capacity across multiple tasks under one shared budget. The paper also connects this setting to predict-then-optimize methods and reinforcement learning for sequential combinatorial optimization.

2 Related Work

  • Chain-of-thought and reinforcement-learning methods can improve reasoning behavior, but their gains often come with longer outputs and greater inference-time computation.The passage names Chain-of-Thought work by Wei et al. (2022) and Kojima et al. (2022), and reinforcement-learning work by Guo et al. (2025), OpenAI (2024), Team et al. (2025), and Yang et al. (2025).Source: §2 · p. 2
  • Efficiency approaches include switching between thinking and no-thinking modes, anytime reasoning under truncation, and explicit control of reasoning length or token budgets.The cited examples are Zhang et al. (2025), Qi et al. (2025), Lin et al. (2025), Huang et al. (2025), Li et al. (2025), and Han et al. (2025).Source: §2 · p. 2
  • Existing adaptive reasoning approaches regulate computation on individual tasks, while this paper’s focus is allocation across multiple tasks under a shared budget.The cited approaches include uncertainty-based token allocation, adaptive controllers, generation-length policies, and predicted-budget conditioning.Source: §2 · p. 2
  • Predict-then-optimize methods estimate quantities such as costs or rewards before solving a downstream optimization problem based on those predictions.The passage cites Mandi et al. (2020) and Mulamba et al. (2020).Source: §2 · p. 2
  • Reinforcement learning is presented as suited to budgeted sequential decisions with large discrete action spaces, delayed rewards, and difficult cost-benefit modeling.The related-work discussion cites Mazyavkina et al. (2021) and Yang and Whinston (2023).Source: §2 · p. 2

3 Problem Formulation

The formulation models inference as sequential allocation of computation across problems revealed together but processed in a fixed order. Actions carry token costs and uncertain rewards, and the objective is to maximize expected reward under a global computation budget.

3 Problem Formulation

  • All problems are revealed upfront, but the model must answer them sequentially in a fixed order without revising earlier decisions.The model’s action for each problem may be to solve it or strategically abstain.Source: §3 · p. 3
  • Each action earns a non-negative reward based on its answer and incurs computational cost measured in generated tokens.Token use is observable during generation, whereas realized reward is revealed only after evaluation.Source: §3 · p. 3
  • The objective is to maximize expected total reward subject to a global computation budget.The expectation covers both problem instances and policy stochasticity.Source: §3 · p. 3
  • Each problem forms a knapsack class, actions are its alternatives, rewards and token costs act as values and weights, and the total token budget is capacity.Unlike classical MCKP, reward and cost are uncertain at decision time; unlike standard S-MCKP, the problem adds a fixed-order constraint.Source: §3 · p. 3
  • The problem is cast as an Ordered Stochastic Multiple-Choice Knapsack Problem, with sequential decisions under uncertainty and delayed rewards motivating a decomposition of the learning problem.The decomposition separates problem-level meta-cognitive assessment from allocation across the full problem order.Source: §3 · pp. 3–4
  • Because supervision is sparse and delayed, directly optimizing answer quality, ordering decisions, and total budget use end to end is difficult.This motivates separating assessment for an individual problem from global allocation across the sequence.Source: §3 · p. 4

4 Methodology

ROI-Reasoning combines pre-generation cost awareness with budgeted sequential learning to help models decide when to reason and how to allocate computation. Its methodology uses Meta-Cognitive Fine-Tuning (MFT) and Rationality-Aware Reinforcement Learning (RARL).

4.1 Overview

  • The overview frames ROI-Reasoning as a two-stage pipeline: MFT builds cost awareness and RARL trains budget allocation, with prompts exposing the global budget and requiring prereasoning prediction.The inference prompt also constrains the model's output format.Source: §4.1 · p. 4

4.2 Meta-Cognitive Fine-Tuning

  • MFT first teaches pre-generation difficulty and cost assessment, then combines Tag Alignment for cost-level predictions with Refusal Learning for strategic abstention.Training moves from single-problem instances to short multi-problem sequences; excessive-cost or repeatedly unsolved problems may receive a standardized NA answer.Source: §4.2 · pp. 4–5
  • A predicted difficulty level is a coarse proxy for expected computational cost, with each categorical level mapped to a predefined token range.The text gives Level-0 as short solutions within 256 tokens and Level-1 as 256–512 tokens.Source: §4.2 · p. 4

4.3 Rationality-Aware Reinforcement Learning

  • RARL trains models through interaction in simulated exam environments to improve problem-solving and long-horizon allocation across problems under a global token budget.It targets the long-horizon credit assignment that imitation-based MFT cannot fully capture.Source: §4.3 · pp. 5–6
  • The exam-instance objective rewards correct answers whose predicted cost levels match actual token use, while a hard global budget limits total generation.No intermediate or per-token rewards are used; the computational constraint is imposed through the hard budget.Source: §4.3 · p. 5
  • RARL optimizes complete exam-instance rollouts with Dr. GRPO, sampling groups of trajectories under the same budget constraint.Table 1 reports Score, Scoreeasy, and Regret on three-question Medium and Hard test papers under 1024- and 512-token constraints.Source: §4.3 · p. 5, Table 1 · p. 6

5 Experiments

The experiments test whether MFT and RARL improve budget-aware allocation across multi-problem mathematical reasoning tasks. The full pipeline achieves higher scores and lower regret across Medium and Hard settings, with analyses indicating adaptive token use and accurate difficulty prediction.

5 Experiments

  • Training uses GSM8K, MATH, and AIME 2022–2024 problems; evaluation uses their corresponding test sets and AIME25 in three-problem papers.Medium papers interleave difficulty levels, while Hard papers put more challenging problems earlier; evaluation uses 1024- and 512-token budgets.Source: §5 · p. 6
  • Figure 4 shows MFT+RARL shortening reasoning or abstaining earlier on harder problems, while Figure 5 analysis reports high accuracy in difficulty prediction.The token-length histograms compare the base model, MFT, and MFT+RARL across four settings; the error-composition figure compares RFT and RFT+RARL across budgets and paper difficulty.Source: Figure 4 · p. 7, §5 · p. 8, Figure 5 · p. 8
  • Baselines include large proprietary models, open-source instruction-tuned models, prompted reasoning scaffolds, and a greedy knapsack method based on MFT cost estimates.The greedy baseline attempts problems in ascending estimated cost order until its budget is exhausted.Source: §5 · p. 7
  • The primary metric is average total score over each paper’s three problems; approximate regret compares the original ordering with the same problems reordered from lower to higher expected difficulty.Exact best-order regret would require evaluating all N! permutations, so the paper uses the easier ordering as a reference.Source: §5 · pp. 7–8
  • The full MFT+RARL model consistently achieves higher scores and lower regret across both Medium and Hard settings.Advanced models GPT-4o-mini and DeepSeek-V3.2 do not exhibit clear advantages in this setting.Source: §5 · p. 8
  • Qwen2.5-Math-7B-Instruct does not consistently outperform its 1.5B counterpart, while Qwen2.5-7B-Instruct often performs better despite lacking specialized mathematical pretraining.A simple explicit prompt is competitive in several settings, underscoring the difficulty of internalizing meta-cognitive decisions without targeted supervision.Source: §5 · p. 8

6 Conclusion

The conclusion frames budgeted multi-task reasoning as a meta-cognitive allocation problem: models must anticipate difficulty and ROI while distributing limited computation. ROI-Reasoning combines MFT and RARL, and experiments report improved scores and reduced regret under tight budgets.

6 Conclusion

  • Budgeted reasoning across multiple tasks requires deciding how to solve a problem and whether it is worth attempting under a strict global token constraint.The paper highlights anticipating difficulty, estimating ROI, and allocating computation across questions to maximize final score.Source: §6 · p. 8
  • ROI-Reasoning combines MFT, which predicts cost levels and supports solve-or-skip decisions, with RARL for long-horizon planning under a hard global budget.Together, the stages form a two-stage framework for budget-aware reasoning.Source: §6 · p. 8
  • Across budgeted mathematical reasoning benchmarks, combining MFT and RARL consistently improves overall score and substantially reduces regret under tight budgets.Source: §6 · p. 8

7 Use of AI Assistants

The manuscript reports that ChatGPT was used only for proofreading and language refinement, not for its scientific content, results, or conclusions.

7 Use of AI Assistants

  • ChatGPT was used solely for proofreading and language refinement; the paper states that AI assistance did not influence its scientific content, experimental results, or conclusions.Source: §7 · p. 8

8 Limitations

The study is limited to budgeted mathematical reasoning under a fixed three-problem prompt, with token count used as a proxy for computational cost.

8 Limitations

  • Evaluation covers fixed three-problem math papers, and token count does not capture latency, tool overhead, verification, or memory/context costs.The authors also note that coarse difficulty levels and the NA refusal protocol may not transfer to finer-grained or domain-specific abstention and calibration.Source: §8 · pp. 8–9
  • Extending token-budget reasoning to heterogeneous agent tasks remains future work, with non-stationary or delayed rewards, stochastic tool outcomes, safety constraints, and task-success evaluation identified as challenges.The current formulation and results are framed as a step toward broader budget-aware meta-cognitive control.Source: §8 · p. 9

9 Ethics

The paper reports no significant work-specific ethical concerns and describes an MFT procedure with Tag Alignment and Refusal Learning stages.

9 Ethics

  • The work uses publicly available benchmarks, no human subjects, personal data, or sensitive information, and identifies no significant ethical concerns specific to the study.The authors encourage responsible deployment of budget-aware reasoning techniques in real-world systems.Source: §9 · p. 9

A Meta-Cognitive Fine-Tuning

  • MFT alternates single-problem instances with short multi-problem sequences to bridge atomic cost estimation and contextual decision-making.Its two stages are Tag Alignment and Refusal Learning.Source: A Meta-Cognitive Fine-Tuning · p. 11

Stage I: Tag Alignment

  • Tag Alignment trains difficulty prediction with a <predicted_level> tag and standardized reasoning format, using retained correct samples from GSM8K and MATH before sequential adaptation.The sequential adaptation retains correctly solved three-problem sequences with correctly formatted tags.Source: Stage I: Tag Alignment · p. 11

Stage II: Refusal Learning

  • Refusal Learning supervises standardized NA abstention for persistently unsolved problems, including in single-problem and three-problem settings.The single-problem data adds AIME 2022–2024; the resulting final MFT model is M3-Q-v2.Source: Stage II: Refusal Learning · p. 11

B Test Paper Construction

The appendix describes three-problem test-paper construction and contrasts prompt-based budgeting controls with a post-generation greedy selection baseline.

B Test Paper Construction

  • Each test paper is a single query with exactly three problems sampled from GSM8K, MATH, and AIME; Medium and Hard templates control problem ordering, and evaluation uses fixed global budgets of 1024 or 512 tokens.Hard templates favor AIME and MATH problems in earlier positions, while easy variants reorder the same three problems from easiest to hardest.Source: B Test Paper Construction · pp. 11–12

C.1 Simple and Explicit Prompt

  • The simple explicit prompt asks the model to solve sequentially under a global budget and output NA when giving up, without a dedicated planning stage.The prompt-based baselines share the evaluation protocol and decoding configuration but differ in how they elicit budget awareness.Source: C Baselines · p. 12, §C.1 · p. 12

C.2 Plan-and-Solve Prompt

  • Plan-and-Solve requires an upfront scan of all problems, a budget sufficiency estimate, and qualitative ROI-based prioritization before solving.The plan can designate problems to attempt or skip.Source: §C.2 · p. 13

C.3 Least-to-Most Prompt

  • Least-to-Most keeps execution in the presented order while allowing look-ahead planning and strategic skipping of costly early problems.Its design emphasizes that sequential allocation decisions are irreversible even when the model assesses all problems first.Source: §C.3 · p. 14

C.4 Greedy Knapsack Baseline

  • Greedy Knapsack predicts difficulty-based costs, generates candidate solutions first, then uses post-hoc selection to retain problems within budget and replace rejected answers with NA.The method does not regulate reasoning length during generation; outputs exceeding the budget may be hard-truncated.Source: §C.4 · pp. 14–15

D Extra Experiment

The extra experiment examines whether budget-allocation gains vary by problem position, particularly for later problems under tight budgets.

D Extra Experiment

  • The analysis records correctness separately for the first, second, and third problems to distinguish overall gains from position-related effects.It asks whether later-problem handling improves under tight budgets rather than accuracy rising uniformly across positions.Source: D Extra Experiment · p. 15
  • Table 2 reports P1–P3 accuracy and total score on Medium and Hard papers at 1024-token and 512-token budgets.The supplied passage describes the table’s metrics and scope, but gives no cell values or comparative outcome.Source: D Extra Experiment · p. 15

E Case Study: Meta-cognitive Resource Allocation

The case study shows ROI-Reasoning skipping a difficult geometry problem under a 512-token limit, then solving two simpler tasks and scoring 2 overall.

E Case Study: Meta-cognitive Resource Allocation

  • Faced with one complex geometry problem and two simpler tasks, ROI-Reasoning skips the first to stay within the strict 512-token budget.The case study presents this as intrinsic difficulty estimation and strategic decision-making.Source: E Case Study: Meta-cognitive Resource… · p. 15
  • The model predicts Level-0 for Problems 2 and 3, then provides a solution for Problem 3: 15 minutes.The supplied solution computes 5 minutes of hopping and 10 minutes of swimming; the Problem 2 solution is truncated in the passages.Source: E Case Study: Meta-cognitive Resource… · p. 16
  • A Level-3 prediction let the model bypass the geometry problem and solve the next two within 512 tokens, earning a total score of 2.The observation describes the geometry task as AIME-level and likely to exceed the remaining token budget.Source: E Case Study: Meta-cognitive Resource… · pp. 15–16
Loading 2601.03822v1…