Title: A Mean Field Theory of Batch Normalization
Authors: Greg Yang
Jeffrey Pennington
Vinay Rao
Jascha Sohl-Dickstein
Samuel S. Schoenholz
Venue: ICLR 2019
Keywords: theory
batch normalization
mean field theory
trainability

Abstract: We develop a mean field theory for batch normalization in fully-connected feedforward neural networks. In so doing, we provide a precise characterization of signal propagation and gradient backpropagation in wide batch-normalized networks at initialization. Our theory shows that gradient signals grow exponentially in depth and that these exploding gradients cannot be eliminated by tuning the initial weight variances or by adjusting the nonlinear activation function. Indeed, batch normalization itself is the cause of gradient explosion. As a result, vanilla batch-normalized networks without skip connections are not trainable at large depths for common initialization schemes, a prediction that we verify with a variety of empirical simulations. While gradient explosion cannot be eliminated, it can be reduced by tuning the network close to the linear regime, which improves the trainability of deep batch-normalized networks without residual connections. Finally, we investigate the learning dynamics of batch-normalized networks and observe that after a single step of optimization the networks achieve a relatively stable equilibrium in which gradients have dramatically smaller dynamic range. Our theory leverages Laplace, Fourier, and Gegenbauer transforms and we derive new identities that may be of independent interest.

TL;DR: Batch normalization causes exploding gradients in vanilla feedforward networks.

Review discussion:
This paper provides a new dynamic perspective on deep neural network. Based on Gaussian weights and biases, the paper investigates the evolution of the covariance matrix along with the layers.
Eventually the matrices achieve a stationary point, i.e., fixed point of the dynamic system. Local performance around the fixed point is explored. Extensions are provided to include the batch normalization. I believe this paper may stimulate some interesting ideas for other researchers.
Two technical questions:
1.
When the layers tends to infinity, the covariance matrix reaches stationary (fixed) point. How to understand this phenomenon? Does this mean that the distribution of the layer outputs will not change too much if the layer is deep enough? This somewhat conflicts the commonsense of "the deeper the better?"
2. Typos: the weight matrix in the end of page 2 should be N_l times N_{l-1}.
Also, the x_i's in the first line of page 3 should be bold.
Thank you for the review and careful reading of our paper! We’re glad that you found it of interest. On revision we will fix the typos that you identified.
Regarding the first point, your intuition is exactly correct and a slightly simpler discussion of this phenomenon can be found in [1].
When the network is deep enough that the covariance matrix has reached its fixed point, the distribution of the outputs of the network will be independent of the inputs. At this point the network becomes untrainable. To reconcile this with the commonsense intuition that “deeper is better”, our answer is twofold.
1) As in [1] and [2] it is often possible to find configurations or architectural modifications where the covariance matrix doesn’t approach its fixed point over depths often considered in machine learning.
When this is the case one can safely increase the depth without sacrificing accuracy.
2) It seems that the role of depth in performance is more subtle than standard intuition would dictate.
For example, in [3] note that although the authors were able to train a 10k hidden layer network, they did not observe any improvement in accuracy.
In the next version of the manuscript (both in response to your review and that of referee 1) we will add a more intuitive discussion of these results which we agree are somewhat technical.
[1] S. S. Schoenholz, J. Gilmer, S. Ganguli, J. Sohl-Dickstein. Deep Information Propagation (https://arxiv.org/abs/1611.01232)
[2] G. Yang and S. S. Schoenholz. Mean Field Residual Networks (https://arxiv.org/abs/1712.08969)
[3] L. Xiao, Y. Bahri, J. Sohl-Dickstein, S. S. Schoenholz, J. Pennington. Dynamical Isometry and a Mean Field Theory of CNNs: How to Train 10,000-Layer Vanilla Convolutional Neural Networks (https://arxiv.org/abs/1806.05393)
This paper investigates the effect of the batch normalization in DNN learning.
The mean field theory in statistical mechanics was employed to analyze the
progress of variance matrices between layers.
As the results, the batch normalization itself is found to be the cause of gradient explosion.
Moreover, the authors pointed out that near-linear activation function can improve such gradient explosion.
Some numerical studies were reported to confirm theoretical findings.
The detailed analysis of the training of DNN with the batch normalization is quite interesting.
There are some minor comments below.
- in page 3, 2line above eq(2): what is delta in the variance of the multivariate normal distribution?
- the notation q appeared in the middle part of page 3 before the definition of q is shown in the last paragraph of p.3.
- The randomized weight is not very practical.
Though it may be the standard approach of mean field,
some comments would be helpful to the readers.
Thank you for your review and very useful comments! We’re happy you found our manuscript interesting.
To address your comments:
1) Thank you for pointing out that we had not defined the delta. Here delta is the Kronecker delta defined so that \delta_{a,b} = 1 if a = b and 0 if a != b. In the context of the variance of the multivariate normal distribution, the delta function indicates that the different neurons in each layer have zero covariance. We’ll add an explicit discussion of this fact to the manuscript.
2) Thanks for pointing this out, we’ll correct it in the next revision.
3) It is true that the extent to which randomized weights describe trained networks is unclear.
However, it is true that most commonly used weight initialization schemes are random.
For example, He initialization [1] and Xavier initialization [2] strategies are both special cases of the setup considered here. We therefore view our theory as a theory of neural networks at initialization. (There are, however, initialization schemes that are not random and that are not described by our theory).
[1] K. He, X. Zhang, S. Ren, J. Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. (http://www.cv-foundation.org/openaccess/content_iccv_2015/html/He_Delving_Deep_into_ICCV_2015_paper.html)
[2] X. Glorot, Y. Bengio, Y. W. Teh, M. Titterington. Understanding the difficulty of training deep feedforward neural networks. (http://proceedings.mlr.press/v9/glorot10a.html)
This paper develops a mean field theory for batch normalization (BN) in fully-connected networks with randomly initialized weights. There are a number of interesting predictions made in this paper on the basis of this analysis. The main technical results of the paper are Theorems 5-8 which compute the statistics of the covariance of the activations and the gradients.
Comments:
1. The observation that gradients explode in spite of BN is quite counter-intuitive. Can you give an intuitive explanation of why this occurs?
2. In a similar vein, there a number of highly technical results in the paper and it would be great if the authors provide an intuitive explanation of their theorems.
3. Can the statistics of activations be controlled using activation functions or operations which break the symmetry?
For instance, are BSB1 fixed points good for training neural networks?
4. Mean field analysis, although it lends an insight into the statistics of the activations, needs to connected with empirical observations.
For instance, when the authors observe that the structure of the fixed point is such that activations are of identical norm equally spread apart in terms of angle, this is quite far from practice. It would be good to mention this in the introduction or the conclusions.
Thank you for your careful review and useful comments!
Overall, in response to your review and that of referee 3 we will include a more intuitive discussion of our results in the next revision of our text.
To reply to your other specific comments,
1) The intuition for batchnorm can be put in a more general setting. If a function f: X -> Y tends to spread out small clusters in the input space almost evenly in the output space, then one can expect that its gradients will be large typically. In our case, a batchnorm network can be understood as a function that sends a batch of inputs to a batch of outputs. In the appendix, we showed that the correlation between two different batches tend to a constant value independent of the input batches. No matter how close two input batches are, the output batches will have the same “distance” from each other -- small movements in the input space leads to large movements in the output space.
Thus we can expect the gradients to be large as well. We have added a new figure to the Appendix to further support this intuition. In it, we pass through a linear batchnorm network 2 minibatches. Both minibatches contain points on the same circle and 1 point off the circle that is unique to each minibatch.
While the circle in each minibatch will remain an ellipse as they are propagated through the network, the angle between the planes spanned by them increasingly becomes chaotic with depth.
3) As observed in [1] and [2], depthwise convergence to covariance fixed points is bad for training, and the best networks are either moderately deep or initialized such that the depthwise convergence rate to the fixed point is as slow as possible. We observe that deep networks whose activation statistics resemble a non-BSB1 fixed point typically feature worse gradient explosion than BSB1 networks. This seems to be because the nonlinearities that induce these fixed points increase rapidly (for example, polynomials with high degrees), so that the corresponding derivatives are also large, causing gradient explosion.
(The reason that rapidly increasing nonlinearities don’t converge to BSB1 fixed points is that, after a spontaneous symmetry-breaking, begins a “winner-take-all” covariance dynamics, in which the activations of a few examples in the batch suddenly dominates those of the others in the batch, and this dominance persists across each layer.)
4) We were a bit confused by what was meant by “practice” here. We have thoroughly verified that for realistic input distributions (MNIST and CIFAR10) and common initialization strategies (weights that are randomly distributed) our theory makes accurate prediction.
Moreover, we have shown that these predictions can be connected to practice in the sense that they predict whether or not the network can be trained.
Having said this, if by practice you meant that the neural network is accurately described by our theory during training then we do not expect this to be true. We are happy to emphasize this in the camera ready.
If this did not properly address your question, please feel free to let us to know and we will improve this response!
[1] S. S. Schoenholz, J. Gilmer, S. Ganguli, J. Sohl-Dickstein. Deep Information Propagation (https://arxiv.org/abs/1611.01232)
[2] L. Xiao, Y. Bahri, J. Sohl-Dickstein, S. S. Schoenholz, J. Pennington. Dynamical Isometry and a Mean Field Theory of CNNs: How to Train 10,000-Layer Vanilla Convolutional Neural Networks (https://arxiv.org/abs/1806.05393)

Meta-review:
role: AC
data: metareview: This paper provides a mean-field-theory analysis of batch normalization. First there is a negative result as to the necessity of gradient explosion when using batch normalization in a fully connected network. They then provide further insights as to what can be done about this, along with experiments to confirm their theoretical predictions.

The reviewers (and random commenters) found this paper very interesting. The reviewers were unanimous in their vote to accept.

Decision: Accept (Poster)