Speculative Knowledge Distillation: Bridging the Teacher-Student Gap Through Interleaved Sampling

TL;DR

This paper introduces Speculative Knowledge Distillation (SKD), a novel method that improves the teacher-student training process by dynamically generating high-quality training data, leading to better performance in text generation tasks.

Contribution

SKD enables adaptive, on-the-fly data generation by combining student proposals with teacher feedback, effectively bridging the knowledge gap in traditional KD methods.

Findings

SKD outperforms existing KD methods across multiple text generation tasks.

SKD maintains high performance across different data sizes and model initializations.

SKD effectively aligns training data with inference-time distributions.

Abstract

Recent advances in knowledge distillation (KD) have enabled smaller student models to approach the performance of larger teacher models. However, popular methods such as supervised KD and on-policy KD, are adversely impacted by the knowledge gaps between teacher-student in practical scenarios. Supervised KD suffers from a distribution mismatch between training with a static dataset and inference over final student-generated outputs. Conversely, on-policy KD, which uses student-generated samples for training, can suffer from low-quality training examples with which teacher models are not familiar, resulting in inaccurate teacher feedback. To address these limitations, we introduce Speculative Knowledge Distillation (SKD), a novel approach that leverages cooperation between student and teacher models to generate high-quality training data on-the-fly while aligning with the student's inference-time distribution. In SKD, the student proposes tokens, and the teacher replaces poorly ranked ones based on its own distribution, transferring high-quality knowledge adaptively. We evaluate SKD on various text generation tasks, including translation, summarization, math, and instruction following, and show that SKD consistently outperforms existing KD methods across different domains, data sizes, and model initialization strategies.