Scaling Test-Time Compute

Before comparing methods, the paper puts them in one frame. Spending compute at test time means changing the model's output distribution for this prompt, so that sampling from it gives better answers than sampling from the raw LLM. There are only two places to make that change.

The first is the input. Add tokens the model conditions on, such as its own previous attempts, and it proposes from a different and hopefully better distribution. This is modifying the proposal distribution. Training methods like STaR and ReST$^{EM}$ also improve the proposer, but they do it by fine-tuning and add no test-time tokens. Self-critique is the test-time version: the model reads and revises its own output. Prompting off-the-shelf models to revise does not work well on reasoning, so the authors fine-tune a revision model, following the recipe of Qu et al. (Recursive Introspection).

The second is the output. Sample many candidates and use a scorer to aggregate or select among them. This is optimizing the verifier. The simplest form is best-of-N against an outcome reward model (ORM), which scores only finished answers. A process reward model (PRM) predicts the correctness of every intermediate step instead. That makes tree search possible: bad partial solutions can be pruned before they are finished.

The authors compare this to Markov chain Monte Carlo, which samples a complicated target distribution by pairing a simple proposal with a score function. The two axes are independent, and the paper studies each separately. That separation is what lets it say which knob to turn for which kind of question.

Every test-time method has knobs. With a revision model and a verifier you can spend the whole budget on $N$ independent samples, on one chain of $N$ revisions, or on something in between. With a PRM you can choose best-of-N, beam search with different widths, or lookahead search. The paper's central object is the setting of those knobs that works best for a given prompt at a given budget.

Formally, let $\text{Target}(\theta, N, q)$ be the output distribution produced for prompt $q$ with hyper-parameters $\theta$ and budget $N$. The test-time compute-optimal scaling strategy is the $\theta$ that maximizes the chance of sampling the correct answer $y^*(q)$ (Equation 1). Solving this per prompt is impossible, because you would need to know the answer to tune it.

So the paper approximates it with one statistic: question difficulty, measured relative to the base model. Following Lightman et al., each test question's pass@1 is estimated from 2048 samples and the questions are split into five quantile bins. On a validation split the best strategy is chosen per bin and per budget, then applied on held-out questions. Difficulty is treated as a sufficient statistic. Model-relative bins predicted the benefit of test-time compute better than MATH's hand-labelled difficulty levels.

Oracle difficulty needs ground-truth correctness, which deployment does not have. The practical version, model-predicted difficulty, bins the same 2048 samples by the verifier's average final-answer score instead. That costs inference compute too, and the paper is open that its experiments do not charge for it. It compares this to the exploration-exploitation trade-off in RL. To avoid choosing and scoring strategies on the same questions, each bin uses two-fold cross-validation: select on one fold, evaluate on the other, then swap and average.

Search is only as good as the verifier it searches against, so the paper first builds one that holds up under optimization. The obvious starting point, Lightman et al.'s human-labelled PRM800k, did not work for them. A PRM trained on it was easy to exploit even with plain best-of-N. The authors put this down to distribution shift: PRM800k's solutions came from GPT-4, and the PRM was being asked to judge PaLM 2's.

Rather than collect new human labels, they follow Math-Shepherd (Wang et al.) and supervise each step with Monte Carlo rollouts: from every step, continue the solution several times with the base model and record how often it reaches the right answer. The PRM's per-step prediction is therefore a value estimate, the expected reward-to-go under the base model's own sampling policy, not a human judgment of whether the step is valid. Concretely they use 16 samples per training question and 16 rollouts per step, with soft labels and a binary cross-entropy loss.

Using the PRM needs two aggregation choices. Step-wise: prior work scored a solution by the minimum or product of its step scores, but here simply taking the last step's prediction worked best (Figure 13). The authors think the soft Monte Carlo labels behave differently from binary correctness labels. Inter-answer: instead of choosing the single top-scoring solution, they use best-of-N weighted, which adds up verifier scores over all solutions that reach the same final answer and returns the answer with the largest total.

One detail stands out. With last-step scoring the PRM is effectively used like an ORM, yet it still beats an ORM trained the same way, and the gap grows with the number of samples (Figure 14). The authors suggest per-step training may be acting mainly as representation learning.

The second axis changes where answers come from. Rather than sampling $N$ independent attempts, the model writes attempts one after another, each conditioned on the ones before, so it can spot and fix its mistakes. Prompting existing LLMs to correct themselves mostly fails on reasoning, so the authors fine-tune PaLM 2-S* into a revision model, building on Qu et al.

The training data has to show the model a run of wrong answers followed by a right one, and the wrong answers should be related to the right one. Otherwise the model learns to ignore its context and start over, which is just parallel sampling with extra steps. Running multi-turn on-policy rollouts was too expensive on their infrastructure, so they build trajectories after the fact: sample 64 answers per question at a higher temperature, then pair each correct answer with 0 to 4 incorrect ones (count drawn uniformly) as context. The incorrect answer placed last is the one closest by character edit distance to the correct target. That is a rough measure of correlation, but it was enough.

At test time the model was trained on at most four previous answers, so longer chains keep only the four most recent attempts in context. Pass@1 still keeps rising past four revisions (Figure 6, left). There is a catch. Training only ever showed incorrect answers in context, but at test time a correct answer can end up there, and the model may revise it into a wrong one. This happened to about 38% of correct answers. So the final answer is chosen across the chain, by sequential majority voting or by a verifier, not taken from the last revision. The verifier is an ORM trained on the revision model's own outputs, because the PaLM 2 PRM was trained on a different distribution (Appendix J).

With selection in place, a fair comparison at equal budget, $N$ revisions in sequence against $N$ samples in parallel, comes out narrowly in favour of sequential under both majority and verifier selection (Figure 6, right).

With a working PRM, the next question is how hard to optimize against it. The paper compares three search methods, all sampling from a few-shot prompted base model that writes one step per line (Figure 2).

Best-of-N weighted samples $N$ complete solutions independently and picks an answer with the PRM's final-step score. Beam search works step by step. It samples $N$ candidate first steps, keeps the top $N/M$ by PRM score, samples $M$ continuations from each to get back to $N$, and repeats, for up to 40 rounds. Best-of-N weighted then picks among the $N$ finished candidates. Lookahead search changes how beam search scores a step. It rolls each candidate $k$ steps further at temperature 0 and uses the PRM's value at the end of that rollout. Beam search is the special case $k = 0$. Lookahead is also MCTS with the randomness taken out. The exploration in MCTS exists to learn a value function, and here the value function (the PRM) is already trained and frozen.

To compare methods at equal cost, a generation is one sampled answer. Best-of-N and beam search cost $N$. Lookahead costs $N(k+1)$. The sweep, up to a budget of 256, covers beam widths $M = \sqrt{N}$ and $M = 4$, lookahead with $k = 3$ at both widths, and $k = 1$ at $M = \sqrt{N}$.

The results in Figure 3 (left) show the problem with strong optimization. At small budgets beam search clearly beats best-of-N. As the budget grows the gap closes, and beam search often falls below best-of-N. Lookahead, the strongest optimizer, does worst. The cause is exploiting the PRM: search produces low-information repeated steps, or solutions just one or two steps long, that the verifier scores highly anyway.

Splitting by difficulty (Figure 3, right) makes the picture clearer. On easy questions (bins 1 to 2) beam search gets worse as the budget grows, because the verifier is mostly right there and more optimization only amplifies its spurious features. On medium questions (bins 3 to 4) beam search consistently beats best-of-N, because the base model rarely samples a correct answer and guidance helps it find one. On bin 5 nothing makes much progress.

Sequential revisions beat parallel sampling on average, but the two do different jobs. Parallel sampling is a global search: independent attempts can try completely different high-level approaches. Sequential revision is a local refinement: it improves an answer that is already roughly on track. If both matter, the best use of a budget is probably a mix, for example $\sqrt{N}$ chains of $\sqrt{N}$ revisions each.

The paper sweeps the sequential-to-parallel ratio at fixed budgets. With a verifier selecting answers (best within each chain, then best across chains, as in Figure 5), Figure 7 (left) shows that more sequential compute usually helps. At larger budgets, though, there is an ideal ratio in between, and accuracy drops off on either side of it.

The ideal ratio depends on the question. At a budget of 128, split by difficulty (Figure 7, right), easy questions do best with purely sequential compute. When the first attempt is close, every extra sample is better spent polishing it. On harder questions there is an interior optimum: some independent attempts are needed to find a viable approach before revising it pays off. This supports the paper's main point. Revisions (the proposer) and parallel sampling with a verifier are complementary axes, and how to split compute between them is a per-prompt decision.

Put the pieces together: bin each question by difficulty, then use whichever configuration won on that bin at the current budget. The paper does this separately for both axes, and both show the same pattern.

PRM search (Figure 4). At low budgets, compute-optimal search nearly beats PRM best-of-N while using up to 4x less test-time compute, for example 16 generations against 64. At higher budgets some of the advantage fades when difficulty is predicted, while oracle bins keep improving. The oracle and predicted curves mostly overlap, which suggests the practical verifier-based binning gives up little.

Revisions (Figure 8). Choosing the sequential-to-parallel ratio per bin beats parallel best-of-N with up to 4x less compute, for example 64 samples against 256. Parallel sampling levels off at high budgets, while the compute-optimal curve keeps rising, with either oracle or predicted bins.

Overall, in the discussion's terms, compute-optimal scaling makes test-time compute 2 to 4x more efficient. Neither the search methods nor the revision model are new here. The gain comes entirely from spending the same compute differently on different prompts. One caveat applies to the numbers: they do not include the compute needed to estimate difficulty.

The title's claim is tested here. Suppose a model was pretrained with $X$ FLOPs and will serve $Y$ FLOPs of inference. You can multiply the total budget by $M$, to $M(X+Y)$. Should the extra FLOPs go into a bigger model or into more test-time compute for the small one? The paper scales pretraining by adding parameters with the data fixed, as in the LLaMA series, and leaves Chinchilla-style joint scaling to future work.

The exchange rate uses the standard approximations: pretraining costs $X = 6ND_{\text{pretrain}}$ and inference costs $Y = 2ND_{\text{inference}}$, where $N$ here means parameter count. Multiplying parameters by $M$ multiplies both terms by $M$, since the bigger model is also more expensive to serve. To match that with the small model, its inference budget can grow by a factor of $M + 3\,(D_{\text{pretrain}}/D_{\text{inference}})(M-1)$. The answer depends on $R = D_{\text{inference}}/D_{\text{pretrain}}$. Large production deployments may have $R \gg 1$. Self-improvement pipelines, which generate relatively few tokens, have $R \ll 1$.

Figure 9 runs the comparison against a model with ~14x more parameters, at R = 0.16, 0.79 and 22. Stars mark the larger model's greedy accuracy. Where a star sits below the small model's test-time compute curve, test-time compute was the better use of the FLOPs. On easy and intermediate questions (bins 1 to 3, sometimes 4), or when inference requirements are low, test-time compute wins. On the hardest questions (bins 4 to 5), or at large $R$, pretraining wins.

Figure 1's bars give the size of the effect. With revisions at $R \ll 1$, test-time compute improves on the larger model by +21.6% / +27.8% / +11.8% (easy / medium / hard). At $R \gg 1$ the figures are +5.4% / −24.3% / −37.2%. PRM search follows the same pattern more sharply: +19.1% on easy questions at $R \ll 1$, and −52.9% on hard questions at $R \gg 1$. The conclusion is limited and still notable. Even with fairly simple methods, test-time compute can already be preferable to scaling parameters, but it does not substitute one-for-one for pretraining on questions beyond the model's reach.

The paper names its gaps clearly, and several of them became research directions.

Hard problems stay hard. Across search and revisions alike, the hardest difficulty bin barely improved. When the base model almost never samples a correct answer, neither a verifier nor a revision chain has much to work with. The authors call for new ways of using test-time compute that get around this. Until then, the FLOPs-matched results say these are the questions where pretraining still pays.

Difficulty is not free. The compute-optimal policy needs a difficulty estimate, which here came from 2048 samples per question scored by the verifier. That is a lot of test-time compute in its own right, and the reported curves do not charge for it. The authors suggest models that predict difficulty directly, or policies that switch between assessing and solving as they go.

The axes were not fully combined. Revisions were paired with an ORM verifier, but PRM tree search was never combined with revisions, and techniques such as critique-and-revise were not studied. Capabilities were fine-tuned in. Both the PRM and the revision model needed task-specific fine-tuning on MATH. The authors expect future models to be pretrained with these abilities. Pretraining was scaled one way only. Parameters grew and data stayed fixed, so the comparison with compute-optimal pretraining remains open.

The longest-term idea is to close the loop: distil the outputs of test-time compute back into the base model, so that thinking longer at inference becomes a way to train. The paper ends with a prediction: in future, fewer FLOPs will go into pretraining and more into inference.