POP-Norm: A Theoretically Justified and More Accelerated Normalization Approach

1 INTRODUCTION . The last decade has witnessed significant advances in deep neural networks ( DNN ) , which brought substantial improvements for many real-world tasks , such as object ( Szegedy et al. , 2015 ; He et al. , 2015 ) , scene ( Khan et al. , 2016 ; Guo et al. , 2017 ) and action recognition ( Feichtenhofer et al. , 2017 ) , object detection ( Ren et al. , 2015 ; He et al. , 2017b ; Redmon & Farhadi , 2018 ) and image segmentation ( Ronneberger et al. , 2015 ; He et al. , 2017a ) . BatchNorm ( Ioffe & Szegedy , 2015 ) is a milestone technique in the development of DNN , and nowadays it has become a default component to construct modern deep networks . BatchNorm is widely-used in DNN due to that it is able to reduce the sensitivity to initialization , significantly raise the learning rate , substantially speed up the training process and considerably improve the performance . Actually , along with this line a variety of works , such as layer normalization ( Ba et al. , 2016 ) , instance normalization ( Ulyanov et al. , 2016 ) , weight normalization ( Salimans & Kingma , 2016 ) and group normalization ( Wu & He , 2018 ) , batch renormalization , ( Ioffe , 2017 ) were sequently proposed and also have made great success . Nowadays , the practical success of BatchNorm is mainly attributed to reduction of internal covariate shift ( ICS ) by controlling the first two moments ( mean and variance ) of the distributions of layer inputs . Recently , this point of view is challenged by ( Santurkar et al. , 2018 ) . It points out that the link between reduction of ICS and performance improvement of BatchNorm is tenuous , at best . However , the success of BatchNorm is indisputable , therefore there may exist a profound mathematical mechanism behind BatchNorm . In fact , from the perspective of loss landscape , ( Santurkar et al. , 2018 ) demonstrates BatchNorm will make the optimization landscape smoother , and then this smoothness induces a more predictive and stable behavior of the gradients . As stated above , the original theoretical foundation of the vanilla BatchNorm is not so solid and the exact reason for its effectiveness is poorly understood , hence BatchNorm maybe not optimal in theory and we might not make better use of it . In light of this , we try to propose a theoretically justified normalization approach for accelerating training deep neural networks , so that we can continuously further algorithmic progress . To achieve this goal , we firstly demonstrate that lowering gradient Lipschitz constant and gradient variance is the key to boost the convergence rate of optimization algorithms with stochastic gradient methods . And then we construct a normalization approach – removing the mean of inputs and dividing by the scaled l2 norm of inputs before conducting the main operation ( convolution or inner-production ) of a layer , which can be theoretically ensured to reduce the gradient Lipschitz constant and gradient variance . The new normalization approach and BatchNorm are similar . Hence we utilize their similarity to interpret the effectiveness of the vanilla BatchNorm , along the line of lowering gradient Lipschitz constant and gradient variance , which is different from the view of the loss landscape ( Santurkar et al. , 2018 ) . According to our theory , we can easily explain why BatchNorm with ReLU is powerful with but ineffective with Sigmoid and Tanh as activation functions , which may help us to better understand the underlying complexity of neural networks . The main contributions of this paper are summarized as follows : • From scratch a theoretically justified normalization approach is deduced for accelerating the convergence speed of training deep neural networks . • We theoretically explain the root of BatchNorm ’ s effectiveness and the weakness from the perspective of lowering the gradient variance and gradient Lipschitz constant . • With the help of the new normalization approach , in extensive experiments the performance for Sigmoid and Tanh achieve massive improvement , and they are even competitive to ReLU . Our paper is organised as follows . In Section 2 , we will deduce the theoretical premises to accelerate stochastic optimization . Then , built on the premise , we will propose a new normalization approach in Section 3 . In Section 4 , we will analyze the exact root of BatchNorm ’ s effectiveness . In Section 5 , we will implement extensive experiments to demonstrate the superiorly of the new normalization approach . We will summarize our work and discuss further work in Section 6 . 2 THEORETICAL PREMISES . In this section , we will theoretically deduce that reduction of gradient Lipschitz constant and gradient variance plays a vitally important role in convergence when applying stochastic gradient descent ( SGD ) algorithms for machine learning tasks . In a machine learning task , given a set of samples { xi } ni=1 , the empirical risk as the average of all the samples loss is typically treated as the optimizing objective , i.e. , F ( w ) = F ( w ; { xi } ni=1 ) = 1 n n∑ i=1 fi ( w ; xi ) , ( 1 ) where fi ( w ; xi ) is the loss incurred by the parameter vector w with respect to the i-th sample or mini-batch . Noth that sometimes F ( w ; { xi } ni=1 ) and fi ( w ; xi ) may be abbreviated as F ( w ) and fi ( w ) . In large-scale machine learning tasks , such as deep learning , we usually choose SGD to minimize the empirical risk in Eq . ( 1 ) since its computational cost is cheap . Applying SGD , the parameter update at k-th iteration is wk+1 ← wk − αk+1∇fik ( wk ; xik ) ( 2 ) where at the k-th iteration xik is randomly chosen from the full sets { xi } ni=1 , and αk+1 is the learning rate . Convergence guarantee of most gradient-based algorithms including SGD is built upon the assumption of Lipschitz-continuous objective gradients , which is state as follow . The function F ( w ) is continuously differentiable and the gradient function of F ( w ) , namely ∇F ( w ) , is Lipschitz continuous with gradient Lipschitz constant L > 0 , i.e. , ‖∇F ( w1 ) −∇F ( w2 ) ‖ ≤ L‖w1 − w2‖ . ( 3 ) An important equality directly induced from Eq . ( 3 ) is F ( w2 ) ≤ F ( w1 ) + 〈∇F ( w1 ) , w2 − w1〉+ L 2 ‖w2 − w1‖2 . ( 4 ) The proof for it can be found in ( Nesterov , 2004 ; Bottou et al. , 2016 ) . Under the assumption in Eq . ( 4 ) , we provide convergence analysis in the following for the general objective function Eq . ( 1 ) when applying SGD , no matter whether it is convex or not . Theorem 1 Suppose the assumption in Eq . ( 4 ) is satisfied and the gradient of F ( wk ) in Eq . ( 1 ) is bounded ( i.e. , ∇F ( wk ) ≤M ≤ ∞ ) , SGD algorithm is applied to optimize F ( wk ) . ( 1 ) If the leaning rate satisfies ∑∞ k=1 αk =∞ and ∑∞ k=1 α 2 k < ∞ , and then∑K k=0 αk+1E [ ‖∇F ( wk ) ‖2 ] ∑K k=0 αk+1 ≤ E [ F ( w0 ) ] − Finf∑K k=0 αk+1 + L 2 ∑K k=0 α 2 k+1E [ ‖∇fik ( wk ) ‖2 ] ∑K k=0 αk+1 K→∞−−−−→ 0 , ( 5 ) where Finf = minw F ( w ) . Then , ‖∇F ( wK ) ‖ K→∞−−−−→ 0 in probability . ( 2 ) If the leaning rate satisfies 1αk ≥ L and ∑∞ k=1 αk =∞ , we have∑K k=0 αk+1E [ ‖∇F ( wk ) ‖2 ] ∑K k=0 αk+1 ≤ E [ F ( w0 ) ] − Finf 1 2 ∑K k=0 αk+1 + ∑K k=0 αk+1V [ ∇fik ( wk ) ] ∑K k=0 αk+1 , ( 6 ) where V [ ∇fik ( wk ) ] = E [ ‖∇fik ( wk ) − E [ ∇fik ( wk ) ] ‖2 ] that is the gradient variance of fik ( wk ) . When K →∞ , inf k ( E [ ‖∇F ( wk ) ‖2 ] ) ≤ sup k ( V [ ∇fik ( wk ) ] ) . ( 7 ) According to the first conclusion of Theorem 1 , when the learning rate satisfies ∑∞ k=1 αk =∞ and∑∞ k=1 α 2 k < ∞ , we know that SGD algorithm will guarantee the empirical risk function in Eq . ( 1 ) converges to a stationary point in probability . The conditions in the second conclusion of Theorem 1 is easier to be met . In this case we should make the the gradient variance V [ ∇fik ( wk ) ] as smallest as possible to ensure the empirical risk function in Eq . ( 1 ) can converge to a good point . Theorem 1 also implies that increasing learning rate will directly speed up convergence . Since larger αk will make the term E [ F ( w0 ) ] −Finf∑K k=0 αk+1 in Eq . ( 5 ) ( or E [ F ( w0 ) ] −Finf1 2 ∑K k=0 αk+1 in Eq . ( 6 ) ) faster to diminish . However , when the learning rate αk becomes large , the term L2 ∑K k=1 α 2 kE [ ‖∇fik ( wk ) ‖2 ] in Eq . ( 5 ) ( or ∑K k=0 αk+1V [ ∇fik ( wk ) ] in Eq . ( 6 ) ) may also become large and the condition 1 αk+1 ≥ L may not be met . Hence reducing the gradient variance V [ ∇fik ( wk ) ] 1and the value of Lipschitz constant L is the key to increase the learning rate and accelerate the convergence . It is worth noting that the conclusion that gradient variance reduction is beneficial to accelerate convergence has been demonstrated in ( Johnson & Zhang , 2013 ) when the objective function is convex . Therefore , if an normalization approach is effective to make Eq . ( 1 ) converge to a good point and accelerate the convergence process , it needs to reduce gradient variance V [ ∇fik ( wk ) ] and gradient Lipschitz constant L . 3 THE PROPOSED APPROACH . In this section , we will present a new normalization approach of which each step will be theoretically proved to be beneficial to reduce the gradient variance V [ ∇fik ( wk ) ] and the gradient Lipschitz constant L simultaneously , therefore the new approach will accelerate the convergence of training and improve the performance . 1It is known E [ ‖∇fik ( wk ) ‖ 2 ] = V [ ∇fik ( wk ) ] + ( E [ ‖∇fik ( wk ) ‖ ] ) 2 , decreasing V [ ∇fik ( wk ) ] will also lower E [ ‖∇fik ( wk ) ‖ 2 ] . When applying mini-batch SGD , at the k-th iteration the empirical risk loss with respect to the weight/bias of the l-th fully-connected layer of DNN can be expressed as 2 , F ( m ) ( w ( l ) k ) = 1 m m∑ i=1 fik ( w ( l ) k ; xik ) = 1 m m∑ i=1 g ( ( x ( l−1 ) ik ) t w ( l ) k + b ( l ) k ) = G ( m ) ( ( X ( l−1 ) ( m ) k ) t wlk +B ( l ) ( m ) k ) ( 8 ) where m is the size of mini-batch , and X ( l−1 ) ( m ) k = [ x ( l−1 ) 1k ; x ( l−1 ) 2k ; ... ; x ( l−1 ) mk ] and B ( l ) ( m ) k = [ b ( l ) 1k , b ( l ) 2k , ... , b ( l ) mk ] . x ( l−1 ) ik is randomly choose from the entire train set { x ( l−1 ) ik } n i=1 with equal probability , hence F ( w ( l ) k ) = 1 n ∑n i=1 fik ( w ( l ) k ; x ( l−1 ) ik ) is the expectation of F ( m ) ( wlk ) ( E [ F ( m ) ( wlk ) ] = F ( wlk ) ) , where n is the total number of train set . Note that Eq . ( 8 ) also implies that F ( wlk ) = G ( ( X l−1k ) twlk +B l k ) where X l−1k = [ x l−1 1k ; xl−12k ; ... ; x l−1 nk ] and Blk = [ b l 1k , bl2k , ... , b l nk ] . For clarity , henceforth we will omit the layer index . We construct a new normalization approach , referred to as Pre-Operation Normalization ( POPNorm ) , of which one main characteristic is putting the normalization technique before operation . For a layer fed with a d-dimensional input xik = ( x [ 1 ] ik , x [ 2 ] ik ... , x [ d ] ik ) , we ideally normalize each dimension with the global estimation of entire samples , i.e. , x̄ [ j ] ik = x [ j ] ik − η [ j ] · 1 n n∑ i=1 x [ j ] ik , ( Mean Removal ) ( 9 ) x̂ [ j ] ik = x̄ [ j ] ik ρ [ j ] · √ 1 n ∑n i=1 ( x̄ [ j ] ik ) 2 = x̄ [ j ] ik ρ [ j ] √ n ‖x [ j ] k ‖2 , ( Division by sclaed l2 Norm ) ( 10 ) where 0 ≤ η [ j ] ≤ 2 , ρ [ j ] √ 1 n ∑n i=1 ( x̄ [ j ] ik ) 2 ≥ 1 and x [ j ] k = ( x [ j ] 1k , x [ j ] 2k , ... , x [ j ] nk ) . The normalization approach consists of two parts – mean removal and division by l2 norm . The both steps are actually not interdependent and can individually take effects for decreasing gradient variance and gradient Lipschitz constant . We will theoretically demonstrate it in the following . Theorem 2 ( Justification of Mean Removal in Eq . ( 9 ) ) Eq . ( 9 ) is adopted to handle each input xik of the model in Eq . ( 8 ) . Suppose assumption in Eq . ( 3 ) is satisfied and the gradient and Hessian function of g ( z ) and G ( z ) in Eq . ( 8 ) is bounded , i.e. , ‖∇g ( z ) ‖2 ≤ M1 < ∞ and ‖∇2G ( z ) ‖2 ≤ M2 < ∞ , we have ( 1 ) the upper bound of the gradient variance V [ ∇fik ( wk ; x̄ik ) ] will be lower than that of V [ ∇fik ( wk ; xik ) ] . ( 2 ) the upper bound of the minimal gradient Lipschitz constant of F ( wk ; { x̄ik } ni=1 ) will be lower than that of F ( wk ; { xik } ni=1 ) ; Theorem 3 ( Justification of Division by l2 Norm in Eq . ( 10 ) ) Eq . ( 10 ) is adopted to handle each input x̂ik from Eq . ( 9 ) . Suppose assumption in Eq . ( 3 ) is satisfied and the gradient and Hessian function of g ( z ) and G ( z ) in Eq . ( 8 ) is bounded , i.e. , ‖∇g ( z ) ‖2 ≤ M1 < ∞ and ‖∇2G ( z ) ‖2 ≤ M2 < ∞ , we have ( 1 ) the upper bound of the gradient variance V [ ∇fik ( wk ; x̂ik ) ] will be lower than that of V [ ∇fik ( wk ; x̄ik ) ] ; ( 2 ) the upper bound of the minimal gradient Lipschitz constant of F ( wk ; { x̂ik } ni=1 ) will be lower than that of F ( wk ; { x̄ik } ni=1 ) ; ( 3 ) if F ( wk ) is locally convex and maxj ( ( ρ [ j ] ) 2 ) ( minj ( ρ [ j ] ) 2 ) ≤ maxj ( ∑n i=1 ( x̄ [ j ] ik ) 2 ) minj ( ∑n i=1 ( x̄ [ j ] ik ) 2 ) , the lower bound of conditional number of ∇2F ( wk ; { x̂ik } ni=1 ) will be lower than that of∇2F ( wk ; { x̄ik } ni=1 ) in the local space . 2 A convolution layer can be also reformulated to be a matrix form , hence the analysis for convolution layers is similar , but it is still somewhat complicated and we discuss it in the appendix . Due to the limited space , the proofs of Theorem 2 and Theorem 3 are provided in the appendix . The computation of 1n ∑n i=1 xik and 1 n ∑n i=1 x̄ 2 ik in Eq . ( 9 - 10 ) over the entire training set is impractical when adopting SGD . Therefore , just like the vanilla BatchNorm , each mini-batch produces 1 m ∑n i=1 xik and 1 m ∑n i=1 x̄ 2 ik as the estimates instead . This implementation will also make all the inputs fully participate in the gradient back-propagation . The vanilla BatchNorm introduces the scale and shift parameter to be learned to restore the presentation power of the network . To make POP-Norm more like the vanilla BatchNorm , we factorize η [ j ] and ρ [ j ] into a fixed parameter and a parameter to be learned . In practice POP-Norm is presented as follows : x̄ [ j ] ik = x [ j ] ik − 1 m m∑ i=1 x [ j ] ik + β [ j ] , ( Mean Removal ) ( 11 ) x̂ [ j ] ik = γ [ j ] κ · x̄ [ j ] ik√ 1 m ∑m i=1 ( x̄ [ j ] ik ) 2 , ( Division by scaled l2 Norm ) ( 12 ) where β [ j ] and γ [ j ] are the shift parameter and the scale parameter to be learned , and κ ( ≥ 1 ) is a hyperparameter . If |β [ j ] | < 1m ∑m i=1 x [ j ] ik and |γ [ j ] | < κ √ 1 m ∑m i=1 ( x̄ [ j ] ik ) 2 are satisfied , Theorem 2 and Theorem 3 will still hold . Remark 1 . Although the main steps of POP-Norm in Eq . ( 11 - 12 ) are similar to the vanilla BatchNorm , they are still slight different , as shown in Figure 1 . For BatchNorm , x̂ [ j ] ik = γ [ j ] ( x [ j ] ik − µ [ j ] + β [ j ] ) / √ 1 m ∑m i=1 ( x [ j ] ik − µ [ j ] ) 2 where µ = 1m∑mi=1 x [ j ] ik , while for POP-Norm , x̂ [ j ] ik = γ [ j ] ( x [ j ] ik − µ [ j ] + β [ j ] ) / √ 1 m ∑m i=1 ( x [ j ] ik − µ [ j ] + β [ j ] ) 2 . Besides , there are other two significant differences between the proposed POP-Norm and the vanilla BatchNorm – placing normalization before the main operation ( convolution and inner-production ) and adding a new scale parameter κ . Actually , we deliberately make POP-Norm more like the the vanilla BatchNorm . POP-Norm has other forms , for example , mean removal in Eq . ( 11 ) can be x̄ [ j ] ik = x [ j ] ik − 1+β [ j ] m ∑m i=1 x [ j ] ik , and the l2 norm to divide can be substituted by the other norm , such as l1 norm and l∞ nrom . Mean removal and division by l2 norm of POP-Norm are actually not interdependent and can individually take effects for decreasing gradient variance and gradient Lipschitz constant , and their order can be also exchanged . From the proof of Theorem 2 and Theorem 3 , we also know the division by l2 norm will lowering the upper bound of V [ ∇fik ( wk ) ] and L more substantially than mean removal . Additionally , the theoretical analysis and construction of POP-Norm are generalized , and we can also apply them to construct a new normalization approach to compete with instance normalization ( Ulyanov et al. , 2016 ) and group normalization ( Wu & He , 2018 ) . Remark 2 . Adding a new parameter κ and setting κ ≥ 1 is more beneficial to decrease the gradient variance and the gradient variance V [ ∇fik ( wk ) ] and the gradient Lipschitz constant L. When looking into the proof of Theorem 2 and Theorem 3 , we can conclude the upper bound of V [ ∇fik ( wk ) ] and L will be proportional to 1κ2 , which means the larger κ will be quickly smaller V [ ∇fik ( wk ) ] and L so that speeds up the convergence . However , κ can not be too large , otherwise the value after the operation will be much small , and then it may harm the representation capability of the layer , since the non-linear activation function will mainly work in the linear regime . Moreover , when κ is too large that will directly make gradient variance much small , and then it is difficult for train loss to escape the trap of early local minimums . In our practices , when κ is moderately large , it will substantially improve the performance , which is demonstrated in the experimental section . Not strictly speaking , adding κ is just make f ( w ) change to f ( wκ ) . It seems the optimization trajectory will be the same if we magnify the initial value ofw by an exact factor κ and scale the corresponding learning rate . However , we do not exactly change the initial value and the learning rate , and the optimum is not unique . On the contrary , the number of local minimums of a deep neural network is much large and they scatter throughout all the space . Therefore , adding κ is more likely to help w in the network quickly converge to a local minimum near the initial value according to Theorem 3 rather than along with a far trajectory converge to the original local minimum . Remark 3 . As proved in Theorem 3 , adopting the POP-Norm will be conducive to reduction of the condition number of∇2F ( wk ) if F ( wk ) is locally convex . Smaller condition number will make the training algorithm avoid the danger of running into a sudden change of the loss such as a flat region or a steep slope , which also enable us to adopt a larger learning rate . In other words , the smaller condition number will make the landscape or convergence trajectory of the loss is more smooth and substantially speed up the training process .

Summary: This paper proposes Pre-OPeration Normalization (POP-norm) as a theoretically grounded modification of batch normalization (Batch-norm). The authors provide an analysis of stochastic gradient descent (SGD) showing that the convergence can be accelerated by lowering the gradient Lipschitz constant and gradient variances. Based on this observation, the authors show that by applying POP-norm which is quite similar to the Batch-norm applied before the activation functions, the gradient Lipschitz constants and the gradient variances can be lowered.