Learning How Hard to Think: Input-Adaptive Allocation of LM Computation (2410.04707v1)

Abstract: Computationally intensive decoding procedures--including search, reranking, and self-critique--can improve the quality of LLM (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 response quality, or improve quality by up to 10% at a fixed computational budget.

• The paper introduces a difficulty prediction model that estimates the marginal quality improvement from added compute.
• The paper proposes a dynamic allocation algorithm that maximizes rewards under fixed computational budgets.
• The paper demonstrates up to 50% computation savings and 10% quality gains across math, programming, and dialog tasks.

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

The paper addresses the challenge of computational efficiency in LLMs (LMs) by proposing an approach for dynamically adjusting the computational allocation based on the difficulty of each input query. This adaptive approach aims to optimize resource usage by providing a mechanism to allocate more computation to inputs predicted to benefit from it, thereby enhancing performance without unnecessary computational expenditure.

Key Contributions and Methodology

The paper introduces a method that hinges on predicting the potential benefit of additional computational resources for a given input. The primary contributions can be outlined as follows:

1. Difficulty Prediction Model: The authors train a lightweight model capable of estimating the marginal improvement in output quality that an increased computational budget would provide. This model leverages pre-trained LM hidden representations without significant computational overhead at inference time.
2. Computation Allocation Algorithm: With predicted difficulty levels, the paper outlines an efficient algorithm to allocate computational resources dynamically. The algorithm operates by maximizing the expected reward across inputs while adhering to an overall computational budget constraint.
3. Two-fold Applications: The proposed framework is evaluated across two adaptive decoding procedures—best-of-$k$ sampling and routing. The best-of-$k$ approach adjusts the number of samples generated for reranking, while routing selectively uses either an expensive, accurate decoding procedure or a less costly alternative.

Experimental Results

The experiments demonstrate the efficacy of this adaptive strategy across diverse domains, including programming, mathematics, and dialog tasks:

• Computational Savings: The adaptive approach significantly reduces the computational load, showing up to 50% computation reduction without quality loss in many cases. Particularly in math and code, the adaptive method matched or exceeded reward outcomes with fewer resources than non-adaptive baselines.
• Quality Improvements: In scenarios with fixed computational budgets, the method improved response quality by up to 10%, showcasing its potential for enhancing LM output without increasing resource expenditure.
• Predictive Accuracy: The marginal reward prediction model proved to be effective, indicating that pre-trained LMs inherently capture aspects of problem difficulty, which can be leveraged for computational efficiency.

Theoretical and Practical Implications

The approach holds both theoretical appeal and practical utility. Theoretically, it advances the understanding of how models encode problem difficulty implicitly and how this can be exploited to make models more efficient. Practically, it offers an implementable framework to optimize LM computations, which is crucial in resource-constrained settings or applications dealing with large-scale user interactions.

Future Directions

The research opens avenues for further exploration in adaptive computation scaling in AI. Enhancing the difficulty prediction models could yield even better resource allocations. Additionally, exploring alternative predictor architectures or integrating this approach with a broader range of decoding procedures could extend its applicability.

In summary, this work presents a robust framework for adaptively managing LM computation, providing a scalable solution for improving efficiency and output quality in AI applications.