Strategic Scaling of Test-Time Compute: A Bandit Learning Approach

arXiv:2506.12721v2 Announce Type: replace Abstract: 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, up to 10.82% (14.44% relative) on the AIME25 dataset, and up to an 11.23% performance improvement (15.29% relative) on the LiveCodeBench dataset.

Related

• Procedural Knowledge at Scale Improves Reasoning
• Inference-Time Scaling of Diffusion Language Models via Trajectory Refinement
• PhysMem: Scaling Test-time Physical Memory for Robot Manipulation
• A Mechanistic Analysis of Looped Reasoning Language Models
• Pruning Long Chain-of-Thought of Large Reasoning Models via Small-Scale Preference Optimization

Source: arXiv cs.AI | 2026-04-24