taid
TAID: Temporally Adaptive Interpolated Distillation for Efficient Knowledge Transfer in Language Models
Shing, Makoto, Misaki, Kou, Bao, Han, Yokoi, Sho, Akiba, Takuya
Causal language models have demonstrated remarkable capabilities, but their size poses significant challenges for deployment in resource-constrained environments. Knowledge distillation, a widely-used technique for transferring knowledge from a large teacher model to a small student model, presents a promising approach for model compression. A significant remaining issue lies in the major differences between teacher and student models, namely the substantial capacity gap, mode averaging, and mode collapse, which pose barriers during distillation.s To address these issues, we introduce Temporally Adaptive Interpolated Distillation (TAID), a novel knowledge distillation approach that dynamically interpolates student and teacher distributions through an adaptive intermediate distribution, gradually shifting from the student's initial distribution towards the teacher's distribution. We provide a theoretical analysis demonstrating TAID's ability to prevent mode collapse and empirically show its effectiveness in addressing the capacity gap while balancing mode averaging and mode collapse. Our comprehensive experiments demonstrate TAID's superior performance across various model sizes and architectures in both instruction tuning and pre-training scenarios. These results demonstrate TAID's effectiveness in creating high-performing and efficient models, advancing the development of more accessible AI technologies. Large language models are too large. Causal language models (LMs) are increasingly becoming essential tools across various sectors (Malinka et al., 2023; Wu et al., 2023; Zhang et al., 2023a; He et al., 2024). Scaling data size, model size, and training steps has been the primary approach to improve LM performance (Kaplan et al., 2020; Hoffmann et al., 2022; OpenAI et al., 2024), leading to rapid advancements in both proprietary and open-source LMs (Touvron et al., 2023; Abdin et al., 2024; Yang et al., 2024). This paradox of scale hinders the widespread deployment and use of LMs despite their potential and high demand. Knowledge distillation offers a promising prescription. One promising approach to developing compact yet high-performing models is knowledge distillation (KD) (Hinton et al., 2015).