On the convergence of gradient descent for two layer neural networks

On the convergence of gradient descent for two layer neural networks

Lei Li leili2010@sjtu.edu.cn School of Mathematical Sciences, Institute of Natural Sciences, MOE-LSC, Shanghai Jiao Tong University, Shanghai, 200240, P. R. China.
Abstract

It has been shown that gradient descent can yield the zero training loss in the over-parametrized regime (the width of the neural networks is much larger than the number of data points). In this work, combining the ideas of some existing works, we investigate the gradient descent method for training two-layer neural networks for approximating some target continuous functions. By making use the generic chaining technique from probability theory, we show that gradient descent can yield an exponential convergence rate, while the width of the neural networks needed is independent of the size of the training data. The result also implies some strong approximation ability of the two-layer neural networks without curse of dimensionality.

1 Introduction

The universal approximation theorem tells us that a two-layer neural networks can approximate a broad class of functions, provided that the width is sufficiently large, i.e. in the over-parametrized regime [1, 2]. Recently, it has been shown that gradient descent can find the optimal parameters for the wide two layer neural networks using optimal transport theory [3, 4, 5], and using direct estimation on particular neural network structure [6, 7, 8]. All these results are in the over-parametrized regime. We are interested to see whether the loss function can converge to zero in the under-parametrized regime, i.e., the number of data points are larger than the width of the two layer neural networks.

Following [4, 6], we consider the loss function given by the quadratic function

 ℓ:=12∫D|f(n)(x)−f(x)|2dν(x), (1.1)

where is the approximation given by the two layer neural network with width , is a compact set and , is the dimension, and is some probability measure on . We aim to show that under suitable setup, with high probability, the convergence rate is exponentially fast, and we desire the width of the neural networks is independent of the training data, by using suitable approximation . Note that this does not contradict with the result of [9], which shows that the under-parametrized networks cannot approximate the function well. The reason is that the loss function here only measures the approximation at the given data points, instead of approximation for the whole function.

Here, we mention the two main works that motivate this work. In [4], the following form of approximations were considered inspired by weak approximation of measures

 ~f(n)(x,t):=1nn∑i=1ai(t)φ(x,wi(t)), (1.2)

where represents any unit function with parameter . We now also regard as a function of since we are considering that are varying. For example, it could be the Gaussian, the one-layer neural networks or deep neural networks. They were able study the dynamics as the interacting particle systems, so that the decaying of the loss function can be understood as gradient flows in Wasserstein-2 space using optimal transport theory. However, the decaying can only be shown in the limiting regime where and the decaying rate is unknown. The mechanism is that the parameters will converge to some minimization regions.

In the work [6] by S. Du et al., the two layer of neural networks were studied, i.e.,

 φ(x,w)=σ(w⋅x), w∈Rd.

Here, is the so-called activation function. In their work, the following form of approximation is used:

 f(n)(x,t):=1√nn∑i=1ai(t)φ(x,wi(t))=1√nn∑i=1ai(t)σ(wi(t)⋅x). (1.3)

Though this seems to be a little change, the dynamics has been changed significantly compared with (1.2). They are actually able to show that the weights ’s will not converge under the gradient descent method: they are close to the initialization. They claim that the loss function converges to zero exponentially in the overparamterized regime. They obtain this result by considering the dynamics of the prediction . In [6], the convergence largely is due to their Assumption 3.1 and Theorem 3.1. The analysis for their matrix needs to exclude the cases that there are two vectors being nearly parallel, which is not desired in the regime, where is the number of training data. We note that if we rely on the dynamics of , the matrix can be positive definite simply due to the positivity of , and this is the strategy we adopt here.

Below, we aim to combine and improve the ideas from these two works to show a convergence rate of the gradient descent independent of the sample size . We use the approximation from [6], i.e. (1.3), to study the dynamics of the training loss directly, as done in [4, 8]. We makes use the key observation from [10, 6] that the weights are close to the initial values, but instead rely on the dynamics on ’s to obtain the positive definiteness of the Gram matrix. Another important technique is the generic chaining [11] from probability theory that kind of guarantees the the Gram matrix is close to some positive definite matrix uniformly in data . In section 2, we set up the mathematical formulation and state the main results. The proof is then performed in section 3. Lastly, we make some discussions in section 4.

2 Mathematical setup and the main results

We focus on the following special setup, and consider two layer neural networks as in (1.3). The generalization to other cases will be done soon in subsequent works.

About the data set and measure for , we assume the following:

Assumption 1.

The domain is a compact set such that there is a constant that is independent of satisfying for all . The measure is a probability measure supported on a countable subset of .

Below, we will denote the support of by :

 T:=suppν. (2.1)

One common example of is the empirical measure

 ν:=1MM∑p=1δ(⋅−xp), (2.2)

where ’s are the training data sampled i.i.d from some underlying distribution. Of course, in our discussion in this work, ’s are fixed.

Moreover, we assume the initial conditions as follows.

Assumption 2.

For the weight , we assume

 wi(0)∼N(0,Id), i.i.d.,

and ’s are i.i.d such that .

The condition is important because it guarantees that the approximation (1.3) is with high probability at to close up the estimation later (Lemma 4). The distribution for ’s are assumed to be Gaussians only for technical convenience. The argument can be extended to other distributions with finite second moments.

Moreover, we assume

Assumption 3.

The activation function satisfies , and . Moreover, , and thus

 ∥σ∥∞+∥σ′∥∞<∞.

The sigmoid function should satisfy the condition. This condition may be too strong, however, we impose this only for simplicity. The generalization to other functions should be doable.

Under the gradient descent with loss function (1.1), the parameters ’s satisfy

 ˙wi(t)=−∇wiℓ=−1√n∫D(f(n)(x)−f(x))ai(t)σ′(wi⋅x)xdν(x), (2.3)

where means the time derivative. Similarly, the parameters ’s satisfy

 ˙ai(t)=−∂ℓ∂ai=−1√n∫D(f(n)(x)−f(x))σ(wi⋅x)dν(x). (2.4)

Instead of studying the prediction as in [6], we study the dynamics of the loss function directly as motivated by [4, 8],

 (2.5)

The matrix

 M(x,y):=1nn∑i=1(σ(wi(t)⋅x)σ(wi(t)⋅y)+a2i(t)σ′(wi(t)⋅x)σ′(wi(t)⋅y)x⋅y) (2.6)

is called the Gram matrix in the terminology from [6]. We aim to use the positive definiteness of this matrix to show the exponential convergence. The second part is positive semi-definite (in fact in [6], it was argued that this part is also positive definite when the data points are not parallel), so we can control easily as

 ddtℓ≤−1nn∑i=1∬D×D(f(n)(x)−f(x))(f(n)(y)−f(y))σ(wi(t)⋅x)σ(wi(t)⋅y)dν(x)dν(y). (2.7)

One key observation from [6] is that the weights do not change much during the dynamics. The choice (1.3) ensures the form , so that one has kind of law of large number convergence here. Moreover, we note that are i.i.d D Gaussian variables regardless the dimension of itself. Using the generic chaining approach [11], this allows us obtain the convergence rate of the loss function independent of and when is large enough ( without exponential dependence in ).

Theorem 1.

Under Assumptions 1-3, for any small, when , with probability , it holds that

 ℓ(t)≤(∥f∥2∞+C(logδ)2)exp(−λ2t), (2.8)

where and are independent of , and .

Remark 1.

The dependence in the lower bound of comes from a constant associated with the geometry of , as in Proposition 1. In our case, the domain is like a ball whose volume goes to zero as . If we consider other domains, like , the dependence in will be changed. On one hand, the bounds of will be changed and also the constant in Proposition 1 will be changed.

We remark that as does not mean approximates well. It only approximates the function values on the . However, there is another implication of the result. In the case of the empirical measures (2.2), the width and convergence rates are independent of the data set . This somehow says that any bounded functions can be approximated by the two-layer neural networks with high probability. In fact, it seems that the following approximation ability of two layer neural networks holds.

Corollary 1.

Suppose satisfying Assumption 1 is the union of closures of some domains (open connected sets) so that the Lebesgue measure can be defined. Then, the two layer neural networks satisfying Assumption 3 with width are able to approximate any functions , under the norm.

Here, means the norm with Lebesgue measure. As commented in section 4, it seems that such good approximations are available ”everywhere” in the parameter space. In Theorem 1, the lower bound of depends on , but in this corollary, there is no dependence in the lower bound of . The reason is that one can approximate first and then scale ’s to be . The dependence of origins from the gradient descent dynamics (2.3) where the derivatives depend on .

3 The proof of the main results

Below, we use to indicate a constant independent of , whose concrete meaning can change from line to line.

Lemma 1.

The derivatives of the parameters satisfy

 12ddtn∑i=1a2i=−∫D(f(n)(x)−f(x))f(n)(x)dν(x)=−2ℓ(t)−∫D(f(n)(x)−f(x))f(x)dν(x)≤∥f∥∞∫D|f(n)(x)−f(x)|dν(x)≤∥f∥∞√ℓ(t), (3.1)

and

 |˙wi|≤C|ai(t)|√n√ℓ(t),   |˙ai|≤1√n√ℓ(t). (3.2)

where is independent of , and .

The proof is straightforward and we omit. Let us just remark that if we consider, for example, . Using (2.3), by Hölder’s inequality, it is straightforward to obtain

 |˙wi(t)|≤1√n∣∣∣∫D(f(n)(x)−f(x))ai(t)σ′(wi⋅x)xdν(x)∣∣∣≤c|ai(t)|√n(∫D(f(n)(x)−f(x))2dν(x))1/2(∫D∥σ′∥2∞dν(x))1/2≤C|ai(t)|√n√ℓ(t).

Using Lemma 1, we are able to establish:

 ddtℓ(t)≤−1nn∑i=1∬D×D(f(n)(x)−f(x))(f(n)(y)−f(y))×σ(wi(0)⋅x)σ(wi(0)⋅y)dν(x)dν(y)+R(t)=:−Dn(t)+R(t), (3.3)

where we have introduced

 Dn(t):=1nn∑i=1∬D×D(f(n)(x)−f(x))(f(n)(y)−f(y))σ(wi(0)⋅x)σ(wi(0)⋅y)dν(x)dν(y), (3.4)

and

 R(t):=Cnn∑i=1∬D×D|f(n)(x)−f(x)||f(n)(y)−f(y)|×∫t0|ai(s)|√n√ℓ(s)dsdν(x)dν(y). (3.5)

We now estimate these two terms separately.

Lemma 2.

The term satisfies

 R(t)≤C(1√nb(t)+1nb(t)2)ℓ(t),

where

 b(t):=∫t0√ℓ(s)ds. (3.6)
Proof.

First of all, one finds

 |ai(s)|≤|ai(0)|+∫s0|˙ai(τ)|dτ.

Using (3.2), one has (note that )

 |ai(s)|≤1+C√n∫s0√ℓ(τ)dτ.

Hence,

 1nn∑i=1|ai(s)|√n≤1√n+Cn∫s0√ℓ(τ)dτ.

Consequently, one has

 1nn∑i=1∫t0|ai(s)|√n√ℓ(s)ds≤1√n∫t0√ℓ(s)ds+Cn(∫t0√ℓ(s)ds)2, (3.7)

where we have used the fact that

 (∫s0√ℓ(τ)dτ)√ℓ(s)=12dds(∫s0√ℓ(τ)dτ)2.

The claim then follows. ∎

We are now in a position to take in . In particular, we need to estimate

 Dn(x,y):=1nn∑i=1σ(wi(0)⋅x)σ(wi(0)⋅y). (3.8)

For fixed , it is clear that this converges to

 D(x,y):=Ew∼N(0,Id)σ(w⋅x)σ(w⋅y), (3.9)

and the rate is independent of . However, to obtain the uniform convergence rate for is challenging. The usual technique is the union bound, which gives a quantity involving , and this further requires the over-parametrized regime. One may also use the convergence of empirical measures in the so called Wasserstein spaces [12], which yields a quantity independent of the training data, but suffers from curse of dimensionality. In our problem here, we are essentially considering the projections of high dimensional Gaussians. Hence, we expect the convergence rate will not suffer from curse of dimensionality. We find that the generic chaining approach in [11, 13] useful.

We introduce the following result essentially derived from [13, Proposition 2.2] and the proof of [13, Theorem 1.1]. The proof is then deferred to Appendix A using generic chaining approach.

Proposition 1.

Let be as in Assumptions 1. Consider

 {Zx:=1nn∑i=1Yix}x∈D.

Suppose that for any given , ’s are i.i.d, and for any , and are independent for but may have different distributions. Moreover, ’s satisfy the following.

1. They are uniformly bounded, have mean zero.

2. There is a constant such that

 E|Yix−Yiy|2≤C1|x−y|2, ∀i,x,y.

Then there is a parameter only depending on with magnitude , such that for any small and that is countable, when , with probability ,

 supx∈T|Zx|≤Clog(1δ)γ2√n. (3.10)

We now move to the estimates of . We have

Lemma 3.

For any small , when , with probability , it holds that

 supx,y∈T|Dn(x,y)−D(x,y)|≤Clog(1δ)γ2√n. (3.11)
Proof.

We only need to verify the conditions in Proposition 1. Here, the domain we work on is while is the countable subset we consider. The distance is defined by .

Clearly, we need to identity

 Yix,y=σ(wi(0)⋅x)σ(wi(0)⋅y)−Ew∼N(0,Id)σ(w⋅x)σ(w⋅y).

Note that even if , are still i.i.d.

Noting that

one has

 |Yix,y−Yix1,y1|2≤C(|wi(0)⋅(x−x1)|2+|wi(0)⋅(y−y1)|2)+CEw(|w⋅(x−x1)|2+|w⋅(y−y1)|2).

Taking expectation, one finds . The claim follows.

Note that in this case, the dimension is , but clearly, . ∎

Note that the boundedness of is also used here for the Bernstein inequality in Appendix A. Consequently, regarding defined in (3.4), when (3.11) holds, one has

 |Dn(t)−D(t)|≤Clog(1δ)γ2√nℓ(t), (3.12)

where we introduced

 D(t):=∬D×D(f(n)(x)−f(x))(f(n)(y)−f(y))×Ew∼N(0,Id)σ(w⋅x)σ(w⋅y)dν(x)dν(y). (3.13)

Now, we move onto the initial values of . If we identify

 Yix:=ai(0)σ(wi(0)⋅x)

and use Proposition 1, we can find that when , with probability , it holds that

 supx∈T∣∣1√nn∑i=1ai(0)σ(wi(0)⋅x)∣∣≤Clog(1δ)γ2. (3.14)

Here, depends on . However, according to [8, Lemma 1], one has the following.

Lemma 4.

For any small , with probability , it holds that

 supx∈T∣∣1√nn∑i=1ai(0)σ(wi(0)⋅x)∣∣≤Clog(1δ). (3.15)

The proof is done via the so-called Rademacher complexity, and we refer the readers to [8].

Consequently, we find that for any small , with probability , it holds that

 ℓ(0)≤(∥f∥2∞+C(logδ)2). (3.16)

Now, we estimate the Gram matrix.

Lemma 5.

There exists independent of such that

 Ew∼N(0,Id)σ(w⋅x)σ(w⋅y)≥λ, ∀x,y∈D. (3.17)
Proof.

If either or is zero, the conclusion is clearly correct.

Now, assume . Then, , which is a 2D Gaussian with covariant matrix

 Σ=(|x|2x⋅yx⋅y|y|2).

Consider the interval where . Then, contains for some .

Now, we should estimate . The expectation can be rewritten as

where and are independent standard normal variables, and

 r=x⋅y|x||y|.

Note that is bounded, hence, we only need and where is a number chosen such that and are in for all . Clearly can be made universal independent of . The probability hence has a lower bound independent of . Then, setting will suffice. ∎

We are now able to prove the main result.

Proof of Theorem 1.

By the previous lemmas, for any small, when , we may pick an event with probability , such that outside this, we have , and also .

Then, on

 dℓ(t)dt≤−λℓ(t)+Clog(1δ)γ2ℓ(t)√n+C(b(t)√n+1nb(t)2)ℓ(t).

We first of all note that is nonincreasing by (2.7). We now consider a fixed . Then, all the quantities below are essentially discussed for such , and we omit this dependence for convenience.

Define

 Tb:=inf{t≥0:λ−Clog(1δ)γ2√n−C√nb(s)−Cnb(s)2>λ2, ∀0≤s≤t}.

Since , if , then . Then, for it holds that

 dℓ(t)dt≤−12λℓ(t).

Then,

 ℓ(t)≤ℓ(0)exp(−λ2t)≤(∥f∥2∞+C(logδ)2)exp(−λ2t).

Hence, for all

 b(t)=∫t0√ℓ(s)ds<4λ√∥f∥2∞+C(logδ)2.

Hence, if is chosen such that

 Clog(1/δ)γ2√n+4Cλ√n√∥f∥2∞+C(logδ)2+16Cnλ2(∥f∥2∞+C(logδ)2)<λ2,

also holds, then we must have . In other words, for some ,

 n≥C1(γ22(logδ)2+∥f∥2∞)

will imply .

Otherwise, suppose . Then,

still holds. Then, by continuity, there exists such that on this still holds. This then contradicts with the definition of .

Now that , holds for all . Lastly, we note , the claim then follows. ∎

We now move to the corollary regarding approximation.

Proof of Corollary 1.

If we can approximate the functions with , we can multiply functions with suitable constants to approximate any functions. Hence, we assume below. Moreover, one can use functions to approximate with small error under norm. Hence, we assume that below.

Fix . One first of all has the following estimates:

 ∫D|f(x)−f(n)(x)|2dx≤∣∣∣∫D|f(x)−f(n)(x)|2dx−∫D|f(x)−f(n)(x)|2dνM∣∣∣+∫D|f(x)−f(n)(x)|2dνM, (3.18)

where

 νM=1MM∑p=1δ(⋅−xp)

is the empirical measure for data .

By the standard result of convergence empirical measures [12], there exists such with such that,

 W1(νM,dx)<ϵ,

where is the Wasserstein distance.

Then, we set . Then, we take the event such that the claims in Theorem 1 hold for , and such that . Hence, we are able to find the parameters for large enough such that

 ∫D|f(x)−f(n)(x,t)|2dνM≤ϵ.

The important consequence of Theorem 1 is that is independent of (and thus ).

Now, we fix such and estimate the Lipschitz constant of . The gradient of this function is

 −2(f(x)−f(n)(x,t))(∇f+1√nn∑i=1ai(t)σ′(wi(t)⋅x)wi(t)).

Hence, we need to estimate and .

By Lemma 3.1,

 n∑i=1a2i(t)≤n∑i=1a2i(0)+C=n+C.

Similarly,

 |wi(t)|≤|wi(0)|+C√n∫t0|ai(s)|√ℓ(s)ds.

Hence,

 n∑i=1|wi|≤√n ⎷n∑i=1|wi(0)|2+C∫t0 ⎷n∑i=1a2i(s)√ℓ(s)ds≤Cn√d.

Hence, and . In other words, the Lipschitz constant of is controlled by a polynomial of and the Lipschitz constant of (but independent of ), since is independent of .

The whole error is thus given by

 Cϵ+Lip(|f(x)−f(n)|2)ϵ≤C(n,d,f)ϵ.

The constant is independent of , and the concrete expression of . Hence, it is a fixed number once are given. As is arbitrary, the claim follows. ∎

4 Conclusion and discussion

Current works show the convergence of loss function in the overparametrized regime for neural networks. By considering the loss function directly, we obtain the decay of the loss function in a rate independent of the number of training data, where the approximation is given specifically by (1.3). The key observation was made by S. Du et al that the weights are close to the initialization. In this sense, the loss function has global minimizations “everywhere” in the parameter space. As we have seen, the proof works for large number of data , and as Corollary 1 indicates, the parameters that can achieve the good approximation seems to be available ”everywhere” in the parameter space, which we strongly believe leads to the good performance of stochastic gradient descent. If we instead use the approximation of the form (1.2) inspired by mean field limit, then one must wait the weights to converge to some specific regions to gain global minimization. This probably means the approximation (1.3) has larger capacity and may be used to explain why SGD behaves well.

The loss function used here is quadratic functions. However, we argue that quadratic functions are general enough since near the minima, all functions look like quadratic functions.

There are many instant subsequent works, for example, one may explore whether the ideas and techniques here can be applied to improve the the generalization error estimates in literature, and to get better explanation of why stochastic gradient descent works well. These should be challenging but exciting problems.

Acknowledgement

L. Li would like to thank Haizhao Yang for helpful comments. The work of L. Li was partially sponsored by NSFC 11901389 and 11971314, and Shanghai Sailing Program 19YF1421300.

Appendix A Proof of Proposition 1

The proof divides into several steps. Below, gain denotes some generic constant that is independent of whose concrete meaning can change from line to line.

Step 1. Basic concentration inequalities.

Define

 φ(v):=√v+1 (A.1)

and define

 Wx:=φ(|Zx|). (A.2)

Hence, using , one has

 |Wx−Wy|≤√||Zx|−|Zy||≤√|Zx−Zy|.

Note that

 Zx−Zy=1nn∑i=1(Yix−Yiy),

where ’s are still i.i.d with mean zero. By Bernstein inequality [14], one has for any

 P(|Zx−Zy|≥t|x−y|)≤2exp(−nt2|x−y|2/σ22(1+Lt|x−y|/σ))≤2exp(−C1nmin(t2,t)), (A.3)

where .

Consequently,

 P(|Wx−Wy|≥t√|x−y|)≤P(|Zx−Zy|≥t2|x−y|)≤2exp(−Cnmin(t4,t2)).

Hence, one defines the metric (note that the maximum of two metrics is still a metric)

 g(x,y)=max(|x−y|,|x−y|1/2), (A.4)

then one has

 P(|Zx−Zy|≥tg(x,y))≤2exp(−Cnmin(t2,t)), (A.5)

while for ,

 P(|Wx−Wy|≥tg(x,y))≤2exp(−Cnt2). (A.6)

Step 2 Generic chaining

Now consider the domain equipped with the metric in (A.4). Define a number

 γα:=infsupx∈D∞∑s=02s/αg(x,Ts), (A.7)

where is a set containing points. The is taken over the collections of such sets . From to , the number of points in is squared. Hence, the distance decays double exponentially if the sequence is chosen suitably. Hence, is well-defined. For a given , for every , there must be , such that

 g(x,Ts)=g(x,y),

and we define

 πs(x):=y. (A.8)

Let be the usual Euclidean distance, and we find that on the suitable sequence of subsets, . Hence, the number is comparable to

 ~γα:=infsupx∈D∞∑s=02s/αge(x,Ts).

Below, we focus only on . The number of for the unit ball is like as shown in [11, sec. 2.2], where is the dimension.

By the definition of , we are able to choose the collection of such that

 ∞∑i=12s/2g(πs(x),πs+1(x))≤2γ2,

which holds for all .

Step 3 Controlling .

Let be such that . Consider any . By step 1, for any fixed , one has

 P⎛⎝|Wx−Wy|≥2C1√log(|Ts||Ts+1|)ng(x,y)⎞⎠≤2exp(−2log(|Ts||Ts+1|))=2|Ts|2|Ts+1|2≤1|Ts|2|Ts+1|,

where is the number of elements in .

Define

 Es:=⋃x∈Ts,y