Mean Field Residual Networks: On the Edge of Chaos

By Greg Yang, Sam Schoenholz. 2017

219 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Materials Chemistry, Statistical and Nonlinear Physics, Model Architectures

How researchers see this paper

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

In Artificial Intelligence

In Materials Chemistry

In Statistical and Nonlinear Physics

In Model Architectures

Citations in the corpus

Wide neural networks of any depth evolve as linear models under gradient descent (2019, background): “[35] Ge Yang and Samuel Schoenholz.”
Initialization and Regularization of Factorized Neural Layers (2021, background): “Initialization is a major focus of deep learning research (He et al., 2015; Mishkin & Matas, 2016; Yang & Schoenholz, 2017).”
ReZero is All You Need: Fast Convergence at Large Depth (2020, methodology): “We speculate that this is because at initialization the variance of the signal is not independent of depth, see [20].”
Fixup Initialization: Residual Learning Without Normalization (2019, background): “Yang & Schoenholz (2017) study the effect of the initialization of residual nets to the test performance and pointed out Xavier or He initialization scheme is not optimal.”
Scaling Limits of Wide Neural Networks with Weight Sharing: Gaussian Process Behavior, Gradient Independence, and Neural Tangent Kernel Derivation (2019, background): “We hope Cor 2.3 will allow future works to optimize initialization hyperparameters and prevent gradient explosion/vanishing problems for modern architectures in the way of Schoenholz et al. (2017); Yang & Schoenholz (2017; 2018); Chen et al. (2018); Xiao et al. (2018); Yang et al. (2018).”
Neural Tangents: Fast and Easy Infinite Neural Networks in Python (2019, background): “…studied in the context of mean ﬁeld theory and initialization (Cho & Saul, 2009; Daniely et al., 2016; Poole et al., 2016; Schoenholz et al., 2016; Yang & Schoenholz, 2017; Xiao et al., 2018; Li & Nguyen, 2019; Pretorius et al., 2018; Hayou et al., 2018; Karakida et al., 2018; Blumenfeld et al.,…”
On the Convergence Rate of Training Recurrent Neural Networks (2018, methodology): “There is huge literature on using the mean-ﬁeld theory to study neural networks [10, 14, 25, 30, 34, 38–40, 42, 47, 56–58].”
How to Start Training: The Effect of Initialization and Architecture (2018, methodology): “The articles [32, 33], in contrast, use mean ﬁeld theory to derive the dependence of both forward and backward dynamics in a randomly initialized ResNet on the residual module weights.”
Bayesian Deep Convolutional Networks with Many Channels are Gaussian Processes (2018, background): “The line of work examining signal propagation in random deep networks [33, 37, 47, 16, 9] is related to the construction of the GPs we consider.”
Dynamical Isometry and a Mean Field Theory of RNNs: Gating Enables Signal Propagation in Recurrent Neural Networks (2018, background): “…theory to explore questions of trainability and expressivity in fully-connected networks and residual networks (Poole et al., 2016; Schoen-holz et al., 2017; Yang & Schoenholz, 2017; Schoenholz et al., 2017; Karakida et al., 2018; Hayou et al., 2018; Hanin & Rolnick, 2018; Yang & Schoenholz, 2018).”
When Vision Transformers Outperform ResNets without Pretraining or Strong Data Augmentations (2021, result): “As shown in Table 1, κ is pretty stable for ResNets, echoing previous results that ResNets enjoy superior trainability regardless of the depth (Yang & Schoenholz, 2017; Li et al., 2018).”
Freeze and Chaos for DNNs: an NTK view of Batch Normalization, Checkerboard and Boundary Effects (2019, background): “Recent theoretical results [29] suggest some relation to the transition from “freeze” (also known as “order”) to “chaos” observed as the depth of the NN goes to infinity [22, 6, 28].”
Characterizing signal propagation to close the performance gap in unnormalized ResNets (2021, background): “These insights from studies of batch-normalized ResNets are also supported by theoretical analyses of unnormalized networks (Taki, 2017; Yang & Schoenholz, 2017; Hanin & Rolnick, 2018).”
Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer (2022, background): “Many previous works proposed different initialization or parametrizations with favorable properties, such as better stability for training deep neural networks [5, 10, 13, 21, 30, 47, 48, 51].”
A Fine-Grained Spectral Perspective on Neural Networks (2019, background): “Closely related to the discussion of CK and NTK is the signal propagation literature, which tries to understand how to prevent pathological behaviors in randomly initialized neural networks when they are deep [48, 52, 71, 70, 23, 24, 14, 72, 44, 25, 47].”
Cerebras-GPT: Open Compute-Optimal Language Models Trained on the Cerebras Wafer-Scale Cluster (2023, methodology): “Various techniques have been developed to control training dynamics and train models stably (Glorot & Bengio, 2010; Yang & Schoenholz, 2017; Schoenholz et al., 2017; Yang & Schoenholz, 2018; Zhang et al., 2019; Bachlechner et al., 2020; Huang et al., 2020; Liu et al., 2020; Li et al., 2022). µP is…”
Mean-field Behaviour of Neural Tangent Kernel for Deep Neural Networks (2019, methodology): “In parallel, information propagation in wide DNNs has been studied in (Hayou et al., 2019; Lee et al., 2018; Schoenholz et al., 2017; Yang and Schoenholz, 2017a).”
Dynamical Isometry and a Mean Field Theory of LSTMs and GRUs (2019, methodology): “Forward signal propagation and the spectral properties of the Jacobian at initialization can be studied using mean field theory and random matrix theory (Poole et al., 2016; Schoenholz et al., 2016; Yang and Schoenholz, 2017; 2018; Xiao et al., 2018; Pennington et al., 2017; 2018; Chen et al., 2018;”