Compute-optimal test-time scaling

A PRM trained on rollouts from each step serves as the verifier. A PRM trained on PRM800Kthe human step-level correctness labels released by Lightman et al. for GPT-4 solutions to MATH problems was easy to exploit even with best-of-N on PaLM 2 samples, which the authors attribute to the difference between GPT-4's solutions and PaLM 2's. Following Math-Shepherd (Wang et al.), each step of 16 sampled solutions per training question is labeled instead with the fraction of 16 continuations from that step that reach the right answer, so the PRM predicts the chance of eventual success under the base model's own sampling, a value estimate. As the score of a whole answer, the PRM's score at the last step works better than the minimum or product over steps, and the PRM beats an ORMoutcome reward model: a verifier trained to score only complete answers, by their final correctness by a margin that grows with $N$; read at its last step the PRM scores whole answers, so the authors suggest the step-level training helps mainly by giving better representations. Answers are chosen by best-of-N weightedselection of the final answer whose samples have the largest total verifier score, so agreement among samples and verifier confidence both count selection: with $a_i$the final answer of sample $i$ and $v_i$the verifier's score of sample $i$, the choice is $$\hat a=\arg\max_a\ \sum_{i:\,a_i=a}v_i,$$ which is majority votingchoosing the final answer that appears most often among the samples when all scores are equal.

Stronger search against the PRM helps at small budgets and hurts at large ones. Beam search keeps $N$ beams with beam width $w$the beam width: each kept partial solution is extended with $w$ sampled next steps, so $N/w$ of the $N$ candidates survive each round: sample $N$ first steps, keep the $N/w$ the PRM scores highest, sample $w$ next steps from each, and repeat for up to 40 rounds, ending with best-of-N weighted over the $N$ final answers. Lookahead searchbeam search in which each candidate step is scored by the PRM at the end of a greedy rollout of up to $k$ further steps; beam search is the case $k=0$ is charged $N(k+1)$ samples. Swept with $w=\sqrt N$ or 4 and $k\in\{1,3\}$ at budgets $N$ up to 256, beam search is clearly ahead of best-of-N at small $N$, its lead shrinks and often reverses as $N$ grows, and lookahead, which optimizes the PRM hardest, is worst at equal cost. The search outputs show what is being exploited: long tails of repetitive, uninformative steps, and solutions of one or two steps, both of which the PRM scores highly.

Difficulty decides which search wins, and choosing by difficulty nearly matches best-of-N on a quarter of the budget. At $N\in\{4,16,64,256\}$, beam search with $w=4$ and best-of-N look alike in aggregate at the largest budget, but on bins 1 and 2 beam search gets worse as $N$ grows while best-of-N does not, on bins 3 and 4 beam search is consistently ahead, and on bin 5 neither makes meaningful progress. The authors' explanation: on easy questions the verifier already judges the base model's samples mostly correctly, so a stronger search gains little and mainly amplifies spurious features the verifier learned; on harder ones a correct answer is rarely sampled at all, and step-level guidance finds it more often. Picking the best search setting per bin and budget nearly matches best-of-N at 64 generations with 16, and the curves for predicted and oracle bins nearly coincide; at the largest budgets the gain shrinks with predicted bins and continues with oracle ones.

A fine-tuned revision model improves its answers along a chain, and it also breaks correct ones. Prompting alone does not give useful self-correction on these problems, so the model is fine-tuned on constructed chains. From 64 samples per training question, each correct answer becomes the target after zero to four incorrect answers placed in context, the last of them the incorrect answer closest to the correct one in character edit distancethe number of character insertions, deletions and substitutions needed to turn one text into another, so the target reads as a correction of what precedes it. At test time the context keeps the four most recent attempts. Pass@1 rises along the chain, even past the four revisions seen in training, but about 38% of correct answers are revised into incorrect ones, since training never showed a correct answer in context. The final answer is therefore selected over the chain by majority vote or by a verifier; the PRM of Step 1, trained on base-model samples, transferred poorly to revision outputs, so an ORM trained on revision outputs, with the earlier attempts in its context, is used. With either selector, $N$ revisions in sequence beat $N$ independent samples from the same model by a small margin, and the margin remains when the verifier does not see the history. Further training of the revision model with ReST-EMa simple RL method: sample attempts, keep the successful ones, fine-tune on them, and repeat on its own revision chains made additional revisions harmful, which the authors attribute to spurious correlations in data the model generated itself.

The best split between revisions and independent samples depends on difficulty. Split a budget $N$ into $N/s$ independent chains of $s$ revisions each, and select with the verifier first within each chain and then across chains. More sequential samples generally help, and at the higher budgets accuracy peaks at an intermediate ratio of sequential to parallel samples. At $N=128$ the peak moves with difficulty: the easiest bins do best with one chain, harder bins with a mix. The authors' reading is that revisions refine an approach that is already roughly right, while independent samples try different approaches. Choosing the ratio per bin beats parallel best-of-N at 256 generations with 64, and keeps improving at budgets where parallel sampling has flattened.

Extra inference replaces a larger model on easy questions, and on harder ones only when inference is light. With $P$the small model's parameter count, $D_{\rm pre}$the number of pretraining tokens, the same for both models and $D_{\rm inf}$the total number of tokens the deployed model generates, one greedy answer per query, pretraining costs $6PD_{\rm pre}$ FLOPs and inference $2PD_{\rm inf}$. A model with $M$the ratio of the larger model's parameter count to the small model's; about 14 here times the parameters, trained on the same data, costs $M$ times both. With the same total FLOPs the small model can run its inference $K$ times over, where $6PD_{\rm pre}+K\cdot2PD_{\rm inf}=M(6PD_{\rm pre}+2PD_{\rm inf})$. Dividing by $2PD_{\rm inf}$ and writing $R$the inference load $D_{\rm inf}/D_{\rm pre}$: tokens generated in deployment per pretraining token for the inference load $D_{\rm inf}/D_{\rm pre}$ gives $$K=\underbrace{M}_{\text{inference}}+\underbrace{3(M-1)/R}_{\text{extra pretraining}}.$$ For $M\approx14$ and the paper's $R=0.16,\ 0.79,\ 22$, $K$ is about 258, 63 and 16 (my computation), budgets of roughly $2^8,2^6,2^4$ generations. Against the larger model's greedy accuracy, compute-optimal revisions (with oracle bins) are 22, 28 and 12% better (the labels of the paper's Figure 1), relative, on the easy, medium and hard groups at $R=0.16$; at $R=22$ they are 5% better on easy questions and 24% and 37% worse on medium and hard ones. PRM search does worse than revisions in every group and at every load, already about 35% behind on medium and hard questions at $R=0.79$. The larger model's advantage is concentrated where the small model rarely produces a correct answer by sampling or by revising.