Scaling LLM Inference with Optimized Sample Compute Allocation

TL;DR

This work addresses efficient LLM inference under tight compute by optimizing how many samples to allocate across multiple sampling configurations. OSCA formulates the problem as learning an allocation that maximizes pass@C, estimates per-configuration success probabilities from a small training budget, and solves a hill-climbing optimization over mixed configurations. Across LiveCodeBench, LiveBench, and SWE-Bench, OSCA substantially outperforms pure and uniform allocations, achieving strong accuracy with far less compute (e.g., up to 128x less for code generation and 25x less for reasoning). The results demonstrate the practical value of mixed sampling allocations for both single-turn problems and complex, agentic workflows, offering a pathway to more efficient and reliable LLM-based inference.

Abstract

Sampling is a basic operation in many inference-time algorithms of large language models (LLMs). To scale up inference efficiently with a limited compute, it is crucial to find an optimal allocation for sample compute budgets: Which sampling configurations (model, temperature, language, etc.) do we use? How many samples do we generate in each configuration? We formulate these choices as a learning problem and propose OSCA, an algorithm that Optimizes Sample Compute Allocation by finding an optimal mix of different inference configurations. Our experiments show that with our learned mixed allocation, we can achieve accuracy better than the best single configuration with 128x less compute on code generation and 25x less compute on 4 reasoning tasks. OSCA is also shown to be effective in agentic workflows beyond single-turn tasks, achieving a better accuracy on SWE-Bench with 3x less compute than the default configuration. Our code and generations are released at https://github.com/LeiLiLab/OSCA.

Figures (6)

• Figure 1: On 2 single-turn benchmarks and 1 agentic benchmark with a total of 6 tasks, our optimized allocations of sample compute are better than both optimal pure allocations and uniform allocations in most cases, especially when the compute budget is small.
• Figure 2: With 4 sampling configurations and a compute budget of 1, pure allocation spends all the budget on one configuration. Uniform allocation evenly distributes the budget across configurations. Optimized allocation learns where to spend more budget.
• Figure 3: Osca's pass rates on LiveBench when trained with different proportions of the training data.
• Figure 4: Osca with different sample compute for estimating pass rates. $C_0=50$.
• Figure 5: Osca's pass rates on LiveBench when it is banned from allocating compute to multiple temperatures or multiple models.