PaperlayerView exact revision on arXiv

Source-linked AI summary

Scalable Power Sampling: Unlocking Efficient, Training-Free Reasoning for LLMs via Distribution Sharpening

Xiaotong Ji, Rasul Tutunov, Matthieu Zimmer, Haitham Bou Ammar

arXiv:2601.21590v1cs.LGcs.AI

TL;DR

The paper addresses the computational cost of MCMC-based power-distribution sampling for obtaining RL-like reasoning improvements without external rewards. It derives a future-aware, token-level approximation and implements it as training- and verifier-free autoregressive generation. Across four LLMs and math, code, and QA benchmarks, the method matches GRPO while achieving over 10× lower latency than MCMC.

  • Problem

    MCMC can sample a sharpened power distribution without external rewards, but its high computational cost limits scalability.

  • Method

    The method approximates the global power distribution with a scaled low-temperature distribution whose token-specific factor reflects future trajectory quality, enabling training- and verifier-free autoregressive sampling.

  • Results

    Across MATH500, HumanEval, and GPQA on four LLMs, the method matches or surpasses one-shot GRPO and reduces MCMC inference latency by over 10×.

  • Takeaways & Limitations

    Distribution sharpening can provide a scalable, inference-time alternative to RL post-training and MCMC-based power sampling within the evaluated tasks.

  • Takeaways & Limitations

    The paper identifies future work in reducing estimator variance, optimizing compute through adaptive budgets, amortizing rollout costs, and extending the method to agentic settings.

Abstract

from arXiv · show

Reinforcement learning (RL) post-training is a dominant approach for improving the reasoning performance of large language models (LLMs), yet growing evidence suggests that its gains arise primarily from distribution sharpening rather than the acquisition of new capabilities. Recent work has shown that sampling from the power distribution of LLMs using Markov chain Monte Carlo (MCMC) can recover performance comparable to RL post-training without relying on external rewards; however, the high computational cost of MCMC makes such approaches impractical for widespread adoption. In this work, we propose a theoretically grounded alternative that eliminates the need for iterative MCMC. We derive a novel formulation showing that the global power distribution can be approximated by a token-level scaled low-temperature one, where the scaling factor captures future trajectory quality. Leveraging this insight, we introduce a training-free and verifier-free algorithm that sharpens the base model's generative distribution autoregressively. Empirically, we evaluate our method on math, QA, and code tasks across four LLMs, and show that our method matches or surpasses one-shot GRPO without relying on any external rewards, while reducing inference latency by over 10x compared to MCMC-based sampling.

1. Introduction

The paper argues that RL gains can reflect distribution sharpening already latent in base models, but existing power-distribution sampling is too costly for scalable use. It therefore derives a direct autoregressive approximation and evaluates it across math, code, and QA tasks.

  • RL post-training may primarily sharpen probability mass over reasoning trajectories rather than introduce fundamentally new capabilities.
  • MCMC sampling from the power distribution can recover reasoning improvements without external rewards, but its iterative forward passes create substantial latency.
  • The proposed formulation approximates the global power distribution with a suitably scaled low-temperature policy, identifying a scaling factor and bounding approximation error.
  • The training- and verifier-free algorithm applies this scaling during standard autoregressive generation, eliminating iterative MCMC overhead.
  • Over 10× lower inference latency than MCMC is reported while preserving the reasoning improvements of the global power distribution.
  • Across four LLMs and MATH500, HumanEval, and GPQA, the method matches or surpasses one-shot GRPO without parameter updates or external rewards.

2. Power & Low-Temperature Distributions

The power distribution globally reweights complete trajectories, making it more expressive but harder to sample than local low-temperature distributions. The paper motivates an efficient approximation that retains future-aware reweighting while avoiding iterative MCMC.

  • Low-temperature sampling transforms current token probabilities using only the generated history, whereas power sampling requires future-completion marginalisation.
  • Direct power-distribution sampling is impractical because summing over future sequences grows as |V|^(T−t).
  • Power sampling is globally trajectory-aware, while low-temperature sampling is local; the two distributions generally differ at every token position.
  • Prior work uses Metropolis–Hastings MCMC over token sequences, including prefix-level variants with bounded MCMC steps during autoregressive generation.
  • Approximately 8.8× higher cost than standard sampling is reported for MCMC on MATH500 sequences.
  • The proposed approximation connects low-temperature sampling and the power distribution while preserving future-aware reweighting without iterative MCMC loops.

3. Approximating Power Distributions

The paper relates the global power distribution to low-temperature sampling through a token-specific, future-aware scaling factor, enabling autoregressive approximation. It then reduces finite-sample bias with jackknife correction and implements scalable candidate-based trajectory sampling.

  • Power-distribution decomposition: Theorem 3.1 expresses power sampling as low-temperature sampling augmented by a token-specific scaling factor based on expected future-completion likelihood.This future-aware correction reweights locally probable tokens according to the downstream sequences they induce.
  • Power-distribution decomposition: The scaling factor can be estimated with autoregressive Monte Carlo rollouts, avoiding trajectory-level MCMC for approximate power-distribution sampling.The resulting procedure uses standard autoregressive generation to estimate future-completion quality.
  • Bias analysis and correction: The plug-in Monte Carlo estimator is biased because substituting estimated scaling factors creates a ratio of expectations rather than an expectation of ratios.The bias arises when the estimated scaling factors are substituted back into the power-distribution decomposition.
  • Bias analysis and correction: Jackknife correction removes the leading-order bias term, yielding inverse-quadratic bias decay in the rollout budget Mt.The correction combines the original estimator with leave-one-out variants.
  • Bias analysis and correction: O(1/√ϵ) samples achieve a target bias ϵ with jackknife correction, compared with O(1/ϵ) samples for the original Monte Carlo estimator.This reduces the number of language-model generations required for a given accuracy.
  • Scalable power sampling: At each autoregressive step, the algorithm evaluates Top-Kt candidate tokens with Mt independent rollouts, applies bias-corrected scaling, and samples the next token.Targeted compute allocation focuses rollout budgets at critical decision points rather than generating complete trajectories indiscriminately.

4. Experiments & Results

The experiments show that rollout-based power sampling captures future trajectory quality more effectively than local sharpening, while preserving strong performance with substantially lower latency than MCMC.

  • Building Intuition: Planning vs. Guessing: The toy example contrasts a locally attractive GUESS shortcut with the higher-quality PLAN → CALC trajectory.Although GUESS has higher initial probability, PLAN leads to correct answers with higher probability.
  • Building Intuition: Planning vs. Guessing: Low-temperature decoding amplifies locally probable shortcuts, whereas power sampling evaluates complete trajectories and favors the structured PLAN branch.The distinction is whether sharpening accounts for downstream consequences.
  • Large-Scale Evaluation: The evaluation covers four models and mathematics, code, and knowledge-intensive QA benchmarks.The models include Qwen2.5-7B, Qwen2.5-Math-7B, DeepSeek-Math-7B, and DeepSeek-Math-7B-RL; benchmarks are MATH500, HumanEval, and GPQA.
  • Large-Scale Evaluation: Power-sampling methods achieve the strongest sampling performance, with the proposed estimator matching or exceeding MCMC by up to +13.4%.The comparison uses equal sampling budgets and evaluates pass@1 across the three benchmarks.
  • Large-Scale Evaluation: Our method consistently outperforms MCMC in speed, reaching up to approximately 10× speedup across models and benchmarks.On Qwen2.5-Math-7B and MATH500, average per-prompt time falls from 2.5 minutes for MCMC to 0.22 minutes.
  • Diversity Analysis & Ablation: For pass@K, GRPO improves pass@1 but its gains diminish as K increases, while the proposed method preserves diversity and is stronger at small K.The method closely tracks base-model performance for K≥8 and is stronger for K∈{1, 2, 4}.
  • Diversity Analysis & Ablation: Performance is robust to candidate and rollout budgets but sensitive to α, with intermediate values such as α∈{4, 5} strongest.The experiments fix α=4 and Kt=Mt=8.

5. Related Works

Prior work frames RL gains as distribution sharpening and uses power-distribution sampling to reproduce them, but MCMC implementation imposes substantial computational cost.

  • Reinforcement Learning and Distribution Sharpening: RL post-training methods optimize LLMs against automated verifiers to improve reasoning performance.Examples include rejection-sampling fine-tuning, GRPO, and related variants.
  • Inference-Time Reasoning and Sampling: Power Distribution Sampling uses the globally sharpened distribution pα to match or exceed RL-trained model performance without additional training.The approach provides the conceptual basis for training-free reasoning improvement.
  • Inference-Time Reasoning and Sampling: MCMC power sampling is effective but suffers high latency from iterative resampling and slow mixing in high-dimensional token spaces.The proposed method targets the same theoretical distribution without this iterative implementation burden.

6. Conclusions and Future Work

The paper concludes that distribution sharpening can expose latent reasoning capabilities in base models through a scalable, training-free sampler, while identifying several efficiency extensions for future work.

  • Conclusions: The training-free sampler bridges global power distributions and local policies to match GRPO on MATH500, HumanEval, and GPQA.The paper presents this result as evidence that latent reasoning capabilities are accessible through distribution sharpening.
  • Conclusions: The method achieves over 10× lower latency than MCMC, offering a scalable alternative to post-training pipelines.This efficiency result is the paper’s main practical consequence.
  • Future Work: Future work targets variance reduction, adaptive compute budgets, rollout-cost amortization, and extension to agentic settings.The listed directions include control variates, adaptive budgets, speculative decoding, and agentic applications.

Impact Statement

The work presents inference-time sampling as a way to make high-performance reasoning more accessible while reducing some costs associated with training and MCMC. It also highlights a safety boundary: sharpening can amplify undesirable behaviors already present in the base model.

  • Democratisation of Advanced Reasoning: Inference-time sampling can let base models match RL-finetuned models without requiring large-scale post-training resources.The authors frame this as lowering the barrier for academic researchers and smaller organisations without large GPU clusters.
  • Environmental Considerations: The autoregressive method is slower than standard decoding but reduces inference energy consumption by an order of magnitude compared to MCMC-based baselines.The paper attributes this trade-off to the added lookahead mechanism and the elimination of iterative RL fine-tuning.
  • Safety and Dual-Use Risks: The method relies on the base model’s safety alignment because it sharpens latent probability mass without external rewards or human feedback.The authors recommend applying it to base models that have undergone rigorous safety filtering or instruction tuning.

A.1. Proof of Theorem 3.1

This appendix develops the theorem and the batch implementation behind scalable power sampling. It describes candidate generation, rollout-based scaling-factor estimation, jackknife correction, and trajectory construction, while also stating concentration and complexity results.

  • A.1. Proof of Theorem 3.1: The theorem is formulated for a pretrained LLM p, prompt q, exponent α > 1, and any partially generated sequence.The proof fixes a generation position and conditions on the existing prefix to relate token-level and sequence-level quantities.
  • A.1. Proof of Theorem 3.1: The proof proceeds by conditioning on a fixed token and separately handling the final-token case before completing the theorem argument.The supplied proof passages introduce the fixed token, future-token notation, conditional probability, and last-token case in sequence.
  • A.1. Proof of Theorem 3.1: The appendix analyzes Monte Carlo and jackknife estimators for token-level power-distribution probabilities, including approximation-error terms and scaling-factor concentration bounds.The estimator passages state an O(1/M_t^2) term for the jackknife expression and introduce precision and confidence parameters for sample complexity.
  • B. Batch Tokens Extension: The batch algorithm generates candidate blocks from the base model, retains promising candidates, estimates future scaling factors with rollouts, and samples the next block from the corrected distribution.The procedure then concatenates sampled blocks and constructs the remaining trajectory using candidate completions.
  • B. Batch Tokens Extension: The implementation supports parallel future rollouts and has inference complexity expressed as a sum over stages involving L_t, K_tM_t, and L_w.The appendix also gives a token-generation complexity formulation for the original algorithm.

C.1. Experimental Setups and Implementation Details

The experiments use fixed prompts, decoding settings, and runtime conventions to compare sampling methods fairly across models and benchmarks. The proposed method is implemented as a sampling procedure with lookahead and jackknife correction rather than backpropagation.

  • Experimental setup: The method runs without backpropagation and is comparatively lightweight in memory, with memory usage depending primarily on model size.All experiments use a single GPU with vLLM v0.6.3.
  • Experimental setup: The evaluation fixes prompt templates, maximum generation length, early EOS termination, random seeds, and scoring functions across methods.Qwen prompts follow the MCMC repository, while DeepSeek prompts use the official instruction format.
  • Baseline settings: Low-temperature sampling uses temperature 1/α with α = 4, while verifier-free Best-of-N generates 32 completions and selects the highest base-model sequence score.These settings define the principal non-power-sampling baselines.
  • Baseline settings: MCMC uses α = 4, maximum length 3072, block size B = 192, and 10 MCMC steps per stage with proposal temperature 0.25.The configuration follows the open-source implementation of Karan & Du (2025).
  • Proposed method: The proposed method uses α = 4, T_max = 3072, B = 192, K_t = M_t = 8, and a jackknife-corrected estimator with token-level lookahead.Its likelihood computations are performed in log-space for numerical stability.
  • Inference measurement: The method is typically 2.5–3.5× slower than standard decoding under the reported single-GPU measurement protocol.The measurements average wall-clock generation time per prompt over each evaluation set with matched batching configurations.

C.2. More Pass@k Results and Analysis

Across out-of-domain HumanEval and GPQA, power sampling improves pass@1 while preserving stronger pass@k scaling than GRPO in the settings where GRPO’s diversity declines. The qualitative cases associate GRPO failures with repeated reasoning or answer-format behavior rather than only missing capability.

  • Pass@K on out-of-domain tasks: On HumanEval, GRPO has pass@k scaling similar to the base model, whereas on GPQA it improves pass@1 but gains less as k increases.The authors interpret the GPQA pattern as consistent with reduced diversity in an out-of-domain task.
  • Power sampling comparison: Power sampling improves pass@1 while both MCMC and the proposed method retain strong pass@k growth on out-of-domain tasks.The result supports inference-time sharpening that raises high-quality sampling probability without the same degradation in scaling.
  • Qualitative examples: GRPO failure cases often repeat near-identical reasoning across 16 samples, limiting pass@16 despite the possibility that the underlying capability is present.The paper contrasts this reduced diversity with broader candidate sets from related base models.
  • Figure 5: Figure 5 compares pass@k for k ∈ {1, 2, 4, 8, 16} on HumanEval and GPQA, showing stronger scaling for power sampling than GRPO on GPQA.On HumanEval, GRPO does not show an obvious diversity collapse.
  • Qualitative examples: On GPQA, GRPO-MATH can produce correct reasoning but select the wrong option letter or stop before giving a final multiple-choice answer.The paper notes that this behavior is less frequent in the corresponding base model and may reflect post-training on the MATH split.

C.3. Ablation Study on Different Models

Across Qwen2.5-7B and DeepSeek-Math-7B, performance is relatively robust to sampling budgets, with the strongest results typically at moderate sharpening exponents. Figure 6 evaluates pass@1 across MATH, HumanEval, and GPQA while varying α and Mt = Kt.

  • Budget sensitivity: Varying Mt = Kt over the tested range only mildly affects pass@1 under the same sharpening exponent α.The ablation indicates that most gains can be obtained with relatively small sampling budgets.
  • Sharpening sensitivity: The strongest region is typically α ≈ 4–5 across the evaluated base models and tasks.This pattern is reported for math, code, and QA validation subsets.
  • Sharpening sensitivity: Overly aggressive sharpening such as α = 8 can degrade performance, while α = 1 performs poorly on most tasks.DeepSeek-Math-7B on GPQA is a notable exception, where α = 1 can be competitive in parts of the budget range.
  • Overall finding: The ablations suggest improved pass@1 across diverse base models without post-training using a small and stable hyperparameter range.The reported trends are consistent across the evaluated models and tasks.
  • Experimental setup: Figure 6 varies α from 1 to 8 and Mt = Kt from 4 to 16, reporting pass@1 for MATH, HumanEval, and GPQA on two models.Qwen2.5-7B appears in the top panels and DeepSeek-Math-7B in the bottom panels.
Loading 2601.21590v1…