SWITCH: Studying with Teacher for Knowledge Distillation of Large Language Models
This work addresses the challenge of high inference costs and memory requirements in large language models for AI practitioners, offering an incremental improvement in knowledge distillation techniques.
The paper tackles the problem of noisy and biased student-generated outputs in knowledge distillation for large language models, which can cause teacher misguidance, especially in long sequences, and proposes SWITCH, a method that selectively incorporates teacher intervention during student generation, achieving superior performance over traditional methods in experiments across multiple models and datasets.
Despite the success of Large Language Models (LLMs), they still face challenges related to high inference costs and memory requirements. To address these issues, Knowledge Distillation (KD) has emerged as a popular method for model compression, with student-generated outputs (SGOs) as training data being particularly notable for reducing the mismatch between training and inference. However, SGOs often produce noisy and biased sequences, which can lead to misguidance from the teacher model, especially in long sequences. To mitigate these challenges, we propose SWITCH (Studying WIth TeaCHer for Knowledge Distillation), a novel approach that strategically incorporates the teacher model during the student's sequence generation. SWITCH identifies discrepancies between the token probabilities of the teacher and student models, allowing the teacher to intervene selectively, particularly in long sequences that are more prone to teacher misguidance. Extensive experimental results across three model families and five instruction-following datasets show that SWITCH surpasses traditional KD methods, particularly excelling in the generation of long sequential data.