[Paper Review] MixKD: Towards Efficient Distillation of Large-scale Language Models

Kevin J Liang, Weituo Hao+5 more|arXiv (Cornell University)|Nov 1, 2020

Topic ModelingComputer Science52 references30 citations

TL;DR

MixKD improves knowledge distillation for large-scale language models by augmenting data with mixup and querying the teacher on interpolated samples, yielding stronger generalization and better performance for compact students on GLUE tasks.

ABSTRACT

Large-scale language models have recently demonstrated impressive empirical performance. Nevertheless, the improved results are attained at the price of bigger models, more power consumption, and slower inference, which hinder their applicability to low-resource (both memory and computation) platforms. Knowledge distillation (KD) has been demonstrated as an effective framework for compressing such big models. However, large-scale neural network systems are prone to memorize training instances, and thus tend to make inconsistent predictions when the data distribution is altered slightly. Moreover, the student model has few opportunities to request useful information from the teacher model when there is limited task-specific data available. To address these issues, we propose MixKD, a data-agnostic distillation framework that leverages mixup, a simple yet efficient data augmentation approach, to endow the resulting model with stronger generalization ability. Concretely, in addition to the original training examples, the student model is encouraged to mimic the teacher's behavior on the linear interpolation of example pairs as well. We prove from a theoretical perspective that under reasonable conditions MixKD gives rise to a smaller gap between the generalization error and the empirical error. To verify its effectiveness, we conduct experiments on the GLUE benchmark, where MixKD consistently leads to significant gains over the standard KD training, and outperforms several competitive baselines. Experiments under a limited-data setting and ablation studies further demonstrate the advantages of the proposed approach.

Motivation & Objective

• Motivate efficient deployment of large Language Models by reducing model size and inference cost without sacrificing accuracy.
• Address overfitting and memorization when task data is limited by enriching training data via augmentation.
• Propose a data-agnostic distillation framework (MixKD) that leverages mixup to generate augmented samples for teacher–student learning.
• Provide theoretical justifications for improved generalization with augmented KD.
• Demonstrate empirical gains on GLUE tasks, especially in low-data regimes, and compare against baseline KD and related methods.

Proposed method

• Perform knowledge distillation from a large teacher (BERT) to a smaller student (BERT-3 or BERT-6).
• Apply mixup to word embeddings to generate augmented inputs for both student and teacher supervision (x' = λx_i + (1-λ)x_j; y' = λy_i + (1-λ)y_j).
• Train with a joint objective combining: standard cross-entropy on original data, cross-entropy on mixup data for the student, and a distillation loss between teacher and student on mixup data (L = L_MLE + α_SM L_SM + α_TMKD L_TMKD).
• Query the teacher on mixup samples and minimize the distillation loss against the student predictions on those samples (L_TMKD).
• Provide theoretical results showing conditions under which augmented data reduces the gap between empirical and generalization error, improving generalization.

Experimental results

Research questions

• RQ1 Does mixup-based augmentation in KD (MixKD) improve generalization and reduce the generalization gap for small Student models compared to standard KD?
• RQ2 How does MixKD perform on GLUE tasks across varying data availability, and how does it compare to DistilBERT, PKD, and other baselines?
• RQ3 What is the impact of combining mixup augmentation with backtranslation and other data augmentation techniques in KD?
• RQ4 How sensitive is MixKD to hyperparameters and mixup ratios, particularly in limited-data regimes?

Key findings

• MixKD consistently improves over standard KD and related baselines on GLUE tasks, with notable gains in limited-data settings.
• A 6-layer MixKD student with SM+TMKD+BT often captures a large fraction of teacher performance while maintaining significantly faster inference and fewer parameters.
• In SST-2, MixKD with SMTKD and backtranslation closes a substantial portion of the gap to the teacher compared to training from scratch or standard KD.
• Theoretical results show that data augmentation via mixup can reduce the generalization gap under several data-generation scenarios, supporting empirical gains.
• Embeddings visualizations indicate MixKD guides the student to align interpolated samples with the data manifold, improving latent space organization.