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

Q1. Please explain the choice of benchmarks and how they compare to HumanEval and MATH in terms of difficulty distribution.

A1. Thanks for the question! To address concerns about generalization to other datasets, we have added a new set of benchmark results (Appendix B in the revised submission) on MATH and GSM8K. Specifically, we evaluate on MATH using our adaptive best-of-k approach and on GSM8K with our routing approach. We find that adaptive compute allocation improves performance on both benchmarks. Notably, adaptive routing on GSM8K improves absolute success rates by up to 5% (a relative increase of nearly 20%) while using the same amount of compute as non-adaptive methods.

We also present results on Anthropic HH [3] (Appendix C), which has been used as a standard benchmark in RLHF. However, HumanEval/MBPP do not have training datasets and thus we are unable to train difficulty models for these benchmarks.

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

A2. Thank you for your question. We would like to emphasize that the contribution of this paper is not a particular decoding method (such as adaptive best-of-k or routing), but instead to show that adaptive test-time compute allocation can be beneficial across a diverse set of existing decoding methods. In this sense, the “baseline” methods are those that allocate a fixed amount of computation per query. Concurrent work by Snell et al.[2], also shows the value of adaptive computation and does not consider method-specific baselines.

Most of the relevant work in the area presents different methods to use test-time computation (chain-of-thought, generate and revise, MCTS, etc) . In this work, we consider a different axis which tries to adaptively allocate this computation, making our method complementary to most test-time methods.

Speculative decoding is only applicable to our routing setting with LM size (which is only 1 out of our 5 experiments). Moreover, speculative decoding is not input-adaptive (our main novelty) and uses a fixed frequency to query the large model, making it combinable with our method. Specifically, speculative decoding uses a fixed, query-invariant frequency to query the larger model. Combining with our framework would imply input-adaptively choosing at what frequency queries should be verified on the larger model.

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

A1. Thank you for the suggestion. To address this concern, we have added a new set of benchmark results (Appendix B in the revised submission) on MATH and GSM8K. Specifically, we evaluate on MATH using our adaptive best-of-k approach and on GSM8K with our routing approach. We find that adaptive compute allocation improves performance on both benchmarks (Figure 7). Notably, adaptive routing on GSM8K improves absolute success rates by up to 5% (a relative increase of nearly 20%) while using the same amount of compute as non-adaptive methods. We also present results on Anthropic HH [1] (Figure 8), which has been used as a standard benchmark in RLHF. However, HumanEval/MBPP do not have training datasets and thus we are unable to train difficulty models for these benchmarks.

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

A2. Thank you for your question! Although the downstream compute allocation is independent of the base LLM, our difficulty model is learned on top of the base LLM’s representations. Our main reason for using different LLMs was to show that we can learn effective difficulty predictors across a variety of base LLM models.

We also provide new results on GSM8K with the Gemma family of models (Appendix B) and find that the performance gains are significant and follow the trend we observed for specialized Math models. Thus, all of our chat experiments (with the exception of value-augmented search for which no Gemma model is openly available) and the new GSM8K experiment use the Gemma family of models.

Q3. The heavy training data resource requirement for learning a good reward predictor have not been explicitly discussed in the context.

A3. Thank you for pointing this out! The probes we train are actually very lightweight and even when we use LoRA, the average training time is 3-4 hours on 1 A100 GPU. Constructing the training dataset requires MC sampling to estimate the targets and using VLLM, we found that this generally took anywhere between 4 hours (Math, 10K examples) to 10 hours (LMSYS, 50K examples). We also acknowledge that we need a training dataset for our probe, and have added this to the limitations section. Let us know if there is anything else we can do to address this.

Q4. The underlying difficulty of this method is to actually train a very good difficulty estimator. I suspect the difficulty of some tasks will not be easy to predict. 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.

A4. While perfectly predicting difficulty is indeed a challenging problem, our experiments show that even somewhat noisy predictions are good enough: on all 3 domains (and our new benchmark results),we are still able to predict difficulty to a degree that makes it useful for the downstream application of adaptive compute allocation.

Developing better difficulty models, although beyond the scope of this work, is definitely worth exploring and can significantly boost performance. We fully agree with the reviewer’s comment that conditioning on the response y can boost the performance of the difficulty model and are actually considering this for future work! We are considering a setting where we allocate some amount of initial computation to each query and use the evaluation of those responses to refine our difficulty estimate. In such a setting, conditioning on the response y can significantly improve performance. However, this will also substantially increase latency and comes with other computational costs.

Thank the author for the response. Since it has mostly addressed my concerns, I have increased my score.

Regarding A5, could you please also report the correlation if you have excluded the extreme regions' points? I am curious to know the moderate region's correlation.

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

A1. Thank you for your suggestion, this is an extremely interesting idea and one we had already starting thinking about as a follow-up study! In a serial setting (and unlike the standard secretary problem), individual decoding results y can provide fine-grained information about the distribution of future outputs from different sampling methods. A good difficulty model can condition on ys, and many more complex decoding strategies are possible. We think this is a great direction for future research.

Q2. Online vs offline

A2. Thanks for pointing this out—we will clarify in our final revision.

Q1. Does the computation allocation solution generalize to new data distributions?

A1. Great question! To answer it, we ran two new experiments (see Appendix C in the revised submission) to evaluate how our difficulty model generalizes to data distributions it was not trained on:

1. Applying the difficulty model trained on the Numina dataset to the popular MATH benchmark (Figure 8): We find that our difficulty model shows strong generalization and that downstream adaptive compute allocation with this probe leads to significant gains over non-adaptive baselines. Interestingly, we also find that our difficulty model matches the performance of a difficulty model trained on MATH. This suggests that our difficulty model is able to capture general features that are applicable across different mathematical datasets.

2. Applying the difficulty model trained on the LMSYS dataset to the popular Anthropic HH dataset (Figure 8): LMSYS and Anthropic HH are both chat datasets but were collected in significantly different ways [1,2]. Despite this, we find that our difficulty model generalizes well and using it to route queries adaptively is significantly better than non-adaptive methods. In particular, we can achieve up to a 40% reduction in calls to the more expensive decoding scheme while maintaining similar levels of reward.

In addition to these results, we would like to highlight that the LMSYS dataset is itself very diverse and captures a large distribution of users. In particular, the LMSYS dataset is composed of real-world user conversations collected from 210K unique IP addresses [1]. Although we do agree and note in the paper that this still maybe somewhat controlled (for example, the website used for collection may have users that are primarily LLM hobbyists), we believe that the chat results on LMSYS (for routing and adaptive BoK) also demonstrate some degree of generalizability of our difficulty model.

Finally, we also present results that show generalization of our difficulty model to decoding procedures that it was not trained for (see Appendix D).

Q2. 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?

A2. Thank you for your question! Although the downstream compute allocation is independent of the base LLM, our difficulty model is learned on top of the base LLM’s representations. Our main reason for using different LLMs was to show that we can learn effective difficulty predictors across a variety of base LLM models.

We also provide new results on GSM8K with the Gemma family of models (see Appendix B) and find that the performance gains are significant and follow the trend we observed for specialized Math models. Thus, all of our chat experiments (with the exception of value-augmented search for which no Gemma model is openly available) and the new GSM8K experiment use the Gemma family of models.

Q3. The implementation of the baselines is weak, with only one effective but not particularly practical baseline (best-of-k and random) in each scenario. There are likely other reasonable approaches that could demonstrate the effectiveness of the proposed framework.

A3. Thank you for your question. We would like to emphasize that the contribution of this paper is not a particular decoding method (such as adaptive best-of-k or routing), but instead to show that adaptive test-time compute allocation can be beneficial across a diverse set of existing decoding methods. In this sense, the “baseline” methods are those that allocate a fixed amount of computation per query. Concurrent work by Snell et al.[3], also shows the value of adaptive computation and does not consider method-specific baselines.

Finally, most of the relevant work in the area presents different methods to use test-time computation (chain-of-thought, generate and revise, MCTS, etc) . In this work, we consider a different axis which tries to adaptively allocate this computation, making our method complementary to most of these test-time methods.

Please let us know if there are specific comparisons you would like to see!