PreSTO: Predictive Subtree Prefetching for Fast LLM Power Sampling

AI summary

Overview

• Research area: Efficient inference and test-time compute for large language models, specifically the acceleration of Markov chain Monte Carlo (MCMC) samplers used for "power sampling"; it sits at the intersection of NLP, LLM serving systems, and Monte Carlo methodology.
• Technical level: Advanced. The problem framing is accessible, but the correctness argument rests on Metropolis–Hastings transition kernels, suffix-resampling proposals, and acceptance-probability algebra. Readers without MCMC background will follow the empirical story more easily than the theory.
• Scope (one sentence): The paper introduces PreSTO, a prefetching framework that batches future Metropolis–Hastings proposals into a tree so that several sequential sampling steps complete per model call, and shows it accelerates PowerMH, EntropyCut, and MultiTryMH across seven datasets and five base LLMs without changing the sampler's target distribution.

Note: the supplied text is truncated partway through Case Study 3, so results reported only in the later appendices are not covered here. Values the paper does not state (for example, the specific power exponent used, proposal temperature, and block size) are marked as not reported.

What This Paper Is About

Power sampling improves LLM reasoning at test time by drawing from a sharpened version of the base model's own distribution, but its sampler, PowerMH, is slow because every accept/reject decision must be resolved before the next proposal can be generated. The paper's goal is to make this sampler faster without altering what it samples: instead of one proposal per model call, PreSTO speculatively generates a whole subtree of possible future proposals in a single batched call, then walks the one path that the original acceptance rule actually realizes.

Key Contributions

1. A prefetching framework for MH-based power sampling (PreSTO). Future accept/reject outcomes are represented as a tree; each round generates the proposals of a budgeted subtree in one batched model call and then follows the realized root-to-leaf path using the original MH acceptance rule. The authors prove (Proposition 1) that PreSTO-PowerMH leaves the MH transition kernel of PowerMH unchanged.
2. A feasibility condition that makes batching valid (Condition 1). Because suffix resampling preserves prefixes, a child request with cut index c′ can be generated in the same batch as its parent with cut index c whenever c′ ≤ c. This lets multiple suffix-resampling requests that share a prefix be issued together, and modern serving-engine prefix caching (vLLM, SGLang) makes the shared KV cache reusable.
3. An empirical characterization of prefetched subtrees, and a predictive traversal rule. The authors observe that under the power-sharpened target, MH reaches only a sparse, nearly unary part of each subtree: 67%–95% of edges are ϵ-certain at ϵ = 0.01 across all seven datasets (about 95% on LCB V6). They replace default breadth-first traversal with a rule that uses a three-class logistic regression on the cut fraction c_j/T to predict branch probabilities and then selects the at-most-N_pf requests that maximize the predicted number of completed transitions, solved by dynamic programming.
4. A broad empirical evaluation. Speedups are measured over 105 sampler–LLM–dataset settings (three sequential samplers, five base LLMs, seven datasets), plus case studies isolating call-count reduction, budget sensitivity, output distribution preservation, and traversal-rule impact.

Main Findings

• Consistent speedups, never slower: Across the 105 settings, PreSTO is never slower than its baseline and reaches at least 1.5× in 96 of them. Median speedups are 1.71× over PowerMH, 1.91× over EntropyCut, and 2.33× over MultiTryMH, with ranges of 1.00×–2.29×, 1.19×–2.85×, and 1.91×–7.01× respectively.
• Fewer calls, not cheaper ones: Over PowerMH and EntropyCut, PreSTO makes a median of 3.9× fewer model calls, each only 2.1× longer. Across settings the call reduction is 1.94×–4.92× while time per call rises 1.90×–3.42×.
• What drives the speedup depends on the sampler: A variance decomposition of log speedup shows that, over PowerMH, the reduction in calls takes a 95% share (R² = 0.78) because time per call rises by a nearly constant 1.90×–2.42×. Over EntropyCut and MultiTryMH, the time-per-call ratio (1.37×–2.74× and 0.26×–1.00×) takes 94% and 99% shares (R² = 0.75 and 0.99), because their call reductions vary less (3.11×–4.52× and 1.77×–2.05×).
• Detailed case study (LCB V6, Qwen3.5-9B, N_pf = 20, BFS): PreSTO-PowerMH is 1.44× faster in mean cumulative model-call time, making 4.92× fewer calls (a ratio of means), each 3.42× as long, averaging 4.96 transitions per call. It completes K = 100 transitions in 375.12 ± 35.08 s versus 539.24 ± 30.92 s; a PreSTO call takes 18.45 ± 0.51 s rather than 5.39 ± 0.31 s. PreSTO-EntropyCut on the same setting is 1.91× faster (4.22 ± 0.58 transitions per call at 10.08 ± 0.47 s rather than 4.64 ± 0.52 s; 4.15× fewer calls, each 2.17× as long).
• Prefetch budget trade-off: Sweeping N_pf from 8 to 20 on LCB V6 with Qwen3.5-9B raises transitions per call from 3.59 to 4.91 and speedup from 1.20× to 1.44×, while per-call time grows from 16.15 s to 18.43 s against 5.39 s for a PowerMH step. Every prefetched proposal is generated whether or not it is realized, so speedup grows more slowly than transitions per call.
• Gains are largest when the sequential step is expensive: On MBPP with Qwen3.5-9B, an EntropyCut step (7.62 s) is slower than a PowerMH step (4.97 s), so a PreSTO-EntropyCut call costs only 1.45× a step and speedup reaches 2.85×, versus 1.61× for PreSTO-PowerMH.
• Gains are smallest when few proposals are realized: Of the five settings below 1.25×, four are PreSTO-PowerMH (MBPP with both Qwen3 models, and HumanEval and MMLU with Qwen3.5-9B), realizing only 1.96–2.72 transitions per call at about 2× the cost of a PowerMH step; PreSTO-EntropyCut on MBPP with Qwen3-4B realizes 3.51 per call at 2.7× the cost of an EntropyCut step.
• Output distributions are largely preserved: Paired TOST equivalence tests at a margin of 0.2 baseline standard deviations on terminal-trace log-likelihood and confidence pass in 29 of 38 comparisons. With Qwen3.5-4B, the two distributions overlap closely on five datasets and equivalence holds in 7 of the 10 comparisons; the three failures are both statistics on AIME and confidence on MATH500, where the 90% intervals extend past the margin.
• The predictive traversal rule beats BFS: On LCB V6 with Qwen3.5-9B it is 1.25×–1.37× faster than the default breadth-first traversal at all five prefetch budgets tested.
• A caveat on MultiTryMH: Part of the MultiTryMH speedup (which reaches 7.01× on Qwen3.5-9B MATH500) comes from the implementation rather than from prefetching, since sequential MultiTryMH rescores every candidate at all three temperatures while PreSTO-MultiTryMH reuses the sampled log-probabilities.
• Wasted work under BFS: Because BFS enqueues rejection children even when a node's own request is deferred, requests placed below a deferred node are generated and scored but never reached; at N_pf = 10 on LCB V6 with Qwen3.5-9B these are 19% of issued requests.

Methodology in Plain English

The starting point is standard Metropolis–Hastings: propose a change to the current text, then accept or reject it with a probability that depends on how much more (or less) likely the proposal is under the sharpened target. PowerMH applies this block by block, resampling the suffix of the sequence after a randomly chosen cut point and refining it with K MH steps per block.

The bottleneck is that step k+1 cannot start until step k's coin flip is resolved. The authors note that all possible futures over the next K steps form a binary tree: each node branches left if the proposal is accepted and right if it is rejected. A normal run traces exactly one root-to-leaf path, but if the proposals on that path were already available, several steps would complete from a single model call.

Two ideas make this practical. First, a feasibility rule: if a parent proposal cuts at position c and a child proposal cuts at position c′ ≤ c, the child's prefix is untouched by the parent's resampling, so the child's suffix can be drawn from the same already-known prefix in the same batch. Requests that cannot yet be issued are deferred with their cut indices retained for later batches. Second, batched serving: the collected requests go to the inference engine as one batch, and because many share long prefixes, prefix caching avoids recomputing the same KV entries.

After generation, the sampler simply walks the tree using the original acceptance probabilities until the block's K transitions are exhausted, and the reached leaf becomes the new state. No approximation is introduced into the kernel.

The refinement concerns which subtree to spend the budget on. Under a sharpened target, accept/reject probabilities are usually pushed close to 0 or 1, so most branches are "ϵ-certain" and effectively unary — the chain almost always takes the same child. BFS, however, expands nodes regardless of reach probability. The predictive traversal rule fits a three-class logistic regression on a cheap feature available before generation (the cut fraction c_j/T) to predict whether a decision will be a certain accept, a certain reject, or uncertain, converts those into branch probabilities, and uses dynamic programming to pick the at-most-N_pf requests that maximize the expected number of completed transitions.

Why This Matters

Research impact. The paper demonstrates that classical MCMC prefetching — long studied when proposals are cheap — can be retargeted to LLM settings where proposals are expensive, by treating the budget explicitly and allocating it by predicted reach probability. Because Proposition 1 guarantees the transition kernel is unchanged, the acceleration is orthogonal to work that changes the sampler (sequential Monte Carlo, replica exchange), approximates the target, or alters proposals; it could in principle be layered on top. The sparsity observation — that power-sharpened targets make most MH decisions near-deterministic — is a reusable insight for anyone designing speculative or batched inference procedures.

Real-world applications (as evaluated in the paper):

• Competition mathematics and reasoning workloads (MATH500, AIME 24&25), where power sampling is used to improve reasoning at test time without RL post-training.
• Code generation (MBPP, HumanEval, LCB V6), where the largest measured speedups appear.
• Graduate-level science question answering (GPQA) and general knowledge evaluation (MMLU).
• Serving-stack deployments that already run vLLM or SGLang, since the method's economics depend on prefix caching across batched requests.

Industry relevance. Power sampling has been reported as competitive with RL post-training without additional training, which makes inference cost the binding constraint on adoption. The paper quantifies

Authors’ abstract

Power sampling improves large language model (LLM) reasoning at test time by sampling from a sharpened version of the base model's distribution. Its sampler, PowerMH, is slow because its Metropolis--Hastings (MH) updates are sequential: each accept/reject decision sets the state from which the next proposal is generated. Prior work reduces this cost by changing the sampler or the model (e.g., sequential Monte Carlo, approximate lookahead, distillation). We propose Predictive SubTree Prefetching (PreSTO), which accelerates PowerMH by prefetching future proposals and provably leaves its MH transition kernel unchanged. PreSTO views future MH outcomes as a tree of accept/reject decisions: in each model call, it generates the proposals of a \textit{prefetched subtree} of this tree in one batch, then follows the realized path with the original MH acceptance rule. We also observe that MH reaches only a sparse, nearly unary part of each prefetched subtree, we further design a \textit{predictive traversal rule} that chooses the subtree to maximize the predicted number of completed transitions. Across seven datasets and five base LLMs, PreSTO achieves median speedups of $1.71\times$, $1.91\times$, and $2.33\times$ on PowerMH and its follow-ups EntropyCut and MultiTryMH, respectively; on the first two, it makes a median of $3.9\times$ fewer model calls, each only $2.1\times$ longer. On LCB V6 with Qwen3.5-9B, the predictive traversal rule is a further $1.25\times$--$1.37\times$ faster than the default breadth-first traversal.