Efficiently Aligning Draft Models via Parameter- and Data-Efficient Adaptation

TL;DR

This paper presents EDA, a novel framework that efficiently adapts draft models for speculative decoding, significantly reducing training costs while maintaining high inference performance on fine-tuned large language models.

Contribution

EDA introduces a decoupled architecture, data regeneration, and sample selection to enable parameter- and data-efficient adaptation of draft models.

Findings

Restores speculative decoding performance on fine-tuned models

Achieves higher average acceptance lengths

Reduces training costs significantly

Abstract

Speculative decoding accelerates LLM inference but suffers from performance degradation when target models are fine-tuned for specific domains. A naive solution is to retrain draft models for every target model, which is costly and inefficient. To address this, we introduce a parameter- and data-efficient framework named Efficient Draft Adaptation, abbreviated as EDA, for efficiently adapting draft models. EDA introduces three innovations: (1) a decoupled architecture that utilizes shared and private components to model the shared and target-specific output distributions separately, enabling parameter-efficient adaptation by updating only the lightweight private component;(2) a data regeneration strategy that utilizes the fine-tuned target model to regenerate training data, thereby improving the alignment between training and speculative decoding, leading to higher average acceptance length;(3) a sample selection mechanism that prioritizes high-value data for efficient adaptation. Our experiments show that EDA effectively restores speculative performance on fine-tuned models, achieving superior average acceptance lengths with significantly reduced training costs compared to full retraining. Code is available at https://github.com/Lyn-Lucy/Efficient-Draft-Adaptation.