Learning How Hard to Think: Input-Adaptive Allocation of LM Computation

Authors: Mehul Damani, Idan Shenfeld, Andi Peng, Andreea Bobu, Jacob Andreas
Primary area: foundation or frontier models, including LLMs
Keywords: LLM, inference, scaling, test-time compute

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.

Decision: accept (poster)
Reviewer scores: 6, 8, 6, 6

Reviews and discussion:
<thread>
creation date:
1730036128110

thread_type:
official_review

<thread_object>
creator: reviewer_amki

<content>
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.

strengths:
1. This paper achieves efficient decoding from a different perspective, showing clear improvements over the original BoK method.
2. 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.

2. The paper's baseline comparison is limited to the BoK method, lacking comparative experiments with other stronger efficient decoding methods, such as Speculative Decoding.

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. 

2. 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730524600963

thread_type:
official_review

<thread_object>
creator: reviewer_iggm

<content>
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.

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.

2. The empirical improvement in efficiency is noticeable, and this work has covered adaptive allocation in representative popular subdomains: sampling, model size, and decoding method.

3. 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.

2. 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. 

3. 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. 

4. 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.

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?

2. Could you please report the Spearman Correlation in Figure 3 (b, middle column)?  

3. Could you provide more clarification on the computing budget? Is it based on the inference calls?

I will be happy to raise my score if the author could address the aforementioned limitations and concerns.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730691102677

thread_type:
official_review

<thread_object>
creator: reviewer_nngm

<content>
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.

The 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.

Experimental 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.

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.

Other 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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731073415417

thread_type:
official_review

<thread_object>
creator: reviewer_46zi

<content>
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.

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.
- 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.
- 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.
- 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.)
- 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.
- 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.

questions:
- Typos:
    - Line 352: "which in an" should be "which is an".
    - "LoRa" should be changed to "LoRA".
    - In Figure 1, "Large LM" should be changed to "large LM" for consistency.
- What's the definition of "N" in equation (10)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732037595276

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Comments and Revision Summary

comment:
We thank the reviewers for their time and effort in reviewing this paper. We are happy that the reviewers found our work timely and comprehensive. The insightful comments received have helped us refine this paper. Below, we summarize the main concerns raised and outline the additional experiments conducted to address them:

1. **Evaluation on standard benchmarks**: We have included new benchmark results on MATH and GSM8K (**Appendix B** in the revised submission). These results confirm that adaptive compute allocation remains beneficial in these settings. 


2. **Generalization of the Learned Difficulty Model to Unseen Data Distributions**: We conducted two additional experiments (**Appendix C**) to assess how well our difficulty model generalizes to data outside its training distribution. The results demonstrate strong generalization: adaptive compute allocation using our model consistently outperforms non-adaptive baselines.

3. **Generalization of the Learned Difficulty Model to New Decoding Procedures**: To evaluate robustness across decoding methods, we ran two new experiments (**Appendix D**) testing our difficulty model on decoding procedures it was not trained for. While there is a slight performance drop, the model still delivers strong results, indicating that our difficulty model has learned a general notion of query difficulty which is applicable to a range of decoding procedures.  its robustness to varied decoding approaches.


4. **Reasoning for using different LLMs for different experiments**: Since our difficulty model is learned on top of an LLM’s representations, our primary reason for using different LLMs was to demonstrate that it is possible to learn effective difficulty models for a variety of base LLM models. Since that raised concern among reviewers, we tried to address that (within the time limit of the rebuttal), and now, with the inclusion of our new results, all of our chat experiments (with one exception) and one Math experiment use the Gemma family of models.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper presents an adaptive computation allocation approach for LM decoding that predicts input difficulty to optimize resource usage. Key strengths include: comprehensive evaluation across diverse domains (code, math, chat), significant compute reduction (up to 50%) without quality loss, and strong generalization results on standard benchmarks (MATH, GSM8K). Main weaknesses are: initial limitation to same-distribution evaluation, use of different LLMs across experiments, and moderate prediction accuracy for medium-difficulty cases. However, authors adequately addressed these through additional experiments showing cross-dataset/model generalization. The paper makes a valuable contribution to efficient LM deployment and is recommended for acceptance.