Task-Agnostic and Adaptive-Size BERT Compression
Authors: Jin Xu
Xu Tan
Renqian Luo
Kaitao Song
Li Jian
Tao Qin
Tie-Yan Liu
Venue:  (2021)
Keywords: BERT compression
neural architecture search
adaptive sizes
across tasks
knowledge distillation

Abstract: While pre-trained language models such as BERT and RoBERTa have achieved impressive results on various natural language processing tasks, they have huge numbers of parameters and suffer from huge computational and memory costs, which make them difficult for real-world deployment. Hence, model compression should be performed in order to reduce the computation and memory cost of pre-trained models. In this work, we aim to compress BERT and address the following two challenging practical issues: (1) The compression algorithm should be able to output multiple compressed models with different sizes and latencies, so as to support devices with different kinds of memory and latency limitations; (2) the algorithm should be downstream task agnostic, so that the compressed models are generally applicable for different downstream tasks.  We leverage techniques in neural architecture search (NAS) and propose NAS-BERT, an efficient method for BERT compression. NAS-BERT trains a big supernet on a carefully designed search space containing various architectures and outputs multiple compressed models with adaptive sizes and latency. Furthermore, the training of NAS-BERT is conducted on standard self-supervised pre-training tasks (e.g., masked language model) and does not depend on specific downstream tasks. Thus, the models it produces can be used across various downstream tasks. The technical challenge of NAS-BERT is that training a big supernet on the pre-training task is extremely costly. We employ several techniques including block-wise search, search space pruning, and performance approximation to improve search efficiency and accuracy. Extensive experiments on GLUE benchmark datasets demonstrate that NAS-BERT can find lightweight models with better accuracy than previous approaches, and can be directly applied to different downstream tasks with adaptive model sizes for different requirements of memory or latency.

Under review as a conference paper at ICLR 2021
TASK-AGNOSTL AND ADAPTIVE-SIZE BERT COM- PRESSION
Anonymous authors Paper under double-blind review
ABSTRACT
While pre-trained language models such as BERT and RoBERTa have achieved impressive results on various natural language processing tasks, they have huge numbers of parameters and suffer from huge computational and memory costs, which make them difficult for real-world deployment. Hence, model compres- sion should be performed in order to reduce the computation and memory cost of pre-trained models. In this work, we aim to compress BERT and address the fol- lowing two challenging practical issues: (I) The compression algorithm should be able to output multiple compressed models with different sizes and latencies, as to support devices with different kinds of memory and latency limitations; (2) the algorithm should be downstream task agnostic, so that the compressed models are generally applicable for different downstream tasks. We leverage techniques neural architecture search (NAS) and propose NAS-BERT, an efficient method for BERT compression. NAS- BERT trains big supernet on carefully designed search space containing various architectures and outputs multiple compressed models with adaptive sizes and latency. Furthermore, the training of NAS-BERT is conducted on standard self-supervised pre-training tasks (e.g.. masked language model) and does not depend on specific downstream tasks. Thus, the models it produces can be used across various downstream tasks. The technical challenge of NAS-BERT is that training a big supernet on the pre-training task is extremely costly. We employ several techniques including block-wise search, search space pruning, and performance approximation to improve search efficiency and ac- curacy. Extensive experiments on GLUE benchmark datasets demonstrate that NAS-BERT can find lightweight models with better accuracy than previous ap- proaches, and can be directly applied to different downstream tasks with adaptive model sizes for different requirements of memory or latency.
INTRODUCTION
Pre-trained Transformer Vaswani al. (017)-based language models like BERT (Devlin et al. 2019), XLNet (Yang et al.) 2019) and RoBERTa (Liu et al. (019) have achieved impressive perfor- mance variety of downstream natural language processing tasks. These models pre-trained on massive language corpus via self-supervised tasks to learn language representation and fine-tuned on specific downstream tasks. Despite their effectiveness, these models are quite expensive in terms of computation and memory cost, which makes them difficult for the deployment on different down- stream tasks and various resource-restricted scenarios such as online servers, mobile phones, and embedded devices. Therefore, it is crucial to compress pre-trained models for practical deployment.
Recently, variety of compression techniques have been adopted to compress pre-trained models, such as pruning (McCarley) (2019; Gordon et al. (020). weight factorization (Lan et al., 2019).
quantization (Shen et al. [2200] /afrirr et al.) (209), and knowledge distillation (Sun et al. 2019
Sanh et al., 2019: Chen et al. 20.02020jjaaa
et al. 2019 Hou et al.) 2020 Song et (000)) Several
existing works (Tsai et al. 2020; McCarley 2019, Gordon et al. 2020 Sanh et al. 2019 Zafrin
et al. 2019 Chen et al. 2020 Lan et al. 2019 Sun et al. 2019) compress large pre-trained model into a small or fast model with fixed size on the pre- pre-tnginig or fine-tuning stage and have achieved good compression efficiency and accuracy. However, from the perspective of practical deployment, a fixed size model cannot be deployed in devices with different memory and latency constraints.
For example, smaller models preferred in embedded devices than in online servers, and faster

---

Under review as a conference paper at ICLR 2021
inference speed is more critical in online services than in offline services. On the other hand, some previous methods (Chen et al. 2020 Hou et al. 2020) compress the models on the fine-tuning stage for each specific downstream task. This can achieve good accuracy due to the dedicated design in each task. However, compressing the model for each task can be laborious and a compressed model for one task may not generalize well on another downstream task.
In this paper, we study the BERT compression in a different setting: the compressed models need to cover different sizes and latencies, in order to support devices with different kinds of memory and latency constraints; the compression is conducted on the pre-training stage so as to be downstream task agnostic. To this end, we propose NAS-BERT, which leverages neural architecture search (NAS) to automatically compress BERT models. We carefully design a search space that contains multi-head attention (Vaswani et al. 2017), separable convolution (Kaiser et al.] 2018), feed-forward network and identity operations with different hidden sizes to find efficient models. In order to search models with adaptive sizes that satisfy diverse requirements of memory and latency constraints in different devices, we train big supernet that contains all the candidate operations and architectures with weight sharing (Bender et al. 2018 Cai et al. 2018 2019 Yu et al., 2020). In order to reduce the laborious compressing on each downstream task, we directly train the big supernet and get the compressed model on the pre-training task to make it applicable across different downstream tasks.
However, it is extremely expensive to directly perform architecture search in a big supernet on the heavy pre-training task. To improve the search efficiency and accuracy, we employ several tech- niques including block-wise search, search space pruning and performance approximation during the search process: (1) We adopt block-wise search (Li et al. (020a) to divide the supernet into blocks so that the size of the search space can be reduced exponentially. To train each block, we leverage a pre-trained teacher model, divide the teacher model into blocks similarly, and use the input and output hidden states of the corresponding teacher block as paired data for training. (2) To further reduce the search cost of each block (even if block-wise search has greatly reduced the search space) due to the heavy burden of the pre-training task, we propose progressive shrinking to dynamically prune the search space according to the validation loss during training. To ensure that architectures with different sizes and latencies can be retained during the pruning process, we further divide all the architectures in each block into several bins according to their model sizes and perform progressive shrinking in each bin. (3) We obtain the compressed models under specific con- straints of memory and latency by assembling the architectures in every block using performance approximation, which can reduce the cost in model selection.
We evaluate the models compressed by NAS-BERT the GLUE benchmark (Wang et al. 2018). The results show that NAS-BERT can find lightweight models with various sizes from 5M to 60M with better accuracy than that achieved by previous work. Our contributions of NAS-BERS can be summarized as follows:
We carefully design a search space that contains various architectures and different sizes, and apply NAS on the pre-training task to search for efficient lightweight models, which is able to deliver adaptive model sizes given different requirements of memory or latency and apply for different downstream tasks.
We further apply block-wise search, progressive shrinking and performance approximation to reduce the huge search cost and improve the search accuracy.
Experiments on the GLUEE benchmark datasets demonstrate the effectiveness of NAS-BERT for model compression.
RELATED WORK
BERT Model Compression Recently, compressing pre-trained language models has been studied extensively and several techniques have been proposed such as knowledge distillation, pruning,
weight factorization, quantization and so on. Existing works (Tsai et al. 2020 Sanh et al. 2019
Sun et al.] 2019; Song et al.) 2020; Jiao al. 2019 Lan et al. 2019 Zafrir et al. 2019 Shen
et al.) 2020; Wang et al. 2019b Lan et al. 2019; Zafrir et al. 2019; Chen et al. 2020) am to
compress the pre- trained model into a fixed size of the model and have achieved trade-off between the small parameter size (usually no more than 66M) and the good performance. However, these compressed models cannot be deployed in devices with different memory and latency constraints.
2

---

Under review as a conference paper at ICLR 2021
Recent works (Hou et al. 2020) can deliver adaptive models for each specific downstream task and demonstrate the effectiveness of the task-oriented compression. For practical deployment, can be laborious to compress models from each task. Other works (Fan et al. 2019) can produce compressed models on the pre-training stage that can directly generalize to downstream tasks, and allow for efficient pruning at inference time. However, they do not explore the potential of different architectures as in our work. Different from existing works, NAS-BERT aims for task-agnostic compression on the pre-training stage which eliminates the laborious compression for each specific downstream task, and carefully designs the search space which is capable to explore the potential of different architectures and deliver various models given diverse memory and latency requirements.
Neural Architecture Search for Efficient Models Many works have leveraged NAS to search
2018 Howard et al. 2019] Tan & Le 2019: Cai et al.) 2019 Yu et al. 2020 Wang et al. 2020a Tsai et al. 2020). Most of them focus computer vision tasks and rely on specific designs on the convolutional layers (e.g., inverted bottleneck con- volution (Howard et al.] 2019) elastic kernel size (Cai et al. 2019 Yu et al.] 20200). Among them, once-for-all (Cai et al. (019) and BigNAS (Yu et al. (020) train a big supernet that contains all the candidate architectures and can get specialized sub-network by selecting from the super- net to support various requirements (e.g., model size and latency). HAT (Wang et al.) 2020a) also trains supernet with the adaptive widths and depths for machine translation tasks. Our proposed NAS-BERS also trains big supernet. However, different from these methods, we target model compression for BERT at the pre-training stage, which is a more challenging task due to the large model size and huge pre-training cost. Therefore, we introduce several techniques including block- wise search, progressive shrinking, and performance approximation to reduce the training cost and improve search efficiency. Tsai al. (2020) apply one-shot NAS to search faster Transformer but they cannot deliver multiple architectures to meet various constraints for deployment. Different from Tsai et al. (2020). NAS-BERT 1) progressively shrinks the search space to allocate more sources to promising architectures and thus can deliver various architectures without adding much computation: 2) designs bins in the shrinking algorithm to guarantee that can search architectures to meet diverse memory and latency constraints. 3) explores novel architectures with convolution layer, multi-head attention, and feed-forward layer, and achieves better performance than previous works for BERT compression.
efficient models (Liu et al. 2018: Cai et al.
3 METHOD
In this section, we describe NAS-BERT, which conducts neural architecture search to find small, novel and accurate BERT models. To meet the requirements of deployment for different memory and latency constraints and across different downstream tasks, we 1) train a supernet with a novel search space that contains different sizes of models for various resource-restricted devices, and 2) directly search the models on the pre-training task to make them generalizable on different down- stream tasks. The method can be divided into three steps: 1) search space design (Section 3.1); 2) supernet training (Section 3.2); 3) model selection (Section (33)). Due to the huge cost to train the big supernet the heavy pre-training task and select compressed models under specific constraints, we introduce several techniques including block-wise search, search space pruning and performance approximation in Section 3.2 and 3.3 to reduce the search space and improve the search efficiency.
3.1 SEARCH SPACE DESIGN
A novel search space allows the potential of combinations of different operations, instead of simply stacking basic Transformer block (multi-head attention and feed- -oedwrd network) as in the original BERT model. We adopt the chain-structured search space (Elsken et al.) 2018), and construct an over-parameterized supernet A with L. layers and each layer contains all candidate operations in
O = {o),- ... ) where is the number of predefined candidate operations. Residual connection is applied to each layer by adding the input to the output. There are C' possible paths (architectures)
in the supernet, and a specific architecture a = (a' .a'....... a²) where a. O is the operation in the [-th layer, mechanism that is widely used in NAS (Bender et al. 2018: Cai et al. 2019) for efficient training, where each architecture (path) shares the same set of operations in each layer.
) is a sub-net (path) in the supernet, shown in Fig. 2 (a). We adopt weight sharing
3

---

Under review as a conference paper at ICLR 2021
We further describe each operation in O follows: 1) Multi-head attention (MHA) and feed- forward network (FFN), which are the two basic operations in Transformer and are popular in pre- training models (in this way convolution (SepConv), whose effectiveness and efficiency in natural language processing tasks have been demonstrated by previous work (Kaiser al. 2018 Karatzoglou et al., (020). 3) Identity operation, which can support architectures with the number of layers less than L. Identity operation regarded as placeholder in candidate architecture and be removed to obtain shallower network. More detailed considerations on choosing the operation set are in Appendix A. Apart from different types of operations, to allow adaptive model sizes, each operation can have different hidden sizes: {128. 192, 256, 384, 512). In this way, architectures in the search space can be of different depths and widths. The complete candidate operation set contains (1+1+3) + 26 operations, where the first product term represents the number of types of operations and 3 represents the SepConv with different kernel size {3, 5, 7}. the second product term represents that there are 5 different hidden sizes. We list 26 operations in Table T
cover BERT model subnet in our supernet). 2) Separable
The detailed structure of separable
convolution shown in Fig.
Hidden Size
128
192
256
384
512
MHA
2 Heads
3 Heads
Heads
Hcads
8 Heads
FFN
512
768
1024
1536
2048
Kerne Kemel 5
Kernel 3
Kernel 3
Kemel 3
Kernel 3
SepConv
Kernel 5
Kernel 5
Kernl 5
Kernel 5
Kernl
Kernel 7
Kernel 7
Kemel 7
Kernel
Identity
Identity
Table 1: Candidate operation set. For each type of operation including multi-head attention (MHA), feed-forward network (FFN) and separable convolution (SepConv) in each row, we list the number of heads in MHA and the size of the intermediate layer in FFN, and kernel size in SepConv under different hidden sizes (in different columns).
Pointaie Conv
Depthwkk Coev
Pontwaae
Layatharm
Corv
GGUU
Copu
Separable Cony
Separable Conv
Figure 1: Structure of separable convolution.
3.2
SUPERNET TRAINING
3.2.1
BLOCK-WISH TRAINING WITH KNOWLEDGE DISTILLATION
Directly training the whole supernet causes huge cost due to its large model size and huge search space. With limited computational resources (total training time, steps, etc.), the amortized training time of each architecture from the huge search space is insufficient for accurate evaluation (Chu al.) 2019 Luo et al. 2019 Li al. 2020b). Inspired by Li al. (2020a), we adopt block-wise search to uniformly divide the supernet A into N blocks (A,, space and improve the efficiency. To train each block independently and effectively, knowledge dis- tillation is applied with prp-preiinn areitrained BERT model. The pre- pre-trained BERT model (teacher) is divided into corresponding N blocks as Fig. 2 The input and output hidden states of the corresponding block in the teacher model are used as the paired data to train the block in the supernet (student). Specifically, the 1ith student block receives the output of the (n 1)-th teacher block as the input and is optimized to predict the output of the n-th teacher block with mean square loss
AN) ) to reduce the search
C, = where f(:: A..) is the mapping function of n-th block A,,, Y,, is the output of the n-th block of the teacher model (Yn is the output of the embedding layer of the teacher model). At each training step,
A,, - Y.
(1)
4

---

Under review as a conference paper at ICLR 2021
Block
Block
Block
Block 4
L
Layer
Teacher (BERT (.)
Layer4 4
Layer 3
Layer
NAS Search Block
NAS Search Block
NAS Search Block
NAS Search Block
Layer
Student
Operation
Transformer Layer
Hidden Size Transformation
Figure 2: (a) an architecture (path) in the supernet. (b) an illustration of block-wise distillation (N 4 blocks). The supernet (student) and the pre-trained teacher model divided into blocks respectively and each student block is trained to mimic the corresponding teacher block.
we randomly sample an architecture from the search space following Guo et al. (2019): Bender et al. 2018); Cal al. (2019). which is memory-efficient due to the single path optimization. Different from Li et al. (2020a), we allow different hidden sizes and incorporate identity layer in each block to support elastic width and depth to derive models that meet various requirements. Besides, the search space within each block in our work larger compared to Li et al. (2020a) (100x larger) which is much more sample in-efficient and requires more techniques (described in Section (.2.2) to improve the training efficiency. Since the hidden sizes of the student block may be different from that in the teacher block, we cannot directly leverage the input and output hidden of the teacher block as the training data of the student block. To solve this problem, we use learnable linear transformation layer at the input and output of student block to transform each hidden size match that of the teacher block, as shown in Fig. 2
3.2.2 PROGRESSIVE SHRINKING
Although block-wise training can largely reduce the search space, the supernet still requires huge time for convergence due to the heavy pre-training task. To further improve the training effectiveness and efficiency, we propose to progressively shrink the search space in each block during training to allocate more training resources to more promising candidates (Wang et al. 2019a Li et al.
2020b
Luo et al. 2020). However, simply pruning the architectures cannot ensure to obtain different sizes of models, since larger models in An. are likely to be pruned the early stage of training due its difficulty of optimization (Chu et al. 2019 Luo et al. 2019 Li al.) 2020b) and smaller models are likely to be pruned the late stage due limited capacity. Therefore,
assign the architectures in An into different bins where each bin represents a short range of model sizes. Besides, we also apply latency constraints in each bin avoid models accepted parameter size but large latency. Denote Ph B "p(a") and I, = B I(a") as the maximum parameter size and latency for the b-th bin, where p(-) and I(.) calculate the parameter size and latency, a' is the largest model in the search space and B is the number of bins. The architecture
b-th bin should meet (1) ps > p(a) Pb-1 and (2) 4 ((a) 4-11 Architectures that cannot satisfy the constraint of latency are removed.
Then we conduct the progressive shrinking algorithm in each bin at the end of each training epoch as follows: 1) Sample E. architectures in each bin and get the validation losses on the dev set; 2) Rank the architectures according to their validation losses in descending order; 3) Remove architectures with the largest losses. The shrinking algorithm terminates when there only architectures left in the search space to avoid all the architectures being deleted. The design of bins ensures the diversity of models when shrinking the search space, which makes it possible to select model given diverse constraints at the model selection stage.
3.3 MODEL SELECTION
After the supernet training with progressive shrinking, each block contains B possible archi- tectures and the whole supernet (N blocks) contains (m B) complete architectures. The model selection procedure is as follows: 1) We build large lookup table T.aarc with (m B) items,
5

---

Under review as a conference paper at ICLR 2021
where each item contains the meta-information of complete architecture: (architecture, parame- ter, latency, loss). Since it is extremely time-consuming to measure the exact latency and loss for (m * B) (c.g., 10 in experiments) architectures, we use performance approximation to obtain the two values as described in the next paragraph. 2) For a given constraint of model size and infer- ence latency, we select the top T architectures with low loss from I.Tarch that meet the parameter and latency constraint. 3) We evaluate the validation loss of the top '" complete architectures on the dev set and select the best one as the final compressed model. The compressed model associ- ated with an embedding layer with adaptive size rather than a fixed size, and the embedding size is determined by the rules according to the requirement of target model size (see Appendix (A.5).
Next we introduce the performance approximation of the latency and loss when building the lookup table L.T a.Ter We measure the latency of each candidate operation (just 26 in our design) on the target device and store in a lookup table I.Tiae in advance, and then approximate the latency of an architecture I(a) by ((a) Σt= I(a') following Cai et al. (2018). where I(a²) is from LT To approximate the loss of an architectures in LT we add up the block-wise distillation loss of the sub-archites arbberrchee in each block on the dev set. Obtaining the dev loss of all sub-architectures in all blocks only involves m * B * N evaluations.
4 EXPERIMENT
4.1 EXPERIMENTAL SETUP
Supernet Training Setup The supernet consists of L 24 layers, which is consistent with BERT base (Devlin et al. 2019) (BERT base has 12 Transformer layers with 24 sub-layers in total, since each Transformer layer has a MHA and FFN). We use a pre-trained BERT (Devlin et al.) 2019) as the teacher model. The detailed configurations of the search space and teacher model train- ing can be found in Appendix A and Appendix A.2 The supernet is divided into N = blocks and the search space in each block is divided into B = 10 bins. We first train the supernet for 3 epochs without progressive shrinking as a warm start, and then begin to shrink at the end of each later epoch. We randomly sample E' 2000 architectures for validation (evaluate all architectures when the number of architectures in search space is less than 2000) and perform the progressive shrinking to remove R E/2 architectures for each bin as in Section 3.2.2 The shrinking process terminates when only m = 10 architectures are left in each bin in each block, and the training also ends. The considerations about how decide these hyper-parameters are described in Appendix AA.5 The supernet is trained on English Wikipedia plus BookCorpus (16GB size), with a batch size of 1024 sentences and each sentence consists of 512 tokens. The training costs 3 days on 32 NVIDIA P40 GPUs while training the BERT base teacher model costs 5 days. Other training configurations re- main the same as the teacher model. The latency used in progressive shrinking and model selection is measured on Intel(R) Xeon(R) CPU E5-2690 v4 @ 2.60 GHz with 12 cores, but our method can be easily applied to other devices (e.g., mobile platforms, embedded devices) by using the lookup table LTiat (described in Section 3.3 measured and built for the corresponding device. We select = 100 models from the table Taach on the model selection stage. In progressive shrinking, to reduce the time of evaluating all the candidates, we only evaluate on 5 batches rather than the whole dev set, which is accurate enough for the pruning according to our preliminary study.
Evaluation Setup on Downstream Tasks We evaluate the effectiveness of NAS-BERT by pre- training the compressed models on the original pre-training task and fine-tuning on the GLUE benchmark (Wang et al. 2018), which includes CoLA (Warstadt et al. 2018). SST-2 (Socher al. 2013). MRPC (Dolan & Brockett (005). STS-B (Cer et al. [017]) QQP (Chen et al. 2018). MNLI (Williams et al.) 2018). QNLI (Rajpurkar et al. 2016) and RTE (Dagan et al. (006). Similar to previous methods (Sanh et al. 2019] Wang et al.) 2020b] Ture et al. 2019 Hou et al.] we
also apply knowledge distillation and conduct it on two stages (i.e., pre-training and fine-tuning) as the default setting for evaluation. The details of two-stage distillation can be found in Appendix A.4 Considering the focus of our work is to compress pre- trained models with novel architectures instead of knowledge distillation, we only use prediction layer distillation and leave the various distillation techniques like layer-by-layer distillation, embedding layer distillation and attention matrix distilla-
tion (Sun et al. 2019 Jiao al. 2019 Sanh et al. 2019 Hou al. 2020 Wang et al. 2020b) that

---

Under review as a conference paper at ICLR 2021
Setting
FLOPs
Speedup
MNLI
QQP
QNLI
CoLA
SST-2
STS-B
RTE
MRPC
AVG
BERT PF
1.3e10
2.2x
82.6
90.3
89.4
52.6
92.1
88.3
75.6
89.2
82.5
NAS- BERToo + PF
1.3e10
2.2x
83.0
90.9
90.8
53.8
92.3
88.7
76.7
88.9
83.2
BERToo KD
1.3e10
2.2%
83.2
90.5
90.2
56.3
91.8
88.8
78.5
88.5
83,5
NAS-BERIE BES-BERToo + KD
1.3e10
2.2x
84.1
91.0
91.3
58.1
92.1
89.4
79.2
88.5
84.2
BERTa + PF
7.1e9
3.6 x
80.0
89.6
87.6
40.6
90.8
87.7
73.2
84.1
79.2
NAS-BERT PF
7.0c9
3.6 x
80.4
90.0
87.8
48.1
90.3
87.3
71.4
84.9
BERT + KD NAS-BERT KD
7.1e9
3.6 x
80.8
89.8
88.7
44.7
90.5
87.6
70.3
85.2
79.7
7.0c9
3.6
81.0
90.2
88.4
48.7
90.5
87.6
71.8
84.6
80.3
BE + PF
2.369
9.1 x
74.0
87.6
84,9
25.7
88.3
85.3
64.1
81.9
74.0
NAS- MASBBERT1o PF
2.3e9
8.7
76.0
88.4
86.3
27.8
88.6
84.3
68.7
81.5
75.2
BERT KD
2.3c9
9.1
74.4
87.8
85.7
32.5
86.6
85.2
66.9
77.9
74.6
1A + KD
2.3e9
8.7 x
76.4
88.5
86.3
34.0
88.6
84.8
66.6
79.1
75.5
BERT, + PF
8.6e8
20.7
67.7
84.1
80.9
10.4
81.6
81.1
62.8
78.6
68.4
NAS-BERT; + PF
8.7e8
23.6 x
74.2
85.7
83.9
19.6
84.9
82.8
67.0
80.0
72.3
BERT, + KD
8.6e8
20.7 x
67.9
83.2
80.6
12.6
82.8
81.0
61.9
78.1
68.5
NAS-RERT; KD
x
85.8
84.9
87.3
Table 2: Comparison of NAS-BERT models and hand-designed BERT models under different sizes (60M, 30M, 10M, 5M) on GLUE dev set. "PF" means pre-training and fine-tuning. "KD" means two-stage knowledge distillation. MNLI is reported in the matched set. Spearman correlation is reported for STS-B. Matthews correlation is reported for CoLA. Accuracy is reported for other tasks.
can further improve the performance as to future work. During fine-tuning on the GLUE benchmark, RTE, MRPC and STS-B are started from the model fine-tuned on MNLJ following Liu et 2019
4.2 RESULTS
Accuracy of NAS-BERT While our NAS-BERT can compress models with adaptive sizes, we only show the results of compressed models with 60M, 30M, 10M and 5M parameter sizes (de- noted as NAS-
NAS-BERT NAS-BERT 1o and NAS-BERT; respectively) on the GLUE benchmark due to the large evaluation cost, and list the model structures with different sizes in Appendix A.6) We compare our NAS-BERT models with hand-designed BERT models under the same parameter size and latency (denoted as BERT BERT 30. BERT 10 and BERT, respectively). We follow several principles when manually designing the BERT models: we just use the origi- nal BERT structure (MHA plus FFN) and keep the parameter, latency, depth and width as close as possible to the corresponding NAS-BERT models. The detailed structures of the hand-designed BERT models are introduced in the next paragraph. To demonstrate the advantages of architectures searched by NAS-BERT, the comparisons are evaluated in two settings: only with pre-training and fine-tuning (PF), and 2) pre- pre-training and fine-tuning with two-stage knowledge distillation (KD). We measure the inference speedup of following Clark et al.] (019)) The configurations are in Appendix A.3 The results are shown in Table from which we can see that NAS- outperforms hand-designed BERT baseline across almost all the tasks under various model sizes. Especially, the smaller the model size is, the larger gap can be observed (c.g., NAS-BERT, vs. BERT,). The results show that NAS-BERT can search for efficient lightweight models that are better than Transformer based models.
compared with BERT base and inference FLOPs
BERT Baselines We follow several principles when manually designing the BERT models in Table 2 1) The size of the embedding layer is the same as that of the corresponding NAS-BERT model; 2) We use the original BERT structure (MHA plus FFN) and keep the parameter, latency, depth and width as close as possible to the corresponding NAS-BERT model. The baseline BERT models in Table are: BERT (L=10, H=512, A=8), BERT (L=6, H=512, A=8), BERT (L=6, H=256, A=4) and BERT; (L=6, H=128, A=2) where L. is the number of layers, H is the hidden size, and A is the number of attention heads.
Comparison with Previous Work Next, we compare the effectiveness of our DAS-BERT to pre- vious methods on BERT compression. Since they usually compress BERT into model size of about
7

---

Under review as a conference paper at ICLR 2021
Model
Params
QQP
QNLI
CoLA
SST-2
MRPC
dev set
DistiBERT
66M
82.2
88.5
89.2
51.3
91.3
86.9
59.9
87.5
79.6
MiniLM
66M
84.0
91.0
91.0
49.2
92.0
71.5
88.4
BERT-offTheseus
66M
82.3
89.6
89.5
51.1
91.5
88.7
68.2
PD-BERT
66M
82.5
90.7
89.4
91.1
66.7
84,9
DynaBERT
60M
84.2
91.2
91.5
56.8
92.7
89.2
72.2
84.1
82.7
NAS-BERT
60M
84.1
91.0
91.3
58.1
92.1
89.4
79.2
88.5
84.2
NAS-BERTT
60M
84.8
91.2
91.9
58.7
93.1
89.9
79.8
88.9
84.8
Jest set
BERT-of-Thesens
66M
82.4
89.3
89.6
47.8
92.2
84.1
66.2
83.2
79.4
PD-BERT
66M
82.8
88.5
88,9
91.8
65.3
81.7
BERT-PKD
66M
81.5
88.9
89.0
92.0
65.5
79.9
TinyBERT* 66M
84.6
89.1
90.4
51.1
93.1
83.7
70.0
82.6
80.6
NAS-BERT 60M
83.5
88.9
90.9
48.4
92.9
86.1
73.7
84.5
81.1
NAS-BERT* 60M
84.1
88.8
91.2
50.5
92.6
86.9
72.7
86.4
81.7
Table 3: Results on the dev and test set of the GLUE benchmark. "s" means using data augmenta- tion. The test set results are obtained from the official GLUE leaderboard.
Setting
MNLI
QQP
QNLI
CoLA
SST-2
STS-B
RTE
MRPC
AVG
w/o PS
83.5
90.6
90.8
55.5
92.0
88.1
77.4
86.5
83.1
w PS
84.1
91.0
91.3
58.1
92.1
89.4
79.2
88.5
84.2
Table 4: The results of NAS-BERT with and without progressive shrinking (PS) on NAS-BERT,
66M or 60M, we use our NAS-BERT, BERT for comparison. We mainly compare our NAS-BERT with
1) DistilBERT (Sanh et al. 2019) which uses knowledge distillation on the pre-training stage; 2) BERT-PKD (Sun et al.) 2019), which distills the knowledge from the intermediate layers and the final output logits on the pre-training stage; 3) BERT-of- -hesess (Xu et al. (020), which uses mod- ule replacement for compression; 4) MiniLM (Wang et al. 2020b), which transfers the knowledge from self-attention module; 5) PD-BERT (Ture et al. 2019), which distills knowledge from the target domain in BERT training: 6) DynaBERT (Hou al. 2020). which uses network rewiring to adjust width and depth of BERT for each downstream task. 7) Tiny TERyBERT (Jiao et al.] 2019), which leverage embedding layer, hidden layer, and attention matrix distillation mimic the teacher model at both the pre-training and fine-tuning stages. To make comparison with DynaBERT and TinyBERT, we also use their data augmentation (Hou al. 2020 Jiao al. (019)) on downstream tasks. Table 3 mentation, reports the results on the dev and test set of the GLUE benchmark. Without data aug- BERT achieves better results on nearly all the tasks compared previous work. Further, with data augmentation, NAS-BERT outperforms DynaBERT and TinyBERT. Different from these methods that leverage advanced knowledge distillation techniques in pre-training and/or fine-tuning, HAS-BERY mainly takes advantage of architectures and achieves better accuracy, which demonstrates the advantages of NAS-BERT in model compression.
4.3
ABLATION STUDY
Setting
MNLI
QQP
QNLI
CoLA
SST-2
STS-B
RTE
MRPC
AVG
PS-arch
84.1
91.0
91.3
58.1
92.1
89.4
79.2
88.5
84.2
PS-op
83.8
91.0
90.2
54,5
92.0
88.8
75.6
86.1
82.8
Table 5: The results of different progressive shrinking approaches on NAS-BERT PS-op denote pruning architectures and operations in progressive shrinking.
PS-arch and
Ablation Study on Progressive Shrinking To verify the effectiveness progressive shrink- ing (PS), train the supernet with the same number of training epochs but without progressive
8

---

Under review as a conference paper at ICLR 2021
shrinking. We follow the same procedure in DAS-BERT for model selection and final evaluation downstream tasks. Due to the huge evaluation cost during the model selection on the whole search space without progressive shrinking, it costs 50 hours (evaluation cost on the shrinked search space only takes 5 minutes since there are only 10 architectures remaining in each bin in each block). The results are shown in Table 4
that NAS-BERT with progressive shrinking searches better architectures, with less total search time.
It be
We further show the training loss curve in Fig. It can be seen that the superset without progressive shrinking suffers from slow convergence. The huge number of architectures in the supemet need long time for suffi- cient training. Given a fixed budget of training time, progressive shrinking can ensure the promising archi- tectures be trained with more resources and result in more accurate evaluation, and thus better architectures can be selected. On the contrary, without progressive shrinking, the amortized training time of each architec- ture is insufficient, resulting in inaccurate evaluation and model selection.
0.75
w/o PS
0.65
0.55
0.45
0.25
Epoch
Figure 3: Loss curve of supernet training. Progressive shrinking starts at the epoch 4.
Different Progressive Shrinking Approaches Instead of pruning architectures (paths) from the search space, we can also prune operations (nodes) from the supernet (Radosavovic et al. 2020 Luo al.) 2020) in progressive shrinking. From the perspective of supernet, the former is to remove paths and the latter to remove nodes from the supernet. To evaluate the performance of operations (nodes) in supernet, the end of each training epoch, evaluate the architectures prune the search space according to the performance (validation loss) of operations. The validation loss of the operation o; in l-layer is estimated by the mean validation losses of all the architectures whose operation in the 1-th layer a' = o, The shrinking algorithm proceeds as follows:
the dev set and
Sample E architectures and get the validation losses on the dev set.
Rank operations according their mean validation losses in descending order.
Prune operations with the largest losses from the supernet repeatedly until removing R (R a hyper-parameter to control the speed of pruning) of the architectures in the search space.
The shrinking algorithm performs at the end of each training epoch, and terminates when only m = 10 architectures are left in each bin in each block, and the training also ends. For the fair comparison, we set m = 10 and E = 1000, which are same as settings in Section [. At the end of each epoch, we perform this algorithm to remove R 30% architectures for each bin. In this way, the algorithm terminates at the same epoch as that in Section 3.2.2 At the end of each training epoch, we evaluate the architectures on the dev set and prune the search space according to the performance of operations. As shown in Table pruning architectures in progressive shrinking achieves better results.
5 CONCLUSION
In this paper, we propose NAS-BERT, which leverages neural architecture search (NAS) to compress BERT models. We carefully design a search space with different operations associated with different hidden sizes, to explore the potential of diverse architectures and derive models with adaptive sizes according to the memory and latency requirements of different devices. The compression con- ducted on the pre-training stage and is downstream task agnostic, where the compressed models can be applicable for different downstream tasks. Experiments on the GLUE benchmark datasets demon- strate the effectiveness of our proposed NAS-BERT compared with both hand-designed BERT base- lines and previous works on BERT compression. For future work, we will explore more advanced search space and NAS methods to achieve better performance.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv: eeeei.4094047777 2014.
Gabriel Bender, Pieter-Jan Kindermans, Barret Zoph, Vijay Vasudevan, and Quoe Le. Understand- ing and simplifying one-shot architecture search. In International Conference on Machine Learn- ing. pp. 550-559, 2018.
Han Cai, Ligeng Zhu, and Song Han. Proxylessnas: Direct neural architecture search on target task and hardware. In International Conference Learning Representations, 2018.
Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, and Song Han. Once-for-all: Train one network and specialize it for efficient deployment. arXiv preprint wwXiv:1908.0999) 2019.
Daniel Cer, Mona Diab, Eneko Agirre, Iñigo Lopez-Gazpio, and Lucia Specia. SemEval-201 task 1: Semantic textual similarity multilingual and crosslingual focused evaluation. In Proceedings of the 11th International Workshop on Semantic Evaluation (SemEval-2017), pp. 1-14, August 2017.
Daoyuan Chen, Yaliang Li, Minghui Qiu, Zhen Wang, Bofang Li, Bolin Ding, Hongbo Deng, Jun Huang, Wei Lin, and Jingren Zhou. Adabert: Task-adaptive bert compression with differentiable neural architecture search. In Christian Bessiere (ed.), Proceedings of the Twenty-Ninth Interna- tional Joint Conference on Artificial Intelligence, IJCAI-20, pp. 2463-2469, International Joint Conferences on Artificial Intelligence Organization, 7 2020. doi: 10.24963/jecii URL https:///oi. org 24963/ /ttp:////caaaa22//33322 2020/341 Main track.
Zihan Chen, Hongbo Zhang, Xiaoji Zhang, and Leqi Zhao. Quora question pairs, 2018.
Xiangxiang Chu, Bo Zhang. Ruijun Xu, and Jixiang Li. Fairnas: Rethinking evaluation faimess of weight sharing neural architecture search. arXiv preprint arXiv:
2019.
Kevin Clark, Minh-Thang Luong, Quoe V Le, and Christopher 1) Manning. Electra: Pre-training text encoders as discriminators rather than generators. In International Conference on Learning Representations, 2019.
Ido Dagan, Oren Glickman, and Bernardo Magnini. The pascal recognising textual entailment challenge. In Machine Learning Challenges. Evaluating Predictive Uncertainty, Visual Object Classification, and Recognising Tectual Entailment, pp. 177-190, 2006.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 18.18653/v 1/N19-1444
William B. Dolan and Chris Brockett. Automatically constructing corpus of sentential paraphrases. In Proceedings of the Third International Workshop on Paraphrasing (/WP2005), 2005.
Thomas Elsken, Jan Hendrik Metzen, and Frank Hutter. Neural architecture search: A survey, arXiv preprint 1rXiv:1808.05377. 2018.
Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth demand with structured dropout. In International Conference on Learning Representations, 2019.
Mitchell A Gordon, Kevin Duh, and Nicholas Andrews. Compressing bert: Studying the effects of weight pruning on transfer learning. arXiv preprint arXiv:2002.0m 08300, 2020.
Zichao Guo, Xiangyu Zhang, Haoyuan MU, Wen Heng. Zechun Liu, Yichen Wei, and Jian Sun. Single path one-shot neural architecture search with uniform sampling. arXiv preprint 1rXiv:1904.004420, 2019.
Lu Hou, Lifeng Shang, Xin Jiang, and Qun Liu. Dynabert: Dynamic bert with adaptive width and depth. arXiv preprint arXiv:2004.0403 040377 2020.
10

---

Under review as a conference paper at ICLR 2021
Andrew Howard, Mark Sandler, Grace Chu, Liang- Chich Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukum Zhu, Ruoming Pang, Vijay Vasudevan, et al. Searching for mobilenetv3. In Pro- ceedings of the IEEE International Conference Computer Vision, pp. 1314-1324, 2019.
Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. Tinybert: Distilling bert for natural language understanding. arXiv preprinl arXiv: /900. 1035,, 2019.
Lukasz Kaiser, Aidan N Gomez, and Francois Chollet. Depthwise separable convolutions for neural machine translation. In International Conference on Learning Representations, 2018.
Antonios Karatzoglou, Nikolai Schnell, and Michael Beigl. Applying depthwise separable and multi-channel convolutional neural networks of varied kernel size semantic trajectories. Neural Computing and Applications, 32(1 1):6685-6698, 2020.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.698) 2014.
Guillaume Lample, Alexandre Sablayrolles, Mare' Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys. In Advances in Neural Information Processing Systems, pp. 8548-8559, 2019.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Sori- cut. Albert: A lite bert for self-supervised learning of language representations. In International Conference on Learning Representations, 2019.
Changlin Li, Jiefeng Peng, Liuchun Yuan, Guangrun Wang, Xiaodan Liang, Liang Lin, and Xi- aojun Chang. Block-wisely supervised neural architecture search with knowledge distillation. In Proceedings of the IEEÉ/CVF Conference on Computer Vision and Pattern Recognition, pp. 1989-1998, 2020a.
Xiang Li, Chen Lin, Chuming Li, Ming Sun, Wei Wu, Junjie Yan, and Wanli Ouyang. Improving one-shot nas suppressing the posterior fading. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 13836-1885,, 2020b.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. In International Conference on Learning Representations, 2018
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoye, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint 197.v 2019.
Rengian Luo, Tao Qin, and Enhong Chen. Balanced one-shot neural architecture optimization, 2019.
Renqian Luo, Xu Tan, Rui Wang. Tao Qin, Enhong Chen, and Tie-Yan Liu. Neural architecture search with gbdt. arXiv preprint arXiv:2007.04783, 2020.
J Scott McCarley. Pruning bert-based question answering model. arXiv preprint arXiv:1910.06360, 2019.
Ilija Radosavovic, Raj Prateek Kosaraju, Ross Girshick, Kaiming He, and Piotr Dollár. Designing network design spaces. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11428-10436, 2020.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. In EMNLP, pp. 2383-2392, 2016.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108, 2019.
Sheng Shen, Zhen Dong, Jiayu Ye, Linjian Ma, Zhewei Yao, Amir Gholami. Michael W Mahoney, and Kurt Keutzer. Q-bert: Hessian based ultra low precision quantization of bert. In AAAI, pp- 8815-8821, 2020.
11

---

Under review as a conference paper at ICLR 2021
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In EMNLP, pp. 1631-16422 October 2013.
Kaitao Song, Hao Sun, Xu Tan, Tao Qin, Jianfeng Lu, Hongzhi Liu, and Tin Liu. Light- paff: A two-stage distillation framework for pre-training and fine-tuning. arXiv preprint arXiv:2004.12817, 2020.
Siqi Sun, Yu Cheng, The Gan, and Jingjing Liu. Patient knowledge distillation for bert model compression. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP- IJCNLP), pp. 4314 4323, 2019.
Mingxing Tan and Quoe Le. Efficientnet: Rethinking model scaling for convolutional neural net- works. In International Conference on Machine Learning, pp. 6105-6114, 2019.
Henry Tsai, Jayden Ooi, Chun-Sung Ferng, Hyung Won Chung, and Jason Riesa. Finding fast transformers: One-shot neural architecture search by component composition. arXiv preprint orXiv:2008.06880 2020.
Julia Ture, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better: On the importance of pre-training compact models. arXiv preprint arXiv:7908.088 08962, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008 2017.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. In International Conference on Learning Representations, 2018.
Hanrui Wang, Zhanghao Wu, Zhijian Liu, Han Cai, Ligeng Zhu, Chuang Gan, and Song Han.
Hat: Hardware-awaae transformers for efficient natural language processing. arXiv preprint arXiv:2005. 14177, 2020a.
Linnan Wang, Saining Xie, Teng Li, Rodrigo Fonseca, and Yuandong Tian. Sample-efficient neural architecture search by learning action space. arXiv preprint 1xXi::996.008833 2019a.
Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou. Minilm: Deep self- attention distillation for task-agnostic compression of pre-trained transformers. arXiv preprint arXiv:2002. 10957, 2020b.
Ziheng Wang. Jeremy Wohlwend, and Tao Lei. Structured pruning of large language models. arXiv 2019b.
preprint arXiv:
Alex Warstadt, Amanpreet Singh, and Samuel R. Bowman. Neural network acceptability judgments. CoRR, abs//805.1 12471, 2018.
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sen- tence understanding through inference. In NAACI, pp. 1112-1122, June 2018.
Felix Wu, Angela Fan, Alexei Baevski, Yann Dauphin, and Michael Auli. Pay less attention with lightweight and dynamic convolutions. International Conference on Learning Representattions 2018.
Canwen Xu, Wangchunshu Zhou, Tao Ge, Furu Wei, and Ming Zhou. Bert-of-theseus: Compressing bert by progressive module replacing. arXiv preprint arXiv:2002.02925 2020.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Lee XInet: Generalized autoregressive pretraining for language understanding. In Advances in neural information processing systems, pp. 5753-5763, 2019.
Jiahui Yu, Pengchong Jin, Hanxiao Liu, Gabriel Bender, Pieter-Jan Kindermans, Mingxing Tan, Thomas Huang, Xiaodan Song, Ruoming Pang, and Quoc Le. Bignas: Scaling up neural archi- tecture search with big single-stage models. arXiv preprint arXiv:2003.1 142, 2020.
12

---

Under review as a conference paper at ICLR 2021
Ofir Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat. Q8bert: Quantized 8bit bert. arXiv preprint arXiv:1910.06188, 2019.
13

---

Under review as a conference paper at ICLR 2021
A
APPENDIX
A.1 OPERATION SET AND SEARCH SPACE
The Design Choice the Operation Set In addition to MHA and FFN, LSTM, convolution and variants of MHA and FFN have achieved good performance in many NLP tasks (Chen et al. 22000] Kaiser et al., 2018; Karatzoglou et al.) 2020 Babdanau et al. 2014). We describe the considerations to choose the operations in Table as follows: I) is not considered due to slow training and inference speed. 2) In our preliminary experiments, we try some variants of MHA and FFN (Lample et al. 2019 Wu et al. 2018), but fail to observe the advantages of small model size and/or better performance. 3) Considering the parameter size of convolution is K * 2 and separable convolution (SepConv) is + K + H where K and H the kernel size and hidden size, instead of convolution, we use SepConv with a larger kernel (with a larger receptive field) without significantly increasing the model size and the latency. Based on these considerations, we add SepConv into the candidate operation set.
To determine the possible hidden sizes for operations, we mainly consider the range of the com-
pressed model sizes. Previous works (Sanh et al. 2019 Sun et al. 2019] [Song et al.. 2020 Jiao
2019 Lan al. 2019 Zafrir et al.) 2019 Shen et al. 2020, Lan et al. 2019; Zafrir et al. Chen et al. 2020)) usually compress pre- pre-tr BERT model into small model (usually no more than 66M) for efficiency and effectiveness. Similarly,
al.
2019
this work, we aim to obtain com-
pressed models less than 66M. Therefore, we choose the hidden sizes between 128 and 512 for the candidate operations, which enables a good trade-off between efficiency and effectiveness
The Complexity of the Search Space The supernet consists of L = 24 layers. If we do not use block-wise search, there would be 26 x 1034 paths (possible candidate architectures) in the super- net. We divide the supernet into N 4 blocks and each block contains 6 layers. Within each block, the hidden size of the 6 layers are required to be consistent, and the hidden sizes across different blocks can be different. So there are 5 * 6 233280 paths (possible candidate sub-architectures in each block, where the first 5 is the number of candidate hidden sizes and 6 represents that there are 6 operations in 6 layers. Due to the identity operation, there is an unnecessary increase in the possible sequence of operations (architecture) as pointed in Li et al. (2020a). For example, the archi- tecture (FFN, identity, FFN, identity) is equivalent to (FFN, FFN, Identity, identity). }. Thus we only keep the architectures that all of the identity operations are the end (e.g.. {FFN, FFN, identity, identity) })) and delete other redundant architectures. After cleaning the redundancy, the search space in each block is reduced from the original 233280 97650, which largely improves the sample efficiency. We can select sub-architectures from each block and ensemble them to get a complete model. Considering NN 4 blocks, there are 97650 (about 10a0 possible combinations). Therefore the number of possible models is reduced from 10 to 1020.
A.2 TRAINING CONFIGURATIONS
Teacher Model We train the BERT base (L=12, H=768, A=12) (Devlin et al. 2019) the teacher model, where L is the number of layers, H is the hidden size, and A is the number of attention heads. Following Devlin et al. (019), we use BookCorpus plus English Wikipedia as pre-training data (16GB in total). We use Adam (Kigmma & Bal 2014) with a learning rate of 1e-4, 81 = 0.9 and B3 0.999. The learning rate is warmed up peak value of 5e-4 for the first 10,000 steps, and then linearly decays. The weight decay is 0.01 and the dropout rate is O.I. We apply the best practices proposed in Liu et al. (2019) train the BERT base on 16 NVIDIA V100 GPUs with large batches leveraging gradient accumulation (2048 samples per batch) and 125000 training steps. We present the performance of the teacher model and compare it with teacher models used in other works in T Our teacher model is better than others, which is mainly caused by the volatility of RTE and CoLA (small dataset). After removing these two datasets, the performance of the teacher model (average score: 89.74) is close to the teacher model of other methods (DistilBERT: 89.73, BERT- of-Theseus: 88.76 and DynaBERT: 89.68). In this way, compared with other approaches, without considering RTE and CoLA in Table 3
can still show its effectiveness
We do not hidden size smaller than 128 since it cannot yield model with enough accuracy.
14

---

Under review as a conference paper at ICLR 2021
Teacher model
MNLI-m
QQP
QNLI
CoLA
SST-2
STS-B
RTE
MRPC
AVG
DistilBERT
86.7
89.6
91.8
56.3
92.7
89.0
69.3
88.6
83.0
MiNiLM
84.5
91.3
91.7
58.9
93.2
68.6
87.3
BERT-off ofRT-of-Theseus Theseus
83.5
89.8
91.2
54.3
91.5
88.9
71.1
89.5
82.3
DynaBERT 84.8
90.9
92.0
58.1
92.9
89.8
71.1
87.7
83.4
MobileBERT
87.0
93.2
94.1
87.3
Ours
85.2
91.0
91.3
61.0
92.9
90.3
76.0
87.7
84.4
Table 6: The accuracy of the teacher models on dev set of the GLUE benchmark. The teacher model of MobileBERT' is IB-BERTIARRG LAngE which reaches the similar accuracy original BERT LERTE.RGE.
A.3 INFERENCE AND FLOPS SETUP
Following Sun al. 2019): Wang et al. (2020b), the inference time evaluated on QNLI training set with the batch size of 128 and the maximum sequence length of 128. The numbers reported in Table are the average 100 batches on Intel(R) (eon(R) CPU ES-2690 v4 2.60GHz with 12 cores. Following Clark et al. (2019), the inference FLOPs are calcuated with single length-128 input.
A.4 TWO-STAGE DISTILLATION
Pre-training Corpus
K distillation
Teacher Model (Pre- trained)
Student Model
(Pre- trained)
fine-tune
fine
Teacher Model (Fine-tuned)
K distillation
Student Model (Fine buned)
Downstreem task
Figure 4: The pipeline of two-stage distillation.
Two-stage distillation means applying knowledge distillation in both the pre-training and the fine- tuning stage. Previous methods (Song et al.) 22200] Sanh al. 2019 Jiao et al., 2019:/Gordoo al., 2020) have proved that using two-stage distillation is superior to the single-stage distillation. The pipeline of two-stage distillation shown in Fig. The procedure of two-stage distillation can be summarized as follows:
1. Pre-train the teacher model on the pre-training corpus.
2. Pre-train the light- student model with knowledge distillation from the pre-trained teacher model in step
3. Fine-tune the pre-trained teacher model in step 1 on the downstream task.
4. Fine-tune the pre-trained student model in step with knowledge distillation from the fine- tuned teacher model in step 3 on the target downstream task.
To simplify our expression, we denote the parameter of the student and the teacher model as Os and Or respectively. We adopt general formulation for knowledge distillation in both stages:
(X,)))
Lir,y;0s,00 0r)
Σ -.-E))LMME(RRpp + (xr (r.
(2)
where LMLE is the maximal likelihood loss of the student model 0g over the training corpus (X', )) (i.e.. masked language model loss on the pre-training stage or clasification/regrssio loss on the fine-tuning stage), and LKL is the KL divergence between the predicted probability distribution fæ; 0x) of the teacher model fr and the distribution fæ; 0g) of the student model 8g. is a hyper-paramette to trade off LMLE and LKL, and is set 0.5 in our experiments. There are other advanced distillation techniques (e.g.. embedding distillation or attention distillation) but here we only consider prediction layer distillation.
15

---

Under review as a conference paper at ICLR 2021
A.5 OTHER DESIGN CHOICES IN NAS-BERT
Why dividing the model into (-= 4) blocks? If N is small, the search space for each block will be huge (e.g.. C² possible architectures when N = and the training is costly. If N is large, candidate architectures in each block will be very limited (e.g., C architectures when N = L) and cannot explore the potential of combinations of different operations. We N for trade-off following Li ct al. 2020a).
Why dividing the search space each block into B(= 10) bins? If B too large, there are at least B architectures during the whole process of supernet training. The amortized training time of each architecture can be insufficient, resulting in inaccurate evaluation and model selection as shown in Section If B is very small, cannot get various architectures at the end of the training. Consequently. we set B = 10 for the trade-off.
Why conducting progressive shrinking until m(= 10) architectures is reached? As introduced in Section 3.3 we have (m B) possible combinations to build the lookup table LTarch. When m 10, the table has 10 architectures for the selection, which enough to select models under diverse requirements. If is too large, storing the big table T.T T.Tarch with many items ((m B)*) costly and the evaluation cost in Section 3.3 will increase exponentially.
Is the hidden size transformation module in Fig. retained in the final model? The answer is No. In the final derived architecture, we re-train the architecture by two-stage distillation without adding hidden size transformation module because the architecture does not need to match the hidden size to the one in the teacher model. However, there may he hidden size mismatch in the architecture itself due to the design of blocks, and we simply add additional linear layer handle it.
How to decide the embedding layer on the model selection stage? In order to save computing resources, we do not search the width of the embedding layer in NAS-BERT. We manually design the embedding layer for the final derived architecture according to the requirement of compressed model size: 1) < <10M, Wg=64, 2) <20M, Wg=128, 3) <35M, Wg=256, 4) <50M, Wg=384 and 5) >50M, ===512, where WE is the width of the embedding layer. For example, we choose the embedding layer with hidden size 512 when that target model size is required to be larger than 50M.
A.6 SEARCHED ARCHITECTURES BY NAS-BERT
Our NAS-BERS can generate different compressed models given specific constraints, as described in Section 333 Besides the several NAS-BERY models we evaluate in Table 2 we further select architectures with various sizes, from 5M to 60M with 5M interval, yielding totally 12 different architectures with different sizes. We present the architectures in the below figures (Figure (a)~(I)).
16

---

Under review as a conference paper at ICLR 2021
Embedding 512
Embeediin 512
:SepCenv
FFN, H 512
1 :Multi-boa,, 11 312
:SepCony 5, H 512
MMulti-hoad,
Multt-head, II 512
3 :SepConv H 512
:EPN,, H 513
PFN, 512
SepCorv 513
7 :Mulli beard 11512
Muiti head, H 512
bead,
SepConv 7, H 512
9 :SepCeny H 512
9. -Mutti-hadd H 512
PER,
FFN,
11 FFN, H H512
-SppCon
12 Multi-hend, 51512
SSpp 512
13 SepConv 5, F 512
13. Multi-haad, 512
head,
14 SepConv 5, 512
15 FFN H 512
15. :SepCovv 5, II 512
SapConv 3, 512
6 .FN, If 512
17 FFN, H 512
17 FFN, 512
18 Multi-hoad, 515
FFN 512
19 Multt-hoad, 512
20 SepConv 5, 512
-Mult-head,
21 FTN. IT 512
H151
22. FTN, 512
21 SepCanv 3, 512
PEA, H 5122
22. -HN, 312
(a) 60M
55M
17

---

Under review as a conference paper at ICLR 2021
Embedding 384
Embeddinn 384
C.N,,
0 :SopConv 3, H 512
SepCCan 5, H12
:Mulli hmad, 11 512
SepConv 512
2 Muutiiheaa, 11 512
N,
.FN, 512
SooCoon 512
:SepConn 512
PFN,
S:FNN H 512
SepCorv 513
6 :SupCeary 5122
Muiti head, H 512
7 Multii head, H 513
SepConv 7, H 512
SupCCenv H
9. -Mutti-hedd, H 512
9 9xFPN, H 512
FFN,
10 PEN, H11
-SeCoane
SupConv H.512
SSpp 512
12 -Multi-hedd 512
13. Multi-haad 512
13 SSpConv 5, 512
14 SepConv 7, 512
14 SepConv 7, 512
15 SSppCavv 3, II 51222
FFN, 512
6 .FN, if 512
16 Multi head, H 512
17 Multthhead, H 512
17 FEN, H 512
FFN 512
18 Multi-houd, 512
-Muti-head,
SepConv 5, 512
SepConv 3, H512
20 hoult-houd, 512
21 SepCanv 7, 512
21 LETN, 5151
22. -HNN 312
(c)
45M
18

---

Under review as a conference paper at ICLR 2021
Embedding 384
Embodding 256
0 :SepCenv's 3, HI 512
:SeeCooy 11
.Multi-head, 11 512
SeppCrrv H 512
:SepCeav
2 Multt-head, H 512
3 Multi Multi-head, H 512
SepCouv 512
:.SpConn 5, H H151
SapConv II 512
11512
6 SeppCenv H 512
Mulli bead, H
7 Mulbi head, H 512
-SapConv 11 512
:SepCenv H
SeepConv H 512
JFNN, H 512
9 SSepoonv H 512
10 FFN, H 512
SepCanv 3, 51512
11 :SepConv 3, 512
:uuti-heed, 512
12 Multi head, 512
:SepCore 5,
13. CoopConv 5, 11 512
13 SepCarv 5, 112512
14 SepConv 5, 512
14 SepCouv 3, H 512
CoppCovv 512
1 .PNN, 512
16 .TN, II 512
16 SepCanv 7, 512
17 SSepConv 7, 512
17 Mulli head, 512
Multi 512
19 Molti-head, 512
512
20 FFN H 512
-FNN, 512
(e)
35M
19

---

Under review as a conference paper at ICLR 2021
Embedding 256
0 SSepCcav 3, H 512
1 MMultihhaa, H 512
2 :SepCenv 5, H 512
3 :SepConv H 512
Emheading 256
IEFNN, H 512
:SepCouv H 304
5 :SepCenv H1512
Multi-head, H 384
6 Muulti head, H 512
SeppCorv H 384
:SepCcav H12
SSepConv H 384
:SepCenv
:SopConv
SSpponr
10 SSeCoovv 3, 512
-SapCory
11 :SopConv 5, 512
SSepConv H
12 Mullti-had, II 512
Multi-head, H 512
13 FTN 512
PFN, H 512
14 Coppoov 3, 512
:0 .FN, 512
15 :SepConv 7, 512
11 :Multihhad, H 512
15 Multi head, II 512
12 SapConv 7, 11512
17 Coppouv 3, 512
13. Mulu head, Il 512
18 FFN, 512
14 Mulli-hed,, 512
19. SepConv 5, 512
.UNN 512
30M
(h) 25M
20

---

Under review as a conference paper at ICLR 2021
Embedding 12B
Embedding 128
:SopCony
:SapConv 3, H
:Multi-boad, 11 304
1 Muuubi head, 11 384
:SepCony
:SepCony H 381
2 :SepCenv 5, H 384
:SppCony
3 :SepCon 3,
5 :SepConv H 381
:Septony
4 :SepCCanv 3, H 384
:SepCony II 384
3 SepCony 5, H 384
head,
:SepConv H 384
6 Multi head, H
:Spppoan 5,
7:FFN, H 512
11 :SepGove 7, H 384
Multiihead, 384
FFFN, 512
13 SepGorv 5, H 364
9 Multi-headd, H 512 SepCorv 5, 384
15 SepCorv 5, H 384
10 SepConv 3, 384
SepCaree 5,
11 :SepConv 5, 384
17 SupCarv 5, H 384
SepConv 3, 384
12 Multi-head, H 384
SapCoav 5, H 364
13 :SepConv 5, 384
20. Multt-had, It 384
21 SepConv 5, [I 364
11 :SepConv 3, 384
22 SepConv 3, 384
15 SepCCov 5, 384
23 SeppCrre (t1-
ii 20M
ii
21

---

Under review as a conference paper at ICLR 2021
Embodding 128
Embedding 64
SSppCouv 3, II 384
SSppConv 3, 192
:SopConv 381
SSepConv 7, 192
Multi-head, 384
Multiihhadd 192
-appConv 384
:SepConv 7, 192
Seppoon 384
SSepCony 7, 192
Muulti-head, Mulli-heed,, 192
Soppoov
-SSpConn 3, 192
bead, 384
Multi bead, 192
SSppoonv 3, 384
SSpCConv 7, 192
Seppoov 7, 11 384
SeppCony 5, 192
:Mult head, Il 384
Muulti head, Il 192
:SepCoav 3, H 364
11 :SepConv 3, H 112
(k) 10M
(1) 5M
22


Reviews:
[{"review_id": "uezBS89xMhk", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "First of all, I believe the paper is looking into a very important question that attracts lots of attention recently. The set of techniques proposed in the work are also reasonable and practical, where the proposed progressive space pruning seems to work very well. Empirically, the obtained models do perform better compared to standard Transformer baselines. As for the comparison with previous methods, since there are too many implementation details that can affect the fairness of the comparison (e.g. length of pretraining, batch size, teacher performances, etc), it's hard to judge the actual scale of the gain.\n\nThere are also a few concerns. \n- Firstly, when the block-wise search is used, it feels like the NAS-BERT is trained in a way that is more similar to a variant of distillation that additionally utilizes intermediate hidden states. As this signal is not used in the standard BERT baseline, some improvement could actually come from this factor besides a better model (architecture+param). A better baseline could be a Transformer trained in a similar way.\n- Secondly, in terms of novelty, this work is more like the combination of existing ideas, namely \"once-for-all\" and \"block-wise search\". One general issue with \"once-for-all\" is after the search, although you obtain multiple models of different sizes, the excellence these models are tied to the (1) specific set of shared parameters obtained and (2) the pre-training task. So, whether this is really a good way to obtain the desired task-agnostic compressed models is still questionable.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "nhZqyL8xXv6", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary\n\nThe paper develops a new method to compress the BERT model with varying model sizes depending on the underlying usage. They use block-wise neural architecture search to choose the best set of submodules for each of the blocks. To reduce the size of the exponential search space they progressively remove the architectural configurations that yields high loss. Over-all the paper is well written and nicely presented.\n\n+ve\n\n- The NAS-BERT can produce pertained models with varying model sizes which is better than DistilBERT and BERT-PKD that requires pre-training every time the number of layers are varied.\n- The paper conducts rigorous experiments to demonstrate the effectiveness of NAS-BERT. The baseline methods used for comparison covers most of the state-of-the-art methods used for model compression for BERT.\n\nConcerns:\n \n- Can the authors clarify if they save the model parameters corresponding to all the possible architectural choices? Or they find out the best configuration matching the model size and latency requirements and then do the pre-training again with those architectural choices.\n\n- In appendix A.7 the paper demonstrates some of the architectures used by NAS-BERT. For lower model size ( < 20M ) it can be observed that the NAS-BERT ends up choosing SepConv layers most of the times. Do the authors do any analysis on why SepConv layer works better than the self-attention layer and the feed forward network. How does the network perform if it is composed of all SepConv layers? Have the authors tried to use only SepConv layers and see if that itself gives good accuracy rather than doing the architecture search.\n\n- In terms of original ideas, although the concepts of block-wise architecture search, using SepConv layer for NLP tasks and using block-wise knowledge distillation are not novel by themselves but this paper has efficiently made use of the available techniques (along with efficient engineering work like progressively reducing the search space) to develop a method that gives good performance on NLP tasks.  ", "extra_scores": {}, "extra_text": {}}, {"review_id": "mMhsFT_vAAw", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThis paper proposes to search architectures of BERT model under various memory and latency contraints. The search algorithm is conducted by pretraining a big supernet that contains the all the sub-network structures, where the optimal models for different requirements are selected from it. Once an architecture is found, it is re-trained through pretraining-finetuning or two-stage distillation for each specific task. Several approaches (block-wise training and search, progressive shrinking, performance approximation) are proposed to improve the search efficiency. Experiments on GLUE benchmark shows the models found by proposed methods can achieve better accuracy than some of the previous compressed BERT models. The paper (together with the appendix) is clearly presented, and the idea is new and interesting to me. The experiments are detailed and comprehensive.\n\nPros:\nThe paper is well presented. The architecture of the superent and the candidate operations are carefully designed and selected. It seems that the SpeConv operation is particularly effective when the model size is small. The search algorithm including the block-wise training, progressive shrinking can remove less-optimal structures quickly and significantly reduce the search space. The performance of NAS-BERT models are generally better than those of the compressed BERT models with similar model size, although the comparisons may not be completely fair.\n\nConcerns:\n1. The organization of the paper can be further improved. The paper may not be easy to follow if the appendix is skipped, especially for the readers who are not familiar with NAS or related work. Many of the important information can only be found in appendix.\n2. The novelty of the paper is unclear to me. Although this work may be new on search BERT-like language model, it seems many of the ideas such as block-wise search and distillation are borrowed from existing work. Please the author clarify the main novelties and technical contribution of this work, especially to the field of neural architecture search or more broadly, AutoML . Moreover, some of the proposed techniques such as progressive shrinking are merely empirical practices and are lack of theory or insight showing how accurate the approximation would be.\n3. It is usually more illustrative (and also space saving) to plot accuracy versus latency/#parameters of different models in the same figure. Some of the well noted models such as MobileBERT and TinyBERT are not included in comparison. For DynaBERT, there are multiple configurations but only one is included. AdaBERT, which adopts NAS for each specific task, should also be included if possible. Again, since there are of many models with different size and latency, it may be better to have a plot for clear comparison. \n4. HAT (Wang et al. HAT: Hardware-Aware Transformers for Efficient Natural Language Processing. ACL 2020.) is not mentioned in the paper, which share similarities (training supernet) and differences (search algorithm) with this work from technical point of view. It will be better if the author can explain and compare the proposed search algorithm to evolutionary search.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "dEmgbf1omy", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents an effective NAS method for pre-trained language models at the pre-training stage, so the selected models can be applied to various downstream tasks with fine-tuning. To achieve better performances the widely used two-stage distillation and data augmentation are applied to the selected models from super-net.  \n\nThe main contribution of this work lies in the designed search space and the proposed three strategies (block-wise search, progressive shrinking and performance approximation) for improving search efficiency and accuracy. Although the novelty of this work is quite limited, training a big supernet for BERT at the pre-training stage is not trivial, which is useful for industry applications. The authors evaluate their approach on GLUE datasets and compare it to other state-of-the-art models.\n\nThe paper is well-written and organized, the experiments are thorough. However, I have several concerns:\n\n1)  In the Table 1 under the KD setting, “two-stage distillation” is conducted on the selected models from supernet to further improve the performances, it would be interesting to add another two settings: a) only conducting the distillation at the pre-training stage, b) continuing to pre-training on large scale unlabeled data to finally obtain better task-agnostic models. \n\n2) The models are evaluated on the GLUE dataset, more experiments on challenging QA tasks should be added. \n\n3) In the Table 1, the comparison to MobileBERT and TinyBERT should be added, and the FLOPs or the inference time on CPU/GPU can be provided.  \n\n4) Some important related work should be included, [1] HAT: Hardware-Aware Transformers for Efficient Natural Language Processing (although this work focuses on the machine translation task) [2] Finding Fast Transformers: One-Shot Neural Architecture Search by Component Composition.\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Compressing BERT is a practically important research direction. Our main concern on this submission is on its practical value. Comparing with MobileBERT in the literature, NAS-BERT does not show advantages on any aspect: latency, prediction performance, or model size (less important), while being much more costly to build because of NAS. MobileBERT just simply narrowed the original BERT models (8x narrower than BERT large). So it is hard to convince the readers that adaptive-size or NAS is interesting or matters. On the research side, this paper have some interesting points on designing the search space, but overall the novelty of this paper is limited, as all of the reviewers pointed out. It is also worth noticing that the claim of "task agonistic" in this paper does not fully hold: in the downstream tasks, the soft labels of the teacher model are required to train the compressed model. To be fully "task agonistic", the results on downstream tasks should be solely based on training with the ground truth labels, as in the MobileBERT paper. Once following the exact task agnostic experimental protocol, the reported performance in this paper may be significantly lower. 
Decision: Reject