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

arXiv: https://arxiv.org/abs/2604.14853 · HTML: https://mlsanity.com/paper/2604.14853/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 42 of its 50 references that are in the map.
- First posted on arXiv 16 Apr 2026.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)

## Builds on (most cited first, 30 of 42)

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [Large Language Models are Zero-Shot Reasoners](https://mlsanity.com/paper/2205.11916/) (2022)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Self-Refine: Iterative Refinement with Self-Feedback](https://mlsanity.com/paper/2303.17651/) (2023)
- [Tree of Thoughts: Deliberate Problem Solving with Large Language Models](https://mlsanity.com/paper/2305.10601/) (2023)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [Fast Inference from Transformers via Speculative Decoding](https://mlsanity.com/paper/2211.17192/) (2022)
- [Qwen2.5-Math Technical Report: Toward Mathematical Expert Model via Self-Improvement](https://mlsanity.com/paper/2409.12122/) (2024)
- [Large Language Monkeys: Scaling Inference Compute with Repeated Sampling](https://mlsanity.com/paper/2407.21787/) (2024)
- [Solving math word problems with process- and outcome-based feedback](https://mlsanity.com/paper/2211.14275/) (2022)
- [Reasoning with Language Model is Planning with World Model](https://mlsanity.com/paper/2305.14992/) (2023)
- [FrugalGPT: How to Use Large Language Models While Reducing Cost and Improving Performance](https://mlsanity.com/paper/2305.05176/) (2023)
- [Selective Classification for Deep Neural Networks](https://mlsanity.com/paper/1705.08500/) (2017)
- [RouteLLM: Learning to Route LLMs with Preference Data](https://mlsanity.com/paper/2406.18665/) (2024)
- [BranchyNet: Fast Inference via Early Exiting from Deep Neural Networks](https://mlsanity.com/paper/1709.01686/) (2017)
- [Adaptive Computation Time for Recurrent Neural Networks](https://mlsanity.com/paper/1603.08983/) (2016)
- [Dynamic Neural Networks: A Survey](https://mlsanity.com/paper/2102.04906/) (2021)
- [Hybrid LLM: Cost-Efficient and Quality-Aware Query Routing](https://mlsanity.com/paper/2404.14618/) (2024)
- [Inference Scaling Laws: An Empirical Analysis of Compute-Optimal Inference for Problem-Solving with Language Models](https://mlsanity.com/paper/2408.00724/) (2024)
- [RouterBench: A Benchmark for Multi-LLM Routing System](https://mlsanity.com/paper/2403.12031/) (2024)

## Built on by (most cited first, 9 of 9)

- [Adaptive Generate-Rank-Verify: Inference-Time Search with Costly Verification](https://mlsanity.com/paper/2605.17609/) (2026)
- [Think Again or Think Longer? Selective Verification for Budget-Aware Reasoning](https://mlsanity.com/paper/2606.19808/) (2026)
- [Thinking Hard, Not Smart: Reasoning Models Fail to Ration Test-Time Compute Across Questions](https://mlsanity.com/paper/2608.07968/) (2026)
- [Steering Recurrent Reasoners at Inference Time with Readout Feedback](https://mlsanity.com/paper/2608.24136/) (2026)
- [Can your AI agent be cheaper? Investigating the effects of task specifications on token spend in agentic coding tasks](https://mlsanity.com/paper/2608.25399/) (2026)
- [BudgetVerify: Budget-Tiered Verification for Financial QA](https://mlsanity.com/paper/2609.33052/) (2026)
- [Solver-Hard Is Not Model-Hard: A Hardness-Controlled Diagnostic for LLM Constraint Reasoning](https://mlsanity.com/paper/2607.17047/) (2026)
- [Sampling headroom is not selection gain: a compute-value audit of test-time scaling for video world models](https://mlsanity.com/paper/2609.13257/) (2026)
- [PAGE: Partition-Aware Gated KV-Cache Eviction](https://mlsanity.com/paper/2609.22157/) (2026)

## Abstract (arXiv)

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.