Free probability as a solution to the problem of tuning neural networks
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Gradient descent during the learning process of a neural network can be subjectto many instabilities.
The spectral density of the Jacobian is a key component foranalyzing stability.
Following the works of Pennington et al., such Jacobians aremodeled using free multiplicative convolutions from Free Probability Theory (FPT).
We present a reliable and very fast method for computing the associated spectraldensities, for given architecture and initialization.
This method has a controlledand proven convergence.
Our technique is based on an homotopy method: it is anadaptative Newton-Raphson scheme which chains basins of attraction.
In orderto demonstrate the relevance of our method we show that certain FPT metricscomputed before training are highly correlated to ﬁnal test accuracies – up to85%.
This nuances the idea that learning happens at the edge of chaos by givingevidence that a very desirable feature for neural networks is the hyperbolicity oftheir Jacobian at initialization.
Introduction
Neural network training and tuning can be wasteful in human and energy resources.
For example[21, Table 1] shows that a single GPU training can have a large carbon footprint.
Often, this isnothing compared to architecture search.
In this context, an obvious moonshot would be estimatingthe performance of an architecture before training.
This paper takes a more realistic and indirectapproach. We aim for a fast and reliable computational method for estimating stability before training– and stability is a good proxy for performance.
Framework: Consider a feed-forward network of depth L ∈ N , with L full-connected layers.
Foreach depth (cid:96) ∈ { 1 , 2 , . . .
, L } , the layer has activation vector x (cid:96) ∈ R N (cid:96) , where N (cid:96) is the currentwidth.
The vector x 0 ∈ R N 0 takes in the neural network’s input, while x L ∈ R N L gives theoutput.
The vector of widths is written N :=
( N 0 , N 1 , . .
. , N L ) and will appear in superscriptto indicate the dependence in any of the N (cid:96) ’s. The following recurrence relation holds betweenlayers x (cid:96) =
φ (cid:96) (cid:16) W ( N ) (cid:96) x (cid:96) − 1 + b ( N ) (cid:96) (cid:17) , where φ (cid:96) is a choice of non-linearity applied entry-wise, W ( N ) (cid:96) ∈
M N (cid:96) ,N (cid:96) − 1 ( R ) is a weight matrix and b ( N ) (cid:96) ∈ R N (cid:96) is the vector of biases.
We writeh (cid:96) : = W ( N ) (cid:96) x (cid:96) [Equation] +
[Equation] N
[Equation]
(cid:96) for the pre-activations. The Jacobian computed during back-propagation can be written explicitly by using the chain rule: [Equation] where D (cid:96) ’s are the diagonal matrices given by
[Equation]
Technically, a step of gradient descent updates weights and biases following
[Equation]
for each (cid:96) = 1 , . . .
, L .
Here α > 0 is the learning rate and L is the loss on a minibatch.
If theminibatch has size B ∈ N , and corresponds a small sample (( X i , Y i ) ; i = 1 , . .
, B ) of the dataset
,we have L = 1 B (cid:80) Bi
=1 d ( x L ( X i ) , Y i ) .
Here d is a real-valued distance or similarity function, the
X i ’s are the input vectors while the Y i ’s are the output vectors (e.g. labels in the case of classiﬁer,
Y i ≈
X i in the case of an autoencoder etc...).
The chain rule dictates:
[Equation]
where
[Equation]
Therefore, for the sake of simplicity, we shall focus on the Jacobian J ( N ) given in Eq.
(1) since it hasexactly the same form as the J ( N ) (cid:96) given in Eq.
The issue is that a large product of (even larger)matrices can easily become unstable.
If many singular values are (cid:28) 1 , we have gradient vanishing.
If many singular values are (cid:29) 1 , we have gradient explosion.
Such a transition can be referred to asthe edge of chaos.
Intuition.
This instability is easily understood thanks to the naive analogy with the one-dimensionalcase.
Indeed, the geometric progression q n with n → ∞ is the archetype of a long product and itconverges extremely fast, to either 0 if | q | < 0 or to ∞ if | q | > 1 .
A less naive intuition consists in recalling that mini-batch sampling in Eq.
(5) is noise source in
Stochastic Gradient Descent (SGD). It is fair to assume that ∂ L ∂x L has a Gaussian behavior with covariance proportional to I N L – either because of the Central Limit Theorem if B is large enough or after time averaging, because of the mixing properties of SGD.
Therefore, each gradient step α ∂ L ∂ ( W ( N ) (cid:96) ,b ( N ) (cid:96) ) in Eq.
(3) should approximately be a Gaussian vector with covariance proportional to:
[Equation]
∂ ( W ( N ) (cid:96) , b ( N ) (cid:96) ) .
Simplifying further, we see the importance of the spectrum of (cid:16) J ( N ) (cid:96) (cid:17) T J ( N ) (cid:96) for stability.
Basically,eigenvectors of (cid:16) J ( N ) (cid:96) (cid:17) T J ( N ) (cid:96) are the directions along which the one-dimensional intuition applies.
Randomness.
Starting from the pioneering works of Glorot and Bengio [8] on random initializations,it was suggested that the spectral properties of J ( N ) are an excellent indicator for stability andlearning performance.
In particular, an appropriate random initialization was suggested and sinceimplemented in all modern ML frameworks [17, 1].
We make classical choices of random initializations.
The biases b ( N ) (cid:96) are taken as random vectorswhich entries are centered i.i.d.
Gaussian random variables with standard deviation σ b (cid:96) .
For theweights, we will consider the following matrix ensembles: the [ W ( N ) (cid:96) ] i,j are drawn from i.i.d.centered random variables
Modeling spectrum thanks to Free Probability Theory (FPT).
Already, notice that the squareroot normalization suggested in [8] is the correct scale for the convergence of spectral measures in
FPT.
This ought to be recognized and provides a new justiﬁcation for that choice.
Now, followingthe works of Pennington et al.
[18], the tools of FPT can be used to more quantitatively analyze thesingular values of J ( N ) in the large width limit.
The large width limit is particularly attractive whenstudying large deep networks, especially because free probability appears at relatively small sizesthanks to strong concentration properties [12].
Indeed, the theory applies to random matrices of sizeas small as a hundred.
For the purposes of this paragraph, we restrict ourselves to square matrices and assume N (cid:96) = N for all (cid:96) = 1 , . . .
, L .
In fact, FPT is concerned with the behavior of spectral measures as N → ∞ .
For any diagonalizable A N ∈ M N ( R ) , the associated spectral measure on the real line is:
[Equation]
with the a ( N ) i ’s being the eigenvalues of A N .
For ease of notation, the spectrum of (squared) singular values is written ν
A ( N ) :
= µ ( A ( N ) )
T A ( N ) .
A fundamental assumption for invoking tools from Free
Probability Theory, is the assumption of asymptotic freeness .
Without deﬁning the notion, which can be found in [14], let us describe the important computation it allows, discovered in the seminal work of Voiculescu [23].
Given two sequences of square matrices A ( N ) , B ( N ) in M N ( R ) , with converging spectral measures:
[Equation]
we have that, under the assumption of asymptotic freeness lim N →∞ ν A
[Equation]
B [Equation] N ) =
ν
A [Equation] ν B , where 66 (cid:2) is a deterministic operation between measures called multiplicative free convolution.
The (cid:2) willbe detailed in Section 2.
The letter A (as well as B ) does not correspond to a limiting matrix but toan abstract operator, with associated spectral measure µ A and measure of squared singular values ν A .
For such limiting operators, we drop the superscript ( N ) .
Under suitable assumptions which are motivated and detailed later following the works of [18, 10,16 , 15, 7], for all (cid:96) = 1 , . . .
, L , the measures ν W ( N ) (cid:96) and ν D ( N ) (cid:96) will respectively converge to ν W (cid:96) andν D (cid:96) .
Again the W (cid:96) ’s and D (cid:96) ’s are abstract operators which only make sense in the inﬁnite widthregime.
In the limit, asymptotic freeness will also hold.
Therefore, we will see that the measure ofinterest is:
[Equation]
The goals of this paper are (1) To give a very fast and stable computation of ν J , in the more generalsetup of rectangular matrices (2) Empirically demonstrate that FTP metrics computed from ν J docorrelate to the ﬁnal test accuracy.
1.1 Contributions
We aim at streamlining the approach of Pennington et al.
by providing the tools for a systematic useof FPT.
The contributions of this paper can be categorized as:
• Theoretical:
In Pennington et al., a constant width is assumed.
We generalize the model to allowfor varying width proﬁles, which is more inline with design practices.
This requires us to develop arectangular multiplicative free convolution.
Then we propose a computational scheme for computing spectral densities named "
Newton
lilypads".
The method relies on adaptative inversions of S -transforms using the Newton-Raphson
algorithm.
If the Newton-Raphson scheme is only local, we achieve a global resolution by chaining basins of attractions, thanks to doubling strategies.
As such, we have theoretical guarantees for the convergence.
Interestingly, even in the FPT community, inverting S -transforms has been considered impossibleto realize in practice.
As we shall see, an S -transform is a holomorphic map with multi-valuedinverse.
In the words of [5, p.218], “the operations which are necessary in this method (the inversionof certain functions, the extension of certain analytic functions) are almost always impossible torealise practically.”
• Numerical:
This misconception led to the use of combinatorial methods based on moments, or
[Figure] ﬁxed-point algorithms via the subordination method [2, 13, 22].
In the ML community, Pennington et al. pioneered the application of FPT to the theoretical foundations of Machine Learning and [Figure] did not shy away from inverting S -transforms.
Their [18, Algorithm 1] is based on a generic root [Figure] ﬁnding procedure, and choosing the root closest to the one found for the problem with one less layer.
A major drawback of their method is that there is no guarantee of ﬁnding the correct value, [Figure] unlike our chaining which always chooses the correct branch.
[Figure] Not only Newton lilypads ﬁnds the correct root, but it is also an order of magnitude faster (Fig. 1). [Figure]
A few standard Cython optimizations allow to gain another order of magnitude, although this can [Figure] certainly be reﬁned
[Figure]
Corr. factor Value p -value Spearman 0. 3ePearson 0.
9eR0. N/A value is based on
hypothesis that and are
• Empirical: We demonstrate that FPT is not a mere abstract tool by designing a simple andunequivoqual experience. Fig.
2 shows the correlation between
the accuracy of several randomlygenerated architectures and the quantiles of ν J , after a ﬁxed number of epochs.
Remarkably, we ﬁnd that accuracy is very correlated to the higher quantiles.
This suggests thatspread-out spectral distributions ν J are more desirable, provided of course we avoid the vanishingand explosive regimes.
In the language of dynamical systems, we say that the Jacobian needs tobe hyperbolic i.e. with both contracting and expanding directions.
This considerably nuances theidea that learning happens at the edge of chaos.
A similar point was made in the conclusion of [6],using the language of hyperbolic attractors.
Also, larger neural networks obviously perform better. But smaller networks with more spread-outν J perform better than larger counterparts.
For reproducibility of the numerical and empirical aspects, a complete implementation is provided ina Github repository https://anonymous.4open.science/r/FreeNN-331E/
1.2 Structure of the paper
We start in Section 2 by stating facts from Free Probability Theory.
Most of it is available in theliterature, except the setup for the product of rectangular free matrices.
To the best of our knowledge,this is novel.
There, we establish in the rectangular setting an analogue of Eq.
(7) in Theorem 2.3.
In Section 3, we explain in detail the FPT model for random neural networks.
Then thanks to theresults of [16, 15] and our rectangular setting, we show that the spectral measure of the Jacobian
J ( N ) converges to ν J
and we encode the limit in explicit generating series in Theorem 3.1.
This giveshow ν J can theoretically be recovered.
Section 4 presents the numerical resolution which inverts the (multi-valued) generating series.
Bychaining different (local) basins of attractions, we obtain a global resolution method.
Our algorithmis detailed in Algorithm 1 and Theorem 4.1 states the theoretical guarantees.
Finally Section 5 presents the experiment leading to Fig.
More details are given in Appendix F,including comments on the benchmark of Fig.
Free Probability
2.1 Deﬁnitions and notations
Free Probability Theory provides a framework to analyze eigenvalues and singular values of large random matrices.
We now introduce various complex-analytic generating series which encode the measures and the basic operations on them.
First, the Cauchy-Stieltjes transform of µ , a probability measure on R + is:
[Equation]
where C ± := { z ∈ C | ± (cid:61) z > 0 } .
The transform G µ encodes the measure µ and reciprocally, themeasure can be recovered thanks to:
Lemma 2.1 (Cauchy-Stieltjes inversion formula – Theorem 6 in [14]) .
We have the weak convergence of probability measures:
[Equation]
The moment generating function is
[Equation]
where for all k ∈ N , m k ( µ ) :
= (cid:82) R x k µ ( dx ) is the k -th moment of µ .
For µ (cid:54) =
δ 0 , M µ is invertiblein the neighborhood of ∞ and the inverse is denoted by M (cid:104)− 1 (cid:105) µ .
The S-transform of µ is deﬁnedas S µ ( m ) =[-
-]
[Equation] m [Equation] ( m ) , and is analytic in a neighborhood of m = 0 .
Furthermore, the variablez will always denote an element of C + , while the variables g and m will denote elements in theimage of G µ and M µ .
For a diagonalizable matrix A ( N ) ∈ M N ( R ) , we write S A ( N ) :
= S µ A ( N ) ,
G A ( N ) :
= G µ A ( N ) , M A ( N ) :
=
M µ A ( N ) .
A landmark result in the ﬁeld introduces free multiplicative convolution in a natural way, and showsthat this operation is linearized by the S -transform:
Theorem 2.2 (Voiculescu, [23]) .
Consider two sequences of positive matrices, each element in
M N ( R )
[Equation]
such that:
[Equation]
Under the assumption of asymptotic freeness for A ( N ) and B ( N ) , there exists a deterministic probability measure µ A (cid:2) µ B such that lim N →∞ µ A ( N ) B ( N ) =
µ A (cid:2) µ B .
The operation (cid:2) is themultiplicative free convolution.
Moreover
[Equation]
This convergence akin to a law of large numbers is the key ingredient which allows to build thedeterministic model for the back-propagation of gradients in Eq.
2.2 Product of rectangular free matrices
As a generalization of Eq.
(9) to rectangular matrices, we state:
Theorem 2.3.
Let ( p N ) N ≥ 1 , ( q N ) N ≥ 1 , ( r N ) N ≥ 1 , be three sequences of integers satisfying
[Equation]
Consider for all N ≥ 1 let A ( N ) , B ( N ) be random matrices of respective sizes p N × q N and q N × r N such that the (squared) singular laws of A ( N ) , B ( N ) converge weakly.
Assuming that B ( N ) (cid:0) B ( N ) (cid:1) T and (cid:0) A ( N ) (cid:1)
T A ( N ) are asymptotically free, we have that in the limit N → [Equation]
[Equation]
Proof.
See the appendix, Subsection C.
Implicitly this deﬁnes a rectangular multiplicative free convolution, which could be denoted (cid:2) c inthe spirit of the rectangular free additive convolution [4].
But, in the current setting, this is not a goodidea.
Indeed, if one deﬁnes µ 1 (cid:2) c µ 2 as the measure whose S -transform is S µ 1 ( c · ) S µ 2 , then a quickcomputation shows that (cid:2) c is not associative i.e. for a triplet ( µ 1 , µ 2 , µ 3 ) of probability measuresand a pair ( c 1 , c 2 )
∈ R ∗ + × R ∗ + , we generically have µ 1 (cid:2)
c 1 ( µ 2 (cid:2) c 2 µ 3 ) (cid:54) = ( µ 1 (cid:2) c 1 µ 2 ) (cid:2) c 2 µ 3 .
A better idea is to treat the dimension ratio c as part of the data via a central extension:
Deﬁnition 2.4.
On the set of pairs ( µ, c ) such that µ is a probability measure on R + and c ∈ R ∗ + , deﬁne the operation (cid:2) as:
[Equation]
where ν is the unique probability measure such that S ν
= S µ 1 ( c 2 · ) S µ 2 .
This extends the classicaldeﬁnition as the usual free convolution is recovered with ( µ 1 , 1) (cid:2) ( µ 2 , 1) :
= ( µ 1 (cid:2) µ 2 , 1) .
Such an operation is associative and will allow a neat formulation of the measure of interest in theupcoming Theorem 3.1, entirely analogous to Eq.
Theoretical resolution of the model
Width proﬁle: Pennington et al.
[18] consider N (cid:96) = N for (cid:96) = 1 , 2 , . . .
, N .
Here, we consider thatthe width of layers is not constant across layers, which is mostly the case in practice.
Indeed, modernarchitectures typically have very sophisticated topologies with layers varying in widths.
Let us assume that we are in the inﬁnite width regime in the sense that N (cid:96) → ∞ , for all (cid:96) =0 , 1 , 2 . . .
, L with: N (cid:96) − 1 N (cid:96) −→ N →∞ λ (cid:96) ≥ 0 .
And let us denote Λ (cid:96) :
= lim N →∞ N 0 N (cid:96) =
(cid:81) (cid:96)k =1 λ k , with theconvention Λ 0 = 1 .
FPT limits: Let N be a standard Gaussian random variable on the real line P ( N ∈ dx )
=
e − x 22 √ 2 π dx .
Here D ( N ) (cid:96) is diagonal with entries φ (cid:48) (cid:96) ( [ h (cid:96) ] i ) (see Eq. (2)), and the pre-activations
[Equation]
clearly depend on the previous layers.
Because of this lack of independence, the standard results of
FPT cannot be applied directly i.e. asymptotic freeness does not obviously hold.
This is an importantsubtlety that is addressed in the upcoming Theorem 3.1.
Based on an information propagationargument, the papers [19, 20] argue that the entries of h (cid:96) behave as the i.i.d.
samples of a Gaussiandistribution with zero mean and variance q (cid:96) .
A basic law of large numbers applied to Eq.
(2) gives alimit for the empirical measure µ D (cid:96) = lim N (cid:96) [Equation] µ D
[Equation]
N )
[Equation]
q (cid:96) N (cid:17) .
Also the recurrence forthe variance is:
[Equation]
10) with initial condition q 1 =
[Equation]
b 1 .
Recently Pastur et al.
completed this heuristic thanks to a swapping trick – see [16, Lemma 3.3] and[15, Remark 3.4].
They proved that, regarding the asymptotical spectral properties of J ( N ) , one canreplace each D ( N ) (cid:96) by a diagonal matrix with independent Gaussian entries (cid:112) q (cid:96) N independent fromthe rest.
By combining that result and our rectangular setup from Section 2, we ﬁnd:
Theorem 3.1.
In terms of the rectangular multiplicative free convolution, the measure of (squared)singular values of J ( N ) converges to
[Equation]
Moreover, the S-transform of J T J in the inﬁnite width regime veriﬁes
[Equation]
In particular, under the assumption that the entries of W (cid:96) are i.i.d. :
[Equation]
m ) .
Proof.
See the appendix, Subsection C.2.
Master equation: In the end, we only need to ﬁx width ratios and non-linearities to form M (cid:104)− 1 (cid:105) J ( m ) ,and get the master equation which we solve numerically thanks to an adaptive Newton-Raphsonscheme.
The non-linearities ReLu, Hard Tanh and Hard Sine yield explicit formulas, which can befound in Table 1 of the appendix.
If W (cid:96) has i.i.d.
entries, one ﬁnds the explicit master equation:
[Equation]
[Equation]
where c (cid:96) = 12 when φ (cid:96) is ReLu, c (cid:96) = C (cid:96) = P
[Equation]
if φ (cid:96) is Hard Tanh and c (cid:96) = 1 if φ (cid:96)is Hard Sine.
Numerical resolution
Here we describe the numerical scheme aimed at computing the spectral density of J T J in Eq.
We use the following steps to compute the spectral density at a ﬁxed x ∈ R + :
• Because of the
Cauchy-Stieltjes inversion formula
given in Lemma pick a small y > 0 inorder to compute: − 1 π (cid:61) G J T J ( z
= x + iy ) . The smaller the better, and in practice our methodworks for up to y = 10 − 9 .
Figure 3 shows the same target distribution but convolved with various
Cauchy distributions y C where y ∈ { 1 , 10 − 1 , 10 − 4 } .
This corresponds to computing the density − 1 π (cid:61) G µ ( + iy ) for different y
’s.
• Because of Eq.
(8), we equivalently need to compute M J T J ( z ) .•
M (cid:104)− 1 (cid:105) J T J ( m ) is available thanks to the master equation in Theorem 3.1.
Therefore, we need to
invert
m (cid:55)→ M (cid:104)− 1 (cid:105) J T J ( m ) .
This step is the crucial part: M (cid:104)− 1 (cid:105) J T J is multi-valued and one needs to choose thecorrect branch.
4.1 Initial setup
We ﬁrst use the classical Newton-Raphson scheme to invert the equation z
= f ( m ) where z ∈ C + isﬁxed and f is rational.
A neat trick which leverages the fact that f is rational and that z ∈ C + is todeﬁne:
[Equation]
As such, we have z = f ( m )
=
P ( m ) Q ( m ) ⇐⇒ ϕ
z ( m ) = 0 .
There are several advantages of doingthat: (1) Inversion is recast into ﬁnding the zero of a polynomial function.
(2) Since we havelim z →∞ M ( z ) = 0 , if z is large in modulus, m = 0 is a natural starting point for the algorithm whenz is large.
It is well-known that the Newton-Raphson scheme fails unless the initial guess m 0 ∈ C belongs to abasin of attraction for the method.
And, provided such a guarantee, the Newton-Raphson scheme isexceptionally fast with a quadratic convergence speed.
Kantorovich’s seminal work in 1948 providessuch a guarantee locally.
For the reader’s convenience, we give in Appendix D the pseudo-code forthe Newton-Raphson algorithm (Algorithm 2), as well as a reference for the optimal form of the
Kantorovich criterion (Theorem D.1).
Therefore, we assume that we have at our disposal a function ( z, m ) (cid:55)→ IS _ IN _ BASIN ( z, m ) whichindicates if the Kantorovich criterion is satisﬁed for ϕ z at any m ∈ C .
It is particularly easy toprogram with ϕ z polynomial.
Algorithm Name: NEWTON _ LILYPADS
Input: Image value: z objective ∈ C + , (Optional) Proxy: ( z 0 , m 0 )
∈ C + × C . Output: M ( z objective )
# Find a proxy ( z 0 , m 0 = 0) using a doubling strategy, if None given if ( z 0 , m 0 ) is None then m ←z ← z objective while not IS _ IN _ BASIN (z, m) do z ← z + i (cid:61) ( z ) = (cid:60) ( z ) + i 2 (cid:61) ( z ) # Double imaginary part end while m ← NEWTON _ RAPHSON ( z, Guess = m ) else ( z, m ) ← ( z 0 , m 0 ) end if # Starts heading towards z objective using dichotomy while | z objective − z | > 0 do ∆ z ← z objective − z while not IS _ IN _ BASIN ( z + ∆ z , m) do ∆ z ← 0 . 5 ∗ ∆ z end while z ← z + ∆ z m ← NEWTON _ RAPHSON ( z, Guess = m ) end while return m
Now we have all the (local) ingredients in order to describe a global strategy which solves in m ∈ Cthe equation ϕ z ( m ) = 0 .
First, one has to notice that this problem is part of a family parametrized by z ∈ C + .
And the solutionis m ≈ 0 for z large.
Therefore, one can ﬁnd a proxy solution for z ∈ C + high enough.
This is donethanks to a doubling strategy until a basin of attraction is reached.
Second, if a proxy ( z, m ) is available, we can use the Newton-Raphson algorithm to ﬁnd a solution( z + ∆ z, m + ∆ m ) starting from m .
To do so, we need ∆ z small enough.
This on the other hand isdone by dichotomy.
Tying the pieces together allows to chain the different basins of attraction and leads to Algorithm 1.
Notice that in the description of the algorithm, we chose to make implicit the dependence in thefunction f , since it is only passed along as a parameter.
Technically, f is a parameter for all threefunctions NEWTON _ RAPHSON , IS _ IN _ BASIN , NEWTON _ LILYPADS .
The discussion leading to thisalgorithm, combined with the Kantorovich criterion yields:
Theorem 4.1.
Given f : m (cid:55)→ M (cid:104)− 1 (cid:105) ( m ) and z ∈ C + , Algorithm 1 has guaranteed convergence.
Moreover it returns m = M ( z ) i.e. the (unique) holomorphic extension of the inverse of f in theneighborhood of 0 .
On the experiment
To leverage the numerical scheme, we designed the following experiment whose results are reportedon Fig. 2.
Consider a classical Multi-Layer Perceptron (MLP) with L = 4 layers, feed-forward andfully-connected with ReLu non-linearities.
The task at hand is classiﬁcation on the classical MNISTdataset.
The MLP’s architecture is determined by the vector λ = ( λ 0 , λ 1 , . .
λ L ) while the gains atinitialization are determined by the vector σ = ( σ 1 , σ 2 , . . .
, σ L ) .
By randomly sampling the vector λ , we explore the space of possible architectures.
In other tohave balanced architectures, we chose independent λ
i ’s with E ( λ i )
= 1 .
Likewise, we also sampledifferent possible gains.
Hence we ﬁnd ourselves with several MLPs architectures each with it’sunique initialization.
The spectral distributions are computed thanks to our Algorithm 1.
As shown in Fig.
2, the parameter count varies by orders of magnitude.
To guarantee convergencedespite the changing architectures, we intentionally chose a low learning rate and 50 epochs oftraining.
For 200 MLPs, the experiment takes 10 hours of compute on a consumer GPU (RTX2080Ti).
This is to be contrasted with less than one minute of CPU compute for the spectral measureν J for all the MLPs.
To control for any stochastic variability, we also trained multiple instances for each MLP, and theshown results are averages between a few runs.
Finally we calculate the correlations between theaccuracy on the test set and the 90th percentile of the spectral distribution ν J .
A posteriori, the FPTregime is justiﬁed beyond initialization because the results are coherent despite various width scales.
Conclusion
In summary, this paper developed FPT in the rectangular setup which is the main tool for a metamodelfor the stability of neural network.
Then we gave a method for the numerical resolution both fastand with theoretical guarantees.
Finally, conﬁrming the initial claim that a metamodel focused onstability yields a good proxy for performance, we empirically demonstrated that the accuracy of afeed-forward neural network after training is highly correlated to higher quantiles of the theoreticallycomputed spectral distribution.
Regarding the importance of hyperbolicity, we surmise that the presence of a very few large singularvalues allows SGD (Stochastic Gradient Descent) to escape local minima without compromisingoverall stability.
Additionally to this conjecture, many other challenges remain in order to build a moreversatile metamodel for stability: for example accomodating for skip-connections and convolutionallayers fall outside the scope of the theory.
A Appendix: More ﬁgures
In this section, we record non-essential yet pedagogical ﬁgures.
B Appendix: Table of non-linearities
Let us now collect the various required formulas, and specialize them to a selection of non-linearities:
ReLu, Hard Tanh and Hard Sine.
These non-linearities are very tractable hence the choice.
Asdiscussed in the introduction, the empirical distribution of D ( N ) (cid:96) converges to the law of φ (cid:48) l (cid:16)(cid:112) q (cid:96) N (cid:17) .
From this observation was deduced in [18] the following formula:
[Equation]
Linear
ReLu
Hard Tanh
Hard sine / Triangle
[Equation]
[Equation]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table] In Table 1.1, the reader shall ﬁnd the formulae of M D 2 (cid:96) , S D 2 (cid:96) and the recurrence relation given by f (cid:96)where
[Equation]
(cid:96)
[Equation]
q [Equation] − 1 ) + σ
[Equation]
b
[Equation]
, with g (cid:96) ( q ) :
= E (cid:104) φ (cid:96) ( √ q N ) 2 (cid:105) .
and
[Equation]
Proof of formulas in Table 1: ReLu and Hard Tanh have already been computed in [18].
Hard Sine: This is the most tricky formula to establish.
If φ ( x ) = [-
-]2 π arcsin ◦ sin (cid:0) π 2 x (cid:1) and (cid:98) f ( ξ ) = 398 (cid:82) dx f ( x ) e iξx is the Fourier transform on f , then the application of the Plancherel formula yields:
[Equation]
But in term of Fourier series:
[Equation]
as φ 2 ( x ) = x 2 on [ − 1 , 1] and extended in order to become 2 -periodic.
In terms of Schwartz distributions:
[Equation]
Hence:
[Equation]
[Table]
We conclude by computing the Fourier coefﬁcients of φ 2 .
[Equation] [Equation]
In the end:
[Equation]
C Appendix: ProofsC.1 Proof of Theorem 2.3
For k ∈ N , we have :
Tr
[Equation]
N
[Equation]
(cid:16) A ( N )
[Equation]
T AB (cid:19) k (cid:35) =
Tr
[Equation]
( B ( N ) (cid:16) B ( N ) (cid:17) T (cid:16) A ( N ) (cid:17) T A )
k (cid:21) .
As (cid:0) A ( N ) B ( N ) (cid:1) T A ( N ) B ( N ) ∈ M r N ( C ) and B ( N ) (cid:0) B ( N ) (cid:1) T (cid:0) A ( N ) (cid:1) T A ( N ) ∈
M q N ( C ) , this showsthat
[Equation]
and then
[Equation]
Consequently,
[Equation]
As (cid:0) A ( N ) (cid:1) T A ( N ) and B ( N ) (cid:0) ( N ) (cid:1) T are asymptotically free, taking the limit N → + ∞ and applying
Voiculescu’s Theorem 2.2 , we get
[Equation]
Moreover, the above equality is true replacing A with the identity I .
S I ( m ) =
1 yields:
[Equation]
Finally we have
[Equation]
This concludes the proof.
C.2 Proof of Theorem 3.1
We assume that W (cid:96) have i.i.d.
entries.
Thanks to a swapping trick justiﬁed in [15], we can assumethat the matrices D (cid:96) have i.i.d.
entries independent from the rest of the network, distributed asφ (cid:48) (cid:96) ( √ q (cid:96) N ) .
Notice that we can also replace the D (cid:96) ’s by deterministic matrices that use the quantilesof the same distribution.
This together with standard results from FPT such as [14] gives asymptoticfreeness – see [3] for a more general result reﬂecting the current state of the art.
Therefore, we canapply Theorem 2.3.
Starting from Eq.
(1) , we get in the inﬁnite width regime and by induction:
[Equation]
with the convention Λ 0 = (cid:81)
k =1 λ k = 1 .
Under the asumption that the entries of W (cid:96) are i.i.d., the Marcenko-Pastur Theorem gives
[Equation]
which leads to
[Equation]
We thus have
[Equation]
D Appendix: On the classical Newton-Raphson scheme
Algorithm scheme for a function f
Name: NEWTON _ RAPHSON
Input:
Numerical precision: ε > 0 (Default: 10 − 12 ),
Image value: z ∈ C + ,
Polynomials: P , Q such that f = PQ , (Optional) Guess: m 0 ∈ C , (Default: m 0 = 0 ).
Code: m ← mwhile True do value ← ϕ z ( m )
# See Eq.
(14) if | value | < ε then return m end if grad ← ϕ (cid:48) z ( m ) m ← m − value/grad end while
Here we give the optimal Kantorovich criterion from [9] adapted to this paper.
Fix z ∈ C + and recallthat ϕ z in Eq.
(14) is the map whose zero we want to ﬁnd.
Theorem D.1 (Kantorovich’s criterion, [11]) .
Consider a starting point m 0 ∈ C , and deﬁne:
[Equation]
If the starting point satisﬁes h := δκλ < 12 , where
[Equation]
Then, the Newton-Raphson scheme, starting from m 0 converges to m ∗ such that ϕ z ( m ∗ ) = 0 .
Furthermore, the convergence at each step is at least quadratic.
E Appendix: Moments of J
We can reach an early understanding of the behavior of J ’s singular values by computing mean and variance.
For ease of notation, we write:
[Equation]
for any operator A , which admits a measure of singular values.
We have under the assumptions thatthe entries of W (cid:96) are i.i.d. :
[Equation]
Under the asumption that W T(cid:96) W
(cid:96) =
σ W (cid:96) I N (cid:96) − 1 , we ﬁnd the same m ( s )
( J ) and :
[Equation]
These formulas need to be interpreted:
• Variance grows with L , showing increased instability with depth.
• Larger layers, relative to N 0 , give larger Λ (cid:96) ’s and thus the same effect.
Proof: Computations of moments.
The following remark is useful in the computation of moments.
Remark E.1 (Moments) .
At the neighborhood of z ∼ ∞ :
[Equation]
By inversion, at the neighborhood of m ∼ 0 :
[Equation]
Hence:
[Equation]
Thanks to this
[Equation]
we can prove Eq.
(18) and (19) [Equation]
By Remark E [Equation] 1 and Theorem 3.
[Equation]
we have as m → 0 :
[Equation]
Identifying the ﬁrst order term, one ﬁnds indeed Eq.
Continuing the previous computation:
[Equation]
Applying Remark E.1 again for S J T J , we get :
[Equation]
which is equivalent to Eq.
[Equation]
We conclude by specializing to classical weight distributions.
Under the assumption that the entriesof W (cid:96) are i
[Equation] i.d. we have m ( s ) {+
+}
( W (cid:96) ) = σ 2 W (cid:96) and m ( s ) 2 ( W (cid:96) ) =
σ 4 W (cid:96) (1 + λ (cid:96) ) which gives
[Equation]
[Table]
F Appendix: More details on the benchmarks and the experiment
Recall that we have provided an anonymized Github repository at the address:https://anonymous.4open.science/r/FreeNN-331E/
F.
On the benchmarks of Fig.
The ﬁgure presents the computational time required for computing the density of ν J .
Let us make thefollowing comments:
• Pennington and al. ’s
algorithm has been implemented using a native root ﬁnding procedure.
Assuch, it is much more optimized than our code.
• A closer examination of the timings shows that Newton lilypads scales sublinearly with the numberof required points.
This is easily understood by the fact that smaller N requires the computation ofmore basins of attraction per point. • Also, the Monte-Carlo method used matrices of size n = 3000 .
Not only Monte-Carlo is imprecisebecause of the noise, but its performance scales very poorly with n since one needs to diagonalize ever larger matrices.
Of course, Monte-Carlo remains the easiest method to implement.
But
F.2 Description of the experiment of Fig
Recall that we considered a classical Multi-Layer Perceptron (MLP) with L = 4 layers, feed-forwardand fully-connected with ReLu non-linearities.
The MLP’s architecture is determined by the vectorλ
= ( λ 0 , λ 1 , . .
. , λ L ) .
The initialization follows the Xavier normal initialization [8] as implementedin Pytorch [17].
Thus the full vector of variances of this initialization σ = ( σ 1 , σ 2 , . . .
, σ L ) isdetermined up to a multiplicative factor called the gain.
First, we sample random architectures.
We chose
• the λ i to be i i
d. and uniform on { 14 , 13 , 12 , 23 , 1 . 0 , 32 , 2 , 3 , 4 } .
As such the mean is E λ i
= 1 in order to obtain relatively balanced architectures. • the single gain is taken as a uniform random variable on { 14 , 12 , 1 , 2 , 4 } .
Secondly, we compute the spectral distributions predicted by FPT for each random architecture.
Some architectures are very unbalanced and are discarded.
Thirdly, we train more than 200 instances of MLPs and record the learning curves and accuracies.
In the end, by considering FPT quantiles and the ﬁnal test accuracy for each MLP, we obtain datawhich amenable to a bivariate statistical analysis. More comments are in order:
• Other initializations than normal were tested, without a noticeable difference. In any case, FPTpredicts the same outcomes.• It was our impression that more stable architectures train better regardless of the SGD ﬂavor. Therefore we ended up using a vanilla SGD without momentum.•
Due to the sampling procedure, the number of parameters for our MLPs can differ by orders ofmagnitudes. As shown in Fig. 2, the range is from 10 4 to 10 8 . The largest ones do beneﬁt fromGPU acceleration.
F.3 Final remarks
Our sampling procedure does not consider large ﬂuctuations in the λ
i ’s and focuses on balancedarchitectures.
Likewise, the gains at initialization do not deviate much from the classical [8].
It isimportant to recognize that without the insight of FPT, the scalings applied to such initializations arealready normalized so that spectral measures do converge.
As such, we never encounter truly problematic gradient vanishing or gradient explosion, whichcompletely sabotage the convergence of the neural network.
Our reﬁned FPT metrics are arguably"second order corrections".
Nevertheless, it is surprising the 90th percentile in Fig. 2 highly correlatesto the ﬁnal test accuracy after training.
In the end, tuning FPT metrics does not amount to secondorder corrections.