Scaling LLM Inference Efficiently with Optimized Sample Compute Allocation (2025)

Abstract: Sampling is a basic operation for large language models (LLMs). In reinforcement learning rollouts and meta generation algorithms such as Best-of-N, it is essential to sample correct trajectories within a given compute budget. To find an optimal allocation for sample compute budgets, several choices need to be made:Which sampling configurations (model, temperature, language, etc.) to use?How many samples to 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.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.

## introduction large language models (llms) solve more problems with more inference compute. different ways of scaling up llm inference include sampling @xcite , self-consistency @xcite , tree search @xcite , etc. among these, sampling is the most basic and serves as an atomic operation needed in all other more complicated methods. sampling is also essential for collecting high-quality rollouts in reinforcement learning. therefore, it is crucial to do it well. previous studies @xcite have investigated how to find the optimal sampling configuration, such as the best temperature. while * equal contribution. correspondence to kexun@cmu.edu. figure 1 : on 4 benchmarks with a total of 7 tasks, our optimized allocations are better than both optimal pure allocations and uniform allocations in most cases. as shown in figure 1 , just by uniformly allocating the compute budget over all possible inference configurations (dubbed "uniform mixed"), llms' accuracy can already surpass the optimal pure allocation on livecodebench. uniform mixed allocation gets 64% accuracy with 4 samples, while optimal pure needs 16x more samples to get a similar accuracy. since uniform allocation is only one of the exponentially many allocations in the search space, it is natural to ask: how do we find the optimal sample budget allocation for llm inference? we formulate this as a learning problem: given a set of different inference configurations, a training problem set, and a compute budget, we need to distribute the budget to maximize the expected accuracy. to solve this problem, we use a hill-climbing algorithm and demonstrate its effectiveness with theoretical justification and experiments. as shown in figure 1 , the learned allocation from osca is significantly better than the optimal pure allocation, especially when the sample size is small. to provide further insights for future adopters of osca, we conduct ablation studies on the effectiveness of different hyperparameters in the mix, the number of problems needed in the training set, and scenarios where mixed allocations do not offer significant improvements. moreover, we show on swe-bench, a benchmark for llm agents, that replacing pure sampling with osca's learned allocation in just one step boosts the entire workflow's performance. this demonstrates that a mixed sampling strategy not only improves single-turn tasks but also enhances agentic workflows, leading to better reasoning and decision-making. our contributions are the following: ## related work inference time algorithms. following the taxonomy of @xcite on inference-time algorithms, chained meta-generators run multiple llm calls sequentially and use the output sample from each call as the input to the next one @xcite . parallel meta-generators samples multiple candidates for a problem and selects the best candidate @xcite @xcite @xcite . step-level search methods regards problem-solving as a multistep process and sample candidate next steps at each intermediate state, using algorithms like tree search @xcite , graph search, and monte-carlo tree search @xcite @xcite . refinement-based methods samples candidate solutions sequentially, relying on some feedback to revise the next candidate @xcite . although these algorithms scale up inference differently, they all need llm sampling as a basic operation. ## evaluation on single-turn tasks we first evaluate osca on single-turn tasks, where a solution to a problem can be generated with a single llm call. ## evaluation on agentic tasks optimizing sample compute allocation is not just useful for its own purpose, it is also useful for more complicated llm inference algorithms, such as tree search and agentic workflows, because sampling is a basic operation in these. to demonstrate such usefulness, we apply osca to an agentic workflow for swe-bench, which needs 2 stages and 7 steps involving numerous llm calls to resolve software engineering issues in repositorylevel code. ## conclusion in conclusion, this paper presents osca, an algorithm designed to optimize sample compute allocation for large language models (llms) during inference. through various experiments on both single-turn and agentic tasks, the study demonstrates that osca significantly improves accuracy with reduced compute resources compared to traditional methods, such as pure or uniform mixed allocations. by leveraging a mixed allocation allocation, osca balances different inference configurations, proving particularly effective in code generation and reasoning tasks, as seen in benchmarks like livecodebench and livebench. this work demonstrates the importance of adapting sampling allocation to the specific characteristics of the problem at hand. furthermore, osca's application to more complex workflows, such as multi-step agentic tasks, shows potential for broader utility in improving the efficiency of llms in various realworld applications. however, future work could explore optimizing additional hyperparameters and testing the scalability of the method with larger compute budgets.