The importance of better models in stochastic optimization

Significance

Sensitivity of optimization algorithms to problem and algorithmic parameters leads to tremendous waste in time and energy, especially in applications with millions of parameters, such as deep learning. We address this by developing stochastic optimization methods demonstrably—both by theory and by experimental evidence—more robust, enjoying optimal convergence guarantees for a variety of stochastic optimization problems. Additionally, we highlight the importance of method sensitivity to problem difficulty and algorithmic parameters.

Abstract

Standard stochastic optimization methods are brittle, sensitive to stepsize choice and other algorithmic parameters, and they exhibit instability outside of well-behaved families of objectives. To address these challenges, we investigate models for stochastic optimization and learning problems that exhibit better robustness to problem families and algorithmic parameters. With appropriately accurate models—which we call the aprox family—stochastic methods can be made stable, provably convergent, and asymptotically optimal; even modeling that the objective is nonnegative is sufficient for this stability. We extend these results beyond convexity to weakly convex objectives, which include compositions of convex losses with smooth functions common in modern machine learning. We highlight the importance of robustness and accurate modeling with experimental evaluation of convergence time and algorithm sensitivity.

A major challenge in stochastic optimization—the algorithmic workhorse for much of modern statistical and machine-learning applications—is in setting algorithm parameters (or hyperparameter tuning). This sensitivity causes multiple issues. It results in thousands to millions of wasted engineer and computational hours. It also leads to a lack of clarity in research and development of algorithms—in claiming that one algorithm is better than another, it is unclear whether this is due to judicious choice of dataset or judicious parameter settings or whether indeed the algorithm does exhibit new desirable behavior. Consequently, in this paper we pursue 2 main thrusts: First, by using models more accurate than the first-order models common in stochastic gradient methods, we develop families of algorithms that are provably more robust to input parameter choices, with several corresponding optimality properties. Second, we argue for a different type of experimental evidence in evaluating stochastic optimization methods, where one jointly evaluates convergence speed and sensitivity of the methods.

The wasted computational and engineering energy is especially pronounced in deep learning, where engineers use models with millions of parameters, requiring days to weeks to train a single model. To get a sense of this energy use, we consider a few recent papers we view as exemplars of this broader trend: In searching for optimal neural network architectures and hyperparameters, the papers (1–3) used approximately 3,150 graphics processing unit (GPU) days, 22,000 GPU days, and 750,000 central processing unit (CPU) days of computation, respectively. To put this in perspective, assuming standard CPU energy use of between 60 and 100 W, the energy (ignoring network interconnect, monitors, etc.) for the paper (3) is roughly between 4 and $6\cdot 10^{12}$ J. At $10^9$ J per tank of gas, this is sufficient to drive 4,000 Toyota Camrys the 380 miles between San Francisco and Los Angeles.

To address these challenges, we develop stochastic optimization procedures that exhibit similar convergence to classical approaches—when the classical approaches have good tuning parameters—while enjoying better robustness, achieving this performance over a range of parameters. We argue too for evaluation of optimization algorithms based not only on convergence time but also on robustness to input choices. Briefly, a fast algorithm that converges for a small range of stepsizes is too brittle; we argue instead for (potentially slightly slower) algorithms that converge for broad ranges of stepsizes and other parameters. Our theory and experiments demonstrate the effectiveness of our methods for applications including phase retrieval, matrix completion, and deep learning.

Problem Setting and Approach

We begin by making our setting concrete. We study the stochastic optimization problem

$$\begin{array}{cc}\hfill minimize\hspace{0.17em}& F\left(x\right)≔{\mathbb{E}}_P\left[f\left(x;S\right)\right]={\int }_{\mathcal{S}}f\left(x;s\right)dP\left(s\right)\hfill \\ x\\ \hfill \mathrm{s}\mathrm{u}\mathrm{b}\mathrm{j}\mathrm{e}\mathrm{c}\mathrm{t}\hspace{0.17em}\mathrm{t}\mathrm{o}\hspace{0.17em}\hspace{0.17em}& x\in \mathcal{X}.\hfill \end{array}$$ [1]

In problem 1, the set $\mathcal{S}$ is a sample space, $\mathcal{X}\subset {\mathbb{R}}^n$ is closed convex, and $f\left(x;s\right)$ is the loss $x$ suffers on sample $s$. In this paper, we move beyond convex optimization by considering $\rho \left(s\right)$-weakly convex functions $f$, meaning (cf. refs. 4 and 5) that $f\left(x;s\right)+\frac{\rho \left(s\right)}{2}{∥x∥}_2^2$ is convex. We recover convexity when $\rho \left(s\right)\le 0$. Examples include linear regression, $f\left(x;\left(a,b\right)\right)={\left(⟨a,x⟩-b\right)}^2$, and phase retrieval, $f\left(x;\left(a,b\right)\right)=|{⟨a,x⟩}^2-b|$, which is $2{∥a∥}_2^2$-weakly convex.

Most optimization methods iterate by making an approximation—a model—of the objective at the current iterate, minimizing this model and reapproximating. Stochastic (sub)gradient methods (6, 7) instantiate this approach using a linear approximation; following initial work of our own and others (5, 8, 9), we study the modeling approach in more depth for stochastic optimization. Thus, the aprox algorithms we develop iterate as follows: For $k=\mathrm{1,2},\dots$, we draw a random $S_k\sim P$ and then update the iterate $x_k$ by minimizing a regularized approximation to $f\left(\cdot ;S_k\right)$, setting

$$x_{k+1}≔\underset{x\in \mathcal{X}}{argmin}\left\{f_{x_k}\left(x;S_k\right)+\frac{1}{2{\alpha }_k}{∥x-x_k∥}_2^2\right\}.$$ [2]

We call $f_x\left(\cdot ;s\right)$ the model of $f$ at $x$, where $f_x$ satisfies 3 conditions (cf. refs. 5, 8, and 9):

• C.i)(Model convexity): The function $y\mapsto f_x\left(y;s\right)$ is convex and subdifferentiable.
• C.ii)(Weak lower bound): The model $f_x$ satisfies

$$f_x\left(y;s\right)\le f\left(y;s\right)+\frac{\rho \left(s\right)}{2}{∥y-x∥}_2^2\hspace{0.17em}\hspace{0.17em}\text{for all}y\in \mathcal{X}.$$

• C.iii)(Local accuracy): We have $f_x\left(x;s\right)=f\left(x;s\right)$.

The containment ${\partial }_y{f}_x\left(y;s\right){\mid }_{y=x}\subset {\partial }_{x}f\left(x;s\right)$ is immediate from condition C.iii. We provide examples presently.

We show that models slightly more accurate than the first-order model used by the stochastic gradient method—sometimes as simple as recognizing that if $f$ is nonnegative, we should truncate the approximation at zero—achieve substantially better theoretical guarantees and practical performance. While the iterates of gradient methods can (superexponentially) diverge for misspecified stepsizes, our methods guarantee the iterates never diverge. Even more, this stability guarantees convergence and, in convex cases, optimal asymptotic normality of the averaged iterates. Finally, we evaluate the performance of our methods, validating our theoretical findings on convergence and robustness for a range of problems, including matrix completion, phase retrieval, and classification with neural networks. We defer proofs to SI Appendix.

In optimization broadly, proximal point methods and their related robust convergence are classical (10–12), and their role in smoothing and Moreau–Yosida regularization is also central in convex and variational analysis (13–15). In signal processing, least-mean squares for adaptive filtering is an important instance of the stochastic proximal point method (16, 17). More recent work in large-scale optimization and machine learning revisits Moreau smoothing and regularization, extending acceleration and stability properties of proximal-point-type methods to finite sum and stochastic problems (18–20).

Notation and Basic Assumptions

For a weakly convex function $f$, we let $\partial f\left(x\right)$ denote its Fréchet subdifferential at the point $x$, and $f\prime \left(x\right)\in \partial f\left(x\right)$ denotes an arbitrary element of the subdifferential. Throughout, we let $x^{\star }$ denote a minimizer of problem 1 and ${\mathcal{X}}^{\star }={argmin}_{x\in \mathcal{X}}F\left(x\right)$ denote the optimal set for problem 1. We let ${\mathcal{F}}_k≔\sigma \left(S_1,\dots ,S_k\right)$ denote the $\sigma$ field generated by the first $k$ random variables $S_i$. Note that $x_k\in {\mathcal{F}}_{k-1}$ for all $k$. Unless stated otherwise, we assume that the function $f\left(x;s\right)$ is $\rho \left(s\right)$ -weakly convex for each $s\in \mathcal{S}$. Finally, the following assumption implicitly holds throughout.

Assumption A1.

The set ${\mathcal{X}}^{\star }≔{argmin}_{x\in \mathcal{X}}\left\{F\left(x\right)\right\}$ is nonempty, and there exists ${\sigma }^2<\infty$ such that for each $x^{\star }\in {\mathcal{X}}^{\star }$ and selection $f\prime \left(x^{\star };s\right)\in \partial f\left(x^{\star };s\right)$, we have $\mathbb{E}\left[{∥f\prime \left(x^{\star };S\right)∥}_2^2\right]\le {\sigma }^2$.

Methods

To make our approach more concrete, we identify several models that fit into our framework. These have appeared in refs. 5, 8, and 9, but we believe a self-contained presentation is beneficial. Each one satisfies our conditions C.i to C.iii. The most common model in stochastic optimization is the first-order model.

Stochastic Subgradient Methods.

The stochastic subgradient method uses the model

$$f_x\left(y;s\right)≔f\left(x;s\right)+⟨f\prime \left(x;s\right),y-x⟩.$$ [3]

Proximal Point Methods.

In the convex setting (8, 20, 21), the stochastic proximal point method uses the model $f_x\left(y;s\right)≔f\left(y;s\right)$; in the weakly convex setting, we regularize and use

$$f_x\left(y;s\right)≔f\left(y;s\right)+\frac{\rho \left(s\right)}{2}{∥y-x∥}_2^2.$$ [4]

Other models require less knowledge than proximal model 4 but preserve structural properties in the original function.

Prox-Linear Model.

Let the function $f$ have the composite structure $f\left(x;s\right)=h\left(c\left(x;s\right);s\right)$, where $h\left(\cdot ;s\right)$ is convex and $c\left(\cdot ;s\right)$ is smooth. The stochastic prox-linear method applies $h$ to a first-order approximation of $c$, using

$$f_x\left(y;s\right)≔h\left(c\left(x;s\right)+\nabla {c\left(x;s\right)}^T\left(y-x\right);s\right).$$ [5]

In the nonstochastic setting, these models are classical (22), while recent work establishes convergence and convergence rates in restrictive stochastic settings (5, 9). When $h$ is $L_h$ Lipschitz and $c$ has an $L_c$-Lipschitz gradient, then $f$ is $\rho =L_h\cdot L_c$-weakly convex.

Example 1 (phase retrieval):

In phase retrieval (23), we wish to recover an object $x^{\star }\in {\mathbb{C}}^n$ from a diffraction pattern $A{x}^{\star }$, where $A\in {\mathbb{C}}^{m\times n}$, but physical sensor limitations mean we observe only amplitudes $b=|A{x}^{\star }{|}^2$. A natural objective is

$$\underset{x\in {\mathbb{C}}^n}{minimize}\frac{1}{m}\sum _{i=1}^{m}f\left(x;\left(a_i,b_i\right)\right),\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}f\left(x;\left(a_i,b_i\right)\right)=\left||⟨a_i,x⟩{|}^2-b_i\right|.$$

This is the composition of $h\left(z\right)=\left|z\right|$ and $c\left(x;\left(a_i,b_i\right)\right)=|⟨<a_i,x⟩{|}^2-b_i$, so $f\left(\cdot ;\left(a_i,b_i\right)\right)$ is $2{∥a_i∥}_2^2$-weakly convex (24).$♢$

Example 2 (matrix completion):

In the matrix completion problem (25), which arises (for example) in the design of recommendation systems, we have a matrix $M\in {\mathbb{R}}^{m\times n}$ with decomposition $M=X_{\star }{Y}_{\star }^T$ for $X_{\star }\in {\mathbb{R}}^{m\times r}$ and $Y_{\star }\in {\mathbb{R}}^{n\times r}$. Based on the incomplete set of known entries $\mathrm{\Omega }\subset \left[m\right]\times \left[n\right]$, our goal is to recover the matrix $M$, giving rise to the objective

$$\underset{X\in {\mathbb{R}}^{m\times r},Y\in {\mathbb{R}}^{n\times r}}{minimize}\frac{1}{\left|\mathrm{\Omega }\right|}\sum _{\left(i,j\right)\in \mathrm{\Omega }}f\left(x_i,y_j;M_{i,j}\right),$$

where $f\left(x,y;z\right)≔|⟨x,y⟩-z|$ and $x_i$, $y_j$ are the $\left(i,j\right)$ rows of $X$ and $Y$. This is the composition of $h\left(z\right)=\left|z\right|$ and $c\left(x,y,z\right)=⟨x,y⟩-z$, so that $f=h○c$ is 1-weakly convex.$♢$

Truncated Models.

The prox-linear model 5 may be challenging to implement for complex compositions (e.g., deep learning). If instead we know a lower bound on $f$, we may incorporate this into the model

$$f_x\left(y;s\right)≔max\left\{f\left(x;s\right)+⟨f\prime \left(x;s\right),y-x⟩,\underset{z\in \mathcal{X}}{inf}f\left(z;s\right)\right\}.$$ [6]

In our examples—linear and logistic regression, phase retrieval, and matrix completion (more generally, typical loss functions in machine learning)—we have ${\text{inf}}_z\hspace{0.17em}f\left(z;s\right)=0$. The assumption that we have a lower bound is thus rarely restrictive. This model satisfies the conditions C.i to C.iii, also satisfying the following condition.

• C.iv)(Lower optimality): For all $s\in \mathcal{S}$ and $x,y\in \mathcal{X}$,

$$f_x\left(y;s\right)\ge \underset{z\in \mathcal{X}}{inf}f\left(z;s\right).$$

As we show, condition C.iv is sufficient to derive several optimality and stability properties.

Stability and Its Consequences

In our initial study of stability in optimization (8), we defined an algorithm as stable if its iterates remain bounded and then showed several consequences of this in convex optimization (which we review presently). Here, we develop 2 important extensions. First, we show that any model satisfying condition C.iv has stable iterates under mild assumptions, in strong contrast to models (e.g., linear) that fail the condition. Second, we develop an analogous stability theory for weakly convex functions, proving that accurate enough models are stable. In parallel to the convex case, stability suffices for more: It implies convergence (with an asymptotic rate) to stationary points for any model-based method on weakly convex functions. Let us formalize stability (8). A pair $\left(\mathcal{F},\mathcal{P}\right)$ is a collection of problems if $\mathcal{P}$ consists of probability measures on a sample space $\mathcal{S}$ and $\mathcal{F}$ of functions $f:\mathcal{X}\times \mathcal{S}\to \mathbb{R}$.

Definition 1.

An algorithm generating iterates $x_k$ according to the model-based update 2 is stable in probability for the class of problems $\left(\mathcal{F},\mathcal{P}\right)$ if for all $f\in \mathcal{F}$, $P\in \mathcal{P}$ defining $F\left(x\right)={\mathbb{E}}_P\left[f\left(x;S\right)\right]$, and ${\mathcal{X}}^{\star }={argmin}_{x\in \mathcal{X}}F\left(x\right)$,

$$\underset{k}{sup}\mathrm{d}\mathrm{i}\mathrm{s}\mathrm{t}\left(x_k,{\mathcal{X}}^{\star }\right)<\infty \hspace{0.17em}\hspace{0.17em}with\hspace{0.17em}probability\hspace{0.17em}1.$$ [7]

Typically, stability 7 requires the standard assumptions

$${\alpha }_k>0,\hspace{0.17em}\hspace{0.17em}\sum _{k\ge 1}{\alpha }_k=\infty ,\hspace{0.17em}\hspace{0.17em}\text{and}\hspace{0.17em}\hspace{0.17em}\sum _{k\ge 1}{\alpha }_k^2<\infty .$$ [8]

Even under these, models such as the linear model 3 and consequent subgradient method are unstable (ref. 8, section 3). They may even cause superexponential divergence.

Example 3 (divergence):

Let $F\left(x\right)=e^x+e^{-x}$, $p<\infty$, and ${\alpha }_0>0$, and let ${\alpha }_k$ satisfy ${\alpha }_k\ge {\alpha }_0{k}^{-p}$. Let $x_{k+1}=x_k-{\alpha }_{k}F\prime \left(x_k\right)=x_k-{\alpha }_k\left(e^{x_k}-e^{-x_k}\right)$ be generated by the gradient method. For large $x_1$, $\log \frac{x_{k+1}}{x_k}\ge 2^k$ for all $k$.$♢$

The Importance of Stability in Stochastic Convex Optimization.

To set the stage for what follows, we begin by motivating the importance of stable procedures. Briefly, any stable aprox model converges for any convex function under weak assumptions, which we now elucidate. First, we make an assumption.

Assumption A2.

There exists ${\mathsf{G}}_{\mathrm{b}\mathrm{i}\mathrm{g}}:{\mathbb{R}}_{+}\to \left[0,\infty \right)$ such that for all $x\in \mathcal{X}$ and each measurable selection $f\prime \left(x;s\right)\in \partial f\left(x;s\right)$,

$$\mathbb{E}\left[{∥f\prime \left(x;S\right)∥}_2^2\right]\le {\mathsf{G}}_{\mathrm{b}\mathrm{i}\mathrm{g}}\left({∥x∥}_2\right).$$

Assumption A2 is equivalent to assuming $\mathbb{E}\left[{∥f\prime \left(x;S\right)∥}_2^2\right]$ is bounded on compact sets; it allows arbitrary growth as long as the subgradients have second moments.

Proposition 1 [Asi and Duchi (8), proposition 1].

Assume that $f\left(\cdot ;s\right)$ is convex for each $s\in \mathcal{S}$ and let Assumption A2 hold. Let the iterates $x_k$ be generated by any method satisfying conditions C.i to C.iii and [8]. On the event $\underset{k}{sup}∥x_k∥\infty$, ${\sum }_k{\alpha }_k\left(F\left(x_k\right)-F\left(x^{\star }\right)\right)<\infty$ and $\mathrm{d}\mathrm{i}\mathrm{s}\mathrm{t}\left(x_k,{\mathcal{X}}^{\star }\right)\hspace{0.17em}\stackrel{a.s.}{\to }\hspace{0.17em}0$.

Proposition 1 establishes convergence of stable procedures and also (via Jensen’s inequality) provides asymptotic rates of convergence for weighted averages ${\sum }_k{\alpha }_k{x}_k/{\sum }_k{\alpha }_k$.

Stability is additionally important when the functions $f$ are smooth: Any stable aprox method achieves asymptotically optimal convergence. In particular, let us assume $F$ is ${\mathcal{C}}^2$ near $x^{\star }={argmin}_{\mathcal{X}}F\left(x\right)$ with ${\nabla }^{2}F\left(x^{\star }\right)\succ 0$, and the $f\left(\cdot ;s\right)$ have an $L\left(s\right)$ -Lipschitz gradient near $x^{\star }$ with $\mathbb{E}\left[L{\left(S\right)}^2\right]<\infty$.

Proposition 2 [Asi and Duchi (8), theorem 2].

In addition to the conditions of Proposition 1, let the conditions of the previous paragraph hold. Then ${\overline{x}}_k=\frac{1}{k}{\sum }_{i=1}^k{x}_i$ satisfies

$$\sqrt{k}\left({\overline{x}}_k-x^{\star }\right)\hspace{0.17em}\stackrel{d}{\to }\hspace{0.17em}\mathsf{N}\left(0,{\nabla }^{2}F{\left(x^{\star }\right)}^{-1}\hspace{0.17em}\mathrm{C}\mathrm{o}\mathrm{v}\left(\nabla f\left(x^{\star };S\right)\right){\nabla }^{2}F{\left(x^{\star }\right)}^{-1}\right).$$

This convergence is optimal for any method (26).

Stability of Lower-Bounded Models for Convex Functions.

With these consequences of stability in hand—convergence and asymptotic optimality—it behooves us to provide conditions sufficient to guarantee stability. To that end, we show that lower-bounded models satisfying condition C.iv are stable in probability (Definition 1) for functions whose (sub)gradients grow at most polynomially. We begin with an assumption.

Assumption A3.

There exist $C<\infty$, $2\le p<\infty$ such that

$$\mathbb{E}\left[{∥f\prime \left(x;S\right)∥}_2^2\right]\le C\left(1+\mathrm{d}\mathrm{i}\mathrm{s}\mathrm{t}{\left(x,{\mathcal{X}}^{\star }\right)}^p\right),\hspace{0.17em}\hspace{0.17em}all\hspace{0.17em}x\in \mathcal{X},$$

and $\mathbb{E}\left[{\left(f\left(x^{\star };S\right)-{\text{inf}}_{z\in \mathcal{X}}f\left(z;S\right)\right)}^{p/2}\right]\le C$ for all $x^{\star }\in {\mathcal{X}}^{\star }$.

The analogous condition (27) for stochastic gradient methods holds for $p=2$, or quadratic growth, without which the method may diverge. In contrast, Assumption A3 allows polynomial growth; for example, the function $f\left(x\right)=x^4$ is permissible, while the gradient method may exponentially diverge even for stepsizes ${\alpha }_k=1/k$. The key consequence of Assumption A3 is that if it holds, truncated models are stable:

Theorem 1.

Assume the function $f\left(\cdot ;s\right)$ is convex for each $s\in \mathcal{S}$. Let Assumption A3 hold and ${\alpha }_k={\alpha }_0{k}^{-\beta }$ with $\frac{p+2}{p+4}<\beta <1$. Let $x_k$ be generated by the iteration 2 with a model satisfying conditions C.i to C.iv. Then

$$\underset{k\in \mathbb{N}}{sup}\mathrm{d}\mathrm{i}\mathrm{s}\mathrm{t}\left(x_k,{\mathcal{X}}^{\star }\right)<\infty \hspace{1em}with\hspace{0.17em}probability\hspace{0.17em}1.$$

Theorem 1 shows that truncated methods enjoy the benefits of stability we outline in Propositions 1 and 2 above. Thus, these models, whose updates are typically as cheap to compute as a stochastic gradient step (especially in the common case that ${\text{inf}}_{z}f\left(z;s\right)=0$) provide substantial advantage over methods using only (sub)gradient approximations.

Stability and Its Consequences for Weakly Convex Functions.

We continue our argument that—if possible—it is beneficial to use more accurate models, even in situations beyond convexity, investigating the stability of proximal models (Eq. 4) for weakly convex functions. Establishing stability in the weakly convex case requires a different approach to the convex case, as the iterates may not make progress toward a fixed optimal set. In this case, to show stability, we require an assumption bounding the size of $f\prime \left(x;S\right)$ relative to the population subgradient $F\prime$.

Assumption A4.

There exist $C_1,C_2<\infty$ such that for all measurable selections $f\prime \left(x;s\right)\in \partial f\left(x;s\right)$ and $F\prime \left(x\right)\in \partial F\left(x\right)$,

$$\mathbb{E}\left[{∥f\prime \left(x;S\right)-F\prime \left(x\right)∥}_2^2\right]\le C_1{∥F\prime \left(x\right)∥}_2^2+C_2.$$

By providing a relative noise condition on $f\prime$, Assumption A4 allows for more than the typical class of functions with global Lipschitz properties (cf. ref. 5), such as the phase retrieval and matrix completion objectives (Examples 1 and 2). It can allow exponential growth, addressing the challenges in Example 3. For example, let $f\left(x;1\right)=e^x$ and $f\left(x;2\right)=e^{-x}$, where $S$ is uniform in $\left\{\mathrm{1,2}\right\}$ so that $F\left(x\right)=\frac{1}{2}\left(e^x+e^{-x}\right)$; then $\mathbb{E}\left[f\prime {\left(x;S\right)}^2\right]=2F\prime {\left(x\right)}^2+1$.

To describe convergence and stability in nonconvex (even nonsmooth) settings, we require appropriate definitions. Finding global minima of nonconvex functions is computationally infeasible (28), so we follow established practice and consider convergence to stationary points via the Moreau envelope (5, 29). To formalize, for $x\in {\mathbb{R}}^n$ and $\lambda \ge 0$, the Moreau envelope and associated proximal map are

$$F_{\lambda }\left(x\right)≔\underset{y\in \mathcal{X}}{inf}\left\{F\left(y\right)+\frac{\lambda }{2}{∥y-x∥}_2^2\right\}\hspace{0.17em}\hspace{0.17em}\text{and}\hspace{0.17em}{\mathrm{p}\mathrm{r}\mathrm{o}\mathrm{x}}_{F/\lambda }\left(x\right)≔\underset{y\in \mathcal{X}}{argmin}\left\{F\left(y\right)+\frac{\lambda }{2}{∥y-x∥}_2^2\right\}.$$

For large $\lambda$, the minimizer $x^{\lambda }≔{\mathrm{p}\mathrm{r}\mathrm{o}\mathrm{x}}_{F/\lambda }\left(x\right)$ is unique whenever $F$ is weakly convex. Adopting the techniques Davis and Drusvyatskiy (5) pioneer for weakly convex problems, we rely on the Moreau envelope’s connections to (near) stationarity:

$$\begin{array}{cc}\hfill & \nabla F_{\lambda }\left(x\right)=\lambda \left(x-x^{\lambda }\right),\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}F\left(x^{\lambda }\right)\le F\left(x\right),\hfill \\ \hfill & \mathrm{d}\mathrm{i}\mathrm{s}\mathrm{t}\left(0,\partial F\left(x^{\lambda }\right)\right)\le {∥\nabla F_{\lambda }\left(x\right)∥}_2.\hfill \end{array}$$ [9]

The 3 properties in [9] imply that any nearly stationary point $x$ of $F_{\lambda }$—when $\Vert {\nabla F_{\lambda }\left(x\right)\Vert }_2$ is small—is close to a nearly stationary point $x^{\lambda }$ of $F$. To prove convergence for weakly convex $F$, then, it suffices to show $\nabla F_{\lambda }\left(x_k\right)\to 0$.

Using full proximal models guarantees convergence.

Theorem 2.

Let Assumption A4 hold, let $\lambda <\infty$ satisfy $\mathbb{E}\left[\rho \left(S\right)\right]<\lambda$, and assume $\underset{x\in \mathcal{X}}{inf}F\left(x\right)>-\infty$ and $\mathbb{E}\left[\rho {\left(S\right)}^2\right]<\infty$. Let $x_k$ follow the iteration 2 with proximal model 4 and stepsizes 8. Then there exists a random variable $G_{\lambda }<\infty$ satisfying

$$F_{\lambda }\left(x_k\right)\to G_{\lambda }\hspace{0.17em}\hspace{0.17em}and\hspace{0.17em}\hspace{0.17em}\sum _k{\alpha }_k{∥\nabla F_{\lambda }\left(x_k\right)∥}_2^2<\infty \hspace{0.17em}\hspace{0.17em}w.p.\hspace{0.17em}1.$$

Theorem 2 shows that $F_{\lambda }\left(x_k\right)$ is bounded almost surely. Thus, if $F$ is coercive, meaning $F\left(x\right)\uparrow \infty$ as $∥x∥\to \infty$, the Moreau envelope $F_{\lambda }$ is coercive, yielding the following.

Corollary 1.

Let the conditions of Theorem 2 hold and let $F$ be coercive. Then

$$\underset{k\in \mathbb{N}}{sup}\mathrm{d}\mathrm{i}\mathrm{s}\mathrm{t}\left(x_k,{\mathcal{X}}^{\star }\right)<\infty \hspace{0.17em}\hspace{0.17em}\hspace{0.17em}with\hspace{0.17em}probability\hspace{0.17em}1.$$

In parallel with our development of the convex case, stability is sufficient to develop convergence results for any aprox method, highlighting its importance. Indeed, we can show that stable methods guarantee convergence, although for probability 1 convergence of the iterates, we require a slightly elaborate assumption (cf. refs. 9 and 30), which rules out pathological limits.

Assumption A5 (Weak Sard).

Let ${\mathcal{X}}_{\mathrm{s}\mathrm{t}\mathrm{a}\mathrm{t}}=\left\{x\mid 0\in \partial F\left(x\right)\right\}$ be the collection of stationary points of $F$ over $\mathcal{X}$. The Lebesgue measure of the image $F\left({\mathcal{X}}_{\mathrm{s}\mathrm{t}\mathrm{a}\mathrm{t}}\right)$ is zero.

Under this assumption, aprox methods converge to stationary points whenever the iterates are stable.

Proposition 3.

Let Assumption A2 hold and the iterates $x_k$ be generated by any method satisfying conditions C.i to C.iii and [8]. Assume that $\lambda$ is large enough that $\mathbb{E}\left[\rho \left(S\right)\right]<\lambda$. There exists a finite random variable $G_{\lambda }$ such that on the event that $\underset{k}{sup}{∥x_k∥}_2<\infty$, with probability 1 we have

$$\sum _k{\alpha }_k{∥\nabla F_{\lambda }\left(x_k\right)∥}_2^2<\infty \hspace{0.17em}\hspace{0.17em}\hspace{0.17em}and\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}{F}_{\lambda }\left(x_k\right)\to G_{\lambda }.$$ [10]

Under Assumption A5, then $\mathrm{d}\mathrm{i}\mathrm{s}\mathrm{t}\left(x_k,{\mathcal{X}}_{\mathrm{s}\mathrm{t}\mathrm{a}\mathrm{t}}\right)\hspace{0.17em}\stackrel{a.s.}{\to }\hspace{0.17em}0$ and ${∥\nabla F_{\lambda }\left(x_k\right)∥}_2\stackrel{a.s.}{\to }\hspace{0.17em}0$.

The condition 10 is enough to develop a conditional ${\ell }_2$-convergence guarantee similar to what stochastic (sub)gradient methods achieve to stationary points for Lipschitz $F$ (5, 31). Indeed, assume ${\alpha }_k={\alpha }_0{k}^{-\beta }$ for some $\beta \in \left(\frac{1}{2},1\right)$ and that the iterates $x_k$ are stable; choose $I_k\in \left\{1,\dots ,k\right\}$ with probability $P\left(I_k=i\right)={\alpha }_i/{\sum }_{j=1}^k{\alpha }_j$. Then inequality 10 shows

$$\underset{k}{lim\; sup}{k}^{1-\beta }\mathbb{E}\left[{∥\nabla F_{\lambda }\left(x_{I_k}\right)∥}_2^2\mid {\mathcal{F}}_k\right]<\infty \hspace{0.17em}\hspace{0.17em}\text{with probability}\hspace{0.17em}1.$$

Fast Convergence for Easy Problems

In many engineering and learning applications, solutions interpolate the data. Consider, for example, signal recovery problems with $b=A{x}^{\star }$ or modern machine-learning applications, where frequently training error is zero (32, 33). We consider such problems here, showing how models that satisfy the lower-bound condition C.iv enjoy linear convergence, extending our earlier results (8) beyond convex optimization.

Definition 2.

Let $F\left(x\right)≔{\mathbb{E}}_P\left[f\left(x;S\right)\right]$. Then $F$ is easy to optimize if for each $x^{\star }\in {\mathcal{X}}^{\star }$ and $P$ almost all $s\in \mathcal{S}$,

$$\underset{x\in \mathcal{X}}{inf}f\left(x;s\right)=f\left(x^{\star };s\right).$$

For such problems, we can guarantee progress toward minimizers for appropriate $f$, as the following lemma shows.

Lemma 1.

Let $F$ be easy to optimize (Definition 2). Let $x_k$ be generated by the updates 2 using a model satisfying conditions C.i to C.iv. Then for any $x^{\star }\in {\mathcal{X}}^{\star }$,

$${∥x_{k+1}-x^{\star }∥}_2^2\le \left(1+{\alpha }_k\rho \left(S_k\right)\right){∥x_k-x^{\star }∥}_2^2-\left[f\left(x_k;S_k\right)-f\left(x^{\star };S_k\right)\right]min\left\{{\alpha }_k,\frac{f\left(x_k;S_k\right)-f\left(x^{\star };S_k\right)}{{∥f\prime \left(x_k;S_k\right)∥}_2^2}\right\}.$$

Lemma 1 allows us to prove fast convergence as long as $f$ grows quickly enough away from $x^{\star }$; a sufficient condition for us is a sharp growth condition away from the optimal set ${\mathcal{X}}^{\star }$. To meld with Lemma 1, we consider the following:

Assumption A6 (Expected Sharp Growth).

There exist constants ${\lambda }_0,{\lambda }_1>0$ such that for $\alpha \in {\mathbb{R}}_{+}$, $x\in \mathcal{X}$, and $x^{\star }\in {\mathcal{X}}^{\star }$,

$$\mathbb{E}\left[min\left\{\alpha ,\frac{f\left(x;S\right)-f\left(x^{\star };S\right)}{\Vert \underset{2}{\overset{2}{f\prime \left(x;S\right)\Vert }}}\right\}\left(f\left(x;S\right)-f\left(x^{\star };S\right)\right)\right]\hspace{1em}\hspace{1em}\hspace{0.17em}\ge \mathrm{d}\mathrm{i}\mathrm{s}\mathrm{t}\left(x,{\mathcal{X}}^{\star }\right)min\left\{{\lambda }_0\alpha ,{\lambda }_1\hspace{0.17em}\mathrm{d}\mathrm{i}\mathrm{s}\mathrm{t}\left(x,{\mathcal{X}}^{\star }\right)\right\}.$$

Assumption A6 is tailored to Lemma 1, so we discuss a few situations where it holds. One sufficient condition is the small-ball condition that there exists $C$ such that $\mathbb{P}\left(f\left(x;S\right)-f\left(x^{\star };S\right)\ge ϵ\hspace{0.17em}\mathrm{d}\mathrm{i}\mathrm{s}\mathrm{t}\left(x,{\mathcal{X}}^{\star }\right)\right)\ge 1-Cϵ$ for $ϵ>0$ and $\mathbb{E}\left[{∥f\prime \left(x;S\right)∥}_2^2\right]\le C\left(1+\mathrm{d}\mathrm{i}\mathrm{s}\mathrm{t}{\left(x,{\mathcal{X}}^{\star }\right)}^2\right)$. We can be more explicit:

Example 4 (Example 1 continued):

Consider the (real-valued) phase retrieval problem with objective $f\left(x;\left(a,b\right)\right)=|{⟨a,x⟩}^2-b|$. Assume the vectors $a_i\in {\mathbb{R}}^n$ are drawn from a distribution satisfying the small-ball condition $P\left(\left|⟨a_i,u⟩\right|\ge ϵ{∥u∥}_2\right)\ge 1-ϵ$ for $ϵ>0$ and any $u\in {\mathbb{R}}^n$ and additionally that $\mathbb{E}\left[{∥a_i∥}_2^2\right]\le M^{2}n$ and $\mathbb{E}\left[{⟨a_i,x⟩}^2\right]\le M^2{∥x∥}_2^2$ for some $M<\infty$. For a sample of size $m$, Assumption A3 holds with high probability for the objective $F\left(x\right)=\frac{1}{m}{\sum }_{i=1}^{m}f\left(x;\left(a_i,b_i\right)\right)$ with ${\lambda }_0=c{∥x^{\star }∥}_2$, and ${\lambda }_1=\frac{c}{16M^{4}n}$, for a numerical constant $c>0$. The full calculation is in SI Appendix.$♢$

The following proposition is our main result in this section, showing lower-bounded models may enjoy linear convergence.

Proposition 4.

Let Assumption A6 hold and $x_k$ be generated by the stochastic iteration 2 using any model satisfying conditions C.i to C.iv, where the stepsizes ${\alpha }_k$ satisfy ${\alpha }_k={\alpha }_0{k}^{-\beta }$ for some $\beta \in \left(\mathrm{0,1}\right)$. If $f\left(\cdot ;S_k\right)$ is $\rho \left(S_k\right)$-weakly convex with $\mathbb{E}\left[\rho \left(S_k\right)\right]=\overline{\rho }$, then for any $m\in \mathbb{N}$and $ϵ>0$, there exists a finite random variable $V_{\infty ,m}<\infty$ such that

$$\frac{\mathrm{d}\mathrm{i}\mathrm{s}\mathrm{t}{\left(x_k,{\mathcal{X}}^{\star }\right)}^2}{{\left(1-{\lambda }_1\right)}^k}\cdot 1\left\{\underset{m\le i\le k-1}{max}\mathrm{d}\mathrm{i}\mathrm{s}\mathrm{t}\left(x_i,{\mathcal{X}}^{\star }\right)\le \frac{{\lambda }_0}{\left(1+ϵ\right)\overline{\rho }}\right\}\stackrel{a.s.}{\to }\hspace{0.17em}{V}_{\infty ,m}.$$

When the functions $f$ are convex, we have $\overline{\rho }=0$, so that Proposition 4 guarantees linear convergence for easy problems. In the case that $\overline{\rho }>0$, the result is conditional: If an aprox method converges to one of the sharp minimizers of $f$, then this convergence is linear (i.e., geometrically fast). In the case of phase retrieval, we can guarantee convergence:

Example 5 (Example 4 continued):

Let $A\in {\mathbb{R}}^{m\times n}$ be a matrix with rows $a_i$ that satisfy the conditions of Example 4. For $F\left(x\right)=\frac{1}{m}{∥|Ax{|}^2-|A{x}^{\star }{|}^2∥}_1$ where $m\gtrsim n$, the truncated model 6 requires overall computation time $O\left(mn\log \frac{1}{ϵ}\right)$ to achieve an $ϵ$-accurate solution to phase retrieval, which is the best-known time complexity. See proof in SI Appendix.$♢$

Experiments

An important question in the development of any optimization method is its sensitivity to algorithm parameters. Consequently, we conclude by experimentally examining convergence time and robustness of each of our optimization methods. We consider each of the models in this paper: the stochastic gradient method (i.e., the linear model 3), the proximal model 4, the prox-linear model 5, and the (lower) truncated model 6.

We test both convergence time and, dovetailing with our focus in this paper, robustness to stepsize for several problems: phase retrieval, matrix completion, and 2 classification problems using deep learning. We consider stepsize sequences of the form ${\alpha }_k={\alpha }_0{k}^{-\beta }$ and perform $K$ iterations over a wide range of different initial stepsizes ${\alpha }_0$. (For brevity, we present results only for the power $\beta =0.6$; experiments with varied $\beta \in \left(\frac{1}{2},1\right)$ were similar.) For a fixed accuracy $ϵ>0$, we record the number of steps $k$ to achieve $F\left(x_k\right)-F\left(x^{\star }\right)\le ϵ$, reporting these times (where we terminate each run at iteration $K$). We perform $T$ experiments for each initial stepsize choice, reporting median time to $ϵ$ accuracy and 90% confidence intervals.

Phase Retrieval.

We start our experiments with the phase retrieval problem in Examples 1 and 4, focusing on the real case for simplicity, where we are given $A\in {\mathbb{R}}^{m\times n}$ with rows $a_i\in {\mathbb{R}}^n$ and $b={\left(A{x}^{\star }\right)}^2\in {\mathbb{R}}_{+}^m$ for some $x^{\star }\in {\mathbb{R}}^n$. Our objective is the nonconvex and nonsmooth function

$$F\left(x\right)=\frac{1}{m}\sum _{i=1}^m\left|{⟨a_i,x⟩}^2-b_i\right|.$$

We sample the entries the vectors $a_i$ and $x^{\star }$ i.i.d. $\mathsf{N}\left(0,I_n\right)$.

We present the results in Fig. 1, comparing the stochastic gradient method 3, the proximal method 4, and the truncated method 6 (whose updates are identical to the prox-linear model 5 in this case). The plots demonstrate the expected result that the stochastic gradient method has good performance in a narrow range of stepsizes, ${\alpha }_1\approx 10$ in this case, while better approximations for aprox yield convergence over a large range of stepsizes. The truncated model 6 exhibits oscillation for large stepsizes, in contrast to the exact model 4.

Fig. 1.

The number of iterations to achieve $ϵ$ accuracy versus initial stepsize ${\alpha }_0$ for phase retrieval with $n=50$, $m=\mathrm{1,000}$. SGM, stochastic gradient method.

Matrix Completion.

For our second experiment, we investigate aprox procedures for the matrix completion problem of Example 2. In this setting, we are given $M=X_{\star }{Y}_{\star }^T$, for $X_{\star }\in {\mathbb{R}}^{m\times r}$ and $Y_{\star }\in {\mathbb{R}}^{n\times r}$, and a set of indexes $\mathrm{\Omega }\subset \left[m\right]\times \left[n\right]$. We aim to recover $M$ observing only ${\left\{M_{ij}\right\}}_{i,j\in \mathrm{\Omega }}$, so our goal is to

$$minimize\hspace{0.17em}F\left(X,Y\right)≔\frac{1}{\left|\mathrm{\Omega }\right|}\sum _{i,j\in \mathrm{\Omega }}\left|X_i^T{Y}_j-M_{i,j}\right|.$$

We optimize over matrices $X\in {\mathbb{R}}^{m\times \widehat{r}}$ and $Y\in {\mathbb{R}}^{n\times \widehat{r}}$, where the estimated rank $\widehat{r}\ge r$. We generate $X_{\star }$ and $Y_{\star }$ by drawing their entries i.i.d. $\mathsf{N}\left(0,1\right)$, choosing $\mathrm{\Omega }$ uniformly at random of size $\left|\mathrm{\Omega }\right|=5\left(nr+mr\right)$. We present the timing results in Fig. 2, which tells a similar story to Fig. 1: Better approximations, such as the truncated models (which again yield identical updates to the prox-linear models 5), are significantly more robust to stepsize specification. The proximal update requires solving a nontrivial quartic, so we omit it.

Fig. 2.

Number of iterations to achieve $ϵ$ accuracy versus initial stepsize ${\alpha }_0$ for matrix completion with $m=\mathrm{2,000}$, $n=\mathrm{2,400}$, $r=5$. Shown are estimated ranks (A) $\widehat{r}=5$ and (B) $\widehat{r}=10$.

Neural Networks.

As one of our main motivations is to address the extraordinary effort—in computational and engineering hours—spent carefully tuning optimization methods, we would be remiss to avoid experiments on deep neural networks. Therefore, in our last set of experiments, we test the performance of our models for training neural networks for classification tasks over the CIFAR10 dataset (34) and the fine-grained 128-class Stanford dog multiclass recognition task (35). For our CIFAR10 experiment, we use the Resnet18 architecture (36); we replace the rectified linear unit (RELU) activations internal to the architecture with exponentiated linear units (ELUs) (37) so that the loss is of composite form $f=h○c$ for $h$ convex and $c$ smooth. For Stanford dogs we use the VGG16 architecture (38) pretrained on Imagenet (39), again substituting ELUs for RELU activations. For this experiment, we also test a modified version of the truncated method, truncadagrad, which uses the truncated model in iteration 2 and a diagonally scaled Euclidean distance (40), updating at iteration $k$ by setting $x_k$ to minimize

$${\left[f\left(x_k;S_k\right)+⟨g_k,x-x_k⟩\right]}_{+}+\frac{1}{2{\alpha }_0}{\left(x-x_k\right)}^T{H}_k\left(x-x_k\right),$$

where $H_k=\mathrm{diag}{\left({\sum }_{i=1}^k{g}_i{g}_i^T\right)}^{1/2}$ for $g_i=f\prime \left(x_i;S_i\right)$. This update requires no more of standard deep-learning software than computing a gradient (backpropagation) and loss. We also compare to adam, the default optimizer in TensorFlow (41).

Figs. 3 and 4 show our results for the CIFAR10 and Stanford dogs datasets, respectively. Fig. 3A and 4A give the number of iterations required to achieve $ϵ$ test-classification error (on the highest or “top-1” predicted class), while Figs. 3B and 4B show the maximal accuracy each procedure achieves for a given initial stepsize ${\alpha }_0$. The plots demonstrate the sensitivity of the standard stochastic gradient method to stepsize choice, which converges only for a small range of stepsizes, in both experiments. ADAM exhibits better robustness for CIFAR10, while it is extremely sensitive in the second experiment (Fig. 4), converging only for a small range of stepsizes—this difference in sensitivities highlights the importance of robustness. In contrast, our procedures using the truncated model are apparently robust for all large enough stepsizes. Figs. 3B and 4B show additionally that the maximal accuracy the 2 truncated methods achieve changes only slightly for ${\alpha }_0\ge 10^{-1}$, again in strong contrast to the other methods, which achieve their best accuracy only for a small range of stepsizes.

Fig. 3.

(A) The number of iterations to achieve $ϵ$ test error versus initial stepsize ${\alpha }_0$ for CIFAR10. (B) The best achieved accuracy after $T=50$ epochs.

Fig. 4.

(A) The number of iterations to achieve $ϵ$ test error versus initial stepsize ${\alpha }_0$ for the Stanford dogs dataset. (B) The best achieved accuracy after $T=30$ epochs.

These results reaffirm the insights from our theoretical results and experiments: It is important and possible to develop methods that enjoy good convergence guarantees and are robust to algorithm parameters.

Data Availability.

All data discussed in this paper are available at GitHub (https://github.com/HilalAsi/APROX-Robust-Stochastic-Optimization-Algorithms) (42).