M ODEL A RCHITECTURE C ONTROLS G RADIENT D ESCENT D YNAMICS : A C OMBINATORIAL P ATH -B ASED F ORMULA
Anonymous authors
Paper under double-blind review
A BSTRACT
Recently, there has been a growing interest in automatically exploring neural network architecture design space with the goal of ﬁnding an architecture that improves performance (characterized as improved accuracy, speed of training, or resource requirements).
However, our theoretical understanding of how model architecture affects performance or accuracy is limited.
In this paper, we study the impact of model architecture on the speed of training in the context of gradient descent optimization.
We use the ideas from prior work that shows gradient descent can be modeled as a ﬁrst-order ODE and use ODE’s coefﬁcient matrix H to characterize the convergence rate.
We introduce a simple analysis technique that enumerates H in terms of all possible “paths” in the network.
We show that changes in model architecture parameters reﬂect as changes in the number of paths and the properties of each path, which jointly control the speed of convergence.
We believe our analysis technique is useful in reasoning about more complex model architecture modiﬁcations.
I NTRODUCTION
Gradient descent and its variants are the cornerstones of deep learning.
The theoretical properties of gradient descent have been widely studied in the literature; study of the convergence bounds and guarantees [6, 12, 15, 13], the characterization of the local geometry of stationary points [21, 11, 10, 19, 25, 30, 24, 23, 13], exploration of better algorithms to optimize the descent process [28, 16, 26, 33, 7] are just a few example of active research areas in this domain.
Another major research area is the exploration of network architecture and its impact on performance.
In recent years, network architecture search and design have shown major performance boosts in various aspects of deep learning [36, 35, 18, 8, 34].
Recent theoretical results are trying to relate model architecture and gradient descent properties.
It has been shown that over-parametrization in width guarantees convergence in deep neural network [15].
It has also been shown that increasing depth has an impact similar to adding momentum optimization and adaptive learning rate to the objective function [4].
The gradient descent process can be described by a system of ﬁrst-order differential equations in the continuous-time limit, in the form of ˙x ( t )
=
H ( t ) x ( t ) .
In this paper, we study the properties of the coefﬁcient matrix H ( t ) , which governs the dynamics of the gradient descent.
Speciﬁcally, we formulate H ( t ) in terms of all possible paths in the network.
This representation is powerful as it enables us to analyze the gradient behavior symbolically and abstracts away the unnecessary bookkeeping of derivatives and weights.
Speciﬁcally, we use this representation to explore the followings.
(1) We study the impact of width on convergence.
Using a simple path counting argument, we show that H ( t ) is a sum of m i.i.d.
terms where m is dictated by the width of the network.
This has been implied in the work of Simon Du, et al.
[12] for the special case of a 2-layer RELU-activated network.
We contrast the ease of arriving at such conclusions using our representation against methods used in previous work [12, 15, 13, 6].
(2) We study the impact of depth on convergence.
While prior work [3] established the relationship between momentum and depth, we show a direct relationship between depth-induced momentum and the scale of the output value. We also discuss how depth-induced momentum is different than explicit momentum. (3) We argue why the number of paths is more important than the number of nodes in a network.
R ELATED W ORK
Related work can be categorized into three groups which we will explain in detail below.
In general, the common theme across previous works is to show that if some conditions, such as bound on models’ width or depth or assumptions on geometrical properties of the loss landscape are met, gradient descent or its variance converges to a global minimum.
Landscape Analysis Researchers have extensively studied the properties of error surface in linear neural network ([10, 21]) and have proved convergence to a global minimum for non-convex optimizations if error surface shows certain properties, such as all saddle points to be strict (i.e. there exists a negative curvature)([11, 19, 25, 30, 24, 23, 13]).
However, it has been shown that the strict saddle property is not guaranteed for deep neural networks.
Width Analysis
A number of recent work have characterized minimum width that guarantees convergence in terms of model architecture parameters and input data.
We on the other hand study how convergence rate varies with width.
Examples of ﬁnding width bounds include [15] where it has been shown that for a 2-layer ReLU-activated neural network with squared loss, as long as every hidden layer is wide enough, gradient descent converges to a global minimum at a linear rate.
In [14, 1], it is extended to fully connected linear networks and ResNets.
Du et.
al.
[12] has also generalized this work to an L-layer fully connected linear neural network and shows the convergence rate as a function of depth, output dimension and least eigenvalue of the Gram matrix for large enough hidden layers.
Ardalani et.
al.
[2] have shown empirically that for a wide range of RNN applications, increasing width will reduce the number of steps to minimum validation loss.
Neural tangent kernel [22] studies convergence in the inﬁnite width limit.
Depth Analysis Bartlett et al.
[6] have proven for a deep linear neural network with isotropic input and identity initialization that the number of steps to (cid:15) -proximity of the best answer is polynomial in the number of layers.
Neural ODE studies the convergence in the inﬁnite depth limit [27, 9, 20, 17].
Work that motivated our depth analysis are [3, 4]. They have shown that over-parametrization accelerate training with these extra restrictions: (1) linear neural network with (2) fully-connected layers and (3) "balanced" initialization where (4) dimensions of hidden layers are at least the minimum of the input and output dimensions. We contribute in the following ways. We characterize depthinduced momentum in value space with an inequality and compare with explicit gradient descent momentum.
We observe that momentum is directly impacted by value of the sub-network.
We lift the structure and activation restrictions.  N OTATIONS AND C ONVENTIONS
Matrix entries are denoted with row and column indices, for instance H
=
( H ij ) is a matrix with entries H ij .
When H is diagonalizable, we denote by λ min ( H ) its smallest eigenvalue.
For a matrix W , W
( i, :) and W (: , j ) denote the row vector in the i -th row and the column vector in the j -th column, respectively.
Write (cid:80) w ∈ W for summing over all elements of W .
We use pairs { X i , y i } , 1 ≤ i ≤ N to represent labeled data where X i is input and y i is the output.
We write x i,k for the k -th component of X i .
We denote loss by L .
We work with feed-forward neural networks, which we sometimes denote abstractly by a function f , or f ( w, x ) as a function of input x and weights w .
We denote network weights individually by w i .
Predictions are denoted by u i =
f
( w, X i ) 1 ≤ i ≤ N .
The prediction on all inputs is denoted by the vector u .
l p loss is deﬁned as (cid:80) ( y − u i ) p for even integers p ≥ 2 .
We use angle brackets to denote inner products, i.e. (cid:104)−→ v , −→ w (cid:105) =
(cid:80) i v i w i represents inner product of two vectors.
i.i.d.
random variables stands for independent and identically distributed random variables.
ODE stands for ordinary differential equations.
NN stands for neural networks.
GD stands for gradient descent.
B ACKGROUND
We build upon prior work [15, 12, 14, 3, 31] which take gradient descent to a continuous time limit and derive the Gram-matrix H governing gradient descent dynamics. In this section, we give a quick derivation for Equation (4) for completeness.
Convergence rate can be characterized as the rate of change in loss during the training process.
If steps are inﬁnitesimally small, the rate of change in loss can be characterized as:
[Equation]
Where u i ( t ) =
f ( w ( t ) , X i ) is the prediction on input X i at time t .
For simplicity we focus on the dynamics of du i dt , which can be expanded as follows:
[Equation]
We use the deﬁnition of gradient descent and chain rule to derive the equation above.
Gradient descent at inﬁnitesimal small step is deﬁned by:
[Equation]
Observe that −→ ∂L∂u is a column vector.
Hence, if we deﬁne H ij as:
[Equation]
Then we can simplify equation (2) in a vector form as:
[Equation]
Assuming L is l 2 loss, equation (5) simpliﬁes to the following.
This result will hold up to a constant for other l p losses (see Appendix A).
[Equation]
Notice that this puts y − u in a system of differential equations whose single variable analogue is:
[Equation]
The above equation’s solution is:
[Equation]
when h < 0 , e h < 1 and let e h = 1 − ξ .
ξ governs the convergence rate in Equation (9).
In the actual system of equations (7), ξ will be determined by the minimum eigenvalue of H ( t ) , denoted λ min ( H ( t )) .
M AIN I DEA : C HARACTERIZING H IN TERMS OF N ETWORK P ATHS
In order to state the path decomposition formula, we ﬁrst establish some notations.
• We denote the neural network abstractly with a function f ( X ) .
• For any neural network, we view the underlying network as a directed graph where edges are in the backward propagation direction.
• We deﬁne path p , as a series of connected nodes and directed edges.
We deﬁne an output path as a path that starts with the output node.
We denote the length of p by l ( p ) .
• We index the nodes as follows.
Index the output node by 0, the one after by 1, and so on.
We denote the j -th node in the i -th path by p ij .
• We allow different activation functions at nodes.
We denote the activation function at p s by σ s .
• An example path that ends in an input node may look like this:
[Equation]
[Equation] σ... σ s X i,k www s • The post-subnetwork at w s is deﬁned as the network formed by the subgraph rooted at the right vertex of w s .
The pre-subnetwork at w s is deﬁned as the union of all output paths ending in w s , denoted by pre ( w s ) .
• For every input
X i , we denote the activation value at the node after w s with f [ w s ]( X i ) .
Note that this is equivalent to the value at the right vertex of w s once X i propagated through the post-subnetwork at w s .
• We denote the local gradient wrt.
input at node s along the path p with τ ip s .
• Given the pre-subnetwork at w s , we write the explicit formula for ∂u i ∂w s :
[Equation]
where path _ gradient p ( w s , X i ) is the gradient value along the path p at the left vertex of w s , and can be expanded as follows using the chain rule:
[Equation]
we refer to path _ gradient p ( w s , X i ) as PG p ( w s , X i ) for brevity from now on.
Theorem 5.1.
For a network f ,
[Equation]
which easily follows by substituting Equation (10) in Equation (4).
In the rest of the paper, we will make many arguments most of which rely on the number of terms in this equation.
Number of terms in Equation 12 is itself controlled by the number of paths in the pre-subnetworks at each weight.
Following examples provide some intuition.
Example 5.2.
In this example, our network has a structure as illustrated in Figure 2a.
The output node is denoted with value one, indicating that the activation function at output layer is an identity function.
As shown in Figure 2b, there are only one path in each weight’s pre-subnetwork.
Therefore, there are only three terms in Equation 12, where each summand is of form PG p ( w, X i )
PG p ( w, X j ) .
We refer to these terms as symmetrical pairs since both multipliers are representing the same path.
Example 5.3.
A more complicated application of the formula can be found in Figure 2d.
There is one path in each of the following weights’ presubnetwork: w 0 , w 1 , w 2 , w 3 , and w 4 , and three paths at w 5 ’s presubnetwork.
Therefore, Equation 12 would have 5 + 3 2 = 14 summands, 8 of which are terms with symmetrical pairs and 6 are cross-terms induced by the paths in w 5 ’s pre-subnetwork.
In general, cross-terms will appear only if there are two paths that share a ﬁnal edge which can only happen if there are three or more number of layers in the network.
σxxwww(a) Example 5.
σ 1 wpre ( w 0 ) σxwwpre ( w 1 ) σxwwpre ( w 2 ) (b) Paths Example 5.
σσσx w 0 www 3 ww(c) Example 5.
σ 0 wpre ( w 0 ) σwpre ( w 1 ) σ 2 wpre ( w 4 ) σσwwpre ( w 2 ) σσ 2 wwpre ( w 3 ) σσx www σx wwpre ( w 5 ) σσx www(d)
Paths Example 5.
A RCHITECTURE I MPLICATIONS 6. C ONVERGENCE R ATE S CALES L INEARLY WITH W IDTH
In this Section, we show that for sufﬁciently wide neural networks, convergence rate has a linear relationship with the network’s width.
Results are drawn theoretically at time 0 assuming random i.i.d initialization.
For these results to generalize across time, we assume that H stays close to its initial value.
In the Appendices D and E, we show empirical results that suggest this assumption is true.
For the extreme case of inﬁnitely-wide networks, it has been shown that H is constant throughout training [22].
Others [5, 15] have also shown that for sufﬁciently-wide neural networks weights remain close to initialization, and so our argument holds true across training.
Proposition 6.1.
For a 2-layer neural network with hidden dimension m , the expected convergence rate scales linearly with m .
To demonstrate the simplicity and expressiveness of our approach, next we will discuss proofs with and without our approach.
Brute-force approach A 2-layer fully-connected linear neural network can be modeled as:
[Equation]
where V ∈ R 1 ,m and W ∈ R m,d are the weights of the second and ﬁrst layer, respectively and the input
X is a d -dimensional column vector.
We begin the analysis with the simple case of a 2-layer linear network where the network’s width is one, i.e. m = 1 and u ( X ) = v 0 · (cid:104) W (: , 1) , X (cid:105) .
H ij can be decomposed into two terms, partial derivatives with respect to weights of the ﬁrst layer and partial derivatives with respect to weights of the second layer.
[Equation]
For a more general case where the width is m , we have:
[Equation]
[Equation]
Adding non-linearity would make this equation even more complex, which is the approach used in previous literature together with induction:
[Equation]
See Appendix B for details of how we derive these equations.
Since we assume all weights ( w lk , v 0 ) are independent and identically distributed, the summands within Equations (15) and (16) are also independent and identically distributed.
This indicates that H ij is composed of m i.i.d components, hence convergence rate scales linearly with m .
Our approach: Using Theorem 5.1 More generally, by applying Theorem 5.1, it is immediate to see every pre-subnetwork will be a path in a 2-layer fully connected network.
There are two types of paths; m length 1 paths and m length 2 paths.
The weights are independent and follow the same distribution, hence, a 2-layer network’s H matrix always decomposes as m i.i.d sums.
Hence convergence rate scales linearly with m .
This proves Proposition 6.1.
(cid:3)
See Section 7 for empirical analysis that motivated this study.
We show for a 2-layer linear and non-linear networks that convergence rate has linear relationship with width.
6. D EPTH - INDUCED F IRST -O RDER M OMENTUM
In this section, we use our formula to explain how depth has a momentum-like effect.
Take a network f .
Construct a new network g by adding a multiplication layer on top, i.e. g ( X ) = µf ( X ) with an extra trainable parameter µ .
We ﬁrst present an inequality (with the proof left to Appendix F), then explain why the inequality implies momentum-like behavior and ﬁnally how they are different. Proposition 6.2. With u g = µu f ,
[Equation]
[Equation]
where cos denotes cosine between two vectors.
Remark 6.3 (Momentum interpreted as directional preference) . In physics, momentum is deﬁned as the tendency of an object to follow its current velocity of motion. In gradient based optimization, momentum can be deﬁned as in Equation (1) of [32],
[Equation] where v ’s are the updates, (cid:15) is the learning rate and m 0 is the momentum coefﬁcient. Quantitatively, cos ( v t +1 , v t ) increases with m 0 .
A bigger m 0 implies a stronger preference in v t +1 for the direction of v t . For network depth, Proposition 6.2 describes the stronger directional preference of the deeper network g in the direction of the current value vector.
In the depth-induced momentum experiment below, we show that this directional preferences can have similar effects on convergence rates as momentum in gradient based optimization.
Remark 6.4 (Difference to momentum) . The key difference between depth-induced momentum and gradient descent momentum is the following. The depth-induced momentum is a condition between u (cid:48) and u , whereas gradient descent momentum is a second order condition between u (cid:48)(cid:48) and u (cid:48)
(see [29] Eq.
(4) and (10)).
Remark 6.5 (Sensitivity to output scale) . Since (cid:104) u (cid:48) , u (cid:105) = 12 ddt (cid:104) u, u (cid:105) = 12 ddt (cid:107) u (cid:107) 2 , depth-induced momentum will be more pronounced when (cid:107) u (cid:107) can increase. In Appendix G, we present an experiment that shows that depth-induced momentum is more pronounced when output scale is larger.
6. O N THE I MPORTANCE OF N UMBER OF P ATHS FOR C ONVERGENCE
Here, we will provide some heuristic arguments, based on our formula, on why over-parametrization through increasing the number of paths can be more effective than increasing the number of nodes (parameters) and conclude that increasing depth is more effective than increasing width for accelerating the training process.
We ﬁrst introduce an important terminology which we will use below.
Deﬁnition 6.6.
Given a collection of vectors { g i } , the Gram matrix G is deﬁned as a matrix whose ij -th entry is (cid:104) g
i , g j (cid:105) under the Euclidean inner product.
Remark 6.7.
Based on the deﬁnition above it is easy to see that Gram matrix is symmetric and positive-semideﬁnite and its rank is exactly the dimension of the subspace spanned by { g i } .
Suppose
H =
H 1 + · · · + H n are all Gram matrices.
Then H is positive deﬁnite if any of the H i is positive deﬁnite.
This is because of supperadditivity of the minimum eigenvalue function,
i
[Equation]
e. λ min ( H ) ≥ λ min ( H 1 ) + · · · + λ min ( H n ) .
Strongly-Gram vs. Weakly-Gram Equation(12) in Theorem 5.1 can be expanded further as follows and the terms can be classiﬁed into two groups.
The ﬁrst group will include all the symmetrical pairs and the second group will include all the cross-terms.
[Equation]
We refer to the ﬁrst group as strongly Gram , and denote it by H Sij ( f ) , and refer to the second group as weakly Gram and denote it by H Wij ( f ) , i.e. we have:
Remark 6.8 (Gramness of the strongly-Gram component) .
Suppose there are q paths in the network, p 1 , p 2 , . . .
, p q .
Then Equation (19) implies that the strongly Gram component is Gramian with vectors
[Equation]
Remark 6.9 (Insigniﬁcance of the weakly Gram component) .
We claim that the weakly Gram component is insigniﬁcant, see Appendix C for heuristics and empirical evidence.
Remark 6.10 (Number paths vs. number of nodes) .
The Gram matrix H S should ideally be of full rank for good convergence.
So we want the subspace spanned by { g i } to have dimension at least
N .
N is the dimension of H and also the number of sample points.
On the other hand, as shown in
Equation 21 g i ’s are embedded as vectors in a q -dimensional vector space where q is the number of paths.
Therefore, to ensure H is full-rank, q should satisfy q ≥ N .
The bigger q is, the more likely they are linearly independent and make the matrix full rank.
In deep learning, N is usually large, therefore, we need to have networks whose number of paths is comparable to N for best convergence.
Remark 6.11 (Depth vs. width) .
Adding a new parameter at a new layer can increase the number of paths more effectively than adding it at an existing layer.
In other words, making models deeper increases the number of paths faster than making models wider.
In this section, we provide experimental results that demonstrate the impact of model structure on convergence.
Setup We employ a teacher-student framework to put our ideas to test.
Our teacher model is a two-layer neural network with hidden dimension 100, input dimension 50 and output dimension 1.
We assume a Gaussian distribution both for the inputs and ground-truth weights.
The student models are also two-layer neural networks with variable hidden dimensions, and l 2 loss.
The convergence rate at step t is deﬁned as 1 − Loss t +1 Loss t .
[Figure] [Figure]
Two-Layer Linear Neural Network As shown in Figure 3a, training curve in log scales look like straight lines.
As shown in Figure 3b, convergence rate ( λ ) grows linearly with width.
Two-Layer Non-Linear Neural Network As shown in Figure 4a, training curve for models with the sigmoid non-linearity begins in a steep region where loss drops very quickly for a short period of time and ends in a steady region where loss drops very slowly.
We are interested in characterizing the curve in the steady region.
As shown in Figure 4b, convergence rate ( λ ) grows linearly with width.
Depth induced momentum Figure 5 compares the convergence rate curves for 2-layer, 3-layer and 4-layer NN using GD, against a 2-layer network with a momentum optimizer.
The deeper networks are constructed from 2-layer networks with additional scalar multiplication layers.
As discusses in Section 4.5, one can observe depth-induced momentum in convergence rates, similar to using momentum optimizer with a 2-layer NN.
C ONCLUSION
In this paper, we provide a path-based combinatorial formula for the coefﬁcient matrix H which governs the gradient descent dynamics.
We advocate analyzing H with this path-based approach by showing the insights we could derive from this.
A l p LOSS
We argue very brieﬂy that without loss of generality, we can assume that we work with l 2 loss instead of other L p norms.
For l p loss,
[Equation]
So for l 2 loss, Equation 5 becomes
[Equation]
Recall the argument for decay rate under Equation 8.
The single variable analogue of the system is:
[Equation]
24) the solution looks like:
[Equation]
and decays like:
[Equation]
with
[Equation]
When h is small, by Taylor expansion
[Equation]
Instead of using the closed form solution, this ODE could equivalently be solved using the series method and we would be getting the Taylor expansion of equation (25) on the right and equation (28) remains the same.
Note that the series method (more fundamentally, Taylor expansion involving only one matrix) works for both ODEs and systems with convergent assumptions on the coefﬁcient matrix
H .
ηH is small by assumption, so we have convergent Taylor expansions, so we continue to argue for l p loss as if we are in the single variable case.
For general l p loss, we have:
[Equation]
[Equation]
then
[Equation]
) and the decay rate would be h p , scaled down by p , so all the convergence rate statements we make for the rest of the paper for l 2 remains true if scaled by p 2 .
B 2 LAYER LINEAR NETWORK
The detailed brute-force calculation of H ij for a 2-layer fully connected linear network with m =is as follows.
Recall m is the width of the hidden layer.
[Equation] [Equation]
[Figure]
For a general m , we have:
[Equation]
C W EAKLY G RAM MATRIX IS EMPIRICALLY INSIGNIFICANT
In terms of network architecture, the weakly Gram term would only appear when two paths share a ﬁnal edge.
That means the network needs to have at least 3 layers.
The number of pairs could grow quadratically with width.
Heuristically, each term E [PG p 1 ( w, X i ) · PG p 2 ( w, X j )] tend to be 0 because of independent weights in PG and the weights having expectation 0.
Hence, we focus on the strongly Gram part for the rest of the section.
As in Equation (19), the number of terms in weakly Gram component grows quadratically in the width of the hidden layer.
One might expect this results in a quadratic growth in convergence, given all the terms are i.i.d.
random variables and strictly positive.
In order to test this hypothesis, we exploit a teacher-student framework as outlined in Section 7, however the student models are 3-layer RELU-activated networks to ensure the number of terms in weakly Gram component is non-zero.
We choose RELU activation function to ensure that the expected value of the activation to be greater than zero.
Our experimental results shows that the convergence rate relationship with width is linear (see Figure 6) and not quadratic.
This implies that weakly Gram components have expectation zero.
D J USTIFY H STAYS CLOSE TO INITIALIZATION
Weights and H are indeed time-dependent and follow their respective stochastic process.
However, theoretical results in the overparametrized regime suggest that every weight vector remains close to initialization (Lemma 5.3 [5] and Lemma 3.3 [15]).
We also observe that the mean and variance stay very close to initialization, N (0 , 0 . 1) (Figure 7).
In Figure 8, we ran the Kolmogorov-Smirnov test to compare the weight distribution against the initialization distribution.
As shown, the p-values are overall stable and far from signiﬁcance levels to invalidate the null hypothesis ( H 0 : trained weights follow N (0 , 0 . 1) ).
[Figure]
E J USTIFY IID ASSUMPTION FOR WEIGHTS
The Turning Point Test is commonly used to check iid-ness.
Figure 9 shows the p-value through the training process.
As shown, p-values do not exhibit any major regime shifts and the null hypothesis (the weights are iid samples) stays far from rejectable with any commonly used signiﬁcance levels (0.1, 0.05 ...).
iid-ness can also be seen as a consequence of closeness to initialization in response 1.
By staying close to initialization, statistical properties are preserved and the weights remain an iid sample of the original distribution.
F P ROOF OF P ROPOSITION 6.
Proof. By deﬁnition of f and g , in g , there’s a new path of length 1 whose edge is µ and all paths in f are annexed by µ at the front.
Hence H ij for the new network g [-
-]can be computed in terms of H ij for the network f as follows:
[Equation]
This can be written in the matrix form as:
[Equation]
We ﬁrst calculate an inner product, [Equation] Now we calculate cosines, [Equation]
Using equation (34) to expand the denominator,
[Equation] It sufﬁces to check [Equation] which after cross-multiplication, reduces to:
[Equation]
This follows from Cauchy–Schwarz inequality. Equality holds precisely when H ( f ) u f is parallel to u f (Equation (47), (44)). Remark F.1. One further remark that could be made after Remark 6.3 and 6.4 is the following.
The proof above tells us that depth-induced momentum comes mainly from the extra term u f in
Equation (34).
Note that u f s are the values of the networks at sub-network f .
Comparing this with our path-based expansion formula at Equation (12) where all sub-networks’ values appear with different scaling factors proportional to their path gradients sum, we can conclude qualitatively that this phenomenon is present across all sub-networks and each sub-network contributes to overall momentum proportional to its path gradients sums.
G T HE I MPACT OF O UTPUT V ALUE D ISTRIBUTION ON D EPTH -I NDUCED M OMENTUM
Figure 10 shows the impact of output distribution on momentum.
The underscore separated strings represent hidden dimensions.
"2000" represents hidden dimension of width 2000 and "2000_1" represents hidden layers of width 2000 and 1.
The number after "std" represents the output’s standard deviation in the teacher network, equivalently scale as the mean is 0.
As we can observe from the graph, momentum acceleration is not present for small standard deviation (0.06 and 0.12) and the behavior is more similar to a shallow 2-layer network of the same width.
When standard deviation is raised to 0.24, then we observe acceleration of the convergence rate.
In all of these experiments, µ is initialized to 1 and empirically stays close to 1 to counteract the effect of µ 2 in equation (34).