#  MicroAdam: Accurate Adaptive Optimization with Low Space Overhead and
Provable Convergence

Ionut-Vlad Modoranu1 Mher Safaryan1 Grigory Malinovsky2 Eldar Kurtic1 Thomas
Robert1 Peter Richtárik2 Dan Alistarh1
1Institute of Science and Technology Austria (ISTA)
2King Abdullah University of Science and Technology (KAUST) Correspondence to
<EMAIL_ADDRESS>
###### Abstract

We propose a new variant of the Adam optimizer (Kingma and Ba, 2014) called
MicroAdam that specifically minimizes memory overheads, while maintaining
theoretical convergence guarantees. We achieve this by compressing the
gradient information before it is fed into the optimizer state, thereby
reducing its memory footprint significantly. We control the resulting
compression error via a novel instance of the classical _error feedback_
mechanism from distributed optimization (Seide et al., 2014; Alistarh et al.,
2018; Karimireddy et al., 2019) in which _the error correction information is
itself compressed_ to allow for practical memory gains. We prove that the
resulting approach maintains theoretical convergence guarantees competitive to
those of AMSGrad, while providing good practical performance. Specifically, we
show that MicroAdam can be implemented efficiently on GPUs: on both million-
scale (BERT) and billion-scale (LLaMA) models, MicroAdam provides practical
convergence competitive to that of the uncompressed Adam baseline, with lower
memory usage and similar running time. Our code is available at
https://github.com/IST-DASLab/MicroAdam.

##  1 Introduction

The Adam (Kingma and Ba, 2014) adaptive optimizer and its variants (Reddi et
al., 2019; Loshchilov and Hutter, 2019) has emerged as a dominant choice for
training deep neural networks (DNNs), especially in the case of large language
models (LLMs) with billions of parameters. Yet, its versatility comes with the
drawback of substantial memory overheads: relative to naive SGD-based
optimization, the Adam optimizer states doubles the memory overhead, as it
requires storing two additional parameters for each variable. For large-scale
models, these memory demands pose a significant challenge. In turn, this has
spurred research into memory-efficient adaptive optimizers, such as AdaFactor
(Shazeer and Stern, 2018), 8-bit Adam (Dettmers et al., 2021), or the very
recent GaLore (Zhao et al., 2024) low-rank projection approach. Despite their
popularity and practical utility, the above methods lack rigorous convergence
guarantees, and often trade off memory reductions with decreased convergence
in practice. This raises the question of whether it is possible to design
adaptive optimizers that are not only memory-efficient, but also maintain
strong theoretical and practical performance metrics.

#### Contributions.

In this paper, we address this gap by introducing MicroAdam, an adaptive
optimizer which guarantees low memory usage but also ensures provable
convergence. We start from the idea that we can allow the (lossy) sparse
projection of gradient information before it enters the optimizer states;
crucially, different from prior work, we ensure convergence by correcting for
the inherent error due to compression by employing a novel variant of _error
correction_ , a mechanism introduced for distributed optimization (Seide et
al., 2014). However, simply using error feedback would not lead to memory
savings, since the size of the error correction buffer is comparable to the
that of the original optimizer state. Instead, our main algorithmic innovation
is in showing that the error feedback _can itself be compressed_ in the
context of adaptive optimization. This renders the memory overhead of error
feedback negligible, while preserving convergence guarantees.

Specifically, on the theoretical side, we provide a new analysis showing that,
under reasonable assumptions on the loss function being optimized and on the
degree of compression, MicroAdam provably guarantees convergence, at
asymptotically the same rate as AMSGrad (Zhou et al., 2024a), i.e. a version
of Adam with general convergence guarantees. The key finding is that our
approach allows for the overhead of compression to be shifted to the higher-
order terms, where it should not impact practical convergence in common cases.
This claim holds both for general smooth non-convex functions, and for non-
convex functions under the Polyak-Lojasiewicz (PL) assumption, highlighting a
trade-off between the degree of compression of the gradients, and that of the
error feedback.

We complement our algorithmic and analytic results with an efficient GPU
implementation of MicroAdam, which we validate for fine-tuning language models
from the BERT (Devlin et al., 2018), OPT (Zhang et al., 2022) and LLaMA
(Touvron et al., 2023) families, with hundreds of millions to billions of
parameters. We show that, in practice, gradients can be projected to very high
sparsity (99%), while the error correction can be stored at 4 bits per
component, without loss of convergence. Concretely, our method can
significantly improve upon the memory footprint of the extremely popular 8bit
Adam (Dettmers et al., 2021) when fine-tuning models such as LLaMA2-7B
(Touvron et al., 2023), at similar or better accuracy. At the same time,
MicroAdam provides better accuracy relative to high-compression heuristics
such as GaLore (Zhao et al., 2024).

In summary, we provide a new theoretically-grounded approach to memory-
efficient adaptive optimization, which has the advantage of providing both
theoretical guarantees and good practical convergence, while being scalable to
billion-parameter models. MicroAdam could therefore serve as a useful new tool
for accurate and memory-efficient optimization of large models.

##  2 Related Work

Due to space constraints, we mainly focus on related work reeducing the cost
of optimizer states. Dettmers et al. (2021) considers this problem,
specifically by performing fine-grained quantization of the optimizer states.
Their work does not alter the Adam algorithm; instead, it deals with the
challenge of accurately compressing the dynamically-changing meta-data
sequence. As the name suggests, the space savings correspond to roughly
halving the memory required by the optimizer states, relative to FP16. In the
same vein, AdaFactor (Shazeer and Stern, 2018) and CAME (Luo et al., 2023)
reduce memory cost by factorizing the second-order statistics, while the
recent GaLore (Zhao et al., 2024) factorizes the gradients themselves before
they enter the optimizer state (but does not use error correction).
Importantly, these methods are _heuristics_ : they do not provide theoretical
guarantees under standard assumptions,111GaLore (Zhao et al., 2024) does state
convergence guarantees for a variant of the algorithm with fixed projections,
but this is under a strong “stable rank” assumption, which may not hold in
practice. and in practice require careful tuning to preserve convergence (Luo
et al., 2023). By contrast, our method is theoretically justified, and
provides good practical convergence. Earlier work by Anil et al. (2019)
provides convergence guarantees for a compressed variant of Adagrad (Duchi et
al., 2010) called SM3, improving upon earlier work by Spring et al. (2019).
However, it is not clear how to extend their approach to the popular Adam
optimizer, and heuristic methods appear to provide superior performance (Luo
et al., 2023).

Conceptually, our work is related to error feedback mechanisms studied in
distributed optimization, e.g. (Seide et al., 2014; Alistarh et al., 2018;
Karimireddy et al., 2019; Richtárik et al., 2021). Specifically, Li et al.
(2022) proved convergence of AdaGrad-like algorithms in conjunction with error
feedback, in a distributed environment. Our focus is different: minimizing
memory costs in the single-node setting: for this, we show that the error
correction buffer can itself be compressed. We provide an analysis for the
resulting new algorithm, and efficient CUDA implementations.

More broadly, scaling adaptive or second-order optimizers to large models is a
very active area. Works such as GGT (Agarwal et al., 2019), Shampoo (Gupta et
al., 2018) and M-FAC (Frantar et al., 2021) provided quadratic-space
algorithms that are still feasible to execute for moderate-sized DNNs, but
will not scale for billion-parameter models. Follow-up work such as AdaHessian
(Yao et al., 2020), Sophia (Liu et al., 2023), Sketchy (Feinberg et al., 2023)
and EFCP (Modoranu et al., 2023), scaled these approaches via additional
approximations. Of these, the closest work to ours is EFCP, which uses
sparsification plus standard error feedback to compress the gradient window
employed in the Fisher approximation of the Hessian. However, EFCP does not
compress the error accumulator, assumes a different optimization algorithm
(Natural Gradient (Amari, 2016)), lacks convergence guarantees, and does not
scale to billion-parameter models.

##  3 The MicroAdam Algorithm

#### Notation.

We consider a standard Adam-type algorithm, which we will augment for memory
savings. We will use $f$ for the loss function, $d$ for the model size, $k$
for the gradient density (sparsity $d-k$), $\theta_{t}$ and $g_{t}$ for the
model parameters and gradient at step $t$ respectively, $\eta_{t}$ for the
learning rate, $\lambda$ for the weight decay parameter, $m_{t}$ and $v_{t}$
for the first and second moment of the gradient, $\epsilon$ for the numerical
stability constant, $\beta_{1}$ and $\beta_{2}$ for the momentum coefficients
for $m_{t}$ and $v_{t}$ respectively. Furthermore, we use $e_{t}$ for the
error feedback (EF) vector, $b$ the number of bits for EF quantization, $m$
for the sliding window size, $\mathcal{G}=(\mathcal{I},\mathcal{V})$ for the
sliding window of size $m\times k$ that stores indices $\mathcal{I}$ and
values $\mathcal{V}$ selected by the Top-K operator.

#### Algorithm Description.

We provide pseudocode in Algorithm 1 and highlight the parts related to error
feedback quantization in blue. The main idea is to compress the gradients via
TopK sparsification before they enter the optimizer state, and to correct for
the inherent error by applying error feedback $e_{t}\in\mathbb{R}^{d}$.
Instead of storing the optimizer state directly, we maintain a “sliding
window” of highly-sparse past gradients and dynamically re-compute the Adam
statistics at each step based on this window. Yet, this alone does not improve
space, as the error buffer partially negates the benefits of gradient
compression. Instead, we prove that the error feedback accumulator can itself
be compressed via quantization.

Algorithm 1 MicroAdam: Implementation View

1:Input: $\beta_{1},\beta_{2},\epsilon,\mathcal{G},T,d$

2:$m_{0},v_{0},e_{1},\delta_{1},\Delta_{1}\leftarrow
0_{d},0_{d},0_{d}^{4b},0_{d},0_{d}$

3:for $t=\\{1,2,...,T\\}$ do

4:
$a_{t}\leftarrow\widetilde{\nabla}_{\theta}f(\theta_{t})+{\color[rgb]{0,0,1}Q^{-1}(e_{t},\delta_{t},\Delta_{t}})$

5: $\mathcal{I}_{t},\mathcal{V}_{t}\leftarrow T_{k}(|a_{t}|)$

6: $a_{t}[\mathcal{I}_{t}]\leftarrow 0$

7: $\delta_{t+1},\Delta_{t+1}\leftarrow min(a_{t}),max(a_{t})$

8: $e_{t+1}\leftarrow Q(a_{t},\delta_{t+1},\Delta_{t+1})$

9: $\mathcal{G}_{i,:}\leftarrow(\mathcal{I}_{t},\mathcal{V}_{t})$

10: $\hat{m_{t}}\leftarrow\textsc{AdamStats}(\beta_{1},\mathcal{G})$

11: $\hat{v_{t}}\leftarrow\textsc{AdamStats}(\beta_{2},\mathcal{G}^{2})$

12:
$\theta_{t+1}\leftarrow\theta_{t}-\eta_{t}\frac{\hat{m_{t}}}{\epsilon+\sqrt{\hat{v_{t}}}}$

13: $i\leftarrow(i+1)\%m$

14:end for

#### Dynamic Statistics.

In Algorithm 3, in line 4 we dynamically re-compute the decay factor
$\beta^{r}$ with exponent $r$ determined based on the current optimization
step $t$, $i^{th}$ row of the buffer $\mathcal{G}$ and the window size $m$.
The most recent row will have $r=0$, while the latest one will have $r=m-1$.
In the end, we will add the values $\beta^{r}\cdot\mathcal{V}_{i}$ to the
buffer $z$ at the corresponding indices $\mathcal{I}_{i}$. We discuss the
efficient CUDA implementation in the Appendix.

Algorithm 2 Symmetric Uniform Quantization and Dequantization

1:procedure $Q$($x,\delta,\Delta,b=4$)

2: $u\leftarrow\frac{\Delta-\delta}{2^{b}-1}$

3: $x_{Q}\leftarrow\lfloor\frac{x-\delta}{u}+\frac{1}{2}\rfloor$

4: return  $x_{Q}$

5:end procedure

1:procedure $Q^{-1}$($x_{Q},\delta,\Delta,b$)

2: $u\leftarrow\frac{\Delta-\delta}{2^{b}-1}$

3: $x\leftarrow x_{Q}\cdot u+\delta$

4: return  $x$

5:end procedure

Algorithm 3 Maintaining Adam Statistics

1:procedure AdamStats($\beta,\mathcal{G},t,m,d$)

2: $z\leftarrow 0_{d}$

3: for $i\in\\{1,2,...,\min(t,m)\\}$ do

4: $r\leftarrow(t-i-1)\%m$

5: $z[\mathcal{I}_{i}]\leftarrow
z[\mathcal{I}_{i}]+\beta^{r}\cdot\mathcal{V}_{i}$

6: end for

7: return  $z/(1-\beta^{t})$

8:end procedure

In detail, at each optimization step $t$ we add up the current gradient
$g_{t}=\nabla_{\theta}f(\theta_{t})$ and the dequantized error (see procedure
$Q^{-1}$ in Algorithm 2) from the previous iterations (feeding back the error)
into the accumulator variable $a_{t}$. The accumulator is compressed using the
Top-K operator before being passed to Adam. The pair
$(\mathcal{I}_{t},\mathcal{V}_{t})$ represents the indices of the largest
components by absolute value and the corresponding values (this step is
equivalent to selecting the outliers from $a_{t}$). In step 6, we remove the
outliers from the accumulator $a_{t}$ to finally obtain the error at the
current step $t$, which is still stored in $a_{t}$. From an algorithmic
perspective, this step is equivalent to writing $e_{t}\leftarrow
a_{t}-T_{k}(a_{t})$.

Next, we compress the error feedback. In our practical implementation, we will
use compression to 4-bit integers, more specifically using symmetric uniform
quantization in large blocks of fixed size $B_{q}$ (step 8). More precisely,
this computes the minimum and maximum values $\delta_{t}$ and $\Delta_{t}$
(step 7) to quantize the accumulator $a_{t}$, as detailed in the procedure
$Q$.

Finally, we update the $i^{th}$ row in the buffer $\mathcal{G}$ by replacing
the oldest pair of indices and values with the one newly computed at step $t$.
The statistics $m_{t}$ and $v_{t}$ are both estimated using the sliding window
$\mathcal{G}$, where $\mathcal{G}^{2}=(\mathcal{I},\mathcal{V}^{2})$ (via
element-wise squaring). We would like to point out that the resulting update
$u_{t}=m_{t}/(\epsilon+\sqrt{v_{t}})$ may be highly sparse when the window
size $m$ is small. For illustration, if we use density $k=d/100$ ($1\%$
density, e.g. $99\%$ sparsity) with $m=10$ and suppose that all rows in the
indices matrix $\mathcal{I}$ are disjoint, then the overall density in the
update $u_{t}$ will be at least $90\%$. (Note that the sparsity increases if
rows in $\mathcal{I}$ have common values.) However, over large windows, the
update sparsity should decrease.

#### Algorithm Intuition.

To gain more intuition, we illustrate the impact of compressing gradients
using TopK compression before they are incorporated into the optimizer state
for Adam, both with and without error feedback (EF). We use the classic 2D
Rosenbrock function $f(x,y)=(1-x)^{2}+100(y-x^{2})^{2}$ to illustrate this
effect. In the Appendix, we provide similar results for another ill-
conditioned problem. The optimization trajectories are discussed in Figure 1.
The results clearly show that EF is essential for fast convergence. Besides,
TopK with EF, which is a surrogate of MicroAdam, allows for competitive
convergence relative to the uncompressed baseline. In Appendix E we discuss
implications of Error Feedback applied to GaLore.

Figure 1:  Optimization trajectories of Adam, TopK-Adam and TopK-Adam with EF applied on the Rosenbrock function starting from $(x_{0},y_{0})=(-\frac{1}{2},1)$. Notice the extremely “jagged” profile of TopK-Adam without EF, and the recovered convergence when EF is added. |  |
---|---|---

###  3.1 Efficient Implementation

A direct implementation (e.g., in Pytorch) of the previous algorithm would not
bring significant benefits, and would in fact might slow down optimization in
terms of wall-clock time. To realize the theoretical gains, we detail a GPU-
aware implementation below.

#### Accumulator $a_{t}$.

First, we do not use an additional accumulator tensor $a_{t}$; instead, we use
a CUDA kernel to dequantize the error buffer, and store the result in the grad
attribute of the model parameters. This allows us to accumulate the error
feedback into the gradients, without any additional memory.

#### Top-K.

Since we run on LLMs with billions of parameters, naive storage of the sparse
indices would require using an int64 type for the indices matrix
$\mathcal{I}$, assuming that the Top-K operator is applied globally to all the
parameters in $a_{t}$. To avoid this cost, we apply Top-K in blocks of fixed
size $B_{d}<2^{16}-1=65535$ and store block-relative indices in int16 format.
Applying Top-K per row to $a_{t}$ reshaped to 2D is not only faster, but
provides the block-relative indices directly.

Computing Top-K in blocks also allows us to allocate and efficiently use CUDA
shared memory blocks to compute the statistics $m_{t}$ and $v_{t}$ for Adam,
as described in the AdamStats procedure below. Note that the block-relative
indices returned by Top-K will be directly used as indices in the shared
memory array of CUDA thread blocks to retrieve values from $\mathcal{I}$ and
$\mathcal{V}$.

#### Quantization metadata.

Our approach also stores two additional vectors $\delta$ and $\Delta$ used for
quantization. Since the quantization block size $B_{q}$ is set to a very large
value (e.g. $100K$), the space required for these two arrays becomes
negligible in comparison to the buffer $\mathcal{G}$ and error feedback $e$.

#### Practical memory usage.

We note that we apply MicroAdam per layer, and that the size of quantization
statistics $\delta$ and $\Delta$ are allocated based on the layer size. Having
many such small tensors may result in slightly sub-optimal memory allocation
from Pytorch. This is why our reported memory usage can be higher than the
theoretical usage for small models, in the 100M parameter range; these effects
disappear for billion-parameter models, where the savings are significant.

###  3.2 Memory footprint analysis and comparison with other methods

We now compare the theoretical memory footprint of MicroAdam AdamW (Loshchilov
and Hutter, 2019) and AdamW-8 bits (Dettmers et al., 2021), and GaLore (Zhao
et al., 2024), focusing on memory usage of the optimizer states $m_{t}$ and
$v_{t}$, each of size $d$, expressed in bytes (B). For concreteness, we report
the practical memory usage for the optimizer state for a Llama-2 7B model for
each optimizer.

  * •

AdamW stores states in float32 format (4 B per component), resulting in a
total memory footprint of $4d+4d=8d$ (B), while using bfloat16 would result in
$4d$ (B) memory. We will refer to these memory footprints as
$M_{AW32}=8d\text{ (B)}=50.21\text{ (GB)}$ and $M_{AW16}=4d\text{
(B)}=25.10\text{ (GB)}$.

  * •

AdamW-8 bit stores states in 8-bit format (1 B per component), both with $d$
components, with memory footprint of $M_{AW8}=d+d=2d\text{ (B)}=12.55\text{
(GB)}$.

  * •

MicroAdam stores the error feedback $e$ in 4-bit format (0.5 B per component)
and the sliding window $\mathcal{G}$ that stores the indices $\mathcal{I}$ in
int16 and $\mathcal{V}$ in bfloat16 format. Both have $m\times k$ components,
each requiring 2 B per component. In the end, for $m=10$ and $k=d/100$, the
memory footprint is $M_{\mu A}(m)=0.5d+4mk\text{ (B)}=0.9d\text{
(B)}=5.65\text{ (GB)}$, that is, half of AdamW-8bit.

  * •

GaLore. Given a neural network with $L$ layers, where each layer has weight
matrix $W_{i}\in\mathbb{R}^{A_{i}\times B_{i}}$ (shaped as a 2D matrix), the
model size $d=\sum_{i=1}^{L}A_{i}B_{i}$. GaLore uses a rank-$r$ compression
via the SVD composition as $W_{i}=USV^{T}$, where $U\in\mathbb{R}^{A_{i}\times
A_{i}}$ and we choose $R_{i}\in\mathbb{R}^{A_{i}\times r}$ to project the
gradient of $W_{i}$ to the $r$-dimensional space. As a consequence, the model
dimension shrinks from $d$ to
$d_{r}=\sum_{i=1}^{L}A_{i}r=r\sum_{i=1}^{L}A_{i}$, which represents the total
number of components to be stored in the GPU memory only for the projection
matrices $R_{i}$. If we suppose they are stored in bfloat16 (2 B per
component), then the entire memory usage for low-rank projection would be
$2d_{r}$. We would like to mention that some layers in the model are rank-$1$
and they do not need low-rank compression, but will still have associated
states in Adam, meaning that we have to take them into consideration when
computing the theoretical memory (we will use $\epsilon_{1}$ for memory of
rank-$1$ layers). In addition, we have to store the states $m_{t}$ and $v_{t}$
for AdamW in bfloat16 format, which adds another $4d_{r}$ bytes. In sum, the
total memory footprint of GaLore is $M_{GLAW16}(r)=6d_{r}+2\epsilon_{1}\text{
(B)}$, while for the 8-bit version we get
$M_{GLAW8}(r)=4d_{r}+2\epsilon_{1}\text{ (B)}$. In the end, the practical
memory usage for Llama-2 7B is $M_{GLAW8}(256)=1.36\text{ (GB)}$,
$M_{GLAW8}(1024)=5.43\text{ (GB)}$, $M_{GLAW16}(256)=2.04\text{ (GB)}$ and
$M_{GLAW8}(1024)=8.15\text{ (GB)}$.

#### Discussion.

Assume that our goal is to obtain a lower memory footprint compared to AdamW-8
bit. Recall that we use sparsification to compress gradients; then, we will
compute the gradient density $k=d/100$ for MicroAdam in order to be
competitive with AdamW-8 bit.

For this, we have to solve the equation $M_{\mu A}(m)=M_{AW8}$ for $m$,
resulting in $m_{\max}=37.5$. Specifically, if we use a gradient history of
$m<m_{\max}$ gradients in $\mathcal{G}$, MicroAdam will have theoretical
memory savings. We will see that, in practice, this history size is more than
enough for good practical results. As entries in the window past this range
are dampened extremely significantly, their influence is negligible. In
Appendix C, we provide Python code to compute the theoretical memory usage for
the three optimizers MicroAdam for Llama-2 7B.

##  4 Convergence Guarantees for MicroAdam

In this section, we present our theoretical framework. We begin by introducing
and discussing the analytical assumptions used in our theory, providing an
analytical view of the proposed MicroAdam algorithm, along with two
theoretical convergence results.

###  4.1 Gradient and Error Compression

We now define two classes of compression operators widely used in the
literature.

######  Assumption 1.

The gradient compressor $\mathcal{C}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ is
$q$-contractive with $0\leq q<1$, i.e.,

$\left\|\mathcal{C}(x)-x\right\|\leq q\left\|x\right\|,\quad\text{for any
}x\in\mathbb{R}^{d}.$

The compression we use in Algorithm 1 is the TopK compressor $T_{k}$ selecting
top $k$ coordinates in absolute value. This is known to be contractive with
$q=\sqrt{1-\nicefrac{{k}}{{d}}}$. Another popular contractive compressor is
the optimal low-rank projection of gradient shaped as a $d\times d$ matrix, in
which case $q=\sqrt{1-R/d}$ where $R$ is the projection rank.

The second class of compressors, which we use for the error feedback, requires
unbiasedness and relaxes the constant in the uniform bound.

######  Assumption 2.

The error compressor $\mathcal{Q}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ is unbiased
and $\omega$-bounded with $\omega\geq 0$, namely,

$\mathbb{E}[\mathcal{Q}(x)]=x,\quad\left\|\mathcal{Q}(x)-x\right\|\leq\omega\left\|x\right\|,\quad\text{for
any }x\in\mathbb{R}^{d}.$

One example of $\omega$-bounded compressor, a version of which is used in
Algorithm 2, is the randomized rounding quantizer we employ, whose properties
we provide below.

######  Lemma 1.

Consider Algorithm 2 with randomized rounding, i.e., for a vector
$x\in\mathbb{R}^{d}$ with $\delta=\min_{i}x_{i}$ and $\Delta=\max_{i}x_{i}$,
let $\hat{x}_{i}\mathrel{\mathop{:}}=\lfloor\frac{x_{i}-\delta}{u}+\xi\rfloor
u+\delta$ be the $i$-th coordinate of the quantized vector $\hat{x}$, where
$\xi\sim\textrm{U}[0,1]$ is the uniform random variable and
$u=\frac{\Delta-\delta}{2^{b}-1}$ is the quantization level. Then

$\mathbb{E}[\hat{x}]=x,\quad\|\hat{x}-x\|\leq\frac{\sqrt{d-2}}{2^{b}-1}\frac{\Delta-\delta}{\sqrt{\Delta^{2}+\delta^{2}}}\|x\|,\quad\text{for
all }x\in\mathbb{R}^{d}.$

Next, we provide an “analytical” view of our method in Algorithm 4.
Essentially, we use the contractive compressor $\mathcal{C}$ for compressing
the error corrected gradient information $g_{t}+e_{t}$, and the unbiased
compressor $\mathcal{Q}$ to compress the remaining compression error
$g_{t}+e_{t}-\mathcal{C}(g_{t}+e_{t})$.

Algorithm 4 MicroAdam: Analytical View

1:Input: parameters $\beta_{1},\,\beta_{2}\in(0,1)$, $\epsilon>0$, step-size
$\eta>0$, $\theta_{1}\in\mathbb{R}^{d}$, $e_{1}=m_{0}=v_{0}=\hat{v}_{0}=0_{d}$

2:for $t=\\{1,2,...,T\\}$ do

3: $g_{t}=\widetilde{\nabla}_{\theta}f(\theta_{t})$ $\diamond$ Compute
unbiased stochastic gradient

4: $\tilde{g}_{t}=\mathcal{C}(g_{t}+e_{t})$ $\diamond$ Add accumulated error
$e_{t}$ and compress

5: $e_{t+1}=\mathcal{Q}(e_{t}+g_{t}-\tilde{g}_{t})$ $\diamond$ Update and
compress the error

6: $m_{t}=\beta_{1}m_{t-1}+(1-\beta_{1})\tilde{g}_{t}$ $\diamond$ Update
first-order gradient moment

7: $v_{t}=\beta_{2}v_{t-1}+(1-\beta_{2})\tilde{g}_{t}^{2}$ $\diamond$ Update
second-order gradient moment

8: $\hat{v}_{t}=\max(v_{t},\hat{v}_{t-1})$ $\diamond$ Apply AMSGrad
normalization

9:
$\theta_{t+1}=\theta_{t}-\eta\frac{\hat{m_{t}}}{\sqrt{\hat{v}_{t}+\epsilon}}$
$\diamond$ Update the model parameters

10:end for

It is clear from this description that our objective with these two
compressors, $\mathcal{C}$ and $\mathcal{Q}$, is to approximate the dense
gradient information $g_{t}+e_{t}$ using two compressed vectors:
$\tilde{g}_{t}=\mathcal{C}(g_{t}+e_{t})$ and
$\mathcal{Q}(g_{t}+e_{t}-\tilde{g}_{t})$. However, in doing so, we inevitably
lose some information about $g_{t}+e_{t}$ depending on the degree of
compression applied to each term. Thus, the condition $(1+\omega)q<1$ required
by our analysis can be seen as preventing excessive loss of information due to
compression.

###  4.2 Convergence Guarantees for General Smooth Non-convex Functions

Next, we state our algorithm’s convergence guarantees under standard
assumptions, stated below:

######  Assumption 3 (Lower bound and smoothness).

The loss function $f\colon\mathbb{R}^{d}\to\mathbb{R}$ is $L$-smooth and lower
bounded by some $f^{*}\in\mathbb{R}$:

$\left\|\nabla f(\theta)-\nabla f(\theta^{\prime})\right\|\leq
L\left\|\theta-\theta^{\prime}\right\|,\quad\text{for any
}\theta,\theta^{\prime}\in\mathbb{R}^{d}.$

######  Assumption 4 (Unbiased and bounded stochastic gradient).

For all iterates $t\geq 1$, the stochastic gradient $g_{t}$ is unbiased and
uniformly bounded by some constant $G\geq 0$:

$\mathbb{E}[g_{t}]=\nabla f(\theta_{t}),\quad\left\|g_{t}\right\|\leq G.$

######  Assumption 5 (Bounded variance).

For all iterates $t\geq 1$, the variance of the stochastic gradient $g_{t}$ is
uniformly bounded by some constant $\sigma^{2}\geq 0$:

$\mathbb{E}[\|g_{t}-\nabla f(\theta_{t})\|^{2}]\leq\sigma^{2}.$

Main Result. The above assumptions are standard in the literature, e.g.
(Défossez et al., 2022; Li et al., 2022; Xie et al., 2023; Zhou et al.,
2024a). Under these conditions, if the two compressors satisfy the basic
condition $(1+\omega)q<1$, we show:

######  Theorem 1.

(Non-convex convergence rate) Let Assumptions 1, 2, 3, 4, 5 hold and
$q_{\omega}\mathrel{\mathop{:}}=(1+\omega)q<1$. Then, choosing
$\eta=\frac{1}{\sqrt{T}}\leq\frac{\epsilon}{4LC_{0}}$, MicroAdam (Algorithm 4)
satisfies

$\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}]\leq
2C_{0}\left(\frac{f(\theta_{1})-f^{*}}{\sqrt{T}}+\frac{L(\sigma^{2}+C_{2}^{2}G^{2})}{\epsilon\sqrt{T}}\right)+\mathcal{O}\left(\frac{G^{3}(G+d)}{T}\right)$

with constants
$C_{0}\mathrel{\mathop{:}}=\sqrt{\frac{4(1+q_{\omega}^{2})^{3}}{(1-q_{\omega}^{2})^{2}}G^{2}+\epsilon}$
and $C_{2}\mathrel{\mathop{:}}=\omega
q(1+\frac{2q_{\omega}}{1-q_{\omega}^{2}})$.

#### Discussion.

First, notice that the leading term $\tfrac{1}{\sqrt{T}}$ of the rate is the
optimal convergence speed for non-convex stochastic gradient methods (Ghadimi
and Lan, 2016). Furthermore, the obtained convergence rate
$\mathcal{O}(\tfrac{1}{\sqrt{T}}+\tfrac{d}{T})$ asymptotically matches the
rate of uncompressed AMSGrad in the stochastic non-convex setup (Zhou et al.,
2024a). Hence, the added compression framework of the MicroAdam together with
error feedback mechanism can slow down the convergence speed only up to some
constants. Evidently, the additional constants $C_{0}$ and $C_{2}$ affected by
compression and appearing in the leading terms can be easily estimated once
the compressors are fixed. Besides, if we store the full error information
without applying $\mathcal{Q}$ compressor (i.e., $\omega=0,\,q_{\omega}=q$),
then MicroAdam reduces to the single-node Comp-AMS method by Li et al. (2022)
recovering the same convergence rate. The full proof is provided in the
Appendix.

###  4.3 Convergence Rate for Non-Convex Functions under the PL Condition

Next, we show that we can obtain even stronger bounds when the objective
satisfies the PL condition:

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

For some $\mu>0$ the loss $f$ satisfies Polyak-Lojasiewicz (PL) inequality

$\|\nabla f(\theta)\|^{2}\geq 2\mu(f(\theta)-f^{*}),\quad\text{for any
}\theta\in\mathbb{R}^{d}.$

In this case, we can show:

######  Theorem 2.

(PL convergence rate) Let Assumptions 1, 2, 3, 4, 5 and 6 hold, and
$q_{\omega}<1$. Then, choosing $\eta=\frac{C_{0}\log T}{\mu
T}\leq\frac{\epsilon}{4LC_{0}}$, MicroAdam (Algorithm 4) satisfies

$\mathbb{E}[f(\theta_{T+1})]-f^{*}\leq\frac{f(\theta_{1})-f^{*}}{T}+\frac{\log
T}{T}\left(\frac{LC_{0}^{2}}{\mu}\frac{\sigma^{2}+(C_{1}+C_{2}^{2})G^{2}}{\mu\epsilon}+\frac{C_{0}(1+C_{1})(1+d)G^{2}}{\mu\sqrt{\epsilon}}\right)\\\
+\widetilde{\mathcal{O}}\left(\frac{G^{4}(G+d)}{T^{2}}\right)$

with constant
$C_{1}\mathrel{\mathop{:}}=\frac{\beta_{1}}{1-\beta_{1}}(1+C_{2})+\frac{2q_{\omega}}{1-q_{\omega}^{2}}$.

#### Discussion.

In contrast to the general non-convex setup, the study of non-convex analysis
under the PL condition for AMSGrad or Adam-type methods is much less
extensive. The only work we found analyzing the PL condition, which claims to
be the first in this direction, focuses on Adam when $\beta_{2}\to 1$,
achieving a convergence rate of $\mathcal{O}(\tfrac{1}{T})$ (He et al., 2023).
However, our MicroAdam is based on AMSGrad normalization, and no constraint on
$\beta_{2}$ is imposed in the analysis. Therefore, similar to the general non-
convex case, we are able to achieve the best-known convergence rate in the
leading term, up to a logarithmic factor. The third, higher-order term has
higher constant dependencies, but they should be negligible as the term is
dampened by $T^{2}$. Hence, in this case as well, the theory predicts that the
convergence rate of the algorithm should be similar to the uncompressed
version, modulo a constant that can be controlled using the compression
parameters.

##  5 Experiments

We now validate our optimizer experimentally. We focus on comparing MicroAdam
with Adam, Adam-8bit, GaLore and CAME in the context of language model
finetuning on different tasks. Concretely, we test our optimizer in full
finetuning (FFT) scenario on BERT-Base/Large (Devlin et al., 2018) and
OPT-1.3B (Zhang et al., 2022) on GLUE/MNLI and Llama2-7B (Touvron et al.,
2023) on the GSM8k math reasoning dataset and on the Open-Platypus instruction
tuning dataset. We train all our models in bfloat16 format and tune the
learning rates on a grid and report the best accuracy among 3 seeds. All
methods use default AdamW parameters, such as
$\beta_{1}=0.9,\beta_{2}=0.999,\epsilon=10^{-8}$. MicroAdam uses a window size
of $m=10$ gradients with $k=1\%$ density and quantization bucket size is set
to 64 for the error feedback. For GaLore we use rank $r=256$ and the SVD
update interval is set to $T=200$, as suggested by the original paper. We run
our experiments on NVidia GPUs A100-SXM4-80GB and on RTX 3090 with 24GB RAM in
single GPU setup. We provide full details regarding hyper-parameters in
Appendix A.

#### Results on GLUE/MNLI.

We first test our integration of MicroAdam in HuggingFace Transformers (Wolf
et al., 2020) on moderate-sized language models such as BERT-Base/Large (110M
and 335M parameters) and OPT-1.3B, comparing with Adam, Adam-8bit, CAME and
GaLore. The results are shown in Table 1. (Certain optimizers, notably CAME
and GaLore, had numerical stability issues across runs; for a fair comparison,
we report the numbers for the run with maximum accuracy. We emphasize that all
methods were tuned using the same protocol.)

The results show that MicroAdam achieves comparable memory usage to the state-
of-the-art heuristics Adam-8bit and GaLore, while being surprisingly lower
than CAME on all tasks. The memory savings for GaLore are more visible when
the model size increases, which follows our analysis of theoretical memory
usage. However, we see that these gains come at a significant _accuracy cost_
for GaLore: across all tasks, it drops at least 1% accuracy relative to
MicroAdam. For BERT-Base we ran GaLore with a higher SVD re-computation
frequency $T=20$ ($10\times$ lower) and the results did not improve, but its
running time was much higher. Relative to 8bit Adam, MicroAdam uses
essentially the same memory, but achieves slightly better accuracy.

From these results, we conclude that MicroAdam can provide better accuracy
relative to other memory-efficient methods on moderate-sized models, at
similar space costs. We show training loss curves in Appendix B.

Table 1:  Finetuning results on GLUE/MNLI. We report the entire memory usage read from the GPU during training, that includes the optimizer state, activations and gradients. The asterisk flags the runs for which one or two seeds did not converge (we report the run with maximum performance). Model | Metric |  MicroAdam $(m=10)$ | Adam | Adam-8b | CAME |  GaLore $r=256$
---|---|---|---|---|---|---
| train loss | 0.2651 | 0.4228 | 0.3402 | 0.6431* | 0.3908*
Base | accuracy | 85.10% | 83.53% | 84.61% | 76.13%* | 83.82%*
(110M) | memory | 2.55 GB | 2.70 GB | 2.53 GB | 2.69 GB | 2.53 GB
| train loss | 0.2509 | 0.3857 | 0.2876 | 0.6658* | 0.3768*
Large | accuracy | 86.17% | 84.79% | 86.18% | 75.23%* | 84.90%*
(335M) | memory | 5.98 GB | 6.64 GB | 6.04 GB | 6.59 GB | 5.85 GB
| train loss | 0.2122 | 0.2066 | 0.2611 | 0.4959 | 0.2831
OPT-1.3B | accuracy | 88.18% | 87.90% | 87.81% | 83.15% | 87.70
(1.3B) | memory | 15.28 GB | 17.66 GB | 15.00 GB | 17.13 GB | 13.66 GB

#### Results on LLaMA2-7B/GSM-8k.

Next, we perform finetuning on Llama-2 7B on GSM-8k, a challenging grade-
school-level mathematical reasoning dataset. The baseline model obtains
extremely low zero-shot accuracy on this task, therefore fine-tuning is
necessary. In this setup, we compare MicroAdam with Adam and Adam-8bit in
terms of evaluation accuracy and memory usage. In Table 2 we show our results
for 3 training epochs, global batch size $32$ with micro-batch (per-device)
size $1$, max sequence length $512$ on a single GPU. These parameters are
standard for this task. We integrated our optimizer with the llm-foundry
repository of MosaicML and tested via lm-evaluation-harness.

Table 2:  FFT results for Llama-2 7B on GSM-8k | Adam | Adam-8b | MicroAdam
---|---|---|---
|  |  | $m=10$ | $m=20$
accuracy | 34.50% | 34.34% | 34.72% | 35.10%
memory | 55.2 GB | 42.5 GB | 37.1 GB | 39.7 GB
runtime | 1h 17m | 1h 18m | 1h 8m | 1h 37m

The results show that MicroAdam can allow accurate full fine-tuning of a 7B
model on this task using a single 40GB GPU. Moreover, MicroAdam preserves
accuracy relative to Adam, with lower memory usage than the well-optimized
implementation of 8bit AdamW, and marginally lower running time for the
shorter gradient window $m=10$. Increasing the window size $m$ to $20$
gradients leads to slightly better accuracy, at the cost of higher runtime and
space. Running GaLore in this setup was infeasible since using SVD
decomposition for all layers in the model was too slow; preliminary
experiments (with high runtimes) did not yield competitive accuracy. We show
training loss curves in Appendix B.

#### Results on LLaMA2-7B/Open-Platypus.

Finally, in Table 3 we present FFT results with various optimizers on the
popular instruction-tuning Open-Platypus dataset (Lee et al., 2023). To ensure
fair comparisons, we perform the same grid search for each optimizer to find
the best performing learning-rate, while keeping all other hyperparameters at
their default values. We use $m=10$ gradients for the sliding window and
gradient density $k=1\%$. Evaluations are conducted following the standard
few-shot setup of the Open LLM Leaderboard (Beeching et al., 2023) on the
following datasets: ARC-c (Clark et al., 2018), HellaSwag (Zellers et al.,
2019), MMLU (Hendrycks et al., 2021), and Winogrande (Sakaguchi et al., 2019).

Table 3:  FFT results on instruction-following Open-Platypus (Lee et al., 2023) dataset. The results show that MicroAdam fully recovers accuracy relative to baseline Adam, and outperforms the 8bit variant, despite using less memory.  Optimizer | Memory |  Average Accuracy |  ARC-c 25-shot |  HellaSwag 10-shot |  MMLU 5-shot |  Winogrande 5-shot
---|---|---|---|---|---|---
AdamW | 67.17 GB | 62.10 | 52.56 | 77.38 | 45.53 | 72.93
Adam-8b | 53.93 GB | 61.84 | 51.96 | 77.51 | 44.11 | 73.79
MicroAdam | 46.63 GB | 62.36 | 53.07 | 77.46 | 45.04 | 73.87

The results show that MicroAdam allows for full accuracy recovery on this task
as well relative to Adam, despite using 50% less memory. (The memory usage and
runtime are very similar to those in Table 2 and are therefore omitted from
Table 3.) Moreover, MicroAdam obtains consistently better accuracy relative to
Adam-8b, especially on the more challenging ARC-c task.

#### Discussion.

In summary, the experimental results have shown that MicroAdam can recover the
state-of-the-art accuracy of the the uncompressed Adam baseline, while
providing significant memory gains and matching wall-clock speed on billion-
parameter models. Specifically, our approach matches and outperforms Adam-8b
and CAME both in terms of memory use and in terms of final accuracy. Relative
to the high-compression GaLore method, MicroAdam provides consistently higher
accuracy, as well as more stable practical convergence. We conclude that
MicroAdam should be a good alternative to 8bit Adam in memory-constrained
settings, and that the empirical results appear to validate our theoretical
predictions.

##  6 Limitations and Broader Impact

The MicroAdam algorithm we propose is designed and tested with fine-tuning
workloads in mind, where the user aims to minimize the memory cost of
optimizing over a powerful pre-trained model. Additional work is needed to
adapt our approach to the case of LLM pretraining, which presents a different
set of challenges, both in terms of implementation and optimization
trajectory. We plan to undertake this study in future work.

Another limitation we aim to address in future work is that we have only
focused on sparsity as a form of gradient projection. However, our theoretical
analysis also applies to low-rank projection of gradients. We believe that our
practical implementation can be extended to this case as well, although
providing a general, accurate, and efficient implementation will require non-
trivial efforts.

Our work introduces a new, accurate, and memory-efficient optimizer for fine-
tuning LLMs. The major positive impact of our approach is its ability to
maintain performance while reducing memory requirements, thereby lowering the
cost of running experiments due to the reduced hardware expenses. It is
important to note that while our optimizer can enhance performance and reduce
costs, we do not have control over the neural network applications trained
with it.

## Acknowledgements

The authors thank Razvan Pascanu, Mahdi Nikdan and Soroush Tabesh for their
valuable feedback, the IT department from Institute of Science and Technology
Austria for the hardware support and Weights and Biases for the infrastructure
to track all our experiments. Mher Safaryan has received funding from the
European Union’s Horizon 2020 research and innovation program under the Marie
Sklodowska-Curie grant agreement No 101034413.

## References

  * Agarwal et al. [2019] N. Agarwal, B. Bullins, X. Chen, E. Hazan, K. Singh, C. Zhang, and Y. Zhang.  Efficient full-matrix adaptive regularization.  In _International Conference on Machine Learning_ , pages 102–110. PMLR, 2019.
  * Alistarh et al. [2018] D. Alistarh, T. Hoefler, M. Johansson, N. Konstantinov, S. Khirirat, and C. Renggli.  The convergence of sparsified gradient methods.  In _Advances in Neural Information Processing Systems_ , pages 5973–5983, 2018.
  * Amari [2016] S.-i. Amari.  _Information geometry and its applications_ , volume 194.  Springer, 2016.
  * Anil et al. [2019] R. Anil, V. Gupta, T. Koren, and Y. Singer.  Memory efficient adaptive optimization.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Beeching et al. [2023] E. Beeching, C. Fourrier, N. Habib, S. Han, N. Lambert, N. Rajani, O. Sanseviero, L. Tunstall, and T. Wolf.  Open llm leaderboard.  https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard, 2023.
  * Clark et al. [2018] P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord.  Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018.
  * Défossez et al. [2022] A. Défossez, L. Bottou, F. Bach, and N. Usunier.  A simple convergence proof of adam and adagrad.  _Transactions on Machine Learning Research_ , 2022.  ISSN 2835-8856.  URL https://openreview.net/forum?id=ZPQhzTSWA7.
  * Dettmers et al. [2021] T. Dettmers, M. Lewis, S. Shleifer, and L. Zettlemoyer.  8-bit optimizers via block-wise quantization.  _arXiv preprint arXiv:2110.02861_ , 2021.
  * Devlin et al. [2018] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding, 2018.
  * Duchi et al. [2010] J. C. Duchi, E. Hazan, and Y. Singer.  Adaptive subgradient methods for online learning and stochastic optimization.  _J. Mach. Learn. Res._ , 12:2121–2159, 2010.
  * Feinberg et al. [2023] V. Feinberg, X. Chen, Y. J. Sun, R. Anil, and E. Hazan.  Sketchy: Memory-efficient adaptive regularization with frequent directions, 2023.
  * Frantar et al. [2021] E. Frantar, E. Kurtic, and D. Alistarh.  M-fac: Efficient matrix-free approximations of second-order information, 2021.
  * Ghadimi and Lan [2016] S. Ghadimi and G. Lan.  Accelerated gradient methods for nonconvex nonlinear and stochastic programming.  _Mathematical Programming_ , 2016.  ISSN 156(1-2):59–99.
  * Gupta et al. [2018] V. Gupta, T. Koren, and Y. Singer.  Shampoo: Preconditioned stochastic tensor optimization.  In _International Conference on Machine Learning_ , pages 1842–1850. PMLR, 2018.
  * He et al. [2023] M. He, Y. Liang, J. Liu, and D. Xu.  Convergence of adam for non-convex objectives: Relaxed hyperparameters and non-ergodic case.  _Journal of Machine Learning Research_ , 2023.
  * Hendrycks et al. [2021] D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt.  Measuring massive multitask language understanding, 2021.
  * Karimireddy et al. [2019] S. P. Karimireddy, Q. Rebjock, S. U. Stich, and M. Jaggi.  Error feedback fixes SignSGD and other gradient compression schemes.  In _Proceedings of the Thirty-sixth International Conference on Machine Learning_ , pages 3252–3261, 2019.
  * Kingma and Ba [2014] D. P. Kingma and J. Ba.  Adam: A method for stochastic optimization, 2014.
  * Lee et al. [2023] A. N. Lee, C. J. Hunter, and N. Ruiz.  Platypus: Quick, cheap, and powerful refinement of llms.  _arXiv preprint arXiv:2308.07317_ , 2023.
  * Li et al. [2022] X. Li, B. Karimi, and P. Li.  On distributed adaptive optimization with gradient compression.  _arXiv preprint arXiv:2205.05632_ , 2022.
  * Liu et al. [2023] H. Liu, Z. Li, D. Hall, P. Liang, and T. Ma.  Sophia: A scalable stochastic second-order optimizer for language model pre-training.  _arXiv preprint arXiv:2305.14342_ , 2023.
  * Loshchilov and Hutter [2019] I. Loshchilov and F. Hutter.  Decoupled weight decay regularization.  In _Proceedings of the Seventh International Conference on Learning Representations_ , 2019.
  * Luo et al. [2023] Y. Luo, X. Ren, Z. Zheng, Z. Jiang, X. Jiang, and Y. You.  Came: Confidence-guided adaptive memory efficient optimization.  _arXiv preprint arXiv:2307.02047_ , 2023.
  * Modoranu et al. [2023] I.-V. Modoranu, A. Kalinov, E. Kurtic, and D. Alistarh.  Error feedback can accurately compress preconditioners.  _arXiv preprint arXiv:2306.06098_ , 2023.
  * Reddi et al. [2019] S. J. Reddi, S. Kale, and S. Kumar.  On the convergence of adam and beyond.  _arXiv preprint arXiv:1904.09237_ , 2019.
  * Richtárik et al. [2021] P. Richtárik, I. Sokolov, and I. Fatkhullin.  Ef21: A new, simpler, theoretically better, and practically faster error feedback.  _Advances in Neural Information Processing Systems_ , 34:4384–4396, 2021.
  * Sakaguchi et al. [2019] K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi.  WINOGRANDE: an adversarial winograd schema challenge at scale, 2019.
  * Seide et al. [2014] F. Seide, H. Fu, J. Droppo, G. Li, and D. Yu.  1-bit stochastic gradient descent and its application to data-parallel distributed training of speech DNNs.  In _Fifteenth Annual Conference of the International Speech Communication Association_ , 2014.
  * Shazeer and Stern [2018] N. Shazeer and M. Stern.  Adafactor: Adaptive learning rates with sublinear memory cost.  In _International Conference on Machine Learning_ , pages 4596–4604. PMLR, 2018.
  * Spring et al. [2019] R. Spring, A. Kyrillidis, V. Mohan, and A. Shrivastava.  Compressing gradient optimizers via count-sketches.  In _International Conference on Machine Learning_ , pages 5946–5955. PMLR, 2019.
  * Touvron et al. [2023] H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample.  Llama: Open and efficient foundation language models, 2023.
  * Wolf et al. [2020] T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, J. Davison, S. Shleifer, P. von Platen, C. Ma, Y. Jernite, J. Plu, C. Xu, T. L. Scao, S. Gugger, M. Drame, Q. Lhoest, and A. M. Rush.  Transformers: State-of-the-art natural language processing.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations_ , pages 38–45, Online, Oct. 2020. Association for Computational Linguistics.  URL https://www.aclweb.org/anthology/2020.emnlp-demos.6.
  * Xie et al. [2023] X. Xie, P. Zhou, H. Li, Z. Lin, and S. Yan.  Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models, 2023.
  * Yao et al. [2020] Z. Yao, A. Gholami, S. Shen, K. Keutzer, and M. W. Mahoney.  Adahessian: An adaptive second order optimizer for machine learning.  _arXiv preprint arXiv:2006.00719_ , 2020.
  * Zellers et al. [2019] R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi.  Hellaswag: Can a machine really finish your sentence?, 2019.
  * Zhang et al. [2022] S. Zhang, S. Roller, N. Goyal, M. Artetxe, M. Chen, S. Chen, C. Dewan, M. Diab, X. Li, X. V. Lin, T. Mihaylov, M. Ott, S. Shleifer, K. Shuster, D. Simig, P. S. Koura, A. Sridhar, T. Wang, and L. Zettlemoyer.  Opt: Open pre-trained transformer language models, 2022.
  * Zhao et al. [2024] J. Zhao, Z. Zhang, B. Chen, Z. Wang, A. Anandkumar, and Y. Tian.  Galore: Memory-efficient llm training by gradient low-rank projection.  _arXiv preprint arXiv:2403.03507_ , 2024.
  * Zhou et al. [2024a] D. Zhou, J. Chen, Y. Cao, Z. Yang, and Q. Gu.  On the convergence of adaptive gradient methods for nonconvex optimization.  _Transactions on Machine Learning Research_ , 2024a.  ISSN 2835-8856.  URL https://openreview.net/forum?id=Gh0cxhbz3c.  Featured Certification.
  * Zhou et al. [2024b] P. Zhou, X. Xie, Z. Lin, and S. Yan.  Towards understanding convergence and generalization of adamw.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , pages 1–8, 2024b.  doi: 10.1109/TPAMI.2024.3382294.

###### Contents

  1. 1 Introduction
  2. 2 Related Work
  3. 3 The MicroAdam Algorithm
    1. 3.1 Efficient Implementation
    2. 3.2 Memory footprint analysis and comparison with other methods
  4. 4 Convergence Guarantees for MicroAdam
    1. 4.1 Gradient and Error Compression
    2. 4.2 Convergence Guarantees for General Smooth Non-convex Functions
    3. 4.3 Convergence Rate for Non-Convex Functions under the PL Condition
  5. 5 Experiments
  6. 6 Limitations and Broader Impact
  7. A Hyper-parameters
    1. A.1 GLUE/MNLI
    2. A.2 GSM-8k.
  8. B Training Graphs
  9. C Memory footprint for the optimizer state
  10. D Deferred Proofs
    1. D.1 Intermediate Lemmas
    2. D.2 Non-convex Analysis
    3. D.3 Analysis Under PL Condition
    4. D.4 Non-convex Analysis with Weight Decay
  11. E Error Feedback applied to GaLore
    1. E.1 Behaviour of the Error Feedback Mechanism
    2. E.2 Consequences on Training

##  Appendix A Hyper-parameters

In this section we provide details about the hyper-parameters that we used for
each model and dataset. All Adam variants use default parameters
$\beta_{1}=0.9,\beta_{2}=0.999,\epsilon=10^{-8},\lambda=0$, where $\lambda$ is
the weight decay parameter. We use $3$ seeds (7, 42 and 1234) for each
learning rate and report the results for the best configuration that
converged.

###  A.1 GLUE/MNLI

For GLUE/MNLI, we used the learning rate grid
$\\{1e-6,3e-6,5e-6,7e-6,1e-5,3e-5,5e-5,7e-5\\}$ for all optimizers and models.
Certain optimizers diverge for specific seeds. Next, we provide some details
about hyper-parameters for each optimizer individually.

####  MicroAdam.

We use $m=10$ gradients in the sliding window, $k=1\%$ density (e.g. $99\%$
sparsity) and quantization bucket size $64$ (we also tried 100 000, but this
didn’t affect performance or memory usage in a meaningful way).

#### Adam and Adam-8bit.

All hyper-parameters mentioned above apply for these two main baseline
optimizers.

#### GaLore.

We use rank $r=256$ and SVD update interval $T\in\\{20,200\\}$. In the
original GaLore paper, the authors tune both learning rate and in our
experiments we keep scale fixed to value 1 and augment the learning rate grid
with the values $\\{1e-4,3e-4,5e-4,7e-4\\}$.

#### CAME.

This optimizer has some additional parameters that we keep to default values,
such as $\beta_{3}=0.9999$. Instead of $\epsilon$, it uses
$\epsilon_{1}=1e-30$ and $\epsilon_{2}=1e-16$. The authors mention that the
learning rate should be much smaller than Adam’s and because of that we
augment the learning rate grid with the values $\\{1e-7,3e-7,5e-7,7e-7\\}$.

###  A.2 GSM-8k.

For GSM-8k, we used the learning rate grid
$\\{1e-5,2e-5,3e-5,4e-5,5e-5,6e-7,7e-5,8e-5,9e-5\\}$ and reported the model
with the best evaluation accuracy. We found that different versions for
PyTorch, lm-eval-harness and llm-foundry have large variance in the results.

####  MicroAdam.

We use similar settings as for GLUE/MNLI above in terms of other hyper-
parameters.

##  Appendix B Training Graphs

In this section we show training loss curves for BERT-Base, BERT-Large and
OPT-1.3b on GLUE/MNLI and Llama-2 7B on GSM-8k.

Figure 2: Training curves for BERT-Base on GLUE/MNLI Figure 3: Training curves
for BERT-Large on GLUE/MNLI Figure 4: Training curves for OPT-1.3B on
GLUE/MNLI Figure 5: Training curves for Llama-2 7B on GSM-8k

##  Appendix C Memory footprint for the optimizer state

In this section we provide a python script to simulate the memory usage for
our optimizer’s state for Llama2-7b model. Note that the theoretical memory
usage will always be slightly lower than the actual allocated memory on the
GPU because PyTorch usually allocates more. To run this script, run the
following commands:

⬇

import math

d = 6_738_415_616 # actual number of parameters for Llama-2 7b

k = math.ceil(d / 100)

m = 10

M_AW32 = 8 * d / (2 ** 30)

M_AW16 = 4 * d / (2 ** 30)

M_AW8 = 2 * d / (2 ** 30)

M_muA = (0.5 * d + 4 * m * k) / (2 ** 30) # B to GB

print(f’{M_AW32=:.2f}␣GB’)

print(f’{M_AW16=:.2f}␣GB’)

print(f’{M_AW8=:.2f}␣GB’)

print(f’{M_muA=:.2f}␣GB’)

GL_sumA = 1_423_872 # sum_Ai from Llama-2 7B

epsilon = 266_240 # sum of sizes for rank-1 layers

for bits, const in [(8, 4), (16, 6)]:

for rank in [256, 1024]:

dr = rank * GL_sumA

M_GLAW_rank = (const * dr + 2 * epsilon) / (2 ** 30)

print(f’M_GLAW{bits}_rank{rank}={M_GLAW_rank:.2f}␣GB’)

##  Appendix D Deferred Proofs

At time step $t$, let the uncompressed stochastic gradient be
$g_{t}=\widetilde{\nabla}_{\theta}f(\theta_{t})$, the error accumulator be
$e_{t}$, and the compressed gradient after the error correction be
$\tilde{g}_{t}=\mathcal{C}(g_{t}+e_{t})$. The second moment computed by the
compressed gradients is denoted as
$v_{t}=\beta_{2}v_{t-1}+(1-\beta_{2})\tilde{g}_{t}^{2}$, and
$\hat{v}_{t}=\max\\{\hat{v}_{t-1},v_{t}\\}$ is the AMSGrad normalization for
the second-order momentum. Besides the first-order gradient momentum $m_{t}$
used in the algorithm description, we define similar running average sequence
$m^{\prime}_{t}$ based on the uncompressed gradients $g_{t}$.

$\displaystyle m_{t}=\beta_{1}m_{t-1}+(1-\beta_{1})\tilde{g}_{t}\quad$
$\displaystyle\textrm{and}\quad
m_{t}^{\prime}=\beta_{1}m_{t-1}^{\prime}+(1-\beta_{1})g_{t},$

Note that $m_{t}^{\prime}$ is used only in the analysis, we do not need to
store or compute it. By construction we have

$m_{t}=(1-\beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}\tilde{g}_{\tau},\quad
m_{t}^{\prime}=(1-\beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}g_{\tau}$

Denote by
$\zeta_{t}=e_{t+1}-(e_{t}+g_{t}-\tilde{g}_{t})=\mathcal{Q}(e_{t}+g_{t}-\tilde{g}_{t})-(e_{t}+g_{t}-\tilde{g}_{t})$
the compression noise from $\mathcal{Q}$. Due to unbiasedness of the
compressor $\mathcal{Q}$ (see Assumption 2), we have
$\mathbb{E}[\zeta_{t}\mid\theta_{t},g_{t},\tilde{g}_{t},e_{t}]=0$. Also, from
the update rule of $e_{t+1}$ we get
$e_{t+1}=e_{t}+g_{t}-\tilde{g}_{t}+\zeta_{t}$. Moreover, we use the following
auxiliary sequences,

$\displaystyle\mathcal{E}_{t+1}\mathrel{\mathop{:}}=\beta_{1}\mathcal{E}_{t}+(1-\beta_{1})e_{t+1}=(1-\beta_{1})\sum_{\tau=1}^{t+1}\beta_{1}^{t+1-\tau}e_{\tau}.$
$\displaystyle\mathcal{Z}_{t+1}\mathrel{\mathop{:}}=\beta_{1}\mathcal{Z}_{t}+(1-\beta_{1})\zeta_{t+1}=(1-\beta_{1})\sum_{\tau=1}^{t+1}\beta_{1}^{t+1-\tau}\zeta_{\tau}.$

###  D.1 Intermediate Lemmas

See 1

###### Proof.

The unbiasedness can be verified directly from the definition for each
coordinate. Without loss of generality assume that $\delta=x_{1}\leq
x_{2}\leq\dots\leq x_{d-1}\leq x_{d}=\Delta$. By construction of the
quantization, we have $|\hat{x}_{1}-x_{1}|=|\hat{x}_{d}-x_{d}|=0$ and
$|\hat{x}_{i}-x_{i}|\leq u$ for the remaining coordinates $2\leq i\leq d-1$.
Then

$\|\hat{x}-x\|^{2}=\sum_{i=1}^{d}|\hat{x}_{i}-x_{i}|^{2}\leq(d-2)u^{2}\leq\frac{(d-2)u^{2}}{\Delta^{2}+\delta^{2}}\|x\|^{2},$

which completes the proof. ∎

######  Lemma 2.

Under Assumptions 1-5, for all iterates $t$ and $T$ we have

$\|m_{t}^{\prime}\|\leq
G,\quad\text{and}\quad\sum_{t=1}^{T}\mathbb{E}[\|m_{t}^{\prime}\|^{2}]\leq
T\sigma^{2}+\sum_{t=1}^{T}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}].$

###### Proof.

The first part follows from triangle inequality and the Assumption 4 on
bounded stochastic gradient:

$\|m_{t}^{\prime}\|=(1-\beta_{1})\left\|\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}g_{\tau}\right\|\leq(1-\beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}\|g_{\tau}\|\leq
G.$

For the second claim, the expected squared norm of average stochastic gradient
can be bounded by

$\mathbb{E}\left[\|g_{t}\|^{2}\right]=\mathbb{E}\left[\|g_{t}-\nabla
f(\theta_{t}))\|^{2}\right]+\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}]\\\
\leq\sigma^{2}+\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}],$ (1)

where we use Assumption 5 that $g_{t}$ is unbiased with bounded variance. Let
$g_{t,j}$ denote the $j$-th coordinate of $g_{t}$. Applying Jensen’s
inequality for the squared norm, we get

$\displaystyle\mathbb{E}[\|m_{t}^{\prime}\|^{2}]$ $\displaystyle=$
$\displaystyle\mathbb{E}\left[\left\|(1-\beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}g_{\tau}\right\|^{2}\right]$
$\displaystyle\leq$
$\displaystyle(1-\beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}\mathbb{E}[\|g_{\tau}\|^{2}]$
$\displaystyle\leq$
$\displaystyle\sigma^{2}+(1-\beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}\mathbb{E}[\|\nabla
f(\theta_{\tau})\|^{2}],$

Summing over $t=1,\dots,T$, we obtain

$\displaystyle\sum_{t=1}^{T}\mathbb{E}[\|m_{t}^{\prime}\|^{2}]\leq
T\sigma^{2}+(1-\beta_{1})\sum_{t=1}^{T}\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}\mathbb{E}[\|\nabla
f(\theta_{\tau})\|^{2}]\leq T\sigma^{2}+\sum_{t=1}^{T}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}],$

which completes the proof. ∎

######  Lemma 3.

Let $q_{\omega}=(1+\omega)q<1$. Under Assumptions 1-5, for all iterates $t$ we
have

$\displaystyle\|e_{t}\|^{2}\leq\frac{4q_{\omega}^{2}}{(1-q_{\omega}^{2})^{2}}G^{2},$
$\displaystyle\mathbb{E}[\|e_{t+1}\|^{2}]\leq\frac{4q_{\omega}^{2}}{(1-q_{\omega}^{2})^{2}}\sigma^{2}+\frac{2q_{\omega}^{2}}{1-q_{\omega}^{2}}\sum_{\tau=1}^{t}\left(\frac{1+q_{\omega}^{2}}{2}\right)^{t-\tau}\mathbb{E}[\|\nabla
f(\theta_{\tau})\|^{2}].$

###### Proof.

We start by using Assumption 1, 2 on compression and Young’s inequality to get

$\displaystyle\|e_{t+1}\|^{2}$
$\displaystyle=\|\mathcal{Q}(g_{t}+e_{t}-\mathcal{C}(g_{t}+e_{t}))\|^{2}$
$\displaystyle\leq(1+\omega)^{2}q^{2}\|g_{t}+e_{t}\|^{2}$ $\displaystyle\leq
q_{\omega}^{2}(1+\rho)\|e_{t}\|^{2}+q_{\omega}^{2}\left(1+\frac{1}{\rho}\right)\|g_{t}\|^{2}$
$\displaystyle\leq\frac{1+q_{\omega}^{2}}{2}\|e_{t}\|^{2}+\frac{2q_{\omega}^{2}}{1-q_{\omega}^{2}}\|g_{t}\|^{2},$
(2)

where (2) is derived by choosing
$\rho=\frac{1-q_{\omega}^{2}}{2q_{\omega}^{2}}$ and the fact that
$q_{\omega}<1$. For the first claim we recursively apply the obtained
inequality and use bounded gradient Assumption 4. For the second claim,
initialization $e_{1}=0$ and the obtained recursion imply

$\displaystyle\mathbb{E}[\|e_{t+1}\|^{2}]$ $\displaystyle\leq$
$\displaystyle\frac{2q_{\omega}^{2}}{1-q_{\omega}^{2}}\sum_{\tau=1}^{t}\left(\frac{1+q_{\omega}^{2}}{2}\right)^{t-\tau}\mathbb{E}[\|g_{\tau}\|^{2}]$
$\displaystyle\overset{(\ref{var-decomp})}{\leq}$
$\displaystyle\frac{4q_{\omega}^{2}}{(1-q_{\omega}^{2})^{2}}\sigma^{2}+\frac{2q_{\omega}^{2}}{1-q_{\omega}^{2}}\sum_{\tau=1}^{t}\left(\frac{1+q_{\omega}^{2}}{2}\right)^{t-\tau}\mathbb{E}[\|\nabla
f(\theta_{\tau})\|^{2}],$

which concludes the lemma. ∎

######  Lemma 4.

Let $q_{\omega}=(1+\omega)q<1$. Under Assumptions 1-5, for all iterates $t$ we
have

$\displaystyle\|\zeta_{t}\|\leq\omega
q\left(1+\frac{2q_{\omega}}{1-q_{\omega}^{2}}\right)G,\quad\text{and}\quad\|\mathcal{Z}_{t}\|\leq\omega
q\left(1+\frac{2q_{\omega}}{1-q_{\omega}^{2}}\right)G.$

###### Proof.

Using the bounds defining compressors and Lemma 3, we get

$\displaystyle\|\zeta_{t}\|$
$\displaystyle=\|\mathcal{Q}(e_{t}+g_{t}-\tilde{g}_{t})-(e_{t}+g_{t}-\tilde{g}_{t})\|$
$\displaystyle\leq\omega\|e_{t}+g_{t}-\tilde{g}_{t}\|=\omega\|e_{t}+g_{t}-\mathcal{C}(e_{t}+g_{t})\|$
$\displaystyle\leq\omega q\|e_{t}+g_{t}\|$ $\displaystyle\leq\omega
q\|e_{t}\|+\omega q\|g_{t}\|$ $\displaystyle\leq\omega
q\left(1+\frac{2q_{\omega}}{1-q_{\omega}^{2}}\right)G.$

For the second claim, recall the definition of $\mathcal{Z}_{t}$ and apply
triangle inequality:

$\|\mathcal{Z}_{t}\|\leq(1-\beta_{1})\sum_{\tau=1}^{t}\beta^{t-\tau}\|\zeta_{\tau}\|\leq\omega
q\left(1+\frac{2q_{\omega}}{1-q_{\omega}^{2}}\right)G.$

∎

######  Lemma 5.

For the moving average error sequence $\mathcal{E}_{t}$, it holds that

$\displaystyle\sum_{t=1}^{T}\mathbb{E}[\|\mathcal{E}_{t}\|^{2}]\leq\frac{4Tq_{\omega}^{2}}{(1-q_{\omega}^{2})^{2}}\sigma^{2}+\frac{4q_{\omega}^{2}}{(1-q_{\omega}^{2})^{2}}\sum_{t=1}^{T}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}].$

###### Proof.

Let $e_{t,j}$ be the $j$-th coordinate of $e_{t}$ and denote

$K_{t}\mathrel{\mathop{:}}=\sum_{\tau=1}^{t}\left(\tfrac{1+q_{\omega}^{2}}{2}\right)^{t-\tau}\mathbb{E}[\|\nabla
f(\theta_{\tau})\|^{2}].$

Applying Jensen’s inequality and Lemma 3, we get

$\displaystyle\mathbb{E}[\|\mathcal{E}_{t}\|^{2}]$
$\displaystyle=\mathbb{E}\left[\left\|(1-\beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}e_{\tau}\right\|^{2}\right]$
$\displaystyle\leq(1-\beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}\mathbb{E}[\|e_{\tau}\|^{2}]$
$\displaystyle\leq\frac{4q_{\omega}^{2}}{(1-q_{\omega}^{2})^{2}}\sigma^{2}+\frac{2q_{\omega}^{2}(1-\beta_{1})}{(1-q_{\omega}^{2})}\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}K_{\tau},$

Summing over $t=1,\dots,T$ and using the technique of geometric series
summation leads to

$\displaystyle\sum_{t=1}^{T}\mathbb{E}[\|\mathcal{E}_{t}\|^{2}]$
$\displaystyle\leq\frac{4Tq_{\omega}^{2}}{(1-q_{\omega}^{2})^{2}}\sigma^{2}+\frac{2q_{\omega}^{2}(1-\beta_{1})}{(1-q_{\omega}^{2})}\sum_{t=1}^{T}\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}K_{\tau}$
$\displaystyle\leq\frac{4Tq_{\omega}^{2}}{(1-q_{\omega}^{2})^{2}}\sigma^{2}+\frac{2q_{\omega}^{2}}{(1-q_{\omega}^{2})}\sum_{t=1}^{T}K_{t}$
$\displaystyle=\frac{4Tq_{\omega}^{2}}{(1-q_{\omega}^{2})^{2}}\sigma^{2}+\frac{2q_{\omega}^{2}}{(1-q_{\omega}^{2})}\sum_{t=1}^{T}\sum_{\tau=1}^{t}\left(\frac{1+q_{\omega}^{2}}{2}\right)^{t-\tau}\mathbb{E}[\|\nabla
f(\theta_{\tau})\|^{2}]$
$\displaystyle\leq\frac{4Tq_{\omega}^{2}}{(1-q_{\omega}^{2})^{2}}\sigma^{2}+\frac{4q_{\omega}^{2}}{(1-q_{\omega}^{2})^{2}}\sum_{t=1}^{T}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}],$

The desired result is obtained. ∎

######  Lemma 6.

Let $q_{\omega}=(1+\omega)q<1$. Under Assumptions 1-5, for all iterates
$t\in[T]$ and coordinates $i\in[d]$, the following bound holds

$\hat{v}_{t,i}\leq\frac{4(1+q_{\omega}^{2})^{3}}{(1-q_{\omega}^{2})^{2}}G^{2}.$

###### Proof.

Lemma 3 and Assumption 4 imply

$\displaystyle\|\tilde{g}_{t}\|^{2}$
$\displaystyle=\|\mathcal{C}(g_{t}+e_{t})\|^{2}$
$\displaystyle\leq\|\mathcal{C}(g_{t}+e_{t})-(g_{t}+e_{t})+(g_{t}+e_{t})\|^{2}$
$\displaystyle\leq 2(q^{2}+1)\|g_{t}+e_{t}\|^{2}$ $\displaystyle\leq
4(q^{2}+1)\left(G^{2}+\frac{4q_{\omega}^{2}}{(1-q_{\omega}^{2})^{2}}G^{2}\right)$
$\displaystyle=\frac{4(1+q^{2})(1+q_{\omega}^{2})^{2}}{(1-q_{\omega}^{2})^{2}}G^{2}.$

It’s then easy to show by the updating rule of $\hat{v}_{t}$, there exists a
$j\in[t]$ such that $\hat{v}_{t,i}=v_{j,i}$. Then

$\displaystyle\hat{v}_{t,i}=(1-\beta_{2})\sum_{\tau=1}^{j}\beta_{2}^{j-\tau}\tilde{g}_{\tau,i}^{2}\leq\frac{4(1+q^{2})(1+q_{\omega}^{2})^{2}}{(1-q_{\omega}^{2})^{2}}G^{2},$

which concludes the claim. ∎

######  Lemma 7.

For
$D_{t}\mathrel{\mathop{:}}=\frac{1}{\sqrt{\hat{v}_{t-1}+\epsilon}}-\frac{1}{\sqrt{\hat{v}_{t}+\epsilon}}$
we have

$\displaystyle\sum_{t=1}^{T}\|D_{t}\|_{1}\leq\frac{d}{\sqrt{\epsilon}},\quad\sum_{t=1}^{T}\|D_{t}\|^{2}\leq\frac{d}{\epsilon}.$

###### Proof.

By the update rule, we have $\hat{v}_{t-1,i}\leq\hat{v}_{t,i}$ for any iterate
$t$ and coordinate $i\in[d]$. Therefore, by the initialization
$\hat{v}_{0}=0$, we get

$\sum_{t=1}^{T}\|D_{t}\|_{1}=\sum_{t=1}^{T}\sum_{i=1}^{d}\left(\frac{1}{\sqrt{\hat{v}_{t-1,i}+\epsilon}}-\frac{1}{\sqrt{\hat{v}_{t,i}+\epsilon}}\right)\\\
=\sum_{i=1}^{d}\left(\frac{1}{\sqrt{\hat{v}_{0,i}+\epsilon}}-\frac{1}{\sqrt{\hat{v}_{T,i}+\epsilon}}\right)\\\
\leq\frac{d}{\sqrt{\epsilon}}.$

For the sum of squared $l_{2}$ norms, note the fact that for $a\geq b>0$, it
holds that

$(a-b)^{2}\leq(a-b)(a+b)=a^{2}-b^{2}.$

Thus,

$\sum_{t=1}^{T}\|D_{t}\|^{2}=\sum_{t=1}^{T}\sum_{i=1}^{d}\left(\frac{1}{\sqrt{\hat{v}_{t-1,i}+\epsilon}}-\frac{1}{\sqrt{\hat{v}_{t,i}+\epsilon}}\right)^{2}\\\
\leq\sum_{t=1}^{T}\sum_{i=1}^{d}\left(\frac{1}{\hat{v}_{t-1,i}+\epsilon}-\frac{1}{\hat{v}_{t,i}+\epsilon}\right)\\\
\leq\frac{d}{\epsilon},$

which gives the desired result. ∎

###  D.2 Non-convex Analysis

Here we derive the convergence rate with fixed step-size $\eta$. The rate
shown in the main part can be obtained by plugging the expression of $\eta$.
Note that the constants in the rate can be slightly improved if we choose
$\eta=\min\\{\frac{\epsilon}{4LC_{0}},\frac{1}{\sqrt{T}}\\}$.

######  Theorem 3.

(Non-convex convergence rate) Let Assumptions 1, 2, 3, 4, 5 hold and
$q_{\omega}\mathrel{\mathop{:}}=(1+\omega)q<1$. Then, choosing any step-size
$\eta\leq\frac{\epsilon}{4LC_{0}}$, MicroAdam (Algorithm 4) satisfies

$\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}]\leq
2C_{0}\left(\frac{f(\theta_{1})-f^{*}}{T\eta}+\frac{\eta
L\sigma^{2}}{\epsilon}+\frac{\eta LC_{2}^{2}G^{2}}{\epsilon}\right.\\\
+\left.\frac{\eta^{2}LC_{0}C_{1}^{2}G^{2}}{\epsilon^{2}}+\frac{(1+C_{1})G^{2}d}{T\sqrt{\epsilon}}+\frac{\eta(1+2C_{1})C_{1}LG^{2}d}{T\epsilon}\right),$

with constants
$C_{0}\mathrel{\mathop{:}}=\sqrt{\frac{4(1+q_{\omega}^{2})^{3}}{(1-q_{\omega}^{2})^{2}}G^{2}+\epsilon},\,C_{1}\mathrel{\mathop{:}}=\frac{\beta_{1}}{1-\beta_{1}}(1+C_{2})+\frac{2q_{\omega}}{1-q_{\omega}^{2}},\,C_{2}\mathrel{\mathop{:}}=\omega
q(1+\frac{2q_{\omega}}{1-q_{\omega}^{2}})$.

###### Proof.

Similar to the proof of Comp-AMS [Li et al., 2022], we define two virtual
iterates $\theta^{\prime}_{t}$ and $x_{t}$.

$\displaystyle\theta_{t+1}^{\prime}$
$\displaystyle\mathrel{\mathop{:}}=\theta_{t+1}-\eta\frac{\mathcal{E}_{t+1}}{\sqrt{\hat{v}_{t}+\epsilon}}$
$\displaystyle x_{t+1}$
$\displaystyle\mathrel{\mathop{:}}=\theta_{t+1}^{\prime}-\eta\frac{\beta_{1}}{1-\beta_{1}}\frac{m_{t}^{\prime}+\mathcal{Z}_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}.$

Then, we derive the recurrence relation for each sequence as follows:

$\displaystyle\theta_{t+1}^{\prime}$
$\displaystyle=\theta_{t+1}-\eta\frac{\mathcal{E}_{t+1}}{\sqrt{\hat{v}_{t}+\epsilon}}$
$\displaystyle=\theta_{t}-\eta\frac{(1-\beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}\tilde{g}_{\tau}+(1-\beta_{1})\sum_{\tau=1}^{t+1}\beta_{1}^{t+1-\tau}e_{\tau}}{\sqrt{\hat{v}_{t}+\epsilon}}$
$\displaystyle=\theta_{t}-\eta\frac{(1-\beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}(\tilde{g}_{\tau}+e_{\tau+1})+(1-\beta)\beta_{1}^{t}e_{1}}{\sqrt{\hat{v}_{t}+\epsilon}}$
$\displaystyle=\theta_{t}-\eta\frac{(1-\beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}(g_{\tau}+e_{\tau}+\zeta_{\tau})}{\sqrt{\hat{v}_{t}+\epsilon}}$
$\displaystyle=\theta_{t}-\eta\frac{(1-\beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}e_{\tau}}{\sqrt{\hat{v}_{t}+\epsilon}}-\eta\frac{m_{t}^{\prime}}{\sqrt{\hat{v}_{t}+\epsilon}}-\eta\frac{\mathcal{Z}_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}$
$\displaystyle=\theta_{t}-\eta\frac{\mathcal{E}_{t}}{\sqrt{\hat{v}_{t-1}+\epsilon}}-\eta\frac{m_{t}^{\prime}}{\sqrt{\hat{v}_{t}+\epsilon}}+\eta\left(\frac{1}{\sqrt{\hat{v}_{t-1}+\epsilon}}-\frac{1}{\sqrt{\hat{v}_{t}+\epsilon}}\right)\mathcal{E}_{t}-\eta\frac{\mathcal{Z}_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}$
$\displaystyle=\theta_{t}^{\prime}-\eta\frac{m_{t}^{\prime}}{\sqrt{\hat{v}_{t}+\epsilon}}+\eta\left(\frac{1}{\sqrt{\hat{v}_{t-1}+\epsilon}}-\frac{1}{\sqrt{\hat{v}_{t}+\epsilon}}\right)\mathcal{E}_{t}-\eta\frac{\mathcal{Z}_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}$
$\displaystyle=\theta_{t}^{\prime}-\eta\frac{m_{t}^{\prime}+\mathcal{Z}_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}+\eta
D_{t}\mathcal{E}_{t},$

where we used the fact that $\tilde{g}_{t}+e_{{t+1}}=g_{t}+e_{t}+\zeta_{t}$
with quantization noise $\zeta_{t}$, and $e_{0}=0$ at initialization. Next,
for the $x_{t}$ iterates we have

$\displaystyle x_{t+1}$
$\displaystyle=\theta_{t+1}^{\prime}-\eta\frac{\beta_{1}}{1-\beta_{1}}\frac{m_{t}^{\prime}+\mathcal{Z}_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}$
$\displaystyle=\theta_{t}^{\prime}-\eta\frac{m_{t}^{\prime}+\mathcal{Z}_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}-\eta\frac{\beta_{1}}{1-\beta_{1}}\frac{m_{t}^{\prime}+\mathcal{Z}_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}+\eta
D_{t}\mathcal{E}_{t}$
$\displaystyle=\theta_{t}^{\prime}-\eta\frac{\beta_{1}(m_{t-1}^{\prime}+\mathcal{Z}_{t-1})+(1-\beta_{1})(g_{t}+\zeta_{t})+\frac{\beta_{1}^{2}}{1-\beta_{1}}(m_{t-1}^{\prime}+\mathcal{Z}_{t-1})+\beta_{1}(g_{t}+\zeta_{t})}{\sqrt{\hat{v}_{t}+\epsilon}}$
$\displaystyle\quad+\eta D_{t}\mathcal{E}_{t}$
$\displaystyle=\theta_{t}^{\prime}-\eta\frac{\beta_{1}}{1-\beta_{1}}\frac{m_{t-1}^{\prime}+\mathcal{Z}_{t-1}}{\sqrt{\hat{v}_{t}+\epsilon}}-\eta\frac{g_{t}+\zeta_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}+\eta
D_{t}\mathcal{E}_{t}$
$\displaystyle=x_{t}-\eta\frac{g_{t}+\zeta_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}+\eta\frac{\beta_{1}}{1-\beta_{1}}D_{t}(m_{t-1}^{\prime}+\mathcal{Z}_{t-1})+\eta
D_{t}\mathcal{E}_{t}.$

Next we apply smoothness (Assumption 3) of the loss function $f$ over the
iterates $x_{t}$. From the gradient Lipschitzness we have

$\displaystyle f(x_{t+1})\leq f(x_{t})+\langle\nabla
f(x_{t}),x_{t+1}-x_{t}\rangle+\frac{L}{2}\|x_{t+1}-x_{t}\|^{2}.$

Due to unbiasedness of the compressor $\mathcal{Q}$ (see Assumption 2), we
have $\mathbb{E}[\zeta_{t}|g_{t},\tilde{g}_{t},e_{t},\hat{v}_{t}]=0$. Taking
expectation, we obtain

$\displaystyle\mathbb{E}[f(x_{t+1})]-\mathbb{E}[f(x_{t})]$ $\displaystyle\leq$
$\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla
f(x_{t}),\frac{g_{t}+\zeta_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}\right\rangle\right]$
(4) $\displaystyle+\eta\mathbb{E}\left[\left\langle\nabla
f(x_{t}),\frac{\beta_{1}}{1-\beta_{1}}D_{t}(m_{t-1}^{\prime}+\mathcal{Z}_{t-1})+D_{t}\mathcal{E}_{t}\right\rangle\right]$
$\displaystyle+\frac{\eta^{2}L}{2}\mathbb{E}\left[\left\|\frac{g_{t}+\zeta_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}-\frac{\beta_{1}}{1-\beta_{1}}D_{t}(m_{t-1}^{\prime}+\mathcal{Z}_{t-1})-D_{t}\mathcal{E}_{t}\right\|^{2}\right]$
$\displaystyle=$
$\displaystyle\underbrace{-\eta\mathbb{E}\left[\left\langle\nabla
f(\theta_{t}),\frac{g_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}\right\rangle\right]}_{I}$
$\displaystyle+\underbrace{\eta\mathbb{E}\left[\left\langle\nabla
f(x_{t}),\frac{\beta_{1}}{1-\beta_{1}}D_{t}(m_{t-1}^{\prime}+\mathcal{Z}_{t-1})+D_{t}\mathcal{E}_{t}\right\rangle\right]}_{II}$
$\displaystyle+\underbrace{\frac{\eta^{2}L}{2}\mathbb{E}\left[\left\|\frac{g_{t}+\zeta_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}-\frac{\beta_{1}}{1-\beta_{1}}D_{t}(m_{t-1}^{\prime}+\mathcal{Z}_{t-1})-D_{t}\mathcal{E}_{t}\right\|^{2}\right]}_{III}$
$\displaystyle+\underbrace{\eta\mathbb{E}\left[\left\langle\nabla
f(\theta_{t})-\nabla
f(x_{t}),\frac{g_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}\right\rangle\right]}_{IV},$

In the following, we bound all the four terms highlighted above.

Bounding term I. We have

$\displaystyle I$ $\displaystyle=$
$\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla
f(\theta_{t}),\frac{g_{t}}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right\rangle\right]-\eta\mathbb{E}\left[\left\langle\nabla
f(\theta_{t}),\left(\frac{1}{\sqrt{\hat{v}_{t}+\epsilon}}-\frac{1}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right)g_{t}\right\rangle\right]$
(5) $\displaystyle\leq$ $\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla
f(\theta_{t}),\frac{\nabla
f(\theta_{t})}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right\rangle\right]+\eta
G^{2}\mathbb{E}[\|D_{t}\|].$ $\displaystyle\leq$
$\displaystyle-\frac{\eta}{\sqrt{\frac{4(1+q_{\omega}^{2})^{3}}{(1-q_{\omega}^{2})^{2}}G^{2}+\epsilon}}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}]+\eta G^{2}\mathbb{E}[\|D_{t}\|_{1}],$

where we use Assumption 4, Lemma 6 and the fact that $l_{2}$ norm is no larger
than $l_{1}$ norm.

Bounding term II. By the definition of $\mathcal{E}_{t}$ and
$\mathcal{Z}_{t}$, we know that

$\displaystyle\|\mathcal{E}_{t}\|$ $\displaystyle\leq$
$\displaystyle(1-\beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}\|e_{t}\|\leq\frac{2q_{\omega}}{1-q_{\omega}^{2}}G,$
$\displaystyle\|\mathcal{Z}_{t}\|$ $\displaystyle\leq$
$\displaystyle(1-\beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}\|\zeta_{t}\|\leq\omega
q\left(1+\frac{2q_{\omega}}{1-q_{\omega}^{2}}\right)G.$

Then we have

$\displaystyle II$ $\displaystyle\leq\eta\mathbb{E}\left[\left\langle\nabla
f(\theta_{t}),\frac{\beta_{1}}{1-\beta_{1}}D_{t}(m_{t-1}^{\prime}+\mathcal{Z}_{t-1})+D_{t}\mathcal{E}_{t}\right\rangle\right]$
$\displaystyle\quad+\eta\mathbb{E}\left[\left\langle\nabla f(x_{t})-\nabla
f(\theta_{t}),\frac{\beta_{1}}{1-\beta_{1}}D_{t}(m_{t-1}^{\prime}+\mathcal{Z}_{t-1})+D_{t}\mathcal{E}_{t}\right\rangle\right]$
$\displaystyle\leq\eta\mathbb{E}\left[\|\nabla
f(\theta_{t})\|\left\|\frac{\beta_{1}}{1-\beta_{1}}D_{t}(m_{t-1}^{\prime}+\mathcal{Z}_{t-1})+D_{t}\mathcal{E}_{t}\right\|\right]$
$\displaystyle\quad+\eta^{2}L\mathbb{E}\left[\left\|\frac{\frac{\beta_{1}}{1-\beta_{1}}m_{t-1}^{\prime}+\frac{\beta_{1}}{1-\beta_{1}}\mathcal{Z}_{t-1}+\mathcal{E}_{t}}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right\|\left\|\frac{\beta_{1}}{1-\beta_{1}}D_{t}(m_{t-1}^{\prime}+\mathcal{Z}_{t-1})+D_{t}\mathcal{E}_{t}\right\|\right]$
$\displaystyle\leq\eta
C_{1}G^{2}\mathbb{E}[\|D_{t}\|_{1}]+\frac{\eta^{2}C_{1}^{2}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|D_{t}\|_{1}],$
(6)

where $C_{1}=\frac{\beta_{1}}{1-\beta_{1}}\left(1+\omega
q\left(1+\frac{2q_{\omega}}{1-q_{\omega}^{2}}\right)\right)+\frac{2q_{\omega}}{1-q_{\omega}^{2}}$.
The second inequality is because of smoothness of $f(\theta)$, and the last
inequality is due to Lemma 3, Assumption 4 and the property of norms.

Bounding term III. This term can be bounded as follows:

$\displaystyle III$
$\displaystyle\leq\eta^{2}L\mathbb{E}\left[\left\|\frac{g_{t}+\zeta_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}\right\|^{2}\right]+\eta^{2}L\mathbb{E}\left[\left\|\frac{\beta_{1}}{1-\beta_{1}}D_{t}(m_{t-1}^{\prime}+\mathcal{Z}_{t-1})-D_{t}\mathcal{E}_{t}\right\|^{2}\right]$
$\displaystyle\leq\frac{2\eta^{2}L}{\epsilon}\mathbb{E}[\|g_{t}-\nabla
f(\theta_{t})+\nabla
f(\theta_{t})\|^{2}]+\frac{2\eta^{2}L}{\epsilon}\mathbb{E}[\|\zeta_{t}\|^{2}]$
(7)
$\displaystyle\quad+\eta^{2}L\mathbb{E}\left[\left\|D_{t}\left(\frac{\beta_{1}}{1-\beta_{1}}m_{t-1}^{\prime}+\frac{\beta_{1}}{1-\beta_{1}}\mathcal{Z}_{t-1}-\mathcal{E}_{t}\right)\right\|^{2}\right]$
$\displaystyle\leq\frac{2\eta^{2}L}{\epsilon}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}]+\frac{2\eta^{2}L\sigma^{2}}{\epsilon}+\frac{2\eta^{2}L}{\epsilon}\omega^{2}q^{2}\left(1+\frac{2q}{1-q^{2}}\right)^{2}G^{2}+\eta^{2}C_{1}^{2}LG^{2}\mathbb{E}[\|D_{t}\|^{2}]$
$\displaystyle\leq\frac{2\eta^{2}L}{\epsilon}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}]+\frac{2\eta^{2}L(\sigma^{2}+C_{2}^{2}G^{2})}{\epsilon}+\eta^{2}C_{1}^{2}LG^{2}\mathbb{E}[\|D_{t}\|^{2}],$
(8)

where $C_{2}=\omega q(1+\frac{2q}{1-q^{2}})$ and we used Assumption 5 that
$g_{t}$ is unbiased with bounded variance $\sigma^{2}$.

Bounding term IV. We have

$\displaystyle IV$ $\displaystyle=\eta\mathbb{E}\left[\left\langle\nabla
f(\theta_{t})-\nabla
f(x_{t}),\frac{g_{t}}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right\rangle\right]$ (9)
$\displaystyle\quad+\eta\mathbb{E}\left[\left\langle\nabla
f(\theta_{t})-\nabla
f(x_{t}),\left(\frac{1}{\sqrt{\hat{v}_{t}+\epsilon}}-\frac{1}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right)g_{t}\right\rangle\right]$
$\displaystyle\leq\eta\mathbb{E}\left[\left\langle\nabla f(\theta_{t})-\nabla
f(x_{t}),\frac{\nabla
f(\theta_{t})}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right\rangle\right]$ (10)
$\displaystyle\quad+\eta^{2}L\mathbb{E}\left[\left\|\frac{\frac{\beta_{1}}{1-\beta_{1}}m_{t-1}^{\prime}+\frac{\beta_{1}}{1-\beta_{1}}\mathcal{Z}_{t-1}+\mathcal{E}_{t}}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right\|\|D_{t}g_{t}\|\right]$
$\displaystyle\overset{(a)}{\leq}\frac{\eta\rho}{2\epsilon}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}]+\frac{\eta}{2\rho}\mathbb{E}[\|\nabla
f(\theta_{t})-\nabla
f(x_{t})\|^{2}]+\frac{\eta^{2}C_{1}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|D_{t}\|]$
$\displaystyle\overset{(b)}{\leq}\frac{\eta\rho}{2\epsilon}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}]+\frac{\eta^{3}L}{2\rho}\mathbb{E}\left[\left\|\frac{\frac{\beta_{1}}{1-\beta_{1}}m_{t-1}^{\prime}+\frac{\beta_{1}}{1-\beta_{1}}\mathcal{Z}_{t-1}+\mathcal{E}_{t}}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right\|^{2}\right]+\frac{\eta^{2}C_{1}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|D_{t}\|_{1}]$
$\displaystyle\leq\frac{\eta\rho}{2\epsilon}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}]+\frac{\eta^{3}L}{2\rho}\frac{C_{1}^{2}G^{2}}{\epsilon}+\frac{\eta^{2}LC_{1}G^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|D_{t}\|_{1}],$
(11)

where (a) is due to Young’s inequality and (b) is based on Assumption 3. Now
integrating (5), (6), (8), (11) into (4),

$\displaystyle I$ $\displaystyle\leq$
$\displaystyle-\frac{\eta}{C_{0}}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}]+\eta
G^{2}\mathbb{E}[\|D_{t}\|_{1}]$ $\displaystyle II$ $\displaystyle\leq$
$\displaystyle\eta
C_{1}G^{2}\mathbb{E}[\|D_{t}\|_{1}]+\frac{\eta^{2}C_{1}^{2}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|D_{t}\|_{1}]$
$\displaystyle III$ $\displaystyle\leq$
$\displaystyle\frac{\eta^{2}L}{\epsilon}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}]+\frac{\eta^{2}L(\sigma^{2}+C_{2}^{2}G^{2})}{\epsilon}+\eta^{2}C_{1}^{2}LG^{2}\mathbb{E}[\|D_{t}\|^{2}]$
$\displaystyle IV$ $\displaystyle\leq$
$\displaystyle\frac{\eta\rho}{2\epsilon}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}]+\frac{\eta^{3}L}{2\rho}\frac{C_{1}^{2}G^{2}}{\epsilon}+\frac{\eta^{2}LC_{1}G^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|D_{t}\|_{1}],$

and taking the telescoping summation over $t=1,\dots,T$, we obtain

$\displaystyle\mathbb{E}[f(x_{T+1})-f(x_{1})]$
$\displaystyle\leq\left(-\frac{\eta}{C_{0}}+\frac{\eta^{2}L}{\epsilon}+\frac{\eta\rho}{2\epsilon}\right)\sum_{t=1}^{T}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}]+\frac{T\eta^{2}L(\sigma^{2}+C_{2}^{2}G^{2})}{\epsilon}+\frac{T\eta^{3}LC_{1}^{2}G^{2}}{2\rho\epsilon}$
$\displaystyle\quad+\left(\eta(1+C_{1})G^{2}+\frac{\eta^{2}(1+C_{1})C_{1}LG^{2}}{\sqrt{\epsilon}}\right)\sum_{t=1}^{T}\mathbb{E}[\|D_{t}\|_{1}]+\eta^{2}C_{1}^{2}LG^{2}\sum_{t=1}^{T}\mathbb{E}[\|D_{t}\|^{2}].$

Setting $\eta\leq\frac{\epsilon}{4LC_{0}}$ and choosing
$\rho=\frac{\epsilon}{2C_{0}}$, we further arrive at

$\displaystyle\mathbb{E}[f(x_{T+1})-f(x_{1})]$
$\displaystyle\leq-\frac{\eta}{2C_{0}}\sum_{t=1}^{T}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}]+\frac{T\eta^{2}L(\sigma^{2}+C_{2}^{2}G^{2})}{\epsilon}$
$\displaystyle\quad+\frac{T\eta^{3}LC_{0}C_{1}^{2}G^{2}}{\epsilon^{2}}+\frac{\eta(1+C_{1})G^{2}d}{\sqrt{\epsilon}}+\frac{\eta^{2}(1+2C_{1})C_{1}LG^{2}d}{\epsilon}.$

where the inequality follows from Lemma 7. Re-arranging terms, we get that

$\displaystyle\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}]$ $\displaystyle\leq
2C_{0}\left(\frac{\mathbb{E}[f(x_{1})-f(x_{T+1})]}{T\eta}+\frac{\eta
L(\sigma^{2}+C_{2}^{2}G^{2})}{\epsilon}+\frac{\eta^{2}LC_{0}C_{1}^{2}G^{2}}{\epsilon^{2}}\right)$
$\displaystyle\quad+2C_{0}\left(\frac{(1+C_{1})G^{2}d}{T\sqrt{\epsilon}}+\frac{\eta(1+2C_{1})C_{1}LG^{2}d}{T\epsilon}\right)$
$\displaystyle\leq 2C_{0}\left(\frac{f(\theta_{1})-f^{*}}{T\eta}+\frac{\eta
L(\sigma^{2}+C_{2}^{2}G^{2})}{\epsilon}+\frac{\eta^{2}LC_{0}C_{1}^{2}G^{2}}{\epsilon^{2}}\right)$
$\displaystyle\quad+2C_{0}\left(\frac{(1+C_{1})G^{2}d}{T\sqrt{\epsilon}}+\frac{\eta(1+2C_{1})C_{1}LG^{2}d}{T\epsilon}\right),$

where in the last inequality we used $x_{1}=\theta_{1}$ and the lower bound
$f^{*}\leq f(\theta)$ for all $\theta\in\mathbb{R}^{d}$. ∎

###  D.3 Analysis Under PL Condition

As in the non-convex analysis, here we derive the convergence rate with fixed
step-size $\eta$. The rate shown in the main part can be obtained by plugging
the expression of $\eta$. Also, the constants in the rate can be slightly
improved if we choose $\eta=\min\\{\frac{\epsilon}{4LC_{0}},\frac{C_{0}\log
T}{\mu T}\\}$.

######  Theorem 4.

(Convergence rate under PL) Let Assumptions 1, 2, 3, 4, 5 and 6 hold, and
$q_{\omega}<1$. Then, choosing any step-size
$\eta\leq\frac{\epsilon}{4LC_{0}}$, MicroAdam (Algorithm 4) satisfies

$\displaystyle\mathbb{E}[f(\theta_{T+1})]-f^{*}\leq\bigl{(}1-\frac{\eta\mu}{C_{0}}\bigr{)}^{T}(f(\theta_{1})-f^{*})+\eta\left(\frac{LC_{0}\sigma^{2}+LC_{0}(C_{1}+C_{2}^{2})G^{2}}{\mu\epsilon}+\frac{(1+C_{1})G^{2}d+C_{1}G^{2}}{\sqrt{\epsilon}}\right)$
$\displaystyle\hfill+\,\eta^{2}\left(\frac{3L^{2}C_{0}C^{2}_{1}G^{2}}{2\mu\epsilon^{3/2}}+\frac{(1+2C_{1})C_{1}LG^{2}d}{\epsilon}+\frac{LC_{1}^{2}G^{2}}{2\epsilon}\right).$

###### Proof.

We start from descent lemma

$\displaystyle\mathbb{E}[f(x_{t+1})]-f(x_{t})$
$\displaystyle\leq-\eta\mathbb{E}\left[\left\langle\nabla
f(x_{t}),\frac{g_{t}+\zeta_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}\right\rangle\right]+\eta\mathbb{E}\left[\left\langle\nabla
f(x_{t}),\frac{\beta_{1}}{1-\beta_{1}}D_{t}(m_{t-1}^{\prime}+\mathcal{Z}_{t-1})+D_{t}\mathcal{E}_{t}\right\rangle\right]$
$\displaystyle\hskip
108.405pt+\frac{\eta^{2}L}{2}\mathbb{E}\left[\left\|\frac{g_{t}+\zeta_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}-\frac{\beta_{1}}{1-\beta_{1}}D_{t}(m_{t-1}^{\prime}+\mathcal{Z}_{t-1})-D_{t}\mathcal{E}_{t}\right\|^{2}\right]$
$\displaystyle=\underbrace{-\eta\mathbb{E}\left[\left\langle\nabla
f(x_{t}),\frac{g_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}\right\rangle\right]}_{I^{\prime}}+\underbrace{\eta\mathbb{E}\left[\left\langle\nabla
f(x_{t}),\frac{\beta_{1}}{1-\beta_{1}}D_{t}(m_{t-1}^{\prime}+\mathcal{Z}_{t-1})+D_{t}\mathcal{E}_{t}\right\rangle\right]}_{II}$
$\displaystyle\hskip
108.405pt+\underbrace{\frac{\eta^{2}L}{2}\mathbb{E}\left[\left\|\frac{g_{t}+\zeta_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}-\frac{\beta_{1}}{1-\beta_{1}}D_{t}(m_{t-1}^{\prime}+\mathcal{Z}_{t-1})-D_{t}\mathcal{E}_{t}\right\|^{2}\right]}_{III}.$
(12)

We bound part $II$ and part $III$ in the same way as it was done in the non-
convex analysis. We now provide a bound for part $I^{\prime}$:

$\displaystyle I^{\prime}=$
$\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla
f(x_{t}),\frac{g_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}\right\rangle\right]$
$\displaystyle=$ $\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla
f(x_{t}),\frac{g_{t}}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right\rangle\right]-\eta\mathbb{E}\left[\left\langle\nabla
f(x_{t}),g_{t}\left(\frac{1}{\sqrt{\hat{v}_{t}+\epsilon}}-\frac{1}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right)\right\rangle\right]$
$\displaystyle=$ $\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla
f(x_{t}),\frac{g_{t}-\nabla f(x_{t})+\nabla
f(x_{t})}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right\rangle\right]$
$\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla
f(x_{t}),g_{t}\left(\frac{1}{\sqrt{\hat{v}_{t}+\epsilon}}-\frac{1}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right)\right\rangle\right]$
$\displaystyle=$ $\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla
f(x_{t}),\frac{\nabla
f(x_{t})}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right\rangle\right]-\eta\mathbb{E}\left[\left\langle\nabla
f(x_{t}),\frac{g_{t}-\nabla
f(x_{t})}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right\rangle\right]$
$\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla
f(x_{t}),g_{t}\left(\frac{1}{\sqrt{\hat{v}_{t}+\epsilon}}-\frac{1}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right)\right\rangle\right].$

We further expand and bound this equation as follows:

$\displaystyle I^{\prime}\leq$
$\displaystyle-\frac{\eta}{C_{0}}\mathbb{E}\left[\left\|\nabla
f(x_{t})\right\|^{2}\right]$
$\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla f(x_{t})-\nabla
f(\theta_{t})+\nabla
f(\theta_{t}),\frac{1}{\sqrt{\hat{v}_{t-1}+\epsilon}}\left(\nabla
f(\theta_{t})-\nabla f(x_{t})\right)\right\rangle\right]$
$\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla f(x_{t})-\nabla
f(\theta_{t})+\nabla
f(\theta_{t}),\left(\frac{1}{\sqrt{\hat{v}_{t}+\epsilon}}-\frac{1}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right)g_{t}\right\rangle\right]$
$\displaystyle=$
$\displaystyle-\frac{\eta}{C_{0}}\mathbb{E}\left[\left\|\nabla
f(x_{t})\right\|^{2}\right]$
$\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla f(x_{t})-\nabla
f(\theta_{t}),\frac{1}{\sqrt{\hat{v}_{t-1}+\epsilon}}\left(\nabla
f(\theta_{t})-\nabla f(x_{t})\right)\right\rangle\right]$
$\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla
f(\theta_{t}),\frac{1}{\sqrt{\hat{v}_{t-1}+\epsilon}}\left(\nabla
f(\theta_{t})-\nabla f(x_{t})\right)\right\rangle\right]$
$\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla f(x_{t})-\nabla
f(\theta_{t}),\left(\frac{1}{\sqrt{\hat{v}_{t}+\epsilon}}-\frac{1}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right)g_{t}\right\rangle\right]$
$\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla
f(\theta_{t}),\left(\frac{1}{\sqrt{\hat{v}_{t}+\epsilon}}-\frac{1}{\sqrt{\hat{v}_{t-1}+\epsilon}}\right)g_{t}\right\rangle\right]$
$\displaystyle\leq$
$\displaystyle-\frac{\eta}{C_{0}}\mathbb{E}\left[\left\|\nabla
f(x_{t})\right\|^{2}\right]+\frac{\eta}{\sqrt{\epsilon}}\mathbb{E}\left[\left\|\nabla
f(x_{t})-f(\theta_{t})\right\|^{2}\right]$
$\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla
f(\theta_{t}),\frac{1}{\sqrt{\hat{v}_{t-1}+\epsilon}}\left(\nabla
f(\theta_{t})-\nabla f(x_{t})\right)\right\rangle\right]$
$\displaystyle+\eta\mathbb{E}\left[\left\langle\nabla f(x_{t})-\nabla
f(\theta_{t}),D_{t}g_{t}\right\rangle\right]+\eta\mathbb{E}\left[\left\langle\nabla
f(\theta_{t}),D_{t}g_{t}\right\rangle\right].$

Next, we use the Cauchy–Schwartz inequality to bound inner products above,
$L$-smoothness inequality to bound $\|\nabla f(x_{t})-\nabla
f(\theta_{t})\|\leq L\|x_{t}-\theta_{t}\|\leq\frac{\eta
LC_{1}G}{\sqrt{\epsilon}}$, and the inequality
$-\|a\|^{2}\leq-\frac{1}{2}\|b\|^{2}+\|a-b\|^{2}$ for the first term:

$\displaystyle I^{\prime}\leq$
$\displaystyle-\frac{\eta}{C_{0}}\mathbb{E}\left[\left\|\nabla
f(x_{t})\right\|^{2}\right]+\frac{\eta}{\sqrt{\epsilon}}\mathbb{E}\left[\left\|\nabla
f(x_{t})-f(\theta_{t})\right\|^{2}\right]+\frac{\eta
G}{\sqrt{\epsilon}}\mathbb{E}\left[\left\|\nabla f(\theta_{t})-\nabla
f(x_{t})\right\|\right]$ $\displaystyle+\eta G\mathbb{E}\left[\|\nabla
f(x_{t})-\nabla f(\theta_{t})\|\|D_{t}\|\right]+\eta
G^{2}\mathbb{E}\left[\|D_{t}\|\right]$ $\displaystyle\leq$
$\displaystyle-\frac{\eta}{2C_{0}}\mathbb{E}\left[\left\|\nabla
f(x_{t})\right\|^{2}\right]-\frac{\eta}{2C_{0}}\mathbb{E}\left[\left\|\nabla
f(x_{t})\right\|^{2}\right]+\frac{\eta}{\sqrt{\epsilon}}\frac{\eta^{2}L^{2}C^{2}_{1}G^{2}}{\epsilon}+\frac{\eta^{2}LC_{1}G^{2}}{\epsilon}$
$\displaystyle+\eta G\frac{\eta
LC_{1}G}{\sqrt{\epsilon}}\mathbb{E}\left[\|D_{t}\|_{1}\right]+\eta
G^{2}\mathbb{E}\left[\|D_{t}\|_{1}\right]$ $\displaystyle\leq$
$\displaystyle-\frac{\eta}{2C_{0}}\mathbb{E}\left[\left\|\nabla
f(x_{t})\right\|^{2}\right]-\frac{\eta}{4C_{0}}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}]+\frac{\eta}{2C_{0}}\mathbb{E}[\|\nabla f(x_{t})-\nabla
f(\theta_{t})\|^{2}]$
$\displaystyle+\frac{\eta^{3}L^{2}C^{2}_{1}G^{2}}{\epsilon^{3/2}}+\frac{\eta^{2}LC_{1}G^{2}}{\epsilon}+\frac{\eta^{2}LC_{1}G^{2}}{\sqrt{\epsilon}}\mathbb{E}\left[\|D_{t}\|_{1}\right]+\eta
G^{2}\mathbb{E}\left[\|D_{t}\|_{1}\right]$ $\displaystyle\leq$
$\displaystyle-\frac{\eta}{2C_{0}}\mathbb{E}\left[\left\|\nabla
f(x_{t})\right\|^{2}\right]-\frac{\eta}{4C_{0}}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}]$
$\displaystyle+\frac{3\eta^{3}L^{2}C^{2}_{1}G^{2}}{2\epsilon^{3/2}}+\frac{\eta^{2}LC_{1}G^{2}}{\epsilon}+\frac{\eta^{2}LC_{1}G^{2}}{\sqrt{\epsilon}}\mathbb{E}\left[\|D_{t}\|_{1}\right]+\eta
G^{2}\mathbb{E}\left[\|D_{t}\|_{1}\right].$

Plugging the obtained bound for $I^{\prime}$ with previously obtained bounds
for $II$ and $III$

$\displaystyle II$ $\displaystyle\leq$ $\displaystyle\eta
C_{1}G^{2}\mathbb{E}[\|D_{t}\|_{1}]+\frac{\eta^{2}C_{1}^{2}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|D_{t}\|_{1}]$
$\displaystyle III$ $\displaystyle\leq$
$\displaystyle\frac{\eta^{2}L}{\epsilon}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}]+\frac{\eta^{2}L(\sigma^{2}+C_{2}^{2}G^{2})}{\epsilon}+\eta^{2}C_{1}^{2}LG^{2}\mathbb{E}[\|D_{t}\|^{2}]$

into (12) and using the step-size bound $\eta\leq\frac{\epsilon}{4LC_{0}}$ we
get

$\displaystyle\mathbb{E}[f(x_{t+1})]-\mathbb{E}[f(x_{t})]\leq$
$\displaystyle-\frac{\eta}{2C_{0}}\mathbb{E}\left[\left\|\nabla
f(x_{t})\right\|^{2}\right]-\frac{\eta}{4C_{0}}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}]$
$\displaystyle+\frac{3\eta^{3}L^{2}C^{2}_{1}G^{2}}{2\epsilon^{3/2}}+\frac{\eta^{2}LC_{1}G^{2}}{\epsilon}+\frac{\eta^{2}LC_{1}G^{2}}{\sqrt{\epsilon}}\mathbb{E}\left[\|D_{t}\|_{1}\right]$
$\displaystyle+\eta
C_{1}G^{2}\mathbb{E}[\|D_{t}\|_{1}]+\frac{\eta^{2}C_{1}^{2}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|D_{t}\|_{1}]+\eta
G^{2}\mathbb{E}\left[\|D_{t}\|_{1}\right]$
$\displaystyle+\frac{\eta^{2}L}{\epsilon}\mathbb{E}[\|\nabla
f(\theta_{t})\|^{2}]+\frac{\eta^{2}L(\sigma^{2}+C_{2}^{2}G^{2})}{\epsilon}+\eta^{2}C_{1}^{2}LG^{2}\mathbb{E}[\|D_{t}\|^{2}]$
$\displaystyle\leq$
$\displaystyle-\frac{\eta}{2C_{0}}\mathbb{E}\left[\left\|\nabla
f(x_{t})\right\|^{2}\right]+\frac{\eta^{2}L\sigma^{2}}{\epsilon}+\frac{\eta^{2}L(C_{1}+C_{2}^{2})G^{2}}{\epsilon}$
$\displaystyle+\eta(1+C_{1})G^{2}\mathbb{E}[\|D_{t}\|_{1}]+\frac{3\eta^{3}L^{2}C^{2}_{1}G^{2}}{2\epsilon^{3/2}}$
$\displaystyle+\frac{\eta^{2}(1+C_{1})C_{1}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|D_{t}\|_{1}]+\eta^{2}C_{1}^{2}LG^{2}\mathbb{E}[\|D_{t}\|^{2}]$
$\displaystyle\leq$
$\displaystyle-\frac{\eta\mu}{C_{0}}(\mathbb{E}[f(x_{t})]-f^{*})+\frac{\eta^{2}L\sigma^{2}}{\epsilon}+\frac{\eta^{2}L(C_{1}+C_{2}^{2})G^{2}}{\epsilon}$
$\displaystyle+\eta(1+C_{1})G^{2}\mathbb{E}[\|D_{t}\|_{1}]+\frac{\eta^{2}(1+C_{1})C_{1}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|D_{t}\|_{1}]$
$\displaystyle+\eta^{2}C_{1}^{2}LG^{2}\mathbb{E}[\|D_{t}\|^{2}]+\frac{3\eta^{3}L^{2}C^{2}_{1}G^{2}}{2\epsilon^{3/2}},$

where in the last inequality we applied PL condition from Assumption 6. After
some reshuffling of the terms, we obtain the following recursion:

$\displaystyle\mathbb{E}[f(x_{t+1})]-f^{*}\leq$
$\displaystyle\left(1-\frac{\eta\mu}{C_{0}}\right)(\mathbb{E}[f(x_{t})]-f^{*})+\frac{\eta^{2}L\sigma^{2}}{\epsilon}+\frac{\eta^{2}L(C_{1}+C_{2}^{2})G^{2}}{\epsilon}+\frac{3\eta^{3}L^{2}C^{2}_{1}G^{2}}{2\epsilon^{3/2}}$
$\displaystyle+\eta(1+C_{1})G^{2}\mathbb{E}[\|D_{t}\|_{1}]+\frac{\eta^{2}(1+C_{1})C_{1}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|D_{t}\|_{1}]$
$\displaystyle+\eta^{2}C_{1}^{2}LG^{2}\mathbb{E}[\|D_{t}\|^{2}].$

Notice that $\eta\leq\frac{\epsilon}{4LC_{0}}\leq\frac{C_{0}}{4\mu}$, so that
the coefficient $1-\frac{\eta\mu}{C_{0}}\in(0,1)$. Unrolling the recursion, we
arrive

$\displaystyle\mathbb{E}[f(x_{T+1})]-f^{*}\leq$
$\displaystyle\left(1-\frac{\eta\mu}{C_{0}}\right)^{T}(\mathbb{E}[f(x_{1})]-f^{*})$
$\displaystyle+\left(\frac{\eta^{2}L\sigma^{2}}{\epsilon}+\frac{\eta^{2}L(C_{1}+C_{2}^{2})G^{2}}{\epsilon}+\frac{3\eta^{3}L^{2}C^{2}_{1}G^{2}}{2\epsilon^{3/2}}\right)\sum_{t=1}^{T}\left(1-\frac{\eta\mu}{C_{0}}\right)^{t}$
$\displaystyle+\eta(1+C_{1})G^{2}\sum_{t=1}^{T}\left(1-\frac{\eta\mu}{C_{0}}\right)^{t}\mathbb{E}[\|D_{t}\|_{1}]$
$\displaystyle+\frac{\eta^{2}(1+C_{1})C_{1}LG^{2}}{\sqrt{\epsilon}}\sum_{t=1}^{T}\left(1-\frac{\eta\mu}{C_{0}}\right)^{t}\mathbb{E}[\|D_{t}\|_{1}]$
$\displaystyle+\eta^{2}C_{1}^{2}LG^{2}\sum_{t=1}^{T}\left(1-\frac{\eta\mu}{C_{0}}\right)^{t}\mathbb{E}[\|D_{t}\|^{2}].$
(13)

For the second sum above we upper bound it by its infinite sum as

$\sum_{t=1}^{T}\left(1-\frac{\eta\mu}{C_{0}}\right)^{t}\leq\sum_{t=0}^{\infty}\left(1-\frac{\eta\mu}{C_{0}}\right)^{t}=\frac{C_{0}}{\eta\mu}.$

For the other three sums we bound $1-\frac{\eta\mu}{C_{0}}\leq 1$ and apply
the bounds in Lemma 7:

$\sum_{t=1}^{T}\left(1-\frac{\eta\mu}{C_{0}}\right)^{t}\mathbb{E}[\|D_{t}\|_{1}]\leq\sum_{t=1}^{T}\mathbb{E}[\|D_{t}\|_{1}]\leq\frac{d}{\sqrt{\epsilon}},$
$\sum_{t=1}^{T}\left(1-\frac{\eta\mu}{C_{0}}\right)^{t}\mathbb{E}[\|D_{t}\|^{2}]\leq\sum_{t=1}^{T}\mathbb{E}[\|D_{t}\|^{2}]\leq\frac{d}{\epsilon}.$

Plugging all this bounds into (13) and noticing that $x_{1}=\theta_{1}$, we
finally get

$\displaystyle\mathbb{E}[f(x_{T+1})]-f^{*}\leq$
$\displaystyle\left(1-\frac{\eta\mu}{C_{0}}\right)^{T}(f(\theta_{1})-f^{*})$
$\displaystyle+\frac{C_{0}}{\eta\mu}\left(\frac{\eta^{2}L\sigma^{2}}{\epsilon}+\frac{\eta^{2}L(C_{1}+C_{2}^{2})G^{2}}{\epsilon}+\frac{3\eta^{3}L^{2}C^{2}_{1}G^{2}}{2\epsilon^{3/2}}\right)$
$\displaystyle+\frac{\eta(1+C_{1})G^{2}d}{\sqrt{\epsilon}}+\frac{\eta^{2}(1+C_{1})C_{1}LG^{2}d}{\epsilon}+\frac{\eta^{2}C_{1}^{2}LG^{2}d}{\epsilon}$
$\displaystyle\leq$
$\displaystyle\left(1-\frac{\eta\mu}{C_{0}}\right)^{T}(f(\theta_{1})-f^{*})$
$\displaystyle+\eta\left(\frac{LC_{0}\sigma^{2}}{\mu\epsilon}+\frac{LC_{0}(C_{1}+C_{2}^{2})G^{2}}{\mu\epsilon}+\frac{(1+C_{1})G^{2}d}{\sqrt{\epsilon}}\right)$
$\displaystyle+\eta^{2}\left(\frac{3L^{2}C_{0}C^{2}_{1}G^{2}}{2\mu\epsilon^{3/2}}+\frac{(1+C_{1})C_{1}LG^{2}d}{\epsilon}+\frac{C_{1}^{2}LG^{2}d}{\epsilon}\right).$

The obtained rate above is with respect to the virtual iterates $x_{t}$ that
we defined for the purposes of analysis. To convert this rate with respect to
the iterates $\theta_{t}$ of the algorithm, we apply $L$-smoothness to bound
the functional difference:

$|f(x_{t})-f(\theta_{t})|\leq|\langle\nabla
f(\theta_{t}),x_{t}-\theta_{t})\rangle|+\frac{L}{2}\|x_{t}-\theta_{t}\|^{2}\leq\frac{\eta
C_{1}G^{2}}{\sqrt{\epsilon}}+\frac{\eta^{2}LC_{1}^{2}G^{2}}{2\epsilon},$

which implies

$\displaystyle\mathbb{E}[f(\theta_{T+1})]-f^{*}\leq$
$\displaystyle\left(1-\frac{\eta\mu}{C_{0}}\right)^{T}(f(\theta_{1})-f^{*})$
$\displaystyle+\eta\left(\frac{LC_{0}\sigma^{2}}{\mu\epsilon}+\frac{LC_{0}(C_{1}+C_{2}^{2})G^{2}}{\mu\epsilon}+\frac{(1+C_{1})G^{2}d}{\sqrt{\epsilon}}+\frac{C_{1}G^{2}}{\sqrt{\epsilon}}\right)$
$\displaystyle+\eta^{2}\left(\frac{3L^{2}C_{0}C^{2}_{1}G^{2}}{2\mu\epsilon^{3/2}}+\frac{(1+C_{1})C_{1}LG^{2}d}{\epsilon}+\frac{C_{1}^{2}LG^{2}d}{\epsilon}+\frac{LC_{1}^{2}G^{2}}{2\epsilon}\right)$

and completes the proof. ∎

###  D.4 Non-convex Analysis with Weight Decay

Algorithm 5 MicroAdamW (MicroAdam with Weight Decay)

1:Input: parameters $\beta_{1},\,\beta_{2}\in(0,1)$, $\epsilon>0$, step-size
$\eta>0$, $\theta_{1}\in\mathbb{R}^{d}$, $e_{1}=m_{0}=v_{0}=\hat{v}_{0}=0_{d}$

2:for $t=\\{1,2,...,T\\}$ do

3: $g_{t}=\widetilde{\nabla}_{\theta}f(\theta_{t})$ $\diamond$ Compute
unbiased stochastic gradient

4: $\tilde{g}_{t}=\mathcal{C}(g_{t}+e_{t})$ $\diamond$ Add accumulated error
$e_{t}$ and compress

5: $e_{t+1}=\mathcal{Q}(e_{t}+g_{t}-\tilde{g}_{t})$ $\diamond$ Update and
compress the error

6: $m_{t}=\beta_{1}m_{t-1}+(1-\beta_{1})\tilde{g}_{t}$ $\diamond$ Update
first-order gradient moment

7: $v_{t}=\beta_{2}v_{t-1}+(1-\beta_{2})\tilde{g}_{t}^{2}$ $\diamond$ Update
second-order gradient moment

8:
$\theta_{t+1}=(1-\eta_{t}\lambda)\theta_{t}-\eta\frac{\hat{m_{t}}}{\sqrt{\hat{v}_{t}+\epsilon}}$
$\diamond$ Update the model parameters with

9:end for

######  Lemma 8.

Under Assumptions 1-5, for all iterates of Algorithm 5 we have

$\displaystyle\mathbb{E}\left[\|m_{t}-\nabla f(\theta_{t})\|^{2}\right]\leq$
$\displaystyle\left(1-\frac{\beta_{1}}{2}\right)\mathbb{E}\left[\|m_{t-1}-\nabla
f(\theta_{t-1})\|^{2}\right]+\frac{2}{\beta_{1}}L^{2}\mathbb{E}\left[\|\theta_{t}-\theta_{t-1}\|^{2}\right]$
$\displaystyle+\beta_{1}\mathbb{E}\|\nabla f(\theta_{t})-\tilde{g}_{t}\|^{2}.$

###### Proof.

We start our proof from

$\displaystyle\mathbb{E}\left[\|m_{t}-\nabla f(\theta_{t})\|^{2}\right]=$
$\displaystyle\mathbb{E}\left[\|(1-\beta_{1})m_{t-1}+\beta_{1}\tilde{g}_{t}-\nabla
f(\theta_{t})\|^{2}\right]$ $\displaystyle=$
$\displaystyle\mathbb{E}\left[\|(1-\beta_{1})m_{t-1}+(1-\beta_{1})\nabla
f(\theta_{t-1})\right.$ $\displaystyle-\left.(1-\beta_{1})\nabla
f(\theta_{t-1})+\beta_{1}\tilde{g}_{t}-\nabla f(\theta_{t})\|^{2}\right]$
$\displaystyle=$
$\displaystyle\mathbb{E}\left[\|(1-\beta_{1})m_{t-1}+(1-\beta_{1})\nabla
f(\theta_{t-1})-(1-\beta_{1})\nabla f(\theta_{t})\right.$
$\displaystyle-\left.(1-\beta_{1})\nabla
f(\theta_{t-1})+\beta_{1}\tilde{g}_{t}-\beta_{1}\nabla
f(\theta_{t})\|^{2}\right].$

Using Jensen’s inequality we have

$\displaystyle\mathbb{E}\left[\|m_{t}-\nabla f(\theta_{t})\|^{2}\right]\leq$
$\displaystyle(1-\beta_{1})\mathbb{E}\left[\|m_{t-1}-\nabla
f(\theta_{t-1})+\nabla f(\theta_{t-1})-\nabla f(\theta_{t})\|^{2}\right]$
$\displaystyle+\beta_{1}\mathbb{E}\left[\|\tilde{g}_{t}-\nabla
f(\theta_{t})\|^{2}\right].$

Using Young’s inequality we have

$\displaystyle\mathbb{E}\left[\|m_{t}-\nabla f(\theta_{t})\|^{2}\right]\leq$
$\displaystyle(1-\beta_{1})(1+b)\mathbb{E}\left[\|m_{t-1}-\nabla
f(\theta_{t-1})\|^{2}\right]$
$\displaystyle+(1-\beta_{1})\left(1+\frac{1}{b}\right)\mathbb{E}\|\nabla
f(\theta_{t-1})-f(\theta_{t})\|^{2}$
$\displaystyle+\beta_{1}\mathbb{E}\left\|\nabla
f(\theta_{t})-\tilde{g}_{t}\right\|^{2}.$

Setting $b=\frac{\beta_{1}}{2}$ we have
$(1-\beta_{1})\left(1+\frac{\beta_{1}}{2}\right)\leq 1-\frac{\beta_{1}}{2}$
and $(1-\beta_{1})\left(1+\frac{2}{\beta_{1}}\right)\leq\frac{2}{\beta_{1}}$:

$\displaystyle\mathbb{E}\left[\|m_{t}-\nabla f(\theta_{t})\|^{2}\right]\leq$
$\displaystyle\left(1-\frac{\beta_{1}}{2}\right)\mathbb{E}\left[\|m_{t-1}-\nabla
f(\theta_{t-1})\|^{2}\right]$
$\displaystyle+\frac{2}{\beta_{1}}\mathbb{E}\|\nabla
f(\theta_{t-1})-f(\theta_{t})\|^{2}+\beta_{1}\mathbb{E}\left\|\nabla
f(\theta_{t})-\tilde{g}_{t}\right\|^{2}.$

Combining this bound with $L$-smoothness we obtain

$\displaystyle\mathbb{E}\left[\|m_{t}-\nabla f(\theta_{t})\|^{2}\right]\leq$
$\displaystyle\left(1-\frac{\beta_{1}}{2}\right)\mathbb{E}\left[\|m_{t-1}-\nabla
f(\theta_{t-1})\|^{2}\right]$
$\displaystyle+\frac{2}{\beta_{1}}L^{2}\mathbb{E}\|\theta_{t-1}-\theta_{t}\|^{2}+\beta_{1}\mathbb{E}\left\|\nabla
f(\theta_{t})-\tilde{g}_{t}\right\|^{2}.$

∎

######  Lemma 9.

Under Assumptions 1-5, for all iterates of Algorithm 5 we have

$\displaystyle\mathbb{E}\left[\|\nabla
f(\theta_{t})-\tilde{g}_{t}\|^{2}\right]$ $\displaystyle\leq
9\mathbb{E}\left[\|e_{t}\|^{2}\right]+6G^{2}+3\sigma^{2}.$

###### Proof.

We start from

$\displaystyle\mathbb{E}\left[\|\nabla
f(\theta_{t})-\tilde{g}_{t}\|^{2}\right]=$
$\displaystyle\mathbb{E}\left[\|\mathcal{C}(e_{t}+g_{t})-\nabla
f(\theta_{t})\|^{2}\right]$ $\displaystyle=$
$\displaystyle\mathbb{E}\left[\|\mathcal{C}(e_{t}+g_{t})-(e_{t}+g_{t})+(e_{t}+g_{t})-\nabla
f(\theta_{t})\|^{2}\right]$ $\displaystyle\leq$ $\displaystyle
3\mathbb{E}\left[\|\mathcal{C}(e_{t}+g_{t})-(e_{t}+g_{t})\|^{2}\right]$
$\displaystyle+3\mathbb{E}\left[\|e_{t}\|^{2}\right]+3\mathbb{E}\left[\|\nabla
f(\theta_{t})-g_{t}\|^{2}\right].$

Using definition of contractive compressor we have

$\displaystyle\mathbb{E}\left[\|\nabla
f(\theta_{t})-\tilde{g}_{t}\|^{2}\right]$ $\displaystyle\leq
3q^{2}\mathbb{E}\left[\|e_{t}+g_{t}\|^{2}\right]$
$\displaystyle+3\mathbb{E}\left[\|e_{t}\|^{2}\right]+3\mathbb{E}\left[\|\nabla
f(\theta_{t})-g_{t}\|^{2}\right].$

Using Young’s inequality we have

$\displaystyle\mathbb{E}\left[\|\nabla
f(\theta_{t})-\tilde{g}_{t}\|^{2}\right]$ $\displaystyle\leq
6q^{2}\mathbb{E}\left[\|e_{t}\|^{2}\right]+6q^{2}\mathbb{E}\left[\|g_{t}\|^{2}\right]$
$\displaystyle+3\mathbb{E}\left[\|e_{t}\|^{2}\right]+3\mathbb{E}\left[\|\nabla
f(\theta_{t})-g_{t}\|^{2}\right]$ $\displaystyle\leq
9\mathbb{E}\left[\|e_{t}\|^{2}\right]+6G^{2}+3\sigma^{2}.$

∎

######  Lemma 10.

Under Assumptions 1-5, for all iterates of Algorithm 5 we have

$\mathbb{E}\left[\|e_{t}\|^{2}\right]\leq\frac{(1+\omega)^{2}q^{2}}{\left(1-(1+\omega)q\right)^{2}}G^{2}$

###### Proof.

Using definition of contractive compressor we have

$\displaystyle\mathbb{E}\left[\|e_{t+1}\|^{2}\right]$
$\displaystyle=\mathbb{E}\left[\left\|\mathcal{Q}\left(e_{t}+g_{t}-\mathcal{C}(e_{t}+g_{t})\right)\right\|^{2}\right]$
$\displaystyle\leq(1+\omega)^{2}q^{2}\mathbb{E}\left[\|e_{t}+g_{t}-\mathcal{C}(e_{t}+g_{t})\|^{2}\right].$

Using Young’s inequality we have

$\displaystyle\mathbb{E}\left[\|e_{t+1}\|^{2}\right]$
$\displaystyle\leq(1+\omega)^{2}q^{2}\left(1+a\right)\mathbb{E}\left[\|e_{t}\|^{2}\right]+(1+\omega)^{2}q^{2}\left(1+\frac{1}{a}\right)\mathbb{E}\left[\|g_{t}\|^{2}\right].$

We need to satisfy the following condition:

$\displaystyle(1+\omega)^{2}q^{2}(1+a)\leq(1+\omega)q.$

It holds for $0\leq\omega,$ $0\leq q<1$, $(1+\omega)q<1$ and
$a=\frac{1}{(1+\omega)q}-1$. Using this parameters we have

$\displaystyle\mathbb{E}\left[\|e_{t+1}\|^{2}\right]$
$\displaystyle\leq(1+\omega)q\mathbb{E}\left[\|e_{t}\|^{2}\right]+(1+\omega)^{2}q^{2}\left(\frac{1}{1-(1+\omega)q}\right)\|g_{t}\|^{2}$
$\displaystyle\leq(1+\omega)q\mathbb{E}\left[\|e_{t}\|^{2}\right]+\frac{(1+\omega)^{2}q^{2}}{1-(1+\omega)q}\mathbb{E}\left[\|g_{t}\|^{2}\right].$

Unrolling this recursion allows us to obtain

$\displaystyle\mathbb{E}\left[\|e_{t+1}\|^{2}\right]$
$\displaystyle\leq\left((1+\omega)q\right)^{t}\mathbb{E}\left[\|e_{1}\|^{2}\right]+\sum^{T}_{t=1}\left((1+\omega)q\right)^{t}\frac{(1+\omega)^{2}q^{2}}{1-(1+\omega)q}G^{2}$
$\displaystyle\leq\left((1+\omega)q\right)^{t}\mathbb{E}\left[\|e_{1}\|^{2}\right]+\frac{1}{1-(1+\omega)q}\frac{(1+\omega)^{2}q^{2}}{1-(1+\omega)q}G^{2}$
$\displaystyle\leq\frac{(1+\omega)^{2}q^{2}}{\left(1-(1+\omega)q\right)^{2}}G^{2},$

last inequality holds because $e_{1}=0.$ ∎

######  Lemma 11.

Under Assumptions 1-5, for all iterates of Algorithm 5 we have

$\displaystyle\mathbb{E}\left[\left\|\tilde{g_{t}}\right\|^{2}\right]\leq
4(1+q^{2})\left(1+\frac{(1+\omega)^{2}q^{2}}{(1-(1+\omega)q)^{2}}\right)G^{2}.$

###### Proof.

$\displaystyle\mathbb{E}\left[\left\|\tilde{g_{t}}\right\|^{2}\right]$
$\displaystyle=\mathbb{E}\left[\left\|\mathcal{C}\left(g_{t}+e_{t}\right)\right\|^{2}\right]$
$\displaystyle=\mathbb{E}\left[\left\|\mathcal{C}\left(g_{t}+e_{t}\right)-\left(g_{t}+e_{t}\right)+\left(g_{t}+e_{t}\right)\right\|^{2}\right]$
$\displaystyle\leq
2\mathbb{E}\left[\left\|\mathcal{C}\left(g_{t}+e_{t}\right)-\left(g_{t}+e_{t}\right)\right\|^{2}\right]+2\mathbb{E}\left[\left\|g_{t}+e_{t}\right\|^{2}\right]$
$\displaystyle\leq
2(1+q^{2})\mathbb{E}\left[\left\|g_{t}+e_{t}\right\|^{2}\right].$

Using Lemma 10 we obtain

$\displaystyle\mathbb{E}\left[\left\|\tilde{g_{t}}\right\|^{2}\right]$
$\displaystyle\leq
4(1+q^{2})\left(1+\frac{(1+\omega)^{2}q^{2}}{(1-(1+\omega)q)^{2}}\right)G^{2}.$

∎

######  Lemma 12 (From paper: Zhou et al. [2024b]).

Let us consider the update rule:

$\theta_{t+1}=(1-\eta\lambda)\theta_{t}-\eta_{t}\frac{m^{t}}{\sqrt{v^{t}+\varepsilon}}.$

For brevity, we denote $\widehat{v}_{t}=\sqrt{v_{t}+\varepsilon}$. Also we
define

$u_{t}:=m_{t}+\lambda\theta_{t}\otimes\widehat{v}_{t},$

where $\otimes$ denotes element-wise product. Moreover, we also define
$\widetilde{f}\left(\theta_{t}\right)$ as follows:

$\widetilde{f}\left(\theta_{t}\right)=f(\theta_{t})+\lambda_{t}\|\theta_{t}\|_{\widehat{v}_{t}}^{2}$

where
$\lambda_{t}=\frac{\lambda}{2}\sum_{i=1}^{t}\left(\frac{1-q}{2}\right)^{i}(t>0),$
$\lambda_{0}=0$ with $0<q<1$ and
$\|\theta_{t}\|_{\hat{v}_{t}}=\sqrt{\left\langle\theta_{t},\hat{v}_{k}t\otimes\theta_{t}\right\rangle}$.
Also let $c_{1}\leq\|\widehat{v}_{t}\|_{\infty}<c_{2}$, then iterates of
Algorithm 5 satisfy

$\displaystyle\widetilde{f}\left(\theta_{t}\right)$
$\displaystyle\leq\widetilde{f}\left(\theta_{t-1}\right)+\frac{\eta_{t}}{2c_{1}}\left\|\nabla
f\left(\theta_{t-1}\right)-m_{t-1}\right\|^{2}$
$\displaystyle-\frac{\eta_{t}}{2c_{2}}\left\|\nabla\widetilde{f}\left(\theta_{t-1}\right)\right\|^{2}-\frac{\eta_{t}}{4c_{2}}\left\|u_{t-1}\right\|^{2}.$

######  Lemma 13 (From paper: Zhou et al. [2024b]).

Assume that $c_{s,\infty}\leq\|\tilde{g}_{t}\|_{\infty}\leq c_{\infty},$ then
we have

$\|m_{t}\|_{\infty}\leq
c_{\infty},\quad\|v_{t}+\epsilon\|_{\infty}c^{2}_{\infty}+\epsilon,\quad\left\|\frac{(v_{t}+\epsilon)^{p}}{(v_{t+1}+\epsilon)^{p}}\right\|_{\infty}\in[1-\mu,1+\mu]\left(\forall
p\in(0,1)\right),$

where $\mu=\frac{\beta_{2}c^{2}_{\infty}}{c^{2}_{s,\infty}+\epsilon}$.

######  Theorem 5.

Let Assumptions 1 to 5 hold. Let us define
$\widetilde{f}\left(\theta_{t}\right)=f(\theta_{t})+\lambda_{t}\|\theta_{t}\|_{\widehat{v}_{t}}^{2}$,
where
$\lambda_{t}=\frac{\lambda}{2}\sum_{i=1}^{t}\left(\frac{1-q}{2}\right)^{i}(t>0),$
$\lambda_{0}=0$ with $0<q<1$ and
$\|\theta_{t}\|_{\hat{v}_{t}}=\sqrt{\left\langle\theta_{t},\hat{v}_{k}\otimes\theta_{t}\right\rangle}$.
Also define
$\Psi_{t}=\widetilde{f}(\theta_{t})+\frac{\eta_{t}}{2c_{1}\beta_{1}}\mathbb{E}\left[\|m_{t}-\nabla
f(\theta_{t})\|^{2}\right].$ With
$\eta_{t}=\eta\leq\frac{\beta_{1}c_{1}}{2L}\sqrt{\frac{c_{1}}{2c_{2}}}$,
Algorithm 5 satisfies

$\displaystyle\frac{1}{T}\sum^{T}_{t=1}\left\|\nabla\widetilde{f}\left(\theta_{t-1}\right)\right\|^{2}$
$\displaystyle\leq\frac{2c_{2}}{\eta T}\Psi^{0}$
$\displaystyle+\frac{c_{2}}{c_{1}}\left(9\left(1+\frac{(1+\omega)^{2}q^{2}}{\left(1-(1+\omega)q\right)^{2}}\right)G^{2}+6G^{2}+3\sigma^{2}\right).$

###### Proof.

We start from main lemma and lemma for momentum, summing inequalities together
we obtain

$\displaystyle\widetilde{f}\left(\theta_{t}\right)+V\mathbb{E}\left[\|m_{t}-\nabla
f(\theta_{t})\|^{2}\right]$
$\displaystyle\leq\widetilde{f}\left(\theta_{t-1}\right)+\frac{\eta_{t}}{2c_{1}}\left\|\nabla
f\left(\theta_{t-1}\right)-m_{t-1}\right\|^{2}$
$\displaystyle-\frac{\eta_{t}}{2c_{2}}\left\|\nabla\widetilde{f}\left(\theta_{t-1}\right)\right\|^{2}-\frac{\eta_{t}}{4c_{2}}\left\|u_{t-1}\right\|^{2}$
$\displaystyle+V\left(1-\frac{\beta_{1}}{2}\right)\mathbb{E}\left[\|m_{t-1}-\nabla
f(\theta_{t-1})\|^{2}\right]$
$\displaystyle+V\frac{2}{\beta_{1}}L^{2}\mathbb{E}\|\theta_{t-1}-\theta_{t}\|^{2}+V\beta_{1}\mathbb{E}\left\|\nabla
f(\theta_{t})-\tilde{g}_{t}\right\|^{2}.$

Using previous lemmas we have

$\displaystyle\widetilde{f}\left(\theta_{t}\right)+V\mathbb{E}\left[\|m_{t}-\nabla
f(\theta_{t})\|^{2}\right]$
$\displaystyle\leq\widetilde{f}\left(\theta_{t-1}\right)+\frac{\eta_{t}}{2c_{1}}\left\|\nabla
f\left(\theta_{t-1}\right)-m_{t-1}\right\|^{2}$
$\displaystyle-\frac{\eta_{t}}{2c_{2}}\left\|\nabla\widetilde{f}\left(\theta_{t-1}\right)\right\|^{2}-\frac{\eta_{t}}{4c_{2}}\left\|u_{t-1}\right\|^{2}$
$\displaystyle+V\left(1-\frac{\beta_{1}}{2}\right)\mathbb{E}\left[\|m_{t-1}-\nabla
f(\theta_{t-1})\|^{2}\right]$
$\displaystyle+V\frac{2}{\beta_{1}}L^{2}\mathbb{E}\|\theta_{t-1}-\theta_{t}\|^{2}$
$\displaystyle+V\beta_{1}\left(9\left(1+\frac{(1+\omega)^{2}q^{2}}{\left(1-(1+\omega)q\right)^{2}}\right)G^{2}+6G^{2}+3\sigma^{2}\right).$

Using $\theta_{t}-\theta_{t-1}=-\eta_{t}\frac{u_{t-1}}{\widehat{v}_{t-1}}$ we
have

$\displaystyle\widetilde{f}\left(\theta_{t}\right)+V\mathbb{E}\left[\|m_{t}-\nabla
f(\theta_{t})\|^{2}\right]$
$\displaystyle\leq\widetilde{f}\left(\theta_{t-1}\right)-\left(\frac{\eta_{t}}{4c_{2}}-\frac{2VL^{2}\eta_{t}^{2}}{\beta_{1}c_{1}^{2}}\right)\left\|u_{t-1}\right\|^{2}$
$\displaystyle-\frac{\eta_{t}}{2c_{2}}\left\|\nabla\widetilde{f}\left(\theta_{t-1}\right)\right\|^{2}$
$\displaystyle+\left(V\left(1-\frac{\beta_{1}}{2}\right)+\frac{\eta_{t}}{2c_{1}}\right)\mathbb{E}\left[\|m_{t-1}-\nabla
f(\theta_{t-1})\|^{2}\right]$
$\displaystyle+V\beta_{1}\left(9\left(1+\frac{(1+\omega)^{2}q^{2}}{\left(1-(1+\omega)q\right)^{2}}\right)G^{2}+6G^{2}+3\sigma^{2}\right).$

Using $V=\frac{\eta_{t}}{2c_{1}\beta_{1}}$ and
$\Psi_{t}=\widetilde{f}(\theta_{t})+\frac{\eta_{t}}{2c_{1}\beta_{1}}\mathbb{E}\left[\|m_{t}-\nabla
f(\theta_{t})\|^{2}\right]$ we have

$\displaystyle\Psi_{t}$
$\displaystyle\leq\Psi_{t-1}-\frac{\eta_{t}}{4c_{2}}\left(1-\frac{4c_{2}L^{2}\eta_{t}^{2}}{\beta_{1}^{2}c_{1}^{3}}\right)\left\|u_{t-1}\right\|^{2}$
$\displaystyle-\frac{\eta_{t}}{2c_{2}}\left\|\nabla\widetilde{f}\left(\theta_{t-1}\right)\right\|^{2}$
$\displaystyle+\frac{\eta_{t}}{2c_{1}}\left(9\left(1+\frac{(1+\omega)^{2}q^{2}}{\left(1-(1+\omega)q\right)^{2}}\right)G^{2}+6G^{2}+3\sigma^{2}\right).$

Using $\eta_{t}=\eta\leq\frac{\beta_{1}c_{1}}{2L}\sqrt{\frac{c_{1}}{2c_{2}}}$
we have

$\displaystyle\Psi_{t}$
$\displaystyle\leq\Psi_{t-1}-\frac{\eta_{t}}{2c_{2}}\left\|\nabla\widetilde{f}\left(\theta_{t-1}\right)\right\|^{2}$
$\displaystyle+\frac{\eta}{2c_{1}}\left(9\left(1+\frac{(1+\omega)^{2}q^{2}}{\left(1-(1+\omega)q\right)^{2}}\right)G^{2}+6G^{2}+3\sigma^{2}\right).$

Summing from $t=1$ to $T$ we have

$\displaystyle\frac{1}{T}\sum^{T}_{t=1}\left\|\nabla\widetilde{f}\left(\theta_{t-1}\right)\right\|^{2}$
$\displaystyle\leq\frac{2c_{2}}{\eta
T}\Psi^{0}+\frac{c_{2}}{c_{1}}\left(9\left(1+\frac{(1+\omega)^{2}q^{2}}{\left(1-(1+\omega)q\right)^{2}}\right)G^{2}+6G^{2}+3\sigma^{2}\right).$

∎

#### Discussion.

This result is similar to the one from Zhou et al. [2024b] in the non-convex
case, where the decay rate for the first term is
$\mathcal{O}\left(\frac{1}{T}\right)$ and the second term is a non-vanishing
$\mathcal{O}\left(\beta_{1}\frac{c_{2}}{c_{1}}\sigma^{2}\right).$ In our
result, the non-vanishing term is proportional to
$\mathcal{O}\left(\frac{c_{2}}{c_{1}}\left(\sigma^{2}+G^{2}\right)\right).$

A key difference is that our term is not proportional to $\beta_{1}.$ It is
important to note that $\beta_{1}$ typically takes a value close to $1$ in
practical applications, meaning its influence on the bound is minimal.
Therefore, even though our result does not directly involve $\beta_{1}$, the
impact on the overall bound is not significantly different.

Moreover, our bound includes an additional term proportional to $G^{2},$ which
represents the gradient norm squared. This makes the bound slightly worse
compared to the result in Zhou et al. [2024b]. However, this degradation is
only by a constant factor, which means that while the theoretical bound may be
worse, the practical implications are often negligible.

In summary, our result aligns closely with previous findings, with differences
primarily in the constant factors and the presence of $G^{2}.$ Despite these
differences, the practical performance remains largely unaffected, ensuring
that the bound remains robust and applicable in a variety of scenarios.

##  Appendix E Error Feedback applied to GaLore

###  E.1 Behaviour of the Error Feedback Mechanism

The GaLore low-rank updates introduced by [Zhao et al., 2024] enable the
compression of optimizer states by performing learning updates on a lower-
dimensional subspace. In this approach, the optimizer receives gradients
projected on a defined learning subspace. Theoretical convergence guarantees
are provided under a “stable rank” assumption, where learning subspace is
fixed during training. However, in practice, convergence is attained by
occasionally updating the learning subspace and allowing full space learning
to better align with the gradient trajectory during training.

Here, it is useful to draw an analogy with the TopK method, as the occasional
updates of the learning subspace resembles working with a fixed mask for many
steps. Using a fixed mask would result in discarding the same coordinates of
the gradient at each step. Similarly, in the case of low-rank updates,
components orthogonal to the same learning subspace are discarded at each
step.

The systematic nature of the information discarded by compression carries
significant implications for error feedback behavior. Over multiple steps, the
error accumulates gradient components belonging to the orthogonal space of the
same learning subspace. Consequently, by linearity, the error itself resides
in the orthogonal space of this learning subspace. As a result, when the error
is passed to the accumulator, its projection onto the learning space is
effectively disregarded until it is potentially utilized at the specific step
when the learning subspace is updated. Therefore, the behavior of error
feedback in the case of low-rank updates is non-standard: it accumulates
gradient components over numerous steps before unloading them all at once.

For a better understanding, we derive analytical evidence for the described
behaviour by induction. Let $L$ be fixed learning subspace and assume that
$e_{t-1}\in L^{\perp}$. Then, gradient passed to the optimizer at step $t$ is:
$\mathcal{C}_{GaLore}(a_{t})=proj_{L}(a_{t})=proj_{L}(e_{t-1}+g_{t})=proj_{L}(g_{t})$
where error feedback is discarded. Thus,
$e_{t}=a_{t}-\mathcal{C}_{GaLore}(a_{t})=e_{t-1}+g_{t}-proj_{L}(g_{t})=e_{t-1}+proj_{L^{\perp}}(g_{t})\in
L^{\perp}$ which completes the induction.

Assume now that learning subspace is updated every $T$ steps, and denote
$L_{t}$ the learning subspace at step $t$. Then, a similar induction leads to:

$e_{t}=\sum_{i=1}^{t}\underset{j=i}{\overset{t}{\circ}}proj_{L_{j}^{\perp}}(g_{i})=\sum_{i=1}^{t}\underset{j=\lfloor\frac{i}{T}\rfloor}{\overset{\lfloor\frac{t}{T}\rfloor}{\circ}}proj_{L_{jT}^{\perp}}(g_{i})$
$a_{kT}=proj_{L_{kT}}(g_{kT}+e_{kT-1})=proj_{L_{kT}}(g_{kT})+\sum_{i=1}^{kT-1}proj_{L_{kT}}\circ(\underset{j=i}{\overset{t}{\circ}}proj_{L_{j}^{\perp}})(g_{i})$

###  E.2 Consequences on Training

Such behaviour of the error feedback mechanism results in the dominance of the
error norm over the gradient norm. Before learning subspace updates, the error
is the sum over past gradient components that belong to the orthogonal of the
current learning subspaces. Since these components represent descent
directions that were not used, they are not expected to compensate each other
on average. Consequently, between learning subspace updates, the error norm is
expected to grow linearly. Figure 6 provides evidence of such linear growth of
the error norm during fine-tuning of RoBERTa-base model on GLUE/MNLI task.

It implies that known analysis techniques [Alistarh et al., 2018, Karimireddy
et al., 2019] of convergence for the error feedback mechanism do not apply to
GaLore. Indeed, such proofs rely on the assumption that the compression
operator is contractive, as it allows the error to be bounded. Given a fixed
vector, low-rank compression based on its singular value decomposition is a
contraction operator. However, in our case, the compression is based on a
previously-computed singular value decomposition and therefore may not be a
contraction operator for newly computed gradients. The extreme case being when
the gradient is orthogonal to the learning subspace, in which case the
compression operator returns the null vector. Figure 6 shows that during
training the error norm is not on the same order of magnitude of the gradient
norm.

Figure 6: Dynamics of the norm of the error compared to norm of the gradient
(of output of the 3rd attention layer) during fine-tuning of RoBERTa-base
model on GLUE/MNLI from surrogate GaLore with error feedback optimizer. We
used hyperparameters from [Zhao et al., 2024], i.e. batch size 16, learning
rate 0.00001, projection update gap 200, rank 4 and GaLore scale 4.

The dominance of the error over the gradient also has effects on space
exploration, as the learning subspaces are computed from the singular value
decomposition of the accumulator (i.e. the sum of the gradient and the error).
Since the main components of the accumulator belong to the orthogonal of
current learning subspaces, successive learning subspaces will tend to be
orthogonal to each other. This allows errors to be effectively passed to the
optimizer, but all at once which can introduce irregularities in the learning
trajectory. However, it also implies that learning is performed on a learning
subspace that is suboptimal in terms of the direction of the gradient, but
this may help convergence by enforcing space exploration. See Figure 7 for
examples of how induced orthogonality of successive learning subspaces affects
the learning trajectory.

Figure 7:  Optimization trajectory for Adam, GaLore-Adam and GaLore-Adam-EF for ill-conditioned function $f(x,y)=cos(\frac{5\pi}{4}x)+sin(\frac{7\pi}{4}y)$ starting from $(x_{0},y_{0})=(-\frac{1}{4},\frac{1}{4})$ (on first row) and for Rosenbrock function starting from $(x_{0},y_{0})=(-\frac{1}{2},1)$ (on second row).  |  |
---|---|---
|  |