Title: MEAN FIELD MODELS FOR NEURAL NETWORKS IN TEACHER-STUDENT SETTING

Abstract: Mean field models have provided a convenient framework for understanding the training dynamics for certain neural networks in the infinite width limit. The resulting mean field equation characterizes the evolution of the time-dependent empirical distribution of the network parameters. Following this line of work, this paper first focuses on the teacher-student setting. For the two-layer networks, we derive the necessary condition of the stationary distributions of the mean field equation and explain an empirical phenomenon concerning training speed differences using the Wasserstein flow description. Second, we apply this approach to two extended ResNet models and characterize the necessary condition of stationary distributions in the teacher-student setting.

1. INTRODUCTION:

In the past several years, neural networks have achieved tremendous successes in many areas of machine learning and artificial intelligence (Devlin et al., 2018; He et al., 2016; Silver et al., 2016; Vaswani et al., 2017) , based on remarkable developments in algorithms (Robbins & Monro, 1951; Duchi et al., 2011; Rumelhart et al., 1988) , software (Paszke et al., 2017; Jia et al., 2014; Abadi et al., 2016) , and hardware (Lindholm et al., 2008; Jouppi et al., 2018) . In spite of the empirical successes, the theoretical understanding of neural networks is rather incomplete. To bridge this gap, tools from different fields, such as statistical mechanics, partial differential equations, dynamics systems, have been brought into use. Recently, a series of remarkable papers (Chizat & Bach, 2018; Mei et al., 2018; Rotskoff & Vanden-Eijnden, 2018; Sirignano & Spiliopoulos, 2018) analyze the two-layer neural networks using mean field models and nonlinear partial differential equations (PDEs). The main result is that, under the large width limit and sufficient small learning rates, stochastic gradient descent (SGD) dynamics of a two-layer fully connected neural network with a convex loss can be modeled as a gradient flow of a convex functional of probability distributions under the Wasserstein metric (Peyré et al., 2019; Santambrogio, 2015; Villani, 2008) . Though only for two-layer networks, this result significantly simplifies our understanding of the training dynamics. In two recent papers (Nguyen, 2019; Sirignano & Spiliopoulos, 2019) , this approach has been extended to multi-layer fully-connected networks by taking the widths to infinity successively from higher to lower layers.

Contributions. This paper follows this line of research, focusing on the teacher-student setting. The teacher network is typically a narrow network that computes the function to be learned, while the student network is a vastly over-parameterized wide network. The goal is to optimize the parameters of the student in order to produce the function represented by the teacher. Though somewhat restricted, the teacher-student setting allows for assessing performance by comparing the parameters of the networks. The main contributions of this paper are listed as follows.

First, for the two-layer networks in the teacher-student setting, we characterize for the first time the necessary condition for a stationary distribution of the mean field PDE and show how it relates to the parameters of the teacher network. We also discuss how different activation functions influence the results. Using the Wasserstein flow description, we also provide a simple explanation for why the convergence to teacher parameters with larger weight is much faster than to the ones with smaller weight.

Second, we extend the mean field analysis to two extended ResNet models: the ensemble-ResNet (an average of a large number of independent ResNets) and the wide-ResNet (a ResNet with wide residual blocks). We derive mean field PDEs for the training dynamics of the parameters in both cases and also characterize the necessary condition for a stationary distribution for the ensemble-ResNet model as well as the first and second order approximations of the wide-ResNet model. Though our derivation remains formal, this is the first time that the mean field model of the ResNet has been considered.

Related work.:

Besides the aforementioned papers on the mean field analysis for neural networks, tools from stochastic analysis (Mandt et al., 2015; Li et al., 2017; 2019; Yaida, 2018; An et al., 2018) and Wasserstein geometry (Arjovsky et al., 2017; Cuturi, 2013; Frogner et al., 2019; Peyré et al., 2019; Solomon et al., 2014) have been playing an increasingly prominent role in machine learning and artificial intelligence.

Another direction that considers the infinite width neural networks is the work (Jacot et al., 2018) of the neural tangent kernel (NTK), which describes under an appropriate scaling the evolution of the function represented by the neural network. At the infinite width limit, the NTK converges to a deterministic limit that depends on the nonlinear activation function and remains unchanged during the training process. Though it is not yet clear how this limit relates to the actual training dynamics of neural networks used in practice (Chizat et al., 2019) , mathematically it is an interesting regime that has inspired a lot of recent work (Arora et al., 2019a; b; Cao & Gu, 2019) .

In recent years, particle-based methods have also been widely applied to many machine learning problems, such as in Bayesian inference. One such recent example is the Stein variational gradient descent (SVGD) method (Liu & Wang, 2016; Liu, 2017) that approximates a target distribution with a group of interacting particles using iterative gradient updates. It has been shown that, in the large particle number limit, the empirical distribution of the parameters can be modeled by a nonlinear PDE, which can be viewed as a gradient flow of a convex loss function under a Wasserstein-type but non-local metric (Liu, 2017; Lu et al., 2019) .

Organization. The rest of the paper is organized as follows. Section 2 concerns the mean field description of two-layer networks and provides necessary conditions for stationary distributions. Section 3 extends the analysis to two extended ResNet models. Finally, Section 4 discusses some future directions.

2. TWO-LAYER NETWORK IN TEACHER-STUDENT SETTING:



2.1. TWO-LAYER NETWORKS:

Setup. Given a function f : R d → R of the independent variable x ∈ R d , consider its approximation by a two-layer network model 1 N N j=1 w j σ(x • θ j ), where N is the number of hidden units, {θ j } are the parameters in a domain Θ ⊂ R d , {w j } are the weights, and σ : R → R is an activation function such as rectified linear unit (ReLU), softplus, etc. The following discussion focuses on the special case where the weights w j are fixed to be one, i.e.,

1 N N j=1 σ(x • θ j ).

(1)

The more general case of trainable w j values are discussed in Appendix 5.1. Assuming that the input x is sampled from a measure µ(x) independently, the learning problem of this two-layer network is to minimize the loss

E(θ 1 , . . . , θ N ) = x 1 N N j=1 σ(x • θ j ) -y(x) µ(x)dx

, where (•) is a loss function. The rest of this paper focuses for example on the quadratic loss case, i.e., (z) = 1 2 |z| 2 :

E(θ 1 , . . . , θ N ) ≡ 1 2 x 1 N N j=1 σ(x • θ j ) -y(x) 2 µ(x)dx ≡ 1 2 1 N N j=1 σ(x • θ j ) -y(x) µ .

When optimized with a standard algorithm such as gradient descent (GD), each parameter θ j for 1 ≤ j ≤ N is updated with

θ i [k+1] = θ i [k]-α ∂E ∂θ i [k], ∂E ∂θ i [k] = x   1 N N j=1 σ(x • θ j [k]) -y(x)   1 N σ (x•θ i [k])xµ(x)dx,

where α is the learning rate of the iterative algorithm. For the SGD, the right hand side is supplemented with a mean-zero noise term. When the learning rate α is sufficiently small, the iterates {θ i [k]} can be modeled by a continuous-time ordinary differential equation (ODE) system {θ i (t)} by identifying t = kα/N

EQUATION

Mean field equation. Viewing θ i (t) as interacting particles and considering the empirical distribution ρ(t, θ)

≡ 1 N N j=1 δ θj (t) (θ), we can rewrite 1 N j σ(x • θ j ) -y(x) as σ(x • θ)ρ(θ)dθ and E(θ 1 , . . . , θ N ) as E(ρ) ≡ 1 2 σ(x • θ)ρ(θ)dθ -y(x) 2 µ .

The mean field evolution equation for ρ(t, θ) associated with ( 2) is

ρ t (θ, t) = div θ ρ(θ, t) x θ σ(x • θ )ρ(θ , t)dθ -y(x) σ (x • θ)xµ(x)dx .

Noticing that the functional derivative of E(ρ) with respect to ρ is given by

δE(ρ) δρ = x θ σ(x • θ )ρ(θ )dθ -y(x) σ(x • θ)µ(x)dx,

one can write the mean field PDE for ρ compactly as

EQUATION

which can be interpreted as the gradient flow of the functional E(ρ) under the the Wasserstein metric (∇ T θ ρ∇ θ ) + (Jordan et al., 1998; Otto, 2001; Otto & Villani, 2000) . Though the above derivation is mostly formal, one can show that the dynamics of θ i (t) converges to the evolution of ρ(θ, t) for any fixed period of time [0, T ] when the number of hidden units N goes to infinity. This is formalized for example in the following theorem. Theorem 1 (Reformulated from (Chizat & Bach, 2018) ). If θ i (0) are sampled independently from a density ρ 0 (x) at time t = 0, then for any fixed T > 0 the measure 1 N N j=1 δ θj (t) (θ) converges weakly to the solution ρ(t, θ) of (3) for any t ∈ [0, T ] as N goes to infinity.

For the case of bounded parameter domain Θ, such as a ball in R d with a fixed radius, one needs to specify the appropriate boundary conditions both for the ODE system (2) and the mean field PDE (3). First, the dynamics of θ i (t) should follow normal reflection when θ i (t) hits the boundary ∂Θ of Θ. Second, for the mean field equation of ρ(t, θ), the Neumann boundary condition ∂ρ(t,θ) ∂n (t, θ) = 0 should be used at θ ∈ ∂Θ. The bounded case is particularly relevant to the discussion in Section 3.2.

Stationary distribution in teacher-student setting. Given the PDE formulation of the parameter evolution, one natural question is to characterize the stationary distribution of (3). In what follows, we address this question in the teacher-student setting, i.e., there exists a parameter distribution ρ * (θ) such that the target function y(x) satisfies

y(x) = θ σ(x • θ)ρ * (θ)dθ.

With this teacher representation, the loss E(ρ) can be written equivalently as

E[ρ] = 1 2 x θ σ(x • θ)(ρ -ρ * )(θ)dθ 2 µ(x)dx.

The following theorem characterizes the stationary distribution in this teacher-student setting.

Theorem 2. Suppose that µ(x) > 0 for x ∈ R d . Then any stationary distribution ρ(θ) that satisfies ρ(θ) > 0 for any θ ∈ Θ has zero loss and satisfies the condition

σ(x • θ)(ρ(θ) -ρ * (θ))dθ = 0. (4)

Proof. By introducing the operator Σ with kernel given by σ(x•θ), one can write the loss as

E(ρ) = 1 2 Σ(ρ -ρ * ) 2

µ and the functional derivative as

δE δρ = Σ T µΣ(ρ -ρ * ),

where µ is considered to the operator of multiplying with µ(x) at each x. At a stationary distribution ρ, since ρ t ≡ 0, div θ ρ∇ θ δE δρ = 0 at any θ ∈ Θ. Taking inner product with δE δρ gives δE δρ

T div θ ρ∇ θ δE δρ = 0.

Using the fact that div θ is the adjoint of ∇ θ , this states ρ(θ)∇ θ δE δρ = 0, which further implies ∇ θ δE δρ = 0 as ρ(θ) > 0. This means that δE δρ = Σ T µΣ(ρ -ρ * ) as a function of θ is a constant. Taking inner-product of δE δρ with (ρ -ρ * )(θ) and using the fact that (ρ

-ρ * )(θ)dθ = 0 gives (ρ -ρ * ) T Σ T µΣ(ρ -ρ * ) = 0.

This implies zero loss and µ(x)Σ(ρ When σ(z) is an odd function such as tanh or arctan, (4) states that ρ is a stationary distribution if and only if the difference (ρ -ρ * )(θ) is an even function. In addition, when σ(z) is the sum of an odd function and a constant (for example, like the Sigmoid function), (4) is equivalent to (ρ -ρ * )(θ) being even, since the integral of (ρ -ρ * )(θ) is zero.

-ρ * ) = 0, which further implies Σ(ρ -ρ * ) = 0 since µ(x) > 0. Writing Σ(ρ -ρ * ) =

Finally, the softplus function σ(z) = ln(1 + e z ), which is a smooth approximation of ReLU. (4) states that ln(1 + e θx )(ρ -ρ * )(θ)dθ = 0.

Let us assume for simplicity that both ρ(θ) and ρ * (θ) decay faster than θ -2 at large θ and define f (θ) as the anti-derivative of (ρ -ρ * )(θ). Integrating by parts for (4) results in

1 1 + e -θx f (θ)dθ = 0.

Noticing that the Sigmoid function 1 1+e -θx is a sum of an odd function and a constant, we conclude that (4) is equivalent to f (θ) being even and mean-zero. In terms of the difference (ρ -ρ * )(θ), these two conditions means that (ρ -ρ * )(θ) is odd and has a vanishing first moment θ • (ρ -ρ * )(θ)dθ = 0.

2.2. CONVERGENCE SPEED ISSUE:

One interesting empirical observation from training neural networks in the teacher-student setting is that the convergence to a teacher parameter with a large weight is much faster than the speed to a teacher parameter with a small weight. More specifically, consider a teacher network given by ρ

* (θ) = α 1 δ c1 (θ) + α 2 δ c2 (θ),

where c i are the parameter values in R d and the weights α i satisfy α 1 α 2 > 0 and α 1 + α 2 = 1. This can either be realized as a general two-layer network of two hidden units with weights α 1 and α 2 or an equal-weight network (1) of α 1 N units with parameter c 1 and α 2 N units with parameter c 2 . After an initial short period, the evolution of the student parameters decouple into two nearly independent problems: an α 1 fraction of the N student parameters converges to c 1 , while an α 2 fraction of the student parameters converges to c 2 . The observed behavior is that the convergence to c 1 is much faster than the convergence to c 2 .

The mean field PDE allows us to understand this phenomenon using scale analysis. After the initial stage, the PDE ρ t = div θ ρ∇ θ Σ T µΣ(ρ -ρ * ) is effectively decoupled into two equations

ρ 1 t = div θ ρ 1 ∇ θ Σ T µΣ(ρ 1 -ρ 1 * ) , where ρ 1 * = α 1 δ c1 (θ), ρ 2 t = div θ ρ 2 ∇ θ Σ T µΣ(ρ 2 -ρ 2 * ) , where ρ 2 * = α 2 δ c2 (θ).

The key observation is that the quadratic dependence on ρ on the right hand side is the reason of different convergence speeds. To see this, let us divide these two equations by α 1 and α 2 , respectively,

ρ 1 α 1 t = α 1 div θ ρ 1 α 1 ∇ θ Σ T µΣ ρ 1 α 1 - ρ 1 * α 1 , where ρ 1 * α 1 = δ c1 (θ), ρ 2 α 2 t = α 2 div θ ρ 2 α 2 ∇ θ Σ T µΣ ρ 2 α 2 - ρ 2 * α 2 , where ρ 2 * α 2 = δ c2 (θ).

These two are the same equations for normalized densities ρ 1 α1 and ρ 2 α2 , except the coefficients α 1 and α 2 at the beginning of the right hand sides. The difference in these two coefficients clearly demonstrates that the first equation converges α 1 /α 2 times faster than the second one.

3. RESNET MODELS IN TEACHER-STUDENT SETTING:

Consider a simple ResNet model with input x ∈ R d (see Figure 1 

left) X 0 = x, X +1 = X + 1 L σ(θ X ), 0 ≤ < L,

where X ∈ R d , the matrix parameter θ ∈ Θ ⊂ R d×d , and the activation σ is applied entry-wise to each component. By introducing θ = (θ 0 , . . . , θ L-1 ), we denote the final result at level L by X L θ in order to emphasize the dependence on θ. Given a function y : R d → R d , the training problem searches for the parameters θ = (θ 0 , . . . , θ L-1 ) ∈ Θ L in order to minimize, say, the quadratic loss

E(θ) = 1 2 x X L θ (x) -y(x) 2 µ(x)dx = 1 2 X L θ (x) -y(x) 2 µ .

When optimized with GD, the parameter vector θ is updated at each step with

θ[k + 1] = θ[k] -α ∂E ∂θ (θ[k])

(or with an extra noise term for SGD). Instead of directly working with this vanilla ResNet, we introduce two extended ResNet models in the rest of this section and derive a mean field equation for each of them. 

E(θ 1 , . . . , θ N ) = 1 2 1 N N j=1 X L θj (x) -y(x) 2 µ .

The GD/SGD algorithm with learning rate α takes the following step for each i

θ i [k + 1] = θ i [k] -α x 1 N ∇ θi X L θi[k] (x) T   1 N N j=1 X L θj [k] (x) -y(x)   µ(x)dx.

When α/N is sufficiently small, the dynamics of θ i [k] can be approximated by an ODE system θ i (t) by identifying t = kα/N

dθ i dt = - x ∇ θi X L θi (x) T   1 N N j=1 X L θj (x) -y(x)   µ(x)dx. Viewing θ i (t) ≡ (θ 0 i (t), . . . , θ L-1 i (t)) as particles in Θ L and introducing the distribution ρ(t, θ) = 1 N N j=1 δ θj (t) (θ), we rewrite E(θ 1 , . . . , θ N ) as E(ρ) ≡ 1 2 θ X L θ (x)ρ(θ)dθ -y(x) 2 µ .

The mean field equation of ρ(t, θ) takes the same form as (3)

ρ t (θ) = div θ ρ∇ θ δE δρ ,

which is a gradient flow of quadratic loss E(ρ) under the Wasserstein metric on the densities ρ(θ) on the product space Θ L .

Stationary distribution in teacher-student setting. In the teacher-student setting, the function y(x) from R d to R d is assumed to take the form y(x) = θ X L θ (x)ρ * (θ)dθ, for a certain distribution ρ * on the space Θ L . This include the special case that y(x) can be represented by a vanilla ResNet described above. More precisely, assuming that

X 0 = x, X +1 = X + 1 L σ(θ * X ), y(x) = X L (x), then y(x) corresponds to ρ * (θ) = δ θ * (θ) with θ * ≡ (θ 0 * , . . . , θ L-1 * ) ∈ Θ L .

Theorem 3. Suppose that µ(x) > 0 for x ∈ R d . Then any stationary distribution ρ(θ) that satisfies ρ(t, θ) > 0 for any t > 0 and θ ∈ Θ L has zero loss and satisfies the condition, for each x,

EQUATION

The proof of the theorem follows the same argument as Theorem 2. One potential criticism of the ensemble-ResNet model is that the space Θ L ⊂ (R d×d ) L can be very high-dimensional even for moderate values of L. Therefore, the number of branches N required in order for the mean field model to be accurate can be quite large.

3.2. WIDE-RESNET:

In order to overcome the exponential grow of the space Θ L in the ensemble-ResNet model, we consider a different ResNet model in this subsection. Instead of allowing for multiple independent branches, the wide-ResNet model sticks to a single path but increases the width of each residual block (see Figure 1 right). For an input x ∈ R d , the wide-ResNet evaluates

EQUATION

where the activation function σ(z) is assumed to be smooth and vanishes at z = 0. By introducing the parameter group θ = (θ 1 , . . . , θ N ) for each level , we write the final result of (6) as X L θ 0 ,...,θ L-1 (x) and the quadratic loss function as

E(θ 0 , . . . , θ L-1 ) = E(θ 0 1 , . . . , θ 0 N , . . . , θ L-1 1 , . . . , θ L-1 N ) = 1 2 X L θ 0 ,...,θ L-1 (x) -y(x) 2 µ .

The GD/SGD algorithm with learning rate α takes the form

θ i [k + 1] = θ i [k] -α x ∇ θ i X L θ 0 [k],...,θ L-1 [k] (x) T X L θ 0 [k],...,θ L-1 [k] (x) -y(x) µ(x)dx,

where

∇ θ i X L θ 0 [k],...,θ L-1 [k] (x) contains a factor 1 LN .

When the stepsize α is sufficiently small, the dynamics of θ i [k] can approximated with an ODE system

dθ i dt = - x ∇ θ i X L θ 0 ,...,θ L-1 (x) T X L θ 0 ,...,θ L-1 (x) -y(x) µ(x)dx.

By viewing θ 1 (t), . . . , θ N (t) as interacting particles at each level and introducing ρ (t, θ ) = 1 N j δ θ j (t) (θ ) for each , one can rewrite the wide-ResNet computation as

EQUATION

By taking the limit N to infinity, the loss becomes E(ρ 0 , . . . , ρ

L-1 ) = 1 2 X L ρ 0 ,...,ρ L-1 (x) -y(x) 2 µ

(with the subscripts denoting the dependence on ρ 0 , . . . , ρ L-1 ) and the mean field equation for ρ (t, θ ) for each level reads

EQUATION

As σ(z) is assumed to be smooth, expanding (7) recursively and applying the Taylor expansion (see Appendix 5.2 for a derivation) gives the following approximation

EQUATION

Furthermore, since σ(z) vanishes at zero by assumption, this expansion has L terms of order O(θ) each with weight 1/L, L 2 terms of order O(θ 2 ) each with weight 1/L 2 , etc. When the domain Θ of the matrix parameter θ is bounded by a sufficiently small radius, the contribution from the k-linear terms is of order O(1/k!) and hence decays rapidly with increasing values of k. Therefore, in the same parameter regime, it is sufficient to consider approximations of the first few orders.

Linear approximation. The linear approximation for X L in (9) leads to the following first-order approximation to the loss E(ρ 0 , . . . , ρ L-1 ) ≈ E (1) (ρ 0 , . . . , ρ L-1 ) ≡

1 2 x + 1 L a σ(θx)ρ a (θ)dθ -y(x) 2 µ .

Now take the limit of L going to infinity and define ρ(θ, ε) = a ρ a (θ) • δ a/L (ε) where ε ∈ [0, 1] is a rescaled continuous variable associated with the depth of the ResNet. In this formal limit, the loss function becomes

E (1) (ρ) ≡ 1 2 x + σ(θx)ρ(θ, ε)dθdε -y(x) 2 µ .

The mean field equation for ρ(t, θ, ε) takes the form

EQUATION

)

Notice that this is a degenerate gradient flow with mobility only in the θ variable. In the teacherstudent setting, i.e. y(x) = x+ σ(θx)ρ * (θ, ε)dθdε for some teacher distribution ρ * (θ, ε), we can write compactly

E (1) (ρ) = 1 2 Σ 1 (ρ -ρ * ) 2 µ

, where Σ 1 is the operator with kernel σ(θx). Notice that ε does not appear in the kernel explicitly. Theorem 4. Suppose that µ(x) > 0 for x ∈ R d . Then any stationary distribution ρ(θ, ε) that satisfies ρ(θ, ε) > 0 for any θ ∈ Θ and ε ∈ [0, 1] has zero loss and satisfies, at each

EQUATION

The proof follows the argument of Theorem 2. It is not surprising that the constraint ( 11) is formulated in terms of the marginal of ρ, since the parameters at different levels become almost commutative in the small θ regime.

Quadratic approximation. The quadratic approximation for X L in (9) leads to the following second-order approximation to the loss

E (2) (ρ 0 , . . . , ρ L-1 ) ≡ 1 2 x + 1 L L-1 a=0 σ(θx)ρ a (θ)dθ -y(x) + 1 L 2 b>a ∇σ(θ b X 0 )θ b σ(θ a X 0 )ρ b (θ b )ρ a (θ a )dθ a dθ b -y(x) 2 µ .

Taking the large L limit and introducing ρ(θ, ε) ≡ a ρ a (θ) • δ a/L (ε) gives rise to the loss function

E (2) (ρ) ≡ 1 2 x + σ(θx)ρ(θ, ε)dθdε + ε >ε ∇σ(θ x)θ σ(θx)ρ(θ , ε )ρ(θ, ε)dθ dθdε dε -y(x) 2 µ

in terms of ρ(θ, ε) and the mean field equation takes the same form as (10) except with E (1) replaced with E (2) . In the teacher-student setting, there exists some ρ * (θ, ε) such that

y(x) = x + σ(θx)ρ * (θ, ε)dθdε + ε >ε ∇σ(θ x)θ σ(θx)ρ * (θ , ε )ρ * (θ, ε)dθ dθdε dε. By introducing a new operator Σ 2 (ρ ⊗ ρ) = ε >ε ∇σ(θ x)θ σ(θx)ρ * (θ , ε )ρ * (θ, ε)dθ dθdε dε,

we can write it more compactly as

EQUATION

µ . Unfortunately, since E (2) (ρ) is quartic in ρ, a stationary distribution might be a local minimum even with the extra assumption ρ(θ, ε) > 0. On the other hand, the explicit and simple form of E (2) (ρ) might allow multiple runs with different initial condition ρ(0, θ) to identify the ρ * more efficiently.

4. DISCUSSIONS:

This paper discusses necessary conditions for the stationary distributions of mean field models for the two-layer neural networks as well as two extended ResNet models. A lot of interesting questions remain to be addressed concerning the ResNet. The first, and probably most important, question is that whether it is possible to obtain a mean field model for the vanilla ResNet by treating the depth L as the large parameter going to infinity. This in fact has been the original motivation of this paper and it still remains unresolved. Second, in the wide-ResNet model, the resulting PDE is degenerate as there is no diffusion in the ε variable. Would it be possible to include the diffusion in ε by allowing parameters to jump between adjacent layers? If this is possible, the next step is to formulate a mean field model with densities over some "architecture" space for neural architecture search (Zoph & Le, 2016; Zoph et al., 2018; Elsken et al., 2018; Wistuba et al., 2019) .

where µ is considered to the operator of multiplying with µ(x) at each x. At a stationary distribution ρ, since ρ t ≡ 0, div θw ρ∇ θw 

5.2. DERIVATION OF (9):

The computation follows the definition in (7) and uses Taylor expansion. At the first level,

X 1 = X 0 + 1 L θ 0 σ(θ 0 X 0 )ρ 0 (θ 0 )dθ 0 .

At the second level,

X 2 = X 1 + 1 L θ 1 σ(θ 1 X 1 )ρ 1 (θ 1 )dθ 1 = X 0 + 1 L θ 0 σ(θ 0 X 0 )ρ 0 (θ 0 )dθ 0 + 1 L θ 1 σ θ 1 X 0 + 1 L θ 0 σ(θ 0 X 0 )ρ 0 (θ 0 )dθ 0 ρ 1 (θ 1 )dθ 1 .

Recalling that σ is assumed to be smooth near the origin, a Taylor expansion for σ in the last term

X 2 ≈X 0 + 1 L θ 0 σ(θ 0 X 0 )ρ 0 (θ 0 )dθ 0 + 1 L θ 1 σ(θ 1 X 0 )ρ 1 (θ 1 )dθ 1 + 1 L 2 θ 1 ∇σ(θ 1 X 0 )θ 1 θ 0

σ(θ 0 X 0 )ρ 0 (θ 0 )dθ 0 ρ 1 (θ 1 )dθ 1 + h.o.t.

Iterating this until the L-th level gives rise to 

X L ≈ X 0 + 1 L L-

:



Figures and Tables:

0 in the integral form gives (4). Below we consider the implication of Theorem 2 for a few commonly-used activation functions. Example 1. Consider the ReLU activation σ(z) = max(z, 0). It is convenient to use the polar coordinates x → (x, |x|) and θ → ( θ, |θ|), where | • | and • denote the radial coordinate and the angular variables, respectively. Then (4) reads θ |θ| |x||θ| d σ(x • θ)(ρ( θ, |θ|) -ρ * ( θ, |θ|))d θd|θ| = 0, where the factor |θ| d comes from the homogeneity of ReLU as well as the Jacobian from the polar coordinate transformation. If the operator with kernel (σ(x • θ)) x, θ∈S d-1 has a trivial null space, then one can conclude that the stationary distribution ρ( θ, |θ|) satisfies for any θ ∈ S d-1 |θ| |θ| d (ρ( θ, |θ|) -ρ * ( θ, |θ|))d|θ| = 0.Example 2. Consider for simplicity the 1D case d = 1. When σ(z) is an even function such as the Gaussian or quadratic activations, (4) states that ρ(θ) is a stationary distribution if and only if the difference (ρ -ρ * )(θ) is an odd function.

Figure 1: ResNet models. Left: vanilla ResNet. Middle: ensemble-ResNet. Right: wide-ResNet. Mean field models for the ensemble-ResNet and wide-ResNet models are discussed in Section 3.

any θ ∈ Θ and w ∈ W . Taking inner product with δE δρ Using the fact that div θ is the adjoint of ∇ θw , this means ρ(θ, w)∇ θw δE δρ = 0, which further implies ∇ θw δE δρ = 0 as ρ(θ, w) > 0. This means that δE δρ = Σ T µΣ(ρ -ρ * ) as a function of θ and w is constant. Taking inner-product of δE δρ with (ρ -ρ * )(θ, w), one arrives at(ρ -ρ * ) T Σ T µΣ(ρ -ρ * ) = 0,using the fact that (ρ -ρ * )(θ, w)dθdw = 0. This implies zero loss andµ(x)Σ(ρ -ρ * ) = 0,which further implies Σ(ρ -ρ * ) = 0 since µ(x) > 0. Writing the last statement explicitly in the integral form gives (13). Example 3. Let us consider the ReLU activation σ(z) = max(z, 0). Using the polar coordinates x → (x, |x|) and θ → ( θ, |θ|), (13) reads θ |θ| w |x||θ| d wσ(x • θ)(ρ( θ, |θ|, w) -ρ * ( θ, |θ|, w))d θd|θ|dw = 0. If the operator with the kernel (σ(x • θ)) x, θ∈S d-1 has a trivial null space, then one can conclude that that for each θ |θ| w |θ| d w ρ( θ, |θ|, w) -ρ * ( θ, |θ|, w) d|θ|dw = 0 that at the stationary distribution ρ.

1 a=0 σ(θX 0 )ρ a (θ)dθ+ 1 L 2 b>a ∇σ(θ b X 0 )θ b σ(θ a X 0 )ρ b (θ b )ρ a (θ a )dθ a dθ b +h.o.t.

Peer review (iclr2020):
This paper studies the mean fields limit of neural networks and extends the mean field limit treatment to resnets. The paper seems to be presenting non-trivial formal results but I could not really understand what are their specific consequences of interest. What are the concrete insights that the results of this paper bring? Even as a theoretician, I did not really understand what should I take from this paper. As such I cannot recommend its acceptance. But it is well plausible that with better explanations I will understand its value and change my mind. Some more concrete questions/issues follow: The introduction several times states: *First, for the two-layer networks in the teacher-student setting, we characterize for the first time the necessary condition for a stationary distribution of the mean field PDE and show how it relates to the parameters of the teacher network. We also discuss how different activation functions influence the results.* The outline states: *Section 2 concerns the mean field description of two-layer networks and provides necessary conditions for stationary distributions.* But where is this discussion of necessary condition? I cannot find a single mention of *necessary* in section 2 that only includes formal derivations and statements. Can the result of Section 2 be translated to somewhat more general audience? What is the take-home message there? Where is the discussion on the role of the activation function? I see the two examples stating the formal result for two different activations, but again I am unable to understand what should one take from this? I could understand the result of section 2.2. about stronger weight-teachers being learned first. But this is completely intuitive in the same way as independent components corresponding to larger eigenvalues would get learned first in PCA. See also similar conclusions in works on so-called *INCREMENTAL LEARNING* in Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120, 2013. And subsequent results. So this alone does not seem very novel result. As I had hard time understanding the section two on the simpler feedforward case, Section 3 was even less clear to me. How does what is done here compare to what is known on resnets and what should somebody interested in resnets (but not specifically in the mean field equations) take out from this. Some non-technical summary of the findings is seriously missing in this paper. Comments that are less fundamental to the overall understanding: ** The mean field treatment was also extended to the multi-layer case in: https://arxiv.org/abs/1906.00193 ** The paper present nice account on previous results involving the mean field limit. The manuscript should also discuss the long line of papers analyzing the teacher student setting on one-hidden-layer neural networks. After-all this seems to be the main object of study here so the results only make sence presented against what was previously known. Notably the case of eq. (1) where the 2nd layer weights are fixed to one is a model called the committee machine widely considered in previous literature (in the case of finitely wide one-hidden-layer network). The (non-exhaustive) list of paper on the teacher-student setting is (more references are in those papers): -- The teacher-student setting was introduced in Garder, Derrida*83 (model B, https://iopscience.iop.org/article/10.1088/0305-4470/22/12/004/pdf). -- In the classical textbook on neural networks Engel, Andreas, and Christian Van den Broeck. Statistical mechanics of learning. Cambridge University Press, 2001, there is a rather detailed account of many results on the setting from 80s and 90s. -- Notably the plain SGD was analyzed via ODEs for the teacher-student setting in classical papers: David Saad and Sara A Solla. On-line learning in soft committee machines. Physical Review E, 52 (4):4225, 1995. David Saad and Sara A Solla. Dynamics of on-line gradient descent learning for multilayer neural networks. In Advances in neural information processing systems, pp. 302–308, 1996. -- This line of work was recently extended with a focus on the overparametrized regime (but not infinitely wide) in: Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup, Sebastian Goldt, Madhu S. Advani, Andrew M. Saxe, Florent Krzakala, Lenka Zdeborová, NeurIPS*19. -- There even in analysis with more than one hidden layer in Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. NeurIPS, 2019a. It would be really interesting to see a discussion of what is similar and different between these works and results and the present work. And what is the added value of the present one about understanding the teacher-student setting.