Dynamically Stable Infinite-Width Limits of Neural Classifiers
Authors: Eugene Golikov
Venue:  (2021)
Keywords: neural tangent kernel
mean field limit

Abstract: Recent research has been focused on two different approaches to studying neural networks training in the limit of infinite width (1) a mean-field (MF) and (2) a constant neural tangent kernel (NTK) approximations. These two approaches have different scaling of hyperparameters with the width of a network layer and as a result, different infinite-width limit models. Restricting ourselves to single hidden layer nets with zero-mean initialization trained for binary classification with SGD, we propose a general framework to study how the limit behavior of neural models depends on the scaling of hyperparameters with network width. Our framework allows us to derive scaling for existing MF and NTK limits, as well as an uncountable number of other scalings that lead to a dynamically stable limit behavior of corresponding models. However, only a finite number of distinct limit models are induced by these scalings. Each distinct limit model corresponds to a unique combination of such properties as boundedness of logits and tangent kernels at initialization or stationarity of tangent kernels. Existing MF and NTK limit models, as well as one novel limit model, satisfy most of the properties demonstrated by finite-width models. We also propose a novel initialization-corrected mean-field limit that satisfies all properties noted above, and its corresponding model is a simple modification for a finite-width model.

Under review as a conference paper at ICLR 2021
DYNAMICALLY STABLE SHENITE-WIDIHH LIMITS OF NEURAL CLASSIFIERS
Anonymous authors Paper under double-blind review
ABSTRACT
Recent research has been focused on two different approaches studying neural networks training in the limit of infinite width (1) a mean-field (MF) and (2) a constant neural tangent kernel (NTK) approximations. These two approaches have different scaling of hyperparameters with the width of network layer and as a result, different infinite-width limit models. Restricting ourselves to single hidden layer nets with zero-mean initialization trained for binary classification with SGD, we propose a general framework to study how the limit behavior of neural models depends on the scaling of hyperparameters with network width. Our framework allows us to derive scaling for existing MF and NTK limits, as well as an uncount- able number of other scalings that lead to dynamically stable limit behavior of corresponding models. However, only a finite number of distinct limit models are induced by these scalings. Each distinct limit model corresponds combination of such properties as boundedness of logits and tangent kernels at initialization or stationarity of tangent kernels. Existing MF and NTK limit models, as well as one novel limit model, satisfy most of the properties demonstrated by Enite-width. models. We also propose novel initializatioo-ccreeett mean-field limit that satisfies all properties noted above, and its corresponding model is a simple modification for a mnite-width model.
unique
INTRODUCTION
For a couple of decades neural networks have proved to useful in a variety of applications. However, their theoretical understanding is still lacking. Several recent works have tried to simplify the object of study by approximating a training dynamics of finite-width neural network with its limit counterpart in the limit of large number of hidden units; refer it as "infinite-width" limit. The exact type of the limit training dynamics depends on how hyperparameters of the training dynamics scale with different types of limit models have been already extensively discussed in the literature: an NTK model (Jacot et al. 2018) and a mean-field field limit model (Mei et al. 2018 2019 Rotskoff & Vanden-FFijndnn 2019; Sirignano & Spiliopoulos 2020 Chizat & Bach 2018 Yarotsky 2018). A recent work (Golikov] 2020) attempted to provide a link between these two different types of limit models by building framework for choosing scaling of hyperparameters that lead to "well-defined" limit model. Our work is the next step in this direction. We study infinite-width limits for networks with single hidden layer trained to minimize cross-entropy loss with gradient descent. Our contributions are following.
width. In particular,
We develop framework for reasoning about scaling of hyperparameters, which allows one to infer scaling parameters that allow for dynamically stable model evolution in the limit of infinite width. This framework allows us to derive both mean-field and NTK limits that have been extensively studied in the literature, well as the "intermediate limit" introduced in Golikov (2020).
2.
Our framework demonstrates that there are only 13 distinct stable model evolution equations in the limit of infinite width that can be induced by scaling hyperparameters of a unite-width model. Each distinet limit model corresponds to region (two-, one-, or zero-dimensional of green band of the Figure left.
3. We consider a list of properties that are statisfied by the evolution of finite-width models, but not generally are for its affinitewwidth limits. We demonstrate that mean-field and NTK

---

Under review as a conference paper at ICLR 2021
moon ting kerrel, Fetto (2)
train size 1024, batch size 1024
meon sts testoot. 4.)na))
MF
100
sym-dcfault
finite logits
10-1
initializaiion
12a
101
12
10'
10a
10'
dynamicall ally stable model cvolution
trening dtep, k+
training step. k+1
train size
1024, batch size 1024
finite tangent kernels initialization
evolving
E.AAAR xij
10.
1.0
kernels
0.5
NTK
logits kernels are of the same order initialization
0.6
10
10a
10" training stop, k+1
10'
102
102
312
112
Uniit stop. ko11
Figure I: A diagram on the left specifies several properties demonstrated by finite-width mod- els. As plots on the right demonstrate, our novel IC-MF limit model satisly all of these prop- erties, while MF and NTK limit models, as well as sym-default limit model presented in the paper violate some of them. Left: A band of scaling exponents (go. that lead to dynamically stable model evolutions in the limit of infinite width, as well as dashed lines of special properties that corresponding limits satisfy. Three colored points correspond these properties. Right: Training dynamics of three models that correspond to color points on the left plot, as well as of wawaaliiaaiinccorreett mean-field model (IC-MF), which does not correspond to any point of the left plot. These models results of scaling of reference model of width d 2 (black line) up to width = 26 (colored lines). Solid lines correspond to the test set, while dashed lines are for the train set. Note that we have added small vertical displacement all curves in order to make them visually distinguishable. See Appendix for details.
limit models that satisfy most of
limit models, as well as "sym-default" limit model which was not discussed in the literature previously, special in the sense that they satisfy most of these properties among all limit models induced by hyperparameter scalings. We propose a model modification that allows for all of these properties in the limit of infinite width and call the corresponding limit iinitaaiiatiionccrreed
mean-field limit (IC-MF)".
4. We discuss the ability of limit models approximate the training dynamics of finite-width ones. We show that our proposed IC-MF limiting model is the best among all other possible limit models.
While our present analysis is restricted to networks with single hidden layer, we discuss high-level plan for generalizing it to deep nets, as well as an expected outcome of this research program, in App.
2 TRAINING ONE HIDDEN LAYER NET WITH SGD
Here we consider training a one hidden layer net fa with d hidden units with SGD. We assume the hyperparameters, namely, initialization variances and learning rates, scaled as power-laws of Each scaling induces a limit model foe lim fa. We present a notion of dynamical stability, which states that the change of logits after a single gradient step comparable to logits themselves. We derive a necessary condition for dynamical stability in terms of the power-law exponents of hyperparameters. We then present list of conditions that divide the class of scalings into 13 subclasses; each subclass corresponds to a unique distinct limit model.
Consider one hidden layer network:
f(x: a. W) - -aFoWWTx Σ c,o(wxx),
where x E Rd, W [w....., wal E Rdxxd xd and a = la,...., ......,,!!! az. E Rd, We assume a nonlinearity be real analytic and asymptotically linear: dbz) Az-co (2). Such nonlinearity
2

---

Under review as a conference paper at ICLR 2021
+ for > O. This is technical assumption introduced to simplify proofs. Note that we have used traditional leaky ReLUs in our experiments: App. Forr details. We assume the loss function Ely, z) to be the standard binary cross-entropy loss: {(y ((y.z)= = In(1 + e v*), where labels E {-1,1}, The data distribution loss is defined L(a, W) E Exxy-Dly fx; W)). We assume that the data distribution 'D does not depend on width d.
can be, e.g. "leaky softplus": d(2) In(1 + +*) a
Weights are initialized with isotropic gaussians with zero means: wi NN o Vr = l...d. The evolution of weights is driven by the stochastic gradient descent (SGD):
o0 ot
A = 0((++)) - ((()) -100 Oelya f(x (1)
(k)
a, W)).
(x8 (1) (k) (x00
D,
00
where 0 is either a or W. We assume that gradients for and W are estimated using independent
data samples (x. ya (lk) and (xx ym) While this assumption is indeed non-standard, we note
that corresponding stochastic gradients still give unbiased estimates for true gradients. Moreover, we have used either full-batch GD or standard mini-batch SGD in our experiments: see App. F for details. Define:
(k)
Or
w
(8)
WT
11a
710
iww
the
(3)
Fa
02
o
Then the dynamics transforms to:
A +jo ae(ye f(x o F W))
(1)
(k),
(xg .39 ~ D,
(4)
00,
while scaled initial conditions become: at ~ N(0, 1), w: ~N(O,I) Vr = vr== 1...d By expanding gradients, we get the following:
A (k) = -, (h) (1) v ola (law (K).Txx x!1))
fa ((x!!
a ~ N(0, 1),
(5)
Aw(h) o
w V (k) ((x! ( al.) (x...)))
w ~ N(..)).
(6)
(k)
La ((x. v)
ak(y,2)
-
(x) 0a
a 0(0 w ( (x)x))))..
Oz
=x (x) 1 + exp (x)))
Without loss of generality assume ae = 1 (we can rescale inputs x otherwise). We shall omit a subscript of aa from now Assume hyperparameters that drive the dynamics obey power-law dependence on d:
(dd =0 (d/d")a, (d( nh (d/d")&,,
- (d/d")0=
Given this, network of width d' has hyperparameters o² and navi- Here and then we write "a V w" meaning "a or w".
This assumption is quite natural: for He initialization (He et al. 2015) commonly used in practice d/1/22 while we keep learning rates in the original parameterization constant while changing width by default: Naxw = const, which implies n, x d and fhv xx d. On the other hand, NTK scaling (Jacot al. 2018 Lee et al. 2019) requires scaled learning rates be constants: ńav x d"
Scaling exponents (gr; Ta, T.)) together with proportionality factors (d", ad" -. Nw) define limit model | (x) lima-voo (())()),))) (x). We call a model "dynamically stable in the limit of large width" if it satisfies the following condition which we state formally in Appendix A
Condition (informal version of Condition in Appendix
Aj. Let A (1) (x) = f1 ((++))()))))))) (x) (x).
(x).
As(k) (k)
3kbalance E N: 4k: kiatance
stays finite for large d.
Proatance
3

---

Under review as a conference paper at ICLR 2021
Roughly speaking, this condition states that the change of logits after a single step is comparable to logits themselves. This means that the model learns.
Note that this condition is weaker than the used in Golikov (2020), because it allows logits to vanish or diverge with width. Such situations are fine, because only logit signs matter for the binary classification.
For simplicity assume ÿu = qw = q. We prove the following in Appendix B.1
Proposition 1. Suppose Ta Tw q and DD is continuous distribution. Then Condition requires 4o + E 121 0] to hold.
This statement gives a necessary condition for growth rates of a and á to lead to a well-defined limit model evolution. This condition corresponds to a band in (4o, -pppne see Figure I left. We refer it as a "band of dynamical stability".
Each point of this band corresponds dynamically stable limit model evolution. We present several conditions that separate the dynamical stability band into regions. We then show that each region corresponds to single limit model evolution.
We start with defining tangent kernels. Since Φ is smooth, we have:
1((+11) (x) 18 (x) = Σ 0fa(x)
f (x) fa
A Oiavv--007 + t/00"
a0, (,=0(()
- (1) ((x!!*) ( K and (x, (1)) Ni (()) El {xx(()) ((2))) K (xc + + =)::))
where we have defined kernels:
(x,(xx)))) (x, -d/d')/ o (H).T ())(
(9)
(8) (x,) = (d/d")
Σ (K) 12 | (w (k). x)(0 (w: (k).. xx)xxxx"xxxx x'.
(10)
Here we deviate from the traditional definition of tangent kernels (e.g. from Jacot al. 2018) in embedding learning rate growth factors into kernels. This is done for avoiding 0 x oc ambiguity when nev grows with while a vanishes that "a learning rate times a kernel" stays finite. This is the case for the mean-fiee field scaling: fav & d, while x d
While for the NTK scaling kernels stop evolving with & in the limit of large d, this is not the case generally. Indeed, for the mean-field scaling mentioned above we have:
K =
Σ
ob (R).T
(R).T x)) ( (x)xx)
(1I)
Similarly the NTK case, the kernel above converges due the Law of Large Numbers, however in contrast to the NTK case the weights evolve in the limit: (() ~ (0) This is due to the fact that
(k)
weight increments are proportional to f' which is x d0 for the mean-field scaling but x d 1/2 for
the NTK one. For this reason, similarly to model increments p (k) we define kernel increments:
AK (k) aVi.d ( x') = d (x,x)) - (() a xx,
(12)
Condition 2 (informal version Condition 5 in Appendix A}. Following conditions separate the band of dynamical stability (Figure 7 left):
1. fd stays finite for large d.
2. K aV stays finite for large
3. k-(0)
a/ (0) stays finite for large d.
4

---

Under review as a conference paper at ICLR 2021
4. AK K We prove the following Appendix B.2
(0) av ava.a stays finite for large d.
Proposition 2 (Separating conditions). Given Condition) 7
Condition reads as, point by point:
1. A limit model initialization is finite: qa 1/2
2. Tangent kernels at initialization are finite: 2qo + 1
3. Tangent kernels and a limit model are of the same order at initialization: 9a + + 1/2 0.
4. Tangent kernels start to evolve: qo O.
We have also checked this Proposition numerically for limit models discussed below: see Figure] D right. Each condition corresponds to a straight line the (gr: ()pplane: see Figure left. These four lines divide the well-definiteness band into 13 regions: three are two-dimensional, seven are one- dimensional, and three are dinddimennion... In Appendix we show that each region corresponds to single distinct limit model evolution; also list corresponding evolution equations. Note that segment (a one-dimensional region) that corresponds to the Condition 2 exactly coincides with a family of "intermediate scalings" introduced in Golikov 2200
CAPTURING THE BEHAVIOR OF FINITE-WIDTH NETS
possible use- for a limit model is being proxy for a given finite-width net, useful for theoretical considerations. For example, number of theoretical properties, including convergence to global minimum and generalization, are already proven for nets near the NTK limit: see Arora et al. (20195).
Note that typical finie-wwdth model satisfies all four statements of Condition (if exclude the word "limit" from them). Indeed, neural nets are typically initialized with He initialization (He et al. 2015) that guarantees finite f(0) even for large width d. Since learning rates of finite nets are finite, the tangent kernels are finite as well. Nevertheless, a neural tangent kernel of a typical finite-width network evolves significantly: AArora et al.) (2019a) have shown that freezing NTK of practical convolutional nets sufficiently reduces their generalization ability; Woodworth et al. 20199 also noticed that evolution of NTK is sufficient for good performance.
Consequently, if we want a limit model to capture the dynamics of finite-width net, we have to satisfy all four statements Condition 2 However, as one can see from Figure ! we cannot satisfy all of them simultaneously. We say that one limit model captures the behavior of a finite-width one satttioonnsssssssiiiiiiisi satisfied by the latter are satisfied by the former too. If we say in this case that "the former dominates the latter" then one can easily notice that there are only three don-ndninateed limit models which we discuss in the upcoming section. After that, we introduce model modification that allows for a limit satisfying all four statements.
better than the other, if all statements of Conditions
3.1 "NON-DOMINATED" LIMIT MODELS: MF, NTK AND "SYM-DEFAULT"
Obviously, the three "non-dominatee limit models exactly three zero-dimensional regions (points) in Figure [ left. First suppose statements 1, and 3 hold, hence tangent kernels are constant throughout training (see Figure [ right). A corresponding point yo -1/2, q reads as d-1/2 and ij = const, which is the case considered in the seminal paper NTK (Jacot et al. C. T.I.Tand and App. C for the general derivation,
2018). The limit dynamics is then given as (see App. App.[C..I.T and eqs. 75) for complete system of evolution equations):
muk,c (x) ouk.oc -fj V (K)
((x!))
(1)) a. x+)) th
(K) (k) fa e(x(() (() K((0) 100 (x, x(())
mi (x) (1)
(0)
(13)
where (x (x)
(0) (Navis ) ~ D and limit tangent kernels K aV and standard deviations at the initializa- tion of (0) (x) can be calculated along the same lines in Lee al. 2019
5

---

Under review as a conference paper at ICLR 2021
Next, suppose statements 2 and hold. In this case K does not coincide with K (see Figure
D right), hence the dynamics analogous to (1) is not closed. However, the limit dynamics can be
expressed as an evolution of a weight-space measure (see Rotskoff & Vanden Siiddenn-ijneenni (2019): Chizat
& Bach (2018) for a similar dynamics for the gradient flow, App. [C.2. and App. for the general derivation, and eqs. (94) for a complete system of evolution equations):
Mo00
11++)
(R)
+ div
(K) A
) Hoc (0(0,
(14)
mf (x) =0" o abl x) (b ( (dā,dw))
(k)
f
(15)
where the vector field (g mf is defined as follows:
Aoth mf (ā, w) = -V fur ((x!") (k) (k) (o("" ()),
(k)
(k) (k) (1) jao' (w" x xxx (x
e(x
(11)
(16)
Far
where we write " [,, v]" meaning a concatenation of two row vectors u and v. Here we have fr
g= 1. hence x d'1 and rj d: this hyperparameter scaling were used in Rotskoff & Vanden
Eijnden (2019): Chizat & Bach (018)) Note that since a measure the initialization foc has
zero mean, a limit model vanishes at the initialization f = 0 (see Figure right) thus violating statements and of Condition 2
Finally, consider point for which statements and 4 hold: qa = - 1/2, q 1/2. This situation is very similar to what we call "default" scaling. Consider He initialization (He et al. 2015). typically used in practice: da < dd1/22 and uw x d=1/2. Assume learning rates (in original parameterization) are not modified with width: 77 = const and f = const. This implies f) x d and fur x 1, or To = 1 and qw 0. We refer the scaling & 1/22, qa 1 and fw 0 as "default", and the scaling 9- = 112, q 1/2 as "sym-default". A limit model evolution for the sym-default scaling looks follows (see App. C.2.2 for equivalent formulation and App. [for for the general derivation, and eqs. (97 101) for a complete system of evolution equations):
(k+1) = (k) + div(u
(k) A0(K)
symn-de)),
(0) = =0. I+dd);
(17)
If aco x) (8) (dā dw) > >0 (18)
de (x) ~ NO0 0² 22 20))
sym 2o (x)
(k)
where the vector field A (K) der der is defined similarly to the MF case 16
(g(k)
sym-der (ā, w) - (V(M) (k)
(k) ya (K)
)o( (k) (()))
(K)
((x
(k)
yie jad WT x x(x))) ((0.TT
famder
Symer
V (k)
fommder ((x, v) ----yly sym (x)
(A)
for k 2 1.
(19)
As we show in Appendix D. the default scaling leads
almost similar limit dynam- ics as the sym-default scaling: eqs. (114 for a complete system of the corresponding
evolution equations. The quantity 2sym-dof.co should be perceived as sign of f
(EF)
o² lim (d=+11 fap(TTx x) (1()) (dá, dw)). The reason why we have to switch from logits
to their signs is that the limit model diverges for k 2 1: lim (x) = ac. Nevertheless the gradient of the cross-entroe cross-yntropy loss is well-defined even for infinite logits: it just degenerates into the a( -ylyz < 0j. For this reason, we redefine the 8f
gradient of hinge-type loss: limf--ooxe loss gradient for k > in terms of logit signs: eq. (19). Note that besides of the fact that logits diverge in the limit of large width, the measure in the parameter space ula stays well-defined.
3.2
MEAN-FIELD (IC-MF) LIMIT
Here we propose dynamics that satisfy all four statements of Condition 2 modify the network training for the finite width in order to ensure that in the limit of the infinite width its training dynamies converge to the proposed limit one. Consider the following:
We then show how to
(k+1) 2
(b( + div(j
(k) AO(A) )
1 = [tdd))
(20)

---

Under review as a conference paper at ICLR 2021
f (x) (x) f 1 (dā, dw) + =m (x),
(21)
where pll)) is defined similarly to above:
mik na (x) ~
(22)
the vector field A is defined analogously to the mean-field case:
(e) lemi (ā, w
(V
(k)
f(x,
(K)
(1)), (k)
f(x!)) y jāo ( (w (w"x()) )x
(k)
(23)
Fane
See App. E and eqs.
[3]
(34) for a complete system of evolution equations. The only difference
between this dynamics and the mean-field dynamics is a bias term ( in the definition of logits.
This bias term does not depend on and stays finite for large d in contrast to mm, which vanishes for large d; it ensures Condition 2 1 to hold. As for Condition 2 4, tangent kernels evolve with
simply because the measure Hlac evolves with k similarly to the mean-field case (see Figure right). Indeed,
K(() oi (x' 02.20" |aa00)p =x (f (k (8).7 x') poo (dā, dw),
(24)
and the limit of K ([(a] is written in a similar way. Kernels initialization K av are finite due to the Law of Large Numbers (Condition 2); this, and the finiteness of (ll) ensures Condition 2
As we show in Appendix E the dynamics (20 is limit for the GD dynamics of the following model with learning rates fav Wlvw(ddd"
fiemt, a W) = (d(dd))
d ( + (1dd"")))
0 W (0). xx).
(25)
Note that fiemt. (x) = 0* Dd= (r x)):
we have not altered the model definition at d'
3.3 EXPERIMENTS
Consider a network of width d' initialized with standard deviation a² and trained with learning rates fav We call this model a "reference". Consider a family of models indexed by width d with hyperparameters specified by the power-law scaling (7). We train a reference network of width d' = 128 for the binary classification with a cross-entrops loss on the CIFAR2 dataset (a subset of first two classes of CIFAR1O). We track the divergence of limit network from the reference one using the following quantity: E Xx-Deest Diogito (f (x)
f (x)), where
Diogits (5 |1 (*) (E,,(EE, Var E) || [E",", Var
(26)
We have also tried other divergence measures; see Appendix [
Results are shown in Figure 2 The NTK limit tracks the reference network well only for the first 20 training steps; a similar observation has been already made by Lee al. (2019). At the same time, the mean-field field limit starts with high divergence (since the initial limit model is zero in this case), however, after the 80-th step. it becomes smaller than that of the NTK limit. This can be the implication of non-stationary kernels. As for the default case, divergence of logits results in blow-up of the Kl.-divergence.
The best overall case is the proposed IC-MF limit, which retains the small KL-divergence related to the reference model throughout the training process. Capturing the behavior of finite-width nets is also possible by introducing finite-width corrections for the NTK (Dyer & Gur-Ari 2019 Huang & Yau 2019). However, this gives us an infinite sequence of equations, which is intractable. We have to cut this sequence; this gives us approximate dynamics, which is still complicated. In contrast, our IC-MF limit is a simple modification of the MF limit, and at the same time, a good proxy for finite-width networks.
7

---

Under review as a conference paper at ICLR 2021
10a
MF
MTK
IC-MF
15a
default
symn-eeault
KL/linit,
10a
10a
10a
136
_0'
10a 10a
training step.
Figure 2: Initfalization-correetee mean-field (IC-MF) limit captures the behavior of a given finite-width network best among other limit models. We plot KL--dvergeence of logits of different infinite-width limits of a fixed widthh reference model relative to logits of this reference model. Setup: we train a one hidden layer network with SGD on CIFAR2 dataset; see Appendix F for details. KL-divergences are estimated using gaussian fits with 10 samples.
RELATED WORK
pioneering work of Jacot et al. (2018) have shown that gradient descent training of a neural net can be viewed as a kernel gradient descent in the space of predictors. The corresponding kernel is called a neural tangent kernel (NTK). ) Generally, NTK is random and non-stationary, however et al. (2018) have shown that in the limit of infinite width it becomes constant given a network is parameterized appropriately. In this case the evolution the model determined by this constant kernel; eq. (3).. The training regime when NTK is hardly varying is coined "lazy training". as opposed to the "rich" training regime, when NTK evolves significantly (Woodworth et al.) 2019). Chizat et al. (2019) noted that the training becomes lazy for finite width if one scales the output of the network appropriately. While being theoretically appealing, "laziness" assumption turns out to have number of limitations in explaining the success of deep learning (Arora al. 2019a Ghorbani
et al. 2019).
Another line of works considers the evolution of weights an evolution of weight-space sure, similar to eq. (14) (Mei et al. 2018 2019 Sirignano & Spiliopoulos 2020; Chizat & Bach
2018 Rotskoff & Vanden-Eijndee
2019 Yarotsky 2018). This weight-space measure becomes deterministic in the limit of infinite width, given the network is parameterized appropriately; the
corresponding limit dynamics is called "mean-fiedd" Note that the parameterization required here for the convergence to a limit dynamics differs from the one used in the NTK literature.
Our framework for reasoning about scaling of hyperparameters in Golikov (2020). However, there are several crucial differences. First, we do not consider weight increments, as well as a model decomposition, and do not try to estimate exponents of the former and for terms of the latter, which arguebly complicates the work of Golikov (2020). Instead, derivations in terms of the limit behavior of logits and kernels which appears to be simpler and clearer. Second, our criterion of "dynamical stability" of scaling is weaker compared to the one of Golikov (2020) and more suitable for classification problems, since it allows for diverging or vanishing logits, long as they give meaningful classification responses. In particular, our dynamical stability condition covers practically important "default" limit for which learning rates are kept constant while width grow up to infinity. Note that "intermediate limits" investigated in Golikov (2020) exactly correspond to limit models which satisfy Condition 2,2. Moreover, both "sym-default" and IC-MF limit models we propose in the present work have not been discussed previously; we present limit evolution equations for both of them (see Appendix(C) (). Finally, our analysis suggests that there are only 13 distinct limit models that can be induced by power-law scaling of hyperparameters.
similar in spirit to the one used
present
5
CONCLUSIONS
The current work follows a direction started in Golikov (0220): we study how one should scale hyperparameters of neural network with single hidden layer in order to converge "dynamically
8

---

Under review as a conference paper at ICLR 2021
stable" limit training dynamics. A weaker dynamical stability condition leads us to richer class of possible limit models as compared to Golikov (2020). In particular, the class of limit models we consider includes "default" limit model that corresponds to network with infinitely large number of nodes and finite learning rates in the original parameterization. This "default" limit model does not satisfy "well-definiteess" condition of Golikov (2020).
Moreover, we show that the class of limit models that can be achieved by scaling hyperparameters of finite-width nets is finite. The space of hyperparameter scalings is divided by regions with certain conditions on the training dynamics, and each region corresponds to single limit model. All of these conditions are satisfied by finite-width networks, but cannot be satisfied by limit models all simultancously. We propose modification of a finite-width model; the limit of this modification corresponds a limit model that satisfy all of the conditions mentioned above and tracks the dynamics of "reference" finite-width net better than other limit models.
REFERENCES
Dyego Araújo, Roberto I Oliveira, and Daniel Yukimura. A mean-field limit for certain deep neural networks. arXiv preprint arXiv: 1 1906.00/9, 2019.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances in Neural Information Processing Systems, pp. 8139-8144 -18,,8148,220 2019a.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. arXiv preprint 2019b.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. In Advances in neural information processing systems, pp. 3036- 3046, 2018.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In Advances Neural Information Processing Systems, pp. 2933-2943, 2019.
Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. arXiv preprint arXiv: 1304, 2019.
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Limitations of lazy training of two-layers neural network. Advances in Neural Information Processing Systems, pp- 9108-911 91 18, 2019.
Eugene A Golikov. Towards general theory of infinite-width limits of neural classifiers. arXiv preprint wxXiv:2003.05884 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance imagenet classification. In Proceedings of the IEEE international conference on computer vision, pp. 1026-1034, 2015.
Jiaoyang Huang and Horn-- Tzer Yau. Dynamics of deep neural networks and neural tangent hierarchy. arXiv preprint arXiv: 190909900000 08156, 2019.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems, pp. 8571-8580, 2018
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. In Advances in neural information processing systems, pp. 8580-8581, 2019.
Song Mei, Andrea Montanari, and Than-Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 115633))7666-E7677 2018.

---

Under review as a conference paper at ICLR 2021
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory, pp. 2019.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Grant Rotskoff and Eric Vanden-Eijnden. Trainability and accuracy of neural networks: an interacting particle system approach. stat, 1050:30, 2019.
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of deep neural networks. arXiv preprint Att3:/9903.44444 2019.
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A law of large numbers. SIAM Journal on Applied Mathematics, 80(2)725-752, 2020.
Blake Woodworth, Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Kernel and deep regimes in overparametrized models. arXiv preprint arXiv: 1906.0887, 2019.
Dmitry Yarotsky. Collective evolution of weights in wide neural networks. arXiv preprint arXiv:1810.03974,20000 2018.
A FORMAL CONDITIONS FOR SECTION
Here we present formal definitions for notions that appear in Section mathematical rigor. First, recall the definition of tangent kernels:
2
they are required for
(k) a.d [d
(x, x') (d/d
(K).T
x x)
(1.Tx)) x')
(27)
_h0 w.d (x,x = (d/d') 0²
(k).T
(28)
The kernels are used to express a model increment:
(x) fa (k+11) - -100 (x) = = 0fa(x)
A@(K) + Ontvv -0 + 2)
aθ,
à g(60
no V (4) (R)
((x!
() K! ,
(x, x (k))
no
v fa
{xx
(k) (1)) (1) (x, x x0) + + fr.2)
(29)
Define the linear part of the model increment with respect to learning rate proportionality factors:
Afa aaf(1) (x)
(k) fa ((x warw) (Kk)
K (x,)
(k)
(30)
We use this quantity to rewrite the model increment:
Af( (x)
(x))
d.a
f A (k)) (x) O + ))
(31)
Let us consider kernel definitions 27 and 28 again. Their increments given by:
(K((k) (x, xx) - (
((w (R), x) (w. (3).Tx))
(1).T xxxxxx))e x) (w
())(W(()).T x))
(k) ((x(1) )
)a! (k) b((w.))) (c). (k) (x + (xx+x))+ (1) (k) + O - a -5 +go (32)
d=
10

---

Under review as a conference paper at ICLR 2021
=
0 Σ |0 12
()) (K), x)+
x)) x x''xx x
yi (
)
(w,
x (R)
+ Oi No -o (
~000
3 (d/ddd
20 (k) 25' (w (x)_)))xx())))))))))))))) x)0 (w, (6)
(w x')x
66 (K),T x (k)) + Oz -0 201-2689-100 d24-110+1).
(33)
d++x0
Similarly to what was done for model increments, we define linear parts of the kernel increments with respect to learning rate proportionality factors:
2D OAK
(k) a.d (x, x')
A A awaa (x, x)=
aii
-ddddd)
(W (6),T x) (1) x') o' (x)("))))))))))))))) x)) x
E(xx0 (1) )a( (((w (k), (x xx)x (()) (34)
OAK (*() w.d
(x, x')
(K(K), ww,d xx, x))
6-=0
(dddd")
0
Σ (a.4) (w(W (K).)) x (w (K).T) + ("(!!h) "')''"""""" (1), x) xxxxxx
r=1
1(x fa
(k)
y (h) )a( 0' (((w (() xxx) (35)
AA K w.d Xx x')
AK(())
wadd(x,
Oña
=1=0
( - Σ 20( ø' ('(w!h) ( (W(lt).T "x)x'"'VVVV (k) ((x!*) (00 (o)) ( xt)) (36)
Note that AK AK(()) aud X x') = since -errs are absent in the definition of Ka,d: eq. 27
Define Perrd P (g.x.gh" (Kk 1) xxx Gik 1) x "y (uf (x) (x) < 0} the probability of
(k)
giving wrong answer the step k. Let kterm.d N U (+o0} be maximal & such that Vk' < k (K)
Pernd > O. Generally, Kterm,d depends on hyperparameters, as well as on the data distribution D.
Scaling exponents (ge, qu. Fam) together with proportionality factors (d" o" na. Nir define limit model ((6) (x) lim ((()) (x). We call a model "dynamically stable in the limit of large width" if it satisfies the following condition:
Condition 3. 3kbalance
E : Vk
(k) g(h)) (x(()) < 0 and [Riatancee kteem, N. ya (Klsalaane))
You Jo (x()) < O imply ((k),
(1)
Add,vv (x) = Oe (f (x) x-a.e. (:k) x Xavw -a.s.
(1k)
This condition puts constraint exponents (go, Jos qw); this constraint generally depends on the train data distribution D and on proportionality factors d' o", and fāve- In order to obtain a constraint, need the condition above to hold for
data- data-independent any value of keerm. and any values of d', o", and y Without loss of generality we can assume ozerm be infinite, which gives the following condition:
11

---

Under review as a conference paper at ICLR 2021
Condition 4 (a formal version of Condition
D. Given Rterm, 20 400,
Ekbalance E N Yo
Vilvw > vk 2 libalance yo (k)
((k)
(R)
0 and you (R)
fob (xm < imply A Sd.avm (x) =
(Xa
Od (fd ( (x)) x-a.e. (Vavmr XVVio) -a.s
(sk)
Condition 5 formal version of Condition 2 Following conditions separate the band of dynamical stability (Figure left):
A limit model at initialization is finite: _2 = (1- x-a.e.
2. Tangent kernels at initialization are finite: K K(0)
d, x') O (1) (x, x-a.
3. Tangent kernels and limit model are of the same order initialization: K K(0) davv (x, x')
(f (x)) (x.x')aa.e.
4. Tangent kernels start to evolve: (0), d,waVio (x, = O (K
d, d,, (x, (x, x') and
A (0)) d,a (x, x²) = Od-roo (K
2nd x x')) (x, )-ae
B PROOFS OF PROPOSITIONS
We restate all necessary definitions here. We assume the non-linearity to be real analytic and asymptotically linear: ((z) 22)x We assume the loss function Ely, z) to be the standard binary cross-entropy loss: Ely, 2) = In(1 + e e==*), where labels y E (-1,1}.
The training dynamics is given (1)
Aa(k) a fa [(x!! (k), ( (((()))XX! x
~
(37)
Aw(K) -m V (b)
{(x!!
(k), y ( agh o' ('(W!!)) ( )x!
w ~((,
Vr E [d].
(38)
fa
V (k) [xx, y) ol(y.2)
(1) (x) Σ
a
(6) ((W(T)Tx))
0z
=== 1 (x0
1 exp (x)v)
r=1
where (x Dava) ) ~ DD for 'D being the data distribution.
We assume hyperparameters to be scaled with width as power-laws:
(d) a² (d/d")a> (d) va (d/dr)ā,, (m the (d/d")/
B.
PROOF OF PROPOSITION
Define:
40 (k)
ig g(k) Od-
L inf Ag(a) 08
(k)
(39)
where 0 should be substituted with w. We define inf +oo. We introduce similar definitions for other quantities:
of (x) = inf(g: (x) = (d))
(x, v) = inf V (8) ((x, v) = Od+)
(40)
(1) = inf(q: 15(00 Od=r)o(dd)) Iafivo (x) inf Aldii = Oa-soa (
(d
(41)
Lemma 1. Assume D is continuous distribution. Then Then following hold:
1. Vk 2 0 VVx 1et (x, v) 5 O. while ly (k) (x) < 0] implies d76 (x, v) 0.
(k)
(k)
2. Lave 0. o3 (x) go + x-aee
(0)
12

---

Under review as a conference paper at ICLR 2021
(8) 3. vk 2 0 0//0/ / -avi + 90 + Twja a + 9o0 (x(1), (y)) (x Nevwlaa..
(k)
(k)
4. Vk 2 0 0g favm (x) 2go + 1 + vav + 2wwa +97 (Xavis Navw) x-a.e (Xavis )avw)-a.s
(k)
(k) (k)
5. 4A 2 0 10 + tw + (k) 1 implies that for sufficiently small ve and the (e) = (M)
max 93%% (x), Ia5% (x)) x-a.e (xa (k) (k) (K) (K)
ya.ss
Yu .w Vou
max Gavis (ba/Dw) (xx(t), ((0)),(()))) y(.)))..s.. of (x++)((x)
(k) (6)
6. 4k max(aj (x). M3f (x)) x-a.e. (xa (M) (K)
qavw
(N)
ya xu ye )-a.s
Proof. (1) follows from the fact that ally, =) is bounded Vy, while (ally, z)/02l [1/2, 1] when y2 < O.
(0) = 0; similar holds for w. ( ~ 1)) which is not zero and does not depend on d, hence qa x) Od-so (d) (2+ due to the Central Limit
For x F we have f/ (x) Σd a p (0)T
Theorem. Hence (2) holds.
Since D is a.c. wrt Lebesgue measure on R Rl and o is real analytic and non-zero, o(w, (K),T xavw) # O and d'(wi (1),T xovw) is well-defined (xav Vavw) ) This implies that (k) da//A lave + to + Dw/a + avv (Xav yavi (x yavw) )-a.s, which is exactly (3).
(k)
(K) (k) (6)
Consider A//
Af(h)((x) = -v/()) (+) l(x!!*) , ) K.t)( (x,, x((1))-
= e(x! (8) ) (d/d") 0² Σ (o x) (WWWWW/Txx (42)
For the same reason as discussed above o(w, (1).T ( Xa (k) * 0 (xm Ya Ja.ss and o(w, (1), x) + 0
(k) (k)
x-a.e Since the summands are distributed identically and are generally non-zero, the sum introduces a factor of d. Indeed, an expectation of the sum is a sum of expectations by a linearity property. Moreover, the absolute value of the k-th moment of the sum of d indentically distributed terms does not exceed d times the k-th moment of each summand. Hence the sum itself scales as d. Since is asymptotically linear, each ó-term scales as d (2) Collecting all terms together, we obtain (k)
(k) (k) (k)) x-a.e. (xa yo (R) )-a.s. Following the same steps fa (x) = 2// + 1 + Ha + (8) + Tel (xa ya
for A A
(k), we get (4).
Let us overview Afd (k)
(x) in detail:
xC
a fx(x
A)(1) (x)
Σ
(wi
Awi
Jwyy
(R)
a,=a(a)
e fx()
Ad, An,
Anitthis)
0úit
Σ
(1) l( (k) y )
((x((x)
(a( (0'
13 na
()))))² () x
3=1
x (((4)))-11 (w( xxx) (())))))))))0)))11(0)))) (10,)xx)) xxxxx(1)xxxxx11 2 (43)
13

---

Under review as a conference paper at ICLR 2021
(8) 0 implies i (0) ( and
Assumption 40 + qv qo
aa 1 (ar ( (0( o0 ( (2)
(8) Since a00 (x, v) I 0 yx,y due to (1), (V (k) f(xu yu )) Od (V f(xe yy and (1)
(k) (K)
(k) (k) (1))
(8) y/a))) e(x)))
(V
fo ((x! "
() - O (V [x(!
(K)
Since 6(2)
(K).T (z), d'( d((w, )
(K).T x(x )) Oa (1) and (d'(w, (w!
Od- (d' (W. x ) for 1.
1.
Hence for small enough th and fi the first term of each which corresponds to j dominates all others, even in the limit of infinite d:
d
Afd (x) =
16) ((x!0) bl (a! V
x ('((W()))))))) (w(t). x) x)x (K).Tx x+
Sa
f(x!
fa
+0 (Od+20 (11v (k) {(x
AN La {(x -0 (a.t)))) )0)))
- (x) + no 22 (k) (x)+
+ Oñi +0 (Oa-1o (1 Va
(k)) Ne V (k) e(x!!
(8) w (a( o²d)): (44)
Note that two summands depend on (xlx -yu on each other. Hence (x) = max(1 (x), " (x)) x-a.e (xoves )-v.)- which is (5).
(k) and (xa (R)
yo (8) respectively, which do not depend (6)
(k)
(K)
Note that the o-term does not alter the exponent. Indeed,
na V (k) fo ((x!1))
(k) 1xx,) + iN V
(8) ((x( v (a(ki)) o²d =
(k)
= (dlt2g + (xe!)) w/th 1. gel (xc1) a ( +22)))))))))
O
(02 (dl+200 +
2004420 au t (xe00) w +
O m (2) (2)) (x))
(45)
One before the last equality holds, because q due to (4).
(k)
due to (2) and (6), while the last equality holds
(k) (k) (k-1) = (k) Aar (k) Since the second term depends on (xa Ma while (k+1) max (8) (k))
By definition we have ar
the first term does not, we get qa which gives (6).
(.a). Similar holds for w,, and f(() (x) x-a.e.
Lemma 2. Assume D is continuous distribution, Kzerm +o0 and qo = qw = g. Then
1. If qr + + S 0 then vk 2 0 quvu 0 (xx ( 1) Vaa 11) xx
(1k-1)
Ve (1k-1)) aass )-a..
2. If qr + q O then vk 2 0 gavu = klga + ②) with positive probability wrr (x( (1) yi 11) 11) you 11)
Proof. Here and in subsequent proofs will write "almost surely" meaning "almost surely wrt (x(()) Na (ik) (ik) yu (ok)y" " for appropriate we apply similar shortening for "with positive probabil-
ity wrt (x1 yo x )"
(1k) (k) (:k)
If go + q S 0 then statements 1, 2, and 6 of ( Vk: 2 0 00g4) 0 a.s.
Assume 9rr q by induction. Induction base is given by Lemma 2
q > O. We will prove that Vk: 2 0 gave max(0, A(9z + 7)) with positive probability
(k)
14

---

Under review as a conference paper at ICLR 2021
Combining the induction assumption and Lemma
1 we get 1sa/Dw (k + 1) + ÿ) yio 1 (xavir (k)
(R)
aab (x((), (()) with positive probability wrt (xm (((-)))
(1)
ya
(ck
X.i
(ck-1)
Jas.
Since kteem to0 > k. Javw-d (k)
(K) ((() (b) < 0 with positive probability wrt (x. ya x vw
and Lemma implies that 4 = (k + 1) + ?) with positive probability wrt
(xc) (RR)
(.k)
Finally, Lemma 6 concludes the proof of the induction step.
Lemma 3. Assume DD is a continuous distribution, keerm, = +00, 9a = gw = q and qo + S
Then vk 2
1. yavw
(R) (() (Xovw) <
0 implies
(k) aDfavo (x) = 24 + + q x-a.e.
(xn (ck-1) (ck-1) 1c
(ck
1 Jass.
-ya
Xw
Vou
2. yo soe (11))
(K)
O and you
(6) 5 (xx < O imply 42f
240 + 1 + ④ x-a.e.
(xp= (c (1 x1x 11) (sk ((--1) 1- for sufficiently small il and ii
Vw
Proof. By Lemma vk 2 0 davi = -Oas. Since (k) (k) (k) < O, Ove O due to Lemma Davar
(k)
1.. Given this, Lemma I implies Ak 2 0 4afvvo (x) = 24. 1 + q x-a.e a.s. Hence by virtue of 2gr + + x-a.e. a.s. for sufficiently small il and iik Proposition 3. Suppose qa que q and D is continuous distribution. Then Condition requires 4a + + [-1/2, 0) to hold.
Lemma 5 vk 2 0 aay
Proof. By Lemma if go + q same time by virtue of Lemma 1111 hiermoo +0 implies AVe (K)
> 0 then lavw k(go + q) with positive probability. At the
(k)
(k) O with positive probability.
Given this, Lemma implies Va/vva (x) 4r + 1 + (2k + 1)(ar + 9) x-a.e with positive
probability. This means that the last quantity cannot be almost surely equal to of (kratance) (x) for (x))
any bbalance independent on k. Since Afdavo (x) A (fd (((hbalonne)) (x)) requires Aafavm (x)
(k)
(kadlance) (x), we conclude that Condition cannot be satisfied if go + q 7>00 O.
4f
(k) < and you (R) ((k)) (x (k) < 0 imply Joo (xo
Hence Go + q I 0. Then by Lemma (k)
vk 2 (sk-1) 1) (sk 1) (=k-
Laf (x) 2g + ② x-a.e
(x.
(sk
Ua
Xi
Yyu
1) -a.s. for sufficiently small ia and
f We will show that Condition 4 requires Oo + q [-1/2, 0] hold already for these sufficiently small n% and niu
(R) 5oo (x(()) 0 and you 3o (x 0. Given this, points and of Lemma imply Suppose No
(k) (k)
(k)
Vkbalance 2 1 9f
Abalance (x) max (x), 2g. + + 7) max(4, + 2. 2.24 + 1 + ②) x-a.e. a.s.
Hence 4a/avm (x) of
(Mientance) (x) x-a.e. a.s. if and only if go + S 2go + +1+ q. which is + 2
-1/2; take kbalance 1 without loss of generality. Having lakvve (x) 9/ of (heatence) ) necessary to have A (x) = O
(A)
(x) is
(1d (f
(Neatance) (x)).
Summing all together, Condition 4 requires do + [-1/2, 0] to hold.
B.2 PROOF OF PROPOSITION
Proposition 4. Let Condition 4 holds: then
1. 1 (x) = Oa- (1) x-a.e. is equivalent to q, + 1/2 0.
2. K-10) (Lava(xxx = Ad-roo (1) (x, xa.e is equivalent to 24o + 1 = 0.
15

---

Under review as a conference paper at ICLR 2021
3. d,vv (x, x') = (Jo
K-(0]
(x2 (x)) (x, x') is equivalent to qr +- 1/2
4. AK (0),
d,wavm (x,x) = @ (K
(0)
d, (x, x')) (x,
and dx, x')
Oo (x,x) x')) (x, x')- e. is equivalent to qo q O.
Proof. Statement (1) directly follows from Lemma 2:
(x) Σ
a 0
00. x)
-dooo (MMot122
(46)
(x)-a.e. due to the Central Limit Theorem. Statement (2) follows from the definition of kernels and the Law of Large Numbers:
K Th (x,x' (d/
Σ ob (O) x))(
(O) x) x))
A (d (d+22++1))
(47)
(x, xa.e.. the same logic holds for the other kernel: [ w. (x, x') O (d(+2g+t1) (x, )-a..
Combining derivations of the two previous statements, we get the statement (3). Now proceed to the last statement. Consider again the kernel k[0); a linear part of this increment with respect to proportionality factors of learning rates is given by, see eq. (34
A AK aw (x, x') = OAK (0)) (x,x') x')
Oij
12=0
d
-ddd/d") 2
(o(w))) () (0). x') + (w (0)T ()((0)))xx // x))
fa ((x
00)
(0)
ye a Ø' (w (0) (0)) )(x x') (0)
(48)
Hence A 1 = 2 ( (K is equivalent to do q 0. Considering the second kernel K (0) aw,d and its increment is equivalent to the same condition.
w.d
C THE NUMBER OF DISTINCT LIMIT MODELS IS FINITE
is easy to see that due to the Proposition 4 Condition divides the well-definiteness band into 13 regions. We now show that when proportionality factors o* and fay are fixed, choosing limit model evolution is equivalent to picking single region from these 13.
Indeed, for any width d a model evolution can be written as follows:
(k) (x) A Afdh)
ni v v (k)
fa (xx(x) y((b))
(K:") wnd (x, xx, (+))+
AK(()))
(k).)
wo.d (x, (K)) + t. A A (K) (x, x((1)))
+0
10000
n v
(1) 2(x!!
Vua (k)) K5! (k)
[x, x (1)) + Og 10 fir AK x 1))
(49)
d=
f1 1) (x) = )((h (x) + A (x) (x), 1/2 (1),)(1))) y) = 1 exp((s[ (x)y))
-
(50)
12 (x) o² d
a 00) (0,,x x), (a.), w(()) ~ ((0da))
(51)
Now we introduce normalized kernels:
(x, (x)x') = 00-2 dt d
R:00) (x, x') -0 (d/d") 11-9")11-9---- -1-9-24. 24 (
Σ
olw ((w())) (6),3
x))( (H).T))
(52)
d
r=1
16

---

Under review as a conference paper at ICLR 2021
R R (x, = (d/d")-1---20)
24 k w.d (x,,') o 0z. 222 d'
1 (2d (wr( x) (
(
x')xxx
(53)
Note that after normalization kernels stay finite in the limit of large width due to the Law of Large Numbers. Similarly, we normalize logits, as well as kernel and logit increments:
A AK (d/d")-1- -----0 AK(
20,
(54)
AJ (1/2")11~--4
24- Ap/h) pam) = (d/d') 11------200 20 s/p/.
(55)
We then rewrite the model evolution as: (x)
A (x)
n V (k) fa
((x(()) yew )
( (M) w.d (x, x )
O nd +0 hh AK ( (2, x (b)) + nh AK ward (X, (b111) 1)
((xx! 0 -ya (k))
(k)
a.d x, (k)) O 0p: fi A (u, (x,xx
(56)
La
f(k+)) (x) = f(h))(x + A (x) Vk 2 0,
(57)
04 (x) o² (d/d")-1-4-3
11 Σ
a Φl w (0),T (0),x)),
( (
NO0 I+dd),
(58)
f4 (d/d)) 24 (x),,
(x),
V (c) ((x,)) 9)
y
Vk #x20 0.
(59)
+ exp(fa (x)
(x)y)
C.I CONSTANT NORMALIZED KERNELS CASE
(R) Kernels K avvodd are either constants (hence A (A) andd - 0 as d - oc) or evolve with k in the limit of
large d. First assume they are constants; in this case do + q < 0 due to Proposition 44, and
A ((k) (x) = ne V f(x V/
(K-0)) w.d (x, x (k)) + Oiyy o
d+x
- fo ((x!_ :yaa) .(aa)) (K.0 x!!)) + - (60)
(k) (b)
d=100
(0) Since normalized kernels av converge to non zero limit kernels (1(0) formula above as:
Ravi.oos we can rewrite the
A 2(K) (x) = -i V (K) fa ((x 1
(000)
0000 (x, x (k) + (0)=
nh fa
(t)
((x!!) (k) , (:(0) (x,,
(x, x²1)) + (1))
(61)
(x) = - (dd/d)))--1 (N(0, . 00) (x)) + (1))
(62)
where (0)) (x) can be calculated in the same manner as in Lee et al. 2019 As required by Proposition
3 1 1/2 + q+ 4o 2 0, hence f(0)) = Oa-sno (1). This implies the following:
V (0) Jm [(xxxx) v) lim
00) V
{(x,))=
La
00
o < 0
for 1/2 + do > 0:
lim
yv
for 1/2 + do 0;
d=000 1 + exp((d/d))
1 2g 7(0) (x)y)
-- (0, (0.000000 (x))))
y/2
for 1/2 + do < 0.
(63)
17

---

Under review as a conference paper at ICLR 2021
On the other hand, A d
00 (x) a (1) with positive probability over (Xavw Navw): ). Hence
(0)
((o) Os (00 and ((1) 2(0)
+ + (1)) For the same reason, fd
f(k+1)
A/ +
A1 +
0 (1) Vk 2 0.
This implies the following:
AA 0
V (ltt [(x, y) = lim v(ttt)((x v/
lim
y
for
d0000
fo
d+ 1 + exp((d/d+)) 1++42 få f(atl)/ (x)+)
ullimd-+o
glk+1) (x)y < 0]
for 1 + q + 290 > 0;
for 1 + q + 24a 0; (64)
1+exp(lims- -//2
11 (x)x)
for + q+ + 24 < 0.
If we define |/00 (x) = lim Sa (1) (x), we get the following limit dynamics:
((k) (x) = -w V (b) f(x!!) (k) u (1)) 1 (x, x
(k) Soe ((x!1)
ya (1)) K )
(x, x(()),
(65)
f0o
K10d (x, (x,x) = 2r
(66)
k10) 10, 2 (a. (a. ((,,,,
xx)xxx
(67)
(x) 00 (0 (x)) for 1/2 + q 9, > 0:
for 1/2+ ++ 9- 0:
(x) = F(h (x)+ (()) (x),
0
(68)
([ (0), (x) y < 0]
for 1/2 4r > 0;
V (0) fx
((x, v)
for 1/2 + 9 = 0;
(69)
14explo-d d 12 (0, (20))))))))
9/2
for 1/2 + go < 0;
v v 11++1) (x)y
< 0]
for 1 + q + 2g > 0;
(6+1)8(x,11 f00
for 1 q + 2g,- = 0;
Ak 2 O.
(70)
1-exp( (1ht)))))))))) (x)x) n/2
for + + 24 < 0:
This dynamics is defined by proportionality factors o", d" Vave and signs of three exponents: 1/2 9o. 1 + q- 24 and 1/2 + q + 9.- Since we assume proportionality factors to be fixed, choosing signs of exponents is equivalent to choosing a limit model. Note that these exponents exactly correspond to those mentioned in Proposition 4 points 1, 2 and 3. One can easily notice from Figure (left) that given le + + < O, there are 8 distinct sign configurations.
Note also that since we interested in binary classification problems, only the sign of logits matters. Since f(()) = (d/d+))+ f(()) signs of f(()) and of j(k) are the same for all d. Hence Vx, y lima-too sign(Sa
(1) (x)) lima sign( ((k) (x)) sign ((1) 150 (x)).
C.I.1 NTK LIMIT MODEL
We state here special case of the NTK scaling (g., 1122, q = 0, see Jacot et al. (2018) explicitly. Since in this case 1 + q + +2g = O, we can omit tildas everywhere. This results in the following limit dynamics:
(A) no
(x) = -1in (k) (K) (1) ( wipo x(0)
J00 ((x!8)
%%
((x!1 (k)
Va
(K) ) x!)),
(71)
[00) = - w.
(72)
k(00 (x,x)) 222T (2.w)
(73)
18

---

Under review as a conference paper at ICLR 2021
1 = -xx + (x))
(x) = 0² d0.1//21(00 (0))((x)))
(74)
(k) xx E(x y)
Ak 0.
(75)
1 1+ expl (k) (x)y)
C.2 NON-STATIONARY NORMALIZED KERNELS CASE
Suppose now do - = 0. In this case (0), A
maviio (x, O (K d, (x,x)) (x,x')-a.e. and
(0)
AK daw xx)x) = u (K d.a x')) (x, x').e. by virtue of the Proposition 444. Hence kernels evolve in the limit of large width (at least, for sufficiently small vav
x')
If we follow the lines of the previous section, we will get a limit dynamics which not closed:
((() A ((H) (x) = - ni 1/ 150 in Die 00,00
(2) (1)) E(x
(x, x (+)
+ OREv 00 The
((),,
(x, (k)) + fl (+).
10,20 (x, x x)(
1001,00
- 1x0
(K) {(x!!a)
(b) ,y(2))
(T(R)
(x, x + Onz 00 (n (((R),)
a10,00 (x, )))
(76)
0,00
f(++1) (x) = F(k) (x) + 7(k (x),
F (x) 0,
(77)
>l σ (0), x))y < 0]
for 1/2 + 4r > 0;
()) fox
((x.v) v))
for 1/2 + Qr 0;
(78)
1+explo (0, (x)))))) 1/2
for 1/2 + ga < 0;
v1 1()+1) (x)y < 0]
for 1 + 9- 0:
((xtt)((x.))
for + G 0:
Vk 2 0.
(79)
L+exp(i ((6+11)) (x)9) -9/2
for + 9e < O;
The reason for this is non-stationarity of kernels. As a workaround we consider a measure in the weight space:
Pd
(b)
Σ
000 3 (e.
(80)
Recall the stochastic gradient descent dynamics:
Aa(h) N
v (k) Sa
l( ((x!!
(8)
ye (3)) (( x((x)))))))))))
a,0)~ ~ ~N 1),
(81)
Aw(h) M o (N) ((x(b)
(k) ) a D' (w!
(k)), x xx{
w.!) ~ NN
(82)
Here we have replaced f with May T" because 9o + q O. Similar to Rotskoff & Vanden
Eijnden (019): Chizat & Bach 20118). this dynamics can be expressed in terms of the measure
defined above:
Nd Mdd + divlul)) (8) so(+),
(k+1) (k)
(0) 1
000~~ ~ ~N(, V [d] (83)
g00 Pa 7
10 (ā, w) =
= E(x! (k) (>][T>x())) nh -V( ((((() yeu Yas' (WT )x)) x (84)
fx) の (d/d / (ao(WTx)) x)x1( (k) (dū dw),
(85)
19

---

Under review as a conference paper at ICLR 2021
(8) fa
((x,
1 +xx(pff (k) (x)y) Ak 2 0.
(86)
We rewrite the last equation in terms of (x) - (d/d))-1- 1- ((h) (x): (k)
(x) o'd' I ccpliw' x) ud (dā dw/),
(87)
(k)) fa 1(xx v) = + )xp f( (x)y)
-y
Vk 2 0.
(88)
This dynamics is closed. Taking the limit d yields:
Mo0 20
(1++)
(R) + div(u
(1) A N = N' I+tde)
(89)
(0 (a, W)
=-1ia0 100'0VV0 (k) (k)
1xc ((x!1) va jo(
) xx! Ne V ((x!))
(b) v yao' (wT x((x()) )x())77
(90)
(x) = o²dd f / x) x) (dā dw),
(91)
-v[ 0())())0()2(())) <0
for 1/2 + do > 0;
((x,))
for 1/2 + Qr = 0;
(92)
I+explo d' (0,00000 11/2
for 1/2 + do < 0;
uI 1(k+1) (x)y < 0]
for + de > 0:
V ((xt))
for + = 0;
Vk 2 0.
(93)
1+exp(f2 9/2
PIN (x)))
for + 9r O:
Since proportionality factors a², d², and fav are assumed to be fixed, choosing d- is sufficient to define the dynamics. Signs exponents 1/2+ + 9o and + lo give distinct limit dynamics. Together with 8 limit dynamics for constant normalized kernels case, this gives 13 distinct limit dynamics, each corresponding to region in the band of a dynamical stability (Figure left).
As was noted earlier, only the sign of logits matters, and our (()) preserve the sign for any d: Vx lim sfgn(fa (f(((())))))))))) (x)) = lim sign(fa ((k) ((x)(x)) (x)) = sign( ((h)) (x)).
C.2.1
MF LIMIT MODEL
We state here special case of the mean-field scaling (q. -1, q 1, see Rotskoff & Vanden Eijnden (2019) or Chizat & Bach (2018) explicitly. Similar to NTK case, since + 4o = O we can omit tildas. This results in the following limit dynamics:
(k+1) 1
(k) + div(u (k) AO(H), 1 - I1td2):
(94)
A0(K) (a, w) =
- V (k)
((x!1) 18) (6( xx()), nio (h) ((x(0) ye ja (W xx)) (x))x( (95)
Ja
100 d'd" / (ab((TT)) x1 (dā, dw),
(k) ((x,1)
y
Wk 2 0. (96)
for
(k) + exp( (x)y)
20

---

Under review as a conference paper at ICLR 2021
C.2.2 SYM-DEFAULT LIMIT MODEL
Another special case which deserves explicit formulation is what we have called "sym-default"
limit model. The corresponding scaling is: 9- = -1/2, q = 1/2. The resulting limit dynamics is the following:
(b++1) = 10 + div(u (1) ) Me = (1
(97)
A w)
= =1
(k)
((x!1)) Va )ol))(((" xx()) o So ((x! (b) (0) ja
(k)
x xx1)xxxxx )x
(98)
for
(x) = d'd' [ x
1oo (dá dw),
(99)
(0) f/ ((xx 5) = + explordd (0. (x)2(
(100)
((+t1)((x y) = -v(f(++)(x)) <0 V4
(101)
fm
D DEFAULT SCALING
Consider the special case of the default scaling: da = 1122, qu =1, -. 0. Then corresponding dynamics can be written as follows:
A (K) = 0 (d/d
fo ((x!!), (R), 2) ((W()))xx(
~ 1),
(102)
Aw(M) = 0n (d/d") -1 1/2 11//111/////
(k) fa f(x! Vu " (((w!").T x xx! (R) w(() ~N Ia.). (103)
d
(00 (x) = o² (d/d") 1/2
Σ (k) (l (4),)
((
fa (
y) + exp(f1 (x)y)
-y
Wk 2 0.
(104)
As one can see, increments of output layer weights Aa diverge with d. We introduce their ) Aa(") Similarly, we normalize output layer weights themselves: a(*²) (d/d*) 1 (k) Then the dynamics transforms to:
normalized versions: Aa(h) (k) = (d/ dd) 1/2
Ad(h) 1 V (k)
((x (2) T. xx())))) (k),, ai N(O,0(((/
(105)
la
Aw(k) (k) -he
(k)
((x
(b) a(th) 60 (w (+). ( x w/o) NO Iae),
(106)
(K) (x)
(k) 2 w (A).T))
x),
V ((x,))
fa
1 exp(s (k) (x)y)
Vk 2 0.
(107)
Similar to Appendix C we have to introduce weight-space measure in order to take a limit of + 0o:
Pd
(K)
6 00) w (e)
(108)
In terms of the measure the dynamics expressed then as follows:
Pd 1td + div(j1 (8) 10(A)) )
(k+1)
(k)
(109)
21

---

Under review as a conference paper at ICLR 2021
d
(0)
/00 ~ N(O, 8 1000
w/o) ~ ~N(0, e (110)
Pd d rm
10)
@
or
A0 (a,, W)
-110000 V Sa
!(x
.va )o())w")) x(")), o' "V/o e(x!())) (k), w/ )āo' (wo" (wTTx!!)) (111)
(x) 12)((x)
o'd f āØ(w" x) u (dá, dw), (k) ((x, y)
44 0. (112)
1 +exxp(ff exp(Fa (x)y)
We rewrite the last equation in terms of fa (x) = d-1 (1-1 (x):
(x) =0" f
(K) x) pd (dá, dw),
V fa ((x,))
44 0. (113)
1 + exp(ddf ((k)
A limit dynamics then takes the following form:
1 1o + div(u so(), _ = N0,
(k+1)
(114)
A0(k) (ā, W)
00
- V (k)
E ((x!x) vaa
(k) (8) )
x (.)), nh - e(x(()
yh jads (ool(wTTT (wTx )x
(115)
100
fo0
1
1z ((x,)) = + exp or o'd...../// (x)))
(116)
(x) = 0' / (oolTTx x) (k) (dä, dw), (x)=
Inc ((ttt) ((x,)) y) = -ylf(++ (x) 0]
vk 2 0.
(117)
As one can notice, the only difference between this limit dynamics and the limit dynamics of sym-default scaling (Appendix C.2.2) is the initial measure.
We now check the Condition 5 First of all, by the Central Limit Theorem, f1)(x)) ~ (1),
hence the first point of Condition 5
holds. As for kernels, we have:
K a.d xx xxx 02 Σ
d( w (R).T x) (WW
(118)
K ((d(()))))
(M), Σ |a(k) 12 (w: (k).T x) (o''(W!!)))) (w x')x x')x""x' x.
(119)
We see that while (0) converges to a constant due to the Law of Large Numbers, (0) diverges as d + co. This violates the second statement of Condition 5 and the third as well, since f( is finite.
Consider now kernel increments:
AR(tid ax,a(xx))) (x, x') -d'd 1/2
/b (1).))))))) (W
(1): x))
x fz ((x!!) (k) (k) )a! (b) (w (k), (k) )x (xxxxxx++ x(0)) (120)
x()
22

---

Under review as a conference paper at ICLR 2021
d +
3/2 (x,,)x = - (d/ddd
Ja
(1).T) x) (w ( x'+
(x xx)xxx)) x (121)
A AR (x,d(x x') = (d/- - 1/2
2ā (k) & (w) (1)T x) (w())T))
) (w())) !
(122)
For h: = 0 terms inside sums of each increment have zero expectations. Hence the Central Limit Theorem be used here. We get: AK( O (1), (
aw,d Od-oc (dt A K (0),
Os-xx( Since [ Osoox (d), =dd Os (1), the last statement of Condition violated as well.
E INITIALFZATION-CORFETTTD MEAN-FIELD (IC-MF) LIMIT
Here we consider the same training dynamics for the mean-field scaling (see Appendix C.2 but with a modified model definition:
Aa (k)
1000 (k)
((x ) of o( (x)), a ~ N 1),
(R)
(123)
Aw( 1i
(k)
f(x!!
Vw 0r
(k)
( N(0, Ta.)-
(124)
Sa
(x) or (d/d') Σ a (t) (( (b). x) o (d/d') 1/2 Σ a Φ (())x) xx,,
(125)
(1) fa
((x,0)
Wk 2 0.
(126)
1 + exp(f / (x)y)
Similar the mean-field case (Appendix C.2 weight-space measure:
we rewrite the dynamics above in terms of the
(k+1) =! + div(u A0()),
(0)
1
0:00 ~ ~N(00 Vr [d], (127) 8g0
Ha
Ha
A0(1) w)
-11i00 (//i) ((x!(t) v doW" x(1)), f (1) w )ad((w (wTx((h)) (x))))))()) ( (128)
(x) = o=d" d" I ( x) pd (2) (dā, dw) + (dd / "0o(""' x) 11 (dā, dw),
(129)
V (1) 1 1 exp(f
Vk O.
(130)
(k) (x)y)
Note that here ((A) stays finite in the limit of d + xo for any k 2 0. Hence taking the limit d + oo yields:
(k+1) p(( + div(u (k) A0()) 1 (O
(131)
AB(R) (R)
(1) w)
= =- V (k)
(xx
(k) (b) (o(W)> ni ( ((x(0)
y.
(k) 1ee "ao' (wT (xxx(() )x(1T
(132)
for
23

---

Under review as a conference paper at ICLR 2021
(x) d'dd) (ubb(""" x) 1 (dū dw) + o2d':1/2N(0,00 (0).
(133)
((),)(()))) foc 1 expl[xx (x)y)
y
(134)
F EXPERIMENTAL DETAILS
We perform our experiments on a feed-forward fully-connected network with single hidden layer with no biases. We learn our network as a binary classifier on a subset of the CIFAR2 dataset (which a dataset of first two classes of CIFARI of size 1024. We report results using test set from the same dataset of size 2000. We do not do hyperparameter search, for this reason, we do not use validation set.
We train our network for 2000 training steps to minimize the binary cross-entropy loss. We use full-batch GD as an optimization algorithm. We repeat our experiments for random seeds and report mean and deviations in plots for logits and kernels (e.g. Figure [ left). For plots of the KL-divergence, we use logits from these 10 random seeds to fit single gaussian. Where necessary, estimate data expectations (e.g. E E (x())) using 10 samples from the test dataset.
We experiment with other setups (i.e. using a mini-batch gradient estimation instead of exact one, a larger train dataset, multi-class classification) in Appendix G All experiments were conducted single NVIDIA GeForce GTX 1080 Ti GPU using the Py PyTocc framework (Paszke et al. 2017). Our code is available online: (suppressed for anonymity).
Although our analysis assumes initializing variables with samples from a gaussian, nothing changes sample of instead, where can be any symmetric random variable with distribution independent on hyperparameters.
In experiments, took network of width d' 2 27 128 with leaky ReLU activation and apply the Kaiming He uniform initialization (He et al., 2015) to its layers; we call this network reference network. According the Kaiming He initialization strategy, initial weights have zero mean and standard deviation o² (d")-1/2 for the output layer, while the standard deviation of the input layer does not depend on the reference width d'. For this network we take learning rates in the original parameterization V/2 "h- 0.02. After that, scale its initial weights and learning rates with width d according to a scaling at hand:
d
9r
4avw
Navi Nave
d
d
Note that we have assumed ow 1. By definition, nav navi / this implies:
d
da
d
Ta 22
The = ni ( d
Naa
1o
d
d
d
G EXPERIMENTS FOR OTHER SETUPS
Although plots provided in the main body represent the full-batch GD on subset of CIFAR2, we have experimented with other setups as well. For instance, we have varied the batch size and the size of the train dataset. Results are shown in Figures Here differences are marginal and not qualitative.
We have also experimented with multi-class classification: see Figure g Here we have trained our network on the full CIFARIO dataset with SGD with batches of size 100. As we see left plot, the IC-MF limit model has the lowest KL-divergence relative to the reference model, however, in terms of the test accuracy, all the limit models are similar.
CIFAR 10 be downloaded httpe///e////////ooooooooeo/o//r html
/ttp:://wwwww..o
24

---

Under review as a conference paper at ICLR 2021
train size = 1024, batch size = 1024
train size 1024, batch size 256
1.0
1.0
accuracy
0.8
accuracy
0.8
0.6
0.6
0.4
0.4
10
10a
102
10a
10a
102
10a
training step, k+1
training step, k+4
train size
1024, batch size
64
train size 8192, batch size 512
1.0
1.0
accuracy
0.8
accuracy
0.8
0.6
0.6
0.4
0.4
102
10a
102
10a
training step, k+16
training step, k+16
Figure 3: Test accuracy of different limit models, as well as of the reference model. Serup: We train one hidden layer network subsets of the CIFAR2 dataset of different sizes with SGD with varying batch sizes.
H GENERALIZATION TO DEEP NETS PROPOSAL
While our present analysis is devoted to networks with a single hidden layer, we discuss possible generalizations to deep nets here.
Consider network with H hidden layers. For simplicity, assume that widths of all hidden layers are equal to d. We thus have consider H +-1 learning rates fo.H1 for each layer, and similarly H initialization variances d..H. Without loss of generality, we may assume the input layer variance to be equal (we can rescale inputs otherwise). This gives 2/1 + hyperparameters in total.
Similarly to what we did for H = 1, we assume that each hyperparameter obeys power-law with respect width. Let refer the of the power-law exponents reason about what the scaling should be in order to converge to dynamically stable limit model: see Condition[]] Moreover, we want derive conditions that separate the domain of "dynamically stable" scalings, such that each region corresponds to a distinct unique dynamically stable limit model: see Condition 2
"scaling". Again, we want
Having that much hyperparameters seems burdening, and this prohibits to draw nice two- dimensional scaling plane as we did for // 1: see Figure [ For this reason, one have to reduce dimensionality of scaling.
First, it tempting to consider a homogeneous activation function: leaky ReLU. This introduces symmetry in the weight space that guarantees that dynamics depends only the product of initialization variances: σ X X do: let us refer this product as a.. This approach was previously used by Golikov (2020), however have note that non-smoothness of the activation function introduces certain mathematical obstacles. Nevertheless, one may consider sacrificing mathematical rigor in favor of reducing the number of hyperparameters from 2H to H + 1.
The next simplification should affect learning rate scaling exponents. Similar to what we have done for shallow net, may assume all learning rate exponents to be equal: dn ... qn q.
25

---

Under review as a conference paper at ICLR 2021
train size 1024, batch size = 1024
mean diag kernel, E,Kix, x)
train size 1024, batch size 256
mean diag kernel, ExKix,
10²
102
10
10
10a
102
training step, k+1
training step, k+4
mean diag kernel, E,Kix, x)
train size 1024, batch size = 64
mean diag kernel, E,K(x,x x)
train size 8192, batch size
512
10a
102
2 10 x 10 10 6 10
10
2 10 x 10 10 6 x 10 training step, k+16
10a
training step, k+16
Figure 4: Mean kernel diagonals E x>>D(f> (f Ka (x, K (x,)xxx))) of different limit models, as well as of the reference model. Setup: We train one hidden layer network on subsets of the CIFAR2 dataset of different sizes with SGD with varying batch sizes. Data expectations are estimated with 10 test data samples.
The NTK limit, which generalizes naturally to deep nets, requires do ... 4H 0. and hence conforms the assumption above. However, possible generalization of the mean-field limit requires 4o O while @ Golikov (202)). This aspect suggests the following alternatives:
qH 2; see (Siiignano & Spiliopoulos/ (2019); Araújo al.)
(2019)
Consider to QH = q. while ` = ... -H qridi this results dhree-aimensional space of scalings: (4a 7, Thid).
2. Consider qe = ĀH = q. while & = =... = - = 24; this results in a two-dimensional space of scalings that covers both of the NTK and the mean-field scalings.
The former class of scalings is richer, but if it does not contain any interesting limit models that are present the second class, it can more expository to tighten the class the latter. By "interesting" we mean limit models that are "non-dominated" in a similar sense as we have specified in Section 3
In order define which limit models are better than others approximating finite-width dominated"), we have to derive conditions that separate the domain of dynamically stable scalings into regions of distinct unique corresponding limit models, similar Condition 2 We hypothesize that these conditions are similar to the shallow case: (1) limit model at initialization is finite, (2) kernels at initialization are finite, (3) limit model and kernels are of the same order, (4) kernels evolve at initialization. Since have decided to consider separate learning rate scalings for hidden layers and for input and output layers, we expect that the above-proposed conditions should consider two distinct families of kernels respectively: hidden kernels and input plus output kernels.
("non-
will be very interesting to check all of the dynamically stable limit models are specified either by an evolution in model space driven by constant kernel, by evolution of weight-space measure, as was the case for = 1; see Appendix C Investigating a non-dominated limit model, different from both the NTK and the mean-field models, should be valuable outcome of the proposed research
26

---

Under review as a conference paper at ICLR 2021
train size 1024, batch size 1024
train size 10224 batch size 256
abs logit, Exf((x)
mean abs logit, Exf((x)
10a
10a
100
10
10-1
10
10-2
10
10
10
10a
10a
10a
training step, k+1
training step, k+44
train size
1024, batch size = 64
train size 8192, batch size
512
mean abs logit, Ex]((x)
10a
mean abs logit, Exf((x)
10
10
100
10-1
10
10-2
10
2 10 3 10
102
2 10a 3 10 10 10* training step, k+16
10a
training step, k+16
Figure 5: Mean absolute logits Ex~~ f [x) of different limit models, as well of the reference model. Serup: We train a one hidden layer network on subsets of the CIFAR2 dataset of different sizes with SGD with varying batch sizes. Data expectations estimated with 10 test data samples.
train size 1024, batch size 1024
train size = 1024, batch size 256
10
10
5 x]]
10-2
10-2
10-2
1
10c
101
102
10'
10a
training step, k+1
training step, k+4
train size
1024, batch size
64
train size
8192, batch size
512
10
10"
Ex/fbo/KKx, x11
10
10
10
10a3
2 10a 3x 10xx 10a 6x 10a
10a
2 3 3 10x 10 10
training step, k+16
training k+16
Figure 6: Mean absolute logits relative kernel diagonals (x>Dlfa(x)/Cl (f) Ka ax,a(x,5 x) fh (x,dxx, x)]] of different limit models, well as of the reference model. Setup: We train one hidden layer network on subsets of the CIFAR2 dataset of different sizes with SGD with varying batch sizes. Data expectations are estimated with 10 test data samples.
program; it will be even more valuable this limit model will not be covered by both mean-firn -ield and constant kernel formalisms.
27

---

Under review as a conference paper at ICLR 2021
train size 1924. batch size 1024
train 1024. betch size 256
10a
10"
NTK
10
defieutt
K. MP
10,0irst ruf)
sym
-07
sym
200
10
10"
10a
10'
30'
10
10'
10' traliing stap. K44
training step. k+1
train size 1024, batch size 64
10"
10-
ICMF
107
NIK
K-NY
defailt
1.07
don-dofa.lt
10:4
15 training step, 16
10
A0t tra ning step. k-16
LUP
Figure 7: KL-divergence of different limit models relative to reference model. Setup: We train one hidden layer network on subsets of the CIFAR2 dataset of different sizes with SGD with varying batch sizes.
10a
NTK
ICMF
0.550
NT
K-AP
enfault
sym-dcfout
0.52%
K4Cimi:
0.500
0.475
0.656
0.455
training step,
ning step.
Figure Left: KI--divergence of different limit models relative to a reference model. Right: Accuracies the test set of different limit models well as of the reference model. Setup: We train a one hidden layer network on the full CIFARIO dataset with SGD with batches size 100.
We also have to note that according to Golikov (020). the mean-field limit vanishes for H > This fact suggests that the analysis for deep nets should be held for H = 2 and for H > separately.
I MEASURING DIVERGENCE BETWEEN A LIMIT MODEL AND A REFERENCE ONE
We track the divergence of a limiting network from reference one, which can be done in two ways.
The first one is tracking divergence directly between logits: L Ex~PPr. Diogittee (x) || f( (x)) for some divergence measure D || -). The second is tracking divergence between probabilities:
E Ex~DD... pros (f (x)) || (f
(x))), where we have overloaded the notation by denoting the
standard sigmoid as a: (x) = (1 +
We choose KL-divergence for the first case. However, measuring KL-divergence between logits is hardly possible, since do not have access the distribution of f()) (x) random variable depending on initialization. For this reason, we a gaussian to its samples:
Diogits (5 || (*) KL(W(ES, Var E) || W[E Var
(135)
This case is depicted in Figure 9 left.
28

---

Under review as a conference paper at ICLR 2021
10"
10"
NTK
ICMM
107
10a
default
sym
KUElmn,, nuf
107
MF
100
syrechfask
10 303
100
Figure 9: Left: plot KL-divergence of logits of different infinite-width limits of fixed finite- width reference model relative to logits of this reference model. KL-divergences are estimated using gaussian fits with 10 samples. Right: same, for probabilities instead of logits. KL-divergences are estimated using beta-distribution fits with 10 samples. Setup: we train one hidden layer network with SGD on CIFAR2 dataset; see Appendix for details.
As for the second case, may want measure KL-divergence between distributions on probabili- ties. Again, this not possible, because the true distribution ([((((() is not known; for this reason, we decide to first fit beta distribution, and then measure the divergence:
Dprob || 5**) KL (Betaaa(emmte (E), (E)) || Beta(a ()))(
(136)
where Ommtt (E) and (Ente(E) are maximum-likellban estimations of hyperparameters of a beta:
[R EEE) -1). (E) = (1 - Rgg)
EE - (EC)-
-1) (137)
(Emle(E) = E
Var E
War
This case plotted in Figure 9 right.
Both cases can be generalized to work with multi-class classification. In the first case, we can simply fit gaussian with diagonal covariance matrix; this was done Figure left. the second case, -okimuum-ikklihood likelihood estimation before.
we can fit Dirichlet random variable using maximum
29


Reviews:
[{"review_id": "f36_I7ZdNTV", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\nThis paper investigates the dynamics of the training of fully connected neural networks with one hidden layer in the infinite-width limit, for classification. Starting from the observation [Golikov, 2020] that the mean-field limit and the NTK regime are only two special cases of continuous families of scalings, it attempts at identifying important or desirable features that an infinite-width limit should display in order to give good insight into the dynamics of large but finite neural networks. \n\nThis is an exciting program, and it looks like the ideas are very good. If / when laid out, they will really bring insight into the training of large neural nets.\n\nUnfortunately, the way the paper is written makes it hard to accept in its present form. While the global structure is quite clear, the wording makes it very hard to extract information. For instance, it is not clear what is mean by \"finite\": does it mean there is a finite limit (what should be expected in principle), that it is bounded from above (what it seems to mean sometimes) or that it is strictly positive (what it appears to mean sometimes)? Also, there is a classification into 13 cases that are promised, but they are not clearly listed in the main, the fact that we are working with only one hidden layer is not clearly said in the abstract (and the discussion on how to extend to more hidden layers is not convincing). The grammar is somehow problematic (there is in particular a lot of missing \"the\" articles in front of nouns), and it sometimes makes it hard to follow. Also, it is not very clear what is proven, where assumptions are used (e.g. analyticity of the nonlinearity), the seemingly most interesting regimes are not defined in the main, etc.  \n\nI am tempted to think that the authors are very lucid about their understanding of what is happening and that they really have interesting something to convey, but the present version makes it very hard to get a reliable information. A (very significantly) revised version of this paper could, I believe, bring much insight to our understanding of neural nets. There is a lot of potential with this paper, just not realized in terms of exposition.", "extra_scores": {}, "extra_text": {}}, {"review_id": "_ZC206F-Kb", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposed a general framework to derive different stable limiting behaviors of the dynamics of two-layers neural networks, under different parameterization of the hyper-parameters. For certain choices of hyper-parameters, this recovers the mean-field limit and the NTK limit. This paper also proposed certain properties of the limiting dynamics and showed that using these properties as the classification criteria, there are only a finite number of distinct models in the limit. This paper also proposed a novel initialization-corrected mean-field limit that satisfies all properties.\nThis paper tells an interesting story. The question is whether the problem solved by the story is important.\nThe main consideration of this paper is to find regimes of hyper-parameters such that the limiting dynamics are stable in some sense. For this purpose, the authors advocate the IC-MF regime, such that the limiting dynamics are stable with respect to all the conditions that the authors proposed. This story seems to be self-contained on its own. However, I believe that two more important criteria for good training algorithms are optimization and generalization efficiency, which are not discussed by the authors.\nThe optimization and generalization efficiency are much more important than the stability condition in practice. There could be some regimes of hyper-parameters that do not satisfy the stability condition, but has good optimization and generalization efficiency. The IC-MF regime proposed in this paper, although seems to satisfy additional stability conditions, but intuitively, it seems that its generalization efficiency is not as good as that of MF regime: it added an additional noisy function $f_{ntk, \\infty}^{0)}$ to the mean-field prediction function, and this additional noise (intuitively will) hurt generalization.\nI believe some of these stability properties should have some connection to optimization and generalization. For example, if some simple stability property is violated, the algorithm cannot generalize well. I feel that the authors should try to build connections of stability properties to optimization and generalization, to justify the importance of condition 1 and condition 2 defined in the paper.\nAbove all, I feel that this paper is interesting in its own criteria. However, it didn't justify that its criteria are important. So I feel that this paper is on the borderline.\n\nMinor issues:\n\t1.\tSome notations are easy to get readers confused. Eq. (7), $\\sigma(d) = \\sigma^*(d / d^*)^{q_\\sigma}$. Here $\\sigma$ is a function of $d$ while $\\sigma^*$ is a scaler (not as a function of $d / d^*$). It takes me while to understand this.\n\t2.\tTypos: page 18: in this case $1 + \\tilde q + 2 q_\\sigma$.\n\t3.\tThe notations of this paper looks very complicated, especially the superscripts and subscripts.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "9t_VGS5ys3G", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper analyzes joint scalings of the parameter initialization and the learning rate, with respect to the limit of infinite width, in the context of two-layer neural networks with stochastic gradient descent and binary logistic loss. It proposes some “dynamically stable” conditions and identifies a range of scalings that satisfy these conditions. This covers the neural tangent kernel (NTK) and mean field (MF) scalings, as well as others. The paper then proposes to add an extra correction function to the initialization of the MF limit and argues experimentally that this correction can give a proxy for standard neural networks.\n\nOverall there are interesting points in the paper, but I think the paper needs a lot more works to make the study rigorous, complete and easier to read. More analyses into why any point in the “dynamically stable model evolution” band leads to stable dynamics at all training time are needed.\n\n\n### Positive points:\nThe recent trend in analysis of large-width neural networks has drawn a lot of attention, so the paper is timely. \n\nThe idea of joint scalings between the initialization and the learning rate, w.r.t. the width d, has been explored before; but here the paper proposes to do so in relation with a new set of criteria. The results make an interesting view on the connection between two distinct regimes NTK and MF, and beyond.\n\n\n### Negative points:\n\n**Presentation:**\nThere are many unnecessary notations, for example, \\tilde{q}_a and \\tilde{q}_w which are equal to each other in the main analysis. It is better to keep all discussions around just the two key scaling exponents. \n\nThe main figure, Figure 1, is very hard to read, and I cannot see where the curve for MF is in this plot. \n\nThe statements also seem to miss several assumptions e.g. assumptions on the data (will things hold if |x| ~ exp(d)??).\n\nSome notations are not explained, for example, \\sigma^{(0)}(x) that appears in Eq (13).\n\n**Incomplete proofs:**\nAside from many derivations which are heuristic (which are fine as long as they are not stated as propositions/theorems), the proofs for the stated lemmas / propositions are unrigorous and may entirely omit the more difficult technical points. For example, in the proof of Lemma 1.4 in Appendix B.1, why is it that the collection of the weights \\hat{w}_r^{(k), r=1…,d, at any time k>1 allows one to apply the law of large numbers? The variables in this collection are not independent; if they are mildly dependent, in what sense are they so that one can apply the law of large numbers?\n\nIn fact, there are simple mathematical mistakes throughout the proofs. For example, to prove Lemma 1.2 in Appendix B.1, the paper only proves the statement for a single fixed x; why does it hold for almost every x?\n\n**Abstract missing key information:**\nThe abstract should be more precise: it should mention logistic loss for binary classification and that the initializations are zero-mean. These are important; without them, a number of key results will fail.\n\n**Dynamical stability criteria are insufficiently justified; most are limited to just initializations:**\nThe paper does not give sufficient justification of why the proposed criteria (Conditions 1 and 2) should be considered. In particular, Condition 2 only concerns with what happens at initialization; does it guarantee stable dynamics well beyond initialization and in what sense? The paper does not discuss this for most points in the “dynamically stable model evolution” band. In fact, compared to the well-studied NTK and MF scalings, the newly identified scaling sym-default is said to have some divergence behavior (in which its prediction function is always infinite, for infinite d, at any positive training time). If anything, the paper should justify why this divergence behavior does not lead to instability; by looking at Figure 1, I think it suffers from numerical instability.\n\nIn fact, scalings other than NTK and MF have been studied well beyond initialization in a rigorous fashion. In particular, [1] shows that with a wide range of different scalings, it’s entirely possible to have the same stable behavior: NTK behavior where weights do not move and the model tracks some random-feature models. Given that such rigorous and thorough study is possible, it is unclear how the proposal of Conditions 1 and 2 leads to novel and significant insights.\n\n**Condition 2.3 is unusual in terms of physical units and not motivated:**\nIn particular, why should we compare the magnitude of the kernel K with the magnitude of the prediction function f, while they should have very different physical interpretation? Why is this condition interesting? The second paragraph in Section 3 seems to give some arguments for Condition 2.1 and 2.2, but none for Condition 2.3.\n\nGiven that the role of Condition 2.3 is unclear, if one removes this condition from the diagram of Figure 1, the band picture disappears and we are effectively left with a half plane, separated by the “evolving kernels” boundary. To the left of this boundary, it is the NTK behavior as expected in [1]. So unless Condition 2.3 signifies some significant and meaningful change in network behaviors, the diagram would not show new significant insights compared to the literature.\n\n**The derivations for the dynamics in Section 3 are incomplete:**\nIn the regime of small learning rate (\\hat{\\eta}* goes to 0), what we should obtain is a continuous time evolution with an expectation over data, not a discrete-time one with stochastically drawn data. Evidently Eq (14) as given in the paper completely contradicts all previous published reports on the MF limit.\n\nIt does not make sense to claim, after Equation (19), that the prediction function diverges from iteration k=1 onwards. Firstly, as said, the analysis has to be done with respect to small learning rate and hence the iteration k has to be determined in relation with the learning rate. Secondly, that claim contradicts Figure 1: the prediction function does not seem to diverge at iteration 1.\n\n**The initialization-corrected mean field (IC-MF) limit lacks justification:**\nIts aim is to correct the MF limit w.r.t. Condition 2.1, but there are very simple alternatives to do this correction. The first way is to add any fixed function, whose magnitude are independent of d in suitable sense, to the MF limit. The second way is to initialize the MF limit with non-zero mean distributions; in fact, it is known that non-zero mean initializations are more typical in MF limit.\n\nThe paper argues that IC-MF limit can be good proxy for standard networks, giving only one simple experiment shown in Figure 1. There is no proof or mathematical heuristics provided. I think this IC-MF idea does not have anything to do with the theory in the previous sections, nor the binary classification problem with logistic loss. As such, it should be at least further tested on more complex experimental tasks, e.g. CIFAR-10. From a mathematical standpoint, I do not foresee a simple argument to show why it can be a good proxy for standard networks.\n\n\n[1] A comparative analysis of the optimization and generalization property of two-layer neural network and random feature models under gradient descent dynamics, E, Ma and Wu, 2019.\n\n---------------------\n\nPost rebuttal: \n\nI thank the authors for their rebuttal. Let me focus my reply on a few important points. I first thank the authors for clarifying the meaning of Condition 2. In this sense, Condition 1 is the key main contribution; however the current proof does not look correct to me, and the revised argument is far from being sufficient. In particular:\n\n- Point 5 of the rebuttal: I think the revised argument here is incomplete. The given argument concerns trivial facts and does not imply the claim. For example, what if the distribution of the terms is symmetric, the expected sum is zero, and hence the quantity might be of order smaller than d? Note that this is an example problem; there are multiple problems with the proof of Lemma 1.4. For example, the paper claims this for all k; but if k is something like d^100, would things hold? What would stop the magnitude of the weights to grow with time?\n\n- Point 6 of the rebuttal: The CLT, when applied w.r.t. the randomness of the weights, says that for a fixed $x$, $\\sum_{r=1}^d \\hat{a}_r \\phi (\\hat{w}_r x) / \\sqrt{d} \\sim N(0,v_x)$ approximately. That is, there is a non-zero probability (w.r.t. the randomness of weights) that the claim in the paper for a fixed $x$ fails. As such, to reason the claim for many $x$, one requires doing probabilistic arguments very carefully.\n\nThe paper should execute the proof very carefully. It is not just a matter of technicality; I suspect some of the claims are actually wrong.\n\nMore importantly, Condition 1 alone is insufficient to claim dynamical stability at any time k. What should qualify for dynamical stability is rather the existence of a well-defined limiting dynamics exists (which is argued heuristically in Appendix C), and its proof. In the current writing, it’s unclear how Condition 2 is crucial; while it studies interesting properties, it is very restrictive.\n\nAs said in my last review, one thing that has been missing is really whether the insight here differs qualitatively from the known NTK and MF limits. Further looking at the limiting dynamics in Appendix C, one sees that they are qualitatively either NTK or MF. There are possible degeneracies due to scalings and the use of logistic loss, but these do not lead to much deviation from NTK or MF behaviors. If one is to use a squared loss for instance, what one would obtain in Figure 1 is just the line connecting NTK and MF; all other points in the band outside this line are degeneracies due to logistic loss. The behavior on this line, again, is qualitatively either NTK or MF, and this is shown (somewhat implicitly) already by a number of past works.\n\nI would imagine a rigorous derivation of the limiting dynamics for each point in the band revolves around the renormalized dynamics in Appendix C. When translating from the renormalized dynamics to the original one, the extra scaling factors will complicate the proof (for instance, they can blow up Lipschitz constants). Again this has to be done very carefully.", "extra_scores": {}, "extra_text": {}}, {"review_id": "3Mno_f9gMj", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nMain topic of the paper is to study various infinite width limits. The paper notices different scaling used in NTK limit and MF limit and proposes a general framework for studying the limit behavior depending on the scaling. This allows the authors to define more general dynamically stable models in the infinite width limit. \n\nStated contribution by authors are:\n\n- Framework for reasoning about scaling that leads to ‘dynamically stable’ model evolution in the infinite width limit. \n- Categorization of 13 distinct dynamically stable models in the infinite width limit.\n- Characterization of “sym-default” model, which along with NTK and MF limit which shows properties most of finite-width network evolution. \n- Model modification “Initialization-corrected mean-field limit”(IC-MF) that satisfy all identified property of finite-width network evolution\n- Demonstration of IC-MF limit approximating finite-width network the best among all other possible models. \n \nReason for score:\n\nThe paper proposes an interesting theoretical framework to capture various infinite width limits. While settings are limited to make theory tractable, there is some empirical validation as well as capacity to broaden the study to different infinite width limits. I believe there are interesting novelties to be shared among ICLR participants who are interested in deep learning theory. \n\nPros:\nGeneral framework encapsulates widely studied infinite width limits and generalizes them which would be useful for theoretical study of large neural networks.\n\nThe paper makes predictions on scaling limits and finds a limit that could agree with finite networks better (IC-MF). For a variety of experiments the fact that this class matches with a reference finite width network is demonstrated well. \n\nWhile the proposed analysis is limited to a single hidden layer case, the authors have described thoughtful possibilities how the framework could potentially generalize to deeper networks. \nAuthors shared code to reproduce their experiments in the paper which is useful for reproduction and clearly understand different proposed scalings.\n\nCons:\n\nBeing single hidden layer analysis is a drawback (while discussion on extension is appreciated and stated in the “pros”). Often multi-layer infinite width is more interesting since the number of hidden layer weights scales quadratically while only the input and output layer weight count scales linearly with width. This can induce different training dynamics and scaling, so limitation to single hidden layer analysis is a drawback. \n\nTo simplify proof non-linearity of the choice is “leaky softplus” which is not widely used in practice. Having general comments on extension or limitation to well-known activations is advised. \n\nAssumption in p3 stating “gradients for a and W are estimated using independent data samples” may be too strong and not realistic. I cannot find no validation that the assumption is a reasonable one to make beyond making the theory simple. This may lead to inconsistencies raised in the “Questions” below. \n\nExperiments demonstrating similar effect for more realistic setting (e.g. full CIFAR-10 classification task achieving reasonable performance) would have been better to convince the validity or generalizability of the theory. While one can see that IC-MF and reference agrees well for the training curve, for test-set since every curve essentially behaves the same it is hard to determine which limit is superior or not. \n\nQuestions:\nIn [Liu et al., 2020], they show non-constancy of NTK for non-linear output including the soft-max layer in the outputs. I believe this fact may be in conflict with this paper’s analysis on binary cross-entropy loss and still having constant NTK. I do suspect it may be due to assumptions on independent gradients for `a` and `W` in page 3. Do you think in a realistic training setting with softmax-cross entropy, you can generalize the analysis and reconcile the fact that NTK is non-constant?\n\nIn general, does the generalization of “dynamical stability” to admit infinite logits valid for multi-class cases? In multi-class simple “sign” can’t be used for classification and I wonder if authors have ideas on generalization to multi-class settings. \n\nIn eq (26), why is Gaussian distribution over logit a good way to measure KL divergence? I understand at initialization prior is distributed over Gaussian, however after training with softmax I expect the logit distribution is no-longer Gaussian. \n\nNits and additional feedback:\n\nNit: use \\citep in places where appropriate.\n\nDivision (/) for denoting “or” is often confusing and would suggest other notation. Especially in the Figures and Condition 2 when sometimes it means “or” and sometimes it means division.\n\np4: ‘grows width’ -> ‘grows with’\n\n[Sohl-Dickstein et al., 2020] show different scaling of weight and bias scaling extending standard parameterization to work well with NTK limit. While this paper mostly works with the NTK parameterization, it would be interesting to discuss how the improved parameterization in [Sohl-Dickstein et al., 2020] can be utilized. I suspect this is quite related to (9) / (10) where definition of NTK deviates from [Jacot et al., 2018]. \n\nIn Section 3: When comparing performance on finite networks and NTK, it can be subtle depending on how one trains finite networks and there’s architecture dependence [Lee et al., 2020] \nIn Section 4: suggest citing [Chizat et al., 2019] for ‘lazy training’\n\nChizat et al., On Lazy Training in Differentiable Programming, NeurIPS 2019\nLee et al., Finite Versus Infinite Neural Networks: an Empirical Study, arXiv: 2007.15801\nLiu et al., On the linearity of large non-linear models: when and why the tangent kernel is constant, arXiv:2010.01092\nSohl-Dickstein et al., On the infinite width limit of neural networks with a standard parameterization, arXiv:2001.07301\n\n----\n[post rebuttal]\nI thank the authors for their revision and clarifying many of my questions and adding results based on that. I have read other reviewers concern and while I agree some room for improvement on clarity, I still think the paper brings in value. Unless there's technical flaws spotted by other reviewers that has not been resolved, I'm still leaning towards accepting (increased score from 6 to 7). ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a general framework to study the limit behavior of neural models with respect to the scaling of hyperparameters in terms of network width, which covers existing mean-field (MF) and neural tangent kernel (NTK) limits, as well as other new limit models that were not discovered before. While the reviewers agree that the study of limiting behavior of neural network models is of great importance and could be a good addition to the current understanding of NTK and MF, there is a technical flaw in the proof (regarding Condition 1) pointed by the reviewer. After reviewer discussion, all reviewers agree that this is a serious issue, and needs to be addressed before publication. I believe it could be a strong paper if the technical flaw can be fixed. I encourage the authors to revise the paper and resubmit it to the next conference. 
Decision: Reject