Regularization in ResNet with Stochastic Depth

1 Introduction . Stochastic Depth ( SD ) is a well-established regularization method that was first introduced by Huang et al . [ 2016 ] . It is similar in principle to Dropout [ Hinton et al. , 2012 , Srivastava et al. , 2014 ] and DropConnect [ Wan et al. , 2013 ] . It belongs to the family of noise-based regularization techniques , which includes other methods such as noise injection in data [ Webb , 1994 , Bishop , 1995 ] and noise injection throughout the network [ Camuto et al. , 2020 ] . While Dropout , resp . DropConnect consists of removing some neurons , resp . weights , at each iteration , SD randomly drops full layers , and only updates the weights of the resulting subnetwork at each training iteration . As a result of this mechanism , SD can be exclusively used with residual neural networks ( ResNets ) . There exists a stream of papers in the literature on the regularization effect of Dropout for linear models [ Wager et al. , 2013 , Mianjy and Arora , 2019 , Helmbold and Long , 2015 , Cavazza et al. , 2017 ] . Recent work by Wei et al . [ 2020 ] extended this analysis to deep neural networks using second-order perturbation analysis . It disentangled the explicit regularization of Dropout on the loss function and the implicit regularization on the gradient . Similarly , Camuto et al . [ 2020 ] studied the explicit regularization effect induced by adding Gaussian Noise to the activations and empirically illustrated the benefits of this regularization scheme . However , to the best of our knowledge , no analytical study of SD exists in the literature . This paper aims to fill this gap by studying the regularization effect of SD from an analytical point of view ; this allows us to derive principled guidelines on the choice of the survival probabilities for network layers . Concretely , our contributions are four-fold : • We show that SD acts as an explicit regularizer on the loss function by penalizing a notion of information discrepancy between keeping and removing certain layers . • We prove that the uniform mode , defined as the choice of constant survival probabilities , is related to maximum regularization using SD . • We study the large depth behaviour of SD and show that in this limit , SD mimics Gaussian Noise Injection by implicitly adding data-adaptive Gaussian noise to the pre-activations . ∗Equal contribution . Correspondence to : < soufiane.hayou @ yahoo.fr ; fadhel.ayed @ huawei.com > 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . • By defining the training budget L̄ as the desired average depth , we show the existence of two different regimes : small budget and large budget regimes . We introduce a new algorithm called SenseMode to compute the survival rates under a fixed training budget and provide a series of experiments that validates our Budget hypothesis introduced in Section 5 . 2 Stochastic Depth Neural Networks . Stochastic depth neural networks were first introduced by Huang et al . [ 2016 ] . They are standard residual neural networks with random depth . In practice , each block in the residual network is multiplied by a random Bernoulli variable δl ( l is the block ’ s index ) that is equal to 1 with some survival probability pl and 0 otherwise . The mask is re-sampled after each training iteration , making the gradient act solely on the subnetwork composed of blocks with δl = 1 . We consider a slightly different version where we apply the binary mask to the pre-activations instead of the activations . We define a depth L stochastic depth ResNet by y0 ( x ; δ ) = Ψ0 ( x , W0 ) , yl ( x ; δ ) = yl−1 ( x ; δ ) + δlΨl ( yl−1 ( x ; δ ) , Wl ) , 1 ≤ l ≤ L , yout ( x ; δ ) = Ψout ( yL ( x ; δ ) , Wout ) , ( 1 ) where Wl are the weights in the lth layer , Ψ is a mapping that defines the nature of the layer , yl are the pre-activations , and δ = ( δl ) 1≤l≤L is a vector of Bernoulli variables with survival parameters p = ( pl ) 1≤l≤L . δ is re-sampled at each iteration . For the sake of simplification , we consider constant width ResNet and we further denote by N the width , i.e . for all l ∈ [ L− 1 ] , yl ∈ RN . The output function of the network is given by s ( yout ) where s is some convenient mapping for the learning task , e.g . the Softmax mapping for classification tasks . We denote by o the dimension of the network output , i.e . s ( yout ) ∈ Ro which is also the dimension of yout . For our theoretical analysis , we consider a Vanilla model with residual blocks composed of a Fully Connected linear layer Ψl ( x , W ) = Wφ ( x ) , where φ ( x ) is the activation function . The weights are initialized with He init [ He et al. , 2015 ] , e.g . for ReLU , W lij ∼ N ( 0 , 2/N ) . There are no principled guidelines on choosing the survival probabilities . However , the original paper by Huang et al . [ 2016 ] proposes two alternatives that appear to make empirical consensus : the uniform and linear modes , described by Uniform : pl = pL , Linear : pl = 1− lL ( 1− pL ) , where we conveniently parameterize both alternatives using pL . Training budget L̄ . We define the training budget L̄ to be the desired average depth of the subnetworks with SD . The user typically fixes this budget , e.g. , a small budget can be necessary when training is conducted on small capacity devices . Depth of the subnetwork . Given the mode p = ( pl ) 1≤l≤L , after each iteration , the subnetwork has a depth Lδ = ∑L l=1 δl with an average Lp : = Eδ [ Lδ ] = ∑L l=1 pl . Given a budget L̄ , there is a infinite number of modes p such that Lp = L̄ . In the next lemma , we provide probabilistic bounds on Lδ using standard concentration inequalities . We also show that with a fixed budget L̄ , the uniform mode is linked to maximal variability . Lemma 1 ( Concentration of Lδ ) . For any β ∈ ( 0 , 1 ) , we have that with probability at least 1− β , |Lδ − Lp| ≤ vp u−1 ( log ( 2/β ) vp ) , ( 2 ) where Lp = E [ Lδ ] = ∑L l=1 pl , vp = Var [ Lδ ] = ∑L l=1 pl ( 1− pl ) , and u ( t ) = ( 1 + t ) log ( 1 + t ) − t. Moreover , for a given average depth Lp = L̄ , the upperbound in Eq . ( 2 ) is maximal for the uniform choice of survival probabilities p = ( L̄ L , ... , L̄ L ) . ( a ) Distributions of Lδ for a Resnet100 with average survival rate L̄/L = 0.5 for the uniform and linear modes . ( b ) Training time of Dropout and SD on CIFAR10 with ResNet56 for 100 epochs . Lemma 1 shows that with high probability , the depth of the subnetwork that we obtain with SD is within an ` 1 error of vp u −1 ( log ( 2/β ) vp ) from the average depth Lp . Given a fixed budget L̄ , this segment is maximized for the uniform mode p = ( L̄/L , . . . , L̄/L ) . Fig . 1a highlights this result . This was expected since the variance of the depth Lδ is also maximized by the uniform mode . The uniform mode corresponds to maximum entropy of the random depth , which would intuitively results in maximum regularization . We depict this behaviour in more details in Section 4 . SD vs Dropout . From a computational point of view , SD has the advantage of reducing the effective depth during training . Depending on the chosen budget , the subnetworks might be significantly shallower than the entire network ( Lemma 1 ) . This depth compression can be effectively leveraged for computational training time gain ( Fig . 1b ) . It is not the case with Dropout . Indeed , assuming that the choice of dropout probabilities is such that we keep the same number of parameters on average compared to SD , we still have to multiply matricesL times during the forward/backward propagation . It is not straightforward to leverage the sparsity obtained by Dropout for computational gain . In practice , Dropout requires an additional step of sampling the mask for every neuron , resulting in longer training times than without Dropout ( Fig . 1b ) . However , there is a trade-off between how small the budget is and the performance of the trained model with SD ( Section 6 ) . 3 Effect of Stochastic Depth at initialization . Empirical evidence strongly suggests that Stochastic Depth allows training deeper models [ Huang et al. , 2016 ] . Intuitively , at each iteration , SD updates only the parameters of a subnetwork with average depth Lp = ∑ l pl < L , which could potentially alleviate any exploding/vanishing gradient issue . This phenomenon is often faced when training ultra deep neural networks . To formalize this intuition , we consider the model ’ s asymptotic properties at initialization in the infinite-width limit N → +∞ . This regime has been the focus of several theoretical studies [ Neal , 1995 , Poole et al. , 2016 , Schoenholz et al. , 2017 , Yang , 2020 , Xiao et al. , 2018 , Hayou et al. , 2019 , 2020 , 2021a ] since it allows to derive analytically the distributions of different quantities of untrained neural networks . Specifically , randomly initialized ResNets , as well as other commonly-used architectures such as Fully connected Feedforward networks , convolutional networks and LSTMs , are equivalent to Gaussian Processes in the infinite-width limit . An important ingredient in this theory is the Gradient Independence assumption . Let us formally state this assumption first . Assumption 1 ( Gradient Independence ) . In the infinite width limit , we assume that the weightsW used for back-propagation are an iid version of the weightsW used for forward propagation . Assumption 1 is ubiquitous in the literature on the signal propagation in deep neural networks . It has been used to derive theoretical results on signal propagation in randomly initialized deep neural network [ Schoenholz et al. , 2017 , Poole et al. , 2016 , Yang and Schoenholz , 2017 , Hayou et al. , 2021b , a ] and is also a key tool in the derivation of the Neural Tangent Kernel [ Jacot et al. , 2018 , Arora et al. , 2019 , Hayou et al. , 2020 ] . Recently , it has been shown by Yang [ 2020 ] that Assumption 1 yields the exact computation for the gradient covariance in the infinite width limit . See Appendix A0.3 for a detailed discussion about this assumption . Throughout the paper , we provide numerical results that substantiate the theoretical results that we derive using this assumption . We show that Assumption 1 yields an excellent match between theoretical results and numerical experiments . Leveraging this assumption , Yang and Schoenholz [ 2017 ] , Hayou et al . [ 2021a ] proved that ResNet suffers from exploding gradient at initialization . We show in the next proposition that SD helps mitigate the exploding gradient behaviour at initialization in infinite width ResNets . Proposition 1 . Let φ = ReLU and L ( x , z ) = ` ( yout ( x ; δ ) , z ) for ( x , z ) ∈ Rd × Ro , where ` ( z , z′ ) is some differentiable loss function . Let q̃l ( x , z ) = EW , δ ‖∇ylL‖ 2 ‖∇yLL‖2 , where the numerator and denomi- nator are respectively the norms of the gradients with respect to the inputs of the lth and Lth layers . Then , in the infinite width limit , under Assumption 1 , for all l ∈ [ L ] and ( x , z ) ∈ Rd × Ro , we have • With Stochastic Depth , q̃l ( x , z ) = ∏L k=l+1 ( 1 + pk ) , • Without Stochastic Depth ( i.e . δ = 1 ) , q̃l ( x , z ) = 2L−l . Proposition 1 indicates that with or without SD , the gradient explodes exponentially at initialization as it backpropagates through the network . However , with SD , the exponential growth is characterized by the mode p. Intuitively , if we choose pl 1 for some layer l , then the contribution of this layer in the exponential growth is negligible since 1 + pl ≈ 1 . From a practical point of view , the choice of pl 1 means that the lth layer is hardly present in any subnetwork during training , thus making its contribution to the gradient negligible on average ( w.r.t δ ) . For a ResNet with L = 50 and uniform mode p = ( 1/2 , . . . , 1/2 ) , SD reduces the gradient exploding by six orders of magnitude . Fig . 2a and Fig . 2b illustrates the exponential growth of the gradient for the uniform and linear modes , as compared to the growth of the gradient without SD . We compare the empirical/theoretical growth rates of the magnitude of the gradient in Table 1 ; the results show a good match between our theoretical result ( under Assumption 1 ) and the empirical ones . Further analysis can be found in Section A4 . Stable ResNet . Hayou et al . [ 2021a ] have shown that introducing the scaling factor 1/ √ L in front of the residual blocks is sufficient to avoid the exploding gradient at initialization , as illustrated in Figure 2c . The hidden layers in Stable ResNet ( with SD ) are given by , yl ( x ; δ ) = yl−1 ( x ; δ ) + δl√ L Ψl ( yl−1 ( x ; δ ) , Wl ) , 1 ≤ l ≤ L. ( 3 ) The intuition behind the choice of the scaling factor 1/ √ L comes for the variance of yl . At initialization , with standard ResNet ( Eq . ( 1 ) ) , we have Var [ yl ] = Var [ yl−1 ] + Θ ( 1 ) , which implies that Var [ yl ] = Θ ( l ) . With Stable ResNet ( Eq . ( 3 ) ) , this becomes Var [ yl ] = Var [ yl−1 ] + Θ ( 1/L ) , resulting in Var [ yl ] = Θ ( 1 ) ( See Hayou et al . [ 2021a ] for more details ) . In the rest of the paper , we restrict our analysis to Stable ResNet ; this will help isolate the regularization effect of SD in the limit of large depth without any variance/gradient exploding issue . Nevertheless , the natural connection between SD and Dropout , coupled with the line of work on the regularization effect induced by the latter [ Wager et al. , 2013 , Mianjy and Arora , 2019 , Helmbold and Long , 2015 , Cavazza et al. , 2017 , Wei et al. , 2020 ] , would indicate that the benefits of SD are not limited to controlling the magnitude of the gradient . Using a second order Taylor expansion , Wei et al . [ 2020 ] have shown that Dropout induces an explicit regularization on the loss function . Intuitively , one should expect a similar effect with SD . In the next section , we elucidate the explicit regularization effect of SD on the loss function , and we shed light on another regularization effect of SD that occurs in the large depth limit .

Summary: The paper focuses on understanding the regularization effect of stochastic depth in residual neural networks. The authors give 3 main results:   (a) Stochastic depth helps to mitigate (to some extent) the gradient explosion phenomena observed at initialization in the Vanilla ResNet model,   (b) Stochastic depth mimics Gaussian Noise Injection (GNI) in deeper models by implicitly adding some non-isotropic Gaussian noise at each layer during training. Thereby, they show that Stochastic depth has the same regularization effects as GNI.  (c) They propose SenseMode, a principled method to choose the survival rates of each layer for training. SenseMode performs better than the existing methods in certain cases.  The authors corroborate all their propositions in an experimental study.