Source-linked AI summary
Mean Field Residual Networks: On the Edge of Chaos
Greg Yang, Samuel S. Schoenholz
TL;DR
The paper asks whether residual skip connections avoid the exponential geometry collapse and gradient instability seen in vanilla networks. It extends mean-field theory and difference-equation analysis to residual networks, showing mostly subexponential or polynomial dynamics and initialization-dependent predictions of trained performance. It also finds that optimal initialization variances depend on depth and derives new α-ReLU kernel identities.
Problem
Vanilla networks exhibit exponential forward and backward dynamics, rapidly collapsing input geometry and causing vanishing or exploding gradients.
Method
The paper extends mean-field theory and difference-equation analysis to randomly initialized residual networks with different nonlinearities and initialization settings.
Results
Residual networks generally have subexponential, often polynomial, forward and backward dynamics; initialization theory accurately predicts relative trained performance across settings.
Takeaways & Limitations
Optimal residual-network initialization must account for depth, and common Xavier or He schemes are not generally optimal.
Takeaways & Limitations
Across the experiments, cosine-distance dynamics did not predict neural-network performance.
Abstract
from arXiv · showhide
We study randomly initialized residual networks using mean field theory and the theory of difference equations. Classical feedforward neural networks, such as those with tanh activations, exhibit exponential behavior on the average when propagating inputs forward or gradients backward. The exponential forward dynamics causes rapid collapsing of the input space geometry, while the exponential backward dynamics causes drastic vanishing or exploding gradients. We show, in contrast, that by adding skip connections, the network will, depending on the nonlinearity, adopt subexponential forward and backward dynamics, and in many cases in fact polynomial. The exponents of these polynomials are obtained through analytic methods and proved and verified empirically to be correct. In terms of the "edge of chaos" hypothesis, these subexponential and polynomial laws allow residual networks to "hover over the boundary between stability and chaos," thus preserving the geometry of the input space and the gradient information flow. In our experiments, for each activation function we study here, we initialize residual networks with different hyperparameters and train them on MNIST. Remarkably, our initialization time theory can accurately predict test time performance of these networks, by tracking either the expected amount of gradient explosion or the expected squared distance between the images of two input vectors. Importantly, we show, theoretically as well as empirically, that common initializations such as the Xavier or the He schemes are not optimal for residual networks, because the optimal initialization variances depend on the depth. Finally, we have made mathematical contributions by deriving several new identities for the kernels of powers of ReLU functions by relating them to the zeroth Bessel function of the second kind.
1 Introduction
Prior mean-field work found exponential forward and backward dynamics in randomly initialized vanilla networks, rapidly distorting input geometry and causing vanishing or exploding gradients. This paper studies whether residual skip connections produce different dynamics.
- Vanilla networks can range from stable to chaotic depending on initialization variance, with input-vector cosine distance converging exponentially to a fixed point.A fixed point of 1 corresponds to stability, whereas 0 corresponds to chaos.
- Residual networks add identity passthrough connections so each nonlinearity computes a residual relative to the previous activation.
- The paper characterizes randomly initialized residual networks and analyzes their forward and backward behavior across nonlinearities.
- The work aims to explain residual-network behavior using theoretical analysis, with formal statements and proofs placed in the appendix.
2 Background
Background work established exponential convergence of activation statistics, input geometry, and gradients in vanilla feedforward networks. The paper contrasts these dynamics with subexponential or polynomial behavior in residual networks.
- Vanilla feedforward networks propagate layer activations through affine transformations followed by coordinatewise nonlinearities such as tanh or ReLU.
- Prior analyses found normalized activation norms and dot products converge exponentially to fixed points, so cosine distance also converges exponentially.
- Exponential cosine-distance convergence means vanilla networks quickly forget the geometry of the input space.
- Vanilla networks can exhibit exponentially vanishing or exploding gradients, with the gradient half-life controlled by initialization variances.
- Residual networks instead show subexponential, and often polynomial, cosine-distance and gradient growth rates.
3 Preliminaries
The paper defines reduced and full residual architectures, assumes Gaussian initialization and mean-field independence conditions, and introduces forward, correlation, expressivity, and gradient quantities for large-width analysis.
- A reduced residual network uses a shortcut connection, while a full residual network additionally uses affine weights and biases from the nonlinearity to the next layer.
- Gradient quantities χ measure the expected squared gradient size at layers under the large-width mean-field setup.
- Weights and biases are sampled independently from Gaussian distributions with layer-independent standard deviations σw, σb, σv, and σa.
- The analysis adopts symmetry, gradient-independence, and related physical assumptions to make mean-field calculations tractable.These assumptions include independence between backpropagated gradients and previous-layer activations.
- The central forward quantities include length statistics p and q, correlation γ, cosine distance e, and metric expressivity s.Under equal initial input norms, e(l)=γ(l)/p(l) and s(l)=(1−e(l))p(l).
- The paper defines asymptotic notation, including a softened polynomial-equivalence notation ˇΘ for growth within arbitrary powers x^ε.
4 Overview
Using central-limit-theorem-based mean-field recurrences, the paper analyzes forward geometry and backward gradients in residual networks. It finds polynomial forward behavior and broadly subexponential backward behavior, while initialization theory predicts trained performance patterns.
- The central limit theorem makes affine-layer outputs approximately Gaussian, enabling mean-field recurrences for p, q, γ, λ, and χ.
- For tanh and α-ReLU, cosine distance converges polynomially rather than exponentially, keeping finite-depth networks away from the stable or chaotic endpoints.This slower convergence preserves geometric information and lets residual networks hover near the edge of chaos.
- For tanh and α-ReLU with α<1, metric expressivity s varies polynomially with depth.
- Residual networks generally avoid vanishing gradients and exhibit subexponential gradient dynamics, with polynomial behavior for α-ReLU when α<1.The full residual network with ReLU is the stated exception to subexponential gradient dynamics.
- Initialization theory predicts relative test performance by tracking gradient explosion for tanh residual networks and expressivity for ReLU-family residual networks.
- The α-ReLU analysis derives new kernel identities relating powers of ReLU kernels to zeroth Bessel functions.
5 Theoretical Results
The paper derives mean-field recurrences and asymptotic dynamics for residual networks with tanh and α-ReLU activations. Skip connections generally produce subexponential or polynomial behavior, with initialization and activation controlling expressivity, gradient growth, and numerical stability.
- Method: Theoretical analysis uses central-limit-based recurrences and solves their high-depth asymptotics for tanh and α-ReLU residual networks.The recurrences track forward quantities, cosine distance, and gradient dynamics.
- Validation: The recurrences predict empirical dynamics very well for tanh and α-ReLU residual networks.Figure 1 compares theory and measurements for tanh/FRN, with similar agreement reported for tanh/RRN and α-ReLU/FRN.
- 5.1 Tanh: For tanh residual networks, forward quantities grow polynomially, while cosine distance converges to a stable fixed point at a polynomial rate.For FRN, the fixed point satisfies e∗ < 1, and metric expressivity scales as Θ(l).
- 5.1 Tanh: For tanh/FRN, the convergence exponent δ∗ depends on ρ = σa/σv and increases from approximately 0.36338 toward 1/2 as ρ increases.The limiting regimes correspond to e∗ approaching 0 for small ρ and 1 for large ρ.
- 5.1 Tanh: Tanh residual networks avoid the exponential gradient behavior of vanilla tanh networks, although gradient explosion still trades off against expressivity through ρ.Large ρ improves trainability by reducing the explosion constant A but makes the network less expressive; small ρ has the opposite tradeoff.
- 5.2 α-ReLU: Standard ReLU residual networks have exponential forward growth and can overflow after about 100 layers, whereas α-ReLU with α < 1 has slower polynomial forward behavior.For standard ReLU, metric expressivity grows exponentially; α-ReLU convergence and exponents depend on α.
- 5.2 α-ReLU: For α-ReLU, the exponent governing one analyzed asymptotic quantity is minimized at α = 3/4 on [3/4, 1], where R = 9/2.On (1/2, 1), the minimum occurs at α = 2/3, where R = 4.
- 5.2 α-ReLU: Experiments show that α-ReLU performance is governed by metric expressivity, with the largest α avoiding numerical issues performing best.Smaller α avoids numerical problems from exponentiating forward and backward dynamics.
6 Experimental Results
Experiments on MNIST show that initialization quality depends on activation and architecture: gradient explosion predicts tanh-resnet performance, whereas expressivity predicts ReLU-family performance.
- Initialization predictors: Initialization quality is governed by gradient explosion for tanh resnets but by random-network expressivity for ReLU resnets.The paper measures these effects using χ(0)/χ(L) and s(L), respectively.
- Tanh resnets: For tanh resnets, deeper networks favor smaller σw, and gradient-ratio level curves identify the largest σw associated with optimal test accuracy.The relationship follows log R ≈ σw^2L.
- Tanh resnets: Tanh-resnet performance is maximized around a fixed depth across ρ, supporting gradient dynamics as the main initialization-quality determinant.Increasing ρ produces a minor unexplained performance increase despite decreasing expressivity.
- ReLU resnets: For ReLU resnets, the best test accuracies occur at depths just below numerical failure from exploding activation and gradient norms.Optimal performance tracks the highest s(L), p(L), and χ(0)/χ(L) before infinities occur.
- α-ReLU resnets: For α-ReLU resnets, χ(0)/χ(L) does not predict accuracy, while s(L) contours predict performance remarkably well.This supports expressivity rather than trainability as the relevant test-time predictor in this setting.
- Additional observation: Across the experiments, e dynamics did not predict neural-network performance.
7 Conclusion
The paper extends mean field theory to residual networks and finds that most studied models avoid exponential collapse of input geometry and gradient information. It also concludes that initialization should depend on depth, while several architectural and theoretical extensions remain open.
- Conclusion: The extended mean field analysis shows that most studied residual networks avoid exponential collapse of input geometry and gradient information in forward and backward passes.
- Conclusion: Theory and experiments indicate that optimal residual-network initialization must account for depth.The paper attributes Xavier and He suboptimality to structural assumptions that break on residual networks, not to fragile statistical assumptions.
- Open Problems: Mean field theory remains incomplete for combinations involving residual connections, batch normalization, convolutional layers, recurrent architectures, and attention mechanisms.The authors also identify the accuracy of their axioms as an unresolved theoretical question.
Notes
These notes record technical qualifications concerning the residual-network analysis, tempered α-ReLU, numerical visualization, and regimes where theoretical gradient predictions lose empirical agreement.
- Related work: The fixed-point existence result cited from Daniely et al. applies to feedforward networks but not directly to residual architectures.
- Implementation: Tempered α-ReLU replaces the divergent behavior near x=0 with a small-ϵ regularized form, and the paper expects similar conclusions.
- Background: The fixed-ρ=1 version of Wφ is called the dual function of φ in Daniely et al.’s terminology.
- Visualization: Floating-point precision motivates plotting e(l)−e* against l−δ* indirectly rather than relying on convergence to exact zero.
- Assumptions: The derivations cover α∈(1/2,1], but empirical agreement is not expected when α≤3/4 because gradient variance diverges.
- Assumptions: The theoretical predictions may nevertheless act as an upper bound on gradient norms in the divergent-variance regime.
- Visualization: The p(l) contour has similar shape to the heatmap contours, but its slopes are slightly different.
A Additional Figures
Additional figures compare theoretical and empirical forward, backward, and gradient dynamics across residual-network nonlinearities. The results generally track empirical averages, while ReLU and low-α regimes exhibit numerical or variance-related limitations.
- Empirical verification: Theoretical predictions closely track empirical averages for α-ReLU forward quantities p(l), γ(l), and e(l) across α.For gradient norms, accuracy holds from α = 1 to α = .7, but variance becomes problematic at lower α.
- Numerical behavior: At α = 1, typical ReLU residual-network activation norms become NaN around layer 90, motivating shorter plotted depth ranges.For α < 1, the corresponding numerical issue is not observed in these experiments.
- Gradient dynamics: For α = .6 to α = .5, theory consistently and significantly overestimates average gradient norms because variance is extremely large.At these α values, one standard deviation below the mean can produce negative plotted values.
- Definitions and recurrences: The appendix reports empirical verification of recurrence-based dynamics for reduced residual networks and defines notation through a symbol glossary.The formal results rely on recurrence equations for network quantities and gradients.
B.1.2 Full Residual Network
For tanh-like nonlinearities in full residual networks, forward quantities grow asymptotically linearly or with logarithmic corrections, while cosine-distance and gradient behavior follow distinct regimes.
- Cosine-distance dynamics: For σw > 0, cosine distance converges to a unique fixed point e* ≠ 1 determined by the stated fixed-point equation.The σw = 0 regime instead has e(l) → 1 with 1 − e(l) = Θ(l^-1).
- Gradient dynamics: The appendix empirically verifies the asymptotic expansion of χ for several small values of σ•s.For larger values, the constant term can dominate and χ may behave like exp(Θ(l)) up to depth 1000.
B.2.1 Full Residual Network
For α-ReLU full residual networks, α = 1 produces exponential dynamics, whereas 0 < α < 1 yields polynomial forward, backward, and gradient behavior with α-dependent exponents.
- Forward dynamics: For α = 1, p(l) has exponential asymptotics Θ((1 + σw^2/2)^l), verified empirically for multiple σv and σw pairs.The corresponding plots use a logarithmic y-axis.
- Forward dynamics: For 0 < α < 1, p(l) grows polynomially as Θ(l^(1/(1−α))), while e(l) converges to a unique nonunit fixed point e*.The convergence rate is |e* − e(l)| = Θ̃(l^-μ), with μ = (1 − J̇α(e*))/(1 − α).
- Open issue: The conjecture μ = 1/(2(1−α)) is unsupported by proof and would imply μ → 0 as α → 1, unlike the α = 1 exponent −2.The text explicitly identifies this as a discontinuity conjecture rather than an established theorem.
- Validity conditions: The variance of ψ̇α(ζ)^2 diverges for α ≤ 3/4 and is finite for α > 3/4, limiting applicability of corresponding theoretical equations.Empirical verification of backward asymptotics is reported for α < 1.
- Backward dynamics: For α ∈ (3/4, 1), gradient dynamics are polynomial, with exponent R = α^2/((1−α)(2α−1)).The displayed recurrence gives χ(l−m) = Θ(1)χ(l)(l/(l−m))^R.
C.5.1 Forward Dynamics
The forward-dynamics analysis derives asymptotic growth and convergence properties for residual-network quantities using recurrence analysis, integral transformations, and smooth expansions of tanh kernels.
- Asymptotic growth: With σw > 0, p(l) and q(l) are asymptotically linear in depth, with tanh expansions including square-root and logarithmic correction terms.The analysis obtains terms such as b0l + b1l^1/2 + b2 log l + O(1).
- Asymptotic expansion: The forward analysis repeatedly substitutes asymptotic forms into V^tanh(q) and sums the resulting recurrence to derive successive expansion terms.This procedure uses the expansion V^tanh(q) = 1 − Cq^-1/2 + Θ(q^-3/2).
- Kernel convergence: For fixed c ∈ (0, 1), the tanh kernel satisfies 2π^-1 arcsin(c) − Wφ(q,cq) = Θ(q^-1).The hidden constants depend on c, but can be uniform when c stays in [ε, 1−ε].
C.6.1 Forward Dynamics
For tanh residual networks with e(0) < 1, the forward correlation either approaches 1 with polynomially decaying error or converges to a unique nonunit fixed point, depending on σw.
- σw = 0 yields e(l) → 1 with 1−e(l) = Θ(l−1), while the scaled quantity s(l) remains Θ(1).
- σw > 0 makes e(l) converge to the unique fixed point e∗ ≠ 1 determined by the governing equation.
- The convergence proof uses interval trapping and a difference-equation lemma to establish limiting behavior for every e(0) ∈ [0, 1).
- |e(l)−e∗| decays between Ω(l−δ∗−ε) and O(l−δ∗+ε) for any ε > 0.
C.7.1 Forward Dynamics
For power-ReLU nonlinearities, forward dynamics vary sharply with α: α = 1 produces exponential-scale behavior, whereas 0 < α < 1 yields polynomial growth and convergence to a nonunit fixed point.
- The α < 1 asymptotic expansion is supported by analysis of Jα and empirically verified for its leading coefficient and exponent.
- For α ∈ [1/2, 1), Jα has an unstable fixed point at 1 and a stable fixed point below 1.
- For α = 1, e(l) → 1 and s(l) = exp(Θ(l)), indicating exponential forward growth.
- For 0 < α < 1, e converges to the unique nonunit fixed point e∗ of Jα and s(l) = Θ(p(l)) = Θ(l1/(1−α)).
- The backward-dynamics variance condition changes at α = 3/4: Var(˙ψα(ζ)^2) diverges for α ≤ 3/4 and is finite for α > 3/4.