OptScale: Probabilistic Optimality for Inference-time Scaling

TL;DR

This paper introduces OptScale, a probabilistic framework and algorithm for principled inference-time scaling of LLMs, optimizing sampling efficiency while maintaining reasoning performance.

Contribution

It provides the first theoretical lower bound for sample requirements and develops OptScale, a practical method for dynamic, optimal sampling in LLM inference.

Findings

OptScale reduces sampling overhead significantly.

It maintains or improves reasoning performance compared to state-of-the-art methods.

Experimental results on reasoning benchmarks validate the effectiveness of OptScale.

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.