DIPA: Difficulty-Informed Probabilistic Allocation of Test-Time Compute via Training-Free Proxies
Abstract
Large language models (LLMs) excel at complex tasks but incur prohibitive computational costs, particularly when using techniques like self-consistency that require multiple generation attempts. This paper addresses the challenge of input-adaptive test-time compute allocation. We propose a framework that leverages training-free difficulty proxies derived directly from LLMs to distribute a fixed compute budget across the test queries, without requiring specialized training for the allocation mechanism. Our objective is to maximize the number of solved instances by dynamically allocating more compute to difficult instances and less to simpler ones, while adhering to a total budget constraint. We first introduce several training-free proxies and empirically demonstrate their effectiveness in estimating instance difficulty. We then design an adaptive allocation strategy guided by these proxies, which is theoretically grounded in a novel bandit formulation. Experiments across math, coding, and Q&A benchmarks demonstrate that our method significantly outperforms both uniform budget allocation and training-based allocation baselines. This work presents a practical and readily deployable approach to enhance the resource efficiency of LLM inference for demanding reasoning tasks.