MEAN-FIELD ANALYSIS OF BATCH NORMALIZATION

1 INTRODUCTION . Deep neural networks have achieved remarkable success in the past decade on tasks that were out of reach prior to the era of deep learning ( Krizhevsky et al. , 2012 ; He et al. , 2016b ) . Amongst the myriad reasons for these successes are powerful computational resources , large datasets , new optimization algorithms , and modern architecture designs ( Russakovsky et al. , 2015 ; Kingma & Ba , 2015 ) . In many modern deep learning architectures , one key component is batch normalization ( BatchNorm ) . BatchNorm is a module that can be introduced in layers of deep neural networks that normalizes hidden layer outputs to have a common first and second moment . Empirically , BatchNorm enables optimization using much larger learning rates , and achieves better convergence ( Ioffe & Szegedy , 2015 ) . Despite significant practical success , a theoretical understanding of BatchNorm is still lacking . A widely held view is that BatchNorm improves training by “ reducing of internal covariate shift ” ( ICF ) ( Ioffe & Szegedy , 2015 ) . Internal covariate shift refers to the change in the input distribution of internal layers of the deep network due to changes of the weights . Recent results ( Santurkar et al. , 2018 ) , however , cast doubt on the ICF expalanation , by demonstrating that noisy BatchNorm increases ICF yet still improves training as in regular BatchNorm . This raises the question of whether the utility of BatchNorm is indeed related to the reduction of ICF . Instead , it is argued by Santurkar et al . ( 2018 ) that BatchNorm actually improves the Lipschitzness of the loss and gradient . Meanwhile , dynamical mean-field theory ( Sompolinsky & Zippelius , 1982 ) , a powerful theoretical technique , has recently been applied by Poole et al . ( 2016 ) to ensembles of multi-layer random neural networks . This theory studies networks with an i.i.d . Gaussian distribution of weights and biases . Most recent work focuses on the analysis of order parameter flows and their fixed points ( Schoenholz et al. , 2017 ; Xiao et al. , 2018 ; Yang & Schoenholz , 2017 ) , including their stability and decay rates . Importantly , Karakida et al . ( 2018 ) also successfully used mean-field analysis to estimate the spectral properties of the Fisher Information Matrix . In this paper , we analytically quantify the impact of BatchNorm on the landscape of the loss function , by using mean-field theory to estimate the spectral properties of the Fisher Information Matrix ( FIM ) for typical batch-normalized neural networks . In particular , it is shown that BatchNorm reduces the maximal eigenvalue of the FIM provided that the normalization coefficient γ is not too large . By drawing on results linking Fisher Information to the geometry of the loss function , we explain how BatchNorm neural networks can be trained with a larger learning rate without leading to parameter explosion , and provide upper bounds on the learning rate in terms of the BatchNorm parameters . As an additional contribution motivated by our theoretical findings , we demonstrate an empirical correlation between the BatchNorm parameter γ and test loss . In particular , networks with smaller γ achieve lower loss after a fixed number of training epochs . 2 PRELIMINARIES . In our theoretical analysis , we employ the recent application of mean-field theory to neural networks which studies an ensemble of random neural networks with pre-defined i.i.d . Gaussian weights and biases . In this section , we provide background information and briefly recall the formalism of Karakida et al . ( 2018 ) which first computes spectral properties of the Fisher Information of a neural network and then relates it to the maximal stable learning rate . 2.1 FISHER INFORMATION MATRIX AND LEARNING DYNAMICS . Given a data distribution D over the set of instance-label pairs X × Y , a family of parametrized functions fθ : X → Y and a loss function l ( f , y ) , our focus will be to ensure convergence of the following gradient descent with momentum update rule : θt+1 = θt − η∇L ( θt ) + µ ( θt − θt−1 ) , ( 1 ) where L ( θ ) is the unobserved population loss , L ( θ ) : = E ( x , y ) ∼D [ l ( fθ ( x ) , y ) ] . ( 2 ) In practice , the parameters are determined by minimizing an empirical estimate of equation 2 using a stochastic generalization ( SGD ) of the update rule equation 1 . We neglect this difference by always working in the asymptotic limit of large sample size and moreover assuming full-batch gradient updates . Suppose the loss function can be expressed in terms of a parametric family of positive densities as l ( fθ ( x ) , y ) = : − log pθ ( x , y ) . This assumption holds true for a large class of losses including squared loss and cross-entropy loss . Let Iθ denote the Fisher Information Matrix ( FIM ) associated with the parametric family induced by the loss , Iθ : = E ( x , y ) ∼Pθ [ ∇θ log pθ ( x , y ) ⊗∇θ log pθ ( x , y ) ] , ( 3 ) where ⊗ denotes Kronecker product and Pθ denotes the probability distribution over X × Y with density pθ ( x , y ) . Recall that under suitable regularity conditions the following identity holds : Iθ = −E ( x , y ) ∼Pθ [ Hess θ log pθ ( x , y ) ] , ( 4 ) where Hessθ denotes the Hessian with respect to θ . The above right-hand side is closely related to the Hessian of the population loss , Hess ( L ( θ ) ) = −E ( x , y ) ∼D [ Hess θ log pθ ( x , y ) ] , ( 5 ) where we interchanged the Hessian with the expectation value . In fact , if we assume that the estimation problem is well-specified so that there exist parameters θ∗ such that the data distribution is generated by Pθ∗ = D , then we obtain the following equality between the Hessian of the population loss and the FIM evaluated at the optimal parameters , Hess ( L ( θ∗ ) ) = Iθ∗ . ( 6 ) If θ is initialized in a sufficiently small neighborhood of θ∗ , then by expanding the population loss L ( θ ) to quadratic order about θ∗ one can show that a necessary condition for convergence is that the step size is bounded from above by ( LeCun et al. , 2012 ; Karakida et al. , 2018 ) 1 , η < η∗ : = 2 ( 1 + µ ) λmax ( Hess ( L ( θ∗ ) ) ) = 2 ( 1 + µ ) λmax ( Iθ∗ ) , ( 7 ) 1In the quadratic approximation to the loss , the optimal learning rate is in fact η∗/2 . where λmax ( M ) denotes the largest eigenvalue of the matrixM . Rather than computing the optimal parameters θ∗ directly , we follow the strategy of Karakida et al . ( 2018 ) by estimating the following quantity and arguing that the distribution of the weights and biases is not significantly impacted by the training dynamics , λ̄max : = E θ [ λmax ( Iθ ) ] , ( 8 ) where Eθ denotes the expectation value with respect to the weights and biases . This heuristic was shown to yield a remarkably accurate prediction of the maximal learning rate in ( Karakida et al. , 2018 ) . In this paper we adopt the data modeling assumption that the joint density factors as pθ ( x , y ) = p ( x ) pθ ( y |x ) where p ( · ) denotes the probability density of the marginal distribution of the covariates , which is independent of θ . Under this factorization assumption , the FIM simplifies to Iθ = E ( x , y ) ∼Pθ [ ∇ log pθ ( y |x ) ⊗∇θ log pθ ( y |x ) ] . ( 9 ) Focusing on the Gaussian conditional model pθ ( y |x ) ∝ exp ( 12‖fθ ( x ) − y‖ 2 2 ) , the FIM further simplifies to Iθ = Ex∼D [ ∇θfθ ( x ) ⊗∇θfθ ( x ) ] . ( 10 ) The family of parametrized functions fθ : RN0 → RNL is chosen to be the family of functions computed by a multi-layer neural network architecture with N0 input nodes , NL output nodes and L ≥ 1 layers . In this paper , we consider neural networks consisting of fully-connected ( FC ) and convolutional ( Conv ) layers , with and without batch normalization . The pointwise activation is denoted by σ , which is taken to be the rectified linear unit ( ReLU ) in this paper . Our analysis can be straightforwardly extended to other architectures and non-linearities . We use hlθ ( x ) to denote the output of layer l and the input to layer l + 1 . Clearly we have h0θ ( x ) = x and h L θ ( x ) = fθ ( x ) . 3 THEORY . In this section we focus on applying dynamical mean-field theory to study the effect of introducing batch normalization modules into a deep neural network by estimating the largest eigenvalue of the FIM . This estimate , in turn , provides an upper bound on the largest learning rate for which the learning dynamics is stable . This section is structured as follows : We first define various thermodynamic quantities ( order parameters , 6 for fully-connected layers and 9 for convolutional layers ) that satisfy recursion relations in the mean-field approximation . Then we present an estimate of λ̄max in terms of these order parameters , generalizing a result of Karakida et al . ( 2018 ) . Using this estimate , we study how λ̄max and η∗ are affected by BatchNorm and calculate their dependence on the BatchNorm coefficient γ . Detailed derivations of the order parameters , their recursions , and the associated eigenvalue bound are deferred to the Supplementary Material . 3.1 FULLY CONNECTED LAYERS . A general fully connected layer with input activation hl ( x ) and output pre-activation zl+1 ( x ) is described by the affine transformation , zl+1 ( x ) : = W l+1hl ( x ) + bl+1 , ( 11 ) where W l+1 ∈ RNl+1×Nl , bl+1 ∈ RNl+1 and Nl denotes the number of units in layer l. In the framework of mean-field theory , we will consider an ensemble of neural networks with Gaussian random weights and biases distributed as follows , [ W l+1 ] ij ∼ N ( 0 , σ2w/Nl ) , bl+1i ∼ N ( 0 , σ 2 b ) . ( 12 ) In the case of a standard fully connected layer , the input activation satisfies the recursions hl ( x ) = σ ( zl ( x ) ) , where σ denotes the pointwise activation . A batch-normalized fully connected layer , in contrast , satisfies the following recursion , hl ( x ) : = σ ( zl ( x ) − µl sl γl + βl ) , ( 13 ) where denotes the elementwise ( Hadamard ) product , µl ∈ RNl and ( sl ) 2 : = sl sl ∈ RNl denote the mean and variance of the pre-activation layers with respect to the data distribution , µl : = E x [ zl ( x ) ] , ( 14 ) ( sl ) 2 : = E x [ ( zl ( x ) − µl ) 2 ] . ( 15 ) The weights and biases are drawn from the same distributions as in the standard , no BatchNorm , case . In addition , we now have the BatchNorm parameters γl+1 , βl+1 ∈ RNl+1 which are assumed to be non-random for simplicity , γl [ i ] = γl , βl [ i ] = 0 . ( 16 )

Summary: In this paper, the effect of batch normalization to the maximum eigenvalue of the Fisher information is analyzed. The techinique is mostly developed by Karakida et al. (2018). The main result is an informal bound of the maximum eigenvalue, which is given without proof. Though, the numerical result corresponds to the derived bound.