Source-linked AI summary
Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters
Charlie Snell, Jaehoon Lee, Kelvin Xu, Aviral Kumar
TL;DR
The paper asks how effectively LLMs can use additional inference-time computation on challenging prompts, addressing limited and conflicting evidence about test-time scaling. It analyzes adaptive proposal updates and verifier-based search, finding that difficulty-dependent compute allocation improves efficiency and can outperform larger models in FLOPs-matched comparisons.
Problem
The paper asks how much a language model can improve a challenging response when given a fixed, non-trivial inference-time compute budget.
Method
The study analyzes adaptive proposal-distribution updates and search against dense process-based verifier reward models, using prompt difficulty to allocate compute.
Results
Compute-optimal scaling improves test-time compute efficiency by 2–4× and can outperform a 14× larger model in FLOPs-matched settings on many easier and intermediate problems.
Takeaways & Limitations
Test-time computation can be preferable to additional pretraining for many problems, but the hardest questions benefit more from pretraining.
Takeaways & Limitations
Model-predicted difficulty adds inference computation, and the experiments largely do not account for that cost.
Abstract
from arXiv · showhide
Enabling LLMs to improve their outputs by using more test-time computation is a critical step towards building generally self-improving agents that can operate on open-ended natural language. In this paper, we study the scaling of inference-time computation in LLMs, with a focus on answering the question: if an LLM is allowed to use a fixed but non-trivial amount of inference-time compute, how much can it improve its performance on a challenging prompt? Answering this question has implications not only on the achievable performance of LLMs, but also on the future of LLM pretraining and how one should tradeoff inference-time and pre-training compute. Despite its importance, little research attempted to understand the scaling behaviors of various test-time inference methods. Moreover, current work largely provides negative results for a number of these strategies. In this work, we analyze two primary mechanisms to scale test-time computation: (1) searching against dense, process-based verifier reward models; and (2) updating the model's distribution over a response adaptively, given the prompt at test time. We find that in both cases, the effectiveness of different approaches to scaling test-time compute critically varies depending on the difficulty of the prompt. This observation motivates applying a "compute-optimal" scaling strategy, which acts to most effectively allocate test-time compute adaptively per prompt. Using this compute-optimal strategy, we can improve the efficiency of test-time compute scaling by more than 4x compared to a best-of-N baseline. Additionally, in a FLOPs-matched evaluation, we find that on problems where a smaller base model attains somewhat non-trivial success rates, test-time compute can be used to outperform a 14x larger model.
1. Introduction
The paper studies how LLMs can use additional test-time computation to improve challenging responses and whether inference compute can substitute for larger models or more pretraining. It finds that the best strategy depends on prompt difficulty, motivating adaptive compute-optimal scaling.
- Motivation: The paper asks how LLMs can most effectively use a fixed test-time compute budget to improve accuracy on challenging prompts.The question matters for self-improving agents and for trading off inference-time computation against pretraining and model size.
- Motivation: Prior research reports mixed results on inference-time computation, especially for complex mathematical reasoning, motivating systematic comparison of scaling methods.Some studies find improvements, while others report limited effectiveness on complex tasks.
- Main results: Compute-optimal scaling adaptively allocates test-time computation per prompt and can improve scaling efficiency by more than 4× over best-of-N.The strategy is designed to select the most effective approach for the problem under a given compute budget.
- Approaches: Test-time strategies include best-of-N sampling, sequential revision, and search using process-based reward models.These methods respectively sample and select candidates, modify the proposal distribution, or verify individual solution steps.
- Core finding: The effectiveness of each strategy varies with prompt difficulty and base-model capability: revisions suit easier problems, whereas parallel sampling or tree search suits harder problems.Easier questions may need refinement of an initially promising answer; harder questions may require exploring different high-level solution approaches.
- Main results: In FLOPs-matched evaluations, additional test-time computation can outperform a 14× larger model on problems where the smaller model has non-trivial success rates.For the most challenging questions, however, additional pretraining is more effective and test-time compute is not fully exchangeable with pretraining.
2. A Unified Perspective on Test-Time Computation: Proposer and Verifier
The paper unifies test-time computation as adaptive modification of a model’s response distribution through proposal changes or verifier-based output modification. It analyzes these as complementary axes involving proposal distributions and verifiers.
- Unified perspective: The framework views test-time computation as adaptively modifying the model’s predicted output distribution conditioned on a prompt.The goal is to generate better outputs than naïve sampling from the language model.
- Proposal distributions: Proposal distributions can be modified at the input level by adding conditioning tokens or at the output level by operating on sampled candidates.These mechanisms alter either what the model conditions on or how its generated responses are transformed.
- Unified perspective: The study treats modifying the proposal distribution and using a verifier as independent axes for analyzing test-time computation.The abstraction supports comparisons among revisions, sampling, and process-based search methods.
- Verifiers: Verifier-based methods use post-hoc scorers to aggregate or select candidates produced by the proposal distribution.Best-of-N samples complete solutions and selects the highest-scoring answer according to a verifier.
3. How to Scale Test-Time Computation Optimally
The paper defines compute-optimal scaling as prompt-dependent selection of test-time hyperparameters and approximates it using model-specific difficulty. This enables different allocations for revision and search strategies while exposing practical inference-cost limitations.
- Problem formulation: The central problem is choosing the most effective use of a fixed test-time compute budget for each prompt and comparing it with a larger pretrained model.Different strategies may have different effectiveness depending on the problem.
- Strategy allocation: Easier problems may benefit more from sequential revisions, while harder problems may benefit from independent parallel sampling or more sophisticated verifier search.The allocation reflects whether the initial answer is likely to be on the right track or whether broad exploration is needed.
- Definition: Compute-optimal scaling selects the test-time hyperparameters that maximize the accuracy of the target output distribution for a given prompt and compute budget.The target distribution is induced by the model under a chosen strategy and its hyperparameters.
- Approximation: The approximation defines strategy choices as a function of a prompt-difficulty statistic and selects the best-performing strategy independently within each difficulty bin.Question difficulty is used as a sufficient statistic when designing the compute-optimal strategy.
- Difficulty estimation: Difficulty is estimated from the base LLM’s pass@1 rate over 2048 samples, then divided into five quantiles with increasing difficulty.This model-specific difficulty measure was more predictive of test-time-compute efficacy than hand-labeled MATH difficulty bins.
- Difficulty estimation: Model-predicted difficulty replaces ground-truth correctness with averaged learned-verifier scores, making deployment feasible without answer labels.Oracle difficulty uses ground-truth correctness, whereas predicted difficulty uses verifier-based estimates from the same 2048 samples.
- Limitation: Difficulty estimation adds inference computation, and the experiments largely exclude this cost while leaving its exploration–exploitation tradeoff for future work.Deployment must balance computation spent assessing difficulty against computation spent applying the selected strategy.
4. Experimental Setup
The experiments evaluate test-time scaling on MATH using PaLM 2-S* and capability-specific fine-tuning for revision and verification. The setup targets problems requiring complex inference across several difficulty levels.
- Setup: The experimental setup analyzes multiple verifier designs and proposal distributions before presenting scaling results.The setup is organized around representative mechanisms for using test-time computation.
- Dataset: Experiments use the MATH benchmark, consisting of high-school competition-level problems spanning a range of difficulty levels.The dataset split contains 12k training questions and 500 test questions.
- Models: The analysis uses PaLM 2-S* (Codey) as the base model because it has non-trivial but unsaturated MATH performance.The authors view it as a test bed representative of many contemporary LLM capabilities.
5. Scaling Test-Time Compute via Verifiers
The paper evaluates PRM-based search methods and finds that their effectiveness depends on compute budget and question difficulty. This motivates compute-optimal allocation, which can nearly match best-of-N with substantially less test-time compute.
- Search methods: The study analyzes best-of-N, beam search, and lookahead search for scaling test-time computation with process reward models.Best-of-N samples complete answers; beam search expands and retains high-scoring prefixes; lookahead search extends beam search with k-step rollouts.
- Search methods: Lookahead-search costs N × (k + 1) sampled answers because it performs k additional rollout steps at each search step.The cost protocol counts one sampled answer as one generation and accounts for lookahead’s extra computation.
- Search results: At low generation budgets, beam search outperforms best-of-N, but its gains diminish at higher budgets and lookahead-search generally underperforms at matched budgets.The authors attribute diminishing returns partly to exploitation of PRM predictions, including repetitive or overly short solutions.
- Search results: Search behavior varies by difficulty: beam search degrades on easy questions as budget increases but consistently outperforms best-of-N on harder questions.On the most difficult questions, neither method makes much meaningful progress.
6. Refining the Proposal Distribution
The paper refines test-time proposal distributions by training models to revise answers sequentially, then studies how to combine sequential revisions with parallel sampling. The best allocation depends on question difficulty and can substantially reduce the compute needed to outperform parallel best-of-N.
- Iteratively Revising Answers at Test-time: Revision models sequentially condition on previous answers to learn from mistakes and dynamically improve their proposal distribution at test time.The models are finetuned on trajectories of incorrect answers followed by a correct answer, with revision data constructed from sampled responses.
- Iteratively Revising Answers at Test-time: Pass@1 gradually improves across revision steps, including beyond the four previous answers used during training.Longer revision chains are sampled by retaining the four most recent revised responses in context.
- Iteratively Revising Answers at Test-time: 38% of correct answers are converted back to incorrect ones by a naïve revision approach because inference can include correct answers in contexts trained to contain only incorrect answers.Sequential majority voting or verifier-based selection is therefore used to select the most correct answer from a revision chain.
- Iteratively Revising Answers at Test-time: Sequential revision sampling narrowly outperforms parallel sampling when selecting answers with either verifier-based or majority-based mechanisms.The comparison matches the number of sequential revisions against the number of parallel attempts.
- Compute Optimal Revisions: At a fixed generation budget, an ideal sequential-to-parallel ratio maximizes accuracy, with the best ratio varying by question difficulty.Easy questions benefit more from sequential revisions, whereas difficult questions perform best with a balance between sequential and parallel computation.
- Compute Optimal Revisions: Up to 4x less test-time compute can outperform the parallel best-of-N baseline when compute is allocated according to question difficulty.This result holds for both oracle and predicted difficulty bins; at higher budgets, parallel sampling plateaus while compute-optimal scaling continues improving.
7. Putting it Together: Exchanging Pretraining and Test-Time Compute
The paper compares allocating a fixed total FLOPs budget to larger model parameters versus adaptive test-time computation. The better allocation depends on question difficulty and the ratio of inference to pretraining tokens.
- The section asks whether additional FLOPs should fund increased pretraining compute or additional test-time computation.
- The comparison scales model parameters while holding training-data quantity fixed, following the approach used for the open-source LLaMA series.
- Pretraining FLOPs are approximated as X = 6ND_pretrain, while inference FLOPs are approximated as Y = 2ND_inference.
- The exchange depends on R, the ratio of inference to pretraining tokens, because this ratio determines how much inference computation can be applied.
- On harder questions or with higher inference requirements, additional pretraining is generally more effective than test-time computation.
- On easy or intermediate questions and with lower inference requirements, compute-optimal test-time computation can outperform scaling parameters by approximately 14x.
8. Discussion and Future Work
The discussion reports that compute-optimal test-time scaling improves efficiency but remains limited by the studied methods, difficulty estimation cost, and the lack of integration with training.
- The study analyzes verifier search and proposal-distribution refinement, finding that effectiveness strongly correlates with difficulty relative to the base LLM’s capabilities.
- An adaptive compute-optimal strategy improves test-time compute-scaling efficiency by 2−4× and can outperform spending equivalent FLOPs on pretraining for some prompts.
- The study does not combine revisions with PRM tree-search techniques or examine other methods such as critique and revise.
- The studied schemes provide small gains on hard problems, motivating methods that can circumvent this limitation.
- Estimating question difficulty requires a non-trivial amount of test-time computation, creating a practical efficiency constraint.
- The work does not interleave test-time computation with training, although future work could distill test-time outputs back into the base LLM.
Appendices
The appendices situate the study among work on language-model reasoning, training–inference tradeoffs, and methods that augment LLMs with test-time computation.
- Recent mathematical-reasoning gains have come from continued pretraining, proposal-distribution improvements, and other training or inference interventions.
- Earlier work studied train-time versus test-time compute using Monte-Carlo tree search for Hex, whereas this paper analyzes full-scale language-model math reasoning.
- Related approaches augment language models with hierarchical hypothesis search or tools at test time to improve downstream-task performance.
B. Additional Revision Results
Additional revision experiments show that majority selection follows trends similar to verifier selection, while predicted difficulty bins reproduce oracle-bin patterns despite costly estimation.
- Majority selection exhibits largely similar trends to verifier selection in the additional revision results.
- Difficulty bins are estimated without oracle correctness by averaging PRM final-answer scores over 2048 samples for each question.
- The resulting predicted-difficulty bins divide test questions into five quintiles using the oracle-bin procedure.
- The predicted-difficulty procedure is extremely costly in practice, and the study does not explore cheaper direct correctness-prediction methods.
- Predicted difficulty produces trends similar to oracle difficulty in both PRM-search and revision settings.
D. PRM Training Details
The paper trains process reward models to score solution steps and compares score aggregation, verifier selection, and revision-model training procedures for test-time inference.
- PRM training: The PRM predicts a value between 0 and 1 at each solution step and is trained as a binary classifier with soft Monte Carlo targets.Training uses binary cross entropy on soft ground-truth values estimated from Monte Carlo rollouts.
- PRM training: The PRM uses 16 samples per question and 16 Monte Carlo rollouts per step, while filtering samples without valid parsable final answers.The same base model and prompt generate the training samples and rollouts.
- Difficulty estimation: Difficulty bins computed without ground-truth correctness show performance trends largely similar to bins based on ground-truth correctness for revision and PRM search.The comparison appears in Figures 11 and 12 for verifier, majority, and PRM-search settings.
- Score aggregation: The last-step PRM score outperforms minimum-score and product aggregation for assigning a final score to the full solution.The compared aggregators are min, prod, and last; prior work reported minimum aggregation as best.
- PRM versus ORM: Using the last-step score makes the PRM functionally similar to an ORM, yet the PRM outperforms the ORM, suggesting value from per-step training as representation learning.The authors propose that per-step supervision may help beyond direct inference-time scoring.
- Revision-model training: Revision-model training uses trajectories of incorrect answers followed by a correct answer, with incorrect examples selected to correlate somewhat with the target correction.The base model is fine-tuned on correct solutions in these generated trajectories.
J. Revision Model Verifier Training
The revision model requires a verifier matched to its output distribution, and revision history provides only a slight additional benefit in verifier scoring.
- Verifier matching: The PRM trained on base-model outputs is less effective on revision-model outputs, likely because of distribution shift.The authors therefore train a separate ORM for the revision model instead of generating costly per-step PRM labels.
- Verifier context: The revision-setting ORM receives previous revision attempts in context so it can score the current answer with the same context as the revision model.This modifies the standard ORM while keeping other training details identical to PRM training.
- Verifier context: Including revision history improves verifier performance slightly, while sequential revisions still slightly outperform parallel sampling without revision history in context.This indicates the sequential-sampling advantage is not solely attributable to verifier access to prior attempts.
K. ReSTEM Revision Model Experiments
The ReSTEM-optimized revision model degrades when more computation is allocated to sequential revisions, unlike the earlier revision setup, and the appendix illustrates revision and beam-search behaviors.
- ReSTEM training: ReSTEM training generates 64 revision trajectories of maximum length 5, stops at the first correct answer, and balances trajectory lengths during fine-tuning.The resulting data consists of correct-answer solutions from MATH training questions.
- ReSTEM results: Additional sequential revisions substantially hurt performance in the ReSTEM-optimized revision model as the sequential-to-parallel ratio increases.Performance is evaluated with majority voting while varying that ratio.
- ReSTEM results: The authors hypothesize that online ReSTEM data amplifies spurious correlations in revision data and prevents the optimized model from learning the revision task effectively.They suggest more offline data collection as a possible alternative for future work.
- Qualitative examples: The appendix includes examples of revision trajectories in which later attempts correct arithmetic, formatting, conversion, or coordinate errors.Examples span Figures 17–23, including corrections across multiple attempts.
- Qualitative examples: Additional appendix examples show PRM beam search trajectories with per-step scores reported for each solution step.These examples appear in Figures 24–29.