L IFTING THE C URSE OF C APACITY G AP IN D ISTILLING L ARGE L ANGUAGE M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 (M O E), we propose a mixture of minimal experts (M INI M O E), 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 M INI M O E to a large extent.
M INI M O E 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 × , M INI M O E preserves ∼ 95% GLUE score of the teacher.
I NTRODUCTION
Pretrained language models (LMs) have become a popular choice for various downstream tasks, e.g., text classification, token classification, and question answering (Devlin et al., 2019; Liu et al., 2019;
Raffel et al., 2020).
Unfortunately, appealing performance comes with a huge cost of inference compute due to the scale of LMs.
Knowledge distillation (Hinton et al., 2015; Sun et al., 2019), as an alternative to model pruning (Han et al., 2015) and quantization (Sung et al., 2015), discovers a way to compress (Bucila et al., 2006) LMs with a teacher-student paradigm.
However, in LM distillation, we recognize a curse of capacity gap as: “ Large teachers, poor students. ”
The curse of capacity gap refers to a deficiency that a larger teacher might unexpectedly result in a poorer student especially when the capacity gap between the teacher and the student is large (Mirzadeh et al., 2020; Cho & Hariharan, 2019), as illustrated in Table 1.
Although a few studies (Wang et al., 2020; Zhang et al., 2022a; Park et al., 2021a) have investigated to fill the gap, the curse is still not yet tackled.
To the demand, we aim at lifting the curse of capacity gap via enlarging the capacity of the student without notably increasing the inference compute.
We propose a mixture of minimal experts (M INI M O E), inspired by the intuition of sparse activation of mixture of experts (M O E)
(Shazeer et al., 2017).
Thanks to that the activation process can be parallel on either single or multiple devices (He et al., 2021; Rajbhandari et al., 2022), M INI M O E on the one hand imposes extra parameters to the student, but on the other hand introduces negligibly additional inference compute brought by routing algorithm.
Experiments are conducted on GLUE (Wang et al., 2019) and CoNLL (Sang & Meulder, 2003).
The results exhibit that M INI M O E largely lifts the curse of the gap as in Table 1.
M INI M O E also achieves state-of-the-art performance compared with a range of competitive baselines, as shown in
[Table]
Our contributions can be summarized as follows:
• We recognize that LM distillation is faced a curse of capacity gap, invoking a deficiency when the capacity gap between the teacher and the student is large.
This is the first verification in LM distillation since previous studies recognize the curse in vision model distillation.
• We propose a M INI M O E to lift the curse of capacity gap.
As M INI M O E enjoys a sparse activation, it enlarges the capacity of the student without notably increasing the inference compute.
To our best knowledge, this is the first work aiming at lifting the curse completely.
• We examine our method on GLUE and CoNLL.
Experimental results show that our method lifts the curse of capacity gap, and realizes new state of the arts at almost all small FLOPs.
Thereby, we state that M INI M O E is a small yet nontrivial magic, making a great difference in circumventing the curse
C URSE OF C APACITY G AP
The curse of capacity gap is not new but is already recognized in studies on vision model distillation (Mirzadeh et al., 2020; Cho & Hariharan, 2019).
While a hit-the-mind drawback of the curse is that the performance of distilling to a small student can be dramatically worse than that of distilling to a slightly larger one, a rather counter-intuitive deficiency is invoked as that the performance of distilling from a large teacher can be unexpectedly worse than that of distilling from a smaller one (i.e., large teacher, poor student ).
We here give a minor theoretical justification on the curse, as a plus to the empirical justification.
Proposition 1 (VC dimension theory, Vapnik, 1998) .
Assuming that the teacher function is f T ∈
F T , the labeling function is f ∈ F , and the data is D , we have:
[Equation]
where r ( · ) is the risk function, | · | c is the function class capacity measure, and | · | is the data scale measure.
It should be highlighted that the approximation error ϵ T is negatively correlated with the capacity of the teacher model while the estimation error o ( · ) is correlated with the learning optimization.
Proposition 2 (Generalized distillation theory, Lopez-Paz et al., 2016) .
Additionally providing that the student function is f S ∈ F S , we have:
[Equation]
where the approximation error ϵ G is positively correlated with the capacity gap between the teacher and the student models, and 1 / 2 ≤ α ≤ 1 is a factor correlated to the learning rate.
Theorem 1. The bound for the student function at a learning rate can be written as:
[Equation]
[-Proof 
-]Proof.
The proof is rather straightforward by combining Proposition 1 and 2.
Remark 1. Under the same distillation setting, we can ignore the estimation error.
When we compare two students of different capacities distilled from a teacher of the same capacity, the student of a smaller capacity has a larger ϵ G thus lower performance.
When we compare two students of the same capacities distilled from teachers of different capacities, the student distilled from the teacher of a larger capacity has a smaller ϵ T yet a larger ϵ G thus a tradeoff.
Remark 1 basically tells that a tradeoff is associated with the increase of teacher capacity, implying that increasing teacher capacity would first lead to improved but then degraded student performance.
This tradeoff naturally corresponds with the curse.
On the other hand, it is accepted that large capacity gap is a pain and is processed in literature of
LM distillation (Wang et al., 2020; Zhang et al., 2022a; Zhou et al., 2022).
Being unaware of the curse of capacity gap, these studies attempt to offer student-friendly teachers by either interpolating teacher assistants (Wang et al., 2020; Zhang et al., 2022a) or adapting teacher knowledge (Zhou et al., 2022).
The unawareness is largely due to a fun fact that they only distil LMs like BERT base , but neglect the scalability to LMs like BERT large especially when the student is small.
Though the performance of student can be boosted in this way, the curse still remains in LM distillation as in
Figure 2.
[Figure]
Embarrassingly, while the curse is claimed to be tackled in vision model distillation (Zhu & Wang, 2021;
Park et al., 2021a; Zhao et al., 2022), our preliminary study (cf.
Table 6) indicates they are either expensive or not capable of LMs.
The potential differences are as follows: tasks (e.g., ImageNet v.s. GLUE), backbones (e.g., ResNets v.s. transformers), and paradigms (e.g., from scratch v.s. pretraining).
M INI M O E 3. M OTIVATION
Enlarging the capacity of the student is an intuitive solution to lift the curse of capacity gap.
However, regarding the inference compute efficiency, the increase of capacity should not introduce much inference compute.
An initial proposal can be using quantized backbones (Zafrir et al., 2019; Bai et al., 2021).
Quantized backbones may decrease the compute precision, therefore maintaining inference compute constant, along the course of enlarging the capacity.
But a vital portion of hardware-specific modifications are needed to do so.
We hence move on to next possibility.
Another alternative is using dynamic networks (Han et al., 2021)
based on the idea of conditional computation (Bengio et al., 2015).
The other commonly used one is depth-adaptive computation (Xin et al., 2020;
Zhou et al., 2020; Goyal et al., 2020; Kim & Cho, 2021) which involves layers into computation adaptively on either example (alias early exiting, Xin et al., 2020;
Zhou et al., 2020) or token (alias token reduction, Goyal et al., 2020; Kim & Cho, 2021) level.
A critical distinction between MoE and depth-adaptive models is that the compute of an MoE model is accurately under control while that of a depth-adaptive model is not.
We are impelled by the merits of MoE, and propose a M INI M O E so that the capacity of the student can be enlarged without much inference overhead increment.
Additionally, we argue that M INI M O E is orthogonal to alternatives mentioned above, and M INIM O E can be incorporated to these alternatives and makes it possible to serve more extreme scenarios.
It is noteworthy that a certain stream of work (Zhang et al., 2022b; Zuo et al., 2022) actually accelerates LMs via precisely converting them into MoE models.
Nonetheless, the moefication process is directly exerted to LMs with limited inference compute improvements (cf.
MoEBERT in Figure 1).
Contrarily, M INI M O E is comprised of minimal experts, each of which can be extremely small.
A comparison between mentioned possibilities and M INI M O E is listed in Table 2.
[Table]
MHA
FFN
MHA
FFN
MHA
FFN
MHA
FFN
MHA
FFN
MHA
FFN
MHA
FFN
MHA
FFN
MHA
FFN FFN FFN
MHA
FFN FFN FFN
MHA
FFN FFN
FFN
MHA
FFN
MHA
FFN
MHA
FFN
MHA
Note that there are emergent work exploring compressing MoE LMs (Xue et al., 2022) to dense students, which is walking down the same street in the opposite side since we instead focus on compressing dense LMs to MoE students.
3. I MPLEMENTATION
Minimal Language Models Typical language models are comprised of a stack of transformers layers (Vaswani et al., 2017), and are pretrained with language modeling tasks such as masked language modeling (Devlin et al., 2019).
A transformer layer can be decomposed to a multi-head self-attention (MHA) block and a feed-forward network (FFN) block.
Concretely, given an n -length sequence of d -dimension input vectors X ∈ R n × d with the i -th vector being x i , the output of the MHA block with A independent heads can be represented as:
[Equation]
where the j -th head is parameterized by W Q j , W K j , W V j ∈ R d × d A , and
W O j ∈ R d A × d .
On the other hand, the output of the FFN block is shown as:
[Equation]
where two fully-connected layers are parameterized by W I ∈
R d × d I and W O ∈ R d I × d respectively.
Details like biases, normalizations of a transformer layer are omitted for brevity.
To reach an acceptable compute budget, pioneering studies either pretrain language models or distil ones of small scales from LMs as in Figure 3.
There are three lines of work in LM distillation: firstly, task-specific distillation (Sun et al., 2019; Li et al., 2020; Sun et al., 2020a; Park et al., 2021b;
Hou et al., 2020; Xia et al., 2022) that conducts distillation on a specific task at finetuning stage; secondly, task-agnostic distillation (Turc et al., 2019; Sanh et al., 2019; Sun et al., 2020b; Wang et al., 2021b) that conducts distillation at pretraining stage; and thirdly, two-stage distillation (Jiao et al., 2020) that combines the power of both task-agnostic and -specific distillation.
Here, the distilled language models only refer to language models distilled with task-agnostic distillation regarding better taskscalability as the number of concerned tasks explodes.
We formally define the distilled language models as minimal language models (MiniLMs, somehow abuse of notation with Wang et al., 2020) notated with S .
In contrast, LMs are notated with T .
The learning objective of MiniLMs can be abstracted as L ( S ; T , D ) , where D denotes the data.
The specific form of L can be adapted to arbitrary alignment strategies.
We adopt a relation alignment strategy (Wang et al., 2021b) as follows:
[Equation]
where KL stands for kullback-leibler divergence.
Essentially, relation heads are derived by merging the original A attention heads and then splitting them to R heads.
T W Q j is the redistributed query parameter of the j -th relation head within totally R heads from the last layer of the LM, likewise T
W K j and T W V j are the key and value parameters.
An auxiliary MHA block is employed as the last layer of the MiniLM for better alignment following Wang et al.
(2021a).
The MiniLM can be then finetuned on any tasks.
Mixture of Minimal Experts Naturally, in order to enlarge the learning capacity gap of the student, we should add more parameters to the student.
However, trivially adding parameters usually leads to a loss of inference compute efficiency.
To remedy this, a mixture of minimal experts is proposed as in Figure 3.
Following prior literature (Shazeer et al., 2017; 2018), if we consider a FFN block in a MiniLM as a minimal expert, then extra parameters are exactly imposed as minimal experts to be added to the FFN block.
The FFN block is enabled as a mixture of m minimal experts FFN MoE in an expert gating tactic as:
[Equation]
where the j -th gate is parameterized by w G j ∈ R d , and correspondingly the j -th minimal expert is denoted as FFN j .
We further follow Fedus et al.
(2021) to only allow top- one gating (i.e., only the expert with highest gating probability is reserved) because we want to keep the inference compute untouched.
There are also diverse designs to achieve the sparse routing, such as hashing (Roller et al., 2021) which we find performs worse (cf.
Figure 5).
Since only one minimal expert is activated during the inference, the compute is only negligibly increased by expert routing.
As a complement, we can also achieve, if necessary, a mixture of experts in an MHA block similarly.
To encourage a balanced load across minimal experts, a differentiable load balancing objective
B ( S ; D ) is added from Lepikhin et al.
(2021) as:
[Equation]
where α is a coefficient that should be manually tune and is kept as 0.01 throughout this work following (Fedus et al., 2021).
While f j depicts the fraction of tokens dispatched to the j -th minimal expert, P j describes the fraction of the routing probability to the j -th minimal expert.
And a multiplier m is used to make the magnitude of the objective invariant to the number of minimal experts.
The load balancing objective basically desires a uniform routing so that the loss can be minimized.
The objective is added to the MiniLM not only at task-agnostic distillation stage but also but also at finetuning stage for practical concerns (cf.
Figure 5).
E XPERIMENTS 4. D ATA AND M ETRICS
We conduct experiments on GLUE (Wang et al., 2019) and CoNLL (Sang & Meulder, 2003).
The
GLUE originally consists of two sequence classification tasks, SST-2 (Socher et al., 2013) and
CoLA (Warstadt et al., 2019), with seven sequence-pair classification tasks, i.e., MRPC (Dolan & Brockett, 2005), STS-B (Cer et al., 2017), QQP, MNLI (Williams et al., 2018), QNLI (Rajpurkar et al., 2016), RTE (Bentivogli et al., 2011) and WNLI (Levesque et al., 2012).
We exclude WNLI and CoLA due to the evaluation inconsistency (in other words, MiniLMs get dramatically worse results while LMs get much better ones as found out in Xia et al., 2022) and use the left tasks.
The CoNLL is a token classification task.
Following BERT (Devlin et al., 2019), we report Accuracy (Acc) on SST-2, MNLI, QNLI, RTE, Spearman Correlation scores (SpCorr) on STS-B, and F1 on
MRPC, QQP, CoNLL.
Average score over tasks from GLUE (GLUE Score) is additionally computed.
Results on development sets are reported.
GFLOPs are also attached as theoretical speedup references.
We adopt Wikipedia data for task-agnostic disitllation.
The detailed statistics, maximum sequence lengths, and metrics of GLUE, CoNLL, and Wikipedia are supplied in Appendix A.
H ANDS - ON D ETAILS
Experiments are conducted upon distilling BERT base and BERT large (Devlin et al., 2019).
The distillation carried out on eight Nvidia A100s.
The number of relation heads is set to 32.
After the distillation, finetuning is carried out on one Nvidia A100.
The number of minimal experts m is default to 4 otherwise specified.
Other details are supplied in Appendix B. All experiments are task-agnostic ones, except those in Table 6.
4. B ASELINES
We compare M INI M O E with several state-of-the-art baselines.
Conventional Distillation FT indicates direct finetuning the student.
KD (Hinton et al., 2015),
PKD (Sun et al., 2019), and CKD (Park et al., 2021b) are methods with different distillation objectives, i.e., KD directly distills logits, PKD distills both logits and hidden states, and CKD distills high-order relations.
While above four methods originally initialize student structures by dropping layers, we enable them with a global pruning so that they can adapt to students of small scales.
DynaBERT (Hou et al., 2020) uses a two-step pruning to regulate student structures and a distillation objective akin to PKD.
MoEBERT (Zuo et al., 2022) moefies LMs by decomposing FFN blocks to
MoE layers.
For these task-specific distillation methods, student structures are denoted either with *L for preserved number of layers in layer-dropping or with *% for preserved portion of parameters in pruning.
As aforementioned methods are task-specific distillation ones, we then introduce task-agnostic ones.
TinyBERT (Jiao et al., 2020) exploits a distillation objective distilled with a combination of various feature alignments.
MiniLM (Wang et al., 2021b) straightforwardly utilizes a distillation objective with a deep relation alignment exactly the same with ours.
Since task-agnostic distillation allows both dropping layers and hidden dimensions, student structures are denoted with *L;*H accordingly.
Capacity-aware Distillation MiniLM w/ TA (Wang et al., 2020) specifically incorporates a teacher assistant to MiniLM.
AutoDisc (Zhang et al., 2022a) argues that the scale of the teacher assistant is crucial for student performance and proposes an automatic teacher assistant scheduler based on properties of pruning.
While MiniLM w/ TA is only inspected under a task-agnostic setting, AutoDisc offers results under both task-specific and task-agnostic settings.
Nevertheless, only task-specific AutoDisc is selected since pruned MiniLMs can be unfair to compare with.
There is scarce work in this direction in which we find these two are the most comparable ones.
4. M AIN R ESULTS
From results in Table 3, we observe that M INI M O E generally outperforms both conventional and capacity-aware baselines and achieves new state-of-the-art performance at all concerned times of compression.
For example, M INI M O E 4L;192H has an absolute 0.8 performance improvement over
MiniLM 4L;192H on GLUE.
Another observation is that the larger times of compression, the larger the performance improvements are.
For example, M INI M O E 4L;384H yields an absolute 0.5 performance improvement over MiniLM 4L;384H in contrast to that M INI M O E 6L;384H only has an absolute 0.2 performance improvement over MiniLM 6L;384H on GLUE.
Two more notes are that, MoEBERT nearly reaches the compression upper bound, and TinyBERT is reproduced without data augmentation for a fair comparison while the results with data augmentation are supplied in Appendix C.
From results in Table 4, we find that M INI M O E also lifts the curse of capacity gap at all concerned times of compression.
For example, M INI M O E 3L;384H disitlled from BERT large has an absolute 0.
performance gain over that distilled from BERT base on GLUE, and the value on CoNLL is 0.9.
On another note, MiniLM is free of the curse only at small times of compression, and MiniLM w/ TA can somewhat saves MiniLM from the curse at intermediate times of compression.
For example, both
MiniLM 3L;384H and MiniLM 3L;384H w/ TA fail to improve the performance via replacing BERT base with BERT large .
Results on larger LMs like BERT xlarge are supplied in Appendix F for scalability check.
[Table]
4. A NALYSES
Practical Inference Compute Since GFLOPs can only measure the theoretical inference compute, we further provide throughput (i.e., tokens per micro second) as a practical inference compute measure.
As in Table 5, 20 × compression can realize a significant inference compute gain in comparing KD 5% to BERT base .
The practical speedup is approximately 6.7 × .
Moreover, M INIM O E 3L;384H can retain most inference compute gain even if the routing algorithm can slightly reduce the gain when compared to MiniLM 3L;384H .
Student Scale Following the behavior of Figure 2, we would like to showcase whether M INI M O E can lift the curse across difference student scales.
From Figure 4, the curse is lifted to a large extent by M INI M O E in comparison with MiniLM and MiniLM w/ TA. However, M INI M O E meets [Table] [Table]
[Table] [Table] [Table] a bottleneck that distilling BERT large makes no difference from distilling BERT base when the FLOPs is at an extreme value 0.04G ( ∼ 273 × compression from BERT base , ∼ 968 × compression from
BERT large ).
We explore the extreme case by plugging a TA to M INI M O E as supplied in Appendix D.
[Figure] [Figure]
Routing Algorithm
Routing algorithm is also a crucial part benefiting from a nice design choice.
We compare our used gating with another fancy choice hashing.
We at the same time show the effect of using load balance at finetuning stage as well.
From the results in Figure 5, we see that gating outperforms hashing, and load balancing at both distillation and finetuning stages is superior to that at only distillation stage.
Expert Number Regarding the expert number m is a parameter of great importance for M INIM O E, we here study its impact on the performance.
The results in Figure 6 reveal a first ascending then descending phenomenon while adding experts at a time.
The phenomenon suggests there is a tradeoff when increasing the number of experts, and we conjecture the tradeoff accords with the famous bias-variance tradeoff (Hastie et al., 2001, Chapter 7).
That is, adding experts grows the parameter scale, thus decreasing bias yet increasing variance.
Another interesting notice is that smaller students favor fewer experts.
Based on the tradeoff conjecture, we hypothesize that smaller students are more sensitive to variance increment, as the biases of smaller students can arrive at a minimum more quickly than those of larger ones.
[Figure]
Failure of Vision Method We examine in a preliminary study the effectiveness of one of the vision model distillation methods (DeKD, Zhao et al., 2022) which can lift the curse of capacity gap.
From the results in Table 6, we unfortunately discover that DeKD can only give comparable performance in distilling BERT base , which even lags behind KD w/ TA.
It hints that vision model distillation methods are not that capable of LMs.
C ONCLUSIONS
In this work, we uncover a curse of capacity gap in LM distillation, which is well discussed in previous studies on vision model distillation but not recognized in distilling LMs.
While there are some studies investigating to fill the gap, we find they can hardly tackle the curse.
Interestingly, existing solutions in large vision language model distillation which are stated to be able to lift the curse fail to achieve so for LMs.
So we aim at lifting the curse by proposing a well-motivated
M INI M O E.
The M INI M O E can essentially enlarge the capacity of the student but leave the inference compute nearly untouched.
Our experimental results indicate that M INI M O E can not only lift the curse but also realize new state of the arts.
A D ATA S UMMARY
The detailed statistics, maximum sequence lengths, and metrics for datasets we use are shown in
Table 7, where the Wikipedia corpus used for distillation is also attached.
[Table]
B M ORE H ANDS - ON D ETAILS
General Guidelines The details of hyperparameters for distillation and finetuning are shown in Table 8.
We will be releasing our code and scripts in the final version for exact reproducibility.
For all cases, students are always randomly initialized following MiniLM.
[Table]
Implementation of MiniMoE We strictly follow the design of SwitchTransformer (Fedus et al., 2021) and extend it to the design of our M
INI M
O E.
We also follow their associated appendices to implement an MoE for multihead attention.
In detail, based on the original design, we treat an
FFN/MHA as an minimal expert, adopt top- one gating with load balancing, and employ a capacity factor of 1.25 for a good tradeoff (where overflowed tokens are dropped). For the parameter effect of adding an expert, we take expanding MiniLM 4L;192H (11.3M) to MiniMoE 4L;192H - 1,2E (14.9M) as an example. The number of parameters for embeddings is not changed (6.0M → 6.0M), but adding an expert ( 1,1E → 1,2E ) results in an increased number of parameters for transformers (5.4M → 9.0M). Further, our design for HashLayer (Roller et al., 2021) also strictly follows the original random hash design, i.e., per-token hash is used.
We strictly follow the best configuration of DeKD as reported in their paper (Zhao et al., 2022), where α is 1.0 and β is 8.0.
C R ESULTS W / D ATA A UGMENTATION
The results with data augmentation are produced from released checkpoints.
The results in Tabledemonstrate that TinyBERT is largely supported with data augmentation for great performance.
Another intriguing observation is that data augmentation only works for distillation but not for finetuning potentially due to the noise-resilience of distillation, so we preferably replace the finetuning stage with a task-specific distillation stage in experimenting with MiniLM.
[Table]
D M INI M O E AT E XTREME
The results in Table D witness that, M INI M O E sometimes struggles with extreme cases but can be enhanced with the help of TA. E R ELATED W ORK
Knowledge Distillation Distillation (Hinton et al., 2015) is a de facto way to compression (Bucila et al., 2006) LMs by transferring the knowledge of LMs to small language models.
During the distillation, a small language model serves as a student and treats a LM as a teacher to learn from.
There are three lines of work in LM distillation: firstly, task-specific distillation (Sun et al.,
[Table] [Table]
Li et al
[Table]
[Table]
Sun et al.
[Table]
2020a
[Table]
et al.
[Table]
2021b
[Table]
Hou et al.
[Table]
[Table]
Xia et al.
[Table]
2022) that conducts distillation on a specific task at finetuning stage; secondly, task-agnostic distillation (Turc et al., 2019; Sanh et al., 2019; Sun et al., 2020b;
Wang et al., 2021b) that conducts distillation at pretraining stage; and thirdly, two-stage distillation (Jiao et al., 2020) that combines the power of both task-agnostic and -specific distillation.
Though these methods realize promising performance when distilling LMs like BERT base , they can come short of scalability to LMs like BERT large especially when the student is of a small scale.
In fact, driven by recent observations (Wang et al., 2020;
Zhang et al., 2022a; Mirzadeh et al., 2020; Cho & Hariharan, 2019), distillation with a small student can be faced with two deficiencies due to the large capacity gap.
A few studies including teacher assistant-based (Mirzadeh et al., 2020; Zhang et al., 2022a) and student-friendly (Park et al., 2021a;
Zhou et al., 2022) distillation can alleviate the first but fail to resolve the second.
It is noteworthy that some work states they can tackle both deficiencies for vision models (Zhu & Wang, 2021; Zhao et al., 2022), but preliminary studies have found that they are either expensive or not capable of
LMs.
In our work, we follow the line of task-agnostic distillation of LMs and aims at lifting both efficiencies for the first time.
Mixture of Experts Based on the idea of conditional computation (Bengio et al., 2015), MoE layer is proposed to scale-up LMs in a sparsely activated fashion (Shazeer et al., 2017).
There are diverse designs to achieve the sparse routing, such as gating (Shazeer et al., 2018) and hashing (Roller et al., 2021), with necessary balance constraints (Lepikhin et al., 2021).
MoE layers are then joined to
LMs in the past one or two years (Fedus et al., 2021; Du et al., 2022).
Owing to the sparse activation property, the scales of LMs are significantly increased with only minor losses in compute efficiency on modern GPU devices so that the underneath scaling laws can be uncovered in a comparably cheap manner (He et al., 2021; Rajbhandari et al., 2022).
In our work, we are impelled by the merits of
MoE, and propose a M INI M O E so that the capacity of the student can be enlarged without much inference overhead increment.
M INI M O E can be similar to a certain stream of methods (Zhang et al., 2022b;
Zuo et al., 2022) that pursue accelerating LMs via precisely moefying them.
Nonetheless, the moefication process is exerted to LMs with limited inference compute improvements compared to those advanced by M INI M O E. Note that there are emergent work exploring compressing MoE
LMs (Xue et al., 2022) to dense students, which is walking down the same street in the opposite side since we instead focus on compressing dense LMs to MoE students.
F R ESULTS ON BERT xlarge
LM distillation, under either the task-agnostic setting as in our paper or the task-specific setting, has seldom been investigated to distil LMs larger than BERT large .
Even worse, there is only little work has been investigated to distil BERT large under the task-agnostic setting.
In the main results, we just follow the paces of the task-agnostic setting, not only due to the huge scales of larger LMs like T5 and GPT3 but also due to that task-agnostic LM distillation requires the access to the original pretraining data of usually vast volume.
What’s more, larger LMs like T5 can be incomparable to BERT owing to the architectural difference, and existing task-agnostic methods including ours may easily fail.
Regarding all the considerations mentioned above, however, we try to check the existence of the curse of capacity gap and examine M INI M O E under a comparably larger-scale setting, i.e., Chinese
BERT base v.s. BERT xlarge on some datasets from CLUE (Xu et al., 2020) (which can be viewed as the Chinese GLUE).
These datasets include a topic classification dataset TNews, a similar question matching dataset AFQMC, and a natural language inference dataset OCNLI.
The preliminary results are shown in Table 11.
As far as we know, while English BERT xlarge with more than one billion parameters trained by Nvidia Megatron (Shoeybi et al., 2019) is not publicly available, Chinese
BERT xlarge can be easily downloaded through huggingface.
1 It is noteworthy that Chinese BERT base is trained on Chinese Wikipedia ( ∼ 15G) while Chinese BERT xlarge is trained on Wudao Corpus ( ∼ 300G)
(Yuan et al., 2021).
We use Wikipedia data as the default choice for distillation, but Wudao data seems to be a more suitable (though not that fair) one for distilling Chinese BERT xlarge as we have found that Wikipedia could not make the distillation converge properly.
Painfully, it consumes around one week to achieve one epoch of distilling Chinese BERT xlarge using Wudao in contrast to five epochs of distilling Chinese BERT base on Wikipedia in one day.
The results show that Chinese
BERT xlarge is cursed to realize better students than Chinese BERT base does, and M INI M O E has the potential to lift the curse under the larger-scale setting.
[Table]
G P OTENTIAL OF M EMORY - EFFICIENT M INI M O E
One may argue that M INI M O E introduces much more memory consumption than MiniLM does, largely limiting the application scenarios for memory-sensitive devices (e.g., mobile devices).
However, there is no free lunch to enlarge the capacity of the student.
We should claim that, in order to increase the capacity, memory/space consumption is a cheaper choice (e.g., more experts) than latency/time consumption (e.g., more operations), and this is potentially the reason why large LMs like PaLM (Chowdhery et al., 2022) and FLAN (Chung et al., 2022) could become so popular.
We should also highlight that scenarios that require rather limited memory consumption (e.g., mobile scenarios) is currently not (though can be in the near future) the main concern of LMs.
In contrast,
LMs are usually served in GPU scenarios, where memory/space is easy to access.
Luckily, we find a potential path to address the memory efficiency concern based on the idea of parameter decomposition (e.g., SVD).
While embedding parameter decomposition is a general way to reduce the number of parameters for embeddings and could not make M INI M O E as memoryefficient as MiniLM.
We uncover that, without much performance sacrifice, transformer parameter decomposition in M INI M O E can be easier in comparison with that in MiniLM owing to the sparse activation property of MoE. That is, transformer parameters in M INI M O E have lower ranks than those in MiniLM, and this can be shown by analyzing the magnitudes of the normalized singular values using SVD.
The preliminary results of the output matrices of the last FFN layers separately from MiniLM 3L;384H and M INI M O E 3L;384H are shown in Table 12.
[Table]
With this finding, M INI M O E can compress more parameters than MiniLM does using parameter decomposition and finally yield a similar memory efficiency to that of MiniLM