A uto T iny BERT : Automatic Hyper-parameter Optimization for Efficient Pre-trained Language Models (2021)

Abstract: Pre-trained language models (PLMs) have achieved great success in natural language processing. Most of PLMs follow the default setting of architecture hyper-parameters (e.g., the hidden dimension is a quarter of the intermediate dimension in feed-forward sub-networks) in BERT. Few studies have been conducted to explore the design of architecture hyper-parameters in BERT, especially for the more efficient PLMs with tiny sizes, which are essential for practical deployment on resource-constrained devices. In this paper, we adopt the one-shot Neural Architecture Search (NAS) to automatically search architecture hyper-parameters. Specifically, we carefully design the techniques of one-shot learning and the search space to provide an adaptive and efficient development way of tiny PLMs for various latency constraints. We name our method AutoTinyBERT and evaluate its effectiveness on the GLUE and SQuAD benchmarks. The extensive experiments show that our method outperforms both the SOTA search-based baseline (NAS-BERT) and the SOTA distillation-based methods (such as DistilBERT, TinyBERT, MiniLM, and MobileBERT). In addition, based on the obtained architectures, we propose a more efficient development method that is even faster than the development of a single PLM. The source code and models will be publicly available upon publication.

## introduction pre-trained language models, such as bert @xcite , roberta @xcite and xlnet @xcite , have become prevalent in natural language processing. to improve model performance, most plms (e.g. electra @xcite and gpt-2/3 @xcite ; under the same speedup constraint, our method outperforms both the default hyper-parameter setting of bert @xcite , pf @xcite ) and nas-bert @xcite . more details are in the section 4.2. brown et al., 2020)) follow the default rule of hyper-parameter setting foot_0 in bert to scale up their model sizes. due to its simplicity, this rule has been widely used and can help large plms obtain promising results @xcite . in many industrial scenarios, we need to deploy plms on resource-constrained devices, such as smartphones and servers with limited computation power. due to the expensive computation and slow inference speed, it is usually difficult to deploy plms such as bert (12/24 layers, 110m/340m parameters) and gpt-2 (48 layers, 1.5b parameters) at their original scales. therefore, there is an urgent need to develop plms with smaller sizes which have lower computation cost and inference latency. in this work, we focus on a specific type of efficient plms, which we define to have inference time less than 1/4 of bert-base. 3 although, there have been quite a few work using knowledge distillation to build small plms @xcite @xcite , all of them focus on the application of distillation techniques @xcite and do not study the effect of architecture hyper-parameter settings on model performance. recently, neural architecture search and hyper-parameter optimization @xcite have been widely explored in machine learning, mostly in computer vision, and have been proven to find better designs than heuristic ones. inspired by this research, one problem that naturally arises is can we find better settings of hyper-parameters foot_3 for efficient plms? in this paper, we argue that the conventional hyper-parameter setting is not best for efficient plms (as shown in figure 1 ) and introduce a method to automatically search for the optimal hyper-parameters for specific latency constraints. pre-training efficient plms is inevitably resourceconsuming @xcite . therefore, it is infeasible to directly evaluate millions of architectures. to tackle this challenge, we introduce the one-shot neural architecture search (nas) @xcite @xcite to perform the automatic hyper-parameter optimization on efficient plms, named as autotinybert. specifically, we first use the one-shot learning to obtain a big superplm, which can act as proxies for all potential sub-architectures. proxy means that when evaluating an architecture, we only need to extract the corresponding sub-model from the superplm, instead of training the model from scratch. super-plm helps avoid the time-consuming pre-training process and makes the search process efficient. to make superplm more effective, we propose practical techniques including the head sub-matrix extraction and efficient batch-wise training, and particularly limit the search space to the models with identical layer structure. furthermore, by using superplm, we leverage search algorithm @xcite ) to find hyperparameters for various latency constraints. in the experiments, in addition to the pre-training setting @xcite , we also consider the setting of task-agnostic bert distillation @xcite ) that pre-trains with the loss of knowledge distillation, to build efficient plms. exten-sive results show that in pre-training setting, au-totinybert not only consistently outperforms the bert with conventional hyper-parameters under different latency constraints, but also outperforms nas-bert based on neural architecture search. in task-agnostic bert distillation, autotinybert outperforms a series of existing sota methods of distilbert, tinybert and mobilebert. our contributions are three-fold: (1) we explore the problem of how to design hyper-parameters for efficient plms and introduce an effective and efficient method: autotinybert; (2) we conduct extensive experiments in both scenarios of pretraining and knowledge distillation, and the results show our method consistently outperforms baselines under different latency constraints; (3) we summarize a fast rule and it develops an autotiny-bert for a specific constraint with even about 50% of the training time of a conventional plm. ## preliminary before presenting our method, we first provide some details about the transformer layer @xcite to introduce the conventional hyperparameter setting. transformer layer includes two sub-structures: the multi-head attention (mha) and the feed-forward network (ffn). for clarity, we show the mha as a decomposable structure, where the mha includes h individual and parallel self-attention modules (called heads). the output of mha is obtained by summing the output of all heads. specifically, each head is represented by four main matrices figure 2 : overview of autotinybert. we first train an effective superplm with one-shot learning, where the objectives of pre-training or task-agnostic bert distillation are used. then, given a specific latency constraint, we perform an evolutionary algorithm on the superplm to search optimal architectures. finally, we extract the corresponding sub-models based on the optimal architectures and further train these models. in the conventional setting of the hyper-parameters in bert, all dimensions of matrices are the same as the dimension of the hidden vector, namely, d q|k|v|@xmath0. in fact, there are only two requirements of d @xmath1. transformer layer also contains an ffn that is stacked on the mha, that is: similarly, there are modules of residual connection and layer normalization on top of ffn. in the original transformer, d @xmath2. thus, we conclude that the conventional hyper-parameter setting follows the rule of {d q|k|v|@xmath3, d @xmath4. ## related work efficient plms with tiny sizes. there are two widely-used methods for building efficient plms: pre-training and model compression. knowledge distillation (kd) @xcite is the most widely studied technique in plm compression, which uses a teacher-student framework. the typical distillation studies include distilbert @xcite , bert-pkd @xcite , minilm @xcite , mobile-bert @xcite , minibert @xcite and etd @xcite . in addition to kd, the techniques of pruning @xcite , quantization @xcite @xcite and parameter sharing @xcite introduced for plm compression. our method is orthogonal to the building method of efficient plm and is trained under the settings of pre-training and task-agnostic bert distillation, which can be used by direct finetuning. nas for nlp. nas is extensively studied in computer vision @xcite , but relatively little studied in the natural language processing. evolved transformer @xcite and hat @xcite search architecture for transformer-based neural machine translation. for bert distillation, adabert @xcite focuses on searching the architecture in the fine-tuning stage and relies on data augmentation to improve its performance. schubert @xcite obtains the optimal structures of plm by a pruning method. a work similar to ours is nas-bert @xcite . it proposes some techniques to tackle the challenging exponential search space of non-identical layer structure and heterogeneous modules. our method adopts a linear search space and introduces several practical techniques for superplm training. moreover, our method is efficient in terms of computation cost and the obtained plms are easy to use. ## conclusion we propose an effective and efficient method au-totinybert to search for the optimal architecture hyper-parameters of efficient plms. we evaluate the proposed method in the scenarios of both the pre-training and task-agnostic bert distillation. the extensive experiments show that autotiny-bert can consistently outperform the baselines under different latency constraints. furthermore, we develop a fast development rule for efficient plms which can build an autotinybert model even with less training time of a conventional one.