PaperlayerView exact revision on arXiv

Source-linked AI summary

Can 1B LLM Surpass 405B LLM? Rethinking Compute-Optimal Test-Time Scaling

Runze Liu, Junqi Gao, Jian Zhao, Kaiyan Zhang, Xiu Li, Biqing Qi, Wanli Ouyang, Bowen Zhou

arXiv:2502.06703v1cs.CL

TL;DR

Studies had not systematically examined how policy models, Process Reward Models (PRMs), and problem difficulty shape optimal test-time scaling. This paper evaluates reward-aware compute-optimal scaling across these factors and reports that a 3B Llama model can outperform a 405B Llama model on MATH-500 and AIME24.

  • Problem

    The optimal allocation of test-time computation remains a challenge, and the effects of policy models, PRMs, and problem difficulty have not been systematically analyzed.

    Source: §1 · p. 2

  • Method

    The paper experimentally analyzes reward-aware compute-optimal test-time scaling across policy models, PRMs, scaling methods, and problem difficulty on MATH-500 and AIME24.

    Source: §1 · p. 2

  • Results

    3B Llama-3.2-3B-Instruct with compute-optimal TTS outperforms 405B Llama-3.1-405B-Instruct on MATH-500 and AIME24.

    Source: §5.1 · p. 11

  • Takeaways & Limitations

    Compute-optimal TTS is a promising approach for enhancing LLM reasoning capabilities on mathematical tasks.

    Source: §5.2 · p. 12

    The evaluation covers mathematical tasks; the authors identify extending TTS to coding and chemistry as future work.

    Source: §7 · p. 15

Abstract

from arXiv · show

Test-Time Scaling (TTS) is an important method for improving the performance of Large Language Models (LLMs) by using additional computation during the inference phase. However, current studies do not systematically analyze how policy models, Process Reward Models (PRMs), and problem difficulty influence TTS. This lack of analysis limits the understanding and practical use of TTS methods. In this paper, we focus on two core questions: (1) What is the optimal approach to scale test-time computation across different policy models, PRMs, and problem difficulty levels? (2) To what extent can extended computation improve the performance of LLMs on complex tasks, and can smaller language models outperform larger ones through this approach? Through comprehensive experiments on MATH-500 and challenging AIME24 tasks, we have the following observations: (1) The compute-optimal TTS strategy is highly dependent on the choice of policy model, PRM, and problem difficulty. (2) With our compute-optimal TTS strategy, extremely small policy models can outperform larger models. For example, a 1B LLM can exceed a 405B LLM on MATH-500. Moreover, on both MATH-500 and AIME24, a 0.5B LLM outperforms GPT-4o, a 3B LLM surpasses a 405B LLM, and a 7B LLM beats o1 and DeepSeek-R1, while with higher inference efficiency. These findings show the significance of adapting TTS strategies to the specific characteristics of each task and model and indicate that TTS is a promising approach for enhancing the reasoning abilities of LLMs.

1. Introduction

The paper studies how policy models, PRMs, and problem difficulty shape compute-optimal external test-time scaling, and reports that reward-aware scaling can let smaller models outperform larger ones on challenging reasoning tasks.

1. Introduction

  • External TTS uses sampling or search with fixed LLMs, and its central challenge is allocating computation optimally for each problem.Existing approaches use PRMs to guide generation and select final answers, but the roles of policy models, PRMs, and difficulty have not been systematically analyzed.Source: §1 · p. 2
  • Compute-optimal TTS strategy varies with the policy model, PRM, and problem difficulty, motivating the paper’s reward-aware approach.Experiments cover PRMs from 1.5B to 72B and policy models from 0.5B to 72B across two model families.Source: §1 · p. 2
  • A 3B LLM can outperform a 405B LLM, and a 7B LLM can surpass o1 and DeepSeek-R1 on MATH-500 and AIME24 with reward-aware compute-optimal TTS.The paper also reports that a 1B model can outperform a 405B model and state-of-the-art reasoning models on challenging reasoning tasks.Source: §1 · p. 2
  • The study evaluates diverse TTS methods, policy models, PRMs, scaling methods, and challenging tasks to characterize performance boundaries.The experiments focus on MATH-500 and AIME24.Source: §1 · p. 2

2. Setup & Preliminaries

The setup models reasoning as a prompt-conditioned MDP and evaluates test-time scaling methods alongside a prompt-specific compute-optimal strategy. The methods vary in how they generate and select responses or search paths.

2.1. Problem Formulation

  • Reasoning is modeled as an MDP in which a policy generates actions from the prompt, receives rewards, and appends actions to the state until an episode-ending condition.Episodes end at the maximum number of steps or when the policy generates an <EOS> token; trajectories are sequences of actions.Source: §2.1 · p. 3

2.2. Test-Time Scaling Method

  • The evaluated TTS methods are Best-of-N, beam search, and Diverse Verifier Tree Search; lookahead methods such as MCTS are excluded as inefficient.Best-of-N generates N responses and selects a final answer using scoring and voting; beam search expands verifier-selected steps, while DVTS explores independent beam-search subtrees.Source: §2.2 · pp. 3–4

2.3. Compute-Optimal Test-Time Scaling

  • Compute-optimal TTS selects strategy hyperparameters for a given prompt and compute budget to maximize performance benefits.The framework denotes the policy model's output distribution for prompt x under budget N as Target(θ, N, x).Source: §2.3 · p. 4

3. Rethinking Compute-Optimal Test-Time Scaling

The section argues that compute-optimal test-time scaling should account for the reward model as well as policy and problem difficulty. It proposes reward-aware scaling and an absolute Pass@1-based difficulty criterion.

3.1. Compute-Optimal Scaling Strategy Should be Reward-Aware

  • Earlier TTS work used either an on-policy PRM trained on the same policy’s responses, which gives more accurate rewards, or an offline PRM that can suffer OOD inaccuracies.Training a PRM for every policy to avoid OOD issues is computationally expensive, motivating study of the more general setting where PRM and TTS policies may differ.Source: §3.1 · p. 4
  • Beam search with Llama-3.1-8B-Instruct showed that RLHFlow-PRM-Mistral-8B favored short, incorrect answers, while RLHFlow-PRM-Deepseek-8B found correct answers using more tokens.The authors report that rewards significantly affect generation and outcomes, and later empirically show effects on TTS performance and output tokens.Source: §3.1 · pp. 4–5
  • Reward-aware TTS adjusts the policy’s output distribution using the reward function, prompt, and compute budget; rewards affect selection across methods and also guide search processes.For sampling-based methods, the target distribution remains unchanged by the reward function.Source: §3.1 · pp. 4–5

3.2. Absolute Problem Difficulty Criterion is More Effective Than Quantiles

  • Qwen2.5-72B-Instruct exceeds 80% Pass@1 on 76.2% of MATH-500 problems, motivating absolute difficulty thresholds rather than model-dependent quantiles.The paper defines easy as 50%–100% Pass@1, medium as 10%–50%, and hard as 0%–10%, because policy models have different reasoning capabilities.Source: §3.2 · p. 5

4. How to Scale Test-Time Compute Optimally?

The compute-optimal test-time scaling strategy varies with the policy model, PRM, and problem difficulty. PRM generalization, response-length bias, and voting-method sensitivity further affect search performance and efficiency.

4.1. Setup

  • Experiments use MATH-500 and AIME24, Llama 3 and Qwen2.5 Instruct policies, open-source PRMs, and three trajectory-scoring and three answer-voting methods.The scoring methods are PRM-Min, PRM-Last, and PRM-Avg; voting methods are Majority Vote, PRM-Max, and PRM-Vote.Source: §4.1 · p. 6

4.2. How does TTS improve with different policy models and PRMs? (Q1)

  • PRM generalization varies across policy models and tasks: search improves with larger budgets for some PRMs on MATH-500, but gains remain moderate on AIME24.For Llama-3.1-8B-Instruct, Skywork and Qwen2.5-Math PRMs scale better than Math-Shepherd and RLHFlow; Qwen2.5-7B-Instruct also scales with Skywork-PRM-7B and Qwen2.5-Math PRMs.Source: Figure 4 · p. 7, Figure 5 · p. 7, §4.2 · p. 7
  • BoN usually works best with Math-Shepherd or RLHFlow PRMs, while search works better with Skywork or Qwen2.5-Math PRMs; small policies favor search and large policies favor BoN.Using an out-of-distribution PRM can yield suboptimal answers or trap step selection in local optima; smaller policies benefit from verifier-guided intermediate-step selection, whereas larger policies need it less.Source: §4.2 · pp. 7–8, Figure 7 · p. 8

4.3. How does TTS improve for problems with different difficulty levels? (Q2)

  • Difficulty groups are defined by Pass@1 accuracy: easy is 50%–100%, medium 10%–50%, and hard 0%–10%; optimal methods vary across policy sizes and difficulty.Policies under 7B favor BoN on easy problems and beam search on harder ones; 7B–32B models favor DVTS on easy and medium problems and beam search on hard ones, while 72B models favor BoN throughout.Source: §4.3 · p. 9, Figure 8 · p. 9

4.4. Does PRMs have bias towards specific response lengths or sensitivity to voting methods? (Q3)

  • RLHFlow-PRM-Deepseek-8B uses nearly 2× as many inference tokens as RLHFlow-PRM-Mistral-8B under the same budget, while their training data also differs in response and step length.Qwen2.5-Math-7B uses more inference tokens than Skywork-PRM-7B for near-equal performance, indicating the latter is more efficient.Source: §4.4 · pp. 9–10, Table 1 · p. 9
  • Skywork-PRM-7B works better with PRM-Vote than PRM-Max, whereas Qwen2.5-Math-PRM-7B is relatively insensitive to voting method.The paper attributes Qwen2.5-Math PRMs’ robustness to training-data processing that removes wrongly positive intermediate-step labels, improving error detection during search.Source: Table 2 · p. 10, §4.4 · p. 10

5. Results for Compute-Optimal Test-Time Scaling

Compute-optimal TTS lets small policy models surpass larger or frontier models in some benchmark comparisons, while its gains vary with policy-model reasoning ability and task difficulty.

5.1. Can smaller policy models outperform larger models with the compute-optimal TTS strategy (Q4)

  • With compute-optimal TTS, 3B beats 405B on both benchmarks, 1B beats it only on MATH-500 at N=512, 0.5B and 3B beat GPT-4o, and 7B beats o1 and DeepSeek-R1.The 3B model outperforms a model 135× larger on MATH-500 and AIME24.Source: Table 3 · p. 11, §5.1 · p. 11
  • DeepSeek-R1-Distill-Qwen-1.5B with TTS beats o1-preview and o1-mini on both benchmarks, while smaller TTS models reduce total FLOPS by 100× ∼1000× versus larger CoT models.The paper reports that small policy models can surpass large ones using less inference FLOPS.Source: §5.1 · pp. 11–12, Table 4 · p. 12

5.2. How does compute-optimal TTS improve compared with CoT and majority voting? (Q5)

  • Compute-optimal TTS is 256× more efficient than majority voting and improves reasoning performance by 154.6% over CoT; gains shrink as policy-model reasoning ability rises.The reported results are for MATH-500; TTS yields larger improvements for weaker-reasoning models and more limited gains for stronger ones.Source: Table 5 · p. 12, §5.2 · p. 12

5.3. Is TTS more effective than long-CoT-based methods? (Q6)

  • On MATH-500 and AIME24, Qwen2.5-7B TTS beats rStar-Math, Eurus-2, SimpleRL, and Satori; it is close to DeepSeek-R1-Distill-Qwen-7B on MATH-500 but drops significantly on AIME24.The authors find TTS more effective than direct RL or SFT on MCTS-generated data, but less effective than distillation from strong reasoning models; it is more effective on simpler tasks.Source: Table 6 · p. 13, §5.3 · p. 13

6. Related Work

The related work spans test-time scaling, mathematical reasoning methods, and process reward models. This paper positions itself as a systematic evaluation of PRM-guided search across policies and verifiers.

LLM Test-Time Scaling

  • Prior test-time scaling uses majority voting, search, refinement, verifier guidance, and self-evaluation, but lacks systematic evaluation across strong verifiers and policies of different capabilities.The paper aims to evaluate up-to-date policies and verifiers on more challenging tasks and provide practical principles for test-time scaling.Source: §6 · p. 13

Improving Mathematical Reasoning Abilities of LLMs

  • Mathematical reasoning training methods include corpus pre-training, supervised fine-tuning, self-training, self-improvement, and long chain-of-thought approaches.The passage notes OpenAI o1’s powerful reasoning capabilities with long thinking.Source: §6 · pp. 13–14
  • Mathematical reasoning test-time methods include chain-of-thought prompting, self-refinement, external tools, and self-verification; this work scales compute with PRM-guided search.These strategies enhance reasoning without changing model parameters and complement training-time enhancements.Source: §6 · p. 14

Process Reward Models

  • PRM research addresses costly data collection through automatic collection methods, explores newer reward-modeling approaches, and releases open-source models and evaluation benchmarks.The passage notes prior findings that PRMs are more effective than ORMs and cites ProcessBench and PRMBench as comprehensive PRM evaluations.Source: §6 · p. 14

7. Conclusion & Discussion

The conclusion finds that compute-optimal TTS depends on the policy model, PRM, and problem difficulty, and can enable smaller models to outperform larger ones. It also identifies supervision and task coverage as priorities for future work.

Conclusion & Discussion

  • A 1B model achieves better performance than a 405B model through compute-optimal TTS.The authors also report that a 7B PRM can achieve strong TTS results when supervising a more capable 72B policy model.Source: §7 · p. 14
  • Compute-optimal TTS strategies depend on the policy model, PRM, and problem difficulty.The authors propose investigating “weak-to-strong” rather than “strong-to-weak” supervision for policy optimization.Source: §7 · p. 14
  • Future work should extend TTS to coding and chemistry, improve compute-optimal methods, and develop more adaptable supervision mechanisms.Both PRM-based and RL-based approaches depend on high-quality supervision; the authors seek more efficient mechanisms to improve small models on complex tasks.Source: §7 · p. 15

A. Prompt Template for Test-Time Scaling

The section specifies system-prompt templates for Llama 3 and Qwen2.5 models, and notes that Llama 3 uses its official evaluation prompt to prevent performance drop.

A. Prompt Template for Test-Time Scaling

B. Full Results of Test-Time Scaling with Different Policy Models, PRMs, and Scaling Methods

This section presents full test-time scaling results across policy models, PRMs, and scaling strategies on MATH-500 and AIME24.

C. Cases

The cases identify four PRM failure modes—over-criticism, error neglect, error localization bias, and scoring bias—that persist on both in-distribution and OOD data. These issues distort reasoning search, degrade performance, and reduce the reliability of PRM-assisted reasoning.

C. Cases

  • PRMs show over-criticism, scoring mathematically correct steps too low, and error neglect, assigning relatively high scores to steps with clear mathematical errors.The paper illustrates over-criticism in Figure 13 and error neglect in Figures 14 and 15.Source: C. Cases · p. 27, Figure 13 · p. 29, Figure 14 · p. 30
  • PRMs exhibit error localization bias, lowering scores on steps that are not the source of critical errors, and scoring bias, including sensitivity to intermediate-step token length.The paper illustrates these cases in Figures 16–18, noting that equally correct reasoning steps can receive large score discrepancies.Source: C. Cases · p. 27, Figure 16 · p. 32, Figure 17 · p. 33
  • These problems persist on both OOD AIME24 data not used for PRM training and in-distribution MATH data used for training, distorting search, degrading performance, and reducing PRM-assisted reasoning reliability.The paper argues that addressing these biases in model architectures and training procedures is necessary to improve PRM robustness and interpretability.Source: C. Cases · p. 27

Related papers

Loading 2502.06703v1…