POSS: Position Specialist Generates Better Draft for Speculative Decoding

2025-06-04

AI Summary

In speculative decoding, draft models suffer from accumulated feature errors, causing a sharp decline in token acceptance rates toward later generation positions. To address this, we propose Position-Specialized Drafting (PosS), a novel mechanism that departs from the conventional single-model paradigm by assigning dedicated draft layers to distinct generation positions. PosS integrates LLM hidden-state enhancement, position-aware modular draft architecture, dynamic position allocation, and parallel verification. Evaluated on six benchmarks using Llama-3-8B-Instruct and Llama-2-13B-Chat, PosS achieves superior average acceptance length and end-to-end speedup over state-of-the-art methods. It is the first work to realize position-specialized draft modeling, effectively mitigating accuracy degradation and significantly improving efficiency for long-sequence generation.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Generation
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG

Abstract

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.

Problem

• Improves draft token prediction accuracy in speculative decoding
• Addresses error accumulation in draft model generated features
• Enhances token acceptance rate at later positions

Innovation

• Position-specialized draft layers improve accuracy
• Reduces error accumulation in draft predictions
• Enhances token acceptance rate at later positions