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

1. Problem Definition

The paper addresses the inefficiency of current Test-Time Compute (TTC) scaling methods for Large Language Models (LLMs).

• Current State: Existing methods (e.g., Best-of-N sampling, consistency checking) typically allocate a uniform compute budget across all queries in a dataset. They generate $N$ responses for every question regardless of difficulty.
• The Flaw: This "one-size-fits-all" approach is suboptimal. Simple queries waste resources by receiving more compute than necessary, while complex queries suffer from insufficient budget.
• The Goal: The authors aim to formulate Strategic Test-Time Compute Allocation: adaptively distributing a fixed total compute budget $B$ across a set of queries $S = \{x_1, \dots, x_n\}$ to maximize the fraction of correctly answered queries. The core challenge is doing this on-the-fly without prior knowledge of query difficulty or the specific compute required to solve each instance.

2. Methodology

The authors propose a novel framework that bridges LLM inference and Multi-Armed Bandit (MAB) learning, specifically focusing on Pure Exploration (Best Arm Identification).

A. Bandit Formulation

• Actions: Each query $x \in S$ is treated as a "bandit arm."
• Feedback: Sampling an arm (querying the LLM) generates a response $y$. A Reward Oracle (e.g., a Process Reward Model or Ground Truth verifier) scores the response.
• Objective: Unlike standard bandits that seek the single best arm, this is a thresholding problem. The goal is to find at least one correct response for every query in the set. A query is considered "solved" (and removed from the active set) if its best observed score exceeds a threshold $\gamma$.

B. Algorithmic Framework (Algorithm 1: Elimination)

The core algorithm operates in rounds with an Active Set $A$ of unsolved queries:

1. Exploration: In each round, the algorithm allocates a small compute budget $K$ to every query currently in $A$.
2. Evaluation: Responses are scored by the reward oracle.
3. Elimination:
   • If a query's best response score $\ge \gamma$, the query is marked as solved and removed from $A$.
   • The algorithm maintains the best response found so far for each query.
4. Termination: The process stops when the total budget $B$ is exhausted or all queries are solved.

C. Extensions & Variants

The framework supports flexible exploration strategies beyond uniform sampling:

• UCB (Upper Confidence Bound): Prioritizes queries with high empirical rewards plus an uncertainty bonus.
• Gap-based: Focuses on queries near the elimination threshold $\gamma$.
• Entropy-based (Novel): Selects queries with the highest response diversity (empirical entropy). This is particularly effective for distinguishing between solvable and unsolvable queries (see Results).
• Adaptations: The framework extends to Token-constrained settings (counting tokens instead of generations) and Streaming settings (queries arriving sequentially).

D. Theoretical Analysis

The authors provide theoretical guarantees comparing their adaptive approach to uniform allocation:

• Uniform Allocation: Requires a budget scaling with $O(|S| / \min \Delta_x)$, where $\Delta_x$ is the probability of a correct answer. It is inefficient because it over-samples easy queries.
• Adaptive Allocation (Theorem 1): Requires a budget scaling with $O(\sum_{x \in S} 1/\Delta_x)$.
• Result: In heterogeneous regimes (where query difficulties vary significantly), the adaptive approach achieves a strict separation in efficiency, theoretically proving it can be orders of magnitude more efficient than uniform allocation.

3. Key Contributions

1. Novel Formulation: First to frame strategic test-time compute allocation as a pure-exploration bandit problem, providing a decision-theoretic foundation for adaptive scaling.
2. Adaptive Algorithms: Developed a general framework (Elimination) with novel exploration rules (specifically Entropy-based sampling) that dynamically allocates compute based on real-time difficulty estimates.
3. Theoretical Proof: Proved that adaptive allocation achieves better compute efficiency than uniform allocation and two-stage "explore-then-commit" baselines.
4. Comprehensive Empirical Validation: Extensive testing across math and code benchmarks, demonstrating significant performance gains without additional model training.

4. Experimental Results

The authors evaluated their methods on MATH-500, AIME25, and LiveCodeBench using various models (Llama-3, Qwen3, Gemini, DeepSeek-R1).

• Performance Gains:
  • MATH-500: Up to 11.10% absolute improvement (15.04% relative) over uniform baselines.
  • AIME25: Up to 10.82% absolute improvement (14.44% relative).
  • LiveCodeBench: Up to 11.23% absolute improvement (15.29% relative).
• Efficiency: The adaptive methods achieved the same accuracy as uniform baselines using 2x to 4x less compute (e.g., achieving the same performance with 8 units of compute that uniform methods needed 32 units for).
• Handling Hard/Unsolvable Queries:
  • On the MATH-500-Hard subset (containing unsolvable queries), the Entropy-based variant outperformed others.
  • Mechanism: Unsolvable queries tended to produce low-entropy, invalid, or repetitive outputs. Solvable queries produced diverse, high-entropy outputs. The Entropy rule successfully identified and prioritized solvable instances, avoiding wasted compute on unsolvable ones.
• Robustness: The method remained effective under token-matched constraints and streaming query scenarios.

5. Significance

This paper represents a paradigm shift in how inference-time compute is managed for LLMs:

• From Static to Dynamic: It moves away from static, uniform sampling strategies toward dynamic, data-driven resource allocation.
• Cost-Effectiveness: By reducing the compute required to solve difficult problems and avoiding waste on easy ones, it offers a pathway to significantly lower inference costs for high-performance reasoning.
• Generalizability: The bandit framework is model-agnostic and can be applied to any task where a reward oracle exists (math, code, logic), making it a versatile tool for future LLM scaling.
• Theoretical Grounding: It provides the first rigorous theoretical proof that strategic allocation outperforms uniform allocation, moving the field beyond purely empirical observations.

In summary, the paper demonstrates that treating test-time compute as a sequential decision-making problem allows LLMs to "think harder" only when necessary, leading to substantial efficiency and accuracy gains.