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):
The paper studies the dynamics of neural network in the Teacher-Student setting, using the approach pionnered in the last few years. Concerning the presentation and the review of other works, I am a bit surprise that the *teacher-student* setting appears out of the blue, without any references to any paper. This is not a new area! Especially when it concerns learning with SGD, these were the subject of intense studies, especially in the 80s, see for instance: * Saad & Sola On-line learning in soft committee machines, *95 * Exact solution for on-line learning in multilayer neural networks *95 or the book *On-Line Learning in Neural Networks*, By David Saad, with contribution of Botou, etc... * Many of these results were proven rigorously recently in *Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup* by Goldt et al, There is a difference with the current formulation: in those papers, both the teacher AND student had a finite-size second (or more) layer, while here, one is working in the mean-field regime where the student is infinity wide. This is indeed a different situation, where the system can be much more *over-parametrized*. But this does not mean that the subject is *terra incognita*. There are three main sections in the paper, discussing the results. * The first result is a theorem that, if I understand correctly, says that the stationary distribution of gradient descent on the population loss (i.e. the test error) a necessary condition for the stationary distribution is that it has zero generalisation error (Eq. 4). That seems like an incremental-step compared to previous results that write down the PDE (the four mean-field papers from last year) and it looks very similar to results that show gradient descent provably converges to global optima etc. I am not sure I see the importance of the result. Also, it is not clear *how long* SGD should run to reach such a point and this regime might be entirely out of interest. * Sec. 2.2 also discusses the difference in learning to teacher nodes with large or small weights, resp. Again, this is well-known in the asymptotic limit and rather unsurprising. This is also discussed in, e.g. in Saxe et al *Exact solutions to the nonlinear dynamics of learning in deep linear neural networks* at least in the linear case. * The extension to ResNets is definitely more interesting. The authors write down mean-field equations for two models, and prove, if I understand correctly, that it is a necessary condition to recover the teacher weights to generalise perfectly, which, as I said above, seems unsurprising. In the end, given the paper is not discussing the relevant literature in teacher-student setting, and that I (perphaps wrongly) do not find the results surprising enough, I would not support acceptance in ICLR.