Empirical Phase Diagram for Three-layer Neural Networks with Inﬁnite Width
Abstract
Substantial work indicates that the dynamics of neural networks (NNs) is closely related to their initialization of parameters.
Inspired by the phase diagram for two-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 for common initialization methods.
With carefully designed experiments and a large computation cost, for both synthetic datasets and real datasets, we ﬁnd that the dynamics of each layer also could be divided into a linear regime and a condensed regime, separated by a critical regime.
The criteria is the relative change of input weights (the input weight of a hidden neuron consists of the weight from its input layer 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, we also demonstrate that different layers can lie in different dynamical regimes in a training process within a deep NN.
In the condensed regime, we also observe the condensation of weights in isolated orientations with low complexity.
Through experiments under three-layer condition, our phase diagram suggests a complicated dynamical regimes consisting of three possible regimes, together with their mixture, for deep NNs and provides a guidance for studying deep NNs in different initialization regimes, which reveals the possibility of completely different dynamics emerging within a deep NN for its different layers.
Introduction
As neural networks (NNs) have been employed to extensive practical and scientiﬁc problems, such as computer vision, natural language processing, and automatic driving, etc.
Studying the rationale behind 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 and generalization 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 the initialization, therefore, the NN after training can be well approximated by the ﬁrst-order Taylor expansion at the initial 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 where parameters signiﬁcantly deviates from the initialization, therefore, the NN dynamics exhibits 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 three distinct regimes, i.e., linear regime, critical regime and condensed regime, based on the relative change of input weights (the input weight of a hidden neuron consists of the weight from its input layer 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 empirical study ﬁnds that orientations of the input weights condense in several isolated orientations, which is a strong 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 that deep networks can learn faster and generalize better than shallow networks in both real data and synthetic data (He et al., 2016; Arora et al., 2018; Xu and Zhou, 2021).
Because of multi-layer structure 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 to obtain clear phenomena as well, where multi-layer NNs could have distinct characteristics compared with shallow two-layer NNs.
For example, due to the existence of other layers, how is the dynamical regime of the weights of a layer in multi-layer NNs different from two-layer ones and can weights of two different layers experience distinct dynamical behaviors in one training process?
Answering these questions is important to understand how the dynamics of multi-layer NNs depend on the initial hyper-parameters, thus, leading to a deeper understanding of deep NNs and providing guidance for tuning parameters in practical training a tool for the theory community.
For instance, if we want to study the nonlinear effect of a hidden layer, we can set this interested hidden layer in the non-linear regime and all other hidden layers in the linear regime.
Further, we can continue to study how the non-linear effect of a hidden layer depends on its depth.
Three-layer network is a good starting point to study multi-layer NNs as it only has two hidden layers for studying the interaction between layers.
Besides, to have a clear boundary of different dynamical regimes, we focus on three-layer ReLU NNs with large-width limit.
A difﬁculty to overcome is to ﬁnd appropriate quantities required for describing a phase diagram, which should be able to identify distinct different regimes, and can also account for the same dynamical behavior when hyper-parameters are different but those quantities are the same.
To this end, we derive a normalized gradient ﬂow for a three-layer ReLU NN.
Since we only care about the training trajectory up 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 initialization scheme as the commonly used in practice, and we obtain two key independent quantities.
With carefully designed 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 a critical regime.
The criteria is the relative change of input weights as the width approaches 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 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 shown in Fig.
For deeper networks, this work suggests that the dynamical regime for the weights of each layer is similar to that of a two-layer NN and different layers can lie in distinct regimes in a training process.
For NNs with not too large width, the boundary between regimes is not sharp but this work still sheds lights on how the initialization affects the dynamics of weights in different layers.
Our study on phase diagram of three-layer ReLU NNs thus lays an important foundation for future work on the dynamical behavior of multi-layer NNs and corresponding implicit regularizations at each of the 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 class of activation functions, when the input dimension is at least the size of the training set, the training loss 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 a few 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 learning 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 (Zhou et al., 2021).
The embedding principle (Zhang et al., 2021a,b) shows that the loss landscape of an NN contains all critical points/functions of all the narrower NNs, thus providing a basis 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]
where x ∈ R d , α is the scaling factor, m is the width of both layers, θ =
vec { a , W [2] , W [1] } with a ∈ 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 x and 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 ∼ √m 1 + d , β 2 ∼ √ 2 m 1 + m 2 , β 3 ∼ √m 2 +1 , for α ∼ O (1) gradient ﬂow of NN exhibits 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 is equal, i.e. m 1 = m
= m .
Since the output is linear with respect to a , a is not the main concern, and the initialization methods used for a neural network in practice often has β 2 =
Bβ 3 , where B is a constant in dependent of m .
So, we take the following assumption.
Assumption 3.2.
β
= Bβ 3 , where B is a constant independent of m .
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 starting point.
Identiﬁcation of the coordinates is important for drawing the phase diagram.
There are some guiding principles for ﬁnding the coordinates of a phase diagram (Luo et al., 2021): (i) They should be effectively independent; (ii) Given a speciﬁc coordinate in the phase diagram, the learning dynamics of all the corresponding NNs statistically should be similar up to a time scaling; (iii) They should well differentiate dynamical differences except for the time scaling.
Guided by above principles, in this section, we perform the following rescaling procedure for a fair comparison between different choices of hyperparameters and obtain a normalized model with two independent 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]
where e i = (α a T σ ( W [2] σ ( W [1] x i )) − y i ) , the operation ⊙ is the Hadamard product.
Then, setting
[Equation]
[Table]
so that all the parameters are picked from standard normal distribution.
By the chain rule of derivative and the homogeneity of ReLU , i.e. σ ( a u ) = aσ ( u ) and σ ′ ( a u )
= σ ′ ( u ) , where a > 0 is a scalar and u is a vector, we obtain that:
[Equation]
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 are in general functions of m , which attains 0 , O (1) , + ∞ at m → ∞ .
To account for such dynamical difference 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]
where our assumptions 3.1 and 3.2 mean κ 1 = B , leading to γ 1 = lim m →∞ − log κ 1 log m = 0 .
So we discuss 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 with their 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 by a
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 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 a diversity of typical cases over the phase diagram regarding W [1] and W [2] respectively.
We use a simple 1-d problem of 4 training points, which allows an easy visualization.
[Figure]
The ﬁrst row in Fig.
2 shows typical learning results over different γ 3 ’s, from a relatively jagged interpolation (NTK scaling) to a smooth interpolation (mean-ﬁeld scaling) and further to a linear spline 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 in Fig.
For γ 3 = 1 . 0 , the initial scatter plot is very close to the one after training.
However, for γ 3 = 2 . 0 , active neurons (i.e., neurons with signiﬁcant amplitude and away from the origin) are condensed at a few orientations, which strongly deviates from 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 group feature, which can also reﬂect the characteristics between different γ 3 and γ 2 .
Considering that W [2] k is m + 1 dimensional, We cannot directly and intuitively show { W [2] k } mk =1 in a ﬁgure.
Here, W [2] k represents the input weight of the k -th neuron in the second hidden layer.
Therefore, to characterize the distance between different vectors in { W [2] k } mk =1 , we introduce D ( u , v ) to denote the inner product of two { W [2] k } mk =1 , as Eq.
As shown in Fig.
3, for γ 3 = 2 . 2 , W [2] k hardly has a strong relationship with each other, while for γ 3 = 2 . 8 , most of the W [2] k are very close to each other, showing strong condensation.
5.3 Phase diagram for three-layer neural networks
In this section, with γ 3 and γ 2 as coordinates, we characterize at m → ∞ the dynamical regimes of NNs and identify their boundaries in the phase diagram through experimental and theoretical approaches.
How to characterize and classify different types of training behaviors of NNs is an important 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 works have proved speciﬁc points in the three-layer phase diagram belong to the linear regime of
[Figure]
both W [1] and W [2] .
However, its exact range in the phase diagram remains unclear.
On the other hand, NN training dynamics can also be highly nonlinear at m → ∞ , such as widely used initialization method, Glorot and Bengio (2010), as a point shown in the phase diagram.
However, whether there are other points in the phase diagram that has similar training behavior is not well understood in three-layer condition.
In addition, it is still not clear if there are other regimes in the 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-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.
[Figure]
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 .
1 , the corresponding slopes grow as well.
[Table]
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.
[Table]
we deﬁne
[Equation] [Equation]
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 show result 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.
The result for the same 1-d problem as in Fig.
2 is presented in Fig.
In the red zone, where S W i is less 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 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 ) } mk =1 for two different regions in Fig.
As is shown in Fig.
7, on the left side of 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 a few orientations, which strongly deviates from the initial scatter plot.
The graphs in the blue boxes present an intermediate state or process state, making the above-mentioned change a continuous 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 the average 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 with phase distribution in Fig. 8.
W [1] k } mk =1 at width 10000.
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 dots in the small ﬁgure correspond to the initialization parameters, and the green dots correspond to the post-training parameters.
Since results with initialization by Gaussian distribution times averaged the results, ﬁnally obtained the
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.
The homogeneity of ReLU is is an important basis for deriving the coordinate of the phase diagram, so it could not be replaced by any non-linearity activation function.
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 having different training dynamics for different layers within a neural network, and serves as a map that provide a guidance for the future study of deep NNs, such as how different dynamical regimes of different hidden layers affect the generalization.
Acknowledgments and Disclosure of Funding
This work is sponsored by the National Key R&D Program of China Grant No. 2019YFA (Z. X.), the Shanghai Sailing Program, the Natural Science Foundation of Shanghai
Grant No.
20ZR1429000 (Z. X.), the National Natural Science Foundation of China Grant No.(Z. X.), the National Natural Science Foundation of China Grant No. 12101401 (T. L.), {+Shanghai Municipal Science and Technology Key Project No. 22JC
 (T. L.),+} the National Natural Science Foundation of China Grant No.
(Y. Z.), Shanghai Municipal of Science and Technology Project Grant
No. 20JC1419500 (Y.Z.), the Lingang Laboratory Grant No.LG-QS-202202- (Y.Z.), Shanghai Municipal of Science and Technology Major Project No. 2021SHZDZX0102, and the HPC of School of Mathematical Sciences and the Student Innovation Center at Shanghai Jiao
Tong University.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on 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 referencing the 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 are proprietary.
• 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 instructions block 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] 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] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] In supplementary.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were 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] 4.
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] 5.
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]