PaperlayerView exact revision on arXiv

Source-linked AI summary

Rollout Roulette: A Probabilistic Inference Approach to Inference-Time Scaling of LLMs using Particle-Based Monte Carlo Methods

Isha Puri, Shivchander Sudalairaj, Guangxuan Xu, Kai Xu, Akash Srivastava

arXiv:2502.01618v5cs.LGcs.AI

TL;DR

PRM-guided deterministic search can discard promising reasoning paths because early scores are imperfect. The paper uses particle filtering for inference-time scaling, reporting 4–16x faster scaling than search-based methods and strong results with small models.

  • Problem

    PRM-guided inference-time search can irreversibly prune promising paths when noisy early scores favor an incorrect trajectory.

    Source: §1 · p. 1

  • Method

    The paper formulates inference-time scaling as posterior inference in an LLM-defined state-space model and uses particle filtering to sample accepted trajectories.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    Across MATH500 and AIME 2024, particle filtering scales 4–16x faster than search-based methods; Qwen2.5-Math-1.5B-Instruct surpasses GPT-4o accuracy with 4 rollouts.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The probabilistic framing offers a principled path toward more robust inference algorithms by integrating uncertainty into LLM generation.

    Source: §6 · p. 9

    Reward-model latency adds computational cost, and budget hyperparameters may require tuning across domains.

    Source: §6 · p. 9

Abstract

from arXiv · show

Large language models (LLMs) have achieved significant performance gains via scaling up model sizes and/or data. However, recent evidence suggests diminishing returns from such approaches, motivating scaling the computation spent at inference time. Existing inference-time scaling methods, usually with reward models, cast the task as a search problem, which tends to be vulnerable to reward hacking as a consequence of approximation errors in reward models. In this paper, we instead cast inference-time scaling as a probabilistic inference task and leverage sampling-based techniques to explore the typical set of the state distribution of a state-space model with an approximate likelihood, rather than optimize for its mode directly. We propose a novel inference-time scaling approach by adapting particle-based Monte Carlo methods to this task. Our empirical evaluation demonstrates that our methods have a 4-16x better scaling rate over our deterministic search counterparts on various challenging mathematical reasoning tasks. Using our approach, we show that Qwen2.5-Math-1.5B-Instruct can surpass GPT-4o accuracy in only 4 rollouts, while Qwen2.5-Math-7B-Instruct scales to o1 level accuracy in only 32 rollouts. Our work not only presents an effective method to inference-time scaling, but also connects the rich literature in probabilistic inference with inference-time scaling of LLMs to develop more robust algorithms in future work. Code, videos, and further information available at https://probabilistic-inference-scaling.github.io.

1 Introduction

Inference-time scaling offers a way to improve reasoning without increasing model size, but PRM-guided deterministic search can discard correct paths because of noisy early scores. The paper proposes particle filtering to preserve uncertainty across candidate trajectories and reports faster scaling than search-based methods.

1 Introduction

  • Inference-time scaling allocates more computation at inference time and can improve performance without increasing model size.Prior work has shown smaller models can match or exceed larger models on complex reasoning tasks.Source: §1 · p. 1
  • PRM-guided search can prune promising solutions early because noisy or miscalibrated partial scores may favor an incorrect path.Figure 1 illustrates a correct solution receiving a lower initial PRM score than an incorrect one.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • Particle filtering maintains a weighted population of candidate sequences, balancing high-PRM paths with stochastic exploration of alternatives.The paper frames inference-time scaling as posterior inference in a state-space model defined by the LLM transition model and PRM emission model.Source: §1 · p. 2
  • The paper studies mathematical and out-of-domain reasoning performance and ablates compute allocation, PRM aggregation, and generation temperature.It proposes a model-based reward aggregation method reported to improve stability and performance.Source: §1 · p. 2
  • 4–16x faster scaling speed than search-based methods is reported on MATH500 and AIME 2024; with 4 rollouts, Qwen2.5-Math-1.5B-Instruct surpasses GPT-4o accuracy.With a budget of 32, Qwen2.5-Math-7B-Instruct reaches o1 accuracy.Source: §1 · p. 2

2 Background

A state-space model describes sequential latent states and observations through transition and emission models. Inference estimates the hidden-state posterior, with particle filtering approximating it using weighted particles.

2 Background

  • A state-space model represents a sequential system through latent states that evolve step by step.The states follow a transition model, while observations are governed by an emission model.Source: §2 · pp. 2–3
  • The joint distribution of states and observations combines the initial-state prior with the emission probabilities across time.The supplied passage gives the joint form as p(x1:T, o1:T) = p(x1)∏ p(ot | xt).Source: §2 · p. 3
  • Particle filtering approximates the hidden-state posterior with weighted particles propagated by the transition model and reweighted using the emission model.Posterior inference is generally intractable in high-dimensional state spaces with dependencies.Source: §2 · p. 3

3 Method

The method casts inference-time scaling as approximate posterior inference over token sequences, using an LLM transition model and a reward-model acceptance likelihood. Particle filtering samples accepted trajectories through iterative generation, scoring, and resampling.

3 Method

  • Inference-time scaling becomes approximate posterior inference over token sequences, with the LLM as transition model and reward-model acceptances as observations.The target is the posterior over trajectories whose steps are all accepted, using an approximate reward likelihood when the true reward is unavailable.Source: §3 · pp. 3–4

3.1 Particle Filtering for Estimating the Posterior

  • Particle filtering repeatedly generates partial trajectories, scores them by reward-induced likelihood, and resamples particles to balance high-reward paths with stochastic exploration.The procedure initializes and propagates particles with the LLM, then uses softmax-weighted scores to form the next population.Source: Figure 2 · p. 3, §3.1 · p. 4
  • With prefix caching, particle filtering's runtime is comparable to generating the same number of full completions independently.Token sampling and reward computation can be parallelized across particles.Source: §3.1 · p. 4
  • The default particle weight aggregates step-level PRM rewards by product, while min, last-step, and model-based aggregation are alternatives.The alternative strategies are compared in an appendix ablation.Source: §3.1 · p. 5
  • Unlike deterministic search, particle filtering preserves a distribution over partial sequences, helping recover from early errors and smoothing noisy reward scores.It samples from the typical set rather than targeting the approximate posterior's mode, and naturally represents multiple valid completions.Source: §3.1 · p. 5
  • Under mild assumptions, particle filtering provides consistent, unbiased estimators, and its posterior samples support unbiased estimates of expected correctness.For final answer selection, the authors score completed outputs with an outcome reward model and choose the top-scoring answer for comparison with prior work.Source: §3.1 · pp. 4–5
  • The particle-filtering approach also defines an MCMC kernel for multiple iterations of complete answers and parallel simulations inspired by parallel tempering.The methods and results for these extensions are detailed in the appendix.Source: §3.1 · p. 5

4 Evaluation

The evaluation compares particle filtering with inference-time scaling baselines on mathematical and broader reasoning tasks, then examines reward-model and compute-allocation choices. PF leads the reported comparisons, including on non-math benchmarks.

4.1 Setup

  • The setup evaluates general Qwen and Llama instruct models alongside Qwen2.5-Math 1.5B and 7B models, guided by Qwen2.5-Math-PRM-7B.The selected PRM outperformed the other tested PRMs.Source: §4.1 · p. 5

Baselines

  • Baselines include greedy decoding, self-consistency, BoN/WBoN, beam search, and DVTS; tasks span MATH500, AIME 2024, NumGLUE Chemistry, and FinanceBench.MATH500 and AIME 2024 are the two benchmarks used for the main fixed-budget comparison.Source: Baselines · pp. 5–6, §4.2 · p. 6

4.2 Results on Mathematical Reasoning Datasets

  • At 32 generations, PF gives Qwen2.5-Math-7B 87.7% on MATH500 and 10/30 on AIME; PF also leads the inference-time scaling baselines across model sizes.The main comparison uses Qwen2.5-Math-PRM-7B as the reward model.Source: §4.2 · pp. 6–7, Table 1 · p. 7

Performance of different PRMs

  • Qwen2.5-Math-PRM-7B consistently outperforms the other tested PRMs, while model-based reward aggregation scales best as budget increases.The aggregation ablation compares last, minimum, product, and model aggregation strategies.Source: Performance of different PRMs · p. 8

4.3 Results on Generalized Reasoning Datasets

  • On FinanceBench and NumGLUE, PF achieves the highest accuracy among inference-time scaling baselines, using Llama 3.1-8B-Instruct with 8 and 32 particles, respectively.The evaluation uses Qwen2.5-Math-PRM-7B as the reward model.Source: §4.3 · p. 7, Table 2 · p. 8

Performance of different PRMs

  • Temperatures from 0.4–1.0 perform well with little accuracy variation across budgets, so experiments use 0.8; the appendix also explores budget allocation across iterations and parallelism.The allocation study varies particle count, iteration count, and parallelism under a fixed budget.Source: Performance of different PRMs · p. 8

5 Related Work

Related work covers reward-guided search and inference-time scaling, process reward models, and particle-based Monte Carlo methods. The paper positions its approach as using a noisy pretrained reward model without additional training, unlike verifier-dependent Twisted SMC.

Process reward models

  • PRMs score intermediate reasoning steps and are trained with process supervision; cited approaches include step-by-step verification, annotated training data, and implicit-reward objectives.PRIME's implicit objective learns a token-level Q-function without process labels.Source: §5 · p. 9

Inference-time scaling

  • Prior inference-time scaling work includes best-of-N, beam-search variants, MCTS-related methods, and DVTS, which runs parallel subtrees to avoid all leaves getting stuck in local minima.The cited work frames inference-time scaling as a training-free strategy for enhancing LLM performance.Source: §5 · p. 9

Particle-based Monte Carlo methods

  • Sequential Monte Carlo and particle filtering approximate state-space posteriors, while particle Gibbs adds MCMC; token-based SMC and Twisted SMC have also been applied to language models.The cited applications include steering generation and self-constrained generation.Source: §5 · p. 9
  • Unlike Twisted SMC, this method uses an off-the-shelf generator and noisy pretrained PRM without extra training or a ground-truth verifier; on MATH500 it reports 75.4% with 128 samples, 14.6 points above TSMC.The comparison uses the same DeepSeekMath7B generator and MATH500 dataset as the TSMC results table.Source: §5 · p. 9

Decision making with uncertainty

  • Softmax uncertainty representation is called Boltzmann exploration; a multi-armed-bandit formulation permits temperature scheduling and regret bounds but loses the unbiasedness of particle filtering.Source: §5 · p. 9

6 Conclusion

The conclusion presents particle filtering as a probabilistic alternative to deterministic inference-time search, designed to preserve candidate diversity under imperfect reward models. It also identifies computational and deployment challenges and points to probabilistic inference as a basis for future work.

6 Conclusion

  • Particle filtering maintains a diverse candidate population to balance exploration and exploitation, addressing early pruning caused by imperfect reward models.The paper frames this as probabilistic inference for inference-time scaling.Source: §6 · p. 9
  • Inference-time scaling can incur high reward-model latency; smaller models may need format-oriented prompt engineering, and budget hyperparameters may require tuning across domains.Source: §6 · p. 9
  • The authors hope the probabilistic modeling connection will support systematic solutions to current limitations and bring advanced inference algorithms into LLM scaling.Source: §6 · p. 9

A Aggregation Strategies for PRM Scores

The appendix compares ways to turn PRM scores into particle weights and studies how to distribute a fixed compute budget across particles, iterations, and parallelism.

A Aggregation Strategies for PRM Scores

  • The default product of step-level PRM rewards matches the factorized likelihood and permits online weight updates as generation proceeds.Alternatives include minimum, last-step, and model-based aggregation.Source: A Aggregation Strategies for PRM Scores · p. 15
  • Minimum aggregation penalizes weak intermediate steps but prevents online updates, whereas last-step aggregation is efficient and follows the scoring mode used in.Source: A Aggregation Strategies for PRM Scores · p. 15
  • Model-based aggregation scores the full partial trajectory with one scalar, smoothing noisy token-level rewards and potentially improving stability when PRMs are inconsistent across steps.The PRM receives the prompt and prefix and returns a scalar reward.Source: A Aggregation Strategies for PRM Scores · p. 15

B Multiple Iterations and Parallel Chains

B.1 Particle Gibbs

Particle Gibbs (PG) extends particle filtering (PF) with repeated particle-filter steps and a reference particle retained during resampling. Results compare PG budget allocations with PF, including a 7B-model run that exceeds o1 performance.

B.1 Particle Gibbs

  • Each PG iteration runs a modified PF step with a reference particle from the preceding iteration that resampling cannot replace, though its partial trajectory can be forked.PG is described as an MCMC algorithm using PF as its transition kernel.Source: §B.1 · p. 15
  • PG configurations showed no marked benefit over equivalently budgeted PF runs, while larger particle counts typically performed better than runs with fewer particles.The comparison included Llama-3.2-1B-Instruct configurations with varying iteration counts.Source: §B.1 · pp. 15–16, Figure 6 · p. 16
  • 87.2% accuracy on MATH500 with 16 particles and 4 iterations beat o1 performance.The run used Qwen 2.5 7B Math Instruct as the policy model.Source: §B.1 · pp. 15–16

B.2 Parallel Tempering

Parallel tempering (PT) runs chains at different temperatures and swaps states to support exploration and exploitation. The reported allocation experiments find improved performance when increasing the number of parallel chains at fixed particle count, and better overall scaling for PT with N = 16 than PF.

B.2 Parallel Tempering

  • PT runs parallel MCMC chains at different temperatures and swaps their states, letting hotter chains explore and cooler chains exploit newly found regions.The experiments explored only the special case of multiple chains with a single iteration.Source: §B.2 · p. 16
  • PT with N = 16 had better overall scaling than PF, and increasing parallel-chain count M improved performance for any fixed N.Figure 7 compares PF and three PT configurations on a 100-question MATH500 subset using Llama-3.2-1B-Instruct.Source: Figure 7 · p. 16, §B.2 · p. 16
  • Jointly optimizing the allocation across N, T, and M is left for future work.N is the particle count, T the iteration count, and M the parallelism.Source: §B.2 · p. 16

C Proof of Theorem 2

The proof establishes unbiasedness of the particle-filter estimate of expected accuracy. The accompanying comparison contrasts particle filtering's uncertainty-aware sampling with beam search's deterministic selection and expansion, while the algorithms outline the particle-filtering and parallel-tempering procedures.

C Proof of Theorem 2

  • Particle filtering's Monte Carlo estimate of expected correctness is unbiased over the randomness of the algorithm.The proof applies particle filtering's unbiasedness for expectations to the binary correctness function.Source: C Proof of Theorem 2 · p. 17

D Side by Side Comparison of Particle Filtering vs Beam Search

  • Particle filtering propagates expansions by sampling while retaining reward uncertainty, unlike beam search's deterministic selection and expansion.The comparison frames deterministic beam search as implicitly assuming the rewards are correct.Source: Figure 8 · p. 17

E Algorithms

  • Algorithm 1 takes a particle count, reward model, language model, and prompt, then returns weighted particles after reward-based resampling.Algorithms 2 and 3 extend this setup to Gibbs iterations and parallel chains with temperature-based resampling and state exchanges.Source: E Algorithms · pp. 17–19

F Particle Filtering Results with More Generator Models

The appendix points to further particle-filtering results with a wider variety of generator models and identifies Table 3 as reporting MATH500 and AIME 2024 performance. Its prompt template distinguishes simple from complex problems and specifies answer formatting.

F Particle Filtering Results with More Generator Models

  • Further particle-filtering results cover a wider variety of generator models, with Table 3 scoped to MATH500 and AIME 2024 results for greedy decoding and particle filtering at 64 generations per problem.The supplied table caption gives the scope but no cell values.Source: F Particle Filtering Results with More… · p. 19, Table 3 · p. 19

G Inference Prompt Template

  • The prompt calls for concise solutions with minimal explanation for problems of two steps or fewer, and a step-by-step format for problems of three steps or more.The complex-problem template uses concise step descriptions with brief explanations and calculations.Source: G Inference Prompt Template · p. 20
  • The prompt instructs every solution to conclude with a boxed final answer, where the answer is only the final number or expression.Source: G Inference Prompt Template · p. 20

H Evaluation details

The paper describes answer parsing and exact-match verification, alongside checklist statements about reproducibility, statistical reporting, compute resources, and limitations. The checklist also records that one theorem is proved in Appendix C.

H Evaluation details

  • Answers are parsed with heuristics, normalized through SymPy, and counted correct when numerically or symbolically equal to the ground truth.Generated and reference expressions are converted from LaTeX to SymPy and back; accuracy is the fraction of exact matches.Source: H Evaluation details · p. 20

NeurIPS Paper Checklist

  • The stated limitations include reward-model latency, prompt-formatting needs for smaller models, and problem-dependent budgets that may require cross-domain tuning.The authors identify hosting and running a reward model as resource-intensive.Source: NeurIPS Paper Checklist · p. 21
  • The paper reports one theorem and says its proof appears in Appendix C.The checklist marks the theoretical assumptions and proof as provided.Source: NeurIPS Paper Checklist · p. 22
  • The authors report full experimental details, including datasets, model versions, hyperparameter selection, reward-model aggregation, and ablations, as support for reproducibility.They state that code is not included at submission to preserve confidentiality and will be open-sourced upon acceptance.Source: NeurIPS Paper Checklist · pp. 22–23
  • The authors did not run every experiment multiple times, citing academic-budget computational limits, and report no error bars across all experiments.They say many experiments were repeated during research and development, but not every experiment received multiple runs.Source: NeurIPS Paper Checklist · p. 24
  • The authors do not provide formal compute-resource details because different experiments used different numbers of GPUs.They note that an off-the-shelf library with basic usage information is planned for release upon acceptance.Source: NeurIPS Paper Checklist · p. 24

Related papers

Loading 2502.01618v5…