Adaptive Test-Time Compute Allocation for Reasoning LLMs via Constrained Policy Optimization

TL;DR

This paper introduces a method for dynamically allocating compute resources during inference of large language models, optimizing accuracy within a fixed compute budget through a two-stage solve-then-learn approach.

Contribution

It formalizes the compute allocation as a constrained optimization problem and develops a practical, real-time policy learned via supervised classification to improve inference efficiency.

Findings

Achieves up to 12.8% relative accuracy improvement on MATH dataset.

Closely tracks the Lagrangian oracle with over 91% imitation accuracy.

Outperforms uniform and heuristic baselines across multiple models.

Abstract

Test-time compute scaling, the practice of spending extra computation during inference via repeated sampling, search, or extended reasoning, has become a powerful lever for improving large language model performance. Yet deploying these techniques under finite inference budgets requires a decision that current systems largely ignore: which inputs deserve more compute, and which can be answered cheaply? We formalize this as a constrained optimization problem (maximize expected accuracy subject to an average compute budget) and solve it with a two-stage Solve-then-Learn pipeline. In the solve stage, Lagrangian relaxation decomposes the global constraint into per-instance sub-problems, each admitting a closed-form oracle action that optimally prices accuracy against cost. We prove that the induced cost is monotone in the dual variable, enabling exact budget targeting via binary search. In the learn stage, a lightweight classifier is trained to predict oracle actions from cheap input features, amortizing the allocation rule for real-time deployment. We establish that the task-level regret of the learned policy is bounded by its imitation error times the worst-case per-instance gap, yielding a clean reduction from constrained inference to supervised classification. Experiments on MATH and GSM8K with three LLMs (DeepSeek-V3, GPT-4o-mini, Qwen2.5-7B) show that our method consistently outperforms uniform and heuristic allocation baselines, achieving up to 12.8% relative accuracy improvement on MATH under matched budget constraints, while closely tracking the Lagrangian oracle upper bound with over 91% imitation accuracy.