Length-Adaptive Transformer: Train Once with Length Drop, Use Anytime with Search

1 INTRODUCTION . Pretrained language models ( Peters et al. , 2018 ; Devlin et al. , 2018 ; Radford et al. , 2019 ; Yang et al. , 2019 ) have achieved notable improvements in various natural language processing ( NLP ) tasks . Most of them rely on transformers ( Vaswani et al. , 2017 ) , and the number of model parameters ranges from hundreds of millions to billions ( Shoeybi et al. , 2019 ; Raffel et al. , 2019 ; Kaplan et al. , 2020 ; Brown et al. , 2020 ) . Despite this high accuracy , excessive computational overhead during inference , both in terms of time and memory , has hindered its use in real applications . This level of excessive computation has further raised the concern over energy consumption as well ( Schwartz et al. , 2019 ; Strubell et al. , 2019 ) . Recent studies have attempted at addressing these concerns regarding large-scale transformers ’ computational and energy efficiency ( see §6 for a more extensive discussion . ) Among these , we focus on PoWER-BERT ( Goyal et al. , 2020 ) which progressively reduces sequence length by eliminating word-vectors based on the attention values as passing layers . PoWER-BERT establishes the superiority of accuracy-time trade-off over earlier approaches ( Sanh et al. , 2019 ; Sun et al. , 2019 ; Michel et al. , 2019 ) . It however requires us to train a separate model for each efficiency constraint . In this paper , we thus develop a framework based on PoWER-BERT such that we can train a single model that can be adapted in the inference time to meet any given efficiency target . In order to train a transformer to cope with a diverse set of computational budgets in the inference time , we propose to train one while reducing the sequence length with a random proportion at each layer . We refer to this procedure as LengthDrop which was motivated by the nested dropout ( Rippel et al. , 2014 ) . We can extract sub-models of shared weights with any length configuration without requiring extra post-processing nor additional finetuning . Once a transformer is trained with the proposed LengthDrop , we search for the length configuration that maximizes the accuracy given a computational budget . Because this search is combinatorial and has multiple objectives ( accuracy and efficiency ) , we use an evolutionary search algorithm , which further allows us to obtain a full Pareto frontier of accuracy-efficiency trade-off of each model . It is not trivial to find an optimal length configuration given the inference-time computational budget , although it is extremely important in order to deploy these large-scale transformers in practice . In this work , we propose to use evolutionary search to find a length configuration that maximizes the accuracy within a given computational budget . We can further compute the Pareto frontier of accuracy-efficiency trade-off to obtain a sequence of length configurations with varying efficiency profiles . PoWER-BERT , which forms the foundation of the proposed two-stage procedure , is only applicable to sequence-level classification , because by design it eliminates some of the word vectors at each layer . In other words , it can not be used for token-level tasks such as span-based question answering ( Rajpurkar et al. , 2016 ) , because these tasks require hidden representations of the entire input sequence at the final layer . We thus propose to extend PoWER-BERT with a novel Drop-and-Restore process ( §3.3 ) , which eliminates this inherent limitation . Word vectors are dropped and set aside , rather than eliminated , in intermediate layers to maintain the saving of computational cost , as was with the original PoWER-BERT . These set-aside vectors are then restored at the final hidden layer and provided as an input to a subsequent task-specific layer , which is unlike the original PoWER-BERT . The main contributions of this work are two-fold . First , we introduce LengthDrop , a structured variant of dropout for training a single Length-Adaptive Transformer model that allows us to automatically derive multiple sub-models with different length configurations in the inference time using evolutionary search , without requiring any re-training . Second , we design Drop-and-Restore process that makes PoWER-BERT applicable beyond classification , which enables PoWER-BERT to be applicable to a wider range of NLP tasks such as span-based question answering . We empirically verify Length-Adaptive Transformer works quite well using the variants of BERT on a diverse set of NLP tasks , including SQuAD 1.1 ( Rajpurkar et al. , 2016 ) and two sequence-level classification tasks in GLUE benchmark ( Wang et al. , 2018 ) . Our experiments reveal that the proposed approach grants us a fine-grained control of computational efficiency and a superior accuracy-efficiency trade-off in the inference time , compared to existing approaches . 2 BACKGROUND : TRANSFORMERS AND POWER-BERT . Before we describe our main approach , we review some of the building blocks in this section . In particular , we review transformers , which are a standard backbone used in natural language processing these days , and PoWER-BERT , which was recently proposed as an effective way to train a large-scale , but highly efficient transformer for sequence-level classification . 2.1 TRANSFORMERS AND BERT . A transformer is a particular neural network that has been designed to work with a variable-length sequence input and is implemented as a stack of self-attention and fully-connected layers ( Vaswani et al. , 2017 ) . Here , we give a brief overview of the transformer which is the basic building block of the proposed approach . Each token xt in a sequence of tokens x = ( x1 , . . . , xN ) , representing input text , is first turned into a continuous vector h0t ∈ RH which is the sum of the token and position embedding vectors . This sequence is fed into the first transformer layer which returns another sequence of the same length h1 ∈ RN×H . We repeat this procedure L times , for a transformer with L layers , to obtain hL = ( hL1 , . . . , h L N ) . We refer to each vector in the hidden sequence at each layer as a word vector to emphasize that there exists a correspondence between each such vector and one of the input words . Although the transformer was first introduced for the problem of machine translation , Devlin et al . ( 2018 ) demonstrated that the transformer can be trained and used as a masked language model . More specifically , Devlin et al . ( 2018 ) showed that the transformer-based masked language model , called BERT , learns a universally useful parameter set that can be finetuned for any downstream task including sequence-level and token-level classification . In the case of sequence-level classification , a softmax classifier is attached to the word vector hL1 associated with the special token [ CLS ] , and the entire network , including the softmax classifier and BERT , is finetuned . For token-level classification , we use each hLt as the final hidden reprsentation of the associated t-th word in the input sequence . This strategy of pretraining followed by finetuning , often referred to as transfer learning , has recently become a dominant approach to classification in natural language processing . 2.2 POWER-BERT . PoWER-BERT keeps only the topmost lj word vectors at each layer j by eliminating redundant ones based on the significance score , which is the total amount of attention imposed by a word on the other words ( Goyal et al. , 2020 ) . lj is the hyper-parameter determines how many vectors to keep at layer j. PoWER-BERT has the same model parameters with BERT , but the extraction layers are interspersed after the self-attention layer in every transformer blocks ( Vaswani et al. , 2017 ) . PoWER-BERT reduces inference time successfully , achieving better accuracy-time trade-off than DistilBERT ( Sanh et al. , 2019 ) , BERT-PKD ( Sun et al. , 2019 ) , and Head-Prune ( Michel et al. , 2019 ) . Despite the original intention of maximizing the inference efficiency with the minimal loss in accuracy , it is possible to set up PoWER-BERT to be both more efficient and more accurate compared to the original BERT , which was observed but largely overlooked by Goyal et al . ( 2020 ) . Training a PoWER-BERT model consists of three steps : ( 1 ) finetuning , ( 2 ) length configuration search , and ( 3 ) re-training . The finetuning step is just like the standard finetuning step of BERT given a target task . A length configuration is a sequence of retention parameters ( l1 , · · · lL ) , each of which corresponds to the number of word vectors that are kept at each layer . These retention parameters are learned along with all the other parameters to minimize the original task loss together with an extra term that approximately measures the number of retained word vectors across layers . In the re-training step , PoWER-BERT is finetuned with the length configuration fixed to its learned one . For each computational budget , we must train a separate model going through all three steps described above . Moreover , the length configuration search step above is only approximate , as it relies on relaxation of retention parameters which are inherently discrete . This leads to the lack of guaranteed correlation between the success of this stage and true run-time . Even worse , it is a delicate act to tune the length configuration given a target computational budget , because trade-off is implicitly made via a regularization coefficient . Furthermore , PoWER-BERT has an inherent limitation in that it only applies to sequence-level classification because it eliminates word vectors in intermediate layers . 3 LENGTH-ADAPTIVE TRANSFORMER . In this section , we explain our proposed framework which results in a transformer that reduces the length of a sequence at each layer with an arbitrary rate . We call such a resulting transformer a Length-Adaptive Transformer . We train Length-Adaptive Transformer with LengthDrop which randomly samples the number of hidden vectors to be dropped at each layer with the goal of making the final model robust to such drop in the inference time . Once the model is trained , we search for the optimal trade-off between accuracy and efficiency using multi-objective evolutionary search , which allows us to use the model for any given computational budget without finetuning nor re-training . At the end of this section , we describe Drop-and-Restore process as a way to greatly increase the applicability of PoWER-BERT which forms a building block of the proposed framework . In short , we train a Length-Adaptive Transformer once with LengthDrop and Drop-and-Restore , and use it with an automatically determined length configuration for inference with any target computational budget , on both sequence-level and token-level tasks . 3.1 LENGTHDROP . Earlier approaches to efficient inference with transformers have focused on a scenario where the target computational budget for inference is known in advance Sanh et al . ( 2019 ) ; Goyal et al . ( 2020 ) . This greatly increases the cost of deploying transformers , as it requires us to train a separate transformer for each scenario . Instead , we propose to train one model that could be used for a diverse set of target computational budgets without retraining . LengthDrop randomly generates a length configuration by sequentially sampling a sequence length li+1 at the ( i+ 1 ) -th layer based on the previous layer ’ s sequence length li , following the uniform distribution U ( d ( 1 − p ) lie , li ) , where l0 is set to the length of the input sequence , and p is the LengthDrop probability . This sequential sampling results in a length configuration ( l1 , · · · , lL ) . Length-Adaptive Transformer can be thought of as consisting of a full model and many sub-models corresponding to different length configuration , similarly to a neural network trained with dropout ( Srivastava et al. , 2014 ) . LayerDrop From the perspective of each word vector , the proposed LengthDrop could be thought of as skipping the layers between when it was set aside and the final layer where it was restored . The word vector however does not have any information based on which it can determine whether it would be dropped at any particular layer . In our preliminary experiments , we found that this greatly hinders optimization . We address this issue by using LayerDrop ( Fan et al. , 2019 ) which skips each layer of a transformer uniformly at random . The LayerDrop encourages each word vector to be agnostic to skipping any number of layers between when it is dropped and when it is restored , just like dropout ( Srivastava et al. , 2014 ) prevents hidden neurons from co-adapting with each other by randomly dropping them . Sandwich Rule and Inplace Distillation We observed that standard supervised training with LengthDrop does not work well in the preliminary experiments . We instead borrow a pair of training techniques developed by Yu & Huang ( 2019 ) which are sandwich rule and inplace distillation , for better optimization as well as final generalization . At each update , we update the full model without LengthDrop as usual to minimize the supervised loss function . We simultaneously update ns randomly-sampled sub-models ( which are called sandwiches ) and the smallest-possible submodel , which corresponds to keeping only d ( 1− p ) lie word vectors at each layer i , using knowledge distillation ( Hinton et al. , 2015 ) from the full model . Here , sub-models mean models with length reduction . They are trained to their prediction close to the full model ’ s prediction ( inplace distillation ) .

Summary: This work introduces a method, called LengthDrop, to train a Length-Adaptive Transformer that supports adaptive model architecture based on different latency constraints. In order to make the model robust to variable input lengths, the method stochastically reduces the length of a sequence at each layer during training. Once the model is trained, the method uses an evolutionary search to find subnetworks that maximize model accuracy under a latency budget.