Draft Model Distillation

The core goal is not to replicate the target model's internal weights, but to train the draft model to produce token probability distributions that closely match those of the target model. A high-fidelity match directly increases the acceptance rate during speculative decoding, as the target model is more likely to agree with the draft's proposals. This is typically achieved using knowledge distillation loss functions like Kullback-Leibler (KL) divergence, which minimizes the statistical distance between the two models' output probabilities.

The draft model is architecturally distinct and significantly smaller than the target model. Common approaches include:

• Reduced Parameters: Fewer layers and a smaller hidden dimension.
• Shallow Transformers: A 2-4 layer network versus the target's 10s or 100s of layers.
• Specialized Designs: Models like Medusa heads are not standalone networks but lightweight prediction heads attached to the target model for self-speculation. This asymmetry ensures the draft model's forward pass is orders of magnitude faster than the target's, which is the prerequisite for net latency reduction.

Effective distillation requires training on data representative of the target model's operational domain.

• Static Datasets: Often uses a subset of the target model's pre-training corpus (e.g., C4, Wikipedia).
• Dynamic/Online Data: Can involve generating outputs from the target model on-the-fly for distillation, ensuring alignment with the target's current behavior and any fine-tuning.
• Context Importance: Training must teach the draft to correctly interpret long contexts, as speculative decoding operates on the full prompt and generation history. Poor context handling leads to low acceptance rates after the first few tokens.

Standard cross-entropy loss is insufficient. Training employs distillation-specific objectives:

• KL Divergence Loss: The primary loss, minimizing KL(Target_Probabilities || Draft_Probabilities).
• Temperature Scaling: A softmax temperature parameter (T) is used to soften the target model's probability distribution, making it easier for the draft to learn.
• Combined Losses: Often combines KL loss with a small amount of standard cross-entropy on ground-truth tokens to stabilize training.
• Sequence-Level Training: Models are trained to predict multi-token sequences, not just single next tokens, to better prepare for generating the candidate sequence in speculation.

The ultimate success metric is the draft model's acceptance rate when paired with its target model in a speculative decoding loop. This is measured empirically and differs from traditional accuracy or perplexity.

• A high acceptance rate (e.g., >80% for 3-5 draft tokens) leads to a high speedup factor.
• Evaluation must be performed on held-out, domain-relevant text to be meaningful.
• The relationship is non-linear: small improvements in acceptance rate can lead to large improvements in latency reduction, as they reduce the frequency of costly rollback mechanisms.

Draft model distillation is a specialized subset of general knowledge distillation, with key differences:

An inference acceleration technique where a smaller, faster draft model proposes a sequence of future tokens. A larger target model then verifies this entire candidate sequence in a single, parallel forward pass, accepting correct tokens and rejecting incorrect ones. This reduces latency by replacing multiple sequential autoregressive steps with fewer, batched verification steps.

• Core Mechanism: Parallel verification of token batches.
• Primary Goal: Reduce wall-clock latency for text generation.
• Key Requirement: The draft model must be significantly faster than the target model to offset verification overhead.

The primary, larger, and more accurate language model in a speculative decoding setup. It is responsible for the verification forward pass, where it scores the candidate sequence proposed by the draft model against its own probability distributions.

• Role: Acts as the authoritative verifier and quality gate.
• Operation: Processes the candidate sequence in a single, batched inference step.
• Output: Decides to accept or reject each draft token, generating a corrected token if a rejection occurs.

The percentage of tokens proposed by the draft model that are accepted by the target model during verification. This is the critical determinant of speedup in speculative decoding.

• Calculation: (Number of Accepted Tokens) / (Total Proposed Tokens).
• Impact: A higher acceptance rate directly increases the speedup factor by maximizing the utility of each parallel verification pass.
• Goal of Distillation: The primary objective of draft model distillation is to maximize this metric by aligning the draft model's output distribution with the target model's.

The single, batched inference step where the target model evaluates the entire candidate sequence from the draft model. This is the computational cost traded for multiple sequential autoregressive steps.

• Process: The target model runs one forward pass on the concatenated context and candidate tokens.
• Output: Produces probability distributions for each position in the sequence.
• Efficiency: The cost of this pass must be less than the cost of the target model generating the same number of tokens autoregressively for a net speedup.

The broader training paradigm where a compact student model (the draft) is trained to mimic the behavior or output distributions of a larger, more complex teacher model (the target). Draft model distillation is a specific application focused on inference-time alignment for speculative decoding.

• Training Objective: Minimize a divergence loss (e.g., KL-divergence) between the student's and teacher's output probabilities.
• Contrast with Fine-Tuning: Focuses on behavioral cloning rather than task-specific performance on a downstream dataset.
• Benefit: Creates a small, fast proxy that retains key predictive characteristics of the large model.

The fundamental balance in speculative decoding between inference speed (latency reduction) and faithfulness to the target model's exact autoregressive output distribution (accuracy).

• Source of Tradeoff: The draft model is an approximation. Its errors cause rollbacks and may slightly alter the generated text flow compared to pure target model generation.
• Engineering Decision: Choosing draft model size and distillation quality involves optimizing this tradeoff for a specific use case.
• Measurement: Evaluated via both speedup factor and output quality metrics (e.g., win rate in human evaluations).