The Emergence of Spectral Universality in Deep Networks
Jeffrey Pennington Samuel S. Schoenholz Surya Ganguli


[

The Emergence of Spectral Universality in Deep Networks

 Jeffrey Pennington Samuel S. Schoenholz Surya Ganguli 

 Google Brain Google Brain Google Brain
Applied Physics, Stanford University   ]


Recent work has shown that tight concentration of the entire spectrum of singular values of a deep network's input-output Jacobian around one at initialization can speed up learning by orders of magnitude. Therefore, to guide important design choices, it is important to build a full theoretical understanding of the spectra of Jacobians at initialization. To this end, we leverage powerful tools from free probability theory to provide a detailed analytic understanding of how a deep network's Jacobian spectrum depends on various hyperparameters including the nonlinearity, the weight and bias distributions, and the depth. For a variety of nonlinearities, our work reveals the emergence of new universal limiting spectral distributions that remain concentrated around one even as the depth goes to infinity.
 



§ INTRODUCTION


A well-conditioned initialization is essential for successfully training neural networks.  Seminal initial work focused on random weight initializations ensuring that the second moment of the spectrum of singular values of the network Jacobian from input to output remained one, thereby preventing exponential explosion or vanishing of gradients <cit.>.  However, recent work has shown that even among different random initializations sharing this property,  those whose entire spectrum tightly concentrates around one can often yield faster learning by orders of magnitude.  For example, deep linear networks with orthogonal initializations, for which the entire spectrum is exactly one, can achieve depth-independent learning speeds, while the corresponding Gaussian initializations cannot <cit.>. 

Recently, it was shown <cit.> that a similarly well-conditioned Jacobian could be constructed for deep non-linear networks using a combination of orthogonal weights and tanh nonlinearities. The result of this improved conditioning was an orders-of-magnitude speedup in learning for tanh networks. However, the same study also proved that a well-conditioned Jacobian could not be achieved with Rectified Linear units (ReLUs). Together these results explained why, historically, in some cases orthogonal weight initialization had been found to improve training efficiency only slightly <cit.>.

These empirical results connecting the conditioning of the Jacobian to a dramatic speedup in learning raise an important theoretical question. Namely, how does the entire shape of this spectrum depend on a network's nonlinearity, weight and bias distribution, and depth?  Here we provide a detailed analytic answer by using powerful tools from free probability theory.  Our answer provides theoretical guidance on how to choose these different network ingredients so as to achieve tight concentration of deep Jacobian spectra even at very large depths.  Along the way, we find several surprises, and we summarize our results in the discussion.      



§ PRELIMINARIES




 §.§ Problem Setup


Consider an L-layer feed-forward neural network of width N with synaptic weight matrices W^l∈ℝ^N× N,  bias vectors b^l, pre-activations h^l, and post-activations x^l, with l=1,…,L.  The forward-propagation dynamics are given by,

    x^l = ϕ(h^l) ,  h^l = W^l x^l-1 + b^l ,

where ϕ : ℝ→ℝ is a pointwise nonlinearity and the input is  x^0 ∈ℝ^N.  Now consider the input-output Jacobian J∈ℝ^N× N given by

    J = ∂x^L/∂x^0 = ∏_l=1^L D^l W^l.

Here D^l is a diagonal matrix with entries D^l_ij = ϕ'(h^l_i)  δ_ij, where δ_ij is the Kronecker delta function.
The input-output Jacobian J is closely related to the backpropagation operator mapping output errors to weight matrices at a given layer, in the sense that if the former is well-conditioned, then the latter tends to be well-conditioned for all weight layers. We are therefore interested in understanding the entire singular value spectrum of J for deep networks with randomly initialized weights and biases.

In particular, we will take the biases b^l_i to be drawn i.i.d. from a zero-mean Gaussian with standard deviation σ_b.  For the weights, we will consider two random matrix ensembles: (1) random Gaussian weights in which each W^l_ij is drawn i.i.d from a Gaussian with variance σ_w^2/N, and (2) random orthogonal weights, drawn from a uniform distribution over scaled orthogonal matrices obeying  (W^l)^T W^l = σ_w^2  I. 




 §.§ Review of Signal Propagation


The random matrices D^l in (<ref>) depend on the empirical distribution of pre-activations h^l_i for i=1,…,N entering the nonlinearity ϕ in (<ref>).  The propagation of this empirical distribution through different layers l was studied in <cit.>. In those works, it was shown that in the large N limit this empirical distribution converges to a Gaussian with zero mean and variance q^l, where q^l obeys a recursion relation induced by the dynamics in (<ref>):

    q^l   =  σ_w^2 ∫ 𝒟h   ϕ( √(q^l-1) h )^2  + σ_b^2 ,

with initial condition q^1 = σ_w^2/N∑_i=1^N (x^0_i)^2+σ_b^2,
and 𝒟h = dh/√(2π) exp(-h^2/2) denoting the standard normal measure. This recursion has a fixed point obeying, 

    q^*   =  σ_w^2 ∫ 𝒟h   ϕ( √(q^*) h )^2  + σ_b^2 .

If the input x^0 is chosen so that q^1 = q^*, then the dynamics start at the fixed point and the distribution of D^l is independent of l. Moreover, even if q^1  q^*, a few layers is often sufficient to approximately converge to the fixed point (see <cit.>). As such, when L is large, it is often a good approximation to assume that q^l = q^* for all depths l when computing the spectrum of J.

Another important quantity governing signal propagation through deep networks <cit.> is 

    χ   = 1/N⟨Tr  (DW)^T DW⟩
       = σ_w^2 ∫ 𝒟h [ ϕ' ( √(q^*) h ) ]^2,

where ϕ' is the derivative of ϕ.  Here χ is second moment of the distribution of squared singular values of the matrix DW, when the pre-activations are at their fixed point distribution with variance q^*.  As shown in <cit.>, χ(σ_w,σ_b) separates the (σ_w,σ_b) plane into two regions: (a) when χ > 1, forward signal propagation expands and folds space in a chaotic manner and back-propagated gradients exponentially explode; and (b) when χ < 1, forward signal propagation contracts space in an ordered manner and back-propagated gradients exponentially vanish.   Thus the constraint χ(σ_w,σ_b)=1 determines a critical line in the (σ_w,σ_b) plane separating the ordered and chaotic regimes.  Moreover, the second moment of the distribution of squared singular values of J was shown simply to be χ^L in <cit.>.  Fig. <ref> shows an example of an order-chaos transition for the tanh nonlinearity.  




 §.§ Review of Free Probability

The previous section revealed that the mean squared singular value of J is χ^L. Indeed when χ≪ 1 or χ≫ 1 the vanishing or explosion of gradients, respectively, dominates the learning dynamics and provide a compelling case for choosing an initialization that is critical with χ = 1. We would like to investigate the question of whether or not all cases where χ = 1 are the same and, in particular, to obtain more detailed information about entire the singular value distribution of J when χ=1. Since (<ref>) consists of a product of random matrices, free probability becomes relevant as a powerful tool to compute the spectrum of J, as we now review. See <cit.> for a pedagogical introduction, and <cit.> for prior work applying free probability to deep learning.

In general, given a random matrix X, its limiting spectral density is defined as

    ρ_X(λ) ≡⟨1/N∑_i=1^N δ(λ - λ_i) ⟩_X,

where ⟨·⟩_X denotes an average w.r.t to the distribution over the random matrix X.  

The Stieltjes transform of ρ_X is defined as,

    G_X(z) ≡∫_ℝρ_X(t)/z-tdt ,     z∈ℂ∖ℝ ,

which can be inverted using,

    ρ_X(λ) = -1/πlim_ϵ→0^+Im  G_X(λ + i ϵ) .

G_X is related to the moment generating function M_X,

    M_X(z) ≡ z G_X(z) -1 = ∑_k=1^∞m_k/z^k ,

where m_k is the kth moment of the distribution ρ_X,

    m_k  = ∫ dλ ρ_X(λ) λ^k = 1/N⟨X^k ⟩_X .

In turn, we denote the functional inverse of M_X by M_X^-1, which by definition satisfies  M_X(M_X^-1(z)) = M_X^-1(M_X(z))=z.  
Finally, the S-transform <cit.> is defined as,

    S_X(z) = 1+z/z  M_X^-1(z) .

The utility of the S-transform arises from its behavior under multiplication. Specifically, if A and B are two freely independent random matrices, then the S-transform of the product random matrix ensemble AB is simply the product of their S-transforms,

    S_AB(z) = S_A(z) S_B(z) .
 



§ MASTER EQUATION FOR SPECTRAL DENSITY




 §.§ S-transform for Jacobians


We can now write down an implicit expression of the spectral density of JJ^T, which is also the distribution of the square of the singular values of J.   In particular, in the supplementary material (SM) Sec. 1, we combine (<ref>) with the facts that the S-transform depends only on traces of moments through (<ref>), and that these traces are invariant under cyclic permutations, to derive a simple expression for the S-transform of JJ^T,

    S_JJ^T =  ∏_l=1^L S_(D^l)^2 S_(W^l)^T W^l = S^L_D^2 S^L_W^T W.

Here the lack of dependence on the layer index l on the RHS  is valid if the input x^0 is such that q^1 = q^*. 

Thus, given expressions for the S-transforms associated with the nonlinearity, S_D^2, and the weights, S^L_W^T W, one can compute the S-transform of the input-output Jacobian S_JJ^T at any network depth L through (<ref>).  Then from S_JJ^T, one can invert the sequence (<ref>), (<ref>), and (<ref>) to obtain ρ_JJ^T(λ). 



 §.§ An Efficient Master Equation

The previous section provides a naive method for computing the spectrum ρ_JJ^T(λ), through a complex sequence of calculations.  One must start from ρ_W^TW(λ) and  ρ_D^2(λ), compute their respective Stieltjes transforms, moment generating functions, inverse moment generating functions, and S-transforms, take the product in (<ref>), and then invert this sequence of steps to finally arrive at ρ_JJ^T(λ).  Here we provide a much simpler “master" equation for extracting information about ρ_JJ^T(λ) and its moments directly from knowledge of the moment generating function of the nonlinearity, M_D^2(z), and the S-transform of the weights, S_W^TW(z).  As we shall see, these latter two functions are the simplest functions to work with for arbitrary nonlinearities.

To derive the master equation, we insert (<ref>), for X = D^2, into (<ref>), and perform some algebraic manipulations (see SM Sec. 3 for details) to obtain implicit functional equations for M_JJ^T(z) and G(z),


    M_JJ^T(z) = M_D^2( z^1/L F(M_JJ^T(z))) ,        
    
    z G(z) - 1 =  M_D^2( z^1/L F(z G(z)-1)),

where,

    F(x) = S_W^TW(x)(1+x/x)^1-1/L .

In principle, a solution to eq. (<ref>) allows us to compute the entire spectrum of J J^T. In practice, when an exact solution in terms of elementary functions is lacking, it is still possible to extract robust numerical solutions, as we describe in the next subsection.





 §.§ Numerical Extraction of Spectra


Here we describe how to solve (<ref>) numerically. The difficulty is that (<ref>) implicitly defines G(z) through an equation of the form ℱ(G,z)=0. Notice that, for any given z, this equation may have multiple roots in G.  The correct branch can be chosen by requiring that z →∞,  G(z) ∼ 1/z <cit.>. Therefore, one point on the correct branch can be found by taking |z| large, and finding the solution to ℱ(G,z)=0 that is closest to G = 1/z.  Recall that to obtain the density ρ_JJ^T(λ) through the inversion formula ((<ref>)), we need to extract the behavior of G(z) near the real axis at a point z = λ + iϵ where ρ_JJ^T(λ) has support. So, practically speaking, for each λ we can walk along the imaginary direction obeying Re(z)=λ from large imaginary values to small, and repeatedly solve ℱ(G,z)=0, always choosing the root that is closest to the previous root. 

A potential pitfall arises if we approach a point z where ℱ(G,z)=0 has a double root in G, which could cause us to leave the correct branch of roots and then traverse an incorrect branch.  However, points in the complex two dimensional plane (G,z) ∈ℂ^2 where ℱ has a double root in G are expected to be a set of measure 0, and in practice they do not seem to be a concern. Algorithm <ref> summarizes our heuristic for computing ρ(λ) for each λ of interest.



In the following sections, we demonstrate through many examples a precise numerical match between the outcome of Algorithm <ref> and direct simulations of various random neural networks, thereby justifying not only (<ref>), but also the efficacy our algorithm. 
 


 §.§ Moments of Deep Spectra


In addition to numerically extracting the spectrum of JJ^T, we can also calculate its moments m_k encoded in the function

    M_JJ^T(z) ≡∑_k=1^∞m_k/z^k .

These moments in turn can be computed in terms of the series expansions of S_W^TW and M_D^2,  which we define as

    S_W^TW(z) ≡   σ_w^-2(1 + ∑_k=1^∞ s_k z^k) 
    
    M_D^2(z) ≡   ∑_k=1^∞μ_k/z^k ,

where the moments μ_k  of D^2 are given by,

    μ_k = ∫𝒟h ϕ'(√(q^*) h)^2k .

Substituting these expansions into (<ref>), we obtain equations for the unknown moments m_k in terms of the known moments μ_k and s_k. We can solve for the low-order moments by expanding (<ref>) in powers of z^-1. By equating the coefficients of z^-1 and z^-2, we find equations for m_1 and m_2 whose solution yields (see SM Sec. 3),

    m_1    = (σ_w^2 μ_1)^L
    
    m_2    = (σ_w^2 μ_1)^2L  L (μ_2/μ_1^2 + 1/L - 1 - s_1 ) .

Note the combination σ_w^2 μ_1 is none other than χ defined in (<ref>), and so (<ref>) recovers the result that the mean squared singular value m_1 of J either exponentially explodes or vanishes unless χ(σ_w, σ_b)=1 on a critical boundary between order and chaos.   However, even on this critical boundary where the mean m_1 of the spectrum of JJ^T is one for any depth L, the variance

    σ_JJ^T^2 = m_2 - m_1^2 = L (μ_2/μ_1^2 - 1 - s_1)

grows linearly with depth L for generic values of μ_1, μ_2 and s_1.
Thus J can be highly ill-conditioned at large depths L for generic choices of nonlinearities and weights, even when σ_w and σ_b  are tuned to criticality.  



§ SPECIAL CASES OF DEEP SPECTRA




Exploiting the master equation (<ref>) requires information about M_D^2(z), and  S_WW^T(z).  We first provide this information and then use it to look at special cases of deep networks.



 §.§ Transforms of Nonlinearities


First, for any nonlinearity ϕ(h), we have, through (<ref>) and (<ref>), 

    M_D^2(z) = ∫𝒟h ϕ'(√(q*) h)^2/z - ϕ'(√(q^*)h)^2 .

The integral over the Gaussian measure 𝒟h reflects a sum over all the activations h^l_i in a layer l, since in the large N limit the empirical distribution of activations converges to a Gaussian with standard deviation √(q^*).   Moreover, an activation h^l_i feels a squared slope ϕ'(h^l_i)^2, which appears as an eigenvalue of the diagonal matrix (D^l)^2.  Thus M_D^2(z) naturally involves an integral over a function of ϕ'(·)^2 against a Gaussian.

Table <ref> provides the moment generating function and moments of D^2 for several nonlinearities. Detailed derivations of the results in Table <ref>, which follow from performing the integral in (<ref>), can be found in the SM Sec. 3. In the Erf case, Φ is a special function known as the Lerch transcendent, which can be defined by its moments μ_k. 



 §.§ Transforms of Weights
 


The S-transforms of the weights can be obtain through the sequence of equations (<ref>), (<ref>), and (<ref>), starting with ρ_W^TW(λ) = δ(λ-1) for an orthogonal random matrix W, and ρ_W^TW(λ) = (2π)^-1√(4-λ)  for λ∈ [0,4], for a Gaussian random matrix W with variance 1/N (see SM Sec. 5). Furthermore, by scaling W→σ_w W, the S-transform scales as S_W^TW→σ_w^-2 S_W^TW, yielding the S-transforms and first moments in Table <ref>.  



 §.§ Exact Properties of Deep Spectra


Now for different randomly initialized deep networks, we insert the appropriate expressions in Tables <ref> and <ref> into our master equations (<ref>) and (<ref>) to obtain information about the spectrum of JJ^T, including its entire shape, through Algorithm  <ref>, and its variance σ_JJ^T^2 through (<ref>) and (<ref>).  We always work at criticality, so that in (<ref>), χ=σ_w^2 μ_1 = 1.  The resulting condition for σ_w^2 at criticality and the value of σ_JJ^T^2 are shown in Table <ref> for different nonlinearities, both for orthogonal (s_1=0) and Gaussian (s_1=-1) weights.






  §.§.§ Linear Networks

For linear networks, the fixed point equation (<ref>) reduces to q^* = σ_w^2 q^* + σ_b^2, and (σ_w, σ_b)=(1,0) is the only critical point.  Moreover, linear Gaussian networks behave very differently from orthogonal ones.  The latter are well conditioned, with σ^2_JJ^T=0 because the product of orthogonal matrices is orthogonal and so ρ_JJ^T(λ) = δ(λ-1) for all L.  However, σ^2_JJ^T=L for Gaussian weights. This radically different behavior of the spectrum of JJ^T is shown in Fig. <ref>A.   



  §.§.§ ReLU Networks

For ReLU networks, the fixed point equation (<ref>) reduces to q^* = 1/2σ_w^2 q^* + σ_b^2, and (σ_w, σ_b)=(√(2),0) is the only critical point. Unlike the linear case, σ_JJ^T^2 becomes L for orthogonal and 2L for Gaussian weights.  In essence, the ReLU nonlinearity destroys the qualitative scaling advantage that linear networks possess for orthogonal weights versus Gaussian.  The qualitative similarity of spectra for ReLU Orthogonal and linear Gaussian is shown in Fig. <ref>AB.



  §.§.§ Hard Tanh and Erf Networks

For Hard Tanh and Erf Networks, the criticality condition σ_w^2  = μ_1^-1 does not determine a unique value of σ_w^2 because μ_1, the mean squared slope ϕ'(h)^2, now depends on the variance q^* of the distribution of pre-activations h.  Since q^* itself is a function of σ_w and σ_b through (<ref>), these networks enjoy an entire critical curve in the (σ_w, σ_b) plane, similar to that shown in Fig. <ref>. As q^* decreases monotonically towards zero, the corresponding point on this curve approaches the point (σ_w,σ_b)=(1,0).

Moreover, Table <ref> shows that σ_JJ^T^2 = L(ℱ(q^*) - 1 - s_1) with lim_q^* → 0ℱ(q^*) = 1.  This implies that for Gaussian weights (s_1=-1), no matter how small one makes σ_w, σ^2_JJ^T∝ L.  However, for orthogonal weights (s_1=0), for any fixed L, one can reduce σ_w and therefore q^*, so as to make  σ_JJ^T^2 arbitrarily small. Thus Hard Tanh and Erf nonlinearities rescue the scaling advantage that orthogonal weights possess over Gaussian, which was present in linear networks, but destroyed in ReLU networks.  Examples of the well-conditioned nature of orthogonal Hard Tanh and Erf networks compared to orthogonal ReLu networks are shown in Fig. <ref>.



§ UNIVERSALITY IN DEEP SPECTRA



Table <ref> shows that for orthogonal Erf and Hard Tanh networks (but not ReLU networks), since σ_JJ^T^2 = L(ℱ(q^*)-1) with lim_q^* → 0ℱ(q^*) = 1, one can always choose q^* to vary inversely with L so as to achieve a desired L-independent constant variance σ^2_JJ^T≡σ^2_0. To achieve this scaling, q^*(L) should satisfy the equation ℱ(q^*(L)) = 1 + σ^2_0/L, which implies σ_w→ 1 and q^* → 0 as L→∞.

Remarkably, in this double scaling limit, not only does the variance of the spectrum of JJ^T remain constant at the fixed value σ_0^2, but the entire shape of the distribution converges to a universal limiting distribution as L →∞. There is more than one possible limiting distribution, but its form depends on ϕ only through the distribution of ϕ'(h)^2 as q^*→ 0 via the expression for M_D^2(z) in (<ref>). Therefore, many qualitatively different activation functions may in fact be members of the same universality class. We identify two universality classes that correspond to many common activation functions: the Bernoulli universality class and the smooth universality class, named based on the distribution of ϕ'(h)^2 as q^*→ 0.

The Bernoulli universality class contains many piecewise linear activation functions, such as Hard Tanh (Fig. <ref>C) and a version of ReLU shifted so as to be linear at the origin, which for concreteness we define as ϕ(x) = [x+1/2]_+ - 1/2 (Fig. <ref>E). While these functions look quite different, their derivatives are both Bernoulli-distributed (Fig. <ref>DF) and the limiting spectra of their corresponding Jacobians are the same (Fig. <ref>AB).

The smooth universality class contains many smooth activation functions, such as Erf (Fig. <ref>G) and a smoothed version of ReLU that we take to be the sigmoid-weighted linear unit (SiLU) <cit.> (Fig. <ref>I). In this case, not only do the activation functions themselves look different, but so too do their derivatives (Fig. <ref>HJ). Nevertheless, in the double scaling limit, the limiting spectra of their corresponding Jacobians are the same (Fig. <ref>CD). The rate of convergence to the limiting distribution is different, because the moments μ_k differ substantially for non-zero q^*.

Unlike the smoothed and shifted versions of ReLU, the vanilla ReLU activation (Fig. <ref>AB) behaves entirely differently and has no limiting distribution because the μ_k are independent of q^* and therefore it is impossible to attain an L-independent constant variance σ^2_JJ^T≡σ^2_0 in this case.

To understand the mechanism behind the emergence of spectral universality, we now examine orthogonal networks whose activation functions have squared derivatives obeying a Bernoulli distribution and show that they all share a universal limiting distribution as L→∞. To this end, we suppose that,

    M_D^2 = p(q^*)1/z-1 ,

for some function p(q^*) that measures the probability of the nonlinearity having slope one as a function of q^*. We will assume that p(q^*)→ 1 as q^*→ 0. The relevant ratio of moments and the weight variance σ_w^2 are given as,

    μ_2/μ_1^2 = 1/μ_1 = σ_w^2 = 1/p(q^*) .

From (<ref>), we have,

    σ_JJ^T^2 = σ_0^2 = L (1/p(q^*)-1) ⇒ p(q^*) = 1 + σ_0^2/L .

Notice that a solution q^*(L) to (<ref>) will exist for large L since we are assuming p(q^*)→ 1 as q^*→ 0. Substituting this solution in (<ref>) and (<ref>) gives for large L,

    M_D^2 = L/L+σ_0^21/z-1  and  μ_1 = L/L+σ_0^2 .

Using these expressions and (<ref>), we find that the S-transform obeys,

    S^Bernoulli_JJ^T = (μ_1 1+z/z M^-1_D^2)^L = (1 + z σ_0^2/L(1+z))^-L .

The large depth limit gives,

    S^Bernoulli_JJ^T  = e^-z σ_0^2/(1+z) .

Using (<ref>) and (<ref>) to solve for G(z) gives,

    G(z) = 1/zσ_0^2/σ_0^2+W(-σ_0^2/z) ,

where W denotes the principal branch of the Lambert-W function <cit.> and solves the transcendental equation,

    W(x) e^W(x) = x .

The spectral density can be extracted from (<ref>) easily using (<ref>). The results are shown in black lines in Fig. <ref>AB. Both Hard Tanh and Shifted ReLU have Bernoulli-distributed ϕ'(h)^2 and, despite being qualitatively different activation functions, have the same limiting spectral distributions. It is evident that the empirical spectral densities converge to this universal limiting distribution as the depth increases.

Next we build some additional understanding of the spectral density implied by (<ref>). Because the spectral density is proportional to the imaginary part of G(z), we expect the locations of the spectral edges to be related to branch points of G(z), or more generally to poles in its derivative. Using the relation,

    W'(x) = 1/x + e^W(x) ,

we can inspect the derivative of G(z). It may be expressed as,

    G'(z) = -σ_0^2(σ_0^2 +W(-σ_0^2/z)(σ_0^2 +W(-σ_0^2/z))/z^2(1+W(-σ_0^2/z))(σ_0^2 + W(-σ_0^2/z))^2 .

By inspection, we find that G'(z) has double poles at,

    z = λ_0 = 0,    z = λ_2 = e^σ_0^2 ,

which are locations where the spectral density diverges, i.e. there are delta function peaks at λ_0 and λ_2. Note that there is only a pole at λ_2 if σ_0 ≤ 1. There is also a single pole at,

    λ_1 = σ_0^2 e ,

which defines the right spectral edge, i.e. the maximum value of the bulk of the density.

The above observations regarding λ_0, λ_1, and λ_2 are evident in Fig. <ref>AB. Noting that in the figure, σ_0 = 1/2, we predict that the bulk of the density to have its right edge located at s = √(λ_1) = √(e)/2 ≈ 0.82 and that there should be a delta function peak at s=√(λ_2) = e^1/8≈ 1.13, both of which are reflected in the figure.

A similar analysis can be carried out for activation functions for which the distribution of ϕ'(h)^2 is smooth and concentrates around one as q^*→ 0. The analysis for Erf is presented in the SM. We find that,

    S^Smooth_JJ^T  = e^-z σ_0^2 ,

and that G(z) can be expressed in terms of a generalized Lambert-W function <cit.>. The locations of the spectral edges are given by s_± = e^-1/4σ_±^2√(1 + 1/2σ_∓^2), where,

    σ_±^2 = σ_0(σ_0±√(σ_0^2 + 4)) .

For σ_0=1/2, these results give s_-≈0.57 and s_+=1.56, which is in excellent agreement with the behavior observed in Fig. <ref>CD. Overall, Fig. <ref> provides strong evidence supporting our predictions that orthogonal Hard Tanh and shifted ReLU networks have the Bernoulli limit distribution, while orthogonal Erf and smoothed Relu networks have the smooth limit distribution.  

Finally, we derived these universal limits assuming orthogonal weights. In the SM we show that orthogonality is in fact necessary for the existence of a stable limiting distribution for the spectrum of JJ^T.  No other random matrix ensemble can yield a stable distribution for any choice of nonlinearity with ϕ'(0)=1. Essentially, any spread in the singular values of W grows in an unbounded way with depth and cannot be nonlinearly damped. 




§ DISCUSSION

In summary, motivated by a lack of theoretical clarity on when and why different weight initializations and nonlinearities combine to yield well-conditioned spectra that speed up deep learning, we developed a calculational framework based on free probability to provide, with unprecedented detail, analytic information about the entire Jacobian spectrum of deep networks with arbitrary nonlinearities. Our results provide a principled framework for the initialization of weights and the choice of nonlinearities in order to produce well-conditioned Jacobians and fast learning. Intriguingly, we find novel universality classes of deep spectra that remain well-conditioned as the depth goes to infinity, as well as theoretical conditions for their existence. Our results lend additional support to the surprising conclusions revealed in <cit.>, namely that using either Gaussian initializations or ReLU nonlinearities precludes the possibility of obtaining stable spectral distributions for very deep networks. Beyond the sigmoidal units advocated in <cit.>, our results suggest that a wide variety of nonlinearities, including shifted and smoothed variants of ReLU, can achieve dynamical isometry, provided the weights are orthogonal.  Interesting future work could involve the discovery of new universality classes of well-conditioned deep spectra for more diverse nonlinearities than considered here.

unsrt





The Emergence of Spectral Universality in Deep Networks:
Supplementary Material












§ REVIEW OF FREE PROBABILITY


For what follows, we define the key objects of free probability.  Given a random matrix X, its limiting spectral density is defined as

    ρ_X(λ) ≡⟨1/N∑_i=1^N δ(λ - λ_i) ⟩_X,

where ⟨·⟩_X denotes an average w.r.t to the distribution over the random matrix X.  For large N, the empirical histogram of eigenvalues of a single realization of X converges to ρ_X.  In turn, the Stieltjes transform of ρ_X is defined as,

    G_X(z) ≡∫_ℝρ_X(t)/z-tdt ,     z∈ℂ∖ℝ ,

which can be inverted using,

    ρ_X(λ) = -1/πlim_ϵ→0^+Im  G_X(λ + i ϵ) .

G_X is related to the moment generating function M_X,

    M_X(z) ≡ z G_X(z) -1 = ∑_k=1^∞m_k/z^k ,

where the m_k is the k'th moment of the distribution ρ_X,

    m_k  = ∫ dλ ρ_X(λ) λ^k = 1/N⟨X^k ⟩_X .

In turn, we denote the functional inverse of M_X by M_X^-1, which by definition satisfies  M_X(M_X^-1(z)) = M_X^-1(M_X(z))=z.  
Finally, the S-transform <cit.> is defined in terms of the functional inverse M_X^-1 as,

    S_X(z) = 1+z/z  M_X^-1(z) .

The utility of the S-transform arises from its behavior under multiplication. Specifically, if A and B are two freely independent random matrices, then the S-transform of the product random matrix ensemble AB is simply the product of their S-transforms,

    S_AB(z) = S_A(z) S_B(z) .
 




§ FREE PROBABILITY AND DEEP NETWORKS


We will now use eqn. (<ref>) to write down an implicit definition of the spectral density of JJ^T, which is also the distribution of the square of the singular values of J.  Here J is the input-output Jacobian of a deep network defined in the main paper.  First notice that, by eqn. (<ref>), M(z) and thus S(z) depend only on the moments of the spectral density. The moments, in turn, can be defined in terms of traces (as in eqn. (<ref>)), which are invariant to cyclic permutations, i.e.,

    (^1 ^2⋯^m)^k =  (^2⋯^m ^1)^k .

Therefore the S-transform is invariant to cyclic permutations. Now define matrices Q^l and Q̃^l as,

    Q^L    ≡JJ^T = ( D^LW^L D^L-1⋯D^1 W^1)(D^LW^L D^L-1⋯D^1 W^1)^T
    Q̃^L    ≡[ (D^LW^L)^T D^LW^L ](W^L-1D^L-2⋯D^1 W^1) (W^L-1D^L-2⋯D^1 W^1)^T 
        =  [ (D^LW^L)^T D^LW^L ] Q^L-1 .

Now Q^L and  Q̃^L are related by a cyclic permutation. Therefore the above argument shows that their S-transforms are equal, i.e. S_Q_L = S_Q̃_L. Furthermore (D^LW^L)^T D^LW^L  and (D^L)^2 (W^L)^T W^L are related by a cyclic permutation, implying their S-transforms are also equal.   Then a recursive application of eqn. (<ref>) and cyclic invariance of S-transforms implies that,

    S_JJ^T = S_Q_L  = S_(D^L)^2 S_(W^L)^T W^L S_Q^L-1 = ∏_l=1^L S_(D^l)^2 S_(W^l)^T W^l = S^L_D^2 S^L_W^T W

where the last equality follows if each term in the Jacobian product identically distributed. 



Given the expression for S_JJ^T, a simple procedure recovers the density of singular values of J: 


  * Use eqn. (<ref>) to obtain the moment generating function M_JJ^T(z)

  * Use eqn. (<ref>) to obtain the Stieltjes transform G_JJ^T(z)

  * Use eqn. (<ref>) to obtain the spectral density ρ_JJ^T(λ)

  * Use the relation λ = σ^2 to obtain the density of singular values of J.

So in order to compute the distribution of singular values of of J, all that remains is to compute the S-transforms of W^TW and of D^2. We will attack this problem for specific activation functions and matrix ensembles in the following sections.



§ DERIVATION OF MASTER EQUATIONS FOR THE SPECTRUM OF THE JACOBIAN


To derive the master equation, we first insert (<ref>), for X = D^2, into (<ref>) to obtain

    S_JJ^T = S_W^TW^L(1+z/z)^L(M_D^2^-1)^-L .

Then we find M_JJ^T^-1 = (1+z)(zS_J^TJ)^-1 by inverting (<ref>), which combined with the above equation yields

    M_JJ^T^-1 = S_W^TW^-L(1+z/z)^1-L(M_D^2^-1)^L .

Then solving for M_D^2^-1 yields

    M_D^2^-1  = (M_JJ^T^-1 S_W^TW^L(1+z/z)^L-1)^1/L .

Applying M_D^2 to both sides gives,

    z  = M_D^2( (M_JJ^T^-1 S_W^TW^L(z)(1+z/z)^L-1)^1/L) .

Finally, evaluating this equation at z = M_JJ^T gives our sought after master equation:

    M_JJ^T(z) = M_D^2( z^1/L S_W^TW(M_JJ^T(z)) (1+1/M_JJ^T(z))^1-1/L) .

This is an implicit functional equation for M_JJ^T(z), an unknown quantity, in terms of the known functions M_D^2(z) and S_W^TW(z).   
Furthermore, by substituting (<ref>), M_JJ^T = zG_JJ^T-1, into (<ref>), we also obtain an implicit functional equation for the Stieltjes transform G of ρ_JJ^T(λ),

    z G - 1  =M_D^2( z^1/L S_W^TW(z G-1)(z G/z G - 1)^1-1/L).




§ DERIVATION OF MOMENTS OF DEEP SPECTRA

The moments m_k of  the spectrum of JJ^T are encoded in the moment generating function

    M_JJ^T(z) ≡∑_k=1^∞m_k/z^k .

These moments in turn can be computed in terms of the series expansions of S_W^TW and M_D^2,  which we define as

    S_W^TW(z) ≡   σ_w^-2(1 + ∑_k=1^∞ s_k z^k) 
    
    M_D^2(z) ≡   ∑_k=1^∞μ_k/z^k ,

where the moments μ_k  of D^2 are given by,

    μ_k = ∫𝒟h ϕ'(√(q^*) h)^2k .



We can substitute these moment expansions into (<ref>) to obtain equations for the unknown moments m_k of the spectrum of JJ^T, in terms of the known moments μ_k and s_k.  We can solve for the low order moments by expanding  (<ref>) in powers of z^-1.  By equating the coefficients of z^-1 and z^-2, we obtain the following equations for m_1 and m_2,

    m_1    = σ_w^2 μ_1 m_1^1-1/L
    
    m_2    = σ_w^4 μ_2 m_1^2-2/L
       + σ_w^2 μ_1 m_1^2-1/L((m_2/m_1^2-1)(1-1/L) - s_1) .

Solving for m_1 and m_2 yields,

    m_1    = (σ_w^2 μ_1)^L
    
    m_2    = (σ_w^2 μ_1)^2L  L (μ_2/μ_1^2 + 1/L - 1 - s_1 ) .





§ TRANSFORMS OF NONLINEARITIES

Here we compute the moment generating functions M_D^2(z) for various choices of the nonlinearity ϕ, some of which are displayed in Table 1 of the main paper. 


 §.§ ϕ(x) = x


    M_D^2(z)    = ∫𝒟x 1/z - 1
       = 1/z-1 .



 §.§ ϕ(x) = [x]_+


    M_D^2(z)    = ∫𝒟x θ(x)^2/z - θ(x)^2
       = 1/2∫𝒟x 1/z - 1
       = 1/21/z - 1 .



 §.§ ϕ(x) = (x)


    M_D^2(z)    = ∫𝒟x θ(1-q_* x)^2/z - θ(1+ q_* x)^2
       = (1/√(2) q_*) ∫𝒟x 1/z - 1
       = (1/√(2) q_*)  1/z - 1 .



 §.§ ϕ(x) = [x]_+ + α [-x]_+


    M_D^2(z)    = ∫𝒟x ϕ'(q_* x)^2/z - ϕ'(q_* x)^2
       =1/2(z-1) + 1/2(z/α^2-1) .



 §.§ ϕ(x) = (√(π)/2x)


    M_D^2(z)    = ∫𝒟x ϕ'(q_* x)^2/z - ϕ'(q_* x)^2
       = ∑_k=1^∞z^-k∫𝒟x  e^-1/2π k q_*^2 x^2
       = ∑_k=1^∞1/z^k √(1 +π k q_*^2)
       = 1/√(π) q_* zΦ(1/z, 1/2, 1 + 1/π q_*^2) ,

where Φ is the special function known as the Lerch transcendent. 


 §.§ ϕ(x) = 2/πarctan(π/2x)


    M_D^2(z)    = ∫𝒟x ϕ'(q_* x)^2/z - ϕ'(q_* x)^2
       = ∫𝒟x 16/(4+π^2 q_*^2 x^2)^2 z + 16
       = -√(2)/π^3/2 q_*^2 √(z)(e^z_+/2/√(z_+)(√(z_+/2)) - e^z_-/2/√(z_-)(√(z_-/2))) ,

where,

    z_± = 4(√(z)± 1)/π^2 q_*^2 √(z) .





§ TRANSFORMS OF WEIGHTS


First consider the case of an orthogonal random matrix satisfying W^TW = I.  Then 

    ρ_W^TW(λ)    = δ(λ-1) 
    
    G_W^TW(z)    = (z-1)^-1
     
    M_W^TW(z)    = (z-1)^-1
    
    M_W^TW^-1(z)    = (1+z)/z 
    
    S_W^TW(z)    = 1.

The case of a random Gaussian random matrix W with zero mean, variance 1/N entries is more complex, but well known:

    ρ_W^TW(λ)    = (2π)^-1√(4-λ)  for λ∈ [0,4] 
    
    G_W^TW(z)    = 1/2(1- √(z-4/2)) 
     
    M_W^TW(z)    =  1/2(z-√(z(z-4))-2) 
    
    M_W^TW^-1(z)    = (1+z)^2/z 
    
    S_W^TW(z)    = (1+z)^-1.

Furthermore, by scaling W→σ_w W, the S-transform scales as S_W^TW→σ_w^-2 S_W^TW, yielding the S-transforms in Table <ref>.  




§ UNIVERSALITY CLASS OF ORTHOGONAL HARD TANH NETWORKS


We consider hard tanh with orthogonal weights. The moment generating function is,

    M_D^2 = (1/√(2q^*))  1/z - 1 ,

so that

    μ_2/μ_1^2 = 1/(1/√(2q^*))   and   g = 1/√(μ_1) = 1/(1/√(2q^*)) .

Also we have,

    σ_JJ^T^2 = L (1/(1/√(2q^*))-1)   ⇒   q_*(L) = 1/√(2 ^-1(L/L+σ_0^2)) .

if we wish to scale q^* with depth L so as to achieve a depth independent constant variance σ_JJ^T^2 = σ_0^2 as L →∞. 
This expression for q^* gives,

    M_D^2 = L/L+σ_0^21/z - 1  and  μ_1 = L/L+σ_0^2 ,

so that,

    S_JJ^T = S_D^2 = (μ_1 1+z/z M^-1_D^2)^L = (L(1+z)/L(1+z) + z σ_0^2)^L = (1 + z σ_0^2/L(1+z))^-L .

The large depth limit gives,

    S_JJ^T  = e^-z σ_0^2/(1+z) .

Solving for G(z) gives,

    G(z) = 1/z1/1+W(-σ_0/z)/σ_0^2 ,

where W is the standard Lambert-W function, or product log. The derivative of this function has double poles at,

    λ_0 = 0,   λ_2 = e^σ_0^2 ,

which are locations where the spectral density diverges. There is also a single pole at,

    λ_1 = σ_0^2 e ,

which is the maximum value of the bulk of the density.



§ UNIVERSALITY CLASS OF ORTHOGONAL  NETWORKS

Consider ϕ(x) = √(π/2)(x/√(2)), which has been scaled so that ϕ'(0) = 1 and ϕ”'(0) = -1. The μ_k are given by,

    μ_k = 1/√(1 + 2 k q^*)

so that

    σ_JJ^T^2 = L (1 + 2 q^*/√(1 +4 q^*) - 1 ) .

If we wish to scale q^* with depth L so as to achieve a depth independent constant variance σ_JJ^T^2 = σ_0^2 as L →∞, then we can choose 


    q^*(L) = (σ_0^2/L+σ_0^4/2 L^2+(σ_0/2L+σ_0^3/2L^2)√(2L+σ_0^2))^1/4 .

Since we also assume the network is critical, we also have that,

    σ_w^2 = (1 + 2 q^*)^1/4 .



To illustrate universality, we next consider an arbitrary activation function, and assume that it has a Taylor expansion around 0. This allows us to expand the μ_k. First we write,

    ϕ(x) = ∑_k=0^∞ϕ_k x^k ,

We will need ϕ_1  0. First we will assume that  ϕ_2  0. Using this expansion we can write,

    μ_k = ϕ_1^2k(1 + k((2k-1)ϕ_2^2/ϕ_1^2 + ϕ_3/ϕ_1)q_*^2 + 𝒪(q_*^4)) .

We also have

    S_JJ^T    = (μ_1 1+z/z M_D^2^-1)^L ,

where we have used the fact that the network is critical so that we have μ_1 = g^-2. Using the Lagrange inversion theorem to expand M_D^2^-1, we find that 

    μ_1 1+z/z M_D^2^-1 = 1 - 4ϕ_2^2/ϕ_1^2z q_*^2 + 𝒪(q_*^4) .

Meanwhile,

    L    = σ_0^2  (μ_2/μ_1^2-1)^-1,
       = σ_0^2 ϕ_1^2/4ϕ_2^2 q_*^2 ,

so that,

    S_JJ^T    = (μ_1 1+z/z M_D^2^-1)^L
       =(1 - σ_0^2/L z)^L 
       = e^-σ_0^2 z + 𝒪(L^-1) ,

Next we will assume that  ϕ_2 = 0 and ϕ_3  0[We suspect these additional assumptions are unnecessary and that the results which follow are valid so long as there exists a k for which ϕ_k  0. It would be interesting to prove this.]. Using the above expansion we can write,

    μ_k = ϕ_1^2k(1 + kϕ_3/ϕ_1q_*^2 + 𝒪(q_*^4)) .

Also we have,

    S_JJ^T    = (μ_1 1+z/z M_D^2^-1)^L ,

where we have used the fact that the network is critical so that we have μ_1 = σ_w^-2. Using the Lagrange inversion theorem to expand M_D^2^-1, we find that 

    μ_1 1+z/z M_D^2^-1 = 1 - 2ϕ_3^2/ϕ_1^2z q_*^4 + 𝒪(q_*^6) .

Meanwhile,

    L    = σ^2  (μ_2/μ_1^2-1)^-1,
       = σ^2 ϕ_1^2/2ϕ_3^2 q_*^4 ,

so that,

    S_JJ^T    = (μ_1 1+z/z M_D^2^-1)^L
       =(1 - σ_0^2/L z)^L 
       = e^-σ_0^2 z + 𝒪(L^-1) ,

establishing a universal limiting S-transform (subject to our assumptions). From this result we can extract the Stieltjes transform and thus the spectral density. The result establishes a universal double scaling limiting spectral distribution.



Next we observe that the Stieltjes transform can be expressed in terms of a generalization of the Lambert -W function called the r-Lambert function, W_r(z), which is defined by

    W_r e^W_r + r W_r = z .

In terms of this function, the Stieltjes transform is,

    G(z) = W_-e^σ_0^2z(-σ_0^2 z e^σ_0^2)/z σ_0^2 .

We can extract the maximum and minumum eigenvalue by finding the branch points of this function. It suffices to look for poles in the derivative of the numerator of G(z). Using r = -σ_0^2 z e^σ_0^2, eqn. (<ref>) and its total derivative with respect to z yields the following equation defining the locations of these poles,

    e^W_r(1 + W_r) = z e^z^2 ,

which is solved by

    W_r = W(e^1+σ_0^2 z)-1 ,

where W is the standard Lambert W function. Next we substite this relation into eqn. (<ref>); zeros in z then define the location of the branch points. Some straightforward algebra yields the maximum and minimum eigenvalue,

    λ_± = 1/2 e^-1/2σ_±^2(2 + σ_∓^2) ,  where  σ_±^2 = σ_0(σ_0±√(σ_0^2 + 4))




§ ORTHOGONAL WEIGHTS ARE REQUIRED FOR STABLE, UNIVERSAL LIMITING DISTRIBUTIONS


We work at criticality so χ = σ_w^2 μ_1 = 1. This implies that 

    μ_JJ^T    = m_1 = 1
    σ_JJ^T^2    = m_2 - m_1^2 = L (μ_2/μ_1^2 - 1 - s_1) .

Observe that Jensen's inequality requires that μ_2 ≥μ_1^2. If we require that σ_JJ^T^2 approach a constant as L→∞, we must have that,

    s_1≥ 0 .

Similarly, writing

    M_W^TW(z) = ∑_k=1^∞𝔪_k/z^k ,

we can relate σ_w and s_1 to 𝔪_1 and 𝔪_2. Specifically, evaluating the relation,

    M_W^TW^-1(z) = 1+z/z S_W^TW(z) ,

at z = M_W^TW(x), gives,

    x = 1 + M_W^TW(x)/M_W^TW(x) S_W^TW(M_W^TW(x)) .

Expanding this equation to second order gives,

    𝔪_1    = g^2
    𝔪_2    = g^2 𝔪_1 (1 - s_1) .

Finally we see that,

    σ_WW^T^2 = 𝔪_2 - 𝔪_1^2 =  -g^4 s_1 .

Positivity of variance gives s_1 ≤ 0, which, together with eqn. (<ref>) implies,

    s_1 = 0 .

Altogether we see that the variance of the distribution of eigenvalues of WW^T must be zero. Since its mean is equal to σ_w^2, we see that the only valid distribution for the eigenvalues of WW^T is a delta function peaked at σ_w^2, i.e. the distribution corresponding to the singular values of an orthogonal matrix scaled by σ_w.

