Skip to main navigation Skip to search Skip to main content

MIXKD: TOWARDS EFFICIENT DISTILLATION OF LARGE-SCALE LANGUAGE MODELS

  • Kevin J. Liang
  • , Weituo Hao
  • , Dinghan Shen
  • , Yufan Zhou
  • , Weizhu Chen
  • , Changyou Chen
  • , Lawrence Carin
  • Duke University
  • Meta
  • Microsoft Dynamics 365 AI
  • SUNY Buffalo

Research output: Contribution to conferencePaperpeer-review

26 Scopus citations

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.

Original languageEnglish
StatePublished - 2021
Event9th International Conference on Learning Representations, ICLR 2021 - Virtual, Online, Austria
Duration: May 3 2021May 7 2021

Conference

Conference9th International Conference on Learning Representations, ICLR 2021
Country/TerritoryAustria
CityVirtual, Online
Period05/3/2105/7/21

Fingerprint

Dive into the research topics of 'MIXKD: TOWARDS EFFICIENT DISTILLATION OF LARGE-SCALE LANGUAGE MODELS'. Together they form a unique fingerprint.

Cite this