Reducing Transformer Depth on Demand with Structured Dropout

1 INTRODUCTION . Transformer architectures ( Vaswani et al. , 2017 ) have become the dominant architecture in natural language processing , with state-of-the-art performance across a variety of tasks , including machine translation ( Vaswani et al. , 2017 ; Ott et al. , 2018 ) , language modeling ( Dai et al. , 2019 ; Baevski & Auli , 2018 ) and sentence representation ( Devlin et al. , 2018 ; Yang et al. , 2019 ) . Each of its layers contains millions of parameters accessed during the forward pass , making it computationally demanding in terms of memory and latency during both training and inference . In an ideal situation , we would be able to extract sub-networks — automatically and without finetuning — from this over-parameterized network , for any given memory or latency constraint , while maintaining good performance . In contrast , standard pruning or distillation methods follow a strategy that often includes a finetuning or retraining step , and the process must be repeated for each desired depth . In this work , we propose a novel approach to extract any sub-network without a post-hoc pruning process from over-parameterized networks . The core of our method is to sample small sub-networks from the larger model during training by randomly dropping model weights as in Dropout ( Hinton et al. , 2012 ) or DropConnect ( Wan et al. , 2013 ) . This has the advantage of making the network robust to subsequent pruning . If well-chosen groups of weights are dropped simultaneously , the resulting small sub-networks can be very efficient . In particular , we drop entire layers to extract shallow models at inference time . Previous work ( Huang et al. , 2016 ) has shown that dropping layers during training can regularize and reduce the training time of very deep convolutional networks . In contrast , we focus on pruning . As illustrated in Figure 1 , an advantage of our layer dropping technique , or LayerDrop , is that from one single deep model , we can extract shallow sub-networks of any desired depth on demand at inference time . We validate our findings on a variety of competitive benchmarks , namely WMT14 EnglishGerman for machine translation , WikiText-103 ( Merity et al. , 2016 ) for language modeling , CNNDailymail ( Hermann et al. , 2015 ) for abstractive summarization , ELI5 ( Fan et al. , 2017 ) for long form question answering , and several natural language understanding tasks ( Wang et al. , 2019a ) for sentence representation . Our approach achieves state of the art on most of these benchmarks as a result of the regularization effect , which stabilizes the training of larger and deeper networks . We also show that we can prune Transformer architectures to much smaller models while maintaining com- petitive performance , outperforming specific model reduction strategies dedicated to BERT ( Devlin et al. , 2018 ; Sanh , 2019 ) as well as training smaller models from scratch . Overall , applying LayerDrop to Transformer networks provides the following key advantages : • LayerDrop regularizes very deep Transformers and stabilizes their training , leading to stateof-the-art performance across a variety of benchmarks . • Small and efficient models of any depth can be extracted automatically at test time from a single large pre-trained model , without the need for finetuning . • LayerDrop is as simple to implement as dropout . 2 RELATED WORK . Our approach is a form of Dropout ( Srivastava et al. , 2014 ) applied to model weights instead of activations , as in DropConnect ( Wan et al. , 2013 ) . Different from DropConnect , we drop groups of weights to induce group redundancy to create models suited for pruning to shallow , efficient models at inference time . Gomez et al . ( 2018 ) propose a targeted Dropout and DropConnect , where they learn the drop rate of the weights to match a targeted pruning scheme . Instead , we adapt the masks to the structures that we are interested in pruning . Closer to our work , the Stochastic Depth approach of Huang et al . ( 2016 ) drops layers randomly during training . As opposed to our work , they are interested in accelerating the training of very deep ResNets ( He et al. , 2016 ) , so their dropping schedule is adapted to this goal . Concurrently to this work , Pham et al . ( 2019 ) applied Stochastic Depth to train very deep Transformers for speech and show the benefits of its regularization effect . More generally , our method is a form of structured pruning ( Liu et al. , 2018b ) . As opposed to weight pruning ( LeCun et al. , 1990 ) , structured pruning removes coherent groups of weights to preserve the original structure of the network . Structured pruning has been used in some NLP applications , such as machine translation ( See et al. , 2016 ) , text classification ( Joulin et al. , 2016 ) and language modeling ( Murray & Chiang , 2015 ) . However , it has been more widely adopted in computer vision and applied to convolutional network to remove filters ( Li et al. , 2016 ; Wen et al. , 2016 ) , channels ( He et al. , 2017 ) , or residual blocks ( Huang et al. , 2018 ; Huang & Wang , 2018 ) . Similar to Mittal et al . ( 2018 ) , we take advantage of the plasticity of neural networks to learn models that are resilient to random pruning or skipping connections Wang et al . ( 2018 ) ; Wu et al . ( 2018 ) ; Liu et al . ( 2018a ) , rather than learning the pruning itself . We refer the reader to Liu et al . ( 2018b ) for an exhaustive study of these approaches and their evaluation in the context of convolutional networks . Reducing the memory footprint of Transformer architectures and BERT in particular is an active subject of research . Several works have compressed BERT as a post-processing step using different forms of distillation ( Turc et al. , 2019 ; Tang et al. , 2019 ; Shulga , 2019 ; Sanh , 2019 ) . Similarly , various papers have shown evidence that Transformers are over-parameterized , especially that most self-attention heads can be dropped at test time ( Michel et al. , 2019 ; Voita et al. , 2019 ) . Different from these , our models are trained to be resilient to pruning , which significantly reduces the performance drop induced by test time pruning . Others have proposed trainable adaptive mechanisms to control their memory footprint ( Jernite et al. , 2016 ; Sukhbaatar et al. , 2019 ; Correia et al. , 2019 ) . These approaches are complementary to ours and should benefit from each other . 3 METHOD . In this section , we briefly introduce the Transformer , then describe our Structured Dropout technique and its application to layers . We also discuss several inference time pruning strategies . 3.1 THE TRANSFORMER ARCHITECTURE . We succinctly review the Transformer architecture and refer the reader to Vaswani et al . ( 2017 ) for additional details . A Transformer is a stack of layers composed of two sub-layers : multi-head self-attention followed by a feedforward sub-layer . The multi-head self-attention sub-layer consists of multiple attention heads applied in parallel . Each attention head takes a matrix X where each row represents an element of the input sequence and updates their representations by gathering information from their context using an Attention mechanism ( Bahdanau et al. , 2014 ) : Y = Softmax ( XTK ( QX+P ) ) VX , where K , V , Q and P are matrices of parameters . The outputs of the heads are then concatenated along the time step into a sequence of vectors . The second sub-layer then applies a fully connected feedforward network to each element of this sequence independently , FFN ( x ) = U ReLU ( Vx ) , where V and U are matrices of parameters . Each sub-layer is followed by a AddNorm operation that is a residual connection ( He et al. , 2016 ) and a layer normalization ( Ba et al. , 2016 ) . 3.2 TRAINING TRANSFORMERS WITH RANDOM STRUCTURED PRUNING . We present a regularization approach that makes Transformers robust to subsequent structured pruning at inference time . We focus in particular on the case where the targeted structure is a layer . 3.2.1 RANDOMLY DROPPING STRUCTURES AT TRAINING TIME . Regularizing networks to be robust to pruning can be achieved by randomly removing weights during its training as in DropConnect ( Wan et al. , 2013 ) . In this approach , each weight is dropped independently following a Bernoulli distribution associated with a parameter p > 0 that controls the drop rate . This is equivalent to a pointwise multiplication of the weight matrix W with a randomly sampled { 0 , 1 } mask matrix M : Wd = M W. DropConnect is a form of random unstructured pruning that leads to smaller , but not necessarily more efficient , models . We propose to add structure to this mechanism to target model efficiency . Random Structured Dropout . The weights of a Transformer network belong to multiple overlapping structures , such as heads , FFN matrices , or layers . Dropping weights using groups that follow some of these inherent structures potentially leads to a significant reduction of the inference time . This is equivalent to constraining the mask M to be constant over some predefined groups of weights . More precisely , given a set G of predefined groups of weights , the { 0 , 1 } mask matrix M is randomly sampled over groups instead of weights : ∀i , M [ i ] ∈ { 0 , 1 } , and ∀G ∈ G , ∀ ( i , j ) ∈ G , M [ i ] = M [ j ] . This structured dropout formulation is general and can be applied to any overlapping groups of weights , whether heads , FFN matrices , or layers . Nonetheless , not all of the structures in a Transformer lead to the same benefits when dropped . For example , dropping attention heads does not reduce runtime as they are usually computed in parallel . For simplicity , we focus on dropping layers , and we name this structured pruning , LayerDrop . This is inspired by the Stochastic Depth approach of Huang et al . ( 2016 ) used to train very deep ResNets ( He et al. , 2015 ) . 3.2.2 PRUNING AT INFERENCE TIME . Selecting Layers to Prune Training with LayerDrop makes the network more robust to predicting with missing layers . However , LayerDrop does not explicitly provide a way to select which groups to prune . We consider several different pruning strategies , described below : • Every Other : A straightforward strategy is to simply drop every other layer . Pruning with a rate p means dropping the layers at a depth d such that d ≡ 0 ( modb 1pc ) . This strategy is intuitive and leads to balanced networks . • Search on Valid : Another possibility is to compute various combinations of layers to form shallower networks using the validation set , then select the best performing for test . This is straightforward but computationally intensive and can lead to overfitting on validation . • Data Driven Pruning : Finally , we propose data driven pruning where we learn the drop rate of each layer . Given a target drop rate p , we learn an individual drop rate pd for the layer at depth d such that the average rate over layers is equal to p. More precisely , we parameterize pd as a non-linear function of the activation of its layer and apply a softmax . At inference time , we forward only the fixed top-k highest scoring layers based on the softmax output ( e.g . chosen layers do not depend on the input features ) . In practice , we observe that the Every Other strategy works surprisingly well across many tasks and configurations . Search on Valid and Data Driven Pruning only offer marginal gains . Note that we do not further finetune any of the pruned networks ( see Appendix for analysis of finetuning ) . Setting the drop rate for optimal pruning . There is a straightforward relationship between the drop rate of groups and the average pruning level that the network should be resilient to . Assuming N groups and a fixed drop ratio p , the average number of groups used by the network during training is N ( 1− p ) . As a consequence , to target a pruning size of r groups , the optimal drop rate is : p∗ = 1− r N In practice , we observe that networks are more robust to pruning than their expected ratio but higher pruning rates leads to better performance for smaller models . We use a LayerDrop rate of p = 0.2 for all our experiments , but we recommend p = 0.5 to target very small inference time models .

Summary: This work explored the effect of LayerDrop training in efficient pruning at inference time. The authors showed that it is possible to have comparable performance from sub-networks of smaller depth selected from one large network without additional finetuning. More encouraging is that the sub-networks are able to perform better than the same network trained from scratch or learned based on distillation.