Title: MEAN-FIELD ANALYSIS OF BATCH NORMALIZATION

Abstract: Batch Normalization (BatchNorm) is an extremely useful component of modern neural network architectures, enabling optimization using higher learning rates and achieving faster convergence. In this paper, we use mean-field theory to analytically quantify the impact of BatchNorm on the geometry of the loss landscape for multi-layer networks consisting of fully-connected and convolutional layers. We show that it has a flattening effect on the loss landscape, as quantified by the maximum eigenvalue of the Fisher Information Matrix. These findings are then used to justify the use of larger learning rates for networks that use BatchNorm, and we provide quantitative characterization of the maximal allowable learning rate to ensure convergence. Experiments support our theoretically predicted maximum learning rate, and furthermore suggest that networks with smaller values of the BatchNorm parameter γ achieve lower loss after the same number of epochs of training.

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:

EQUATION

where L(θ) is the unobserved population loss,

EQUATION

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,

EQUATION

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:

EQUATION

where Hess θ denotes the Hessian with respect to θ. The above right-hand side is closely related to the Hessian of the population loss,

EQUATION

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,

EQUATION

)

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) 

EQUATION

where λ max (M ) denotes the largest eigenvalue of the matrix M . 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,

EQUATION

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

EQUATION

Focusing on the Gaussian conditional model p θ (y |x) ∝ exp( 1 2 f θ (x) -y 2 2 ), the FIM further simplifies to

EQUATION

)

The family of parametrized functions f θ : R N0 → R N L is chosen to be the family of functions computed by a multi-layer neural network architecture with N 0 input nodes, N L 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 h l θ (x) to denote the output of layer l and the input to layer l + 1. Clearly we have h 0 θ (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 Batch-Norm 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 h l (x) and output pre-activation z l+1 (x) is described by the affine transformation,

EQUATION

where

W l+1 ∈ R N l+1 ×N l , b l+1 ∈ R N l+1

and N l 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,

EQUATION

In the case of a standard fully connected layer, the input activation satisfies the recursions h l (x) = σ(z l (x)), where σ denotes the pointwise activation.

A batch-normalized fully connected layer, in contrast, satisfies the following recursion,

EQUATION

where denotes the elementwise (Hadamard) product, µ l ∈ R N l and (s l ) 2 := s l s l ∈ R N l denote the mean and variance of the pre-activation layers with respect to the data distribution,

EQUATION

EQUATION

)

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 ∈ R N l+1 which are assumed to be non-random for simplicity,

EQUATION

3.1.1. ORDER PARAMETERS AND THEIR RECURSIONS:

To investigate the spectral properties of the FIM, we define the following order parameters,

EQUATION

)

EQUATION

EQUATION

where • denotes the Euclidean norm and δ l (x) := ∂f θ ∂z l (x). Here we assume that the data x are drawn i.i.d. from a distribution with mean 0 and variance 1, and also that the last layer is linear for classification. We then have the base cases: q0 = 0, q0 xy = 1, qL = qL xy = 1. The order parameters in the absence of BatchNorm satisfy the following recursions derived in Karakida et al. (2018) ,

EQUATION

EQUATION

EQUATION

where c xy := q l xy /q l . In the case of batch normalization we find the following recursions, which are derived in the Supplementary Material,

EQUATION

EQUATION

EQUATION

)

3.2. CONVOLUTIONAL LAYER:

The mean field theory of convolutional layer was first studied by Xiao et al. (2018) . In this paper, the results of the preceding section also apply to structured affine transformations including convolutional layers. Let K l denote the set of allowable spatial locations of the the lth layer feature map and let F l+1 index the sites of the convolutional kernel applied to that layer. Let C l denote the number of input channels. The output of a general convolutional layer is of the form,

EQUATION

where

α ∈ K l+1 , W l+1 β ∈ R C l+1 ×C l and b l+1 ∈ R C l+1 .

The weights and biases are now distributed as

EQUATION

)

where now N l := C l |F l+1 |. As in the fully connected case, we consider convolutional layers with both vanilla activation functions of the form h l α (x) := σ(z l α (x)) as well as batch normalized convolutional layers, for which the input activations satisfy the recursive identity,

EQUATION

3.2.1 ORDER PARAMETERS AND THEIR RECURSIONS Similar to the definitions for fully connected layer, we define the following set of order parameters:

EQUATION

)

EQUATION

EQUATION

EQUATION

)

EQUATION

where now δ l α := ∂f θ /∂z l α in analogy with the fully connected layer. The expectations over α and β are with respect to the uniform measure over the set of allowed indices. For a standard convolutional layer without BatchNorm, the order parameters can be shown to satisfy the following recursion relations:

EQUATION

)

EQUATION

EQUATION

EQUATION

EQUATION

In the case of convolutional layers with BatchNorm, the following recursions hold:

EQUATION

)

EQUATION

EQUATION

EQUATION

)

EQUATION

)

where c xy := q l xy /q l and c αβ := q l αβ,xy /q l . The derivations of the recursion relations for both vanilla and batch-normalized convolutional layers are deferred to the Supplementary Material.

3.3. EIGENVALUE BOUND AND THERMODYNAMIC VARIABLES:

The order parameters derived in the previous section are useful because they allow us to gain information about the maximal eigenvalue λmax of the FIM. We derived a generalization of (Karakida et al., 2018, Theorem 6) to allow for the inclusion of batch normalization and convolutional layers. In particular, we obtain a lower bound on the maximal eigenvalue λmax in terms of the previously introduced order parameters which satisfy the stated recursion relations in the mean-field approximation.

Claim 3.1. If the layer dimension N l of the fully connected layers and the number of channels C l of the convolutional layers satisfy N l 1 and C l 1 for 0 < l < L, we have,

EQUATION

where

f l = N l-1 ql-1 xy ql

xy for fully connected layers and

EQUATION

for convolutional layers, where recall

N l-1 = C l-1 |F l |.

The index sets F l and K l are defined in section 3.2. The order parameters are defined in the previous subsection.

Now we are ready to calculate the lower bound on λmax for a given model architecture by calculating the order parameters using their recursions. In the next section, we will focus on the numerical analysis of these recursion relations as well as present experiments that support our calculation.

4. NUMERICAL ANALYSIS AND EXPERIMENTS:

In order to understand the effect of BatchNorm on the loss landscape, we theoretically compute λmax as a function of the BatchNorm parameter γ, for both fully connected and convolutional architectures (Fig. 1 ) with and without BatchNorm. For γ 3 (typical for deep network initialization (Ioffe & Szegedy, 2015) ) BatchNorm significantly reduces λmax compared to the vanilla networks. As a direct consequence of this, the theory predicts that batch normalized networks can be trained using significantly higher learning rates than their vanilla counterparts.

We tested the above theoretical prediction by training the same architectures on MNIST and CIFAR-10 datasets, for different values of η and γ, starting from randomly initialized networks with same variances employed in the mean-field theory calculations. As shown in Fig. 2 , the (γ, log 10 η)-plane clearly partitions into distinct phases characterized by convergent and non-convergent optimization dynamics, and our theoretically predicted upper bound η * closely agrees with the experimentally determined phase boundary. The experiment of vanilla network is shown in 6.4 as a baseline.

In addition to the striking match between our theoretical prediction and the experimentally determined phase boundaries, the experimental results also suggest a tendency for smaller γ-initiations to produce lower values of test loss after a fixed number of epochs, i.e. faster convergence. We leave detailed investigation of this initialization scheme to future work. Also, dark strips can be observed in the heatmaps indicate the optimal learning rates for optimization, which is around η * /2 and consistent with LeCun et al. (2012) in the quadratic approximation to the loss. Our analysis also suggests that small γ initialization benefits the convergence of training. Additional experiments supporting this intuition can be found in Section 6.5 of Supplementary Material.

The architectural design for our experiments is as follows. In the fully connected architecture, we choose L = 4 layers with N l = 1000 hidden units per layer except the final (linear) layer which has N L = 10 outputs. Batch normalization is applied after each linear operation except for the final linear output layer. The convolutional network has a similar structure with L = 4 layers. The first three are convolutional layers with filter size 3, stride 2, and number of channels C 1 = 30, C 2 = 60, C 3 = 90. The final layer is a fully connected output layer to perform classification. The other architectural/optimization hyperparameters were chosen to be σ 2 w = 2, σ 2 b = 0.5, β = 0 and µ = 0.9. Momentum µ here was set to be 0.9 to match the value frequently used in practice, which only affects the dependency of η * on FIM. The white region indicates parameter explosion for at least one of the runs. The red line shows the theoretical prediction for the maximal learning rate η * . The dark band on the heatmaps for CIFAR-10 approximately tracks the optimal learning rate η * /2 in the quadratic approximation to the loss. Note the log scale for the learning rate, so the theory matches the experiments over three orders of magnitude for η.

5. CONCLUSION AND FUTURE WORK:

In this paper, we studied the impact of BatchNorm on the loss surface of multi-layer neural networks and its implication for training dynamics. By developing recursion relations for the relevant order parameters, the maximum eigenvalue of the Fisher Information matrix λmax can be estimated and related to the maximal learning rate. The theory correctly predicts that adding BatchNorm with small γ allows the training algorithm to exploit much larger learning rates, which speeds up convergence.

The experiments also suggest that using a smaller γ results in a lower test loss for a fixed number of training epochs. This suggests that initialization with smaller γ may help the optimization process in deep learning models, which will be interesting for future study.

The close agreement between theoretical predictions and the experimentally determined phase boundaries strongly supports the validity of our analysis, despite the non-rigorous nature of the derivations. Although similar approaches have been used in other work (Poole et al., 2016; Schoenholz et al., 2017; Yang & Schoenholz, 2017; Xiao et al., 2018; Karakida et al., 2018) , we hope that future work will place these results on a firmer mathematical footing. Furthermore, our BatchNorm analysis is not limited to the convolutional and fully-connected architectures we considered in this paper and can be extended to arbitrary feedforward architectures such as ResNets.

6. SUPPLEMENTARY MATERIAL:

This section provides non-rigorous derivations of the order parameters, their recursions, and the associated eigenvalue bound. Despite the non-rigorous nature of these calculations, we remark that similar reasoning has been successfully used in a number of related works on mean-field theory, demonstrating impressive agreement with experiments (Poole et al., 2016; Schoenholz et al., 2017; Yang & Schoenholz, 2017; Xiao et al., 2018; Karakida et al., 2018) . Derivation. In general, we have

EQUATION

)

Thus, setting l = 0 we obtain q 1 = σ 2 w + σ 2 b if we assume E

x

x 2 = N 0 , and

EQUATION

Therefore, setting x = y and taking expectation values over x gives,

EQUATION

EQUATION

EQUATION

EQUATION

EQUATION

EQUATION

EQUATION

)

where we have approximated each component of the random vector u l (x) as a standard Gaussian.

Similarly, taking expectations over x, y gives

EQUATION

)

Consider the approximation in which the random pair (u l (x)[1], u l (y) [1] ) is Gaussian distributed with zero mean and covariance,

EQUATION

)

Then the recursion becomes

EQUATION

)

where c l xy := Σ l xy / Σ l xx Σ l yy . Observe that by independence of x and y,

EQUATION

EQUATION

)

Thus Σ l xy = 0 and consequently,

EQUATION

)

EQUATION

)

The derivation here assumes an infinitely large dataset. For limited dataset size m, an error is introduced from the non-zero ratio of m/m x =y where m x =y denotes the total number of sample pairs (x, y) where x = y. We have m x =y = m 2 -m and therefore the error is O(1/m) which is negligible for most of the frequently-used datasets.

Claim 6.2. The backward recursions for

l ∈ [1, L -1] are ql = γ 2 l σ 2 w 2 ql+1 q l and ql xy = γ 2 l σ 2 w 4 ql+1 xy q l

with base cases qL = qL xy = 1.

Derivation. For each layer l ∈ [L] and for each unit i ∈

EQUATION

In particular, since we assume linear output δ L [i] = e i and thus

EQUATION

where 1 ∈ R N L is the vector of ones. For ease of presentation and without loss of generality we restrict to the first component qL [1] = 1 and abuse notation by writing qL = 1. For l < L we have by the chain rule,

EQUATION

EQUATION

Hence,

δ l (x)[i]δ l (y) = γ l [i] 2 s l [i] 2 σ l u l (x)[i] γ l [i] σ l u l (y)[i] γ l [i] k,k ∈[N l+1 ] [W l+1 ] ki [W l+1 ] k i δ l+1 [k]δ l+1 [k ] .

(67) Following the usual assumption that the back-propagated gradient is independent of the forward signal we obtain,

EQUATION

Setting x = y and taking expectation values over x we obtain,

EQUATION

EQUATION

Similarly, taking expectation values over x, y we obtain,

EQUATION

EQUATION

6.2. RECURSIONS FOR CONVOLUTIONAL LAYERS:

Recall that the output of a general convolutional layer is of the form,

EQUATION

where α ∈ K l+1 . As a concrete example, consider CIFAR-10 input of dimension 32 × 32 × 3 which is mapped by a convolutional layer with 3 × 3 kernels, stride 2 and 20 output channels and no padding. Then We begin by deriving some useful identities for convolutional layers, before specializing to the batch-normalized and vanilla networks. For each channel i ∈ [C l ], we have,

EQUATION

Hence,

EQUATION

where

N l := C l |F l+1 |.

Moreover, let us introduce the following shorthand,

EQUATION

EQUATION

EQUATION

Then we can write the recursion relations as,

EQUATION

EQUATION

EQUATION

EQUATION

)

Furthermore we have,

EQUATION

EQUATION

EQUATION

EQUATION

Finally, we have,

EQUATION

EQUATION

= E

x,y

EQUATION

EQUATION

EQUATION

where we used, .

EQUATION

EQUATION

EQUATION

EQUATION

EQUATION

EQUATION

In order to derive the backward recursion, define (for each l ∈

EQUATION

Then by the chain rule,

EQUATION

Now,

EQUATION

EQUATION

EQUATION

Thus,

EQUATION

By the distributional assumption on the weights, for each (α, β) ∈ K l × K l we have,

EQUATION

Thus, under the usual independence assumptions,

EQUATION

Under review as a conference paper at ICLR 2019 Setting α = β, x = y, averaging over α and taking the expectation value over x,

EQUATION

EQUATION

EQUATION

)

Similarly, setting α = β, averaging over α and taking expectation values over x, y gives

EQUATION

EQUATION

EQUATION

EQUATION

EQUATION

Let us further assume that

EQUATION

Thus, taking expectation values over (x, y) and averaging over the allowable indices such that α = β we obtain, 

EQUATION

EQUATION

EQUATION

EQUATION

and c xy := q l xy /q l , c αβ,xy := q l αβ,xy /q l . Derivation. Substituting into equation 85, equation 89 and equation 94 we obtain, Derivation. In the backward direction, we have

EQUATION

EQUATION

EQUATION

EQUATION

EQUATION

EQUATION

= σ l z l α [i] Recall that the Fisher information matrix I θ ∈ R n×n is given by

EQUATION

)

The claim is that the maximum eigenvalue of the Fisher Information Matrix is bounded as follows

EQUATION

The upper bound follows from convexity of λ max (•). Karakida et al. (2018) obtain the lower bound by considering the empirical estimate of the Fisher Information Matrix,

EQUATION

EQUATION

)

where we have defined the matrix B m ∈ R m×n with components [B m ] ij := ∂f θ (xi) ∂θj . We can then define a symmetric matrix F θ ∈ R m×m with the same eigenvalues as I θ ,

EQUATION

The maximal eigenvalue of I θ is thus computed by the Rayleigh quotient,

EQUATION

Letting 1 ∈ R m denote the vector of ones,

EQUATION

which is the the plug-in estimator (V-statistic). Thus, for a layered neural network we have 

EQUATION

where N l-1 = C l-1 |F l |. In the first approximation we use the fact that the terms with respect to b, γ, β are c l times smaller than the term with respect to W . The second approximation uses the assumption that forward and backward order parameters are independent.

6.4. BASELINE:

Baseline was an experiment of vanilla fully-connected networks trained on MNIST, with various σ w weight initializations. Result is shown in Fig. 3 .

6.5. ADDITIONAL EXPERIMENTS:

Our theory predicts that smaller γ has the effect of greatly reducing the λ max of the FIM, and empirically networks with BatchNorm converge faster. Following this intuition, we performed additional

:

In the quadratic approximation to the loss, the optimal learning rate is in fact η * /2.

Figures and Tables:

Figure 1: The maximum eigenvalue λmax and associated critical learning rate η * for vanilla (blue) and BatchNorm networks (red) as a function of the BatchNorm parameter γ for different choices of architecture (fully-connected and convolutional), calculated by theory. (a, c) shows the flattening effect of BatchNorm on the loss function for a wide range of hyperparameters and (b, d) furthershow that for sufficiently small γ BatchNorm enables optimization with much higher learning rate than vanilla networks.

RECURSIONS FOR FULLY CONNECTED LAYERSClaim 6.1. The forward recursions for 0 ≤ l ≤ L-1 are q l+1 = σ 2 ql xy = γ 2 l /(2π) for l ∈ [L -1] whereas q0 = 1 and q0 xy = 0.

C 0 = 3, C 1 = 20, |K 0 | = 1024 and |F 1 | = 9 and |K 1 | = 225.

) 6.2.1 BATCH NORMALIZATION Claim 6.3. The forward recursions for 0 ≤ l ≤ L -1 are q l+1 = σ 2 b + σ 2 w ql , q l+1 xy = σ 2 b + σ 2 w ql xy and q l+1 αβ,xy = σ 2 b + σ 2 w ql αβ,xy where for each l ∈ [L -1] we have ql = γ 2 l /2 and ql xy = ql αβ,xy = γ 2 l /(2π).Derivation. Recalling that h l α (x) = σ l u l α (x) γ l where u l α

Claim 6.4. The backward recursions are ql =

,xy sin -1 (c αβ,xy ) . (136) Claim 6.6. The backward recursions are ql = -1 (c αβ,xy ) .

l α (x) , σ l z l β (y) δ∈F l+1 δ l+1 α-δ (x), δ l+1 β-δ (y) .

|K l |(|K l | -1)q l αβ,xy + |K l |q l xy C l-1 |F l ||K l |(|K l | -1) |K l | 2 ql-1 αβ,xy + C l-1 |K l ||F l | |K l N l-1 (|K l | -1)q l αβ,xy + ql xy (|K l | -1)q l-1 αβ,xy + ql-1 xy ,

Peer review (iclr2019):
 Interesting application of MFT on FIM to understand Batch Normalization This paper applies mean field analysis to networks with batch normalization layers. Analyzing maximum eigenvalue of the Fisher Information Matrix, the authors provide theoretical evidence of allowing higher learning rates and faster convergence of networks with batch normalization. The analysis reduces to providing lower bound for maximum eigenvalue of FIM using mean-field approximation. Authors provide lower bound of the maximum eigenvalue in the case of fully-connected and convolutional networks with batch normalization layers. Lastly authors observe empirical correlation between smaller .gamma and lower test loss. Pro: - Clear result providing theoretical ground for commonly observed effects. - Experiments are simple but illustrative. It is quite surprising how well the maximum learning rate prediction matches with actual training performance curve. Con: - While mean field analysis a-priori works in the limit where networks width goes to infinity for fixed dataset size, the analysis of Fisher and Batch normalization need asymptotic limit of dataset size. - Although some interesting results are provided. The content could be expanded further for conference submission. The prediction on maximum learning rate is interesting and the concrete result from mean field analysis - While correlation between batch norm .gamma parameter and test loss is also interesting, the provided theory does not seem to provide good intuition about the phenomenon. Comments: - The theory provides the means to compute lower bound of maximum eigenvalue of FIM using mean-field theory. In Figure 1, is .bar .lambda_{max} computed using the theory or empirically computed on the actual network? It would be nice to make this clear. - In Figure 2, the observed .eta_*/2 of dark bands in heatmap is interesting. While most of networks without Batch Norm, performance is maximized using learning rates very close to maximal value, often networks using batch norm the learning rate with maximal performance is not the maximal one and it would be interesting to provide theoretical - I feel like section 3.2 should cite Xiao et al (2018). Although this paper is cited in the intro, the mean field analysis of convolutional layers was first worked out in this paper and should be credited.