R-Drop: Regularized Dropout for Neural Networks

1 Introduction . In recent years , deep learning has achieved remarkable success in various areas , e.g. , natural language processing , computer vision , speech/audio processing , etc . When training a deep neural network , regularization techniques [ 57 , 60 , 27 , 3 , 67 , 58 , 23 , 71 ] are indispensable to prevent over-fitting and improve the generalization ability of deep models . Among them , the dropout technique [ 24 ] , the most widely used one , aims to prevent co-adaptation and performs implicit ensemble by simply dropping a certain proportion of hidden units from the neural network during training . Existing literature [ 40 , 77 ] has revealed the possible side effect of dropout that there is an unnegligible inconsistency between training and inference stage of dropout models , i.e. , the randomly sampled sub model ( caused by dropout ) during training is inconsistent with the full model ( without dropout ) during inference . Through imposing L2 regularization on the inconsistent hidden states [ 40 , 77 ] , current methods can mitigate the inconsistency problem to some extent but are far from being widely used . In this paper , we introduce a simple yet more effective alternative to regularize the training inconsistency induced by dropout , named as R-Drop . Concretely , in each mini-batch training , each data sample goes through the forward pass twice , and each pass is processed by a different sub model by randomly dropping out some hidden units . R-Drop forces the two distributions for the same ∗Equal contribution ( listed in alphabetical order ) . 2https : //github.com/dropreg/R-Drop 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . data sample outputted by the two sub models to be consistent with each other , through minimizing the bidirectional Kullback-Leibler ( KL ) divergence between the two distributions . That is , R-Drop regularizes the outputs of two sub models randomly sampled from dropout for each data sample in training . In this way , the inconsistency between the training and inference stage can be alleviated . Compared with the dropout strategy in conventional neural network training , R-Drop only adds a KL-divergence loss without any structural modifications . From the perspective of deep neural network regularization , our proposed R-Drop can be treated as a new variation of dropout . Different from most of the previous methods that merely work on the hidden units of each layer ( e.g. , the standard dropout [ 24 ] ) or model parameters ( e.g. , dropconnect [ 60 ] ) , R-Drop works on both the hidden units and the output of sub models sampled by dropout , which is much more effective . We theoretically analyze the regularization effect of R-Drop , where the result shows that R-Drop can reduce the inconsistency existed in the training and inference . Though R-Drop regularization is simple , we find it is surprisingly effective through extensive experiments on 5 tasks with 18 datasets , spanning from natural language processing , including language modeling , neural machine translation , abstractive summarization , and language understanding , to computer vision , i.e. , image classification . It creates new records on multiple datasets , such as 30.91 BLEU score on WMT14 English→German and 43.95 on WMT14 English→French translation tasks while only be simply applied to the training of the vanilla Transformer , and also achieves SOTA results on the CNN/DailyMail summarization dataset . These universal improvements clearly demonstrate the effectiveness of R-Drop . Our main contributions are summarized as follows : • We propose R-Drop , a simple yet effective regularization method built upon dropout , which can be universally applied to train different kinds of deep models . • We theoretically show that our R-Drop can reduce the inconsistency between training and inference of the dropout based models . • Through extensive experiments on 4 NLP and 1 CV tasks with a total of 18 datasets , we show that R-Drop achieves extremely strong performances , including multiple SOTA results . 2 Approach . The overall framework of our R-Drop regularization method is shown in Figure 1 . Before elaborating on the details , we first present some necessary notations . Given the training datasetD = { ( xi , yi ) } ni=1 , the goal of the training is to learn a model Pw ( y|x ) , where n is the number of the training samples , ( xi , yi ) is the labeled data pair . xi is input data and yi is the label . For example , in NLP , xi can be the source language sentence in machine translation , and yi is the corresponding target language sentence . In CV , xi can be one image , and yi is the categorical class label . The probability distribution of the mapping function is also denoted as Pw ( y|x ) , and the Kullback-Leibler ( KL ) divergence between two distributions P1 and P2 is represented by DKL ( P1||P2 ) . In the following , we will explain our proposed R-Drop , training algorithm , and theoretical analysis , respectively . 2.1 R-Drop Regularization . We introduce our simple regularization method in this part . Given the training dataD = { ( xi , yi ) } ni=1 , the main learning objective for a deep learning model is to minimize the negative log-likelihood loss function , which is as follow : Lnll = 1 n n∑ i=1 − logPw ( yi|xi ) . ( 1 ) Since the deep neural networks are prone to over-fitting , regularization methods such as dropout [ 57 ] are usually adopted during training to reduce the generalization error of the model . Specifically , dropout randomly drops part of units in each layer of the neural network to avoid co-adapting and over-fitting . Besides , dropout also approximately performs to combine exponentially many different neural network architectures efficiently [ 57 ] , while model combination can always improve the model performance . Though simple and effective , there is a huge inconsistency between training and inference that hinders the model performance . That is , the training stage takes the sub model with randomly dropped units , while the inference phase adopts the full model without dropout . Also , the sub models caused by randomly sampled dropout units are also different without any constraints . Based on above observations and the randomness of the structure brought by dropout , we propose our R-Drop to regularize the output predictions of sub models from dropout . Concretely , given the input data xi at each training step , we feed xi to go through the forward pass of the network twice . Therefore , we can obtain two distributions of the model predictions , denoted as Pw1 ( yi|xi ) and Pw2 ( yi|xi ) . As discussed above , since the dropout operator randomly drops units in a model , the two forward passes are indeed based on two different sub models ( though in the same model ) . As shown in the right part of Figure 1 , the dropped units in each layer of the left path for the output prediction Pw1 ( yi|xi ) are different from that of the right path for output distribution Pw2 ( yi|xi ) . Thus the distributions of Pw1 ( yi|xi ) and Pw2 ( yi|xi ) are different for the same input data pair ( xi , yi ) . Then , at this training step , our R-Drop method tries to regularize on the model predictions by minimizing the bidirectional Kullback-Leibler ( KL ) divergence between these two output distributions for the same sample , which is : LiKL = 1 2 ( DKL ( Pw1 ( yi|xi ) ||Pw2 ( yi|xi ) ) +DKL ( Pw2 ( yi|xi ) ||Pw1 ( yi|xi ) ) ) . ( 2 ) With the basic negative log-likelihood learning objective LiNLL of the two forward passes : LiNLL = − logPw1 ( yi|xi ) − logPw2 ( yi|xi ) , ( 3 ) the final training objective is to minimize Li for data ( xi , yi ) : Li = LiNLL + α · LiKL =− logPw1 ( yi|xi ) − logPw2 ( yi|xi ) ( 4 ) + α 2 [ DKL ( Pw1 ( yi|xi ) ||Pw2 ( yi|xi ) ) +DKL ( Pw2 ( yi|xi ) ||Pw1 ( yi|xi ) ) ] , where α is the coefficient weight to control LiKL . In this way , our R-Drop further regularizes the model space beyond dropout and improves the generalization ability of a model . Compared Equation ( 1 ) with Equation ( 4 ) , our R-Drop only adds a KL-divergence loss LiKL based on two forward passes in training . Note that our regularization methodology can be universally applied on different model structures if there exists randomness in a model ( e.g. , dropout ) that can produce different sub models or outputs . We leave further explorations as future work . 2.2 Training Algorithm . The overall training algorithm based on our R-Drop is presented in Algorithm 1 . As introduced before , at each training step , Line 3-4 show that we go forward the model and obtain output distributions Pw1 ( y|x ) and Pw2 ( y|x ) , then Line 5-6 calculate the negative log-likelihood and the KL-divergence Algorithm 1 R-Drop Training Algorithm Input : Training data D = { ( xi , yi ) } ni=1 . Output : model parameter w. 1 : Initialize model with parameters w. 2 : while not converged do 3 : randomly sample data pair ( xi , yi ) ∼ D , 4 : repeat input data twice as [ xi ; xi ] and obtain the output distribution [ Pw1 ( yi|xi ) , Pw2 ( yi|xi ) ] , 5 : calculate the negative log-likelihood loss LiNLL by Equation ( 3 ) , 6 : calculate the KL-divergence loss LiKL by Equation ( 2 ) , 7 : update the model parameters by minimizing loss Li of Equation ( 4 ) . 8 : end while between the two distributions . It is worth nothing that we do not forward the input data twice , instead , we repeat the input data x and concatenate them ( [ x ; x ] ) in batch-size dimension , which can make forward procedure happen in the same mini-batch to save the training cost . Finally , the model parameters are updated ( Line 7 ) according to the loss of Equation ( 4 ) . The training will continue over the data epochs till convergence . Compared to the conventional training , our implementation is similar to enlarge the batch size to be double , and one potential limitation is that the computational cost of R-Drop increases at each step . As we show in Section 4.1 , similar to other regularization methods ( e.g. , training w/ or w/o dropout ) , though R-Drop needs more training to converge , the final optimum is much better with a superior performance . We also show another study of baseline with doubled batch size in Appendix C.1 . 2.3 Theoretical Analysis . We analyze the regularization effect of R-Drop in this subsection . Let hl ( x ) ∈ Rd denote the output of the l-th layer of a neural network with input vector x , and let ξl ∈ Rd denote a random vector , each dimension of which is independently sampled from a Bernoulli distribution B ( p ) : ξli = { 1 , with probability p , 0 , with probability 1-p. Then the dropout operation on hl ( x ) can be represented by hlξl ( x ) = 1 p ξl hl ( x ) , where denotes the element-wised product . Hence , the output distribution of the neural network with parameter w after applying dropout is Pwξ ( y|x ) : = softmax ( linear ( hLξL ( · · · ( h 1 ξ1 ( xξ0 ) ) ) ) ) , where ξ = ( ξ L , · · · , ξ0 ) . The objective for R-Drop enhanced training can be formulated as solving the following constrained optimization problem : min w 1 n n∑ i=1 Eξ [ − logPwξ ( yi|xi ) ] , ( 5 ) s.t . 1 n n∑ i=1 Eξ ( 1 ) , ξ ( 2 ) [ DKL ( Pwξ ( 1 ) ( yi|xi ) ||P w ξ ( 2 ) ( yi|xi ) ) ) ] ≤ . ( 6 ) More precisely , R-Drop optimizes the constrained optimization problem in Equation ( 5 ) and Equation ( 6 ) in a stochastic manner , i.e. , it samples two random vectors ξ ( 1 ) and ξ ( 2 ) ( corresponding to two dropout instantiations ) from Bernoulli distribution and one training instance ( xi , yi ) , and updates the parameter w according to the stochastic gradient∇wLi from Equation ( 4 ) . As we presented , one problem for dropout is the inconsistency between the training and inference models . Specifically , the training objective for dropout is the average loss of the sub models , i.e. , minw 1 n ∑n i=1 Eξ [ − logPwξ ( yi|xi ) ] , while the full model ( denoted as P̃w ( y|x ) ) is used for inference . Our proposed R-Drop enhanced training reduces this inconsistency by forcing the sub structures to be similar . The following proposition uses a linear model to demonstrate that with the constraint in Equation ( 6 ) , the inconsistency gap between the average loss of sub structures and the loss of the full model can be bounded ( detailed proof can be found in Appendix B ) . Proposition 2.1 . For a linear model Pw ( y|x ) = softmax ( Norm ( wTx ) ) where Norm ( · ) denotes the normalization layer and x ∈ Rd , with the constraint in Equation ( 6 ) in the main paper , we have |Lnll ( w ) − Eξ [ Lnll ( w , ξ ) ] | ≤ c √ , where Lnll ( w ) , Lnll ( w , ξ ) are the empirical loss calculated by the full model and a random sub model respectively , c is a constant related to the Liptschtz constant of the softmax operator .

Summary: This paper proposes a simple yet effective regularization method for deep learning models. The model performs two forward passes with different dropout masks for the same input and minimizes the difference between the two output distributions by the bidirectional KL divergence. The proposed method has been sufficiently validated on both CV and NLP benchmarks.