Phase diagram of Stochastic Gradient Descent in high-dimensional two-layer neural networks
Anonymous Author(s)
AffiliationAddress email
Abstract
Despite the non-convex optimization landscape, over-parametrized shallow networks are able to achieve global convergence under gradient descent.
The picturecan be radically different for narrow networks, which tend to get stuck in badlygeneralizing local minima.
Here we investigate the cross-over between these tworegimes in the high-dimensional setting, and in particular investigate the connectionbetween the so-called mean-field/hydrodynamic regime and the seminal approachof Saad & Solla.
Focusing on the case of Gaussian data, we study the interplaybetween the learning rate, the time scale, and the number of hidden units in thehigh-dimensional dynamics of stochastic gradient descent (SGD).
Our work buildson a deterministic description of SGD in high-dimensions from statistical physics,which we extend and for which we provide rigorous convergence rates.
Descent-based algorithms such as stochastic gradient descent (SGD) and its variants are the workhorseof modern machine learning.
They are simple to implement, efficient to run and most importantly:they work well in practice.
A detailed understanding of the performance of SGD is a major topicin machine learning.
Quite recently, significant progress was achieved in the context of learning inshallow neural networks.
In a series of works, it was shown that the optimisation of wide two-layerneural networks can be mapped to a convex problem in the space of probability distributions overthe weights [1–4].
This remarkable result implies global convergence of two-layer networks towardsperfect learning provided that the number of hidden neurons is large, the learning rate is sufficientlysmall and enough data is at disposition.
This line of work is commonly referred to as the mean-fieldor the hydrodynamic limit of neural networks.
Mathematically, these works showed that one coulddescribe the entire dynamics using a partial differential equation (PDE) in d dimensions.
In a different, and older, line of work, one-pass SGD for two-layer neural networks with a finitenumber p of hidden units, synthetic Gaussian input data and teacher-generated labels has been widelystudied starting with the seminal work of [5].
These works consider the limit of high-dimensional dataand show, in particular, that the stochastic process driven by gradient updates converge to a set of p 2deterministic ordinary differential equations (ODEs) as the input dimension d → ∞ and the learningrate is proportional to 1 /d .
The validity of these ODEs in this limit was proven by [6].
However, thepicture drawn from the analysis of these ODEs is slightly different from the mean-field/hydrodynamicpicture: in this case SGD can get stuck for long time in minima associated to no specialization of thehidden units to the teacher hidden units, and even when it converges to specializing minima, it fails toperfectly learn (i.e. to achieve zero population risk).
In fact, in this analysis, the interplay betweenthe limit of the learning rate going to zero and d → ∞ appeared to be fundamental.
One should naturally wonder about the link between these two sets of works with, on the one hand ad -dimensional PDE (with large p ), and on the other a p 2 -dimensional ODE (with large d ).
In thiswork we aim to build a bridge between these two approaches for studying one-pass SGD.
[Figure]
∞ 𝜅
Perfect learning
B a d l e a r n i n g
N o O
D E s
P l a t e a u
[Figure]
Our starting point is the framework from [5], which we build upon and expand to a much broaderrange of choices of learning rate, time scales, and hidden layer width.
This allows us to provide asharp characterisation of the performance of SGD for two-layer neural networks in high-dimensions.
We show they depends on the precise way in which the limit is taken, and in particular on how thequantity of data, the hidden layer width, & the learning rate scales as d → ∞ .
For different choicesof scaling, we can observe scenarios such as perfect learning, imperfect learning with an unavoidableerror, or even no learning at all.
As a consequence of our analysis, we provide a phase diagram(see Fig.1) 1 describing the possible scenarios arising in the high-dimensional setting.
Our maincontributions are as follow:
C1 We rigorously show that the dynamics of SGD can be captured by a set of deterministic ODEs,considerably extending the proof of [6] to accommodate for general time scalings defined by anarbitrary learning rate, and a general range of hidden layer width.
We provide much finer nonasymptotic guarantees which are crucial for our subsequent analysis.
C2 From the analysis of the ODEs, we derive a phase diagram of SGD for two-layer neural networksin the high-dimensional input layer limit d → ∞ .
In particular, scaling both the learning rate γ andhidden layer width p with the input dimension d as
[Equation]
we identify 4 different learning regimes which are summarized in Fig.
1(left):
• Perfect learning (green region, κ > − δ ): we show that perfect learning (zero population risk
can
be asymptotically achieved with n ∼ d 1+ κ + δ samples even for tasks with additive noise.
• Plateau (blue line κ = − δ ): learning reaches a plateau related to the noise strength.
The pointκ = δ = 0 goes back to the classical work of [5].
• Bad learning (orange region − 1 / 2 < κ + δ < 0 ): here the noise dominates the learning process.•
No ODEs (red region κ
+ δ < − 1 / 2 ): the stochastic process associated to SGD is not guaranteedto converge to a set of deterministic ODEs.
This region is thus outside the scope of our analysis.
Relation to previous work – Deterministic dynamical descriptions of one-pass stochastic gradientdescent in high-dimensions have a long tradition in the statistical physics community, starting withsingle- and two-layer neural networks with few hidden units [7–11].
The seminal work by [5]overcame previous limitations by constructing a set of deterministic ODEs for two-layer networkswith any finite number of hidden units, paving the way for a series of important contributions [12–14, 6].
This line of work corresponds to the κ = δ = 0 case of Fig.
1 (left).
One of our goal is togeneralize this picture beyond fixed hidden layer size and learning rate
A more recent line of work investigating the dynamics of SGD is the so-called mean-field limit[1, 15, 2–4], which connects the SGD dynamics of large-width two-layer neural networks to adiffusion equation in the hidden layer weight density.
In particular, [15] provide non-asymptoticconvergence bounds for sufficiently small learning rates.
The mean-field approach computes theempirical distribution (in R d ) of the hidden layer weights, while we focus on the macroscopic overlapsbetween the teacher and student weights.
Setting
Consider a supervised learning regression task.
The data set is composed of n pairs ( x ν , y ν ) ν ∈ [ n ] ∈
R d +1 identically and independently sampled from P ( x , y ) .
The probability P ( x ) is assumed to beknown and P ( y | x ) is modelled by a two layer neural network called the teacher .
Given a featurevector x ν ∈ R d , the respective label y ν ∈ R is defined as the output of a network with k hidden units,fixed weights W ∗ ∈ R k × d and an activation function σ : R → R :
[Equation]
with w ∗ r ≡ [ W ∗ ] r ∈ R d as the r -th row of the matrix W ∗ and λ ∗ νr ≡ w ∗⊤ r x ν / √ d ∈ R as the r -thcomponent of the teacher local field vector λ ∗ ν ∈ R k .
The parameter ∆ ≥ 0 controls the strength ofadditive label noise: ζ ν ∼ P ( ζ ν ) such that E ζ ∼ P ( ζ ) [ ζ ] = 0 and E ζ ∼ P ( ζ ) [ ζ 2 ] = 1 .
Given a new sample x ∼ P ( x ) outside the training data, the goal is to obtain an estimation ˆ f ( x )for the respective label y .
The error is quantified by a loss function
L ( y, ˆ f ( x , Θ )) , where Θ is anarbitrary set of parameters to be learned from data.
In this manuscript we are interested in the problem of estimating W ∗ with another two-layer neuralnetwork with the same activation function, which we will refer to as the student .
The student networkhas p hidden units and a matrix of weights W ∈ R p ×
d to be learned from the data.
Given a featurevector x ∼ P ( x ) the student prediction for the respective label is given as
[Equation]
where w j ≡ [ W ] j ∈ R d is the j -th row of the matrix W and λ j ≡ w ⊤ j x / √ d ∈ R is defined as thej -th component of the student local field vector λ ∈ R p .
One-pass gradient descent – Typically, one minimizes the empirical risk over the full data set.
Instead, learning with one-pass gradient descent minimizes directly the population risk :
[Equation]
Given a single sample ( x ν , y ν ) the weights are updated sequentially by the gradient descent rule:
[Equation]
[Equation]
with ν ∈ [ n ] and j ∈ [ p ] .
The parameter γ > 0 is the learning rate.
Despite being a simplificationwith respect to batch learning, one-pass gradient descent is an amenable surrogate for the theoreticalanalysis of non-convex optimization, since at each step the gradient is computed with a fresh datasample, which is equivalent to performing SGD directly on the population risk.
In particular, inthis manuscript we assume realizability p ≥ k , and focus our analysis on the square loss
L ( y, ˆ y ) =
12 ( y − ˆ y ) 2 , leading to
[Equation]
where
[Equation]
Therefore, from the above expression we can see that to monitor the population risk along the learningdynamics it is sufficient to track the joint distribution of the local fields ( λ , λ ∗ ) .
For Gaussian data 103 P ( x ) =
N ( x | 0 , 1 ) , one can replace the expectation E x ,y ∼ P ( x ,y )
[ · ] by E λ , λ ∗ ∼N ( λ , λ ∗ | 0 , Ω ) [ · ] andfully describe the dynamics through the following sufficient statistics, known in the statistical physicsliterature as macroscopic variables :
[Equation]
with matrix elements, called order parameters in the statistical physics literature, denoted by q νjl ≡[ Q ν ] jl , m νjr ≡ [ M ν ] jr and ρ rs ≡ [ P ] rs .
The macroscopic state of the system at the learning step νis given by the overlap matrix Ω ν ∈ R ( p + k ) × ( p + k ) :
[Equation]
and the population risk is completely determined by the macroscopic state:
[Equation]
[Equation]
The training dynamics (5) defines a discrete-time stochastic process for the evolution of the overlapmatrix
(cid:110) Ω ν ∈ R ( p + k ) ×
( p + k ) , ν ∈ [ n ] (cid:111) , with P fixed and Q ν and M ν updated as:
[Equation]
with ν ∈ [ n ] , j, l ∈ [ p ] , r ∈ [ k ] and E νj ≡ σ ′ ( λ νj ) E ν .
Albeit our work particularizes to Gaussian input data, we believe our conclusion, and the phasediagram discussed in Fig.
1, to hold beyond this restricted case.
Indeed, while the Gaussian assumptionis crucial to reach a particular set of ODEs and their analytic expression, the approach can be appliedto more complex data distribution, as long as one can track the sufficient statistics required to have aclosed set of equations.
For instance, [16] obtained very similar equations for an arbitrary mixture of
Gaussians —that would obbey the same scaling analysis as ours– while [17–19] proved that manycomplex distributions behave as Gaussians in high-dimensional setting, including, e.g. realistic
GAN-generated data.
We thus expect our conclusions to be robust in this respect.
Main results
Although t 0 = ν/d seems to be the most natural time scaling in the high-dimensional limit d → ∞ ,if γ and p are allowed to vary with d the right-hand side (RHS) of Eqs.
(11) can diverge and renderthe ODE approximation obsolete.
Instead, for a given time scaling δt , we can rewrite eq (11) as
[Equation]
In Theorem 2.1 we prove that as d → ∞ , Ω ν converges to the solution of the ODE:
[Equation]
where ψ : R ( p + k ) ×
( p + k ) → R ( p + k ) ×
( p + k ) is the expected value of the RHS of Eqs.
(12), providedthat this solution stays bounded.
This enhances the result of [6] by providing convergence rates to the
ODEs encompassing all scalings adopted hereafter:
Theorem 2.1 (Deterministic scaling limit of stochastic processes) .
Let τ ∈ R be the continuous timehorizon and δt =
δt ( d ) be a time scaling factor such that the following assumptions hold:
the activation function σ is L -Lipschitz,
the time scaling δt satisfies for some constant c ,
[Equation]
the ODE trajectory ¯ Ω is bounded away from zero: for some ϵ > 0 and any 0 ≤ t ≤ τ , ¯
[Equation]
[Equation]
) ⪰ ϵI p + k .
Then, for any 0 ≤ ν ≤ ⌊ τ/δt ⌋ , the following inequality holds:
[Equation]
Our proof is based on techniques introduced in [20] (namely, their Lemma 2) which studies a differentproblem with related proof techniques.
The proof involves decomposing Ω ν +1 as
[Equation]
where the two first terms can be considered as a deterministic discrete process, and the last term isa martingale increment.
The main challenge lies in showing that the martingale contribution staysbounded throughout the considered time period.
Although the method is similar to [6], there are a number of differences between the two approaches.
First, our proof fixes a number of holes in [6], in particular bounding q νjj by a sufficiently slowlydiverging function of ν .
Additionally, the techniques used in this paper yield a dependency in p thatis nearly negligible, while the previous methods imply bounds that are much too coarse for our needs.
Assumption 3 can seem out of place, since it is a condition on the trajectory itself, and difficult tocheck a priori .
It can however be checked self-consistently and explicitly during the integration ofthe ODE.
Although it is hard to prove directly, we never encountered a situation where this conditionis not verified.
We believe that this condition is actually not important in practice, and that it is anartefact of our proof technique when two neurons learn the exact same weights (which is anywayrather unlikely with a sufficiently spread out initial condition), a situation that would necessitate aseparate treatment.
Alternatively, one could also get rid of this assumption by adding an infinitesimalnoise at each steps to the gradient iteration.
Theorem 2.1 allows us to safely navigate through Fig.
1 by keeping track of convergence rates ofthe discrete process to a set ODEs.
The interplay between learning rate and hidden layer widthdefines the time scaling δt and the trade-off between the linear contribution on E j and the quadraticone, playing a central role on whether the network achieves perfect learning or not.
Specifically,consider the following learning rate and hidden layer width scaling with d : γ = γ
/d
δ , p = p
d κwhere γ 0 ∈ R + and p 0 ∈ N are constants.
The exponent δ ∈ R can be either greater or smaller thanzero, while κ ∈ R + .
Replacing these scalings on Eqs.
(11), we find (where we have chosen γ 0 = p 0without loss of generality):
[Equation]
learning
[Equation]
16) Since the distribution of the label noise P ( ζ ) is such that E ζ ∼ P ( ζ ) [ ζ ] = 0 , the linear contribution in
E j is noiseless in the high-dimensional limit d → ∞ , and therefore we will refer to it as the learningterm .
The noise enters in the equations through the variance computed on the quadratic contribution
E j E l , which we will refer to as the noise term ; intuitively, it is a high-dimensional variance correctionwhich hinders learning.
In order to satisfy (14), we shall take
[Equation]
When κ + δ ̸ = 0 , this implies that either the learning term or the noise term scale like a negativepower of d , and is negligible with respect to the other term.
It is then easy to check that at a finitetime horizon τ , the resulting ODE behaves as if the negligible term was not present.
We refer to
Theorem B.1 in the appendix for a quantitative proof of this phenomenon.
Let us now describe thedifferent regimes depicted in Fig.
[Figure]
Δ
[Figure]
101010101010 5 𝑡 = 𝜈𝑑 10 −10 −10 −10 −
R 𝑑 1 /𝛾 = 𝑝
𝛾 = 𝑝 0 𝑑 1 / 4 𝛾 = 𝑝 0 𝑑 3 /𝛾
= 𝑝
𝑑 1 /
dq jl
[Equation]
18a
[Equation]
[Equation]
(18b) with δt 0 ≡ 1 /d .
This regime is an extension of [5] for which κ = δ = 0 .
The convergence rate tothe ODEs scales with d − 1 / 2 log( d ) , and the phenomenology we observe for κ = δ = 0 is consistentwith previous works studying the setting κ = δ = 0 ; namely the existence of an asymptotic plateauproportional to the noise level.
For instance, the asymptotic population risk R ∞ is known to beproportional to γ ∆ [6] when κ = δ = 0 and the dynamics is driven by a rescaled version of Eqs.
Since the noise term does not vanish under this scaling, perfect learning to zero population risk is notpossible.
There is always an asymptotic plateau related to the noise level ∆ , and the learning rate γ .
Green region (perfect learning) – If κ > − δ we can define the time scaling δt κ + δ ≡ 1 /d
+ κ + δ .
By Theorem 2.1, Eqs.
(16) converge to the following deterministic set of ODEs:
[Equation]
[Equation]
at a rate proportional to d − (1+ κ + δ ) / 2 log( d ) , where we have highlighted that the noise term vanisheswith d − ( κ + δ ) .
Hence, as long as κ > − δ the noise does not play any role on the dynamics.
Thissetting could be understood by taking an effect learning rate γ eff ∝ d − κ − δ on R ∞ ∝ γ ∆ , whichleads to zero population risk, i.e. perfect learning, in the high dimensional limit d → ∞ .
We validatethis claim by a finite size analysis in the next section.
As discussed, the time scaling determines the number of data samples required to complete onelearning step on the continuous scale.
The bigger κ + δ , the more attenuated the noise term, thus thecloser to perfect learning.
The trade-off is that the bigger κ + δ , the larger the number of samplesneeded is, since n = τd 1+ κ + δ .
Given a realizable learning task, one would thus rather choose theparameters to attain the perfect learning region, but being as close as possible to the plateau line fornot increasing too much the needed number of samples.
We remark that [15] provides non-asymptoticconvergence bounds for p ≫ 1 and γ ≪ 1 /d , corresponding to the region δ ≥ 1 in Fig.
Orange region (bad learning) – We now step in the unusual situation where the learning rate growsfaster with d than the hidden layer width: κ < − δ .
In this case, by (17) the noise term dominatesover the dynamics.
Defining the time scaling δt 2( κ + δ ) ≡ 1 /d 1+2( κ + δ ) , we have
[Equation]
20a
[Equation] [Equation] 20b) [Figure]
According to Theorem 2.1 the convergence rate of Eqs.
(16) to Eqs.
(20) scales withd − (1 / 2+ κ + δ ) log( d ) .
Therefore the existence of the noisy ODEs above is circumscribed to theregion − 1 / 2 < κ + δ < 0 and presents a convergence trade-off absent in the other regimes: the fasterof the contributions of Eqs.
(16) goes to zero, the worse is the convergence rate.
In the present case,the more the learning term is attenuated, i.e. the more negative is κ + δ , the worse the dynamics isdescribed by Eqs.
Although the weights are updated, the correlation between the teacher and thestudent weights parametrized by the overlap matrix M remains fixed on its initial value M 0 , whichis a fixed point of the dynamics under this scaling.
Not surprisingly, this leas to poor generalizationcapacity.
Red region (no ODEs) – If κ + δ < − 1 / 2 , the stochastic process driven by the weight dynamicsdoes not converge to deterministic ODEs under the assumptions of Theorem 2.1.
We are then notable to state any claim about this regime.
Initialization and convergence – There are two additional features worth commenting on thehigh-dimensional dynamics and its connection to the mean-field/hydrodynamic approach, regardinginitialization and the specialization transition.
In the ODE approach we discuss here, we alwaysobserve a first plateau where the teacher-student overlaps are all the same.
This means all the hiddenlayer neurons learned the same linear separator.
At this point, the two-layer network is essentiallylinear.
This is called a unspecialized network in [21, 5].
In fact, this is a perfectly normal phenomenon,as with few samples even the Bayes-optimal solution would be unspecialized [22].
Only by runningthe dynamics long enough the student hidden neurons start to specialize , each of them learning adifferent sub-function so that the two-layer network can learn the non-trivial teacher.
Let us make two comments on this phenomenon: (i) while the "linear" learning in the unspecializedregime may remind the reader of the linear learning in the lazy regime [23, 24] of neural nets, the twophenomena are completely different.
In lazy training, the learning is linear because weights changevery little, so that the effective network is a linear approximation of the initial one.
Here, instead, theweights are changing considerably , but each hidden neuron learns essentially the same function.
(ii)
If the ODEs are initialized with weights uncorrelated with the teacher, then the unspecialized regimeis a fixed point of the ODEs: the student thus never specializes, at any time.
Strikingly, such conditionarises as well in the analysis of mean-field equations (see e.g. Theorem 2 in [25] that discussesthe need to have spread initial conditions with a non-zero overlap with the teacher) to guaranteeglobal convergence.
This raises the question about the precise dependence of the learning on theinitialization condition in the high-dimensional regime, where a random start gets a vanishing ( 1 / √ d )overlap.
This is a challenging problem that only recently has been studied (though in a simplersetting) in [26–28] who showed it yields an additional log d time-dependence.
Generalizing theseresults for high-dimensional two-layer nets is an open question which we leave for future work.
[Figure]
Discussion, special cases, and simulations
To illustrate the phase diagram of Fig.
1, we present now several special cases for which we canperform simulations or numerically solve the set of ODEs.
Henceforth, we take σ ( x )= erf
( x/ √ 2) ,for which the expectations of the ODEs and of the population risk, Eq.
(10), can be calculatedanalytically [5].
The explicit expressions are presented in Appendix C. Teacher weights are such thatρ rs = δ rs .
The initial student weights are chosen such that the dimension d can be varied withoutchanging the initial conditions Q 0 , M 0 , P and consequently the initial population risk R 0 .
Adetailed discussion can be found in Appendix D.
Saad & Solla scaling κ = δ =
— We start by recalling the well-known setting characterized bythe point κ = δ = 0 .
The convergence of the stochastic process for fixed learning rate and hiddenlayer width to Eqs.
(18) was first obtained heuristically by [5].
In Fig.
2 (left) we recall this classicalresult by plotting the population risk dynamics for different noise levels.
Dots represent simulations,while solid lines are obtained by integration of the ODEs, Eqs.
Learning is characterized by two phases after the initial decay.
The first is the unspecialized plateauwhere all the teacher-student overlaps are approximately the same: m jr ≈ m .
Waiting long enough,the dynamics reaches the specialization phase, where the student neurons start to specialize , i.e., theiroverlaps with one of the teacher neurons increase and consequently the population risk decreases.
This specialization is discussed extensively in [5].
If ∆ = 0 , the population risk goes asymptoticallyto zero.
Instead, if ∆ ̸ = 0 , the specialization phase presents a second plateau related to the noise ∆ .
The asymptotic population risk R ∞ related to the 2nd plateau is proportional to γ ∆ [6] in thehigh-dimensional limit d → ∞ with p finite.
As mentioned in the previous section, the expectationover E j E l in Eq.
(18a) prevents one from obtaining zero population risk for a noisy teacher.
Perfect learning for κ =
In this section we study the line κ = 0 with δ > 0 of Fig.
1, forwhich Eqs.
(19) with κ = 0 hold.
We show that perfect learning can be asymptotically achieved inthe realizable setting for any finite hidden layer width p = p 0 .
Keeping δ and ∆ fixed, we have donesimulations increasing the input layer dimension d .
In Fig.
3 (left) we set δ = 1 / 2 , ∆ = 10 − 3 andvary the input layer dimension.
The bigger d is, the closer we are to the ODE-derived noiseless result.
Gathering the asymptotic population risk from simulations for varying d and ∆ we perform a finitesize analysis to study the dependence of R ∞ with d .
This shows that the noise term goes to zerounder this setting.
In Fig.
3 (right) we plot R ∞ versus d from simulations (dots) for different noiselevels.
We fit lines under the log-log scale showing that R ∞ ∝ d − δ , as expected.
Figure 4 draws thesame conclusion for δ = 1 / 4 .
As already stated, the interplay between the exponents directly affects the time scale.
We end thissubsection by graphically illustrating this fact through simulations.
Setting the noise to ∆ = 10 − {+
 [Figure] we+} compare the cases δ =0 , 1 / 4 , 3 / 8 , 1 / 2 in Fig.
2 (right).
All simulations are rendered on the scaleδt 0 = 1 /d to illustrate the trade-off between asymptotic performance and training time.
Bad learning for κ =
— We now quickly discuss the uncommon case of γ growing with dwithin the orange region.
In Fig.
5 (left) we compare simulations varying d with the solution of the
ODEs given by Eqs.
Both lead to poor results compared to the green & blue regions.
Moreover,this regime present strong finite-size effects, making it harder to observe the asymptotic ODEs atsmall sizes.
However, the trend as d increases is very clear from the simulations.
As discussed in Sec.2, the more the learning term is attenuated on the ODEs, the worse they describe the dynamics.
Large hidden layer: κ > 0 — Finishing our voyage through Fig. 1 with examples, we brieflydiscuss the case where both input and hidden layer widths are large.
Although Theorem 2.1 providesnon-asymptotic guarantees for κ > 0 , the number of coupled ODEs grows quadratically with p ,making the task of solving them rather challenging.
Thus, we present simulations that illustrate theregions of Fig.
Fixing d = 100 we show in Fig.
5 (right) learning curves for different values of κand δ .
The colors are chosen to match their respective regions in the phase diagram.
Due to the relatively small sizes used in Fig.
5 (right), the green dots seem to decrease towards perfectlearning, even when δ < 0 , provided that κ is large enough, as is predicted by the phase diagram in
Fig.
Moreover, since d is not large enough, when the parameters are within the orange region thefinite-size effects actually dominates, similarly to Fig.
The learning contribution still plays a roleand the asymptotic population risk is similar to the case κ = δ =0 .
Within the red region, which is outof scope of our theory, the simulation gets stuck on a plateau with larger population risk.
Conclusion
Building up on classical statistical physics approaches and extending them to a broad range oflearning rate, time scales, and hidden layer width, we rendered a sharp characterisation of theperformance of SGD for two-layer neural networks in high-dimensions.
Our phase diagram describesthe possible learning scenarios, characterizing learning regimes which had not been addressed byprevious classical works using ODEs.
Crucially, our key conclusions do not rely on an explicitsolution, as our theory allows the characterization of the learning dynamics without solving thesystem of ODEs.
The introduction of scaling factors is non-trivial and has deep implications.
Ourgeneralized description enlightens the trade-off between learning rate and hidden layer width, whichhas also been crucial in the mean-field theories.