A UTO D ISC : A UTOMATIC D ISTILLATION S CHEDULE FOR L ARGE
L ANGUAGE M ODEL C OMPRESSION
Anonymous authors
Paper under double-blind review
A BSTRACT
Driven by the teacher-student paradigm, knowledge distillation is one of the de facto ways for language model compression.
Recent studies have uncovered that conventional distillation is less effective when facing a large capacity gap between the teacher and the student, and introduced teacher assistant-based distillation to bridge the gap.
As a connection, the scale and the performance of the teacher assistant is crucial for transferring the knowledge from the teacher to the student.
However, existing teacher assistant-based methods manually select the teacher assistant, requiring many trials before identifying the optimal teacher assistant.
To this end, we propose an Automatic Distillation Schedule (A UTO D ISC ) for large language model compression to discover the optimal teacher assistant in only one trial.
In particular, motivated by the finding that the performance of the student is positively correlated to the scale-performance tradeoff of the teacher assistant,
A UTO D ISC designs a λ -tradeoff to measure the optimality of the teacher assistant.
A UTO D ISC then yields the λ -tradeoff s of all teacher assistant candidates in a once-for-all optimization with two approximations.
The optimal teacher assistant can be automatically selected by uncovering the best λ -tradeoff .
A UTO D ISC is evaluated with an extensive set of experiments on a language understanding benchmark GLUE.
Experimental results demonstrate the improved efficiency with similar or even better effectiveness of our A UTO D ISC compared to several stateof-the-art baselines.
We further apply A UTO D ISC to a language model with over one billion parameters and show the scalability of A UTO D ISC .
I NTRODUCTION
Large language models (LLMs) (Devlin et al., 2019; Liu et al., 2019; Radford et al., 2019; Brown et al., 2020; Raffel et al., 2020) achieve promising results in various downstream tasks (Wang et al., 2019; Rajpurkar et al., 2018), but are inapplicable to those requiring limited compute or low latency (Liu et al., 2021b).
To fulfill the computational requirement, LLMs can be compressed via a range of strategies such as model quantization (Zafrir et al., 2019; Bai et al., 2021), pruning (Michel et al., 2019;
Hou et al., 2020), etc., among which knowledge distillation (Sun et al., 2019; Wang et al., 2020) is an appealing choice under the teacher-student paradigm.
In knowledge distillation,
LLMs serve as teachers and are distilled to small students.
Recent advances (Mirzadeh et al., 2020) have shown that conventional distillation suffers from severe performance decline when facing a large capacity gap between the teacher and the student.
To alleviate the shortcoming, teacher assistant-based distillation (Son et al., 2021) has been proposed, where the teacher is first distilled into an teacher assistant of an intermediate scale.
This teacher assistant then serves as an alternative teacher to better transfer the knowledge to the student.
While teacher assistant-based distillation generally lifts the performance of the student (Wang et al., 2020;
Wu et al., 2021), teacher assistants of different scales and performance may lead to different students and it is thus conjectured careful tuning is required to identify the teacher assistant that yields the best student.
To verify this, we conduct a pilot study and illustrate in Figure 1 that the performance of the student is largely impacted by the choice of the teacher assistant.
In fact, we observe in Figure 1 that the performance of the student is impacted by the teacher assistant from a perspective of scale-performance tradeoff.
However, existing studies only manually schedule the teacher assistant ( manual distillation schedule , in short M AN D ISC ), resulting in an inefficient solution that requires many trials to meet the optimal teacher assistant.
[Figure]
To this demand we propose an automatic distillation schedule (A UTO D ISC ) that automatically schedules the optimal teacher assistant in only one trial.
A quantitative measure, λ -tradeoff , is defined to empirically measure the scale-performance tradeoff for a given teacher assistant so that the optimality of the teacher assistant in terms of the student performance can be determined without trial distillation to the student.
Specifically, A UTO D ISC is implemented in three stages.
Specification , we leverage gridding and pruning techniques (Li et al., 2017; Frankle & Carbin, 2019) to generate a series of teacher assistant candidates of different scales and structures.
Optimization , we demonstrate that the generated candidates satisfy the incremental property and the sandwich rule , and present two optimization approximations that yields the λ -tradeoff of the teacher assistant candidate of each scale in a one-run optimization.
Selection , we select the optimal teacher assistant with the highest λ -tradeoff .
It is noteworthy that our method can be flexibly extended to the case with multiple teacher assistants by recursively applying A UTO D ISC .
However, this work only focuses on one teacher assistant between the teacher and student due to its sufficiency.
To verify the effectiveness of A UTO D ISC , we conduct experiments on a language understanding benchmark
GLUE (Wang et al., 2019) with both task-specific and task-agnostic distillation settings.
Experimental results exhibit the competitive performance of A UTO D ISC compared with an array of stateof-the-art baselines, with significantly improved efficiency ( ∼ 5-6 × ) of A UTO D ISC compared with that of M AN D ISC .
Further, we apply A UTO D ISC to a LLM EncT5 xl
(Liu et al., 2021a) and show the scalability of A UTO D ISC .
Our main contributions are summarized as follows:
• We investigate the impact of teacher assistants with different scales on the performance of the student, and introduce a quantitative scale-performance tradeoff measure, λ -tradeoff , on the teacher assistant that is positively correlated with the student performance.
• We show two properties of the specified candidates.
These properties lead to a novel optimization framework that jointly achieves yielding λ -tradeoff s of teacher assistants of all scales in one run.
The optimization framework, together with the λ -tradeoff , enables a one-trial distillation schedule.
• We validate the effectiveness and efficiency of the A UTO D ISC .
Our results of a language model with over one billion parameters show the scalability of our approach.
To our best knowledge, our work is the first one exploring the distillation of true LLMs.
M ETHODOLOGY 2. P ROBLEM F ORMULATION
We formally define the problem of A UTO D ISC .
Assuming we have a teacher ( T , s t , m t ) that should be distilled to a student ( S , s s , m s ), the goal is to find a teacher assistant ( A , s a , m a ) such that the student performance can be maximized when distilling the teacher to the student via the teacher
[Figure]
Selection : The assistant (i.e., T → A → S ).
Here the first, second and third elements in a tuple denote the model structure, the scale, and the performance respectively.
It is straightforward that the scale and the performance of the teacher assistant are bounded by the teacher and the student, i.e., s s ≤ s a ≤ s t and m s ≤ m a ≤ m t .
Ideally, we are seeking a teacher assistant with small scale and high performance so that the best student can be realized.
However, the scale-performance tradeoff is not easy to measure.
To empirically quantify scale-performance balance, we first introduce a new tradeoff measure below:
Definition 1 ( λ -tradeoff )
The λ scale-performance tradeoff measure of a teacher assistant ( A , s a , m a ) is defined as t a = m a /m t + λ · (1 − s a /s t ) , where λ ∈ [0 , 1] .
It is clear that the value of λ -tradeoff is bounded by 1 + λ · (1 − s s /s t ) when the teacher assistant can achieve the performance of the teacher ( m
a = m t ) with the student scale ( s a
= s s ).
However, in practice, this is impossible as smaller models usually lead to lower performance as shown by the blue curves in Figure 1.
We further observe that the λ -tradeoff (red curves) of the teacher assistant is positively correlated with the performance of the student (green curves).
Theoretically, due to the linear property of the λ -tradeoff and the concave property of the teacher assistant scale-performance correlation, there should always be one and only one maximum value of λ -tradeoff .
Therefore, the problem can be reformulated as finding an optimal teacher assistant with the maximum value of λ -tradeoff :
[Equation]
Based on the above reformulation, our methodology can be decomposed into three main stages: specification , optimization , and selection .
Essentially, during specification , a set of teacher assistant candidates are generated of different scales.
Then the performance metric of teacher assistant of each scale is obtained through a one-run optimization .
These two stages form a feasible region for the above reformulation.
Finally, the optimal teacher assistant A ∗ is selected with a linear programming of the feasible region during selection .
After the discovery of the optimal teacher assistant, the teacher assistant can subsequently be distilled to the expected student.
An overview of the methodology is given in Figure 2.
2. S PECIFICATION
Gridding .
Theoretically, one needs to generate candidates at every possible scale to find the optimal solution.
However, it is impossible to enumerate all possibilities in a continuous space.
Therefore, we discretize the candidate scales into n discrete values, {A =
( A k , s a k , m a k )
| ∆ s a =
( s t − s s ) /n } , with equal slicing between the teacher scale and student scale.
Pruning .
For candidates at each scale, there are still an infinite number of possible structures, e.g., different combinations of width and depth.
A number of approaches have been proposed to identify a good structure at a scale, including dynamic search (Hou et al., 2020), layer dropping (Fan et al., 2020) and pruning (Michel et al., 2019).
In this work, we adopt pruning to assign structures A k to the candidates due to its known advantages in knowledge distillation (Xia et al., 2022).
Concretely, following previous work (Michel et al., 2019), the pruning starts with the least important parameters/features based on their importance scores, which are approximated by masking the parameterized structures.
The technical details of our pruning are supplied in Appendix A.
Essentially, gridding positions the scales of candidates between the scales of the teacher and student with equal intervals and pruning assigns candidates with pruned structures.
2. O PTIMIZATION
A straightforward solution to unearth the optimality of each candidate is exhaustively measuring the student performance distilled from each, e.g., M AN D ISC .
λ -tradeoff offers a chance to measure the optimality without actual distillation.
However, the memory footprints and the computational costs apparently can also be extremely large considering the number of candidates.
To reduce the memory overhead and the computation complexity, we introduce two effective approximations, parametersharing and sandwich-optimization , so that the λ -tradeoff s of all candidates at different scales can be yielded in a once-for-all optimization.
The feasibility of the approximations are guarded by the following two Properties.
Property 1 (Incremental Property)
For two candidates A i and A j in the teacher assistant candidate set A , if s i < s j , then we have A i ⊂
A j .
This incremental property is an outcome of the pruning approach (Li et al., 2017; Frankle & Carbin, 2019), which essentially tells that among all candidates obtained from the specification, the structure of a candidate at a smaller scale is a subset of the structure for a candidate at a larger scale.
Remark 1 The incremental property affirms that a larger candidate can result in a smaller one by continuously pruning less significant parameters, which enables these candidates to be assembled into one sandwich-like model in a parameter-sharing fashion.
The memory scale of the sandwichlike model is exactly that of the largest candidate.
Property 2 (Sandwich Rule)
For two candidates A i and A j from candidate set
A , if s i < s j , then we have m s ≤ m i ≤ m j ≤ m t .
The sandwich rule (Yu & Huang, 2019; Cai et al., 2020) states that the performance of a candidate is bounded by the best performance of a larger candidate and a smaller one, due to the subset structure.
Therefore, a candidate can be optimized by alternatively distilling its larger and smaller candidates, without direct distillation.
Remark
The sandwich rule allows us to sub-sample η out of all n ( η ≤ n ) filling-like candidates and conduct sandwich-optimization over the sampled candidates, which substantially reduces the computational cost.
With the two approximations, we reduce the memory footprints of all candidates to a distinguished one with the largest scale.
The computational costs are also largely reduced with the sub-sampling.
Finally, we formulate the distillation objectives for task-specific distillation (TSD) and task-agnostic distillation (TAD) respectively as below:
[Equation]
where MSE, CE and KL stand for mean squared error, cross entropy and kullback-leibler divergence respectively.
X j is the intermediate output of the j -th layer within totally l layers, y is the final prediction.
As is taken from MiniLMv2 (Wang et al., 2021), Q R j is the query relation matrix of the j -th head within totally h attention heads from the last layer, likewise K R j and V R j are the key and value relation matrices.
Since heads can be pruned for a teacher assistant candidate, an additional self-attention module is employed as the last layer for TAD.
The teacher assistants with the best performance at different scales can be obtained after the above sandwich-optimization .
The unsampled teacher assistants can be retrieved based on the larger teacher assistant from the sampled pool using the shared parameters.
2. S ELECTION
The intuition for imposing the teacher assistant is to adequately sacrifice teacher scale for student performance, where the adequacy should be assured by retaining teacher assistant performance as much as possible.
In other words, a good teacher assistant is the one at a small scale yet with nice performance.
As observed in the pilot study, the λ -tradeoff measure is positively correlated with the final student performance and thus is directly used as the selection criterion.
The optimal teacher assistant can be identified by selecting the candidate with the best tradeoff measure.
The optimal teacher assistant is then distilled to the expected student again following above distillation objectives.
Note that the tradeoff measure is also dependent on λ .
However, we empirically find that the optimal solution of A UTO D ISC is relatively stable with a wide range of λ , and we fix λ to 0.2 in all our experiments.
More discussion on the impact of λ is provided in the experiments.
E XPERIMENTS 3. S ETUP
Datasets and Metrics We conduct experiments on a language understanding benchmark
GLUE (Wang et al., 2019).
The GLUE originally consists of two sequence classification tasks, SST2 (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., 2009) and WNLI (Levesque et al., 2012).
We exclude WNLI and CoLA due to the evaluation inconsistency (in other words, compressed LMs get dramatically worse results while original LMs get much better ones as found out in (Xia et al., 2022)) and use the other seven tasks for evaluation.
Following the work in BERT (Devlin et al., 2019), we report F1 on MRPC and QQP, Spearman Correlation scores (Sp Corr) on
STS-B, and Accuracy (Acc) on other tasks.
Macro average scores (Average) over these seven tasks are computed for overall performance.
Results on development sets are reported.
We also adopt
Wikipedia for pretraining in task-agnostic disitllation.
The detailed statistics, maximum sequence lengths, and metrics of GLUE and Wikipeida are supplied in Appendix B.
Implementation Details Experiments are carried out on BERT base (Devlin et al., 2019) and
EncT5 xl (Liu et al., 2021a).
EncT5 is a language model which achieves competitive performance as T5 3b (Raffel et al., 2020) on GLUE with a nearly encoder-only T5 (incorporated with a decoder layer).
Our task-specific experiments are carried out on either one Nvidia A100 for EncT5 xl or one Nvidia V100 for BERT base , and η is set to 9 according to our empirical investigation.
On the other hand, the task-agnostic experiments are carried out on eight Nvidia A100s with BERT base only
(Liu al is to
[Table]
1 to substantially reduce computational burden.
The number of relation heads is set to 32 since we use deep relation distillation as the task-agnostic distillation objective.
Other implementation details are supplied in Appendix C. Generally, the sampling is performed from candidates at scales { 100% , 95% , 90% , . . .
Baselines We compare our model with several state-of-the-art baselines.
BERT *L , MiniLMv *L;*H , and TinyBERT *L;*H denote methods via dropping layers and hidden dimensions, while DynaBERT *% , BERT *% , and EncT5 *% represent structured pruning with either local ranking or our global ranking (see Appendix A).
• Conventional Distillation: FT (Li et al., 2017) indicates direct fine-tuning after pruning.
KD (Hinton et al., 2015), PKD (Sun et al., 2019) and CKD (Park et al., 2021) are methods with different objectives,
i
e., KD directly distills logits, PKD distills both logits and hidden states and CKD distills token and layer relations.
DynaBERT (Hou et al., 2020) uses structured pruning with a local ranking in each layer.
MiniLMv2 (Wang et al., 2021) is distilled with the deep relation alignment.
TinyBERT (Jiao et al., 2020) is distilled with a combination of various feature distillations.
• Teacher Assistant-based Distillation: M AN D ISC manually selects the best teacher assistant among available trials.
TA (Wang et al., 2020) is specifically incorporated with MiniLMv2 for task-agnostic distillation.
3. M AIN R ESULTS
Results of Task-specific Distillation Table 1 presents the comparison results of different methods on task-specific distillation at four student scales.
The highlighted rows are the results from this work, where L TSD denotes the student directly obtained from distillation and A UTO D ISC denotes the student with additional distillation using the optimal teacher assistant.
There are several observations: First , L TSD achieves the best performance among all conventional distillation methods, while A UTO D ISC further improves the model and obtains similar or even better results compared to
M AN D ISC .
This validates the improved performance of A UTO D ISC for automatically identifying a
[Table]
good teacher assistant.
Notably, for further smaller BERT 3% , the improvement still holds, as supplied in Appendix D. Second
[Table]
conventional distillations generate reasonable results at large student scale 30% but fail to maintain the student performance at small scale 15% (with extremely worse results at scales 10% and 5% which are not included).
Nonetheless, A UTO D ISC consistently outperforms the baselines at all scales.
Additional comparisons of practical inference measurement are supplied in Appendix E. Third , pruning based models perform much better compared to the layer dropping methods, e.g., BERT 15% -KD achieves much higher score than FLOPs-matched BERT 2LKD, which verifies the effectiveness of pruning approach in knowledge distillation.
Moreover, we discover the global ranking strategy surpasses the local ranking one by comparing BERT 15% - L TSD to FLOPs-matched DynaBERT 15% .
We speculate the narrow structures induced by the local ranking strategy are not effective for rather small students.
The distributions of example pruned structures are supplied in Appendix F.
Results of Task-agnostic Distillation We also apply A UTO D ISC to task-agnostic distillation and report the results in Table 2.
The first glimpse is that L TAD surpasses L TSD , indicating the deep relation alignment is more suitable for task-agnostic distillation.
Surprisingly, we discover that the pruned structures can boost the performance of MiniLMv2, i.e., BERT- L TAD , and establish a new state-of-the-art for conventional task-agnostic distillation.
Another interesting observation is that teacher assistant-based distillations do not improve the performance over conventional distillations until the scale is reduced to 5%, indicating that conventional distillations are already promising choices on task-agnostic distillations.
Nonetheless, we still argue the applicability of A UTO D ISC to task-agnostic distillation for a performance guarantee.
Note that TinyBERT is less effective without data augmentation, and its results with data augmentation are supplied in Appendix G.
Results of Large-scale Distillation As is shown in Table 3, we conduct a similar comparison on a true LLM, EncT5 xl , with over one billion parameters.
The very first results of a true LLM also exhibit an akin trend as the one in BERT base .
And the results on a moderate BERT large are supplied in
Appendix H. We therefore conclude that the scalability of A UTO D ISC is also compelling.
Reversely, the results of A UTO D ISC on small LMs are supplied in Appendix I.
3. A NALYSES
Scale-performance Tradeoff To validate the existence of scale-performance tradeoff, we use teacher assistants at different scales for M AN D ISC and plot performance variations of these schedules upon BERT base in Figure 3(a).
It can be seen that reducing the scale can lead to performance improvement until a certain scale, after which performance degradation is witnessed.
Almost all manual schedules underperform the automatic one.
We attribute the inferiority of manual schedules to improper scale-performance tradeoffs, as concentrating only on either scale or performance will
[Table]
[Figure] [Figure] [Figure] [Figure]
[Figure] [Figure] [Figure] [Figure]
[Figure]
[Figure]
[Figure]
[Figure]
M
[Figure]
D ISC
[Figure] [Figure]
) 7.8G
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
A UTO D ISC
[Figure] [Figure]
three teacher assistants.
give rise to a trivial solution with pareto optimality (Sener & Koltun, 2018; Lin et al., 2019).
The overall phenomenon implies the existence of scale-performance tradeoff.
Similar phenomenon is also observed in EncT5, which is supplied in Appendix J. One may note that A UTO D ISC does not achieve the best performance, which is expected as M AN D ISC exhaustively searches for the teacher assistant at all scales.
However, A UTO D ISC is able to automatically identify a reasonably good teacher assistant in a much more efficient manner compared to M AN D ISC .
Sufficiency of One Teacher Assistant To examine whether one teacher assistant is sufficient, we insert more than one teacher assistant to A UTO D ISC and present the results in Figure 3(b).
It is clear that there is no obvious performance gain when applying more than one teacher assistant (two and three) in schedules.
Therefore, we alternatively choose to use only one teacher assistant in
A UTO D ISC for training efficiency based on the sufficiency.
The conclusion still holds for EncT5, which is supplied in Appendix J.
Impact of Candidate
Sampling We then study the impact of the sandwich-optimization in A UTO D ISC by varying the number of sampled candidates η , and measuring the training cost and the student performance.
From Table 4, we show the assembled sandwich together with sub-sampled fillings brings acceptable performance detriment and efficiency gain.
In comparison with M AN D ISC which we conduct 10 trials with different teacher assistants, A UTO D ISC with just one candidate is able to achieve similar performance with much less training time.
Impact of λ To show λ -tradeoff is robust on the value of λ , we vary λ within { 0.1,0.2,0.3 } .
It can be seen from Table 5 that the performance of A UTO D ISC is relatively stable with different values of λ .
Moreover, we offer a λ -independent solution using a negative derivative of performance to scale as the tradeoff measure, which yields slightly worse results, as supplied in Appendix K.
[Table] [Table]
R ELATED W ORK
Language Model Language models (LMs) (Devlin et al., 2019; Raffel et al., 2020) are widely adopted in various natural language tasks (Bao et al., 2020; Zhang et al., 2020).
Typical LMs consist of a stack of transformer (Vaswani et al., 2017) encoder/decoder layers.
Each encoder layer has two modules.
The first is a self-attention module
[Table]
and the second is a feed-forward module.
A residual connection is employed around each of these modules, with a layer normalization placed either in (pre-norm) or out of (post-norm) the connection (Xiong et al., 2020).
Each decoder layer additionally has a cross-attention module between the self-attention and feed-forward modules.
While LMs exhibit excellent performance in various downstream tasks, their scales impede the deployment in real-world applications.
Therefore, it is an important research problem of learning compact language models from the large ones.
In our work, we aim to make LMs deployable via model compression.
Model Pruning Inspired by the idea that not all parameters contribute equally to the overall performance of a model, model pruning (Han et al., 2015) is widely adopted to waive the parameters with little impact.
Model pruning spans from unstructured pruning (Frankle & Carbin, 2019; Louizos et al., 2018; Sanh et al., 2020; Chen et al., 2020) to structured pruning (Michel et al., 2019;
Hou et al., 2020; Li et al., 2017; Xia et al., 2022; Lagunas et al., 2021).
Unstructured pruning prunes parameters at neuron level referring to parameter magnitude (Han et al., 2015;
Louizos et al., 2018) or learning dynamics (Sanh et al., 2020), while structured pruning (Michel et al., 2019; Xia et al., 2022) prunes parameters at module level relying on parameter sensitivity to performance.
Although unstructured pruning enjoys a finer-grained pruning, it can only fit specialized devices.
In contrast, structured pruning generally fits modern acceleration devices.
In our work, we adopt structured pruning for deriving the structures of candidates for its benefits for distillation.
Pruning also offers an opportunity to optimize the efficiency of our method due to its merits (Li et al., 2017; Frankle &
Carbin, 2019; Yu & Huang, 2019; Cai et al., 2020).
C ONCLUSIONS
In this paper, we propose A UTO D ISC to automatically identify an effective teacher assistant for teacher assistant-based distillation, bridging the large capacity gap between the teacher and student in only on trial in contrast to M AN D ISC .
Based on the largely-neglected observation that scaleperformance tradeoff of the teacher assistant is of great importance to the performance of the student, we introduce a λ -tradeoff measure that quantifies the scale-performance tradeoff of the teacher assistant, and show that it is positively correlated with the student performance.
To compute the measures for possible teacher assistant candidates, we leverage gridding and pruning to specify these candidates and achieve a once-for-all optimization for these candidates based on two properties.
The best teacher assistant is selected according to the λ -tradeoff value.
Comprehensive experiments demonstrate the improved performance and applicability of A UTO D ISC .
Experimental results of a language model over one billion parameters show the scalability of A UTO D ISC .
Learning sparse neural networks
forum id neurips cc/paper/ 2019/hash/2c601ad9d2ff9bc8b282670cdd54f69f-Abstract.html .
org https:/
Attention
A T ECHNICAL D ETAILS OF P RUNING
Concretely, following previous work (Michel et al., 2019), the pruning always starts with the least important parameters/features, which are identified according to importance scores .
The importance scores are approximated by first masking the parameterized structures.
µ ( i ) , ν ( i ) , and ξ ( j ) denote the mask variables respectively for a self-attention head, optionally a cross-attention head, and a feed-forward neuron, such that for an intermediate input X and potentially an encoder-produced input E :
[Equation]
where potential bias terms (e.g., linear bias and position bias) are omitted, i means i -th head among h heads, j means j -th intermediate neuron among d neurons, and g is an activation function.
We initialize all mask variables to ones to preserve the original structure at the very beginning.
Then expected absolute gradients over either fine-tuning or pre-training data gives the important scores:
[Equation]
7) where ( x, y ) is a data point and L is the task-specific loss for task-specific models or the language modeling loss for pre-trained models.
E represents expectation.
The absolute value of gradient for a mask indicates how large the impact of pruning the corresponding structure is, thus implying how important the structure is.
Intuitively, we take a global ranking, in contrast to a local one as in other literature (Hou et al., 2020), for the structures of the same type (i.e., attention head or feed-forward element) from all stacking layers for pruning preference, before which we also normalize the importance scores for same-type structures in a layer with ℓ 2 norm, as suggested by Molchanov et al.
(2017), for a balanced pruning.
Therefore, for each candidate, we separately prune attention heads and feed-forward elements to the scale so that we reach a qualified structure.
For the sake of a corner case that all structures in a module are pruned, we skip the module by feeding the input as the output.
While we can alternate to an quite recent pruning method (Xia et al., 2022) exploiting both coarse-grained and fine-grained strategies for state-of-the-art performance, we argue that our framework is agnostic to pruning methods and keep the pruning method simple.
B D ATASET S TATISTICS
We conduct experiments on seven datasets.
The detailed statistics, maximum sequence lengths, and metrics for datasets we use are shown in Table 6, where the Wikipedia corpus used for pretraining is also attached.
C A DDITIONAL I MPLEMENTATION D ETAILS
The summary of hyperparameters for both task-specific and task-agnostic distillation is shown in
Table 7.
We will be releasing our code and scripts in the final version for exact reproducibility.
D A DDITIONAL R ESULTS UPON BERT BASE
We further conduct experiments on extremely small scale student model, i.e., BERT 3% .
The results are shown in Table 8.
[Table] [Table]
E I NFERENCE M EASUREMENT
Since FLOPs only offers theoretical inference compute, we additionally provide throughput for empirical inference compute of each model with throughput (i
[Table]
e., processed tokens per micro second) in Table 9.
The test environment is established by feeding 32 × 128 tokens to models.
The amount of decomposed parameters is also attached for a reference
[Table]
F P RUNED S TRUCTURE D ISTRIBUTIONS
We give the distributions of example pruned structures in Figure 4, which exactly show what pruned
LMs consist of.
While pruned BERT base tends to preserve bottom and middle layers, pruned EncT5 xl tends to preserve bottom layers.
Meanwhile, neurons in feed-forward layers are more likely to be pruned than heads in attention layers, owing to the centrality of the attention module within an transformer layer.
G D ATA A UGMENTATION FOR T INY BERT
We compare TinyBERT with and without data augmentation as in Table 10.
The results with data augmentation are retrieved from the original paper, since the augmented data is not publicly available.
The results demonstrate that TinyBERT is largely supported with data augmentation for good performance.
H R ESULTS UPON BERT LARGE
We show extended results of A UTO D ISC on BERT large for readers’ interest in Table 11.
Consistent patterns have been observed as in BERT base .
[Table] [Table]
I R ESULTS OF S MALL - SCALE D ISTILLATION
When A UTO D ISC is applied to small MiniLM 12;384H and BERT mini as shown in Table 12, A UTO D ISC can reversely affect the performance of conventional distillation.
Contrarily, M AN D ISC can still improve or at least retain the performance.
However, it is less necessary to compress small
LMs.
J V ARYING S CHEDULES FOR E NC T
Performance variations among possible schedules for EncT5 are displayed in Figure 5, where the existence of scale-performance tradeoff and sufficiency of one teacher assistant can be verified.
K N EGATIVE D ERIVATIVE -T RADEOFF
As mentioned in the main paper, although λ -tradeoff is able to provide stable tradeoff measurement, it is dependent on the value of λ .
To eliminate this dependency, we design a new measure, negative derivative-tradeoff, which computes the negative derivative of performance to scale at each candidate scale as: t a = lim δ → 0 − ( m a + δ − m a ) s a + δ − s a .
In the discrete case, t a i = − ( m a i +1 − m a i ) ∆ s a .
The idea of the measure is basically derived from saving the performance from a potentially significant drop.
However, first-order estimation can lead to a high estimation variance and can be further tuned with second-order or so for better performance.
The comparison results using λ -tradeoff and ND-tradeoff are shown in Table 13.
It can be seen from the table that A UTO D ISC -ND also achieves comparable results. [Figure] [Figure] [Table]
[Figure] [Figure] [Table] [Table]
[Table] [Figure] TSD
[Figure]
M
[Figure]
D
[Figure]
A UTO D
[Figure]
[Table]