Guided Speculative Inference for Efficient Test-Time Alignment of LLMs

TL;DR

Guided Speculative Inference (GSI) is a new algorithm that improves the efficiency and accuracy of reward-guided decoding in large language models, reducing latency and outperforming existing methods.

Contribution

GSI combines test-time scaling with a small auxiliary model to approximate optimal reward-guided policies, enhancing decoding efficiency and accuracy.

Findings

GSI achieves higher accuracy than standard soft best-of-$n$ methods.

GSI reduces end-to-end latency by up to 28%.

GSI outperforms some existing reward-guided decoding approaches.

Abstract

We propose Guided Speculative Inference (GSI), a novel algorithm for efficient reward-guided decoding in large language models. GSI combines soft best-of-$n$ test-time scaling with a reward model $r(x,y)$ and speculative samples from a small auxiliary model $\pi_S(y\mid x)$. We provably approximate both the optimal tilted policy $\pi_{\beta,B}(y\mid x) \propto \pi_B(y\mid x)\exp(\beta\,r(x,y))$ of soft best-of-$n$ under the base model $\pi_B$, as well as the expected reward under the optimal policy. In experiments on reasoning benchmarks (MATH500, OlympiadBench, Minerva Math, MMLU-STEM, GSM8K) and across different model families, our method achieves higher accuracy than standard soft best-of-$n$ with $\pi_S$ and reward-guided speculative decoding (Liao et al., 2025), and in certain settings even outperforms soft best-of-$n$ with $\pi_B$, while reducing end-to-end latency by up to $28\%$. The code is available at https://github.com/j-geuter/GSI .