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. Deep neural networks have been enormously successful across a broad range of disciplines. These successes are often driven by architectural innovations. For example, the combination of convolutions BID15 , residual connections BID12 , and batch normalization BID13 has allowed for the training of very deep networks and these components have become essential parts of models in vision (Zoph et al.) , language BID4 , and reinforcement learning (Silver et al., 2017) . However, a fundamental problem that has accompanied this rapid progress is a lack of theoretical clarity. An important consequence of this gap between theory and experiment is that two important issues become conflated. In particular, it is generally unclear whether novel neural network components improve generalization or whether they merely increase the fraction of hyperparameter configurations where good generalization can be achieved. Resolving this confusion has the promise of allowing researchers to more effectively and deliberately design neural networks.Recently, progress has been made (Poole et al., 2016; Schoenholz et al., 2016; BID6 BID16 BID11 in this direction by considering neural networks at initialization, before any training has occurred. In this case, the parameters of the network are random variables which induces a distribution of the activations of the network as well as the gradients. Studying these distributions is equivalent to understanding the prior over functions that these random neural networks compute. Picking hyperparameters that correspond to well-conditioned priors ensures that the neural network will be trainable and this fact has been extensively verified experimentally. However, to fulfill its promise of making neural network design less of a black box, these techniques must be applied to neural network architectures that are used in practice. Over the past year, this gap has closed significantly and theory for networks with skip connections (Yang & Schoenholz, 2017; , convolutional networks (Xiao et al., 2018) , and gated recurrent networks BID3 BID9 have been developed. More recently, Yang (2019) devised a formalism that extends this approach to include an even wider range of architectures.Before state-of-the-art models can be analyzed in this framework, a slowly-decreasing number of architectural innovations must be studied. One particularly important component that has thus-far remained elusive is batch normalization.Our Contributions. In this paper, we develop a theory of fully-connected networks with batch normalization whose weights and biases are randomly distributed. A significant complication in the case of batch normalization (compared to e.g. layer normalization or weight normalization) is that the statistics of the network depend non-locally on the entire batch. Thus, our first main result is to recast the theory for random fully-connected networks so that it can be applied to batches of data. We then extend the theory to include batch normalization explicitly and validate this theory against Monte-Carlo simulations. We show that as in previous cases we can leverage our theory to predict valid hyperparameter configurations.In the process of our investigation, we identify a number of previously unknown properties of batch normalization that make training unstable. In particular, for most nonlinearities used in practice, batchnorm in a deep, randomly initialized network induces high degree of symmetry in the embeddings of the samples in the batch (Thm 3.4) . Whenever this symmetry takes hold, we show that for any choice of nonlinearity, gradients of fully-connected networks with batch normalization explode exponentially in the depth of the network (Thm 3.9) . This imposes strong limits on the maximum trainable depth of batch normalized networks. This limit can be lifted partially but not completely by pushing activation functions to be more linear at initialization. It might seem that such gradient explosion ought to lead to learning dynamics that are unfavorable. However, we show that networks with batch normalization causes the scale of the gradients to naturally equilibrate after a single step of gradient descent (provided the initial gradients are not so large as to cause numerical instabilities). For shallower networks, this equilibrating effect is sufficient to allow adequate training.Finally, we note that there is a related vein of research that has emerged that leverages the prior over functions induced by random networks to perform exact Bayesian inference BID16 BID7 BID18 BID8 . One of the natural consequences of this work is that the prior for networks with batch normalization can be computed exactly in the wide network limit. As such, it is now possible to perform exact Bayesian inference in the case of wide neural networks with batch normalization. In this work we have presented a theory for neural networks with batch normalization at initialization. In the process of doing so, we have uncovered a number of counterintuitive aspects of batch normalization and -in particular -the fact that at initialization it unavoidably causes gradients to explode with depth. We have introduced several methods to reduce the degree of gradient explosion, enabling the training of significantly deeper networks in the presence of batch normalization. Finally, this work paves the way for future work on more advanced, state-of-the-art, network architectures and topologies. : Batch norm leads to a chaotic input-output map with increasing depth. A linear network with batch norm is shown acting on two minibatches of size 64 after random orthogonal initialization. The datapoints in the minibatch are chosen to form a 2d circle in input space, except for one datapoint that is perturbed separately in each minibatch (leftmost datapoint at input layer 0). Because the network is linear, for a given minibatch it performs an affine transformation on its inputs -a circle in input space remains an ellipse throughout the network. However, due to batch norm the coefficients of that affine transformation change nonlinearly as the datapoints in the minibatch are changed. (a) Each pane shows a scatterplot of activations at a given layer for all datapoints in the minibatch, projected onto the top two PCA directions. PCA directions are computed using the concatenation of the two minibatches. Due to the batch norm nonlinearity, minibatches that are nearly identical in input space grow increasingly dissimilar with depth. Intuitively, this chaotic input-output map can be understood as the source of exploding gradients when batch norm is applied to very deep networks, since very small changes in an input correspond to very large movements in network outputs. (b) The correlation between the two minibatches, as a function of layer, for the same network. Despite having a correlation near one at the input layer, the two minibatches rapidly decorrelate with depth. See Appendix H for a theoretical treatment.A VGG19 WITH BATCHNORM ON CIFAR100Even though at initialization time batchnorm causes gradient explosion, after the first few epochs, the relative gradient norms ∇ θ L / θ for weight parameters θ = W or BN scale parameter θ = γ, equilibrate to about the same magnitude. See Fig. 7 . with batchnorm on CIFAR100 with data augmentation. We use 8 random seeds for each combination, and assign to each combination the median training/validation accuracy over all runs. We then aggregate these scores here. In the first row we look at training accuracy with different learning rate vs β initialization at different epochs of training, presenting the max over . In the second row we do the same for validation accuracy. In the third row, we look at the matrix of training accuracy for learning rate vs , taking max over β. In the fourth row, we do the same for validation accuracy.
Batch normalization causes exploding gradients in vanilla feedforward networks. Develops a mean field theory for batch normalization (BN) in fully-connected networks with randomly initialized weights. Provides a dynamic perspective on deep neural network using the evolution of the covariance matrix along with the layers.