Title: Leveraging the two-timescale regime to demonstrate convergence of neural networks
Venue: NeurIPS 2023
Accepted: yes
Keywords: neural networks

non-convex optimization

gradient flow

convergence proof

two-timescale algorithm
Abstract: We study the training dynamics of shallow neural networks, in a two-timescale regime in which the stepsizes for the inner layer are much smaller than those for the outer layer. In this regime, we prove convergence of the gradient flow to a global optimum of the non-convex optimization problem in a simple univariate setting. The number of neurons need not be asymptotically large for our result to hold, distinguishing our result from popular recent approaches such as the neural tangent kernel or mean-field regimes. Experimental illustration is provided, showing that the stochastic gradient descent behaves according to our description of the gradient flow and thus converges to a global optimum in the two-timescale regime, but can fail outside of this regime.

Leveraging the two-timescale regime to demonstrate convergence of neural networks

Abstract

We study the training dynamics of shallow neural networks, in a two-timescale regime in which the stepsizes for the inner layer are much smaller than those for the outer layer. In this regime, we prove convergence of the gradient flow to a global optimum of the non-convex optimization problem in a simple univariate setting. The number of neurons need not be asymptotically large for our result to hold, distinguishing our result from popular recent approaches such as the neural tangent kernel or mean-field regimes. Experimental illustration is provided, showing that the stochastic gradient descent behaves according to our description of the gradient flow and thus converges to a global optimum in the two-timescale regime, but can fail outside of this regime.

1 Introduction

Artificial neural networks are among the most successful modern machine learning methods, in particular because their non-linear parametrization provides a flexible way to implement feature learning (see, e.g., Goodfellow et al., 2016, chapter 15). Following this empirical success, a large body of work has been dedicated to understanding their theoretical properties, and in particular to analyzing the optimization algorithm used to tune their parameters. It usually consists in minimizing a loss function through stochastic gradient descent (SGD) or a variant (Bottou et al., 2018). However, the non-linearity of the parametrization implies that the loss function is non-convex, breaking the standard convexity assumption that ensures global convergence of gradient descent algorithms.

In this paper, we study the training dynamics of shallow neural networks, i.e., of the form

where m denotes the number of hidden neurons, a = ( a 0 , . . . , a m ) and u = ( u 1 , . . . , u m ) denote respectively the outer and inner layer parameters, and g ( x ; u ) denotes a non-linear function of x and u . The novelty of this work lies in the use of a so-called two-timescale regime (Borkar, 1997) to train the neural network: we set stepsizes for the inner layer u to be an order of magnitude smaller than the stepsizes of the outer layer a . This ratio is controlled by a parameter ε . In the regime ε ≪ 1 , the neural network can be thought of as a fitted linear regression with slowly evolving features

g ( x ; u j ) , j = 1 , . . . , m : this reduction enables us to precisely describe the movement of the inner layer parameters u j .

Our approach proves convergence of the gradient flow to a global optimum of the non-convex landscape with a fixed number m of neurons. The gradient flow can be seen as the simplifying yet insightful limit of the SGD dynamics as the stepsize h vanishes. Proving convergence with a fixed number of neurons contrasts with two other popular approaches that require to take the limit m →∞ : the neural tangent kernel (Jacot et al., 2018; Allen-Zhu et al., 2019; Du et al., 2019; Zou et al., 2020) and the mean-field approach (Chizat and Bach, 2018; Mei et al., 2018; Rotskoff and Vanden-Eijnden, 2018; Sirignano and Spiliopoulos, 2020). As a consequence, this paper is intended as a step towards understanding feature learning with a moderate number of neurons.

While our approach through the two-timescale regime is general, our description of the solution of the two-timescale dynamics and our convergence results are specific to a simple example showcasing the approach. More precisely, we consider univariate data x ∈ [0 , 1] and non-linearities of the form g ( x ; u j ) = σ ( η -1 ( x -u j )) , where u j is a variable translation parameter, η is a fixed dilatation parameter, and σ is a sigmoid-like non-linearity. Finally, we restrict ourselves to the approximation of piecewise constant functions.

Organization of this paper. In Section 2, we detail our setting and state our main theorem on the convergence of the gradient flow to a global optimum. Section 3 articulates this paper with related work. Section 4 provides a self-contained introduction to the two-timescale limit ε → 0 . We explain how it simplifies the analysis of neural networks, and provides heuristic predictions for the movement of neurons in our setting. Section 5 gives a rigorous derivation of our result. We prove convergence first in the two-timescale limit ε → 0 , then in the two-timescale regime with ε small but positive. Section 6 presents numerical experiments showing that the SGD dynamics follow closely those of the gradient flow in the two-timescale regime, and therefore exhibit convergence to a global optimum. On the contrary, SGD can fail to reach a global optimum outside of the two-timescale regime.

2 Setting and main result

We present a setting in which a piecewise constant univariate function f ∗ : [0 , 1] → R is learned with gradient flow on a shallow neural network. Our notations are summarized on Figure 1. We begin by introducing our class of functions of interest.

Definition 1. Let n ⩾ 2 , ∆ v ∈ (0 , 1) , ∆ f > 0 and M ⩾ 1 . We denote F n, ∆ v, ∆ f,M the class of functions f ∗ : [0 , 1] → R satisfying the following conditions:

f ∗ is piecewise constant: there exists

for all i ∈ { 1 , . . . , n } , v i -v i -1 ⩾ ∆ v ,

for all i ∈ { 1 , . . . , n -1 } , | f ∗ i -f ∗ i -1 | ⩾ ∆ f ,

for all i ∈ { 0 , . . . , n -1 } , | f ∗ i | ⩽ M .

f(-,a,u)

△f

f

f

0

In

I

1

U3

1

U2

I

1

Uu2

Let us now define our class of neural networks. Consider σ : R → R an increasing, twice continuously differentiable non-linearity such that σ ( x ) = 0 if x ⩽ -1 / 2 , σ ( x ) = 1 if x ⩾ 1 / 2 , and σ -1 / 2 is odd. Then, our class of shallow neural networks is defined by

where 0 < η ⩽ 1 measure the sharpness of the non-linearity σ η . Note that that inner layer parameter u j determines the translation of the non-linearity; no parameterized multiplicative operation on x is performed in this layer. We refer to the parameter u as the 'positions' of the neurons (or, sometimes, simply as the 'neurons') and to the parameter a as the 'weights' of the neurons. We define the quadratic loss as

We use gradient flow on L to fit the parameters a and u : they evolve according to the dynamics

where ε corresponds to the ratio of the stepsizes of the two iterations.

Main result. By leveraging the two-timescale regime where ε is small, our theorem shows that, with high probability, a neural network trained with gradient flow is able to recover an arbitrary piecewise constant function to an arbitrary precision. The proof is relegated to the Appendix.

Theorem 1. Let ξ, δ > 0 , and f ∗ a piecewise constant function from F n, ∆ v, ∆ f,M . Assume that the neural network has m neurons with

Assume that, at initialization, the positions u 1 , . . . , u m of the neurons are i.i.d. uniformly distributed on [0 , 1] and their weights a 0 , . . . , a m are equal to zero.

Then there exists Q 1 > 0 and Q 2 > 0 depending on ξ, δ, m, ∆ f, M such that, if

then, with probability at least 1 -δ , the solution to the gradient flow (1) is defined at least until T = 6 ε (∆ f ) 2 , and

Further, Q 1 = C 1 M 2 ( m +1) min ( δ 2 (∆ f ) 2 ( m +1) 4 , ξ ) and Q 2 = C 2 δ 2 M 4 ( m +1) 17 / 2 min ( δ (∆ f ) 2 m +1 , ξ ) for some universal constants C 1 , C 2 > 0 .

For this result to hold, the inequality (2) requires the number of neurons in the neural network to be large enough. Note that the minimum number of neurons required to approximate the n pieces of f ∗ is equal to n . If the length of all the intervals is of the same order of magnitude, then ∆ v = Θ(1 /n ) and thus the condition is m = Ω( n (1 + log n +log1 /δ )) . In this case, condition (2) only adds a logarithmic factor in n and δ . Moreover, the lower bound on m does not depend on the target precision ξ . Thus we observe some non-linear feature learning phenomenon: with a fixed number m of neurons, gradient flow on a neural network can approximate any element from the infinite-dimensional space of piecewise constant functions to an arbitrary precision.

The recovery result of Theorem 1 is provided under two conditions (3). The first one should not surprise the reader: the condition on η enables the non-linearity to be sharp enough in order to approximate well the jumps of the piecewise constant function f ∗ . The novelty of our work lies in the condition on ε , that we refer to as the two-timescale regime . This condition ensures that the stepsizes taken in the positions u are much smaller than the stepsizes taken in the weights a . As a consequence, the weights a are constantly close to the best linear fit given the current positions u . This property decouples the dynamics of the two layers of the neural network; this enables a sharp description of the gradient flow trajectories and thus the recovery result shown above. This intuition is detailed in Section 4.

3 Related work

Two-timescale regime. Systems with two timescales, or slow-fast systems , have a long history in physics and mathematics, see Berglund and Gentz (2006, Chapter 2) for an introduction. In particular, iterative algorithms with two timescales have been used in stochastic approximation and optimization, see Borkar (1997) or Borkar (2009, Section 6). For instance, they are used in the training of generative adversarial networks, to decouple the dynamics of the generator from those of the discriminator (Heusel et al., 2017), in reinforcement learning, to decouple the value function estimation from the temporal difference learning (Szepesvári, 2010), or more generally in bilevel optimization, to decouple the outer problem dynamics from the inner problem dynamics (Hong et al., 2023). However, to the best of our knowledge, the two-timescale regime has not been used to show convergence results for neural networks.

Layer-wise learning rates. Practitioners are interested in choosing learning rates that depend on the layer index to speed up training or improve performance. Using smaller learning rates for the first layers and higher learning rates for the last layer(s) improves performance for fine-tuning (Howard and Ruder, 2018; Ro and Choi, 2021) and is a common practice for transfer learning (see, e.g., Li et al., 2022). Another line of work proposes to update layer-wise learning rates depending on the norm of the gradients on each layer (Singh et al., 2015; You et al., 2017; Ko et al., 2022). However, they aim to compensate the differences across gradient norms in order to learn all the parameters at the same speed, while on the contrary we enjoy the theoretical benefits of learning different speeds.

Theory of neural networks. A key novelty of the analysis of this paper is that we show recovery with a fixed number of neurons. We now detail the comparison with other analyses.

The neural tangent kernel regime (Jacot et al., 2018; Allen-Zhu et al., 2019; Du et al., 2019; Zou et al., 2020) corresponds to small movements of the parameters of the neural network. In this case, the neural network can be linearized around its initial point, and thus behaves like a linear regression. However, in this regime, the neural network can approximate only a finite dimensional space of functions, and thus it is necessary to take m →∞ to be able to approximate the infinite-dimensional space of piecewise constant functions to an arbitrary precision.

The mean-field regime (Chizat and Bach, 2018; Mei et al., 2018; Rotskoff and Vanden-Eijnden, 2018; Sirignano and Spiliopoulos, 2020) describes the dynamics of two-layer neural networks in the regime m ≫ 1 through a partial differential equation on the density of neurons. This regime is able to describe some non-linear feature learning phenomena, but does not explained the observed behavior with a moderate number of neurons. In this paper, we show that in the two-timescale regime, only a single neuron aligns with each of the discontinuities of the function f ∗ . However, it should be noted that the neural tangent kernel and mean-field regimes have been applied to show recovery in a wide range of settings, while our work is restricted to the recovery of piecewise constant functions. Extending the application of the two-timescale regime is left for future work.

Our work includes a detailed analysis of the alignment of the positions of the neurons with the discontinuities of the target function f ∗ . This is analogous to a line of work (see, e.g., Saad and Solla (1995); Goldt et al. (2020); Veiga et al. (2022)) interested in the alignment of a 'student' neural network with the features of a 'teacher' neural network that generated the data, for high-dimensional Gaussian input. In general, the non-linear evolution equations describing this alignment are hard to study theoretically. On the contrary, thanks to the two-timescale regime and to the simple setting of this paper, we are able to give a precise description of the movement of the neurons.

Our study bears high-level similarities with the recent work of Safran et al. (2022). In a univariate classification setting, they show that a two-layer neural network achieves recovery with a number of neurons analogous to (2): inversely proportional to the length of the smallest constant interval of the target, up to logarithmic terms in the number of constant intervals and in the failure probability. However, the two papers have different settings: Safran et al. (2022) consider classification with ReLU activations while we consider regression with sigmoid-like activations. More importantly, the authors do not use the two-timescale regime. Instead, by a specific initialization scale, they ensure that the neural network has a first lazy phase where the positions of the neurons do not move significantly. For the second rich phase, they describe the implicit bias of the limiting point; this approach does not lead to an estimate of the convergence time while the fine description of the two-timescale limit does.

Finally, a related technique is the so-called layerwise training, which consists in first training the inner layer with the outer layer fixed, and then doing the reverse. This setup has been used in theoretical works to show convergence of (stochastic) gradient descent in a feature learning regime with a moderate number of neurons (Abbe et al., 2023; Damian et al., 2022). The two-timescale regime can be seen as a refinement of this technique since we allow both layers to move simultaneously instead of sequentially, which is closer to practical setups.

4 A non-rigorous introduction to the two-timescale limit

This section introduces the core ideas of our analysis in a non-rigorous way. Section 4.1 introduces the limit of the dynamics when ε → 0 , called the two-timescale limit . Section 4.2 applies the two-timescale limit to predict the movement of the neurons.

4.1 Introduction to the two-timescale limit

Let us consider the gradient flow equations (1) and perform the change of variables τ = εt :

In the two-timescale regime ε ≪ 1 , the rate of the gradient flow in the weights a is much larger than then the rate in the positions u . Note that L is marginally convex in a , and thus, for a fixed u , the gradient flow in a must converge to a global minimizer of a ↦→ L ( a, u ) . More precisely, assume that { σ η ( . -u 1 ) , . . . , σ η ( . -u m ) } forms an independent set of functions in L 2 ([0 , 1]) . Then the global minimizer of a ↦→ L ( a, u ) is unique; we denote it as a ∗ ( u ) . In the limit ε → 0 , we expect that a evolves sufficiently quickly with respect to u so that it converges instantaneously to a ∗ ( u ) . In other words, the gradient flow system (4) reduces to its so-called two-timescale limit when ε → 0 :

The two-timescale limit considerably simplifies the study of the gradient flow system because it substitutes the weights a , determined to be equal to a ∗ ( u ) . However, showing that (4) reduces to (5) requires some mathematical care, including checking that a ∗ ( u ) is well-defined.

Remark 1 (abuse of notation) . Equation (5) contains an ambiguous notation: does ∇ u L ( a ∗ ( u ) , u ) denote the gradient in u of the map L ∗ : u ↦→ L ( a ∗ ( u ) , u ) or the gradient ( ∇ u L )( a, u ) taken at a = a ∗ ( u ) ? In fact, by definition of a ∗ ( u ) , both quantities coincide:

where d a ∗ d u denotes the differential of a ∗ in u and we use that, by definition of a ∗ ( u ) , ( ∇ a L )( a ∗ ( u ) , u ) = 0 . This is a special case of the envelope theorem (Border, 2015, Sec. 5.10).

The discussion in this section is not specific to the setting of Section 2. Using the two-timescale limit to decouple the dynamics of the outer layer a and the inner layer u is a general tool that may be used in the study of any two-layer neural network. We chose the specific setting of this paper so that the two-timescale limit (5) can be easily studied, thereby showcasing the approach. The next section is devoted to a sketch of this study.

4.2 Sketch of the dynamics of the two-timescale limit

In this section, in order to simplify the exposition of the behavior of the two-timescale limit (5), we consider the limiting case η → 0 . Note that this is coherent with Theorem 1 that requires η to be small. This limit is a neural network with a non-linearity equal to the Heaviside function

Note that σ 0 would be a poor choice of non-linearity in practice: as its derivative is 0 almost everywhere, the positions u would not move. However, it is a relevant tool to get an intuition about the dynamics of our system for a small η . Moreover, as we will see in Section 6, the dynamics sketched here match closely those of the SGD (with η > 0 ).

f*

f(,a*(u),u)

f

f

0

I

U1

I

u

I

U1

R

U4U5

I

u

I

U2

I

R

1

U

The set { 1 , σ 0 ( . -u 1 ) , . . . , σ 0 ( . -u m ) } generates the space of functions that are piecewise constant with respect to the subdivision { u 1 , . . . , u m } . Furthermore, if u 1 , . . . , u m are distinct and in (0 , 1) , then this set is an independent set of functions in L 2 ([0 , 1]) . Thus a ∗ ( u ) is well defined and represents the coefficients of the best piecewise constant approximation of f ∗ with subdivision { u 1 , . . . , u m } .

This quantity is straightforward to describe under the mild additional assumption that there are at least two neurons u j in each interval ( v i -1 , v i ) between two points of discontinuity of f ∗ . For each 1 ⩽ i ⩽ n , let u L i denote the largest position of neurons below v i and u R i denote the smallest position above v i (with convention u R 0 = 0 and u L n +1 = 1 ). By assumption, 0 = u R 0 < u L 1 < u R 1 < · · · < u L n < u R n < u L n +1 = 1 are distinct. A simple computation then shows the following identities:

and for all i ∈ { 1 , . . . , n +1 } , for all x ∈ ( u R i -1 , u L i ) , f ( x ; a ∗ ( u ) , u ) = f ∗ i -1 ,

where we recall that f ∗ i denotes the value of f ∗ on the interval ( v i , v i +1 ) . Figure 2 illustrates the situation. Moreover, the loss L ( a ∗ ( u ) , u ) , which is half of the square L 2 -error of this optimal approximation, can be written

The dynamics of the two-timescale limit (5) corresponds to the local optimization of the subdivision u in order to minimize the loss (6). A remarkable property of this loss is that it decomposes as a sum of local losses around the jump points v i for i ∈ { 1 , . . . , n -1 } . Each element of the sum involves only the two neurons located at u L i and u R i . As a consequence, the dynamics of the two-timescale limit (5) decompose as n independent systems of two neurons u L i and u R i : for all i ∈ { 1 , . . . , n -1 } ,

All neurons other than u L 1 , u R 1 , . . . , u L n -1 , u R n -1 do not play a role in the expression (6), thus they do not move in the two-timescale limit (5). The position u L i moves right and u R i moves left, until one of them hits the point v i . This shows that the positions of the neurons eventually align with the jumps of the function f ∗ , and thus that the function f ∗ is recovered.

5 Convergence of the gradient flow

In this section, we give precise mathematical statements leading to the convergence of the gradient flow to a global optimum, first in the two-timescale limit ε → 0 , then in the two-timescale regime with ε small but positive. All proofs are relegated to the Appendix.

5.1 In the two-timescale limit

This section analyses rigorously the two-timescale limit (5), which we recall for convenience:

We start by giving a rigorous meaning to these equations. First, for L to be differentiable in u , we require the parameter η of the non-linearity to be positive. Second, for a ∗ ( u ) to be well-defined, we need u ↦→ L ( a, u ) to have a unique minimum. Obviously, if the u i are not distinct, then the features { σ η ( . -u 1 ) , . . . , σ η ( . -u m ) } are not independent and thus the minimum can not be unique. We restrict the state space of our dynamics to circumvent this issue. For u ∈ [0 , 1] m , we denote

̸

with the convention that u 0 = -η/ 2 and u m +1 = 1 + η/ 2 . Further, we define U = { u ∈ [0 , 1] m | ∆( u ) > 2 η } . The proposition below shows that U gives a good candidate for a set supporting solutions of (8).

Proposition 1. For u ∈ U , the Hessian H ( u ) of the quadratic function L ( ., u ) is positive definite and its smallest eigenvalue is greater than ∆( u ) / 8 . In particular, L ( ., u ) has a unique minimum a ∗ ( u ) .

The bound on the Hessian is useful in the following, in particular in the proof of the following result. Proposition 2. Let G ( u ) = ∇ u L ( a ∗ ( u ) , u ) for u ∈ U . Then G : U → R m is Lipschitz-continuous.

Then, the Picard-Lindelöf theorem (see, e.g., Luk, 2017 for a self-contained presentation and Arnold, 1992 for a textbook) guarantees, for any initialization u (0) ∈ U , the existence and unicity of a maximal solution of (8) taking values in U . This solution is defined on a maximal interval [0 , T max ) where it could be that T max < ∞ if u hits the boundary of U . However, the results below show that the target function f ∗ is recovered before this happens (with high probability over the initialization), and thus that this notion of solution is sufficient for our purposes. To this aim, we first define some sufficient conditions that the initialization should satify.

Definition 2. Let D be a positive real. We say that a vector of positions u ∈ [0 , 1] m is D -good if

(a) for all i ∈ { 0 , . . . , n -1 } , there are at least 6 positions u j in each interval [ v i , v i +1 ] ,

(b) ∆( u ) ⩾ D , and

(c) for all i ∈ { 1 , . . . , n -1 } , denoting u L i the position closest to the left of v i and u R i the position closest to the right, we have | u R i + u L i -2 v i | ⩾ D .

Condition (a) is related to the fact that the derivation in Section 4.2 is valid only if there are at least two neurons per piece. This requirement that the neurons be distributed on every piece of the target seems to be necessary for our result to hold, and we provide in Appendix D.1 a counter-example where recovery fails otherwise. Condition (b) indicates that the neurons have to be sufficiently spaced at initialization, which is not surprising since we have to guarantee that ∆( u ( τ )) > 2 η , that is, u ( τ ) ∈ U , for all τ until the recovery of f ∗ happens. Finally, condition (c) also helps to control the distance between neurons: although u L i and u R i move towards each other, as shown by (7), their distance can be controlled throughout the dynamics as a function of | u R i + u L i -2 v i | .

We can now state the Proposition showing the recovery in finite time. The proof resembles the sketch of Section 4.2 with additional technical details since we need to control the distance between neurons, and the fact that η > 0 makes the dynamics more delicate to describe.

Proposition 3. Let f ∗ ∈ F n, ∆ v, ∆ f,M . Assume that the initialization u (0) is D -good with D = 2 13 / 2 ( m +1) 1 / 2 Mη 1 / 2 (∆ f ) -1 . Then the maximal solution of (8) taking values in U is defined at least on [0 , T ] for T = 6 / (∆ f ) 2 , and at the end of this time interval, there is a neuron at distance less than η from each discontinuity of f ∗ .

This Proposition is the main building block to show recovery in the next Theorem, along with some high-probability bounds to ensure that an i.i.d. uniform initialization is D -good.

Theorem 2. Let ξ, δ > 0 , and f ∗ a piecewise constant function from F n, ∆ v, ∆ f,M . Assume that the neural network has m neurons with

Assume that, at initialization, the positions u 1 , . . . , u m of the neurons are i.i.d. uniformly distributed on [0 , 1] . Then there exists Q depending on ξ, δ, m, ∆ f, M such that, if

then, with probability at least 1 -δ , the maximal solution to the two-timescale limit (8) is defined at least until T = 6 (∆ f ) 2 , and

Furthermore, we have Q = C M 2 min ( δ 2 (∆ f ) 2 ( m +1) 5 , ξ n ) for some universal constant C > 0 .

5.2 From the two-timescale limit to the two-timescale regime

We now briefly explain how the proof for the two-timescale limit can be adapted for the gradient flow problem in the two-timescale regime (1), that is with a small but non-vanishing ε . First note that the existence and uniqueness of the maximal solution to the dynamics (1) follow from the local Lipschitz-continuity of ∇ a L and ∇ u L with respect to both their variables.

The heuristics of Section 4.1 indicate that, for ε small enough, at any time t , the weights a ( t ) are close to a ∗ ( u ( t )) , the global minimizer of L ( · , u ( t )) . The next Proposition formalizes this intuition.

Proposition 4. Assume that a (0) = 0 and that, for all s ∈ [0 , t ] , there at least 2 positions u j ( s ) in each interval [ v i , v i +1 ] and ∆( u ( s )) ⩾ D / 2 for some D ⩾ 32 η . Finally, assume that ε ⩽ 2 -16 D 2 M -2 ( m +1) -5 / 2 . Then

The crucial condition in the Proposition is ∆( u ( s )) ⩾ D / 2 ; it is useful to control the conditioning of the quadratic form L ( · , u ( s )) . The Proposition shows that ∥ a ( t ) -a ∗ ( u ( t )) ∥ is upper bounded by the sum of two terms; the first term is a consequence to the initial gap between a (0) and a ∗ ( u (0)) and decays exponentially quickly. The second term is negligible in the regime ε ≪ 1 .

Armed with this Proposition, we show that the two-timescale regime has the same behavior as the two-timescale limit and thereby prove Theorem 1.

6 Numerical experiments and discussion

Numerical illustration in the setting of Section 2. We first compare the dynamics of the gradient flow in the two-timescale limit presented in Section 4.2 with the dynamics of SGD. To simulate the SGD dynamics, we assume that we have access to noisy observations of the value of f ∗ ∈ F n, ∆ v, ∆ f,M : let ( X p , Y p ) p ⩾ 1 be i.i.d. random variables such that X p is uniformly distributed on [0 , 1] , and Y p = f ∗ ( X p ) + N p where N p is additive noise. The (one-pass) SGD updates are then given by

with ℓ ( X,Y ; a, u ) = 1 2 ( Y -f ( X ; a, u )) 2 . The experimental settings, as well as additional results, are given in the Appendix.

Remarkably, the dynamics of SGD in the two-timescale regime with η small match closely the gradient flow in the two-timescale limit with η = 0 , as illustrated in Figure 3. This validates the use

of the gradient flow to understand the training dynamics with SGD. Both dynamics are close until the two-timescale limit achieves perfect recovery of the target function, at which point the SGD stabilizes to a small non-zero error. The fact that SGD does not achieve perfect recovery is not surprising, since SGD is performed with η > 0 and f ∗ is not in the span of { 1 , σ η ( x -u 1 ) , . . . , σ η ( x -u m ) } for any u 1 , . . . , u m and for η > 0 . On the contrary, we simulated the dynamics of gradient flow for η = 0 , as presented in Section 4.2, enabling perfect recovery in that case.

(a) Loss L as a function of the number of steps p

(b) Plot of the functions after

p

= 2

.

7

·

10

7

steps

101

SGD

5

Target

Two-timescale limit

SGD

4

TT limit

100

3

10-1

2

10-2

1

0.0

0.5

1.0

1.5

0

p

×108

0.0

0.2

0.4

0.6

0.8

1.0

Next, we compare the SGD dynamics in the two-timescale regime ( ε ≪ 1 ) and outside of this regime ( ε ≈ 1 ). In Figure 4, we see that the network trained by SGD (in orange) in the two-timescale regime ε = 2 · 10 -5 , achieves near-perfect recovery. If we change ε to 1 , while keeping all other parameters equal, the algorithm fails to recover the target function (Figure 5). This shows that, in our setting with a moderate number of neurons m , recovery can fail away from the two-timescale regime. It could seem that we are favouring the two-timescale regime by running it for more steps. In fact, it is not the case since both regimes are run until convergence. We refer to Appendix C for details.

Note that the dynamics of the neurons in Figures 4 and 5 are different. In the two-timescale regime, only the neurons closest to a discontinuity move significantly, while the others do not. These dynamics correspond to the sketch of Section 4. Interestingly, it means that in this regime, the neural network learns a sparse representation of the target function, meaning that only n out of the m neurons are active after training. On the contrary, when ε = 1 , all neurons move to align with discontinuities of the target function, thus the learned representation is not sparse. Furthermore, since the number of neurons is moderate, one of the discontinuities is left without any neuron.

(a) At initialization

p

= 0

(b) After

p

= 5

.

4

·

10

6

steps

(c) After

p

= 1

.

8

·

10

8

steps

5

Target

L

5

SGD

4

4

3

3

II

1

3

2

2

2

1

1

1

0

11

1

0

0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

(a) At initialization

p

= 0

(b) After

p

= 10

4

steps

(c) After

p

= 10

6

steps

5

Target

5

5

SGD

4

4

4

3

3

3

2

2

2

1

0

11

11

1

0

1

0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

Discussion. The two-timescale regime decouples the dynamics of the two layers of the neural network. As a consequence, it is a useful theoretical tool to simplify the evolution equations. In this paper showcasing the approach, the two-timescale regime enables to show the alignment of the neurons with the discontinuities of the target function in the piecewise constant 1D case, and thus to prove recovery. A full general understanding of the impact of the two-timescale regime is an open question, which is left for future work. We provide in the following some practical evidence of the applicability of this regime to other settings (higher-dimensional problems, ReLU networks, finite sample size). Additional technical difficulties significantly complicate the proof in these settings, but we believe that there is no fundamental reason that our mathematical approach should not apply.

Higher dimensions. We consider piecewise constant functions on R d with pieces that are cuboids aligned with the axes of the space (see Figure 6). Neural networks are of the form f ( x ; a, u ) = a 0 + ∑ m j =1 ∑ d k =1 a jk σ η ( x k -u jk ) , where the j -th neuron has d -dimensional position ( u jk ) 1 ⩽ k ⩽ d and weight ( a jk ) 1 ⩽ k ⩽ d . The results are similar to the 1D case: convergence to a global minimum is obtained in the two-timescale regime but not in the standard regime.

(a) After training in the standard regime

(b) After training in the two-timescale regime

1.0

0.7

Target

0.5

Neurons

0.3

0.0

0.0

0.3

0.5

0.7

1.0

1.0

0.7

Target

0.5

Neurons

0.3

0.0

0.0

0.3

0.5

0.7

1.0

ReLU networks. Appendix C reports the case of using ReLU activations to approximate piecewiseaffine targets. A similar conclusion holds.

Finite sample size. We believe that it should be possible to generalize our results to finite sample sizes (say, for single-pass SGD), following a perturbation analysis similar to the one of Section 5.2, with additional terms due to random sampling. Numerically, Figure 3 shows that SGD indeed closely follows the behavior of the two-timescale limit. Some ideas about the proof are provided in Appendix D.2.

Acknowledgments and Disclosure of Funding

The authors thank Emmanuel Abbé for many discussions and suggestions on this project, and Quentin Berthet for pointing them towards the envelope theorem. The authors are grateful towards the Simons Institute for the Theory of Computing where this work was initiated. P.M. is supported by a grant from Région Île-de-France, by a Google PhD Fellowship and by Mines Paris - PSL.

References

Emmanuel Abbe, Enric Boix Adserà, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. In Gergely Neu and Lorenzo Rosasco, editors, Proceedings of Thirty Sixth Conference on Learning Theory , volume 195 of Proceedings of Machine Learning Research , pages 2552-2623. PMLR, 2023.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via overparameterization. In Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pages 242-252, 2019.

Vladimir Arnold. Ordinary Differential Equations . Springer Berlin Heidelberg, 1992.

Nils Berglund and Barbara Gentz. Noise-induced phenomena in slow-fast dynamical systems: a sample-paths approach . Springer Science & Business Media, 2006.

Kim Border. Miscellaneous notes on optimization theory and related topics, 2015. URL: https:// healy.econ.ohio-state.edu/kcb/AddedByPJ/Maximization.pdf (version: 2023-02-08).

Vivek Borkar. Stochastic approximation with two time scales. Systems & Control Letters , 29(5): 291-294, 1997.

Vivek Borkar. Stochastic approximation: a dynamical systems viewpoint , volume 48. Springer, 2009.

Léon Bottou, Frank Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM review , 60(2):223-311, 2018.

Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. Advances in Neural Information Processing Systems , 31, 2018.

Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Po-Ling Loh and Maxim Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory , volume 178 of Proceedings of Machine Learning Research , pages 5413-5452. PMLR, 2022.

Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pages 1675-1685, 2019.

Sebastian Goldt, Madhu Advani, Andrew Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. Journal of Statistical Mechanics: Theory and Experiment , 2020(12):124010, 2020.

Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning . MIT Press, 2016.

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems , 30, 2017.

Mingyi Hong, Hoi-To Wai, Zhaoran Wang, and Zhuoran Yang. A two-timescale stochastic algorithm framework for bilevel optimization: Complexity analysis and application to actor-critic. SIAM Journal on Optimization , 33(1):147-180, 2023.

Jeremy Howard and Sebastian Ruder. Universal language model fine-tuning for text classification. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 328-339, 2018.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in Neural Information Processing Systems , 31, 2018.

Yunyong Ko, Dongwon Lee, and Sang-Wook Kim. Not all layers are equal: A layer-wise adaptive approach toward large-scale DNN training. In Proceedings of the ACM Web Conference 2022 , page 1851-1859, 2022.

Fei-Fei Li, Jiajun Wu, and Ruohan Gao. Deep learning for computer vision course, 2022. URL: https://cs231n.github.io/transfer-learning/ (version: 2023-02-02).

Jonathan Luk. Notes on existence and uniqueness theorems for ODEs, 2017. URL: http://web. stanford.edu/~jluk/math63CMspring17/Existence.170408.pdf (version: 2023-01-19).

Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of twolayer neural networks. Proceedings of the National Academy of Sciences , 115(33):E7665-E7671, 2018.

Youngmin Ro and Jin Young Choi. AutoLR: Layer-wise pruning and auto-tuning of learning rates in fine-tuning of deep networks. Proceedings of the AAAI Conference on Artificial Intelligence , 35 (3):2486-2494, 2021.

Grant Rotskoff and Eric Vanden-Eijnden. Trainability and accuracy of neural networks: An interacting particle system approach. arXiv preprint arXiv:1805.00915 , 2018.

David Saad and Sara Solla. On-line learning in soft committee machines. Phys. Rev. E , 52:4225-4243, Oct 1995.

Itay Safran, Gal Vardi, and Jason D Lee. On the effective number of linear regions in shallow univariate relu networks: Convergence guarantees and implicit bias. Advances in Neural Information Processing Systems , 35, 2022.

Bharat Singh, Soham De, Yangmuzi Zhang, Thomas Goldstein, and Gavin Taylor. Layer-specific adaptive learning rates for deep networks. In 2015 IEEE 14th International Conference on Machine Learning and Applications (ICMLA) , pages 364-368, 2015.

Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: a central limit theorem. Stochastic Processes and their Applications , 130(3):1820-1852, 2020.

Csaba Szepesvári. Algorithms for reinforcement learning . Springer, 2010.

Rodrigo Veiga, Ludovic Stephan, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Phase diagram of stochastic gradient descent in high-dimensional two-layer neural networks. Advances in Neural Information Processing Systems , 35, 2022.

Yang You, Igor Gitman, and Boris Ginsburg. Scaling SGD batch size to 32k for imagenet training. arXiv preprint arXiv:1708.03888 , 2017.

Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes overparameterized deep ReLU networks. Machine Learning , 109(3):467-492, 2020.

Appendix

Organization of the Appendix In Section A, we introduce additional notations that will be used throughout the Appendix, then proceed to prove useful technical lemmas. We proceed in Section B to prove the results presented in the main text. Section C contains details about our experimental settings as well as some additional simulations. Section D gives a couple of additional insights about our results.

A Additional notations and technical lemmas

For a vector a , we denote ∥ a ∥ its ℓ 2 -norm, ∥ a ∥ 1 its ℓ 1 -norm and ∥ a ∥ ∞ its ℓ ∞ -norm. For matrices H , ∥ H ∥ denotes the operator norm associated to the ℓ 2 norm and ∥ H ∥ F denotes the Frobenius norm. For real-valued functions f , ∥ f ∥ ∞ denotes the supremum norm.

In all of the Appendix, we denote u 0 = -η/ 2 and u m +1 = 1 + η/ 2 . Note that σ η ( x -u 0 ) = 1 for all x ∈ [0 , 1] , meaning that σ η ( · -u 0 ) corresponds to the bias term. This notation allows to treat the bias term in a unified fashion with respect to the other terms of f ( x ; a, u ) . Since u i ∈ (0 , 1) for i ∈ { 1 , . . . , m } , we assume in the following w.l.o.g. that the ( u i ) 0 ⩽ i ⩽ m +1 are ordered in increasing order. Note that we prove in the following that the ( u i ) 1 ⩽ i ⩽ m do not cross during the dynamics, so they remain ordered throughout the dynamics.

The proofs involve comparisons of some quantities when η > 0 and when η = 0 . To avoid confusion, we make explicit the dependency of L on η ⩾ 0 , i.e., we let L η ( a, u ) in place of L ( a, u ) of the main paper, and similarly, when the arg min is well-defined and unique,

in place of a ∗ ( u ) . Similarly, we now make explicit the dependence of f on η ⩾ 0 , i.e., we denote

The Hessian of the quadratic function L η ( · , u ) is denoted H η ( u ) ∈ R ( m +1) × ( m +1) (in place of H ( u ) ), and satisties that, for i, j ∈ { 0 , . . . , m } ,

Also let, for η ⩾ 0 and u ∈ R m , b η ( u ) ∈ R m +1 such that, for j ∈ { 0 , . . . , m } ,

Finally, we let U η in place of U in the paper.

With these notations, we have, for η ⩾ 0 and a, u ∈ R m ,

and

We now move on to a series to lemmas that will be helpful in the proofs of Appendix B.

Lemma 1. For η ⩾ 0 and u ∈ R m , we have

Proof. For any j ∈ { 0 , . . . , m } ,

where in the last step we use that ∥ f ∗ ∥ ∞ ⩽ M and that σ η ( x ) = 0 for x ⩽ -η/ 2 , σ η ( x ) ∈ [0 , 1] for -η/ 2 < x < η/ 2 and σ η ( x ) = 1 for x ⩾ η/ 2 . Similarly,

Lemma 2. For η ⩾ 0 and u ∈ U η , H η ( u ) = H 0 ( u ) + D η , where D η is a diagonal matrix whose elements are independent of u and bounded in absolute value by η/ 2 .

Proof. Let i, j ∈ { 0 , . . . , m } , and denote c = max( u i , u j , 0) . Then

If i = j = 0 , max( u i , u j ) = -η/ 2 , and H η,ij ( u ) = 1 = H 0 ,ij ( u ) . If i = j = 0 ,

Note that the last integral is non-negative and less than 1 , hence | H η,ij ( u ) -H 0 ,ij ( u ) | ⩽ η/ 2 . Finally, if i = j , since | u i -u j | > η ,

̸

Furthermore, 0 < max( u i , u j ) < 1 -η 2 , thus

where the last equality comes from the oddness of σ -1 / 2 .

Lemma 3. For η > 0 , let a ∗ η : u ∈ U η ↦→ a ∗ η ( u ) . Then a ∗ η is differentiable and for any u ∈ U η ,

Proof. By Proposition 1 (whose proof does not rely on this lemma), for u ∈ U η , L η ( · , u ) has a unique minimizer a ∗ η ( u ) , which is equal to H η ( u ) -1 b η ( u ) by (11). Furthermore, H η and b η are differentiable with respect to u , hence a ∗ η is also differentiable with respect to u , and we have

̸

Denote w k ( u ) := ∂H η ∂u k ( u ) a ∗ η ( u ) and W ( u ) the ( m + 1) × ( m + 1) matrix formed by stacking column-wise the vectors ( w k ( u )) 0 ⩽ k ⩽ m . Then

We now estimate the Frobenius norm of the matrix W ( u ) . By Lemma 2, for u ∈ U η , H η ( u ) = H 0 ( u ) + D η . Take i, j ∈ { 0 , . . . , m } , then

̸

̸

Hence ∂H η,ij ∂u k = 0 if i, j = k . Further, if i = k and j = k ,

̸

Of course, the bound | ∂H η,ij ∂u k ( u ) | ⩽ 1 also holds when j = k and i = j . Finally, a similar bound shows that | ∂H η,ij ∂u k ( u ) | ⩽ 2 when i = j = k .

As a consequence, for k, i ∈ { 0 , . . . , m } ,

Thus

With a reasoning similar to the above, note that ∂b η ∂u ( u ) is a diagonal matrix with diagonal entries in [ -M,M ] . Finally, putting these elements together, using Proposition 1 and that ∥ W ( u ) ∥ ⩽ ∥ W ( u ) ∥ F, we obtain

The following lemma gives exact formulae for the derivative of the loss L η with respect to the positions of the neurons, evaluated for a = a ∗ 0 ( u ) , that is the best piecewise constant approximation of f ∗ with subdivision { u 1 , . . . , u m } . Note that the formulae are the same as in Section 4.2, but the derivation is slightly more intricate since we consider here the loss L η and not L 0 .

Lemma 4. Take η > 0 and u ∈ U η such that there are at least two neurons on each piece [ v i , v i +1 ] of f ∗ . Then, if u j does not flank a discontinuity of f ∗ ,

Furthermore, for a discontinuity v i , denote u L i is the closest neuron to its left and u R i the closest neuron to its right. If v i -u L i ⩾ η 2 and u R i -v i ⩾ η 2 , then

̸

Proof. In this proof, let us denote for simplicity a = a ∗ 0 ( u ) . At the condition that there is at least two neurons on each piece of f ∗ , Section 4.2 gives the optimal approximation f 0 ( x ; a, u ) of f ∗ that is piecewise constant with respect to the subdivision { u 1 , . . . , u m } . As a consequence, we easily get the value of a . Namely, if u j does not flank a discontinuity of f ∗ , the value of f 0 ( x ; a, u ) is locally constant around u j , thus a j = 0 . Plugging into (10), we obtain

Further, for a discontinuity v i , denote respectively a L i and a R i the coefficients associated to u L i and u R i . At u L i , the value of f 0 ( x ; a, u ) jumps from f ∗ i -1 to v i -u L i u R i -u L i f ∗ i -1 + u R i -v i u R i -u L i f ∗ i , thus

Similarly, we have

We now compute, using (10),

Using that ∆( u ) > 2 η and that there are at least two neurons on each piece of f ∗ , we have that u L i -v i -1 ⩾ 2 η . Since, in addition, by assumption, v i -u L i ⩾ η 2 , we get that for x ∈ [ u L i -η 2 , u L i + η 2 ] , f ∗ ( x ) = f ∗ i -1 . Moreover, using again ∆( u ) ⩾ 2 η that σ η is equal to σ 0 on ( -∞ , -η/ 2] and [ η/ 2 , ∞ ) , we have for x ∈ [ u L i -η 2 , u L i + η 2 ] ,

Thus we obtain

The computation of ∂L η ∂u R i ( a, u ) is similar.

Lemma 5. Consider η ⩾ 0 and u ∈ U η such that there are at least two neurons on each piece [ v i , v i +1 ] of f ∗ . Then, for all x ∈ [0 , 1] , | f η ( x ; a ∗ 0 ( u ) , u ) | ⩽ M .

Proof. In the case where η = 0 , the result easily follows from the expressions for f 0 ( x ; a ∗ 0 ( u ) , u ) provided in Section 4.2. We now assume η > 0 .

Denote A ∗ k ( u ) = ∑ k j =0 a ∗ 0 ,j ( u ) (with the convention A ∗ -1 ( u ) = 0 ). Recall the convention u 0 = -η/ 2 . We compute

Note that A ∗ k ( u ) = lim x → u k + f 0 ( x ; a ∗ 0 ( u ) , u ) , and thus, from the case η = 0 , we have | A ∗ k ( u ) | ⩽ M . Moreover, σ η is increasing and the u k are in increasing order. We thus get

Lemma 6. Consider η > 0 and u ∈ U η such that there are at least two neurons on each piece [ v i , v i +1 ] of f ∗ . Then, for j ∈ { 0 , . . . , m } ,

and, for any a ∈ R m +1 ,

Proof. The first statement of the Lemma comes from the explicit formulae for a ∗ 0 ( u ) given in the proof of Lemma 4, namely each a ∗ 0 ,j ( u ) is either zero or less in magnitude than the gap between two pieces of f ∗ that is less than 2 M .

By (10), we have

We bound the two terms separately. For the first term, we use Lemma 5.

We now continue with the second term.

and thus

Returning to our initial upper bound, we obtain, using the first statement of the Lemma,

Lemma 7. For η ⩾ 0 and u ∈ U η ,

Proof. By (11), and by (11) and by Lemma 2,

According to Proposition 1 (whose proof does not rely on this lemma), H η ( u ) is invertible with ∥ H η ( u ) -1 ∥ ⩽ 8 / ∆( u ) . We thus have

The result then unfolds from Lemmas 1 and 6.

Lemma 8. Let η > 0 , u ∈ R m and a, a ′ ∈ R m +1 . Then

As a consequence of the second inequality, by the fundamental theorem of calculus for line integrals,

Proof. Recall that, for all j ∈ { 1 , . . . , m } , and for all a, u ∈ R m ,

From the first equality, we have

As a consequence,

Similarly, from the second equality, we have

As a consequence,

Lemma 9. Consider η ⩾ 0 and u ∈ U η such that there is a neuron at distance less than η from each discontinuity of f ∗ and 3 η ⩽ ∆ v . Then

Proof. By definition of a ∗ η ( u ) ,

Thus it is enough to exhibit some a for which the latter integral is smaller than 6 M 2 ηn to conclude.

We construct such an a as follows: set a 0 = f ∗ (0) , and for each discontinuity v i , set the coefficient of a neuron at distance less than η to the value f ∗ i -f ∗ i -1 and set all other neurons to zero. Note that the active neurons are distinct since 3 η ⩽ ∆ v .

Then the neural network is equal to the target function everywhere except on an interval of size 3 η/ 2 around each discontinuity, where they disagree (in infinite norm) by at most 2 M .

Lemma 10. Let m be a positive integer and u 1 , . . . , u m be i.i.d. uniform random variables in [0 , 1] . Assume that

Then, with probability at least 1 -δ , the vector u is D -good with D = δ 6( m +1) 2 .

Proof. We define the following events:

(a) A is the event 'there are at least 6 positions u j in each interval [ v i , v i +1 ] for i ∈ { 0 , . . . , n -1 } ',

(b) B is the event ' ∆( u ) ⩾ D ',

(c) for all i ∈ { 1 , . . . , n -1 } , E i is the event 'there are at least one neuron on the left and on the right of v i ' and C i is the event ' E i holds and | u R i + u L i -2 v i | ⩾ D '.

Note that by Definition 2, u is D -good if and only if the event A ∩ B ∩ ( ⋂ i C i ) holds. To show that this holds with high probability, we bound the probability of the complement

Thus

Below, we bound separately the three terms of the right hand side.

(a) Denote m ′ = ⌊ m/ 6 ⌋ . For any i ∈ { 0 , . . . , n -1 } , the set A i = { j ∈ { 1 , . . . , m ′ } | u j ∈ [ v i , v i +1 ] } is empty with probability (1 -( v i +1 -v i )) m ′ ⩽ (1 -∆ v ) m ′ . Thus by the union bound, the probability that at least one of A 1 , . . . , A n is empty is upper bounded by n (1 -∆ v ) m ′ .

We now check that n (1 -∆ v ) m ′ ⩽ δ/ 18 . Indeed,

where we use ∆ v ⩽ 1 , 3 ⩾ log(18) , and log(1 -∆ v ) ⩽ -∆ v < 0 . This gives the desired inequality.

In other words, the probability that at least one of the intervals [ v i , v i +1 ] contains none of the u 1 , . . . , u m ′ is bounded by δ/ 18 . As a consequence, by the union bound, the probability that at least one of the intervals [ v i , v i +1 ] contains strictly less than 6 of the u 1 , . . . , u m is bounded by δ/ 3 , i.e., P ( A c ) ⩽ δ/ 3 .

̸

(b) Recall that by convention, u 0 = -η 2 and u m +1 = 1+ η 2 . For all i ∈ { 0 , . . . , m +1 } , denote I i = ( u i -D,u i + D ) . Denote F j the event ' u j ∈ I i for some i ∈ { 0 , . . . , m +1 } , i = j '. Note that B c = ∪ j m =1 F j .

̸

Fix j = 1 , . . . , m . By conditioning on u i for all i ∈ { 0 , . . . , m +1 } , i = j , we see that P ( F j ) ⩽ 2( m +1) D . By the union bound,

(c) Take i ∈ { 1 , . . . , n -1 } . For convenience, we define the random variable u L i (resp. u R i ) on the full probability space by setting u L i = 0 (resp. u R i = 1 ) when there is no neuron on the left (resp. the right) of v i . We compute the joint cumulative distribution function of ( u L i , u R i ) (with a convenient change of inequality): for all 0 ⩽ y ⩽ v i ⩽ z ⩽ 1 ,

We observe that the joint cumulative distribution function of ( u L i , u R i ) is a smooth function of ( y, z ) when ( y, z ) ∈ (0 , v i ) × ( v i , 1) . Note that the events E i and { ( u L i , u R i ) ∈ (0 , v i ) × ( v i , 1) } are equal up to a null set. Therefore, on this event, ( u L i , u R i ) is an absolutely continuous random variable with density g : (0 , v i ) × ( v i , 1) → R ,

We compute

We make the change of variables θ = z -y , ν = z + y .

Using m ⩾ 24 / ∆ v ⩾ 24 n , we have

This concludes the proof.

B Proofs of the results of the main text

B.1 Proof of Proposition 1

Let us lower-bound the smallest eigenvalue of H η ( u ) which is equal to

Now for a ∈ R m +1 such that ∥ a ∥ = 1 ,

Since ∆ u > 2 η (because u ∈ U ) and u 0 = -η / 2 , u m +1 = 1+ η / 2 , the intervals [ u i + η/ 2 , u i +1 -η/ 2] for i ∈ { 0 , . . . , m } are disjoint and included in [0 , 1] . Thus

Since σ ( x ) = 0 if x < -1 / 2 and σ ( x ) = 1 if x > 1 / 2 , we have that σ η ( x ) = 0 if x < -η/ 2 and σ η ( x ) = 1 if x > η/ 2 . Further recall that the u i are ordered in increasing order. As a consequence,

where in the last step, we used that ∆( u ) > 2 η and thus u i +1 -u i -η ⩾ ∆( u ) -η ⩾ ∆( u ) -∆( u ) / 2 = ∆( u ) / 2 . Now, denote c 0 = 0 and c i = ∑ i -1 k =0 a k . Then ∥ a ∥ = 1 writes

Furthermore,

Hence

which shows in conjunction with (12) that the smallest eigenvalue of H η ( u ) is lower-bounded by ∆ u 8 .

B.2 Proof of Proposition 2

To show that G ( u ) = ( ∇ u L η )( a ∗ η ( u ) , u ) is Lipschitz-continuous on U η , we show that it is differentiable on U η and that its derivatives are uniformly bounded. The chain rule gives

From (10), using that σ is twice continuously differentiable, it can be checked that ∂L η ∂u j is differentiable in both its arguments and its derivatives are uniformly upper-bounded when a is bounded. Furthermore, for u ∈ U η ,

by Lemma 1 and Proposition 1. Finally, according to Lemma 3, a ∗ η is differentiable with derivatives uniformly upper-bounded on U η . This concludes the proof.

B.3 Proof of Proposition 3

In this proof, we denote u L i ( τ ) (resp. u R i ( τ ) ) the position at time τ of the neuron that is at initialization closest to v i to the left (resp. the right). Note that because of the movement of the neurons, it could be that u L i (resp. u R i ) does not remain the neuron closest to the left (resp. the right) throughout the dynamics. Our proof discusses when this phenomenon occurs. Similarly, denote u LL i (resp. u RR i ) the neuron second closest to the left (resp. the right) of v i . Since the initialization is D -good, note that all these neurons are distinct.

Denote T the minimal time τ ∈ [0 , T max ) such that ∆( u ( τ )) ⩽ D / 2 or there are less than two neurons in some piece [ v i , v i +1 ] of f ∗ . Note that by assumption, ∆( u (0)) ⩾ D > D/ 2 and there are at least 6 neurons in each interval at initialization, thus T > 0 . Furthermore, using the trivial inequalities M ⩾ ∆ f/ 2 , m + 1 ⩾ 1 and η 1 / 2 ⩾ η , we have D 2 = 2 11 / 2 M √ m +1 √ η ∆ f ⩾ 8 η > 2 η . Recall that 2 η is the quantity defining the set U η supporting the maximal solution of the equation (8). As a consequence, we do have T < T max . At the end of the proof, we check that T < T , by controlling carefully the movement of each neuron.

Let us first bound the difference between the dynamics of u and the dynamics that we would have if at each time τ , the weights a were given by a ∗ 0 ( u ( τ )) , the best approximation of f ∗ by a piecewise constant function with subdivision u ( τ ) . For any τ < T and j ∈ { 1 , . . . , m } , by Lemma 6, we have

We are therefore led to bounding ∥ a ∗ η ( u ( τ )) -a ∗ 0 ( u ( τ )) ∥ , as follows:

Then the first term in (13) is less than

and the second term in (13) is less than

Hence we obtain, for any τ < T and j ∈ { 1 , . . . , m } ,

Now, let us examine how the neurons move, by leveraging Lemma 4 that gives exact formulae for ∂L η ∂u j ( a ∗ 0 ( u ( τ )) , u ( τ )) . First, if u j is not next to a discontinuity, ∂L η ∂u j ( a ∗ 0 ( u ( τ )) , u ( τ )) = 0 , hence

Let us now study what happens next to a discontinuity v i . Denote ( δf ) i = f ∗ i -f ∗ i -1 . W.l.o.g., assume that

In the reverse case, the proof is the same by swapping the roles of u L i and u R i , and of u LL i and u RR i . We are going to show that the dynamics of u L i are divided into two phases. Define T i as the minimal τ ∈ [0 , T ] such that u L i ( τ ) = v i -η/ 2 . In the first phase [0 , T i ] , we have u L i ( τ ) < v i -η 2 and u L i moves towards v i . In the second phase [ T i , T ] , we show below that u L i ( τ ) ∈ [ v i -η, v i + η ] . Note that we can have T i = 0 if u L i (0) ⩾ v i -η 2 . It is also possible that T i = ∞ a priori; this means that the second phase does not exist. We show below that this case does not happen. Figure 7 depicts the two phases.

(a) Phase

1

(b) Phase

2

a,u

In

LL

u

L

Ui

urui

RRR

U6U7

U1

1

u

LL

Viu'

L

Ui

R

RR

U6U7

Beginning by the first phase, we have, while u L i ( τ ) < v i -η 2 and u R i ( τ ) > v i + η 2 , according to Lemma 4,

For ease of computation, let d L i ( τ ) = v i -u L i ( τ ) and d R i ( τ ) = u R i ( τ ) -v i be the distances between the neurons and v i . Then, by (14),

since D ⩽ ∆( u (0)) ⩽ 1 . Thus, in some time less than T = 6 (∆ f ) 2 , d R i ( τ ) + d L i ( τ ) ⩽ η , that is, either u L i reaches v i -η 2 or u R i reaches v i + η 2 . Let us check that the second event cannot actually

happen: while u L i ( τ ) < v i -η 2 and u R i ( τ ) > v i + η 2 , we also have

By condition (c) of Definition 2 and by (14), we have d R i (0) -d L i (0) ⩾ D = 60∆ g (∆ f ) 2 ⩾ 60∆ g ( δf ) 2 i , and furthermore d L i ( τ ) + d R i ( τ ) ⩽ 1 . An easy reasoning then shows that d R i -d L i is increasing. Therefore u R i must remain further away from v i than u L i .

In summary, we showed that there exists some time T i ⩽ T when u L i ( T i ) = v i -η 2 , which marks the end of the first phase, and we also have

Moving on to the study of the second phase, let us show that u L i ( τ ) stays in the interval [ v i -η, v i + η ] for τ ∈ [ T i , T ) . Consider any τ ⩽ T such that u L i ( τ ) = v i -η . Then we have by (14) and Lemma 4

where the second upper bound comes from the fact that we have u R i ( τ ) -v i ⩾ D 2 -η since ∆( u ( τ )) ⩾ D / 2 , and furthermore, x ↦→ x 2 ( x + η ) 2 is increasing, hence

Equation (15) implies that u L i ( τ ) ⩾ v i -η for all τ ∈ [ T i , T ) . Similarly, consider any τ ⩽ T such that u L i ( τ ) = v i + η . Note that, for such a τ , u L i ( τ ) is now on the right of v i , and the neurons flanking v i are u LL i and u L i . Thus we have by (14) and Lemma 4

where the second lower bound unfolds similarly as previously. This shows that u L i ( τ ) ⩽ v i + η for all τ ∈ [ T i , T ) .

We now check that T < T , that is, for all τ ⩽ T , ∆( u ( τ )) > D/ 2 and there are at least two neurons in each interval [ v i , v i +1 ] . Starting with the first condition, we say that neurons u j and u k collide if | u j ( τ ) -u k ( τ ) | = D/ 2 for some τ ⩽ T . Let us show that no pair of neurons collide.

We start by showing that there is no collision between u LL i and u L i . In the first phase [0 , T i ] , we have d u LL i d τ ( τ ) ⩽ ∆ g . Recall that we also have d u L i d τ ( τ ) ⩾ -∆ g and thus for τ ⩽ T i ,

since u L i (0) -u LL i (0) ⩾ D and T ∆ g = D / 10 by definition of T and ∆ g . As a consequence, u LL i and u L i do not collide during the first phase, and we have

In the second phase, we can have u L i ∈ [ v i , v i + η ] in which case u LL i becomes the neuron flanking v i to the left and u L i the neuron flanking to the right. Then (14) and Lemma 4 give

Of course, this bound also holds when u L i ∈ [ v i -η, v i ] , because then d u LL i d τ ⩽ ∆ g . Thus, in the second phase τ ∈ [ T i , T ] , by the previous upperbound and the fact that u L i ( τ ) ⩾ v i -η 2 ,

by (16). Let us now upper-bound each of the last two terms by D / 10 to conclude. By definition of D ,

Thus using the definition of T , D ⩽ ∆( u (0)) ⩽ 1 , ∆ f ⩽ 2 M and m +1 ⩾ 1 . Finally, T ∆ g = D / 10 . Thus u LL i and u L i do not collide.

We now show that u L i and u R i do not collide. In the first phase τ ∈ [0 , T i ] , we have

As a consequence, u L i and u R i do not collide during the first phase, and we have

In the second phase, u R i plays a role symmetric to u LL i : it can be, or not, the neuron closest to the right of v i , depending on whether u L i ∈ [ v i -η, v i ] or u L i ∈ [ v i , v i + η ] . As for u LL i , we can show that in any case, for τ ∈ [ T i , T ] ,

Thus one concludes as before: for τ ∈ [ T i , T ] , by the previous lowerbound and the fact that u L i ( τ ) ⩽ v i + η 2 ,

Then, by (17),

where the last lower-bound unfolds similarly as for u LL i and u L i . Thus there is no collision between u L i and u R i .

The reader can check that all other pairs of neurons do not collide, including those involving u 0 = -η/ 2 and u m +1 = 1 + η/ 2 . In fact, the proof is easier than for u LL i , u L i and u L i , u R i because the discontinuity at v i attracts these neurons together.

Furthermore, we proved that before time T at most one neuron can escape on each side of a piece [ v i , v i +1 ] of f . Since we start with at least four (and even six) neurons per piece, there is always before T at least two neurons per piece.

This shows that T < T , and we also proved that at time T , all discontinuities have finished their first phase, hence there is a neuron at distance less than η from each discontinuity of the target function.

B.4 Proof of Theorem 2

Take C = 2 -19 . Then by assumption of Theorem 2,

Moreover, by the definition of D from Proposition 3,

Then Lemma 10 shows that the initialization is D -good with probability at least 1 -δ (since the D -good property is monotonous in D ).

Hence, with probability at least 1 -δ , according to Proposition 3, the maximal solution to (8) is defined at least until T and at that time, there is a neuron at distance less than η from each discontinuity of the target function. Furthermore, 3 η ⩽ 1 m +1 ⩽ 1 n ⩽ ∆ v , hence Lemma 9 applies. This implies that

The assumption on η allow to conclude that the upper-bound is less than ξ .

Remark 2. We did not try to optimize the value of C since our goal was to show convergence to a global optimum and the dependency of the dynamics on the parameters (for instance, it is remarkable that T does not depend on ξ ).

B.5 Proof of Proposition 4

For s ⩽ t , Proposition 1 holds since for ∆( u ( s )) ⩾ 16 η > 2 η . Thus a ∗ η ( u ( s )) is well-defined and verifies

Let, for s ⩽ t , V ( s ) = ∥ a ( s ) -a ∗ η ( u ( s )) ∥ . Recall that, by (11),

This implies that and in consequence

Hence, for s ⩽ t ,

where the first lower bound is a consequence of Proposition 1. Then we have, for any s ⩽ t ,

Let us now upper bound the norm appearing in the second term. We first have by the chain rule

By Lemma 3 (which holds since for ∆( u ( s )) ⩾ 16 η > 2 η ),

Besides,

Furthermore,

By Lemmas 6 and 7, which apply since ∆( u ( s )) > 2 η and since there are at least two positions u j ( s ) in each interval [ v i , v i +1 ] for s ⩽ t ,

where the last upper bound is implied by the assumption D ⩾ 32 η .

Now define T max = inf { s ⩾ 0 , V ( s ) > 3 M √ m +1 } and assume s ⩽ min( t, T max ) so that V ( s ) ⩽ 3 M √ m +1 . Then we proved that ∥ a ( s ) ∥ ⩽ 6 M √ m +1 . Going back to (18), we deduce that

Putting everything together, we obtain

All in all,

Hence

By Grönwall's inequality, for all s ⩽ min( t, T max ) ,

Finally note that V (0) = ∥ a ∗ η (0) ∥ ⩽ 2 M √ m +1 and 2 17 M 3 ( m +1) 3 ε D 2 ⩽ 2 M √ m +1 by the assumption of the Proposition on ε . Hence (20) implies that for all s ⩽ min( t, T max ) , V ( s ) is a (weighted) average of two terms less than 2 M √ m +1 hence it is less than 2 M √ m +1 . This shows that T max ⩾ t , which concludes the proof since (21) is then valid for s = t .

B.6 Proof of Theorem 1

In the proof, we take C 1 = 2 -21 and C 2 = 2 -36 . Denote

Lemma 10 shows that the initialization is D -good with probability at least 1 -δ . In the following, we study the case where this event happens.

Denote T the minimal time t > 0 such that ∆( u ( t )) ⩽ D / 2 or there are less than two neurons in some piece [ v i , v i +1 ] of f ∗ or ∥ a ( t ) ∥ > 7 M √ m +1 . Note that T > 0 since the initialization is D -good. By Lemma 8, ∇ u L η and ∇ a L η are Lipschitz-continuous on compacts, hence the solution of the gradient flow is well defined for t < T since T defines a compact set of parameters.

Then all the assumptions of Proposition 4 are satisfied on the time interval [0 , t ] for any t < T . More precisely, the assumptions that do not come directly from the definition of T are the lower bound for D and the upper bound for ε . The lower bound for D come from

by (3) and the simple bounds δ ⩽ 1 , ∆ f ⩽ 2 M , m +1 ⩾ 1 . The upper bound for ε comes from (3) since

where the second upper bound uses m ⩾ 0 , δ ⩽ 1 and ∆ f ⩽ 2 M . Therefore, according to Proposition 4,

Furthermore, the proof of Proposition 4 actually implies that

The second bound implies that the condition ∥ a ( t ) ∥ > 7 M √ m +1 in the definition of T is actually never active. Let us distinguish between two phases: letting

then the first phase corresponds to t ⩽ T 0 and the second phase for t ⩾ T 0 .

Analysis of the first phase. In the first phase, each neuron moves at most by

where the second upper bound comes from (19) in the proof of Proposition 4. This quantity is less than D 8 if

Let us check this condition: we have

since m +1 ⩾ 1 , δ ⩽ 1 , and 2 M/ ∆ f ⩾ 1 , hence (2 M/ ∆ f ) 1 / 8 ⩽ 2 M/ ∆ f . Next, upper-bounding log( x ) by x , we have, by (3),

using ∆ f ⩽ 2 M and m ⩾ 0 . Note that the upper bound 2 6 εT 0 M 2 ( m +1) 3 / 2 ⩽ D / 8 also implies that

since m ⩾ 0 , D ⩽ 1 and ∆ f ⩽ 2 M . Since each neuron moves by at most D/ 8 in the time interval [0 , T 0 ] and since ∆( u (0)) ⩾ D , we deduce that

Similarly, by condition (c) of the definition of a D -good vector, for all discontinuities v i ,

Furthermore, there at least four neurons on each piece of f at T 0 , because at most one neuron can move out of each piece by either side between 0 and T 0 .

Analysis of the second phase. Let

In the second phase t ⩾ T 0 , we are able to control by ∆ a the distance between a ( t ) and the weights a ∗ 0 ( u ( t )) that are the best approximation of f ∗ by a piecewise affine function with subdivision u ( t ) . To show this, first note that the first term in (23) is smaller than ∆ a 4 when

which is equivalent to which is implied by t ⩾ T 0 . Furthermore, the second term in (23) is smaller than ∆ a 4 because, by definition of D and by (3),

Hence, for all T 0 ⩽ t < T , thus

Furthermore, note that the assumption of Lemma 7 applies for t < T since ∆( u ( t )) ⩾ D 2 > 2 η by (22). Therefore, by Lemma 7 and by (3),

By the triangular inequality, we deduce the upper bound that we were after, that is

As in the proof of Proposition 3, we can now control the distance between the true dynamics and the one that we would have if the weights were equal to a ∗ 0 ( u ) . Namely, for any T 0 ⩽ t ⩽ T and j ∈ { 1 , . . . , m } , by Lemma 6 (which applies since ∆( u ( t )) > 2 η by (22)), we have

The first term is less than

and the second term is less than

using D ⩽ ∆( u (0)) ⩽ 1 , ∆ f ⩽ 2 M and m +1 ⩾ 1 . Hence we obtain, for any T 0 ⩽ t ⩽ T and j ∈ { 1 , . . . , m } ,

We are therefore in a situation very similar to the proof of Proposition 3, starting from (14). One can check that all the arguments used in the proof also apply here. On top of the estimate above that resembles (14), the crucial facts that make the argument of Proposition 3 work here are the bounds (26) and (27) as well as the fact that there are at least four neurons on each piece f at T 0 , which together are very similar to the conditions ensuring that u (0) is D -good in the proof of Proposition 3. Another key point is (25), ensuring that a time at least equal to 11 T/ 12 remains after the first phase of this proof, which is enough time for the dynamics described in the proof of Proposition 3 to unfold.

This yields that T < T , and that at time T , there is a neuron at distance less than η from each discontinuity of f ∗ . Furthermore, 3 η ⩽ 1 m +1 ⩽ 1 n ⩽ ∆ v , hence Lemma 9 applies. Thus

where the second upper bound comes from n ⩽ m +1 and from (3). Furthermore, by (24) and by Lemma 8,

Let us show that this term is less than ξ / 4 . Recall that, by (23),

By definition of D and T , by using exp( -x ) ⩽ 1 / x for x ⩾ 1 and by (3),

using ∆ f ⩽ 2 M , δ ⩽ 1 , and m +1 ⩾ 1 . Furthermore, by (3), we get that

We therefore obtain the sought ξ / 4 upper-bound and can conclude that

C Experimental details and additional experiments

Our code is available at https://github.com/PierreMarion23/two-timescale-nn .

Numerical illustration in the setting of Section 2. To obtain Figures 3 and 4, we use the parameters of Table 1. For Figure 5, we use the parameters of Table 2.

Our target function is defined by f ∗ = 1 on [0 ., 0 . 2] , [0 . 35 , 0 . 5] , [0 . 65 , 0 . 8] , f ∗ = 2 on [0 . 5 , 0 . 65] and f ∗ = 4 elsewhere. The activation function is a piecewise cubic polynomial defined by x ↦→ min(max(4( x +0 . 5) 3 , 0) , 0 . 5) + min(max(0 . 5 -4(0 . 5 -x ) 3 , 0) , 0 . 5) .

We re-run the same SGD experiment as above twenty times, and plot in Figure 8 the average L 2 distance to the target as a function of ε , averaging over the initialization randomness and SGD randomness. This confirms that, in our setting, the SGD is able to recover the target function in the two-timescale regime ( ε ≪ 1 ), but fails outside of the two-timescale regime ( ε = 1 ). The transition between the two regimes seems to occur for ε ≈ 0 . 1 .

The number of iterations in Table 1 is much larger than the one in Table 2. There are two levels of analysis to explain this difference. The most straightforward reason is that the positions u evolve at a speed εh , which is much smaller in Table 1. However, one may note that it is not larger by a

100

error

10-1

10-2

0.01

0.1

1.0

E

factor 1 /ε = 50 , 000 but 'only' by a factor ≃ 200 . To understand why, we have to get more into the details of the dynamics, to understand the order of magnitude of the number of steps required before convergence. In the two-timescale regime, the limiting factor for convergence is the movement of the positions u . At each step, they move by an order of εh ≃ 10 -10 . The positions need to move on a scale of 5 · 10 -2 to align with the discontinuities of the target. Hence the required number of steps is 0 . 05 / (2 · 10 -10 ) ≃ 2 · 10 8 . In the standard regime, on the contrary, the limiting factor for convergence is the movement of the weights a . They move by an order of h ≃ 10 -5 at each step, and they need to move by a distance of ≃ 5 , necessitating ≃ 5 · 10 5 steps to achieve convergence.

Finally, note that it is possible to increase h in Table 1 while keeping the same behavior (in our experiment, h is kept to the same value as in Table 2 in order to facilitate the comparison). More precisely, taking h = 10 -3 in Table 1 yields similar results while dividing the computational cost by 100 .

Higher dimensions. In 2D, we use the parameters of Table 3 and 4. The positions are initialized uniformly over [0 , 1] and the weights uniformly over [0 , 3] . The target is

The activation function is the same as in the one-dimensional case. We use SGD with batch size 10 3 .

As in the one-dimensional case, we re-run the experiment twenty times and report the results in Figure 9a.

We perform a similar experiment in 10D, using the same setup. The target is

We use SGD with batch size 10 4 . The results in terms of L 2 distance are reported in Figure 9b.

(a) In 2D

(b) In 10D

0.025

0.020

0.020

0.015

0.015

error

error

0.010

0.010

0.005

0.005

0.000

0.000

0.01

1.0

0.01

1.0

E

E

ReLU networks. The goal of this experiment is to use ReLU activations to approximate piecewiseaffine targets in 1D. The target is a ReLU network with positions u = [0 . 3 , 0 . 5 , 0 . 7] and weights a = [1 . 0 , -2 . 0 , 3 . 0] . We use the same setup as for the 2D experiment described above. The results are reported in Figure 10.

0.6

Target

SGD

II

0.6

SGD

Target

0.0020

0.5

0.5

0.0015

0.4

0.4

0.3

0.3

0.0010

0.2

0.2

0.0005

0.1

0.1

0.0

0.0

0.0000

0.01

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

3

(a)

After regime

(b)

After training

timescale regime

(c)

L

2

distance with the target as a function of

ε

,

20

repeats

D Additional remarks

D.1 Counter-example in the case of non-uniform initialization

It is easy to craft an example where gradient flow in the two-timescale limit does not converge to the global minimum if the positions of the neurons are not drawn uniformly at random at initialization. Take for instance the case where there are three pieces in the target, the leftmost piece at level 1 , the second one at level -1 and the third one at level 0 . Consider the case where all neurons are in the third piece at initialization. Then (in the two-timescale limit) the bias and the weights are all instantly equal to zero, which is the solution of the optimization problem in a . Looking at (10), this shows that the positions do not move, and thus that the neurons remain in this local minimum.

This example may seem artificial, but in fact, more generally, the case where two consecutive pieces are not covered by any neuron often corresponds to a local minimum of the two-timescale dynamics. This is why we require neurons on every piece at initialization to avoid falling in this local minimum.

D.2 Ideas of proof for SGD

The proof of Theorem 1 consists in bounding the difference between the actual dynamics and the dynamics that we would have if the weights were given by a ∗ 0 ( u ( t )) , that is, by the weights in the training

in the

standard in

the two-

two-timescale limit with η = 0 . This requires to bound two errors terms, one coming from the fact that η > 0 , and one coming from the fact that ε > 0 . By making η and ϵ small enough, we ensure that both terms are small, which in turn allows to describe the trajectories of the weights and of the positions throughout the dynamics.

Moving from gradient flow to stochastic gradient descent induces two additional error terms: a discretization error, and a noise error. However, taking small enough step size allows to bound both error terms and thus ensures that (with high probability) the SGD dynamics are close to the gradient flow dynamics. For this reason, we believe that our proof structure (namely, bounding the difference with the two-timescale limit as in Proposition 4, then showing that this bound allows to describe the trajectories of the weights and of the positions as in Theorem 1) should adapt to the SGD case, at the condition that the step size is small enough.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the training of 2-layer neural networks and proposes a two-timescale limit/regime. In this \nlimit/regime, the learning rate of the first layer is much smaller than the learning rate of the second layer. \nAs a result, the training of the network can be viewed as training the first layer and performing linear regression \nover the first layer outputs, which is simpler/more structured than training both layers at the same rate. To \ndemonstrate the usefulness of this strategy, this paper considers a toy model and shows that a 2-layer network can fit \na certain family of 1D piecewise constant functions. The authors also empirically show that SGD can fail to fit this \nfamily of functions outside the two-timescale regime.\n\nSTRENGTHS:\n* The presentation is clear.\n* This paper not only contains results in the $\\varepsilon \\to 0$ limit, but also non-asymptotic results for small but \n  non-vanishing $\\varepsilon$. The derivation in the $\\eta \\to 0, \\varepsilon \\to 0$ limit is clean (Sec. 4.2), and \n  the asymptotic-to-non-asymptotic parts are also easy to follow.\n* This paper reminds me of a technique that is gaining popularity in the theory community: training the network for \n  one (large) step, freezing the first layer, and then performing linear regression in the second layer using the \n  features learned in the first layer (cf. [AAM22], [DLS22]). The limitation of this technique is that since the \n  first layer is fixed after the first layer, it lacks the ability to refine the learned features. I feel the \n  strategy introduced in this paper is a potential remedy to this problem, as here we also have the linear regression \n  part of the argument, but the first layer can be trained for multiple steps. \n\n[AAM22] Abbe, Emmanuel, Enric Boix Adsera, and Theodor Misiakiewicz. “The Merged-Staircase Property: A Necessary and Nearly Sufficient Condition for SGD Learning of Sparse Functions on Two-Layer Neural Networks.” In Proceedings of Thirty Fifth Conference on Learning Theory, 4782–4887. PMLR, 2022. https://proceedings.mlr.press/v178/abbe22a.html.\n\n[DLS22] Damian, Alex, Jason D. Lee, and Mahdi Soltanolkotabi. “Neural Networks Can Learn Representations with Gradient Descent.” arXiv, June 30, 2022. http://arxiv.org/abs/2206.15144.\n\nWEAKNESSES:\n* Almost all things in this paper are 1D and somewhat tailored to this specific piecewise constant function class. \n  I wonder whether/how this can be generalized to higher dimensions, other network architectures, and more general \n  function classes. \n* It seems that the dynamics are still relatively local. That is, we need some neurons in each target interval at \n  initialization, which may not be reasonable when the dimension is high.\n* The authors should probably add some discussion on the training-for-one-large-step-type technique (see the Strengths \n  part of the review).\n\nQUESTIONS:\n* How general is this strategy? Is it possible to apply this strategy to problems (even toy ones) with dimension higher \n  than $1$ and more standard 2-layer networks?\n* Can we relax the requirement on the initialization? For example, can we remove the requirement that each interval \n  contains at least $6$ $u_j$?\n  Intuitively, when some interval contains no neurons, it is still possible that some neurons from those fitted \n  intervals can move to that interval and fit the target function on that interval. \n  I know the current requirement is reasonable in your setting, I am asking this because I feel this type of global \n  feature learning process is important when the dimension is high.\n\nLIMITATIONS:\nThis is a theoretical work and, as far as I can see, has no potential negative societal impacts.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-26"}

comments: SUMMARY:
This paper studied the problem of fitting a piecewise constant univariate function with a shallow neural network. Specifically, the authors consider a gradient flow with a time-scale difference between the dynamics of the first and the second layer weights. It is shown that the trained shallow network can be arbitrarily close to the target function in $\mathcal{L}_2$, as long as the weight update on the first layer is much slower than one on the second layer.

STRENGTHS:
The two-time-scale regime seems novel in training neural networks. The results are well presented and their proofs are clearly explained.

WEAKNESSES:
Only a very special problem is studied. 

1. The target function is univariate, piece-wise constant. This is very restrictive. 

2. The loss is a population loss, thus the case of finite samples (fixed design or random samples) is not considered.

3. Network is quite different from those studied in other works. The activation function is a smoothed version of a stair function; the first layer is only parametrized by the bias at every neuron. 

With this many assumptions/restrictions, even though the main results are well presented and explained, it is hard to see whether those observations can give insights into the usefulness of implementing two-time-scale training for practical networks. The author did not discuss the limitations of these assumptions, nor did they show how two-time-scale training can be used in practice, even empirically.

QUESTIONS:
See "Weaknesses"

LIMITATIONS:
See "Weaknesses"
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

comments: SUMMARY:
In this paper, the authors considered the problem of learning piece-wise linear function in 1d using two-layer neural network. They considered gradient flow on mean-square loss with different learning rates for 2 layers (two-timescale). Specifically, the outer layer weights are moving much faster than the inner layer weights. The activation used in 2-layer network is similar to a rescale version of sigmoid activation. This paper shows that under proper choice of the parameters, GF converges to small loss within polynomial time in relevant parameters. Experiments are provided to support the results.

STRENGTHS:
1.	Understanding the training dynamics and convergence of neural networks is an important problem.
2.	The paper is overall easy to follow and clearly written. The proof sketch is given so that the reader can understand the main proof idea.
3.	The idea of this two-timescale/two different learning rates in analyzing the training dynamics for neural networks seems to be interesting.

WEAKNESSES:
1.	The problem considered is only in 1d and it would be interesting to see if the analysis could be generalized to multi-dimension.

QUESTIONS:
1.	In the experiments, I was wondering if one could elaborate in Figure 5 that why ‘the dynamics create a zone with no neuron’. It seems to me that there are neurons near every discontinuity of $f^*$. Also, it seems in Figure 4(c) and Figure 5(c), the number of steps for training are different, what is the reason for that?

LIMITATIONS:
The limitation is discussed in the paper. This is a theoretical work and therefore no foresee negative societal impact.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

comments: SUMMARY:
The paper studies the training dynamics of fitting a one hidden layer shallow network with heaviside activation to a piecewise ground truth function with one-dimensional input. It proves that gradient flow always recovers the ground truth in finite time with only mild over-parametrization.

STRENGTHS:
This paper is well-written, and intuitions behind technical proofs are well-presented, and theoretical results are well-supported by numerical experiments. The theoretical result itself is a nice observation, despite being in the simple one-dimensional input case.

WEAKNESSES:
The paper only considers the one-dimensional case. Based on the proof techniques, it is not clear if it is extendable to high dimension, which is of ultimate interests in the deep learning theory community, since the derivation in section 4.2 would not hold any more. While it is understandable that such a result would be difficult to obtain in high dimension, the paper didn't present any experimental results in the high-dimensional case either.

QUESTIONS:
The paper would greatly benefit from a few comments on whether the authors would expect the same phenomenon to be observed in high dimension and what the difficulty would be. Some numerical experiments in the high dimensional case would be a great addition as well.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-08-02