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: This paper studies the regularization effects involved in the stochastic depth training of ResNets. Specifically, by applying the second-order Taylor approximation to the loss function and the intermediate features, the authors show that the stochastic depth training involves a regularization term that promotes flatness and acts like Gaussian noise to the features. Based on this theoretical insight, the authors suggest a strategy for setting the survival probabilities of sub-networks depending on the computational budget, which corresponds to controlling the regularization effect. This strategy developed upon the new insight in the paper is validated by empirical experiments.