OptScale: Probabilistic Optimality for Inference-time Scaling

TL;DR

This work introduces a probabilistic framework for inference-time scaling in LLMs, deriving a fundamental lower bound on the number of samples needed to meet a performance target under i.i.d. assumptions. It then proposes OptScale, a practical adaptive sampling algorithm with a trainable-predictor variant and a training-free variant, to minimize token usage while achieving predefined performance and confidence levels. Extensive experiments across MATH-500, GSM8K, AIME, and AMC show that OptScale reduces inference tokens significantly without sacrificing, and often improving, reasoning accuracy. The combination of theoretical guarantees and empirical validation provides a solid foundation for principled, compute-efficient inference-time scaling in diverse reasoning tasks.

Abstract

Inference-time scaling has emerged as a powerful technique for enhancing the reasoning performance of Large Language Models (LLMs). However, existing approaches often rely on heuristic strategies for parallel sampling, lacking a principled foundation. To address this gap, we propose a probabilistic framework that formalizes the optimality of inference-time scaling under the assumption that parallel samples are independently and identically distributed (i.i.d.), and where the Best-of-$N$ selection strategy follows a probability distribution that can be estimated. Within this framework, we derive a theoretical lower bound on the required number of samples to achieve a target performance level, providing the first principled guidance for compute-efficient scaling. Leveraging this insight, we develop \textsc{OptScale}, a practical algorithm that dynamically determines the optimal number of sampled responses. \textsc{OptScale} employs a language model-based predictor to estimate probabilistic prior parameters, enabling the decision of the minimal number of samples needed that satisfy predefined performance thresholds and confidence levels. Extensive experiments on representative reasoning benchmarks (including MATH-500, GSM8K, AIME, and AMC) demonstrate that \textsc{OptScale} significantly reduces sampling overhead while remaining better or on par with state-of-the-art reasoning performance. Our work offers both a theoretical foundation and a practical solution for principled inference-time scaling, addressing a critical gap in the efficient deployment of LLMs for complex reasoning.

Figures (11)

• Figure 1: Scaling efficiency comparisons (accuracy vs. average token consumption): Both OptScale$^0$ and OptScale$^t$ show consistently faster convergence and optimal accuracy-token tradeoff over the compared baseline methods.
• Figure 2: Token consumption of different methods when scaling across $N$: OptScale consistently achieves reduced completion tokens when scaling to higher $N$ over the compared baseline methods.
• Figure 3: Samples of verifier score distribution: Real vs. Estimated. OptScale accurately fits most distributions under the truncated normal distribution assumption.
• Figure 4: Sensitivity analysis of quality threshold $s_{\min}$: Model performance across target scores.
• Figure 5: Sensitivity analysis of confidence level $\alpha$: Model performance across target levels.