Source-linked AI summary
Recipes for Steering and Scaling LLMs via Sampling
Jiajun He, Zongyu Guo, José Miguel Hernández-Lobato, Yuanqi Du
TL;DR
Sampling strategies for richer LLM target distributions remain inefficient. The paper develops a framework using SMC and RE to sample powered, product, and tilted distributions, and reports more favorable scaling than Best-of-N and standard MCMC while noting computational and model-quality limits.
Problem
Sampling richer target distributions beyond the base LLM remains highly inefficient.
Method
The paper formulates autoregressive generation as time-dependent transport and introduces SMC and RE for powered, product, and tilted target distributions.
Results
SMC and RE improve over naive low-temperature sampling and scale more favorably than standard MCMC on reasoning benchmarks, while entropy tilting further improves performance without an external verifier or reward model.
Takeaways & Limitations
SMC is particularly effective for strong single answers, whereas RE better preserves diversity and delivers stronger pass@k gains.
Takeaways & Limitations
The method adds computational overhead, depends on base-model quality, and may be less useful when the exact target distribution is unimportant.
Abstract
from arXiv · showhide
Large Language Models (LLMs) are probabilistic models, typically defined by an autoregressive factorization. While recent work has begun to study richer target distributions beyond the base model, the sampling strategies remain highly inefficient. In this paper, we present a flexible and theoretically grounded framework for steering and scaling autoregressive LLMs with sampling. Within this framework, we describe two algorithms -- one based on Sequential Monte Carlo (SMC) and one based on Replica Exchange (RE) -- that steer generation toward powering, product or tilting of the base model distribution. We illustrate this framework through scaling the generation quality of LLMs without external supervision or reward models. Experimental results demonstrate our methods scale more favorably than Best-of-N and standard MCMC baselines. Overall, this paper offers a systematic recipe for probabilistic inference with LLMs via sampling.
1 INTRODUCTION
LLMs are probabilistic autoregressive models, motivating inference methods that exploit sampling to improve generation. The paper introduces a framework for steering and scaling LLMs through flexible target distributions and principled sampling algorithms.
- LLMs define distributions over sequences by autoregressively sampling each token conditional on the previously generated prefix.
- Existing inference approaches include speculative decoding, reinforcement-learning-based inference, controlled decoding, and stochastic search methods such as Best-of-N.
- Probabilistic inference studies have examined fine-tuning as distribution sharpening, MCMC sampling from sharpened distributions, and base-model coverage for post-training.
- The paper introduces powered, product, and tilted target distributions and derives sequential Monte Carlo and replica exchange algorithms for inference.
- The framework enhances base-LLM capability on mathematical reasoning by altering inference alone, without external reward supervision.
2 BACKGROUND
Sampling from intractable targets can be approached through importance sampling or MCMC, but both face efficiency challenges when proposals are poorly calibrated or targets are multimodal. Annealing addresses these challenges with intermediate distributions, supporting SMC and RE variants that progressively correct or exchange samples across the path.
- Autoregressive LLM sampling draws tokens sequentially until the absorbing empty token is generated.
- Importance Sampling: Importance sampling draws from a tractable proposal, weights samples by target plausibility, and reweights or resamples them accordingly.
- Importance Sampling: Importance sampling can fail when the proposal does not cover the target, leaving most samples with negligible weight.
- Markov Chain Monte Carlo: MCMC constructs a chain targeting π by proposing local moves and accepting or rejecting them according to target plausibility.
- Markov Chain Monte Carlo: MCMC may mix slowly on multimodal targets because the chain can become stuck in a local mode.
- Accelerated Sampling: Annealing creates intermediate distributions between a tractable prior and a complicated target, yielding sequential Monte Carlo and replica exchange variants.
- Sequential Monte Carlo: SMC propagates weighted particles along the annealing path and corrects proposal-target mismatch with incremental importance weights.
- Replica Exchange: RE assigns one replica to each intermediate distribution, updates replicas locally, and occasionally swaps neighboring replicas.
3 METHODS
The paper formulates flexible target distributions for autoregressive LLMs and develops SMC and replica-exchange methods to sample from them efficiently. These methods operate over token-generation trajectories, including chunked variable-length sequences.
- Transport formulation: The proposed methodology treats autoregressive generation as a time-dependent transport with a forward token-generation kernel and a deterministic backward kernel that removes the last token.
- Flexible target distributions: The framework steers an autoregressive base model toward powered, tilted, product or quotient target distributions, including compositions of these objectives.
- Sampling challenge: Sampling these targets is difficult because they generally lack the base LLM’s autoregressive factorization, making direct MCMC less effective for complicated targets.
- Sequential Monte Carlo: SMC propagates particles along intermediate targets and corrects proposal mismatch with incremental weights derived from forward and backward kernels.
- Replica exchange: The framework also derives replica-exchange swap ratios and applies replica exchange directly to autoregressive generation rather than to replicas sharing one state space.
- Chunked implementation: Both algorithms can operate on contiguous token chunks; for replica exchange, replicas evolve at different lengths and alternate swaps between adjacent pairs.
4 RELATED WORKS
Related work studies inference-time control and scaling across prompting, representation interventions, and guidance-based generation. These approaches provide context for sampling-based steering of LLM behavior.
- Inference-time control: Prompting methods elicit intermediate reasoning or structured decompositions, while optimized prompts further control generated text.
- Inference-time control: Other approaches intervene directly in internal LLM representations to control model behavior.
- Generative-model control: Guidance methods in diffusion models modify generation to satisfy constraints or desired properties, but can introduce bias and fail under strong constraints.
5 EXPERIMENTS
The experiments evaluate SMC and RE for inference-time LLM improvement, comparing them with sampling and Best-of-N baselines across mathematical reasoning benchmarks. Results show favorable scaling, gains from powering and entropy tilting without external verification, and distinct diversity–quality trade-offs between SMC and RE.
- Experimental setup: The evaluation uses MATH500 with Qwen2.5-Math-7B and Qwen2.5-7B, plus GPQA-Diamond with Qwen2.5-32B-Instruct.The same data split and answer format as prior work are used; the GPQA model choice addresses unreliable answer formatting from the 7B base model.
- Baseline comparisons: SMC and RE substantially outperform naive low-temperature autoregressive sampling on MATH500.The paper characterizes both methods as global search, contrasting with greedy standard low-temperature sampling.
- Baseline comparisons: SMC and RE scale more favorably than MCMC because annealed intermediate targets act as checkpoints, avoiding MCMC’s large-revision acceptance-rate dilemma.MCMC randomly removes and refills trajectories, whereas SMC and RE anneal through intermediate targets.
- Steering experiments: Powering and entropy tilting improve base-model accuracy without an external verifier, while combining them improves results further.The experiments cover MATH500 and GPQA-Diamond across different power and tilting strengths.
- Steering experiments: Entropy tilting is non-monotone: very large strengths can hurt performance, with RE tending to drop and SMC tending to plateau.The paper attributes this possibility to increasing target–proposal mismatch and lower acceptance rates, especially for RE.
- Baseline comparisons: Best-of-N outperforms MCMC, exceeds RE, and performs on par with SMC, but unlike SMC and RE it searches for one best sample rather than preserving a sample distribution.SMC and RE are therefore positioned as simple choices when sample diversity matters.
- Population benefits: With multiple samples, RE eventually surpasses SMC in pass@k, while SMC has stronger first-sample performance and mildly improving curves.The pattern indicates more collapsed SMC samples and more diverse RE candidates; both methods use matched sampling budgets.
6 CONCLUSION
The paper presents a probabilistic sampling framework for inference-time steering and scaling, combining flexible target distributions with practical SMC and RE algorithms. The methods improve over naive low-temperature sampling, scale more favorably than standard MCMC, and have complementary strengths, while adding computational overhead and depending on base-model quality.
- The framework derives powered, tilted, and product-form target distributions for steering and scaling autoregressive LLMs at inference time.It views autoregressive generation as a time-dependent transport.
- Practical SMC and RE algorithms sample from these target distributions and improve over naive low-temperature sampling on reasoning benchmarks.The methods also scale more favorably than standard MCMC.
- SMC is particularly effective for obtaining a strong single answer, whereas RE better preserves sample diversity and gives stronger pass@k gains.Best-of-N remains a strong search baseline, while SMC and RE sample from the full target distribution.
- The method introduces additional computational overhead, depends on underlying-model quality, and may be unnecessary when the exact target distribution is unimportant.In such cases, BON is presented as a simple alternative with strong performance.
A EXPERIMENTAL HYPERPARAMETERS
The experiments use specified datasets, models, prompting, and sampling limits to evaluate the methods on mathematical reasoning and question answering.
- The experiments use MATH500 with Qwen2.5-Math-7B and Qwen2.5-7B, and GPQA-Diamond with Qwen2.5-32B-Instruct.All experiments use chain-of-thought prompting.
- The sampling configuration uses maximum generation length Tmax = 3072, proposal temperature τ = 0.25, and 16 chunks.Each chunk contains at most 192 newly generated tokens.
B BROADER IMPACT
The paper notes that its broader societal implications are associated with LLM development and deployment generally rather than being specific to this method.
- The authors advocate responsible and appropriate LLM use with careful consideration of potential risks and impacts.
C DECLARATION OF EXPERIMENTS COMPUTE RESOURCES
Experiments run on NVIDIA A100 40G and 80G GPUs, with complete dataset-and-method result collection requiring several days to weeks.
- The experiments use NVIDIA A100 40G and NVIDIA A100 80G GPUs.
- Collecting results for all datasets and methods can take several days to weeks.Each question is fast to run.
D DECLARATION OF LLM USAGE
The authors used LLMs for coding, manuscript polishing, and proofreading, and studied relatively lightweight Qwen-7B/32B models.
- LLMs assisted with coding, manuscript polishing, and proofreading.
- The paper studies relatively lightweight Qwen-7B/32B LLMs.
E LICENSES FOR EXISTING ASSETS
The paper uses Qwen-2.5 model weights and code from Karan and Du, while GPQA and MATH500 are released under MIT licenses.
- Qwen-2.5 model weights are used under the Qwen LICENSE AGREEMENT.
- The paper uses code from Karan and Du [2025], released without a license.
- GPQA and MATH500 are released under MIT licenses.