Theoretical Analysis of Auto Rate-Tuning by Batch Normalization

1 INTRODUCTION . Batch Normalization ( abbreviated as BatchNorm or BN ) ( Ioffe & Szegedy , 2015 ) is one of the most important innovation in deep learning , widely used in modern neural network architectures such as ResNet ( He et al. , 2016 ) , Inception ( Szegedy et al. , 2017 ) , and DenseNet ( Huang et al. , 2017 ) . It also inspired a series of other normalization methods ( Ulyanov et al. , 2016 ; Ba et al. , 2016 ; Ioffe , 2017 ; Wu & He , 2018 ) . BatchNorm consists of standardizing the output of each layer to have zero mean and unit variance . For a single neuron , if x1 , . . . , xB is the original outputs in a mini-batch , then it adds a BatchNorm layer which modifies the outputs to BN ( xi ) = γ xi − µ σ + β , ( 1 ) where µ = 1B ∑B i=1 xi and σ 2 = 1B ∑B i=1 ( xi − µ ) 2 are the mean and variance within the minibatch , and γ , β are two learnable parameters . BN appears to stabilize and speed up training , and improve generalization . The inventors suggested ( Ioffe & Szegedy , 2015 ) that these benefits derive from the following : 1 . By stabilizing layer outputs it reduces a phenomenon called Internal Covariate Shift , whereby the training of a higher layer is continuously undermined or undone by changes in the distribution of its inputs due to parameter changes in previous layers. , 2 . Making the weights invariant to scaling , appears to reduce the dependence of training on the scale of parameters and enables us to use a higher learning rate ; 3 . By implictly regularizing the model it improves generalization . ∗Work done while visiting Princeton University . But these three benefits are not fully understood in theory . Understanding generalization for deep models remains an open problem ( with or without BN ) . Furthermore , in demonstration that intuition can sometimes mislead , recent experimental results suggest that BN does not reduce internal covariate shift either ( Santurkar et al. , 2018 ) , and the authors of that study suggest that the true explanation for BN ’ s effectiveness may lie in a smoothening effect ( i.e. , lowering of the Hessian norm ) on the objective . Another recent paper ( Kohler et al. , 2018 ) tries to quantify the benefits of BN for simple machine learning problems such as regression but does not analyze deep models . Provable quantification of Effect 2 ( learning rates ) . Our study consists of quantifying the effect of BN on learning rates . Ioffe & Szegedy ( 2015 ) observed that without BatchNorm , a large learning rate leads to a rapid growth of the parameter scale . Introducing BatchNorm usually stabilizes the growth of weights and appears to implicitly tune the learning rate so that the effective learning rate adapts during the course of the algorithm . They explained this intuitively as follows . After BN the output of a neuron z = BN ( w > x ) is unaffected when the weight w is scaled , i.e. , for any scalar c > 0 , BN ( w > x ) = BN ( ( cw ) > x ) . Taking derivatives one finds that the gradient at cw equals to the gradient atw multiplied by a factor 1/c . Thus , even though the scale of weight parameters of a linear layer proceeding a BatchNorm no longer means anything to the function represented by the neural network , their growth has an effect of reducing the learning rate . Our paper considers the following question : Can we rigorously capture the above intuitive behavior ? Theoretical analyses of speed of gradient descent algorithms in nonconvex settings study the number of iterations required for convergence to a stationary point ( i.e. , where gradient vanishes ) . But they need to assume that the learning rate has been set ( magically ) to a small enough number determined by the smoothness constant of the loss function — which in practice are of course unknown . With this tuned learning rate , the norm of the gradient reduces asymptotically as T−1/2 in T iterations . In case of stochastic gradient descent , the reduction is like T−1/4 . Thus a potential way to quantify the rate-tuning behavior of BN would be to show that even when the learning rate is fixed to a suitable constant , say 0.1 , from the start , after introducing BN the convergence to stationary point is asymptotically just as fast ( essentially ) as it would be with a hand-tuned learning rate required by earlier analyses . The current paper rigorously establishes such auto-tuning behavior of BN ( See below for an important clarification about scale-invariance ) . We note that a recent paper ( Wu et al. , 2018 ) introduced a new algorithm WNgrad that is motivated by BN and provably has the above auto-tuning behavior as well . That paper did not establish such behavior for BN itself , but it was a clear inspiration for our analysis of BN . Scale-invariant and scale-variant parameters . The intuition of Ioffe & Szegedy ( 2015 ) applies for all scale-invariant parameters , but the actual algorithm also involves other parameters such as γ and β whose scale does matter . Our analysis partitions the parameters in the neural networks into two groups W ( scale-invariant ) and g ( scale-variant ) . The first group , W = { w ( 1 ) , . . . , w ( m ) } , consists of all the parameters whose scales does not affect the loss , i.e. , scaling w ( i ) to cw ( i ) for any c > 0 does not change the loss ( see Definition 2.1 for a formal definition ) ; the second group , g , consists of all other parameters that are not scale-invariant . In a feedforward neural network with BN added at each layer , the layer weights are all scale-invariant . This is also true for BN with ` p normalization strategies ( Santurkar et al. , 2018 ; Hoffer et al. , 2018 ) and other normalization layers , such as Weight Normalization ( Salimans & Kingma , 2016 ) , Layer Normalization ( Ba et al. , 2016 ) , Group Normalization ( Wu & He , 2018 ) ( see Table 1 in Ba et al . ( 2016 ) for a summary ) . 1.1 OUR CONTRIBUTIONS . In this paper , we show that the scale-invariant parameters do not require rate tuning for lowering the training loss . To illustrate this , we consider the case in which we set learning rates separately for scale-invariant parameters W and scale-variant parameters g. Under some assumptions on the smoothness of the loss and the boundedness of the noise , we show that 1 . In full-batch gradient descent , if the learning rate for g is set optimally , then no matter how the learning rates for W is set , ( W ; g ) converges to a first-order stationary point in the rate O ( T−1/2 ) , which asymptotically matches with the convergence rate of gradient descent with optimal choice of learning rates for all parameters ( Theorem 3.1 ) ; 2 . In stochastic gradient descent , if the learning rate for g is set optimally , then no matter how the learning rate for W is set , ( W ; g ) converges to a first-order stationary point in the rate O ( T−1/4 polylog ( T ) ) , which asymptotically matches with the convergence rate of gradient descent with optimal choice of learning rates for all parameters ( up to a polylog ( T ) factor ) ( Theorem 4.2 ) . In the usual case where we set a unified learning rate for all parameters , our results imply that we only need to set a learning rate that is suitable for g. This means introducing scale-invariance into neural networks potentially reduces the efforts to tune learning rates , since there are less number of parameters we need to concern in order to guarantee an asymptotically fastest convergence . In our study , the loss function is assumed to be smooth . However , BN introduces non-smoothness in extreme cases due to division by zero when the input variance is zero ( see equation 1 ) . Note that the suggested implementation of BN by Ioffe & Szegedy ( 2015 ) uses a smoothening constant in the whitening step , but it does not preserve scale-invariance . In order to avoid this issue , we describe a simple modification of the smoothening that maintains scale-invariance . Also , our result can not be applied to neural networks with ReLU , but it is applicable for its smooth approximation softplus ( Dugas et al. , 2001 ) . We include some experiments in Appendix D , showing that it is indeed the auto-tuning behavior we analysed in this paper empowers BN to have such convergence with arbitrary learning rate for scale-invariant parameters . In the generalization aspect , a tuned learning rate is still needed for the best test accuracy , and we showed in the experiments that the auto-tuning behavior of BN also leads to a wider range of suitable learning rate for good generalization . 1.2 RELATED WORKS . Previous work for understanding Batch Normalization . Only a few recent works tried to theoretically understand BatchNorm . Santurkar et al . ( 2018 ) was described earlier . Kohler et al . ( 2018 ) aims to find theoretical setting such that training neural networks with BatchNorm is faster than without BatchNorm . In particular , the authors analyzed three types of shallow neural networks , but rather than consider gradient descent , the authors designed task-specific training methods when discussing neural networks with BatchNorm . Bjorck et al . ( 2018 ) observes that the higher learning rates enabled by BatchNorm improves generalization . Convergence of adaptive algorithms . Our analysis is inspired by the proof for WNGrad ( Wu et al. , 2018 ) , where the author analyzed an adaptive algorithm , WNGrad , motivated by Weight Normalization ( Salimans & Kingma , 2016 ) . Other works analyzing the convergence of adaptive methods are ( Ward et al. , 2018 ; Li & Orabona , 2018 ; Zou & Shen , 2018 ; Zhou et al. , 2018 ) . Invariance by Batch Normalization . Cho & Lee ( 2017 ) proposed to run riemmanian gradient descent on Grassmann manifold G ( 1 , n ) since the weight matrix is scaling invariant to the loss function . Hoffer et al . ( 2018 ) observed that the effective stepsize is proportional to ηw‖wt‖2 . 2 GENERAL FRAMEWORK . In this section , we introduce our general framework in order to study the benefits of scale-invariance . 2.1 MOTIVATING EXAMPLES OF NEURAL NETWORKS . Scale-invariance is common in neural networks with BatchNorm . We formally state the definition of scale-invariance below : Definition 2.1 . ( Scale-invariance ) LetF ( w , θ′ ) be a loss function . We say thatw is a scale-invariant parameter of F if for all c > 0 , F ( w , θ′ ) = F ( cw , θ′ ) ; if w is not scale-invariant , then we say w is a scale-variant parameter of F . We consider the following L-layer “ fully-batch-normalized ” feedforward network Φ for illustration : L ( θ ) = EZ∼DB [ 1 B B∑ b=1 fyb ( BN ( W ( L ) σ ( BN ( W ( L−1 ) · · ·σ ( BN ( W ( 1 ) xb ) ) ) ) ) ) ] . ( 2 ) Z = { ( x1 , y1 ) , . . . , ( xB , yB ) } is a mini-batch ofB pairs of input data and ground-truth label from a data set D. fy is an objective function depending on the label , e.g. , fy could be a cross-entropy loss in classification tasks . W ( 1 ) , . . . , W ( L ) are weight matrices of each layer . σ : R→ R is a nonlinear activation function which processes its input elementwise ( such as ReLU , sigmoid ) . Given a batch of inputs z1 , . . . , zB ∈ Rm , BN ( zb ) outputs a vector z̃b defined as z̃b , k : = γk zb , k − µk σk + βk , ( 3 ) where µk = Eb∈ [ B ] [ zb , k ] and σ2k = Eb∈ [ B ] [ ( zb , k − µk ) 2 ] are the mean and variance of zb , γk and βk are two learnable parameters which rescale and offset the normalized outputs to retain the representation power . The neural network Φ is thus parameterized by weight matrices W ( i ) in each layer and learnable parameters γk , βk in each BN . BN has the property that the output is unchanged when the batch inputs z1 , k , . . . , zB , k are scaled or shifted simultaneously . For zb , k = w > k x̂b being the output of a linear layer , it is easy to see that wk is scale-invariant , and thus each row vector of weight matrices W ( 1 ) , . . . , W ( L ) in Φ are scaleinvariant parameters of L ( θ ) . In convolutional neural networks with BatchNorm , a similar argument can be done . In particular , each filter of convolutional layer normalized by BN is scale-invariant . With a general nonlinear activation , other parameters in Φ , the scale and shift parameters γk and βk in each BN , are scale-variant . When ReLU or Leaky ReLU ( Maas et al. , 2013 ) are used as the activation σ , the vector ( γ1 , . . . , γm , β1 , . . . , βm ) of each BN at layer 1 ≤ i < L ( except the last one ) is indeed scale-invariant . This can be deduced by using the the ( positive ) homogeneity of these two types of activations and noticing that the output of internal activations is processed by a BN in the next layer . Nevertheless , we are not able to analyse either ReLU or Leaky ReLU activations because we need the loss to be smooth in our analysis . We can instead analyse smooth activations , such as sigmoid , tanh , softplus ( Dugas et al. , 2001 ) , etc .

Summary: The work is motivated by the empirical performance of Batch Normalization and in particular the observed better robustness of the choice of the learning rate.  Authors analyze theoretically the asymptotic convergence rate for objectives involving normalization, not necessarily BN, and show that for scale-invariant groups of parameters (appearing as a result of normalization) the initial learning rate may be set arbitrary while still asymptotic convergence is guaranteed with the same rate as the best known in the general case. Offline gradient descent and stochastic gradient descent cases are considered.