Year: 2020

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.

heading: 1 INTRODUCTION
text: For a couple of decades neural networks have proved to be 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 a finite-width neural network with its limit counterpart in the limit of a large number of hidden units; we refer it as an "infinite-width" limit. The exact type of the limit training dynamics depends on how hyperparameters of the training dynamics scale with width. In particular, two different types of limit models have been already extensively discussed in the literature: an NTK model (Jacot et al., 2018) and a mean-field limit model (Mei et al., 2018; 2019; Rotskoff & Vanden-Eijnden, 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 a framework for choosing a scaling of hyperparameters that lead to a "well-defined" limit model. Our work is the next step in this direction. We study infinite-width limits for networks with a single hidden layer trained to minimize cross-entropy loss with gradient descent. Our contributions are following.
1. We develop a framework for reasoning about scaling of hyperparameters, which allows one to infer scaling parameters that allow for a 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, as 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 finite-width model. Each distinct limit model corresponds to a region (two-, one-, or zero-dimensional) of a green band of the Figure 1, left.
3. We consider a list of properties that are statisfied by the evolution of finite-width models, but not generally are for its infinite-width limits. We demonstrate that mean-field and NTK
limit models, as well as "sym-default" limit model which was not discussed in the literature previously, are 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 "initialization-corrected mean-field limit (IC-MF)".
4. We discuss the ability of limit models to 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 a single hidden layer, we discuss a high-level plan for generalizing it to deep nets, as well as an expected outcome of this research program, in App. H.

heading: 2 TRAINING A ONE HIDDEN LAYER NET WITH SGD
text: Here we consider training a one hidden layer net fd with d hidden units with SGD. We assume the hyperparameters, namely, initialization variances and learning rates, are scaled as power-laws of d. Each scaling induces a limit model f∞ = limd→∞ fd. We present a notion of dynamical stability, which states that the change of logits after a single gradient step is comparable to logits themselves. We derive a necessary condition for dynamical stability in terms of the power-law exponents of hyperparameters. We then present a list of conditions that divide the class of scalings into 13 subclasses; each subclass corresponds to a unique distinct limit model.
Consider a one hidden layer network:
f(x;a,W ) = aTφ(WTx) = d∑ r=1 arφ(w T r x), (1)
where x ∈ Rdx , W = [w1, . . . ,wd] ∈ Rdx×d, and a = [a1, . . . , ad]T ∈ Rd. We assume a nonlinearity to be real analytic and asymptotically linear: φ(z) = Θz→∞(z). Such a nonlinearity
can be, e.g. "leaky softplus": φ(z) = ln(1 + ez) − α ln(1 + e−z) for α > 0. This is a technical assumption introduced to simplify proofs. Note that we have used traditional leaky ReLUs in our experiments: see App. F for details. We assume the loss function `(y, z) to be the standard binary cross-entropy loss: `(y, z) = ln(1 + e−yz), where labels y ∈ {−1, 1}. The data distribution loss is defined as L(a,W ) = E x,y∼D`(y, f(x;a,W )). We assume that the data distribution D does not depend on width d.
Weights are initialized with isotropic gaussians with zero means: w(0)r ∼ N (0, σ2wI), a (0) r ∼ N (0, σ2a) ∀r = 1 . . . d. The evolution of weights is driven by the stochastic gradient descent (SGD):
∆θ(k) = θ(k+1) − θ(k) = −ηθ ∂`(y
(k) θ , f(x (k) θ ;a,W ))
∂θ , (x
(k) θ , y (k) θ ) ∼ D, (2)
where θ is either a or W . We assume that gradients for a and W are estimated using independent data samples (x(k)a , y (k) a ) and (x (k) w , y (k) w ). 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:
â(k)r = a
(k) r
σa , ŵ(k)r =
w (k) r
σw , η̂a = ηa σ2a , η̂w = ηw σ2w . (3)
Then the dynamics transforms to:
∆θ̂(k)r = η̂θ ∂`(y
(k) θ , f(x (k) θ ;σaâ, σwŴ ))
∂θ̂r , (x
(k) θ , y (k) θ ) ∼ D, (4)
while scaled initial conditions become: â(0)r ∼ N (0, 1), ŵ(0)r ∼ N (0, I) ∀r = 1 . . . d. By expanding gradients, we get the following:
∆â(k)r = −η̂aσa∇ (k) fd `(x(k)a , y (k) a ) φ(σwŵ (k),T r x (k) a ), â (0) r ∼ N (0, 1), (5)
∆ŵ(k)r = −η̂wσaσw∇ (k) fd `(x(k)w , y (k) w ) â (k) r φ ′(. . .)x(k)w , ŵ (0) r ∼ N (0, I), (6)
∇(k)fd `(x, y) = ∂`(y, z)
∂z ∣∣∣∣ z=f
(k) d (x)
= −y
1 + exp(f (k) d (x)y)
, f (k) d (x) = σa d∑ r=1 â(k)r φ(σwŵ (k),T r x).
Without loss of generality assume σw = 1 (we can rescale inputs x otherwise). We shall omit a subscript of σa from now on. Assume hyperparameters that drive the dynamics obey power-law dependence on d:
σ(d) = σ∗ × (d/d∗)qσ , η̂a(d) = η̂∗a × (d/d∗)q̃a , η̂w(d) = η̂∗w × (d/d∗)q̃w . (7)
Given this, a network of width d∗ has hyperparameters σ∗ and η̂∗a∨w. Here and then we write "a∨w" meaning "a or w".
This assumption is quite natural: for He initialization (He et al., 2015) commonly used in practice σ ∝ d−1/2, while we keep learning rates in the original parameterization constant while changing width by default: ηa∨w = const, which implies η̂a ∝ d and η̂w ∝ d0. On the other hand, NTK scaling (Jacot et al., 2018; Lee et al., 2019) requires scaled learning rates to be constants: η̂a∨w ∝ d0. Scaling exponents (qσ, q̃a, q̃w) together with proportionality factors (d∗, σ∗, η̂∗a, η̂ ∗ w) define a limit model f (k)∞ (x) = limd→∞ f (k) d (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 1 (informal version of Condition 4 in Appendix A). Let ∆f (k)d (x) = f (k+1) d (x)−f (k) d (x).
∃kbalance ∈ N : ∀k ≥ kbalance ∆f
(k) d
fkbalanced stays finite for large d.
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 one 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 q̃a = q̃w = q̃. We prove the following in Appendix B.1: Proposition 1. Suppose q̃a = q̃w = q̃ and D is a continuous distribution. Then Condition 1 requires qσ + q̃ ∈ [−1/2, 0] to hold.
This statement gives a necessary condition for growth rates of σ and η̂ to lead to a well-defined limit model evolution. This condition corresponds to a band in (qσ, q̃)-plane: see Figure 1, left. We refer it as a "band of dynamical stability".
Each point of this band corresponds to a 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 a single limit model evolution.
We start with defining tangent kernels. Since φ is smooth, we have:
∆f (k) d (x) = f (k+1) d (x)− f (k) d (x) = d∑ r=1 ∂fd(x) ∂θ̂r ∣∣∣∣ θ̂r=θ̂ (k) r ∆θ̂(k)r +Oη̂∗a∨w→0(η̂ ∗ aη̂ ∗ w + η̂ ∗,2 w ) =
= −η̂∗a∇ (k) fd `(x(k)a , y (k) a ) K (k) a,d(x,x (k) a )− η̂∗w∇ (k) fd `(x(k)w , y (k) w ) K (k) w,d(x,x (k) w )+O(η̂ ∗ aη̂ ∗ w+ η̂ ∗,2 w ),
(8)
where we have defined kernels:
K (k) a,d(x,x ′) = (d/d∗)q̃aσ2 d∑ r=1 φ(ŵ(k),Tr x)φ(ŵ (k),T r x ′), (9)
K (k) w,d(x,x ′) = (d/d∗)q̃wσ2 d∑ r=1 |â(k)r |2φ′(ŵ(k),Tr x)φ′(ŵ(k),Tr x′)xTx′. (10)
Here we deviate from the traditional definition of tangent kernels (e.g. from Jacot et al. (2018)) in embedding learning rate growth factors into kernels. This is done for avoiding 0 ×∞ ambiguity when η̂a∨w grows with d while σ vanishes so that "a learning rate times a kernel" stays finite. This is the case for the mean-field scaling: η̂a∨w ∝ d, while σ ∝ d−1. While for the NTK scaling kernels stop evolving with k in the limit of large d, this is not the case generally. Indeed, for the mean-field scaling mentioned above we have:
K (k) a,d(x,x ′) = σ∗,2(d/d∗)−1 d∑ r=1 φ(ŵ(k),Tr x)φ(ŵ (k),T r x ′). (11)
Similarly to the NTK case, the kernel above converges due to the Law of Large Numbers, however in contrast to the NTK case the weights evolve in the limit: ŵ(k)r 9 ŵ(0)r . This is due to the fact that weight increments are proportional to η̂wσ which is ∝ d0 for the mean-field scaling but ∝ d−1/2 for the NTK one. For this reason, similarly to model increments ∆f (k)d we define kernel increments:
∆K (k) a∨w,d(x,x ′) = K (k+1) a∨w,d(x,x ′)−K(k)a∨w,d(x,x ′). (12)
Condition 2 (informal version of Condition 5 in Appendix A). Following conditions separate the band of dynamical stability (Figure 1, left):
1. f (0)d stays finite for large d.
2. K(0)a∨w,d stays finite for large d.
3. K(0)a∨w,d/f (0) d stays finite for large d.
4. ∆K(0)a∨w,d/K (0) a∨w,d stays finite for large d.
We prove the following in Appendix B.2:
Proposition 2 (Separating conditions). Given Condition 1, Condition 2 reads as, point by point:
1. A limit model at initialization is finite: qσ + 1/2 = 0.
2. Tangent kernels at initialization are finite: 2qσ + q̃ + 1 = 0.
3. Tangent kernels and a limit model are of the same order at initialization: qσ + q̃ + 1/2 = 0.
4. Tangent kernels start to evolve: qσ + q̃ = 0.
We have also checked this Proposition numerically for limit models discussed below: see Figure 1, right. Each condition corresponds to a straight line in the (qσ, q̃)-plane: see Figure 1, left. These four lines divide the well-definiteness band into 13 regions: three are two-dimensional, seven are onedimensional, and three are zero-dimensional. In Appendix C we show that each region corresponds to a single distinct limit model evolution; we also list corresponding evolution equations. Note that a segment (a one-dimensional region) that corresponds to the Condition 2-2 exactly coincides with a family of "intermediate scalings" introduced in Golikov (2020).

heading: 3 CAPTURING THE BEHAVIOR OF FINITE-WIDTH NETS
text: A possible use-case for a limit model is being a proxy for a given finite-width net, useful for theoretical considerations. For example, a number of theoretical properties, including convergence to a global minimum and generalization, are already proven for nets near the NTK limit: see Arora et al. (2019b).
Note that a typical finite-width model satisfies all four statements of Condition 2 (if we exclude the word "limit" from them). Indeed, neural nets are typically initialized with He initialization (He et al., 2015) that guarantees finite f (0)d 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: Arora et al. (2019a) have shown that freezing NTK of practical convolutional nets sufficiently reduces their generalization ability; Woodworth et al. (2019) also noticed that evolution of NTK is sufficient for good performance.
Consequently, if we want a limit model to capture the dynamics of a finite-width net, we have to satisfy all four statements of Condition 2. However, as one can see from Figure 1, we cannot satisfy all of them simultaneously. We say that one limit model captures the behavior of a finite-width one better than the other, if all statements of Conditions 2 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 "non-dominated" limit models which we discuss in the upcoming section. After that, we introduce a model modification that allows for a limit satisfying all four statements.

heading: 3.1 "NON-DOMINATED" LIMIT MODELS: MF, NTK AND "SYM-DEFAULT"
text: Obviously, the three "non-dominated" limit models are exactly three zero-dimensional regions (points) in Figure 1, left. First suppose statements 1, 2 and 3 hold, hence tangent kernels are constant throughout training (see Figure 1, right). A corresponding point qσ = −1/2, q̃ = 0 reads as σ ∝ d−1/2 and η̂ = const, which is the case considered in the seminal paper on NTK (Jacot et al., 2018). The limit dynamics is then given as (see App. C.1.1 and App. C for the general derivation, and eqs. (71-75) for a complete system of evolution equations):
f (k+1) ntk,∞ (x) = f (k) ntk,∞(x)−η̂ ∗ a∇ (k) fntk `(x(k)a , y (k) a ) K (0) a,∞(x,x (k) a )−η̂∗w∇ (k) fntk `(x(k)w , y (k) w ) K (0) w,∞(x,x (k) w ),
f (0) ntk,∞(x) ∼ N (0, σ ∗,2σ(0),2(x)), (13)
where (x(k)a∨w, y (k) a∨w) ∼ D and limit tangent kernels K (0) a∨w,∞ and standard deviations at the initialization σ(0)(x) can be calculated along the same lines as in Lee et al. (2019).
Next, suppose statements 2 and 4 hold. In this case K(k)∞ does not coincide with K (0) ∞ (see Figure 1, right), hence the dynamics analogous to (13) is not closed. However, the limit dynamics can be expressed as an evolution of a weight-space measure (see Rotskoff & Vanden-Eijnden (2019); Chizat & Bach (2018) for a similar dynamics for the gradient flow, App. C.2.1 and App. C for the general derivation, and eqs. (94-96) for a complete system of evolution equations):
µ(k+1)∞ = µ (k) ∞ + div(µ (k) ∞ ∆θ (k) mf ), µ (0) ∞ = N (0, I1+dx), (14)
f (k) mf,∞(x) = σ ∗ ∫ âφ(ŵTx)µ(k)∞ (dâ, dŵ), (15)
where the vector field ∆θ(k)mf is defined as follows:
∆θ (k) mf (â, ŵ) = −[∇ (k) fmf `(x(k)a , y (k) a )φ(ŵ Tx(k)a ),∇ (k) fmf `(x(k)w , y (k) w )âφ ′(ŵTx(k)w )x (k),T w ] T , (16)
where we write "[u,v]" meaning a concatenation of two row vectors u and v. Here we have qσ = −1, q̃ = 1, hence σ ∝ d−1 and η̂ ∝ d; this hyperparameter scaling were used in Rotskoff & VandenEijnden (2019); Chizat & Bach (2018). Note that since a measure at the initialization µ(0)∞ has a zero mean, a limit model vanishes at the initialization f (0)mf,∞ = 0 (see Figure 1, right) thus violating statements 1 and 3 of Condition 2.
Finally, consider a point for which statements 1 and 4 hold: qσ = −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: σa ∝ d−1/2 and σw ∝ d−1/2x . Assume learning rates (in original parameterization) are not modified with width: ηa = const and ηw = const. This implies η̂a ∝ d and η̂w ∝ 1, or q̃a = 1 and q̃w = 0. We refer the scaling qσ = −1/2, q̃a = 1 and q̃w = 0 as "default", and the scaling qσ = −1/2, q̃ = 1/2 as "sym-default". A limit model evolution for the sym-default scaling looks as follows (see App. C.2.2 for an equivalent formulation and App. C for the general derivation, and eqs. (97-101) for a complete system of evolution equations):
µ(k+1)∞ = µ (k) ∞ + div(µ (k) ∞ ∆θ (k) sym-def), µ (0) ∞ = N (0, I1+dx), (17)
f (0) sym-def,∞(x) ∼ N (0, σ ∗,2σ(0),2(x)), z (k) sym-def,∞(x) = [∫ âφ(ŵTx)µ(k)∞ (dâ, dŵ) > 0 ] , (18)
where the vector field ∆θ(k)sym-def is defined similarly to the MF case (16):
∆θ (k) sym-def(â, ŵ) = −[∇ (k) fsym-def `(x(k)a , y (k) a )φ(ŵ Tx(k)a ),∇ (k) fsym-def `(x(k)w , y (k) w )âφ ′(ŵTx(k)w )x (k),T w ] T ,
∇(k)fsym-def`(x, y) = −y[yz (k) sym-def,∞(x) < 0] for k ≥ 1. (19)
As we show in Appendix D, the default scaling leads to an almost similar limit dynamics as the sym-default scaling: eqs. (114-117) for a complete system of the corresponding evolution equations. The quantity z(k)sym-def,∞ should be perceived as a sign of f (k) sym-def,∞ =
σ∗ limd→∞ ( dqσ+1 ∫ âφ(ŵTx)µ (k) d (dâ, dŵ) ) . The reason why we have to switch from logits
to their signs is that the limit model diverges for k ≥ 1: limd→∞ f (k)d (x) = ∞. Nevertheless the gradient of the cross-entropy loss is well-defined even for infinite logits: it just degenerates into the gradient of a hinge-type loss: limf→+∞×z ∂`(y,f) ∂f = −y[yz < 0]. For this reason, we redefine the loss gradient for k ≥ 1 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 µ(k)d stays well-defined.

heading: 3.2 INITIALIZATION-CORRECTED MEAN-FIELD (IC-MF) LIMIT
text: Here we propose a dynamics that satisfy all four statements of Condition 2. We then show how to modify the network training for the finite width in order to ensure that in the limit of the infinite width its training dynamics converge to the proposed limit one. Consider the following:
µ(k+1)∞ = µ (k) ∞ + div(µ (k) ∞ ∆θ (k) icmf), µ (0) ∞ = N (0, I1+dx), (20)
f (k) icmf,∞(x) = σ ∗ ∫ âφ(ŵTx)µ(k)∞ (dâ, dŵ) + f (0) ntk,∞(x), (21)
where f (0)ntk,∞ is defined similarly to above:
f (0) ntk,∞(x) ∼ N (0, σ ∗,2σ(0),2(x)), (22)
the vector field ∆θ(k)icmf is defined analogously to the mean-field case:
∆θ (k) icmf(â, ŵ) = −[∇ (k) ficmf `(x(k)a , y (k) a )φ(ŵ Tx(k)a ),∇ (k) ficmf `(x(k)w , y (k) w )âφ ′(ŵTx(k)w )x (k),T w ] T , (23)
See App. E and eqs. (131-134) for a complete system of evolution equations. The only difference between this dynamics and the mean-field dynamics is a bias term f (0)ntk,∞ in the definition of logits. This bias term does not depend on k and stays finite for large d in contrast to f (0)mf,∞ which vanishes for large d; it ensures Condition 2-1 to hold. As for Condition 2-4, tangent kernels evolve with k simply because the measure µ(k)∞ evolves with k similarly to the mean-field case (see Figure 1, right). Indeed,
K(k)w,∞(x ′,x) = σ∗,2d∗ ∫ |â(k)|2φ′(ŵ(k),Tx)φ′(ŵ(k),Tx′)µ(k)∞ (dâ, dŵ), (24)
and the limit of K(k)a,d is written in a similar way. Kernels at initialization K (0) a∨w,∞ are finite due to the Law of Large Numbers (Condition 2-2); this, and the finiteness of f (0)ntk ensures Condition 2-3.
As we show in Appendix E the dynamics (20) is a limit for the GD dynamics of the following model with learning rates η̂a∨w = η̂∗a∨w(d/d ∗)1:
ficmf,d(x; â, Ŵ ) = σ ∗(d/d∗)−1 d∑ r=1 ârφ(ŵ T r x) +σ ∗((d/d∗)−1/2− (d/d∗)−1) d∑ r=1 â(0)r φ(ŵ (0),T r x). (25) Note that ficmf,d∗(x) = σ∗ ∑d∗ r=1 ârφ(ŵ T r x): we have not altered the model definition at d = d ∗.

heading: 3.3 EXPERIMENTS
text: Consider a network of width d∗ initialized with a standard deviation σ∗ and trained with learning rates η̂∗a∨w. We call this model a "reference". Consider a family of models indexed by a 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-entropy loss on the CIFAR2 dataset (a subset of first two classes of CIFAR10). We track the divergence of a limit network from the reference one using the following quantity: E x∼DtestDlogits(f (k) ∞ (x) || f (k)d∗ (x)), where
Dlogits(ξ || ξ∗) = KL(N (E ξ,Var ξ) || N (E ξ∗,Var ξ∗)). (26)
We have also tried other divergence measures; see Appendix I.
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 et al. (2019). At the same time, the mean-field limit starts with a 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 a 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 an 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.

heading: 4 RELATED WORK
text: A pioneering work of Jacot et al. (2018) have shown that a 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 Jacot 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 of the model is determined by this constant kernel; see eq. (13). The training regime when NTK is hardly varying is coined as "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 a finite width if one scales the output of the network appropriately. While being theoretically appealing, "laziness" assumption turns out to have a number of limitations in explaining the success of deep learning (Arora et al., 2019a; Ghorbani et al., 2019).
Another line of works considers the evolution of weights as an evolution of a weight-space measure, similar to eq. (14) (Mei et al., 2018; 2019; Sirignano & Spiliopoulos, 2020; Chizat & Bach, 2018; Rotskoff & Vanden-Eijnden, 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-field". 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 is similar in spirit to the one used 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, we present 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, as 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.

heading: 5 CONCLUSIONS
text: The current work follows a direction started in Golikov (2020): we study how one should scale hyperparameters of a neural network with a single hidden layer in order to converge to a "dynamically
stable" limit training dynamics. A weaker dynamical stability condition leads us to a richer class of possible limit models as compared to Golikov (2020). In particular, the class of limit models we consider includes a "default" limit model that corresponds to a network with infinitely large number of nodes and finite learning rates in the original parameterization. This "default" limit model does not satisfy a "well-definiteness" 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 a single limit model. All of these conditions are satisfied by finite-width networks, but cannot be satisfied by limit models all simultaneously. We propose a modification of a finite-width model; the limit of this modification corresponds to a limit model that satisfy all of the conditions mentioned above and tracks the dynamics of a "reference" finite-width net better than other limit models.

heading: A FORMAL CONDITIONS FOR SECTION 2
text: Here we present formal definitions for notions that appear in Section 2; they are required for mathematical rigor. First, recall the definition of tangent kernels:
K (k) a,d(x,x ′) = (d/d∗)q̃aσ2 d∑ r=1 φ(ŵ(k),Tr x)φ(ŵ (k),T r x ′), (27)
K (k) w,d(x,x ′) = (d/d∗)q̃wσ2 d∑ r=1 |â(k)r |2φ′(ŵ(k),Tr x)φ′(ŵ(k),Tr x′)xTx′. (28)
The kernels are used to express a model increment:
∆f (k) d (x) = f (k+1) d (x)− f (k) d (x) = d∑ r=1 ∂fd(x) ∂θ̂r ∣∣∣∣ θ̂r=θ̂ (k) r ∆θ̂(k)r +Oη̂∗a∨w→0(η̂ ∗ aη̂ ∗ w + η̂ ∗,2 w ) =
= −η̂∗a∇ (k) fd `(x(k)a , y (k) a ) K (k) a,d(x,x (k) a )− η̂∗w∇ (k) fd `(x(k)w , y (k) w ) K (k) w,d(x,x (k) w )+O(η̂ ∗ aη̂ ∗ w+ η̂ ∗,2 w ),
(29)
Define the linear part of the model increment with respect to learning rate proportionality factors:
∆f (k),′ d,a∨w(x) =
∂∆f (k) d (x)
∂η̂∗a∨w ∣∣∣∣∣ η̂∗a=0 η̂∗w=0 = −∇(k)fd `(x (k) a∨w, y (k) a∨w) K (k) a∨w,d(x,x (k) a∨w). (30)
We use this quantity to rewrite the model increment:
∆f (k) d (x) = η̂ ∗ a∆f (k),′ d,a (x) + η̂ ∗ w∆f (k),′ d,w (x) +O(η̂ ∗ aη̂ ∗ w + η̂ ∗,2 w ). (31)
Let us consider kernel definitions (27) and (28) again. Their increments are given by:
∆K (k) a,d(x,x ′) = −η̂∗w(d/d∗)2q̃σ3 d∑ r=1 ( φ(ŵ(k),Tr x)φ ′(ŵ(k),Tr x ′) + φ′(ŵ(k),Tr x)φ(ŵ (k),T r x ′) ) ×
×∇(k)fd `(x (k) w , y (k) w )â (k) r φ ′(ŵ(k),Tr x (k) w )(x + x ′)Tx(k)w +Oη̂∗w→0 d→∞ (η̂∗,2w d 3q̃+4qσ+1), (32)
∆K (k) w,d(x,x ′) = −η̂∗w(d/d∗)2q̃σ3 d∑ r=1 |â(k)r |2 ( φ′(ŵ(k),Tr x)φ ′′(ŵ(k),Tr x ′)+
+ φ′′(ŵ(k),Tr x)φ ′(ŵ(k),Tr x
′) ) xTx′×
×∇(k)fd `(x (k) w , y (k) w )â (k) r φ ′(ŵ(k),Tr x (k) w )(x + x ′)Tx(k)w +Oη̂∗w→0 d→∞ (η̂∗,2w d 3q̃+4qσ+1)−
− η̂∗a(d/d∗)2q̃σ3 d∑ r=1 2â(k)r φ ′(ŵ(k),Tr x)φ ′(ŵ(k),Tr x ′)×
×∇(k)f `(x (k) a , y (k) a )φ(ŵ (k),T r x (k) a ) +Oη̂∗a→0 d→∞ (η̂∗,2a d 3q̃+4qσ+1). (33)
Similarly to what was done for model increments, we define linear parts of the kernel increments with respect to learning rate proportionality factors:
∆K (k),′ aw,d(x,x
′) = ∂∆K
(k) a,d(x,x ′) ∂η̂∗w ∣∣∣∣∣ η̂∗w=0 =
= −(d/d∗)2q̃σ3 d∑ r=1 ( φ(ŵ(k),Tr x)φ ′(ŵ(k),Tr x ′) + φ′(ŵ(k),Tr x)φ(ŵ (k),T r x ′) ) ×
×∇(k)fd `(x (k) w , y (k) w )â (k) r φ ′(ŵ(k),Tr x (k) w )(x + x ′)Tx(k)w , (34)
∆K (k),′ ww,d(x,x
′) = ∂∆K
(k) w,d(x,x ′) ∂η̂∗w ∣∣∣∣∣ η̂∗w=0 =
= −(d/d∗)2q̃σ3 d∑ r=1 |â(k)r |2 ( φ′(ŵ(k),Tr x)φ ′′(ŵ(k),Tr x ′) + φ′′(ŵ(k),Tr x)φ ′(ŵ(k),Tr x ′) ) xTx′×
×∇(k)fd `(x (k) w , y (k) w )â (k) r φ ′(ŵ(k),Tr x (k) w )(x + x ′)Tx(k)w , (35)
∆K (k),′ wa,d(x,x
′) = ∂∆K
(k) w,d(x,x ′) ∂η̂∗a ∣∣∣∣∣ η̂∗a=0 =
= −(d/d∗)2q̃σ3 d∑ r=1 2â(k)r φ ′(ŵ(k),Tr x)φ ′(ŵ(k),Tr x ′)∇(k)fd `(x (k) a , y (k) a )φ(ŵ (k),T r x (k) a ). (36)
Note that ∆K(k),′aa,d (x,x ′) = 0 since âr-terms are absent in the definition of Ka,d, eq. (27).
Define p(k)err,d = P(y,x,y(:k−1)a ,x(:k−1)a ,y(:k−1)w ,x(:k−1)w )∼D2k−1{yf (k) d (x) < 0} — the probability of giving a wrong answer on the step k. Let kterm,d ∈ N ∪ {+∞} be a maximal k such that ∀k′ < k p
(k′) err,d > 0. Generally, kterm,d depends on hyperparameters, as well as on the data distribution D.
Scaling exponents (qσ, q̃a, q̃w) together with proportionality factors (d∗, σ∗, η̂∗a, η̂ ∗ w) define a limit model f (k)∞ (x) = limd→∞ f (k) d (x). We call a model "dynamically stable in the limit of large width" if it satisfies the following condition:
Condition 3. ∃kbalance ∈ N : ∀k ∈ [kbalance, kterm,∞) ∩ N y(k)a f (k)∞ (x(k)a ) < 0 and y (k) w f (k) ∞ (x (k) w ) < 0 imply ∆f (k),′ d,a∨w(x) = Θd→∞(f (kbalance) d (x)) x-a.e. (y (:k) a∨w,x (:k) a∨w)-a.s.
This condition puts a constraint on exponents (qσ, q̃a, q̃w); this constraint generally depends on the train data distribution D and on proportionality factors d∗, σ∗, and η̂∗a∨w. In order to obtain a data-independent hyperparameter-independent constraint, we need the condition above to hold for any value of kterm,∞ and any values of d∗, σ∗, and η̂∗a∨w. Without loss of generality we can assume kterm,∞ to be infinite, which gives the following condition:
Condition 4 (a formal version of Condition 1). Given kterm,∞ = +∞, ∃kbalance ∈ N : ∀σ∗ > 0 ∀η̂∗a∨w > 0 ∀k ≥ kbalance y (k) a f (k) ∞ (x (k) a ) < 0 and y (k) w f (k) ∞ (x (k) w ) < 0 imply ∆f (k),′ d,a∨w(x) = Θd→∞(f (kbalance) d (x)) x-a.e. (y (:k) a∨w,x (:k) a∨w)-a.s. Condition 5 (a formal version of Condition 2). Following conditions separate the band of dynamical stability (Figure 1, left):
1. A limit model at initialization is finite: f (0)d (x) = Θd→∞(1) x-a.e.
2. Tangent kernels at initialization are finite: K(0)d,a∨w(x,x ′) = Θd→∞(1) (x,x ′)-a.e.
3. Tangent kernels and a limit model are of the same order at initialization: K(0)d,a∨w(x,x ′) =
Θd→∞(f (0) d (x)) (x,x ′)-a.e.
4. Tangent kernels start to evolve: ∆K(0),′d,wa∨w(x,x ′) = Θd→∞(K (0) d,w(x,x ′)) (x,x′)-a.e. and
∆K (0),′ d,aw(x,x ′) = Θd→∞(K (0) d,a(x,x ′)) (x,x′)-a.e.

heading: B PROOFS OF PROPOSITIONS
text: We restate all necessary definitions here. We assume the non-linearity φ to be real analytic and asymptotically linear: φ(z) = Θz→∞(z). We assume the loss function `(y, z) to be the standard binary cross-entropy loss: `(y, z) = ln(1 + e−yz), where labels y ∈ {−1, 1}. The training dynamics is given as:
∆â(k)r = −η̂aσ∇ (k) fd `(x(k)a , y (k) a ) φ(ŵ (k),T r x (k) a ), â (0) r ∼ N (0, 1), (37)
∆ŵ(k)r = −η̂wσ∇ (k) fd `(x(k)w , y (k) w ) â (k) r φ ′(ŵ(k),Tr x (k) w )x (k) w , ŵ (0) r ∼ N (0, I) ∀r ∈ [d], (38)
∇(k)fd `(x, y) = ∂`(y, z)
∂z ∣∣∣∣ z=f
(k) d (x)
= −y
1 + exp(f (k) d (x)y)
, f (k) d (x) = σ d∑ r=1 â(k)r φ(ŵ (k),T r x),
where (x(k)a∨w, y (k) a∨w) ∼ D for D being the data distribution.
We assume hyperparameters to be scaled with width as power-laws:
σ(d) = σ∗ × (d/d∗)qσ , η̂a(d) = η̂∗a × (d/d∗)q̃a , η̂w(d) = η̂∗w × (d/d∗)q̃w .
B.1 PROOF OF PROPOSITION 1
Define:
q (k) θ = inf{q : θ (k) = Od→∞(d q)}, q(k)∆θ = inf{q : ∆θ (k) = Od→∞(d q)}, (39)
where θ should be substituted with a or w. We define inf(∅) = +∞. We introduce similar definitions for other quantities:
q (k) f (x) = inf{q : f (k) d (x) = Od→∞(d q)}, q(k)∇` (x, y) = inf{q : ∇ (k) fd `(x, y) = Od→∞(d q)}, (40)
q (k) ∆f (x) = inf{q : ∆f (k) d (x) = Od→∞(d q)}, q(k)∆f ′a∨w(x) = inf{q : ∆f (k),′ d,a∨w(x) = Od→∞(d q)}. (41)
Lemma 1. Assume D is a continuous distribution. Then following hold:
1. ∀k ≥ 0 ∀x, y q(k)∇` (x, y) ≤ 0, while [yf (k) ∞ (x) < 0] implies q (k) ∇` (x, y) = 0.
2. q(0)a∨w = 0, q (0) f (x) = qσ + 1 2 x-a.e.
3. ∀k ≥ 0 q(k)∆a/∆w = q̃a∨w + qσ + q (k) w/a + q (k) ∇` (x (k), y(k)) (x (k) a∨w, y (k) a∨w)-a.s.
4. ∀k ≥ 0 q(k)∆f ′a∨w(x) = 2qσ+1+ q̃a∨w+2q (k) w/a+q (k) ∇` (x (k) a∨w, y (k) a∨w) x-a.e. (x (k) a∨w, y (k) a∨w)-a.s.
5. ∀k ≥ 0 qσ + q̃w + q(k)a ≤ 0 implies that for sufficiently small η̂∗a and η̂∗w q (k) ∆f (x) =
max(q (k) ∆f ′a (x), q (k) ∆f ′w (x)) x-a.e. (x(k)a , y (k) a ,x (k) w , y (k) w )-a.s.
6. ∀k ≥ 0 q(k+1)a∨w = max(q (k) a∨w, q (k) ∆a/∆w) (x (k), y(k))-a.s., q(k+1)f (x) =
max(q (k) f (x), q (k) ∆f (x)) x-a.e. (x (k) a , y (k) a ,x (k) w , y (k) w )-a.s.
Proof. (1) follows from the fact that ∂`(y, z)/∂z is bounded ∀y, while |∂`(y, z)/∂z| ∈ [1/2, 1] when yz < 0.
â (0) r ∼ N (0, 1) which is not zero and does not depend on d, hence q(0)a = 0; similar holds for w. For x 6= 0 we have f (0)d (x) = σ ∑d r=1 â (0) r φ(ŵ (0),T r x) = Θd→∞(d
1/2+qσ ) due to the Central Limit Theorem. Hence (2) holds.
Since D is a.c. wrt Lebesgue measure on R1+dx , and φ is real analytic and non-zero, φ(ŵ (k),T r x (k) a∨w) 6= 0 and φ′(ŵ (k),T r x (k) a∨w) is well-defined (x (k) a∨w, y (k) a∨w)-a.s. This implies that q (k) ∆a/∆w = q̃a∨w + qσ + q (k) w/a + q (k) ∇` (x (k) a∨w, y (k) a∨w) (x (k) a∨w, y (k) a∨w)-a.s., which is exactly (3).
Consider ∆f (k),′d,a :
∆f (k),′ d,a (x) = −∇ (k) f `(x (k) a , y (k) a ) K (k) a,d(x,x (k) a ) =
= −∇(k)f `(x (k) a , y (k) a ) (d/d ∗)q̃aσ2 d∑ r=1 φ(ŵ(k),Tr x)φ(ŵ (k),T r x (k) a ). (42)
For the same reason as discussed above φ(ŵ(k),Tr x (k) a ) 6= 0 (x(k)a , y(k)a )-a.s., and φ(ŵ(k),Tr x) 6= 0 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 dk times the k-th moment of each summand. Hence the sum itself scales as d. Since φ is asymptotically linear, each φ-term scales as dq (k) w . Collecting all terms together, we obtain q (k) ∆f ′a (x) = 2qσ + 1 + q̃a + 2q (k) w + q (k) ∇` (x (k) a , y (k) a ) x-a.e. (x (k) a , y (k) a )-a.s. Following the same steps for ∆f (k),′w , we get (4).
Let us overview ∆f (k)d (x) in detail:
∆f (k) d (x) = d∑ r=1 ( ∞∑ j=1 1 j!
∂jfd(x)
∂ŵi1r . . . ∂ŵ ij r ∣∣∣∣ŵr=ŵ(k)r âr=â (k) r ∆ŵ(k),i1r . . .∆ŵ (k),ij r +
+ ∞∑ j=1 1 j!
∂jfd(x)
∂âr∂ŵ i2 r . . . ∂ŵ ij r ∣∣∣∣ŵr=ŵ(k)r âr=â (k) r ∆âr∆ŵ (k),i2 r . . .∆ŵ (k),ij r ) =
= d∑ r=1 ( ∞∑ j=1 1 j! (−1)j η̂jwσj+1(∇ (k) fd `(x(k)w , y (k) w )) j×
× (â(k)r )j+1(φ′(ŵ(k),Tr x(k)w ))jφ(j)(ŵ(k),Tr x)(x(k),Tw x)j+
+ ∞∑ j=1 1 j! (−1)j η̂aη̂j−1w σj+1∇ (k) fd `(x(k)a , y (k) a )(∇ (k) fd `(x(k)w , y (k) w )) j−1×
× (â(k)r )j−1φ(ŵ(k),Tr x(k)a )(φ′(ŵ(k),Tr x(k)w ))j−1φ(j−1)(ŵ(k),Tr x)(x(k),Tw x)j−1 ) . (43)
Assumption qσ + q̃w + q (k) a ≤ 0 implies η̂jwσj+1(â (k) r )j+1 = Od→∞(η̂wσ 2(â (k) r )2) and η̂aη̂ j−1 w σ j+1(â (k) r )j−1 = Od→∞(η̂aσ 2).
Since q(k)∇` (x, y) ≤ 0 ∀x, y due to (1), (∇ (k) fd `(x (k) w , y (k) w ))j = Od→∞(∇(k)fd `(x (k) w , y (k) w )) and ∇(k)fd `(x (k) a , y (k) a )(∇(k)fd `(x (k) w , y (k) w ))j−1 = Od→∞(∇(k)fd `(x (k) a , y (k) a )).
Since φ(z) = Θz→∞(z), φ′(ŵ (k),T r x (k) w ) = Od→∞(1) and (φ′(ŵ (k),T r x (k) w ))j = Od→∞(φ ′(ŵ (k),T r x (k) w )) for j ≥ 1.
Hence for small enough η̂∗a and η̂ ∗ w the first term of each sum which corresponds to j = 1 dominates all others, even in the limit of infinite d:
∆f (k) d (x) = − d∑ r=1
( η̂wσ
2∇(k)fd `(x (k) w , y (k) w )(â (k) r ) 2φ′(ŵ(k),Tr x (k) w )φ ′(ŵ(k),Tr x)x (k),T w x+
+ η̂aσ 2∇(k)fd `(x (k) a , y (k) a )φ(ŵ (k),T r x (k) a )φ(ŵ (k),T r x)+
+ oη̂∗a∨w→0
( Od→∞ (( η̂a∇(k)fd `(x (k) a , y (k) a ) + η̂w∇ (k) fd `(x(k)w , y (k) w )(â (k) r ) 2 ) σ2 ))) =
= η̂∗w∆f (k),′ d,w (x) + η̂ ∗ a∆f (k),′ d,a (x)+
+ oη̂∗a∨w→0
( Od→∞ (( η̂a∇(k)fd `(x (k) a , y (k) a ) + η̂w∇ (k) fd `(x(k)w , y (k) w )(â (k) r ) 2 ) σ2d )) . (44)
Note that two summands depend on (x(k)w , y (k) w ) and (x (k) a , y (k) a ) respectively, which do not depend on each other. Hence q(k)∆f (x) = max(q (k) ∆f ′a (x), q (k) ∆f ′w (x)) x-a.e. (x(k)a∨w, y (k) a∨w)-a.s., which is (5). Note that the o-term does not alter the exponent. Indeed,( η̂a∇(k)fd `(x (k) a , y (k) a ) + η̂w∇ (k) fd `(x(k)w , y (k) w )(â (k) r ) 2 ) σ2d =
= Od→∞
( d1+2qσ+max(q̃a+q (k) ∇` (x (k) a ,y (k) a ),q̃w+q (k) ∇` (x (k) w ,y (k) w )+2q (k) a ) ) =
= Od→∞
( d1+2qσ+max(q̃a+q (k) ∇` (x (k) a ,y (k) a )+2q (k) w ,q̃w+q (k) ∇` (x (k) w ,y (k) w )+2q (k) a ) ) =
= Od→∞
( d max(q (k) ∆f′a (x),q (k) ∆f′w (x)) ) . (45)
One before the last equality holds, because q(k)w ≥ 0 due to (2) and (6), while the last equality holds due to (4).
By definition we have â(k+1)r = â (k) r + ∆â (k) r . Since the second term depends on (x (k) a , y (k) a ), while the first term does not, we get q(k+1)a = max(q (k) a , q (k) ∆a). Similar holds for ŵr and f (k) d (x) x-a.e., which gives (6).
Lemma 2. Assume D is a continuous distribution, kterm = +∞ and q̃a = q̃w = q̃. Then
1. If qσ + q̃ ≤ 0 then ∀k ≥ 0 q(k)a∨w = 0 (x (:k−1) a , y (:k−1) a ,x (:k−1) w , y (:k−1) w )-a.s.
2. If qσ + q̃ > 0 then ∀k ≥ 0 q(k)a∨w = k(qσ + q̃) with positive probability wrt (x (:k−1) a , y (:k−1) a ,x (:k−1) w , y (:k−1) w ).
Proof. Here and in subsequent proofs we will write "almost surely" meaning "almost surely wrt (x (:k) a , y (:k) a ,x (:k) w , y (:k) w )" for appropriate k; we apply a similar shortening for "with positive probability wrt (x(:k)a , y (:k) a ,x (:k) w , y (:k) w )".
If qσ + q̃ ≤ 0 then statements 1, 2, 3 and 6 of Lemma 1 imply ∀k ≥ 0 q(k)a∨w = 0 a.s.
Assume qσ + q̃ > 0. We will prove that ∀k ≥ 0 q(k)a∨w = max(0, k(qσ + q̃)) with positive probability by induction. Induction base is given by Lemma 1-2.
Combining the induction assumption and Lemma 1-3 we get q(k)∆a/∆w = (k + 1)(qσ + q̃) + q (k) ∇` (x (k), y(k)) with positive probability wrt (x(:k−1)a , y (:k−1) a ,x (:k−1) w , y (:k−1) w ) (x (k) a∨w, y (k) a∨w)-a.s.
Since kterm = +∞ > k, y(k)a∨wf (k) d (x (k) a∨w) < 0 with positive probability wrt (x (k) a , y (k) a ,x (k) w , y (k) w ), and Lemma 1-1 implies that q(k)∆a/∆w = (k + 1)(qσ + q̃) with positive probability wrt (x (:k) a , y (:k) a ,x (:k) w , y (:k) w ).
Finally, Lemma 1-6 concludes the proof of the induction step.
Lemma 3. Assume D is a continuous distribution, kterm,∞ = +∞, q̃a = q̃w = q̃ and qσ + q̃ ≤ 0. Then ∀k ≥ 0
1. y(k)a∨wf (k) ∞ (x (k) a∨w) < 0 implies q (k) ∆f ′a∨w (x) = 2qσ + 1 + q̃ x-a.e.
(x (:k−1) a , y (:k−1) a ,x (:k−1) w , y (:k−1) w )-a.s.
2. y(k)a f (k) ∞ (x (k) a ) < 0 and y (k) w f (k) ∞ (x (k) w ) < 0 imply q (k) ∆f (x) = 2qσ + 1 + q̃ x-a.e.
(x (:k−1) a , y (:k−1) a ,x (:k−1) w , y (:k−1) w )-a.s. for sufficiently small η̂∗a and η̂ ∗ w.
Proof. By Lemma 2 ∀k ≥ 0 q(k)a∨w = 0 a.s. Since y (k) a∨wf (k) ∞ (x (k) a∨w) < 0, q (k) ∇` = 0 due to Lemma 1-1. Given this, Lemma 1-4 implies ∀k ≥ 0 q(k)∆f ′a∨w(x) = 2qσ + 1 + q̃ x-a.e. a.s. Hence by virtue of Lemma 1-5 ∀k ≥ 0 q(k)∆f (x) = 2qσ + 1 + q̃ x-a.e. a.s. for sufficiently small η̂∗a and η̂∗w.
Proposition 3. Suppose q̃a = q̃w = q̃ and D is a continuous distribution. Then Condition 4 requires qσ + q̃ ∈ [−1/2, 0] to hold.
Proof. By Lemma 2 if qσ + q̃ > 0 then q (k) a∨w = k(qσ + q̃) with positive probability. At the same time by virtue of Lemma 1-1 kterm,∞ = +∞ implies q(k)∇` = 0 with positive probability. Given this, Lemma 1-4 implies q(k)∆f ′a∨w(x) = qσ + 1 + (2k + 1)(qσ + q̃) x-a.e. with positive probability. This means that the last quantity cannot be almost surely equal to q(kbalance)f (x) for any kbalance independent on k. Since ∆f (k),′ d,a∨w(x) = Θd→∞(f (kbalance) d (x)) requires q (k) ∆f ′a∨w (x) = q (kbalance) f (x), we conclude that Condition 4 cannot be satisfied if qσ + q̃ > 0.
Hence qσ + q̃ ≤ 0. Then by Lemma 3 ∀k ≥ 0 y(k)a f (k)∞ (x(k)a ) < 0 and y(k)w f (k)∞ (x(k)w ) < 0 imply q
(k) ∆f (x) = 2qσ + 1 + q̃ x-a.e. (x (:k−1) a , y (:k−1) a ,x (:k−1) w , y (:k−1) w )-a.s. for sufficiently small η̂∗a and η̂∗w. We will show that Condition 4 requires qσ + q̃ ∈ [−1/2, 0] to hold already for these sufficiently small η̂∗a and η̂ ∗ w.
Suppose y(k)a f (k) ∞ (x (k) a ) < 0 and y (k) w f (k) ∞ (x (k) w ) < 0. Given this, points 1 and 6 of Lemma 1 imply ∀kbalance ≥ 1 qkbalancef (x) = max(q (0) f (x), 2qσ + 1 + q̃) = max(qσ + 1 2 , 2qσ + 1 + q̃) x-a.e. a.s. Hence q(k)∆f ′a∨w(x) = q (kbalance) f (x) x-a.e. a.s. if and only if qσ + 1 2 ≤ 2qσ + 1 + q̃, which is qσ + q̃ ≥ −1/2; we can take kbalance = 1 without loss of generality. Having q(k)∆f ′a∨w(x) = q (kbalance) f (x) is necessary to have ∆f (k),′d,a∨w(x) = Θd→∞(f (kbalance) d (x)).
Summing all together, Condition 4 requires qσ + q̃ ∈ [−1/2, 0] to hold.
B.2 PROOF OF PROPOSITION 2
Proposition 4. Let Condition 4 holds; then
1. f (0)d (x) = Θd→∞(1) x-a.e. is equivalent to qσ + 1/2 = 0.
2. K(0)d,a∨w(x,x ′) = Θd→∞(1) (x,x ′)-a.e. is equivalent to 2qσ + q̃ + 1 = 0.
3. K(0)d,a∨w(x,x ′) = Θd→∞(f (0) d (x)) (x,x ′)-a.e. is equivalent to qσ + q̃ + 1/2 = 0.
4. ∆K(0),′d,wa∨w(x,x ′) = Θd→∞(K (0) d,w(x,x ′)) (x,x′)-a.e. and ∆K(0),′d,aw(x,x ′) =
Θd→∞(K (0) d,a(x,x ′)) (x,x′)-a.e. is equivalent to qσ + q̃ = 0.
Proof. Statement (1) directly follows from Lemma 1-2:
f (0) d (x) = σ d∑ r=1 â(0)r φ(ŵ (0),T r x) = Θd→∞(d qσ+1/2) (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 (0) a,d(x,x ′) = (d/d∗)q̃σ2 d∑ r=1 φ(ŵ(0),Tr x)φ(ŵ (0),T r x ′) = Θd→∞(d q̃+2qσ+1) (47)
(x,x′)-a.e.; the same logic holds for the other kernel: K(0)w,d(x,x ′) = Θd→∞(d q̃+2qσ+1) (x,x′)-a.e.
Combining derivations of the two previous statements, we get the statement (3). Now we proceed to the last statement. Consider again the kernel K(0)a,d; a linear part of this increment with respect to proportionality factors of learning rates is given by, see eq. (34):
∆K (0),′ aw,d(x,x
′) = ∂∆K
(0) a,d(x,x ′) ∂η̂∗w ∣∣∣∣∣ η̂∗w=0 =
= −(d/d∗)2q̃σ3 d∑ r=1 ( φ(ŵ(0),Tr x)φ ′(ŵ(0),Tr x ′) + φ′(ŵ(0),Tr x)φ(ŵ (0),T r x ′) ) ×
×∇(0)fd `(x (0) w , y (0) w )â (0) r φ ′(ŵ(0),Tr x (0) w )(x + x ′)Tx(0)w , (48)
Hence ∆K(0),′aw,d = Θd→∞(K (0) a,d) is equivalent to qσ + q̃ = 0. Considering the second kernel K (0) w,d and its increment is equivalent to the same condition.

heading: C THE NUMBER OF DISTINCT LIMIT MODELS IS FINITE
text: It is easy to see that due to the Proposition 4 Condition 5 divides the well-definiteness band into 13 regions. We now show that when proportionality factors σ∗ and η̂∗a∨w are fixed, choosing a limit model evolution is equivalent to picking a single region from these 13.
Indeed, for any width d a model evolution can be written as follows:
∆f (k) d (x) = −η̂ ∗ w∇ (k) fd `(x(k)w , y (k) w ) ( K (k) w,d(x,x (k) w )+
+Oη̂∗a∨w→0 d→∞
(η̂∗w∆K (k),′ ww,d(x,x (k) w ) + η̂ ∗ a∆K (k),′ wa,d(x,x (k) w ))
) −
− η̂∗a∇ (k) fd `(x(k)a , y (k) a ) ( K (k) a,d(x,x (k) a ) +Oη̂∗w→0
d→∞ (η̂∗w∆K (k),′ aw,d(x,x (k) a ))
) . (49)
f (k+1) d (x) = f (k) d (x) + ∆f (k) d (x), ∇ (k) fd `(x, y) = −y 1 + exp(f
(k) d (x)y)
, (50)
f (0) d (x) = σ ∗dqσ d∑ r=1 â(0)r φ(ŵ (0),T r x), (â (0) r , ŵ (0) r ) ∼ N (0, I1+dx). (51)
Now we introduce normalized kernels:
K̃ (k) a,d(x,x ′) = (d/d∗)−1−q̃−2qσK (k) a,d(x,x
′) = σ∗,2 d∗
d d∑ r=1 φ(ŵ(k),Tr x)φ(ŵ (k),T r x ′), (52)
K̃ (k) w,d(x,x ′) = (d/d∗)−1−q̃−2qσK (k) w,d(x,x
′) = σ∗,2 d∗
d d∑ r=1 |â(k)r |2φ′(ŵ(k),Tr x)φ′(ŵ(k),Tr x′)xTx′.
(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:
∆K̃ (k),′ ∗∗,d = (d/d ∗)−1−q̃−2qσ∆K (k),′ ∗∗,d , (54)
∆f̃ (k) d = (d/d ∗)−1−q̃−2qσ∆f (k) d , f̃ (k) d = (d/d ∗)−1−q̃−2qσf (k) d . (55)
We then rewrite the model evolution as:
∆f̃ (k) d (x) = −η̂ ∗ w∇ (k) fd `(x(k)w , y (k) w ) ( K̃ (k) w,d(x,x (k) w )+
+Oη̂∗a∨w→0 d→∞
(η̂∗w∆K̃ (k),′ ww,d(x,x (k) w ) + η̂ ∗ a∆K̃ (k),′ wa,d(x,x (k) w ))
) −
− η̂∗a∇ (k) fd `(x(k)a , y (k) a ) ( K̃ (k) a,d(x,x (k) a ) +Oη̂∗w→0
d→∞ (η̂∗w∆K̃ (k),′ aw,d(x,x (k) a ))
) . (56)
f̃ (k+1) d (x) = f̃ (k) d (x) + ∆f̃ (k) d (x) ∀k ≥ 0, (57)
f̃ (0) d (x) = σ ∗(d/d∗)−1−q̃−qσ d∑ r=1 â(0)r φ(ŵ (0),T r x), (â (0) r , ŵ (0) r ) ∼ N (0, I1+dx), (58)
f (k) d (x) = (d/d ∗)1+q̃+2qσ f̃ (k) d (x), ∇ (k) fd `(x, y) = −y 1 + exp(f
(k) d (x)y)
∀k ≥ 0. (59)
C.1 CONSTANT NORMALIZED KERNELS CASE
Kernels K̃(k)a∨w,d are either constants (hence ∆K̃ (k),′ ∗∗,d → 0 as d→∞) or evolve with k in the limit of large d. First assume they are constants; in this case qσ + q̃ < 0 due to Proposition 4-4, and
∆f̃ (k) d (x) = −η̂ ∗ w∇ (k) fd `(x(k)w , y (k) w ) ( K̃ (0) w,d(x,x (k) w ) + oη̂∗a∨w→0 d→∞ (1) ) −
− η̂∗a∇ (k) fd `(x(k)a , y (k) a ) ( K̃ (0) a,d(x,x (k) a ) + oη̂∗w→0 d→∞ (1) ) . (60)
Since normalized kernels K̃(0)a∨w,d converge to non-zero limit kernels K̃ (0) a∨w,∞, we can rewrite the formula above as:
∆f̃ (k) d (x) = −η̂ ∗ w∇ (k) fd `(x(k)w , y (k) w ) ( K̃(0)w,∞(x,x (k) w ) + od→∞(1) ) −
− η̂∗a∇ (k) fd `(x(k)a , y (k) a ) ( K̃(0)a,∞(x,x (k) a ) + od→∞(1) ) . (61)
f̃ (0) d (x) = σ ∗(d/d∗)−1/2−q̃−qσ (N (0, σ(0),2(x)) + od→∞(1)), (62) where σ(0)(x) can be calculated in the same manner as in Lee et al. (2019). As required by Proposition 3 1/2 + q̃ + qσ ≥ 0, hence f̃ (0)d (x) = Od→∞(1). This implies the following:
∇(0)f∞`(x, y) = limd→∞∇ (0) fd `(x, y) =
= lim d→∞ −y 1 + exp((d/d∗)1+q̃+2qσ f̃
(0) d (x)y)
=  −y[N (0, σ(0),2(x))y < 0] for 1/2 + qσ > 0;
−y 1+exp(σ∗d∗,1/2N (0,σ(0),2(x))y) for 1/2 + qσ = 0; −y/2 for 1/2 + qσ < 0.
(63)
On the other hand, ∆f̃ (0)d (x) = Θd→∞(1) with positive probability over (x (0) a∨w, y (0) a∨w). Hence f̃ (0) d = Od→∞(∆f̃ (0) d ) and f̃ (1) d = f̃ (0) d + ∆f̃ (0) d = Θd→∞(1). For the same reason, f̃ (k+1) d = f̃ (k) d + ∆f̃ (k) d = Θd→∞(1) ∀k ≥ 0.
This implies the following:
∀k ≥ 0 ∇(k+1)f∞ `(x, y) = limd→∞∇ (k+1) fd `(x, y) = lim d→∞ −y 1 + exp((d/d∗)1+q̃+2qσ f̃
(k+1) d (x)y)
=
=  −y[limd→∞ f̃ (k+1)d (x)y < 0] for 1 + q̃ + 2qσ > 0; −y 1+exp(limd→∞ f̃ (k+1) d (x)y) for 1 + q̃ + 2qσ = 0;
−y/2 for 1 + q̃ + 2qσ < 0. (64)
If we define f (k)∞ (x) = limd→∞ f (k) d (x), we get the following limit dynamics:
∆f̃ (k)∞ (x) = −η̂∗w∇ (k) f∞ `(x(k)w , y (k) w )K̃ (0) w,∞(x,x (k) w )− η̂∗a∇ (k) f∞ `(x(k)a , y (k) a )K̃ (0) a,∞(x,x (k) a ), (65)
K̃(0)a,∞(x,x ′) = σ∗,2d∗E ŵ∼N (0,Idx )φ(ŵ Tx)φ(ŵTx′), (66)
K̃(0)w,∞(x,x ′) = σ∗,2d∗E (â,ŵ)∼N (0,I1+dx )|â| 2φ′(ŵTx)φ′(ŵTx′)xTx′, (67)
f̃ (k+1)∞ (x) = f̃ (k) ∞ (x)+∆f̃ (k) ∞ (x), f̃ (0) ∞ (x) = { σ∗d∗,1/2N (0, σ(0),2(x)) for 1/2 + q̃ + qσ = 0; 0 for 1/2 + q̃ + qσ > 0;
(68)
∇(0)f∞`(x, y) =  −y[N (0, σ(0),2(x))y < 0] for 1/2 + qσ > 0;
−y 1+exp(σ∗d∗,1/2N (0,σ(0),2(x))y) for 1/2 + qσ = 0; −y/2 for 1/2 + qσ < 0;
(69)
∇(k+1)f∞ `(x, y) =  −y[f̃ (k+1)∞ (x)y < 0] for 1 + q̃ + 2qσ > 0; −y 1+exp(f̃ (k+1) ∞ (x)y) for 1 + q̃ + 2qσ = 0;
−y/2 for 1 + q̃ + 2qσ < 0; ∀k ≥ 0. (70)
This dynamics is defined by proportionality factors σ∗, d∗, η̂∗a∨w, and signs of three exponents: 1/2 + qσ, 1 + q̃ + 2qσ and 1/2 + q̃ + qσ. 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 1 (left) that given qσ + q̃ < 0, there are 8 distinct sign configurations.
Note also that since we are interested in binary classification problems, only the sign of logits matters. Since f (k)d = (d/d ∗)1+q̃+2qσ f̃ (k) d , signs of f (k) d and of f̃ (k) d are the same for all d. Hence ∀x, y limd→∞ sign(f (k) d (x)) = limd→∞ sign(f̃ (k) d (x)) = sign(f̃ (k) ∞ (x)).
C.1.1 NTK LIMIT MODEL
We state here a special case of the NTK scaling (qσ = −1/2, q̃ = 0, see Jacot et al. (2018)) explicitly. Since in this case 1 + q̃ + 2qσ = 0, we can omit tildas everywhere. This results in the following limit dynamics:
∆f (k)∞ (x) = −η̂∗w∇ (k) f∞ `(x(k)w , y (k) w )K (0) w,∞(x,x (k) w )− η̂∗a∇ (k) f∞ `(x(k)a , y (k) a )K (0) a,∞(x,x (k) a ), (71)
K(0)a,∞(x,x ′) = σ∗,2d∗E ŵ∼N (0,Idx )φ(ŵ Tx)φ(ŵTx′), (72)
K(0)w,∞(x,x ′) = σ∗,2d∗E (â,ŵ)∼N (0,I1+dx )|â| 2φ′(ŵTx)φ′(ŵTx′)xTx′, (73)
f (k+1)∞ (x) = f (k) ∞ (x) + ∆f (k) ∞ (x), f (0) ∞ (x) = σ ∗d∗,1/2N (0, σ(0),2(x)), (74)
∇(k)f∞`(x, y) = −y
1 + exp(f (k) ∞ (x)y)
∀k ≥ 0. (75)
C.2 NON-STATIONARY NORMALIZED KERNELS CASE
Suppose now qσ + q̃ = 0. In this case ∆K (0),′ d,wa∨w(x,x ′) = Θd→∞(K (0) d,w(x,x ′)) (x,x′)-a.e. and ∆K (0),′ d,aw(x,x ′) = Θd→∞(K (0) d,a(x,x
′)) (x,x′)-a.e. by virtue of the Proposition 4-4. Hence kernels evolve in the limit of large width (at least, for sufficiently small η∗a∨w).
If we follow the lines of the previous section, we will get a limit dynamics which is not closed:
∆f̃ (k)∞ (x) = −η̂∗w∇ (k) f∞ `(x(k)w , y (k) w ) ( K̃(k)w,∞(x,x (k) w )+
+Oη̂∗a∨w→0(η̂ ∗ w∆K̃ (k),′ ww,∞(x,x (k) w ) + η̂ ∗ a∆K̃ (k),′ wa,∞(x,x (k) w ))
) −
− η̂∗a∇ (k) f∞ `(x(k)a , y (k) a ) ( K̃(k)a,∞(x,x (k) a ) +Oη̂∗w→0(η̂ ∗ w∆K̃ (k),′ aw,∞(x,x (k) a )) ) , (76)
f̃ (k+1)∞ (x) = f̃ (k) ∞ (x) + ∆f̃ (k) ∞ (x), f̃ (0) ∞ (x) = 0, (77)
∇(0)f∞`(x, y) =  −y[N (0, σ(0),2(x))y < 0] for 1/2 + qσ > 0;
−y 1+exp(σ∗d∗,1/2N (0,σ(0),2(x))y) for 1/2 + qσ = 0; −y/2 for 1/2 + qσ < 0;
(78)
∇(k+1)f∞ `(x, y) =  −y[f̃ (k+1)∞ (x)y < 0] for 1 + qσ > 0; −y 1+exp(f̃ (k+1) ∞ (x)y) for 1 + qσ = 0;
−y/2 for 1 + qσ < 0; ∀k ≥ 0. (79)
The reason for this is non-stationarity of kernels. As a workaround we consider a measure in the weight space:
µ (k) d =
1
d d∑ r=1 δ â (k) r ⊗ δ ŵ (k) r . (80)
Recall the stochastic gradient descent dynamics:
∆â(k)r = −η̂∗aσ∗∇ (k) fd `(x(k)a , y (k) a ) φ(ŵ (k),T r x (k) a ), â (0) r ∼ N (0, 1), (81)
∆ŵ(k)r = −η̂∗wσ∗∇ (k) fd `(x(k)w , y (k) w ) â (k) r φ ′(ŵ(k),Tr x (k) w )x (k) w , ŵ (0) r ∼ N (0, Idx). (82)
Here we have replaced η̂a∨wσ with η̂∗a∨wσ ∗, because qσ + q̃ = 0. Similar to Rotskoff & VandenEijnden (2019); Chizat & Bach (2018), this dynamics can be expressed in terms of the measure defined above:
µ (k+1) d = µ (k) d + div(µ (k) d ∆θ (k) d ), µ (0) d =
1
d d∑ r=1 δ θ̂ (0) r , θ̂(0)r ∼ N (0, I1+dx) ∀r ∈ [d], (83)
∆θ (k) d (â, ŵ) =
= −[η̂∗aσ∗∇ (k) fd `(x(k)a , y (k) a )φ(ŵ Tx(k)a ), η̂ ∗ wσ ∗∇(k)fd `(x (k) w , y (k) w )âφ ′(ŵTx(k)w )x (k),T w ] T , (84)
f (k) d (x) = σ
∗(d/d∗)1+qσ ∫ âφ(ŵTx)µ
(k) d (dâ, dŵ), (85)
∇(k)fd `(x, y) = −y
1 + exp(f (k) d (x)y)
∀k ≥ 0. (86)
We rewrite the last equation in terms of f̃ (k)d (x) = (d/d ∗)−1−qσf (k) d (x):
f̃ (k) d (x) = σ
∗d∗ ∫ âφ(ŵTx)µ
(k) d (dâ, dŵ), (87)
∇(k)fd `(x, y) = −y
1 + exp((d/d∗)1+qσ f̃ (k) d (x)y)
∀k ≥ 0. (88)
This dynamics is closed. Taking the limit d→∞ yields:
µ(k+1)∞ = µ (k) ∞ + div(µ (k) ∞ ∆θ (k) ∞ ), µ (0) ∞ = N (0, I1+dx), (89)
∆θ(k)∞ (â, ŵ) =
= −[η̂∗aσ∗∇ (k) f∞ `(x(k)a , y (k) a )φ(ŵ Tx(k)a ), η̂ ∗ wσ ∗∇(k)f∞`(x (k) w , y (k) w )âφ ′(ŵTx(k)w )x (k),T w ] T , (90)
f̃ (k)∞ (x) = σ ∗d∗ ∫ âφ(ŵTx)µ(k)∞ (dâ, dŵ), (91)
∇(0)f∞`(x, y) =  −y[N (0, σ(0),2(x))y < 0] for 1/2 + qσ > 0;
−y 1+exp(σ∗d∗,1/2N (0,σ(0),2(x))y) for 1/2 + qσ = 0; −y/2 for 1/2 + qσ < 0;
(92)
∇(k+1)f∞ `(x, y) =  −y[f̃ (k+1)∞ (x)y < 0] for 1 + qσ > 0; −y 1+exp(f̃ (k+1) ∞ (x)y) for 1 + qσ = 0;
−y/2 for 1 + qσ < 0; ∀k ≥ 0. (93)
Since proportionality factors σ∗, d∗, and η̂∗a∨w are assumed to be fixed, choosing qσ is sufficient to define the dynamics. Signs of exponents 1/2 + qσ and 1 + qσ give 5 distinct limit dynamics. Together with 8 limit dynamics for constant normalized kernels case, this gives 13 distinct limit dynamics, each corresponding to a region in the band of a dynamical stability (Figure 1, left).
As was noted earlier, only the sign of logits matters, and our f̃ (k)d preserve the sign for any d: ∀x limd→∞ sign(f (k) d (x)) = limd→∞ sign(f̃ (k) d (x)) = sign(f̃ (k) ∞ (x)).
C.2.1 MF LIMIT MODEL
We state here a special case of the mean-field scaling (qσ = −1, q̃ = 1, see Rotskoff & VandenEijnden (2019) or Chizat & Bach (2018)) explicitly. Similar to NTK case, since 1 + qσ = 0 we can omit tildas. This results in the following limit dynamics:
µ(k+1)∞ = µ (k) ∞ + div(µ (k) ∞ ∆θ (k) ∞ ), µ (0) ∞ = N (0, I1+dx), (94)
∆θ(k)∞ (â, ŵ) =
= −[η̂∗aσ∗∇ (k) f∞ `(x(k)a , y (k) a )φ(ŵ Tx(k)a ), η̂ ∗ wσ ∗∇(k)f∞`(x (k) w , y (k) w )âφ ′(ŵTx(k)w )x (k),T w ] T , (95)
f (k)∞ (x) = σ ∗d∗ ∫ âφ(ŵTx)µ(k)∞ (dâ, dŵ), ∇ (k) f∞ `(x, y) =
−y 1 + exp(f (k) ∞ (x)y) ∀k ≥ 0. (96)
C.2.2 SYM-DEFAULT LIMIT MODEL
Another special case which deserves explicit formulation is what we have called a "sym-default" limit model. The corresponding scaling is: qσ = −1/2, q̃ = 1/2. The resulting limit dynamics is the following:
µ(k+1)∞ = µ (k) ∞ + div(µ (k) ∞ ∆θ (k) ∞ ), µ (0) ∞ = N (0, I1+dx), (97)
∆θ(k)∞ (â, ŵ) =
= −[η̂∗aσ∗∇ (k) f∞ `(x(k)a , y (k) a )φ(ŵ Tx(k)a ), η̂ ∗ wσ ∗∇(k)f∞`(x (k) w , y (k) w )âφ ′(ŵTx(k)w )x (k),T w ] T , (98)
f̃ (k)∞ (x) = σ ∗d∗ ∫ âφ(ŵTx)µ(k)∞ (dâ, dŵ), (99)
∇(0)f∞`(x, y) = −y
1 + exp(σ∗d∗,1/2N (0, σ(0),2(x))y) , (100)
∇(k+1)f∞ `(x, y) = −y[f̃ (k+1) ∞ (x)y < 0] ∀k ≥ 0. (101)

heading: D DEFAULT SCALING
text: Consider the special case of the default scaling: qσ = −1/2, q̃a = 1, q̃w = 0. Then corresponding dynamics can be written as follows:
∆â(k)r = −η̂∗aσ∗(d/d∗)1/2∇ (k) fd `(x(k)a , y (k) a ) φ(ŵ (k),T r x (k) a ), â (0) r ∼ N (0, 1), (102)
∆ŵ(k)r = −η̂∗wσ∗(d/d∗)−1/2∇ (k) fd `(x(k)w , y (k) w ) â (k) r φ ′(ŵ(k),Tr x (k) w )x (k) w , ŵ (0) r ∼ N (0, Idx),
(103)
f (k) d (x) = σ ∗(d/d∗)−1/2 d∑ r=1 â(k)r φ(ŵ (k),T r x), ∇ (k) fd `(x, y) = −y 1 + exp(f (k) d (x)y) ∀k ≥ 0.
(104)
As one can see, increments of output layer weights ∆â(k)r diverge with d. We introduce their normalized versions: ∆ã(k)r = (d/d∗)−1/2∆â (k) r . Similarly, we normalize output layer weights themselves: ã(k)r = (d/d∗)−1/2â (k) r . Then the dynamics transforms to:
∆ã(k)r = −η̂∗aσ∗∇ (k) fd `(x(k)a , y (k) a ) φ(ŵ (k),T r x (k) a ), ã (0) r ∼ N (0, (d/d∗)−1), (105)
∆ŵ(k)r = −η̂∗wσ∗∇ (k) fd `(x(k)w , y (k) w ) ã (k) r φ ′(ŵ(k),Tr x (k) w )x (k) w , ŵ (0) r ∼ N (0, Idx), (106)
f (k) d (x) = σ ∗ d∑ r=1 ã(k)r φ(ŵ (k),T r x), ∇ (k) fd `(x, y) = −y 1 + exp(f (k) d (x)y) ∀k ≥ 0. (107)
Similar to Appendix C.2, we have to introduce a weight-space measure in order to take a limit of d→∞:
µ (k) d =
1
d d∑ r=1 δ ã (k) r ⊗ δ ŵ (k) r . (108)
In terms of the measure the dynamics is expressed then as follows:
µ (k+1) d = µ (k) d + div(µ (k) d ∆θ (k) d ), (109)
µ (0) d =
1
d d∑ r=1 δ ã (0) r ⊗ δ ŵ (0) r , ã(0)r ∼ N (0, (d/d∗)−1), ŵ(0)r ∼ N (0, Idx) ∀r ∈ [d], (110)
∆θ (k) d (ã, ŵ) =
= −[η̂∗aσ∗∇ (k) fd `(x(k)a , y (k) a )φ(ŵ Tx(k)a ), η̂ ∗ wσ ∗∇(k)fd `(x (k) w , y (k) w )ãφ ′(ŵTx(k)w )x (k),T w ] T , (111)
f (k) d (x) = σ ∗d
∫ ãφ(ŵTx)µ
(k) d (dã, dŵ), ∇ (k) fd `(x, y) = −y 1 + exp(f
(k) d (x)y)
∀k ≥ 0. (112)
We rewrite the last equation in terms of f̃ (k)d (x) = d −1f (k) d (x):
f̃ (k) d (x) = σ
∗ ∫ âφ(ŵTx)µ
(k) d (dâ, dŵ), ∇ (k) fd `(x, y) = −y 1 + exp(df̃
(k) d (x)y)
∀k ≥ 0. (113)
A limit dynamics then takes the following form:
µ(k+1)∞ = µ (k) ∞ + div(µ (k) ∞ ∆θ (k) d ), µ (0) ∞ = δ ⊗N (0, Idx) (114)
∆θ(k)∞ (ã, ŵ) =
= −[η̂∗aσ∗∇ (k) f∞ `(x(k)a , y (k) a )φ(ŵ Tx(k)a ), η̂ ∗ wσ ∗∇(k)f∞`(x (k) w , y (k) w )ãφ ′(ŵTx(k)w )x (k),T w ] T , (115)
∇(0)f∞`(x, y) = −y
1 + exp(σ∗d∗,1/2N (0, σ(0),2(x))y) , (116)
f̃ (k)∞ (x) = σ ∗ ∫ ãφ(ŵTx)µ(k)∞ (dã, dŵ), ∇ (k+1) f∞
`(x, y) = −y[f̃ (k+1)∞ (x)y < 0] ∀k ≥ 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, f (0)d (x) = Θd→∞(1), hence the first point of Condition 5 holds. As for kernels, we have:
K (k) a,d(x,x ′) = σ∗,2 d∑ r=1 φ(ŵ(k),Tr x)φ(ŵ (k),T r x ′), (118)
K (k) w,d(x,x ′) = σ∗,2(d/d∗)−1 d∑ r=1 |â(k)r |2φ′(ŵ(k),Tr x)φ′(ŵ(k),Tr x′)xTx′. (119)
We see that while K(0)w,d converges to a constant due to the Law of Large Numbers, K (0) a,d diverges as d→∞. This violates the second statement of Condition 5, and the third as well, since f (0)∞ is finite. Consider now kernel increments:
∆K (k),′ aw,d(x,x ′) = −σ∗,3(d/d∗)−1/2 d∑ r=1 ( φ(ŵ(k),Tr x)φ ′(ŵ(k),Tr x ′)+
+ φ′(ŵ(k),Tr x)φ(ŵ (k),T r x
′) ) ×
×∇(k)fd `(x (k) w , y (k) w )â (k) r φ ′(ŵ(k),Tr x (k) w )(x + x ′)Tx(k)w , (120)
∆K (k),′ ww,d(x,x ′) = −σ∗,3(d/d∗)−3/2 d∑ r=1 |â(k)r |2 ( φ′(ŵ(k),Tr x)φ ′′(ŵ(k),Tr x ′)+
+ φ′′(ŵ(k),Tr x)φ ′(ŵ(k),Tr x
′) ) xTx′×
×∇(k)fd `(x (k) w , y (k) w )â (k) r φ ′(ŵ(k),Tr x (k) w )(x + x ′)Tx(k)w , (121)
∆K (k),′ wa,d(x,x ′) = −σ∗,3(d/d∗)−1/2 d∑ r=1 2â(k)r φ ′(ŵ(k),Tr x)φ ′(ŵ(k),Tr x ′)×
×∇(k)fd `(x (k) a , y (k) a )φ(ŵ (k),T r x (k) a ). (122)
For k = 0 terms inside sums of each increment have zero expectations. Hence the Central Limit Theorem can be used here. We get: ∆K(0),′aw,d = Θd→∞(1), ∆K (0),′ ww,d = Θd→∞(d −1), ∆K(0),′wa,d = Θd→∞(1). Since K (0) a,d = Θd→∞(d), K (0) w,d = Θd→∞(1), the last statement of Condition 5 is violated as well.
E INITIALIZATION-CORRECTED MEAN-FIELD (IC-MF) LIMIT
Here we consider the same training dynamics as for the mean-field scaling (see Appendix C.2), but with a modified model definition:
∆â(k)r = −η̂∗aσ∗∇ (k) fd `(x(k)a , y (k) a ) φ(ŵ (k),T r x (k) a ), â (0) r ∼ N (0, 1), (123)
∆ŵ(k)r = −η̂∗wσ∗∇ (k) fd `(x(k)w , y (k) w ) â (k) r φ ′(ŵ(k),Tr x (k) w )x (k) w , ŵ (0) r ∼ N (0, Idx). (124)
f (k) d (x) = σ ∗(d/d∗)−1 d∑ r=1 â(k)r φ(ŵ (k),T r x) + σ ∗(d/d∗)−1/2 d∑ r=1 â(0)r φ(ŵ (0),T r x), (125)
∇(k)fd `(x, y) = −y
1 + exp(f (k) d (x)y)
∀k ≥ 0. (126)
Similar to the mean-field case (Appendix C.2), we rewrite the dynamics above in terms of the weight-space measure:
µ (k+1) d = µ (k) d + div(µ (k) d ∆θ (k) d ), µ (0) d =
1
d d∑ r=1 δ θ̂ (0) r , θ̂(0)r ∼ N (0, I1+dx) ∀r ∈ [d], (127)
∆θ (k) d (â, ŵ) =
= −[η̂∗aσ∗∇ (k) fd `(x(k)a , y (k) a )φ(ŵ Tx(k)a ), η̂ ∗ wσ ∗∇(k)fd `(x (k) w , y (k) w )âφ ′(ŵTx(k)w )x (k),T w ] T , (128)
f (k) d (x) = σ
∗d∗ ∫ âφ(ŵTx)µ
(k) d (dâ, dŵ) + σ
∗(dd∗)1/2 ∫ âφ(ŵTx)µ
(0) d (dâ, dŵ), (129)
∇(k)fd `(x, y) = −y
1 + exp(f (k) d (x)y)
∀k ≥ 0. (130)
Note that here f (k)d stays finite in the limit of d→∞ for any k ≥ 0. Hence taking the limit d→∞ yields:
µ(k+1)∞ = µ (k) ∞ + div(µ (k) ∞ ∆θ (k) ∞ ), µ (0) ∞ = N (0, I1+dx), (131)
∆θ(k)∞ (â, ŵ) =
= −[η̂∗aσ∗∇ (k) f∞ `(x(k)a , y (k) a )φ(ŵ Tx(k)a ), η̂ ∗ wσ ∗∇(k)f∞`(x (k) w , y (k) w )âφ ′(ŵTx(k)w )x (k),T w ] T , (132)
f (k)∞ (x) = σ ∗d∗ ∫ âφ(ŵTx)µ(k)∞ (dâ, dŵ) + σ ∗d∗,1/2N (0, σ(0),2(x)), (133)
∇(k)f∞`(x, y) = −y
1 + exp(f (k) ∞ (x)y)
∀k ≥ 0. (134)

heading: F EXPERIMENTAL DETAILS
text: We perform our experiments on a feed-forward fully-connected network with a single hidden layer with no biases. We learn our network as a binary classifier on a subset of the CIFAR2 dataset (which is a dataset of first two classes of CIFAR101) of size 1024. We report results using a test set from the same dataset of size 2000. We do not do a hyperparameter search, for this reason, we do not use a validation set.
We train our network for 2000 training steps to minimize the binary cross-entropy loss. We use a full-batch GD as an optimization algorithm. We repeat our experiments for 10 random seeds and report mean and deviations in plots for logits and kernels (e.g. Figure 1, left). For plots of the KL-divergence, we use logits from these 10 random seeds to fit a single gaussian. Where necessary, we estimate data expectations (e.g. E x∼D|f(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, a multi-class classification) in Appendix G. All experiments were conducted on a single NVIDIA GeForce GTX 1080 Ti GPU using the PyTorch 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 if we sample σξ instead, where ξ can be any symmetric random variable with a distribution independent on hyperparameters.
In our experiments, we took a network of width d∗ = 27 = 128 with leaky ReLU activation and apply the Kaiming He uniform initialization (He et al., 2015) to its layers; we call this network a reference network. According to the Kaiming He initialization strategy, initial weights have a zero mean and a standard deviation σ∗ ∝ (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 η∗a = η ∗ w = 0.02. After that, we scale its initial weights and learning rates with width d according to a scaling at hand:
σ = σ∗ ( d
d∗
)qσ , η̂a∨w = η̂ ∗ a∨w ( d
d∗
)q̃a∨w .
Note that we have assumed σw = 1. By definition, η̂a∨w = ηa∨w/σ2a∨w; this implies:
ηa = η ∗ a ( σ σ∗ )2( d d∗ )q̃a = η∗a ( d d∗ )q̃a+2qσ , ηw = η ∗ w ( d d∗ )q̃w .

heading: G EXPERIMENTS FOR OTHER SETUPS
text: Although plots provided in the main body represent the full-batch GD on a 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 3-7. Here differences are marginal and not qualitative.
We have also experimented with multi-class classification: see Figure 8. Here we have trained our network on the full CIFAR10 dataset with SGD with batches of size 100. As we see on 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.
1CIFAR10 can be downloaded at https://www.cs.toronto.edu/~kriz/cifar.html

heading: H GENERALIZATION TO DEEP NETS PROPOSAL
text: While our present analysis is devoted to networks with a single hidden layer, we discuss possible generalizations to deep nets here.
Consider a network with H hidden layers. For simplicity, assume that widths of all hidden layers are equal to d. We thus have to considerH+1 learning rates q̃0:H , one for each layer, and similarlyH+1 initialization variances σ0:H . Without loss of generality, we may assume the input layer variance to be equal to 1 (we can rescale inputs otherwise). This gives 2H + 1 hyperparameters in total.
Similarly to what we did for H = 1, we assume that each hyperparameter obeys a power-law with respect to width. Let us refer the set of the power-law exponents as a "scaling". Again, we want to reason about what the scaling should be in order to converge to a dynamically stable limit model: see Condition 1. Moreover, we want to 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.
Having that much hyperparameters seems burdening, and this prohibits us to draw a nice twodimensional scaling plane as we did for H = 1: see Figure 1. For this reason, one have to reduce the dimensionality of a scaling.
First, it is tempting to consider a homogeneous activation function: a leaky ReLU. This introduces a symmetry in the weight space that guarantees that dynamics depends only on the product of initialization variances: σH × . . .× σ0; let us refer this product as σ. This approach was previously used by Golikov (2020), however we have to 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 + 1 to H + 1.
The next simplification should affect learning rate scaling exponents. Similar to what we have done for a shallow net, we may assume all learning rate exponents to be equal: q̃0 = . . . = q̃H = q̃.
The NTK limit, which generalizes naturally to deep nets, requires q̃0 = . . . = q̃H = 0, and hence conforms the assumption above. However, a possible generalization of the mean-field limit requires q̃0 = q̃H = 1, while q̃1 = . . . = q̃H−1 = 2; see Sirignano & Spiliopoulos (2019); Araújo et al. (2019); Golikov (2020). This aspect suggests the following alternatives:
1. Consider q̃0 = q̃H = q̃, while q̃1 = . . . = q̃H−1 = q̃hid; this results in a three-dimensional space of scalings: (qσ, q̃, q̃hid).
2. Consider q̃0 = q̃H = q̃, while q̃1 = . . . = q̃H−1 = 2q̃; 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 in the second class, it can be more expository to tighten the class to 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 to define which limit models are better than others in approximating finite-width nets ("nondominated"), we have to derive conditions that separate the domain of dynamically stable scalings into regions of distinct unique corresponding limit models, similar to Condition 2. We hypothesize that these conditions are similar to the shallow case: (1) a limit model at initialization is finite, (2) kernels at initialization are finite, (3) a limit model and kernels are of the same order, (4) kernels evolve at initialization. Since we 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.
It will be very interesting to check if all of the dynamically stable limit models are specified either by an evolution in a model space driven by constant kernel, or by an evolution of a weight-space measure, as was the case for H = 1; see Appendix C. Investigating a non-dominated limit model, different from both the NTK and the mean-field models, should be a valuable outcome of the proposed research
program; it will be even more valuable if this limit model will not be covered by both mean-field and constant kernel formalisms.
We also have to note that according to Golikov (2020), the mean-field limit vanishes for H > 2. This fact suggests that the analysis for deep nets should be held for H = 2 and for H > 2 separately.

heading: I MEASURING DIVERGENCE BETWEEN A LIMIT MODEL AND A REFERENCE ONE
text: We track the divergence of a limiting network from a reference one, which can be done in two ways. The first one is tracking divergence directly between logits: E x∼DtestDlogits(f (k) ∞ (x) || f (k)d∗ (x)) for some divergence measure D(· || ·). The second one is tracking divergence between probabilities: E x∼DtestDprob(σ(f (k) ∞ (x)) || σ(f (k)d∗ (x))), where we have overloaded the notation by denoting the standard sigmoid as σ: σ(x) = (1 + exp(−x))−1. We choose a KL-divergence for the first case. However, measuring a KL-divergence between logits is hardly possible, since we do not have an access to the distribution of f (k)(x) as a random variable depending on initialization. For this reason, we fit a gaussian to its samples:
Dlogits(ξ || ξ∗) = KL(N (E ξ,Var ξ) || N (E ξ∗,Var ξ∗)). (135)
This case is depicted in Figure 9, left.
As for the second case, we may want to measure a KL-divergence between distributions on probabilities. Again, this is not possible, because the true distribution of σ(f (k)∞ (x)) is not known; for this reason, we decide to first fit a beta distribution, and then measure the divergence:
Dprob(ξ || ξ∗) = KL(Beta(αmle(ξ), βmle(ξ)) || Beta(αmle(ξ∗), βmle(ξ∗))), (136)
where αmle(ξ) and βmle(ξ) are maximum-likelihood estimations of hyperparameters of a beta: αmle(ξ) = E ξ ( E ξ(1− E ξ) Var ξ − 1 ) , βmle(ξ) = (1− E ξ) ( E ξ(1− E ξ) Var ξ − 1 ) . (137)
This case is plotted in Figure 9, right.
Both cases can be generalized to work with multi-class classification. In the first case, we can simply fit a gaussian with a diagonal covariance matrix; this was done in Figure 8, left. In the second case, we can fit a Dirichlet random variable using a maximum-likelihood estimation as before.

Review summaries:
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.

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.