Input Adaptive Allocation of Language Model Computation
Name
damani-mehul42-sm-eecs-2025-thesis.pdf
Description
Thesis PDF
Size
2.56 MB
Format
Adobe PDF
Checksum (MD5)
72ea0e9bc6cdb7618e56506e1e7f8e9c
Author(s)
Damani, Mehul
Advisor(s)
Andreas, Jacob
Date Issued
February 2025
Publisher
Massachusetts Institute of Technology
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.
MIT Department
Massachusetts Institute of Technology. Department of Electrical Engineering and Computer Science
Terms of Use
In Copyright - Educational Use Permitted
Copyright retained by author(s)
Persistent DSpace Link
