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

TL;DR

This paper introduces an input-adaptive method for allocating computational resources in language model decoding, improving efficiency or quality by dynamically adjusting the amount of computation based on input difficulty.

Contribution

It proposes a novel approach to predict reward distributions and allocate decoding resources adaptively, reducing computation or enhancing output quality across various tasks.

Findings

Reduced computation by up to 50% without quality loss

Improved output quality by up to 10% within fixed budgets

Effective in programming, mathematics, and dialog tasks

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 response quality, or improve quality by up to 10% at a fixed computational budget.