PaperlayerView exact revision on arXiv

Source-linked AI summary

A Mean Field Theory of Batch Normalization

Greg Yang, Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, Samuel S. Schoenholz

arXiv:1902.08129v2cs.NEcond-mat.dis-nncs.LGmath.DS

TL;DR

The paper addresses the lack of a precise theory for signal and gradient propagation in wide batch-normalized networks at initialization. It develops a mean field theory for fully connected networks with batch-dependent normalization and shows that gradients inevitably explode exponentially with depth, while linearization can reduce the effect and one optimization step can equilibrate gradient scales.

  • Problem

    Batch normalization lacked a precise theoretical characterization of signal propagation and gradient backpropagation in wide networks at initialization.

  • Method

    The paper develops a mean field theory for random fully connected networks that handles batch-dependent statistics and batch normalization, using covariance dynamics and transform-based analysis.

  • Results

    Gradients grow exponentially with depth for any nonlinearity at BSB1 fixed points, while more linear activations reduce the growth and gradients equilibrate after one optimization step.

  • Takeaways & Limitations

    Batch normalization itself causes unavoidable initialization-time gradient explosion, but tuning toward the linear regime enables training of significantly deeper networks.

  • Takeaways & Limitations

    The analysis assumes B ≥4 unless stated otherwise and treats small-batch behavior separately; it also sets γ = 1 and β = 0 without loss of generality.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper addresses the unresolved theoretical status of batch normalization by developing and validating a mean field theory for fully connected networks on data batches. It identifies gradient explosion and connects the resulting theory to hyperparameter selection and exact Bayesian inference.

  • 1 INTRODUCTION: Batch normalization had remained elusive in frameworks for analyzing random neural networks at initialization.The paper positions batch normalization as a missing architectural component in a broader effort to understand signal and gradient distributions before training.
  • 1 INTRODUCTION: The theory recasts random fully connected networks to handle batch-dependent statistics, incorporates batch normalization, and is validated against Monte Carlo simulations.It also uses the theory to predict valid hyperparameter configurations.
  • 1 INTRODUCTION: Gradients explode exponentially with depth for any nonlinearity when the network converges to a BSB1 fixed point.The paper further reports that this instability limits trainable depth, although more linear activations can reduce the explosion without eliminating it.
  • 1 INTRODUCTION: After one gradient-descent step, batch-normalized networks naturally equilibrate gradient scales, allowing adequate training in shallower networks.This equilibration occurs provided the initial gradients are not large enough to cause numerical instabilities.
  • 1 INTRODUCTION: The wide-network prior induced by batch normalization can be computed exactly, making exact Bayesian inference possible for wide batch-normalized networks.This follows from work connecting random-network function priors with exact Bayesian inference.

2 RELATED WORK

Related work has studied batch normalization through competing normalization schemes, empirical gradient behavior, and initialization-time analyses. This paper distinguishes itself by providing a precise theoretical characterization of gradient statistics.

  • 2 RELATED WORK: Performance comparisons among batch, layer, and weight normalization have been challenging and inconclusive.The literature also proposed internal covariate shift as an explanation for batch normalization’s effectiveness.
  • 2 RELATED WORK: Initialization-time analysis of batch-normalized networks yields insights that experimentally predict both training-time and test-time behavior.The analysis concerns networks of infinite width.
  • 2 RELATED WORK: Prior empirical work observed gradient explosion in deep batch-normalized networks across various nonlinearities.This paper contrasts that observation with a precise theoretical characterization of gradient statistics and stronger conclusions.

3 THEORY

The paper develops covariance-based mean field dynamics for wide batch-normalized networks, then uses them to characterize signal propagation, symmetry breaking, and gradient behavior. The theory shows that batch normalization produces exponential gradient explosion and rapid cross-batch information loss, while more linear activations reduce—but do not remove—the instability.

  • 3.1 BATCH NORMALIZATION: Wide batch-normalized networks are analyzed through analytic recurrence relations over batch covariance matrices.The framework reformulates fully-connected mean field theory for batches and extends it to explicit batch normalization.
  • 3.1 BATCH NORMALIZATION: For most practical nonlinearities, deep batch normalization induces a highly symmetric BSB1 embedding fixed point.When activations grow quickly, a winner-take-all transition can instead produce BSB2 block structure, as observed for α-ReLU.
  • 3.1.2 GRADIENT BACKPROPAGATION: Gradients explode exponentially for any nonlinearity whose covariance dynamics converge to a BSB1 fixed point.The result applies to fully connected batch-normalized networks and extends to weight gradients at the same rate.
  • 3.1.2 GRADIENT BACKPROPAGATION: Batch normalization forces layer-Jacobian singular values away from 1, so weight variances and activation choices cannot eliminate gradient explosion.Larger epsilon values can ameliorate, but not eliminate, the explosion; activations closer to linear reduce its rate.
  • 3.1.2 GRADIENT BACKPROPAGATION: ReLU batch-normalized gradient norms explode exponentially, whereas the linear case has the smaller rate B^-2.The ReLU rate decreases toward approximately 1.467 as batch size grows, but remains greater than one.
  • 3.2 CROSS-BATCH DYNAMICS: Cross-batch correlation decays exponentially with depth, causing deep batch-normalized networks to lose input correlation information and become increasingly chaotic.The cross-batch covariance decay rate equals the cross-batch gradient-covariance decay rate, while the identity activation maximizes that rate for fixed batch size.

4 EXPERIMENTS

Experiments link batch-normalized network trainability to initialization-time gradient explosion and test theoretical predictions across depths, optimizers, datasets, and stabilization strategies.

  • ξ = 1/log λ↓G closely predicts the maximum trainable depth across experiments.The theory defines ξ from the largest Jacobian eigenvalue and reports strong agreement with observed trainability limits.
  • Networks deeper than about 50 layers are untrainable on MNIST regardless of batch size.
  • Nearly identical trainable regions appear across modified batch statistics, RMSProp, and CIFAR10 experiments.The modified normalization computes statistics over subsets of size B while keeping the overall batch size fixed.
  • During the first 10 training steps, initially exploding gradients drive exponential weight-norm growth before gradients scale down commensurately.This produces a stable dynamical equilibrium, although deeper networks can experience severe relative gradient vanishing in lower layers.
  • Decreasing γ for tanh or increasing β for rectified linear activations increases maximum trainable depth by moving toward the linear regime.
  • Large ϵ values produce a critical point where gradients do not explode and very deep networks become trainable.The experiments vary ϵ for both tanh and rectified linear networks, with theoretical predictions overlaid.

5 CONCLUSION

The paper presents a mean field theory of batch normalization at initialization, identifies unavoidable depth-dependent gradient explosion, and evaluates methods that reduce it.

  • The theory characterizes fully connected networks with batch normalization at initialization.
  • Batch normalization unavoidably causes gradients to explode with network depth at initialization.
  • Several methods reduce gradient explosion and enable training of significantly deeper batch-normalized networks.

A VGG19 WITH BATCHNORM ON CIFAR100

VGG19 experiments show that batch-normalized parameter gradients initially differ but rapidly equilibrate, while positive β initialization or ϵ can accelerate training.

  • After 10 epochs, relative gradient norms for γ and W become approximately equal across all layers despite initial gradient explosion.
  • Figure 7 orders parameters from input to output with interleaved γ and W curves separated by 3, 5, or 10 epochs.
  • Setting ϵ > 0 or initializing β > 0 produces acceleration effects, especially during initial training.
  • Figure 8 aggregates median training and validation accuracy across 8 random seeds while sweeping learning rate, β initialization, and ϵ.Rows compare learning rate with β initialization or ϵ, using maxima over the remaining hyperparameter.

B GRADIENT INDEPENDENCE ASSUMPTION

The analysis assumes independent weight copies during backpropagation and formalizes batch-normalized covariance dynamics using matrix and integral-transform notation.

  • Gradient independence assumption: During backpropagation, each multiplication by W^T uses an iid copy of the corresponding weight matrix.This is Assumption 2, adopted from prior analyses of random neural networks.
  • Gradient independence assumption: The assumption agrees closely with simulations, although a rigorous justification does not currently extend to batch normalization because its Jacobian is singular at 0.The authors expect that an extension of prior proofs may be possible.
  • Notation: V_Φ maps a batch covariance matrix Σ to the expected outer product of Φ(h) for h drawn from N(0, Σ).For scalar φ, the same transform notation applies to coordinatewise activation functions.
  • Notation: Batch normalization followed by coordinatewise φ is denoted B_φ, with B also used for the identity activation.The notation is defined for batches of size B.
  • Notation: G projects batch vectors onto the zero-sum subspace, and S_G^B contains positive semidefinite matrices with zero row and column means.This captures the batch-centering structure induced by batch normalization.
  • Notation: The appendix also fixes conventions for Gamma, Beta, Pochhammer, dot-product, Jacobian, and related matrix notation.These conventions support the subsequent transform and covariance calculations.

D A GUIDE TO THE REST OF THE APPENDIX

The appendix develops transform-based tools for analyzing forward, backward, and cross-batch dynamics, then applies them to fixed points, convergence, and gradient explosion.

  • Dynamics: Forward and backward propagation are represented as dynamical systems over covariance matrices, simplifying analysis relative to tracking individual pre-activations.Fixed-point behavior provides a central lens for studying these dynamics.
  • Analytical techniques: The Laplace method gives simplified closed forms for positive-homogeneous nonlinearities, while the Fourier method handles broader polynomially bounded continuous nonlinearities.The Laplace method uses Schwinger parametrization; the Fourier method incurs an additional complex integral.
  • Analytical techniques: Spherical integration reduces batch-normalization expectations to one- or two-dimensional angular integrals, enabling numerical evaluation for general nonlinearities.This follows from batch normalization’s interpretation as projection followed by projection onto a sphere.
  • Analytical techniques: The Gegenbauer method expresses nonlinearities in an orthogonal polynomial basis, making forward and backward eigenvalues ratios of quadratic forms.Its structure exposes why gradient explosion occurs under general conditions.
  • Analytical techniques: Ultrasymmetric operators are analyzed through structural eigendecomposition results that support asymptotic calculations for convergence and gradient explosion.These operators arise because the relevant matrix dynamics have many symmetries.
  • Results and scope: The backward analysis concerns hidden preactivation gradients, while Appendix L computes the corresponding moments of weight gradients used in gradient descent.The analysis also studies simultaneous propagation across multiple batches.
  • Results and scope: For φ = id, global convergence to a BSB1 fixed point is established under nondegenerate initial covariance, whereas general nonlinearities receive only local convergence analysis.The main appendix then derives BSB1 fixed points, convergence rates, gradient-explosion rates, and cross-batch dynamics.
  • Laplace method details: The Laplace master equation defines a continuous Gaussian expectation transform when rank Σ > 2k and rewrites it as an integral over normalized covariances.The proof uses integrability near zero, continuity, and Fubini–Tonelli interchange.

E.2 FOURIER METHOD

The Fourier method extends the analysis beyond positive-homogeneous activations, while rigorous justification requires mollifiers and integrability conditions because formal integration steps can fail.

  • Fourier method: For general φ, the Fourier method replaces the positive-homogeneity step that enables the Laplace method’s norm-factor extraction.The Laplace method can use Schwinger parametrization for positive-homogeneous φ, but this simplification is unavailable for general φ.
  • Fourier method: The Fourier representation can be reduced to Gaussian expectations with complex covariance matrices when the relevant covariance is nonsingular and the integral exists.This extension supports the transformed expressions, although the formal derivation requires care around complex Gaussian integration.
  • Fourier method: The derivation uses mollifiers to rigorously justify Fourier expansions involving the delta function and subsequent limiting arguments.The authors state that the formal derivation is not mathematically correct as written, but its result can be justified by expressing the delta function as a mollifier limit.
  • Fourier method: Integrability assumptions on φ and related expressions ensure that the relevant integrals exist and that dominated-convergence arguments apply.The assumptions include local integrability in the radial variable and additional boundedness or polynomial-growth conditions; common nonlinearities such as ReLU, identity, and tanh satisfy them.
  • Fourier method: The paper also uses spherical coordinates and Gegenbauer expansions to analyze batch-normalized covariance and gradient dynamics.The Gegenbauer identity diagonalizes a derivative-dependent quadratic form and is used to establish the necessity of gradient explosion under the BSB1 fixed-point assumption.
  • Fourier method: Under the analyzed dynamics, covariance convergence is exponential, while batch normalization produces exponential gradient-norm growth and linear activations minimize the information loss.The fixed-point analysis gives exponential eigenvalue-gap convergence, and the gradient norm for a linear batch-normalized network explodes at rate B−2.

F.2.2 LAPLACE METHOD

The Laplace method gives a closed-form description of the BSB1 fixed point for positive-homogeneous nonlinearities, while the analysis reduces batch-normalization operators to structured eigenspaces and lower-dimensional integrals.

  • Fixed point: Positive-homogeneous nonlinearities yield a closed-form BSB1 fixed point expressed through the J function.The fixed point exists and is unique under the stated conditions.
  • Fixed point: The BSB1 transform preserves the BSB1 structure, with equal diagonal entries and equal off-diagonal entries.For VBid, the off-diagonal entries are −1/B and the projected result is B/(B −1)G.
  • Operator analysis: The Jacobian analysis uses structured operators whose eigenvalues can be recovered through related operator products.The relevant eigenspaces are studied after projection by G⊗2.
  • Operator analysis: Batchnorm is decomposed into mean-centering, division by standard deviation, and coordinate-wise nonlinearity.This decomposition defines the operators analyzed throughout the Laplace method.

F.3.2 GEGENBAUER EXPANSION

The Gegenbauer expansion expresses local convergence and eigenspace eigenvalues as functions of the activation’s Gegenbauer coefficients. It also clarifies how coefficient structure affects convergence rates.

  • Eigenvalue structure: The local convergence eigenvalue is represented as a quadratic form in the Gegenbauer coefficients of the nonlinearity.The analysis compares the eigenspaces associated with different perturbation types.
  • Coefficient invariance: Batchnorm is invariant to additive shifts of the activation because the zeroth Gegenbauer coefficient does not affect the eigenvalue.The coefficients vB−1,0, wB−1,0, and uB−1,0 are all zero.
  • Convergence behavior: Higher-degree Gegenbauer components grow more strongly near ±1, consistent with the claim that more explosive nonlinearities are less likely to converge to a BSB1 fixed point.The text connects endpoint growth with the convergence behavior of the fixed-point dynamics.
  • Coefficient comparison: For B ≥10, numerical plots suggest wB−1,l−vB−1,l ≥ 0 for all l ≥1.This comparison is shown in Figure 10.
  • Eigenspace comparison: The eigenvalue is typically smaller for the M eigenspace than for the L eigenspace, although counterexamples such as sin exist.The stated comparison applies for degree larger than 0.

F.3.3 LAPLACE METHOD

The Laplace analysis characterizes forward and backward batch-normalization dynamics through eigendecompositions at the BSB1 fixed point, identifying structured eigenspaces and the linear activation as optimal for minimizing one backward eigenvalue.

  • Local convergence: The Laplace method differentiates the fixed-point equations along smooth covariance paths to obtain local convergence rates.The calculation uses determinant derivatives and Gaussian integral identities.
  • Forward eigendecomposition: For positive-homogeneous activations, the relevant operators admit eigendecompositions over the structured eigenspaces RG, L, and M.Several theorems give the corresponding eigenvalues for the projected Jacobian and related operators.
  • Backward dynamics: The backward dynamics are analyzed by studying the adjoint operator at the BSB1 fixed point, whose eigendecomposition governs large-depth behavior.After one backpropagation step, the dynamics are represented by the corresponding adjoint operator.
  • Backward dynamics: The backward operator has structured eigenspaces whose eigenvalues can be computed using Gaussian and spherical integral reductions.The analysis reduces selected expectations to two-dimensional numerical integrals and applies the eigenspace decomposition.
  • Optimization over activations: The eigenvalue associated with RG is minimized over nonlinearities when the activation is linear.This identifies linearity as the minimizing choice for that backward eigenvalue.

G.2 LAPLACE METHOD

The Laplace analysis extends from single-batch dynamics to derivative operators and multiple batches, using positive homogeneity to obtain structured eigendecompositions and cross-batch fixed-point forms.

  • Derivative operator: For positive-homogeneous activations, the derivative operator can be decomposed into the same structured eigenspaces RG, M, and L.Its eigenvalues are expressed through quantities built from the activation and its derivative.
  • Homogeneity reduction: Positive homogeneity gives ψ(y)=yφ′(y)=αφ(y), simplifying the derivative-operator calculations.This identity is used to relate derivative expectations back to Vφ.
  • Derivative operator: The eigenspaces RG, L, and M receive explicit eigenvalue formulas for the analyzed derivative operator.The formulas are stated separately for each eigenspace.
  • Multiple batches: The multi-batch dynamics generalize the single-batch covariance recursion by applying batchnorm separately to each batch block.Restricting the extended dynamics to a diagonal block recovers the original single-batch dynamics.
  • Multiple batches: CBSB1 matrices provide a structured form for studying cross-batch fixed points, with common BSB1 diagonal blocks and constant off-diagonal blocks.The fixed-point analysis first recovers the single-batch fixed point on diagonal blocks and then analyzes cross-batch blocks.

H.1.2 LAPLACE METHOD

The Laplace-method appendix develops integral identities and operator decompositions used to analyze batch-normalized covariance and gradient dynamics. It then characterizes fixed-point linearization through blockwise operator structure and associated eigenspaces.

  • Laplace-method identities: The Laplace cross-batch master equation establishes continuity and an integral representation for expectations involving Gaussian covariance matrices.The construction applies when the covariance rank exceeds 2(a + b).
  • Operator decomposition: Blockwise diagonal-off-diagonal semidirect operators provide the structural form used to analyze linearized covariance dynamics around CBSB1 fixed points.The appendix defines BDOS operators and derives their eigenspaces and eigenvectors under nonsingularity conditions.
  • Operator decomposition: The eigendecomposition separates block-diagonal and off-diagonal components, assigning distinct eigenvalues to invariant matrix subspaces.Permutation-symmetric variants preserve the corresponding eigenvalue structure.
  • Positive-homogeneous activations: For positive-homogeneous nonlinearities, the new cross-batch eigenvalue can be expressed directly using the J function of the activation.The result is stated for degree-α positive-homogeneous φ.

I CROSS BATCH: BACKWARD DYNAMICS

This section analyzes backward cross-batch dynamics by decomposing the relevant operator into invariant blocks and their eigenvalues. The resulting eigenvalue controls how cross-batch gradient correlations evolve.

  • Operator structure: The backward cross-batch operator acts independently on diagonal and off-diagonal batch blocks, allowing the dynamics to be analyzed block by block.The diagonal blocks follow the previously studied single-batch dynamics, while off-diagonal blocks receive a separate scalar evolution.
  • Cross-batch eigenspace: The cross-batch backward operator has a distinguished eigenspace on which its action reduces to multiplication by a scalar.The relevant eigenspace is M_G under the projected operator.
  • Eigenvalue characterization: For positive-homogeneous activations, the cross-batch eigenvalue admits an explicit characterization through the activation’s homogeneity parameters.The theorem specializes the operator result to functions represented using positive and negative ReLU components.
  • Multi-batch dynamics: The backward dynamics over multiple batches is governed by the eigendecomposition of the corresponding block operator.Each off-diagonal block evolves by the cross-batch eigenvalue, while diagonal blocks retain the single-batch eigenvalues.
  • Interpretation: The linear activation maximizes the cross-batch backward eigenvalue for fixed batch size, attaining λ↓eM = 1 and approaching it from below.The paper concludes that linearity maximally ameliorates exponential information loss across batches.

J UNCOMMON REGIMES

The paper examines batch sizes and nonlinearities outside the main BSB1 regime. Small batches can produce either immediate gradient vanishing or explosion, while BSB2 behavior remains analytically unresolved.

  • Small Batch Size: B = 2 causes deterministic zero gradients, whereas B ≥ 4 exhibits gradient explosion; B = 3 is empirically observed to behave similarly to B = 4.The authors conjecture that the main gradient-explosion theorem extends to B = 3.
  • Batch Symmetry Breaking: For nonlinearities whose fixed point is not BSB1, the dynamics can converge to BSB2 fixed points with a dominant diagonal block.Finite-width weight randomness can break the permutation symmetry needed for selecting among such fixed points.
  • Batch Symmetry Breaking: Small input fluctuations can direct dynamics toward one BSB1 fixed point over others, producing large output changes and intuitively large gradients.The dominant BSB2 block is expected to undergo dynamics resembling a smaller-batch BSB1 fixed point.
  • Finite width effect: Width can change which fixed-point structure is favored: for B = 10, width 100 favors BSB2 while widths 1000 and above favor BSB1.This comparison is reported for the fixed-point selection criterion in Eq. (43).
  • BSB2 analysis: The BSB2 forward dynamics reduces to a two-dimensional system over eigenspace eigenvalues, but the Gegenbauer argument for unavoidable gradient explosion does not directly generalize.The paper states that it has no rigorous result for the general BSB2 case.

K.1 \ BSB2

This section analyzes rank-1 BSB2 fixed points and their backward behavior. Depending on the gradient direction, the resulting dynamics can exhibit severe explosion or pathological vanishing, with numerical effects complicating observations.

  • Rank-1 fixed points: Rapidly growing activations can produce BSB2 fixed points with a rank-1 projected covariance structure.The section extends the analysis to the B′ = 1 case and identifies the projected fixed point as λq̂⊗2.
  • Backward dynamics: At a rank-1 BSB2 fixed point, one backpropagation step depends on the gradient component outside the constant direction.The projected covariance determines the normalized Gaussian variables entering the derivative calculation.
  • Gradient behavior: A zero second gradient component yields deterministic zero output, while a nonzero component produces a variable with no finite mean because of the |r|^-1 factor.The paper attributes this divergence specifically to the rank-1 fixed point.
  • Gradient behavior: Generic gradients are therefore expected to show severe explosion, whereas constant gradients are expected to show severe vanishing.The associated numerical behavior can differ from the mathematical prediction.
  • Numerical caveat: Finite-precision computation can make gradients vanish rapidly because projection fails to eliminate one gradient component exactly.This numerical issue is reported for rapidly growing activations such as relu(x)^30.
Loading 1902.08129v2…