Strategic Scaling of Test-Time Compute: A Bandit Learning Approach
This addresses the problem of compute waste for users of large language models by adaptively allocating resources based on query difficulty, though it is an incremental improvement over existing scaling methods.
The paper tackles the inefficiency of uniformly allocating test-time compute across all queries in large language models by formulating it as a bandit learning problem, resulting in up to 11.10% performance improvement on MATH-500 and 7.41% on LiveCodeBench.
Scaling test-time compute has emerged as an effective strategy for improving the performance of large language models. However, existing methods typically allocate compute uniformly across all queries, overlooking variation in query difficulty. To address this inefficiency, we formulate test-time compute allocation as a novel bandit learning problem and propose adaptive algorithms that estimate query difficulty on the fly and allocate compute accordingly. Compared to uniform allocation, our algorithms allocate more compute to challenging queries while maintaining accuracy on easier ones. Among challenging queries, our algorithms further learn to prioritize solvable instances, effectively reducing excessive computing on unsolvable queries. We theoretically prove that our algorithms achieve better compute efficiency than uniform allocation and empirically validate their effectiveness on math and code benchmarks. Specifically, our algorithms achieve up to an 11.10% performance improvement (15.04% relative) on the MATH-500 dataset and up to a 7.41% performance improvement (14.40% relative) on LiveCodeBench.