A Mean Field Theory of Batch Normalization

By Greg Yang, Jeffrey Pennington, Vinay Rao, Jascha Narain Sohl-Dickstein, Sam Schoenholz. 2019

192 citations recorded by the source corpus.

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.

Subfields: Control and Systems Engineering, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Control and Systems Engineering

In Optimization & Training

Citations in the corpus

Wide neural networks of any depth evolve as linear models under gradient descent (2019, background): “…the wide network limit of fully-connected models (Poole et al., 2016; Schoenholz et al., 2017) extends naturally to residual networks (Yang & Schoenholz, 2017), CNNs (Xiao et al., 2018), RNNs (Chen et al., 2018), batch normalization (Yang et al., 2019), and general architectures (Yang, 2019).”
Weight Standardization (2019, background): “[45] proposes a mean ﬁeld theory for BN to study the gradient explosion issues.”
Feature Learning in Infinite-Width Neural Networks (2020, background): “(e.g. see Lemma F.18 of (Yang et al., 2019)) We then easily have Theorem N.39.”
High-Performance Large-Scale Image Recognition Without Normalization (2021, background): “This biases the signal towards the skip path, which ensures that the network has well-behaved gradients early in training, enabling efficient optimization (Balduzzi et al., 2017; Hanin & Rolnick, 2018; Yang et al., 2019).”
On Layer Normalization in the Transformer Architecture (2020, methodology): “As the warm-up stage happens in the ﬁrst several iterations, we investigate the optimization behavior at initialization using mean ﬁeld theory (Lee et al., 2017; Xiao et al., 2018; Yang et al., 2019a; Yang, 2019; Lee et al., 2019; Zhang et al., 2019).”
Batch Normalization is a Cause of Adversarial Vulnerability (2019, methodology): “We use the notation of Yang et al. (2019), where α is the index for a neuron, l for the layer, and i for a mini-batch of B samples from the dataset; N l denotes the number of neurons in layer l , W l is the matrix of weights and b l is the vector of biases that parametrize layer l .”
Information Geometry of Orthogonal Initializations and Training (2018, unknown): “In particular, we observe that the advantage conferred by optimizing over manifolds cannot be explained by the improvement of the gradient smoothness as measured by λmax(Ḡ), which argues against the proposed role of Batch Normalization recently put forward in [27, 36].”
Training BatchNorm and Only BatchNorm: On the Expressive Power of Random Features in CNNs (2020, methodology): “…and direction of weights (Kohler et al., 2019), which weight normalization (an alternative to BatchNorm; Salimans & Kingma, 2016) does explicitly; that it causes gradient magnitudes to enter an equilibrium (Yang et al., 2019); and that it leads to a novel form of regularization (Luo et al., 2019).”
Tensor Programs II: Neural Tangent Kernel for Any Architecture (2020, background): “For example, most of the semirigorous computations made in the signal propagation literature [55, 58, 61, 67, 68] can now be justified rigorously.”
A Signal Propagation Perspective for Pruning Neural Networks at Initialization (2019, unknown): “Notice, perfect layerwise dynamical isometry cannot always be ensured in the modern networks that have components such as ReLU nonlinearities [8] and/or batch normalization [15].”
The Break-Even Point on Optimization Trajectories of Deep Neural Networks (2020, result): “This is consistent with prior work showing that networks with batch normalization layers can exhibit gradient explosion in the ﬁrst iteration (Yang et al., 2019).”
Tensor Programs III: Neural Matrix Laws (2020, background): “[54] Greg Yang, Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, and Samuel S.”
Characterizing signal propagation to close the performance gap in unnormalized ResNets (2021, background): “To understand this effect, a number of papers have analyzed signal propagation in normalized ResNets at initialization (Balduzzi et al., 2017; Yang et al., 2019).”
Is normalization indispensable for training deep neural network? (2020, background): “Other studies [29, 22, 4, 38] also show several other advantages provided by normalization in terms of loss landscape, regularization, etc.”