# Sketchy: Memory-efficient Adaptive Regularization with Frequent Directions

Vladimir Feinberg
Google DeepMind
<EMAIL_ADDRESS>&Xinyi Chen
Princeton University
Google DeepMind &Y. Jennifer Sun
Princeton University
Google DeepMind &Rohan Anil
Google DeepMind &Elad Hazan
Princeton University
Google DeepMind

###### Abstract

Adaptive regularization methods that exploit more than the diagonal entries
exhibit state of the art performance for many tasks, but can be prohibitive in
terms of memory and running time. We find the spectra of the Kronecker-
factored gradient covariance matrix in deep learning (DL) training tasks are
concentrated on a small leading eigenspace that changes throughout training,
motivating a low-rank sketching approach. We describe a generic method for
reducing memory and compute requirements of maintaining a matrix
preconditioner using the Frequent Directions (FD) sketch. While previous
approaches have explored applying FD for second-order optimization, we present
a novel analysis which allows efficient interpolation between resource
requirements and the degradation in regret guarantees with rank $k$: in the
online convex optimization (OCO) setting over dimension $d$, we match full-
matrix $d^{2}$ memory regret using only $dk$ memory up to additive error in
the bottom $d-k$ eigenvalues of the gradient covariance. Further, we show
extensions of our work to Shampoo, resulting in a method competitive in
quality with Shampoo and Adam, yet requiring only sub-linear memory for
tracking second moments.

##  1 Introduction

DL optimization commonly relies on adaptive gradient methods, namely the Adam
optimizer [1]. It differs from stochastic gradient descent in that the
learning rate is a structured diagonal matrix built from previous gradients
rather than a scalar. In full matrix AdaGrad [2], the inverse matrix square
root of the sum of outer products of previous gradients is the learning rate.

Full matrix preconditioning is impractical for modern deep learning
architectures: for instance, the ResNet-50 architecture [3] has over 23
million parameters, requiring more than 2 petabytes to represent its gradient
covariance. Thus, diagonal preconditioning methods remain popular. However,
previous work has demonstrated state-of-the-art results in some settings, such
as large-batch data parallel training, for nondiagonal forms of
preconditioning [4, 5, 6, 7, 8, 9]. In particular, Shampoo [5, 9] introduces a
factorization of full matrix preconditioning method with adoption in large-
scale industrial applications such as training Google’s ads click-through-rate
model [10]. Furthermore, as hardware evolves, memory efficiency becomes an
increasing concern, as “logic improves much faster than wires and SRAM, so
logic is relatively free” [11]: from TPUv2 to TPUv3, per-chip bfloat16
operations per second improved $2.67\times$ but memory bandwidth only improved
$1.29\times$. GPUs exhibit a similar pattern for compute and memory increase,
at $5\times$ and $2.2\times$, for V100 to A100 [12].

Investigation into the Kronecker-factored gradient covariance matrix reveals a
concentrated, but changing, spectrum (Fig. 3), suggesting the majority of the
spectral mass can be represented by a low-rank matrix, albeit rotating over
time. The Frequent Directions (FD) sketch provides a mechanism to track the
top eigenvectors without materializing the full covariance matrix, as proposed
in [13]. Is a large portion of the spectral mass sufficient to retain the
performance of adaptive regularization in theory and practice? In this work,
we investigate this hypothesis.

  * •

In the setting of online convex optimization, by applying a dynamic diagonal
regularization to the FD sketch, we can recover full-matrix AdaGrad regret up
to additive spectral terms under a memory constraint, providing a novel
guarantee without curvature assumptions (Sec. 4.1). Rigorously composing our
approach with Shampoo (Sec. 4.2) unlocks a second-order algorithm which
requires sub-linear memory for its accumulators.

  * •

By modifying FD for exponential moving averages (Sec. 4.3), we demonstrate a
practical algorithm competitive with at-least-linear memory Shampoo and Adam
in three modern DL settings (Sec. 5.1). While previous work [8, 14] shows
rank-1 preconditioners are effective for trading off quality for memory, these
results demonstrate a Pareto improvement by using higher-rank approximations.

  * •

We explain the competitive performance of Sketchy with observations of fast
spectral decay in the moving average of Kronecker-factored gradient covariance
in DL settings (Sec. 5.2).

##  2 Setting and Definitions

##### Regret and Optimization.

The optimization problem of training a deep neural network has a non-convex
objective loss function $f$. Since finding the global optimum is
computationally intractable in general, theoretical guarantees focus on
convergence to an $\varepsilon$-approximate first-order optimum: a point $x$
such that $\|\nabla\mkern-2.5muf(x)\|\leq\varepsilon$. A smooth non-convex
problem can be reduced to solving a series of offline convex problems [6]. The
convex sub-problems have form $f_{t}(x)=f(x)+c\|x-x_{t}\|^{2}$, where $c$ is a
constant and $x_{t}$ is an iterate in the optimization process. Using online-
to-batch conversion, we can translate the regret bound of an online convex
optimization (OCO) [15] algorithm to convergence guarantees for offline
optimization. For more details of this reduction, see Appendix B.1. Therefore,
non-convex optimization guarantees can be obtained from regret bounds, and we
focus on the latter in this paper.

In this setting, an OCO algorithm chooses a point $x_{t}\in\mathcal{K}$
iteratively, where $\mathcal{K}\subseteq\mathbb{R}^{d}$ is a convex decision
set (take $\mathcal{K}=\mathbb{R}^{d}$ if unconstrained). After the decision
is made, the adversary reveals a convex loss function $f_{t}$, to which the
algorithm suffers cost $f_{t}(x_{t})$. Upon receiving the cost, the algorithm
updates its decision for the next iteration. The regret for an online
algorithm is given by

$\displaystyle\mbox{{Regret}}_{T}=\sum_{t=1}^{T}f_{t}(x_{t})-\min_{x\in\mathcal{K}}\sum_{t=1}^{T}f_{t}(x)\,\,.$

##### Sketching and the Frequent Directions Method.

Given a stream of vectors $g_{t}\in\mathbb{R}^{d}$, $t\in[T]$, we utilize the
FD sketch [13] given in Alg. 1 which maintains a low-rank approximation of the
true running covariance $G_{t}=\sum_{s\leq t}g_{s}g_{s}^{\top}$. At each time
$t$, it maintains a matrix $B_{t}$ of size $d\times\ell$ whose last column is
$0$ and whose square is the sketch $B_{t}B_{t}^{\top}=\bar{G}_{t}$. After
seeing $g_{t}$ from the stream, we update the previous matrix using Alg. 1,
which updates $B_{t+1}$ of size $d\times\ell$ whose last column remains $0$;
take $B_{0}=0$. $B_{t+1}$ is obtained by decomposing the sum of $\bar{G}_{t}$
and the newly observed matrix, keeping only the top eigendirections, and
reducing the eigenvalues uniformly by $\ell$-th eigenvalue. At every iteration
$t$, denote $\rho_{t}:=\lambda_{\ell}^{(t)}$ be the removed eigenvalue from
the covariance update in Alg. 1.

For convenience, let
$\rho_{1:t}\stackrel{{\scriptstyle\text{def}}}{{=}}\sum_{s=1}^{t}\rho_{s}$ be
the cumulative escaped mass. For a matrix $X$, we denote its $i$-th leading
eigenvalue by $\lambda_{i}(X)$. Let $\|\cdot\|_{F}$ denote the Frobenius norm
of a matrix.

Algorithm 1 Frequent Directions Update (FD-update)

0: Invariant that last column of $B_{t-1}$ is 0.

0: The last column of $B_{t}$ is 0.

1: Input: Previous state
$\bar{G}_{t-1}=B_{t-1}B_{t-1}^{\top}\in\mathbb{R}^{d\times d}$

2: Input: New symmetric PSD matrix $M_{t}\in\mathbb{R}^{d\times d}$.

3: Eigendecompose
$\bar{U}_{t}\operatorname{diag}\lambda^{(t)}\bar{U}_{t}^{\top}=\bar{G}_{t-1}+M_{t}$
where $\lambda^{(t)}$ contains descending eigenvalues.

4: Define $U_{t}$ as the matrix whose columns are the first $\ell$ columns of
$\bar{U}_{t}$, and $\lambda^{(t)}_{[1:\ell]}$ be its eigenvalues.

5: Update
$B_{t}=U_{t}\operatorname{diag}\left({\lambda^{(t)}_{[1:\ell]}-\lambda^{(t)}_{\ell}}\right)^{1/2}$.
$\lambda^{(t)}_{\ell}$, $B_{t}B_{t}^{\top}$.

The fundamental property of FD is that applying Alg. 1 over a stream of
vectors $g_{t}$, with $B_{0}=0$, the sum of escaped mass
$\rho_{t}=\lambda^{(t)}_{\ell}$ can be bounded by the bottom eigenvalues of
$G_{T}$, formally given by the following lemma:

######  Lemma 1 (Liberty [16]).

The cumulative escaped mass $\rho_{1:T}$ can be upper bounded as

$\rho_{1:T}\leq\min_{k=0,\ldots,\ell-1}\frac{\sum_{i=k+1}^{d}\lambda_{i}(G_{T})}{\ell-k}\leq\sum_{i=\ell}^{d}\lambda_{i}(G_{T})\stackrel{{\scriptstyle\text{def}}}{{=}}\lambda_{\ell:d}\,\,.$

###### Proof.

See Sec. B.2. ∎

##  3 Related Work

###  3.1 Spectral Analysis of DL Training

Denote the loss function of the $i$-th example for weights $x$ as $f_{i}(x)$.
The spectrum of the Hessian matrix $\sum_{i}\nabla\mkern-2.5mu^{2}f_{i}$ has
been the subject of intensive investigation in DL [17, 18, 19, 20] and its
properties have been used to devise training methods [4, 21].

Recent papers [22, 23, 24] inspect the covariance matrix,
$\sum_{i}(\nabla\mkern-2.5muf_{i})(\nabla\mkern-2.5muf_{i})^{\top}$. In small
models, where its computation is feasible, these works identify fast spectral
decay.

Agarwal et al. [6] take advantage of this observation by using a low-rank
approximation of the whole covariance matrix, based on a limited history of
the gradients,
$\sum_{i=t-r}^{r}(\nabla\mkern-2.5muf_{i})(\nabla\mkern-2.5muf_{i})^{\top}$.
This approach still requires $r$ copies of the model gradients in memory,
where typically $r$ should scale with $\beta_{2}^{-1}$, with $\beta_{2}$ the
exponential moving average for second order statistics (the authors set
$r=200$). Fundamentally, approximating the whole covariance matrix constrains
Agarwal et al. [6] application to small models.

In our work, we validate the decay hypothesis holds across the per-layer
factored covariance matrices in several modern neural networks. For a layer’s
gradient matrix $G_{i}$ at the $i$-th example and a second moment decay term
$\beta_{2}$, our work inspects spectral decay for
$L_{t}=\sum_{i}\beta_{2}^{t-i}G_{i}G_{i}^{\top}$ and
$R_{t}=\sum_{i}\beta_{2}^{t-i}G_{i}^{\top}G_{i}$; the spectral structure for
these outer products is not well-documented. Furthermore, as described in Sec.
3.4, approximating the factored covariance $L_{t}\otimes R_{t}$ requires less
memory than the full covariance and explains why our method can scale to large
modern architectures whereas Agarwal et al. [6] cannot.

###  3.2 Sublinear Memory Methods

Extreme Tensoring [7], AdaFactor [14], and SM3 [8] are methods that require
sublinear memory relative to the number of parameters, at the other end of the
memory-quality tradeoff beyond methods that rely on the diagonal of the
gradient covariance such as Adam. Owing to different structural assumptions on
the set of feasible preconditioners, comparison with these methods is out of
scope. However, these methods may compose with our approach. One may apply
Extreme Tensoring first, then sketch the resulting reshaped tensor covariances
with our method to further reduce memory consumption. Similarly, an SM3-like
approach which reduces the indices for each dimension to be preconditioned can
be applied before Sketchy is applied to remaining indices.

Crucially, Adam, which uses linear memory for second moment representations,
compares favorably in terms of quality to all of these sublinear-memory
methods. By increasing rank in factored covariance representation, Sketchy is
competitive with Adam, despite sublinear memory for second moment
representation. Thus, for simplicity, we compare to only Adam, which dominates
the alternative sublinear approaches in terms of quality.

###  3.3 Sketching-based Approaches

Table 1: Memory-efficient adaptive gradient methods, in the OCO setting with dimension $d$ (Sec. 2). We describe the worst-case regret bounds without exp-concavity assumptions, asymptotically, hiding logarithmic factors, treating the decision set diameter as a constant, and assume optimally-tuned hyperparameters. $\ell$ refers to the controllable preconditioner rank. Note $\operatorname{tr}G_{T}^{1/2}=\sqrt{\min_{H\in\mathcal{H}}\sum_{t}\left\lVert\nabla\mkern-2.5mu_{t}\right\rVert_{H}^{2}}$ is the optimal preconditioner’s regret among the class of positive semi-definite, unit-trace matrices, $\mathcal{H}$, and $G_{T}$ is the sum of gradient outer products. We let eigenvalues $\lambda_{i}=\lambda_{i}(G_{T})$ with $\lambda_{i:j}=\sum_{m=i}^{j}\lambda_{m}$.  Reference | Regret (general convex) | Memory
---|---|---
Full Matrix AdaGrad [2] | $\operatorname{tr}G_{T}^{1/2}$ | $d^{2}$
Ada-LR [25] | $\operatorname{tr}G_{T}^{1/2}+\lambda_{\ell+1}^{1/2}\ell^{3/4}d^{1/4}$ | $d^{2}$
Ada-FD [26] |  $\Omega\left({T^{3/4}}\right)$111The regret of Ada-FD is expressed in terms of dynamic run-time quantities which do not admit a universal bound in terms of $G_{T}$; we display its regret for the specific case of Observation 2 instead (a detailed look at its regret is given in Appendix B.3). | $d\ell$
SON [27] | $\sqrt{Td}$ | $d^{2}$
FD-SON [27] | $\sqrt{\ell\lambda_{\ell:d}T}$ | $d\ell$
This paper | $\operatorname{tr}(G_{T}^{1/2})+\sqrt{d(d-\ell)\lambda_{\ell:d}}$ | $d\ell$

Several works have explored sketching-like approximations to the gradient
covariance matrix, but none provide an adaptive bound exploiting fast spectral
decay in gradient covariance without additional assumptions (Tbl. 1). In this
section, we consider the OCO setting over dimension $d$ (Sec. 2).

Random projection (Ada-LR) is most spiritually similar to our work [25].
Although it does not reduce memory usage, it relies on random projections to
lower dimension $\ell\leq d$ to reduce inverse matrix computation costs. An
alternative without formal guarantees, RadaGrad, reduces memory consumption to
$O(d\ell)$; however, as with all Johnson-Lindenstraus projection methods, it
suffers a probabilistic failure rate scaling as $O(\ell^{-1})$ (in comparison,
our method inherits FD’s determinism).

Frequent Directions (FD) [13, 16], provides an alternative matrix sketching
approach from the data streaming literature. As an adaptive sketch, it
dominates random projection in terms of matrix recovery, and lower bounds show
its memory usage is optimal in the sense that any equally-accurate
approximation to an adversarially-chosen true covariance $G_{T}$ in operator
norm constructed from the corresponding gradients must use $O(d\ell)$ bits.

In the context of exp-concave cost functions, Luo et al. [27] provide an FD
sketched version of Online Newton Step (ONS), FD-SON. In this setting, their
approach nearly recovers classical ONS regret, up to logarithmic error in
$\sum_{i=1}^{\ell-1}\lambda_{i}(G_{T})$ and additive error in
$\sum_{i=\ell}^{d}\lambda_{i}(G_{T})$. However, without the exp-concave
assumption, FD-SON falls back to a gradient-descent-like default regret of
$O\left({\lambda_{\ell:d}\sqrt{T}}\right)$, which can be $\Omega(T)$ without
spectral decay. In the context of linear bandits, Chen et al. [28] uses FD for
memory reduction. The resulting algorithm, SCFD, is similar to Alg. 2, but
SCFD lacks a projection step and does not handle general domains
$\mathcal{K}$. We emphasize that our contribution is in the novelty of our
regret analysis for general OCO settings, in contrast to linear bandits.

The main prior work exploring FD for general online convex settings, Wan and
Zhang [26], extends the FD-SON approach by adding a fixed diagonal
perturbation $\delta I$ to an FD-based preconditioner, in Ada-FD. However,
this approach does not achieve $\sqrt{T}$ regret even in a non-adversarial
setting with stochastic linear cost functions (Observation 2), where learning
rate and $\delta$ are tuned. Dynamically changing diagonal regularization is
essential for worst-case $O\left({\sqrt{T}}\right)$ performance.

######  Observation 2.

Suppose we receive linear cost functions $f_{t}(x)=\left\langle
x,g_{t}\right\rangle$, where $g_{t}\in\mathbb{R}^{d}$ is a random vector drawn
iid from any distribution over $r\leq d$ orthonormal vectors $W$. For any
sketch size $\ell\leq r$, the bound on the expected regret of Ada-FD is
$\Omega(T^{3/4})$.

###### Proof.

See Sec. B.3. ∎

Wan and Zhang [26] remark Ada-FD has $\sqrt{T}$ regret when $G_{T}$ is low
rank with rank below $k$, so $d-k$ of its eigenvalues are precisely zero.
However, this setting does not require any sketching in the first place. By
tracking the column space of observed gradients (e.g., with a reduced QR
decomposition, rank-1-updated every step), the full matrix AdaGrad algorithm
can be perfectly recovered without using more than $O(dk)$ memory.

In concrete convex examples, Sketchy compares favorably to these approaches
(Appendix A).

###  3.4 Shampoo

Figure 1: Asymptotic memory consumption for representing gradient covariance
in adaptive regularization approaches for a single matrix parameter of size
$n\times m$. Here, $r$ refers to the GGT history buffer size and $k$ to the
approximation rank of FD (both typically set to hundreds). Past sketching
approaches like Ada-FD and Radagrad take memory similar to GGT, with $r$ being
sketch size; these are all asymptotically superlinear. This figure
demonstrates optimizer memory usage in the theoretical OCO setting; in
practice for deep learning workloads there are additive $O(mn)$ factors for
momentum, the parameters themselves, and grafting parameters for Shampoo and
Sketchy.

Perhaps our most compelling application is reducing the memory of Shampoo [5,
9, 10]. Shampoo is an adaptive preconditioning method that takes into account
the structure of the parameter space, and thus is more efficient than full
matrix AdaGrad. For example, if the parameter is a weight matrix $W$ of size
$m\times n$, AdaGrad treats the matrix-shaped parameters as a vector of size
$mn$, and the preconditioner has size $m^{2}n^{2}$; Shampoo instead has left
and right preconditioners $L,R$ of size $n\times n$ and $m\times m$,
respectively, with the preconditioned update $L^{-1/4}WR^{-1/4}$. Write
$\overline{\text{vec}}\left({W}\right)$ as the vectorized weights, then it is
equivalent to $(L\otimes
R)\overline{\text{vec}}\left({W}\right)=\overline{\text{vec}}\left({LWR}\right)$,
where $\otimes$ denotes the Kronecker product. Figure 1 illustrates the
updates of AdaGrad and Shampoo, where AdaGrad update uses the entire matrix
instead of the diagonal, which is shown in the figure. In other words, Shampoo
uses a Kronecker-factored preconditioner, and the factorization preserves the
matrix structure of the parameters. Since in DL optimization parameters often
have matrix structure, Shampoo has strong empirical performance, and has
improved upon state-of-the-art results in large-scale tasks such as language
modeling with BERT-Large and image classification on ImageNet in [9].

Figure 1 elaborates why the composition of FD and Shampoo is essential to
avoid memory consumption asymptotically greater than parameter count for
approximate full matrix regularization.

However, Shampoo memory costs may still be prohibitive for rectangular weight
matrices. In BERT-Large [29], most parameters are in the feed-forward network
layers, which consist of $4096\times 1024$ dense kernels; other transformers
follow similar narrow-to-wide patterns. For large models, occupying even
$4\times$ memory for the left preconditioner can frequently result in OOM in
memory-constrained settings; this was in fact one of the practical motivations
for our proposed approach.

Anil et al. [9] introduces two workarounds for the problem of rectangular
matrices based on limiting covariance modelling. Furthermore, both
approximations can be applied to our method, so we do not compare against
them. First, the authors propose Blocked Shampoo, which views each weight
matrix $W$ of shape $m\times n$ as $mn/b^{2}$ blocks of size $b\times b$ for
some block size $b<\min(m,n)$ (in the limit $b=1$, this recovers diagonal
AdaGrad). This approach is dependent on the ordering of neurons in hidden
layers. Another approximation relies on only one-sided covariance upper
bounds, $L_{t}\otimes I$ or $I\otimes R_{t}$. Note, however, that the one-
sided approximation doesn’t help with vector parameters, such as those that
appear for the bias terms in dense layers or layer norms [30]. For 3D weights,
such as those which appear in homogeneous Mixtures of Experts [31], blocking
increases memory consumption. These approaches do not take into account the
fast decay of the preconditioner’s spectrum, which is the focus of our work.

##  4 Algorithms and Main Theorems

In this section, we introduce the adaptation of Frequent Directions (FD) to
AdaGrad (Sec. 4.1) and Shampoo (Sec. 4.2), the corresponding algorithms and
regret guarantees. Additionally, in Sec. 4.3, we modify FD to support
exponential moving averages.

The main technical novelty in incorporating Alg. 1 to AdaGrad (Alg. 2) and
Shampoo (Alg. 3) is the construction of preconditioning matrices with FD-
sketched matrices compensated by the cumulative escaped masses. The insight of
such construction lies in the observation that while the FD sketch lower
bounds the full preconditioning matrix, the FD sketch compensated with the
cumulative escaped masses upper bounds the full preconditioning matrix, as
demonstrated in Lemma 10 for AdaGrad and Lemma 14 for Shampoo. The regret
guarantee for AdaGrad ([2]) and Shampoo ([5]) directly depends on the trace of
the preconditioning matrices, therefore obtaining upper and lower bounds on
the preconditioning matrices allows explicit additive dependence on the
cumulative escaped mass. We expect this approach to be reusable for
alternative approximation schemes.

###  4.1 FD for AdaGrad

Our main algorithm in this section is Alg. 2 run with FD (Alg. 1) as the
sketching method. $\tilde{G}_{t}^{-1/2}$ in Alg. 2 denotes the Moore-Penrose
pseudoinverse of the matrix $\tilde{G}_{t}^{1/2}$. Our main algorithm, Sketchy
AdaGrad, in this section exploits the FD approach outlined in Alg. 1 as the
sketching method in AdaGrad. In particular, at every time step, we pass the
newly received subgradient $g_{t}$ into Alg. 1, which updates and maintains a
low-rank sketch $\bar{G}_{t}$ of the AdaGrad preconditioning matrix $G_{t}$.
We keep track of the cumulative escaped mass $\rho_{1:t}$, which we add back
to the low-rank sketch to create the Sketchy preconditioner $\tilde{G}_{t}$,
with which we perform the regular AdaGrad descent and projection.

1: Input: constraint set $\mathcal{K}$, step size $\eta$, time horizon $T$.

2: Initialize $x_{1}\in\mathcal{K}$, $\bar{G}_{0}=\tilde{G}_{0}=0$.

3: for $t=1,\ldots,T$ do

4: Play $x_{t}$, receive $g_{t}\in\partial f_{t}(x_{t})$, suffer cost
$f_{t}(x_{t})$.

5: Sketch $(\rho_{t},\bar{G}_{t})=\texttt{FD-
update}(\bar{G}_{t-1},g_{t}g_{t}^{\top})$.

6: Update $\tilde{G}_{t}=\bar{G}_{t}+\rho_{1:t}I$,
$y_{t+1}=x_{t}-\eta\tilde{G}_{t}^{-1/2}g_{t}$, and
$x_{t+1}=\underset{x\in\mathcal{K}}{\operatorname*{argmin}}\
\|y_{t+1}-x\|_{\tilde{G}_{t}^{1/2}}^{2}$.

7: end for

Algorithm 2 Sketchy AdaGrad (S-AdaGrad)

######  Theorem 3.

Define
$\Omega_{\ell}=\min_{k<\ell}(\ell-k)^{-1}\sum_{i=k+1}^{d}\lambda_{i}(G_{T})$,
then with $\eta=\frac{D}{\sqrt{2}}$, Alg. 2 guarantees the following additive
regret bound:

$\mbox{{Regret}}_{T}(\texttt{S-AdaGrad})\leq
D\left(\sqrt{2}\operatorname{tr}G_{T}^{1/2}+d\sqrt{\frac{\Omega_{\ell}}{2}}\right)\,\,,$

where $D$ is the diameter of the constraint set $\mathcal{K}$ if $\mathcal{K}$
is bounded and $\max_{t\in[T]}\|x_{t}-x^{*}\|_{2}$ otherwise.

###### Proof.

See Sec. B.4.1. ∎

Notably in Thm. 3, $\mbox{{Regret}}_{T}=O\left({\sqrt{T}}\right)$ and the
lower eigenvalue dependence $\Omega_{\ell}$ is additive.

######  Corollary 4.

We can improve Theorem 3 slightly to

$\mbox{{Regret}}_{T}(\texttt{S-AdaGrad})\leq
D\left(\sqrt{2}\operatorname{tr}G_{T}^{1/2}+\sqrt{\frac{d(d-\ell)\Omega_{\ell}}{2}}\right)\,\,.$

###### Proof.

See Sec. B.4.2. ∎

The regret bound above holds under the optimal tuning of the learning rate,
which depends on problem quantities that can be unknown a priori. It is
possible to design a parameter-free variant of Alg. 2 by using the norm
$\|x\|_{t}=(x^{\top}(\tilde{G}_{t}+I)^{1/2}x)^{1/2}$ in the projection step of
Alg. 2, as seen in [32].

###  4.2 FD for Shampoo

In this section, we adapt FD-update to Shampoo [5]. For simplicity, we
optimize over $\mathbb{R}^{m\times n}$ in Alg. 3; projection may be handled as
in Alg. 2. Similar to Sketchy AdaGrad, Sketchy Shampoo uses the FD approach
outlined in Alg. 1 to sketch the left and right preconditioning matrices for
Shampoo. In particular, we maintain two parallel sketching streams using Alg.
1 to produce sketches $\bar{L}_{t},\bar{R}_{t}$ for the left and right
preconditioning matrices. We keep track of the cumulative escaped masses
$\rho_{1:t}^{L}$ and $\rho_{1:t}^{R}$ from sketching the left and right
preconditioning matrices, respectively, and compensate the cumulative escaped
mass to create the left and right Sketchy preconditioning matrices
$\tilde{L}_{t},\tilde{R}_{t}$.

1: Input: step size $\eta$, time horizon $T$.

2: Initialize $X_{0}=0_{m\times n}$, $\tilde{L}_{0}=\varepsilon I_{m}$,
$\tilde{R}_{0}=\varepsilon I_{n}$, $\bar{L}_{0}=0_{m}$, $\bar{R}_{0}=0_{n}$.

3: for $t=1,\ldots,T$ do

4: Play $X_{t}$, suffer $f_{t}(X_{t})$, receive $G_{t}\in\partial
f_{t}(X_{t})$.

5: Sketch $(\rho_{t}^{L},\bar{L}_{t})=\texttt{FD-
update}(\bar{L}_{t-1},G_{t}G_{t}^{\top})$,
$(\rho_{t}^{R},\bar{R}_{t})=\texttt{FD-
update}(\bar{R}_{t-1},G_{t}^{\top}G_{t})$.

6: Update $\tilde{L}_{t}=\bar{L}_{t}+\rho^{L}_{1:t}I_{m}$,
$\tilde{R}_{t}=\bar{R}_{t}+\rho^{R}_{1:t}I_{n}$ and
$X_{t+1}=X_{t}-\eta\tilde{L}_{t}^{-1/4}G_{t}\tilde{R}_{t}^{-1/4}$.

7: end for

Algorithm 3 Sketchy Shampoo (S-Shampoo)

Denote
$L_{T}\stackrel{{\scriptstyle\text{def}}}{{=}}\sum_{t=1}^{T}G_{t}G_{t}^{\top}+\varepsilon
I$ and
$R_{T}\stackrel{{\scriptstyle\text{def}}}{{=}}\sum_{t=1}^{T}G_{t}^{\top}G_{t}+\varepsilon
I$.

######  Theorem 5.

Suppose $G_{1},\ldots G_{T}$ have rank at most $r$. Then Alg. 3 run with
$\eta=D/\sqrt{2r}$ guarantees the following regret bound:

$\displaystyle\mbox{{Regret}}_{T}(\texttt{S-Shampoo})\leq\sqrt{2r}D\left({\operatorname{tr}(L_{T}^{1/4})+m\Omega_{L,\ell}^{1/4}}\right)\left({\operatorname{tr}(R_{T}^{1/4})+n\Omega_{R,\ell}^{1/4}}\right)\,\,,$

where $D=\max_{t\in[T]}\|X_{t}-X^{*}\|_{F}$ and
$\Omega_{L,\ell},\Omega_{R,\ell}$ are analogous bounds for
$\rho_{1:T}^{L},\rho_{1:T}^{R}$ from Lem. 1.

###### Proof.

See Sec. B.5.1. ∎

Bounds may be improved analogous to Cor. 4 for Alg. 3, but we omit the similar
statement due to space.

###  4.3 Exponentially Weighted FD

This section discusses the modification of Alg. 1 to support exponential
moving averages. Early in algorithm development, we noticed that attempting to
approximate the unweighted sum of factored gradient covariances
$\sum_{t}G_{t}G_{t}^{\top}$ and $\sum_{t}G_{t}^{\top}G_{t}$ with FD tended to
an estimate of covariance that was roughly $0$, creating numerical
instabilities. Note that FD guarantee (Lem. 1) still holds—but the error term
$\rho_{1:T}$ becomes greater than $\left\lVert G_{T}\right\rVert$, resulting
in a vacuous bound due to lack of spectral decay.

Indeed, Fig. 3 motivating this work only confirmed that the exponential moving
average $L_{t}(\beta_{2})=\sum_{t}\beta_{2}^{T-t}G_{t}G_{t}^{\top}$ exhibits
fast spectral decay (and analogously for $R_{t}$). Luckily, thanks to the
recursion $L_{t+1}(\beta_{2})=\beta_{2}L_{t}+G_{t+1}G_{t+1}^{\top}$, the FD
sketch may easily be adopted for this setting.

######  Observation 6.

Given a stream $g_{t}$ of vectors for $t\in[T]$, sketch size $\ell$, updates
$(\rho_{t}^{(\beta_{2})},\bar{G}_{t}^{(\beta_{2})})=\texttt{FD-
update}(\beta_{2}\bar{G}_{t-1}^{(\beta_{2})},g_{t}g_{t}^{\top})$, and
$G_{T}^{(\beta_{2})}=\sum_{t=1}^{T}\beta_{2}^{T-t}g_{t}g_{t}^{\top}$, we have

$\left\lVert\bar{G}_{T}^{(\beta_{2})}-G_{T}^{(\beta_{2})}\right\rVert\leq\rho_{1:T}^{(\beta_{2})}\leq\min_{k<\ell}\frac{\sum_{i=k+1}^{d}\lambda_{i}\left({G_{T}^{(\beta_{2})}}\right)}{\ell-k}\,\,.$

##  5 Experiments

We investigate how much of Shampoo’s quality our low-memory approach can
recover (Sec. 5.1) and whether the factored covariance exhibits spectral decay
amenable to sketching (Sec. 5.2). We relegate convex examples to Appendix A
due to space constraints, which check that Sketchy compares favorably to
related work in a setting directly related to the theory.

###  5.1 Deep Neural Networks

We evaluate the effectiveness of S-Shampoo as a practical second-order
algorithm for training networks, including

  * •

ResNet-50 [3] for ImageNet image classification task of ImageNet [33] with
random cropping and flipping augmentations.

  * •

A 16-layer Conformer model [34] for the audio transcription task, Librispeech
[35].

  * •

A GNN with 5 message-passing steps [36] on ogbg-molpcba [37], which classifies
structural properties of graphically encoded molecule inputs.

Our FD variant of Shampoo introduces only one new hyperparameter, the rank
$\ell$, which we do not tune, but set to $\ell=256$, which translates to
$4\times$ memory savings for Shampoo blocks of size $1024$ for the
accumulators. Shampoo, Adam, and the underlying architectures introduce their
own hyperparameters. S-Shampoo inherits those of Shampoo. We tune only common
parameters between the three optimizers with the same budgets, selecting based
on validation set accuracy. In the ImageNet case, we evaluate final test set
performance using ImageNet v2 [38], as the ImageNet test set is unavailable.
Additinoal training information is available in Appendix C.

Figure 2: Test metrics are classification error rate for top-1 accuracy for
ImageNet v2, word error rate for Librispeech, and one minus average precision
for OGBG-MolPCBA. We plot the mean of 5 random seeds, with 1.96 times the
standard error as error bars. For readers familiar with ImageNet v1, final
validation accuracy for Shampoo was 77.69% (0.03%), S-Shampoo having 77.18%
(0.04%), and Adam having 76.76% (0.03%), but we emphasize that due to tuning,
the test set performance pictured above should be of primary concern.

As Fig. 2 demonstrates, the second-order information leveraged by Shampoo
results in improvements over Adam, a first-order method. Our method performs
at least as well as Adam in all cases, despite using asympotically less memory
to represent covariance (as Fig. 1 shows, Adam uses $O(mn)$ for a rectangular
weight matrix’s diagonal accumulators, whereas S-Shampoo uses $O(mk+nk)$). In
the GNN case (OBGB-MolPCBA), Shampoo does not perform as well; its training
curves indicate overfitting, but we note that S-Shampoo was less susceptible
to overfit like Adam.

###  5.2 Spectral Analysis

To explain Sketchy’s strong performance in Sec. 5.1, we inspect the
exponential moving average of Kronecker-factored gradient covariance for fast
spectral decay. We find that this is indeed the case in practice, so Sketchy’s
low-rank plus diagonal covariance is representative of true training
statistics.

For all our architectures, we tune Shampoo and extract the intermediate
gradient covariances over the course of training. To make our curves
comparable across architectures, we fix the parameter for the second moment,
$\beta_{2}=0.999$ for these runs. Furthermore, ResNet-50 has a few parameters
with dimension 2048, but the largest dimension for any parameter from the
other two architectures is 1024, so we use the Blocked Shampoo variant
discussed in Sec. 3.4 with block size 1024. In other words, weights containing
a dimension 2048 are split into two. We tune other Shampoo parameters for each
architecture, and plot statistics of Kronecker factors
$L_{t}=\sum_{i}\beta_{2}^{t-i}G_{t}G_{t}^{\top}$ and
$R_{t}=\sum_{i}\beta_{2}^{t-i}G_{t}^{\top}G_{t}$.

Figure 3: Two measures of spectral decay. As described in Sec. 5.2, we
demonstrate spectral decay of $L_{t}$ and $R_{t}$ covariance factors of shape
$1024\times 1024$. On the left, we take the factors $C$ from the first layer
of each network and plot the proportion of spectral mass captured by its top
$256$ eigenvalues throughout training, i.e.,
$\sum_{i=1}^{256}\lambda_{i}(C)/\sum_{i=1}^{1024}\lambda_{i}(C)$. On the
right, we plot a continuous measure of eigenvalue concentration, the intrinsic
dimension $\operatorname{tr}C/\lambda_{\max}(C)$, typically $10\times$ smaller
than nominal dimension. We plot the average across all covariances for each
network’s weight’s dimensions, with the shaded regions capturing the
interquartile range.

In Fig. 3, we plot the intrinsic dimension of Kronecker covariance factors
over training for our three settings. The intrinsic dimension determines the
rate at which empirical covariance estimates concentrate to their expectation,
rather than a random vector’s actual dimension, up to logarithmic factors
(Vershynin [39], Remark 5.6.3). Despite actual dimensionality being over 1024,
intrinsic dimension across all architectures stays below 105. A conspicuous
phase shift 10% of the way through training may be the result of a change from
linear learning rate warmup to a learning rate decay, starting at roughly 5%
of the way into training.

Given $\beta_{2}=0.999$, we emphasize that the behavior in Fig. 3 is an
emergent property of DL training. Though surely a lower $\beta_{2}$ would
naturally result in lower intrinsic dimension (which can still be taken
advantage of by Alg. 2 and 3), we would still expect higher intrinsic
dimension if covariances were near-isometries. If we observe some large number
$n=10000$ draws $x_{i}$ of $1024\times d$ matrices with iid $N(0,1)$ entries,
then numerical experiments show that the average intrinsic dimension of
$\sum_{i=0}^{n-1}\beta_{2}^{i}x_{i}x_{i}^{\top}$ is $324.63$ ($0.52$) and
$862.13$ ($0.25$) for $d=1,64$, respectively, with parenthesized numbers
denoting standard error across 20 trials. Values generated this way are larger
than the average intrinsic dimension of roughly 10, 30, 50 observed in Fig. 3.

##  6 Discussion

Up to spectral error, Alg. 2 achieves full-matrix AdaGrad regret despite
approximating the _smallest_ part of the spectrum of $G_{t}^{-1/2}$ at each
step. Remarkably, these eigenvectors correspond to the _most_ easily
discernible signals of the covariance for the stream $g_{t}$. This apparent
(and fortuitous) coincidence is resolved by considering the covariance of
$\tilde{G}_{t}^{-1/2}g_{t}$: whitening the gradient to facilitate optimization
best reflects on regret; as a result, approximating top eigenvectors of
$G_{T}$ helps more than the bottom ones.

Our initial implementation focused on correctness rather than physical speed
or memory reduction. Engineering optimizers competitive with existing
industrial-strength implementations of Adam and Shampoo was out of scope. In
implementing FD, we performed updates via the factored SVD of
$[\beta_{2}^{1/2}B_{t};G_{t}]$ rather than the eigendecomposition depicted in
Alg. 1; this avoids squaring, which is unavoidable in Shampoo. For speed,
Shampoo subsamples gradients for its covariance estimation and updates its
inverse matrix roots intermittently, every fixed number of steps. A tuning
script provided by Anil et al. [9] included gradients from every step, but
updated roots every 10 steps. Since FD does not separate sampling from its
computation of estimated covariance eigendecomposition, we took the more
difficult setting for S-Shampoo, only allowing it to simultaneously observe
every 10th gradient and update its covariance inverse roots (see Appendix G
for a theoretical justification).

Though step-skipping makes Shampoo and S-Shampoo tractable, future work may
explore further speedups: since FD only requires the top $\ell$ eigenvalues,
iterative Lanczos-like routines which are accelerator-friendly, such as LOBPCG
[40], may allow incremental updates to $\tilde{G}_{t}^{-1/2}$ in factored form
with only a few matrix multiplies, S-Shampoo may be able to update more
frequently than its non-sketched counterpart, further improving quality.

##  7 Conclusion

In this work, we address a gap in the OCO literature for low-memory
optimization with the novel Alg. 2 and demonstrate its relevance to practical
non-convex problems such as neural net training (Sec. 5.1) by leveraging a new
observation about gradient covariance (Sec. 5.2).

The growing disparity between compute capability and memory bandwidth [11]
underscores the need for further research in this direction. Further, large-
batch settings reduce the performance gap between first and Shampoo-based
second order methods, since the batch-size independent runtime of the
optimizer is amortized per example used for the gradient calculation. Even in
performing experiments for this work, we would frequently find that faster
accelerators were unavailable, but many previous-generation ones were,
encouraging us to leverage data-parallel training. For datasets such as
Imagenet, we notice the advantage of second order methods in dealing with
large batches even at relatively modest sizes, such as 1024; many works on
explore several larger multiples of this [41].

Potential for future work includes numerical methods outlined in the previous
section as well optimizing the rank $\ell$ across the many tensors in a
network, as the spread in Fig. 3 highlights the large variance in covariance
intrinsic dimension. Furthermore, the inductive biases conferred by the minima
which different-rank representations of curvature reach may have problem-
dependent generalization implications, a question which we leave for future
work. For a comparison of full rank preconditioning’s effect versus first-
order minima, see Amari et al. [42].

## References

  * Kingma and Ba [2015] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In _ICLR (Poster)_ , 2015.
  * 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.
  * 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.
  * Martens and Grosse [2015] James Martens and Roger Grosse.  Optimizing neural networks with kronecker-factored approximate curvature.  In _International conference on machine learning_ , pages 2408–2417. PMLR, 2015.
  * 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.
  * Agarwal et al. [2019] Naman Agarwal, Brian Bullins, Xinyi Chen, Elad Hazan, Karan Singh, Cyril Zhang, and Yi Zhang.  Efficient full-matrix adaptive regularization.  In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pages 102–110. PMLR, 09–15 Jun 2019.
  * Chen et al. [2019] Xinyi Chen, Naman Agarwal, Elad Hazan, Cyril Zhang, and Yi Zhang.  Extreme tensoring for low-memory preconditioning.  In _International Conference on Learning Representations_ , 2019.
  * Anil et al. [2019] Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer.  Memory efficient adaptive optimization.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Anil et al. [2020] Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer.  Scalable second order optimization for deep learning.  _arXiv preprint arXiv:2002.09018_ , 2020.
  * Anil et al. [2022] Rohan Anil, Sandra Gadanho, Da Huang, Nijith Jacob, Zhuoshu Li, Dong Lin, Todd Phillips, Cristina Pop, Kevin Regan, Gil I. Shamir, Rakesh Shivanna, and Qiqi Yan.  On the factory floor: Ml engineering for industrial-scale ads recommendation models, 2022.
  * Jouppi et al. [2021] Norman P Jouppi, Doe Hyun Yoon, Matthew Ashcraft, Mark Gottscho, Thomas B Jablin, George Kurian, James Laudon, Sheng Li, Peter Ma, Xiaoyu Ma, et al.  Ten lessons from three generations shaped google’s tpuv4i: Industrial product.  In _2021 ACM/IEEE 48th Annual International Symposium on Computer Architecture (ISCA)_ , pages 1–14. IEEE, 2021.
  * Dally et al. [2021] William J Dally, Stephen W Keckler, and David B Kirk.  Evolution of the graphics processing unit (gpu).  _IEEE Micro_ , 41(6):42–51, 2021.
  * Ghashami et al. [2016] Mina Ghashami, Edo Liberty, Jeff M Phillips, and David P Woodruff.  Frequent directions: Simple and deterministic matrix sketching.  _SIAM Journal on Computing_ , 45(5):1762–1792, 2016.
  * Shazeer and Stern [2018] Noam Shazeer and Mitchell Stern.  Adafactor: Adaptive learning rates with sublinear memory cost.  In _International Conference on Machine Learning_ , pages 4596–4604. PMLR, 2018.
  * Hazan et al. [2016] Elad Hazan et al.  Introduction to online convex optimization.  _Foundations and Trends® in Optimization_ , 2(3-4):157–325, 2016.
  * Liberty [2022] Edo Liberty.  Even simpler deterministic matrix sketching.  _arXiv preprint arXiv:2202.01780_ , 2022.
  * Sagun et al. [2016] Levent Sagun, Leon Bottou, and Yann LeCun.  Eigenvalues of the hessian in deep learning: Singularity and beyond.  _arXiv preprint arXiv:1611.07476_ , 2016.
  * Sagun et al. [2017] Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou.  Empirical analysis of the hessian of over-parametrized neural networks.  _arXiv preprint arXiv:1706.04454_ , 2017.
  * 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.
  * Sankar et al. [2021] Adepu Ravi Sankar, Yash Khasbage, Rahul Vigneswaran, and Vineeth N Balasubramanian.  A deeper look at the hessian eigenspectrum of deep neural networks and its applications to regularization.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 35, pages 9481–9488, 2021.
  * Agarwal et al. [2017] Naman Agarwal, Zeyuan Allen-Zhu, Brian Bullins, Elad Hazan, and Tengyu Ma.  Finding approximate local minima faster than gradient descent.  In _Proceedings of the 49th Annual ACM SIGACT Symposium on Theory of Computing_ , pages 1195–1199, 2017.
  * Gur-Ari et al. [2018] Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer.  Gradient descent happens in a tiny subspace.  _arXiv preprint arXiv:1812.04754_ , 2018.
  * Bakker et al. [2018] Craig Bakker, Michael J Henry, and Nathan O Hodas.  Understanding and exploiting the low-rank structure of deep networks.  2018\.
  * Xie et al. [2022] Zeke Xie, Qian-Yuan Tang, Zheng He, Mingming Sun, and Ping Li.  Rethinking the structure of stochastic gradients: Empirical and statistical evidence.  _arXiv preprint arXiv:2212.02083_ , 2022.
  * Krummenacher et al. [2016] Gabriel Krummenacher, Brian McWilliams, Yannic Kilcher, Joachim M Buhmann, and Nicolai Meinshausen.  Scalable adaptive stochastic optimization using random projections.  _Advances in Neural Information Processing Systems_ , 29, 2016.
  * Wan and Zhang [2021] Yuanyu Wan and Lijun Zhang.  Efficient adaptive online learning via frequent directions.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 2021.
  * Luo et al. [2016] Haipeng Luo, Alekh Agarwal, Nicolo Cesa-Bianchi, and John Langford.  Efficient second order online learning by sketching.  _Advances in Neural Information Processing Systems_ , 29, 2016.
  * Chen et al. [2021] Cheng Chen, Luo Luo, Weinan Zhang, Yong Yu, and Yijiang Lian.  Efficient and robust high-dimensional linear contextual bandits.  In _Proceedings of the Twenty-Ninth International Conference on International Joint Conferences on Artificial Intelligence_ , pages 4259–4265, 2021.
  * Devlin et al. [2019] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 4171–4186, 2019.
  * Ba et al. [2016] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton.  Layer normalization.  _arXiv preprint arXiv:1607.06450_ , 2016.
  * Shazeer et al. [2017] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean.  Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.  _arXiv preprint arXiv:1701.06538_ , 2017.
  * Cutkosky [2020] Ashok Cutkosky.  Better full-matrix regret via parameter-free online learning.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, _Advances in Neural Information Processing Systems_ , volume 33, pages 8836–8846. Curran Associates, Inc., 2020.
  * Russakovsky et al. [2015] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei.  ImageNet Large Scale Visual Recognition Challenge.  _International Journal of Computer Vision (IJCV)_ , 115(3):211–252, 2015.  doi: 10.1007/s11263-015-0816-y.
  * Gulati et al. [2020] Anmol Gulati, James Qin, Chung-Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, et al.  Conformer: Convolution-augmented transformer for speech recognition.  _arXiv preprint arXiv:2005.08100_ , 2020.
  * Panayotov et al. [2015] Vassil Panayotov, Guoguo Chen, Daniel Povey, and Sanjeev Khudanpur.  Librispeech: an asr corpus based on public domain audio books.  In _Acoustics, Speech and Signal Processing (ICASSP), 2015 IEEE International Conference on_ , pages 5206–5210. IEEE, 2015.
  * Battaglia et al. [2018] Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al.  Relational inductive biases, deep learning, and graph networks.  _arXiv preprint arXiv:1806.01261_ , 2018.
  * Hu et al. [2020] Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec.  Open graph benchmark: Datasets for machine learning on graphs.  In Hugo Larochelle, Marc Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, _Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual_ , 2020.
  * Recht et al. [2019] Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar.  Do imagenet classifiers generalize to imagenet?  In _International Conference on Machine Learning_ , pages 5389–5400, 2019.
  * Vershynin [2018] Roman Vershynin.  _High-dimensional probability: An introduction with applications in data science_ , volume 47.  Cambridge university press, 2018.
  * Knyazev [2001] Andrew V Knyazev.  Toward the optimal preconditioned eigensolver: Locally optimal block preconditioned conjugate gradient method.  _SIAM journal on scientific computing_ , 23(2):517–541, 2001.
  * Keskar et al. [2016] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang.  On large-batch training for deep learning: Generalization gap and sharp minima.  _arXiv preprint arXiv:1609.04836_ , 2016.
  * Amari et al. [2020] Shun-ichi Amari, Jimmy Ba, Roger Baker Grosse, Xuechen Li, Atsushi Nitanda, Taiji Suzuki, Denny Wu, and Ji Xu.  When does preconditioning help or hurt generalization?  In _International Conference on Learning Representations_ , 2020.
  * Luo et al. [2019] Luo Luo, Cheng Chen, Zhihua Zhang, Wu-Jun Li, and Tong Zhang.  Robust frequent directions with application in online learning.  _The Journal of Machine Learning Research_ , 20(1):1697–1737, 2019.
  * Chang and Lin [2011] Chih-Chung Chang and Chih-Jen Lin.  Libsvm: a library for support vector machines.  _ACM transactions on intelligent systems and technology (TIST)_ , 2(3):1–27, 2011.
  * Hazan et al. [2014] Elad Hazan, Tomer Koren, and Kfir Y Levy.  Logistic regression: Tight bounds for stochastic and online optimization.  In _Conference on Learning Theory_ , pages 197–209. PMLR, 2014.
  * Kalai and Vempala [2005] Adam Kalai and Santosh Vempala.  Efficient algorithms for online decision problems.  _Journal of Computer and System Sciences_ , 71(3):291–307, 2005.
  * Audenaert [2014] Koenraad MR Audenaert.  A generalisation of mirsky’s singular value inequalities.  _arXiv preprint arXiv:1410.4941_ , 2014.
  * Gilmer et al. [2021] Justin M. Gilmer, George E. Dahl, and Zachary Nado.  init2winit: a jax codebase for initialization, optimization, and tuning research, 2021.  URL http://github.com/google/init2winit.
  * Bradbury et al. [2018] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang.  JAX: composable transformations of Python+NumPy programs, 2018\.
  * Heek et al. [2020] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee.  Flax: A neural network library and ecosystem for JAX, 2020.
  * TFD [2023] Tensorflow datasets, a collection of ready-to-use datasets.  https://www.tensorflow.org/datasets, 2023.
  * Waskom [2021] Michael L. Waskom.  seaborn: statistical data visualization.  _Journal of Open Source Software_ , 6(60):3021, 2021.  doi: 10.21105/joss.03021.
  * Hunter [2007] J. D. Hunter.  Matplotlib: A 2d graphics environment.  _Computing in Science & Engineering_, 9(3):90–95, 2007.  doi: 10.1109/MCSE.2007.55.
  * Harris et al. [2020] Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Rio, Mark Wiebe, Pearu Peterson, Pierre Gérard-Marchant, Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant.  Array programming with NumPy.  _Nature_ , 585(7825):357–362, September 2020\.  doi: 10.1038/s41586-020-2649-2.
  * Virtanen et al. [2020] Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C J Carey, İlhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors.  SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python.  _Nature Methods_ , 17:261–272, 2020.  doi: 10.1038/s41592-019-0686-2.
  * MLCommons® open engineering consortium [2023] MLCommons® open engineering consortium.  MLCommons Algorithmic Efficiency.  https://github.com/mlcommons/algorithmic-efficiency, 2023.
  * Reddi et al. [2019] Vijay Janapa Reddi, Christine Cheng, David Kanter, Peter Mattson, Guenther Schmuelling, Carole-Jean Wu, Brian Anderson, Maximilien Breughe, Mark Charlebois, William Chou, Ramesh Chukka, Cody Coleman, Sam Davis, Pan Deng, Greg Diamos, Jared Duke, Dave Fick, J. Scott Gardner, Itay Hubara, Sachin Idgunji, Thomas B. Jablin, Jeff Jiao, Tom St. John, Pankaj Kanwar, David Lee, Jeffery Liao, Anton Lokhmotov, Francisco Massa, Peng Meng, Paulius Micikevicius, Colin Osborne, Gennady Pekhimenko, Arun Tejusve Raghunath Rajan, Dilip Sequeira, Ashish Sirasao, Fei Sun, Hanlin Tang, Michael Thomson, Frank Wei, Ephrem Wu, Lingjie Xu, Koichi Yamada, Bing Yu, George Yuan, Aaron Zhong, Peizhao Zhang, and Yuchen Zhou.  Mlperf inference benchmark, 2019.
  * Loshchilov and Hutter [2017] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.
  * Dayma and Anil [2022] Boris Dayma and Rohan Anil.  Evaluation of Distributed Shampoo: Comparison of optimizers: Distributed Shampoo, Adam & Adafactor.  Weights & Biases Report, 2022.
  * Hinton et al. [2012] Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky.  Neural networks for machine learning lecture 6a overview of mini-batch gradient descent.  _Cited on_ , 14(8):2, 2012.
  * Agarwal et al. [2020] Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang.  Disentangling adaptive gradient methods from learning rates.  _arXiv preprint arXiv:2002.11803_ , 2020.
  * Ando [1979] Tsuyoshi Ando.  Concavity of certain maps on positive definite matrices and applications to hadamard products.  _Linear algebra and its applications_ , 26:203–241, 1979\.
  * Bhatia [1997] Rajendra Bhatia.  _Matrix analysis_.  Springer, 1997.
  * Petersen et al. [2008] Kaare Brandt Petersen, Michael Syskind Pedersen, et al.  The matrix cookbook.  _Technical University of Denmark_ , 7(15):510, 2008.
  * Brockett [2015] Roger W Brockett.  _Finite dimensional linear systems_.  SIAM, 2015.

##  Appendix A Online Convex Examples

In this section, we evaluate the performance of S-Adagrad in the classical
online convex optimization setting.

We consider three different datasets, summarized in Tbl. 2. For each dataset,
we evaluate several works based on the frequent directions sketch, Ada-FD
[26], FD-SON [27], RFD-SON [43]. As baselines we also add online gradient
descent (OGD) and diagonal Adagrad [2]. For each dataset, we consider the
online convex loss of a logistic binary classification loss over a linear
learner on the features in each dataset augmented with a constant feature for
the intercept.

Table 2: Dataset statistics for simple online convex examples. We stream through each example, providing a logistic loss over the binary label and linear predictor generated by an OCO learner. The feature count includes an all-constant intercept column as well. Datasets were retrieved from Chang and Lin [44]. Dataset Name | Number of Examples | Number of Features
---|---|---
gisette_scale | 6000 | 5001
a9a | 32561 | 124
cifar10 | 50000 | 3073

For methods which require a nonzero diagonal initial regularizer $\delta I$,
namely FD-SON and Ada-FD, we tune $\delta$ in $10^{-6},10^{-5},\cdots 1$ and
the learning rate $\eta$ over the same range as well, for 49 hyperparameters
total. For methods which have $\delta=0$ (Adagrad, OGD, S-Adagrad, RFD-SON),
we instead tune $\eta$ on 49 points spaced evenly on the same logarithmic
scale, $[10^{-6},1]$ to fairly allocate hyperparameter training budget. Note
that the variant of RFD-SON, $\mathrm{RFD}_{0}$, which sets $\delta=0$ is the
main variant evaluated by Luo et al. [43]. The sketch size was fixed to be 10
throughout.

Table 3: Average cumulative online loss across datasets and algorithms, ranked from lowest (1st place) to highest (6th place). Our proposal, S-Adagrad, is in bold. | Place | 1 | 2 | 3 | 4 | 5 | 6
---|---|---|---|---|---|---|---
cifar10 | Alg. | RFD-SON | S-Adagrad | Adagrad | OGD | Ada-FD | FD-SON
| Loss | 0.297 | 0.297 | 0.303 | 0.308 | 2.999 | $6\times 10^{6}$
gisette | Alg. | S-Adagrad | RFD-SON | Ada-FD | Adagrad | OGD | FD-SON
| Loss | 0.158 | 0.167 | 0.196 | 0.209 | 0.224 | 2.432
a9a | Alg. | Adagrad | S-Adagrad | OGD | RFD-SON | Ada-FD | FD-SON
| Loss | 0.332 | 0.333 | 0.335 | 0.335 | 0.354 | 0.539

We sort and display cumulative average regret in Tbl. 3 and Fig. 4. S-Adagrad
is the only method to consistently place among the top three across all
datasets. We suspect that a combination of allowing $\delta=0$ (removing
inductive bias about regularization; notice in Tbl. 3 that Ada-FD and FD-SON,
the two methods with $\delta>0$, routinely place last) and the ability to deal
with the effectively zero exp-concavity constant of the logistic loss [45]
explain S-Adagrad’s performance.

Figure 4: Cumulative average regret curves for logistic loss as a function of
percent of dataset completion in a single online pass, resummarizing Tbl. 3
visually.

##  Appendix B Proof details

##### Notation.

Let $\|\cdot\|$ denote the $\ell_{2}$ norm for a vector. Let $\|\cdot\|_{op}$,
$\|\cdot\|_{F}$ denote the operator norm and the Frobenius norm of a matrix,
respectively. For a positive definite matrix $A$, we use
$\|x\|_{A}=\sqrt{x^{\top}Ax}$ to denote the matrix norm induced by $A$, and
$\|x\|_{A}^{*}=\sqrt{x^{\top}A^{-1}x}$ to denote the dual norm of the induced
matrix norm. For a matrix $A$, $A^{-1}$ is the inverse of $A$ if $A$ is full
rank; otherwise, $A^{-1}$ is taken to be the Moore-Penrose pseudoinverse.
Finally, $\overline{\text{vec}}\left({\cdot}\right)$ denotes the row-major
vectorization of a given matrix, and $\otimes$ denotes the Kronecker product
between two matrices

###  B.1 Reduction from Non-convex Optimization to Online Convex Optimization

In this section, we give more details for the reduction of non-convex
optimization to online convex optimization for completeness. We use the
framework of [6], though many related results exist in the optimization
literature. The algorithm is stated in Algorithm 4, where we optimize the non-
convex function $f$ by iteratively optimizing subproblems $f_{t}$ that are
strongly convex. Given any OCO algorithm $\mathcal{A}$, in each episode, we
first construct $f_{t}$ and then use $\mathcal{A}$ to optimize it for $N$ time
steps.

1: Input: initial $x_{1}$, time horizon $T$, episode length $N$, smoothness
parameter $L$, online convex optimization algorithm $\mathcal{A}$.

2: for $t=1,\ldots,T$ do

3: Construct $f_{t}(x)=f(x)+L\|x-x_{t}\|^{2}$.

4: Initialize $\mathcal{A}$ to start at $x_{t}$, and set $x_{t}^{1}=x_{t}$.

5: for $n=1,\ldots,N$ do

6: Play $x_{t}^{n}$, receive stochastic gradient
$\tilde{\nabla\mkern-2.5mu}f_{t}(x_{t}^{n})$, construct
$g_{t}^{n}(x)=\tilde{\nabla\mkern-2.5mu}f_{t}(x_{t}^{n})^{\top}(x)$.

7: Update $x_{t}^{n+1}=\mathcal{A}(g_{t}^{1},\ldots,g_{t}^{n})$.

8: end for

9: Update $x_{t+1}=\frac{1}{N}\sum_{n=1}^{N}x_{t}^{n}$.

10: end for

11: Output iterate
$x_{t^{*}}=\operatorname*{argmin}_{t\in[T+1]}\|\nabla\mkern-2.5muf(x_{t})\|$.

Algorithm 4 Reduction from Non-convex optimization to Online Convex
Optimization

We state the convergence guarantees in terms of the adaptive ratio, defined
below.

######  Definition 7.

(Adaptive ratio.) Let $\mathcal{A}$ be an algorithm, and consider a convex
function $f$. Given a stochastic gradient oracle with variance bounded by
$\sigma^{2}$, let $x_{\mathcal{A}}$ be the output of $\mathcal{A}$ with at
most $T$ oracle calls, and let $x^{*}\in\operatorname*{argmin}_{x}f(x)$.
Define the adaptive ratio of $\mathcal{A}$ as

$\mu_{\mathcal{A}}(f)=\frac{f(x_{\mathcal{A}})-f(x^{*})}{\|x_{1}-x^{*}\|\frac{\sigma}{T}}.$

The adaptive ratio captures the performance of $\mathcal{A}$ relative to SGD.
For certain algorithms, such as AdaGrad [2], the adaptive ratio can be as
small as $\frac{1}{\sqrt{d}}$. For more discussions on this notion see [6].

######  Theorem 8.

(Theorem A.2 in [6]) Consider a non-convex function $f$, and suppose $f$ is
$L$-smooth and bounded: $\|\nabla\mkern-2.5mu^{2}f(x)\|_{2}\leq L$ and
$\max_{x,y}f(x)-f(y)\leq M$. Additionally suppose we have access to a
stochastic gradient oracle with variance bounded by $\sigma^{2}$. Let
$\mu=\max_{t}\mu_{\mathcal{A}}(f_{t})$. Then Algorithm 4 run with
$T=\frac{12ML}{\epsilon^{2}}$ and $N=\frac{48\mu^{2}\sigma^{2}}{\epsilon^{2}}$
returns a point $x_{t^{*}}$ such that

$\operatorname{\mathbb{E}}[\|\nabla\mkern-2.5muf(x_{t^{*}})\|]\leq\varepsilon.$

The total number of calls to the stochastic gradient oracle is bounded by
$T\cdot N=O(\mu^{2}\sigma^{2}/\epsilon^{4})$.

###  B.2 Proof of Lem. 1

###### Proof.

Let $H_{T}\in\mathbb{R}^{T\times d}$ denote the matrix of stacked gradients,
where the $t$-th row of $H_{T}$ is $g_{t}$. Then $H_{T}^{\top}H_{T}=G_{T}$,
and FD iteratively sketches $H_{T}$. Let $H_{T}=U\Sigma V^{\top}$ be the SVD
of $H_{T}$, and let $H_{T,k}=U_{k}\Sigma_{k}V_{k}^{\top}$ denote the best
rank-$k$ approximation of $H_{T}$, where $U_{k},V_{k}$ are the first $k$
columns of the matrices, and $\Sigma_{k}$ is the upper left $k\times k$
submatrix of $\Sigma$. By the proof of Theorem 1.1 in [13], we have

$\displaystyle\rho_{1:T}\leq\min_{k<\ell}\frac{\|H_{T}-H_{T,k}\|_{F}^{2}}{\ell-k}=\min_{k<\ell}\frac{\sum_{i=k+1}^{d}\lambda_{i}(H_{T}^{\top}H_{T})}{\ell-k}=\min_{k<\ell}\frac{\sum_{i=k+1}^{d}\lambda_{i}(G_{T})}{\ell-k}\leq\sum_{i=\ell}^{d}\lambda_{i}(G_{T})\,\,,$

where the last inequality follows by choosing $k=\ell-1$. ∎

###  B.3 Proof of Observation 2

###### Proof.

Let $\Sigma=\operatorname{\mathbb{E}}\left[{g_{t}g_{t}^{\top}}\right]$ denote
the covariance of the gradients, and $\lambda_{i}=\lambda_{i}(\Sigma)$ denote
its $i$-th eigenvalue. By definition, $g_{t}$ has the following distribution:
$g_{t}=w_{i}$ with probability $\lambda_{i}$. At iteration $t$, we have the
current sketch $\bar{G}_{t-1}\in\mathbb{R}^{\ell\times d}$, and we receive the
new gradient $g_{t}$. Ada-FD uses $\bar{G}_{t-1}+\delta I$ as their
preconditioner.

We first show that under the distribution of the cost functions, if
$\bar{G}_{t-1}$ has rank $\ell-1$, then
$\operatorname{\mathbb{E}}\left[{\rho_{t}|\bar{G}_{t-1}}\right]\geq\sum_{i=\ell}^{r}\lambda_{i}$.
Let $\bar{G}_{t-1}=U\Sigma V^{\top}$ be the SVD of $\bar{G}_{t-1}$, and
$v_{i}$ be the $i$-th row of $V\in\mathbb{R}^{\ell-1\times d}$. Let
$N_{t-1}=W\setminus\\{v_{1},\ldots,v_{\ell-1}\\}$ be the set of basis vectors
not in the row space of $\bar{G}_{t-1}$, then $|N_{t-1}|=r-\ell+1$. If
$g_{t}\in\text{span}(v_{1},\ldots,v_{\ell})$, then $\rho_{t}=0$; otherwise
$\rho_{t}=1$, with probability $\sum_{i:w_{i}\in
N_{t-1}}\lambda_{i}\geq\sum_{i=\ell}^{r}\lambda_{i}$.

We proceed to bound the probability that $\bar{G}_{t-1}$ has rank
$\ell^{\prime}\leq\ell-2$. Note that this event is equivalent to having fewer
than $\ell-1$ distinct vectors drawn from $W$. Let $I_{i}$ be the indicator
variable for drawing $w_{i}$ in the first $t-1$ rounds, then we can obtain the
expected number of distinct vectors as follows

$\displaystyle\operatorname{\mathbb{E}}\left[{\sum_{i=1}^{r}I_{i}}\right]=\sum_{i=1}^{r}\operatorname{\mathbb{E}}\left[{I_{i}}\right]=\sum_{i=1}^{r}1-(1-\lambda_{i})^{t-1}.$

We consider the random variable $r-\sum_{i=1}^{r}I_{i}$, and by Markov’s
inequality,

$\displaystyle\operatorname{\mathbb{P}}\left[r-\sum_{i=1}^{r}I_{i}\geq
r-\ell+2\right]\leq\frac{\sum_{i=1}^{r}(1-\lambda_{i})^{t-1}}{r-\ell+2}\leq\frac{r(1-\lambda_{r})^{t-1}}{2}\,\,.$

Note that this is exactly the probability of having fewer than $\ell-1$
distinct vectors in the first $t-1$ draws. We conclude that for $t\geq\log
r/\lambda_{r}+1$,
$\operatorname{\mathbb{P}}\left[{\text{rank}(\bar{G}_{t-1})=\ell-1}\right]\geq\frac{1}{2}$.
This implies that
$\operatorname{\mathbb{E}}\left[{\rho_{t}}\right]\geq\sum_{i=\ell}^{r}\lambda_{i}/2$
after an initial log number of rounds, and assuming $T\geq 2\log
r/\lambda_{r}$,

$\displaystyle\operatorname{\mathbb{E}}\left[{\sum_{t=1}^{T}\rho_{t}}\right]$
$\displaystyle\geq\operatorname{\mathbb{E}}\left[{\sum_{t=T/2+1}^{T}\rho_{t}}\right]\geq\frac{T}{4}\sum_{i=\ell}^{r}\lambda_{i}.\,\,$

Similarly,

$\displaystyle\operatorname{\mathbb{E}}\left[{\sum_{t=1}^{T}\sqrt{\rho_{t}}}\right]$
$\displaystyle\geq\operatorname{\mathbb{E}}\left[{\sum_{t=T/2+1}^{T}\sqrt{\rho_{t}}}\right]\geq\frac{T}{4}\sum_{i=\ell}^{r}\lambda_{i},$

where the second inequality holds because $\rho_{t}=0$ or $1$, so
$\sqrt{\rho_{t}}=\rho_{t}$ for all $t$. The quantities
$\sum_{t=1}^{T}\rho_{t}$ and $\sum_{t=1}^{T}\sqrt{\rho_{t}}$ correspond to
$\Delta_{T}$ and $\sum_{t=1}^{T}\sqrt{\sigma_{t}}$ in Theorem 1 of [26],
respectively. Therefore, under our setting, the expectation of the upper bound
in Theorem 1 is at least

$\eta\operatorname{\mathbb{E}}\left[{\max\left\\{1,\frac{1+\sqrt{\sum_{t=1}^{T}\rho_{t}}}{\delta}\right\\}\operatorname{tr}(G_{T}^{1/2})}\right]+\frac{D^{2}}{2\eta}\operatorname{\mathbb{E}}\left[{\sum_{t=1}^{T}\sqrt{\rho_{t}}}\right].$
(1)

If we can tune $\delta$, then the $\max$ function evaluates to at least $1$,
and

$\displaystyle(\ref{eq:upper_bound})\geq\eta\operatorname{\mathbb{E}}\left[{\operatorname{tr}(G_{T}^{1/2})}\right]+\frac{D^{2}}{2\eta}\frac{T}{4}\sum_{i=\ell}^{r}\lambda_{i}\geq\eta\operatorname{\mathbb{E}}\left[{\sqrt{\operatorname{tr}(G_{T})}}\right]+\frac{D^{2}T}{8\eta}\sum_{i=\ell}^{r}\lambda_{i}=\eta\sqrt{T}+\frac{D^{2}T}{8\eta}\sum_{i=\ell}^{r}\lambda_{i},$

where the last equality holds because
$\operatorname{tr}(G_{T})=\sum_{t=1}^{T}\|g_{t}\|_{2}^{2}=T$. Optimizing
$\eta$, we conclude that the regret upper bound for Ada-FD is
$\Omega(T^{3/4})$ in expectation.

∎

###  B.4 Proof details for Section 4.1, S-AdaGrad

####  B.4.1 Proof of Theorem 3

The following observations are made of Algorithm 1:

######  Observation 9.

Denote by
$\bar{U}_{t}\stackrel{{\scriptstyle\text{def}}}{{=}}\left[{U_{t};U_{t}^{\bot}}\right]$.
If each $M_{t}$ is of rank $1$, then

$\bar{G}_{t}+\lambda_{\ell}^{(t)}I=\bar{G}_{t-1}+M_{t}+\lambda_{\ell}^{(t)}N_{t}\,\,,$

where $N_{t}=U_{t}^{\bot}\left({U_{t}^{\bot}}\right)^{\top}$.

###### Proof.

By definition of Algorithm 1, $\bar{G}_{t-1}=B_{t-1}B_{t-1}^{\top}$ is of rank
at most $\ell-1$. Under the assumption that $M_{t}$ is of rank $1$,
$\bar{G}_{t-1}+M_{t}$ is of rank at most $\ell$. Therefore,
$\lambda_{\ell+1:d}=0$. Then, we have the following:

$\displaystyle\bar{G}_{t}+\lambda_{\ell}^{(t)}I=U_{t}\operatorname{diag}\left(\lambda_{[1:\ell]}^{(t)}-\lambda_{\ell}^{(t)}\right)U_{t}^{\top}+\lambda_{\ell}^{(t)}I$
$\displaystyle=U_{t}\operatorname{diag}\left(\lambda_{[1:\ell]}^{(t)}-\lambda_{\ell}^{(t)}\right)U_{t}^{\top}+\lambda_{\ell}^{(t)}(U_{t}U_{t}^{\top}+N_{t})$
$\displaystyle=U_{t}\operatorname{diag}\lambda_{[1:\ell]}^{(t)}U_{t}^{\top}+\lambda_{\ell}^{(t)}N_{t}$
$\displaystyle=\bar{U}_{t}\operatorname{diag}\lambda^{(t)}\bar{U}_{t}^{\top}+\lambda_{\ell}^{(t)}N_{t}$
$\displaystyle=\bar{G}_{t-1}+M_{t}+\lambda_{\ell}^{(t)}N_{t}\,\,.$

∎

######  Lemma 10.

Let $\lambda_{\ell}^{(s:t)}$ denote $\sum_{j=s}^{t}\lambda_{\ell}^{(j)}$. Let
$\tilde{G}_{t}\stackrel{{\scriptstyle\text{def}}}{{=}}\bar{G}_{t}+\lambda_{\ell}^{(1:t)}I$,
$G_{t}=\sum_{s=1}^{t}M_{s}$, where each $M_{t}$ is of rank $1$. Let the
initial $G_{0}=\bar{G}_{0}=0$, then the following relation between
$\tilde{G}_{t}$ and $G_{t}$ holds for all $t$:

$\displaystyle\tilde{G}_{t}=G_{t}+\sum_{s=1}^{t}\lambda_{\ell}^{(s)}N_{s}\,\,.$

###### Proof.

The lemma follows from induction on $t$. Base case $\tilde{G}_{0}=G_{0}$ holds
by definition. Suppose the above equation holds for $t-1$. Then,

$\displaystyle\tilde{G}_{t}=\bar{G}_{t}+\lambda_{\ell}^{(1:t)}I$
$\displaystyle=_{1}\bar{G}_{t-1}+M_{t}+\lambda_{\ell}^{(t)}N_{t}+\lambda_{\ell}^{(1:t-1)}I$
$\displaystyle=\tilde{G}_{t-1}+M_{t}+\lambda_{\ell}^{(t)}N_{t}$
$\displaystyle=_{2}G_{t-1}+\sum_{s=1}^{t-1}\lambda_{\ell}^{(s)}N_{s}+M_{t}+\lambda_{\ell}^{(t)}N_{t}$
$\displaystyle=G_{t}+\sum_{s=1}^{t}\lambda_{\ell}^{(s)}N_{s}\,\,,$

where $=_{1}$ follows from Observation 9 and $=_{2}$ follows from induction
hypothesis. ∎

######  Remark 11.

Note that the above lemma immediately provides an approximate isometry
$\bar{G}_{t}\preceq G_{t}\preceq\tilde{G}_{t}$.

Now, we return to the proof of Thm. 3.

###### Proof.

First, we make the following observation of Algorithm 2:

######  Observation 12.

By specification of Algorithm 1,2, $g_{t}\in\text{span}(\tilde{G}_{t})$,
$\forall t$.

We follow the standard AdaGrad [2, 15] analysis. By algorithm specification,

$\displaystyle y_{t+1}-x^{*}$
$\displaystyle=x_{t}-x^{*}-\eta\tilde{G}_{t}^{-1/2}g_{t}\,\,,$
$\displaystyle\tilde{G}_{t}^{1/2}(y_{t+1}-x^{*})$
$\displaystyle=\tilde{G}_{t}^{1/2}(x_{t}-x^{*})-\eta\tilde{G}_{t}^{1/2}\tilde{G}_{t}^{-1/2}g_{t}=_{1}\tilde{G}_{t}^{1/2}(x_{t}-x^{*})-\eta
g_{t}\,\,,$

where $=_{1}$ follows from Observation 12.

With standard AdaGrad analysis, we can bound regret $\mbox{{Regret}}_{T}$
above by the sum of the diameter bound and the gradient bound:

$\displaystyle\underbrace{\frac{1}{2\eta}\sum_{t=1}^{T}\|x_{t}-x_{*}\|^{2}_{\tilde{G}_{t}^{1/2}-\tilde{G}_{t-1}^{1/2}}}_{R_{D}}+\underbrace{\frac{\eta}{2}\sum_{t=1}^{T}\|g_{t}\|_{F_{t}^{-1/2}}}_{R_{G}}\,\,.$

Note that by algorithm specification, we have $\forall t$,

$\displaystyle\tilde{G}_{t}=\bar{G}_{t}+\rho_{1:t}I=_{1}\tilde{G}_{t-1}+g_{t}g_{t}^{\top}+\rho_{t}N_{t}\succeq\tilde{G}_{t-1}+g_{t}g_{t}^{\top},$

where $=_{1}$ follows from the proof of Lemma 10. In particular,
$\tilde{G}_{t}\succeq\tilde{G}_{t-1}$.

Using Remark 11, the gradient norm term in the regret bound can be further
bounded by

$\displaystyle
R_{G}=\frac{\eta}{2}\sum_{t=1}^{T}g_{t}^{\top}\tilde{G}_{t}^{-1/2}g_{t}\leq\frac{\eta}{2}\sum_{t=1}^{T}g_{t}^{\top}G_{t}^{-1/2}g_{t}\leq\eta\operatorname{tr}\left(G_{T}^{1/2}\right)\,\,,$

where the last inequality follows from Lemma 10 of Duchi et al. [2].222The
FTL-BTL lemma [46] alone is not sufficient to justify this inequality, at
least interpreting $G^{-1/2}$ as $\left({G^{1/2}}\right)^{+}$. However, Duchi
et al. [2] rely on concavity of $X\mapsto\operatorname{tr}X^{1/2}$ to show a
semidefinite version of the statement. The diameter norm term in the regret
bound can be bounded by

$\displaystyle R_{D}$
$\displaystyle=\frac{1}{2\eta}\sum_{t=1}^{T}\|x_{t}-x_{*}\|^{2}_{\tilde{G}_{t}^{1/2}-\tilde{G}_{t-1}^{1/2}}\leq_{1}\frac{D^{2}}{2\eta}\operatorname{tr}\left(\tilde{G}_{T}^{1/2}\right)$
$\displaystyle\leq\frac{D^{2}}{2\eta}\operatorname{tr}\left(\left(G_{T}+\rho_{1:T}I\right)^{1/2}\right)$
$\displaystyle\leq_{2}\frac{D^{2}}{2\eta}\left(\operatorname{tr}G_{T}^{1/2}+\operatorname{tr}(\rho_{1:T}I)^{1/2}\right)\,\,,$

where $\leq_{1}$ follows from monotonicity of $\tilde{G}_{t}$’s,
$\|\cdot\|_{op}\leq\operatorname{tr}(\cdot)$ for positive semidefinite
matrices, and linearity of $\operatorname{tr}(\cdot)$, and $\leq_{2}$ follows
from that for $X\in\mathbb{R}^{d},X\succeq 0$, $\operatorname{tr}(X+\sigma
I_{d})^{1/2}\leq\operatorname{tr}(X^{1/2})+d\sqrt{\sigma}$. Combining, we have

$\displaystyle\mbox{{Regret}}_{T}\leq\frac{d\sqrt{\rho_{1:T}}+\operatorname{tr}G_{T}^{1/2}}{2\eta}D^{2}+\eta\operatorname{tr}\left(G_{T}^{1/2}\right)=D\left(\sqrt{2}\operatorname{tr}G_{T}^{1/2}+d\sqrt{\frac{\rho_{1:T}}{2}}\right)\,\,,$

where the last equality is established by choosing $\eta=\frac{D}{\sqrt{2}}$.
∎

####  B.4.2 Proof of Corollary 4

###### Proof.

Following the proof of Theorem 3, we have

$\displaystyle\mbox{{Regret}}_{T}\leq\frac{D^{2}\operatorname{tr}\tilde{G}_{t}^{1/2}}{2\eta}+\eta\operatorname{tr}G_{T}^{1/2}\,\,,$

Denote the accumulated error term

$\displaystyle E=\sum_{t=1}^{T}\rho_{t}N_{t}\,\,.$

Then, by Lemma 10 and sub-additivity of
$\operatorname{tr}\left(\left(\cdot\right)^{1/2}\right)$ [47],

$\displaystyle\mbox{{Regret}}_{T}\leq\left(\frac{D^{2}}{2\eta}+\eta\right)\operatorname{tr}G_{T}^{1/2}+\frac{D^{2}}{2\eta}\operatorname{tr}E^{1/2},$

where it remains to bound the last term. Let $Q$ be a matrix with column
vectors $q_{i}$ that forms an eigenbasis of $E^{1/2}$; this diagonalizes $E$
as well. Notice that

$\lambda_{i}\left({E}\right)=\sum_{t=1}^{T}\rho_{t}q_{i}^{\top}N_{t}q_{i}\,\,,$

and since

$\lambda_{i}^{1/2}\left({E}\right)=\lambda_{i}\left({E^{1/2}}\right)\,\,,$

that we can characterize

$\operatorname{tr}E^{1/2}=\sum_{i=1}^{d}\lambda_{i}^{1/2}(E)=\sum_{i=1}^{d}\left({\sum_{t=1}^{T}\rho_{t}q_{i}^{\top}N_{t}q_{i}}\right)^{1/2}\,\,.$

Denote $u_{t,i}=q_{i}^{\top}N_{t}q_{i}$, since $N_{t}$ is a rank-$(d-\ell)$
projection, $\left\lVert u_{t}\right\rVert_{1}=d-\ell$. Then
$\operatorname{tr}E^{1/2}$ is upper bounded by the value of the program

$\displaystyle\max_{u_{t,i}}$
$\displaystyle\quad\sum_{i=1}^{d}\left({\sum_{t=1}^{T}\rho_{t}u_{t,i}}\right)^{1/2}$
$\displaystyle\mathrm{s.t.}$
$\displaystyle\quad\sum_{i=1}^{d}u_{t,i}=d-\ell\quad\quad\forall t\in[T]\,\,.$

Note that

$\displaystyle\sum_{i=1}^{d}\left(\sum_{t=1}^{T}\rho_{t}u_{t,i}\right)^{1/2}\leq_{1}\sqrt{d}\sqrt{\sum_{i=1}^{d}\sum_{t=1}^{T}\rho_{t}u_{t,i}}=\sqrt{d}\sqrt{\sum_{t=1}^{T}\rho_{t}\sum_{i=1}^{d}u_{t,i}}=_{2}\sqrt{d\rho_{1:T}(d-\ell)}\,\,,$

where $\leq_{1}$ follows from Cauchy-Schwarz, and $=_{2}$ follows from the
constraint on $\sum_{i=1}^{d}u_{t,i}$.

Combining, we have

$\displaystyle\mbox{{Regret}}_{T}\leq\frac{\sqrt{d(d-\ell)\rho_{1:T}}+\operatorname{tr}G_{T}^{1/2}}{2\eta}D^{2}+\eta\operatorname{tr}G_{T}^{1/2}=D\left(\sqrt{2}\operatorname{tr}G_{T}^{1/2}+\sqrt{\frac{d(d-\ell)\rho_{1:T}}{2}}\right)\,\,,$

where the last equality follows from the choice of step size
$\eta=\frac{D}{\sqrt{2}}$.

∎

###  B.5 Proof details for Section 4.2, S-Shampoo

####  B.5.1 Proof of Theorem 5

###### Proof.

First, we establish the following observation and lemma analogous to
Observation 9 and Lemma 10:

######  Observation 13 (Analogous to Observation 9).

Let $V_{t}\Sigma_{t}^{L}V_{t}^{\top}=\bar{L}_{t-1}+G_{t}G_{t}^{\top}$ be the
eigendecomposition of the un-deflated sketch, where
$V_{t}\in\mathbb{R}^{m\times m}$. Suppose
$\operatorname{rank}(\Sigma_{t}^{L})=k$, where $k\in[\ell-1,\ell-1+r]$. Write
$V_{t}=[V_{t}^{\parallel}\ V_{t}^{\perp}]$, where $V_{t}^{\parallel}$ contain
the first $k$ columns of $V_{t}$. Then by definition

$\bar{L}_{t}+\rho_{t}^{L}I\succeq\bar{L}_{t-1}+G_{t}G_{t}^{\top}+\rho_{t}^{L}V_{t}^{\perp}\left({V_{t}^{\perp}}\right)^{\top}.$

Analogously for the right conditioner, let
$W_{t}\Sigma_{t}^{R}W_{t}^{\top}=\bar{R}_{t-1}+G_{t}^{\top}G_{t}$, and write
$W_{t}=[W_{t}^{\parallel}\ W_{t}^{\perp}]$, then

$\bar{R}_{t}+\rho_{t}^{R}I\succeq\bar{R}_{t-1}+G_{t}^{\top}G_{t}+\rho_{t}^{R}W_{t}^{\perp}\left({W_{t}^{\perp}}\right)^{\top}.$

######  Lemma 14.

(Analogous to Lemma 10) Define
$N^{L}_{t}=V_{t}^{\perp}\left({V_{t}^{\perp}}\right)^{\top},N^{R}_{t}=W_{t}^{\perp}\left({W_{t}^{\perp}}\right)^{\top}$,
then

$\displaystyle\tilde{L}_{t}\succeq\sum_{s=1}^{t}G_{s}G_{s}^{\top}+\sum_{s=1}^{t}\rho_{s}^{L}N_{s}^{L}+\varepsilon
I_{m}\,\,,\ \ \
\tilde{R}_{t}\succeq\sum_{s=1}^{t}G_{s}^{\top}G_{s}+\sum_{s=1}^{t}\rho_{s}^{R}N_{s}^{R}+\varepsilon
I_{n}\,\,.$

We follow the shampoo proof in [5]. Let
$x_{t}=\overline{\text{vec}}\left({X_{t}}\right)$,
$g_{t}=\overline{\text{vec}}\left({G_{t}}\right)$, where
$\overline{\text{vec}}\left({\cdot}\right)$ denote the row-major vectorization
of a given matrix.

Kronecker product $\otimes$ obeys the following properties as shown in [5]:

######  Lemma 15 (Lemma 3,4 in Gupta et al. [5]).

For matrices $A,A^{\prime},B,B^{\prime}$ of appropriate dimensions and vectors
$u,v$, $L\in\mathbb{R}^{m\times m},R\in\mathbb{R}^{n\times
n},G\in\mathbb{R}^{m\times n}$, the following properties hold:

  1. 1.

$(A\otimes B)(A^{\prime}\otimes
B^{\prime})=(AA^{\prime})\otimes(BB^{\prime})$.

  2. 2.

$(A\otimes B)^{\top}=A^{\top}\otimes B^{\top}$.

  3. 3.

$A,B\succeq 0$, $(A\otimes B)^{-1}=A^{-1}\otimes B^{-1}$.

  4. 4.

$A\succeq A^{\prime},B\succeq B^{\prime}$, then $A\otimes B\succeq
A^{\prime}\otimes B^{\prime}$.

  5. 5.

$\mathrm{tr}(A\otimes B)=\mathrm{tr}(A)+\mathrm{tr}(B)$.

  6. 6.

$\overline{\text{vec}}\left({uv^{\top}}\right)=u\otimes v$.

  7. 7.

$(L\otimes
R^{\top})\overline{\text{vec}}\left({G}\right)=\overline{\text{vec}}\left({LGR}\right)$.

Then the shampoo update is

$x_{t+1}=x_{t}-\eta(\tilde{L}_{t}^{1/4}\otimes\tilde{R}_{t}^{1/4})^{-1}g_{t}.$

Let
$\tilde{H}_{t}\stackrel{{\scriptstyle\text{def}}}{{=}}\tilde{L}_{t}^{1/4}\otimes\tilde{R}_{t}^{1/4}$,
then by Lemma 15, $\tilde{H}_{t}$ is monotone increasing with $t$, since
$\tilde{L}_{t}$ and $\tilde{R}_{t}$ are monotone by Observation 13. Thus, by
standard analysis [15] for Online Mirror Descent (OMD), we can break down the
regret into the diameter bound and the gradient bound:

$\displaystyle\mbox{{Regret}}_{T}$ $\displaystyle\leq R_{D}+R_{G},\ \ \
\text{where}$

f

$\displaystyle
R_{D}=\frac{1}{2\eta}\sum_{t=1}^{T}\left(\|x_{t}-x^{*}\|_{\tilde{H}_{t}}^{2}-\|x_{t+1}-x^{*}\|_{\tilde{H}_{t}}^{2}\right),\
\ \
R_{G}=\frac{\eta}{2}\sum_{t=1}^{T}\left(\|g_{t}\|_{\tilde{H}_{t}}^{*}\right)^{2}.$

We proceed to bound $R_{D}$ and $R_{G}$ separately. For $R_{D}$,

$\displaystyle R_{D}$
$\displaystyle~{}\leq~{}\frac{1}{2\eta}\sum_{t=1}^{T}\|x_{t}-x^{*}\|_{\tilde{H}_{t}-\tilde{H}_{t-1}}^{2}+\|x_{1}-x^{*}\|_{\tilde{H}_{0}}^{2}$
$\displaystyle\leq\frac{1}{2\eta}\sum_{t=1}^{T}\|\tilde{H}_{t}-\tilde{H}_{t-1}\|_{op}\|x_{t}-x^{*}\|_{2}^{2}+\|x_{1}-x^{*}\|_{\tilde{H}_{0}}^{2}$
$\displaystyle\leq_{1}\frac{D^{2}}{2\eta}\sum_{t=1}^{T}\operatorname{tr}(\tilde{H}_{t}-\tilde{H}_{t-1})+\|x_{1}-x^{*}\|_{\tilde{H}_{0}}^{2}$
$\displaystyle\leq\frac{D^{2}}{2\eta}\operatorname{tr}(\tilde{H}_{T})\,\,,$

where $\leq_{1}$ holds since $\tilde{H}_{t}$’s are increasing in $t$, and we
have for positive semidefinite matrices
$\mathrm{tr}(\cdot)\geq\|\cdot\|_{op}$.

Now we try to bound $R_{G}$. First, we have that

######  Lemma 16 (Lemma 8 in Gupta et al. [5]).

If $G\in\mathbb{R}^{m\times n}$ with rank at most $r$, and
$g=\overline{\text{vec}}\left({G}\right)$, then $\forall\varepsilon\geq 0$,
$\forall t$,

$\displaystyle\varepsilon
I_{mn}+\frac{1}{r}\sum_{s=1}^{t}g_{s}g_{s}^{\top}\preceq\left(\varepsilon
I_{m}+\sum_{s=1}^{t}G_{s}G_{s}^{\top}\right)^{1/2}\otimes\left(\varepsilon
I_{n}+\sum_{s=1}^{t}G_{s}^{\top}G_{s}\right)^{1/2}.$

Define $M_{t}^{L}\in\mathbb{R}^{m\times m},M_{t}^{R}\in\mathbb{R}^{n\times n}$
by

$\displaystyle
M^{L}_{t}\stackrel{{\scriptstyle\text{def}}}{{=}}\sum_{s=1}^{t}G_{s}G_{s}^{\top}+\sum_{s=1}^{t}\rho_{s}^{L}N_{s}^{L}+\varepsilon
I_{m}\,\,,\ \ \
M^{R}_{t}\stackrel{{\scriptstyle\text{def}}}{{=}}\sum_{s=1}^{t}G_{s}^{\top}G_{s}+\sum_{s=1}^{t}\rho_{s}^{R}N_{s}^{R}+\varepsilon
I_{n}\,\,,$

then by Lemma 14,

$\tilde{L}_{t}\succeq M_{t}^{L}\,\,,\ \ \tilde{R}_{t}\succeq M_{t}^{R}\,\,.$

Observe that in addition,

$M_{t}^{L}\succeq\varepsilon I_{m}+\sum_{s=1}^{t}G_{s}G_{s}^{\top}\,\,,\ \
M_{t}^{R}\succeq\varepsilon I_{n}+\sum_{s=1}^{t}G_{s}^{\top}G_{s}\,\,.$

Again by Lemma 15,

$\displaystyle I_{m}\otimes\left(\varepsilon
I_{n}+\sum_{s=1}^{t}G_{s}^{\top}G_{s}\right)\preceq I_{m}\otimes
M_{t}^{R}\,\,,\ \ \ \left({\varepsilon
I_{m}+\sum_{s=1}^{t}G_{s}G_{s}^{\top}}\right)\otimes I_{n}\preceq
M_{t}^{L}\otimes I_{n}\,\,.$

Combining, we have

$\displaystyle\varepsilon
I_{mn}+\frac{1}{r}\sum_{s=1}^{t}g_{s}g_{s}^{\top}\preceq\left({M_{t}^{L}}\right)^{1/2}\otimes\left({M_{t}^{R}}\right)^{1/2}\preceq\tilde{L}_{t}^{1/2}\otimes\tilde{R}_{t}^{1/2}\,\,.$

Define $\hat{H}_{t}\succ 0$ $\forall t\in[T]$ by

$\hat{H}_{t}\stackrel{{\scriptstyle\text{def}}}{{=}}\left(r\varepsilon
I_{mn}+\sum_{s=1}^{t}g_{s}g_{s}^{\top}\right)^{1/2}\preceq\sqrt{r}\tilde{H}_{t}\,\,.$

The bound on $R_{G}$ depends on the following lemma:

######  Lemma 17 (Lemma 2 in Gupta et al. [5]).

Consider a sequence of vectors $\\{g_{t}\\}_{t=1}^{T}$. Given a function
$\Phi(\cdot)$ over positive semidefinite matrices,

$\displaystyle\sum_{t=1}^{T}\left(\|g_{t}\|_{H_{t}}^{*}\right)^{2}\leq\sum_{t=1}^{T}\left(\|g_{t}\|_{H_{T}}^{*}\right)^{2}+\Phi(H_{T})-\Phi(H_{0})\,\,,$

where

$\displaystyle H_{t}=\operatorname*{argmin}_{H\succ
0}\left\\{\left(\sum_{s=1}^{t}g_{s}g_{s}^{\top}\right)\cdot
H^{-1}+\Phi(H)\right\\}\,\,.$

Let
$\Phi(H)\stackrel{{\scriptstyle\text{def}}}{{=}}\mathrm{tr}(H)+r\varepsilon\mathrm{tr}(H^{-1})$
and since

$\displaystyle\operatorname*{argmin}_{H\succ
0}\left\\{\left(\sum_{s=1}^{t}g_{s}g_{s}^{\top}\right)\cdot
H^{-1}+\Phi(H)\right\\}=\operatorname*{argmin}_{H\succ
0}\left\\{\mathrm{tr}{\left(\hat{H}_{t}^{2}H^{-1}+H\right)}\right\\}=\hat{H}_{t}\,\,,$

the above lemma gives

$\displaystyle\sum_{t=1}^{T}\left(\|g_{t}\|_{\hat{H}_{t}}^{*}\right)^{2}\leq\sum_{t=1}^{T}\left(\|g_{t}\|_{\hat{H}_{T}}^{*}\right)^{2}+\Phi(\hat{H}_{T})-\Phi(\hat{H}_{0})\leq
2\operatorname{tr}(\hat{H}_{T})\,\,,$

which by inequality of $\hat{H}_{t}$ and $\tilde{H}_{t}$ established above,
gives

$\displaystyle
R_{G}\stackrel{{\scriptstyle\text{def}}}{{=}}\frac{\eta}{2}\sum_{t=1}^{T}\left(\|g_{t}\|_{\tilde{H}_{t}}^{*}\right)^{2}\leq\frac{\eta\sqrt{r}}{2}\sum_{t=1}^{T}\left(\|g_{t}\|_{\hat{H}_{t}}^{*}\right)^{2}\leq\eta\sqrt{r}\operatorname{tr}(\hat{H}_{T})\leq\eta
r\operatorname{tr}(\tilde{H}_{T})\,\,.$

Combining the bound on $R_{D}$ and $R_{G}$, the overall regret is

$\displaystyle\mbox{{Regret}}_{T}\leq
R_{D}+R_{G}\leq\left({\frac{D^{2}}{2\eta}+\eta
r}\right)\operatorname{tr}(\tilde{H}_{T})=\sqrt{2r}D\operatorname{tr}(\tilde{H}_{T})=\sqrt{2r}D\operatorname{tr}(\tilde{L}_{T}^{1/4})\operatorname{tr}(\tilde{R}_{T}^{1/4})\,\,.$

by the choice of $\eta=\frac{D}{\sqrt{2r}}$ and trace multplicative equality
in Lemma 15. Finally, we have

$\displaystyle\operatorname{tr}\left({\tilde{L}_{T}^{1/4}}\right)$
$\displaystyle\leq_{1}\operatorname{tr}\left(\bar{L}_{T}^{1/4}\right)+\operatorname{tr}\left(\left(\rho_{1:T}^{L}I_{m}\right)^{1/4}\right)$
$\displaystyle\leq_{2}\operatorname{tr}\left({\left({\sum_{t=1}^{T}G_{t}G_{t}^{\top}+\varepsilon
I}\right)^{1/4}}\right)+m\left({\rho_{1:T}^{L}}\right)^{1/4}$
$\displaystyle=\operatorname{tr}\left({L_{T}^{1/4}}\right)+m\left({\rho_{1:T}^{L}}\right)^{1/4},$

where $\leq_{1}$ follows from definition of $\tilde{L}_{T}$ in Algorithm 3 and
subadditivity of $\operatorname{tr}\left(\left(\cdot\right)^{1/4}\right)$
[47], $\leq_{2}$ follows from Remark 11. Similarly,

$\displaystyle\operatorname{tr}\left({\tilde{R}_{T}^{1/4}}\right)\leq\operatorname{tr}\left({R_{T}^{1/4}}\right)+n\left({\rho_{1:T}^{R}}\right)^{1/4}.$

∎

####  B.5.2 Proof of Lemma 14

###### Proof.

We will show the first inequality as the second inequality holds analogously.
For $t=0$, $\tilde{L}_{0}=\varepsilon I_{m}$ by definition of algorithm.
Suppose the first inequality holds for $t$. Consider $t+1$:

$\displaystyle\tilde{L}_{t+1}$
$\displaystyle=\bar{L}_{t+1}+\rho_{1:t+1}^{L}I_{m}$
$\displaystyle\succeq_{1}\bar{L}_{t}+G_{t+1}G_{t+1}^{\top}+\rho_{t+1}^{L}N_{t+1}+\rho_{1:t}^{L}I_{m}$
$\displaystyle=\tilde{L}_{t}+G_{t+1}G_{t+1}^{\top}+\rho_{t+1}^{L}N_{t+1}$
$\displaystyle\succeq_{2}\sum_{s=1}^{t+1}G_{s}G_{s}^{\top}+\sum_{s=1}^{t+1}\rho_{s}^{L}N_{s}^{L}+\varepsilon
I_{m}\,\,,$

where $\succeq_{1}$ follows from Observation 13 and $\succeq_{2}$ follows from
induction hypothesis. ∎

##  Appendix C Training Settings

For repeatable, standard evaluation on modern, competitive tasks we use
init2winit [48] for Jax [49] implementations of architectures in Flax [50] and
standard dataset preprocessing built on top of TFDS [51]. We rely on standard
scientific packages for conducting our work [52, 53, 54, 55]. Our source code
will be released after publication.

Neural net architecture settings are taken from the default settings of the
init2winit library at hash e337ffe [48], which reference the MLCommons
specifications provided at MLCommons® open engineering consortium [56],
including the MLPerf ResNet-50 variant [57], Conformer, and GNN. The
Distributed Shampoo implementation was run at hash 83e6e62 in the repository
referenced by Anil et al. [9].

Throughout, weight decay is applied using its decopled variant [58].

We requested a Shampoo tuning script from Dayma and Anil [59], Anil et al.
[9], which fixed several parameters for Shampoo outside the usual defaults. We
tuned on a cluster of TPUv4s, with 16 TPUv4s per trial in data-parallel mode.

  * •

Block size was already set to 1024. As mentioned in Sec. 5.2, we kept this
change for consistency in covariance factor size across architectures.

  * •

Preconditioning was set to start 101 steps into training
(start_preconditioning_step).

  * •

Preconditioners were updated every 10 steps instead of every step for speed
(preconditioning_compute_steps is 10).

  * •

The grafting type, which controls the per-tensor learning rate schedule, was
set to RMSPROP_NORMALIZED, which applies RMSProp [60] over unit-normalized
gradients.

  * •

moving_average_for_momentum was activated (so the final updates are computed
as $\beta_{1}\mu_{t}+(1-\beta_{1})g_{t}$, where $\mu_{t}$ is the momentum term
and $g_{t}$ is the preconditioned update.

  * •

The virtual batch size, used to compute batch norm statistics, was set to $32$
(the full per-step minibatch size was $1024$, but this enables data-parallel
training).

Also from the provided script, we used a linear warmup rampup starting from 0
to the nominal learning rate hyperparameter setting, followed by a cosine
decay schedule, with the transition happening 5% of the way into training (the
learning rate monotonically increases, then montonically decreases, as the
cosine schedule has a quarter-period set to the number of training steps).

Then, we performed tuning using random hyperparameter search over the space
defined in Tbl. 4. We ran the batch sizes and number of steps provided in the
scripts, which were 256, 512, 1024 for Conformer, GNN, and ResNet-50,
respectively, for about 162, 117, 199 epochs, respectively.

Shampoo is automatically configured with grafting parameters, which we search
over [61].

Table 4: The search space for hyperparameters for tuning Shampoo on our NN architectures for the Kronecker-factored covariance optimization. Note that the search space explores one less momentum, not momentum directly. Label smoothing was only applied to ImageNet. We sample uniformly either from linear or logscale among the ranges specified with 100 trials, and select the best one according to validation accuracy.  Hyperparameter | Range | Log scale?
---|---|---
Learning rate $\eta$ | $[10^{-4},10^{-2}]$ | ✓
Weight decay $\gamma$ | $[10^{-2},1]$ | ✓
Momentum $1-\beta_{1}$ | $[10^{-2},10^{-1}]$ | ✓
Label smoothing | $[0,0.2]$ |

##  Appendix D ResNet-50 Settings

For training ImageNet, we mostly inherited the settings of Appendix C for
Shampoo tuning, but made some minor modifications, namely adding of the second
moment decay ($\beta_{2}$), widening the search space, and, for computational
reasons, performing a shortened run of only 66 epochs for tuning trials. The
architecture details remain the same. The learning rate schedule was stretched
to this interval, so warmup was still 5% of the duration, and cosine decay
ended learning rate at 0 by the end of the 66 epochs of training. The full
search space is elaborated on in Tbl. 5.

Table 5: The search space for hyperparameters for tuning Shampoo on our architectures for ImageNet hyperparameters. The same space was applied to S-Shampoo with a fixed sketch rank $\ell=256$ for all tensors. Note that we search $1-\beta_{1},1-\beta_{2}$, and not the original hyperparameter. We sample uniformly either from linear or logscale among the ranges specified with 256 trials, and select the best one according to validation accuracy. $(*)$ stands for a discrete uniform choice over four different grafting rates, based on AdaGrad, RMSProp, and normalized versions of the two. The gradient clipping norm is similarly discrete.  Hyperparameter | Range | Log?
---|---|---
Learning rate $\eta$ | $[10^{-4},10^{-2}]$ | ✓
Weight decay $\gamma$ | $[10^{-3},0.1]$ | ✓
Momentum $1-\beta_{1}$ | $[10^{-4},10^{-1}]$ | ✓
2nd moment $1-\beta_{2}$ | $[10^{-4},10^{-1}]$ | ✓
Label smoothing | $[0,0.2]$ |
Dropout Rate | $[0,0.2]$ |
Grafting Type | $(*)$ |
Gradient Clip $L_{2}$ | $\left\\{{1,10,10^{2},10^{3}}\right\\}$ |

To tune Adam, a first order method, we considered mostly the same nominaly
hyperparameters (where $\beta_{2}$ refers to second moment momentum now),
except grafting, which instead we replaced with a search over the warmup
duration, summarized in Tbl. 6.

Table 6: The search space for hyperparameters for tuning Adam on our architectures for ImageNet hyperparameters. The same caveats as in Tbl. 5 apply. Also tuned with 256 trials.  Hyperparameter | Range | Log?
---|---|---
Learning rate $\eta$ | $[10^{-4},10^{-2}]$ | ✓
Weight decay $\gamma$ | $[10^{-3},0.1]$ | ✓
Momentum $1-\beta_{1}$ | $[10^{-4},10^{-1}]$ | ✓
2nd moment $1-\beta_{2}$ | $[10^{-4},10^{-1}]$ | ✓
Label smoothing | $[0,0.2]$ |
Dropout Rate | $[0,0.2]$ |
Warmup Duration |  $[2\%,10\%]$ of training |
Gradient Clip $L_{2}$ | $\left\\{{1,10,10^{2},10^{3}}\right\\}$ |

Full evaluation of the selected best hyperparameters for each setting was
performed with the classical 90-epoch setting, with the learning rate schedule
correspondingly stretched.

We provide the full training curves in Fig. 5.

Figure 5: Full test set curves for imagenet.

##  Appendix E Conformer Settings

The Conformer architecture was used from the MLCommons specification as
described in Appendix C, with the following fixed additional settings: 1024
batch size, 100 epochs of training, 5% of training used for linear warmup with
cosine decay of learning rate. We fixed gradient clipping at a value of $10$,
without which we noticed Adam curves were very volatile. We set the eigh
parameter for Shampoo to true (we found that it did not make a difference in a
few sample runs’ loss curves, as an alternative to the iterative $p$-th
inverse root routine in Shampoo, but used it instead since we believe it has
better numerical stability). The hyperparameters we searched over for all
optimizers are described in Tbl. 7.

Table 7: The search space for hyperparameters for tuning Shampoo, Adam, and S-Shampoo for Conformer with fixed $256$ rank. Here we fixed the grafting type to be RMSProp. During initial runs of the baselines, we noticed that Adam preferred larger learning rates, so we changed and reran its space for $\eta$ to be $10\times$ that of Shampoo, namely $[10^{-4},10^{-2}$, still searching over logspace. We also stopped any hyperparameter trial which did not go below 0.875 WER after 5000 training steps.  Hyperparameter | Range | Log?
---|---|---
Learning rate $\eta$ | $[10^{-5},10^{-3}]$ | ✓
Momentum $1-\beta_{1}$ | $[10^{-3},10^{-1}]$ | ✓
2nd moment $1-\beta_{2}$ | $[10^{-3},10^{-1}]$ | ✓
Weight decay $\gamma$ | $[10^{-4},10^{-2}]$ | ✓
Dropout Rate | $[0,0.2]$ | ✓

##  Appendix F GNN Settings

The GNN architecture was used from the MLCommons specification as described in
Appendix C, with the following fixed additional settings: 1024 batch size, 30
epochs of training, 5% of training used for linear warmup with cosine decay of
learning rate, 0.05 dropout, and we set the eigh parameter for Shampoo to true
as in Appendix E.

Then we searched a hyperparameter space for Shampoo, Adam, and S-Shampoo as
described in Tbl. 8.

Table 8: The search space for hyperparameters for tuning Shampoo and S-Shampoo for GNN with fixed regularization settings (due to resource constraints, we only ran 128 samples from the grid here). Here we fixed the grafting type to be RMSProp and did not use gradient clipping, unlike Tbl. 5, based on a few trial runs of the Shampoo baseline from which we determined we could reduce the hyperparameter space.  Hyperparameter | Range | Log?
---|---|---
Learning rate $\eta$ | $[10^{-4},10^{-2}]$ | ✓
Momentum $1-\beta_{1}$ | $[10^{-3},0.5]$ | ✓
2nd moment $1-\beta_{2}$ | $[10^{-3},0.5]$ | ✓
Weight decay $\gamma$ | $[10^{-3},0.5]$ | ✓

##  Appendix G Step-skipping

In this section, we provide some theoretical justification for step-skipping.
We first derive the regret bound of AdaGrad with step skipping, named Generic
Epoch AdaGrad. The additional regret incurred is expressed as an error term.
Then, we describe a setting where the error term admits a simple bound,
showing that step-skipping incurs at most an extra $\log T$ time dependence on
the regret.

###  G.1 Adversarial losses

Consider a generalized epoching AdaGrad with $K$ fixed update points $t_{k}$,
such that $t_{1}=0$ and $t_{K}=T$.

Algorithm 5 Generic Epoch AdaGrad

1: Input: $\eta,T,\left\\{{t_{k}}\right\\}_{k=1}^{K},G_{0}\succ 0$, convex
closed set $\mathcal{K}$.

2: Initialize: $x_{1}$.

3: for $k=1,\dots,K-1$ do

4: for $t=t_{k}+1,\cdots,t_{k+1}$ do

5: Play $x_{t}$, receive $f_{t}$ loss with gradient $g_{t}$.

6: Update $G_{t}=G_{t-1}+g_{t}g_{t}^{\top}$.

7: Update $x_{t+1}=\Pi_{\mathcal{K}}[x_{t}-\eta G_{t_{k}}^{-1/2}g_{t}]$.

8: end for

9: end for

######  Theorem 18.

Generic Epoch AdaGrad (Alg. 5) with fixed update points
$\left\\{{t_{k}}\right\\}_{k=1}^{K}$ satisfies

$\displaystyle
R_{T}\leq\frac{D^{2}}{\eta}\operatorname{tr}G_{T}^{1/2}+\frac{\eta}{2}\left({2\operatorname{tr}G_{T}^{1/2}-2\operatorname{tr}G_{0}^{1/2}+\sum_{k=1}^{K}\epsilon_{k}}\right)\,\,,$

where the error terms $\epsilon_{k}$ are given by

$\displaystyle\epsilon_{k}$
$\displaystyle=\operatorname{tr}\left({G_{t_{k}}^{-1/2}S_{k}G_{t_{k}}^{-1/2}A_{k}}\right)\,\,,$
$\displaystyle S_{k}$ $\displaystyle=\int_{0}^{\infty}\exp\left({-\tau
G_{t_{k}}^{1/2}}\right)A_{k}\exp\left({-\tau
G_{t_{k}}^{1/2}}\right)\mathop{\mathrm{d}\tau}\,\,,$ $\displaystyle A_{k}$
$\displaystyle=G_{t_{k+1}}-G_{t_{k}}\,\,.$

###### Proof of Theorem 18.

First we start with the usual decomposition.

######  Lemma 19.

Consider arbitrary adversarial convex losses $f_{t}$. Without projection, the
regret $R_{T}$ relative to a comparator $x_{*}$ with $D=\max_{t}\left\lVert
x_{t}-x_{*}\right\rVert_{2}$, for generic epoch AdaGrad with fixed update
points $t_{k}$ is given by

$R_{T}\leq\frac{D^{2}}{\eta}\operatorname{tr}G_{T}^{1/2}+\frac{\eta}{2}\sum_{k=1}^{K}\sum_{t=t_{k}+1}^{t_{k+1}}g_{t}^{\top}G_{t_{k}}^{-1/2}g_{t}\,\,.$

###### Proof of Lemma 19.

This follows from the usual AdaGrad analysis since our preconditioners are
monotone $G_{t_{k}}\succeq G_{t_{k+1}}$. ∎

So we must turn our attention to the gradient bound. We start by noting the
following lemmas established in matrix analysis.

######  Lemma 20 (Corollary 4.1 in [62]).

The map $f(X)=X^{-1/2}$ is matrix convex over the positive definite domain;
i.e., for any two matrices $A,B\succ 0$ and any $\theta\in[0,1]$, we have

$\theta f(A)+(1-\theta)f(B)\succeq f(\theta A+(1-\theta)B)\,\,.$

######  Lemma 21 (Theorem V.3.3, Exercise V.3.15 in [63]).

Suppose a matrix convex function $F(X)$ is induced by applying $f$ pointwise
to its spectrum $F(X)=U\operatorname{diag}[f(\Lambda_{ii})]U^{\dagger}$ with
$f\in\mathcal{C}^{1}(I)$ for some $I\subset\mathbb{R}_{+}$. Then

$F(X)+\partial F(X)(\Delta)\preceq F(X+\Delta)\,\,,$

if and only if $F(X)$ is matrix convex, and the linear transformation
$\partial F(X)$ is the derivative of $F$ at $X$.

Matrix derivative computation [64, 65] shows that if $F(X)=X^{-1/2}$ then

$\partial F(X)(\Delta)=-X^{-1/2}\left[(X^{1/2}\oplus
X^{1/2})^{-1}\Delta\right]X^{-1/2}\,\,,$

where $(X^{1/2}\oplus X^{1/2})^{-1}\Delta$ is the solution $S$ to the
continuous Lyapunov equation $\sqrt{X}S+S\sqrt{X}=\Delta$ as $X\oplus
X=I\otimes X+X\otimes I$. For $X\succ 0$, it is known from generic results
about Sylvester’s equation that the solution $S$ is unique. Since $-X$ is
asymptotically stable in the Lyapunov sense,

$S(X,\Delta)=\int_{0}^{\infty}\exp(-\tau\sqrt{X})\Delta\exp(-\tau\sqrt{X})\mathop{\mathrm{d}\tau}\,\,.$

With these results from matrix analysis and linear systems, we are ready to
bound the gradient term in Lemma 19. Consider a single term from the gradient
bound in Lemma 19,
$\sum_{t=t_{k}+1}^{t_{k+1}}g_{t}^{\top}G_{t_{k}}^{-1/2}g_{t}$ for fixed $k$.

With $X=G_{t_{k}}$, $\Delta=A_{k}=G_{t_{k+1}}-G_{t_{k}}$, and $f(X)=X^{-1/2}$
consider applying Lemma 21. $F(X)\preceq F(X+A_{k})-\partial F(X)(A_{k})$, so
overall

$\displaystyle\sum_{t=t_{k}+1}^{t_{k+1}}g_{t}^{\top}G_{t_{k}}^{-1/2}g_{t}$
$\displaystyle\leq\sum_{t=t_{k}+1}^{t_{k+1}}g_{t}^{\top}\left[G_{t_{k+1}}^{-1/2}-\partial
F(G_{t_{k}})(A_{k})\right]g_{t}$
$\displaystyle=\sum_{t=t_{k}+1}^{t_{k+1}}g_{t}^{\top}G_{t_{k+1}}^{-1/2}g_{t}-\operatorname{tr}\left({\partial
F(G_{t_{k}})(A_{k})\sum_{t=t_{k}+1}^{t_{k+1}}g_{t}g_{t}^{\top}}\right)$
$\displaystyle=\sum_{t=t_{k}+1}^{t_{k+1}}g_{t}^{\top}G_{t_{k+1}}^{-1/2}g_{t}-\operatorname{tr}\left({\partial
F(G_{t_{k}})(A_{k})A_{k}}\right)$
$\displaystyle=\sum_{t=t_{k}+1}^{t_{k+1}}g_{t}^{\top}G_{t_{k+1}}^{-1/2}g_{t}+\operatorname{tr}\left(G_{t_{k}}^{-1/2}S_{k}G_{t_{k}}^{-1/2}A_{k}\right)$
$\displaystyle=\sum_{t=t_{k}+1}^{t_{k+1}}g_{t}^{\top}G_{t_{k+1}}^{-1/2}g_{t}+\epsilon_{k}\,\,.$

######  Lemma 22 (FTL-BTL with errors).

Consider arbitrary $\phi_{k}$ for $k\in 0,1,\cdots,K$. Let
$x_{k}\in\operatorname*{argmin}\sum_{j=0}^{k}\phi_{j}$ and suppose
$\phi_{k}(x_{k-1})\leq\phi_{k}(x_{k})+\delta_{k}$. Then $\forall K$,

$\sum_{k=0}^{K}\phi_{k}(x_{k-1})\leq\sum_{k=0}^{K}\phi_{k}(x_{K})+\delta_{k}\,\,,$

with $\delta_{0}=0$ and $x_{-1}=x_{0}$.

Assume Lemma 22 and take $\phi_{k}(X)=\left\langle A_{k},X\right\rangle$ for
$X\succeq 0$ and $\phi_{0}(X)=\left\langle
G_{0},X\right\rangle+\operatorname{tr}X^{-1}$. Note that

$\displaystyle\sum_{j=0}^{k}\phi_{j}(X)$
$\displaystyle=\operatorname{tr}X^{-1}+\sum_{j=0}^{k}\left\langle
A_{j},X\right\rangle\,\,.$

In particular, $G_{t_{k+1}}^{-1/2}=\operatorname*{argmin}_{X\succeq
0}\sum_{j=0}^{k}\phi_{j}(X)$.

Furthermore, with $\delta_{k}=\epsilon_{k}$, the condition
$\phi_{k}(G_{t_{k}})\leq\phi_{k}(G_{t_{k+1}})+\delta_{k}$ is satisfied. Lemma
22 implies

$\displaystyle\sum_{k=1}^{K}\sum_{t=t_{k}+1}^{t_{k+1}}g_{t}^{\top}G_{t_{k}}^{-1/2}g_{t}=\sum_{k=1}^{K-1}\phi_{k}\left(G_{t_{k}}^{-1/2}\right)$
$\displaystyle=-\phi_{0}(G_{0}^{-1/2})+\sum_{k=0}^{K-1}\phi_{k}\left(G_{t_{k}}^{-1/2}\right)$
$\displaystyle\leq-\phi_{0}(G_{0}^{-1/2})+\sum_{k=0}^{K-1}\phi_{k}\left(G_{T}^{-1/2}\right)+\epsilon_{k}\,\,,$

where $G_{t_{0}}\stackrel{{\scriptstyle\text{def}}}{{=}}G_{0}$. Lastly, since

$\displaystyle\sum_{k=0}^{K-1}\phi_{k}\left(G_{T}^{-1/2}\right)=\operatorname{tr}G_{T}^{1/2}+\operatorname{tr}\left({G_{T}^{-1/2}\sum_{k=0}^{K-1}A_{k}}\right)=2\operatorname{tr}G_{T}^{1/2},$

we conclude with the desired bound for $R_{T}$. ∎

###### Proof of Lemma 22.

By induction on $K$. For $K=0$, $\phi_{0}(x_{-1})=\phi_{0}(x_{0})$, holding by
definition. Suppose that the hypothesis now holds for $K$; it then holds for
$K+1$.

$\displaystyle\sum_{k=0}^{K+1}\phi_{k}(x_{K+1})$
$\displaystyle=\sum_{k=0}^{K}\phi_{k}(x_{K+1})+\phi_{K+1}(x_{K+1})$
$\displaystyle\geq\phi_{K+1}(x_{K+1})+\sum_{k=0}^{K}\phi_{k}(x_{K})$
$\displaystyle\geq\phi_{K+1}(x_{K})-\epsilon_{K+1}+\sum_{k=0}^{K}\phi_{k}(x_{K})$
$\displaystyle\geq\phi_{K+1}(x_{K})-\epsilon_{K+1}+\sum_{k=0}^{K}\phi_{k}(x_{k-1})-\epsilon_{k}$
$\displaystyle\geq\sum_{k=0}^{K+1}\phi_{k}(x_{k-1})-\epsilon_{k}\,\,.$

∎

###  G.2 Simplifying the error

We want to simplify the term $\epsilon_{k}$, which is given by

$\displaystyle\epsilon_{k}$
$\displaystyle=\operatorname{tr}\left({G_{t_{k}}^{-1/2}S_{k}G_{t_{k}}^{-1/2}A_{k}}\right)\,\,,$
$\displaystyle S_{k}$ $\displaystyle=\int_{0}^{\infty}\exp\left({-\tau
G_{t_{k}}^{1/2}}\right)A_{k}\exp\left({-\tau
G_{t_{k}}^{1/2}}\right)\mathop{\mathrm{d}\tau}\,\,,$ $\displaystyle A_{k}$
$\displaystyle=G_{t_{k+1}}-G_{t_{k}}\,\,.$

Next, notice that $X$ and $\exp\left({-\alpha X^{-1}}\right)$ commute. Then
along with linearity of trace, we can established that

$\displaystyle\epsilon_{k}=\int_{0}^{\infty}\operatorname{tr}\left[{\left({\exp\left({-\tau
G_{t_{k}}^{1/2}}\right)G_{t_{k}}^{-1/2}A_{k}}\right)^{2}}\right]\mathop{\mathrm{d}\tau}\,\,.$

###  G.3 Towards simpler error

$\epsilon_{k}$ can be further simplified and bounded under additional
assumptions. Namely,

######  Assumption 1.

Suppose that w.p. at least $1-\delta/2K$ for some fixed, universal $\beta>0$,
we have the inequality $A_{k}\preceq\beta G_{t_{k}}$ where
$A_{k}\stackrel{{\scriptstyle\text{def}}}{{=}}G_{t_{k+1}}-G_{t_{k}}$.

######  Assumption 2.

Suppose that w.p. at least $1-\delta/2K$, $G_{t_{k}}$’s are
$(\sigma_{\min},\sigma_{\max})$-well-conditioned, i.e.

$\lambda_{d}(G_{t_{k}})\geq\sigma_{\min}t_{k}\ \ \ \text{and}\ \ \
\lambda_{1}(G_{t_{k}})\leq\sigma_{\max}t_{k}.$

######  Remark 23.

As an example, consider the stochastic linear setting where at each iteration
we receive a loss function $\langle g_{t},x\rangle$, and $g_{t}$’s are
independent, though not necessarily identically distributed, and satisfies
that
$2\sigma_{\min}I\preceq\operatorname{\mathbb{E}}[g_{t}g_{t}^{\top}]\preceq\frac{\sigma_{\max}}{2}I$
and $\|g_{t}\|_{2}\leq\sqrt{\frac{\sigma_{\max}}{2}}$ almost surely. Then, for
$T$ sufficiently large and $t_{k+1}-t_{k}=O(\log T)$, by matrix Chernoff
bounds Assumption 1 and 2 are satisfied.

First, $\forall X,Y\succeq 0$, the following inequality hold:

######  Lemma 24.

If $X\preceq Y$ and $A\succeq 0$, then
$\operatorname{tr}\left[{\left({AX}\right)^{2}}\right]\leq\operatorname{tr}\left[{\left({AY}\right)^{2}}\right]$.

With Lemma 24, we can bound $\epsilon_{k}$. With probability at least
$1-\delta/2K$,

$\displaystyle\epsilon_{k}$
$\displaystyle=\int_{0}^{\infty}\operatorname{tr}\left[{\left({\exp\left({-\tau
G_{t_{k}}^{1/2}}\right)G_{t_{k}}^{-1/2}A_{k}}\right)^{2}}\right]\mathop{\mathrm{d}\tau}$
$\displaystyle\leq\beta^{2}\int_{0}^{\infty}\operatorname{tr}\left[{\left({\exp\left({-\tau
G_{t_{k}}^{1/2}}\right)G_{t_{k}}^{-1/2}G_{t_{k}}}\right)^{2}}\right]\mathop{\mathrm{d}\tau}$
$\displaystyle=\beta^{2}\int_{0}^{\infty}\operatorname{tr}\left[{\left({\exp\left({-\tau
G_{t_{k}}^{1/2}}\right)G_{t_{k}}^{1/2}}\right)^{2}}\right]\mathop{\mathrm{d}\tau}$
$\displaystyle=\beta^{2}\int_{0}^{\infty}\operatorname{tr}\left({\exp\left({-2\tau
G_{t_{k}}^{1/2}}\right)G_{t_{k}}}\right)\mathop{\mathrm{d}\tau}\,\,,$

where the last step only holds since $X$ and $\exp\left(-\alpha
X^{-{1/2}}\right)$ commute.

Next, let $\lambda_{i}$ denote the $i$-th largest eigenvalue and
$\lambda_{-i}$ be the $i$-th smallest. Notice since $\exp\left({-\tau
G_{t_{k}}^{1/2}}\right)$, $G_{t_{k}}^{1/2}$, and $G_{t_{k}}$ are
simultaneously diagonalizable, and montonic matrix functions preserve
eigenvalue ordering, we have

$\displaystyle\lambda_{i}\left({\exp\left({-2\tau
G_{t_{k}}^{1/2}}\right)}\right)$
$\displaystyle=\exp\left({-2\tau\lambda_{-i}\left(\left({G_{t_{k}}}\right)^{1/2}\right)}\right)\,\,,$
$\displaystyle\lambda_{i}\left({\exp\left({-2\tau
G_{t_{k}}^{1/2}}\right)G_{t_{k}}}\right)$
$\displaystyle=\lambda_{i}\left({\exp\left({-2\tau
G_{t_{k}}^{1/2}}\right)}\right)\lambda_{i}\left({G_{t_{k}}}\right)\,\,.$

Returning to our $\epsilon_{k}$ bound, rewriting the trace with eigenvalues,
we have w.p. at least $1-\delta/2K$,

$\displaystyle\epsilon_{k}$
$\displaystyle\leq_{1}\beta^{2}\int_{0}^{\infty}\sum_{i}\lambda_{i}\left({\exp\left({-2\tau
G_{t_{k}}^{1/2}}\right)G_{t_{k}}}\right)\mathop{\mathrm{d}\tau}$
$\displaystyle=\beta^{2}\sum_{i}\lambda_{i}(G_{t_{k}})\int_{0}^{\infty}\exp\left({-2\tau\lambda_{-i}(G_{t_{k}})^{1/2}}\right)\mathop{\mathrm{d}\tau}$
$\displaystyle=\beta^{2}\sum_{i}\frac{\lambda_{i}(G_{t_{k}})}{2\lambda_{-i}(G_{t_{k}})^{1/2}}\,\,,$

where $\leq_{1}$ follows from Tonelli’s Theorem. At this point, we apply
Assumption 2 and get that w.p. at least $1-\delta/K$,

$\displaystyle\epsilon_{k}$
$\displaystyle\leq\frac{\beta^{2}}{\sqrt{t_{k}\sigma_{\min}}}\sum_{i}\lambda_{i}(G_{t_{k}})$
$\displaystyle~{}\leq~{}\frac{\beta^{2}}{\sqrt{t_{k}\sigma_{\min}}}\sum_{i}(t_{k}\sigma_{\max})^{1/2}\lambda_{i}(G_{t_{k}})^{1/2}$
$\displaystyle=\beta^{2}\sqrt{\frac{\sigma_{\max}}{\sigma_{\min}}}\sum_{i}\lambda_{i}(G_{t_{k}})^{1/2}$
$\displaystyle=\beta^{2}\sqrt{\frac{\sigma_{\max}}{\sigma_{\min}}}\operatorname{tr}G_{t_{k}}^{1/2}\,\,.$

Across all epochs, we then have w.p. at least $1-\delta$,

$\frac{1}{\beta^{2}}\sqrt{\frac{\sigma_{\min}}{\sigma_{\max}}}\sum_{k=1}^{K}\epsilon_{k}\leq\sum_{k=1}^{K}\operatorname{tr}G_{t_{k}}^{1/2}\leq\log
T\operatorname{tr}G_{T}^{1/2}\,\,.$

Altogether, since $\beta$ is a universal constant, w.p. at least $1-\delta$,

$R_{T}\lesssim\frac{D^{2}}{\eta}\operatorname{tr}G_{T}^{1/2}+\eta\sqrt{\frac{\sigma_{\max}}{\sigma_{\min}}}\log
T\operatorname{tr}G_{T}^{1/2}\,\,.$

We conclude that in this case, the time dependency of Epoch AdaGrad’s regret
is only $\log T$ factor worse than that of the original AdaGrad regret.

####  G.3.1 Proof of Lemma 24

First, for $0\preceq X\preceq Y$, $BXB\preceq BYB$, $\forall B$, since
$(Bx)^{\top}(Y-X)(Bx)\geq 0$, $\forall x$. By cyclic property of trace and
taking $B=A^{1/2}X$,

$\displaystyle\operatorname{tr}\left({AXAX}\right)=\operatorname{tr}\left({X^{1/2}AXAX^{1/2}}\right)\leq\operatorname{tr}\left({X^{1/2}AYAX^{1/2}}\right)\,\,.$

Continuing,

$\displaystyle\operatorname{tr}\left[{\left({AX}\right)^{2}}\right]$
$\displaystyle\leq\operatorname{tr}\left({X^{1/2}AYAX^{1/2}}\right)$
$\displaystyle=\operatorname{tr}\left({Y^{1/2}AXAY^{1/2}}\right)$
$\displaystyle\leq\operatorname{tr}\left({Y^{1/2}AYAY^{1/2}}\right)$
$\displaystyle=\operatorname{tr}\left[{\left({AY}\right)^{2}}\right]\,\,.$