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: The paper proposes a method, LayerDrop, for pruning layers in Transformer based models. The goal is to explore the stochastic depth of transformer models during training in order to do efficient layer pruning at inference time. The key idea is simple and easy to understand: randomly dropping transformer layers during training to make the model robust to subsequent pruning. The authors perform empirical studies on several sequence modeling task to conclude that the proposed approach allows efficient pruning of deeper models into shallow ones without fine-tuning on downstream tasks. There are also empirical experiments done to demonstrate that the proposed approach outperforms recent model pruning techniques such as DistillBERT under comparable configurations.