Year: 2018

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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 PRELIMINARIES
text: 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.

heading: 2.1 FISHER INFORMATION MATRIX AND LEARNING DYNAMICS
text: 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).

heading: 3 THEORY
text: 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.

heading: 3.1 FULLY CONNECTED LAYERS
text: 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)

heading: 3.1.1 ORDER PARAMETERS AND THEIR RECURSIONS
text: To investigate the spectral properties of the FIM, we define the following order parameters,
ql := 1
Nl E x,θ
[ ‖zl(x)‖2 ] , q̂l := 1
Nl E x,θ
[ ‖hl(x)‖2 ] , (17)
qlxy := 1
Nl E x,y,θ
[ 〈zl(x), zl(y)〉 ] , q̂lxy := 1
Nl E x,y,θ
[ 〈hl(x), hl(y)〉 ] , (18)
q̃l := E x,θ
[ ‖δl(x)‖2 ] , q̃lxy := E
x,y,θ
[ 〈δl(x), δl(y)〉 ] , (19)
where ‖ • ‖ denotes the Euclidean norm and δl(x) := ∂fθ ∂zl (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: q̂0 = 0, q̂0xy = 1, q̃
L = q̃Lxy = 1. The order parameters in the absence of BatchNorm satisfy the following recursions derived in Karakida et al. (2018),
ql = σ2b + σ 2 wq̂
l−1 , q̂l = ql
2 , (20)
qlxy = σ 2 b + σ 2 wq̂ l−1 xy , q̂ l xy =
ql
2π
(√ 1− c2xy + cxyπ
2 + cxy sin
−1 cxy ) , (21)
q̃l = σ2w 2 q̃l+1 , q̃lxy = σ2wq̃ l+1 xy 2π (π 2 + sin−1 cxy ) , (22)
where cxy := qlxy/q l. In the case of batch normalization we find the following recursions, which are derived in the Supplementary Material,
ql = σ2b + σ 2 wq̂ l−1 , q̂l = γ2l 2 , (23)
qlxy = σ 2 b + σ 2 wq̂ l−1 xy , q̂ l xy = γ2l 2π , (24)
q̃l = γ2l σ 2 w
2
q̃l+1
ql , q̃lxy =
γ2l σ 2 w
4 q̃l+1xy ql . (25)

heading: 3.2 CONVOLUTIONAL LAYER
text: 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 Kl denote the set of allowable spatial locations of the the lth layer feature map and let Fl+1 index the sites of the convolutional kernel applied to that layer. Let Cl denote the number of input channels. The output of a general convolutional layer is of the form,
zl+1α (x) = ∑
β∈Fl+1
W l+1β h l α+β(x) + b l+1 , (26)
where α ∈ Kl+1, W l+1β ∈ RCl+1×Cl and bl+1 ∈ RCl+1 . The weights and biases are now distributed as
[W l+1α ]ij ∼ N(0, σ2w/Nl) , bl+1i ∼ N(0, σ 2 b) . (27)
where now Nl := Cl|Fl+1|. As in the fully connected case, we consider convolutional layers with both vanilla activation functions of the form hlα(x) := σ(z l α(x)) as well as batch normalized convolutional layers, for which the input activations satisfy the recursive identity,
hlα(x) := σ
( zlα(x)− µlα
slα γl + βl
) , (28)

heading: 3.2.1 ORDER PARAMETERS AND THEIR RECURSIONS
text: Similar to the definitions for fully connected layer, we define the following set of order parameters:
ql := 1
Cl E α E x,θ
[ ‖zlα(x)‖2 ] , q̂l := 1
Cl E α E x,θ
[ ‖hlα(x)‖2 ] , (29)
qlxy := 1
Cl E α E x,y,θ
[ 〈zlα(x), zlα(y)〉 ] , q̂lxy := 1
Cl E α E x,y,θ
[ 〈hlα(x), hlα(y)〉 ] , (30)
qlαβ,xy := 1
Cl E α 6=β
[ E
x,y,θ 〈zlα(x), zlβ(y)〉
] , q̂lαβ,xy := 1
Cl E α 6=β
[ E
x,y,θ 〈hlα(x), hlβ(y)〉
] , (31)
q̃l := E α E x,θ
[ ‖δlα(x)‖2 ] , q̃lxy := E
α
[ E
x,y,θ 〈δlα(x), δlα(y)〉
] , (32)
q̃lαβ,xy := E α6=β
[ E
x,y,θ 〈δlα(x), δlβ(y)〉
] , (33)
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:
ql = σ2b + σ 2 wq̂l−1 , q̂ l = ql 2 , (34)
qlxy = σ 2 b + σ 2 wq̂ l−1 xy , q̂ l xy =
ql
2π
(√ 1− c2xy + cxyπ
2 + cxy sin
−1 cxy ) , (35)
qlαβ,xy = σ 2 b + σ 2 wq̂ l−1 αβ,xy , q̂ l αβ,xy =
ql
2π
(√ 1− c2αβ + cαβπ
2 + cαβ sin
−1 cαβ ) , (36)
q̃l = σ2w 2 q̃l+1 , q̃lxy = σ2wq̃ l+1 xy 2π (π 2 + sin−1 cxy ) , (37)
q̃lαβ,xy = σ2wq̃ l+1 αβ,xy
2π
(π 2 + sin−1 cαβ ) . (38)
In the case of convolutional layers with BatchNorm, the following recursions hold:
ql = σ2b + σ 2 wq̂l−1 , q̂ l = γ2l 2 , (39)
qlxy = σ 2 b + σ 2 wq̂ l−1 xy , q̂ l xy = γ2l 2π , (40)
qlαβ,xy = σ 2 b + σ 2 wq̂ l−1 αβ,xy , q̂ l αβ,xy = γ2l 2π , (41)
q̃l = γ2l σ 2 w
2
q̃l+1
ql , q̃lxy =
γ2l σ 2 w
4 q̃l+1xy ql , (42)
q̃lαβ,xy = γ2l σ 2 w
4 q̃l+1αβ,xy ql , (43)
where cxy := qlxy/q l and cαβ := qlαβ,xy/q l. The derivations of the recursion relations for both vanilla and batch-normalized convolutional layers are deferred to the Supplementary Material.

heading: 3.3 EIGENVALUE BOUND AND THERMODYNAMIC VARIABLES
text: 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 Nl of the fully connected layers and the number of channels Cl of the convolutional layers satisfy Nl 1 and Cl 1 for 0 < l < L, we have,
λ̄max ≥ ∑ l∈[L] fl , (44)
where fl = Nl−1q̂l−1xy q̃ l xy for fully connected layers and
fl = Nl−1 [ (|Kl| − 1)q̃lαβ,xy + q̃lxy ] [ (|Kl| − 1)q̂l−1αβ,xy + q̂ l−1 xy ] , (45)
for convolutional layers, where recall Nl−1 = Cl−1|Fl|. The index sets Fl and Kl 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.

heading: 4 NUMERICAL ANALYSIS AND EXPERIMENTS
text: 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 CIFAR10 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 (γ, log10 η)-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 Nl = 1000 hidden units per layer except the final (linear) layer which has NL = 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 C1 = 30, C2 = 60, C3 = 90. The final layer is a fully connected output layer to perform classification. The other architectural/optimization hyperparameters were chosen to be σ2w = 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.

heading: 5 CONCLUSION AND FUTURE WORK
text: 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.

heading: 6 SUPPLEMENTARY MATERIAL
text: 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).

heading: 6.1 RECURSIONS FOR FULLY CONNECTED LAYERS
text: Claim 6.1. The forward recursions for 0 ≤ l ≤ L−1 are ql+1 = σ2b+σ2wq̂l and ql+1xy = σ2b+σ2wq̂lxy where q̂l = γ2l /2 and q̂ l xy = γ 2 l /(2π) for l ∈ [L− 1] whereas q̂0 = 1 and q̂0xy = 0.
Derivation. In general, we have
1
Nl+1 E θ 〈zl+1(x), zl+1(y)〉 = σ2b + σ2w Nl 〈hlθ(x), hlθ(y)〉 . (46)
Thus, setting l = 0 we obtain q1 = σ2w + σ 2 b if we assume Ex‖x‖ 2 = N0, and q1xy = σ 2 b since E x,y 〈x, y〉 = 〈Ex,E y〉 = 0.
Recall (for l > 0) zl+1(x) = W l+1σl ( ul(x) γl ) + bl+1 where ul(x) = [zl(x)− µl]/sl so
1
Nl+1 E θ 〈zl+1(x), zl+1(y)〉 = σ2b + σ2w Nl
〈 σl ( ul(x) γl ) , σl ( ul(y) γl )〉 . (47)
Therefore, setting x = y and taking expectation values over x gives,
ql+1 := 1
Nl+1 E x,θ ‖zl+1(x)‖2 , (48)
= σ2b + σ2w Nl E x,θ ‖hl(x)‖2 , (49)
= σ2b + σ 2 w q̂ l , (50)
q̂l := 1
Nl E x,θ ‖hl(x)‖2 , (51) = σ2w E x,θ σ2l ( γl u l(x)[1] ) , (52)
' σ2w ∫ Dz σ2l (γl z) , (53)
= γ2l 2 , (54)
where we have approximated each component of the random vector ul(x) as a standard Gaussian.
Similarly, taking expectations over x, y gives
ql+1xy = σ 2 b + σ 2 w E x,y,θ
σl ( ul(x)[1] γl ) σl ( ul(y)[1] γl ) . (55)
Consider the approximation in which the random pair (ul(x)[1], ul(y)[1]) is Gaussian distributed with zero mean and covariance,
Σl :=
( Σlxx Σ l xy
Σlxy Σ l yy
) := ( Eul(x)[1]2 Eul(x)[1]ul(y)[1] Eul(x)[1]ul(y)[1] Eul(y)[1]2 ) . (56)
Then the recursion becomes ql+1xy ' σ2b + σ2w ∫ Dz1Dz2 σl (√ Σlxx z1γl ) σl [√ Σlyy ( clxy z1 + √ 1− (clxy)2 z2 ) γl ] , (57)
where clxy := Σ l xy/ √ ΣlxxΣ l yy. Observe that by independence of x and y,
E x,y
[ ul(x)ul(y) ] = E x,y
[ zl(x)zl(y) + µ 2 l − µl ( zl(x) + zl(y) ) s2l ] , (58)
= Exzl(x)Eyzl(y) + µ2l − µl
( Exzl(x) + Eyzl(y) ) s2l , (59)
= 0 . (60)
Thus Σlxy = 0 and consequently,
ql+1xy = σ 2 b + σ 2 wq̂ l xy , (61)
q̂lxy = γ2l 2π . (62)
The derivation here assumes an infinitely large dataset. For limited dataset size m, an error is introduced from the non-zero ratio of m/mx 6=y where mx 6=y denotes the total number of sample pairs (x, y) where x 6= y. We havemx6=y = m2−m and therefore the error isO(1/m) which is negligible for most of the frequently-used datasets.
Claim 6.2. The backward recursions for l ∈ [1, L − 1] are q̃l = γ 2 l σ 2 w 2 q̃l+1 ql and q̃lxy = γ2l σ 2 w 4 q̃l+1xy ql with base cases q̃L = q̃Lxy = 1.
Derivation. For each layer l ∈ [L] and for each unit i ∈ [Nl] define δl[i] ∈ RNL by
δl[i] := ∂hLθ ∂zl[i] . (63)
In particular, since we assume linear output δL[i] = ei and thus
q̃L = E x,θ ‖δL[·]‖2 = 1 , (64)
where 1 ∈ RNL is the vector of ones. For ease of presentation and without loss of generality we restrict to the first component q̃L[1] = 1 and abuse notation by writing q̃L = 1. For l < L we have by the chain rule,
δl[i] = ∑
k∈[Nl+1]
∂zl+1[k]
∂zl[i] δl+1[k] , (65)
= γl[i] sl[i] σ′l ( ul[i] γl[i] ) ∑ k∈[Nl+1] [W l+1]kiδ l+1[k] . (66)
Hence,
δl(x)[i]δl(y) = γl[i]
2 sl[i]2 σ′l ( ul(x)[i] γl[i] ) σ′l ( ul(y)[i] γl[i] ) ∑ k,k′∈[Nl+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,
E θ 〈δl(x), δl(y)〉 = σ
2 w
Nl ∑ i∈[Nl] E θ [ γl[i] 2 sl[i]2 σ′l ( ul(x)[i] γl[i] ) σ′l ( ul(y)[i] γl[i] )] E θ 〈δl+1(x), δl+1(y)〉
(68)
Setting x = y and taking expectation values over x we obtain,
q̃l := E x,θ ‖δl(x)‖2 , (69)
' q̃l+1 ql γ2l σ 2 w
∫ Dz σ′l(γlz) 2 , (70)
= γ2l σ 2 w
2 q̃l+1 ql . (71)
Similarly, taking expectation values over x, y we obtain,
q̃lxy = E x,y,θ 〈δl(x), δl(y)〉 , (72)
= σ2wγ 2 l q̃l+1xy ql
∫ Dz1Dz2σ ′ l (√ Σlxxz1 γl ) σ′l [√ Σlyy ( clxy z1 + √ 1− (clxy)2 z2 ) γl ] , (73)
= σ2wγ 2 l
4 q̃l+1xy ql . (74)

heading: 6.2 RECURSIONS FOR CONVOLUTIONAL LAYERS
text: Recall that the output of a general convolutional layer is of the form, zl+1α (x) = ∑
β∈Fl+1
W l+1β h l α+β(x) + b l+1 , (75)
where α ∈ Kl+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 C0 = 3, C1 = 20, |K0| = 1024 and |F1| = 9 and |K1| = 225. We begin by deriving some useful identities for convolutional layers, before specializing to the batch-normalized and vanilla networks. For each channel i ∈ [Cl], we have,
E θ zl+1α (x)[i]z l+1 β (y)[i] = σ 2 b + E θ ∑ (β1,j1)∈Fl+1×[Cl] (β2,j2)∈Fl+1×[Cl] [W l+1β1 ]ij1 [W l+1 β2 ]ij2h l α+β1(x)[j1]h l β+β2(y)[j2] ,
= σ2b + σ2w Nl ∑ δ∈Fl+1 E θ 〈 hlα+δ(x), h l β+δ(y) 〉 . (76)
Hence, 1
Cl+1 E θ
〈 zl+1α (x), z l+1 β (y) 〉 = σ2b +
σ2w Nl ∑ δ∈Fl+1 E θ 〈 hlα+δ(x), h l β+δ(y) 〉 , (77)
where Nl := Cl|Fl+1|. Moreover, let us introduce the following shorthand,
E α ‖hlα(x)‖2 :=
1 |Kl| ∑ α∈Kl ‖hlα(x)‖2 , (78)
E α 〈hlα(x), hlα(y)〉 :=
1 |Kl| ∑ α∈Kl 〈 hlα(x), h l α(y) 〉 , (79)
E α 6=β
〈 hlα(x), h l β(y) 〉 :=
1 |Kl|(|Kl| − 1) ∑
α,β∈Kl×Kl:α 6=β
〈 hlα(x), h l β(y) 〉 (80)
Then we can write the recursion relations as,
ql+1 := 1
Cl+1 E α E x,θ ‖zl+1α (x)‖2 , (81)
= E x σ2b + σ2wNl ∑ δ∈Fl+1 E θ E α ‖hlα+δ(x)‖2  , (82) = σ2b +
σ2w Cl E x,θ E α ‖hlα(x)‖2 , (83)
= σ2b + σ 2 w q̂ l , (84)
q̂l := 1
Cl E x,θ E α ‖hlα(x)‖2 . (85)
Furthermore we have,
ql+1xy = 1
Cl+1 E α E x,y,θ 〈zl+1α (x), zl+1α (y)〉 , (86)
= E x,y σ2b + σ2wNl ∑ δ∈Fl+1 E θ E α 〈hlα+δ(x), hlα+δ(y)〉  , (87) = σ2b + σ 2 w q̂ l xy , (88)
q̂lxy := 1
Cl E α E x,y,θ 〈hlα(x), hlα(y)〉 . (89)
Finally, we have,
ql+1αβ,xy := E α6=β
[ E
x,y,θ 〈zl+1α (x), zl+1β (y)〉
] , (90)
= E x,y σ2b + σ2wNl ∑ δ∈Fl+1 E θ E α6=β 〈hlβ+δ(x), hlα+δ(y)〉  , (91) = E x,y [ σ2b + σ2w Cl E θ E α6=β 〈hlβ(x), hlα(y)〉 ] , (92)
= σ2b + σ 2 wq̂ l αβ,xy , (93)
q̂lαβ,xy := 1
Cl E α 6=β E x,y,θ 〈hlα(x), hlβ(y)〉 , (94)
where we used,∑ δ∈Fl+1 ∑ α,β∈Kl×Kl:α6=β 〈 hlα+δ(x), h l β+δ(y) 〉 = (95)
∑ δ∈Fl+1 〈∑ α∈Kl hlα+δ(x), ∑ β∈Kl hlβ+δ(y) 〉 − ∑ α∈Kl 〈 hlα+δ(x), h l α+δ(y) 〉 (96) = |Fl+1|
〈∑ α∈Kl hlα(x), ∑ β∈Kl hlβ(y) 〉 − ∑ α∈Kl 〈 hlα(x), h l α(y) 〉 (97) = |Fl+1|
∑ α,β∈Kl×Kl:α 6=β 〈 hlα(x), h l β(y) 〉 . (98)

heading: 6.2.1 BATCH NORMALIZATION
text: Claim 6.3. The forward recursions for 0 ≤ l ≤ L− 1 are ql+1 = σ2b + σ2wq̂l, ql+1xy = σ2b + σ2wq̂lxy and ql+1αβ,xy = σ 2 b + σ 2 wq̂ l αβ,xy where for each l ∈ [L − 1] we have q̂l = γ2l /2 and q̂lxy = q̂lαβ,xy = γ2l /(2π).
Derivation. Recalling that hlα(x) = σl ( ulα(x) γl ) where ulα(x) := zlα(x)−µ l α
slα and substituting
into equation 85, equation 89 and equation 94 we obtain, q̂l ' ∫ Dz σ2l (γlz) , (99)
= γ2l 2 , (100)
q̂lxy ' ∫ Dz1Dz2 σl(γlz1)σl(γ1z2) , (101)
= γ2l 2π , (102)
q̂αβ,xy ' ∫ Dz1Dz2 σl(γlz1)σl(γ1z2) , (103)
= γ2l 2π , (104)
Claim 6.4. The backward recursions are q̃l = σ 2 wγ 2 l q̃ l+1
2ql , q̃lxy =
σ2wγ 2 l q̃ l+1 xy
4ql , and q̃lαβ,xy =
σ2wγ 2 l q̃ l+1 αβ,xy
4ql .
In order to derive the backward recursion, define (for each l ∈ [L], i ∈ [Cl], α ∈ Kl)
δlα[i] = ∂hLθ ∂zlα[i] . (105)
Then by the chain rule,
δlα[i] = ∑
(β,k)∈Kl+1×[Cl+1]
∂zl+1β [k]
∂zlα[i] δl+1β [k] . (106)
Now,
zl+1β [k] = ∑
(β′,j)∈Fl+1×[Cl]
[W l+1β′ ]kjσl ( ulβ+β′ [j] γl[j] ) + bl+1[k] , (107)
∂zl+1β [k]
∂zlα[i] = ∑ (β′,j)∈Fl+1×[Cl] [W l+1β′ ]kjσ ′ l ( ulβ+β′ [j] γl[j] ) γl[j] slβ+β′ [j] δijδα,β+β′ , (108)
= [W l+1α−β ]kiσ ′ l ( ulα[i] γl[i] ) γl[i] slα[i] . (109)
Thus,
δlα[i] = γl[i] slα[i] σ′l ( ulα[i] γl[i] ) ∑ (β,k)∈Fl+1×[Cl+1] [W l+1β ]kiδ l+1 α−β [k] . (110)
By the distributional assumption on the weights, for each (α, β) ∈ Kl ×Kl we have,
E θ ∑ (β1,k1)∈Fl+1×[Cl+1] (β2,k2)∈Fl+1×[Cl+1] [W l+1β1 ]k1i[W l+1 β2 ]k2iδ l+1 α−β1 [k1]δ l+1 β−β2 [k2] = σ2w Nl ∑ δ∈Fl+1 E θ 〈δl+1α−δ(x), δ l+1 β−δ(y)〉 .
(111)
Thus, under the usual independence assumptions,
E θ 〈δlα(x), δlβ(y)〉 = σ2w Nl ∑ i∈[Cl] E θ
[ γl[i] 2
slα[i]s l β [i]
σ′l ( ulα(x)[i] γl[i] ) σ′l ( ulβ(y)[i] γl[i] )] ∑ δ∈Fl+1 E θ 〈δl+1α−δ(x), δ l+1 β−δ(y)〉 .
(112)
Setting α = β, x = y, averaging over α and taking the expectation value over x,
q̃l := E α E x,θ ‖δlα(x)‖2 , (113)
= σ2w Nl γ2l 2ql Cl E x,θ ∑ δ∈Fl+1 E α ‖δlα−δ(x)‖2 , (114)
= σ2w Nl γ2l 2ql Cl|Fl+1|E α E x,θ ‖δlα(x)‖2 , (115)
= σ2wγ 2 l q̃ l+1
2ql . (116)
Similarly, setting α = β, averaging over α and taking expectation values over x, y gives
q̃lxy := E α E x,y,θ 〈δlα(x), δlα(y)〉 , (117)
= σ2wγ 2 l q̃ l+1
4ql . (118)
Now, ∑ α,β∈Kl+1×Kl+1:α6=β
δ∈Fl+1
〈δl+1α−δ(x), δ l+1 β−δ(y)〉 (119)
= ∑
δ∈Fl+1 〈 ∑ α∈Kl+1 δl+1α−δ(x), ∑ β∈Kl+1 δl+1β−δ(y) 〉 − ∑ α∈Kl+1 〈δl+1α−δ(x), δ l+1 α−δ(y)〉  , (120) = |Fl+1|
〈 ∑ α∈Kl+1 δl+1α (x), ∑ β∈Kl δl+1β (y) 〉 − ∑ α∈Kl+1 〈δl+1α (x), δl+1α (y)〉  , (121) = |Fl+1|
∑ α,β∈Kl+1×Kl+1:α 6=β 〈 δl+1α (x), δ l+1 β (y) 〉 . (122)
Let us further assume that
1 |Kl|(|Kl| − 1) ∑
α,β∈Kl×Kl:α 6=β
E x,y,θ
〈 δl+1α (x), δ l+1 β (y) 〉 (123)
= 1 |Kl+1|(|Kl+1| − 1) ∑
α,β∈Kl+1×Kl+1:α6=β
E x,y,θ
〈 δl+1α (x), δ l+1 β (y) 〉 . (124)
Thus, taking expectation values over (x, y) and averaging over the allowable indices such that α 6= β we obtain,
E α 6=β
[ E
x,y,θ 〈δlα(x), δlβ(y)〉 ] = Cl|Fl+1| Nl σ2wγ 2 l 4ql E α6=β [ E x,y,θ 〈δl+1α (x), δl+1β (y)〉 ] . (125)
It follows that
q̃lαβ,xy := E α6=β
[ E
x,y,θ
〈 δlα(x), δ l β(y) 〉] , (126)
= σ2wγ 2 l q̃ l+1 αβ,xy
4ql . (127)

heading: 6.2.2 VANILLA CNN
text: Claim 6.5. The forward recursions are ql+1 = σ2b + σ2wq̂l, ql+1xy = σ2b + σ2wq̂lxy and q l+1 αβ,xy = σ2b + σ 2 wq̂ l αβ,xy where
q̂l = 1
2 ql , (128)
q̂lxy = ql
2π
[√ 1− c2xy + cxyπ
2 + cxy sin
−1(cxy) ] , (129)
q̂lαβ,xy = ql
2π
[√ 1− c2αβ,xy + cαβ,xyπ
2 + cαβ,xy sin
−1(cαβ,xy) ] , (130)
and cxy := qlxy/q l, cαβ,xy := qlαβ,xy/q l.
Derivation. Substituting into equation 85, equation 89 and equation 94 we obtain,
q̂l := ∫ Dz σ2l ( √ qlz) , (131)
= ql
2 , (132)
q̂lxy = ∫ Dz1Dz2 σl (√ ql z1γl ) σl [√ ql ( clxy z1 + √ 1− (clxy)2 z2 ) γl ] , (133)
= ql
2π
[√ 1− c2xy + cxyπ
2 + cxy sin
−1(cxy) ] , (134)
q̂lαβ,xy = ∫ Dz1Dz2 σl (√ ql z1γl ) σl [√ ql ( clαβ,xy z1 + √ 1− (clαβ,xy)2 z2 ) γl ] , (135)
= ql
2π
[√ 1− c2αβ,xy + cαβ,xyπ
2 + cαβ,xy sin
−1(cαβ,xy) ] . (136)
Claim 6.6. The backward recursions are q̃l = σ 2 w q̃ l+1
2 , q̃ l xy =
σ2w q̃ l+1 xy
2π [ π 2 + sin −1(cxy) ] and
q̃lαβ,xy = σ2w q̃ l+1 αβ,xy 2π [ π 2 + sin −1(cαβ,xy) ] .
Derivation. In the backward direction, we have
δlα[i] = ∑
(β,k)∈Kl+1×[Cl+1]
∂zl+1β [k]
∂zlα[i] δl+1β [k] , (137)
= σ′l ( zlα[i] ) ∑ (β,k)∈Fl+1×[Cl+1] [W l+1β ]kiδ l+1 α−β [k] (138)
Under the usual independence assumptions,
E θ 〈δlα(x), δlβ(y)〉 = σ2w Nl E θ
〈 σ′l ( zlα(x) ) , σ′l ( zlβ(y) )〉 ∑ δ∈Fl+1 〈δl+1α−δ(x), δ l+1 β−δ(y)〉 . (139)
Thus,
q̃l := E α E x,θ ‖δlα(x)‖2 , (140)
= σ2wq̃ l+1
2 (141)
q̃lxy := E α E x,y,θ 〈δlα(x), δlα(y)〉 , (142)
= σ2wq̃ l+1 xy
2π (π 2 + sin−1 cxy ) (143)
q̃lαβ,xy := E α6=β
[ E
x,y,θ
〈 δlα(x), δ l β(y) 〉] , (144)
= σ2wq̃ l+1 αβ,xy
2π
(π 2 + sin−1 cαβ,xy ) . (145)

heading: 6.3 DERIVATION OF CLAIM 3.1
text: Recall that the Fisher information matrix Iθ ∈ Rn×n is given by
Iθ = E x∼D
[ ∇θfθ(x)⊗∇θfθ(x) ] . (146)
The claim is that the maximum eigenvalue of the Fisher Information Matrix is bounded as follows
E(x,y)∼D〈∇θfθ(x),∇θfθ(y)〉 ≤ λmax(Iθ) ≤ Ex∼D〈∇θfθ(x),∇θfθ(x)〉 . (147)
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,
Îθ = 1
m m∑ i=1 ∇θfθ(xi)⊗∇θfθ(xi) , (148)
= 1
m B>mBm . (149)
where we have defined the matrix Bm ∈ Rm×n with components [Bm]ij := ∂fθ(xi)∂θj . We can then define a symmetric matrix F̂θ ∈ Rm×m with the same eigenvalues as Îθ,
F̂θ = 1
m BmB
> m , (150)
The maximal eigenvalue of Îθ is thus computed by the Rayleigh quotient,
λmax(Îθ) = λmax(F̂θ) = max v:‖v‖=1
〈v, F̂θ v〉 . (151)
Letting 1 ∈ Rm denote the vector of ones, λmax(Îθ) ≥ 〈
1√ m 1, F̂θ 1√ m 1
〉 , (152)
= 1
m2 ∑ i,j∈[m] 〈 ∇θfθ(xi),∇θfθ(xj) 〉 , (153)
which is the the plug-in estimator (V-statistic). Thus, for a layered neural network we have
λ̄max(Iθ) ≥ 1
m2 ∑ (x,y) ∑ l∈[L] 〈 ∇θlfθ(x),∇θlfθ(y) 〉 , (154)
=: ∑ l∈[L] fl. (155)
Specializing to a fully-connected neural network we obtain,
fl = 1
m2 ∑ (x,y) [〈 ∂fθ(x) ∂bl , ∂fθ(y) ∂bl 〉 + 〈 ∂fθ(x) ∂Wl , ∂fθ(y) ∂Wl 〉 (156)
+
〈 ∂fθ(x)
∂γl , ∂fθ(y) ∂γl
〉 + 〈 ∂fθ(x)
∂βl , ∂fθ(y) ∂βl
〉] , (157)
' 1 m2 ∑ (x,y) 〈 ∂fθ(x) ∂Wl , ∂fθ(y) ∂Wl 〉 , (158)
= 1
m2 ∑ (x,y) 〈 hl−1(x), hl−1(y) 〉〈 δl(x), δl(y) 〉 , (159)
'  1 m2 ∑ (x,y) 〈 hl−1(x), hl−1(y) 〉 1 m2 ∑ (x,y) 〈 δl(x), δl(y) 〉 , (160) = Nl−1q̂ l−1 xy q̃ l xy . (161)
In the first approximation we use the fact that the terms with respect to b, γ, β are Nl times smaller than the term with respect to W . The second approximation uses the assumption that forward and backward order parameters are independent. The last approximation is for m 1. For convolutional layers we have
fl = 1
m2 ∑ (x,y) [〈 ∂fθ(x) ∂bl , ∂fθ(y) ∂bl 〉 + 〈 ∂fθ(x) ∂Wl , ∂fθ(y) ∂Wl 〉 (162)
+
〈 ∂fθ(x)
∂γl , ∂fθ(y) ∂γl
〉 + 〈 ∂fθ(x)
∂βl , ∂fθ(y) ∂βl
〉] , (163)
' 1 m2 ∑ (x,y) 〈 ∂fθ(x) ∂Wl , ∂fθ(y) ∂Wl 〉 , (164)
= 1
m2 ∑ (x,y) ∑ α∈Fl ∑ β,β′∈Kl 〈δlβ(x), δlβ′(y)〉〈hl−1α+β(x), h l−1 α+β′(y)〉 , (165)
' 1 |Kl|2  1 m2 ∑ (x,y) ∑ β,β′∈Kl 〈δlβ(x), δlβ′(y)〉  1 m2 ∑ (x,y) ∑ α∈Fl ∑ β,β′∈Kl 〈hl−1α+β(x), h l−1 α+β′(y)〉  , (166)
= [ |Kl|(|Kl| − 1)q̃lαβ,xy + |Kl|q̃lxy ] [Cl−1|Fl||Kl|(|Kl| − 1) |Kl|2 q̂l−1αβ,xy + Cl−1|Kl||Fl| |Kl|2 q̂l−1xy ] ,
(167) = Nl−1 [ (|Kl| − 1)q̃lαβ,xy + q̃lxy ] [ (|Kl| − 1)q̂l−1αβ,xy + q̂ l−1 xy ] , (168)
where Nl−1 = Cl−1|Fl|. In the first approximation we use the fact that the terms with respect to b, γ, β are cl times smaller than the term with respect to W . The second approximation uses the assumption that forward and backward order parameters are independent.

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

heading: 6.5 ADDITIONAL EXPERIMENTS
text: 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
experiments with VGG16 Simonyan & Zisserman (2014) and Preact-Resnet18 He et al. (2016a), with various γ initializations, fixed learning rate 0.1, momentum 0.9 and weight decay 0.0005, trained on CIFAR-10. The result is average over 5 different independent trainings. We find that smaller γ initialization indeed increases the speed of convergence, as shown below in Fig. 4.

Review summaries:
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.