Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
Authors: Zishun Yu
Tengyu Xu
Di Jin
Karthik Abinav Sankararaman
Yun He
Wenxuan Zhou
Zhouhao Zeng
Eryk Helenowski
Chen Zhu
Sinong Wang
Hao Ma
Han Fang
Venue: ICML 2025 poster (2025)
Keywords: large language models
reasoning
inference compute management
reinforcement learning
Primary area: deep_learning->large_language_models

Abstract: Solving mathematics problems has been an intriguing capability of large language models, and many efforts have been made to improve reasoning by extending reasoning length, such as through self-correction and extensive long chain-of-thoughts. While promising in problem-solving, advanced long reasoning chain models exhibit an undesired single-modal behavior, where trivial questions require unnecessarily tedious long chains of thought. In this work, we propose a way to allow models to be aware of inference budgets by formulating it as utility maximization with respect to an inference budget constraint, hence naming our algorithm Inference Budget-Constrained Policy Optimization (IBPO). In a nutshell, models fine-tuned through IBPO learn to ``understand'' the difficulty of queries and allocate inference budgets to harder ones. With different inference budgets, our best models are able to have a $4.14$\% and $5.74$\% absolute improvement ($8.08$\% and $11.2$\% relative improvement) on MATH500 using $2.16$x and $4.32$x inference budgets respectively, relative to LLaMA3.1 8B Instruct. These improvements are approximately $2$x those of self-consistency under the same budgets.

Think Smarter not Harder:
Adaptive Reasoning with Inference Aware Optimization
Zishun Yu' Tengyu Xu² Di Jin Karthik Abinav Sankararaman' 2 Yun He Wenxuan Zhou'
Zhouhao Zeng? Eryk Helenowski Chen Zhu Sinong Wang² Hao Ma² 2 Han Fang?
Abstract
observation merged in the LLM reasoning literature is that longer reasoning traces often leads to improved reasoning soundness and correctness. The seminal work of chain-of- thought (CoT) (Wei et al., 2022) is an excellent example of how enriching reasoning details, by decomposing rea- soning traces into steps, improves its problem-solving ca- pability. CoT has been considered standard technique in reasoning, recent works extend CoT by allow LLMs to expand its reasoning steps, by for example CoT with more steps (Jin et al., 2024) (as explicitly required by in- struction), self-reelletiion/correttio (Madaan et al., 2024; Zelikman al., 2022; Yan et al., 2024; Qu et al., 2024), multi-turn reasoning (Kumar al., 2024) or multi-agent debate (Liang et al., 2023; Pham al., 2023) (as hetero- geneous case of multi-turn). It was conjectured that scaling the test-time compute or the reasoning length unleashes LLMs' potential for reasoning (Snell et al., 2024), which has been empirically verified by recent hype of ultra-long reasoning models, such as OpenAT-ol (Jaech et al., 2024) and DeepSeek-RI (DeepSeek-Ar
Solving mathematics problems has been an in- triguing capability of language models, and many efforts have been made to improve reasoning by extending reasoning length, such through self- correction and extensive long chain-of-thoughts. While promising in problem-solving, advanced long reasoning chain models exhibit an undesired uni-modal behavior, where trivial questions re- quire unnecessarily tedious long chains of thought. In this work, we propose a way to allow models to be aware of inference budgets by formulating it as utility maximization with respect an infer- ence budget constraint, hence naming our algo- rithm Inference Consttt-onstraiiee Policy Opti- mization (IBPO). In a nutshell, models fine-tuned through IBPO learn to "understand" the difficulty of queries and allocate inference budgets to harder ones. With different inference budgets, our best models are able to have a 4.14% and 5.74% ab- solute improvement (8.08% and 11.2% relative) on MATH500 using 2.16x and 4.32x inference budgets respectively, relative to LLaMA3.1 8B In- struct. These improvements are approximately 2x those of self-consistency under the same budgets.
al., 2025). We'll later categorizes these type of responses as (standard) CoT re- sponses, extended responses, and (ultra-)long responses, respectively, given the nature of their reasoning lengths.
While scaling reasoning length is promising in problem- solving, advanced long reasoning-chaie models show an undesired uni-modal behavior that trivial questions may re- quire unnecessarily tedious long reasoning trace, an example is shown in Appendix A. This uni-modal behavior creates unnecessarily higher inference costs and increased carbon footprints (Henderson et al., 2020; Anthony et al., 2020). To partially address this, we study how to enable multi-modal behavior for reasoning models in a way the length of reason- ing traces are automatically adjusted according to the hard- ness level of the queries. From the aspect of query-adaptive reasoning length, some heuristic methods (e.g. Aggarwal et al., 2023; Xu et al., 2024a; Wang al., 2024) have been making effort towards better token efficiency, by which is meant better accuracy with (hopefully) less token overhead. We take a reinforcement learning (RL) perspective, where the accuracy gain over the token overhead nothing but a non-differentiable objective to be optimized. One could, for instance, take the negative response length or metric of this
1. Introduction
Complex reasoning has been an intriguing ability of large language models (LLMs), with application in for exam- ple mathematical problem-solving (Cobbe et al., 2021; Hendrycks et al., 2021b; Lightman et al., 2023) or cod- ing (Chen et al., 2021; Austin et al., 2021; Hendrycks al., 2021a), which does not only require nature language com- prehending but also logical and critical "thinking". An
DDpaarmment of Computer Science, The University of Illinois Chicago, USA Metaall, Menlo Park, USA. Correspondence to: Zishun Yu <zyu32@uicc udu..
Proceedings of the 4dnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
sort an intrinsic reward (Chentanez et al., 2004; Pathak et al., 2017). However, balancing the intrinsic and extrin- sic (accuracy) rewards might also be challenging (Liu et al., 2021), and might be vulnerable to reward hacking (Pan et al., 2022; Skalse et al., 2022; Karwowski et al., 2023).
the sequence-level notation (or the bandit notation), com- monly used in LLMs (Ziegler et al., 2019; Rafailov et al., 2024), especially in preference modeling, that suppresses the transition probabilities and intermediate rewards, and see a response as a whole. In particular, a policy æ : X - A()) takes a prompt z E X and draw a response a1 od2 ... .0ar =: y E y from the produced probability simplex A (J)), where denotes concatenation, a, E V corresponds to the i-th token drawn from the vocabulary V, and T is the maxi- mum length. A LLM is a parametric policy тH E Ill C II, where TT and IT are the parametric and non-parametric policy space respectively. Let [T:; r) or sometimes I (TT) be general objective function, defined by a prompt distribution μ E A(X), and a bounded reward function : X y + [-Rmax: Rraxs) Also, we define pn as an empirical distribution induced from Ω, a set of prompts.
Instead explicit modeling the length of responses, we take more abstract formulation, where we consider labeling each response y with an unique group label G; i E [G] for total number of G groups, so that the union of these disjoint groups exactly form the response space U:G: = ) For example G could be the group of CoT responses (with standard length) or extended responses. We could then im- pose an density constraint on each or a set of groups, by caping Eyxx(_)))) [1 [veg.)] S q: for some prompt distri- bution μ and some response distribution T (2),, induced from LLMs, conditioned a prompt x. This naturally formu- lates a constrained RI. (Garcıa& & Fernández, 2015; Altman, 2021) problem. Also this group definition is motivated by the resource-allocation literature (Chenery& & Kretschmer, 1956; Ibaraki & Katoh, 1988; Karlin, 2003), from the opti- mization and econometric communities, which have been later applied in many machine learning applications (e.g. Zemel et al., 2013; Badanidiyuru et al., 2018). This general- ization allows potential broader application of our algorithm as discussed in Section 6. We've now set our goal:
As aforementioned, we define G disjoint groups 9, such that U.G. = 3 and G; n 9; 0 for all i * j. Each re- sponse y E y is attached to exactly one group, in the sense that y E G, for some i. In the context of LLM reason- ing, without loss of generality, we consider two groups for brevity: Go and G+, corresponding to regular-length CoT responses and extended responses (with low and high in- ference costs), respectively. To conclude the formulation of constrained RL with resource allocation constraints, we could in general define the feasible set as a convex polytope, EyE (vec,)] )] S q: for all i), that caps
Dg := { the total density mass of each group G, by qs In our setting, we only need to cap the total mass of extended responses to optimize the inference efficiency, posing a half-space 4 := { (x ) S 4+} } for some q+ >> 0.
A constrained RL framework controlling how response groups } are distributed.
Therefore, one could control how responses of different lengths (which are supposed belong to different groups) are distributed. Our rationale of algorithm design is given in Section 2, derived from an optimization perspective but ended as a very simple generalization of iterative super- vised fine-tuning (SFT) methods such as reward-ranking fine-tuning (RAFT) (Dong et al., 2023) and rejection sam- pling fine-tuning (RFT) (Ouyang et al., 2022; Touvron et al., 2023), see details in Section 3. Given the motivation of our algorithm, we call the resulted algorithm Inference Budget-Constie Conetrrnnnd Policy Optimization (IBPO).
Background. With the bandit setup, our setting draw lot connections to the online learning literature, especially on- line/bandit convex optimization (Hazan et al., 2016; Slivkins et al., 2019), although we optimize a fixed function the train- ing data however collected in an online fashion. In the sense of distributing resources across groups, it is connected to for example knapsack bandits (Badanidiyuru et al., 2018) and statistical parity (Zemel et al., 2013), aka group fairness. Although the definition of groups and optimization pro- grams could be different. This allocation optimization point of view allows us to further extend our method to broader LLM applications. To the end of policy optimization with constraints, common techniques include projection (Zinke- vich, 2003; Flaxman et al., 2004; Bubeck et al., 2015; Yang et al., 2020), Lyapunov-based (Chow et al., 2018; Cayci et al., 2022), or Lagrangian methods (Ray et al., 2019).
Paper structure. Section 2 present the derivation of our algorithm from an optimization view, resulting in a simple weighted iterative SFT update. Section 3 provides further details on practical implementation, including the base al- gorithm and the reward design. Section 4 introduces some experimental settings. The empirical results of our IBPO are presented in Section 5. And Section 6 concludes our work with limitations, broader impact, and further discussions.
Non Non-parametric space II. In our case, our goal is to solve:
max J(ro) s.t. n. $_
(1)
2. A Agorithm Design
where Hg is the parameterized policy space. Solving Eq. (1) is however intractable due to the LLM parameterized pol- icy space Nl. A common practice is alternating gradients
Problem setup. To make the notation compact, we take
2

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
between reward maximization and constraint satisfaction. For example in Lagrangian methods such as TRPO/PPO- Lagrangian (Ray et al., 2019), one could do alternating up- date Te and the Lagrangian multiplier. Another workaround is to first obtain a solution n² in the non-parametric pol- icy space 11 ;= A(X' x })), aka tabular representations, and project n" onto the parameteric one Hlg, as a technique used in many (constrained) RL works (Peters et al., 2010; Montgomery & Levine, 2016; Zhang et al., 2020).
ing the aforementioned projection step, could write the program as bi-level stochastic optimization:
T = arg min E [KL Xx. [[]]]]]]]]]]]]]]]]]]]
(3)
s.t. * Y. E arg re Y,) I (T;
where Y (), hence Rt is indirectly a function of 0.
Practical update. For general bi-level optimization prob- lems, iteratively solving the upper and lower-level problems by alternatielyy fixing one while optimizing the other could be expensive (Zhang et al., 2024). We've however already setup a manageable inner problem, making it easy to solve for example using convex solvers. One could therefore do iterative gradient updates on the upper level while directly solving the lower-level at each iteration:
The advantage of working in the non-parametric space II is: solving maX.enl I (T) s.t. π E 4+ is easy, on the conditions that (i) [Tr) ) is concave in π so that it is a con- vex program, and (ii) sampling and evaluating the reward function r): and the cost indicator 1 (.) are cheap. The condition (i) is sometimes true, for instances: bandit ob- jective (Slivkins et al., 2019); the LP formulation (Manne, 1960; Denardo, 1970; Nachum & Dai, 2020; Nachum et al., 2019) and (relative) entropy regularized RL (Ziebart, 2010; Haarnoja et al., 2017; 2018) are often concave in occupancy measure p. However, in rare case condition (ii) holds, RL works (Haarnoja et al., 2018; Peng et al., 2019; Zhang et al., 2020) often resort value function approximation, making is easy, for discrete action space, evaluate Q-values (or alternatively advantages) for all actions for a specific state. It is then tractable to obtain closed- form solution (optimal w.r.t. the value/advantage approximations) in II. Once an optimal policy x² is found, one could then project it onto II through (reverse) information projection = arg ming (LT ([TT) (often done approximately by taking gradient steps).
8' = 0 - a V V R
|KL xx Y. (T [x]]
(4)
s.t. #* X.Y. arg m (X. (x) I (x; Yo)
where 0 and 0' are the parameters of current and next itera- tion, respectively; and the projection step is also done with samples (X, Yo) of the current iteration.
Note that A* is indirectly a function of 0 through the sam- ples (X, Yo). The gradient (f Y, [(T)) hence re- quires differentiation through Y,,,, meaning differentiate through an arg max operator, which can in principle be achieved through implicit differentiation (Amos & Kolter, 2017; Lorraine et al., 2020). However, to avoid additional implementation and computation overhead, instead we use the semi aomi-gradient
Stochastic optimization. With L.M,, it is obviously in- tractable to sample and evaluate the reward function r(r, y) for all (z. 8) E X' x D, similarly for the cost indicator 1.3-- To avoid the training of additional value models for LLMs (Snell et al., 2022; Yu et al., 2023), which can cre- ate significant overhead terms of memory usage, imple- mentation complexity, and training stability, we consider stochastic optimization. The stochastic counterpart described in Eq. (2) solves an approximate #* using man- ageable number of samples rather than directly solving for the global optimum x², still, in the non-paramctric space I1:
(sa f, ]))o)) where sGf-) is stop gradient operator. This stop-gradient trick is quite common in many ML applications (Sutton, 2018; Foerster et al., 2018; Chen & He, 2021), leading to the update:
O' 0 - a o Enxpx (SG fxx, Y. H|m [=]]
approximate projectionn weighted SPT
(5)
s.t. AX,Y, E arg max xxn I [TT; (T; X, Y,)
optimization for weight
** Y) arg max I X, Y) s.t. TE T ( Y) [T;; X. Y) = 1 E E ( vct]]
The semi-gradies leads to weighted SFT update (via #*) as
V (|] = -E, a* (Myllz) @ ag logie (Vky
This observation creates an extremely simple update with negligible implementation overhead.
(x,Y) ( (May)(;)) (0ojee+) - 4+)1 s 0}
(2)
Discussion. The update rule ended up aligning many it- erative weighted SFT algorithms, such as RAFT (Dong et al., 2023) and RFT (Ouyang et al., 2022; Touvron et al., 2023). In hindsight, our algorithm is motivated by the ob- servation that these extremely successful algorithms can be interpreted as projecting empirical solutions onto a paramet- ric space. Consequently, it is reasonable to use the empirical estimate in Eq. (2), as RAFT and RFT have demonstrated
where X E X is a vector of n sampled prompts and Y E
mm is a matrix of responses, with m responses for each of the n prompts; we explicitly write I with the conventional expected reward maximization for notational convenience, though alternative objectives are not restricted.
Since the empirical problem (2) is a convex program with relative small sample size, it now manageable. Comb-
3

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
Table 1: SFT methods from our optimization point-of-view. (w)-XE denotes (weighted) cross-entropy loss. RM stands for reward model. Binary reward indicates correctness. ft are (unnormalized) weights for the subsequent XE update.
optimization of LLMs. The choice is driven by implementa- tion considerations: CGPO's modular constraint-handling design makes it straightforward to incorporate additional constraints, such as the group density constraint in our case. CGPO in a nutshell can be viewed as a generalized Best-of- N (BoN), though depending on the specific CGPO settings, It operates by defining a feasible set over the sample space X' x Y. which are designed to capture constraint adherence for, e.g. correctness, factuality, and safety (Xu et al., 2024b). In short, (z,y) & E will be rejected. In the context of math reasoning, the objective of CGPO can be summarized as:
Teas. set +
acceptance weight #t'
ALGGO
loss L
reward r
SFT
XE
constant binary
II
constant
RFT
W-XE
II
1 1r((a,))=1 = 1fr(z,,)) -)- r'}.
(Best-af-N)
RAFT
W-XE
RM
I
r" "-=may ===axxr(pss (45)
Ours
W-XE
Sec. 3 Eq. (2)
Eq. (2)
max (r s.t. Eyl ( [x) (bEE,) ( 2 Vx:
Exi= (y: (y:thhhhh (z, 11 KL Tret) S Kllmax)
strong practical performance despite the inherent bias in- troduced by the non-linearity of these estimations. Since it is essentially generalizes SFT by re-weighting a sample pair (I;, Mij) by 4* (Wiyl 1z)) at each iteration based on the solution of an optimization problem f", we create Table to outline the corresponding interpretation for some iterative SFT methods, from this optimization point-of-view.
correctness
empirical KI.
(7)
This objective essentially optimizes reward over the feasible sets ES such that feasible response y is correct and within an KL range of KLmax. where the KL constraint is mea- sured using point estimate of the (forward) KL defined as KL(y; (y; x, Tref) == log (Ty - log (v)
The optimization problems for RFT and RAFT are trivial as they assign # (ull:) = 1 to accepted responses and to the response with the highest reward model score, respectively. Formulating these methods as optimization does not offer much advantages. However, this perspective provides flexi- bility for future work to extend our framework, allowing for feasible sets and weighting tailored to specific applications.
Resulted update (IBPO w/ J.B). Recap our update defined as 0' 0 a V p
v, s.t. # Y.Y, E arg max, rellnn (X.Y.) [TT: Yo). And we have now defined the RL objective J and the margin reward function ra. We are now ready to put everything together:
Y. E Earg max, 7 JA(T; X, Yo)
1 Σ. Ej m (TVs (x)ra vis)
3. Practical Implementation
s.t. π T 2 (X, Yo) and VzeX
(8)
Yet, as we are working in an algorithm-agnostic fashion, we are now ready to select a specific RI. algorithm, define its corresponding objective J, and specify a reward function r.
(g)
shorthamled OPTInB (inference under hudget)
In addition, as CGPO is a generalization of BoN (with tie breaking randomly). #* will be a pure strategy instead of stochastic one, meaning at most one y will be accepted for each z. for subsequent projection (SFT).
Reward function. Since we are working on mathematical problem-solving, ground-truth reward could be obtained through string matching (Cobbe et al., 2021; Hendrycks et al., 2021b) of the model's solution against the ground truth solution, yielding a binary reward function r mmatch X x V/ + {0, 1} that indicates correctness. On top of the binary reward, we define our reward r' as the reward margin. To formally construct the margin, we first define the expected reward of a set & such that F, (2., 9) Ey--n/r ["match (Ir y) y E G]. We then define the reward margin rs the reward advantage of group G against all other groups ) \ G:
Intuition. Behind this update our intuition can be verbally interpreted as: If prompt is hard, the margin of extended responses r' (1 y E G+) = (z 9+) - -=x(1, G.) is likely to be large for y E G+. Therefore ** will more likely to assign positive weight to an extended response y E G so that the objective receive larger margin reward. In contrast, if query x is simple, r (I. y E G+) is likely to be small. Hence ** will possibly assign positive weight to regular re- sponses y E Go. so that one could save some density budget for harder queries. See concrete example in Table 2.
L (x y E 9) := (x, 9) F, (z. (
(6)
We have AA (z. y ye -4)- (Ir, Go), in our case,
and similarly for G.. r'A may appear odd for not counting the correctness of individual y. This is handled subsequently.
Reward models. Note that the original CGPO implemen- tation has reward models for BoN ranking. We however intentionally excludes reward models, in our OPT formu- lation as shown in Eq. (8), to highlight our methodologi- cal contributions, by decoupling reward modeling efforts.
RL objective. For the learning algorithm, our choice is Con- straint Generative Policy Optimization (CGPO) (Xu et al., 2024b), which is designed for mujti-oheective constrained
4

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
An example of OPTiuB compared to oPTcgoo-Eq. (7)- with tie broken randomly, resulting in polentially non-unique Æ* This table shows several such #* solutions but not all. KL is omitted for brevity. Given the rewards below we have F(Z1, Go) = F( G4) = F(22, G++ = 1 and F(#2,9a = 0.5. Suppose the cap 7+ = 0.5, allowing at most 50% of accepted responses are extended y E G4. For a solution matrix ft,, and 0 represent accepted and rejected response, respectively.
Table 2:
Algorithm 1 Inference Budget-Constrained PO (IBPO)
Require: prompt set D, batch size n, number of responses
m, init policy To Tref. num of iters T', budgets q+
1: for t = 1, 1,.., T do
2: prompt sampling: X ~ u4B
3: response generation: Y ~ TT, (X)
4: evaluate correctness: R Raxm matchh (X, Y)
empirical KL: KL
KL X, Tref)
prompt
I1 (easy)
T2 (hard)
6:
if CGPO (i.e. w/o IuB) then
7: solve BoN (7): y, oPPoogoo (X, RL,
111. 3/2 E Go
121. U22 E 90
8: end if
responses
V13. //4 9+
U23, /24 9+
9:
if IBPO (i.e. w/ luB) then
1
0
10: solve Eq. (8): v, EOPTIB Rc,
r(s,y) v)
1
1
11: end if
8 12: grad step: -Σ:E, x (Mis x (z)) 00
log To (B [z: 8=0.
0
0
To (z, v)
0
0)
13:
end for
2
0 0%
1
0
| = oPTcgpo
0)
0
0
Due to the intricate details and the space limit, we defer the full version of this section to Appendix B, including details of prompting, dataset construction, training pipelines, etc.
0
0
0
0
A = OPTegpo
0 1
1
0)
n' = 1PT
0' 0 0)
1 0)
0
0)
Figure provides examples of our constructions, which should be visually telling. In nutshell, we use the step CoT (SCoT) format from Dubey et al. (2024) as Go. and construct a sequential voting (SV) response as G4. SV serves as a unimodal baseline to demonstrate the performance of our construction. Adaptive SV (ASV) generates a mixture of SCoT and SV responses, allowing us optimize its ability to adaptively choose between the two. For further details about the constructions, refer to Figure and Appendix B.
CGPO case I: for za (hard), Y21 E Go is accepted even though G+ has higher expected reward F F(r2, G+); case 2: y14 and y23
are both accepted, which exceeds the density budget Q+ = 0.5. Ours: accepts y and y23 maximize margin, ra (222, y2s)
0.5, while adhering to density cap q+ 0.5.
Nonetheless, using reward models remains possible. To avoid introducing additional notation, we elaborate verbally: OPTΙиB essentially selects either group G4 or Go for query J. Within a group G, all responses receive the same re- ward ra (I, 9), leaving it possible to further rank responses within each group using reward models.
Notations., we use LLaMA to refer to instruction- LLaMA 3.1 8B (Dubey et al., 2024). ASV-SFT-a denotes ASV models that are supervise fine-tuned with a coefficient a (see Appendix B). ASV-LuB-q+ refers to ASV models optimized by our Algorithm I with budget constraint @+-
Implementation & solvers. A pseudo-code of our IBPO with OPTIuB is listed in Algorithm 1. The essential change to replace the constrained reward ranking OPTearo with general optimization problem, in our case the mar- gin maximization under budget denoted as OPTIuB- The ОРТьь problem is (integer) linear programming prob- lem that could be solved by off-the-shelf solvers, such as CPLEX (Cplex, 2009). Gurobi (Gurobi Optimization, I.I., 2024), or SciPy (Virtanen et al., 2020) which is our choice.
5. Evaluation of IBPO w/ OOTTT. TuB
In this section, we show that: (i) SV is a valid construction with reasonable absolute improvement (Table 3) and scales as good as majority vote (MV, aka self-consistency (Wang et al., 2022)) (Figure 2a & 2b); (ii) ASV-SFT does not achieve good efficiency (Figure 2a & 2b) as SFT does not optimize efficiency as an optimization objective; (iii) ASV- IuB optimized by IBPO achieves better efficiency (Figure 2a & 2b), adherence to constraints (Figure 2c), and allocation of inference to harder quires (Figure 2c & Figure 2f).
4. Naïve Construction of g+
Yet we work on abstract groups Go and G+- In this sec- tion, we present the details of our constructions of extended length responses, i.c. the group G4. However developing long reasoning models is beyond the scope of this work, as our focus is on the constrained optimization of LLMs. Our constructions are for demonstrative purpose only.
5.1. Absolute Improvement (Table 3)
"Baselines". To put SV/ASV in comparison with other baselines in literature, we gather several baselines from the self-correction literature as essentially these methods in-
5

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
SV Prompt (Simplified)
Table 3: Comparison of approaches on their improvements versus base models. 1, 4, it indicate results duplicated from Qu al. (2024); Yan et al. (2024); Kumar et al. (2024), respectively. indicates our methods/constrrutionss and the improvements are relative LLaMA model.
You are asked to give at most eight diverse solutions different way, without referencing to the previous trials. If solution occurs three times. it is considered as consensus and will be used the final answer.
turns/trials per response
ASV Prompt (Simplified)
approach
pass
improv.
SFT//rompting- based
For medium and hard level problems, you are asked to give at most eight diverse solutions in different way. with- referencing to the previous trials.
SV-SFT"
LLaMA
56.8
5.54
5.67x
For easy level prohlems, you are allowed only attempt, which will be considered your final answer.
ASV-SFT-1
LLaMA
55.6
4.43
5.74x
SFT-RISE'
setting I (table of setting 2 (table of f)
5.5
-0.3
5x
Voting Response G+ (SV,ASV ASV Case I)
5.0
0.0
5x
[TRIAL] #B Step 1:
steps omitted
SFT-SCoRe
setting (table of tf) setting 2 (table of tt)
1.8
The final answer is:
A1
[/TRIAL]]
55.0
0.0
[TRIAL] ## Step 1: steps omitted
RISE
The final answer is:
A2
[/TRILL]
LLaMA2 Base
1.4
-0.5
5x
[TRIAL] ##lh Step 1:
steps omitted
boosting
5.5
0.0
5x
The final answer is:
A1 [/TRIL]]
S
[TRIAL] ## Step 1: steps omitted
LLaMA3-8B
33.14
2.56
The final answer is: A
[/TRILL]
Mistral-7B
25.48
1.44
not studied
The answer A1 has occurred three times, and is consid- ered as a consensus.
DappSck-Matt-Bawo-P
41.40
3.18
Qwen2-Math-7B
51.76
0.44
The final answer is A1 hope it is correct.
Self-Refine'
Base
1.9
0.0
3x
GPT-3.5
36.5
-3.2
3x
Non- Non-Voting Response Go (ASV Case 2)
Mistral-7B
7.1
-0.4
3x
Eurus-7B-SFT
9.0
-3.3
3x
[TRIAL]] ## Step 1:
steps omitted
STaR
The final answer is: A1 [/TRILL]
setting 1 (table of +1)
54.0
0.4
2x
Terminated due to difficulty level.
setting 2 (table of tt)
41.2
-14.2
Online Iterative/RL.
Figure Example of prompts and responses. For Go, we use standard step CoT (SCoT). To construct G4, the model is prompted sequentially generate up trials encapsulated within the special tokens [TRIAL] and [/TRIAL], followed by consensus answer. The SV prompt instructs the model to output only voting responses, allowing us to later evaluate its performance and show that is a reasonable construction. The ASV prompt asks the model to decide whether to output voting or non-voting response, facilitating our further optimization of budget allocation in Section
ASV-IuB-g-
25%
54.2
2.94
2.24x
2+ 50%
55.4
4.14
2.16x
4+ 75%
57.0
5.74
4.32x
RISE
Iteration
3.4
Iteration
10.4
4.6
SCoRe
Gemini 1.5 Flash
64.4
4.4
2x
more turss (fig. 8ft)
266
26
5x- 10x
likman et al., 2022), S$C (Yan et al., 2024), SFT-RI and SFT-SCoRe, where SFT-RI and SFT-SCoRe are SFT com- parators implemented in Recursive Introspection (RI) (Qu et al., 2024) and SCoRe (Kumar et al., 2024) respectively.
crease inference length, though not extensively as Jaech al. (2024). The results in Table are mainly gathered from Qu al. (2024); Kumar al. (2024). The self-correction base- lines often admit multi-turn structure, similarly to the multi-trial construction of SV. We therefore include a col- umn of inference cost measured by number of turns/trials for a rough comparison. The SFT comparators we include are (reproduced) Self-Refim (Madaan et al., 2024), STaR (Ze-
ASV experiments. SV-SFT follows our training pipeline of Exp 1.2 in Table 7 (Appendix B) and an voting only baseline. ASV-SFT-o follows our Exp 2.1 setup and an adaptive baseline, meaning model decide whether vote or
6

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
1.0
ASV-UU lUB 75%
140
SCoT responses
54%
SU-SFT
ASV-LU-50%%
vetting respanses
48%
0.56
ASW 30%
AS N---
0.8
AS--
120
"passa1"
ASV-U8-255
vote fraction
0.6
count
180
31%
STV-5TT-2
0.52
0.4
ASV-SFT-4
0.2
OSS to SV interputtion
0.48
0.0
6
7
100
150
200
250
average of trials per problem
training step
difficulty level
performance of last checkpoints
(c) voting ratio of correct responses (train)
(e) voting ratio by difficulty (4+ = 50%)
ASV-UB-75% (224)-IT
voting resporses fraction (teat)
140
SGaT reppores
resporiws
41N
0.30
ASV-UB-BO 50%
(240)
ASV-SFT-3
1144).
100
17%
"pass1""
count
SFV-2FT-2
0.52
40-
ASV-SFT-4
OSS to SV interplation
4SV-UUB-50%
0.08
0.0
ASV
7
50
100
150
200
250
0
average of trials per problem
training step
aitfiiulity
(b) best checkpoints (steps parentheses)
(d) voting response ratio (test set)
(1) voting ratio by difficulty (4+ 25%)
Figure 2: Column I: Comparison of pass@1 @1 (maj @ N for MV) against the average number of trials per response (x-axis). OSS refers to LLaMA. The interpolation between OSS and SV-SFT (aligning with MV) serves as a hypothetical efficiency boundary. ASV-SFT shows lower efficiency relative to this boundary, whereas ASV-luB consistently achieves better efficiency (i.e. above the boundary). Column 2: Voting response ratio versus training steps. Dashed line denotes the budget 4+- On the training set, IuB formulation follows the budget constraints almost exactly. Due to distribution shift, the constraint on testing set is not entirely exact, but still it is noticeable that the voting ratio follows the order of 75% > 50% > 25%. Column 3: The luB formulation enables the model to dynamically allocate voting budget to harder problems.
not for a query. We report ASV-SFT-1 only as is empiri- cally the best ASV-SFT-a baseline, though still fall short in terms of efficiency. The ASV-luB-q+ experiments, initial- ized from ASV-SFT-1, are our RL-tuned adaptive models, optimized by IBPO with 1PT as inner optimization.
improvements. As somewhat tangential yet potentially intriguing observation, it is evident that prompting- bbsed and SFT-based methods struggle with both absolute im- provement (Table 3) and efficiency (Figure 2), support- ing the conjecture that SFT alone does not enable self- correction capabilities (Huang et al., 2023; Kumar et al., 2024). This observation is also partially supported by con- current work (DeepSeek-Al et al., 2025), which suggests that such self-correction emerges automatically during RL rather than manually created by prompting or SFT.
Observations. We would like to first emphasize that the comparison in Table 3 is not intended to demonstrate that SV outperforms SFT-based self-correction or that our ASV. luB surpasses RL-based self-corrections. These efforts are orthogonal,
focus is constrained optimization. As and ASV-SFT-1
observed, our SFT Seanstructions-S--ee. achieve clear improvement in pass@1 1 with high inference costs (5+ times the number trials). The ASV-luB-qe formulation, particularly with 4+ = (50%, 75%}, shows significant improvement while reducing costs by 4.14% at 2.16x and 5.74% 4.32x. This performance on par with SCoRe, state-of-the-an RL-based self-correction method. Note that the performance of ASV-IuB-q.. is re- ported using the best checkpoints. Results from the last checkpoint are shown in Figure 2a. Additionally, training curves are presented in Appendix E, which shows consistent
5.2. Efficiency, Budget Adherence & Allocation
Table shows that our ASV-luB-q+ models achieve results comparable to RL-baaed self-correction models, simply by inference budget management. In this subsection, extend our discussion on: (i) performance-budee efficiency, (ii) constraint satisfaction, and (iii) inference budget allocation.
Performance- efficiency. In Figure 2a and 2b, we vi- sually assess the performance-budge efficiency, compared to hypothetical efficiency boundary. This boundary is

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
an interpolation between OSS LLaMA model and SV-SFT. It is reasonable to consider it a hypothetical boundary for two reasons: (i) OSS and SV-SFT are two extremes of ASV-IuB-q++ corresponding to the cases of q+ = 0 and 4+ = 1 respectively; and (ii) this interpolation achieves an increase comparable to MV, if not slightly better. The SFT version of ASV is generally much worse than the bound- ary, as SFT alone does not optimize resource allocation as mathematical objective. For ASV-IuB-q+ optimized by IBPO, we report both the last and best checkpoint results in Figure 2a and Figure 2b, respectively. Our formulation achieves, in general, better performance-budget efficiency, except q+ = 75% in Figure 2a. We will extend our discus- sion on this unsuccessful case in subsequent paragraphs.
thereby alleviate the distribution shift issue.
6. Conclusion & Discussions
We derived a constrained policy optimization framework, IBPO, from an optimization perspective, resulting in a sim- ple weighted SFT update that resembles successful iterative SFT algorithms such as RFT and RAFT. In each iteration, the optimal weight is obtained by solving an (integer) linear program. The practical implementation of IBPO is build on top of CGPO, and is evaluated on math reasoning task with inference budget constraints. Empirical evaluations show that our framework enables the model to adhere to constraints and dynamically allocate the inference budget.
Constraint satisfaction. We then evaluate how effectively the constraints are enforced. In Figure 2c, the budget constraints are successfully maintained during training for a+ = (25%, 50%, 75%}. Due to distribution shifts, exact adherence to these constraints is not expected on the test set. Nevertheless, Figure 2d demonstrates that constraints are still upheld at the end of training for g+ = (25%, 50%), and the ratio of voting responses follows the order 75% > 50%+ > 25%. (Figure 2c illustrates that our model meets the budget constraints for the set of correct responses. The constraints also hold for all responses, see Appendix E.)
Batch optimization & solver time. Since we solve an opti- mization problem per iteration (i.e. per mini-batch), limited computational resources result in smaller sample sizes for the inner optimization problem, leading to larger vari- ance, This issue can be mitigated through "sample accumu- lation", accumulating samples across multiple consecutive steps, similar to gradient accumulation practice in LLMs. A pseudo-code for sample accumulation can be found in Appendix H. In addition, though integer linear program- ming is NP-hard (Vazirami, 1997), the number of variables in our batch-level optimization is typically small, resulting in minimal computational overhead. Refer to the wall-time plot for the SciPy solver in Appendix H for details.
Didieie allocation. We have show improved efficiency and adherence to constraints, we would like to fur- ther validate the intuition of our design: that more challeng- ing problems may require longer reasoning steps, whereas simpler problems be resolved with just SCoT responses. Ideally, the model should allocate more voting responses to problems with higher difficulty levels. To this end, we use the difficulty levels from the metadata of Hendrycks MATH and plot the ratio of voting responses for each difficulty level. Figure 2 and 2e illustrates that for both g+ = (25%, 50%), more challenging problems, such as those at levels and 5, receive a higher allocation of budgets. This allocation pattern is particularly evident for the case of g+ 25%, where only 2% of level 1 problems receive voting responses.
Broader applications. Our framework has only been evalu- ated with inference in math, the resource allocation problem however has far-reaching implications within the ML com- munity, As result, our framework can be potentially ex- tended to further LLM applications. For example, promis- ing application is statistical parity (Zemel et al., 2013), aka group fairness. In this context, one could consider attribut- ing responses to their respective social groups, and cap the density of responses that correspond to socially privi- leged groups, to encourage more inclusive and equitable responses across different demographies. Another poten- tial application is the balanced expert activation in mixture of experts (Jacobs et al., 1991; Shazeer et al., 2017: Lep- ikhin et al., 2020) systems, which is sometimes achieved by adding an auxiliary balancing loss (Wei et al., 2024). Alter- natively, this balance can be possibly achieved by enforeing minimal total activation density for each expert. This may help to prevent over-reliance and hence enhancing the overall robustness and efficiency. We leave the exploration of broader applications and their implementations as future directions.
Discussion on the unsuccessful case. There is one ob- served unsuccessful case in Figure 2a: the last checkpoint of ASV-IUB-75% which falls on the hypothetical boundary rather than above it. This outcome arguably expected, as observed in Figure 2d, where ASV-IuB-7551 outputs almost exclusively voting responses at the end of training. As result, this model is not adaptively allocating resources, and thus no improvement in efficiency is anticipated. This unsuc- cessful case is hence caused by the distribution shift between the training set and the testing set. It possible that scaling the training set-given that our training set DR.. contains approximately 10k prompts, which is relatively small make the testing set more likely to be in-distribution and
subset of experts,
Impact Statement
This paper presents work whose goal is to advance the field LLMs. There are many potential societal consequences of our work, none of which we feel must be highlighted here.
8

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
References
Chow, Y., Nachum, O., Duenez-Guzman, E., and Ghavamzadeh, M. A lyapunov-based approach to safe reinforcement learning. Advances in neural information processing systems, 31, 2018.
Aggarwal, P., Madaan, A., Yang, Y., et al. Let's sample step by step: -daptive-consistency for efficient reasoning and coding with llms. arXiv preprint arXiv:2305.11888 2023.
Cobbe, K.,, Kosaraju. V., Bavarian, M., Chen, M., Jun, H., Kaiser, L.,, Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10. 14 2021.
Altman, E. Constrained Markov decision processes. Rout- ledge, 2021.
Amos, B. and Kolter, J. Z. Optnet: Differentiable opti- mization as layer in neural networks. In International conference machine learning, pp- 136-145. PMLR, 2017.
Cplex, I. V12. 1: User's manual for eplex. International Business Machines Corporation, 46(53):157, 2009.
DeepSeck- AI, Guo, D., Yang, D., Zhang, H., Song, J., Zhang. R., Xu, R., Zhu, Q., Ma, S., Wang, P., Bi, X., Zhang, X., Yu, X., Wu, Y., Wu, Z. F, Gou, Z., Shao, Z., Li, Z., Gao, Z., Liu, A., Xue, B., Wang, B., Wu, B., Feng, B., Lu, C., Zhao, C., Deng, C., Zhang, C.,, Ruan, C., Dai, D., Chen, D., Ji, D., Li, E., Lin, F., Dai, F., Luo, E. Hao, G., Chen, G., Li, G., Zhang, H., Bao, H., Xu, H., Wang, H., Ding. H., Xin, H., Gao. H., Qu, H., Li, H., Guo, J., Li, J., Wang, J., Chen, J., Yuan, J., Qiu, J., Li, J., Cai, J. L., Ni, J., Liang, J., Chen, J., Dong, K., Hu, K., Gao, K., Guan, K., Huang, K. Yu, K., Wang. I.,, Zhang, L., Zhao, L., Wang, L., Zhang, L., Xu, L., Xia, L., Zhang. M., Zhang, M., Tang, M., Li. M., Wang, M., Li, M., Tian, N., Huang, P., Zhang, P., Wang, Q., Chen, Q., Du, Q. Ge, R., Zhang, R., Pan, R., Wang, R., Chen, R. J., Jin, R. L., Chen, R., Lu, S., Zhou, S., Chen, S., Ye, S.. Wang, S., Yu, S., Zhou, S., Pan. S., Li, S. S., Zhou, S., Wu, S., Ye, S., Yun, T., Pei, T., Sun, T., Wang, T. Zeng, W., Zhao, W., Liu, W., Liang, W., Gao, W. Yu, W., Zhang, W., Xiao, W. L.., An, W., Liu, X., Wang, X., Chen, X., Nic, X., Cheng, X., Liu, X. Xie, X., Liu, X., Yang, X., Li, X., Su, X., Lin, X., Li, X. Q., Jin, X., Shen, X., Chen, X., Sun, X., Wang, X., Song, X., Zhou, X., Wang, X., Shan, X., Li, Y. K., Wang, Y. Q.. Wei, X., Zhang, Y., Xu, Y., Li, Y., Zhao, Y., Sun, Y., Wang, Y., Yu, Y., Zhang, Y., Shi, Y., Xiong, Y., He, Y., Piao, Y., Wang. Y., Tan, Y., Ma, Y., Liu, Y., Guo, Y., Ou, Y., Wang. Y., Gong, Y., Zou, Y., He, Y., Xiong, Y., Luo, Y., You, Y., Liu, Y.. Zhou, Y., Zhu, Y. X., Xu, Y., Huang, Y., Li, Y., Zheng, Y., Zhu, Y., Ma, Y., Tang, Y., Zha, Y., Yan, Y.. Ren, Z. Z., Ren, Z., Sha, Z., Fu, Z., Xu, Z., Xie, Z., Zhang, Z., Hao,
Anthony, L. F. W., Kanding, B., and Selvan, R. Car- bontracker: Tracking and predicting the carbon foot- print of training deep learning models. arXiv preprint vrXiv:2007.0305/1 2020.
Austin, J., Odena, A., Nye, M., Bosma, M., Michalewski, H., Dohan, D., Jiang E., Cai, C., Terry, M., Le, Q., et al. Program synthesis with large language models. arXiv preprint 2rXiv:2108.07732, 2021.
Badanidiyuru, A., Kleinberg, R., and Slivkins, A. Bandits with knapsacks. Journal of the ACM (JACM), 65(3):1-55, 2018.
Bubeck, S., Dekel, O., Koren, T., and Peres, Y. Bandit convex optimization: ssqrtt sqqrt regret in one dimension. In Conference on Learning Theory, pp. 266-278. PMLR, 2015.
Cayci, S., Zheng. Y., and Eryilmaz, A. A lyapunov-based methodology for constrained optimization with bandit feedback. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 3716-3723, 2022.
Chen, M., Tworek. J., Jun, H., Yuan, Q., Pinto, H. D. O., Kaplan, J., Edwards, H., Burda, Y., Joseph, N., Brockman, G., et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374. 2021.
Chen, X. and He, K. Exploring simple siamese represen- tation learning. In Proceedings of the IEFE/CVF con- ference on computer vision and pattern recognition, pp. 15750-15758, 2021.
Z., Ma, Z., Yan, Z., Wu, Z., Gu, Z., Zhu, Z., Liu, Z., Li,
Z. Xie, Z., Song, Z., Pan, Z., Huang, Xu, Z., Zhang, Z., and Zhang, Z. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning, 2025. URL https:/ / wwaaaiivoorgr. abs 2500 12948.
Chenery, H. B. and Kretschmer, K. S. Resource allocation for economic development. Econometrica, Journal of the Econometric Society, pp. 365-399, 1956.
Denardo, V. On linear programming in a markov decision problem. Management Science, 16(5):281-28,, 1970.
Chentanez, N., Barto, A., and Singh, S. Intrinsically moti- vated reinforcement learning. Advances in neural infor- mation processing systems, 17, 2004.
Dong. H., Xiong, W., Goyal, D., Zhang, Y., Chow, W., Pan, R., Diao, S., Zhang, J., Shum, K., and Zhang, T. Raft: Re-
9

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
ward ranked finetuning for generative foundation model alignment. arXiv preprint arXiv:2304. 2023.
Ibaraki, T. and Katoh, N. Resource allocation problems: algorithmic approaches. MIT press, 1988.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The Ilama herd of models. arXiv preprint arXiv:2407.21788. 2024.
Jacobs, R. A., Jordan, M. I., Nowlan, J., and Hinton, G. E. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 87, 1991.
Jaech, A., Kalai, A., Lerer, A., Richardson, A., El-Kishky, A., Low, A., Helyar. A., Madry, A., Beutel, A., Car- ney, A., et al. Openai ol system card. arXiv preprint MrXi::2412166720 2024.
Flaxman, A. D., Kalai, A. T., and McMahan, H. B. On- line convex optimization in the bandit setting: gradient descent without a gradient. arXiv preprint cx/040807, 2004.
Jin, M., Yu, Q., Shu, D., Zhao, H., Hua, W., Meng, Y., Zhang, Y., and Du, M. The impact of reasoning step length on large language models. arXiv preprint arXiv:2401.0492 00495,, 2024.
Foerster, J., Farquhar, G., Al-Shedivat, M., Rocktäschel, T., Xing, E.,, and Whiteson, S. Dice: The infinitely differen- tiable monte carlo estimator. In International Conference on Machine Learning, pp. 1-29-1538, PMLR, 2018.
Karlin, S. Mathematical methods and theory in games, programming. and economics, volume 2. Courier Corpo- ration, 2003.
Garcia, J. and Fernández, F. A comprehensive survey on safe reinforcement learning. Journal of Machine Learning Research,
2015.
Gurobi Optimization, LLC. Gurobi Optimizer Reference Manual, 2024. URL https:tp//.. /ww. hwtp://...........omooo. com.
Karwowski, J., Hayman, O., Bai, X., Kiendlhofer, K., Grif- fin, C., and Skalse, J. Goodhart's law in reinforcement learning. arXiv preprint arXiv:2310.0914 2023.
Haarnoja, T., Tang, H., Abbeel, P., and Levine, S. Rein- forcement learning with deep energy-based policies. In International conference on machine learning, pp. 1352- 1361. PMLR, 2017.
Kumar, A., Zhuang, V., Agarwal, R., Su, Y., Co-Reyes, J. D., Singh, A., Baumli, K., lqbal, S., Bishop, C., Roelofs, R., et al. Training language models to self-correct via reinforcement learning, arXiv preprint arXiv:2409.12911 2024.
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor-critic: Off-policy maximum entropy deep reinforee- ment learning with stochastic actor. In International conference machine learning, pp. 1861-1870. PMLR, 2018.
Lai, X., Tian, Z., Chen, Y., Yang, S., Peng, X., and Jia, J. Step-dpo: Step-wise preference optimiza- tion for long-chain reasoning of 1lms. arXiv preprint arXiv:2406. 18629, 2024.
Hazan, E. et al. Introduction to online convex optimization. Foundations and Trendso in Optimization, 2(3-4):157- 325, 2016.
Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., and Chen, Z. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006.1666 2020.
Henderson, P., Hu, J., Romoff, J., Brunskill, E., Jurafsky, D., and Pineau, J. Towards the systematic reporting of the energy and carbon footprints of machine learning. Journal of Machine Learning Research, 21(248):1-43, 2020.
Liang, T., He, Z., Jiao, W., Wang. X., Wang, Y., Wang, R., Yang, Y., Shi, S., and Tu, Z. Encouraging divergent thinking in large language models through multi-agent debate. arXiv preprint arXiv: 18/ 2023.
Hendrycks, D., Basart, S., Kadavath, S., Mazeika, M., Arora,
A., Guo, E., Burns, C., Puranik, S., He, H., Song, D., et al. Mcasuring coding challenge competence with apps. arXiv preprint arXiv:2 10..0 (09.99938,... 2021a.
Lightman, H., Kosaraju, V., Burda, Y., Edwards, H., Baker, B., Lee, T., Leike, J., Schulman, J., Sutskever, I,, and Cobbe, K. Let's verify step by step. arXiv preprint urXiv:2305.00550 2023.
Hendrycks, D., Burns, C., Kadavath, S., Arora, A., Basart, S., Tang, F., Song, D., and Steinhardt, J. Measuring math- ematical problem solving with the math dataset. arXiv preprint arXiv:2103.03874 2021b.
Liu, E. 7 Raghunathan, A., Liang, P., and Finn, C. Decou- pling exploration and exploitation for meta-reinfooreement learning without sacrifices. In International conference machine learning, pp. 6925-6935. PMLR, 2021.
Huang, J., Chen, X., Mishra, S., Zheng, H. S., Yu, A. W., Song, X., and Zhou, D. Large language mod- els cannot celf-correct reasoning yet. arXiv preprint 10.0077 2023.
Lorraine, J., Vicol, P., and Duvenaud, D. Optimizing mil- lions of hyperparameters by implicit differentiation. In
10

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Er- Madaan, A., Tandon, N., Gupta, P., Hallinan, S., Gao, L., Wiegreffe, S., Alon, U., Dziri, N., Prabhumoye. S., Yang. Y., et al. Self-refine: Iterative refinement with self- feedback. Advances in Neural Information Processing Systems, 36, 2024.
International conference on artificial intelligence and statistics, pp. 1540-1552. PMLR, 2020.
mon, S., and Finn, C. Direct preference optimization: Your language model is secretly a reward model. Ad- vances in Neural Information Processing Systems, 36, 2024.
Ray, A., Achiam, J., and Amodei, D. Benchmarking safe ex- ploration in deep reinforcement learning. arXiv preprint arXiv: 1910.01708, 7(1):2, 2019.
Manne, A. S. Linear programming and sequential decisions. Management Science, 6(3):259-267, 1960.
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q. Hinton, G., and Dean, J. Outrageously large neural networks: The sparsely-gater mixture-of-exppee layer. arXiv preprint 3oXiv:/701.06533 2017.
Montgomery, W. H. and Levine, S. Guided policy search via approximate mirror descent. Advances in Neural Information Processing Systems, 29, 2016.
Skalse, J., Howe, N., Krasheninnikov, D., and Krueger, D. Defining and characterizing reward gaming. Advances in Neural Information Processing Systems, 35:9460-9471, 2022.
Nachum, O. and Dai, B. Reinforcement learning via fenchel- rockafellar duality. arXiv preprint arXiv:2001.01880 2020.
Slivkins, A. et al. Introduction to multi-armed bandits. Foundations and Trends@ in Machine Learning, 12(1- 2):1-286, 2019.
Nachum, O., Chow, Y., Dai, B., and Li, L. Dualdice: Behavior-agnostic estimation of discounted stationary distribution corrections. Advances in neural information processing systems, 32, 2019.
Snell, C.,, Kostrikov, I., Su, Y., Yang, M., and Levine, S. Offline rl for natural language generation with implicit language learning. arXiv preprint arXiv:2206. 11871, 2022.
Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730-27744, 2022.
Snell, C., Lee, J., Xu, K., and Kumar, A. Scaling Ilm test- time compute optimally can be more effective than scal- ing model parameters. arXiv preprint axXiv22408.03314, 2024.
Pan, A., Bhatia, K., and Steinhardt, J. The effects of reward misspecification: Mapping and mitigating misaligned models. arXiv preprint arXiv:2201.035-4 2022.
Sutton, R. S. Reinforcement learning: An introduction. A Bradford Book, 2018.
Pathak, D., Agrawal, P., Efros, A. A., and Darrell, T. Curiosity-drine -riven exploration by self-supervised predic- tion. International conference on machine learning, pp. 2778-2787, PMLR, 2017.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi,
A., Babaei, Y.. Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine- tuned chat models. arXiv preprint arXiv:2307.09288 2023.
Peng, X. B., Kumar, A., Zhang, G., and Levine, S. Advantage-weighter regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177 2019.
Vazirani, V. V. Approximation algorithms. Georgia Inst. Tech, 1997.
Virtanen, P., Gommers, R., Oliphant, T. E., Haberland, M.,
Peters, J., Mulling, K.. and Altun, Y. Relative entropy policy search. Proceedings of the AAAI Conference on Artificial Intelligence, volume 24, pp. 1607-1612, 2010.
Reddy, T., Cournapeau, D., Burovski, E., Peterson, P., Weckesser, W., Bright, J., van der Walt, S. J., Brett, M., Wilson, J., Millman, K. J., Mayorov, N., Nelson, A. R. J., Jones, E., Kern, R., Larson, E., Carey, C. J., Polat, i.. Feng, Y., Moore, E. W., VanderPlas, J., Laxalde, D.. Perktold, J., Cimrman, R., Henriksen, I,, Quintero, E. A., Harris, C. R., Archibald, A. M., Ribeiro, A. H., Pedregosa, F., van Mulbregt, P., and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-272, 2020. doi: 10. 10.10038811999 296---99666--222
Pham, C., Liu, B., Yang, Y., Chen, Z., Liu, T., Yuan, J., Plummer, B. A., Wang. Z., and Yang, H. Let models speak ciphers: Multiagent debate through embeddings. arXiv preprint arXiv:2310.0627 2023.
Qu, Y., Zhang, T., Garg, N., and Kumar, A. Recursive introspection: Teaching language model agents how to self-improve. arXiv preprint arXiv:2407. 18219, 2024.

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
Wang. X., Wei, J., Schuurmans, D., Le, Q. Chi, E., Narang, S., Chowdhery, A., and Zhou, D. Self-consistency im- proves chain of thought reasoning in language models. arXiv preprint arXiv:2203.11171 2022.
Zhang, Y., Khanduri, P., Tsaknakis, I., Yao, Y., Hong, M., and Liu, S. An introduction bilevel optimization: Foun- dations and applications in signal processing and machine learning. IEEE Signal Processing Magazine, 41(1):38- 59, 2024.
Wang, X., Feng, S., Li, Y., Yuan, P., Zhang. Y., Pan, B., Wang, H., Hu, Y., and Li, K. Make every penny count: Dilficult-adaativve self-consistency for cost-efficient rea- soning. arXiv preprint arXiv:2408. 13457, 2024.
Ziebart, B. D. Modeling purposeful adaptive behavior with the principle of maximum causal entropy. Carnegie Mel- lon University, 2010.
Wei, J., Wang, X., Schuurmans, D., Bosma, M., Xia, F., Chi, E., Le, Q. V., Zhou, D., al. Chain-of-though prompting clicits reasoning in large language models. Advances in neural information processing systems, 01224224-24877 2022.
Ziegler, D. M., Stiennon, N., Wu, J., Brown, T. B., Radford, A., Amodei, D., Christiano, P., and Irving, G. Fine-tuning language models from human preferences, arXiv preprint arXiv: 005999,,,, 2019.
Zinkevich, M. Online convex programming and generalized infinitesimal gradient ascent. In Proceedings of the 20th international conference machine learning (icml-03), pp. 928-936, 2003.
Wei, T., Zhu, B., Zhao, L., Cheng, C., Li, B., Lü, W., Cheng, P. Zhang. J., Zhang, X., Zeng, L., et al. Skywork-moe: A deep dive into training techniques for mixture-of-experts language models. arXiv preprint arXiv:2406.065633, 2024.
Xu, M., Li, Y., Sun, K., and Qian, T. Adaption-of- of-thoughhtt Learning question difficulty improves large language models for reasoning. In Proceedings of the 2024 Con- ference on Empirical Methods in Natural Language Pro- cessing, pp. 10665,545,.2 2024a.
Xu, T., Helenowski, E., Sankararaman, K. A. Jin, D., Peng, K., Han, F.,, Nie, S., Zhu, C., Zhang, H., Zhou, W., et al. The perfect blend: Redefining rlhf with mixture of judges. arXiv preprint arXiv:2409.20370 2024b.
Yan, Y., Jiang, J., Liu, Y., Cao, Y., Xu, X., Cai, X., Shao, J., et al. c-c-math: Spontaneous step-level self-correction makes large language models better mathematical reason- ers. arXiv preprint 4rX0v:2409.01524, 2024.
Yang, T.-Y., Rosca, J., Narasimhan, K., and Ramadge, P. J. Projection-based constrained policy optimization. arXiv preprint arXiv:2010.03152, 2020.
Yu, Z., Tao, Y., Chen, L., Sun, T., and Yang, H. B-coder: On value-based deep reinforcement learning for program synthesis. In The Twelfth International Conference Learning Representations, 2023.
Zelikman, E., Wu, Y., Mu, J., and Goodman, N. Star: Boot- strapping reasoning with reasoning. Advances in Neural Information Processing Systems, 35:15476-15488 2022.
Zemel, R., Wu, Y., Swersky, K., Pitassi, T., and Dwork, C. Learning fair representations. In International conference on machine learning, pp. 325-333. PMLR, 2013.
Zhang, Y., Vuong, Q., and Ross, K. First order constrained optimization in policy space. Advances in Neural Infor- mation Processing Systems, 33:15338-155 5338-15349,,,,1 2020.
12

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
A. Motivating Example
Figure is an example that advanced reasoning model spent more than enough time on a trivial problem.
what is plus
Thought for seconds
plus equals
Figure A long reasoning-chan -hain model spent more than enough inference time on trivial problem.
B. Full Version of Section 4: Acronyms, Naïve Construction of g & Training Pipelines
Yet we work on abstract groups Go and G+- In this section, we present the details of our constructions of extended length responses, i.c. the extended group G4. However developing long reasoning models is beyond the scope of this work, as our focus is on the constrained optimization of LLMs. Our constructions are for demonstrative purpose only. Due to the intricate details involved in prompts, datasets, and training pipelines, this section may appear somewhat dense. To make more approachable, we have structured our writing in way that readers can, if they wish, focus on the broader ideas without delving deeply into the specifies of constructions. A TL;DR version of this section is provided below.
TL;DR. We construct two types of illustrative extended responses: Sequential Voting (SV) and Adaptive Sequential Voting (ASV). Figure 4 visually explains how these constructions are implemented. The goal of the SV is to establish a baseline that generates only responses in G.., thereby serving as an uni-modal comparator. SV scales roughly as well as vanilla majority voting (MV), aka coniisonssy (Wang et al., 2022). In contrast, ASV outputs a mixture of responses of y E G. and y E G4. This allows the model to adaptively decide which type of response to produce based on the query. The goal of ASV to further enable IBPO optimization, as IBPO implicitly assumes the model generates both regular and extended-length responses. In Section 5, we show that ASV optimized by IBPO, achieves better allocation of the inference budget.
B.1. Construction of Sequential Vote
Acronyms. For clarity, we explicitly define key terms to hopefully resolve any potential ambiguities. Response: response refers to a sequence generated until a terminal token is encountered. For precision, we sometimes refer to these as voting responses SCoT responses, as illustrated in Figure 4, after introducing our sequential voting baselines. Trial: A trial denotes a solution instance, which is demarcated by the special tokens [TRIAL] ... [/TRIAL], as shown in the voting response example in Figure 4. While voting response contains multiple trials, a SCoT response or a non-voting response contains exactly one trial, as also depicted in Figure 4.
For the description of training and testing details, we use LLaMA and LLaMA-b to denote the instruction and base versions of the LLaMA 3.1 8B models (Dubey et al., 2024), respectively. MATH refers specifically to the training split of the Hendrycks MATH dataset (Hendrycks al., 2021b), while the 500-sample subset of the testing split is referred to as MATH500 (Lightman et al., 2023). SDPO stands for step-DPO (Lai al., 2024) dataset, a curated step-annotated dataset from which we retain only the prompts and positive responses (ground truth solutions), excluding any step signals (see Appendix G for details). It is important to note that while we leverage their curated dataset, the original SDPO method is not relevant to this work. The SDPO dataset was chosen because its ground truth responses follow the SCoT format of L.LMMA responses, making it convenient to run supervised fine-tuning (SFT) mixed with LLaMA samples.
Construction details. To be more specific, the naive sequential voting baselines, as the "expensive" group G4. have increased inference costs by simply sequentially output multiple trials and find the sequential majority vote until stop condition met. In particular, the early-stopping sequential voting (SV) baseline is created to show that such naïve baseline could achieve performance gain, on par with vanilla majority voting (MV). The adaptive sequential voting (ASV) allows model to output both SCoT response yo and sequential voting response y+, allowing us to further conduct the budget controll experiments in Section 5.
Early Stopping Sequential Vote (SV): For single response, model are allowed to output at most 8 trials, and conclude
13

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
SCoT Prompt + Cond.
ASV Cond.
Solve following math problem efficiently and clearly:
For medium and hard level problems, you asked to give most eight diverse solutions in different way, without referencing to previous trials.
For simple problems (2 steps or fewer):
Provide concise solution with minimal explanation.
For complex problems (3 steps more):
Each trial should be contained a separate [TRIAL] trial solution [/TRIAL] block.
Use this step-by-step format:
## Step I: [Concise description]
If solution occurs three times, it is considered consen- and will be used the final answer.
[Brief explanation and calculations]
## Step 2: [Concise description] [Brief explanation and calculations]
there no consensus, use solution from most plausible trial.
For easy level problems. you are allowed only one attempt, which will be considered your final answer.
Regardless of the approach, always conclude with: Therefore, the final answer is: $1 (anxeedaansserr answer]S. hope it is correct.
Where [answer] is just the final number expression that solves the problem.
Voting Response (SV, ASV Case 1)
[Condition Description]
[TRIAL] ## Step steps omitted
Problem:
{{ problem }}
The final answer is:
A1
[/TRIAL.]
[TRIAL] ## Step
steps omitted
SCoT Response
The final answer is:
A2 [/TRIAL]
[TRIAL] ## Step
steps omitted
## Step I: steps omitted
The final answer is:
A1 [/TRIAL]
The final answer is:
Al
hope is correct.
[TRIAL] ## Step
steps omitted
The final answer is:
A1 [/TRIAL]
SV Cond.
The answer A1 has occurred three times, and is considered as consensus.
The final answer is A1 hope it is correct.
You are asked to give most eight diverse solutions in different way, without refereneing to the previous trials. Each trial should be contained in a separate [TRIAL] trial solution [/TRALL] block.
Non- Non-Voting Response (ASV Case 2)
If solution occurs three times, it is considered consen- sus and will be used the final answer.
[TRIAL] ## Step
steps omitted
there no consensus, use the solution from most plausible trial.
The final answer is: Al [/TRIAL] Terminated due to difficulty level.
Figure 4: Prompt templates. For Go. we use the SCoT prompt without including an additional [condition description] and generate standard SCoT response. For 94, simply insert the corresponding condition into the [condition description] placeholder to create new prompt. In the case of SV, the SV condition is integrated into the prompt, and the model is asked to perform repeated trials to reach consensus. The ASV prompt instructs the model to output either a voting response (case 1) or a non-voting response (case 2), with the decision made by the model itself.
with majority of trials. In addition to the terminal condition of maximum 8 trials, model will early stop if an answer appears times and this answer will be considered as the majority answer.
Adaptive Sequential Vote (ASV): Model is allowed to choose either vote (case 1, i.e. 3/) or not (case 2, i.e. yo). This baseline is created to further allow model-driven resource allocation. We later, in Section 5, show that one could optimize the ability of resource allocation with our IBPO.
Dataset. We define our construction of dataset as a product of problem set Q, format template of question T, and answer Ta, and response set A, subjected to some filtration F. Formally, dataset ID is defined as: DD (F Tr o (Q A), where F removes undesired question-answer pairs, such as incorrect responses when A are model generated; The templates Tr and T. collectively transform each question-answer pair into specific text format, as shown in Figure 4; Slightly abusing the notation, the Cartesian product Q x A is used to pair each response with its corresponding question, defined as: Q x A= = {(g, azz): Vi,j).
Specifically, we summarize the datasets used in subsequent sections Table For instance, for SV training, we construct
14

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
Table 4: Summary of constructed datasets for different experimental purposes. Dav. Esss and Asv uses the same set of prompts but different prompt/response tem-
Table 5: Stopping conditions.
stopping conditions
plates. QSDEO and ASDrO Agoedee are from Lai al. (2024) with details deferred to Appendix G. ASVI and ASV2 correspond ASV case and (see Section B.1) re- spectively.
SV
(i) max trials; (ii) answer occurs times.
ASV
voting (case 1): (i) max 8 trials; (ii) if answer
occurs times; non- non-voting (case 2): exact trial.
Set
Ty
Ta
F
Q
A
Table 6: Prompt and response sources.
Dsy
SV
SV
G
QMATH
SAMPLE MATH
ID Asvi
ASV
ASVI
Set
Query-Respont
Source
Dscor
D Asv2
ASV
SCoT
SCoT
ASV2
QSPPO
A GOLDEN SDPO
Training
Qsoro
QMATT & A
ASSPPO
MARATEE MATH & LLaMA samples
GOLDEN
(Lai et al., 2024)
Testing
QMARE1O0
MATH500
DRL
ASV
QSDPO
0
Table 7: Training pipelines. For Sec. B.2, we aim to create a demonstration experiment showing that with the same model (of roughly same math knowledge) SV can achieve reasonable performance-cost efficiency on par with MV. optimize performance-cost efficiency through our luB generalization of CGPO, where in row 2.1 is coefficient of Asv2.
Sec. 5 further show
Exp.
Sec.
Type
Init. Ckpt.
Dataset
Purpose
1.1
Sec. B.2
SFT
LLaMA-b
Dsy U Dssorr
Allow model follow both SV and SCoT prompt.
1.2
Sec. 5
SFT
LLaMA
2.1
SFT
LLaMA Exp. 2.1
TAsVIi U UaD
Follow the ASV instruction to let model decide vote or not. Optimize the capability of dynamic budget allocation.
Sec. 5
2.2
RL
DRR.
MV (SCoT)
MV (SCoT)
0.45
0.45
Nested MV (SV)
Nested MV (SV)
(nested) malority@x
0.40
(nested) majority@@
0.40
0.35
0.35
0.30
0.30
0.25
0.25
103.0
1025
10
1045
10
number of responses (K) per problem
average number of tokens per problem
(a) (Nested) MV measured with number of responses
(Nested) MV measured with number tokens
Figure 5: SV tested MATH500. MV stands for vanilla majority voting, aka self-consistency (Wang et al., 2022), with SCoT responses. Nested MV refers to majority voting with our (early-stopping) SV responses. It "nested" each SV response is already voting, as shown in Figure 4. The SV method has "clear" gains when performance is measured with the number of responses. When measured with the number of tokens, SV aligns the performance-cos efficiency of vanilla MV. This is another indicator that one should worry about token efficiency when measuring reasoning performance.
Dyy with question set QMATH from Hendrycks (2021b) and LLaMA generated responses SAMPLE using the SV templates MATH
15

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
defined in Figure 4, subjected to data selection described in Appendix G.
Training pipelines. We summarize our training pipelines for toy experiments in Section B.2 and our IuB experiments in Section 5. For instance, experiment 2.2 (Section 5) in Table 7 is the summary of our IBPO training, using dataset Dar as constructed in Table 4 and initialized from ASV models from experiment 2.2. Further details of each training can be found in Appendix G and F.
B.2. Performance of Naive Sequential Vote
We start by evaluating the performance of SV compared to vanilla MV, following the setup of Exp. 1.1 in Table 7. This toy experiment is designed to demonstrate: (i) SV scales approximately as well as MV, thereby qualifying it as an example of 94: and (ii) measuring performance based on the number responses inadequate, therefore we later measure performance relative to the number of tokens/trials in Section
are pass@k & and majority @k both of which widely used in the literature (Hendrycks et al., 2021b; Wang al., 2022). We occasionally refer to pass rate as pass 1. In both metrics, k specifically denotes the number of responses, regardless of the number of trials per response. Since our voting methods may involve multiple trials, we may use the average number of trials (as illustrated in Figure 2) on the x-axis for cost-aware comparisons. In addition, we evaluate "performance-cost efficiency by comparing each method's scaling efficiency to that of MV, following the comparison Snell et al. (2024).
Metries, The metrics
As shown in Figure 5a, SV exhibits significant inprovementt--nearyy 10%-in terms majority @k when performance is measured by the number of responses. However, this improvement is misleading since each SV response effectively consists of multiple SCoT responses. Therefore, measuring performance with taking length into consideration provides more reasonable assessment. As demonstrated in Figure 5b, SV's scaling performance aligns with that of vanilla MV, supporting its role as a suitable construction example for G+.
C. Batch Accumulation
This is batch accumulation implementation of Algorithm 1. For brevity, we use OPT as example. One could increase the optimization problem size of to (n ls)) x (m ky)) using Algorithm 2, where superseripts indicate matrix shape, left subscripts denote accumulation indices (distinguishing them from element indices).
Algorithm 2 IBPO with Sample Accumulation
Require: prompt set D, batch size n, number of responses m, init policy To = Tref, num of iters T, budgets g+ for 1, ......, do
2: for 1, ,,... Kg do
3: prompt sampling: X pmm
4: for = L...,..... ky do
5: response generation: Yaxm To, To, (xX)
6: end for
7: end for
8: prompt accumulation: (n) 1 [x,,, 2 X]
9: response accumulation: (mm-k)) In Y, 12
1k. Y: kg Y k22 Y lak. Y]
10: evaluate correctness and empirical KL: Rmatch = r Tmatch X Y and KL KL(Y;; X, Trof)
11:
margin maximization: # E OPTTIE X, Y Rmachh, KL, (+) as defined in Eq. (8)
12: gradient update: with EEh E5=1 # y (Mis (x;) 00 log No ( le=e.
ks m.
end for
16

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
D. Sample Responses
show couple of example responses of ASV-luB-50% the omitted prompt could be found in Figure 4.
Voting Example
Prompt omitted
Problem: Compute arcsin (-1) Express your answer radians.
[TRIAL] ## Step I: Recall the definition of aresine The function aresin is the inverse of the sine function. It returns the angle whose sine is x. ## Step 2: Use the unit circle to find the angle From the unit circle, we know that sin (- ; 1 we have arcsin 11) =
2
## Step 3: Determine the value aresin (-4) Since sin (-5) =
The final answer is:
[/TRIAL]
[TRIAL] ## Step Recall the sine function values for common angles. The sine is v3
##l Step 2: Use the relationship between sine and aresine. Given (0) 2. are looking for aresin (-1).
## Step 3: Identify the angle whose sine is - 3 From the unit circle or trigonometric identities, we know that sin (-a)
## Step 4: Express the answer in radians. Therefore, aresin (-1)
Therefore, the final answer is: π
[/TRIAL]
[TRIAL] ## Step I: Recall the sine function's 's relationship with the unit circle. The sine of an angle in right triangle is the ratio of the length of the side opposite the angle to the length of the hypotenuse. For angles in the unit circle (a circle with radius 1), the sinc of angle is the y-coordinate of the point where the terminal side of the angle intersects the circle.
## Step Identify the angle with sinc 2 From the unit circle, know that sin (g) third and fourth quadrants, and are looking for angle whose sine is the angle must in either the third fourth quadrant.
Since sine negative in the
## Step 3: Determine the angle the correct quadrant. Considering the unit circle, the reference angle for 5 is itself. Since we looking for angle with sine of - and knowing that sine is negative in the third and fourth quadrants, the angle are looking for is in the fourth quadrant hecause sin ( (TT) would indeed be negative. ## Step 4; Calculate the angle. The angle whose sine is is be 3
5 3
Therefore, the final answer is:
[/TRIAL]
[TRIAL] ## Step Recall the definition aresine The function aresin is the inverse of the sine function. returns the angle whose sine is x.
## Step 2: Use the sine function to find the angle We want to find the angle whose sine is - We can use the sine function to find this angle.
## Step 3: Recall common angles and their sines From the unit circle trigonometric identities, know that sin (:)
Since sine is an odd function, sin (-ē) =
## Step 4: Determine the angle Therefore, the angle whose sine is z is
Therefore, the final answer is:
[/TRIAL]
[TRIAL] ## Step Recall the sine function's values for common angles. The sine function has well-known values for certain angles. Specifically, sin
## Step 2: Use the symmetry and periodicity of the sine function. Given that sin (?) 2. and knowing that sine is an odd function, we have (- 7) -.. This step utilizes the property
functions where f( [(-xr = =-(rr).
## Step 3: Solve for arcsin (- (- Given that sin (-ă) =) =-g.
z. the angle whose sine is - This because the inverse sine function. arcsin, returns angle whose a given value.
Therefore, the final answer is:
[/TRAL]
The has occured three times and is considered consensus. The final answer
hope it is correct.
17

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
Non-voting Example
Prompt omitted
Problem: If fx) 3x-2 what is the value of f(-2) + f(-1) f(0)? Express your answer as a common fraction. z-2
[TRIAL.] ## Step Substitute -2 into the function fx). find f(-2) substitute x -2 into function:
f(-2) = 3(-2) -
6 -
8 2.
-2-2 -
-4
## Step 2: Substitute x = -1 into the function f(z). To find f(-1), substitute x = -1 into the function:
S(-1) 3(-1)
-3
-5 3
-1-2
-3
## Step 3: Substitute 0 into the function f(x). To find f(0), substitute = 0 into the function:
3(0) -
0-2
f(0)
0 0 -2 -2
## Step 4: Sum the values of f(-2). f(-1), and f()). Sum the calculated values:
f(-2) + f(-1) f(0) 2
## Step 5: Simplify the sum to express it as common fraction. To combine the terms into a common fraction, convert the whole numbers to fractions with common denominator:
6
5
6
3
14
2
+
3
3
3
3
3
Therefore, the value of f(-2) + f[-1) + f(0) is 14
3.
14
Therefore, the final answer is:
[/TRIL]]
3
Terminated due difficulty level.
E. Training Curves & Further Discussions
Budgets constraints. In Section 5, we show that the constraints hold exactly on the training set, for responses that are correct. Figure 6b further shows the ratio of voting responses for online samples. It's clear that the constraints also hold.
ASV-luB-75%
1.0
ASV-luB-25%
ASV-luB-50%
ASV-luB-50%
voting responses fraction
0.8
ASV-luB-25%
0.8
ASV-luB-75%
WW
Vote fraction
A
0.6
0.6
m
0.4
0.2
0.0
0.0
50
100
150
200
250
50
100
150
200
250
training step
training step
(a) Voting ratio of correct responses (training set)
(b) Voting ratio of all responses (training set)
Figure 6: Voting response ratio versus training steps. Dashed line denotes the budget 9+. On the training set, IuB formulation follows the budget constraints almost exactly for both: (a) correct responses; (b) all responses.
18

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
CGPO DRL w/ LLaMA. To further understand the results presented in Section is benfit from our budget-aware formulation from the prompt set of Lai al. (2024). We further run CGPO with open-sourced LLaMA model and the SDPO dataset in the SCoT format. Figure 7 compares the training dynamics of CGPO with LLaMA versus our ASU-luB-q4 experiments. In general, the SDPO prompt set does not provide much additional knowledge as suggested by the OSS w/ vanilla CGPO experiment, but ASV-luB-0+ experiments are able to achieve noticable gain.
0.56
0.54
pass@1
0.52
0.50
ASV-luB-75%
ASV-lu--50%
0.48
ASV-IuB-25%
0.46
LLaMA 8B Instruct
50
100
150
200
250
training step
Figure 7: Training curves: each point corresponds evaluation of MATH500 test set, and dashed line is the pass @ 1 of LLaMA 8B Instruct.
Controlled setting, Our experiment setting is minimal: we do not use reward models; we use only an 8B model to generate any set of sampled responses, A SAMPLE Our RI. training set, DR.. contains only 10k prompts, leaving the setting quite controlled and providing room for improvement through engineering efforts.
Limitations. Our work is limited its choice of RL algorithms and applications. While should be straightforward apply framework to different RI. frameworks, additional effort is required to derive an optimization problem for different application tasks. Future work may include extending our constrained policy optimization framework to broader range of LLM applications or scaling up our experiments.
19

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
F. Hyperparameters
We list the hyperparameters used for experiments setup 1.2, 2.1, with NVIDIA-A100-80GS. (Please refer Xu al. (2024b) for the definition of some RL-specific
2.2, as described Table 7. And we conduct our experiments
Table 8: Hyperparameters for Experiment setups 1.2, 2.1, and 2.2
Hyperparameter
Sctup 1.2
Setup 2.1
Setup 2.2
prompt size
11295
11295
10795
number of nodes learning rate batch size (per node)
4
8
le-6
le-6
5e-7
16
4
num of steps optimizer
1024
2048
240
AdamW
A
A
scheduler
constant
constant
constant
packing
yes
yes
max sequence length
32768
32768
6144
gradient accumulation
2
1
RL-specifi params
num generation per prompt
8
max generation length temperature
4096
1.0
top-p
0,9
KL-threashold
1024
batch accumulation ks response accumulation k.-
4
1
G. Data Construction
We elaborate the details of @SDRO, ADDO QMATH A AMNT and the filteration function F.
SDPO dataset. Lai et al. (2024) create dataset (see their Section 3.2) contains 10,795 step-wise preference pairs, where the winning response is a correct solution, the losing response is incorrect starting from certain step. For our purpose of GOLDEN with only wining (correct) responses. The SFT and RL training, we do not need losing responses, hence we create AgDED step-formatted response (see an example below) is the reason we choose the dataset from Lai et al. (2024), as adhere the SCoT format of LLaMA models. Therefore, do not need to worry about potential format changes.
SDPO question-response example
Problem:
Sides of a rectangular park are in the ratio 3: 2 and its area 3750 sq the cost of fencing it at 40 ps per meter is?
##ll Step 1: Let's assume the sides of the rectangle are 3r and 2r. The area of the rectangle is the product of its sides, so we have
3x 2x = 6z 3750. Solving for v, we get x² 3750 625, so x = / 625- 25.
## Step 2: So the sides of the rectangle are 3r = 3 25 = 75 meters and 2x = 2 25 = 50 meters.
## Step 3: The perimeter of the rectangle is 2 x (75 50) 2 x 125 250 meters.
## Step 4: The cost of fencing is 40 paise per meter, the total cost of fencing is 250 40 10000 paise. Since rupee 100 rupees.
equal to 100 paise, the cost in rupees is 100b0
100
Therefore, the final answer is: 100
SV dataset for MATH. We also created datasets, Dsy and DASVI, for SFT so that a model could follow the SV instructions. We take D&v as an example and DASvI could created similarly. To do so, we first generate 32 responses per prompt for the entire MATH training split with temperature of 1.2 and top-p of 0.9. We then apply our SV templates Ty and Ta to create corresponding SV question and answer pairs. The procedure of creating SV response is given by Algorithm 3. While we have include an example of SV responses in Figure 8, we make it more concrete the SV response template below,
We then apply some filtration F to the created dataset. We remove SV responses whose final answers are incorrect. From the remaining set, we sub-sample 500 question- response pairs. These pairs are selected from problems that have between 4 and distinct answers out of 32 samples, ensuring that we construct sequential responses with a diverse trials. The distribution of trial counts and distinct answer counts are shown Figure It can be observed that the filtered data is diversely distributed.
20

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
SV template: (A Ag, final answer) (if con- sensus found)
SV template: To(A1, ... As final answer) (if no consensus found)
[TRIAL] } [/TRIAL]
[TRIAL] { [/TRIAL]
[TRIAL] {A2} [/TRIAL]
[TRIAL] {Aa} [/TRIAL]
[TRIAL] [Ax] [/TRIAL]
[TRIAL] {Aa} [/TRIAL]
answer final answer has occurred three times, and is considered as consensus.
Maximum trials reached no consensus found to tie; the most plausible answer is final answer
The final answer is final answer
hope it is correct.
The final answer final answer
hope it correct.
Figure 8: SV response templates. Left: suppose consensus is found k-th answer; Right: no consensus found. Note the subscript of A, only denotes index of answer, A, and A, could still have same final answer for + j.
Algorithm 3 Creating SV Response for SFT.
Require: template Tg. Tas problem Q, set of shuffled responses (A, 1,2,
.K}
1: create prompt: [Q)))
replace problem placeholder with Q
2: responses = []: final final_answer= ansswer= INNAALID found = False
3: for = 1, 2. K do
4: responses.append(AA)
5: majority = find majority(responsss))
6: majority_counl = count manommmaorittyeeeppnnees majority)
if majority
3 then
8: found = True; imal-answer majority; break
9: end if
10: end for
11: if found == False and responses contain correct solution then
12: final answer = random pick a correct solution
13:
end if
create SV response: (responses, final answer)
frequency
frequency
60
40
20
6
7
8
1
3 4
7
number of trials
number of distinct answers
(a) Distribution number of trials per response
(b) Distribution of number of distinct answers
Figure 9: The distribution of filtered prompt-responses subset, which suggests that the construct data is generally diversely distributed.
21

---

Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
H. Integer Linear Program Solving
use SciPy MILP solver, available here, to solve an integer LP every iter. Figure a Pythonic pseudo-code. And wall-time consumed could be found in Figure 10. As the problem size small, so the computational overhead negligible,
import numpy as np From optiiiizz import milp, LinearConstraint, Bounds
def solve_iut (acceptance, is_vote, b
=** solves Inference under Budget (TuB) problem. Parameters:
acceptance: an n x m array where each element is 1 if accepted, 0 otherwise.
is_vote: an x array indicating whether response is voting (2) or non-voting (C)
budget: the fractional budget (q+) constraint for the problem. n, acceptance.. shape
10
calculate pass rates for vote-based and non-vote-based responses vote_pass_rrae mp. (acceptance + is_vote, axis=1, keepdims-Truet
11
sample_pass_rate _atte mean (acceptance (1 - is_vote), axis-1, keepdims-True)
margin vote_pass_rate
sample_pass_ra
(n, 1)
14
flattern acceptance and vote indicator, tile the margin acceptance np. reshape (acceptance, -1)
15
# (o x m
16
vote = np. reshape (is (is_vote, -1)
# (n
margin ap. reshape (np. tile (margin, 11, m]), -1)
m,
18
19
define the objective function coefficients
20
= -1 margin is_vote margin * (1 -s__vte)
# acceptance constraints: ensure each prompt meets acceptance criteria A_acceptance np.eye (len (acceptance))
22
23
b_ p_acceptance acceptance
24
one response per problem constraint (BoN)
25
problem np. zeros (n, m))
for in range (n)
A_problemii, b_problem oneo
mi
m]
28
29
voting responses budget constraint
30
A vote budget whewheee (is_vote vote == 1, 1, 0) reshape (1, -1)
vote budget np. round (budget len (accepttance))
32
33
combine all constraints into a single matrix
34
A np. vstacckkk/ aA_cccptaacc,, A_problem, A_voto_budget)
b_lower inf np. opeones (A. [0]) lower bounds for constraints
36
upper np. hstack ( acceptance, b_problem, cote_budget]) solve the MILP problem using the defined objective and constraints
upper bounds
37
38
result milp(e, Integrality=np.oom (len (c)), ) bounds=Bounds (0, 1), (A, b_lewer, b_upper))
constraintssznnercossttte
return result
Listing I: Pythonic code snippet for solving luB with SciPy: Note that this is for demonstration purposes, and error-free execution is not guaranteed, due to omitted corner cases.
0.010
ASV-luB-75%
wall time (sec)
ASV-IuB-50% ASV-luB-25%
0.008
0.006
0.004
0.002
100
150
200
250
training step
Figure 10: Wall time (averaged across ranks) spent by solver seconds.
22


Reviews:
[{"review_id": "gnlN3Y2KYQ", "reviewer": "Reviewer_toYu", "summary": "With the prevalence of Chain of Thought (CoT) in complex reasoning and the emergence of ultra-long reasoning models such as OpenAI-o1 and DeepSeek-R1, unnecessarily tedious and long generations for trivial problems are increasingly becoming a problem. The paper approaches this problem from a RL perspective, proposing IBPO (Inference Budget-Constrained Policy Optimization) that, rather than simply taking a metric of context length as reward, forms a constraint-RL by controlling how response lengths are distributed. The algorithm is then deduced to be a generalization of iterative SFT. Experiments and positive results are shown.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your thorough review and valuable feedback! We greatly appreciate your insightful suggestions and positive assessment!\n\n---\n\nWe address these three comments collectively, given their relevance to one another.\n\n> However, I'm still hesitant to accept that responses constructed this way are actually long/short-context responses -- the long ones look like a mere aggregation of short responses.\n\n> One major confusion is whether prompts generated via SV sufficiently represent the \"unnecessarily tedious long reasoning trace\" that the paper sets out to improve. This is rather important as SV seems to be the only experiment scenario.\n\n> On the paper structure, maybe it's a better idea to include more explanation on prompt generation (the proposed SV, ASV), not only to clarify the experiment setting more, but to also convince readers that SV is adequate in representing long response situations that the paper initially tries to improve.\n\nWe understand the concern that SV may not fully capture how o1/R1 works at a fundamental level. This construction was chosen due to the unavailability of R1, which was released on Jan 20—roughly 10 days before the submission deadline.\n\nReproducing o1/R1-style longCoT, or anything more similar, was beyond the scope of our focus on constrained LLM inference. Moreover, it was unclear to the community how to reproduce such methods until the R1 technical report [1] was released on Jan 22.\n\nThe SV construction resembles two key characteristics of longCoT: (i) improved accuracy, and (ii) extended response length. It therefore serves an illustrative purpose to highlight our core RL contribution.\n\nAdditionally, our paper (i) casts the problem in terms of distribution constraints, rather than directly optimizing over the long-response group, and (ii) derives an RL method that is theoretically agnostic to the specific type of long response. These efforts together further support the use of a constructed alternative.\n\nGiven these considerations—and the fact that o1/R1-style longCoT was not available at the time—we hope you find the use of this construction reasonable.\n\nWe completely agree that devoting more space to explaining the design decisions and clarifying the role of SV would significantly improve the paper's clarity. Thank you for pointing this out—we will revise the manuscript accordingly!\n\n\n[1] DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. https://arxiv.org/abs/2501.12948\n\n---\n\n> revealing effectiveness of the proposed Sequential Voting and IBPO optimization paired with Adaptive Sequential Voting, although improvements seem a little marginal.\n\n> One potential weakness is that many of the baseline results are not run but transcribed from others' works. This harms the solidity of the experiment comparisons.\n\nWe believe these are relatively minor points, if our understanding is correct. As such, we've kept our response concise to avoid unnecessary distraction. (Further details can be found in our response to Point 3 for Reviewer hh45.)\n\nAfter all, our work on constrained LLM inference and the self-correction literature are essentially two orthogonal directions. Transcribing their results is both convenient and helps avoid potential implementation discrepancies.\n\n\n---\n\nOnce again, we sincerely appreciate the reviewer’s insightful feedback and encouraging assessment! We hope our clarifications sufficiently address your questions and clarify our choice of construction. We deeply appreciate your support and will carefully revise the manuscript to further improve its clarity!", "claims_and_evidence": "A few claims are made during the deduction of main algorithm, but they're all quite well-grounded.\n\nWhen finding  a workaround for solving the parametric objective function with constraints, authors claim that the non-parametric workaround is superior to alternating between gradients of the policy and Lagrangian multipliers. Although no ablation studies are provided, solid literature on other methods utilizing this workaround in similar problems supports the claim well.\n\nIt's also claimed early on in the paper that the proposed IBPO ends up becoming a generalization of SFT methods which is confusing at first but natural after the table comparing SFT, RFT, RAFT and IBPO is provided.\n\nThe choice of some techniques, however, weren't explicitly analyzed/compared, such as the use of semi-gradient and the choice of CGPO, but they are either intuitive choices or self-explanatory in positive experiment results.", "methods_and_evaluation_criteria": "The main method IBPO is framed initially as a constraint RL problem, but after practical adaptations with respect to the policy update (solving for approximate optimal policy via sampling, using semi-gradient) and the reward calculation (introducing a marginal reward, implementing CGPO), the method is deduced to a SFT where samples take the form of weights assigned to long or short context responses. This is in theory sound and intuitive, except that dichotomy between long and short responses is a little counter-intuitive.\n\nThe evaluation is largely sound: authors developed Sequential Voting prompting to generate long-context responses, SCoT for solely short-context responses, and Adaptive Sequential Voting prompting for mixed responses. These are used separately as SFT sets and results are compared. However, I'm still hesitant to accept that responses constructed this way are actually long/short-context responses -- the long ones look like a mere aggregation of short responses.", "theoretical_claims": "The theoretical proportion is mainly definitions, intuitive deductions and plugging-ins. They seem correct to me.", "experimental_designs_or_analyses": "Extensive experiments were conducted, comparing prompt constructing methods and optimizing methods. Experiments consist of the SFT/Prompting-based comparisons and Online Iterative SFT/RL comparisons, respectively revealing effectiveness of the proposed Sequential Voting and IBPO optimization paired with Adaptive Sequential Voting, although improvements seem a little marginal.\n\nAs for the setting, extensive baselines were introduced/re-implemented to show the proposed methods' superiority, among which adapting self-correction methods as baseline is a creative choice, as it generates similar long context responses to SV.", "relation_to_broader_scientific_literature": "This paper not only sheds new lights on the problem of budget-aware/adaptive reasoning but also provides a valuable algorithmic contribution to constrained RL as a whole. During parameter update, a method that is essentially a generalization of RFT and RAFT is proposed; the algorithm is a modified version of CGPO, although the alteration is simple.", "essential_references_not_discussed": "None that I'm aware of.", "other_strengths_and_weaknesses": "Aside from previous comments, the paper is thorough and detailed in introducing the algorithm, providing sufficient citation and overview of related literature for most readers to understand. \n\nOne potential weakness is that many of the baseline results are not run but transcribed from others' works. This harms the solidity of the experiment comparisons.", "other_comments_or_suggestions": "On the paper structure, maybe it's a better idea to include more explanation on prompt generation (the proposed SV, ASV), not only to clarify the experiment setting more, but to also convince readers that SV is adequate in representing long response situations that the paper initially tries to improve. The algorithm and experiment parts can be cut down a bit.", "questions_for_authors": "One major confusion is whether prompts generated via SV sufficiently represent the \"unnecessarily tedious long reasoning trace\" that the paper sets out to improve. This is rather important as SV seems to be the only experiment scenario.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tlz2PAdhT5", "reviewer": "Reviewer_JCKH", "summary": "This paper discusses a method for scaling LLM test time compute on adaptive basis based on prompt difficulty. The proposed approach is a novel reinforcement learning technique that allocates more inference to difficult problems (adaptive number of votes - where each vote requires an inference) and fewer votes to easy questions. The proposed method, called Inference Budget-Constrained Policy Optimization (IBPO) is experimentally verified to improve performance on MATH500 dataset - though at a higher test time compute budget (2x).\n\n## update after rebuttal\n\nI thank the authors for their response.\n\nWhile the response partially addresses my concerns, I am still not confident that the approach generalizes beyond the specific datasets and architectures.\n\nAs the authors mentioned, the method is applicable for reasoning settings that are neither easy (85% or more accuracy per model) nor hard (AIME where model accuracy is low). Normally, this would not be a problem, one could adjust the capacity of the model (use smaller / larger LLMs) to show that the approach can be performant at these settings, but the authors mentioned that their approach cannot scale well due to training complexity.\n\nOn the other hand, while the authors presented an argument as to why the technique should not be compared with evolved SC variants, I am not convinced. If their approach is orthogonal to the reasoning method, then I would like to see some experiments where more modern reasoning backbones are used when combined with their method.\n\nI am certain that more evolved reasoning baselines such as the ones I recommended need to be compared to, to quantify the method's impact on total token usage versus a simple no fine tuning method.\n\nFor this reason I am maintaining my current score.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful feedback! \n\n---\nWe extend some discussion on **accuracy** to clarify our evaluation design, which also helps explain our choice of baseline later on.\nWhile important, accuracy depends on **two orthogonal axes**: (i) the RL axis; and (ii) the reasoning axis.\n$$\n\\underbrace{\\text{constraint satisfaction} \\to \\text{adaptive allocation}}_{\\text{RL}} \\to \\text{accuracy} \\leftarrow \\text{SV (reasoning)}\n$$\n- **RL axis**: We're glad the reviewer agrees the adaptive allocation is ''**theoretically expected**'' and ''**verified**'';\n- **Reasoning axis**: (A)SV was intended as a **simple** alternative to longCoT with **reasonable** accuracy, as R1 was unavailable until 10 days before the submission deadline. Empirically, (A)SV performs **comparably** to SC (Fig 2), enabling:\n     -  a sufficient margin for optimization;\n     -  a fair comparison w/ SC, a baseline more familiar to the community.\n\n(SC is used here to match the reviewer’s terminology, though we use MV in paper.)\n\nOne key observation is: (i) SV is an all-voting case (can be seen as $q_+ = 1$); (ii) ours are adaptive with voting capped by $q_+$. This suggests, roughly speaking, that ours should be theoretically **upper bounded** by SV.\n\nSince SV performs comparably to SC, we consider SC a reasonable baseline—based on the reasons above, not chosen ad hoc.\n\nWhile further engineering on the reasoning axis (e.g. better alternative than SV) could improve overall accuracy, we **refrained** from over-optimizing it, as it's **not** the core insight of this work and could distract from the RL contributions.\n\n---\n- **Baseline**: Continuing from the discussion above \n\n  - **Interpolation baseline**: We compare ours against a **hypothetical efficiency boundary (HEB)** (gray line Fig 2), defined by two **non-adaptive** extremes:\n    - **$q_+ = 0$**: All short responses;\n    - **$q_+ = 1$**: All long responses.\n\n    A model **above** the HEB can be interpreted as having **improved efficiency**.\n\n  - **Why SC**: SC was chosen because it closely aligns with the HEB (Fig 2). We mention SC in the abstract, as readers are likely more familiar with it.\n\n  - **[1, 2]**: While [1, 2] are more recent, our goal is to evaluate against the HEB, which reflects non-adaptive cases. As we refrained from over-optimizing the reasoning axis, the upper bound (SV ≈ SC) is relatively low—but sufficient for testing our RL part—and is likely outperformed by [1, 2]. Thus, we don’t find a direct comparison with them entirely fair/necessary.\n\n  We hope it clarifies our baseline choice and will revise accordingly, including additional discussion of [1, 2].\n\n- **Dataset**: MATH500: (i) it has difficulty metadata; (ii) LLaMA 3.1 8B Inst. has moderate accuracy. All numbers hereafter refer to this model.\n\n    Apart from metadata: Easier ones (GSM8K, SVAMP, ASDIV w/ pass@1 of ~85% [3,4,5]) have small reward margins, as most queries are solved in one attempt; \n    For harder ones like AIME (2/30 solved [6]), efficiency makes little sense as the model rarely succeeds.\n\n    We hope our grounded derivation and constraint satisfaction curves offer more insight and confidence than accuracy numbers alone.\n\n- **Model**: We don't particularly find 8B ''small'', especially given the cost of online training. For reference, [7]—an online self-correction—uses 7B models.\n  While we didn't explore other families, we're optimistic about the method's generalizability, thanks to its (i) model-agnostic derivation; (2) simple update (Eq. 5). \n\n---\n- **2x efficiency**: Apologies for the confusion. It comes from Tab 3 and Fig 2. At $q_+ = 50%$, ours has 55.4% (a 4.14% gain) w/ 2.16x trials (Tab 3), while **the HEB** (aligned with SC, if not better) in Fig 2 requires ~4.5x trials for similar accuracy—implying ~2x efficiency.\n\n- **Table 3**: We didn't re-run; the numbers were taken from the cited papers to avoid implementation discrepancies.\n\n---\n- Online training is **inherently** expensive: For 70B, under our setup, we estimate that a single epoch of generation alone could take ~2,600 H100 hours at 300 output tokens/sec per node. In practice, training throughput is significantly lower due to memory overhead from model optimization—and this estimate excludes other potential costs. These limited us from exploring larger models and harder datasets (which likely require larger training sets).\n\n---\nWe quote reviewer toYu:\n> not only sheds new lights on adaptive reasoning but also provides a valuable algorithmic contribution to constrained RL as a whole.\n\nWe are not aware of RL-for-LLM works that explicitly impose linear constraints to control response distributions.\n\nWe believe our RL contribution merits stronger recognition—especially given your positive evaluation of it.\n\n---\nAgain, we appreciate your valuable suggestions! We will revise accordingly and we hope our responses addressed your questions.\n\n---\n[3] arxiv.org/abs/2407.21783\n\n[4] 2410.06638\n\n[5] 2502.12134\n\n[6] 2410.01560\n\n[7] 2407.18219", "claims_and_evidence": "- IBPO can correctly discriminate easy and hard problems (verified in experiments section e.g. Fig. 2)\n- IBPO can allocate resources efficiently by dedicating more inference compute towards hard questions (this is theoretically expected from the way the RL objective is set up and empirically verified to do so.\n- The paper claims that its proposed method has an improvement that is in relative terms twice that of vanilla Self Consistency (SC) (I am not actually sure where this is verified. The claim is plausible, but I cannot pinpoint where the 2x efficiency is shown - perhaps I am missing it as the paper is quite dense).", "methods_and_evaluation_criteria": "1. The RL method seems correct and novel. I would like to note that while I have understanding of RL, I am not an RL expert thus I would defer to other reviewers to validate my positive evaluation of the proposed method.\n2. Evaluation metrics are reasonable. The experimental pipeline is a fair way test for the method.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "1. There does not seem to be a sufficient amount of experiments. Evaluating on MATH 500 is decent choice, but why not on other datasets? Other datasets with mathematical reasoning problems include AIME 2024 (or other years of AIME), SVAMP, ASDIV, AQUA, some big bench hard datasets, etc.\n2. I am curious as to why the selected LLMs seems to be small. The experiments go up to an 8B model. Do the performance gains generalize across different parameter sizes? What about other LLM families besides LLAMA? \n3. Baseline selection: In the abstract the authors make a point of improving over vanilla self consistency (SC). I am curious if the proposed method competes well against modern SC variants that have been demonstrated to significantly reduce inference costs. Given that vanilla SC is reasonably dated I think it would be fair to compare against [1] and [2].\n\n[1] Let’s Sample Step by Step: Adaptive-Consistency for Efficient Reasoning and Coding with LLMs\n \n[2] Escape Sky-high Cost: Early-stopping Self-Consistency for Multi-step Reasoning", "relation_to_broader_scientific_literature": "The work relates to the test time compute literature. This is a major area of research for LLMs. The paper's problem is well motivate.", "essential_references_not_discussed": "I understand that [1] was mentioned in the introduction section of the paper but I believe that it merits experimental comparison.\n\n[1] Let’s Sample Step by Step: Adaptive-Consistency for Efficient Reasoning and Coding with LLMs\n\n[2] Escape Sky-high Cost: Early-stopping Self-Consistency for Multi-step Reasoning", "other_strengths_and_weaknesses": "This is polished paper but I personally find it dense to read with a lot of heavy notations. Perhaps it could be made sparser with more references to appendix to make it easier? For example, it was not easy to understand the dataset preparation section.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "- In table 3 it is stated that a lot of the numbers are duplicated from Qu et al. (2024); Yan et al. (2024); Kumar et al. (2024). Did the authors run the code for some of these methods in their setting to make sure that the results reproduce? Is the setting reported from these works identical? For example, are the exact same prompts used in all works in this table?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hOpd9u3SzW", "reviewer": "Reviewer_hh45", "summary": "This paper proposes IBPO to optimize reasoning length allocation in large language models (LLMs). While extended reasoning chains improve accuracy, they often lead to inefficiencies by applying unnecessary long reasoning to trivial queries. IBPO formulates this as a constrained reinforcement learning (RL) problem, categorizing responses into different length groups and imposing density constraints to allocate inference budgets adaptively.  \n\nEmpirical results show that IBPO improves efficiency, achieving accuracy gains on MATH500 over LLaMA3.1 8B Instruct, with efficiency gains compares to self-consistency. The paper details IBPO’s derivation, implementation and experiments.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your time in assessing our paper and your thoughtful feedback.\n\n--- \n\nWe’d like to first clarify a possible **misunderstanding** regarding the **role of sequential voting** (SV): it is **not** intended as a significant contribution of this work.\n\n**Role of SV:** It is a **simple, constructive alternative** to o1/R1-style long CoT, with **extended length** and **reasonable performance gains** over conventional CoT. \n\n- **Why an alternative**: As o1 is proprietary and R1 was released on Jan 20, ~10 days before the submission deadline.\n\n- **Why SV**: \n  - **Simple**: A minimal construction that allows us to focus on our RL contribution without the distraction of engineering long responses; \n  - **Characteristics**: (i) gains over standard CoT; (ii) extended length. \n    - SV resembles longCoT in these characteristics (though its gains over CoT are moderate), enabling us to demonstrate the effectiveness of our RL algorithm.\n\n- **What do we expect from SV?**:\n\n  - **Superior accuracy?**: No. As shown in Fig 2, SV performs comparably to parallel majority voting (MV). We expect advanced reasoning methods to outperform MV, and hence SV.\n\n  - **Practical value?**: No. The main insights we convey do **not** lie in SV. After all, it is a dummy construction used to highlight our RL algorithm, which carries the core insights of this work.\n\nWe understand the concern that SV may not fully capture how o1/R1 works fundementally. Due to the **unavailability** of R1, we construct SV that at least resembles longCoT in terms of the aforementioned characteristics.\n\n\nGiven the considerations above, we believe our construction is a fair and reasonable choice for **illustrative** purposes.\n\n\n\n\n---\n\nTo reiterate, our core contributions are:\n- (i) casting adaptive reasoning as constrained inference; \n- (ii) designing a constrained RL algorithm — which merits particular emphasis — that is simple and grounded.\n\nIn addition, the final **accuracy** stems from **two orthogonal axes**: (i) **the reasoning axis**; and (ii) **the RL axis**.\n\n$$\n\\underbrace{\\text{constraint satisfaction} \\to \\text{adaptive allocation}}_{\\text{the RL axis}} \\to \\text{final accuracy} \\leftarrow \\text{SV (the reasoning axis)}\n$$\n\nThe derivation of our algorithm and the ''upstream'' metric (constraint satisfaction) are theoretically **agnostic** to the choice of long-response.\n\nWe intentionally kept the reasoning axis simple and did **not** attempt to boost performance through it, so as to isolate and highlight the RL component.\n\nThese points collectively suggest that a reasonable construction is sufficient for emphasizing our RL contribution.\n\n---\nWe hope these discussions clarify our contributions. We believe our RL contribution—unanimously considered as **sound and novel** by other reviewers—merits recognition.\n\nWe quote reviewer toYu:\n\n> not only sheds new lights on adaptive reasoning but also provides a valuable algorithmic contribution to constrained RL as a whole.\n\n----\nFor specific comments:\n\n1. **R1/O1 type models**: We hope the above clarifies our choice of SV, given the unavailability of such models.\n\n2. **Trials as metric**: \n   1. - This choice was mainly motivated by Table 3. Since the self-correction works don't report token counts, we use the number of trials/turns as a proxy.  \n\n      - Besides, SV-SFT performs similarly to MV in trials (Figures 2a/2b); and scales comparably in tokens (Figure 5b, Appendix). Hence, trials serve as a reasonable proxy.\n  \n   2. **Practical value considering inference time**: This is an inherent limitation of any long-form reasoning method like longCoT, not a consequence of our RL algorithm. Substituting SV with O1/R1-style responses would result in similar inference-time overhead.   \n    \n      After all, SV is a dummy alternative to longCoT, and was **not** intended to offer practical value beyond its illustrative purpose.\n\n3. **Table 3**: As noted in Sec. 5.1, Table 3 is not intended to suggest our method outperforms self-correction. These are two **orthogonal** research directions: constrained inference and self-correction.\n\n    Its purpose is simply to illustrate that constrained inference can achieve **comparable** performance to **a well-established line of work**. For this purpose, we believe transcribing their results is reasonable—it is convenient and avoids potential discrepancies from re-implementation.\n\n    A more informative comparison is Fig. 2, where we evaluate against an **efficiency boundary** interpolating between two extremes of non-adaptive (homogeneous) cases. See our response to reviewer JCKH for details.\n\n---\nAgain, we sincerely appreciate your suggestions and your assessment that our claims are supported, our theoretical and experimental designs have no issues, and our framework is rigorous.\n\nWe will carefully revise the manuscript to further improve clarity based on your suggestions, and we hope our responses have adequately addressed your questions.", "claims_and_evidence": "Yes. The claims made in the submission are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes, I think they have no issues.", "experimental_designs_or_analyses": "Yes, I checked. I think they have no issues.", "relation_to_broader_scientific_literature": "This work contributes to the broader scientific discourse on reasoning in large language models (LLMs), particularly in mathematical problem-solving and efficient inference. It builds on foundational research in chain-of-thought (CoT) prompting, which demonstrates that decomposing reasoning into explicit steps improves problem-solving accuracy. While CoT and its extensions—such as self-correction and multi-turn reasoning.", "essential_references_not_discussed": "To the best of my knowledge, there are no essential related works that are missing from the citations or discussion in the paper.", "other_strengths_and_weaknesses": "Strengths: \n\n(1) The method is rooted in constrained reinforcement learning (RL) and resource allocation theory, providing a rigorous framework for optimizing inference budgets.      \n(2) IBPO aims to balancing inference cost and reasoning accuracy by dynamically adjusting reasoning length based on problem difficulty, reducing unnecessary computational overhead.      \n\nWeaknesses:\n\n(1) The paper frequently refers to extended CoT (or long CoT) but selects SV as the representative for extended CoT. This deviates significantly from the standard academic understanding, where extended CoT is typically represented by models like O1 and R1, which focus on deeper rather than broader reasoning (e.g., majority voting). Since the proposed method has not been evaluated on R1 (O1)-like extended CoT, its effectiveness in such settings remains uncertain.        \n(2) According to my understanding, SV generate different trails sequentially, considering the quadratic complexity in transformer-based models, SV can result in longer inference times compared to Majority Voting (MV), where trails are generated in parallel. What’s more, the paper uses trails per response as the inference budget metric, which has two key issues:      \n1. Trail count does not directly correspond to token count, whereas token budget would be a more practical and meaningful measure.    \n2. Despite a lower token budget, the sequential nature of SV may lead to longer inference times compared to parallel methods like MV. Additional runtime experiments comparing MV, SV, and ASV should be conducted to verify whether the proposed method increases inference time. If the approach actually extends inference time, its practical value is questionable.    \n(3) While the core idea of IBPO is concise, the formulation in the paper contains substantial redundancy, which hinders readers from extracting the key information efficiently.      \n(4) In Table 3, the IBPO experiments are conducted on LLaMA3.1-8B, whereas S3C and Self-Refine utilize multiple models different from LLaMA3.1-8B. This inconsistency weakens the rigor of the comparison, making the evaluation less reliable.", "other_comments_or_suggestions": "(1) The authors should simplify the formulation appropriately to improve the readability of the paper.      \n(2) The authors should supplement the paper with experiments comparing the inference time overhead of MV, SV, and ASV to provide a more comprehensive evaluation.", "questions_for_authors": "No other question for authors, please see the weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eFzUGiDQO1", "reviewer": "Reviewer_8P46", "summary": "This paper targets the problem of solving mathematical problems with LLMs. While the currently prevalent method LongCoT brings promising improvements for mathematical reasoning, they are sometimes unnecessary and cause token waste. To alleviate this issue, the authors propose an algorithm Inference Budget-Constrained Policy Optimization (IBPO), where the problem is formulated as a resource allocation scenario. The budget is assigned with respect to the difficulty level of the problem. IBPO is implemented based on the RL objective of constraint generative policy optimization by replacing it with margin maximization under budget. Experiments on MATH dataset verify the effectiveness of the proposed method. Analyses also prove its design motivation, that difficult problems will get more budget to solve.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your thorough review and valuable feedback! We greatly appreciate your insightful suggestions and positive assessment!\n\n---\n\n> The authors demonstrate (specifically in Figure 2 Column 3) that the voting budget can adaptively change w.r.t. the difficulty level of the problem. Apart from the token budget, I am wondering if the authors could also offer accuracy on different difficulty levels. Now there is only overall accuracy in Table 3.\n\n| Level               | 1             | 2              | 3              | 4              | 5              | Overall       |\n|--------------------|---------------|----------------|----------------|----------------|----------------|----------------|\n| # of problems | 43            | 90             | 105            | 128            | 134            | 500            |\n| q₊ = 0.25          | 93.02% (40)   | 76.67% (69)    | 64.76% (68)    | 48.44% (62)    | 23.88% (32)    | 54.2% (271)    |\n| q₊ = 0.5           | 83.72% (36)   | 80.00% (72)       | 66.67% (70)    | 52.34% (67)    | 23.88% (32)    | 55.4% (277)    |\n| q₊ = 0.75          | 90.70% (39)    | 77.78% (70)    | 74.29% (78)    | 51.56% (66)    | 23.88% (32)    | 57.0% (285)      |\n\nAbsolutely. We break down the overall accuracies reported in Table 3 by difficulty level in the table above (numbers in parentheses indicate the number of problems solved). Some observations are:\n\n- With larger budgets, we observe improved performance on the harder levels (3 and 4), as more voting budget becomes available. \n\n- It is surprising that level 5 performance remains the same across all budgets. We conjecture that this is due to the limitations of the 8B model, which may only solve a subset of the hardest problems.\n\n---\n\n> The design choices of the dataset construction in Section 4 (Appendix B) need more justification. For example, why at most \"8 trials\" and early stop if an answer appears \"3 times\"? Would the design choices have a significant impact on the final performance (sensitivity)?\n\nIn general, these choices do impact the performance of sequential voting (SV) itself. These choices were made to ensure that SV performs **comparably** to majority voting (MV), a baseline that is more familiar to the community.\n\n- **Reason of chocies**: \n  - In our observations, SV without early stopping tends to underperform compared to parallel MV. With the current setup, SV achieves performance similar to MV (Figures 2a/2b), which allows for a meaningful comparison between our ASV-IuB-$q$ models and MV.\n\n  - We set the trial cap at 8 to avoid unbounded responses. This limit, combined with early stopping, empirically matched MV performance, so we kept it.\n\nSince our key insights lie in the RL component, we found that aligning SV with parallel MV is a reasonable choice, as it allows for: (i) a sufficient reward margin for optimization, and (ii) a fair comparison to a community-standard baseline (MV).\n\n  \n---\n\nOnce again, we gratefully thank the reviewer for the insightful suggestions and positive evaluation! We hope our clarifications have adequately addressed your questions. We sincerely appreciate your support and will carefully revise the manuscript to further improve its clarity!", "claims_and_evidence": "The major claims in the paper, specifically regarding the motivation and experiment results, are good to me.", "methods_and_evaluation_criteria": "The paper is evaluated mainly on MATH dataset, which contains annotations on the difficulty levels of different problems. This setting is natural and could well validate the proposed method, as it claims to allocate different budgets on different questions.", "theoretical_claims": "I checked the proposed algorithm and the corresponding equations, all seem sound to me.", "experimental_designs_or_analyses": "The overall experiment design makes sense. I have some minor issues or comments:\n- The authors demonstrate (specifically in Figure 2 Column 3) that the voting budget can adaptively change w.r.t. the difficulty level of the problem. Apart from the token budget, I am wondering if the authors could also offer accuracy on different difficulty levels. Now there is only overall accuracy in Table 3.\n- The design choices of the dataset construction in Section 4 (Appendix B) need more justification. For example, why at most \"8 trials\" and early stop if an answer appears \"3 times\"? Would the design choices have a significant impact on the final performance (sensitivity)?", "relation_to_broader_scientific_literature": "This paper studies a very important and interesting research question of the current community, allowing LLMs to allocate different token budgets for different problems of difficulty levels. This is timely research, especially when o1 and Deepseek styles begin to dominate in reasoning tasks. \n\nIn terms of ideas and findings, many concurrent works try to decrease the number of tokens as much as possible while maintaining performance, such as TokenSkip, this paper takes the approach of resource allocation and designs a novel algorithm to do this. While a lot of previous works noticed the correlation between token budget and problem difficulty, they are mostly doing direct preference alignment. The algorithm proposed in this paper is valuable.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "See the above review.", "other_comments_or_suggestions": "See the above review.", "questions_for_authors": "See the above review.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces IBPO, a reinforcement learning algorithm that helps LLMs assign longer reasoning chains to harder problems and shorter ones to easier ones. The approach frames this as a constrained optimization problem, where the model learns to balance accuracy and computation cost. The method is theoretically sound and the use of constrained RL is interesting. However, the experiments could be more complete by comparing to modern variants of self-consistency methods. Another issue is to use SV to represent models such as o1 and R1 where SV focuses on breath and in fact, o1 and R1 focuses on the depth of the reasoning. The experiments are also limited to the math500 dataset.
Decision: Accept (poster)