POSS: Position Specialist Generates Better Draft for Speculative Decoding
This work addresses a specific bottleneck in speculative decoding for faster LLM inference, offering an incremental improvement over existing methods.
The paper tackles the problem of degrading draft token prediction quality at later positions in speculative decoding for LLM inference by proposing Position Specialists (PosS), which uses multiple position-specialized draft layers to improve token acceptance rates, resulting in enhanced average acceptance length and speed-up ratios across datasets.
Speculative decoding accelerates Large Language Model (LLM) inference by using a small draft model to predict multiple tokens, and a large target model to verify these tokens in parallel. Recent studies leverage the hidden state of the target model to enhance draft model prediction accuracy. However, existing methods suffer from the degrading quality of draft token predictions at later positions, due to error accumulation in draft model generated features. In this paper, we propose Position Specialists (PosS), which consist of multiple position-specialized draft layers to generate tokens at assigned position(s). Position specialists greatly improve token acceptance rate at later positions per drafting round, as each specialist only needs to focus on handling a certain level of draft model feature deviation. Experiment results on Llama-3-8B-Instruct and Llama-2-13B-chat across six datasets demonstrate that PosS effectively improves over baselines on average acceptance length and speed-up ratio. Our codebase is available at https://github.com/shrango/PosS.