MixKD: Towards Efficient Distillation of Large-scale Language ModelsDownload PDF

28 Sep 2020 (modified: 25 Jan 2021)ICLR 2021 PosterReaders: Everyone
  • Keywords: Natural Language Processing, Representation Learning
  • Abstract: Large-scale language models have demonstrated impressive empirical performance in recent years. Nevertheless, the improved results are attained at the price of bigger size, more power consumption, and slower inference, which hinder their applicability to low-resource (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 memorizing training instances, and thus tend to make inconsistent predictions when the data distribution is slightly altered. Moreover, the student model has few opportunities to request useful information from 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 teacher's behaviour on the linear interpolations of example pairs as well. We prove, from a theoretical perspective, that MixKD gives rise to a smaller gap between the generalization error and the empirical error. To verify its effectiveness, we conduct extensive 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.
  • Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics
  • One-sentence Summary: We propose MixKD, a distillation framework leveraging mixup for large-scale language models.
22 Replies