Task-Agnostic and Adaptive-Size BERT Compression

1 INTRODUCTION . Pre-trained Transformer ( Vaswani et al. , 2017 ) -based language models like BERT ( Devlin et al. , 2019 ) , XLNet ( Yang et al. , 2019 ) and RoBERTa ( Liu et al. , 2019 ) have achieved impressive performance on a variety of downstream natural language processing tasks . These models are 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 downstream 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 , a variety of compression techniques have been adopted to compress pre-trained models , such as pruning ( McCarley , 2019 ; Gordon et al. , 2020 ) , weight factorization ( Lan et al. , 2019 ) , quantization ( Shen et al. , 2020 ; Zafrir et al. , 2019 ) , and knowledge distillation ( Sun et al. , 2019 ; Sanh et al. , 2019 ; Chen et al. , 2020 ; Jiao et al. , 2019 ; Hou et al. , 2020 ; Song et al. , 2020 ) . Several existing works ( Tsai et al. , 2020 ; McCarley , 2019 ; Gordon et al. , 2020 ; Sanh et al. , 2019 ; Zafrir et al. , 2019 ; Chen et al. , 2020 ; Lan et al. , 2019 ; Sun et al. , 2019 ) compress a large pre-trained model into a small or fast model with fixed size on the pre-training or fine-tuning stage and have achieved good compression efficiency and accuracy . However , from the perspective of practical deployment , a fixed size model can not be deployed in devices with different memory and latency constraints . For example , smaller models are preferred in embedded devices than in online servers , and faster 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 a 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 techniques including block-wise search , search space pruning and performance approximation during the search process : ( 1 ) We adopt block-wise search ( Li et al. , 2020a ) 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 constraints 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 on 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-BERT 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 GLUE benchmark datasets demonstrate the effectiveness of NAS-BERT for model compression . 2 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 et 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 ) aim to compress the pre-trained model into a fixed size of the model and have achieved a trade-off between the small parameter size ( usually no more than 66M ) and the good performance . However , these compressed models can not be deployed in devices with different memory and latency constraints . 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 , it 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 efficient models ( Liu et al. , 2018 ; Cai et al. , 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 on computer vision tasks and rely on specific designs on the convolutional layers ( e.g. , inverted bottleneck convolution ( Howard et al. , 2019 ) or elastic kernel size ( Cai et al. , 2019 ; Yu et al. , 2020 ) ) . Among them , once-for-all ( Cai et al. , 2019 ) and BigNAS ( Yu et al. , 2020 ) train a big supernet that contains all the candidate architectures and can get a specialized sub-network by selecting from the supernet to support various requirements ( e.g. , model size and latency ) . HAT ( Wang et al. , 2020a ) also trains a supernet with the adaptive widths and depths for machine translation tasks . Our proposed NAS-BERT also trains a 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 blockwise search , progressive shrinking , and performance approximation to reduce the training cost and improve search efficiency . Tsai et al . ( 2020 ) apply one-shot NAS to search a faster Transformer but they can not 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 resources to promising architectures and thus can deliver various architectures without adding much computation ; 2 ) designs bins in the shrinking algorithm to guarantee that we 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 . 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 downstream 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 3.3 ) . Due to the huge cost to train the big supernet on 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-forward 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 = { o1 , · · · , oC } , where C is the number of predefined candidate operations . Residual connection is applied to each layer by adding the input to the output . There areCL possible paths ( architectures ) in the supernet , and a specific architecture a = ( a1 , · · · , aL ) is a sub-net ( path ) in the supernet , where al ∈ O is the operation in the l-th layer , as shown in Fig . 2 ( a ) . We adopt weight sharing 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 . We further describe each operation in O as follows : 1 ) Multi-head attention ( MHA ) and feedforward network ( FFN ) , which are the two basic operations in Transformer and are popular in pretraining models ( in this way we can cover BERT model as a subnet in our supernet ) . 2 ) Separable convolution ( SepConv ) , whose effectiveness and efficiency in natural language processing tasks have been demonstrated by previous work ( Kaiser et al. , 2018 ; Karatzoglou et al. , 2020 ) . 3 ) Identity operation , which can support architectures with the number of layers less than L. Identity operation is regarded as a placeholder in a candidate architecture and can be removed to obtain a shallower network . More detailed considerations on choosing the operation set are in Appendix A.1 . 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 setO contains ( 1+1+3 ) ∗5+1 = 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 1 . The detailed structure of separable convolution is shown in Fig . 1 .

Summary: This 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.