Free Probability for predicting the performance of feed-forward neural networks
Abstract
Gradient descent during the learning process of a neural network can be subject to many instabilities.
The spectral density of the Jacobian is a key component for analyzing stability.
Following the works of Pennington et al., such Jacobians are modeled using free multiplicative convolutions from Free Probability Theory (FPT).
We present a reliable and very fast method for computing the associated spectral densities, for given architecture and initialization.
This method has a controlled and proven convergence.
Our technique is based on an homotopy method: it is an adaptative Newton-Raphson scheme which chains basins of attraction.
In order to demonstrate the relevance of our method we show that the relevant FPT metrics computed 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 giving evidence that a very desirable feature for neural networks is the hyperbolicity of their 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 nothing compared to architecture search.
In this context, an obvious moonshot is estimating the performance of an architecture before training.
The goal of this paper is more realistic: providing a fast and reliable computational method for estimating stability before training – and stability is a good proxy for performance [12].
Framework: Consider a feed-forward network of depth L ∈ N , with L full-connected layers.
For each depth ℓ ∈ { 1 , 2 , . . .
, L } , the layer has activation vector x ℓ ∈ R N ℓ , where N ℓ is the current width.
The vector x 0 ∈ R N 0 takes in the neural network’s input, while x L ∈ R N L gives the output.
The vector of widths is written N :=
( N 0 , N 1 , . .
. , N L ) and will appear in superscript to indicate the dependence in any of the N ℓ ’s. The following recurrence relation holds between layers x ℓ = ϕ ℓ (cid:16) W ( N )
ℓ x ℓ − 1 + b ( N ) ℓ (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 write h ℓ := W ( N ) ℓ [Equation] 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 the minibatch 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 it has exactly the same form as the J ( N ) ℓ given in Eq.
The issue is that a large product of (even larger) matrices can easily become unstable.
If many singular values are ≪ 1 , we have gradient vanishing.
If many singular values are ≫ 1 , we have gradient explosion.
Such a transition can be referred to as the edge of chaos [26, 27].
Intuition.
This instability is easily understood thanks to the naive analogy with the one-dimensional case.
Indeed, the geometric progression q n with n → ∞ is the archetype of a long product and it converges extremely fast, to either {+
+} 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 α ∂ L ∂ ( W ( N ) ℓ ,b ( N ) ℓ ) 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 ) ℓ [Equation] 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 and learning performance.
In particular, an appropriate random initialization was suggested and since implemented in all modern ML frameworks [19, 1].
We make classical choices of random initializations.
The biases b ( N ) ℓ are taken as random vectors which entries are centered i.i.d.
Gaussian random variables with standard deviation σ b ℓ .
For the weights, 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 et al.
[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 studying large deep networks, especially because free probability appears at relatively small sizes because of strong 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 →∞ ν A ( N ) B ( N ) =
ν
A ⊠ ν B , where ⊠ is a deterministic operation between measures called multiplicative free convolution.
The ⊠ will be detailed in Section 2.
The letter A (as well as B ) does not correspond to a limiting matrix but to an 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 width regime.
In the limit, asymptotic freeness will also hold.
Therefore, we will see that the measure of interest is:
[Equation]
The goals of this paper are (1) To give a very fast and stable computation of ν J , in the more general setup of rectangular matrices (2) Empirically demonstrate that FTP metrics computed from ν J do correlate 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 use of 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 allow for varying width profiles, which is more inline with design practices.
This requires us to develop a rectangular 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 impossible to realize in practice.
As we shall see, an S -transform is a holomorphic map with multi-valued inverse.
In the words of [5, p.218], “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, or fixed-point algorithms via the subordination method [2, 15, 24].
In the ML community, Pennington et al. pioneered the application of FPT to the theoretical foundations of Machine Learning and did 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 less layer.
A major drawback of this method is that there is no guarantee to find the correct value, unlike our chaining which always chooses the correct branch.
Not only Newton lilypads has theoretical guarantees of convergence, but it is also an order of magnitude faster (Fig. 1.1).
A few standard Cython optimizations allow to gain another order of magnitude, although this can certainly be refined.
• 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
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 1.2 for 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.
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.
For reproducibility of the numerical and empirical aspects, a complete implementation is provided in a Github repository https://github.com/redachhaibi/FreeNN
1.2 Structure of the paper
We start in Section 2 by stating facts from Free Probability Theory.
Most of it is available in the literature, except the definition of 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.
(1.7) in
Theorem 2.3.
[Table]
[Table]
For each dataset and each quantile, the table reports the R factor in a bivariate linear regression.
In Section 3, we explain in detail the FPT model for random neural networks.
Then thanks to the results 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 gives how ν J can theoretically be recovered.
Section 4 presents the numerical resolution which inverts the (multi-valued) generating series.
By chaining different (local) basins of attractions, we obtain a global resolution method.
Our algorithm is 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 Probability 2.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, the measure 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]
where for all k ∈ N , m k ( µ ) :
= (cid:82) R x k µ ( dx ) is the k -th moment of µ .
For µ ̸ = δ 0 , M µ is invertible in the neighborhood of ∞ and the inverse is denoted by M ⟨− 1 ⟩ µ .
The S-transform of µ is defined as S µ ( m ) =[-
-]
[Equation] ( m ) , and is analytic in a neighborhood of m = 0 .
Furthermore, the variable z will always denote an element of C + , while the variables g and m will denote elements in the image 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 shows that this operation is linearized by the S -transform:
Theorem 2.2 (Voiculescu, [25]) .
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 ⊠ µ B such that lim N →∞ µ A ( N ) B ( N ) =
µ A ⊠ µ B .
The operation ⊠ is the multiplicative free convolution.
Moreover
[Equation]
This convergence akin to a law of large numbers is the key ingredient which allows to build the deterministic 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.
-]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 classical definition 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 the upcoming 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 that the width of layers is not constant across layers, which is mostly the case in practice.
Indeed, modern architectures 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 the convention Λ 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 ℓ ] i ) (see Eq.
(1.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 important subtlety that is addressed in the upcoming Theorem 3.1.
Based on an information propagation argument, the papers [21, 22] argue that the entries of h ℓ behave as the i.i.d.
samples of a Gaussian distribution with zero mean and variance q ℓ .
A basic law of large numbers applied to Eq.
(1.2) gives a limit for the empirical measure µ D ℓ = lim N ℓ →∞ µ D ( N ) ℓ [Equation] ϕ
[Equation]
N
[Equation]
Also the recurrence for the 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 1 .
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 can replace each D ( N ) ℓ by a diagonal matrix with independent Gaussian entries (cid:112) q ℓ N independent from the rest.
In that setting, one can apply the results on products of asymptotically free matrices which were 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-Raphson scheme.
The non-linearities ReLU, Hard Tanh and Hard Sine yield explicit formulas, which can be found 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 order to compute : − 1 π ℑ G J T J ( z = x + iy ) .
The smaller the better, and in practice our method works for 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 −π ℑ G µ ( · + iy ) for different y ’s.
Because of Eq.
(2.1) we equivalently need to compute M J T J ( z ) .
M ⟨− J T J m ) is available thanks to the master equation in Theorem 3.1.
Therefore, we need to invert m (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 the correct 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 + is fixed and f is rational.
A neat trick which leverages the fact that f is rational and that z ∈ C + is to define:
[Equation]
As such, we have z = f ( m )
=
P ( m ) Q ( m )
⇐⇒ φ z ( m )
= 0 .
There are several advantages of doing that: (1) Inversion is recast into finding the zero of a polynomial function.
(2) Since we have lim z →∞ M ( z ) = , if z is large in modulus, m = 0 is a natural starting point for the algorithm when z is large.
It is well-known that the Newton-Raphson scheme fails unless the initial guess m 0 ∈ C belongs to a basin of attraction for the method.
And, provided such a guarantee, the Newton-Raphson scheme is exceptionally fast with a quadratic convergence speed.
Kantorovich’s seminal work in 1948 provides such a guarantee locally.
For the reader’s convenience, we give in Appendix D the pseudo-code for the 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 ) which indicates if the Kantorovich criterion is satisfied for φ z at any m ∈ C .
It is particularly easy to program 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 ∈ C the equation φ z ( m ) = 0 .
First, one has to notice that this problem is part of a family parametrized by z ∈ C + .
And the solution is m ≈ 0 for z large.
Therefore, one can find a proxy solution for z ∈ C + high enough.
This is done thanks 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 is done 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 the function f , since it is only passed along as a parameter.
Technically, f is a parameter for all three functions 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 the neighborhood of 0 .
On the experiment
To leverage the numerical scheme, we designed the following experiment whose results are reported on Table 1.1.
Consider a classical Multi-Layer Perceptron (MLP) with L = 4 layers, feed-forward and fully-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 to have balanced architectures, we chose independent λ
i ’s with E ( λ i )
= 1 .
Likewise, we also sample different possible gains.
Hence we find ourselves with several MLPs architectures each with it’s unique initialization.
The spectral distributions are computed thanks to our Algorithm 1.
As shown by the parameter count in Fig.
1.2, some architectures are quite large for task at hand.
To guarantee convergence despite the changing architectures, we intentionally chose a low learning rate and did 50 epochs of training.
For 200 MLPs, the experiment takes 10 hours of compute on a consumer GPU (RTX 2080Ti).
This is to be contrasted with less than one minute of CPU compute for the spectral measure ν J .
To control for any stochastic variability, we also trained multiple instances for each MLP, and the shown results are averages between a few runs.
Finally we calculate the correlations between the accuracy 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 metamodel for the stability of neural network.
Then we gave a method for the numerical resolution both fast and with theoretical guarantees.
Finally, confirming the initial claim that stability is a good proxy for performance, we empirically demonstrated that the accuracy of a feed-forward neural network after training is highly correlated to higher quantiles of the theoretically computed spectral distribution using 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 rule used for back-propagation changes in a fundamental way.
The Jacobian cannot be approximated by a simple product of free matrices as the same free variable will appear at multiple locations.
Acknowledgements
R.C. recognizes the support of French ANR STARS held by Guillaume Cébron as well as support from the ANR-3IA Artificial and Natural Intelligence Toulouse Institute.
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.
As discussed 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 ( ξ ) =
(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]
In the end:
[Equation]
C Appendix: Proofs C.1 Proof of Theorem 2.
For k ∈ N , we have :
Tr (cid:34)(cid:18)(cid:16)
A ( N ) B ( N ) (cid:17) T A ( N ) B ( N ) (cid:19) k
(cid:35) =
Tr (cid:34)(cid:18)(cid:16) 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 shows that
[Equation]
and then
M ⟨−
[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.
We assume that W ℓ have i.i.d.
entries.
Thanks to a swapping trick justified in [17], we can assume that the matrices D ℓ have i.i [Equation]
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 quantiles of the same distribution.
This together with standard results from FPT such as [16] gives asymptotic freeness – see [3] for a more general result reflecting the current state of the art.
Therefore, we can apply Theorem 2.3.
Starting from Eq.
[Equation] 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 + and recall 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]
(cid:12) (cid:12)(cid:12)(cid:12) ,
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12) [Equation]
[Equation]
) (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 that the entries of W ℓ are i.i.d. :
[Equation]
Under the asumption that W Tℓ W ℓ = σ W ℓ I N ℓ −
[Equation]
we find the same m ( s )
( J ) and :
[Equation]
(cid:19)(cid:19) L (cid:89) ℓ
= (cid:0) c
[Equation]
4 W ℓ (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 as m → 0 :
[Equation]
Identifying the first order term
[Equation]
one finds indeed Eq.
(E.1).
Continuing the previous computation:
[Equation]
Applying Remark E [Equation] 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 entries of W ℓ are i.i.d.
we have m ( s ) 1 ( W ℓ ) =
σ 2 W ℓ and m ( s ) 2 ( W ℓ ) =
σ 4 W ℓ (1 + λ ℓ ) which gives
[Equation]
F Appendix: More details on the benchmarks and the experiment
Recall that we have provided an anonymized Github repository at the address: https://github.com/redachhaibi/FreeNN/
F.
On the benchmarks of Fig. 1.
The figure presents the computational time required for computing the density of ν J .
Let us make the following 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 number of required points.
This is easily understood by the fact that smaller N requires the computation of more 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 diagonalize ever 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-forward and 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 implemented in Pytorch [19].
Thus the full vector of variances of this initialization σ = ( σ 1 , σ 2 , . . .
, σ L ) is determined 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 data which amenable to a statistical analysis.
Finally, MNIST and FashionMNIST both have the same format: grayscale images of 28x28 pixels.
To be able to train exactly the same MLPs on all three datasets, we applied a grayscale transformation and resizing to 28x28 to CIFAR10.
F.3 Final remarks
Our sampling procedure does not consider large fluctuations in the λ
i ’s and focuses on balanced architectures.
Likewise, the gains at initialization do not deviate much from the classical [9].
It is important to recognize that without the insight of FPT, the scalings applied to such initializations are already normalized so that spectral measures do converge.
As such, we never encounter truly problematic gradient vanishing or gradient explosion, which completely 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 highly correlates to the final test accuracy after training.
In the end, tuning FPT metrics does not amount to second order corrections.