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

• URL: http://arxiv.org/abs/2604.14853v1
• Date: Thu, 16 Apr 2026 10:39:22 GMT
• Title: Adaptive Test-Time Compute Allocation for Reasoning LLMs via Constrained Policy Optimization
• Abstract summary: Test-time compute scaling has become a powerful lever for improving large language model performance.<n>Yet deploying these techniques under finite inference budgets requires a decision that current systems largely ignore.<n>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.
• Score: 18.737087162461563
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Cost-Aware Learning [72.31444819326795]
  We consider the problem of Cost-Aware Learning, where sampling different component functions incur different costs.<n>We propose the Cost-Aware Descent algorithm for convex functions, and derive its cost complexity to attain an error of $$.<n>We introduce Cost-Aware GRPO, an algorithm designed to reduce the cost of policy optimization while preserving performance.
  arXiv  Detail & Related papers  (2026-04-30T15:39:09Z)
• Granularity-Regulated Adaptive Computational Efficiency for Optimal Verification in Test-Time Scaling [0.0]
  Test-time scaling (TTS) has emerged as a powerful paradigm for improving the reasoning performance of large language models.<n>A central component of TTS is the emphverifier, which selects or scores candidate solutions to guide the search process.<n>We prove that there exists a phase transition: fine-grained verification dominates when either the compute budget is large or the problem is hard, whereas coarse-grained verification is preferred in the low-budget, easy-problem regime.
  arXiv  Detail & Related papers  (2026-04-28T18:19:16Z)
• $\nabla$-Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space [71.23672814629448]
  $nabla$-Reasoner is an iterative generation framework that integrates differentiable optimization over token logits into the decoding loop.<n>$nabla$-Reasoner achieves over 20% accuracy improvement on a challenging mathematical reasoning benchmark.
  arXiv  Detail & Related papers  (2026-03-05T08:42:54Z)
• ODAR: Principled Adaptive Routing for LLM Reasoning via Active Inference [60.958331943869126]
  ODAR-Expert is an adaptive routing framework that optimize the accuracy-efficiency trade-off via principled resource allocation.<n>We show strong and consistent gains, including 98.2% accuracy on MATH and 54.8% on Humanity's Last Exam.
  arXiv  Detail & Related papers  (2026-02-27T05:22:01Z)
• LLM-as-Judge on a Budget [35.393598355979385]
  We present a principled variance-adaptive approach leveraging multi-armed bandit theory and concentration inequalities.<n>Our algorithm is shown to achieve a worst-case score-estimation error of $tildeOleft(sqrtfracsum_i=1K_i2Bright)$.<n>Experiments on emphSummarize-From-Feedback and emphHelpSteer2 demonstrate that our method significantly outperforms uniform allocation.
  arXiv  Detail & Related papers  (2026-02-17T10:35:41Z)
• Predictive Scheduling for Efficient Inference-Time Reasoning in Large Language Models [6.002670452103349]
  Large language models (LLMs) achieve state-of-the-art accuracy on complex reasoning tasks.<n>But using a fixed token budget per query leads to over-computation on easy inputs and under-computation on hard ones.<n>We introduce Predictive Scheduling, a plug-and-play framework that pre-runs lightweight predictors to estimate each query's optimal reasoning length or difficulty before any full generation.
  arXiv  Detail & Related papers  (2026-02-01T13:58:23Z)
• Closing the Approximation Gap of Partial AUC Optimization: A Tale of Two Formulations [121.39938773554523]
  The Area Under the ROC Curve (AUC) is a pivotal evaluation metric in real-world scenarios with both class imbalance and decision constraints.<n>We present two simple instance-wise minimax reformulations to close the approximation gap of PAUC optimization.<n>The resulting algorithms enjoy a linear per-iteration computational complexity w.r.t. the sample size and a convergence rate of $O(-2/3)$ for typical one-way and two-way PAUCs.
  arXiv  Detail & Related papers  (2025-12-01T02:52:33Z)
• $\ exttt{SPECS}$: Faster Test-Time Scaling through Speculative Drafts [55.231201692232894]
  $textttSPECS$ is a latency-aware test-time scaling method inspired by speculative decoding.<n>Our results show that $textttSPECS$matches or surpasses beam search accuracy while reducing latency by up to $sim$19.1%.
  arXiv  Detail & Related papers  (2025-06-15T05:50:05Z)
• Fast Controlled Generation from Language Models with Adaptive Weighted Rejection Sampling [90.86991492288487]
  evaluating constraint on every token can be prohibitively expensive.<n> LCD can distort the global distribution over strings, sampling tokens based only on local information.<n>We show that our approach is superior to state-of-the-art baselines.
  arXiv  Detail & Related papers  (2025-04-07T18:30:18Z)
• Semiparametric Double Reinforcement Learning with Applications to Long-Term Causal Inference [33.14076284663493]
  Long-term causal effects must be estimated from short-term data.<n>MDPs provide a natural framework for capturing such long-term dynamics.<n>Nonparametric implementations require strong intertemporal overlap assumptions.<n>We introduce a novel plug-in estimator based on isotonic Bellman calibration.
  arXiv  Detail & Related papers  (2025-01-12T20:35:28Z)
• Minimax and Communication-Efficient Distributed Best Subset Selection with Oracle Property [0.358439716487063]
  The explosion of large-scale data has outstripped the processing capabilities of single-machine systems. Traditional approaches to distributed inference often struggle with achieving true sparsity in high-dimensional datasets. We propose a novel, two-stage, distributed best subset selection algorithm to address these issues.
  arXiv  Detail & Related papers  (2024-08-30T13:22:08Z)
• Self-Evaluation Guided Beam Search for Reasoning [61.523627290397556]
  We introduce a stepwise self-evaluation mechanism to guide and calibrate the reasoning process of Large Language Model (LLM) We propose a decoding algorithm integrating the self-evaluation guidance via beam search. Our approach surpasses the corresponding Codex-backboned baselines in few-shot accuracy by $6.34%$, $9.56%$, and $5.46%$ on the GSM8K, AQuA, and StrategyQA.
  arXiv  Detail & Related papers  (2023-05-01T02:37:59Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.