PreSTO: Predictive Subtree Prefetching for Fast LLM Power Sampling

PreSTO. Predictive Subtree Prefetching for Fast LLM Power Sampling. By Nan Jiang and two others, from the University of Texas at El Paso, the Georgia Institute of Technology, and one other institution. Published on the ninth of October. Its code repository is worked out: about five hundred source files, with a readme.

Power sampling improves the reasoning of a large language model at test time. Instead of drawing from the base model's distribution directly, it samples from a sharpened version of that distribution. The sampler that does this, PowerMH, is slow, and the reason is structural. It works through MH updates, and those updates are sequential. Each accept or reject decision sets the state from which the next proposal is generated, so one step cannot begin until the previous one has been settled. Earlier attempts to reduce the cost have changed something about the setup: the sampler itself, through sequential Monte Carlo or approximate lookahead, or the model, through distillation. Each of those buys speed by altering what is being run.

Predictive SubTree Prefetching, or PreSTO, takes a different line. It accelerates PowerMH by prefetching future proposals, and it provably leaves the MH transition kernel unchanged. The sampler that results is the same sampler, only faster. The way in is to look at the future differently. The outcomes of upcoming steps form a tree of accept and reject decisions. Any one run follows a single path through that tree, but the proposals along many possible paths can be generated before anyone knows which path will be taken.

In each model call, PreSTO generates the proposals of a prefetched subtree of this tree together, in one batch. It then follows the path that is actually realized, using the original acceptance rule, so the decisions are made exactly as before. A second observation shapes which subtree to fetch. The sampler reaches only a sparse part of each prefetched subtree, one that is nearly unary, close to a single chain. So rather than fetching a broad, even subtree, the authors design a predictive traversal rule that chooses the subtree expected to maximize the number of completed transitions.

The evaluation spans seven datasets and five base language models. PreSTO reaches a median speedup of one point seven one times on PowerMH. On its follow-up EntropyCut, the median speedup is one point nine one times. On MultiTryMH, it is two point three three times.

For the first two samplers, the source of the gain is visible. PreSTO makes a median of three point nine times fewer model calls, and each call is only two point one times longer. A further comparison isolates the predictive traversal rule. On LCB V6 with a nine billion parameter Qwen model, that rule is between one point two five and one point three seven times faster than the default breadth-first traversal.

What PreSTO trades on is the gap between sequential decisions and parallel computation. The decisions stay strictly in order, under the original rule, which is why the kernel is untouched. Only the expensive part, producing the proposals, moves ahead in batches, guided by a prediction of where the chain is likely to go.