Free probability as a solution to the problem of tuning neural networks
Anonymous Author(s)
AffiliationAddress 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 order todemonstrate the relevance of our method we show that the relevant FPT metricscomputed before training are highly correlated to final test accuracies – up to 85%.
We also nuance 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[23, Table 1] show that a single GPU training can have high energetic costs.
Often, this is nothingcompared to architecture search.
In this context, an obvious moonshot is estimating the performanceof an architecture before training.
The goal of this paper is more realistic: providing a fast and reliablecomputational method for estimating stability before training – and stability is a good proxy forperformance [12].
Framework: Consider a feed-forward network of depth L ∈ N , with L full-connected layers.
Foreach depth ℓ ∈ { 1 , 2 , . . .
, L } , the layer has activation vector x ℓ ∈ R N ℓ , where N ℓ 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 ℓ ’s. The following recurrence relation holds betweenlayers x ℓ = ϕ ℓ (cid:16) W ( N
[Equation]
(cid:17) , where ϕ ℓ is a choice of non-linearity applied entry-wise,
W ( N ) ℓ ∈ M N ℓ ,N ℓ − 1 ( R ) is a weight matrix and b ( N ) ℓ ∈ R N ℓ is the vector of biases.
We writeh ℓ := W ( N ) ℓ x ℓ − 1 + b ( N ) ℓ for the pre-activations.
The Jacobian computed during back-propagation can be written explicitly by using the chain rule:
[Equation]
where D ℓ ’s are the diagonal matrices given by
[Equation]
Technically, a step of gradient descent updates weights and biases following
[Equation]
for each ℓ = 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 classifier,
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.1) since ithas exactly the same form as the J ( N ) ℓ given in Eq.
The issue is that a large product of (evenlarger) matrices can easily become unstable.
If many singular values are ≪ 1 , we have gradientvanishing.
If many singular values are ≫ 1 , we have gradient explosion.
Such a transition can bereferred to as the edge of chaos [26, 27].
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 observing that mini-batch sampling in Eq.
(1.5) is very noisy.
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 [6, 8].
Therefore, each gradient step α
[Equation]
b
[Equation] [Equation]
) in Eq.
(1.3) is approximately a Gaussian vector with covariance proportional to:
[Equation]
Simplifying further, we see the importance of the spectrum of (cid:16) J ( N ) ℓ (cid:17) T J ( N ) ℓ for stability.
Basically,eigenvectors of (cid:16) J ( N ) ℓ (cid:17) T J ( N ) ℓ are the directions along which the one-dimensional intuition applies.
Randomness.
Starting from the pioneering works of Glorot and Bengio [9] 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 [19, 1].
We make classical choices of random initializations.
The biases b ( N ) ℓ are taken as random vectorswhich entries are centered i.i.d.
Gaussian random variables with standard deviation σ b ℓ .
For theweights, we will consider the following matrix ensembles: the [ W ( N ) ℓ ] i,j are drawn from i.i.d.centered random variables with variance σ 2 W ℓ /N ℓ and finite fourth moment as in [17].
Modeling spectrum thanks to Free Probability Theory.
Now, following the works of Pennington etal.
[20], the tools of Free Probability Theory (FPT) can be used to quantitatively analyze the singular values of J ( N ) in the large width limit.
The large width limit is particularly attractive when studyinglarge deep networks, especially because free probability appears at relatively small sizes because ofstrong concentration properties [14].
Indeed, random matrices of size 100 exhibit freeness.
For the purposes of this paragraph, we restrict ourselves to square matrices and assume N ℓ = N for all ℓ = 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 defining the notion, which can be found in [16], let us describe the important computation it allows, discovered in the seminal work of Voiculescu [25].
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 →∞ ν
[Equation]
=
ν
A ⊠ ν B , where 63 ⊠ is a deterministic operation between measures called multiplicative free convolution.
The ⊠ 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 [20, 11,18, 17, 7], for all ℓ = 1 , . . .
, L , the measures ν W ( N ) ℓ and ν D ( N ) ℓ will respectively converge to ν W ℓ andν D ℓ .
Again the W ℓ ’s and D ℓ ’s are abstract operators which only make sense in the infinite 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 final 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 profiles, 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 chainingbasins of attractions, thanks to doubling strategies.
As such, we have theoretical guarantees for theconvergence.
Interestingly,
even in the FPT community inverting S -transforms has been considered impossible
to
realize
in practice.
As
we shall see
, an
S -
transform is a holomorphic map with multi-valuedinverse.
In the words of [5,
p
], “the operations which
are
necessary in this
method
(the inversion
of certain functions, the extension of certain
analytic functions
) are almost always impossible to
realise practically.”
• Numerical:
This misconception led to the use of combinatorial methods based on moments, orfixed-
point algorithms via the subordination method [2, 15, 24].
In the ML community, Penningtonet al. pioneered the application of FPT to the theoretical foundations of Machine Learning anddid not shy away from inverting S -transforms.
Their [20, Algorithm 1] is based on a generic root finding procedure, and choosing the root closest to the one found for the problem with one lesslayer.
A major drawback of this method is that there is no guarantee to find the correct value, unlikeour chaining which always chooses the correct branch.
Not only Newton lilypads has theoretical guarantees of convergence, but it is also an order ofmagnitude faster (Fig. 1.1).
A few standard Cython optimizations allow to gain another order ofmagnitude, although this can certainly be refined.
[Figure]
• Empirical: We analyze the correlation between the test accuracy of several randomly generated
MLP architectures and the quantiles of ν J , after a fixed number of epochs.
The same architectures
were independently trained on the MNIST, FashionMNIST and CIFAR10 datasets.
We find that
accuracy is strongly correlated to FTP quantiles of ν J (see Table 1.1).
Remarkably, the correlation
is almost entirely captured by the higher quantiles – see Table A.1 individual R factors.
Scatter
plots showing the distribution of test accuracy, log 10 of the 90th quantile and number of parameters
can be seen in Fig.
A.3.
Interestingly, smaller networks can perform better than larger one provided
they have more spread-out ν J distribution.
This suggests that spread-out spectral distributions
ν J are more desirable, provided of course we avoid the vanishing and explosive regimes.
In the
language of dynamical systems, we say that the Jacobian needs to be hyperbolic i.e. with both
contracting and expanding directions.
This considerably nuances the idea that learning happens
at the edge of chaos.
A similar point was made in the conclusion of [6], using the language of
hyperbolic
attractors.
[Table]
All quantiles of ν J
[Table] [Table]
Github repository https://anonymous.4open.science/r/FreeNN-7ED4/
1.2 Structure of the paper
We start in Section 2 by stating facts from Free Probability Theory.
Most of it is available inthe literature, except the definition of the product of rectangular free matrices.
To the best of ourknowledge, this is novel.
There, we establish in the rectangular setting an analogue of Eq.
(1.7) in
Theorem 2.3.
In Section 3, we explain in detail the FPT model for random neural networks.
Then thanks to theresults of [18, 17] 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 Table 1.1.
More details are given in Appendix F,including comments on the benchmark of Fig.
Free Probability2.1 Definitions 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 | ± ℑ 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 [16]) .
We have the weak convergence of probability measures:
[Equation]
The moment generating function is
[Equation]
zG
[Equation]
−
[Equation]
(cid:88) k
[Equation]
m
[Equation] [Equation]
, (2.1) where for all k ∈ N , m k ( µ ) :
= (cid:82) R x k µ ( dx ) is the k -th moment of µ .
For µ ̸ = δ 0 , M µ is invertiblein the neighborhood of ∞ and the inverse is denoted by M ⟨− 1 ⟩ µ .
The S-transform of µ is definedas S µ ( m ) =
+ m mM ⟨− 1 ⟩ µ ( m ) , and is analytic in a neighborhood of m =
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 field introduces free multiplicative convolution in a natural way, and showsthat this operation is linearized by the S -transform:
Theorem 2.2 (Voiculescu, [25]) .
Consider two sequences of positive matrices, each element in
[Equation]
such that:
[Equation]
Under the assumption of asymptotic freeness for A ( N ) and B ( N ) , there exists a deterministic probability measure µ A ⊠ µ B such that lim N →∞ µ A ( N ) B ( N ) =
µ A ⊠ µ B .
The operation ⊠ 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.
(2.2) 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]
Proof.
See the appendix, Subsection C.
Implicitly this defines a rectangular multiplicative free convolution, which could be denoted ⊠ c in the spirit of the rectangular free additive convolution [4].
But, in the current setting, this is not a good idea.
Indeed, if one defines µ 1 ⊠ c µ 2 as the measure whose S -transform is S µ 1 ( c · ) S µ 2 , then a quick computation shows that ⊠ c is not associative, i.e., for a triplet ( µ 1 , µ 2 , µ 3 ) of probability measures and a pair ( c 1 , c 2 )
∈ R ∗ + × R ∗ + , we generically have:
[Equation]
A better idea is to treat the dimension ratio c as part of the data via a central extension:
Definition 2.4.
On the set of pairs ( µ, c ) such that µ is a probability measure on R + and c ∈ R ∗ + , define the operation ⊠ as:
[Equation]
where ν is the unique probability measure such that S ν
= S µ 1 ( c 2 · ) S µ 2 .
This extends the classicaldefinition as the usual free convolution is recovered with ( µ 1 , 1) ⊠ ( µ 2 , 1) :
= ( µ 1 ⊠ µ 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 profile: Pennington et al.
[20] consider N ℓ = N for ℓ = 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 infinite width regime in the sense that N ℓ → ∞ , for all ℓ =0 , 1 , 2 . . .
, L with: N ℓ − 1 N ℓ −→ N →∞ λ ℓ > 0 .
And let us denote Λ ℓ :=
lim N →∞ N 0 N ℓ = (cid:81) ℓ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 ) ℓ is diagonal with entries ϕ ′ ℓ ([ h ℓ
[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 [21, 22] argue that the entries of h ℓ behave as the i.i.d.
samples of a Gaussiandistribution with zero mean and variance q ℓ .
A basic law of large numbers applied to Eq.
(1.2) givesa limit for the empirical measure µ D ℓ = lim N ℓ →∞ µ D ( N ) ℓ
[Equation]
N
[Equation]
Also the recurrence forthe variance is:
[Equation]
with initial condition q 1 = σ 2 W 1 N 1 (cid:80) N 1 i =1 ( x i
[Equation]
2 + σ 2 b
[Equation]
Recently Pastur et al.
completed this heuristic thanks to a swapping trick – see [18, Lemma 3.3] and[17, Remark 3.4].
They proved that, regarding the asymptotical spectral properties of J ( N ) , one canreplace each D ( N ) ℓ by a diagonal matrix with independent Gaussian entries (cid:112) q ℓ N independent fromthe rest.
In that setting, one can apply the results on products of asymptotically free matrices whichwere given in Section 2.
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 infinite width regime verifies
[Equation]
In particular, under the assumption that the entries of W ℓ are i.i.d. :
[Equation]
Proof.
See the appendix, Subsection C.2.
Master equation: In the end, we only need to fix width ratios and non-linearities to form M ⟨− 1 ⟩ J T 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 B.1 of the appendix.
If W ℓ has i.i.d.
entries, one finds the explicit master equation:
[Equation]
where c ℓ = 12 when ϕ ℓ is ReLU, c ℓ = C ℓ = P
[Equation]
if ϕ ℓ is Hard Tanh and c ℓ = 1 if ϕ ℓ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 fixed x ∈ R + :
• Because of the Cauchy
-Stieltjes inversion formula given in Lemma 2.1, pick a small y > 0 in orderto compute: − 1 π ℑ G J T J ( z = x + iy ) .
The smaller the better, and in practice our method worksfor up to y = 10 − 9 .
Figure A.1 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 π ℑ G µ ( · + iy ) for different y ’s.
• Because of Eq.
(2.1), we equivalently need to compute M J T J ( z ) .•
M ⟨− 1 ⟩ J T J ( m ) is available thanks to the master equation in Theorem 3.1.
Therefore, we need to invertm (cid:55)→ M ⟨− 1 ⟩ J T J ( m ) .
This step is the crucial part: M ⟨− 1 ⟩ J T J is multi-valued and one needs to choose thecorrect branch.
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 ℑ ( z ) =
ℜ ( z ) + i 2 ℑ ( 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
4.1 Initial setup
We first use the classical Newton-Raphson scheme to invert the equation z
= f ( m ) where z ∈ C + isfixed and f is rational.
A neat trick which leverages the fact that f is rational and that z ∈ C + is todefine:
[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 finding 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 satisfied for φ z at any m ∈ C .
It is particularly easy toprogram with φ z polynomial.
4.2 Newton lilypads: Doubling strategies and chaining
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 find 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 find 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 this algorithm, combined with the Kantorovich criterion yields:
Theorem 4.1.
Given f : m (cid:55)→ M ⟨− 1 ⟩ ( 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 Table 1.1.
Consider a classical Multi-Layer Perceptron (MLP) with L = 4 layers, feed-forward andfully-connected with ReLU non-linearities.
The MLP’s architecture is determined by the vector λ =( λ 0 , λ 1 , . . .
, λ L ) while the gains at initialization 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 find ourselves with several MLPs architectures each with it’sunique initialization.
The spectral distributions are computed thanks to our Algorithm 1.
As shown by the parameter count in Fig.
A.3, some architectures are quite large for task at hand.
To guarantee convergence despite the changing architectures, we intentionally chose a low learningrate and did 50 epochs of training.
For 200 MLPs, the experiment takes 10 hours of compute on aconsumer GPU (RTX 2080Ti).
This is to be contrasted with less than one minute of CPU computefor the spectral measure ν J .
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 percentiles of the spectral distribution ν J .
Notice that, a posteriori, the
FPT regime is justified 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, confirming the initial claim that stability is a good proxy forperformance, we empirically demonstrated that the accuracy of a feed-forward neural network aftertraining is highly correlated to higher quantiles of the theoretically computed spectral distributionusing this method.
Regarding the importance of hyperbolicity, we surmise that a few large singular values allow the
SGD to escape local minima without compromising overall stability.
A challenging problem would be to accommodate for skip-connections.
In this case the chain ruleused for back-propagation changes in a fundamental way.
The Jacobian cannot be approximated by asimple product of free matrices as the same free variable will appear at multiple locations.
A.2 Experiments
[Table]
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 ) ℓ converges to the law of ϕ ′ l (cid:16)(cid:112) q ℓ N (cid:17) .
From this observation was deduced in [20] the following formula:
[Equation]
Linear
ReLU
Hard Tanh
Hard sine / Triangle
[Equation]
In Table 1.1, the reader shall find the formulae of M D 2 ℓ , S D 2 ℓ and the recurrence relation given by f ℓwhere
[Equation]
with
[Equation]
and
[Equation]
Proof of formulas in Table B.1: ReLU and Hard Tanh have already been computed in [20].
Hard Sine: This is the most tricky formula to establish.
If ϕ ( x ) =
π arcsin ◦ sin (cid:0) π 2 x (cid:1) and (cid:98) f ( ξ ) =
343 (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]
We conclude by computing the Fourier coefficients of ϕ 2 .
[Equation] [Equation]
In the end:
[Equation]
C Appendix: ProofsC.1 Proof of Theorem 2.3
For k ∈ N , we have :
Tr
[Equation]
A (
[Equation] [Equation]
[Equation]
B ( N ) (cid:17) T (cid:16) A ( N ) (cid:17) T AB (cid:19) k (cid:35) =
Tr (cid:20) ( 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 ℓ have i.i.d.
entries.
Thanks to a swapping trick justified in [17], we can assumethat the matrices D ℓ have i.i.d.
entries independent from the rest of the network, distributed asϕ ′ ℓ ( √ q ℓ N ) .
Notice that we can also replace the D ℓ ’s by deterministic matrices that use the quantilesof the same distribution.
This together with standard results from FPT such as [16] gives asymptoticfreeness – see [3] for a more general result reflecting the current state of the art.
Therefore, we canapply Theorem 2.3.
Starting from Eq.
(1.1) , we get in the infinite width regime and by induction:
[Equation]
with the convention Λ 0 = (cid:81)
k =1 λ k = 1 .
Under the asumption that the entries of W ℓ 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
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.
(4.1) if | value | < ε then return m end if grad ← φ ′ z ( m ) m ← m − value/grad end while
Here we give the optimal Kantorovich criterion from [10] adapted to this paper.
Fix z ∈ C
+ andrecall that φ z in Eq.
(4.1) is the map whose zero we want to find.
Theorem D.1 (Kantorovich’s criterion, [13]) .
Consider a starting point m 0 ∈ C , and define:
[Equation]
φ ′ z
[Equation]
m 0 ) (cid:12) (cid:12)(cid:12)(cid:12) , κ := (cid:12)(cid:12)(cid:12)(cid:12)
φ ′
z ( m 0 ) (cid:12) (cid:12)(cid:12)(cid:12) .
If the starting point satisfies 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 ℓ are i.i.d. :
[Equation]
[Equation]
Under the asumption that W Tℓ W ℓ = σ W ℓ I N ℓ −
[Equation]
we find the same m ( s )
( J ) and :
[Equation]
c
[Equation]
−
[Equation]
(cid:19)(cid:19) L (cid:89) ℓ
= (cid:0)
[Equation]
(cid:1) .
These formulas need to be interpreted:
• Variance grows with L , showing increased instability with depth.
• Larger layers, relative to N 0 , give larger Λ ℓ ’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.
(E.1) and (E.2).
By Remark E.1 and Theorem 3.
[Equation]
we have asm → 0 :
[Equation]
Identifying the first order term
[Equation]
one finds indeed Eq.
(E.1).
Continuing the previous computation:
[Equation]
Applying Remark E.1 again for S J T J
[Equation]
we get :
[Equation]
which is equivalent to Eq.
(E.2).
We conclude by specializing to classical weight distributions.
Under the assumption that the entriesof W ℓ are i.i.d.
we have m ( s ) 1 ( W ℓ ) =
σ 2 W ℓ and m ( s ) 2 ( W ℓ ) =
σ 4 W ℓ (1 + λ ℓ ) which gives
[Equation]
σ 4 W ℓ (cid:17) .
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-7ED4/
F.
On the benchmarks of Fig.
The figure 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 finding procedure.
As
such, 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 imprecise
because of the noise, but its performance scales very poorly with n since one needs to diagonalizeever larger matrices.
Of course, Monte-Carlo remains the easiest method to implement.
But the
deterministic methods compute the density up to machine precision.
F.2 Description of the experiment
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 [9] as implementedin Pytorch [19].
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 multiple instances of MLPs and record the learning curves.
In the end, by considering FPT quantiles and the final test accuracy for each MLP, we obtain datawhich amenable to a statistical analysis.
Finally, MNIST and FashionMNIST both having the same format: greyscale images of 28x28 pixels.
To be able to train exactly the same MLPs on all three datasets, we applied a grayscale transformationand resizing to 28x28 to CIFAR10.
F.3 Final remarks
Our sampling procedure does not consider large fluctuations in the λ
i ’s and focuses on balancedarchitectures.
Likewise, the gains at initialization do not deviate much from the classical [9].
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 refined FPT metrics are arguably"second order corrections".
Nevertheless, it is surprising the 90th percentile in Table 1.1 highlycorrelates to the final test accuracy after training.
In the end, tuning FPT metrics does not amount tosecond order corrections.