# SketchySGD: Reliable Stochastic Optimization via Robust Curvature Estimates

Zachary Frangella
Stanford University
<EMAIL_ADDRESS>
&Pratik Rathore
Stanford University
<EMAIL_ADDRESS>
&Shipu Zhao
Cornell University
<EMAIL_ADDRESS>
&Madeleine Udell
Stanford University
<EMAIL_ADDRESS>

###### Abstract

We introduce SketchySGD, a stochastic quasi-Newton method that uses sketching
to approximate the curvature of the loss function. Quasi-Newton methods are
among the most effective algorithms in traditional optimization, where they
converge much faster than first-order methods such as SGD. However, for
contemporary deep learning, quasi-Newton methods are considered inferior to
first-order methods like SGD and Adam owing to higher per-iteration complexity
and fragility due to inexact gradients. SketchySGD circumvents these issues by
a novel combination of subsampling, randomized low-rank approximation, and
dynamic regularization. In the convex case, we show SketchySGD with a fixed
stepsize converges to a small ball around the optimum at a faster rate than
SGD for ill-conditioned problems. In the non-convex case, SketchySGD converges
linearly under two additional assumptions, interpolation and the Polyak-
Lojaciewicz condition, the latter of which holds with high probability for
wide neural networks. Numerical experiments on image and tabular data
demonstrate the improved reliability and speed of SketchySGD for deep
learning, compared to standard optimizers such as SGD and Adam and existing
quasi-Newton methods.

_Keywords_ Stochastic optimization $\cdot$ Quasi-Newton methods $\cdot$
Preconditioning $\cdot$ Randomized numerical linear algebra

##  1 Introduction

Stochastic first-order methods such as SGD and Adam reign supreme for training
deep learning models (Davis et al., 2020; Duchi et al., 2011; Johnson and
Zhang, 2013; Kingma and Ba, 2014; Sun, 2019) due to their robustness to
nonconvex landscapes and efficient iterations. However, SGD and its variants
alone typically do not suffice to ensure successful training. Indeed, in order
to successfully reach a point with good generalization properties extensive
hyperparameter tuning, e.g., adjusting learning rates, momentum, and decay
schedules, is required. As described by Yao et al. (2021), “…one has to
babysit the optimizer to make sure that training converges to an acceptable
training loss…”

The need to carefully tune the optimizer can be expensive in terms of human
capital, computational costs, and environmental impact. Manual tuning of the
optimizer is costly in terms of man-hours and computational resources, as one
may need to train with several different learning rates, while repeatedly
checking progress, to find a model that yields good performance. In fact, “the
computational cost of [AI] research is increasing exponentially, at a pace
that far exceeds Moore’s Law" (Schwartz et al., 2020), and this computational
burden is compounded further by having to train multiple models with different
hyperparameters. Additionally, heavy computational cost directly correlates
with increased environmental burden, as large-scale deep learning models come
with heavy carbon footprints (Schwartz et al., 2020).

In convex optimization, second-order optimizers based on the Hessian, such as
Newton’s method and quasi-Newton methods converge at super-linear rates under
mild assumptions, making them much faster than first-order methods (Nocedal
and Wright, 1999; Boyd and Vandenberghe, 2004). Given the success of second-
order methods in convex settings, we would hope that these methods outperform
first-order methods such as SGD in deep learning. Alas, second order methods
come with three significant limitations that forbid their use in deep
learning. First, the per-iteration cost of Newton and quasi-Newton methods is
much higher than gradient based methods. These costs are high but surmountable
with an indirect solver like conjugate gradients (CG) with efficient
(backprop-based) Hessian-vector products; we elaborate in Section 2.1. Second,
these methods are not robust to stochastic gradients in theory and in practice
their performance is unreliable; see experiments in Section 4. Third, although
these methods can drive the training loss to a relatively small value in a few
iterations, they often stall at a saddle point that does not generalize well
to the test dataset (Dauphin et al., 2014). The major innovation in this paper
is the use of sketching to produce a new robust stochastic quasi-Newton
framework we call SketchySGD. We show that SketchySGD gives comparable
performance to popular optimizers such as SGD and Adam in the deep learning
setting, while having robust performance with respect to the learning rate. In
addition, we prove that SketchySGD converges at a faster rate than SGD in the
convex setting, and is robust to ill-conditioning.

###  1.1 SketchySGD

SketchySGD solves empirical risk minimization problems of the form

$\textrm{minimize}_{w\in\mathbb{R}^{p}}f(w):=\frac{1}{n}\sum_{i=1}^{n}f_{i}(w)$

given access to a subgradient oracle (computed, e.g., with automatic
differentiation (AD)) for each $f_{i}$. Given iterate $w_{k}$ at iteration
$k$, SketchySGD samples two independent batches $S_{k}$ and
$B_{k}\subseteq[n]$ of data and computes a stochastic gradient
$g_{B_{k}}(w_{k})$ and a randomized low-rank approximation $\hat{H}_{S_{k}}$
to the subsampled Hessian $H_{S_{k}}(w_{k})$. The approximate Hessian may be
constructed from a sketch of the Hessian computed via AD as detailed in
Section 2 at a cost only a constant factor larger than computing a stochastic
gradient. SketchySGD forms a regularized stochastic approximation to the
Hessian $(\hat{H}_{S_{k}}+\rho_{k}I)$ and updates the parameters as

$w_{k+1}=w_{k}-\eta_{k}(\hat{H}_{S_{k}}+\rho_{k}I)^{-1}g_{B_{k}}(w_{k}),$ (1)

where $\eta_{k}$ is the learning rate and $\rho_{k}>0$ is a regularization
parameter. The SketchySGD update may be viewed as a preconditioned stochastic
gradient step with Levenberg-Marquardt regularization (Levenberg, 1944;
Marquardt, 1963) or as a stochastic quasi-Newton method with inexact
gradients.

##### Contributions.

  1. 1.

We develop a new robust stochastic quasi-Newton algorithm that is fast and
generalizes well by accessing only a subsampled Hessian and stochastic
gradient.

  2. 2.

In the convex case, we show that SketchySGD converges to the optimum at a
faster rate than SGD for ill-conditioned problems. In the non-convex case, we
show that SketchySGD converges linearly to the optimum for wide neural
networks under the Polyak-Lojaciewicz condition and interpolation.

  3. 3.

We present logistic regression experiments showing that SketchySGD with a
theoretically-motivated learning rate can match or outperform SGD, Adam, SAGA,
SVRG, and L-BFGS. These experiments corroborate our theory that SketchySGD
converges at a faster rate than SGD.

  4. 4.

We present deep learning experiments showing that SketchySGD reliably yields
models that generalize well across a variety of learning rates on both tabular
and image tasks, while SGD and Adam often yield poor performance without
careful selection of the learning rate.

  5. 5.

We compare SketchySGD to several quasi-Newton methods for deep learning. We
find SketchySGD yields comparable or better performance, is the most robust,
and is either faster or comparable to its competitors in terms of runtime.

##### Outline.

Section 2 describes the SketchySGD algorithm in detail, explaining how to
compute $\hat{H}_{S_{k}}$ and the update in (1) efficiently. Section 2.1
surveys previous work on stochastic quasi-Newton methods, particularly in the
context of machine learning. Section 3 establishes convergence of SketchySGD
in both convex and non-convex settings. Section 4 provides numerical
experiments showing the superiority of SketchySGD relative to competing
optimizers.

##  2 SketchySGD

##### Notation.

Throughout the paper $B_{k}$ and $S_{k}$ denote subsets of $\\{1,\ldots,n\\}$
that are sampled independently and uniformly without replacement. The
corresponding stochastic gradient and minibatch Hessian are given by

$\displaystyle g_{B_{k}}(w)=\frac{1}{b_{g_{k}}}\sum_{i\in B_{k}}g_{i}(w),$
$\displaystyle H_{S_{k}}(w)=\frac{1}{b_{h_{k}}}\sum_{i\in
S_{k}}\nabla^{2}f_{i}(w),$

where $b_{g_{k}}=|B_{k}|,b_{h_{k}}=|S_{k}|.$ For shorthand we often omit the
dependence upon $w$ and simply write $g_{B_{k}}$ and $H_{S_{k}}$. We also
define $H(w)$ as the Hessian of the objective $f$ at $w$, and we set
$M(w)=\max_{1\leq i\leq n}\|\nabla^{2}f_{i}(w)\|$. Given any $\tau>0$, we use
the notation $H^{\tau}_{S_{k}}$ to denote $H_{S_{k}}+\tau I$. We abbreviate
positive-semidefinite as psd. We denote the Loewner order on the convex cone
of psd matrices by $\preceq$, where $A\preceq B$ means $B-A$ is psd. Given a
psd matrix $A\in\mathbb{R}^{p\times p}$, we enumerate its eigenvalues in
descending order,
$\lambda_{1}(A)\geq\lambda_{2}(A)\geq\cdots\geq\lambda_{p}(A)$. Finally given
a psd matrix $A$ and $\tau>0$ we define the effective dimension by
$d_{\textrm{eff}}(\tau)=\textrm{tr}(A(A+\tau I)^{-1})$, which provides a
smoothed measure of the eigenvalues greater than or equal to $\tau$.

##### Hessian vector product via autodiff.

SketchySGD relies on one main computational primitive, a (minibatch) Hessian
vector product (hvp), to compute a low-rank approximation of the (minibatch)
Hessian. An hvp can be computed by automatic differentiation (AD); it requires
roughly twice the computation of one (minibatch) gradient, or $O(p)$ with a
constant batch size. In contrast, explicitly instantiating a Hessian entails a
heavy $O(p^{2})$ storage and $O(np^{2})$ computational cost.

Let $g_{S_{k}}\in\mathbb{R}^{p}$ denote the gradient with respect to the
sampled minibatch $S_{k}\subseteq[n]$, and let $\varphi\in\mathbb{R}^{p}$ be a
vector. Then calculus yields an expression for the minibatch hvp
$H_{S_{k}}\varphi$ as the gradient of $g^{\top}_{S_{k}}\varphi$:

$\frac{\partial}{\partial
w}(g^{T}_{S_{k}}\varphi)=g^{T}_{S_{k}}\frac{\partial\varphi}{\partial
w}+\frac{\partial g_{S_{k}}}{\partial
w}\varphi=0+H_{S_{k}}\varphi=H_{S_{k}}\varphi.$ (2)

That is, given the gradient $g_{S_{k}}$, the hvp $H_{S_{k}}\varphi$ can be
computed by AD on the scalar function $g_{S_{k}}^{T}\varphi$ with respect to
$w$. This technique has been used in several recent papers (Ghorbani et al.,
2019; Yao et al., 2020, 2021) that use Hessian information, and goes back to
Pearlmutter (1994).

##### Randomized low-rank approximation.

The hvp primitive allows for efficient randomized low-rank approximation to
the minibatch Hessian by sketching. Sketching reduces the cost of fundamental
numerical linear algebra operations without much loss in accuracy (Woodruff,
2014; Martinsson and Tropp, 2020) by computing the quantity of interest from a
_sketch_ , or randomized linear image, of a matrix. In particular, sketching
enables efficient computation of a near-optimal low-rank approximation to
$H_{S_{k}}$ (Halko et al., 2011; Cohen et al., 2015; Tropp et al., 2017a).

SketchySGD computes a sketch of the subsampled Hessian using hvps and returns
a randomized low-rank approximation $\hat{H}_{S_{k}}$ of $H_{S_{k}}$ in the
form of an eigendecomposition $\hat{V}\hat{\Lambda}\hat{V}^{T}$, where
$\hat{V}\in\mathbb{R}^{p\times r}$ and $\hat{\Lambda}\in\mathbb{R}^{r\times
r}$. Many algorithms are available to construct a low-rank approximation from
the sketch, including the randomized SVD (Halko et al., 2011). In this paper,
we use the randomized Nyström approximation, following the stable
implementation in Tropp et al. (2017b).

When the loss is convex, $H_{S_{k}}$ is symmetric psd; otherwise, $H_{S_{k}}$
is not guaranteed to be psd. To handle nonconvex problems for which the
subsampled Hessian may not be psd, SketchySGD uses _dynamic regularization_ to
ensure all eigenvalues of the approximation are nonnegative. The resulting
algorithm appears in the supplement as RandPDLowRankApprox.

Algorithm 1 SketchySGD

initialization $w_{0}$, ranks $\\{r_{j}\\}$, regularization $\\{\rho_{k}\\}$,
learning rates $\\{\eta_{k}\\}$ , Hessian update frequency $u$

repeat

Sample independent batches $S_{k}$ and $B_{k}$

Compute stochastic gradient $g_{B_{k}}(w_{k})$ $\triangleright$ via AD

if $k\equiv 0\pmod{u}$ then

$\Phi=\frac{1}{\sqrt{p}}\text{randn}(p,r_{k/u})$ $\triangleright$ Gaussian
test matrix

Compute sketch $Y=H_{S_{k}}(w_{k})\Phi$ $\triangleright$ $r$ hvps

$[\hat{V},\hat{\Lambda}]$ = RandPDLowRankApprox$(Y,\Phi,r)$

Compute $v_{k}=(\hat{H}_{S_{k}}+\rho_{k}I)^{-1}g_{B_{k}}(w_{k})$ via (3)

$w_{k+1}=w_{k}-\eta_{k}v_{k}$ $\triangleright$ Update parameters

until convergence

##### SketchySGD.

As shown in Algorithm 1, the SketchySGD method uses a low-rank approximation
to a minibatch Hessian to perform a stochastic quasi-Newton update with the
minibatch gradient. SketchySGD reestimates the approximation to the minibatch
Hessian at frequency $u$ to ensure fast convergence; empirically, reestimating
every epoch (i.e., one pass through the training data) or two suffices. In
practice, a constant rank $r$ between 50 and 100 works well. While several
hyperparameters appear in the algorithm, including the sequence of ranks
$\\{r_{j}\\}$, regularizations $\\{\rho_{k}\\}$, and learning rates
$\\{\eta_{k}\\}$, one major benefit of SketchySGD is robustness to these
choices. We present default values, which work well and robustly, in the
supplement.

Given the rank-$r$ approximation
$\hat{H}_{S_{k}}=\hat{V}\hat{\Lambda}\hat{V}^{\top}$ to the minibatch Hessian
$H_{S_{k}}$, the main cost of SketchySGD relative to standard SGD is computing
the search direction $v_{k}=(\hat{H}_{S_{k}}+\rho_{k}I)^{-1}g_{B_{k}}$. This
(parallelizable) computation requires $O(pr)$ flops, by the Woodbury formula
(Higham, 2002):

$v_{k}=\hat{V}\left(\hat{\Lambda}+\rho_{k}I\right)^{-1}\hat{V}^{T}g_{B_{k}}+\frac{1}{\rho_{k}}(g_{B_{k}}-\hat{V}\hat{V}^{T}g_{B_{k}}).$
(3)

##### And it’s fast!

The cost of a fresh low-rank Hessian approximation is $O((b_{h_{k}}+r^{2})p)$.
Given this approximation, the cost of each iteration (i.e., quasi-Newton step)
is only $O\left((b_{g_{k}}+r)p\right)$. With a typical batch size of $256$ or
$512$ for SGD and rank $r=100$ for SketchySGD, the complexity per iteration is
only about $20$–$40\%$ larger than standard SGD. With weight-sharing, the time
to compute a backprop step increases relative to the number of parameters,
which makes the additional matrix computations required to implement
SketchySGD even cheaper in comparison.

###  2.1 Comparison to previous work

Many authors have sought to scale Newton’s method and its variants to large-
scale machine learning tasks. Classical approaches include BFGS, its low-
memory variant L-BFGS, and inexact Newton methods, where the sub-problem is
solved approximately (Broyden, 1970; Liu and Nocedal, 1989; Dembo et al.,
1982). Given a function and gradient oracle, BFGS and L-BFGS yield
satisfactory solutions to most optimization problems. However, for very large-
scale problems, the cost of computing an exact gradient or function evaluation
at each iteration is too large; hence methods for these problems must work
with sampled function values and stochastic gradients.

_Stochastic quasi-Newton methods_ address both problems by taking a
generalized Newton step using a stochastic approximation to the Hessian (and,
sometimes, to the gradient as well). Stochastic quasi-Newton methods that use
exact gradients with a stochastic Hessian approximation constructed via
sketching or subsampling include Byrd et al. (2011); Erdogdu and Montanari
(2015); Pilanci and Wainwright (2017); Gower et al. (2019). Others, such as
Roosta-Khorasani and Mahoney (2019); Bollapragada et al. (2019); Meng et al.
(2020), can handle stochastic approximations to both the gradient and Hessian.
Among this work, our convergence analysis in the non-convex case is inspired
by Meng et al. (2020), who identify the key role of interpolation in
establishing convergence without requiring increasingly accurate
approximations to the gradient.

In recent years there has been a flurry of work focusing on quasi-Newton
methods for deep learning. Grosse and Martens (2016) uses a Kroenecker-
factored approximation to the Fisher information matrix, which leads to a
scalable approximate natural gradient method called K-FAC. Yang et al. (2020)
develops another approximate natural gradient method called the Sketchy
Empirical Natural Gradient (SENG) method. The authors find SENG delivers
better performance than K-FAC and first-order methods such as Adam. K-FAC and
SENG are not the only methods that approximate the natural gradient; others
include Roux et al. (2007); Yang et al. (2021a). Gupta et al. (2018)
introduces Shampoo, which preconditions stochastic tensor iterations using the
second moments of accumulated gradients. Taking a more classical approach,
Goldfarb et al. (2020); Yang et al. (2021b) design versions of L-BFGS tailored
to neural networks. Another recent method is AdaHessian (Yao et al., 2021).
AdaHessian reduces the complexity of Newton’s method by replacing the inverse
Hessian with a randomized diagonal approximation. All of these methods have
been observed to improve performance relative to standard optimizers such as
SGD and Adam, and are comparable in cost, excluding K-FAC, and Shampoo, which
are much slower than SGD, Adam, and our method, SketchySGD.

SketchySGD improves on these previous stochastic quasi-Newton methods and
generalizes many of them. In particular, subsampled Newton methods (Roosta-
Khorasani and Mahoney, 2019; Bollapragada et al., 2019; Meng et al., 2020) may
be viewed as a special case of SketchySGD that uses the true minibatch Hessian
instead of its low-rank approximation. We recover the method of Erdogdu and
Montanari (2015) by using exact gradients and a regularized exact low-rank
approximation to the subsampled Hessian. SketchySGD is theoretically much
faster than or comparable to prior work on stochastic quasi-Newton methods.
Table 1 compares the complexity per iteration of SketchySGD to prior methods
that work in the deep learning setting, assuming the Hessian approximation of
SketchySGD is updated once per epoch. Although SketchySGD’s theoretical
complexity is larger than AdaHessian’s, our method generally requires fewer
hvps, since SketchySGD computes $r$ hvps per epoch, while AdaHessian computes
$O(N)$ hvps per epoch, where $N$ is the number of minibatches. For large-scale
problems, $N$ is typically much larger than $r=100$. In our experiments, we
see that SketchySGD is significantly cheaper than all previous stochastic
(quasi-)Newton methods and is comparable to SGD and Adam.

Table 1:  Per-iteration Complexity of Stochastic Quasi-Newton Methods Method | Iteration Complexity | References
---|---|---
Subsampled Newton-CG (inexact gradients) | $\tilde{O}\left(\left(\sqrt{\kappa(H_{S_{k}})}b_{h_{k}}+b_{g_{k}}\right)p\right)$ |  Roosta-Khorasani and Mahoney (2019), Bollapragada et al. (2019), Meng et al. (2020)
AdaHessian | $O(b_{g_{k}}p)$ | Yao et al. (2021)
SENG | $O((b_{g_{k}}+b_{h_{k}}^{2})p)$ | Yang et al. (2020)
SketchySGD | $O\left((b_{g_{k}}+r)p\right)$ | This paper

##  3 Theory

We now prove convergence of SketchySGD both when $f$ is convex and when $f$ is
non-convex but interpolating and satisfies the Polyak-Lojaciewicz (PL)
condition. We begin with an assumption needed throughout.

######  Assumption 1 (Differentiability and smoothness).

The function $f$ is twice differentiable and $L$-smooth. Further, each $f_{i}$
is $L_{i}$-smooth with $L_{i}\leq L_{\textup{max}}$ for every $i=1,\ldots,n$.

###  3.1 Convex case

Our analysis of the convex case requires the bounded iterates, strong
convexity, and the notions of relative smoothness and relative convexity.

######  Assumption 2 (Bounded iterates).

The iterates of Algorithm 1 lie in a compact set $\mathcal{X}$ containing
$w_{\star}$.

######  Assumption 3 (Strong convexity).

The function $f$ is $\mu$-strongly convex for some $\mu>0$.

######  Assumption 4 (Relative smoothness and relative convexity).

There exist constants $\hat{L}\geq\hat{\mu}>0$ such that for all
$w,w^{\prime}\in\mathcal{X}:$

$\displaystyle f(w^{\prime})\leq f(w)+\langle
g(w),w^{\prime}-w\rangle+\frac{\hat{L}}{2}\|w^{\prime}-w\|^{2}_{H(w)},$ (4)
$\displaystyle f(w^{\prime})\geq f(w)+\langle
g(w),w^{\prime}-w\rangle+\frac{\hat{\mu}}{2}\|w^{\prime}-w\|^{2}_{H(w)}.$ (5)

Relative smoothness and relative convexity were introduced in Gower et al.
(2019). These assumptions are weaker than assuming $f$ is $L$-smooth and
$\mu$-strongly convex. Indeed, it is easy to show that if $f$ is $L$-smooth
and $\mu$-strongly convex then (4) and (5) hold, see Gower et al. (2019) for
details. More importantly, (4) and (5) hold with non-vacuous values of
$\hat{L}$ and $\hat{\mu}$. In the case of least-squares $\hat{L}=\hat{\mu}=1$,
and for generalized linear models $\hat{L}$ and $\hat{\mu}$ are $O(1)$ and
independent of the condition number $\kappa=L/\mu$ (Gower et al., 2019). Thus,
for many popular machine learning problems the _relative condition number_
$\hat{\kappa}=\hat{L}/\hat{\mu}$ is a constant independent of the conditioning
of the data.

The worst case $(r+1)$th eigenvalue of the Hessian plays an important role in
our convergence result.

######  Definition 1.

Fix $1\leq r\leq p$, and set

$\lambda^{\star}_{r+1}=\sup_{w\in\mathcal{X}}\lambda_{r+1}(H(w)).$

Observe $\lambda^{\star}_{r+1}\leq L$ and is often significantly smaller. For
the least-squares loss, $\lambda_{r+1}^{\star}=\lambda_{r+1}(H)$ and
$L=\lambda_{1}(H)$. As most machine learning data matrices exhibit polynomial
or exponential spectral decay, (Zhao et al., 2022; Derezinski et al., 2020) we
generally have $\lambda_{r+1}(H)\ll\lambda_{1}(H)$ and hence
$\lambda_{r+1}^{\star}\ll L$.

######  Theorem 5 (SketchySGD convex convergence – informal).

Instate Assumptions 1, 2, 3, and 4. Run Algorithm 1 with batch sizes $b_{g}$,
$b_{h_{k}}=\tilde{O}\left(\frac{M(w_{k})}{\lambda^{\star}_{r+1}}\right)$,
stepsize $\eta=O\left(\hat{\kappa}^{-1}\right)$, fixed regularization
$\rho=\lambda^{\star}_{r+1}$, and at each iteration constructing the
randomized Nyström approximation with rank
$r_{k}=\tilde{O}(d_{\textup{eff}}(\zeta\lambda^{\star}_{r+1}))$ where $r>0$ is
fixed. Then with high probability, for sufficiently small $\epsilon$,

$\mathbb{E}[f(w_{T})]-f(w_{\star})\leq\epsilon$

after
$T=O\left(\frac{\lambda^{\star}_{r+1}\hat{\kappa}}{\mu}\frac{\log\left(\frac{1}{\epsilon}\right)}{\epsilon}\right)$
iterations.

######  Remark 1.

The formal version of Theorem 5 appears in the supplement.

We have the following corollary that gives the improvement ratio of SketchySGD
relative to SGD.

######  Corollary 1.

Let $T_{\textup{SketchySGD}}$ denote the iteration complexity of SketchySGD
from Theorem 5 and $T_{\textup{SGD}}$ denote the iteration complexity of SGD
given from Theorem 4.6 in Gower et al. (2021). Then

$\frac{T_{\textup{SGD}}}{T_{\textup{SketchySGD}}}\geq\frac{\hat{\mu}}{\hat{\kappa}}\frac{L}{30\lambda^{\star}_{r+1}}.$

In particular, in the case of the least-squares loss we have

$\frac{T_{\textup{SGD}}}{T_{\textup{SketchySGD}}}\geq\frac{L}{30\lambda^{\star}_{r+1}}=\frac{\lambda_{1}(H)}{30\lambda_{r+1}(H)}.$

For ill-conditioned problems, we expect the ratio $L/\lambda^{\star}_{r+1}$ to
satisfy $L/\lambda^{\star}_{r+1}\gg 1$, in which case SketchySGD will be
significantly faster than SGD. This is exactly the setting where we would
expect SketchySGD to be useful. Indeed, if the problem is very well-
conditioned than SGD is sufficient as there is no need to precondition.

###  3.2 SketchySGD converges for wide neural networks

In this subsection, we show SketchySGD converges for wide neural networks. The
analysis requires two more assumptions.

######  Assumption 6 (PL condition).

Function $f$ satisfies the PL inequality if

$\frac{1}{2}\|g(w)\|^{2}\geq\theta(f(w)-f(w_{\star})),\;\forall w.$

A function that is $L$-smooth and $\theta$-PL always has $\theta\leq L$; see
more details in the supplement.

######  Assumption 7 (Interpolation).

Any optimal classifier $w_{\star}\in\mathcal{W}_{\star}$ interpolates the
training set: $\|g_{i}(w_{\star})\|=0$ for each $i\in\\{1,\ldots n\\}$.

In the supplement, we show that under Assumptions 6 and 7 SketchySGD converges
linearly to the optimum. We then invoke existing theory that guarantees
interpolation and the PL condition for wide neural nets to show that
SketchySGD converges linearly on wide neural networks.

The PL condition and interpolation are common assumptions in convergence
analysis for non-convex optimization, especially for deep learning (Vaswani et
al., 2019a; Gower et al., 2021). Liu et al. (2022) showed that sufficiently
wide neural networks satisfy the PL condition. This result holds for
multilayer perceptrons (MLPs) as well as more sophisticated architectures such
as ConvNets and ResNets. Further, apriori, over-parameterized neural networks
are believed to interpolate the training dataset (Vaswani et al., 2019b, a;
Meng et al., 2020; Ma et al., 2018), and it is now standard practice for many
problems to train until the model interpolates the training set (or longer)
(Belkin, 2021).

Using the results in Liu et al. (2022), we establish convergence theory for
SketchySGD on wide neural networks. As an example, we state the following
theorem for MLPs. As in Liu et al. (2022), define a scalar-valued $L$-layer
MLP $F_{\textrm{NN}}(w;x)$ with weights $w\in\mathbb{R}^{p}$, width $m^{\ell}$
for each layer $\ell\in\\{1,\ldots,L\\}$, and input $x\in\mathbb{R}^{d}$.
Given a dataset $\\{x_{1},\ldots,x_{n}\\}$, stack the output of the MLP at
each data point into a vector and write the result as a function of the
parameters $w$ as $\mathcal{F}_{\textrm{NN}}(w)\in\mathbb{R}^{n}$. Define the
_tangent kernel_ of $F_{\textrm{NN}}$ at $w$ as
$K(w)=\mathcal{DF}_{\textrm{NN}}(w)\mathcal{DF}_{\textrm{NN}}(w)^{T}$, where
$\mathcal{DF}_{\textrm{NN}}(w)$ denotes the Jacobian.

######  Theorem 8 (SketchySGD converges on wide MLPs).

Consider an MLP $F_{\textup{NN}}(w,x)$ with weights $w_{0}$ initialized
according to a standard Gaussian distribution and fit on a dataset of size $n$
using the square-loss $f(w)=\|\mathcal{F}_{\textup{NN}}(w)-y\|^{2}$. Suppose
the last layer activations $\sigma^{L+1}$ satisfy
$|\nabla\sigma^{L+1}|\geq\gamma>0$, $K(w_{0})$ has positive smallest
eigenvalue $\lambda_{0}$, and at each layer $\ell=1,\ldots,L$ the width
$m^{\ell}$ of the network satisfies

$m^{l}=\tilde{\Omega}\left(\frac{nR^{6L+2}}{(\lambda_{0}-\theta\gamma^{-2})^{2}}\right)$

for any $\theta\in(0,\lambda_{0}\gamma^{2})$. Instate Assumptions 1, 6, 7, 18.
Run SketchySGD from any initialization $w_{0}\in\mathbb{R}^{p}$, with constant
batch sizes for $B_{k}$ and $S_{k}$, and appropriate step size and
regularization sequences. Then with high probability SketchySGD converges
linearly to a global optimizer $w_{\star}$ within the ball $B(w_{0},R)$ for
some $R>0$.

######  Remark 2.

Theorem 8 pertains only to the square loss function. Convergence for other
popular loss functions can be established using ideas from Scaman et al.
(2022).

##  4 Numerical experiments

In this section, we evaluate the performance of SketchySGD through four groups
of experiments. The first group of experiments is in the logistic regression
setting (Section 4.1), where we compare the performance of SketchySGD to SGD,
SVRG (Johnson and Zhang, 2013), SAGA (Defazio et al., 2014), Adam (Kingma and
Ba, 2014), and L-BFGS (Liu and Nocedal, 1989). The second, third, and fourth
groups of experiments are in the deep learning setting. The second and third
groups train classifiers on tabular data (Section 4.2.1) and image data
(Section 4.2.2), respectively. These experiments compare SketchySGD with Adam
and SGD to show its effectiveness and reliability in practice compared to
commonly used methods across different types of tasks; the image
classification experiments also add AdamW (Loshchilov and Hutter, 2019) as a
point of comparison. The fourth group of experiments (Section 4.2.3) compares
SketchySGD to other quasi-Newton methods including AdaHessian (Yao et al.,
2021), Shampoo (Gupta et al., 2018), K-FAC (Grosse and Martens, 2016), and
SENG (Yang et al., 2020) to show the significant speedup and robustness
SketchySGD offers compared to other quasi-Newton methods.

Additional details appear in the appendix and code to reproduce our
experiments is forthcoming on GitHub.

###  4.1 Logistic regression

In the logistic regression setting, we initialize SGD, SVRG, and Adam with a
fixed learning rate $\eta$. For SVRG we update the variance-reduced gradient
at every epoch. The learning rate for L-BFGS and SAGA are automatically
selected by scikit-learn, and the learning rate for SketchySGD is given by
$1/(1+100\hat{\lambda}_{r})$ 111This is a theoretically-motivated choice; see
the supplement for details., where $\hat{\lambda}_{r}$ is the $r$th eigenvalue
of the current approximation $\hat{H}$ to the subsampled Hessian. We update
the approximation $\hat{H}$ every epoch. We run SketchySGD with regularization
$\rho=10^{-3}$ and rank $r=50$. We also use a L2-regularization of
$10^{-2}/n_{\textrm{tr}}$, where $n_{\textrm{tr}}$ is the number of training
samples.

We train logistic regression classifiers on the IJCNN (Prokhorov, ), News20
(Lang, 1995), and RCV1 (Lewis et al., 2004) datasets. For IJCNN, we apply a
random features transformation corresponding to the Gaussian kernel (Rahimi
and Recht, 2007). We run SGD, SVRG, and Adam with learning rates varying from
0.001 to 1, and select the run with the highest peak test accuracy (the
highest test accuracy attained across all epochs and across all learning
rates).

The results are presented in Figure 1. For all four tasks, we observe that
SketchySGD converges much faster than SGD to a model with good test accuracy.
This improved convergence is an empirical confirmation that SketchySGD
converges at a faster rate than SGD for convex problems, as shown in Theorem
5. In addition, SketchySGD outperforms SVRG and L-BFGS and performs comparably
to Adam and SAGA on on all tasks.

(a) IJCNN

(b) News20

(c) RCV1

Figure 1: Performance of SketchySGD, SGD, Adam, SVRG, SAGA, and L-BFGS in the
logistic regression setting. The test accuracy curves for SGD, Adam, and SVRG
are selected according to the peak test accuracy.

###  4.2 Deep learning

In the deep learning setting, we initialize each optimizer with a learning
rate $\eta$, which is decreased by a factor of 10 at a fixed number of epochs
(80, 120, and 160 for image tasks, 10 and 20 for tabular tasks). We use a
weight decay of $5\cdot 10^{-4}$ for all optimizers. We run SketchySGD with
regularization $\rho=\eta$ that decays with the learning rate, and use rank
$r=100$ for all the tasks. (These parameters seem not to require tuning: a
remarkable testament to the reliability of SketchySGD!) We also update the
Hessian approximation $\hat{H}$ every two epochs. All the methods are
implemented using PyTorch, which uses a BSD licence (Paszke et al., 2019).

####  4.2.1 Tabular data

The goal of the tabular tasks is to train an MLP to classify test instances in
the Higgs (Baldi et al., 2014), MiniBooNE (min, ), and Volkert (vol, )
datasets. To assess the robustness of SketchySGD, SGD, and Adam, we train the
MLP models with learning rates varying from 0.0001 to 3 and plot the test
accuracy curves to visualize the robustness of each optimizer to this
hyperparameter.

The results are presented in Figure 2. In each subfigure, the solid lines
indicate the median test accuracy curve for each method (over each learning
rate), and the _75% quantile interval_ is represented by the shaded area and
bounded by the 12.5% and 87.5% quantiles. The 75% quantile interval of
SketchySGD is much narrower than that of Adam and SGD: that is, SketchySGD is
much more robust to different learning rate settings. Second, for all the
tasks, SketchySGD matches or exceeds the peak test accuracy reached by SGD and
Adam, as shown in Table 2.

Table 2:  Peak Test Accuracy for Tabular Tasks Accuracy \ Method | SketchySGD | SGD | Adam
---|---|---|---
Higgs (%) | $75.45$ | $75.47$ | $75.28$
MiniBooNE (%) | $93.40$ | $92.08$ | $92.90$
Volkert (%) | $65.37$ | $58.97$ | $64.24$

(a) Higgs

(b) MiniBooNE

(c) Volkert

Figure 2: SketchySGD converges robustly and nearly as fast or faster than SGD
and Adam on tabular data tasks. The solid line indicates the median of test
accuracy among different learning rate settings, and the shaded area
represents the 75% quantile interval.

####  4.2.2 Image classification

Our image classification tasks train ResNet20 on CIFAR-10 and ResNet32 on SVHN
(Krizhevsky et al., 2009; Netzer et al., 2011; He et al., 2016). Again, we
train the ResNet models with different learning rates sampled from 0.0001 to 3
and report the median and 75% quantile interval of the test accuracy curves.

Figure 3 and Table 3 present the results using the same conventions as in
Section 4.2.1. SketchySGD converges nearly as fast as SGD, Adam, and AdamW and
is more robust to different learning rate settings. SketchySGD reaches peak
test accuracy comparable to SGD, Adam, and AdamW.

Table 3:  Peak Test Accuracy for Image Classification Tasks

Accuracy \ Method | SketchySGD | SGD | Adam | AdamW
---|---|---|---|---
CIFAR-10 (%) | $91.90$ | $92.51$ | $91.31$ | $92.06$
SVHN (%) | $95.97$ | $96.61$ | $96.12$ | $96.58$

(a) CIFAR-10

(b) SVHN

Figure 3: SketchySGD converges robustly and nearly as fast as SGD, Adam, and
AdamW on image classification tasks. The solid line indicates the median of
test accuracy among different learning rate settings, and the shaded area
represents the 75% quantile interval.

(a) Higgs

(b) MiniBooNE

(c) Volkert

Figure 4: SketchySGD performs best among quasi-Newton methods in terms of
robustness, and matches or outperforms the peak test accuracy achieved by
other quasi-Newton methods. The solid line indicates the median of test
accuracy among different learning rate settings, and the shaded area
represents the 75% quantile interval.

####  4.2.3 Quasi-Newton methods comparison

Table 4:  Peak Test Accuracy of Quasi-Newton Methods for Tabular Tasks with 30
Epochs

Accuracy \ Method | SketchySGD | AdaHessian | Shampoo | K-FAC | SENG
---|---|---|---|---|---
Higgs (%) | $75.21$ | $75.07$ | $73.73$ | $72.32$ | $75.33$
MiniBooNE (%) | $92.99$ | $92.59$ | $89.57$ | $92.67$ | $91.70$
Volkert (%) | $62.05$ | $64.21$ | $54.56$ | $67.36$ | $58.75$

Table 5:  Average Per-epoch Time of Quasi-Newton Methods for Tabular Tasks

Time \ Method | SketchySGD | AdaHessian | Shampoo | K-FAC | SENG
---|---|---|---|---|---
Higgs (s) | $16.61$ | $39.90$ | $82.46$ | $31.81$ | $14.40$
MinibooNE (s) | $2.31$ | $15.56$ | $26.35$ | $10.58$ | $1.88$
Volkert (s) | $1.32$ | $13.23$ | $46.20$ | $8.66$ | $1.01$

Experiments in this section compare SketchySGD with the quasi-Newton methods
AdaHessian, Shampoo, K-FAC, and SENG on tabular tasks.

Figure 4 and Table 4 present the results of quasi-Newton methods using the
same conventions as in Section 4.2.1. They show SketchySGD performs best among
quasi-Newton methods in terms of robustness, and matches or outperforms the
peak test accuracy achieved by other quasi-Newton methods. Table 5 provides
the average per-epoch time of quasi-Newton methods on each tabular task, which
shows SketchySGD is the second fastest among all quasi-Newton methods. While
K-FAC sometimes outperforms SketchySGD on the Volkert dataset, it is
unreliable: for example, it performs the worst of all methods on Higgs and for
certain hyperparameter settings it obtains around 40% test accuracy even on
Volkert. K-FAC’s per epoch runtime is also much slower than SketchySGD.
Although SENG is slightly faster than SketchySGD, it has poor robustness and
peak test accuracy.

##  5 Conclusion

In this paper, we have presented SketchySGD, a fast and robust stochastic
quasi-Newton method for machine learning and for deep learning in particular.
SketchySGD uses subsampling, randomized low-rank approximation, and dynamic
regularization to improve conditioning by approximating the curvature of the
loss.

SketchySGD has strong benefits both in theory and in practice. In the convex
case, our theory shows that SketchySGD converges to a neighborhood of the
optimum at a faster rate than SGD. Our logistic regression experiments
validate this claim. In addition, SketchySGD, using a theoretically-motivated
learning rate, outperforms SGD, SVRG, and L-BFGS and performs about as well as
Adam and SAGA in this setting, even when optimizing the learning rate for SGD,
SVRG, and Adam using grid search.

In the non-convex case, SketchySGD converges linearly to the optimum for wide
neural networks. Our deep learning experiments show that SketchySGD reliably
yields models that generalize well across a variety of learning rates on both
tabular and image tasks, while SGD, Adam, and AdamW often fail without
extensive hyperparameter tuning. Furthermore, SketchySGD achieves comparable
performance to quasi-Newton methods for deep learning, such as AdaHessian,
Shampoo, K-FAC, and SENG and runs much faster. Again, SketchySGD is more
robust to the learning rate than these quasi-Newton methods.

We believe widespread adoption of SketchySGD will improve the efficiency of
large-scale deep learning by reducing the need for hyperparameter tuning.
These efficiency gains entail the following benefits: less human effort spent
on hyperparameter tuning, a reduction in computational costs, and lower
greenhouse gas emissions. Thus, SketchySGD has the potential for a significant
positive societal impact that furthers the sustainability goals of green AI.

##  Appendix A RandPDLowRankApprox

In this section we give the pseudocode for the RandPDLowRankApprox algorithm
mentioned in Section 2, which SketchySGD uses to construct the low-rank
approximation $\hat{H}_{S_{k}}$.

Algorithm 2 RandPDLowRankApprox

sketch $Y$ of $H_{S_{k}}$, test matrix $\Phi$, rank $r_{k}$

$\nu=\text{eps}(\text{norm}(Y,2))$ $\triangleright$ compute shift

$Y_{\nu}=Y+\nu\Phi$ $\triangleright$ add shift for stability

$C=\text{chol}(\Phi^{T}Y_{\nu})$ $\triangleright$ Cholesky decomposition

if chol fails then $\triangleright$ Hessian is indefinite

Compute $[W,\Gamma]=\text{eig}(\Phi^{T}Y_{\nu})$ $\triangleright$
$\Phi^{T}Y_{\nu}$ is small and square

Set $\lambda_{k}=\lambda_{\textrm{min}}(\Phi^{T}Y_{\nu})$

$R=W(\Gamma+|\lambda_{k}|I_{k})^{-1/2}W^{T}$ $\triangleright$ $R$ is psd

$B=YR$

$[\hat{V},\Sigma,\sim]=\text{svd}(B,0)$ $\triangleright$ thin SVD

$\hat{\Lambda}=\text{max}\\{0,\Sigma^{2}-(\nu+|\lambda_{k}|)I_{k}\\}$
$\triangleright$ remove shift, compute eigs

return

$B=Y/C$ $\triangleright$ triangular solve

$[\hat{V},\Sigma,\sim]=\text{svd}(B,0)$ $\triangleright$ thin SVD

$\hat{\Lambda}=\text{max}\\{0,\Sigma^{2}-\nu I\\}$ $\triangleright$ remove
shift, compute eigs

$\hat{V},\hat{\Lambda}$

Algorithm 2 gives the pseudocode. $\text{eps}(x)$ is defined as the positive
distance between $x$ and the next largest floating point number of the same
precision as $x$. The test matrix $\Phi$ is the same test matrix used to
generate the sketch $Y$ of $H_{S_{k}}$. The resulting Nyström approximation
$\hat{H}_{S_{k}}$ is given by $\hat{V}\hat{\Lambda}\hat{V}^{T}$. In computing
$R$, the negative square root should be interpreted in the sense of the Moore-
Penrose pseudoinverse, so the inverse of a 0 eigenvalue is 0. The resulting
Nyström approximation is psd but may have 0 eigenvalues. In our algorithms,
this approximation is always used in conjunction with a regularizer to ensure
positive definiteness.

Algorithm 2 follows Tropp et al. (2017b), except that it comes with a dynamic-
regularization fail-safe step in the case the subsampled Hessian is indefinite
and Cholesky fails. Note in the convex case, Cholesky never fails as the
subsampled Hessian is always psd. The fail-safe step is useful in the non-
convex case where $H_{S_{k}}$ may be indefinite, ensuring the resulting
Hessian approximation is always psd, so the regularized Hessian approximation
is always positive definite and invertible.

##  Appendix B Controlling minibatch gradient variance

Central to our analysis is the ability to control the variance of a gradient
minibatch. A key quantity in this regards is the supremum of the variance of
the gradient over the set of minimizers $\mathcal{W}_{\star}$:

$\sigma^{2}_{\star}=\sup_{w_{\star}\in\mathcal{W}_{\star}}\frac{1}{n}\sum^{n}_{i=1}\|g_{i}(w_{\star})\|^{2}.$

We may view $\sigma^{2}_{\star}$ as measure of how well the model fits the
training set. Indeed, in the extreme case when the model interpolates the
dataset, $\sigma^{2}_{\star}=0$. Thus the smaller $\sigma_{\star}^{2}$, the
better the model fit. The fact that $\sigma_{\star}^{2}$ helps control the
variance of minibatch gradient is shown in the following proposition due to
(Gower et al., 2021).

######  Proposition 9 (Proposition 3.2 (Gower et al., 2021)).

Suppose each $f_{i}$ is $L_{i}$-smooth and at least one of the following
conditions holds:

  1. 1.

Each $f_{i}$ is convex

  2. 2.

$f$ interpolates the training set.

Suppose the gradient sample $g_{B_{k}}$ uses batch-size $b_{g}$. Then for
every $w\in\mathbb{R}^{p}$,

$\mathbb{E}_{B_{k}}\|g_{B_{k}}(w)\|^{2}\leq\|g(w)\|^{2}+4\alpha(f(w)-f(w_{\star}))+2\sigma^{2},$

where $\alpha=L_{\textup{max}}\frac{n-b_{g}}{(n-1)b_{g}}$ and
$\sigma^{2}=\frac{1}{b_{g}}\frac{n-b_{g}}{n-1}\sigma_{\star}^{2}$.

In particular, under interpolation $\sigma_{\star}^{2}=0$, so

$\mathbb{E}_{B_{k}}\|g_{B_{k}}(w)\|^{2}\leq\|g(w)\|^{2}+4\alpha(f(w)-f(w_{\star})).$

Proposition 9 shows that the variance of a minibatch gradient is controlled by
the distance to the optimum, the batch-size $b_{g}$, and $\sigma_{\star}^{2}$.
We see that as we increase the batch-size $b_{g}$, the variance of the
minibatch gradient decreases. More interestingly, as $\sigma^{2}_{\star}$
decreases the minibatch gradient variance decreases, and the variance vanishes
when the model class perfectly fits the data. As items 1. or 2. in Proposition
9 always hold in this paper, we may use this proposition to control the
minibatch gradient variance, which is crucial for establishing Theorem 5 and
Theorem 19.

##  Appendix C SketchySGD convergence – convex case

###  C.1 Preliminary technical lemmas

Recall the SketchySGD update is given by

$x_{k+1}=x_{k}-\eta_{k}P_{k}^{-1}g_{B_{k}},$

where $\mathbb{E}_{B_{k}}[g_{B_{k}}]=g_{k}$. We start by making the following
observation about the SketchySGD update.

######  Lemma 10 (SketchySGD is SGD in preconditioned space).

At iteration $k$ define $f_{P_{k}}(z)=f(P_{k}^{-1/2}z)$, that is define the
change of variable $x=P_{k}^{-1/2}z$. Then,

$\displaystyle g_{P_{k}}(z)=P_{k}^{-1/2}g(P_{k}^{-1/2}z)$ $\displaystyle
H_{P_{k}}(z)=P_{k}^{-1/2}H(x)P^{-1/2}_{k}.$

Hence the SketchySGD update may be realized as

$\displaystyle z_{k+1}=z_{k}-\eta_{k}\bar{g}_{P_{k}}(z_{k})$ $\displaystyle
x_{k+1}=P_{k}^{-1/2}z_{k+1},$

where $\bar{g}_{P_{k}}(z_{k})=P_{k}^{-1/2}g_{B_{k}}(P_{k}^{-1/2}z_{k}).$

###### Proof.

The first display of equations follow from the definition of the change of
variable and the chain rule, while the last display follows from definition of
the SketchySGD update and the first display. ∎

The preceding lemma shows SketchySGD update may be realized by performing a
stochastic gradient step in preconditioned space, and then converting back to
the original space. This view will help in establishing the convergence rate
of SketchySGD in the convex case.

######  Definition 2.

We define
$\sigma^{2}_{P}\coloneqq\frac{1}{b_{g}}\frac{n-b_{g}}{n-1}\sup_{k\geq
0}\frac{1}{n}\sum_{i=1}^{n}\|P_{k}^{-1/2}\nabla f_{i}(w_{\star})\|^{2}$.

The quantity $\sigma^{2}_{P}$ is the analogue of $\sigma^{2}$ from Proposition
9 in preconditioned space. $\sigma^{2}_{P}$ will be a critical quantity in
showing convergence of SketchySGD in the convex case.

##### Controlling the quality of the preconditioner

In what follows we show that $\hat{H}_{S_{k}}^{\rho_{k}}$ is close to the
exact Hessian $H$ in the Loewner ordering, which shows it makes excellent
preconditioner. Below, we state Lemma 12 from Ye et al. (2021) that shows if
we sample appropriately then the subsampled Hessian is close to the full
Hessian.

######  Lemma 11 (Closeness in Loewner ordering between $H$ and $H_{S}$).

Let $H$ be the Hessian of $f$, $M(w)=\max_{1\leq i\leq
n}\|\nabla^{2}f_{i}(w)\|$ and let $\xi>0$ and $\delta>0$ be given. Suppose we
construct the subsampled Hessian $H_{S}$ with
$|S|=O(\frac{M(w)}{\xi}\log(\frac{p}{\delta}))$ samples. Then for
$\zeta_{0}=\max\left\\{\frac{3\xi+\mu}{3\xi+3\mu},\frac{1}{2}\right\\}\in(0,1)$
the event

$\mathcal{E}^{(1)}=\\{(1-\zeta_{0})(H_{S}+\xi I)\preceq
H\preceq(1+\zeta_{0})(H_{S}+\xi I)\\}$

holds with probability at least $1-\delta$.

Similarly, we may control the approximation quality of the randomized low-rank
approximation used in SketchySGD via the following technical result from Zhao
et al. (2022).

######  Lemma 12 (Controlling low-rank approximation error).

Let $\tau>0$ and $E_{k}=H_{S_{k}}-\hat{H}_{S_{k}}$. Construct a randomized
Nyström approximation from a standard Gaussian random matrix $\Omega$ with
rank $r_{k}=O(d_{\textup{eff}}(\tau)+\log(\frac{1}{\delta}))$. Then the event

$\mathcal{E}^{(2)}_{k}=\\{\|E_{k}\|\leq\tau\\}$ (6)

holds with probability at least $1-\delta$.

Lemma 12 allows us to establish the following result which shows our low-rank
approximation is close to the regularized subsampled Hessian in the Loewner
ordering.

######  Lemma 13 (Closeness in Loewner ordering between $H_{S_{k}}^{\xi}$ and
$\hat{H}_{S_{k}}^{\xi}$).

Let $0<\zeta<1$. Suppose at iteration $k$ SketchySGD uses a low-rank
approximation $\hat{H}_{S_{k}}$ to $H_{S_{k}}$ with rank
$r_{k}=O(d_{\textup{eff}}(\zeta\xi)+\log(\frac{1}{\delta}))$. Then with
probability at least $1-\delta$,

$(1-\zeta_{0})\hat{H}_{S_{k}}^{\xi}\preceq
H_{S_{k}}^{\xi}\preceq(1+\zeta_{0})(1+\zeta)\hat{H}_{S_{k}}^{\xi}$ (7)

###### Proof.

Let $E_{k}=H_{S_{k}}-\hat{H}_{S_{k}}$, and note by the properties of the
Nyström approximation that $E_{k}\succeq 0$ (Tropp et al., 2017b). Now, notice
the event

$\mathcal{E}^{(2)}_{k}=\\{\|E_{k}\|\leq\zeta\xi\\}$

holds with probability at least $1-\delta$ by Lemma 12. Henceforth, all
analysis is conditioned on $\mathcal{E}_{k}$, and therefore the conclusions
hold with probability at least $1-\delta$. Now, the regularized Hessian and
approximate Hessian satisfy

$H_{S_{k}}^{\xi}=\hat{H}^{\xi}_{S_{k}}+E_{k},$ (8)

Let $P=\hat{H}^{\xi}_{S_{k}}$. Then combining (8) with Weyl’s inequalities
yields

$\displaystyle\lambda_{1}(P^{-1/2}H_{S_{k}}^{\xi}P^{-1/2})\leq\lambda_{1}\left(P^{-1/2}\hat{H}^{\xi}_{S_{k}}P^{-1/2}\right)+\lambda_{1}\left(P^{-1/2}E_{k}P^{-1/2}\right)$
$\displaystyle=$ $\displaystyle 1+\|P^{-1/2}E_{k}P^{-1/2}\|\leq
1+\|P^{-1}\|\|E_{k}\|\leq 1+\frac{\|E_{k}\|}{\xi}\leq 1+\zeta.$

To bound the smallest eigenvalue, observe that

$H_{S_{k}}^{\xi}\succeq\hat{H}^{\xi}_{S_{k}}\quad\implies\quad
H_{S_{k}}^{\xi}\succeq P.$

Thus, conjugating the preceding relation by $P^{-1/2}$ we obtain

$P^{-1/2}H_{S_{k}}^{\xi}P^{-1/2}\succeq I_{p},$

where $I_{p}$ is the $p\times p$ identity matrix. The preceding inequality
immediately yields

$\lambda_{p}(P^{-1/2}H_{S_{k}}^{\xi}P^{-1/2})\geq 1.$

Hence,

$1\leq\lambda_{p}(P^{-1/2}H_{S_{k}}^{\xi}P^{-1/2})\leq\lambda_{1}(P^{-1/2}H_{S_{k}}^{\xi}P^{-1/2})\leq
1+\zeta.$

As an immediate consequence, we obtain the Loewner ordering relation

$I_{p}\preceq P^{-1/2}H_{S_{k}}^{\xi}P^{-1/2}\preceq(1+\zeta)I_{p}$

which we conjugate by $P^{1/2}$ to show

$\hat{H}^{\xi}_{S_{k}}\preceq
H_{S_{k}}^{\xi}\preceq(1+\zeta)\hat{H}^{\xi}_{S_{k}}.$

The claimed result follows. ∎

Combining Lemma 11 and Lemma 13 immediately yields the following corollary.
This corollary shows our regularized low-rank approximation is close to the
true Hessian and so makes an excellent preconditioner.

######  Corollary 2 (Closeness in Loewner ordering between $H$ and
$\hat{H}^{\xi}$).

Instate the hypotheses of Lemma 11 and Lemma 13 with
$\delta^{\prime}=\frac{\delta}{2}$, where $\delta\in(0,1)$ is given. Then with
probability at least $1-\delta$

$(1-\zeta_{0})\hat{H}_{S_{k}}^{\xi}\preceq
H\preceq(1+\zeta_{0})(1+\zeta)\hat{H}_{S_{k}}^{\xi}.$

Hence

$1-\zeta_{0}\leq\lambda_{p}\left((\hat{H}^{\xi}_{S_{k}})^{-1/2}H(\hat{H}^{\xi}_{S_{k}})^{-1/2}\right)\leq\lambda_{1}\left((\hat{H}^{\xi}_{S_{k}})^{-1/2}H(\hat{H}^{\xi}_{S_{k}})^{-1/2}\right)\leq(1+\zeta_{0})(1+\zeta).$

######  Corollary 3 (Union bound).

Let $\mathcal{E}=\bigcap_{k=0}^{t}\mathcal{E}_{k}$, where
$\mathcal{E}_{k}=\left(\mathcal{E}^{(1)}_{k}\cap\mathcal{E}^{(2)}_{k}\right)$
and suppose that at iteration $k$ we construct $\hat{H}_{S_{k}}$ with rank
$r_{k}=\tilde{O}(d_{\textup{eff}}(\zeta\rho_{k}))$, where $\\{\rho_{k}\\}$ is
the regularization sequence. Then

${\mathbb{P}}(\mathcal{E})\geq 1-\delta.$

###### Proof.

Instating the hypotheses of Lemma 11 and Lemma 13 with
$\delta^{\prime}=\frac{\delta}{2(t+1)}$, we obtain for each
$k\in\\{0,\ldots,t\\}$ that

${\mathbb{P}}(\mathcal{E}_{k}^{C})\leq\frac{\delta}{t+1}.$

Consequently,

${\mathbb{P}}(\mathcal{E})=1-{\mathbb{P}}(\mathcal{E}^{C})\geq
1-\sum^{t}_{k=0}\frac{\delta}{t+1}=1-\delta.$

∎

######  Lemma 14 (Gradient-based suboptimality bound).

Assume the conclusion of Corollary 2 holds. Then

$2(1-\zeta_{0})\hat{\mu}\left(f(w_{k})-f(w_{\star})\right)\leq\|P_{k}^{-1/2}g_{k}\|^{2}.$

###### Proof.

Recall $f_{P_{k}}(z)=f(P_{k}^{-1/2}w)$ so Lemma 10 and relative convexity of
$f$ (5) yield

$\displaystyle f_{P_{k}}(z_{\star})$ $\displaystyle\geq
f_{P_{k}}(z_{k})+\langle
g_{P_{k}}(z_{k}),z_{\star}-z_{k}\rangle+\frac{\hat{\mu}}{2}\|z_{\star}-z_{k}\|_{H_{P_{k}}(z_{k})}^{2}$
$\displaystyle\geq f_{P_{k}}(z_{k})+\langle
g_{P_{k}}(z_{k}),z_{\star}-z_{k}\rangle+\frac{(1-\zeta_{0})\hat{\mu}}{2}\|z_{\star}-z_{k}\|^{2}$
$\displaystyle\geq\textrm{min}_{z^{\prime}}\left\\{f_{P_{k}}(z_{k})+\langle
g_{P_{k}}(z_{k}),z^{\prime}-z_{k}\rangle+\frac{(1-\zeta_{0})\hat{\mu}}{2}\|z^{\prime}-z_{k}\|^{2}\right\\}$
$\displaystyle=f_{P_{k}}(z_{k})-\frac{1}{2(1-\zeta_{0})\hat{\mu}}\|g_{P_{k}}(z_{k})\|^{2}.$

Rearranging, we find

$2(1-\zeta_{0})\hat{\mu}\left(f_{P_{k}}(z_{k})-f_{P_{k}}(z_{\star})\right)\leq\|g_{P_{k}}(z_{k})\|^{2}.$

The desired inequality now follows from the preceding display and Lemma 10. ∎

###  C.2 Proof of Theorem 5

Before giving the proof we provide a more detailed statement of Theorem 5,
which includes the exact value of $\eta$.

######  Theorem 15 (SketchySGD convex convergence).

Instate Assumptions 1, 2 and 4. Let
$\zeta_{0}=\frac{3\lambda^{\star}_{r+1}+\mu}{3\lambda^{\star}_{r+1}+3\mu}$.
Run Algorithm 1 with batch sizes $b_{g}$,
$b_{h_{k}}=O\left(\frac{M(w_{k})}{\lambda^{\star}_{r+1}}\log(\frac{p}{\delta})\right)$,
stepsize
$\eta=\frac{\hat{\mu}}{(1+\zeta)(1+\zeta_{0})\hat{L}}\min\left\\{\frac{(1-\zeta_{0})\epsilon}{2\sigma_{P}^{2}},\frac{(1-\zeta_{0})}{(1-\zeta_{0})\hat{\mu}+2\alpha}\right\\}$,
fixed regularization $\rho=\lambda^{\star}_{r+1}$, and at each iteration
constructing the randomized Nyström approximation with rank
$r_{k}=O(d_{\textup{eff}}(\zeta\lambda^{\star}_{r+1})+\log(\frac{1}{\delta}))$
where $r>0$ is fixed. Further, condition on the event
$\mathcal{E}=\cap_{t=1}^{T}\mathcal{E}_{t}$, which holds with probability at
least $1-\delta$. Additionally, suppose
$\epsilon\leq\frac{2\sigma_{P}^{2}}{\hat{\mu}_{P}+2\alpha}$. Then after
$t\geq(1+\zeta)\left(2+\frac{3\lambda_{r+1}^{\star}}{\mu}\right)\frac{\hat{L}}{\hat{\mu}}\frac{2\sigma_{P}^{2}}{(1-\zeta_{0})\hat{\mu}\epsilon}\log\left(\frac{2(f(w_{0})-f(w_{\star}))}{\epsilon}\right)$
iterations,

$\mathbb{E}[f(w_{t})]-f(w_{\star})\leq\epsilon.$

###### Proof.

For convenience in the proof below we define the following quantities:

$P_{k}\coloneqq\hat{H}_{S_{k}}^{\lambda^{\star}_{r+1}},$
$\hat{L}_{P}\coloneqq(1+\zeta)(1+\zeta_{0})\hat{L},$
$\hat{\mu}_{P}\coloneqq(1-\zeta_{0})\hat{\mu}.$

We also observe $\sigma_{P}^{2}$ from Definition 2 is finite as for any $k$ we
have,

$\frac{1}{n}\sum_{i=1}^{n}\|P_{k}^{-1/2}\nabla
f_{i}(w_{\star})\|^{2}\leq\frac{1}{\lambda^{\star}_{r+1}}\frac{1}{n}\sum_{i=1}^{n}\|\nabla
f_{i}(w_{\star})\|^{2}.$

With these preliminaries, we may now begin the proof.

Observe that relative smoothness implies the inequality

$f(w_{k+1})\leq f(w_{k})-\eta\langle
g_{k},P^{-1}_{k}g_{B_{k}}\rangle+\frac{\eta^{2}\hat{L}}{2}\|P_{k}^{-1}g_{B_{k}}\|^{2}_{H_{k}}.$

Hence

$\displaystyle f(w_{k+1})$ $\displaystyle\leq f(w_{k})-\eta\langle
g_{k},P^{-1}_{k}g_{B_{k}}\rangle+\frac{\eta^{2}\hat{L}}{2}\|P_{k}^{-1/2}g_{B_{k}}\|^{2}_{P_{k}^{-1/2}H_{k}P_{k}^{-1/2}}$
$\displaystyle\overset{(i)}{\leq}f(w_{k})-\eta\langle
g_{k},P^{-1}_{k}g_{B_{k}}\rangle+\frac{\eta^{2}(1+\zeta_{0})(1+\zeta)\hat{L}}{2}\|P_{k}^{-1/2}g_{B_{k}}\|^{2}$
$\displaystyle\overset{(ii)}{=}f(w_{k})-\eta\langle
g_{k},P^{-1}_{k}g_{B_{k}}\rangle+\frac{\eta^{2}\hat{L}_{P}}{2}\|P_{k}^{-1/2}g_{B_{k}}\|^{2}$

where (i) follows from Corollary 2 whose conclusions hold as we are
conditioned on $\mathcal{E}$, and (ii) follows from definition of
$\hat{L}_{P}$. Next we take the expectation with respect to $B_{k}$ to obtain

$\displaystyle\mathbb{E}_{B_{k}}[f(w_{k+1})]\leq
f(w_{k})-\eta\|P^{-1/2}g_{k}\|^{2}+\frac{\hat{L}_{P}\eta^{2}}{2}\mathbb{E}_{B_{k}}\left[\|P_{k}^{-1/2}g_{B_{k}}\|^{2}\right].$

We now bound $\mathbb{E}_{B_{k}}\left[\|P_{k}^{-1/2}g_{B_{k}}\|^{2}\right]$.
By Lemma 10 we have
$\mathbb{E}_{B_{k}}\left[\|P_{k}^{-1/2}g_{B_{k}}(w_{k})\|^{2}\right]=\mathbb{E}_{B_{k}}\left[\|P_{k}^{-1/2}g_{B_{k}}(P_{k}^{-1/2}z_{k})\|^{2}\right]=\mathbb{E}_{B_{k}}\left[\|\bar{g}_{P_{k}}(z_{k})\|^{2}\right]$.
Thus we may apply the expected residual condition (Proposition 9) on the
stochastic gradient of the preconditioned objective to conclude

$\mathbb{E}_{B_{k}}\left[\|P_{k}^{-1/2}g_{B_{k}}\|^{2}\right]\leq\|P_{k}^{-1/2}g_{k}\|^{2}+4\alpha(f_{P_{k}}(z_{k})-f_{P_{k}}(z_{\star}))+2\sigma_{P}^{2}.$

Hence

$\displaystyle\mathbb{E}_{B_{k}}[f(w_{k+1})]$ $\displaystyle\leq
f(w_{k})-\eta\|P_{k}^{-1/2}g_{k}\|^{2}+\frac{\hat{L}_{P}\eta^{2}}{2}\left(\|P_{k}^{-1/2}g_{k}\|^{2}+4\alpha(f_{P_{k}}(z_{k})-f_{P_{k}}(z_{\star}))+2\sigma_{P}^{2}\right)$
$\displaystyle=f(w_{k})-\eta\left(1-\frac{\eta\hat{L}_{P}}{2}\right)\|P_{k}^{-1/2}g_{k}\|^{2}+\frac{\hat{L}_{P}\eta^{2}}{2}\left(4\alpha(f(w_{k})-f(w_{\star}))+2\sigma_{P}^{2}\right),$

where in the last equality we have used that
$f_{P_{k}}(z_{k})-f_{P_{k}}(z_{\star})=f(w_{k})-f(w_{\star})$ by Lemma 10. As
$\eta\leq 1/\hat{L}_{P}$, we see $1-\frac{\eta\hat{L}_{P}}{2}\geq 0$, so
invoking Lemma 14 we obtain

$\displaystyle\mathbb{E}_{B_{k}}[f(x_{k+1})]$ $\displaystyle\leq
f(w_{k})-2\hat{\mu}_{P}\eta\left(1-\frac{\hat{L}_{P}\eta}{2}\right)(f(w_{k})-f(w_{\star}))$
$\displaystyle+\frac{\hat{L}_{P}\eta^{2}}{2}\left(4\alpha(f(w_{k})-f(w_{\star}))+2\sigma_{P}^{2}\right)$
$\displaystyle=f(w_{k})-\eta\hat{\mu}_{P}\left(2-\hat{L}_{P}\eta\left(1+\frac{2\alpha}{\hat{\mu}_{P}}\right)\right)(f(w_{k})-f(w_{\star}))+\eta^{2}\hat{L}_{P}\sigma_{P}^{2},$

where in the last equality we have gathered like terms. Now as
$\eta=\frac{1}{\hat{L}_{P}}\min\left\\{\frac{\epsilon\hat{\mu}_{P}}{2\sigma_{P}^{2}},\frac{\hat{\mu}_{P}}{\hat{\mu}_{P}+2\alpha}\right\\}$,
we see

$2-\hat{L}_{P}\eta\left(1+\frac{2\alpha}{\hat{\mu}_{P}}\right)\geq 1.$

Hence we find (after subtracting $f(w_{\star})$ from both sides),

$\displaystyle\mathbb{E}_{B_{k}}[f(w_{k+1})]-f(w_{\star})$
$\displaystyle\leq\left(1-\eta\hat{\mu}_{P}\right)(f(w_{k})-f(w_{\star}))+\hat{L}_{P}\eta^{2}\sigma_{P}^{2}.$

Now, taking total expectation over all timesteps and recursing yields

$\displaystyle\mathbb{E}[f(w_{t})]-f(w_{\star})$
$\displaystyle\leq\left(1-\eta\hat{\mu}_{P}\right)^{t}(f(w_{0})-f(w_{\star}))+\eta\frac{\hat{L}_{P}}{\hat{\mu}_{P}}\sigma_{P}^{2}.$

Using that $\eta\leq\epsilon\hat{\mu}_{P}/(2\hat{L}_{P}\sigma_{P}^{2})$, the
last display simplifies to

$\displaystyle\mathbb{E}[f(w_{t})]-f(w_{\star})$
$\displaystyle\leq\left(1-\eta\hat{\mu}_{P}\right)^{t}(f(w_{0})-f(w_{\star}))+\frac{\epsilon}{2}.$

From this last display it immediately follows that

$\mathbb{E}[f(w_{t})]-f(w_{\star})\leq\epsilon$

after
$t\geq(1+\zeta)\frac{1+\zeta_{0}}{1-\zeta_{0}}\frac{\hat{L}}{\hat{\mu}}\max\left\\{\frac{2\sigma_{P}^{2}}{(1-\zeta_{0})\hat{\mu}\epsilon},1+\frac{2\alpha}{(1-\zeta_{0})\hat{\mu}}\right\\}\log\left(\frac{2(f(w_{0})-f(w_{\star}))}{\epsilon}\right)$
iterations. After some algebra we find

$\frac{1+\zeta_{0}}{1-\zeta_{0}}=2+\frac{3\lambda_{r+1}^{\star}}{\mu}.$

Furthermore, by our hypothesis on $\epsilon$ the maximum in the bound on the
iteration complexity is attained by the first term. Thus, after
$t\geq(1+\zeta)\left(2+\frac{3\lambda_{r+1}^{\star}}{\mu}\right)\frac{\hat{L}}{\hat{\mu}}\frac{2\sigma_{P}^{2}}{(1-\zeta_{0})\hat{\mu}\epsilon}\log\left(\frac{2(f(w_{0})-f(w_{\star}))}{\epsilon}\right)$
iterations

$\mathbb{E}[f(w_{t})]-f(w_{\star})\leq\epsilon,$

yielding the complexity statement in the theorem. ∎

###  C.3 Proof of Corollary 1

We start by stating an abbreviated version of Theorem 4.6 in Gower et al.
(2021) for convex functions.

######  Theorem 16 (Theorem 4.6 in Gower et al. (2021)).

Let $f$ be $L$-smooth and $\mu$-strongly convex and let $\alpha$ and
$\sigma^{2}$ be defined as in Proposition 9. Let $\epsilon>0$ and suppose we
run SGD for $T$ iterations with step size
$\gamma=\frac{1}{L}\min\\{\frac{\mu\epsilon}{2\sigma^{2}},\frac{1}{1+2\alpha/\mu}\\}$.
Then

$T\geq\frac{L}{\mu}\max\left\\{\frac{2\sigma^{2}}{\mu\epsilon},1+\frac{2\alpha}{\mu}\right\\}\log\left(\frac{2(f(w_{0})-f(w_{\star}))}{\epsilon}\right)\implies{\mathbb{E}}[f(w_{T})]-f(w_{\star})\leq\epsilon.$

Now we provide the proof of Corollary 1. The idea of the proof is to compare
the number of iterations that guarantee SketchySGD and SGD to reach an
$\epsilon$-suboptimal point in expectation.

###### Proof.

From the proof of Theorem 15, for sufficiently small $\epsilon$, after

$(1+\zeta)\frac{1+\zeta_{0}}{1-\zeta_{0}}\frac{\hat{L}}{\hat{\mu}}\frac{2\sigma_{P}^{2}}{(1-\zeta_{0})\hat{\mu}\epsilon}\log\left(\frac{2(f(w_{0})-f(w_{\star}))}{\epsilon}\right)\eqqcolon
T_{\textrm{SketchySGD}}$

iterations, SketchySGD is guaranteed to reach an $\epsilon$-suboptimal point
in expectation. Similarly, Theorem 16 guarantees that for sufficiently small
$\epsilon$, SGD will reach an $\epsilon$-suboptimal point in expectation after

$\frac{L}{\mu}\frac{2\sigma^{2}}{\mu\epsilon}\log\left(\frac{2(f(w_{0})-f(w_{\star}))}{\epsilon}\right)\eqqcolon
T_{\textrm{SGD}}$

iterations. Taking the ratio of $T_{\textrm{SGD}}$ to
$T_{\textrm{SketchySGD}}$, we find

$\frac{T_{\textrm{SGD}}}{T_{\textrm{SketchySGD}}}=\frac{\sigma^{2}}{\sigma_{P}^{2}}\frac{\kappa}{\hat{\kappa}}\frac{\hat{\mu}}{\mu}\frac{1-\zeta_{0}}{1+\zeta}\left(\frac{3\lambda_{r+1}^{\star}}{\mu}+2\right)^{-1},$

where we used $\kappa=\frac{L}{\mu},\hat{\kappa}=\frac{\hat{L}}{\hat{\mu}}$,
and $\frac{1+\zeta_{0}}{1-\zeta_{0}}=\frac{3\lambda_{r+1}^{\star}}{\mu}+2$.
Since $\sigma_{P}^{2}\leq\frac{1}{\lambda_{r+1}^{\star}}\sigma^{2}$, we have

$\frac{T_{\textrm{SGD}}}{T_{\textrm{SketchySGD}}}\geq\lambda_{r+1}^{\star}\frac{\kappa}{\hat{\kappa}}\frac{\hat{\mu}}{\mu}\frac{1-\zeta_{0}}{1+\zeta}\left(\frac{3\lambda_{r+1}^{\star}}{\mu}+2\right)^{-1}.$

We now develop a lower bound on the term $\frac{1-\zeta_{0}}{1+\zeta}$. Since
$0<\zeta<1$, $\frac{1}{1+\zeta}>\frac{1}{2}$. In addition,
$1-\zeta_{0}=\frac{2\mu}{3\zeta+3\mu}>\frac{\mu}{3\lambda_{r+1}^{\star}}$,
since $\lambda_{r+1}^{\star}\geq\mu$. Therefore

$\displaystyle\frac{T_{\textrm{SGD}}}{T_{\textrm{SketchySGD}}}$
$\displaystyle\geq\lambda_{r+1}^{\star}\frac{\kappa}{\hat{\kappa}}\frac{\hat{\mu}}{\mu}\frac{\mu}{3\lambda_{r+1}^{\star}}\frac{1}{2}\left(\frac{3\lambda_{r+1}^{\star}}{\mu}+2\right)^{-1}$
$\displaystyle=\frac{\hat{\mu}}{\hat{\kappa}}\frac{\kappa}{6}\left(\frac{3\lambda_{r+1}^{\star}}{\mu}+2\right)^{-1}.$

Since $\lambda_{r+1}^{\star}\geq\mu$,
$\frac{3\lambda_{r+1}^{\star}}{\mu}+2\leq\frac{3\lambda_{r+1}^{\star}}{\mu}+\frac{2\lambda_{r+1}^{\star}}{\mu}=\frac{5\lambda_{r+1}^{\star}}{\mu}$.
Substituting this inequality into the previous display gives

$\displaystyle\frac{T_{\textrm{SGD}}}{T_{\textrm{SketchySGD}}}$
$\displaystyle\geq\frac{\hat{\mu}}{\hat{\kappa}}\frac{\kappa}{6}\left(\frac{5\lambda_{r+1}^{\star}}{\mu}\right)^{-1}$
$\displaystyle=\frac{\hat{\mu}}{\hat{\kappa}}\frac{L}{30\lambda_{r+1}^{\star}},$

which is the claimed ratio. For least-squares problems,
$\hat{\mu}=\hat{\kappa}=1$, implying
$\frac{T_{\textrm{SGD}}}{T_{\textrm{SketchySGD}}}\geq\frac{L}{30\lambda_{r+1}^{\star}}=\frac{\lambda_{1}(H)}{30\lambda_{r+1}(H)}$.
∎

##  Appendix D Convergence of SketchySGD – non-convex case

###  D.1 Relationship between $L$-smoothness and $\theta$-PL

Here we show that just as in the case of $\mu$-strong convexity where $\mu\leq
L$ , we must have that $\theta\leq L$. Thus, for a $L$-smooth and $\theta$-PL
function, $L/\theta$ may be viewed as the problem condition number.

######  Lemma 17.

A function that is $L$-smooth and $\theta$-PL always has $L\geq\theta$.

###### Proof.

Recall that an $L$-smooth function satisfies the inequality (see for instance
Theorem 2.1.5 in (Nesterov, 2018))

$\frac{\|\nabla f(w)\|^{2}}{2L}\leq f(w)-f(w_{\star}).$

Thus

$\|\nabla f(w)\|^{2}\leq 2L(f(w)-f(w_{\star})).$

Now, combining the preceding display with the $\theta$-PL condition, we obtain

$2\theta(f(w)-f(w_{\star}))\leq\|\nabla f(w)\|^{2}\leq 2L(f(w)-f(w_{\star})).$

Dividing by $2(f(w)-f(w_{\star}))$ we conclude

$\theta\leq L.$

∎

###  D.2 SketchySGD converges linearly under interpolation and the PL-
condition

In this section we show SketchySGD exhibits linear convergence under the
hypotheses of interpolation and the $\theta$-PL condition. We require the
following additional hypothesis,

######  Assumption 18.

The maximum eigenvalue remains of $\hat{H}^{\rho_{k}}_{S_{k}}$ remains bounded
throughout the optimization. That is,

$\lambda_{\star}=\sup_{k\geq
0}\lambda_{\text{max}}(\hat{H}^{\rho_{k}}_{S_{k}})<\infty.$

If $H_{S_{k}}(w_{k})$ were psd for all $k$, Assumption 18 would hold trivially
by the properties of the Nyström approximation. Indeed, we would have
$\lambda_{\star}\leq L_{\textrm{max}}+\rho_{\textrm{max}}$. Unfortunately, in
the non-convex setting the subsampled Hessian need not be psd, which
necessitates Assumption 18. The convergence rate here depends on a bound
$\rho_{\text{min}}/\lambda_{\star}$ on the condition number of the approximate
Hessian, where $\rho_{\text{min}}$ lower bounds the minimum eigenvalue.

######  Theorem 19 (SketchySGD non-convex convergence).

Instate Assumptions 1, 6, 7, 18. Run SketchySGD from any initialization
$w_{0}\in\mathbb{R}^{p}$ with constant batch sizes $|B_{k}|=b_{g}$ and
$|S_{k}|=b_{h}$, fixed rank $r$, regularization seqeuence
$\\{\rho_{k}\\}\in[\rho_{\textup{min}},\rho_{\textup{max}}]$, and fixed step-
size $\eta=\frac{\theta\lambda_{\star}}{(\theta+2\alpha)L}$. Then the iterate
$w_{t}$ satisfies

${\mathbb{E}}[f(w_{t})]-f(w_{\star})\leq\left(1-\frac{\theta}{(\theta+2\alpha)}\frac{\rho_{\textup{min}}^{2}}{\lambda_{\star}^{2}}\frac{\theta}{L}\right)^{t}(f(w_{0})-f(w_{\star})).$

###### Proof.

We start by noting the following bound holds $\forall k$,

$\rho_{\textrm{min}}I\preceq\hat{H}^{\rho}_{S_{k}}\preceq\lambda_{\star}I.$
(9)

Indeed, the lower bound holds by definition of $\hat{H}_{S_{k}}^{\rho}$, while
the upper bound follows by assumpution. We now move onto the proof of the main
claim. We start by observing that $L$-smoothness of $f$ and the definition of
the SketchySGD update (1) imply

$\displaystyle f(w_{k+1})\leq f(w_{k})+\langle
g(w_{k}),w_{k+1}-w_{k}\rangle+\frac{L}{2}\|w_{k+1}-w_{k}\|^{2}$
$\displaystyle=f(w_{k})-\eta_{k}\langle
g(w_{k}),[\hat{H}^{\rho_{k}}_{S_{k}}]^{-1}g_{B_{k}}(w_{k})\rangle+\frac{L}{2}\eta^{2}_{k}\|[\hat{H}^{\rho_{k}}_{S_{k}}]^{-1}g_{B_{k}}(w_{k})\|^{2}.$

Now taking expectations with respect to $B_{k}$ and using (9) with Cauchy-
Schwarz the above display becomes

$\displaystyle\mathbb{E}_{B_{k}}[f(w_{k+1})]\leq
f(w_{k})-\frac{\eta_{k}}{\lambda_{\star}}\|g(w_{k})\|^{2}+\frac{L}{2}\frac{\eta^{2}_{k}}{\rho_{\textrm{min}}^{2}}\mathbb{E}_{B_{k}}\left[\|g_{B_{k}}(w_{k})\|^{2}\right].$

Now, invoking Proposition 9 yields

$\displaystyle{\mathbb{E}}_{B_{k}}[f(w_{k+1})]\leq$ $\displaystyle
f(w_{k})-\frac{\eta_{k}}{\lambda_{\star}}\|g(w_{k})\|^{2}+\frac{L}{2}\frac{\eta^{2}_{k}}{\rho_{\textrm{min}}^{2}}\|g(w_{k})\|^{2}$
$\displaystyle+$ $\displaystyle 2\alpha
L\frac{\eta^{2}_{k}}{\rho_{\textrm{min}}^{2}}(f(w_{k})-f(w_{\star})).$

We can bound the last term using the $\theta$-PL condition, which gives

$\displaystyle{\mathbb{E}}_{B_{k}}[f(w_{k+1})]\leq$ $\displaystyle
f(w_{k})-\frac{\eta_{k}}{\lambda_{\star}}\|g(w_{k})\|^{2}+\frac{L}{2}\frac{\eta^{2}_{k}}{\rho_{\textrm{min}}^{2}}\|g(w_{k})\|^{2}$
$\displaystyle+$ $\displaystyle\alpha
L\frac{\eta^{2}_{k}}{\theta\rho_{\textrm{min}}^{2}}\|g(w_{k})\|^{2}.$

Collecting like terms, we obtain

$\displaystyle{\mathbb{E}}_{B_{k}}[f(w_{k+1})]\leq$ $\displaystyle
f(w_{k})+\left(-\frac{\eta_{k}}{\lambda_{\star}}+\frac{L\eta_{k}^{2}}{\rho_{\textrm{min}}^{2}}\left(\frac{1}{2}+\frac{\alpha}{\theta}\right)\right)\|g(w_{k})\|^{2}$
$\displaystyle=$ $\displaystyle
f(w_{k})+\left(-\frac{\eta_{k}}{\lambda_{\star}}+\frac{L\eta_{k}^{2}}{\rho_{\textrm{min}}^{2}}\left(\frac{\theta+2\alpha}{2\theta}\right)\right)\|g(w_{k})\|^{2}.$

Setting
$\eta_{k}=\eta=\frac{\theta\rho_{\textrm{min}}^{2}}{(\theta+2\alpha)\lambda_{\star}L}$
yields

$\displaystyle{\mathbb{E}}_{B_{k}}[f(w_{k+1})]\leq$ $\displaystyle
f(w_{k})-\frac{\theta}{(\theta+2\alpha)}\frac{\rho_{\textrm{min}}^{2}\|g(w_{k})\|^{2}}{2L\lambda_{\star}^{2}}.$

Invoking the $\theta$-PL condition again and subtracting off $f(w_{\star})$,
we obtain the bound

$\displaystyle{\mathbb{E}}_{B_{k}}[f(w_{k+1})]-f(w_{\star})\leq$
$\displaystyle\left(1-\frac{\theta}{(\theta+2\alpha)}\frac{\rho_{\textrm{min}}^{2}}{\lambda_{\star}^{2}}\frac{\theta}{L}\right)(f(w_{k})-f(w_{\star})).$

Taking the expectation over all timesteps and applying recursion in the
preceding display yields

$\displaystyle{\mathbb{E}}[f(w_{t})]-f(w_{\star})\leq$
$\displaystyle\left(1-\frac{\theta}{(\theta+2\alpha)}\frac{\rho_{\textrm{min}}^{2}}{\lambda_{\star}^{2}}\frac{\theta}{L}\right)^{t}(f(w_{0})-f(w_{\star})).$

∎

###  D.3 Proof of Theorem 8

###### Proof.

Let
$\beta^{\prime}=\frac{\theta}{(\theta+2\alpha)}\frac{\rho_{\textrm{min}}^{2}}{\lambda_{\star}^{2}}\frac{\theta}{L}$.
Theorem 19 shows that if the $\theta$-PL condition is satisfied and we run
SketchySGD with appropriate step sizes, then the iterates of SketchySGD
satisfy

${\mathbb{E}}[f(w_{k})]-f(w_{\star})\leq\left(1-\beta^{\prime}\right)^{k}(f(w_{0})-f(w_{\star})).$

In the case of a MLP with square-loss, Theorem 4 of Liu et al. (2022) shows
the $\theta$-PL condition holds in a ball of radius $R$ about the
initialization provided the width of the network at each layer satisfies

$m^{l}=\tilde{\Omega}\left(\frac{NR^{6L+2}}{(\lambda_{0}-\theta\gamma^{-2})^{2}}\right).$

Consequently, to establish the claimed convergence for wide MLPs, we just need
to verify for some $R>0$ that the iterates $\\{w_{k}\\}_{k\geq 0}$ remain in
the ball $B(w_{0},R)$ throughout the optimization with high probability. This
verification is straightforward as the algorithm converges linearly.

Let $\delta>0$ and set
$R=\frac{\eta\sqrt{2L+4\alpha}}{\rho_{\textrm{min}}\delta}\frac{f(w_{0})}{1-\sqrt{1-\beta^{\prime}}}$.
Note that the definition of the radius is not circular even though it depends
on $f(w_{0})$. To see this, observe
$f(w_{0})=\|\mathcal{F}_{\textrm{NN}}(w_{0})-y\|^{2}$ and recall Proposition 1
in Jacot et al. (2018) shows with high probability that for wide NNs
$\|\mathcal{F}_{\textrm{NN}}(w_{0})-y\|=O(1)$. Hence $f(w_{0})=O(1)$ and is
independent of the model, and thus so is $R$. We shall show the iterates
$\\{w_{k}\\}_{k\geq 0}$ lie in the ball $B(w_{0},R)$ with with probability at
least $1-\delta$. To see this, observe that

$\displaystyle\|w_{k+1}-w_{k}\|=\eta\left\|\left(\hat{H}^{\rho_{k}}_{S_{k}}\right)^{-1}g_{B_{k}}\right\|\leq\frac{\eta}{\rho_{\textrm{min}}}\|g_{B_{k}}\|.$
(10)

Taking the expectation with respect to $B_{k}$, using
${\mathbb{E}}\|g_{B_{k}}\|\leq\sqrt{{\mathbb{E}}\|g_{B_{k}}\|^{2}}$, and
Proposition 9 gives

${\mathbb{E}}_{B_{k}}\|w_{k+1}-w_{k}\|\leq\sqrt{\|g(w_{k})\|^{2}+4\alpha((f(w_{k})-f(w_{\star})))}.$

Invoking smoothness of $f$, the preceding display becomes

${\mathbb{E}}_{B_{k}}\|w_{k+1}-w_{k}\|\leq\frac{\eta}{\rho_{\textrm{min}}}\sqrt{(2L+4\alpha)(f(w_{k})-f(w_{\star}))}.$

Now taking the total expectation and using observation (10), we obtain

$\displaystyle{\mathbb{E}}\|w_{k+1}-w_{k}\|\leq\frac{\eta\sqrt{2L+4\alpha}}{\rho_{\textrm{min}}}\sqrt{{\mathbb{E}}[f(w_{k})]-f(w_{\star})}\leq\frac{\eta\sqrt{2L+4\alpha}}{\rho_{\textrm{min}}}\left(1-\beta^{\prime}\right)^{k/2}(f(w_{0})-f(w_{\star})).$

Adding up the distance traversed with every step, we find

$\displaystyle{\mathbb{E}}\left[\sum^{\infty}_{k=0}\|w_{k+1}-w_{k}\|\right]\leq$
$\displaystyle\sum^{\infty}_{k=0}\frac{\eta\sqrt{4\alpha+2L}}{\rho_{\textrm{min}}}\left(1-\beta^{\prime}\right)^{k/2}(f(w_{0})-f(w_{\star}))$
$\displaystyle=$
$\displaystyle\frac{\eta\sqrt{4\alpha+2L}}{\rho_{\textrm{min}}}\frac{(f(w_{0})-f(w_{\star}))}{1-\sqrt{1-\beta^{\prime}}}$
$\displaystyle=$ $\displaystyle\leq\delta R,$

where in the last step we used $f(w_{0})-f(w_{\star})\leq f(w_{0})$ as the
square-loss is non-negative. Hence by Markov’s inequality

$\displaystyle{\mathbb{P}}\left(\left\\{\sum^{\infty}_{k=0}\|w_{k+1}-w_{k}\|\geq
R\right\\}\right)\leq\delta.$

Now,

$w_{k+1}-w_{0}=\sum^{k+1}_{j=1}(w_{j}-w_{j-1}).$

Thus with probability at least $1-\delta$,

$\|w_{k+1}-w_{0}\|\leq\sum^{k}_{j=0}\|w_{j+1}-w_{j}\|\leq\sum^{\infty}_{j=0}\|w_{j+1}-w_{j}\|\leq
R.$

Consequently, with probability at least $1-\delta$, the optimization
trajectory $\\{w_{k}\\}_{k\geq 0}$ remains in the ball $B(w_{0},R)$. ∎

##  Appendix E Experimental Details

Here we provide more details regarding the logistic regression and deep
learning experiments.

##### Logistic regression

We compare SketchySGD against SGD, SVRG, Adam, SAGA, and L-BFGS. For SGD,
SVRG, and Adam we use a learning rate grid of $[10^{-3},2.5\cdot
10^{-3},5\cdot 10^{-3},10^{-2},2.5\cdot 10^{-2},5\cdot
10^{-2},10^{-1},2.5\cdot 10^{-1},5\cdot 10^{-1},10^{0}]$. The learning rate
for SAGA and L-BFGS is automatically selected by scikit-learn. For SketchySGD,
we set the learning rate equal to $1/(1+100\hat{\lambda}_{r})$. In SVRG, we
use $m=\textrm{\\# of minibatches}$ (i.e. we update the "snapshot" once every
epoch). In Adam, we use the default hyperparameters $\beta_{1}=0.9$ and
$\beta_{2}=0.999$. We train for 20 epochs, and use a batch size of 256.

Table 6:  Dimensions of Logistic Regression Datasets Dataset | $n_{\textrm{tr}}$ | $n_{\textrm{test}}$ | $d$
---|---|---|---
IJCNN | $49990$ | $91701$ | $800$
News20 | $15999$ | $3997$ | $1355191$
RCV1 | $20242$ | $677399$ | $47236$

The logistic regression experiments are run on the datasets described in the
main text. For News20, we use a random 80-20 split to form a training and test
set. For IJCNN, we use a random features transformation that gives us 800
features in total. In Table 6, we provide the dimensions of the datasets,
where $n_{\textrm{tr}}$ is the number of training samples, $n_{\textrm{test}}$
is the number of testing samples, and $d$ is the number of features.

##### Deep learning on tabular data

We compare SketchySGD against SGD and Adam. All three methods are used to
train MLPs. For all three methods, we use a learning rate grid of
$[10^{-4},3\cdot 10^{-4},10^{-3},3\cdot 10^{-3},10^{-2},3\cdot
10^{-2},10^{-1},3\cdot 10^{-1},10^{0},3\cdot 10^{0}]$. For all three
optimizers, we use weight decay = $5\cdot 10^{-4}$. In Adam, we use the
default hyperparameters $\beta_{1}=0.9$ and $\beta_{2}=0.999$. We train for 50
epochs, decaying the learning rate by a factor of 10 at 10 and 20 epochs, and
use a batch size of 128. We run SketchySGD with regularization $\rho=\eta$
that decays with the learning rate, and use rank $r=100$ for all the tasks. We
also update the Hessian approximation $\hat{H}$ every two epochs.

Table 7:  Dimensions of Tabular Datasets Dataset \ Size | $n_{\textrm{tr}}$ | $n_{\textrm{test}}$ | $d$ | Network architecture
---|---|---|---|---
Higgs | $800000$ | $200000$ | $28$ | $[28,64,64,64,2]$
MiniBooNE | $104051$ | $26013$ | $50$ | $[50,128,128,128,2]$
Volkert | $46648$ | $11662$ | $180$ | $[180,256,256,256,10]$

The tabular data experiments are run on the datasets described in the main
text. For all three datasets, we use a random 80-20 split to form a training
and test set. In Table 7, we provide the dimensions of the datasets and the
network architectures used in training. In Figure 5, we show peak test
accuracies across learning rates for all three datasets.

(a) Higgs

(b) MiniBooNE

(c) Volkert

Figure 5: For tabular data, the peak test accuracy produced by SketchySGD does
not vary significantly with the learning rate. However, the peak test accuracy
produced by SGD and Adam is sensitive to the learning rate.

##### Deep learning on image classification

We compare SketchySGD against SGD, Adam, and AdamW. We train ResNet20 on
CIFAR10 and ResNet32 on SVHN. We use the train-test splits for CIFAR10 and
SVHN that are provided by PyTorch. For all four optimizers, we use a learning
rate grid of $[10^{-4},3\cdot 10^{-4},10^{-3},3\cdot 10^{-3},10^{-2},3\cdot
10^{-2},10^{-1},3\cdot 10^{-1},10^{0},3\cdot 10^{0}]$. For all four
optimizers, we use weight decay = $5\cdot 10^{-4}$. In Adam and AdamW, we use
the default hyperparameters $\beta_{1}=0.9$ and $\beta_{2}=0.999$. We train
for 200 epochs, decaying the learning rate by a factor of 10 at 80, 120, and
160 epochs, and use a batch size of 128. We run SketchySGD with regularization
$\rho=\eta$ that decays with the learning rate, and use rank $r=100$ for all
the tasks. We also update the Hessian approximation $\hat{H}$ every two
epochs. In Figure 6, we show peak test accuracies across learning rates for
both CIFAR10 and SVHN.

(a) CIFAR-10

(b) SVHN

Figure 6: For image classification, the peak test accuracy produced by
SketchySGD does not vary significantly with the learning rate. However, the
peak test accuracy produced by SGD, Adam, and AdamW is sensitive to the
learning rate.

##### Deep learning: comparing quasi-Newton methods

We compare SketchySGD against AdaHessian, Shampoo, K-FAC, and SENG. For all
five methods, we use a learning rate grid of $[10^{-4},3\cdot
10^{-4},10^{-3},3\cdot 10^{-3},10^{-2},3\cdot 10^{-2},10^{-1},3\cdot
10^{-1},10^{0},3\cdot 10^{0}]$ and weight decay = $5\cdot 10^{-4}$. For
AdaHessian, Shampoo, K-FAC, and SENG, we use all the default hyperparameters
provided by the authors except learning rate and weight decay. We run
SketchySGD with regularization $\rho=\eta$ that decays with the learning rate,
and use rank $r=100$ for all the tasks. We also update the Hessian
approximation $\hat{H}$ every two epochs. All the times provided in Table 5
are evaluated on a NVIDIA GeForce RTX 3090 GPU.

For tabular datasets, Higgs, MiniBooNE, and Volkert, we use the same network
structures as shown in Table 7. We train for 30 epochs, decaying the learning
rate by a factor of 10 at 10 and 20 epochs, and use a batch size of 128. We
provide the peak test accuracies across learning rates for all quasi-Newton
methods on three tabular datasets in Figure 7. For image classification on
CIFAR-10 dataset, we use ResNet20 model. We train for 80 epochs, decaying the
learning rate by a factor of 10 at 30, 50, and 70 epochs, and use a batch size
of 128. For CIFAR-10 dataset, Figure 8(a) shows the robustness of quasi-Newton
methods in the same conventions as Figure 4, and Figure 8(b) shows the peak
test accuracies across learning rates for all quasi-Newton methods.

(a) Higgs

(b) MiniBooNE

(c) Volkert

Figure 7: For tabular data, the peak test accuracy produced by SketchySGD does
not vary significantly with the learning rate. However, the peak test accuracy
produced by other quasi-Newton methods is sensitive to the learning rate.

(a) CIFAR-10 robustness

(b) CIFAR-10 peak

Figure 8: For image classification, SketchySGD performs best among quasi-
Newton methods in terms of robustness, and the peak test accuracy produced by
SketchySGD does not vary significantly with the learning rate. Only AdaHessian
and SENG are capable of producing a comparable test accuracy; however, their
performance is less robust to the learning rate and AdaHessian requires full
Hessian-vector-product and gradient evaluations.

##  Appendix F Default Hyperparameters for SketchySGD

In this section we give guidelines for selecting SketchySGD’s hyperparameters.

##### Hyperparameters for convex problems

For convex problems, the following hyperparameter settings work well:

  * •

Batch sizes set equal $b_{g}=b_{h}$. We used 256 in this paper.

  * •

Fixed rank $r=50$. This value performed well in all our logistic regression
experiments.

  * •

Fixed regularization $\rho\in\\{10^{-1},10^{-2},10^{-3}\\}$ at every
iteration. For the logistic regression experiments, we used $\rho=10^{-3}.$

  * •

Learning rate $\eta=\frac{1}{1+100\hat{\lambda}_{r}}$, where
$\hat{\lambda}_{r}$ is the $r$th eigenvalue of the current subsampled Hessian
approximation $\hat{H}$.

  * •

We recommend computing a fresh approximation $\hat{H}$ after each epoch or
two.

Our learning rate recommendation is motivated by the following theoretical
observation. At each iteration we expect to be able to take a step-size on the
order of $L^{-1}_{P}$, where
$L_{P}=\|(\hat{H}_{S_{k}}^{\rho_{k}})^{-1/2}H(\hat{H}_{S_{k}}^{\rho_{k}})^{-1/2}\|$.
By the triangle inequality and the properties of the Nyström approximation, it
is easy to see that

$L_{P}\leq
1+\frac{\|H-H_{S_{k}}\|}{\rho_{k}}=1+\frac{\|E_{S_{k}}\|}{\rho_{k}}.$

Thus, we expect to be able to take a step-size satisfying

$\eta=O\left(\frac{1}{1+\|E_{S_{k}}\|/\rho_{k}}\right).$

Our choice follows by approximating the ratio as
$\|E_{S_{k}}\|/\rho_{k}\approx 100\hat{\lambda}_{r}$. Empirically we find this
choice works well in practice, having used this selection for all the logistic
regression experiments.

##### Hyperparameters for deep learning

For deep learning we find the following hyperparameter settings work well:

  * •

Batch sizes set equal $b_{g}=b_{h}$. We used 128 in this paper.

  * •

Fixed rank $r=100$. This value performed well in all our deep learning
experiments.

  * •

Fixed learning rate $\eta=10^{-2}$.

  * •

Fixed regularization $\rho=\eta$ at every iteration.

  * •

We recommend computing a fresh approximation $\hat{H}$ after each epoch or
two. Our experiments compute a new approximation every two epochs.

## References

  * Davis et al. [2020] Damek Davis, Dmitriy Drusvyatskiy, Sham Kakade, and Jason D Lee.  Stochastic subgradient method converges on tame functions.  _Foundations of Computational Mathematics_ , 20(1):119–154, 2020.
  * Duchi et al. [2011] John Duchi, Elad Hazan, and Yoram Singer.  Adaptive subgradient methods for online learning and stochastic optimization.  _Journal of Machine Learning Research_ , 12(7), 2011.
  * Johnson and Zhang [2013] Rie Johnson and Tong Zhang.  Accelerating stochastic gradient descent using predictive variance reduction.  _Advances in Neural Information Processing Systems_ , 26, 2013.
  * Kingma and Ba [2014] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Sun [2019] Ruoyu Sun.  Optimization for deep learning: theory and algorithms.  _arXiv preprint arXiv:1912.08957_ , 2019.
  * Yao et al. [2021] Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoney.  AdaHessian: An adaptive second order optimizer for machine learning.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 35, pages 10665–10673, 2021.
  * Schwartz et al. [2020] Roy Schwartz, Jesse Dodge, Noah A Smith, and Oren Etzioni.  Green AI.  _Communications of the ACM_ , 63(12):54–63, 2020\.
  * Nocedal and Wright [1999] Jorge Nocedal and Stephen J Wright.  _Numerical optimization_.  Springer, 1999.
  * Boyd and Vandenberghe [2004] Stephen P Boyd and Lieven Vandenberghe.  _Convex optimization_.  Cambridge University Press, 2004.
  * Dauphin et al. [2014] Yann N Dauphin, Razvan Pascanu, Caglar Gulcehre, Kyunghyun Cho, Surya Ganguli, and Yoshua Bengio.  Identifying and attacking the saddle point problem in high-dimensional non-convex optimization.  _Advances in neural information processing systems_ , 27, 2014.
  * Levenberg [1944] Kenneth Levenberg.  A method for the solution of certain non-linear problems in least squares.  _Quarterly of Applied Mathematics_ , 2(2):164–168, 1944.
  * Marquardt [1963] Donald W Marquardt.  An algorithm for least-squares estimation of nonlinear parameters.  _Journal of the Society for Industrial and Applied Mathematics_ , 11(2):431–441, 1963.
  * Ghorbani et al. [2019] Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao.  An investigation into neural net optimization via Hessian eigenvalue density.  In _International Conference on Machine Learning_ , pages 2232–2241. PMLR, 2019.
  * Yao et al. [2020] Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney.  PyHessian: Neural networks through the lens of the Hessian.  In _2020 IEEE international conference on big data (Big data)_ , pages 581–590. IEEE, 2020.
  * Pearlmutter [1994] Barak A Pearlmutter.  Fast exact multiplication by the Hessian.  _Neural Computation_ , 6(1):147–160, 1994.
  * Woodruff [2014] David P Woodruff.  Sketching as a tool for numerical linear algebra.  _Foundations and Trends® in Theoretical Computer Science_ , 10(1–2):1–157, 2014.
  * Martinsson and Tropp [2020] Per-Gunnar Martinsson and Joel A Tropp.  Randomized numerical linear algebra: Foundations and algorithms.  _Acta Numerica_ , 29:403–572, 2020.
  * Halko et al. [2011] Nathan Halko, Per-Gunnar Martinsson, and Joel A Tropp.  Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions.  _SIAM Review_ , 53(2):217–288, 2011.
  * Cohen et al. [2015] Michael B Cohen, Sam Elder, Cameron Musco, Christopher Musco, and Madalina Persu.  Dimensionality reduction for k-means clustering and low rank approximation.  In _Proceedings of the Forty-seventh Annual ACM Symposium on Theory of Computing_ , pages 163–172, 2015.
  * Tropp et al. [2017a] Joel A Tropp, Alp Yurtsever, Madeleine Udell, and Volkan Cevher.  Practical sketching algorithms for low-rank matrix approximation.  _SIAM Journal on Matrix Analysis and Applications_ , 38(4):1454–1485, 2017a.
  * Tropp et al. [2017b] Joel A Tropp, Alp Yurtsever, Madeleine Udell, and Volkan Cevher.  Fixed-rank approximation of a positive-semidefinite matrix from streaming data.  _Advances in Neural Information Processing Systems_ , 30, 2017b.
  * Higham [2002] Nicholas J Higham.  _Accuracy and stability of numerical algorithms_.  SIAM, 2002.
  * Broyden [1970] Charles George Broyden.  The convergence of a class of double-rank minimization algorithms 1. general considerations.  _IMA Journal of Applied Mathematics_ , 6(1):76–90, 1970.
  * Liu and Nocedal [1989] Dong C Liu and Jorge Nocedal.  On the limited memory BFGS method for large scale optimization.  _Mathematical Programming_ , 45(1):503–528, 1989\.
  * Dembo et al. [1982] Ron S Dembo, Stanley C Eisenstat, and Trond Steihaug.  Inexact Newton methods.  _SIAM Journal on Numerical Analysis_ , 19(2):400–408, 1982.
  * Byrd et al. [2011] Richard H Byrd, Gillian M Chin, Will Neveitt, and Jorge Nocedal.  On the use of stochastic Hessian information in optimization methods for machine learning.  _SIAM Journal on Optimization_ , 21(3):977–995, 2011.
  * Erdogdu and Montanari [2015] Murat A Erdogdu and Andrea Montanari.  Convergence rates of sub-sampled Newton methods.  _Advances in Neural Information Processing Systems_ , 28, 2015.
  * Pilanci and Wainwright [2017] Mert Pilanci and Martin J Wainwright.  Newton sketch: A near linear-time optimization algorithm with linear-quadratic convergence.  _SIAM Journal on Optimization_ , 27(1):205–245, 2017.
  * Gower et al. [2019] Robert Gower, Dmitry Kovalev, Felix Lieder, and Peter Richtárik.  RSN: Randomized subspace Newton.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Roosta-Khorasani and Mahoney [2019] Farbod Roosta-Khorasani and Michael W Mahoney.  Sub-sampled Newton methods.  _Mathematical Programming_ , 174(1):293–326, 2019\.
  * Bollapragada et al. [2019] Raghu Bollapragada, Richard H Byrd, and Jorge Nocedal.  Exact and inexact subsampled Newton methods for optimization.  _IMA Journal of Numerical Analysis_ , 39(2):545–578, 2019.
  * Meng et al. [2020] Si Yi Meng, Sharan Vaswani, Issam Hadj Laradji, Mark Schmidt, and Simon Lacoste-Julien.  Fast and furious convergence: Stochastic second order methods under interpolation.  In _International Conference on Artificial Intelligence and Statistics_ , pages 1375–1386. PMLR, 2020.
  * Grosse and Martens [2016] Roger Grosse and James Martens.  A Kronecker-factored approximate Fisher matrix for convolution layers.  In _International Conference on Machine Learning_ , pages 573–582. PMLR, 2016.
  * Yang et al. [2020] Minghan Yang, Dong Xu, Zaiwen Wen, Mengyun Chen, and Pengxiang Xu.  Sketchy empirical natural gradient methods for deep learning.  _arXiv preprint arXiv:2006.05924_ , 2020.
  * Roux et al. [2007] Nicolas Roux, Pierre-Antoine Manzagol, and Yoshua Bengio.  Topmoumoute online natural gradient algorithm.  _Advances in Neural Information Processing Systems_ , 20, 2007.
  * Yang et al. [2021a] Minghan Yang, Dong Xu, Qiwen Cui, Zaiwen Wen, and Pengxiang Xu.  NG+: A multi-step matrix-product natural gradient method for deep learning.  _arXiv preprint arXiv:2106.07454_ , 2021a.
  * Gupta et al. [2018] Vineet Gupta, Tomer Koren, and Yoram Singer.  Shampoo: Preconditioned stochastic tensor optimization.  In _International Conference on Machine Learning_ , pages 1842–1850. PMLR, 2018.
  * Goldfarb et al. [2020] Donald Goldfarb, Yi Ren, and Achraf Bahamou.  Practical quasi-Newton methods for training deep neural networks.  _Advances in Neural Information Processing Systems_ , 33:2386–2396, 2020.
  * Yang et al. [2021b] Minghan Yang, Dong Xu, Hongyu Chen, Zaiwen Wen, and Mengyun Chen.  Enhance curvature information by structured stochastic quasi-Newton methods.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 10654–10663, 2021b.
  * Zhao et al. [2022] Shipu Zhao, Zachary Frangella, and Madeleine Udell.  NysADMM: faster composite convex optimization via low-rank approximation.  _arXiv preprint arXiv:2202.11599_ , 2022.
  * Derezinski et al. [2020] Michal Derezinski, Feynman T Liang, Zhenyu Liao, and Michael W Mahoney.  Precise expressions for random projections: Low-rank approximation and randomized newton.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, _Advances in Neural Information Processing Systems_ , volume 33, pages 18272–18283. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper/2020/file/d40d35b3063c11244fbf38e9b55074be-Paper.pdf.
  * Gower et al. [2021] Robert Gower, Othmane Sebbouh, and Nicolas Loizou.  SGD for structured nonconvex functions: Learning rates, minibatching and interpolation.  In _International Conference on Artificial Intelligence and Statistics_ , pages 1315–1323. PMLR, 2021.
  * Vaswani et al. [2019a] Sharan Vaswani, Aaron Mishkin, Issam Laradji, Mark Schmidt, Gauthier Gidel, and Simon Lacoste-Julien.  Painless stochastic gradient: Interpolation, line-search, and convergence rates.  _Advances in Neural Information Processing Systems_ , 32, 2019a.
  * Liu et al. [2022] Chaoyue Liu, Libin Zhu, and Mikhail Belkin.  Loss landscapes and optimization in over-parameterized non-linear systems and neural networks.  _Applied and Computational Harmonic Analysis_ , 2022.
  * Vaswani et al. [2019b] Sharan Vaswani, Francis Bach, and Mark Schmidt.  Fast and faster convergence of SGD for over-parameterized models and an accelerated perceptron.  In _The 22nd International Conference on Artificial Intelligence and Statistics_ , pages 1195–1204. PMLR, 2019b.
  * Ma et al. [2018] Siyuan Ma, Raef Bassily, and Mikhail Belkin.  The power of interpolation: Understanding the effectiveness of SGD in modern over-parametrized learning.  In _International Conference on Machine Learning_ , pages 3325–3334. PMLR, 2018.
  * Belkin [2021] Mikhail Belkin.  Fit without fear: remarkable mathematical phenomena of deep learning through the prism of interpolation.  _Acta Numerica_ , 30:203–248, 2021.
  * Scaman et al. [2022] Kevin Scaman, Cédric Malherbe, and Ludovic Dos Santos.  Convergence rates of non-convex stochastic gradient descent under a generic Lojasiewicz condition and local smoothness.  In _International Conference on Machine Learning_ , pages 19310–19327. PMLR, 2022.
  * Defazio et al. [2014] Aaron Defazio, Francis Bach, and Simon Lacoste-Julien.  SAGA: A fast incremental gradient method with support for non-strongly convex composite objectives.  _Advances in Neural Information Processing Systems_ , 27, 2014.
  * Loshchilov and Hutter [2019] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=Bkg6RiCqY7.
  * [51] Danil Prokhorov.  IJCNN 2001 neural network competition.
  * Lang [1995] Ken Lang.  Newsweeder: Learning to filter netnews.  In Armand Prieditis and Stuart Russell, editors, _Machine Learning Proceedings 1995_ , pages 331–339. Morgan Kaufmann, San Francisco (CA), 1995.  ISBN 978-1-55860-377-6.  doi:https://doi.org/10.1016/B978-1-55860-377-6.50048-7.  URL https://www.sciencedirect.com/science/article/pii/B9781558603776500487.
  * Lewis et al. [2004] David D. Lewis, Yiming Yang, Tony G. Rose, and Fan Li.  RCV1: A new benchmark collection for text categorization research.  _J. Mach. Learn. Res._ , 5:361–397, dec 2004.  ISSN 1532-4435.
  * Rahimi and Recht [2007] Ali Rahimi and Benjamin Recht.  Random features for large-scale kernel machines.  _Advances in Neural Information Processing Systems_ , 20, 2007.
  * Paszke et al. [2019] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  Pytorch: An imperative style, high-performance deep learning library.  In _Advances in Neural Information Processing Systems 32_ , pages 8024–8035. 2019.
  * Baldi et al. [2014] P. Baldi, P. Sadowski, and D. Whiteson.  Searching for exotic particles in high-energy physics with deep learning.  _Nature Communications_ , 5(1), July 2014.  doi:10.1038/ncomms5308.  URL https://doi.org/10.1038/ncomms5308.
  * [57] OpenML, miniboone.  https://www.openml.org/search?type=data&status=active&id=41150.  Accessed: 2022-10-07.
  * [58] OpenML, volkert.  https://www.openml.org/search?type=data&sort=runs&id=41166&status=active.  Accessed: 2022-05-01.
  * Krizhevsky et al. [2009] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * Netzer et al. [2011] Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y. Ng.  Reading digits in natural images with unsupervised feature learning.  In _NIPS Workshop on Deep Learning and Unsupervised Feature Learning 2011_ , 2011.  URL http://ufldl.stanford.edu/housenumbers/nips2011_housenumbers.pdf.
  * He et al. [2016] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 770–778, 2016.
  * Ye et al. [2021] Haishan Ye, Luo Luo, and Zhihua Zhang.  Approximate Newton methods.  _Journal of Machine Learning Research_ , 22(66):1–41, 2021.  URL http://jmlr.org/papers/v22/19-870.html.
  * Nesterov [2018] Yurii Nesterov.  _Lectures on convex optimization_ , volume 137.  Springer, 2018.
  * Jacot et al. [2018] Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural Tangent Kernel: Convergence and generalization in neural networks.  _Advances in Neural Information Processing Systems_ , 31, 2018.