Steering Pretrained Drafters during Speculative Decoding

TL;DR

This paper proposes a lightweight dynamic alignment method to improve token acceptance rates in speculative decoding by steering pretrained drafters using verifier information, achieving significant gains with minimal overhead.

Contribution

Introduces a novel steering vector mechanism that enhances pretrained drafters' acceptance rates in speculative decoding, compatible with existing models and architectures.

Findings

Boosts accepted tokens by up to 35% under standard sampling.

Increases accepted tokens by 22% under greedy sampling.

Achieves these improvements with negligible computational overhead.

Abstract

Speculative decoding accelerates language model inference by separating generation into fast drafting and parallel verification. Its main limitation is drafter-verifier misalignment, which limits token acceptance and reduces overall effectiveness. While small drafting heads trained from scratch compensate with speed, they struggle when verification dominates latency or when inputs are out of distribution. In contrast, pretrained drafters, though slower, achieve higher acceptance rates thanks to stronger standalone generation capabilities, making them competitive when drafting latency is negligible relative to verification or communication overhead. In this work, we aim to improve the acceptance rates of pretrained drafters by introducing a lightweight dynamic alignment mechanism: a steering vector computed from the verifier's hidden states and injected into the pretrained drafter. Compared to existing offline alignment methods such as distillation, our approach boosts the number of accepted tokens by up to 35\% under standard sampling and 22\% under greedy sampling, all while incurring negligible computational overhead. Importantly, our approach can be retrofitted to existing architectures and pretrained models, enabling rapid adoption.