On exponential convergence of SGD in non-convex over-parametrized learning

Introduction

Stochastic Gradient Descent and its variants have become a staple of the algorithmic foundations of machine learning. Yet many of its properties are not fully understood, particularly in non-convex settings common in modern practice.

In this note, we study convergence of Stochastic Gradient Descent (SGD) for the class of functions satisfying the Polyak-Lojasiewicz (PL) condition. This class contains all strongly-convex functions as well as a broad range of non-convex functions including those used in machine learning applications (see the discussion below).

The primary purpose of this note is to show that in the interpolation setting (common in modern over-parametrized machine learning and studied in our previous work ) SGD with fixed step size has exponential convergence for the functions satisfying the PL condition. To the best of our knowledge, this is the first such exponential convergence result for a class of non-convex functions.

Most analyses for optimization in machine learning have concentrated on convex or, commonly, strongly convex setting. These settings are amenable to theoretical analyses and describe many important special cases of ML, such as linear and kernel methods. Still, a large class of modern models, notably neural networks, are non-convex. Even for kernel machines, many of the arising optimization problems are poorly conditioned and not well-described by the traditional strongly convex analysis. Below we list some properties of the PL-type setting which make it particularly attractive and relevant to the requirements of machine learning, especially in the interpolated and over-parametrized setting.

To verify the PL condition in the interpolated setting we need access to the norm of the gradient $\|\nabla\mathcal{L}({\boldsymbol{w}})\|$ and the value of the objective function $\mathcal{L}({\boldsymbol{w}})$. These quantities are typically easily accessible empiricallyIn general we need to evaluate $\mathcal{L}({\boldsymbol{w}})-\mathcal{L}({\boldsymbol{w}}^{*})$. Since $\mathcal{L}({\boldsymbol{w}}^{*})=0$, no further knowledge about ${\boldsymbol{w}}^{*}$ is required., can be accurately estimated from a sub-sample of the data, and are often tractable analytically. On the other hand, verifying convexity requires the cumbersome positive definiteness of the Hessian matrix requiring accurate estimation of its smallest eigenvalue $\lambda_{\mathsf{min}}$. Verifying this empirically is often difficult and cannot always be based on a sub-sample due to the required precision of the estimator when $\lambda_{\mathsf{min}}$ is close to zero (as is frequently the case in practice).

The norm of the gradient is much more resilient to perturbation of the objective function than the smallest eigenvalue of the Hessian (for convexity).

Many modern machine learning methods are over-parametrized and result in manifolds of global minima . This is not compatible with strict convexity and, in most circumstancesUnless those manifolds are convex domains in lower-dimensional affine sub-spaces., not compatible with convexity. However, manifolds of solutions are compatible with the PL condition.

Nearly every application of machine learning employs techniques for feature extraction or feature transformation. Global minima and the property of interpolation (shared global minima for the individual loss functions) are preserved under coordinate transformations. Yet convexity is generally not, thus not allowing for a unified analysis of optimization under feature transforms. In contrast, as discussed in Section 3, the PL condition is invariant under a broad class of non-linear coordinate transformations.

Many problems of interest in machine learning involve optimization on manifolds. While geodesic convexity allows for efficient optimization, it is a parametrization dependent notion and is generally difficult to establish, as it requires explicit knowledge of the geodesic coordinates on the manifold. In contrast, the PL condition also allows for efficient optimization, while invariant under the choice of coordinates and far easier to verify. See for some recent applications.

Most convergence analyses in convex optimization rely on the distance to the minimizer. Yet, this distance is often difficult or impossible to bound empirically. Furthermore, the distance to minimizer can be infinite in many important settings, including optimization via logistic loss or inverse problems over Hilbert spaces, as in kernel methods . In contrast, PL-type analyses directly involve the value of the loss function, an empirically observable quantity of practical significance.

As originally observed by Polyak , the PL condition is sufficient for exponential convergence of gradient descent. As we establish in this note, it also allows for exponential convergence of stochastic gradient descent with fixed step size in the interpolated setting.

Technical contributions: The main technical contribution of this note is to show the exponential convergence of mini-batch SGD in the interpolated setting. The proof is simple and is reminiscent of the original observation by Polyak of exponential convergence of gradient descent. It also extends our previous work on the exponential convergence of mini-batch SGD to a non-convex setting. Interestingly, the step size arising from the PL condition in our analysis depends on the parameter $\alpha$ and is potentially much smaller than that in the strongly convex case, where no such dependence is needed. At this point it is an open question whether this dependence is necessary in the PL setting. As an additional contribution, in Section 4, we show that for a special class of PL functions obtained by a composition of a strictly convex function and a linear transformationThese functions are convex but not necessarily strictly convex., we obtain exponential convergence without such dependence on $\alpha$ in the step size. However, this result requires a different type of analysis than that for the general PL setting. In Section 3, we provide a formal statement capturing the transformation invariance property of the PL condition.

The PL condition has recently become popular in optimization and machine learning starting with the work . In fact, as discussed in , several other conditions proposed for convergence analysis are special cases of the PL condition. One such condition is Restricted secant inequality (RSI) proposed in . Another set of conditions that are special cases of the PL condition was referred to as “one-point convexity” in . The two variations of one-point convexity discussed there are special cases of RSI and PL, respectively, and hence are in the PL class. The same reference points out several examples of “one-point convexity” in previous works. Some notable examples satisfying RSI include two-layer neural networks , matrix completion , dictionary learning , and phase retrieval . It has also been observed empirically that neural networks satisfy the PL condition . In particular, we note the recent work which considers a class of neural networks that attain zero quadratic loss implying interpolation. In their proof it is shown that this class of neural nets satisfies the PL condition. Hence our results imply exponential convergence of SGD for this class. To the best of our knowledge this is the first time that exponential convergence of SGD has been established for a class of multi-layer neural networks.

Exponential Convergence of SGD for PL Losses

We start by formally stating the Polyak- Lojasiewicz (PL) Condition.

We do not assume compact parameter space; that is, a parameter vector ${\boldsymbol{w}}\in\mathcal{H}$ can have unbounded norm, however $\mathcal{L}$ is assumed to be bounded. In particular, a global minimizer may not exist, however, we assume the existence of global infimum for $\mathcal{L}$ (which is equal to zero w.o.l.g.).

To elaborate, we assume the existence of a sequence ${\boldsymbol{w}}_{1},{\boldsymbol{w}}_{2},\ldots$ such that

Consider the SGD algorithm that starts at an arbitrary ${\boldsymbol{w}}_{0}\in\mathcal{H}$, and at each iteration $t$ makes an update with a constant step size $\eta$:

where $m$ is the size of a mini-batch of data points whose indices $\{i_{t}^{(1)},\ldots,i_{t}^{(m)}\}$ are drawn uniformly with replacement at each iteration $t$ from $\{1,\ldots,n\}$.

The theorem below establishes the exponential convergence of mini-batch SGD for any smooth, PL loss $\mathcal{L}$ in the interpolated regime.

where the expectation is taken w.r.t. the randomness in the choice of the mini-batch.

From $\lambda$-smoothness of $\mathcal{L}$, it follows that

Fixing ${\boldsymbol{w}}_{t}$ and taking expectation with respect to the randomness in the choice of the batch $i_{t}^{(1)},\ldots,i_{t}^{(m)}$ (and using the fact that those indices are i.i.d.), we get

By invoking $\alpha$-PL condition of $\mathcal{L}$ and assuming that $\eta\leq\frac{2}{\lambda}$, we get

By optimizing the quadratic term in the upper bound 4 with respect to $\eta$, we get $\eta=\frac{\alpha m}{\lambda\left(\beta+\lambda(m-1)\right)}$, which is $\eta^{*}(m)$ in the theorem statement. Hence, (4) becomes

which gives the desired convergence rate. ∎

A Transformation-Invariance Property of PL Functions and Its Implications

Note that the condition that $d\geq k$ is necessary for ${\lambda_{\mathsf{m}in}}\left(J_{\Phi}({\boldsymbol{w}})^{T}\,J_{\Phi}({\boldsymbol{w}})\right)$ to be positive. The condition on ${\lambda_{\mathsf{m}ax}}\left(J_{\Phi}^{T}({\boldsymbol{w}})\,J_{\Phi}({\boldsymbol{w}})\right)$ holds when $\Phi$ is differentiable and Lipschitz-continuous. The above claim follows easily from the chain rule and the PL condition.

Given this property of PL functions and our result in Theorem 1, we can argue that for smooth, PL losses, the exponential convergence rate of SGD is preserved under any transformation that satisfies the conditions in the above claim. We formalize this conclusion below.

Faster Convergence for a Class of Convex Losses

where $\widehat{{\boldsymbol{w}}}_{0}=A^{\dagger}{\boldsymbol{w}}_{0}$ and $\widehat{{\boldsymbol{w}}}^{*}=A^{\dagger}{\boldsymbol{w}}^{*}$ where $A^{\dagger}$ is the pseudo-inverse of $A$.

Using these observations, we can make the following claim.

$\mathcal{L}$ is $\alpha\,\sigma_{\mathsf{min}}^{2}$-strongly convex over $\mathsf{Range}(A^{\dagger})$.

where, as above, $\widehat{{\boldsymbol{w}}}_{t+1}=A^{\dagger}A{\boldsymbol{w}}_{t+1}$ is the projection of ${\boldsymbol{w}}_{t+1}$ onto $\mathsf{Range}(A^{\dagger})$. Similarly, $\widehat{{\boldsymbol{w}}}^{*}$ is the projection of ${\boldsymbol{w}}^{*}$ onto $\mathsf{Range}(A^{\dagger})$. Now, consider $\|\widehat{{\boldsymbol{w}}}_{t+1}-\widehat{{\boldsymbol{w}}}^{*})\|^{2}$. From the update step (2) of the mini-batch SGD and the linearity of the projection operator $A^{\dagger}A$, we have

Using the above inequality together with (7), we have

References