Dynamical Isometry and a Mean Field Theory of CNNs: How to Train 10, 000-Layer Vanilla Convolutional Neural Networks

By Lechao Xiao, Yasaman Bahri, Jascha Narain Sohl-Dickstein, Sam Schoenholz, Jeffrey Pennington. 2018

392 citations recorded by the source corpus.

In recent years, state-of-the-art methods in computer vision have utilized increasingly deep convolutional neural network architectures (CNNs), with some of the most successful models employing hundreds or even thousands of layers. A variety of pathologies such as vanishing/exploding gradients make training such deep networks challenging. While residual connections and batch normalization do enable training at these depths, it has remained unclear whether such specialized architecture designs are truly necessary to train deep CNNs. In this work, we demonstrate that it is possible to train vanilla CNNs with ten thousand layers or more simply by using an appropriate initialization scheme. We derive this initialization scheme theoretically by developing a mean field theory for signal propagation and by characterizing the conditions for dynamical isometry, the equilibration of singular values of the input-output Jacobian matrix. These conditions require that the convolution operator be an orthogonal transformation in the sense that it is norm-preserving. We present an algorithm for generating such random initial orthogonal convolution kernels and demonstrate empirically that they enable efficient training of extremely deep architectures.

Subfields: Computer Vision and Pattern Recognition, Optimization & Training, Scaling & Learning Theory

How researchers see this paper

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

In Computer Vision and Pattern Recognition

In Optimization & Training

In Scaling & Learning Theory

Citations in the corpus

ReZero is All You Need: Fast Convergence at Large Depth (2020, methodology): “Researchers have relied on careful initialization [7, 8] and normalization techniques such as BatchNorm [9] and LayerNorm [10] to mitigate this issue, but these techniques can be costly and limited.”
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).”
Fixup Initialization: Residual Learning Without Normalization (2019, methodology): “, 2017; Hanin, 2018), enabling training 10,000-layer CNNs from scratch (Xiao et al., 2018).”
Sorting out Lipschitz function approximation (2018, methodology): “This property has been shown to speed up training by orders of magnitude when enforced during weight initialization [37, 43], and explored in the contexts of training RNNs [9] and very deep convolutional neural networks [54].”
Scaling Limits of Wide Neural Networks with Weight Sharing: Gaussian Process Behavior, Gradient Independence, and Neural Tangent Kernel Derivation (2019, methodology): “It has been used to optimize initialization hyperparameters to prevent gradient explosion/vanishing, even to allow training of a 10,000 layer CNN without batchnorm or skip connections (Xiao et al., 2018).”
Neural Tangents: Fast and Easy Infinite Neural Networks in Python (2019, methodology): “…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., 2019; Hayou et al., 2019).”
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].”
A Mean Field Theory of Batch Normalization (2019, methodology): “…justified this assumption as used in the computation of moments, for a wide variety of architectures, like multilayer perceptron (Schoenholz et al., 2016), residual networks (Yang & Schoenholz, 2017), and convolutional networks (Xiao et al., 2018) studied previously, but without batchnorm.”
Bayesian Deep Convolutional Networks with Many Channels are Gaussian Processes (2018, methodology): “They apply a mean field approximation in which the pre-activation signal is replaced with a Gaussian, and the derivation of the covariance function with depth is the same as for the kernel function of a corresponding GP. Recently, Xiao et al. (2017b; 2018) extended this to convolutional architecture”
Mean-field Behaviour of Neural Tangent Kernel for Deep Neural Networks (2019, methodology): “In Xiao et al. (2018), authors studied only the limiting behaviour of correlations cα,α′(x, x) (same input x).”
Batch Normalization Biases Deep Residual Networks Towards Shallow Paths (2020, neutral): “While some authors have succeeded in training very deep networks without normalization layers or skip connections (Saxe et al., 2013; Xiao et al., 2018), these papers required orthogonal initialization schemes which are not compatible with ReLU activation functions.”
Disentangling trainability and generalization in deep learning (2019, methodology): “We will focus on the NTK here since Schoenholz et al. (2017); Xiao et al. (2018) contains a detailed description of the NNGP in this case.”