Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Lifting the Curse of Capacity Gap in Distilling Language Models

Chen Zhang, Yang Yang, Jiahao Liu, Jingang Wang, Yunsen Xian, Benyou Wang, Dawei Song

TL;DR

Problem: Distilling large pretrained language models is hindered by capacity gap between teacher and small student. Approach: Propose MiniMoE, a mixture of minimal experts that enlarges the student's capacity with negligible inference cost, via top-one gating and relation-aligned MiniLM distillation. Findings: On GLUE and CoNLL, MiniMoE substantially reduces the curse, achieving state-of-the-art results at low FLOPs and preserving about 95% of the teacher's GLUE score at 50x compression; gating-based routing, load balancing, and optimal expert counts are key for gains. Significance: Enables highly efficient, scalable LM distillation suitable for deployment where compute is constrained, with memory footprint identified as a limitation and potential mitigations explored.

Abstract

Pretrained language models (LMs) have shown compelling performance on various downstream tasks, but unfortunately they require a tremendous amount of inference compute. Knowledge distillation finds a path to compress LMs to small ones with a teacher-student paradigm. However, when the capacity gap between the teacher and the student is large, a curse of capacity gap appears, invoking a deficiency in distilling LMs. While a few studies have been carried out to fill the gap, the curse is not yet well tackled. In this paper, we aim at lifting the curse of capacity gap via enlarging the capacity of the student without notably increasing the inference compute. Largely motivated by sparse activation regime of mixture of experts (MoE), we propose a mixture of minimal experts (MiniMoE), which imposes extra parameters to the student but introduces almost no additional inference compute. Experimental results on GLUE and CoNLL demonstrate the curse of capacity gap is lifted by the magic of MiniMoE to a large extent. MiniMoE also achieves the state-of-the-art performance at small FLOPs compared with a range of competitive baselines. With a compression rate as much as $\sim$50$\times$, MiniMoE preserves $\sim$95\% GLUE score of the teacher.

Lifting the Curse of Capacity Gap in Distilling Language Models

TL;DR

Problem: Distilling large pretrained language models is hindered by capacity gap between teacher and small student. Approach: Propose MiniMoE, a mixture of minimal experts that enlarges the student's capacity with negligible inference cost, via top-one gating and relation-aligned MiniLM distillation. Findings: On GLUE and CoNLL, MiniMoE substantially reduces the curse, achieving state-of-the-art results at low FLOPs and preserving about 95% of the teacher's GLUE score at 50x compression; gating-based routing, load balancing, and optimal expert counts are key for gains. Significance: Enables highly efficient, scalable LM distillation suitable for deployment where compute is constrained, with memory footprint identified as a limitation and potential mitigations explored.

Abstract

Pretrained language models (LMs) have shown compelling performance on various downstream tasks, but unfortunately they require a tremendous amount of inference compute. Knowledge distillation finds a path to compress LMs to small ones with a teacher-student paradigm. However, when the capacity gap between the teacher and the student is large, a curse of capacity gap appears, invoking a deficiency in distilling LMs. While a few studies have been carried out to fill the gap, the curse is not yet well tackled. In this paper, we aim at lifting the curse of capacity gap via enlarging the capacity of the student without notably increasing the inference compute. Largely motivated by sparse activation regime of mixture of experts (MoE), we propose a mixture of minimal experts (MiniMoE), which imposes extra parameters to the student but introduces almost no additional inference compute. Experimental results on GLUE and CoNLL demonstrate the curse of capacity gap is lifted by the magic of MiniMoE to a large extent. MiniMoE also achieves the state-of-the-art performance at small FLOPs compared with a range of competitive baselines. With a compression rate as much as 50, MiniMoE preserves 95\% GLUE score of the teacher.
Paper Structure (32 sections, 3 theorems, 14 equations, 6 figures, 13 tables)

This paper contains 32 sections, 3 theorems, 14 equations, 6 figures, 13 tables.

Table of Contents

  1. Introduction
  2. Curse of Capacity Gap
  3. MiniMoE
  4. Motivation
  5. Implementation
  6. Minimal Language Models
  7. Mixture of Minimal Experts
  8. Experiments
  9. Data and Metrics
  10. Hands-on Details
  11. Baselines
  12. Conventional Distillation
  13. Capacity-aware Distillation
  14. Main Results
  15. Analyses
  16. ...and 17 more sections

Key Result

Proposition 1

Assuming that the teacher function is $f_{\mathcal{T}}\in\mathcal{F}_{\mathcal{T}}$, the labeling function is $f\in\mathcal{F}$, and the data is $\mathcal{D}$, we have: where $r(\cdot)$ is the risk function, $|\cdot|_{c}$ is the function class capacity measure, and $|\cdot|$ is the data scale measure. It should be highlighted that the approximation error $\epsilon_{\mathcal{T}}$ is negatively cor

Figures (6)

  • Figure 1: GLUE v.s. GFLOPs.
  • Figure 2: The performance of MiniLM and MiniLM w/ TA across different student scales upon distilling BERT base. We are glad to share checkpoints of an array of scales, together with those of MiniMoE, to facilitate the development of related research. It should be noted the unit of a vertical grid is comparably large.
  • Figure 3: Implementation of MiniMoE.
  • Figure 4: The performance of MiniMoE across different student scales upon distilling BERT base.
  • Figure 5: The performance of different routing choices with MiniMoE 4L;384H upon distilling BERT base.
  • ...and 1 more figures

Theorems & Definitions (5)

  • Proposition 1: VC dimension theory, Vapnik98
  • Proposition 2: Generalized distillation theory, Lopez-PazBSV15
  • Theorem 1
  • proof
  • Remark 1