Stochastic Cubic Regularization for Fast Nonconvex Optimization

Introduction

We consider the problem of nonconvex optimization in the stochastic approximation framework (Robbins and Monro, 1951):

In this setting, we only have access to the stochastic function $f(\mathbf{x};\xi)$, where the random variable $\xi$ is sampled from an underlying distribution $\mathcal{D}$. The task is to optimize the expected function $f(\mathbf{x})$, which in general may be nonconvex. This framework covers a wide range of problems, including the offline setting where we minimize the empirical loss over a fixed amount of data, and the online setting where data arrives sequentially. One of the most prominent applications of stochastic optimization has been in large-scale statistics and machine learning problems, such as the optimization of deep neural networks.

Classical analysis in nonconvex optimization only guarantees convergence to a first-order stationary point (i.e., a point $\mathbf{x}$ satisfying $\norm{\nabla f(\mathbf{x})}=0$), which can be a local minimum, a local maximum, or a saddle point. This paper goes further, proposing an algorithm that escapes saddle points and converges to a local minimum. A local minimum is defined as a point $\mathbf{x}$ satisfying $\norm{\nabla f(\mathbf{x})}=0$ and $\nabla^{2}f(\mathbf{x})\succeq 0$. Finding such a point is of special interest for a large class of statistical learning problems where local minima are global or near-global solutions (e.g. Choromanska et al. (2015); Sun et al. (2016a, b); Ge et al. (2017)).

Among first-order stochastic optimization algorithms, stochastic gradient descent (SGD) is perhaps the simplest and most versatile. While SGD is computationally inexpensive, the best current guarantee for finding an $\epsilon$-approximate local minimum (see Definition 1) requires $\mathcal{O}(\epsilon^{-4}\text{poly}(d))$ iterations (Ge et al., 2015), which is inefficient in the high-dimensional regime.

In contrast, second-order methods which have access to the Hessian of $f$ can exploit curvature to more effectively escape saddles and arrive at local minima. Empirically, second-order methods are known to perform better than first-order methods on a variety of nonconvex problems (Rattray et al., 1998; Martens, 2010; Regier et al., 2017). However, many classical second-order algorithms need to construct the full Hessian, which can be prohibitively expensive when working with large models. Recent work has therefore explored the use of Hessian-vector products $\nabla^{2}f(\mathbf{x})\cdot\mathbf{v}$, which can be computed as efficiently as gradients in many cases including neural networks (Pearlmutter, 1994).

Several algorithms incorporating Hessian-vector products (Carmon et al., 2016; Agarwal et al., 2017) have been shown to achieve faster convergence rates than gradient descent in the non-stochastic setting. However, in the stochastic setting where we only have access to stochastic Hessian-vector products, significantly less progress has been made. This leads us to ask the central question of this paper: Can stochastic Hessian-vector products help to speed up nonconvex optimization?

The proposed algorithm in this paper is based on a classic algorithm in the non-stochastic setting—the cubic regularized Newton method (or cubic regularization) (Nesterov and Polyak, 2006). This algorithm is a natural extension of gradient descent that incorporates Hessian information. Whereas gradient descent finds the minimizer of a local second-order Taylor expansion,

the cubic regularized Newton method finds the minimizer of a local third-order Taylor expansion,

We provide a stochastic variant of this classic algorithm, bridging the gap between its use in the non-stochastic and stochastic settings.

In contrast to prior work, we present a fully stochastic cubic-regularized Newton method: both gradients and Hessian-vector products are observed with noise. Additionally, we provide a non-asymptotic analysis of its complexity.

There has been a recent surge of interest in optimization methods that can escape saddle points and find $\epsilon$-approximate local minima (see Definition 1) in various settings. We provide a brief summary of these results. All iteration complexities in this section are stated in terms of finding approximate local minima, and only highlight the dependency on $\epsilon$ and $d$.

This line of work optimizes over a general function $f$ without any special structural assumptions. In this setting, the optimization algorithm has direct access to the gradient or Hessian oracles at each iteration. The work of Nesterov and Polyak (2006) first proposed the cubic-regularized Newton method, which requires $\mathcal{O}(\epsilon^{-1.5})$ gradient and Hessian oracle calls to find an $\epsilon$-second-order stationary point. Later, the ARC algorithm (Cartis et al., 2011) and trust-region methods (Curtis et al., 2017) were also shown to achieve the same guarantee with similar Hessian oracle access. However, these algorithms rely on having access to the full Hessian at each iteration, which is prohibitive in high dimensions.

1.2 Finite-Sum Setting

1.3 Stochastic Approximation

Preliminaries

We are interested in stochastic optimization problems of the form

where $\xi$ is a random variable with distribution $\mathcal{D}$. In general, the function $f(\mathbf{x})$ may be nonconvex. This formulation covers both the standard offline setting where the objective function can be expressed as a finite sum of $n$ individual functions $f(\mathbf{x},\xi_{i})$, as well as the online setting where data arrives sequentially.

Our goal is to minimize the function $f(\mathbf{x})$ using only stochastic gradients $\nabla f(\mathbf{x};\xi)$ and stochastic Hessian-vector products $\nabla^{2}f(\mathbf{x};\xi)\cdot\mathbf{v}$, where $\mathbf{v}$ is a vector of our choosing. Although it is expensive and often intractable in practice to form the entire Hessian, computing a Hessian-vector product is as cheap as computing a gradient when our function is represented as an arithmetic circuit (Pearlmutter, 1994), as is the case for neural networks.

Throughout the paper, we assume that the function $f(\mathbf{x})$ is bounded below by some optimal value $f^{*}$. We also make following assumptions about function smoothness:

$\rho$-Lipschitz Hessians: for all $\mathbf{x}_{1}$ and $\mathbf{x}_{2}$,

The above assumptions state that the gradient and the Hessian cannot change dramatically in a small local area, and are standard in prior work on escaping saddle points and finding local minima.

Next, we make the following variance assumptions about stochastic gradients and stochastic Hessians:

2 Cubic Regularization

Our target in this paper is to find an $\epsilon$-second-order stationary point, which we define as follows:

For a $\rho$-Hessian Lipschitz function $f$, we say that $\mathbf{x}$ is an $\epsilon$-second-order stationary point (or $\epsilon$-approximate local minimum) if

An $\epsilon$-second-order stationary point not only has a small gradient, but also has a Hessian which is close to positive semi-definite. Thus it is often also referred to as an $\epsilon$-approximate local minimum.

In the deterministic setting, cubic regularization (Nesterov and Polyak, 2006) is a classic algorithm for finding a second-order stationary point of a $\rho$-Hessian-Lipschitz function $f(\mathbf{x})$. In each iteration, it first forms a local upper bound on the function using a third-order Taylor expansion around the current iterate $\mathbf{x}_{t}$:

This is called the cubic submodel. Then, cubic regularization minimizes this cubic submodel to obtain the next iterate: $\mathbf{x}_{t+1}=\operatorname*{argmin}_{\mathbf{x}}m_{t}(\mathbf{x})$. When the cubic submodel can be solved exactly, cubic regularization requires $\mathcal{O}\left(\frac{\sqrt{\rho}(f(\mathbf{x}_{0})-f^{*})}{\epsilon^{1.5}}\right)$ iterations to find an $\epsilon$-second-order stationary point.

To apply this algorithm in the stochastic setting, three issues need to be addressed: (1) we only have access to stochastic gradients and Hessians, not the true gradient and Hessian; (2) our only means of interaction with the Hessian is through Hessian-vector products; (3) the cubic submodel cannot be solved exactly in practice, only up to some tolerance. We discuss how to overcome each of these obstacles in our paper.

Main Results

We begin with a general-purpose stochastic cubic regularization meta-algorithm in Algorithm 1, which employs a black-box subroutine to solve stochastic cubic subproblems. At a high level, in order to deal with stochastic gradients and Hessians, we sample two independent minibatches $S_{1}$ and $S_{2}$ at each iteration. Denoting the average gradient by

this implies a stochastic cubic submodel:

After sampling minibatches for the gradient and the Hessian, Algorithm 1 makes a call to a black-box cubic subsolver to optimize the stochastic submodel $m_{t}(\mathbf{x})$. The subsolver returns a parameter change $\mathbf{\Delta}$, i.e., an approximate minimizer of the submodel, along with the corresponding change in submodel value, $\Delta_{m}\coloneqq m_{t}(\mathbf{x}_{t}+\mathbf{\Delta})-m_{t}(\mathbf{x}_{t})$. The algorithm then updates the parameters by adding $\mathbf{\Delta}$ to the current iterate, and checks whether $\Delta_{m}$ satisfies a stopping condition.

In more detail, the Cubic-Subsolver subroutine takes a vector $\mathbf{g}$ and a function for computing Hessian-vector products $\mathbf{B}[\cdot]$, then optimizes the third-order polynomial

For any fixed, small constant $c$, $\text{Cubic-Subsolver}(\mathbf{g},\mathbf{B}[\cdot],\epsilon)$ terminates within $\mathcal{T}(\epsilon)$ gradient iterations (which may depend on $c$), and returns a $\mathbf{\Delta}$ satisfying at least one of the following:

The first condition is satisfied if the parameter change $\mathbf{\Delta}$ results in submodel and function decreases that are both sufficiently large (Case 1). If that fails to hold, the second condition ensures that $\mathbf{\Delta}$ is not too large relative to the true solution $\mathbf{\Delta}^{\star}$, and that the cubic submodel is solved to precision $c\cdot\rho\norm{\mathbf{\Delta}^{\star}}^{3}$ when $\norm{\mathbf{\Delta}^{\star}}$ is large (Case 2).

As mentioned above, we assume the subsolver uses gradient-based optimization to solve the subproblem so that it will only access the Hessian only through Hessian-vector products. Accordingly, it can be any standard first-order algorithm such as gradient descent, Nesterov’s accelerated gradient descent, etc. Gradient descent is of particular interest as it can be shown to satisfy Condition 1 (see Lemma 2).

Having given an overview of our meta-algorithm and verified the existence of a suitable subsolver, we are ready to present our main theorem:

There exists an absolute constant $c$ such that if $f(\mathbf{x})$ satisfies Assumptions 1, 2, CubicSubsolver satisfies Condition 1 with $c$, $n_{1}\geq\max(\frac{M_{1}}{c\epsilon},\frac{\sigma_{1}^{2}}{c^{2}\epsilon^{2}})\log(\frac{d\sqrt{\rho}\Delta_{f}}{\epsilon^{1.5}\delta c})$, and $n_{2}\geq\max(\frac{M_{2}}{c\sqrt{\rho\epsilon}},\frac{\sigma_{2}^{2}}{c^{2}\rho\epsilon})\log(\frac{d\sqrt{\rho}\Delta_{f}}{\epsilon^{1.5}\delta c})$, then for all $\delta>0$ and $\Delta_{f}\geq f(\mathbf{x}_{0})-f^{*}$, Algorithm 1 will output an $\epsilon$-second-order stationary point of $f$ with probability at least $1-\delta$ within

total stochastic gradient and Hessian-vector product evaluations.

In the limit where $\epsilon$ is small the result simplifies:

If $\epsilon\leq\min\left\{\frac{\sigma_{1}^{2}}{c_{1}M_{1}},\frac{\sigma_{2}^{4}}{c_{2}^{2}M_{2}^{2}\rho}\right\}$, then under the settings of Theorem 1 we can conclude that Algorithm 1 will output an $\epsilon$-second-order stationary point of $f$ with probability at least $1-\delta$ within

total stochastic gradient and Hessian-vector product evaluations.

Finally, we note that lines 8-11 of Algorithm 1 give the termination condition of our meta-algorithm. When the decrease in submodel value $\Delta_{m}$ is too small, our theory guarantees $\mathbf{x}_{t}+\mathbf{\Delta}^{\star}$ is an $\epsilon$-second-order stationary point, where $\mathbf{\Delta}^{\star}$ is the optimal solution of the cubic submodel. However, Cubic-Subsolver may only produce an inexact solution $\mathbf{\Delta}$ satisfying Condition 1, which is not sufficient for $\mathbf{x}_{t}+\mathbf{\Delta}$ to be an $\epsilon$-second-order stationary point. We therefore call Cubic-Finalsolver to solve the subproblem with higher precision. Since Cubic-Finalsolver is invoked only once at the end of the algorithm, we can just use gradient descent, and its runtime will always be dominated by the rest of the algorithm.

One concrete example of a cubic subsolver is a simple variant of gradient descent (Algorithm 3) as studied in Carmon and Duchi (2016). The two main differences relative to standard gradient descent are: (1) lines 1–3: when $\mathbf{g}$ is large, the submodel (Equation 7) may be ill-conditioned, so instead of doing gradient descent, the iterate only moves one step in the $\mathbf{g}$ direction, which already guarantees sufficient descent; (2) line 6: the algorithm adds a small perturbation to $\mathbf{g}$ to avoid a certain “hard” case for the cubic submodel. We refer readers to Carmon and Duchi (2016) for more details about Algorithm 3.

Adapting their result for our setting, we obtain the following lemma, which states that the gradient descent subsolver satisfies our Condition 1.

There exists an absolute constant $c^{\prime}$, such that under the same assumptions on $f(\mathbf{x})$ and the same choice of parameters $n_{1},n_{2}$ as in Theorem 1, Algorithm 3 satisfies Condition 1 with probability at least $1-\delta^{\prime}$ with

Our next corollary applies gradient descent (Algorithm 3) as the approximate cubic subsolver in our meta-algorithm (Algorithm 1), which immediately gives the total number of gradient and Hessian-vector evaluations for the full algorithm.

Under the same settings as Theorem 1, if $\epsilon\leq\min\left\{\frac{\sigma_{1}^{2}}{c_{1}M_{1}},\frac{\sigma_{2}^{4}}{c_{2}^{2}M_{2}^{2}\rho}\right\}$, and if we instantiate the Cubic-Subsolver subroutine with Algorithm 3, then with probability greater than $1-\delta$, Algorithm 1 will output an $\epsilon$-second-order stationary point of $f(\mathbf{x})$ within

total stochastic gradient and Hessian-vector product evaluations.

The overall runtime of Algorithm 1 in Corollary 3 is the number of stochastic gradient and Hessian-vector product evaluations multiplied by the time to compute a gradient or Hessian-vector product. For neural networks, the latter takes $\mathcal{O}(d)$ time.

From Corollary 3, we observe that the dominant term in solving the submodel is $\frac{\sigma_{1}^{2}}{\epsilon^{2}}$ when $\epsilon$ is sufficiently small, giving a total iteration complexity of $\mathcal{O}(\epsilon^{-3.5})$ when other problem-dependent parameters are constant. This improves on the $\mathcal{O}(\epsilon^{-4}\text{poly}(d))$ complexity attained by SGD.

It is reasonable to believe there may be another cubic subsolver which is faster than gradient descent and which satisfies Condition 1 with a smaller $\mathcal{T}(\epsilon)$, for instance a variant of Nesterov’s accelerated gradient descent. However, since the dominating term in our subsolver complexity is $\frac{\sigma_{1}^{2}}{\epsilon^{2}}$ due to gradient averaging, and this is independent of $\mathcal{T}(\epsilon)$, a faster cubic subsolver cannot improve the overall number of gradient and Hessian-vector product evaluations. This means that the gradient descent subsolver already achieves the optimal asymptotic rate for finding an $\epsilon$-second-order stationary point under our stochastic cubic regularization framework.

Proof Sketch

This section sketches the crucial steps needed to understand and prove our main theorem (Theorem 1). We begin by describing our high-level approach, then show how to instantiate this high-level approach in the stochastic setting, assuming oracle access to an exact subsolver. For the case of an inexact subsolver and other proof details, we defer to the Appendix.

Recall that at iteration $t$ of Algorithm 1, a stochastic cubic submodel $m_{t}$ is constructed around the current iterate $\mathbf{x}_{t}$ with the form given in Equation (6):

where $\mathbf{g}_{t}$ and $\mathbf{B}_{t}$ are the averaged stochastic gradients and Hessians. At a high level, we will show that for each iteration, the following two claims hold:

Claim 1. If $\mathbf{x}_{t+1}$ is not an $\epsilon$-second-order stationary point of $f(\mathbf{x})$, the cubic submodel has large descent $m_{t}(\mathbf{x}_{t+1})-m_{t}(\mathbf{x}_{t})$.

Claim 2. If the cubic submodel has large descent $m_{t}(\mathbf{x}_{t+1})-m_{t}(\mathbf{x}_{t})$, then the true function also has large descent $f(\mathbf{x}_{t+1})-f(\mathbf{x}_{t})$.

Given these claims, it is straightforward to argue for the correctness of our approach. We know that if we observe a large decrease in the cubic submodel value $m_{t}(\mathbf{x}_{t+1})-m_{t}(\mathbf{x}_{t})$ during Algorithm 1, then by Claim 2 the function will also have large descent. But since $f$ is bounded below, this cannot happen indefinitely, so we must eventually encounter an iteration with small cubic submodel descent. When that happens, we can conclude by Claim 1 that $\mathbf{x}_{t+1}$ is an $\epsilon$-second-order stationary point.

We note that Claim 2 is especially important in the stochastic setting, as we no longer have access to the true function but only the submodel. Claim 2 ensures that progress in $m_{t}$ still indicates progress in $f$, allowing the algorithm to terminate at the correct time.

In the remaining parts of this section, we discuss why the above two claims hold for an exact solver.

In this setting, $\mathbf{g}_{t}$ and $\mathbf{B}_{t}$ are the averaged gradient and Hessian with sample sizes $n_{1}$ and $n_{2}$, respectively. To ensure the stochastic cubic submodel approximates the exact cubic submodel well, we need large enough sample sizes so that both $\mathbf{g}_{t}$ and $\mathbf{B}_{t}$ are close to the exact gradient and Hessian at $\mathbf{x}_{t}$ up to some tolerance:

We need to ensure that the random vectors/matrices concentrate along an arbitrary direction (depending on $\mathbf{g}_{t}$ and $\mathbf{B}_{t}$). In order to guarantee the uniform concentration in Lemma 4, we can directly apply results from matrix concentration to obtain the desired result (Tropp et al., 2015).

Let $\mathbf{\Delta}_{t}^{\star}=\operatorname*{argmin}_{\mathbf{\Delta}}m_{t}(\mathbf{x}_{t}+\mathbf{\Delta})$, i.e. $\mathbf{x}_{t}+\mathbf{\Delta}_{t}^{\star}$ is a global minimizer of the cubic submodel $m_{t}$. If we use an exact oracle solver, we have $\mathbf{x}_{t+1}=\mathbf{x}_{t}+\mathbf{\Delta}_{t}^{\star}$. In order to show Claim 1 and Claim 2, one important quantity to study is the decrease in the cubic submodel $m_{t}$:

Let $m_{t}$ and $\mathbf{\Delta}_{t}^{\star}$ be defined as above. Then for all $t$,

Lemma 5 implies that in order to prove submodel $m_{t}$ has sufficient function value decrease, we only need to lower bound the norm of optimal solution, i.e. $\norm{\mathbf{\Delta}_{t}^{\star}}$.

Proof sketch for claim 1: Our strategy is to lower bound the norm of $\mathbf{\Delta}_{t}^{\star}$ when $\mathbf{x}_{t+1}=\mathbf{x}_{t}+\mathbf{\Delta}_{t}^{\star}$ is not an $\epsilon$-second-order stationary point. In the non-stochastic setting, Nesterov and Polyak (2006) prove

which gives the desired result. In the stochastic setting, a similar statement holds up to some tolerance:

Under the setting of Lemma 4 with sufficiently small constants $c_{1},c_{2}$,

That is, when $\mathbf{x}_{t+1}$ is not an $\epsilon$-second-order stationary point, we have $\norm{\mathbf{\Delta}_{t}^{\star}}\geq\Omega(\sqrt{\frac{\epsilon}{\rho}})$. In other words, we have sufficient movement. It follows by Lemma 5 that we have sufficient cubic submodel descent.

Proof sketch for claim 2: In the non-stochastic case, $m_{t}(\mathbf{x})$ is by construction an upper bound on $f(\mathbf{x})$. Together with the fact $f(\mathbf{x}_{t})=m_{t}(\mathbf{x}_{t})$, we have:

showing Claim 2 is always true. For the stochastic case, this inequality may no longer be true. Instead, under the setting of Lemma 4, via Lemma 5, we can upper bound the function decrease with an additional error term:

for some sufficiently small constant $c$. Then when $m_{t}(\mathbf{x}_{t+1})-m_{t}(\mathbf{x}_{t})\leq-4c\sqrt{\epsilon^{3}/\rho}$, we have $f(\mathbf{x}_{t+1})-f(\mathbf{x}_{t})\leq\frac{1}{4}[m_{t}(\mathbf{x}_{t+1})-m_{t}(\mathbf{x}_{t})]\leq-c\sqrt{\epsilon^{3}/\rho}$, which proves Claim 2.

Finally, for an approximate cubic subsolver, the story becomes more elaborate. Claim 1 is only “approximately” true, while Claim 2 still holds but for more complicated reasons. We defer to the Appendix for the full proof.

Experiments

In this section, we provide empirical results on synthetic and real-world data sets to demonstrate the efficacy of our approach. All experiments are implemented using TensorFlow (Abadi et al., 2016), which allows for efficient computation of Hessian-vector products using the method described by Pearlmutter (1994).

We begin by constructing a nonconvex problem with a saddle point to compare our proposed approach against stochastic gradient descent. Let $w(x)$ be the W-shaped scalar function depicted in Figure 1, with a local maximum at the origin and two local minima on either side. While we defer the exact form of $w(x)$ to Appendix C, we note here that it has small negative curvature at the origin, $w^{\prime\prime}(0)=-0.2$, and that it has a 1-Lipschitz second derivative. We aim to solve the problem

with independent noise drawn from $\mathcal{N}(0,1)$ added separately to each component of every gradient and Hessian-vector product evaluation. By construction, the objective function has a saddle point at the origin with Hessian eigenvalues of -0.2 and 20, providing a simple but challenging test case where the negative curvature is two orders of magnitude smaller than the positive curvature and is comparable in magnitude to the noise.

We ran our method and SGD on this problem, plotting the objective value versus the number of oracle calls in Figure 3. The batch sizes and learning rates for each method are tuned separately to ensure a fair comparison; see Appendix C for details. We observe that our method is able to escape the saddle point at the origin and converge to one of the global minima faster than SGD, offering empirical evidence in support of our method’s theoretical advantage.

2 Deep Autoencoder

Results on this autoencoding task are presented in Figure 3. In addition to training the model with our method and SGD, we also include results using AdaGrad, a popular adaptive first-order method with strong empirical performance (Duchi et al., 2011). Since the standard MNIST split does not include a validation set, we separate the original training set into 55,000 training images and 5,000 validation images, plotting training error on the former and using the latter to select hyperparameters for each method. More details about our experimental setup can be found in Appendix C.

We observe that stochastic cubic regularization quickly escapes two saddle points and descends toward a local optimum, while AdaGrad takes two to three times longer to escape each saddle point, and SGD is slower still. This demonstrates that incorporating curvature information via Hessian-vector products can assist in escaping saddle points in practice. However, it is worth noting that AdaGrad makes slightly faster progress than our approach after reaching a basin around a local optimum, indicating that adaptivity may provide complementary benefits to second-order information. We leave the investigation of a hybrid method combining both techniques as an exciting direction for future work.

Conclusion

References

Appendix A Proof of Main Results

In this section, we give formal proofs of Theorems 1 and 3. We start by providing proofs of several useful auxiliary lemmas.

Here we remind the reader of the relevant notation and provide further background from Nesterov and Polyak on the cubic-regularized Newton method. We denote the stochastic gradient as

both for iteration $t$. We draw a sufficient number of samples $|S_{1}|$ and $|S_{2}|$ so that the concentration conditions

are satisfied for sufficiently small $c_{1},c_{2}$ (see Lemma 4 for more details). The cubic-regularized Newton subproblem is to minimize

We denote the global optimizer of $m_{t}(\cdot)$ as $\mathbf{x}_{t}+\mathbf{\Delta}_{t}^{\star}$, that is $\mathbf{\Delta}_{t}^{\star}=\operatorname*{argmin}_{z}m_{k}(\mathbf{z}+\mathbf{x}_{k})$.

As shown in Nesterov and Polyak a global optima of Equation (14) satisfies:

Equation (15) is the first-order stationary condition, while Equation (16) follows from a duality argument. In practice, we will not be able to directly compute $\mathbf{\Delta}_{t}^{\star}$ so will instead use a Cubic-Subsolver routine which must satisfy:

For any fixed, small constant $c_{3},c_{4}$, $\text{Cubic-Subsolver}(\mathbf{g},\mathbf{B}[\cdot],\epsilon)$ terminates within $\mathcal{T}(\epsilon)$ gradient iterations (which may depend on $c_{3},c_{4}$), and returns a $\mathbf{\Delta}$ satisfying at least one of the following:

A.2 Auxiliary Lemmas

We begin by providing the proof of several useful auxiliary lemmas. First we provide the proof of Lemma 4 which characterize the concentration conditions.

We can use the matrix Bernstein inequality from Tropp et al. to control both the fluctuations in the stochastic gradients and stochastic Hessians under Assumption 2.

using the triangle inequality and Jensens inequality. A direct application of the matrix Bernstein inequality gives:

Taking $t=c_{1}\epsilon$ gives the result.

once again using the triangle inequality and Jensens inequality. Another application of the matrix Bernstein inequality gives that:

Taking $t=c_{2}\sqrt{\rho\epsilon}$ ensures that the stochastic Hessian-vector products are controlled uniformly over $\mathbf{v}$:

using $n_{2}$ samples with probability $1-\delta^{\prime}$.

Next we show Lemma 5 which will relate the change in the cubic function value to the norm $\norm{\mathbf{\Delta}_{t}^{\star}}$. See 5

Using the global optimality conditions for Equation (14) from Nesterov and Polyak , we have the global optima $\mathbf{x}_{t}+\mathbf{\Delta}_{t}^{\star}$, satisfies:

Together these conditions also imply that:

Now immediately from the definition of stochastic cubic submodel model and the aforementioned conditions we have that:

An identical statement appears as Lemma 10 in Nesterov and Polyak , so this is merely restated here for completeness. ∎

Thus to guarantee sufficient descent it suffices to lower bound the $\norm{\mathbf{\Delta}_{t}^{\star}}$. We now prove Lemma 6, which guarantees the sufficient “movement” for the exact update: $\norm{\mathbf{\Delta}_{t}^{\star}}$. In particular this will allow us to show that when $\mathbf{x}_{t}+\mathbf{\Delta}_{t}^{\star}$ is not an $\epsilon$-second-order stationary point then $\norm{\mathbf{\Delta}_{t}^{\star}}\geq\frac{1}{2}\sqrt{\frac{\epsilon}{\rho}}$. See 6

As a consequence of the global optimality condition, given in Equation (15), we have that:

Moreover, from the Hessian-Lipschitz condition it follows that:

Combining the concentration assumptions with Equation (18) and Inequality (19), we obtain:

An application of the Fenchel-Young inequality to the final term in Equation (20) then yields:

which lower bounds $\norm{\mathbf{\Delta}_{t}^{\star}}$ with respect to the gradient at $\mathbf{x}_{t}$. For the corresponding Hessian lower bound we first utilize the Hessian Lipschitz condition:

followed by the concentration condition and the optimality condition (16). This immediately implies

We consider the case of large gradient and large Hessian in turn (one of which must hold since $\mathbf{x}_{t}+\mathbf{\Delta}_{t}^{\star}$ is not an $\epsilon$-second-order stationary point). There exist $c_{1},c_{2}$ in the following so that we can obtain:

If $\norm{\nabla f(\mathbf{x}_{t}+\mathbf{\Delta}_{t}^{\star})}>\epsilon$, then we have that

If $-\lambda_{n}(\nabla^{2}f(\mathbf{x}_{t}+\mathbf{\Delta}_{t}^{\star}))>\sqrt{\rho\epsilon}$, then we have that $\norm{\mathbf{\Delta}_{t}^{\star}}>\frac{2}{3}\sqrt{\frac{\epsilon}{\rho}}-\frac{2c_{2}}{3}\sqrt{\frac{\epsilon}{\rho}}=\frac{2}{3}(1-c_{2})\sqrt{\frac{\epsilon}{\rho}}>\frac{1}{2}\sqrt{\frac{\epsilon}{\rho}}.$

We can similarly check the lower bounds directly stated are true. Choosing $c_{1}=\frac{1}{200}$ and $c_{2}=\frac{1}{200}$ will verify these inequalities for example. ∎

A.3 Proof of Claim 1

Here we provide a proof of statement equivalent to Claim 1 in the full, non-stochastic setting with approximate model minimization. We focus on the case when the Cubic-Subsolver routine executes Case 2, since the result is vacuously true when the routine executes Case 1. Our first lemma will both help show sufficient descent and provide a stopping condition for Algorithm 1. For context, recall that when $\mathbf{x}_{t}+\mathbf{\Delta}_{t}^{\star}$ is not an $\epsilon$-second-order stationary point then $\norm{\mathbf{\Delta}_{t}^{\star}}\geq\frac{1}{2}\sqrt{\frac{\epsilon}{\rho}}$ by Lemma 6.

If the routine Cubic-Subsolver uses Case 2, and if $\norm{\mathbf{\Delta}_{t}^{\star}}\geq\frac{1}{2}\sqrt{\frac{\epsilon}{\rho}}$, then it will return a point $\mathbf{\Delta}$ satisfying $m_{t}(\mathbf{x}_{t}+\mathbf{\Delta}_{t})\leq m_{t}(\mathbf{x}_{t})-\frac{1-c_{3}}{12}\rho\norm{\mathbf{\Delta}_{t}^{\star}}^{3}\leq\frac{1-c_{3}}{96}\sqrt{\frac{\epsilon^{3}}{\rho}}$.

In the case when $\norm{\mathbf{\Delta}_{t}^{\star}}\geq\frac{1}{2}\sqrt{\frac{\epsilon}{\rho}}$, by the definition of the routine Cubic-Subsolver we can ensure that $m_{t}(\mathbf{x}_{t}+\mathbf{\Delta}_{t})\leq m_{t}(\mathbf{x}_{t}+\mathbf{\Delta}_{t}^{\star})+\frac{c_{3}}{12}\rho\norm{\mathbf{\Delta}_{t}^{\star}}^{3}$ for arbitarily small $c_{3}$ using $\mathcal{T}(\epsilon)$ iterations. We can now combine the aforementioned display with Lemma 5 (recalling that $m_{t}(\mathbf{x}_{t})=f(\mathbf{x}_{t})$) to conclude that:

for suitable choice of $c_{3}$ which can be made arbitarily small. ∎

Assume we are in the setting of Lemma 4 with sufficiently small constants $c_{1},c_{2}$. If $\mathbf{\Delta}$ is the output of the routine Cubic-Subsolver when executing Case 2 and if $\mathbf{x}_{t}+\mathbf{\Delta}_{t}^{\star}$ is not an $\epsilon$-second-order stationary point of $f$, then $m_{t}(\mathbf{x}_{t}+\mathbf{\Delta}_{t})-m_{t}(\mathbf{x}_{t})\leq-\frac{1-c_{3}}{96}\sqrt{\frac{\epsilon^{3}}{\rho}}$.

This is an immediate consequence of Lemmas 6 and 7. ∎

If we do not observe sufficient descent in the cubic submodel (which is not possible in Case 1 by definition) then as a consequence of Claim 1 and Lemma 7 we can conclude that $\norm{\mathbf{\Delta}_{t}^{\star}}\leq\frac{1}{2}\sqrt{\frac{\epsilon}{\rho}}$ and that $\mathbf{x}_{t}+\mathbf{\Delta}_{t}^{\star}$ is an $\epsilon$-second-order stationary point. However, we cannot compute $\mathbf{\Delta}_{t}^{\star}$ directly. So instead we use a final gradient descent loop in Algorithm 2, to ensure the final point returned in this scenario will be an $\epsilon$-second-order stationary point up to a rescaling.

Assume we are in the setting of Lemma 4 with sufficiently small constants $c_{1},c_{2}$. If $\mathbf{x}_{t}+\mathbf{\Delta}_{t}^{\star}$ is an $\epsilon$-second-order stationary point of $f$, and $\norm{\mathbf{\Delta}_{t}^{\star}}\leq\frac{1}{2}\sqrt{\frac{\epsilon}{\rho}}$, then Algorithm 2 will output a point $\mathbf{\Delta}$ such that $\mathbf{x}_{t+1}=\mathbf{x}_{t}+\mathbf{\Delta}$ is a $4\epsilon$-second-order stationary point of $f$.

We first show that $-\lambda_{\min}(\nabla^{2}f(\mathbf{x}_{t+1}))\lesssim\sqrt{\rho\epsilon}$. Since $f$ is $\rho$-Hessian-Lipschitz we have that:

We now show that $\norm{\nabla f(\mathbf{x}_{t+1})}\lesssim\epsilon$ and thus also small. Once again using that $f$ is $\rho$-Hessian-Lipschitz (Lemma 1 in Nesterov and Polyak ) we have that:

Now, by the termination condition in Algorithm 2 we have that $\norm{\mathbf{g}+\mathbf{B}\mathbf{\Delta}+\frac{\rho}{2}\norm{\mathbf{\Delta}}\mathbf{\Delta}}<\frac{\epsilon}{2}$. So,

Using gradient/Hessian concentration with the previous displays we also obtain that:

for sufficiently small $c_{1}$ and $c_{2}$.

A.4 Proof of Claim 2

We now prove our main descent lemma equivalent to Claim 2—this will show if the cubic submodel has a large decrease, then the underlying true function must also have large decrease. As before we focus on the case when the Cubic-Subsolver routine executes Case 2 since the result is vacuously true in Case 1.

Assume we are in the setting of Lemma 4 with sufficiently small constants $c_{1},c_{2}$. If the Cubic-Subsolver routine uses Case 2, and if $m_{t}(\mathbf{x}_{t}+\mathbf{\Delta}_{t})-m_{t}(\mathbf{x}_{t})\leq-(\frac{1-c_{3}}{96})\sqrt{\frac{\epsilon^{3}}{\rho}}$, then $f(\mathbf{x}_{t}+\mathbf{\Delta}_{t})-f(\mathbf{x}_{t})\leq-\left(\frac{1-c_{3}-c_{5}}{96}\right)\sqrt{\frac{\epsilon^{3}}{\rho}}$.

Using that $f$ is $\rho$-Hessian Lipschitz (and hence admits a cubic majorizer by Lemma 1 in Nesterov and Polyak for example) as well as the concentration conditions we have that:

since by the definition the Cubic-Subsolver routine, when we use Case 2 we have that $\norm{\mathbf{\Delta}_{t}}\leq\norm{\mathbf{\Delta}_{t}^{\star}}+c_{4}\sqrt{\frac{\epsilon}{\rho}}$. We now consider two different situations – when $\norm{\mathbf{\Delta}_{t}^{\star}}\geq\frac{1}{2}\sqrt{\frac{\epsilon}{\rho}}$ and when $\norm{\mathbf{\Delta}_{t}^{\star}}\leq\frac{1}{2}\sqrt{\frac{\epsilon}{\rho}}$.

First, if $\norm{\mathbf{\Delta}_{t}^{\star}}\geq\frac{1}{2}\sqrt{\frac{\epsilon}{\rho}}$ then by Lemma 7 we may assume the stronger guarantee that $m_{t}(\mathbf{x}_{t}+\mathbf{\Delta}_{t})-m_{t}(\mathbf{x}_{t})\leq-(\frac{1-c_{3}}{12})\rho\norm{\mathbf{\Delta}_{t}^{\star}}^{3}$. So by considering the above display in Equation (24) we can conclude that:

since the numerical constants $c_{1},c_{2},c_{3}$ can be made arbitrarily small.

Now, if $\norm{\mathbf{\Delta}_{t}^{\star}}\leq\frac{1}{2}\sqrt{\frac{\epsilon}{\rho}}$, we directly use the assumption that $m_{t}(\mathbf{x}_{t}+\mathbf{\Delta}_{t})-m_{t}(\mathbf{x}_{t})\leq-(\frac{1-c_{3}}{96})\sqrt{\frac{\epsilon^{3}}{\rho}}$. Combining with the display in in Equation (24) we can conclude that:

since the numerical constants $c_{1},c_{2},c_{3}$ can be made arbitrarily small. Indeed, recall that $c_{1}$ is the gradient concentration constant, $c_{2}$ is the Hessian-vector product concentration constant, and $c_{3}$ is the tolerance of the Cubic-Subsolver routine when using Case 2. Thus, in both situations, we have that:

denoting $c_{5}=48c_{1}-48c_{2}c_{4}-96c_{1}c_{4}-60c_{2}c_{4}^{2}$ for notational convenience (which can also be made arbitrarily small for sufficiently small $c_{1},c_{2}$). ∎

A.5 Proof of Theorem 1

We now prove the correctness of Algorithm 1. We assume, as usual, the underlying function $f(x)$ possesses a lower bound $f^{*}$.

For notational convenience let Case 1 of the routine Cubic Subsolver satisfy:

and use $K_{2}=\frac{1-c_{3}}{96}$ to denote the descent constant of the cubic submodel in the assumption of Claim 2. Further, let $K_{\text{prog}}=\min\{\frac{1-c_{3}-c_{5}}{96},K_{1}\}$ which we will use as the progress constant corresponding to descent in the underlying function $f$. Without loss of generality, we assume that $-K_{1}\leq-K_{2}$ for convenience in the proof. If $-K_{1}\geq-K_{2}$, we can simply rescale the descent constant corresponding to Case 2 for the cubic submodel, $\frac{1-c_{3}}{96}$, to be equal to $-K_{1}$, which will require shrinking $c_{1},c_{2}$ proportionally to ensure that the rescaled version of the function descent constant, $\frac{1-c_{3}-c_{5}}{96}$, is positive.

Now, we choose $c_{1},c_{2},c_{3}$ so that $K_{2}>0$, $K_{\text{prog}}>0$, and Lemma 6 holds in the aforementioned form. For the correctness of Algorithm 1 we choose the numerical constant in Line 7 as $K_{2}$ – so the “if statement” checks the condition $\Delta m=m_{t}(\mathbf{x}_{t+1})-m_{t}(\mathbf{x}_{t})\geq-K_{2}\sqrt{\frac{\epsilon^{\prime 3}}{\rho}}$. Here we use a rescaled $\epsilon^{\prime}=\frac{1}{4}\epsilon$ for the duration of the proof.

At each iteration the event that the setting of Lemma 4 hold has probability greater then $1-2\delta^{\prime}$. Conditioned on this event let the routine Cubic-Subsolver have a further probability of at most $\delta^{\prime}$ of failure. We now proceed with our analysis deterministically conditioned on the event $E$ – that at each iteration the concentration conditions hold and the routine Cubic-Subsolver succeeds – which has probability greater then $1-3\delta^{\prime}T_{\text{outer}}\geq 1-\delta$ by a union bound for $\delta^{\prime}=\frac{\delta}{3T_{\text{outer}}}$.

Let us now bound the iteration complexity of Algorithm 1 as $T_{\text{outer}}$. We cannot have the “if statement” in Line 7 fail indefinitely. At a given iteration, if the routine Cubic-Subsolver outputs a point $\mathbf{\Delta}$ that satisfies

then by Claim 2 and the definition of Case 1 of the Cubic-Subsolver we also have that:

Note if the Cubic-Subsolver uses Case 1 in this iteration then we will vacuously achieve descent in both the underlying function $f$, and descent in the cubic submodel greater $-K_{1}\sqrt{\frac{\epsilon^{\prime 3}}{\rho}}$. Since $-K_{1}\sqrt{\frac{\epsilon^{\prime 3}}{\rho}}\leq-K_{2}\sqrt{\frac{\epsilon^{\prime 3}}{\rho}}$ by assumption, the algorithm will not terminate early at this iteration. Since the function $f$ is bounded below by $f^{*}$, the event $m_{t}(\mathbf{x}_{t}+\mathbf{\Delta}_{t})-m_{t}(\mathbf{x}_{t})\leq-K_{2}\sqrt{\frac{\epsilon^{\prime 3}}{\rho}}$ which implies $f(\mathbf{x}_{t+1})-f(\mathbf{x}_{t})\leq-K_{\text{prog}}\sqrt{\frac{\epsilon^{\prime 3}}{\rho}}$ can happen at most $T_{\text{outer}}=\lceil\frac{\sqrt{\rho}(f(x_{0})-f^{*})}{K_{\text{prog}}\epsilon^{\prime 1.5}}\rceil$ times.

Thus in the $T_{\text{outer}}$ iterations of Algorithm 1 it must be the case that there is at least one iteration $T$, for which

By the definition of the Cubic-Subsolver procedure and assumption that $-K_{1}\leq-K_{2}$, it must be the case at iteration $T$ the routine Cubic-Subsolver used Case 2. Now by appealing to Claim 1 and Lemma 7 we must have that $\norm{\mathbf{\Delta}_{T}^{\star}}\leq\frac{1}{2}\sqrt{\frac{\epsilon^{\prime}}{\rho}}$ and that $\mathbf{x}_{T}+\mathbf{\Delta}_{T}^{\star}$ is an $\epsilon^{\prime}$-second-order stationary point of $f$. As we can see in Line 7 of Algorithm 1, at iteration $T$ the “if statement” will be true. Hence Algorithm 1 will run the final gradient descent loop (Algorithm 2) at iteration $T$, return the final point and proceed to exit via the break statement. Since the hypotheses of Lemma 8 are satisfiedWe can also see with this rescaled $\epsilon^{\prime}$ the step-size requirement in Lemma 8 will be satisfied. at iteration $T$, Algorithm 2 will return a final point that is an $\epsilon$-second-order stationary point of $f$ as desired. We can verify the global constant $c=\min\{\frac{K_{\text{prog}}}{8},c_{1},c_{2}\}$ satisfies the conditions of the theorem.

We can also now do a careful count of the complexity of Algorithm 1. First, note at each outer iteration of Algorithm 1 we require $n_{1}\geq\max\left(\frac{M_{1}}{c_{1}\epsilon},\frac{\sigma_{1}^{2}}{c_{1}^{2}\epsilon^{2}}\right)\frac{8}{3}\log\frac{2d}{\delta^{\prime}}$ samples to approximate the gradient and and $n_{2}\geq\max(\frac{M_{2}}{c_{2}\sqrt{\rho\epsilon}},\frac{\sigma_{2}^{2}}{c_{2}^{2}\rho\epsilon})\frac{8}{3}\log\frac{2d}{\delta^{\prime}}$ to approximate the Hessian. The union bound stipulates we should take $\delta^{\prime}(\epsilon)=\frac{\delta}{3T_{\text{outer}}}$ to control the total failure probability of Algorithm 1. Then as we can see in the Proof of Theorem 1, Algorithm 1 will terminate in at most

iterations. The inner iteration complexity of the Cubic-Subsolver routine is $\mathcal{T}(\epsilon)$. The routine only requires computing the gradient vector once, but recomputes Hessian-vector products at each iteration.

The total complexity of Hessian-vector product evaluations is:

Finally, recall the proof of Lemma 8 which shows total complexity of the final gradient descent loop, in Algorithm 2, will be subleading in overall gradient and Hessian-vector product complexity. As before, we can verify the global constant $c=\min\{\frac{K_{\text{prog}}}{8},c_{1},c_{2}\}$ satisfies the conditions of the Theorem 1.

Appendix B Gradient Descent as a Cubic Subsolver

Now, at iteration $t$, using the $\rho$-Hessian Lipschitz condition and concentration conditions we obtain:

This establishes sufficient decrease with respect to the true function $f$. We can verify choosing $c_{1}=c_{2}=1/200$ satisfies all the inequalities in this section. Lastly, note the complexity of computing the Cauchy step is in fact $\mathcal{O}(1)$. ∎

B.2 Gradient Descent Loop in Algorithm 3

Assumption A: The step size for the gradient descent scheme satisfies

Assumption B: The initialization for the gradient descent scheme $\mathbf{\Delta}$, satisfies $\mathbf{\Delta}=-r\frac{\mathbf{g}}{\norm{\mathbf{g}}}$, with $0\leq r\leq R_{c}$.

Then, Theorem 3.2 in Carmon and Duchi (restated here for convenience) gives that:

Note that the iterates $\mathbf{\Delta}_{t}$ are iterates generated from the solving the perturbed cubic subproblem with $\mathbf{g}\to\mathbf{g}+\sigma q$—not from solving the original cubic subproblem. This step is necessary to avoid the “hard case” of the non-convex quadratic problems. See Carmon and Duchi for more details.

To apply this bound we will never have access to $\norm{\mathbf{\Delta}^{\star}}$ apriori. However, in the present we need only to use this lemma to conclude sufficient descent when $\mathbf{\Delta}^{\star}$ is not an $\epsilon$-second-order stationary point and hence when $\norm{\mathbf{\Delta}}^{\star}\geq\frac{1}{2}\sqrt{\frac{\epsilon}{\rho}}$.

for sufficiently small $c_{2}$. This completes the proof of the first statement of the Lemma.

Note that we will eventually choose $\delta^{\prime}\sim\mathcal{O}(\frac{1}{\epsilon^{1.5}})$ for our final guarantee. However, this will only contribute logarithmic dependence in $\epsilon$ to our upper bound.

where we define $c_{4}=\sqrt{\frac{c_{3}}{576}}$. ∎

B.3 Proofs of Lemma 2 and Corollary 3

Here we conclude by showing the correctness of Algorithm 3 which follows easily using our previous results. See 2

Finally, assembling all of our results we can conclude that: Using Lemma 2 and Theorem 1 we can immediately see that: See 3

Appendix C Experimental Details

The W-shaped function used in our synthetic experiment is a piecewise cubic function defined in terms of a slope parameter $\epsilon$ and a length parameter $L$:

We set $\epsilon=0.01$ and $L=5$ in our experiment.

For stochastic cubic regularization, we fix $\rho=1$ at the analytic Hessian Lipschitz constant for this problem, and we use 10 inner iterations for each invocation of the cubic subsolver, finding that this yields a good trade-off between progress and accuracy. Then for each method, we perform a grid search over the following hyperparameters:

Step size: $\{c\cdot 10^{-i}:c\in\{1,3\},i\in\{1,2,3,4,5\}\}$

Gradient and Hessian batch sizes are tuned separately for our method. We select the configuration for each method that converges to a global optimum the fastest, provided the objective value stays within 5% of the optimal value after convergence. Since the global optima are located at $(\pm\frac{3}{5},0)$ and each has objective value $-\frac{2}{375}$, this is equivalent to an absolute tolerance of $\frac{1}{3750}=0.0002666\cdots$.

C.2 Deep Autoencoder

Due to computational constraints, we were unable to perform a full grid search over all hyperparameters for every method. As a compromise, we fix all batch sizes and tune the remaining hyperparameters. In particular, we fix the gradient batch size for all methods at 100, a typical value in deep learning applications, and use a Hessian batch size of 10 for stochastic cubic regularization as motivated by the theoretical scaling. Then we perform a grid search over step sizes for each method using the same set of values from the synthetic experiment. We additionally select $\rho$ from $\{0.01,0.1,1\}$ for stochastic cubic regularization, but find that the choice of this value had little effect on final performance. As in the synthetic experiments, we carry out 10 inner iterations per invocation of the cubic subsolver.