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 step towards 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 obtain two 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 of inputweights (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 , + ∞ and O (1) , respectively.
In addition, wealso demonstrate that different layers can lie in different dynamical regimes in atraining process within a deep NN.
In the condensed regime, we also observe thecondensation of weights in isolated orientations with low complexity.
Throughexperiments under three-layer condition, our phase diagram suggests a complicated dynamical regimes consisting of three possible regimes, together with theirmixture, for deep NNs and provides a guidance for studying deep NNs in different 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; Zhanget al., 2020) leads to a training where parameters stay at a small neighborhood of the initialization,therefore, the NN after training can be well approximated by the ﬁrst-order Taylor expansion at theinitial point, i.e., linear dynamics.
The mean-ﬁeld scaling (Mei et al., 2019; Rotskoff and VandenEijnden, 2018; Chizat and Bach, 2018; Sirignano and Spiliopoulos, 2020) leads to a training whereparameters signiﬁcantly deviates from the initialization, therefore, the NN dynamics exhibits clearlynon-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 three distinct regimes,i.e., linear regime, critical regime and condensed regime, based on the relative change of inputweights (the input weight of a hidden neuron consists of the weight from its input layer to thehidden neuron and its bias term) as the width approaches inﬁnity during the training, which tends to0 , O (1) and + ∞ , respectively.
The condensed regime is named so, because empirical study ﬁndsthat orientations of the input weights condense in several isolated orientations, which is a strongfeature 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 hiddenlayers for studying the interaction between layers.
Besides, to have a clear boundary of differentdynamical regimes, we focus on three-layer ReLU NNs with large-width limit.
A difﬁculty toovercome is to ﬁnd appropriate quantities required for describing a phase diagram, which shouldbe able to identify distinct different regimes, and can also account for the same dynamical behaviorwhen hyper-parameters are different but those quantities are the same.
To this end, we derive anormalized gradient ﬂow for a three-layer ReLU NN.
Since we only care about the training trajectoryup to a time scaling, we obtain three independent quantities that meet our requirements.
And then,we further simplify the setting by assuming weights of the two outer layers use a same initializationscheme as the commonly used in practice, and we obtain two key independent quantities.
Withcarefully designed experiments and a large computation cost, for both synthetic datasets and realdatasets, we ﬁnd that weights of each layer also have a linear regime and a condensed regime,separated by a critical regime.
The criteria is the relative change of input weights as the widthapproaches inﬁnity during the training, which tends to 0 , + ∞ and O (1) , respectively.
In addition,we also ﬁnd that different layers can lie in different dynamical regimes in a training process.
In thecondensed 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ﬁc initializations 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).
Geiger et al.
(2020)probe the dynamic crossover between the NTK limit and the Mean-Feild limit, and Yang (2020);
Yang and Littwin (2021) show that Tangent Kernel of a randomly initialized neural network with‘arbitrary architecture’ converges to a deterministic limit, as the network width goes to inﬁnity.
Chen et al.
(2021) prove that for wide shallow NNs under the mean-ﬁeld scaling and with a general classof activation functions, when the input dimension is at least the size of the training set, the trainingloss converges to 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 featurelearning in the condensation regimes.
To actively learn features, Lyu et al.
(2021) use small initialization to study 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 activation function (Maennel et al., 2018; Pellegrini and Biroli, 2020) and general differential activation functions(Xu et al., 2021).
The embedding principle (Zhang et al., 2021a,b) shows that the loss landscape ofan NN contains all critical points/functions of all the narrower NNs, thus providing a basis for whya large network can condense during the training.
Preliminary
For convenience, we consider a three-layer NN with m hidden neurons for each layer,f
[Equation]
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 ′ ∼ N (0 , β 22 ) , W [1] kk ′ ∼ 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 ∼
[Equation] , for α ∼ O [Equation] 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 .because the output is linear with respect to a , thus, a is not the main concern, and the initializationmethods used for a neural network in practice often has β 2 = Bβ 3 , where B is a constant independent of m . So, we have the following assumption.
Assumption 3.2.
β
= Bβ 3 , where B is a constant in dependent of m .
Cosine similarity: The cosine similarity for two vectors u and v is deﬁned as
[Equation]
[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 aresome guiding principles for ﬁnding the coordinates of a phase diagram (Luo et al., 2021): (i) Theyshould be effectively independent; (ii) Given a speciﬁc coordinate in the phase diagram, the learningdynamics of all the corresponding NNs statistically should be similar up to a time scaling; (iii) Theyshould well 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 following gradient ﬂow of θ ,
[Equation]
More precisely, we obtain that:
[Equation]
5) where e i = (α a T σ ( W [2] σ ( W [1] x i )) − y i ) , the operation ⊙ is the Hadamard product.
Then,setting
[Equation]
so that all the parameters are picked from Standard Normal Distribution.
By the chain rule ofderivative and the homogeneity of ReLU , i.e. σ ( a u ) = aσ ( u ) and σ ′ ( a u )
= σ ′ ( u ) , where a > 0is a scalar and u is a vector, we obtain that:
[Equation]
[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]
(9) where our assumptions 3.1 and 3.2 mean κ 1 = B , leading to γ 1 = lim m →∞ − log
[Equation]
1 log m = 0 .
So wediscuss the phase diagram of three layer network under the coordinates:
( γ 1 , γ 2 , γ 3 ) = (0 , γ 2 , γ 3 ) ,and focus on ( γ 2 , γ 3 ) .
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 ∗ 28 for MNIST.
The output dimension is d out = 1 for synthetic data and d out = 10 for MNIST.
The number of hidden neurons m is speciﬁed in each experiment.
All parameters are initialized bya Gaussian distribution N (0 , var ) , where var depends on β 1 , β 2 and β 3 .
The total data size is n .
The training method is gradient descent with full batch, learning rate lr and MSE loss.
For syntheticdata, 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 ofa diversity 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 displaythe { (( W [1] k ) 1 , ( W [1] k ) 2 ) } mk =1 of the trained network in Fig.
For γ 3 = 1 . 0 , the initial scatter plotis very close to the one after training.
However, for γ 3 = 2 . 0 , active neurons (i.e., neurons withsigniﬁcant amplitude and away from the origin)
[Figure]
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] kis m + 1 dimensional, We cannot directly and intuitively show { W [2] k } mk =1 in a ﬁgure.
Here, W [2] krepresents the input weight of the k -th neuron in the second hidden layer.
Therefore, to characterizethe distance between different vectors in { W [2] k } mk =1 , we introduce D ( u , v ) to denote the innerproduct of two { W [2] k } mk =1 , as Equ. 2.
As shown in Fig.
3, for γ 3 = 2 . 2 , W [2] k hardly has a strongrelationship with each other, while for γ 3 = 2 . 8 , most of the W [2] k are very close to each other,showing strong condensation.
[Figure]
navy-blue block, D ( u, v ) ∼ − 1 ) , their input weights have the same (opposite) direction. Considering that not all neurons are activated for ReLU activation function, we only select half of the neurons with the largest modulus length to draw D ( 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.
However, there is little work on the three-layer phase diagram or the multilayer one.
For the three-layer inﬁnitely wide ReLU neural networks, as shown in Fig. 1, several
[Table]
works have proved speciﬁc points in the three-layer phase diagram belong to the linear regime ofboth W [1] and W [2] .
However, its exact range in the phase diagram remains unclear.
On theother hand, NN training dynamics can also be highly nonlinear at m → ∞ , such as widely used and a the
However,whether there are other points in the phase diagram that has similar training behavior is not wellunderstood in three-layer condition.
In addition, it is still not clear if there are other regimes inthe phase diagram that are nonlinear but behaves distinctively comparing to the Glorot and Bengio(2010) initialization method.
In the following, we will address these problems and draw a three-layerphase 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.
[Figure]
We quantify the growth of RD( W [ i ] ) , i = 1 , 2 as m → ∞ .
As shown in Fig.
4(a-c), they approximately have a power-law relation.
Therefore, similar to Luo et al. (2021), 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 iis less than zero, RD( θ ∗ W i ) → 0 as m → ∞ , indicating a linear regime.
In contrast, in the bluezone, where S W i is greater than zero, RD( θ ∗ W i ) → ∞ as m → ∞ , indicating a highly nonlinearbehavior.
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 large amount of experiments .
In order to verify the rationality of the three-layer phase diagram,we also verify 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 leftside of the blue squares, the parameters after training and the initialization parameters are veryclose, while on the right side of the blue squares, weights with not too small magnitude are condensed at a few orientations, which strongly deviates from the initial scatter plot.
The graphs inthe blue boxes present an intermediate state or process state, making the above-mentioned change acontinuous process.
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 qualitatively measure the degree of condensation.
We deﬁne ζ = 1 ( m/ 2)∑ m/ 2 i,j =1 ∥ D ( u i , u j ) ∥ , 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.
The network structure and the four points of the objective function are consistent with Fig.
3 and γ 2 are indicated by the abscissa and the ordinate, respectively.
The red the small corre spond to the initialization parameters, and the green dots to the post- parameters ensure and ﬁnally
Conclusion
In this paper, we empirically characterized the linear, critical, and condensed regimes with distinctive features and draw the phase diagram for the three-layer ReLU NN at the inﬁnite-width limit.
Through experiments, we further identify the condensation as the signature behavior in the condensed regime of strong non-linearity.
Our phase diagram ﬁgures out the relation between the training dynamics of multi-layer neural networks and their initialization, reveals the possibility of havingdifferent training dynamics for different layers within a neural network, and serves as a map thatprovide a guidance for the future study of deep NNs, such as how different dynamical regimes ofdifferent hidden layers affect the generalization.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for informationon how 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?
[Yes]
(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]
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, if applicable?
[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 amount spent on participant compensation?
[N/A]