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 proposed to forward an input to neural networks twice with two different dropout masks, and apply a consistent loss as regularization, named R-Drop, between the two outputs to penalize the discrepancy. KL divergence of the two output probabilities is used as the consistent loss. Experiments on different NLP tasks demonstrate the effectiveness of R-Drop.