Source-linked AI summary

Inference Scaling Laws: An Empirical Analysis of Compute-Optimal Inference for Problem-Solving with Language Models

Yangzhen Wu, Zhiqing Sun, Shanda Li, Sean Welleck, Yiming Yang

arXiv:2408.00724v3cs.AI

TL;DR

Training scaling laws are well studied, but optimal inference configurations under compute constraints remain less understood. The paper analyzes model sizes and inference strategies under fixed budgets, finding that smaller models with advanced algorithms can deliver favorable cost-performance trade-offs. It also characterizes sampling limits and proposes REBASE as a tree-search method.

  • Problem

    Optimal configurations of model size and inference strategy under fixed inference compute remain insufficiently understood.

  • Method

    The paper evaluates sampling, voting, and tree-search strategies across model sizes and budgets, formulates compute-optimal inference, and proposes REBASE.

  • Results

    Smaller models with advanced inference algorithms often achieve favorable cost-performance trade-offs, while sampling and voting exhibit convergence limits and diminishing returns.

  • Takeaways & Limitations

    Inference compute can be used more efficiently by combining smaller models with sophisticated inference strategies, particularly when deployment compute is constrained.

  • Takeaways & Limitations

    The formal setup assumes a finite vocabulary and bounded answer length, with answers ending at a designated reasoning-end token.

Abstract

from arXiv · show

While the scaling laws of large language models (LLMs) training have been extensively studied, optimal inference configurations of LLMs remain underexplored. We study inference scaling laws (aka test-time scaling laws) and compute-optimal inference, focusing on the trade-offs between model sizes and generating additional tokens with different inference strategies. As a first step towards understanding and designing compute-optimal inference methods, we studied cost-performance trade-offs for inference strategies such as greedy search, majority voting, best-of-$n$, weighted voting, and two different tree search algorithms, using different model sizes and compute budgets. Our findings suggest that scaling inference compute with inference strategies can be more computationally efficient than scaling model parameters. Additionally, smaller models combined with advanced inference algorithms offer Pareto-optimal trade-offs in cost and performance. For example, the Llemma-7B model, when paired with our novel tree search algorithm, consistently outperforms the Llemma-34B model across all tested inference strategies on the MATH benchmark. We hope these insights contribute to a deeper understanding of inference scaling laws (test-time scaling laws) for LLMs.

1 INTRODUCTION

The paper studies how model size and inference strategy should be selected under fixed inference compute. It finds that additional inference computation, especially with smaller models and advanced search, can improve cost-performance trade-offs.

  • Limited knowledge remains about how varying inference compute affects trained-model performance, despite extensive training scaling-law research.
  • The study evaluates greedy search, majority voting, best-of-n, weighted voting, and tree-search variants across model sizes and inference FLOPs.
  • Increasing inference compute usually raises accuracy until saturation; smaller models can lead at lower budgets, while larger models become favorable after saturation.
  • Sampling and voting converge exponentially toward limits determined by model output probabilities, producing diminishing returns without an oracle verifier.
  • REBASE controls node expansion with node-quality rewards, avoiding explicit rollouts while maintaining enough candidate solutions for voting.
  • REBASE consistently outperforms sampling and MCTS across tested settings, while smaller models paired with advanced inference achieve favorable cost-performance trade-offs.
  • The paper formulates compute-optimal inference as selecting model size and strategy to maximize accuracy under a fixed FLOPs budget.

2 RELATED WORKS

Prior work established scaling relationships between model performance, parameters, data, and training compute, alongside diverse inference strategies for generating and selecting solutions. This paper extends the scaling-law perspective to inference-time model and strategy selection.

  • Scaling-law studies have examined model performance across language, image, video, reward, and board-game domains as a function of model size and training compute.
  • Inference methods range from deterministic decoding to sampling and candidate-selection procedures that trade off sequence quality, diversity, performance, and cost.
  • Mathematical problem solving is used to measure LLM reasoning abilities, with later work adding verifiers, majority voting, and weighted majority voting.
  • Figure 2 contrasts training scaling, which selects model size and training tokens, with inference scaling, which selects model size and inference strategy under compute budgets.

3 COMPUTE-OPTIMAL INFERENCE FOR PROBLEM-SOLVING

The paper formulates compute-optimal inference as choosing model size and inference strategy under a fixed FLOPs budget. It analyzes sampling and tree-search approaches, including REBASE, and studies how accuracy scales with additional inference compute.

  • The compute-optimal problem selects model size and inference strategy to maximize accuracy under a fixed FLOPs budget.The formulation represents error as a function of model parameters, generated tokens, and strategy.
  • 3.1 INFERENCE STRATEGIES: Inference compute can be increased by generating more tokens, sampling additional candidate solutions, and ranking or voting over them.The considered strategies include greedy search, majority voting, best-of-n, weighted voting, and tree-search variants.
  • Sampling and voting accuracy converges with more samples, so fixed models eventually exhibit saturation and diminishing returns.The limiting accuracy depends on the language-model distribution and, for weighted voting, the reward-model distribution.
  • 3.1 INFERENCE STRATEGIES: Weighted majority voting can have a higher accuracy limit than majority voting when the reward model is better than random.The experiments consistently find weighted majority voting dominates majority voting, motivating emphasis on best-of-n and weighted voting.
  • 3.1.2 REWARD BALANCED SEARCH (REBASE): MCTS can have unfavorable compute-performance trade-offs because many search paths are used for node-quality estimation.The paper motivates a cheaper tree-search method with costs closer to weighted voting and best-of-n.
  • 3.1.2 REWARD BALANCED SEARCH (REBASE): REBASE uses reward scores to control node expansion without explicit rollouts, preserving tree-search pruning while reducing inference compute.It expands nodes at each depth according to softmax-normalized rewards under a total expansion budget.
  • 3. COMPUTE-OPTIMAL INFERENCE FOR PROBLEM-SOLVING: REBASE is compute-optimal across tested budgets on MATH, with 7B typically the optimal model size.The figure compares error rates for weighted-majority and best-of-n configurations across inference strategies and model sizes.

4 EXPERIMENTS

The experiments examine how model size, inference strategy, and inference compute affect mathematical problem-solving performance. They find that compute-optimal choices vary with budget, task difficulty, and strategy, with REBASE providing especially favorable cost-performance trade-offs.

  • Experimental questions: The experiments ask how performance scales with inference compute when varying model size and inference strategy.
  • Datasets: Experiments use MATH and GSM8K to evaluate inference scaling on challenging and simpler mathematical reasoning problems.
  • Experimental setup: The study evaluates Pythia, Llemma-7B, Llemma-34B, and Mistral-7B with sampling, voting, and tree-search strategies.
  • Compute-optimal model size: The optimal model size varies with inference compute: smaller models are initially preferable, while larger models become favorable after smaller-model accuracy saturates.
  • Model-size comparison: Llemma-7B requires around 2× less total FLOPs than Llemma-34B to achieve comparable accuracy across tested strategies and tasks.
  • Inference strategies: REBASE consistently achieves the best cost-performance trade-offs across tested models, tasks, and inference budgets, whereas MCTS underperforms sampling methods.
  • Inference strategies: With the 7B model, REBASE achieves higher accuracy with 7 times less compute than sampling-based weighted voting.
  • Problem difficulty: REBASE and sampling perform comparably on MATH-easy problems, but REBASE has a significant advantage on MATH-hard problems.

5 CONCLUSIONS

The paper develops empirical inference scaling laws and a compute-optimal inference framework for choosing model sizes and strategies under fixed budgets. Its results favor smaller models with more tokens and advanced inference methods, especially the Pareto-optimal REBASE algorithm.

  • The paper relates task performance to inference compute across model sizes, model families, and inference strategies to characterize compute-optimal configurations.
  • At fixed compute budgets, smaller models generating more tokens often outperform larger models, motivating smaller models with sophisticated inference strategies.
  • Sampling-based majority voting inevitably saturates to a distribution determined by the underlying generation policy as compute approaches infinity.
  • REBASE is Pareto optimal across tested compute budgets and outperforms commonly used weighted majority voting and MCTS methods.

A.1 PROOF OF THEOREM 1

The proof analyzes majority voting by characterizing its error event and showing that voting converges almost surely to the most probable answer as sample count grows.

  • Majority-voting failure occurs only if some alternative answer receives at least as many samples as y* among the first n draws.
  • The count difference between y* and any alternative is modeled as a sum of n i.i.d. variables with expectation δ equal to their probability gap.
  • Borel–Cantelli then implies almost-sure convergence, so majority voting eventually outputs y* for a single example and, by combining examples, for finite datasets.
  • The proof defines E_n(x_i) as the event that majority voting with n samples does not output y*(x_i).
  • Hoeffding-style concentration yields an exponentially decreasing error bound, P(E_n(x_i)) = O(c^-n), for a constant c > 1.

A.2 PROOF OF THEOREM 2

This section introduces the proof’s supporting inference procedures and MCTS mechanics, including node selection, expansion, evaluation, and backpropagation.

  • The proof applies the same argument as Theorem 1 after setting the relevant notation for the second asymptotic result.
  • MCTS selects nodes recursively using the highest Upper Confidence Bound applied to Trees value until reaching an unexpanded node.
  • After selection, MCTS expands a non-terminal node into multiple children and evaluates each child with a value function V(c).
  • Backpropagation updates UCT values and visit counts along the path from the selected node back to the root.
  • Inference uses temperature 1.0, a 1024-token output limit, and a 256-token per-step limit, with method-specific expansion settings.

D.1 MAJORITY VOTING EXPERIMENT RESULTS

The experiments compare majority voting with weighted majority voting and find that REBASE substantially narrows their sampling gap.

  • The sampling gap between majority voting and weighted majority voting is large, but becomes much smaller when REBASE is applied.
  • The authors attribute this narrowing to REBASE’s tree-search selection ability, which samples solutions with high rewards before weighted voting.

D.2 ADDITIONAL EXPERIMENTS ON LLAMA3 MODELS

Additional experiments test the conclusions across Llama3 models, mathematical reasoning, and code generation. They report improved performance with more inference compute until saturation and identify REBASE as the best compute-performance trade-off.

  • D.2 ADDITIONAL EXPERIMENTS ON LLAMA3 MODELS: Experiments with Llama3-8B-Instruct on MATH and GSM8K and with MBPP code generation test whether the conclusions extend across architectures and tasks.
  • D.2 ADDITIONAL EXPERIMENTS ON LLAMA3 MODELS: Across tested settings, increased computational effort improves performance until saturation, while REBASE reaches the optimal performance-compute trade-off.
  • D.2 ADDITIONAL EXPERIMENTS ON LLAMA3 MODELS: REBASE consistently outperforms sampling across best-of-n, majority voting, and weighted voting on mathematical reasoning tasks.
  • D.2 ADDITIONAL EXPERIMENTS ON LLAMA3 MODELS: 90.2% accuracy is achieved by REBASE with weighted majority voting using 128 samples on GSM8K, versus 89.7% for sampling with 256-sample best-of-n.
  • D.2 ADDITIONAL EXPERIMENTS ON LLAMA3 MODELS: 47.4% accuracy is reported for REBASE, compared with 41.9% for sampling with 256 samples using best-of-n.
  • D.2 ADDITIONAL EXPERIMENTS ON LLAMA3 MODELS: For MBPP, REBASE is more compute-efficient than sampling because its reward model prunes suboptimal partial code solutions while exploring promising ones.
Loading 2408.00724v3…