Learning How Hard to Think: Input-Adaptive Allocation of LM Computation
Authors: Mehul Damani
Idan Shenfeld
Andi Peng
Andreea Bobu
Jacob Andreas
Venue: ICLR 2025 Poster (2025)
Keywords: LLM
inference
scaling
test-time compute
Primary area: foundation or frontier models, including LLMs
TL;DR: We introduce a method for input-adaptive allocation of test-time computation

Abstract: Computationally intensive decoding procedures---including search, reranking, and self-critique---can improve the quality of language model (LM) outputs in problems spanning code generation, numerical reasoning, and dialog.
Existing work typically applies the same decoding procedure for every input to an LM. But not all inputs require the same amount of computation to process. Can we allocate decoding computation adaptively, using more resources to answer questions whose answers will be harder to compute? We present an approach that predicts the distribution of rewards given an input and computation budget, then allocates additional computation to inputs for which it is predicted to be most useful. We apply this approach in two decoding procedures: first, an adaptive best-of-$k$ procedure that dynamically selects the number of samples to generate as input to a reranker; second, a routing procedure that dynamically responds to a query using a decoding procedure that is expensive but accurate, or one that is cheaper but less capable. Across a suite of programming, mathematics, and dialog tasks, we show that accurate computation-allocation procedures can be learned, and reduce computation by up to 50% at no cost to quality.

Published as a conference paper at ICLR 2025
LEARNING How HARD TO THINK: INPUT- SEPUT-ADAPTIVE ALLOCATION OF LM COMPUTATION
Mehul Damani" Idan Shenfeld
Andi Peng
Andreea Bobu
Jacob Andreas
Massachusetts Institute Technology
ABSTRACT
Computationally intensive decoding procedures-inclue and self-critique-can improve the quality of language model (LM) outputs problems spanning code generation, numerical reasoning, and dialog. Existing work typically applies the same decoding procedure for every input to LM. But not all inputs require the same amount of computation to process. Can we allo- cate decoding computation adaptively, using more resources answer questions whose answers will be harder to compute? We present an approach that predicts the distribution of rewards given an input and computation budget, then allocates additional computation to inputs for which it is predicted to be most useful. We apply this approach two decoding procedures: first, an adaptive best-of-k pro- cedure that dynamically selects the number of samples to generate input to reranker; second, a roufing procedure that dynamically responds to a query using decoding procedure that is expensive but accurate, less capable. Across suite of programming, mathematics, and dialog tasks, we show that accurate computation-alloaaat
search, reranking,
that is cheaper but
procedures can be learned, and reduce computation by to 50% at no cost to response quality, or improve quality by up 10% a fixed computational budget.
INTRODUCTION
Recent improvements in language models (LMs) have dramatically enhanced their ability to tackle complex tasks in mathematics, cod- ing, and reasoning. However, as with natural and artificial intelligent agents of all kinds (Sil- ver et al. 2016). LMs cannot solve all problems on the first try: they benefit from the ability to perform search (Yao et al. 2024), sampling (Brown al. (024)) or more sophisticated de- coding procedures like chain-of -hhtuuhht (Wei al. 2022) and self-critique (Wang et al. 2023
Hejp me balonce a Bixory neet
Whut & your nome?
Prove Riemsnn lypottissts
Extimatee query dificuutyy
Easy
Impossible
Allocate extra computation hard
LM
7TV
manyy
small
lorge
greedy
dacocing
Figure 1: Overview of our approach. Given a set of input queries to a language model, we train a lightweight model to estimate the difficulty of these queries (more precisely, model that esti- mates how much each query would benefit from more computationally intensive decoding proce- dure.) We then allocate extra computation those queries for which it would be most beneficial.
Importantly, computationally intensive prob- lem domains may exhibit considerable varia- tion in the difficulty of individual problem in- stances: not all problems solve. For example, even a novice program- mer likely write code to to test if an integer is even. Balancing a binary tree might require multiple attempts, and finding polynomial- time set cover algorithm is likely impossible, and not be worth attempting at all. Maximally efficient use of computational resources thus requires identifying, priori, the inputs for which additional computation will improve outputs. In LMs, recent work has shown that significant gains from adap- tive choice of decoding procedures are theoretically achievable (Snell et al. 2024). However, past
equally hard to
"Correspondence mehul445 1iluum t. edu.

---

Published as a conference paper at ICLR 2025
work has relied on measures of problem difficulty that could only be estimated posteriori, using sampling procedure significantly more expensive than the one used to generate model outputs.
In this paper, show that it is possible to perform inference-time scaling of computation before responding to set of user queries. by predicting how much computation will be required to respond successfully each query, and allocating resources accordingly (Fig. D. Our approach has two parts. First, we describe how to train a difficulty model that estimates the marginal improvement in response quality that would result from allocating an additional unit of computing power when responding to query. We show that effective difficulty models can be constructed by training lightweight "probes" on top of a pre-trained LM's hidden representations, suggesting that LMs may already learn to encode problem difficulty as a result of pre- rre-training Second, we describe an efficient allocation algorithm that, given a collection of queries and computation budget, uses predictions from the difficulty model to allocate computation to individual queries. We show how to perform this optimization online (exactly satisfying the budget for batch of queries) or offline (computing a fixed mapping from queries to compute allocations so that constraints are satisfied in expectation).
This approach is flexible with regard to the choice of compute-scaling procedure, and we demon- strate its effectiveness with two different inference procedures. In best-of-k experiments, we choose how many samples to generate from an LM before re-ranking with reward model. In routing ex- periments, choose whether to respond to queries using a computationally expensive but capable decoding procedure, or a less powerful but more cost-effective decoder. For the best-of-k setting, present results across 3 domains: Math, Code and Chat. Adaptive sampling outperforms non- adaptive strategies on all three domains across a range of possible compute budgets. On Math and Code, we achieve the same performance as non-adaptive methods using up to 50% less compute; on Chat, we are able to match reward using up to 10% less compute than base models. For routing, we experiment with both a pair of models (Gemma-2B. Gemma-7F 7B) and a pair of decoding procedures (ordinary sampling and value- augmented sampling; Han et al. (024)) In both cases, we match the performance of the more expensive decoder while calling it only 50-75% of the time.
In summary, this work presents (1) a generally applicable framework for adaptively scaling test-time LM computation; (2) a learned model for estimating the marginal benefit of additional computation in LM decoding; (3) an efficient algorithm for allocating computation queries; and (4) experiments demonstrating significant improvements in LM efficiency and output quality.
2 PRELIMINARIES
Suppose we have an LM-based agent that will interact with a large number of users in parallel. At any moment, each of a set of users has issued a query Ij, for which we wish to produce response y/- We have acccess to both a language model (LM) p(y x)) capable of generating candidate responses, as well as a reward model r(z, v) capable of assessing the quality of candidate responses. In the absence of any computational constraints, we might wish to find the best response to every user query by optimizing arg max, r(z,, y) independently for each query 24.
In practice, however, we generally do not have the ability to perform exhaustive search over can- didate responses y. Instead, response subject to some constrained computation budget b. (In general, increasing the budget query should increase the quality of the response.) Many such decoding procedures are in wide use; our experiments will focus on two of the most widely used.
a decoding procedure f(x, b) that (stochastically) generates
In best-of-k, we generate finite number of samples, then rerank them:
f(x
arg max D: (v.......)..pt
r(x, 1/5)
(I)
In routing, we have access to a strong but expensive decoding scheme pd (which could involve search, reasoning, or even just a larger model). For every query we can either use this decoding scheme to fall back to weak but cheaper decoding scheme D" (e.g. ordinary decoding smaller model). Then:
y D (: x) y ps (.1x)
i DW
f(z, b)
(2)
if = 6s
2

---

Published as a conference paper at ICLR 2025
Here 6" and by denote the cost of calling the weak and strong decoders respectively. Decoding procedures including consensus (Jacob et al. 2024), thain-of-thought (Wei et al. 2022), self-critique (Luo et al. 2024) and debate (Du et al. 2024) may all be expressed in this form, with the budget b controlling the number of generated samples, tokens, or rounds of revision.
In current practice, it is typical to set a fixed budget B, and allocate this uniformly for all queries returning yi ~ f B) for each Ty. In this case, the expected reward for a set of queries simply:
Σ Eyy.s(ou. [r( 34)]
But noted in Section [ different 2, may in practice benefit differentially from increased budgets. In this paper, we consider a more flexible version of the problem- if we are willing to set a fixed av- erage computational cost per query, can we improve overall accuracy by allocating this computation adaptively? In this case, responding to collection of queries {z1,.... In corresponds to solving:
Σ Burs(a brlr (1))));:)) u:)] s.t. brn..bb
max
Σ B.n
(4)
Below, we develop method for solving this adaptive comptuation scaling problem efficiently.
3 METHOD
In Eq. we would like to allocate computation b, to queries I; to maximize expected reward, which we will denote q(zi, b;) Byys/(s (,),))) 3:)]- But there is a problem: without querying the LM, we cannot know the value of q(x,, b,) for some new 24. To make matters worse, reliable Monte Carlo estimation of q(z,, b;) (as in Snell et al. 2024) may require more computation (c.g. more LM samples) than we eventually wish to allocate xy! To efficiently allocate computation, we need some way to predict q(>;, b,) from x; alone, and then use these predictions to solve Eq. (4)
When estimating problem difficulty and allocating decoding budget, it will be useful to reason about the benefits of incremental changes in compute budgets. Formally, let us define the marginal re- ward A.₃ g( - j - 1) (with A(, 0) 0). Under this definition, g(z;, d,) = Σラ= Ay; intuitively, A; represents the expected gain from allocating one more "unit" of computation to 2i, given that we have already allocated j units. We may then re-write Eq.
as:
BRax
max
CCGAB s.t. Σog EossB.m B.B.;
Ci S C.j-1 Vi,j
(5)
Here we have introduced auxiliary variables Cij E {0. } to represent budget increments: intuitively, setting Gi means we have allocated a jth unit of computation to Ti and constraints enforce that
these units are allocated in sequence, shown in Figure 2 To determine how to allocate decoding computation across a set of queries z;, then suffices to predict the value of each term Diy and then solve Eq. 55) to obtain c. Each of these steps is described in more detail below.
3.1 ESTIMATING PROBLEM DIFFICULTY BY PREDICTING MARGINAL REWARDS
Given a training set of queries T4, we collect empirical estimates of A,y by decoding from [ax b) at values of b up to some B Bmax- We then train a model Á (I,; 0) that predicts a vector of marginal rewards for all budgets j simultaneously by optimizing the mean squared error:
arg min
Σ - [z -(z; 0)11
(6)
A.A;;
where A, |A...... A. the vector of empirical marginal rewards for the query 2;, While training this model requires sampling from the LM at many budgets to obtain supervision, it can be called during inference without generating any outputs at all.
We explore two simple parameterizations of Á: (I) a two-layer MLP that takes last the hidden state of the base LM p (obtained by encoding the query) input, and (2) parameter-efficient fine-tuning of the base LM with LoRA (Hu et al., 2022p. During inference, the MLP variant adds extremely little overhead as its input hidden states that are already computed as part of decoding procedure. The LoRA variant is slightly more expensive; however, its overhead is also negligible, as the primary bottleneck during inference occurs when decoding outputs rather than encoding queries.
3

---

Published as a conference paper at ICLR 2025
Incremental computation Gi3
Marginal reward predictions
S1
Allocated budget
Help balance a binary tree.
0.3
0.2
Agy
0.1
0.1
0.3
0.2
0.1
0.1
bj -2
4= What your name?
1.0
0.0
Prove the Riemann hypothesis.
Language
0.0
0.0
by m0
Model
X, Write an essay about Jane Fyre.
bolance binury
0.1
0.1
b4
DIFFICULTY ESTIMATION
COMPUTATION ALLOCATION
Figure 2: Method overview. the first step, the marginal reward predictor is used estimate the marginal reward gains for a batch of queries. In the second step, the allocation algorithm uses these predictions to assign compute budgets to individual queries. The simple query what is your name? is allocated budget of 1, while the harder query balance binary tree is allocated budget of 2.
3.2 ALLOCATING COMPUTATION
Eq. (5) is integer linear program. But under the assumption of decreasing marginal benefits from additional computation, it has a special form that allows it to be solved in O n) time using greedy algorithm that incrementally "turns on" eıj for which Aij largest. We exploit this property in two ways:
Online allocation: If queries { are known priori, simply replace each D.j in Eq. 5 with the corresponding estimate from (xy; 0), solve for Cij, then finally set each 6; = max, Cij-
One drawback of this procedure is that responses must be processed in batch. However, in some cases it is also possible to effectively set budgets independently for each z;:
Offline allocation: If allocations must be made without access full batch of samples, struct a fixed allocation policy as follows:
(1) Hold out a subset of the data used for training the reward estimator Å, then use it to label queries in this held-out set (e.g. based on the first-sample prediction (xx:)). Divide these queries into a fixed set of bins according their predicted marginal rewards.
(2) Solve the allocation problem for this held-out as in Eq. 5), with the additional constraint that all queries in bin receive the same budget allocation. For each bin, store the assigned budget.
(3) During deployment, compute reward prediction for each Ii, map it to a bin, and return the budget associated with that bin.
During deployment, all queries can be processed independently (at the risk of a budget violation the distribution of queries is significantly different from those used to compute the allocation policy).
3.3 INTERESTING SPECIAL CASES
In domains such as coding and math, rewards are often binary, indi- cating success or failure. For instance, an outcome reward model can be used to assess correctness in math while unit tests indicate correctness in coding. In such settings, the probability that single sample from the model succeeds may be used to analytically compute all marginal rewards.
Binary Reward Best-of-k:
Let À |x) [1 y)] denote the probability of obtaining successful result from single sample (for rar, y) E 1)). Then the reward estimate g(2,, b) is simply the probability of getting least one success attempts: q(x, b) - (1 - (.. Then A(x,, b,) A, x,)
In this case, rather than training the reward predictor Å using a squared loss as in Eq. 6), we obtain empirical estimates of X; training, then minimize the cross-entropy:
Σ A
log((((rs))) (1 Al)gg)(1 A(z;; 0)]
with appropriately parameterized {xx;)).
4

---

Published as a conference paper at ICLR 2025
For the routing setting, we learn A in special form that models the probability of outputs from the strong model p¹ being preferred over the weak model p as:
Routing:
Á(z;; bs = p(p" > pW (x) = By ~p w2>pl' (o( v1) - r2),
(8)
4 EXPERIMENTS
We apply our method to several adaptive decoding procedures: best-of-k and routing (to either large model or a sophisticated search algorithm). We evaluate improvements over standard de- coding and procedures that allocate computation uniformly across problem instances, as well as performance relative to the theoretical upper bound. addition, we provide intrinsic evaluations of the accuracy and calibration of reward predictors A.
4.1 ADAPTIVE BEST-OF-k
We use best-of-k reranking with reward model in three settings: Math, Code and Chat.
Methods: We evaluate the following methods:
1. Online Ada-BoK (ours): The online variant of our method that solves a joint optimization
problem, as detailed in Section 3.2
2. Offline Ada-BoK (ours): The offline variant that solves the allocation problem on a held-
out dataset, detailed in Section 3.2 This method only used in Math and Code domains.
3. Best-of-k: This baseline allocates the number of samples k = B to every query.
4. Oracle: Oracle is a non- non-implementable method that uses the ground truth marginal rewards to solve the allocation problem. This method solves the allocation problem by plugging in the true marginal rewards. It is unrealizable in practice, but provides upper bound on the reward that could be obtained if the learned marginal reward predictor Å were perfect.
Evaluation Metrics: Our main evaluation metric for Math and Code is the expected success rate under an oracle verification procedure:
Expected Success Rate
Σ i=1
(a (or is correct]]
(9)
where n is the number queries. Similarly, for Chat, is the expected reward:
Expected Reward
Berseee. (r yi)
(10)
71
To estimate these in practice, sample large number of generations Bmaz for each query and then use bootstrapping to approximate the expectation for different b,
The compute budget B is the average number of samples that may be drawn per query. Best-of-k with = uses the same number of total samples as our adaptive method. We plot the performance of different methods for different compute budgets, where budgets specified in terms of B.
Note that, in the binary reward-basee Math and Code domains, it may efficient to set b, = 0 for some queries, many problems have 0% success rate. In these cases, a default response such as / don't know may be returned, allowing the sample budget to be allocated more effectively elsewhere. For chat experiments, we require all b; 2
CODE
Setup: For the coding experiments, adopt subset of TACO, dataset focused algorith- mic code generation with problems sourced from various programming contest platforms (Li et al. 2023). We use custom off-the-shelf Starcoder-15B model, which was released by the creators of the TACO dataset after fine-tuning on the dataset. We use the official open- evaluation frame- work released with TACO. A generation is classified success if it passes all test cases. Finally,
5

---

Published as a conference paper at ICLR 2025
CODE
0.50
Ground Thuch Probsailties
R.8
4
2.6
0.6
0.35
Oracle
0.30
Offine AduDOs
reduetinne
0.0
0.2
0.4
0.5
0.8
1.0
0.0
02
04
0.5
0,8
1.0
2
2
2
24
2
Ground In Sucewas Probatiiiiien
(Compute Budget)
MATH
Ground Turth Trobabllnte
0.9
0.8
0.7
Chuci
0.2
0.6
Orfirc Adaide
0.0
0.2
0.4 Preductions
2"
8 (CCoppute Budger)
Figure 3: Results on Code and Math. The left column shows the distribution of LM success probabilities across queries. A significant number of problems on Code have 0% success rate, while Math has a flatter distribution difficulty. The middle column compares predicted marginal reward predictors to ground-truth values, alongside calibration curves (our predictors are well calibrated with the ground truth). Finally, the right column shows the performance curves for tested methods. Here, observe that adaptive computation generally outperforms the best-of-k baseline.
as unit tests serve as a verifier, a method is considered successful on a query if at least one of its k generations passes all test cases. We use the MLP variant to learn the marginal rewards.
Results: Present success rates for the code dataset. We first observe that while our methods perform significantly better in the low-budget regime (B 8), the results in the moderate-to-hin budget regime are mixed. Counter-intuitively, while our offline variant always outperforms best- of-k, our online variant actually falls below the best-of-k curve in the high-budget regime. This is because small errors in the learned marginal reward predictor can hugely impact solutions to the allocation problem (if a problem has a true success rate of 0%, but its success rate predicted to be 1%, it becomes an attractive candidate for a large budget). 50% of problems in the coding dataset have 0 success probability (for Math, this is only 5%). The offline variant avoids this pathology by binning impossible and low-probability solutions together, effectively "regularizing" allocations.
In summary, the coding results highlight the fact that prediction errors in marginal reward can ad- versely impact online allocation. At the same time, the consistent outperformance of the offline variant over best-of-lk underscores the value of adaptive compute budget allocation.
MATH
We use subset of the Numina-COw Dataset Li et al. 2024). which contains Math prob- lems obtained from diverse sources. We use Mistral AI's Mathstral-77 model, which is specialized for mathematical tasks, and is based on Mistral 7B (Jiang et al. 2023 We use this model off-the- shelf and do not perform any fine-tuning. We use an oracle verifier select the best answer out of k. This implies that if the model generates at least one correct answer out of k, it will be successful. The verifier uses a 2-stage pipeline that employs the evaluation framework of Hendrycks et al. in the first stage and custom LLM-verifier in the second stage (Hendrycks et al. 202T). This pipeline is detailed in Appendix A We used the LoRA variant to learn the marginal rewards.
Setup:
Figure illustrates the success rate across different compute budgets. Our online and offline methods consistently outperform best-of-kk for all compute budgets. While the improvements marginal in the low-budget regime (B < 8), adaptive computation shines in the moderate-to-higg
Results:

---

Published as a conference paper at ICLR 2025
Gradle Best K
Dracle
-3.8
Best-ofk
-4.0
Plaward
Reward
4.0
-4.2
411
-42
-4.4
-.3
-4.b
-4.4
2
2 B (Compute Budget)
2
2
2
2
B (Compute Budget)
(a) Full variant
(b) Tranches variant
Figure 4: Results of adaptive best-of-k in the Chat domain. Our full variant achieves marginal improvements, enabling a 0-10% reduction in compute budget while maintaining equivalent reward levels. More notably, the tranches variant exhibits substantial gains, achieving the same reward levels with a 25-40% reduction in compute budget.
budget settings (B 2 8), where they can achieve the same success rate as best-of-k while using 25- 50% less computation. This efficiency likely stems from allocation of most of the compute resources to medium and hard problems, which require more samples, while easy problems need only a few to be solved. Both the online and offline variants perform nearly identically, indicating that solving the optimization problem offline does not lead to any performance degradation in this setting. Finally, although significantly outperform best-of-k;, the oracle curve suggests that further improvements in the marginal reward predictor Á could provide even greater gains.
CHAT
Setup: For chat, we use subset of the LMSYS S-Chat dataset, which contains real-world conver- sations with 25 LLMs (Zheng al. 2024). We use the popular Gemma-7b-it model, which is an instruction-tunel model for chat (Team et al. 2024). We use reward as our primary evaluation met- ric. We use an off-the-shelf reward model called was ranked amongst the top 10 on RewardBench at the time writing. We used the MLP variant to learn the marginal rewards. Finally, we conduct evaluation on different subsets of the dataset-
-aa RM-8B, which
1. Full: The vanilla experiment uses the the entire test set.
2. Tranches:
The tranches experiment uses subset of 20% from the full test set. This subset is composed of queries that fall in the lowest 10% highest 10% of variance in rewards. The goal of the tranches experiment is simulate more extreme user distribution. The main goal of tranches is to simulate a more extreme user distribution compared to typical datasets, which are often collected in somewhat controlled settings and might not capture true user diversity. Full details in Appendix A.3
Figures 4a and 4b present the results for the full and tranches variant respectively. We focus on the small compute budget regime as chat requires much lesser search and we empirically observed rewards saturating quickly. the case of the fill variant, the gains are relatively mod- est. The Oracle curve shows 15-25% reduction in compute budget while maintaining the same average reward. Although our adaptive method consistently outperforms best-of-k everywhere, the reductions in compute budget are relatively small, ranging from 0-10%. This indicates that while adaptive allocation can provide some benefit, an equitable allocation performs almost at par. For the tranches variant, the results are notably different. Our adaptive method achieves substantial gains, reducing the budget by 25-40% while matching the rewards of best-of-k.
Results:
4.2 ROUTING
Through this set of experiments, we aim to evaluate the effectiveness of our learned preference predictor in routing queries between weak decoding procedure " and strong procedure n²

---

Published as a conference paper at ICLR 2025
Model Size
7.00
1.0
1.75
-3.0
0.0
1.50
Dendry
0.6
0.4
0.25
0.2
Adapiin Pooding
Drace
0.00
0.0
0.2
0,4
0.6
0.8
10
0.0
0.0
0.2
0,4
0.6
0.8
3.0
-4.
0.2
Dround Trut Phofarence Photali
Prodictinns
(Irraction queries roulted to
Value Augmented Sampling
1.0
-16
-1.2
0.8
-1.
-22
-3.4
1.0
-28
Dracle Adsctive Souting Rancom Rouitio
0.0
0.2 08
1.
Gomnns Trith Prnfernece Profati
Prodictions
(Fraction
p*)
Figure 5: Results for different routing procedures. The left column shows the distribution of ground-truth preference probabilities for model size and value-augmented sampling. The middle column compares our learned preference predictors to ground-truth values, alongside calibration curves (our predictors well calibrated). Finally, the right column shows the performance curves for tested methods. We observe that adaptive routing outperforms the random routing baseline, demonstrating that learned predictor is able effectively route more challenging queries the stronger model while leaving simpler queries to the weaker model.
Settings: We present results for two different (p", .p²) pairs which are based on:
Model Size: this setting, p W and D are models from the same family but with different model sizes. We use the instruction-tuned Gemma-2b-it and Gemma-7b-it models p and pS. We use the full variant of the LMSYS dataset described in Section 4.1
2. Value-Augmented Sampling: In this setting, n and D have the same base LLM model. During decoding, a value function used to guide search, improving performance but with significant computational overhead. In this experiment L.M and the value function. In each decoding step, we compute the value 10 possible tokens, increasing the cost of decoding by factor of 10. We the harmless subset of the popular Anthropic HH dataset (Bai et al. 2022)
use Ilama-2 7B for both the
Methods: We present results for the following methods:
1. Online Routing (ours): Å is used to predict the preference probabilities for set of queries. These predictions are then routed using the online allocation procedure.
2. Random: A simple baseline that randomly routes fixed fraction of queries to ~ Any target may obtained by changing this fraction.
3. Oracle: As above, non-realizable skyline that uses ground-truth information about the reward distribution " (- (x4) and D xx)) for routing.
Setup: We use expected reward, defined in Eq 10 as our main evaluation metric. For reward prediction in experiments with variable model size, we use the smaller nd as the base LLM. That is, we train using the hidden states of p or perform LoRA fine-tuning of p This ensures that computational overhead during inference is minimal. and p does not even have to be called at all for some queries. The compute budget B is the fraction of the total queries that can routed to p*. For example, B 0.7 allows routing 70% of the queries n
8

---

Published as a conference paper at ICLR 2025
Code
70
Easy Medium Hard
Easy Haad
Budget Allecation 180
Budget Allccation (%)
40-
50-
20
10-
2-2
2-4
2+0
2-4
(Compute Budget)
(Compute Budget)
Figure 6: Allocation of compute at different budgets. The "Easy", "Medium", and "Hard" cat- egories are obtained by binning queries according their predicted success probability. In the low-budget regime, the majority of the budget is allocated to easy and medium queries; as budget increases, primarily allocated to hard queries.
Results: shows that for both experiments, the learned predictor effectively routes the more challenging queries to the stronger model, while simpler queries are handled by the weaker model. This leads to substantial compute savings without significant loss overall performance. For example, in Value augmentation, we achieve up to 25-40% reduction in calls to the more expensive decoding scheme while maintaining similar levels of reward.
Strikingly, for some budgets, routing scheme actually outperforms the strong decoder because the weak decoder sometimes outperforms the strong one, routing can achieve better overall perfor- mance. Our reward predictor learns and exploits this pattern.
4.3 ANALYSIS
How DO THE LEARNED MARGINAL REWARD PREDICTORS PERFORM?
So far, we have seen that our learned marginal reward predictors can be used effectively in adaptive compute allocation. However, we also evaluate their performance independently of the adaptive compute allocation. To this end, we introduce three evaluation metrics:
1. Avg. (Average Loss): The empirical loss when the prediction for every query is the average marginal reward, i.e., Á + E. In this case, the model predicts the same marginal reward for each imput. If the representations of the language model do not carry any mean- ingful information, we expect the performance to approximately equal this average loss.
2. Opt.* (Oracle Loss): The loss that a perfect predictor would achieve. Since we use soft labels rather than binary (0-1) labels, the minimum possible loss is positive value.
3. Acc. (Accuracy): The accuracy predictions the median of A values is used threshold. That is, if A; > median, the label for that query is 1, and otherwise 0. The accuracy of random predictor would 50%.
Table summarizes the results. In all cases, the achieved loss is lower than the baseline, indi- cating that the queries contain meaningful sig- nals about the model's response distributions. In routing, the test loss closer to the baseline, which can be attributed to the low entropy in the reward distribution (Figure [5). Across settings, accuracy exceeds 70%, showing that difficulty is predictable but highlighting scope for improvement. In Appendices Cand and D test how our predictors generalize to different data distributions and decoding methods.
Setting
Ours
Avg.
Opt.
Acc
Code
0.33
0.58
74%
Math
0.48
0.69
0.34
84%
Chat (Model routing)
0.64
0.67
0.58
72%
Chat (VAS routing)
0.55
0.57
0.51
72%
Table 1: Learned predictors achieve lower loss than fixed baselines (Avg.), approach optimal val- ues (Min.), and accurately discriminate queries of above- and below-median difficulty (Acc.).

---

Published as a conference paper at ICLR 2025
HOW DOES ALLOCATION VARY WITH PREDICTED DIFFICULTY?
Having established the performance benefits of adaptive allocation, we now explore what the dis- tribution of compute budget for different marginal reward predictions looks like. To investigate this, we focus on the Code and Math domains, where allocation is performed using A2, the success probability of a query. We stratify the predicted success probabilities into three evenly-sized bins according to their predicted success probability.
Figure 6 illustrates how the compute allocation changes across bins as the budget increases. At the lowest budget, the majority of the allocation goes to queries predicted to be easy or medium in difficulty. At the highest compute budget, by contrast, most of the compute is allocated to the hard bin. This shift can be understood intuitively; easy-to-moderate queries typically require only a few samples to solve, beyond which the marginal gain of additional samples decays rapidly. In contrast, for queries with low success probabilities, the marginal gain remains high even with a large number of samples and decays extremely slowly. Finally, the distinct allocation patterns between the two domains highlight how underlying difficulty distributions significantly impact allocation strategies.
RELATED WORK
Decoding Procedures in LLMs. There has been extensive research into utilizing different decod- ing schemes to enhance LLM capabilities, usually by expending inference-time compute. One of the most straightforward approaches is best-of-k sampling. where k different model responses are
generated per query, and the final model output is selected using reward model (Gao et al. [003]
Beirami et al. 2024). majority voting (Wang et al., 2023). or verifiers (Li et al. 2022). Another
line of work has shown that allowing LLMs to generate intermediate "reasoning" steps, or "chain-
of-thoughts" (CoTs) can significantly improve their final answer (Wei et al. 2022] Nye et al. [02T]
Zhou et al.) 2023). Extending this, some recent works have focused training models to generate thought tokens, enabling improved reasoning with longer generations (Zelikman et al. 2024: Goyal et al., 2023). However, while this can boost performance, it incurs additional computational costs due to the need to decode (potentially many) reasoning tokens. LLM decoding can also be framed sampling from tree of possible sequences, inspiring research that uses inference-time compute
to search through this tree YYo et al. 2024 Liu et al. 2024 Han et al. 2024).
Adaptive Computation Time in NN. Several works explored the idea of learning how to adapt inference-tine compute in neural networks (Graves) 2016); Dehghani et al.) 2019 [Banino et al. 2021]. These works focus on networks with recurring components, where the decision is the number of times to pass the input through these components. Our work, however, is architecture-agnstie and focuses on adaptively selecting decoding procedure. recent line of works has demonstrated the compute-optimal inference can outperform simply using larger models (Snell et al.) 2024| Wu al. (024)). Here, we show how to realize these improvements in practice using learned difficulty predictors. Concurrent to our work, Manvi et al.] (2024) demonstrated that language models can self-assess their response quality, enabling adaptive compute allocation.
CONCLUSION
We have introduced an approach for adaptively scaling test-time computation by predicting which queries would benefit the most from additional computation. We first showed that it is possible to learn lightweight models that predict marginal rewards on top of pre-trained that LMs encode usable information about the reward distribution of their responses to inputs. We then presented an allocation algorithm that adaptively allocates computation to queries. Results in programming, mathematics, and chat show our approach gives significant reductions in computation at a target level of output quality, or improvements in quality given a fixed computation budget.
Limitations: In this work, we assumed that we have training dataset of reasonable size (> 3K queries) to train our marginal reward probe. We also assumed access to a verifier in the Math experiments. However, such verifiers will generally not be available, and using specialized reward models might be the more practical implementation of our method.
The gap between our performance and the oracle indicates potential for improvement in marginal reward prediction. This could be addressed by exploring more advanced prediction models or by allocating additional inference-time computation for better marginal reward estimation.
Future Work:
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by the MIT-IBM Watson AI Lab and the National Science Foundation and Intel under grants IIS-2212310 and CCF-2217064. Jacob Andreas supported by a Sloan fellowship; Andi Peng is supported by the NSF GRFP and Open Philanthropy; and Idan Shenfeld is supported by the Qualcomm Innovation Fellowship.
REFERENCES
Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, et al. Training helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05862 2022.
Andrea Banino, Jan Balaguer, and Charles Blundell. Pondernet: Learning to ponder. In 8th ICML Workshop on Automated Machine Learning (AutoML), 2021.
Ahmad Beirami, Alekh Agarwal, Jonathan Berant, Alexander D' Amour, Jacob Eisenstein, Chirag Nagpal, and Ananda Theertha Suresh. Theoretical guarantees on the best-of-n alignment policy. arXiv preprint ar arXiv:2401.01879 2024.
Bradley Brown, Jordan Juravsky, Ryan Ehrlich, Ronald Clark, Quoc V Le, Christopher Ré, and Azalia Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling. arXiv preprint urXiv:2407.21778 2024.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10. 14168, 2021.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations, 2019.
Yilun Du, Shuang Li, Antonio Torralba, Joshua B. Tenenbaum, and Igor Mordatch, Improving fac- tuality and reasoning in language models through multiagent debate. In Forty-first International Conference on Machine Learning, 2024.
Leo Gao, John Schulman, and Jacob Hilton. Scaling laws for reward model overoptimization. In International Conference on Machine Learning, pp. 10835-10866. PMLR, 2023.
Sachin Goyal, Ziwei Ji, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, and Vaishnavh Nagarajan. Think before you speak: Training language models with pause tokens. arXiv preprint arXiv:2310.02226, 2023.
Alex Graves. Adaptive computation time for recurrent neural networks. arXiv preprint arXiv:1603.08983. 2016.
Seungwook Han, Idan Shenfeld, Akash Srivastava, Yoon Kim, and Pulkit Agrawal. Value augmented sampling for language model alignment and personalization. arXiv preprint arXiv:2405.06639, 2024.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the MATH dataset. In Thirty-fifth Conference Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation large language models. In International Con- ference Learning Representations, 2022.
Athul Paul Jacob, Yikang Shen, Gabriele Farina, and Jacob Andreas. The consensus game: Lan- guage model generation via equilibrium search. The Twelfth International Conference on Learning Representations, 2024.
11

---

Published as a conference paper at ICLR 2025
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.06825, 2023.
Jia Li, Edward Beeching, Lewis Tunstall, Ben Lipkin, Roman Soletskyi, Shengyi Costa Huang. Kashif Rasul, Longhui Yu. Albert Jiang, Ziju Shen, Zihan Qin, Bin Dong, Li Zhou, Yann Fleureau, Guillaume Lample, and Stanislas Polu. Numinamath. [Chttps: http:///w........o///////aaaaaa. http://s///giiuuuu..oooereeeee com/projject-nuinna ainneppoggeessssprrr//////////ap
wat.aat.aaaatttpaaaaaaaa pdf) 2024.
Rongao Li, Jie Fu, Bo-Wen Zhang, Tao Huang, Zhihong Sun, Chen Lyu, Guang Liu, Zhi Jin, and Ge Li. Taco: Topics in algorithmic code generation dataset. arXiv preprint arXiv:2312. (in::23/2.44880 14852, 2023.
Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, et al. Competition-level code generation with alphacode. Science, 378(6624):1092-1009 2022.
Jiacheng Liu, Andrew Cohen, Ramakanth Pasunuru, Yejin Choi, Hannaneh Hajishirzi, and Asli Celikyilmaz. Don't 't throw away your value model! generating more preferable text with value- guided monte-carlo tree search decoding. In First Conference on Language Modeling. 2024.
Ziqin Luo. Haixia Han, Haokun Zhao, Guochao Jiang, Chengyu Du, Tingyun Li, Jiaqing Liang, Deqing Yang, and Yanghua Xiao. Sed: Self-evaluation decoding enhances large language models for better generation. arXiv preprint arXiv:2405. /6552, 2024.
Rohin Manvi, Anikait Singh, and Stefano Ermon. Adaptive inference-time compute: Llms can predict if they can do better, even mid mid-generattio. arXiv preprint arXiv:2410.02725, 2024.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2/ 2021.
Show
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game go with deep neural networks and tree search. nature, 529(7587):484-48) 2016.
Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling IIm test-time compute optimally can be more effective than scaling model parameters. arXiv preprint 4:X1::2408.03311,, 2024,
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based gemini research and technology. arXiv preprint arXiv:2403.08295 2024.
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc V Le. Ed H. Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain thought reasoning language models. In The Eleventh International Conference on Learning Representations, 2023.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi. Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in 2022.
neural information processing systems,
Yangzhen Wu, Zhiqing Sun, Shanda Li. Sean Welleck, and Yiming Yang. Inference scaling laws: An empirical analysis of compute-optimal inference for problem-solving with language models. arXiv preprint ar arXiv:2408.0072 2024.
Shunyu Yao. Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. Ad- vances in Neural Information Processing Systems, 36, 2024.
Eric Zelikman, Georges Harik, Yijia Shao, Varuna Jayasiri, Nick Haber, and Noah D Goodman. Quiet-star: Language models can teach themselves to think before speaking. arXiv preprint arXiv:2403.080 099299 2024.
12

---

Published as a conference paper at ICLR 2025
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Tianle Li, Siyuan Zhuang, Zhanghao Wu, Yong- hao Zhuang, Zhuohan Li, Zi Lin, Eric Xing, Joseph E. Gonzalez, lon Stoica, and Hao Zhang. L.MSYS-chat- Im: A large-scale real-world LLM conversation dataset. In The Twelfth Interna- tional Conference on Learning Representations, 2024.
Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuur- mans, Claire Cui, Olivier Bousquet, Quoc V Le, and Ed H. Chi. Least-to-moot -ost prompting enables complex reasoning in large language models. In The Eleventh International Conference on Learn- ing Representations, 2023.
13

---

Published as a conference paper at ICLR 2025
A EXPERIMENT DETAILS
A.1 MATH
We use a subset of the Numina-COT Dataset (Li et al., 2024),w which contains math prob- lems from various sources. We filtered out multiple choice questions from the dataset. Our final dataset contained 15K training samples, 2K validation samples and 3K test samples.
Dataset:
Model: We Mistral Al's Mathstral-7B model, which is specialized for mathematical tasks, and is based on Mistral 7B (Jiang et al. 2023 We use this model directly off-the-shei and do not perform any fine-tuning.
Verification Pipeline: We used 2 stage pipeline to verify the correctness of an answer. The first stage used the evaluation framework of Hendrycks al. (2021). This framework expects the LM to enclose its answer within boxed {} and then runs an automated evaluation of the extracted answer against the ground truth answer. However, we found that the automated evaluation had a high false negative rate for various reasons:
1. Answer not enclosed within boxed O For some queries, the model had the right answer but did not enclose it inside boxed ()- This led to the responses being classified as incorrect.
2. Errors in Checking: We found that due to errors in checking, many correct responses were marked incorrect. Some examples of false negatives are:
1. Model Response: frac (2c) Ground Truth Answer: 12.5
2. Model Response: x=2 Ground Truth Answer: 2
3. Model Response: -0.5 Ground Truth Answer: -frac 1 [2]
4. Model Response: 11.00 Ground Truth Answer: 11
These errors particularly harmful for the training of our marginal reward predictors. because queries predicted to be (and empirically) easy were being marked incorrect. Furthermore, our online allocation exacerbated this issue, by assigning very few samples to such queries.
To address this issue, we introduced a second stage in the verification pipeline that used a LM as evaluator. Because stage does not have false positives, only used stage 2 for responses that were marked incorrect in stage 1. We used Llama-3. 1-B---..-----nttrree with temperature set to 0.1 as our evaluation LM, and prompted follows:
You are a math evaluation agent. You are tasked with evaluating if the final answer from an excerpt of the response matches the given gold truth answer. The format or units of the response and gold truth answer might be different. However, you must evaluate if the answers are numerically equivalent/identical. they must simplify to the same value Your response should be single word followed by an explanation. 'YES' if the unswers are equivalent and 'NO" if they are not. Examples:
extra careful when evaluating fractions,
A) 7% and 7 are equivalent
B) frac frac 10)2} and frac {4) are equivalent. 3,5,7 and 3,8,9 are not equivalent.
Ground Truth Answer: < >
Response: <
Note that if the model's response had an answer enclosed inside boxed {. we only provided that extracted answer. If this not the case, then provided the model's 's entire response to the LLM evaluation agent.
The 2-stage pipeline helped us significantly improve the evaluation process, well as stabilize the training process for our marginal reward predictors. However, this pipeline is not without flaws, did notice examples small number of false positives in the second stage.
14

---

Published as a conference paper at ICLR 2025
We generate 128 responses for every query (lemperature=0.7) and label them using the verification pipeline. We use the labels to compute the empirical mean success probabilities X₃. The estimated A are then used as targets to train (z; f), the marginal reward predictor.
Training:
Evaluation: Our adaptive allocation algorithms assigned maximum budget of Bmax 128 samples per query (that is the allocation for any query is capped at 128 samples).
A.2 CODE
Dataset: We use subset of TACO, dataset focused on algorithmic code generation with problems sourced from various programming contest platforms such as CodeChef, CodeForces and Hacker- Rank. Our primary reason for selecting TACO was the availability of test cases for most problems, which is required to train a predictor of success probability. We filtered out problems that did not have single unit test. We also filtered out problems that were sourced from geeksforgeeks and aizu, as the official evaluation framework did not support those. Our final dataset consisted of 10K training samples, 1K validation samples and 1K test samples. Note that these three splits in our experiments were extracted from official TACO training set. This is because the TACO test set has a very different distribution of problem difficulties, while our method aims to produce accurate distribution allocation of computation. We believe producing distributionally robust predictors is an important topic for future work.
Model: We used the Starcoder-155 15B model, which was finetuned and open-sourced by the authors of Taco.
Training: We generate 100 responses for every query (lemperature=00.) and label them using the unit-test verifier to obtain the mean success probability A, The estimated A; are then used as targets to train À [x; f), the marginal reward predictor.
Evaluation: We used the official evaluation framework released by the creators of the TACO dataset. response was considered a success passed all available test cases. Our adaptive allocation algorithms are assigned a maximum budget of Biaa 100 samples per query (that is the allocation for any query was capped at 100 samples).
A.3 CHAT
Dataset: We use subset of the LMSYS- dataset, which contains one million real-world versations with 25 LLMs. We filter the dataset to only select samples in English and with less than 10 turns. We also filter out samples which were labelled redacted, as these modified. Our final dataset consisted of 50K training samples and 5K test samples.
often artificially
subset of 20% from the full test set. To create this subset, we first generated multiple responses for each query in the test set and labeled them using the reward model. We then calculated the reward variance for each query and selected only those queries that fall in the lowest 10% or highest 10% of variance. In essence, the tranches test set is composed of queries with the lowest and highest variance in rewards, representing the two extremes. The main goal of the tranches experiment is to evaluate our method when the distribution of queries differs from typical datasets, which are often collected in somewhat controlled settings. Such datasets may not fully capture the true diversity of user queries that are encountered by general chatbots. While understanding the true distribution of user queries is beyond the scope this work, the tranches experiment simulates more extreme distribution to provide performance insights.
Tranches Variant:
The tranches experiment
Model: We used the Gemma-2Ba 2B-it model. We did not perform any fine-tuning.
Training: We generate 8 responses (temperature=0.7) for every query and label them using 30OG
which was ranked amongst the top 10 on RewardBench at the time of writing. We then use bootstrapping approximate A, The approximated A, are then used as targets train A (z; 8), the marginal reward predictor.
Evaluation: We evaluate using the same reward model as above. Our adaptive allocation algorithm was assigned maximum budget of B 8 samples per query.
15

---

Published as a conference paper at ICLR 2025
A.4 ROUTING: MODEL SIZE
Dataset: We used the same LLMSYS-Cha dataset that we used for the best-of-k: chat experiments.
Our weak model Gemma-2B-it. Our strong model was Gemma-7B-it. We did perform any fine-tuning on either these models.
Model:
Training: We generate 8 responses (at temperature of 0.7) for every query with both the models. We labelled them using (OODFT//laaaa3-OOtttBBiee RM--B as the reward model. Supervision for the reward model was then computed using Monte Carlo estimate of:
r
By up? w w (a(x,, !/) r(z, y2))]
(11)
Evaluation:
The adaptive allocation procedure uses the predictions of the marginal reward predic- tor to route the top 8th percentile of queries to p We evaluated using the same reward model mentioned above.
A.5 ROUTING: VAS
We use the harmless subset of the popular Anthropic-HiP dataset. We sampled 8K samples randomly from the train set, and 400 samples from the test set.
Dataset:
Model: Our weak decoding procedure was fine-tuned Llama-7B model. Our strong decoding procedure used an additional value function that was also based on the Llama-7B model.
Training: We generated 4 responses per query. The preference probability was then computed using a Monte Carlo estimate of:
p(p -p"1>) -p v1) r
(12)
We labelled the responses using Opunksssaatteeaadddddddetttttree as the reward model.
Evaluation: Our adaptive allocation used the predictions of the marginal reward predictor to route the top Tthh percentile of queries to ps. We evaluated using the same reward model mentioned above.
16

---

Published as a conference paper at ICLR 2025
MATH
GSMBK
0.40
0.9
0.35
Success Rate
0.8
0.30
Success
0.25
0.7
Cracle
0.20
Aoline-AdaBcNN
Oracle
0.6
Offline AdaBoN
0.15
Adaptive Routing
Best-of-K
Random Routing
0.10
20
21
2
2
24
2
2
0.0
0.2
0.4
0.6
0.8
1.0
B (Compute Budget)
(Fraction of queries routed p²)
Figure 7: Results on MATH and GSM8K. We evaluate the MATH dataset using our adaptive best- of-k approach and the GSM8K dataset using routing approach. The results show that adaptive compute allocation improves performance on both benchmarks, demonstrating the ability our learned predictor to effectively assess difficulty and the advantages of allocating compute adaptively.
B ADDITIONAL BENCHMARK RESULTS
B.1 SETUP
MATH We implement adaptive best-of-k pipeline the popular MATH bench- mark (Hendrycks et al. 2021). Similar to our Numina Math experiments, we use Mistral Al's Mathstral-7B model. We use this model directly off-the-shelf and do not perform any fine-tuning. To train our probe, we generate 64 responses per example for a total of 7, 500 training examples from the MATH dataset. We employed the 2-stage verifier pipeline used in the Numina experiments, de- scribed in Appendix A.1 We used the LoRA variant to learn the marginal rewards.
GSM8K We implement our routing setting on the popular GSM8K dataset (Cobbe et al. 2021). We use the instruction-ture Gemma-2b-it and Gemma-7h-it models n" and p² We use these models directly off-the-shelf and do not perform any fine-tuning. We use a simple binary reward structure, where correct answer is given reward 1. To train probe, we generate 16 responses per example for a total of 7, 500 training examples from the GSM8K dataset. We perform LoRA fine-tuning of n to learn the probe. As all problems in GSM8K have numerical answers, simply extract the last 2 numbers in the model's answer and classify an answer as correct if either of them match the ground truth answer.
B.2 RESULTS
Figure illlustrate the results for both benchmarks, highlighting the effectiveness of adaptive com- pute allocation in both scenarios. In GSM8K, adaptive routing improves absolute success rates by up 5% compared non-adaptive approaches while using same amount of compute. Similarly, in MATH, while the performance gains are more modest, adaptive compute allocation achieves com- parable success rates while reducing compute usage by up to 2x compared the baseline best-of-k: approach. This demonstrates that our learned predictor effectively captures query difficulty and facilitates efficient compute allocation.
17

---

Published as a conference paper at ICLR 2025
MATH
Anthrop c HH
-4.0
-4.1
0.9
Success Rate
-42
0.8
Reward
0.7
-4.4
Oradle
-4.5
MATH Probe
Oracle
0.6
Numina Probe
Adaptive Routing ILMSYS probel
-4.6
Best-of-K
Random Routing
20
2
2
2
2
0.0
0.2
0.1
0.6
0.8
1.0
B (Compute Budget)
(Fraction queries routed p²)
Figure 8: Generalization results on MATH and Anthropic HH. We evaluate the generalization of probes trained Numina and LMSYS by applying them to MATH and HH, respectively. The results demonstrate that our probes maintain high effectiveness for adaptive compute allocation, even when used outside their training data distribution. Remarkably, the Numina probe closely matches the performance of the MATH probe, suggesting that it is able to capture general features that are applicable across different mathematical datasets.
C GENERALIZATION TO DIFFERENT DATA DISTRIBUTIONS
In this section, we evaluate the effectiveness of our learned marginal reward predictors (probes) adaptively allocating compute queries outside their training distribution. Specifically, conduct two generalization experiments: a) applying the probe trained on Numina dataset to the MATH dataset and b) applying the probe trained LMSYS Anthropic HH.
C.I SETUP
Numina + MATH We use best-of-k probe, which was trained on the Numina dataset, and apply to the MATH dataset. Although the original Numina dataset includes MATH as a subset, excluded all examples sourced from MATH for this experiment.
LMSYS - Anthropic HH We use our routing probe, which was trained on the LMSYS dataset and apply it to the Anthropic HH dataset (Bai al. 2022). While LMSYS dataset containing real-world user conversations collected in the wild, data in HH was collected by instructing vetted master-qualified US-based MTurk workers with specific instructions. Thus, these datasets have significant differences in their user composition. We use the probe's predictions to route queries between Gemma-2b-it (p") and Gemma-7b-it (pS). Thus, while the same decoding procedures are used for both datasets, the probe we use has only been trained on the LMSYS dataset.
C.2 RESULTS
Figure 8 presents the results for our dataset generalization experiments. Our probes remain effective for adaptive compute allocation even when applied to data outside their training distribution, demon- strating strong generalization capabilities. Notably, the Numina probe matches the performance of the probe trained on MATH, suggesting that it successfully extracts general features related to math- ematical problem difficulty from the language model's 's representations, independent of the dataset.
18

---

Published as a conference paper at ICLR 2025
LMSYS
-3.6
-3.8
Reward
-40
Oratle
-4,2
BoK Probe
Random Routing
Routing Probe (Original)
-4.4
0.0
0.2
0.4
0.8 1.0
B (Fraction of queries routed to p*)
Figure 9: Results of applying best-of-k probe to Routing, We use the probe trained for the best of- decoding method and apply it to routing. The results indicate that while there some reduction in performance compared to probe specifically trained for routing, the best-of-k probe demonstrates effective generalization and is still able to deliver substantial gains.
D GENERALIZATION TO DIFFERENT DECODING METHODS
Our marginal reward predictors are typically trained for specific decoding procedure. In this sec- tion, we explore whether these predictors can generalize decoding procedures they were not ex- plicitly trained on. Specifically, we conduct two generalization experiments: (a) evaluating probe trained on the best-of-k approach when applied routing, and (b) assessing the generalization of probe trained on specific decoding temperature MATH across different decoding temperatures.
D.1 MOTIVATION
1. Applying BoK probe to Routing Through this experiment, we aim to understand if our adap- tive best-of-k: (BoK) probe be applied routing. Intuitively, we want to extract a proxy of query difficulty from our BoK probe and use that for routing. In the best-of-k setting, queries that have higher difficulty benefit from additional search (more samples). The difficult queries should thus have a high difference between the expected reward of best-of-k and the reward of taking a single sample. This range is given by:
Difficulty(z) Bo-Boo (r( v)] Exy v)l
(13)
Using this, we can now approximate the preference probability (similar to Eqn. 8 of using the stronger decoding procedure rd as:
plpS p0))) (Difficulty(x))
(14)
2. Generalization across temperatures The targets for training our probes are generally obtained by using decoding with a fixed temperature (generally 0.7). In this experiment, we want to test the generalization of our probe and the effectiveness of adaptive compute allocation if different decoding temperatures are used.
D.2 SETUP
Applying BoK probe Routing We use the probe trained for the adaptive best-of-k ex- periments on LMSYS. (Section 4.1) and apply it routing on LMSYS. We use the instruction-tunel Gemma-2b-it and Gemma-7b-if models as p 4 and p We compare performance against the probe
that was trained specifically for routing in Section 4.2
Generalization across temperatures We evaluate the performance of our probe on the MATH dataset, focusing on its adaptability across different temperature settings. The probe is trained
19

---

Published as a conference paper at ICLR 2025
Temperature [Onginall
Temperature 0.2
Temperature
Dsacke
Dracke
Bust of
Orins-AcatoN
Driins AdaDoN
of
0.9
Success Rlato
0.0
Succass
0.7
0.6
0.6
20
22
21
27 (CCompute Budgget)
2
27
2
[Commpute Budipet)
[Cooppuee Budgee)
Figure 10: Results on decoding with different temperatures. We assess the performance of our probe, trained at decoding temperature 0.7, across various decoding temperatures. Despite being trained for specific temperature, the probe remains effective across varying decoding temperatures. Intuitively, query difficulty is an intrinsic property and should independent (or only weakly de- pendent) on the decoding temperature.
a temperature of 0.7, but we evaluate its generalization on two temperatures: 0.2 and 1. For each temperature, we perform decoding estimate the empirical ground truth success probabilities and then test performance effectiveness of adaptive compute allocation. Since the scaling curves for Oracle and Best-of-K are temperature-sensitive,, we present three separate plots to illustrate the results for each temperature.
D.3 RESULTS
Figures 9 and 10 present the results for the BoK BoK-routing and temperature experiments respectively. We find that our BoK probe applied routing continues deliver significant gains compared to non-adaptive baselines. However, we do observe a drop in performance when compared to the probe that has been trained specifically for routing, indicating the existence of some generalization gap. In the temperature generalization experiments, we find that our probe remains effective across multiple decoding temperatures. These results suggest that our probes effectively capture notion of query difficulty that shows some degree of consistency for different decoding procedures i.e., query considered difficult for procedure is also likely to considered difficult for procedure B.
20


Reviews:
[{"review_id": "n0lyTxeTcz", "reviewer": "Reviewer_46zi", "summary": "This paper presents an approach to adaptively allocate computational resources for language model (LM) decoding based on input difficulty. The authors propose a framework that predicts the marginal benefit of additional computation for each query, enabling dynamic adjustment of decoding procedures to maximize efficiency without sacrificing output quality. They demonstrate their method across tasks in math, code generation, and dialogue, achieving up to 50% reduction in compute usage in some cases. The paper also introduces two adaptive procedures—best-of-k sampling and routing between models of varying complexity—and provides a thorough evaluation using both online and offline allocation strategies.", "questions": "- Typos:\n    - Line 352: \"which in an\" should be \"which is an\".\n    - \"LoRa\" should be changed to \"LoRA\".\n    - In Figure 1, \"Large LM\" should be changed to \"large LM\" for consistency.\n- What's the definition of \"N\" in equation (10)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The paper adeptly formulates the \"adaptive computation scaling allocation\" in the context of LM decoding, addressing a topic that is both timely and relevant.\n- The proposed computation-allocation framework is comprehensive, covering various cases and scenarios, including binary reward, pairwise optimization in routing, and both online and offline design considerations.\n- The experiments conducted on three diverse and representative domains demonstrate the efficiency and efficacy of the proposed computation-allocation strategies.", "weaknesses": "- The main concern is that the current computation-allocation solution is only evaluated in scenarios with identical distributions (i.e., the training data used to train the difficulty model comes from the same distribution as the test set). It is unclear whether the trained difficulty model generalizes to other distributions. The generalizability of the difficulty model is crucial for determining the practicality of the proposed computation-allocation framework.\n- Following from the above, since the choice of LLMs does not seem to affect the evaluation of the proposed method’s efficacy, why not select a single fixed LLM, such as Llama3-7b-Instruct? By doing so, it might be easier to assess the generalizability of the method. (Please correct me if there is an issue with my understanding.)\n- The implementation of the baselines is weak, with only one effective but not particularly practical baseline (best-of-k and random) in each scenario. Between the proposed method and these baselines, there are likely other reasonable approaches that could better demonstrate the effectiveness of the proposed framework.\n- The related work section is too concise and lacks comprehensiveness, especially in the discussion of relevant adaptive computing research. Only one recently published paper is mentioned, which undermines the paper's completeness and contextual grounding.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "sAVkNlhF3e", "reviewer": "Reviewer_NnGM", "summary": "Presents an input-adaptive method for test-time compute-allocation. Decoding methods apply either sequential (eg weak vs strong model) or parallel compute (eg more samples in best-of-n). For a given method, this paper proposes to predict the marginal utility of every unit of computation, then use these predictions to optimize compute allocation. The paper proposes to predict these utilities given only the input.\n\nThe resource allocation problem can be solved in an offline manner given a fully observed dataset, referred to as online allocation in the paper, or solved via online access to only a partially observed dataset, referred to as offline allocation in the paper.\n\nExperimental results across coding, math, and chat indicate that utility prediction is difficult at the extremes, and allocation decisions are sensitive to utility errors. Overall, adaptive allocation outperforms uniform or random allocations. The partially-observed strategy empirically often does better than the fully-observed case, possibly due to coarsening effects that hide errors in utility prediction.", "questions": "Drawing inspiration from the online secretary problem, it would be interesting to see how online estimation of pass rates for coding can aid utility estimation. For example, one could increase the total computation budget and, for each problem, reserving some of that budget to utility estimation. This would alleviate some of the burden from the prediction model.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper tackles a novel and timely problem, and offers a reasonable approach. The paper is clearly written.", "weaknesses": "A small criticism is the naming convention of online versus offline. Online optimization refers to \"optimization problems having no or incomplete knowledge of the future (online),\" which is not how online is used in this paper.\n\nOther than that, this paper is a good step in improving adaptive test-time compute, identifying the importance of accurate utility estimation in problems with very low success rates.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "TTY7W3L4Ao", "reviewer": "Reviewer_igGM", "summary": "This work proposes an input-adaptive computation allocation mechanism for improving the efficiency of test-time computation. The core idea is to train a model that predicts the distribution of rewards given a query and a budget. It incorporates training an MLP LM head and LoRA as the reward predictor that estimates the difficulty of a batch of queries. The proposed adaptive best-of-k outperforms the efficiency of standard best-of-k baselines in math, code, and chat domains. In addition, the author demonstrates the improvement in routing in terms of different model sizes and decoding schemes. The additional case study in inspecting the allocation of computation at different budgets is intriguing.", "questions": "1. Though I understand using a query only to predict the reward should incur less latency, will $y$ be more informative and easier to train the predictor?\n\n2. Could you please report the Spearman Correlation in Figure 3 (b, middle column)?  \n\n3. Could you provide more clarification on the computing budget? Is it based on the inference calls?\n\nI will be happy to raise my score if the author could address the aforementioned limitations and concerns.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. Scaling the test-time compute is effective but costly, this work contributes to a timely direction with a smart input-adaptive allocation scheme improving test-time efficiency.\n\n2. The empirical improvement in efficiency is noticeable, and this work has covered adaptive allocation in representative popular subdomains: sampling, model size, and decoding method.\n\n3. The presented analysis in Figure 6 is intuitive.", "weaknesses": "1. The selection of datasets and backbone language models may be questionable. I suspect this method should be ideally generalizable across tasks, however, only a single data in each domain is selected. I expect to see more tasks like HumanEval, MBPP for coding, Hendrycks MATH, and GSM for math. Meanwhile, for each domain, the author selects a specific backbone LM rather than the same choice across all tasks. This may raise concerns about the generalization of the proposed method.\n\n2. The underlying difficulty of this method is to actually train a very good difficulty estimator. However, the training difficulty, and the heavy training data resource requirement for learning a good reward predictor have not been explicitly discussed in the context. Moreover, it is highly dependent on the task, and I suspect the difficulty of some tasks will not be easy to predict. \n\n3. The proposed method only considers the query for training the reward predictor. However, though there is a latency for querying the model, I suspect introducing $y$ will be more informative to reflect the difficulty of a task. \n\n4. In Figure 3 (middle), besides the left bottom and right top clusters, the rest correlation appears to be relatively poor. Therefore, I suspect the efficiency gain could be mostly coming from predicting “unanswerable” for the queries in the left bottom regions and putting 0 costs there, also assigning a minimum budget to always correct questions. However, the middle region is actually the region that should benefit from a smart computation allocation scheme, and the correlation is not convincing here.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "mmtkLGRSuJ", "reviewer": "Reviewer_amKi", "summary": "This paper proposes a method for the adaptive allocation of decoding computation. By employing an LLM-based probe to predict the difficulty of a given query, the approach dynamically adjusts the allocation of decoding resources. The authors validate the method’s effectiveness across coding, math, and chat tasks. Results demonstrate that, under computational constraints, this approach outperforms the baseline BoK method.", "questions": "1. Please explain the choice of benchmarks and how they compare to HumanEval and MATH in terms of difficulty distribution. Or please add more benchmarks like HumanEval and MATH. \n\n2. If a more complex decoding method, such as MCTS, is employed, would it necessitate retraining the probe? This could suggest a mismatch between the model's capabilities when using more advanced decoding methods and the probe's predictions. Additionally, it raises the question of whether the probe's prediction accuracy may be affected by factors such as varying prompts or decoding methods, and whether the probe demonstrates robustness under these conditions. Please explain the discuss the generalizability of the probe.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper achieves efficient decoding from a different perspective, showing clear improvements over the original BoK method.\n2. The authors apply their method across three distinct domains—code, math, and chat—demonstrating generalizability of their method.", "weaknesses": "1. In the experiments for code and math, the authors employ less-used benchmarks rather than widely adopted ones like HumanEval and MATH, raising concerns about the method’s applicability to broader tasks.\n\n2. The paper's baseline comparison is limited to the BoK method, lacking comparative experiments with other stronger efficient decoding methods, such as Speculative Decoding.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)