A MEAN FIELD THEORY OF BATCH NORMALIZATION
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.
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.