Convergence of gradient descent for deep neural networks
Sourav Chatterjee
A convergence criterion for gradient descent
with . Gradient descent and its many variants are indispensable tools in all branches of science and engineering, and particularly in modern machine learning and data science. The convergence properties of gradient descent are well-understood when the objective function is convex , and it is known that finding local minima of nonconvex functions by gradient descent is an NP-complete problem . In spite of this, gradient descent is widely used in practice to find local and global minima in highly nonconvex problems, especially in high dimensions. For example, it has been observed that gradient descent can often find global minima of training loss in deep learning , which is one of the reasons behind great success of the ‘deep learning revolution’ .
This article presents a novel criterion for convergence of gradient descent to a global minimum. The criterion is related to (and maybe seen as a strengthening of) the classical Kurdyka–Łojasiewicz inequality . It is also related to results from nonsmooth analysis, such as those in . Indeed, our proof idea bears close resemblance with the classical ‘Łojasiewicz trapping argument’ . Nevertheless, the convergence criterion stated below has not appeared in this exact form previously in the literature.
where is the Euclidean norm of . If for all , then we let . Our main assumption is that for some ,
Under the above assumption, we have two results. The first result, stated below, shows that the gradient flow started at converges exponentially fast to a global minimum of in where is zero. The existence of a global minimum in is a part of the conclusion, and not an assumption.
Let , , and be as above. Assume that (1.1) holds for some , and let . Then there is a unique solution of the gradient flow equation
on with , and this flow stays in for all time, and converges to some where . Moreover, for each , we have
Our second result is the analogue of Theorem 1.1 for gradient descent. It says that under the condition (1.1), gradient descent started at , with a small enough step size, converges to a global minimum of in . Again, the existence of a global minimum in is a part of the conclusion.
Let , , and be as above. Assume that (1.1) holds for some , and let . Choose such that
which is possible since (1.1) holds. Let be a uniform upper bound on the magnitudes of the first-order derivatives of in , and let be a uniform upper bound on the magnitudes of the second-order derivatives of in . Choose any such that
for each . Then for all , and as , converges to a point where . Moreover, with , we have that for each ,
Note that the above results have nothing to say about variants of gradient descent, such as stochastic gradient descent. Adding a stochastic component to the gradient descent algorithm has various benefits, such as helping it escape saddle points . Since it is known that stochastic gradient methods often asymptotically follow the path of a differential equation , it would be interesting to see if analogues of Theorems 1.1 and 1.2 can be proved for stochastic gradient descent. We also do not have anything to say about algorithms that aim to find critical points instead of global minima in nonconvex problems, such as the ones surveyed in . For a recent survey of the many variants of stochastic gradient descent and their applications in machine learning, see . For a comprehensive account of all variants of gradient descent, see . For some essential limitations of nonconvex optimization, see .
It is possible that Theorems 1.1 and 1.2 may be generalizable to what are variously called lower functions, or proximal-regular functions, or weakly convex functions. However, the generalizations are not obvious (especially for Theorem 1.2), and are therefore left for future investigation.
Application to deep neural networks
A feedforward neural network consists of the following components:
A positive integer , which denotes the number of layers. It is sometimes called the depth of the network. The number denotes the ‘number of hidden layers’. To avoid trivialities, we will assume that — that is, there is at least one hidden layer.
A sequence of positive integers , denoting the dimensions of layers . The maximum of is called the width of the network. The dimension of the ‘output layer’, , is often taken to be . We will henceforth take .
where the activation functions act componentwise on vectors of dimensions .
for , where is the step size, and denotes the gradient of (assuming that the activation functions are differentiable).
There is an enormous body of literature on convergence properties of gradient descent for neural networks. The following are some of the most important contributions. Further references can be found in the review sections of these papers and also in the recent comprehensive survey .
Convergence for convex neural networks was studied in , and for linear networks in . Convergence in the absence of convexity and linearity has remained an open problem, except in one particular scenario — when the dimensions of the hidden layers are extremely large, where ‘extremely large’ may mean either tending to infinity, or larger than some large enough power of the sample size. This is now called the ‘infinite width’ or ‘overparametrized’ regime. Following some early results in , this approach was fully developed independently in the concurrent papers . The key idea here is that in the overparametrized regime, gradient descent for the neural net can be approximated by gradient descent in a linear setting. Since then, this idea has been widely applied in a variety of settings, for example, in . For some recent advances beyond the overparametrized regime, see and references therein.
We have two results about convergence of gradient descent for feedforward neural networks of bounded width and depth. The first theorem shows that under fairly general conditions, it is possible to find an exact fit to the data (i.e., a point where ) via gradient descent with suitable initialization and step size. The main requirement is that the input data have to be linearly independent, which necessarily means that the dimension of the input space has to be . This condition is inevitable, because if this is not true, then there may not exist any point where even for linear activation.
Recall that the width of the network is the number . It is important to note that this excludes the input dimension . Theorem 2.1 implicitly needs , but there is no requirement on the width . Most previous works need the width to grow with . For example, Soltanolkotabi et al. require (their result is only for networks with one hidden layer), while both Du et al. and Allen-Zhu et al. — who deal with networks of arbitrary depth — require to be at least as large as some polynomial in . One existing result result that might imply Theorem 2.1 is [19, Theorem 2.4], but it is completely clear if it does.
The class of activation functions allowed by Theorem 2.1 includes many functions used in common practice, such as linear activation (), bipolar sigmoid activation (), and tanh activation (). Moreover, the condition is not a serious restriction, since the presence of the bias vectors implies that the class of models remains the same if we subtract off some constants from our activation functions to make them zero at the origin. In that sense, Theorem 2.1 also allows the sigmoid activation (), smoothed ReLU activation (), and complementary log-log activation (). It does not, however, allow activation functions that are not twice continuously differentiable, such as ReLU activation (), step activation ( if and if ), and piecewise linear activation.
Theorem 2.1 gives a criterion for convergence of gradient descent to a solution that perfectly interpolates the data. It does not, however, say anything about why such interpolating solutions sometimes have good generalization errors, or conditions under which deep learning performs well (or poorly). These are some of the other great mysteries of deep neural networks that have attracted much attention. For more about these, see and references therein.
The proof of Theorem 2.1 yields formulas for and , but they are quite complicated and possibly sub-optimal, and are therefore omitted from the discussion. In practice, it may be easiest to just choose and by trial and error after choosing with arbitrary positive entries, by increasing and decreasing until convergence is achieved.
Not many activation functions in common use satisfy the condition that the slope is uniformly bounded below by a positive constant. One prominent example that satisfies this condition is the leaky ReLU activation function ( if and if , where is some positive number less than ). However, leaky ReLU activation is not twice continuously differentiable. We propose the following smooth modification of the leaky ReLU:
Here , as in the definition of leaky ReLU. Note that as , smooth leaky ReLU has the same asymptotic behavior as leaky ReLU. Although , this can be easily fixed by subtracting a constant (which does not affect anything since we have bias vectors in our model).
Proof of Theorem 1.1
To avoid trivialities, let us assume that is not everywhere zero in , and hence . Throughout this proof, we will denote the closed ball by and its interior by .
Take any . Define for . Then
Thus, satisfies the integral equation for the gradient flow starting from in the interval . Thus, by the uniqueness assumption for this flow in , we get that in this interval. Consequently, in . Thus, the gradient flow starting from has a unique solution in for every . Since , this contradicts the definition of . ∎
Let be the set of all points that are within distance from . Note that and is compact. Since is in , its first and second order derivatives are uniformly bounded on . Thus, there is some such that for any ,
Let be the subset of consisting of all such that and for all . It is easy to see that is a closed subset of . Define a map as
where the second inequality holds because for all , and the third inequality holds because . Since all points at distance from are in , this shows that . Next, note that for any , and any ,
We claim that is the only such map. To prove this, suppose that there exists another map with the above properties. If maps into , then the uniqueness of the fixed point implies that . So, suppose that ventures outside . Let
Since and goes outside , is well-defined and finite. Moreover, since is closed, . But note that since for all ,
But this implies that the ball of radius centered at is completely contained in , and hence, . Thus, cannot venture outside . We conclude that . ∎
The above lemma has several useful corollaries.
Suppose that . Then there is a compact set such that for any , we can find with . By Lemma 3.2, we can find . Take any such that . Then, by Lemma 3.1, we have . But this contradicts the fact that . ∎
If , then since is a nonnegative function, must also be zero. Thus, it suffices to prove the result under the assumption that . Taking in Lemma 3.2, we get that . Suppose that is finite. Then let . Note that since , is a solution of the gradient flow equation starting from in the interval . By uniqueness, this shows that for all . In particular, . By Lemma 3.1, this implies that , which contradicts the fact that is nonzero and finite. Thus, . Again, the function is a solution of the flow equation in . Thus, by uniqueness, for all . ∎
Take any such that or . Let . Then by Lemma 3.1, . But by Corollary 3.4, . Thus, . Also by Corollary 3.4, for all , and by Lemma 3.1, for all . Thus, for all . ∎
In the following, let us fix and as in Theorem 1.1 and write and instead of and , for simplicity of notation.
If , then must visit the boundary of .
Suppose that and remains in throughout. By Lemma 3.2, there is some such that for all . Choose . Let . Then , and therefore . This shows that the gradient flow starting from exists up to time at least . Since , this gives a contradiction which proves that cannot remain in throughout. ∎
Let be any number such that for all . Then for all ,
Note that by the flow equation for ,
By the definition of , the right side is bounded above by if . It is now a standard exercise to deduce the claimed inequality. ∎
The flow cannot visit the boundary of .
Suppose that does visit the boundary of . Let . Then and for all . By the flow equation,
Now, since for all , Corollary 3.5 shows that for all . Thus, the map
is differentiable in , with continuous derivative
where the second identity follows from (3.1). Thus, for any compact interval ,
Since is continuous on , we can take and , and apply the monotone convergence theorem on the left, to get
Therefore, by (3.2) and the Cauchy–Schwarz inequality,
On the other hand, since for all , Lemma 3.7 shows that
for all . Plugging this bound into the previous display, we get
But the last quantity is strictly less than , by assumption (1.1). Since , this gives a contradiction, which proves the lemma. ∎
Combining Lemma 3.6 and Lemma 3.8, we see that . Moreover by Lemma 3.8, stays in forever. Therefore, by Lemma 3.7, it follows that for all . It remains to establish the convergence of the flow and the rate of convergence. Let
with the understanding that if for all . Then note that for any ,
where , as in the proof of Lemma 3.8. As in that proof, note that
Combining the last three displays, and invoking assumption (1.1), we get
Note that the bound does not depend on . If , then by Corollary 3.5, for all . Thus, in this case converges to . The rate of convergence is established by taking in (3.3). If , then (3.3) proves the Cauchy property of the flow , which shows that converges to some as . Again, taking in (3.3) proves the rate of convergence. ∎
Proof of Theorem 1.2
If , then and therefore for all , and there is nothing to prove. So, let us assume that . The following lemma is the key step in the proof of Theorem 1.2.
For all , .
The proof of this lemma will be carried out via induction on . We have . Suppose that for some . We will use this hypothesis to show that , with remaining fixed henceforth.
Then for all , .
Since , the assumed upper bound on implies that
Thus, , and therefore, . Consequently, the line segment joining and lies entirely in . Since , the line segments joining and lies in for all . Thus, by Taylor expansion, we have that for any ,
where is a point on the line segment joining and , and is the Hessian matrix of at . Since is an upper bound on the magnitudes of all second order derivatives of in , and , this gives
which proves that . Since , this proves the claim. ∎
We have , and for any ,
Since for , Lemma 4.2 and the definition of imply that for ,
Since and , we can take above and divide both sides by to get that . Iterating the above inequality gives the desired upper bound for . ∎
Rearranging terms, we get the desired inequality. ∎
Let . By Lemma 4.4,
Also by Lemma 4.4, for each . Thus, for any , we use the Cauchy–Schwarz inequality to get
Using Lemma 4.3 to bound the terms on the right side, we get
Now, for , we have the inequality . This gives
Plugging this into the previous display completes the proof. ∎
We are now ready to complete the proof of Lemma 4.1 and then use it prove Theorem 1.2.
Applying Lemma 4.5 with , and recalling the criterion (1.2) used for choosing , we get
This proves that , completing the induction step. ∎
By Lemma 4.1, we know that for all . Thus, Lemma 4.5 holds for all and all . In particular,
Note that the bound goes to zero as , and has no dependence on . Thus, is a Cauchy sequence in , and therefore, converges to a limit . Moreover, the above bound is also a bound for , since it has no dependence on . Lastly, since for all , Lemma 4.3 also holds for any . This shows that and gives the required bound for . ∎
Proof of Theorem 2.1
Using this relation and the fact that , we get
Then the definition of shows that for any where ,
We obtain a lower bound on the above term by simply considering those ’s that correspond to the entries of . By (5.1), this gives
Now take any as in the statement of Theorem 2.1, that is,
the entries of are all strictly positive, and
irrespective of the values of . Let be the minimum of all the entries of and be the maximum. Take any such that . Then the entries of are all bounded below by and bounded above by
Thus, by (5.3) and (5.4), we see that for any where ,
Since this holds for every , and the numbers have no dependence on , it follows that if we fix and , and take sufficiently large, then by (5.5), we can ensure that
which is the criterion (1.1) for this problem. By Theorems 1.1 and 1.2, this completes the proof of Theorem 2.1.
Proof of Theorem 2.2
In this proof, will denote arbitrary positive constants whose values depend only on , , , , , , and . (The important thing is that these constants do not depend on the input dimension .)
Thus, if happens, and we also have that
then (1.1) holds for in the ball . Looking at the above inequality, it is clear that we can choose so large that the above event is implied by the event
for some suitably defined . Thus, if happens, then (1.1) holds for in the ball . Now note that since the entries of are i.i.d. random variables,
Acknowledgements
I thank Persi Diaconis, Dmitriy Drusvyatskiy, John Duchi, and Lexing Ying for useful feedback, references, and suggestions.