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

Existing TTS studies do not systematically explain how policy models, PRMs, and problem difficulty affect compute allocation, limiting practical understanding of the method. This paper analyzes compute-optimal TTS across mathematical tasks and model configurations, finding that strategy choice is highly dependent on these factors and that small models can outperform much larger ones.

  • Problem

    Existing TTS research does not systematically analyze the effects of policy models, PRMs, and problem difficulty on test-time computation.

  • Method

    The paper conducts experiments on MATH-500 and AIME24 across policy models, PRMs, and external TTS strategies to identify compute-optimal scaling.

  • Results

    Compute-optimal TTS depends on policy models, PRMs, and problem difficulty, while a 1B model can outperform a 405B model through TTS.

  • Takeaways & Limitations

    Adapting TTS strategies to task and model characteristics can enhance reasoning and enable smaller models to outperform larger models.

  • Takeaways & Limitations

    The evaluation focuses on mathematical tasks, and the authors identify extending TTS to coding and chemistry as future work.

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

Test-time scaling improves LLM reasoning by allocating additional inference computation, but its effectiveness depends on policy models, PRMs, and problem difficulty. The paper evaluates these factors and reports that compute-optimal scaling can let smaller models outperform much larger and frontier reasoning models.

  • TTS motivation and scope: External TTS uses sampling or search with fixed LLMs, while PRMs guide generation and final-answer selection.The paper distinguishes external TTS from internal TTS, which trains models to reason with longer chain-of-thought.
  • Study design: The paper evaluates TTS across current policy models, PRMs, scaling methods, and challenging mathematical tasks.Experiments use policy models from 0.5B to 72B parameters and PRMs from 1.5B to 72B on MATH-500 and AIME24.
  • Key finding: Compute-optimal TTS varies with the policy model, PRM, and problem difficulty level.The analysis specifically examines how these factors influence the optimal allocation of test-time computation.
  • Key finding: A 3B LLM can outperform a 405B LLM, while a 7B LLM can surpass o1 and DeepSeek-R1 on MATH-500 and AIME24.These comparisons use the reward-aware compute-optimal TTS strategy.

2. Setup & Preliminaries

The paper formulates reasoning as sequential decision-making and studies external TTS methods that generate, evaluate, and select solution trajectories. Compute-optimal scaling chooses strategy settings for a prompt and budget to maximize performance.

  • Problem formulation: Reasoning is formulated as a Markov Decision Process with states, actions, transitions, rewards, and a discount factor.A policy generates actions from the initial prompt state, receives rewards, and transitions by concatenating prior states and actions.
  • TTS methods: The evaluated TTS methods are Best-of-N, beam search, and Diverse Verifier Tree Search.Lookahead methods such as Monte Carlo Tree Search are excluded because multi-step sampling is considered inefficient.
  • TTS methods: Best-of-N generates N responses and applies scoring and voting to select the final answer.This is a sampling-based method with response-level selection.
  • TTS methods: Beam search expands selected steps, while DVTS independently explores multiple subtrees using beam search.Beam search uses a verifier to retain steps for subsequent search; DVTS adds diversity through independent subtree exploration.
  • Compute-optimal scaling: Compute-optimal scaling selects test-time strategy hyperparameters for each prompt and compute budget to maximize performance.The strategy is defined relative to the policy output distribution and the prompt’s ground-truth response.

3. Rethinking Compute-Optimal Test-Time Scaling

Compute-optimal TTS should account for the reward function, policy model, prompt, and problem difficulty. The paper motivates reward-aware scaling because PRMs can produce different outcomes and token costs, while absolute difficulty thresholds better reflect variation across policy models.

  • Reward-aware scaling: RLHFlow-PRM-Mistral-8B rewards short responses that are incorrect, whereas RLHFlow-PRM-Deepseek-8B produces correct answers with more tokens.
  • Reward-aware scaling: Rewards should be integrated into compute-optimal TTS so scaling adapts to the policy model, prompt, and reward function.For sampling-based methods, the reward-adjusted target distribution equals the original target distribution.
  • Absolute problem difficulty: Absolute Pass@1 thresholds define easy, medium, and hard problems because policy models differ in reasoning capability.The thresholds are easy: 50% ∼100%, medium: 10% ∼50%, and hard: 0% ∼10%.

4. How to Scale Test-Time Compute Optimally?

The paper evaluates how policy models, PRMs, voting methods, and problem difficulty affect compute-optimal TTS on MATH-500 and AIME24. It finds that the best scaling strategy varies across these factors, while PRM generalization, length bias, and voting sensitivity further shape outcomes.

  • Policy Models and PRMs: PRM generalization is difficult across policy models and tasks, especially on more complex tasks.On MATH-500, Skywork and Qwen2.5-Math PRMs scale better for the tested policies, whereas Math-Shepherd and RLHFlow can remain poor; on AIME24, TTS gains remain moderate despite higher Pass@k.
  • Policy Models and PRMs: The optimal TTS method depends on the PRM: BoN usually performs best with Math-Shepherd and RLHFlow, whereas search-based methods perform better with Skywork and Qwen2.5-Math PRMs.The paper attributes this pattern to limited out-of-distribution generalization, with stepwise selection by an OOD PRM potentially trapping answers in local optima.
  • Policy Models and PRMs: The optimal TTS method also depends on policy-model size: search-based methods outperform BoN for small models, while BoN is more effective for large models.The experiments cover Qwen2.5 models from 0.5B through 72B parameters; the paper links the difference to smaller models’ greater reliance on verifier-guided intermediate-step selection.
  • Problem Difficulty: Difficulty changes the preferred method: small models favor BoN on easy problems and beam search on hard ones, while 72B models favor BoN at every difficulty level.For models between 7B and 32B, DVTS performs well on easy and medium problems, while beam search is preferable for hard problems.
  • PRM Properties: PRMs can bias compute allocation toward response length, and their effectiveness can depend on the voting method.Under the same budget, RLHFlow-PRM-Deepseek-8B uses nearly 2× as many inference tokens as RLHFlow-PRM-Mistral-8B; Skywork-PRM-7B works better with PRM-Vote than PRM-Max, while Qwen2.5-Math-PRM-7B is less sensitive.

5. Results for Compute-Optimal Test-Time Scaling

The experiments examine compute-optimal TTS across smaller policy models, frontier-model comparisons, efficiency, and long-CoT baselines on MATH-500 and AIME24. Results show that carefully scaled inference can let smaller models outperform larger ones while using substantially less computation.

  • Compute-optimal TTS is evaluated against CoT and long-CoT methods on MATH-500 and AIME24.
  • Llama-3.2-3B-Instruct with compute-optimal TTS outperforms Llama-3.1-405B-Instruct on both MATH-500 and AIME24.This corresponds to outperforming a model 135× larger.
  • Llama-3.2-1B-Instruct beats Llama-3.1-405B-Instruct on MATH-500 at N=512 but underperforms it on AIME24.
  • Qwen2.5-0.5B-Instruct and Llama-3.2-3B-Instruct with compute-optimal TTS outperform GPT-4o.
  • DeepSeek-R1-Distill-Qwen-7B with compute-optimal TTS outperforms o1 and DeepSeek-R1 on MATH-500 and AIME24.
  • Small policy models surpass larger ones with less inference FLOPS, reducing total FLOPS by 100× ∼1000×.
  • On MATH-500, compute-optimal TTS is 256× more efficient than majority voting and improves reasoning performance by 154.6% over CoT.The improvement decreases as policy-model parameter counts increase, with weaker reasoning models gaining more.
  • Qwen2.5-7B TTS outperforms rStar-Math, Eurus-2, SimpleRL, and Satori on both benchmarks but drops significantly on AIME24 relative to DeepSeek-R1-Distill-Qwen-7B.

6. Related Work

Related work covers test-time scaling, mathematical-reasoning improvement, and process reward models. Prior TTS studies explore sampling, search, refinement, verification, and self-evaluation, but lack evaluation across strong verifiers and policies with varied capabilities.

  • Prior TTS methods include majority voting, search-based methods, refinement, verification-guided inference, and self-evaluation.
  • Existing TTS work lacks evaluation with strong verifiers and policy models of different sizes and capabilities.
  • Mathematical-reasoning methods span training-time pretraining, supervised fine-tuning, self-training, and self-improvement.
  • Test-time reasoning methods include Chain-of-Thought prompting and self-refinement without changing model parameters.
  • Process reward models are generally more effective than outcome reward models, but collecting high-quality PRM data is costly.

7. Conclusion & Discussion

The discussion concludes that compute-optimal TTS depends on policy models, PRMs, and task difficulty, and can enable smaller models to outperform larger ones. The evaluation remains focused on mathematical tasks, while broader tasks and improved scaling methods are future directions.

  • Compute-optimal TTS strategies depend on policy models, PRMs, and problem difficulty.
  • A 1B model can achieve better performance than a 405B model through TTS.
  • A 7B PRM can achieve strong TTS results while supervising a more capable 72B policy model.This motivates investigating weak-to-strong supervision rather than only strong-to-weak supervision.
  • PRM-based and RL-based approaches remain limited by their dependence on high-quality supervision.
  • The study evaluates TTS on mathematical tasks, with extensions to coding and chemistry identified as future work.

A. Prompt Template for Test-Time Scaling

The appendix lists the system prompts used for the Llama 3 and Qwen2.5 model series in Tables 7 and 8.

  • The Llama 3 series system prompt is provided in Table 7.
  • The official Llama 3 evaluation system prompt is used to prevent performance drop.
  • The Qwen2.5 series system prompt is provided in Table 8.

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

Figures 9–11 present TTS results across policy models, PRMs, scaling strategies, and problem-difficulty levels on MATH-500 and AIME24.

  • The full results are shown in Figures 10 and 11.
  • Figure 9 compares three Llama policy models on MATH-500 across different difficulty levels.
  • Figure 10 reports TTS performance on MATH-500 across different policy models, PRMs, and scaling strategies.
  • Figure 11 reports TTS performance on AIME24 across different policy models, PRMs, and scaling strategies.

C. Cases

The case studies examine PRM behavior during TTS, including scoring failures on reasoning steps and worked mathematical examples. These examples illustrate how reward signals can favor incorrect or inefficient search outcomes.

  • PRM issues: PRMs exhibit over-criticism, error neglect, error localization bias, and scoring bias during TTS.The reported problems include low scores for correct steps, high scores for erroneous steps, misplaced penalties, and length-sensitive discrepancies.
  • PRM issues: These PRM issues persist on both OOD AIME24 data and in-distribution MATH data, reducing the reliability of PRM-assisted reasoning.
  • Mathematical examples: One worked case asks for the least positive multiple of 30 expressible using only digits 0 and 2.
  • Mathematical examples: Another worked case simplifies sqrt(242) to 11sqrt(2) through prime factorization.
  • Mathematical examples: A trigonometric worked case computes DE = 10 from sin D = 0.7 and EF = 7 in a right triangle.
Loading 2502.06703v1…