PaperlayerView exact revision on arXiv

Source-linked AI summary

PreSTO: Predictive Subtree Prefetching for Fast LLM Power Sampling

Nan Jiang, Panagiotis Theodoropoulos, Wei Deng

arXiv:2610.13171v1cs.CL

TL;DR

PowerMH is costly because each MH proposal depends on the preceding accept/reject outcome. PreSTO batches proposals for future states and follows the realized path with the original acceptance rule. Across five base LLMs and seven datasets, it is never slower than baseline and achieves median speedups of 1.71×, 1.91×, and 2.33× over PowerMH, EntropyCut, and MultiTryMH, respectively.

  • Problem

    PowerMH’s sequential dependence between MH transitions makes sharpened-distribution sampling costly, because each next proposal must wait for the previous accept/reject outcome.

    Source: §1 · p. 2

  • Method

    PreSTO generates feasible proposals for a prefetched subtree in one batched call, then follows the realized path with the original MH acceptance rule.

    Source: §1 · p. 2

  • Results

    Across five base LLMs and seven datasets, median speedups are 1.71× over PowerMH, 1.91× over EntropyCut, and 2.33× over MultiTryMH; PreSTO is never slower than baseline.

    Source: §6 · p. 12

  • Takeaways & Limitations

    PreSTO trades fewer serial calls for costlier batched calls, with latency falling when saved-call time outweighs added batch cost.

    Source: §1 · p. 2

    A proposal for an acceptance child is feasible in the current batch only when its cut index c′ is no greater than its parent’s cut index c.

    Source: §3.2 · p. 5

Abstract

from arXiv · show

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.

1. Introduction

Power sampling sharpens an LLM’s distribution for test-time reasoning, but PowerMH’s sequential accept/reject updates make sampling costly. PreSTO batches future proposals and follows the original MH rule, with median speedups across three samplers and 105 settings.

1. Introduction

  • Across 105 settings, PreSTO is never slower and has median speedups of 1.71× over PowerMH, 1.91× over EntropyCut, and 2.33× over MultiTryMH.The figure plots one sampler–LLM–dataset setting per point; faster settings lie toward the upper right.Source: Figure 1 · p. 1, §1 · p. 2
  • PowerMH’s bottleneck is that each proposal must await the prior accept/reject outcome, which determines the state for the next proposal.Prefetching makes proposals for future transitions available in advance, allowing several transitions to complete before another model call.Source: §1 · p. 2
  • PreSTO batches feasible proposals in a prefetched subtree, then follows the realized path using the original MH acceptance rule.For PowerMH, this leaves the MH transition kernel unchanged.Source: §1 · p. 2
  • Because MH reaches a sparse, nearly unary part of the prefetched tree, predictive traversal allocates the budget to maximize expected completed transitions per call.The method trades fewer serial calls for costlier batched calls; latency falls when saved-call time outweighs added batch cost.Source: §1 · p. 2
  • Across seven datasets and five base LLMs, the experiments compare PreSTO on PowerMH and its follow-ups EntropyCut and MultiTryMH.Power sampling itself samples directly from a sharpened base-model distribution and has been reported as competitive with RL post-training without additional training.Source: §1 · pp. 1–2

2. Preliminaries

Metropolis–Hastings samples from a target distribution through proposal-and-acceptance steps. Power sharpening raises the base LLM distribution to an exponent, while PowerMH approaches that target through blockwise intermediate distributions.

2. Preliminaries

  • In MH, a candidate is drawn from a proposal distribution given the current state and accepted with a probability; rejection retains the current state.The chain targets a distribution known only up to a normalizing constant.Source: §2 · pp. 2–3
  • Power sampling raises the base LLM distribution to an exponent α; when α > 1, it amplifies high-likelihood continuations and suppresses low-likelihood ones.Prior empirical results report performance competitive with RL-based post-training methods.Source: §2 · p. 3
  • Because exact sampling from the power-sharpened target is intractable, PowerMH proceeds through blockwise intermediate targets, ending at the full target.At each stage it extends the previous sample by one block, then refines it with K MH steps.Source: §2 · p. 3

3. Methodology

PreSTO batches feasible future MH proposals in a subtree, then follows the realized path using the original acceptance rule. Its predictive traversal uses the sparse, nearly unary structure of reached decisions to allocate the prefetch budget toward more completed transitions.

3.1. Motivation

  • PowerMH’s sequential dependence requires each proposal to wait for the preceding accept/reject decision, while prefetching represents future outcomes as a tree of possible paths.Each realized run follows one root-to-leaf path through the tree.Source: §3.1 · pp. 3–4

3.2. Subtree Prefetching for MH-Based Power Sharpening

  • PreSTO builds a subtree of possible future MH states, collects feasible suffix-resampling requests, and selects a root-to-leaf path using the original acceptance rule.Requests are batched when the preserved prefix makes them feasible; a single batched call returns proposals and log-probabilities for scoring.Source: §3.1 · p. 4, §3.2 · pp. 4–6
  • PreSTO-PowerMH preserves PowerMH’s MH transition kernel while batching proposals and completing one or more transitions per prefetch round.The prefetch tree is capped by the number of transitions remaining in the block.Source: §3.2 · p. 6

3.3. Predictive Subtree Allocation for Efficient Prefetching

  • The 𝜖-rounded subtree is predominantly unary, with about 95% of decision edges 𝜖-certain on LCB V6 at 𝜖=0.01.Rounding each 𝜖-certain decision to acceptance or rejection leaves only one reachable child.Source: §3.3 · p. 7, Figure 3 · p. 7
  • The predictive traversal selects up to Npf requests to maximize predicted completed MH transitions, using cut fractions to predict decision-certainty probabilities.Dynamic programming performs the selection over acceptance and rejection subtrees.Source: §3.3 · p. 7

4. Related Work

PreSTO uses a budgeted, irregular subtree to prefetch future MH proposals, unlike earlier MCMC approaches that prefetch many candidates or use regular trees. It accelerates PowerMH itself, rather than changing its target, proposal, or acceptance rule.

MCMC prefetching

  • PreSTO budgets speculative LLM proposal work in an irregular subtree, unlike prior MCMC methods that prefetch many candidates or expand fixed regular trees.Its feasible suffix requests are prioritized by predicted reach probability.Source: §4 · p. 8

Tree-based speculative decoding

  • TreeSD batches alternative draft-token continuations, whereas PreSTO branches on complete-sequence MH accept/reject outcomes and advances one MH transition per expanded node.Following PreSTO’s realized decisions reproduces the sequential PowerMH trajectory.Source: §4 · p. 8

Power-sharpening methods

  • Unlike methods that change samplers, approximate the target, refine proposals, or alter training, PreSTO accelerates PowerMH while retaining its target, suffix proposal, and acceptance rule.Source: §4 · p. 8

5. Experiments

The experiments compare PreSTO with sequential samplers across base-LLM and dataset settings, then examine speedup sources and case studies of prefetch budgets, output equivalence, and traversal rules.

5.1. Acceleration Across Base LLMs and Datasets

  • Across 105 settings, PreSTO was never slower and achieved median speedups of 1.71× over PowerMH, 1.91× over EntropyCut, and 2.33× over MultiTryMH.Table 1 reports speedups as mean cumulative model-call time for each sequential sampler divided by that of its PreSTO variant.Source: §5.1 · p. 8, Table 1 · p. 9

5.2. Case Studies

  • On LCB V6 with Qwen3.5-9B, PreSTO-PowerMH was 1.44× faster, completing 100 transitions in 375.12±35.08 rather than 539.24±30.92 seconds.It realized 4.96±0.47 transitions per call versus one, while each call took 18.45±0.51 rather than 5.39±0.31 seconds.Source: §5.2 · p. 9, Figure 5 · p. 10

5.1. Acceleration Across Base LLMs and Datasets

  • Call reduction accounts for a 95% share of PowerMH speedup variance, while time per call accounts for 94% and 99% for EntropyCut and MultiTryMH.The respective R2 values are 0.78, 0.75, and 0.99.Source: §5.2 · p. 10

5.2. Case Studies

  • Raising the prefetch budget from 8 to 20 increased transitions per call from 3.59 to 4.91 and speedup from 1.20× to 1.44×.Calls also slowed from 16.2 to 18.4 seconds, versus 5.4 seconds for a PowerMH step.Source: §5.2 · p. 10, Table 2 · p. 10
  • Across PowerMH and EntropyCut, 29 of 38 tests established equivalence, and pooled log-likelihood and confidence differences lay within the ±0.2-SD margin.For PowerMH with Qwen3.5-4B, equivalence held in 7 of 10 comparisons; the other three intervals extended past the margin.Source: §5.2 · p. 11, Figure 6 · p. 11
  • The predictive traversal rule was 1.25×–1.37× faster than BFS at every budget and faster on 98 of 100 prompt–budget pairs.It realized 19%–33% more transitions per call, reducing calls by 1.19×–1.33×.Source: §5.2 · p. 11, Figure 7 · p. 12

6. Conclusion

Power-sampling methods differ in whether they change the sampler, approximate the target, or move inference cost into training. The section contrasts these routes with prefetching for sequential MH and derives the suffix-resampling acceptance calculation.

6. Conclusion

  • PowerMH’s sequential transitions require each proposal to await the preceding accept/reject decision, motivating ways to batch future work.The conclusion frames the paper’s focus as accelerating MH-based power sampling, whose sequential decisions force one model call per transition.Source: §6 · p. 12, §A.1 · p. 19

A.1. Sampling from the Power Distribution

  • The exact power-target conditional includes a continuation term that token-level temperature scaling drops, so power-sampling methods must avoid, approximate, or change it.For α > 1, the continuation term favors prefixes with a few high-likelihood completions; computing it directly sums over exponentially many continuations.Source: §A.1 · p. 19
  • PowerMH uses suffix regeneration with a blockwise schedule; EntropyCut changes cut placement, while MultiTryMH considers several suffix candidates per transition.These methods share serial transitions because each proposal depends on the current state.Source: §A.1 · p. 19

A.2. Tree-based Speculative Decoding

  • Tree-based speculative decoding verifies alternative draft-token continuations together, committing an accepted prefix while spending more computation on candidate branches.Under a fixed node budget, tree construction allocates branches toward nodes expected to contribute most to the committed prefix.Source: §A.2 · p. 20

A.3. MCMC Prefetching

  • MCMC prefetching evaluates possible future MH proposals in parallel, then follows the realized accept/reject path with the usual rule to preserve the sequential chain’s law.PreSTO’s tree branches encode complete-sequence MH accept/reject outcomes, unlike TreeSD’s alternative token continuations.Source: §A.3 · p. 21, §A.2 · p. 21

B.1. Computing Log-Acceptance Probabilities

  • The suffix-resampling log-acceptance expression combines proposal and target log-ratios, computed from token-level conditional log-probabilities in the log domain.Log-domain evaluation avoids underflow from multiplying many small token probabilities.Source: §B.1 · pp. 22–23

B.2. Exactness of Subtree Prefetching

PreSTO-PowerMH batches feasible suffix proposals while preserving the sequential PowerMH trajectory under stated random-input and acceptance-rule assumptions. The appendices define the prefetch tree and develop an expected-transition objective with predicted branch probabilities.

B.2. Exactness of Subtree Prefetching

  • A coupling with shared per-decision random inputs shows that PreSTO-PowerMH and PowerMH have the same states and transition kernel after each committed transition.The result assumes feasible proposals use the required prefixes, preserve PowerMH input laws, and apply the exact MH acceptance rule.Source: §B.2 · pp. 23–24

B.3. Examples of Prefix Sharing and MH Transitions through the Prefetched Subtree

  • In the illustrated example, shared cached prefixes support two batched suffix-resampling requests, realizing two MH transitions in one call instead of two serial calls.The acceptance probabilities label tree edges, and the realized path accepts one proposal and rejects the next.Source: §B.3 · pp. 24–25, Figure 8 · p. 25

C.1. Setup and Notation

  • A selection is an ancestor-closed set of expandable nodes; expanding each selected node issues one proposal request and retains both outcome children.The resulting prefetched subtree has 2|I| + 1 nodes for a selection I.Source: §C.1 · p. 26, §C.2 · p. 28

C.2. The Transition-Count Objective

  • The allocation objective counts expected completed transitions, with each reached expanded node contributing one whether its proposal is accepted or rejected.The expected count can be expressed as a sum of reaching weights over selected nodes.Source: §C.2 · p. 27
  • The expected count recurses as one transition at an expanded node plus branch-probability-weighted counts from its acceptance and rejection children.This recursive form supports combining subtree selections in the dynamic program.Source: §C.2 · pp. 27–28

C.3. Predictive Subtree Allocation

  • Before proposals are generated, a classifier predicts certain-reject, certain-accept, and uncertain classes to estimate branch probabilities for subtree allocation.Rounding uses the midpoint for uncertain decisions, with bias bounded per branch probability; classifier error remains on every decision.Source: §C.3 · pp. 28–29

C.4. Optimal Budget Allocation via Dynamic Programming

The dynamic program allocates a request budget across a feasible prefetch tree to maximize predicted completed transitions. Its recurrence expands a node and splits the remaining budget between its two child subtrees.

C.4. Optimal Budget Allocation via Dynamic Programming

  • A feasible selection below a node must use candidate nodes and include every candidate ancestor on the path to each selected node.This ancestor-closure condition ensures the selected requests form an allowable prefetch subtree.Source: §C.4 · p. 30
  • The dynamic program’s budget split assigns one request to expand the current node, then divides the rest between acceptance and rejection subtrees before the outcome is known.The optimal subtrees below both children attain the recurrence’s maximum within the remaining budget.Source: §C.4 · p. 30
  • For a node and request budget, the recurrence adds one transition and maximizes the predicted child-subtree counts over all splits of the remaining requests.The predicted count at the root with budget N_pf gives the optimum of the allocation objective.Source: §C.4 · p. 30, Figure 10 · p. 30

D.1. Baselines, Datasets, and Base LLMs

The evaluation compares PreSTO variants of three MH-based samplers across seven datasets and five base LLMs. Each dataset run uses only its first 20 prompts.

D.1. Baselines, Datasets, and Base LLMs

  • The baselines are PowerMH, EntropyCut, and MultiTryMH, with the latter evaluated using random cut positions rather than entropy-guided selection.PowerMH resamples suffixes from uniformly drawn cuts; EntropyCut selects cuts at entropy increases; MultiTryMH considers several candidate suffixes per transition.Source: §D.1 · p. 32
  • The seven datasets cover mathematical reasoning, code generation, and scientific or general knowledge; runs use only the first 20 prompts from each dataset.The suite includes MATH500, AIME 24&25, HumanEval, MBPP, LiveCodeBench V6, GPQA Diamond, and MMLU.Source: §D.1 · p. 32, Table 3 · p. 32
  • Five base LLMs span 4B to 12B parameters, comprising Qwen3 and Qwen3.5 models at two scales and Gemma4-12B-it.None is a pretrained-only checkpoint; the Qwen releases are not the -Base variants, and Gemma4-12B-it is instruction-tuned.Source: §D.1 · pp. 32–33, Table 4 · p. 32

D.2. Evaluation Metrics

The section defines execution, hardware-resource, and sample-quality metrics, describes paired equivalence testing, and outlines the implementation and default sampling settings.

D.2. Evaluation Metrics

  • Execution efficiency compares cumulative model-call time with the sequential sampler; speedup decomposes into fewer calls versus longer calls.Call time is the number of model calls multiplied by mean call duration.Source: §D.2 · p. 33
  • Sample quality compares length-normalized base-model likelihood and confidence after the same number of MH transitions.Confidence is mean negative entropy of the base model’s next-token distributions along the sequence.Source: §D.2 · p. 33
  • Paired TOST tests matched prompts and seeds at equal transition counts, declaring equivalence when peq < 0.05 and the 90% confidence interval lies within (−δ, δ).The margin δ is 0.2 standard deviations of the statistic across prompts under the sequential sampler.Source: §D.2 · pp. 33–34
  • Hardware-resource efficiency tracks peak KV-pool occupancy, prefix-cache hit rate, and uncached prefix tokens per realized transition.The hit rate measures the fraction of queried prefix tokens already cached.Source: §D.2 · p. 34

D.3. Implementation Environment

  • The implementation uses vLLM with custom engine patches, and SGLang for PowerMH and PreSTO-PowerMH in Case Study 6.The vLLM patch returns proposal and base-model log-probabilities, plus base-model entropy where needed.Source: §D.3 · p. 34

D.4. Hyperparameter Configurations

  • Proposal temperature is set to τ = 1/α except for MultiTryMH and PreSTO-MultiTryMH, which use candidate temperatures listed in Table 6(a).All runs use seed 10086.Source: §D.4 · p. 34

E. Additional Experimental Results

The appendix lists software and hardware configurations, hyperparameter settings, and a case-study configuration.

E. Additional Experimental Results

  • Table 5 summarizes the experimental software and hardware environment.The reported hardware includes one NVIDIA H100 96 GB GPU and a 72-core NVIDIA Grace CPU.Source: Table 5 · p. 35, E. Additional Experimental Results · p. 35
  • Table 6 is identified as presenting hyperparameter configurations; the cited passage provides only that table description and states no specific values or distinction between default and case-study settings.Source: Table 6 · p. 35
  • A separate case-study configuration is identified in the appendix.The supplied passage labels it as the case-study configuration without giving its settings.Source: E. Additional Experimental Results · p. 35

E.1. Full results: Acceleration across base LLMs and datasets

The full-results section reports execution metrics for three samplers across base LLMs and datasets, and includes MATH500 comparisons for two sampler pairs.

E.1. Full results: Acceleration across base LLMs and datasets

  • Across five base LLMs and seven datasets, Tables 7–9 report means and standard deviations of metrics for PowerMH, EntropyCut, and MultiTryMH.The experiments pair each sampler with its PreSTO variant on all five base LLMs.Source: §E.1 · p. 35, Table 7 · p. 36, Table 8 · p. 38
  • Tables 7, 8, and 9 respectively cover PowerMH, EntropyCut, and MultiTryMH execution metrics across base LLMs.The supplied table captions establish the scope, but not individual metric values.Source: Table 7 · p. 36, Table 8 · p. 38, Table 9 · p. 40
  • On MATH500, Figure 11 compares PowerMH with PreSTO-PowerMH, while Figure 12 compares EntropyCut with PreSTO-EntropyCut.The captions identify the comparisons but do not state their outcomes.Source: Figure 11 · p. 37, Figure 12 · p. 39

E.2. Case Study 2: Effect of prefetch budget

On LCB V6 with Qwen3.5-9B, the prefetch-budget sweep shows more realized transitions per call as the budget increases, while call duration also rises.

E.2. Case Study 2: Effect of prefetch budget

  • The comparison uses LCB V6 with Qwen3.5-9B at K = 100 under default breadth-first traversal; no LCB V6 PreSTO-EntropyCut sweep is available.The runs did not record KV-pool use or prefix-cache statistics.Source: §E.2 · p. 41
  • 1.20× to 1.44×: PreSTO-PowerMH speedup over PowerMH rises from Npf = 8 to Npf = 20, while each PreSTO call takes 16.2–18.4 seconds versus 5.4 seconds for a PowerMH step.The largest tested budget is Npf = 20.Source: §E.2 · p. 41
  • 3.59 to 4.91: realized MH transitions per call increase as the prefetch budget grows from Npf = 8 to the default Npf = 20.Figure 13 plots per-transition traces for budgets from Npf = 8 through 20 against PowerMH, with ±1 SD bands.Source: §E.2 · p. 41, Figure 13 · p. 41

E.3. Case Study 3: Likelihood and confidence

Most tests find PreSTO outputs equivalent to their baselines in likelihood and confidence, while most MH decisions are ε-certain under the reported threshold.

E.3. Case Study 3: Likelihood and confidence

  • 29 of 38 tests establish equivalence: 12 of 16 for PowerMH and 17 of 22 for EntropyCut.Figure 14 shows the matched-pair likelihood and confidence comparisons; the tests use a margin of 0.2 baseline standard deviations.Source: §E.3 · p. 41, Figure 14 · p. 42
  • The median absolute mean gap is 0.002 in log-likelihood for both samplers, and 0.007 for PowerMH and 0.004 for EntropyCut in confidence.These gaps are against across-prompt standard deviations of about 0.05 and 0.10, respectively.Source: Figure 14 · p. 42, §E.3 · p. 43

E.4. Case Study 4: Distribution of 𝜖-certainty

  • 56%–95% of edges are ε-certain for Qwen models under PreSTO-PowerMH and 60%–92% under PreSTO-EntropyCut; Gemma4-12B-it ranges are 67%–78% and 60%–93%, respectively.Figures 15 and 16 show edge certainty across models and datasets; ε is 0.01.Source: §E.4 · p. 43, Figure 15 · p. 44, Figure 16 · p. 45

E.5. Case Study 5: Effect of traversal rules

The case studies find that predictive traversal improves average performance, then examine its behavior across serving engines, certainty prediction, memory use, and a realized prefetched trajectory.

E.5. Case Study 5: Effect of traversal rules

  • 1.25×–1.37× faster than BFS in model-call time, the predictive rule is faster on 98 of 100 prompt–budget pairs.It completes 19%–33% more transitions per call; gains average across trees, since it realizes more transitions in 11 of 20 prompt–budget pairs and fewer in 8.Source: §E.5 · p. 46, Table 10 · p. 46

E.6. Case Study 6: Effect of different serving engines

  • 1.97× speedup under SGLang versus 1.40× under vLLM; PreSTO-PowerMH is faster than PowerMH under both engines.The engines reduce calls by similar factors, 4.7× and 4.9×, while a PreSTO call costs 2.4× a PowerMH step under SGLang versus 3.5× under vLLM.Source: §E.6 · p. 48, Table 11 · p. 48

E.7. Case Study 7: Predicting 𝜖-certainty before generation

  • 0.643 for PreSTO-PowerMH and 0.625 for PreSTO-EntropyCut held-out accuracy, above majority-class rates of 0.609 and 0.596.The cut-fraction-only classifier lowers log loss from 0.940 to 0.819 for PowerMH and from 0.954 to 0.884 for EntropyCut.Source: §E.7 · p. 49, Figure 18 · p. 50

E.8. Case Study 8: Comparison with test-time baselines

  • PreSTO-PowerMH peaks at 2.9%–3.1% KV-pool use, versus 0.6% for PowerMH and 3.3% for best-of-8.Its peak use is below the larger parallel baselines, but prefetching narrows rather than closes the cost gap.Source: §E.8 · p. 51, Figure 20 · p. 51

E.9. Example prefetched-subtree trajectory

  • In one LCB V6/Qwen3.5-9B run, 19 BFS subtrees cover 100 MH transitions; the ε-rounded tree reduces to the sampled path after two ε-certain accepts.The example uses a prefetch budget of 20 and ε=0.01.Source: §E.9 · pp. 52–53, Figure 21 · p. 53
Loading 2610.13171v1…