[Paper Review] Spectrum concentration in deep residual learning: a free probability approach

Zenan Ling, Xing He+1 more|arXiv (Cornell University)|Jul 31, 2018

Statistical Mechanics and Entropy27 references5 citations

TL;DR

This paper proposes a free probability-based analysis of weight initialization in deep residual networks, showing that the variance of initial weights must scale as $O(1/L)$ with network depth $L$ to ensure spectrum concentration in the input-output Jacobian. By analyzing the limiting spectral distribution of non-Hermitian random matrices, the authors derive that this scaling prevents gradient vanishing or exploding, enabling training speeds orders of magnitude faster than classical initialization schemes.

ABSTRACT

We revisit the initialization of deep residual networks (ResNets) by introducing a novel analytical tool in free probability to the community of deep learning. This tool deals with non-Hermitian random matrices, rather than their conventional Hermitian counterparts in the literature. As a consequence, this new tool enables us to evaluate the singular value spectrum of the input-output Jacobian of a fully-connected deep ResNet for both linear and nonlinear cases. With the powerful tool of free probability, we conduct an asymptotic analysis of the spectrum on the single-layer case, and then extend this analysis to the multi-layer case of an arbitrary number of layers. In particular, we propose to rescale the classical random initialization by the number of residual units, so that the spectrum has the order of $O(1)$, when compared with the large width and depth of the network. We empirically demonstrate that the proposed initialization scheme learns at a speed of orders of magnitudes faster than the classical ones, and thus attests a strong practical relevance of this investigation.

Motivation & Objective

• To address the lack of theoretical understanding behind effective weight initialization in deep residual networks.
• To investigate why classical initialization schemes like He's fail to prevent ill-conditioned Jacobians in deep ResNets.
• To establish a theoretical framework using free probability to analyze the spectral properties of the input-output Jacobian in deep residual networks.
• To derive a depth-dependent initialization rule that ensures stable gradient flow during backpropagation.
• To empirically validate that $\sigma_w^2 = O(1/L)$ enables significantly faster learning than classical schemes.

Proposed method

• The authors apply non-Hermitian free probability theory to analyze the limiting spectral distribution of the input-output Jacobian matrix in deep fully-connected ResNets.
• They derive the asymptotic spectral density of the Jacobian $\mathbf{J}\mathbf{J}^T$ in both linear and nonlinear cases as network width $N$ and depth $L$ tend to infinity.
• The analysis focuses on the squared singular values of the Jacobian, showing they remain $O(1)$ only when $\sigma_w^2 = O(1/L)$, not just $O(1/N)$ as in classical initialization.
• The method extends from single-layer to multi-layer networks, proving the universality of the $O(1/L)$ scaling across different nonlinearities and weight initialization types.
• The framework is validated empirically on CIFAR-10 using both fully-connected and convolutional ResNets with varying initialization variances.
• The authors compare learning dynamics under different $\sigma_w^2$ settings, including $\sigma_w^2 = 1/\sqrt{L}$, and demonstrate superior convergence speed.

Experimental results

Research questions

• RQ1 Why do classical weight initialization schemes like He's fail to ensure stable gradient flow in deep residual networks despite proper width scaling?
• RQ2 What is the necessary and sufficient condition on weight variance $\sigma_w^2$ to achieve spectrum concentration in the input-output Jacobian of deep ResNets?
• RQ3 How does the depth $L$ of a residual network affect the spectral properties of its input-output Jacobian, and what scaling of $\sigma_w^2$ ensures $O(1)$ squared singular values?
• RQ4 Can free probability theory provide a rigorous asymptotic analysis of the Jacobian spectrum in deep residual networks with arbitrary depth and width?
• RQ5 Is the $O(1/L)$ scaling of weight variance universally effective across different nonlinearities and initialization methods (Gaussian, orthogonal)?

Key findings

• The squared singular values of the input-output Jacobian remain $O(1)$ only when the weight variance is scaled as $\sigma_w^2 = O(1/L)$, not just $O(1/N)$, ensuring spectrum concentration.
• The $O(1/L)$ scaling is necessary and sufficient for spectrum concentration across a wide range of nonlinearities, including ReLU, in both Gaussian and orthogonal weight initialization.
• Empirical results show that $\sigma_w^2 = 1/\sqrt{L}$ enables training speeds orders of magnitude faster than classical initialization, even without batch normalization.
• For residual units with $m > 1$ layers, the optimal scaling becomes $\sigma_w^2 = O(L^{-1/m})$ to maintain global Jacobian conditioning, though this may cause internal gradient vanishing within the residual unit.
• The trade-off between global Jacobian conditioning and local gradient flow within residual units is confirmed, with extremely small $\sigma_w^2$ causing internal vanishing gradients.
• The theoretical framework based on free probability accurately predicts the empirical behavior of Jacobian spectra in deep ResNets, even for finite $N$ and $L$, validating its practical relevance.