Empirical Phase Diagram for Three-layer Neural Networks with Inﬁnite Width
Abstract
Substantial work indicates that the dynamics of neural networks (NNs) is closelyrelated to their initialization of parameters.
Inspired by the phase diagram fortwo-layer ReLU NNs with inﬁnite width (Luo et al., 2021), we make a steptowards drawing a phase diagram for three-layer ReLU NNs with inﬁnite width.
First, we derive a normalized gradient ﬂow for three-layer ReLU NNs and obtaintwo key independent quantities to distinguish different dynamical regimes forcommon initialization methods.
With carefully designed experiments and a largecomputation cost, for both synthetic datasets and real datasets, we ﬁnd that thedynamics of each layer also could be divided into a linear regime and a condensedregime, separated by a critical regime.
The criteria is the relative change ofinput weights (the input weight of a hidden neuron consists of the weight fromits input layer to the hidden neuron and its bias term) as the width approachesinﬁnity during the training, which tends to 0 , + ∞ and O (1) , respectively.
Inaddition, we also demonstrate that different layers can lie in different dynamicalregimes in a training process within a deep NN.
In the condensed regime, we alsoobserve the condensation of weights in isolated orientations with low complexity.
Through experiments under three-layer condition, our phase diagram suggests acomplicated dynamical regimes consisting of three possible regimes, together withtheir mixture, for deep NNs and provides a guidance for studying deep NNs indifferent initialization regimes, which reveals the possibility of completely differentdynamics emerging within a deep NN for its different layers.
Introduction
As neural networks (NNs) have been employed to extensive practical and scientiﬁc problems, suchas computer vision, natural language processing and automatic driving etc, studying the rationalebehind them can help us comprehend the principles behind NNs and better use NNs.
For example,the initialization of neural network parameters plays an important role in the training dynamics andgeneralization of NN ﬁtting.
Several types of initialization schemes are often used in practical trainings, such as He initialization(He et al., 2015), LeCun initialization (LeCun et al., 2012), Xavier initialization (Glorot and Bengio,2010), etc.
Theoretical study further explores how initialization affects the training dynamics of NNs.
For example, the NTK scaling (Jacot et al., 2018; Chizat and Bach, 2019; Arora et al. ,2019; Zhang et al., 2020) leads to a training where parameters stay at a small neighborhood of theinitialization, therefore, the NN after training can be well approximated by the ﬁrst-order Taylorexpansion at the initial point, i.e., linear dynamics.
The mean-ﬁeld scaling (Mei et al., 2019; Rotskoffand Vanden-Eijnden, 2018; Chizat and Bach, 2018; Sirignano and Spiliopoulos, 2020) leads to atraining where parameters signiﬁcantly deviates from the initialization, therefore, the NN dynamicsexhibits clearly non-linear behavior.
Luo et al.
(2021) systematically study the effect of initialization for two-layer ReLU NN with inﬁnite width by establishing a phase diagram, which shows threedistinct regimes, i.e., linear regime, critical regime and condensed regime, based on the relativechange of input weights (the input weight of a hidden neuron consists of the weight from its inputlayer to the hidden neuron and its bias term) as the width approaches inﬁnity during the training,which tends to 0 , O (1) and + ∞ , respectively.
The condensed regime is named so, because empiricalstudy ﬁnds that orientations of the input weights condense in several isolated orientations, which is astrong feature learning behavior, an important characteristic of deep learning.
The study in Luo et al.
(2021) is limited in two-layer NNs, however, empirical studies show thatdeep networks can learn faster and generalize better than shallow networks in both real data andsynthetic data (He et al., 2016; Arora et al., 2018; Xu and Zhou, 2021).
Because of multi-layerstructure and non-linearity, the study of multi-layer nonlinear NNs is often challenging.
For example,the mean-ﬁeld theory is established on two-layer NNs and difﬁcult to be extended to multi-layer
NNs (Mei et al., 2019; Rotskoff and Vanden-Eijnden, 2018; Chizat and Bach, 2018; Sirignano and
Spiliopoulos, 2020).
Even for empirical study, excessive hyper-parameters often set an obstacle toobtain clear phenomena as well, where multi-layer NNs could have distinct characteristics comparedwith shallow two-layer NNs.
For example, due to the existence of other layers, how is the dynamicalregime of the weights of a layer in multi-layer NNs different from two-layer ones and can weightsof two different layers experience distinct dynamical behaviors in one training process?
Answeringthese questions is important to understand how the dynamics of multi-layer NNs depend on the initialhyper-parameters, thus, leading to a deeper understanding of deep NNs and providing guidance fortuning parameters in practical training.
For instance, if we want to study the nonlinear effect of ahidden layer, we can set this interested hidden layer in the non-linear regime and all other hiddenlayers in the linear regime.
Further, we can continue to study how the non-linear effect of a hiddenlayer depends on its depth.
Three-layer network is a good starting point to study multi-layer NNs as it only has two hidden layersfor studying the interaction between layers.
Besides, to have a clear boundary of different dynamicalregimes, we focus on three-layer ReLU NNs with large-width limit.
A difﬁculty to overcome isto ﬁnd appropriate quantities required for describing a phase diagram, which should be able toidentify distinct different regimes, and can also account for the same dynamical behavior whenhyper-parameters are different but those quantities are the same.
To this end, we derive a normalizedgradient ﬂow for a three-layer ReLU NN.
Since we only care about the training trajectory up to atime scaling, we obtain three independent quantities that meet our requirements.
And then, we furthersimplify the setting by assuming weights of the two outer layers use a same initialization schemeas the commonly used in practice, and we obtain two key independent quantities.
With carefullydesigned experiments and a large computation cost, for both synthetic datasets and real datasets, weﬁnd that weights of each layer also have a linear regime and a condensed regime, separated by acritical regime.
The criteria is the relative change of input weights as the width approaches inﬁnityduring the training, which tends to 0 , + ∞ and O (1) , respectively.
In addition, we also ﬁnd thatdifferent layers can lie in different dynamical regimes in a training process.
In the condensed regime,we similarly observe the condensation of weights in isolated orientations.
This work empirically draws a phase diagram for three-layer ReLU NNs with large width, as shownin Fig.
For deeper networks, this work suggests that the dynamical regime for the weights of eachlayer is similar to that of a two-layer NN and different layers can lie in distinct regimes in a trainingprocess.
For NNs with not too large width, the boundary between regimes is not sharp but this workstill sheds lights on how the initialization affects the dynamics of weights in different layers.
Ourstudy on phase diagram of three-layer ReLU NNs thus lays an important foundation for future workon the dynamical behavior of multi-layer NNs and corresponding implicit regularizations at each ofthe identiﬁed regimes and guides the parameter initialization tuning in the practice.
Related works
A series of works have shown that global minima can be found for NNs with several speciﬁcinitializations in linear regime (Zou et al., 2018; Allen-Zhu et al., 2019; Du et al., 2019; E et al.,
2020)or for the two-layer ReLU NNs in the whole linear regime (Luo et al., 2021).
Chen et al.
(2021)prove that for wide shallow NNs under the mean-ﬁeld scaling and with a general class of activationfunctions, when the input dimension is at least the size of the training set, the training loss convergesto zero at a linear rate under gradient ﬂow.
In the condensation regime, a large network with condensation is effectively a network of only afew effective neurons, leading to an output function with low complexity (Bartlett and Mendelson,2002), thus, may provide a possible explanation for good generalization performance of large NNs(Breiman, 1995; Zhang et al., 2021).
The condensation suggests a bias towards simple function,which is consistent with the frequency principle that NNs tend to learn data by low-frequency function(Xu et al., 2019; Rahaman et al., 2019; Xu et al., 2020, 2022).
NNs have signiﬁcant feature learningin the condensation regimes.
To actively learn features, Lyu et al.
(2021) use small initialization tostudy how gradient descent can lead a NN to a solution of “max-margin” solution.
How the condensation happens during the initial training stage is explored for ReLU activationfunction (Maennel et al., 2018; Pellegrini and Biroli, 2020) and general differential activationfunctions (Xu et al., 2021).
The embedding principle (Zhang et al., 2021a,b) shows that the losslandscape of an NN contains all critical points/functions of all the narrower NNs, thus providing abasis for why a large network can condense during the training.
Preliminary
For convenience, we consider a three-layer NN with m hidden neurons for each layer,
[Equation]
(1) where x ∈ R d , α is the scaling factor, m is the width of both layers, θ =
vec { a , W [2] , W [1] } witha ∈ R m , W [2] ∈ R m × m , W [1] ∈ R m × d is the set of parameters initialized by a k ∼ N (0 , β 23 ) ,
W [2] kk (cid:48) ∼ N (0 , β 22 ) , W [1] kk (cid:48) ∼ N (0 , β 21 ) , i.i.d .
The bias term b [1] k can be incorporated by expanding xand W [1] to ( x T , 1) T and ( W [1] , b [1] k ) T , while b [2] k can be incorporated by expanding ¯ x = σ ( W [1] x )and W [2] to (¯ x T , 1) T and ( W [2] , b [2] k ) T .
At the inﬁnite-width limit m → ∞ , given β 1 , β 2 , β 3 ∼ O (1) , for α ∼ √ m 1 m 2 , the gradient ﬂow of
NNs can be approximated by a linear dynamics of neural tangent kernel (NTK)
(Jacot et al.
(2018)) ,whereas given β 1 ∼ (cid:113) 2 m 1 + d , β 2 ∼ (cid:113) 2 m 1 + m 2 , β 3 ∼ (cid:113) 2 m 2 +1 , for α ∼ O (1) gradient ﬂow of NNexhibits highly nonlinear dynamics of Xavier (Glorot and Bengio (2010)).
Assumption 3.1.
For convenience, we assume that the number of neurons in each hidden layer isequal, i.e. m
= m
= m .
Assumption 3.2.
Considering that we are mainly concerned with the dynamics of W [1] and W [2] ,i.e., the dynamics of a is not the main concern, and the initialization methods used for W [2] and awithin a neural network in practice are always the same, which leads to the initialization parametersβ 2 and β 3 to be equal or multiplied, so we assume that β
= Bβ 3 .
Cosine similarity: The cosine similarity for two vectors u and v is deﬁned as
[Equation]
Rescaling and the normalized model of the three layer neural network
To identify sharply distinctive regimes/states, we take the inﬁnite-width limit m → ∞ as our startingpoint.
Identiﬁcation of the coordinates is important for drawing the phase diagram.
There are someguiding principles for ﬁnding the coordinates of a phase diagram (Luo et al., 2021): (i) They should beeffectively independent; (ii) Given a speciﬁc coordinate in the phase diagram, the learning dynamicsof all the corresponding NNs statistically should be similar up to a time scaling; (iii) They shouldwell differentiate dynamical differences except for the time scaling.
Guided by above principles, in this section, we perform the following rescaling procedure for a faircomparison between different choices of hyperparameters and obtain a normalized model with twoindependent quantities irrespective of the time scaling of the gradient ﬂow dynamics.
The empirical risk is
[Equation]
Then the training dynamics based on gradient descent (GD) at the continuous limit obeys the followinggradient ﬂow of θ ,
[Equation]
More precisely, we obtain that:
[Equation]
) where e i = (cid:0) 1 α a T σ ( W [2] σ ( W [1] x i )) − y i (cid:1) , the operation (cid:12) is the Hadamard product.
Then,setting
[Equation]
(6) so that all the parameters are picked from Standard Normal Distribution, by the Chain Rule ofderivative, we obtain that:
[Equation] − (cid:32)n n (cid:88) i = κ 3 σ ( W [2] σ ( W [1] x i )) (cid:33) e i , d W [2] d t = − κ(cid:32)n n (cid:88) i = κ 3 a (cid:12) σ (cid:48) ( W [2] σ ( W [1] x i ))
σ ( W [1] x i ) T (cid:33) e i , d W [1] d t = − κ(cid:32)n n (cid:88) i = κ 3 W [2] T ( a (cid:12) σ (cid:48) ( W [2] σ ( W [1] x i ))) (cid:12) σ (cid:48) ( W [1] x i ) x T i (cid:33) e i , (7)
[Table]
[Table]
where, we introduce four scaling parameters
[Equation]
Note that κ 1 , κ 2 and κ 3 do not follow principle (ii) and (iii) above at inﬁnite-width limit.
They arein general functions of m , which attains 0 , O (1) , + ∞ at m → ∞ .
To account for such dynamicaldifference under different widely considered power-law scalings of α, β 1 , β 2 and β 3 shown in Table.1 and based on our assumption 3.1 and 3.2, we arrive at
[Equation]
Remark 4.1 .
Here we list some commonly-used initialization methods and/or related works withtheir scaling parameters as shown in Table 1.
Empirical phase diagram
5.1 Experimental setup
Throughout this section, we use three-layer fully-connected neural networks with size, d - m - m - d out .
The input dimension d is determined by the training data, i.e., d = 1 for synthetic data and d = 28 ∗ 28for MNIST.
The output dimension is d out = 1 for synthetic data and d out = 10 for MNIST.
Thenumber of hidden neurons m is speciﬁed in each experiment.
All parameters are initialized by a Gaussian distribution N (0 , var ) , where var depends on β 1 , β 2 and β 3 .
The total data size is n .
Thetraining method is gradient descent with full batch, learning rate lr and MSE loss.
For synthetic data,we use a simple 1-d problem of 4 training points.
5.2 Intuitive experiments of synthetic data
With γ 3 and γ 2 as coordinates, in this subsection, we illustrate through experiments the behavior of adiversity of typical cases over the phase diagram regarding W [1] and W [2] respectively.
We use asimple 1-d problem of 4 training points, which allows an easy visualization.
The ﬁrst row in Fig.
2 shows typical learning results over different γ 3 ’s, from a relatively jaggedinterpolation (NTK scaling) to a smooth interpolation (mean-ﬁeld scaling) and further to a linearspline interpolation.
With input weight and bias, W [1] k is two-dimensional for d = 1 .
We display the{ (( W [1] k ) 1 , ( W [1] k ) 2 ) } mk =1 of the trained network Fig.
For γ 3 = 1 . 0 , the initial scatter plot is veryclose to the one after training.
However, for γ 3 = 2 . 0 , active neurons (i.e., neurons with signiﬁcantamplitude and away from the origin) are condensed at a few orientations, which strongly deviatesfrom the initial scatter plot.
For γ 3 = 1 . 5 , the scatter points present an intermediate state.
Inspired by W [1] , we found that the directions between the input weights is a very important groupfeature, which can also reﬂect the characteristics between different γ 3 and γ 2 .
Considering that W [2] k [Figure] is m + 1 dimensional, We cannot directly and intuitively show { W [2] k } mk =1 in a ﬁgure.
Here, W [2] krepresent the input weight of the k -th neuron in the second hidden layer.
Therefore, to characterizethe distance between different { W [2] k } mk =1 , we introduce D ( u , v ) to denote the inner product of two{ W [2] k } mk =1 , as Equ. 2.
As shown in Fig.
3, for γ 3 = 2 . 2 , W [2] k hardly have a strong relationshipwith each other, while for γ 3 = 2 . 8 , most of the W [2] k are very close to each other, showing strongcondensation.
[Figure]
u, v ) .
5.3 Phase diagram for three-layer neural networks
In this section, with γ 3 and γ 2 as coordinates, we characterize at m → ∞ the dynamical regimesof NNs and identify their boundaries in the phase diagram through experimental and theoreticalapproaches.
How to characterize and classify different types of training behaviors of NNs is animportant open question.
Recently, Luo et al.
(2021)’s
work has made clear the research on thetraining behavior of the initialization parameters of the two-layer inﬁnitely wide ReLU neural network.
Many works, like Jacot et al.
; Lee et al.
Arora et al. ; E et al. (2019), are allspeciﬁc points or line in the phase diagram belong to the linear regime in Luo et al. However,there is very little work on the three-layer phase diagram or the multi-layer one.
For the three-layer [Table] inﬁnitely wide ReLU neural networks, as shown in Fig.
[Table]
several works have proved speciﬁc pointsin the three-layer phase diagram belong to the linear regime of both W [1] and W [2] .
However, itsexact range in the phase diagram remains unclear.
On the other hand, NN training dynamics can alsobe highly nonlinear at m → ∞ as widely studied initialization method, Glorot and Bengio (2010), asa point shown in the phase diagram.
However, whether there are other points in the phase diagramthat has similar training behavior is not well understood in three-layer condition.
In addition, it is stillnot clear if there are other regimes in the phase diagram that are nonlinear but behaves distinctivelycomparing to the Glorot and Bengio (2010) initialization method.
In the following, we will addressthese problems and draw the three-layer phase diagram.
5.3.1 Regime identiﬁcation and separation
To characterize different dynamical regimes, we deﬁne the relative change of weights as
[Equation]
where θ ∗ W 1 and θ ∗ W 2 are the weights after training for the ﬁrst and the second hidden layer, andθ W 1 (0) and θ W 2 (0) for initial weights.
[Equation]
The slopes in (a)-(c) are ﬁtted by RD( W [1] )
w.r.t m = 100 , 1000 , 2000 , 5000 , 10000 in a log-log scale.
As γ 3 grows from 0 .
9 to 2 .
the corresponding slopes grow as well.
We quantify the growth of RD( W [ i ] ) , i = 1 , 2 as m → ∞ .
As shown in Fig.
4(a-c), theyapproximately have a power-law relation.
Therefore we deﬁne
[Equation]
(11) which is empirically obtained by estimating the slope in the log-log plot.
As shown in Table 2,
NNs with the same pair of γ 2 and γ 3 , but different α, β 1 , β 2 and β 3 , have very similar S W i , which validates the effectiveness of the normalized model.
In the following experiments, we only showresult of one combination of α, β 1 , β 2 and β 3 for a pair of γ 2 and γ 3 .
Then, we explore the phase diagram by experimentally scanning S W i over the phase space.
Theresult for the same 1-d problem as in Fig.
2 is presented in Fig.
In the red zone, where S W i isless than zero, RD( θ ∗ W i ) → 0 as m → ∞ , indicating a linear regime.
In contrast, in the blue zone,where S W i is greater than zero, RD( θ ∗ W i ) → ∞ as m → ∞ , indicating a highly nonlinear behavior.
Their boundary are experimentally identiﬁed through interpolation indicated by stars in Fig.
5, where RD( θ ∗ W i ) ∼ O (1) .
The dashed lines are our speculative boundary, i e., critical regime, by a largeamount of experiments .
In order to verify the rationality of the three-layer phase diagram, we alsoverify the results on the MNIST dataset, as shown in Fig. 6, and obtain similar results.
[Figure] [Figure]
5.4 Condensation
We then examine the condensation over the phase diagram.
For W [1] k , we display the{ (( W [1] k ) 1 , ( W [1] k ) 2 ) } m k =1 for two different regions in Fig.
As is shown in Fig.
7, on the left sideof the blue squares, the parameters after training and the initialization parameters are very close,while on the right side of the blue squares, weights with not too small magnitude are condensed at afew orientations, which strongly deviates from the initial scatter plot.
The graphs in the blue boxespresent an intermediate state or process state, making the above-mentioned change a continuousprocess.
Next, we need to deﬁne a quantity to describe this phenomenon.
For the second hidden layer, the weights is very high-dimensional.
We deﬁne a quantity to qualitativelymeasure the degree of condensation.
We deﬁne ζ = 1 ( m/ 2) 2 (cid:80) m/ 2 i,j =1 (cid:107) D ( u i , u j ) (cid:107) , which is theaverage inner product distance between different W [2][ k ] ’s and subscript [ k ] indicates that the row of
W [2] which has a magnitude with ranking k among all rows.
More condensation yields to a larger ζ .
As is shown in Fig.
8, we could clearly observe that the condensation tendency is consistent withphase distribution in Fig. 8.
Conclusion
In this paper, we empirically characterized the linear, critical, and condensed regimes with distinctivefeatures and draw the phase diagram for the three-layer ReLU NN at the inﬁnite-width limit.
Throughexperiments, we further identify the condensation as the signature behavior in the condensed regimeof strong non-linearity.
Our phase diagram ﬁgures out the relation between the training dynamics ofmulti-layer neural networks and their initialization, reveals the possibility of having different trainingdynamics for different layers within a neural network, and serves as a map that provide a guidancefor the future study of deep NNs, such as how different dynamical regimes of different hidden layersaffect the generalization.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section ?? .
• Did you include the license to the code and datasets?
[No] The code and the data areproprietary.•
Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[N/A]
(c) Did you discuss any potential negative societal impacts of your work?
[N/A](d)
Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]2.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] In supplemen-
tary.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[No]
(d) Did you include the total amount of compute and the type of resources used (e.g., type
of GPUs, internal cluster, or cloud provider)?
[No]
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[N/A]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]