LDAdam: Adaptive Optimization from Low-Dimensional Gradient Statistics

## Abstract

Abstract We introduce LDAdam 1 1 1 Code available at: https://github.com/IST-DASLab/LDAdam , a memory-efficient optimizer for training large models, that performs adaptive optimization steps within lower dimensional subspaces, while consistently exploring the full parameter space during training. This strategy keeps the optimizer’s memory footprint to a fraction of the model size. LDAdam relies on a new projection-aware update rule for the optimizer states that allows for transitioning between subspaces, i.e., estimation of the statistics of the projected gradients.
To mitigate the errors due to low-rank projection, LDAdam integrates a new generalized error feedback mechanism, which explicitly accounts for both gradient and optimizer state compression.
We prove the convergence of LDAdam under standard assumptions, and show that LDAdam allows for accurate and efficient fine-tuning and pre-training of language models.

## 1 Introduction

Scaling up deep neural networks leads to ever-improving models (Kaplan et al., ), but also to ever-increasing memory and compute requirements. In this work, we focus on the memory cost of *optimizer states*, i.e., the additional first- and second-order gradient statistics required during the adaptive optimization process, which in the case of the baseline algorithm Adam (Kingma & Ba, ; Loshchilov & Hutter, ) reaches up to twice the size of the model. Our goal is to replicate Adam’s performance and theoretical guarantees while significantly reducing its memory footprint.

We start from the observation that the optimization landscape of deep networks exhibits low intrinsic dimensionality (Li et al., ). Low-rank learning has been successfully leveraged for model *fine-tuning*, with the LoRA technique (Hu et al., ) and its many extensions (Dettmers et al., ; Liu et al., ; Nikdan et al., ) becoming standard. However, such adapter-based methods do not extend to end-to-end accurate training (Lialin et al., ; Zhao et al., ), nor even to fine-tuning on more demanding tasks (Zhang et al., ; Nikdan et al., ; Liu et al., ). Despite promising results such as GaLore (Zhao et al., ), designing memory-efficient full-parameter optimizers which leverage low-dimensional gradient statistics
with provable convergence and good practical performance remains an open challenge.

#### Contributions.

We propose *Low-Dimensional Adam (LDAdam)*, a memory-efficient adaptive optimizer which, for the first time, incorporates low-rank compression of *both gradients and optimizer states*, with convergence guarantees under standard assumptions. The goal of the algorithm is to perform optimization steps within subspaces of lower dimension, while introducing mechanisms to 1) maintain relevant information from the last seen subspaces and 2) recover information lost due to low-rank compression. To reach this goal, our work is based on three main techniques. First, we formulate projection-aware update rules that allow us to perform adaptive optimization on dynamically changing coordinate systems and subspaces. Second, we leverage *block power iteration* (Bentbib & Kanber, ) at each step to efficiently and accurately project the gradient and optimizer states. Third, we use a generalized error feedback mechanism for both gradients and optimizer states. Moreover, we prove the convergence of LDAdam under standard assumptions and provide empirical evidence supporting its training efficiency.

When performing adaptive optimization *within low-dimensional subspaces*, a key challenge is to adapt the optimizer states to changes in the projection map. Adapting gradient statistics to new subspaces is not obvious and comes with two challenges. First, standard adaptive optimization algorithms are anisotropic: that is, the parameter-wise learning rate tuning implies that the choice of coordinate system used to express the gradient matrix impacts the optimization process. Second, the second-order statistics are nonlinear, which complicates estimation of the gradient statistics and error compensation. We address the first challenge by carefully selecting the coordinate system used for low-rank compression, i.e. the basis given by the most significant singular vectors. The second challenge is addressed by approximating the gradient statistics in the current subspace given the raw moments estimates in the previous subspace. Thus, in LDAdam, optimizer step is performed from the projection-aware intermediate optimizer states and the projected gradient, rather than from the optimizer states at the previous iteration and the full-rank gradient (Kingma & Ba, ). In turn, this leads to stronger convergence guarantees, and better practical performance.

The choice of projection map at each step is key to low-rank compression. Performing singular value decomposition (SVD) at each step would add prohibitive time overhead; instead, we make use of *block power iteration with warm initialization* (Bentbib & Kanber, ; Vogels et al., ), which can achieve a tight approximation of the most significant singular vectors of successive gradients, at low computational cost. To preserve the information from the last iterations kept within the optimizer states, we show that one can perform the block power iteration on a barycenter of the newly computed gradient and an average of the previous gradients. Finally, to mitigate the errors due to low-rank projection, we introduce a new variant of the error feedback mechanism (Seide et al., ; Karimireddy et al., ; Alistarh et al., ) that supports compression of *both gradient and optimizer states*. Moreover, we propose a memory efficient implementation of it by reusing the space allocated to gradient accumulation.

The resulting LDAdam method is memory-efficient: for a weight layer of shape $n\times m$, the space cost of optimizer states is $nr+2rm$, where $r$ is the projection rank, while the cost of Adam’s state is $2nm$. This matches the efficiency of the popular GaLore algorithm (Zhao et al., ), the closest prior work in this area; and can bring major gains in practice, as usually $r\ll\min(m,n)$. Relative to GaLore, LDAdam has the advantages of tracking gradient evolution in the adapted subspace at each step, and correcting for the projection error. This leads to better practical performance, but also to theoretical guarantees of convergence under standard assumptions. Specifically, for smooth non-convex objectives, LDAdam can preserve the asymptotic convergence rate of AMSGrad (Reddi et al., ), the provably-convergent version of Adam (see Theorem 1). Moreover, we show even faster rates for objectives that obey the Polyak-Łojasiewicz condition (see Theorem 2).

We validate the practical performance of LDAdam via an efficient PyTorch implementation. We apply LDAdam for *fine-tuning* RoBERTa (Liu et al., ) and Llama-family (Touvron et al., ) models on the GLUE (Wang et al., ) and Grade-School Math (GSM) (Cobbe et al., ) benchmarks, respectively. In addition, we provide *pre-training* results for Llama-type models with 130M and 350M parameters. Across all tasks and models, LDAdam shows competitive accuracy relative to Adam, but with much smaller optimizer states relative to the baseline, and comparable runtime. Relative to GaLore, LDAdam shows consistently higher accuracy, for matching memory costs and similar runtime. Ablations show that this is a consequence of both our projection-aware intermediate updates and error correction.

In summary, we show that it is possible to achieve both theoretical guarantees and strong practical performance when performing adaptive optimization from low-dimensional gradient statistics. Our techniques may also be relevant to areas such as distributed optimization, while our implementation is relevant in reducing the practical memory overheads of training large language models (LLMs).

## 2 Related work

#### Memory-Efficient First-Order Methods in Optimization.

Gradient statistics are used to determine the descent direction (Rumelhart et al., ) and to adapt the learning rate per parameter (Duchi et al., ; Hinton, ). Adam (Kingma & Ba, ) combines both ideas in a bias-corrected design, and AdamW (Loshchilov & Hutter, ) adds weight decay. Although it is the de facto optimizer for deep neural networks, Adam/AdamW has been the subject of several extensions aimed at reducing its memory footprint. For example, Adafactor (Shazeer & Stern, ), CAME (Luo et al., ) and Adam-mini (Zhang et al., ) rely on a factorized representation of the second-order statistic to reduce memory. Quantization has also been proposed to reduce the footprint of optimizer states (Dettmers et al., ), and fusing computation of the backward pass and the optimizer step (Lv et al., ; ) has also proven effective.

#### Error Feedback in Distributed Optimization.

Several methods have been proposed to reduce the communication overhead in distributed optimization: quantization (Alistarh et al., ), sparsification (Alistarh et al., ; Lin et al., ) and low-rank approximation (Vogels et al., ). Error feedback (Seide et al., ) has been shown to be effective in improving the empirical performance of biased gradient compression methods (Karimireddy et al., ; Vogels et al., ) methods and key to ensuring their convergence (Stich et al., ; Alistarh et al., ).
MicroAdam (Modoranu et al., ) only stores sparse gradients from each optimization step, and uses error feedback for correction, but reconstructs the full optimizer state at each step, instead of compressing it.

#### Parameter-Efficient Fine-Tuning (PEFT).

A natural approach for reducing memory overheads is to train less parameters (Han et al., ). This idea led to additive PEFT, where new trainable modules are added to the model architecture (Houlsby et al., ; Li & Liang, ); selective PEFT, where only certain parameters of the original model are retrained (Guo et al., ); and reparametrized PEFT, where new modules are trained and merged with the original parameters (Hu et al., ; Nikdan et al., ; Liu et al., ).
However, methods that train significantly less parameters tend to underperform for pre-training and fine-tuning on challenging tasks.

#### Optimization via Low-Rank Gradient Projection.

GaLore (Zhao et al., ), which partly motivated our work, performs adaptive optimization in a low-dimensional subspace that is updated periodically. The adaptive optimizer is fed with a low-rank representation of the gradients obtained via SVD, thus saving memory, and produces a low-rank representation of the descent direction, which is then up-projected to update the model parameters. The idea is similar to PEFT in that low-rank gradient projection yields selective training in the coordinate system induced by SVD. The addition of low-rank components to achieve full-rank training is also proposed in ReLoRA (Lialin et al., ) or COLA (Xia et al., ). However, unlike PEFT methods, GaLore computes and exploits the full gradient of the original model. The final algorithm allows pre-training of 7B parameter language models on consumer GPUs. However, it can only do so with a very small batch size and without the ability to accumulate gradients over multiple mini-batches.

#### GaLore Comparison.

Relative to GaLore, we propose the following new techniques: (1) we introduce a new projection-aware update rule (see Sections
3.1 and 3.2), which avoids accumulating low-rank gradients from *different* subspaces into the optimizer states, and performs frequent updates; (2) to correct for projection errors in *both* gradients and optimizer states, we introduce a new generalized error feedback mechanism (see Section 3.4).
The convergence of GaLore is only guaranteed under a very strong “stable-rank” assumption, whereas we show convergence under standard conditions.
Experimentally, we show that these improvements lead to consistently improved accuracy, and that they can be implemented efficiently in terms of both memory and runtime.

Several orthogonal improvements such as quantization (Dettmers et al., ; Zhang et al., ), per-layer weight update (Lv et al., ), and weight factorization (Jaiswal et al., ) have been proposed to further improve projection-based methods. Many of these improvements are compatible with LDAdam, and we plan to investigate them in future work.

## 3 The Low-Dimensional Adam (LDAdam) Algorithm

### 3.1 The Need for a Projection-Aware Update Rule

#### Standard Adaptive Optimization.

We denote $f$ the loss function to minimize, $\theta_{t}$ the model parameters at step $t$, $G_{t}=\nabla_{\theta}f(\theta_{t})\in\mathbb{R}^{n\times m}$ the gradient at $\theta_{t}$, $g_{t}$ its mini-batch stochastic counterpart, $\eta_{t}$ the learning rate, and ${\epsilon}$ a positive scalar used for numerical stability. The standard adaptive optimization relies on the gradient statistics estimates $m_{t}$ and $v_{t}$, which are obtained from the exponential moving average of the gradients at rate $\beta_{1}$ and the squared gradients at rate $\beta_{2}$, respectively. With $\hat{m}_{t}=\tfrac{m_{t}}{1-\beta_{1}^{t}}$ and $\hat{v}_{t}=\tfrac{v_{t}}{1-\beta_{2}^{t}}$ their unbiased counterpart, Adam’s update is: $\theta_{t+1}=\theta_{t}-\eta_{t}\tfrac{\hat{m}_{t}}{\sqrt{\hat{v}_{t}}+
\epsilon}.$

#### The Challenges Behind Low-Rank Updates.

When adapting the learning subspaces to capture the low-rank gradient structure, one expects the optimizer states to retain information from previous iterations. However, frequent updates of the low-rank projection causes significant challenges: first, projections are lossy and gradient information is lost; second, since the projections may be different between steps, subspace adaptation may alter the representation of the low-dimensional optimizer states. To illustrate this second point, we briefly discuss the GaLore solution  (Zhao et al., ).

Let $r$ be the compression rank, and $\mathcal{U}^{\top}_{t}$ and $\mathcal{U}_{t}$ be the projection and back-projection matrices we use at step $t$. $\mathcal{U}^{\top}_{t}\in\mathbb{R}^{r\times n}$ yields a truncation to its first $r$ rows of the gradient matrix written in an adapted coordinate system, leading to a low-dimensional representation of the gradient with shape $r\times m$. Multiplication with the matrix $\mathcal{U}_{t}$ is then used to express the compressed gradient in the canonical high-dimensional coordinate system. From now on, let $m_{t}$ and $v_{t}$ (resp. $\hat{m}_{t}$ and $\hat{v}_{t}$) be the first and second moment estimates (resp. unbiased estimates) in the *low-dimensional space* induced at step $t$ by $\mathcal{U}^{\top}_{t}$, and let $M_{t}=\mathcal{U}_{t}\cdot m_{t}$ denote the back-projected low-rank counterpart of $m_{t}$.

In this context, Equation 1 describes the GaLore dynamics (Zhao et al., , Algorithm 2):

$$ $\theta_{t+1}=\theta_{t}-\eta_{t}\bm{{\color[rgb]{0,0,1}\mathcal{U}_{t}}}\cdot \frac{\hat{m}_{t}}{\sqrt{\hat{v}_{t}}+{\epsilon}}=\theta_{t}-\frac{\eta_{t}}{ \sqrt{\bm{\hat{v}_{t}}}+{\epsilon}}\frac{1-\beta_{1}}{1-\beta_{1}^{t}}\sum_{ \tau=1}^{t}\beta_{1}^{t-\tau}\bm{{\color[rgb]{0,0,1}\mathcal{U}_{t}}}\cdot\bm{ \mathcal{U}_{\tau}^{\top}}\cdot g_{\tau}$ (1) $$

Notice that, in the above equation, the change in the low-dimensional projection basis between steps is ignored: as can be seen in the right-hand sum, gradients are projected and back-projected using *different maps*, i.e. $\bm{{\mathcal{U}_{t}}}\textnormal{ and }\bm{\mathcal{U}_{\tau}^{\top}}$, respectively. Moreover, the statistics of gradients projected into different subspaces are accumulated on *the same momentum buffers*. The issue we raise here (highlighted in bold in Equation 1) is partially mitigated in practice by heuristics: occasional updates of the projection map and exponential decay rates lead to only a fraction of the steps being effectively resulting from accumulation of compressed gradients in different subspaces. However, we will see experimentally that this leads to a drop in accuracy relative to the “correct” projections.

### 3.2 The LDAdam Update Rule

#### Overview.

In LDAdam, we address two key challenges in low-rank updates: we introduce a projection-aware update rule that accounts for the change in the low-dimensional representation of the optimizer states, and we provide a generalized error feedback mechanism that accounts for the loss of gradient accuracy due to low-rank compression. Compared to standard adaptive optimization step, optimizer states are replaced by low-dimensional intermediate optimizer states, denoted by $m_{t-1/2}$, $v_{t-1/2}$; and the gradient is replaced by the accumulator $a_{t}=\mathcal{U}^{\top}_{t}\cdot A_{t}$ obtained by low-rank projection of the high-dimensional accumulator $A_{t}=g_{t}+\xi_{t}$, where $\xi_{t}$ is the (generalized) error feedback buffer, storing projection errors. In the following, we build the LDAdam algorithm step-by-step. The end-to-end construction is provided in Algorithm 1.

Our first objective is to find a mechanism to transfer gradient information from one lower dimensional vector space to another. For this, we take inspiration from orthogonal projections, which provide the best approximation in terms of $\ell_{2}$ error. The projection-aware update rule for the first moment estimate (Algorithm 1 Line 14) follows from the multiplication with the matrix $\mathcal{U}^{\top}_{t}\cdot\mathcal{U}_{t-1}$, which stands for both the projection matrix and the change of basis matrix. The matrix $\mathcal{U}^{\top}_{t}\cdot\mathcal{U}_{t-1}\in\mathbb{R}^{r\times r}$ is low-dimensional and thus allows efficient transition between subspaces, as no intermediate high-dimensional vector is created. We obtain the following update:

$$ $m_{t}=\beta_{1}\textsc{Proj}_{\mathcal{U}_{t}}(M_{t-1})+(1-\beta_{1})a_{t}= \beta_{1}\underbrace{\mathcal{U}^{\top}_{t}\cdot\mathcal{U}_{t-1}\cdot m_{t-1} }_{m_{t-1/2}}+(1-\beta_{1})a_{t}.$ (2) $$

Optimizer States Compression.
Note already that projection yields compression. Yet, the discrepancy between successive subspaces leads to a loss of information contained in the optimizer states. We address this by enforcing slight shifts between subspaces, and by generalizing error feedback to account for the updating of optimizer states. The information lost on gradient momentum due to compression, and that needs to be reintroduced via the error feedback mechanism, is given by:

$$ $M_{t-1}-\mathcal{U}_{t}\mathcal{U}^{\top}_{t}\cdot M_{t-1}=\mathcal{U}_{t-1} \cdot m_{t-1}-\mathcal{U}_{t}\cdot m_{t-1/2}.$ (3) $$

#### Estimating the Statistics of Projected Gradients.

Although reliable, the orthogonal projection of the optimizer state does not apply to the estimation of the second-order statistic. In particular, Adam does not rely solely on linear operations and is anisotropic (see Appendix A). To overcome these challenges, we interpret Adam’s optimizer states as statistical estimates of the first two moments of each gradient coordinate. Indeed, Adam’s optimization step $t$ holds for an approximation of $(\mathbb{E}_{t,\beta_{1}}[G^{e_{i}}]/\sqrt{\mathbb{E}_{t,\beta_{2}}[(G^{e_{i}}
)^{2}]})_{i\leq n}$, where $\mathcal{E}_{t}=(e_{1},\ldots,e_{n})$ is the basis provided by the canonical parameterization of the model, $G^{e_{i}}=\langle G,e_{i}\rangle$ a column of the gradient matrix, and $\mathbb{E}_{t,\beta}[\cdot]$ is the exponential time-weighted expectation from time $t$ with decay-rate $\beta$.

We rely on this observation to rewrite Adam’s update in any given coordinate system $\tilde{\mathcal{E}}_{t+1}$: we build $\hat{m}_{t+1/2}$ and $\hat{v}_{t+1/2}$ as statistical estimates of $(\mathbb{E}_{t,\beta_{1}}[G^{\tilde{e_{i}}}])_{i\leq n}$ and $(\mathbb{E}_{t,\beta_{2}}[(G^{\tilde{e_{i}}})^{2}])_{i\leq n}$. It follows from $G^{\tilde{e_{i}}}=\sum_{j\leq n}\langle\tilde{e_{i}},e_{j}\rangle G^{e_{j}}$ that the matrix of change of basis $(\langle\tilde{e_{i}},e_{j}\rangle)_{i,j\leq n}$ provides transition between coordinate systems. Its truncation to $i,j\leq r$ induces a rank-$r$ orthogonal projection of the gradients, and thus holds for the transition between subspaces:

$$ $\mathbb{E}_{t,\beta_{1}}[G^{\tilde{e_{i}}}]=\sum_{j=1}^{n}\langle\tilde{e_{i}} ,e_{j}\rangle\mathbb{E}_{t,\beta_{1}}[G^{e_{j}}]\approx\sum_{j=1}^{r}\langle \tilde{e_{i}},e_{j}\rangle(\hat{m}_{t})_{j}=(\mathcal{U}^{\top}_{t+1}\cdot \mathcal{U}_{t}\cdot\hat{m}_{t})_{i}.$ (4) $$

Linearity of expectation leads to an autoregressive update rule for the first moment estimate (see Equation 4). With $(\langle\tilde{e_{i}},e_{j}\rangle)_{i,j}=(\mathcal{U}^{\top}_{t+1}\cdot
\mathcal{U}_{t})_{i,j}$ the estimation of the first-order statistic in the coordinate system $\tilde{\mathcal{E}}_{t+1}$ resolves to an orthogonal projection when considering only the first $r$ components (i.e., after truncation to $i,j\leq r$), and the above discussion on compression of optimizer states remains valid.
Following our statistical approach, we obtain the projection-aware update rule for the second moment estimates (Algorithm 1 Line 15) from the approximation below:

$$ $\displaystyle\mathbb{E}_{t,\beta_{2}}[(G^{\tilde{e_{i}}})^{2}]$ $\displaystyle=\sum_{j=1}^{n}\langle\tilde{e_{i}},e_{j}\rangle^{2}\mathbb{E}_{t ,\beta_{2}}[(G^{e_{j}})^{2}]+\sum_{k\neq l}^{n}\langle\tilde{e_{i}},e_{k} \rangle\langle\tilde{e_{i}},e_{l}\rangle\mathbb{E}_{t,\beta_{2}}[G^{\tilde{e_{ k}}}G^{\tilde{e_{l}}}]$ $\displaystyle\approx\sum_{j=1}^{r}\langle\tilde{e_{i}},e_{j}\rangle^{2}(\hat{v }_{t})_{j}+\sum_{k\neq l}^{r}\langle\tilde{e_{i}},e_{k}\rangle\langle\tilde{e_ {i}},e_{l}\rangle(\hat{m}_{t})_{k}(\hat{m}_{t})_{l}.$ (5) $$

The approximation of the second-order statistics in a new coordinate system requires the estimation of the covariance between the gradient coordinates (see Equation 5). Although too large to store, the covariance can be approximated by the product of first-order moment estimates, assuming independence between gradient coordinates. The latter assumption is enforced on average over the gradient matrix columns by the nature of the coordinate system inherited from SVD. In practice, we guarantee a positive estimate of the second-order statistics by clipping value to zero if necessary.

### 3.3 Learning Subspace Adaptation

LDAdam’s memory savings stem from the low-dimensional structure of the moments estimate. The compression strategy, defined by a choice of truncated coordinate systems, sets the exploration of the parameter space, and its design meets two criteria: (i) the learning subspace must be adapted at each step to integrate error feedback; (ii) the compression operator must be set to approximate the current gradient and to preserve the information stored in the optimizer states.

Projection Map Computation.
Although singular value decomposition (SVD) provides an optimal low-rank approximation, its computational cost makes it prohibitive in our case. We follow PowerSGD (Vogels et al., ) to efficiently approximate the most significant singular vectors: we perform a single block power iteration (Bentbib & Kanber, ) initialized with the approximation from the previous optimization step; and apply the Gram-Schmidt process to derive the projection map (Algorithm 1 Line 11). In addition, we propose to start the whole process by computing SVD of the first gradient.

Learning A Smooth Subspace Transition.
LDAdam relies at each step on compression of the gradient and the optimizer states via a unique low-rank projection map. Instead of fitting two projection maps and interpolate between them, we propose to fit a single projection map to the interpolation of the gradient and an average of the previous gradients (Algorithm 1 Line 10). We introduce an interpolation factor $\rho$ and perform a block power iteration to approximate SVD of the matrix: $B_{t}=\rho\;\mathcal{U}_{t-1}\cdot\hat{m}_{t-1}+(1-\rho)A_{t}$. The interpolation factor helps to balance the accuracy of the gradient compression with the preservation of the optimizer states. We suggest setting $\rho=\beta_{1}$, i.e. computing the optimal subspace for gradient momentum compression.

### 3.4 The Generalized Error Feedback Mechanism

Error Feedback was introduced for distributed optimization (Seide et al., ) to account for gradient compression errors by reintroducing them into the next iteration gradient. We extend this mechanism to account for loss of information on both the gradient and the optimizer states. We keep the same structure of the error feedback mechanism: the optimizer is fed by a unique accumulator that is a sum of the gradient and the error buffer.

Unbiased Error Buffer Loading.
With gradient and optimizer states decayed at different rates within the optimizer step, we must adjust the error buffer loading strategy of the optimizer states projection error by the factor $\beta_{1}/(1-\beta_{1})$ to recover the first moment estimate. Furthermore, the condition under which the error buffer loading strategy is consistent with the second moment estimate is given by: $\beta_{2}=(1-\beta_{2})(\tfrac{\beta_{1}}{1-\beta_{1}})^{2}$. We suggest setting $\beta_{2}=0.99$ and $\beta_{1}\approx 0.908$, as the slight decrease in the decay rate of the second moment estimate helps to adapt faster to a new learning subspace. Equation 6 describes the loading strategy of the generalized error feedback mechanism (Algorithm 1 Line 25).

$$ $\xi_{t+1}=[A_{t}-\mathcal{U}_{t}\cdot a_{t}]+\frac{\beta_{1}(\beta_{2})}{1- \beta_{1}(\beta_{2})}[\mathcal{U}_{t-1}\cdot m_{t-1}-\mathcal{U}_{t}\cdot m_{t -1/2}].$ (6) $$

#### Implementing Error Feedback.

A naive implementation of the error feedback mechanism would add memory overhead equal to the model size. Since the gradient and the error buffer are added together before the optimization step is performed, we store the error buffer in the variable used for gradient accumulation (Algorithm 1 Lines 25 and 7)(e.g. in PyTorch, modify the optimizer.zero_grad function to store the error buffer in the p.grad variable). This results in an error feedback mechanism that does not require any further memory for optimizer states.

Figure: Algorithm 1 LDAdam ($\diamond$ Practical View Only, $g_{t}\in\mathbb{R}^{n\times m}$ / $\diamond$ Analytical View Only, $g_{t}\in\mathbb{R}^{d}$ )

## 4 Convergence Guarantees for LDAdam

#### Analytical Overview.

We now present our theoretical convergence results for LDAdam. We first provide/discuss an “analytical” perspective on our approach in Algorithm 1. Then we introduce and discuss the analytical assumptions used in the theory, along with two theoretical convergence rates.

The outlined steps closely resemble those in the practical view of the algorithm. One notable difference is the introduction of a new AMSGrad-type normalization, $\tilde{v}_{t}=\max(v_{t},||\tilde{v}_{t-1}||_{\max})$ in line 20. This represents uniform version of the original AMSGrad technique (Reddi et al., ), $\tilde{v}_{t}=\max(v_{t},\tilde{v}_{t-1})$, which enforces coordinate-wise monotonicity for the diagonal preconditioning. In our setup, as we transition between different low-dimensional spaces, the preconditioning matrix is no longer diagonal, and scaling factors may change after such a transition. To ensure that adaptive preconditioning remains monotonic, we enforce a monotonic spectrum of the preconditioning.

To avoid unnecessary complications in the notations and derivations, in this section we assume that the stochastic gradients $g_{t}$ and other variables derived from the gradients, such as momentum, are *vectors*. We remind the reader that, in practice and in our implementation, gradients are viewed/represented as 2D matrices. In our analysis, vector gradients should be treated as the columns of the gradient matrix.
Nevertheless, all steps in Algorithm 1 are still applicable when dealing with matrix gradients. Specifically, given a matrix gradient $g_{t}$, we obtain a matrix $B_{t}$, for which we assume that, regardless of the low-rank compression method used (e.g., SVD or power iterations), the contractive inequality $\|(I-\mathcal{U}^{\top}_{t}\mathcal{U}_{t})B_{t}\|_{\rm F}\leq q_{r}\|B_{t}\|_
{\rm F}$ holds with respect to the Frobenious norm. We also incorporate the debiasing factor $\sqrt{1-\beta_{2}^{t}}/(1-\beta_{1}^{t})$ in $\eta_{t}$.

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

We first outline the analytical assumptions under which we establish LDAdam’s convergence guarantees.

###### Assumption 1 (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}$:

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

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

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

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

###### Assumption 3 (Bounded variance) .

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

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

All three assumptions are commonly used in adaptive optimization literature (Chen et al., ; Défossez et al., ; Li et al., ; Xie et al., ; Zhou et al., ). It is worth noting that the bounded gradient condition in Assumption 2 directly implies a bounded variance condition with a constant $\sigma^{2}=2G^{2}$. Therefore, for the purpose of asymptotic analysis, one can omit Assumption 3. However, we will show that in the non-convex convergence rate, the $\sigma^{2}$-term decays at a rate of $1/\sqrt{T}$, while all other $G$-terms decay at a faster rate of $1/T$. Theorem 1 states our general non-convex convergence result.

###### Theorem 1 ( Non-convex convergence rate).

Let Assumptions 1, 2 and 3 hold. Then, choosing step-size $\eta=\min(\frac{\epsilon}{4LC_{0}\sqrt{1+C_{2}}},\frac{1}{\sqrt{T}})$, LDAdam (Algorithm 1) satisfies

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

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

Discussion. Compared to the standard stochastic gradient descent (Ghadimi & Lan, ), the leading term $1/\sqrt{T}$ of the obtained rate recovers the optimal non-convex convergence speed. Additionally, the asymptotic rate $\mathcal{O}(\frac{1+\sigma^{2}}{\sqrt{T}}+\frac{1}{T})$ aligns with the rate of uncompressed AMSGrad in the stochastic non-convex setup (Zhou et al., ). Therefore, the introduced low-rank compression framework along with the error feedback mechanism does not hurt the asymptotic convergence with respect to $T$.
The slowdown caused by compression in the leading term is only through the constant $C_{0}=\mathcal{O}(\frac{1}{1-q_{r}})$, where the factor $1-q_{r}\in(0,1]$ measures the aggressiveness of the compression. Note that our theory applies to any low-rank compression with $r\geq 1$. More compression corresponds to a smaller rank $r$ for the projections and a smaller factor $1-q_{r}$. For instance, when using low-rank compression via SVD decomposition, one can prove $1-q_{r}=\Theta(\frac{r}{d})$ for any rank $r\geq 1$. In this case, the leading term of the rate becomes $\mathcal{O}(\frac{d}{r}\frac{1}{\sqrt{T}})$, indicating that the convergence slows down in the worst case by a factor of $d/r$ due to compression. The proof can be found in Appendix E.2.

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

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

###### Assumption 3 (Bounded variance) .

###### Theorem 1 ( Non-convex convergence rate).

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

Next, we show faster convergence for LDAdam when the loss function satisfies the PL condition.

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

For some constant $\mu>0$ the loss function $f$ satisfies the Polyak-Łojasiewicz (PL) inequality:

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

###### Theorem 2 .

(PL convergence rate)
Let Assumptions 1, 2, 3 and 4 hold. Then, choosing step-size $\eta=\min(\eta_{0},\frac{2C_{0}\log T}{\mu T})$, LDAdam (Algorithm 1) satisfies

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

with constants $C_{1}\mathrel{\mathop{:}}=\frac{\beta_{1}+(1-\beta_{1})q_{r}}{(1-\beta_{1})(1-
q_{r})}$ and $\eta_{0}\mathrel{\mathop{:}}=\min(\frac{\epsilon}{16LC_{0}},\frac{C_{0}(1-
\beta_{1})(1-q_{r})}{2\mu},\frac{\epsilon^{3/4}}{6L\sqrt{C_{0}C_{2}}})$.

Discussion. Notably, the leading term of the convergence rate under the PL condition improves to $\mathcal{O}(\frac{\log T}{T})$, compared to $\mathcal{O}(\frac{1}{\sqrt{T}})$ in the general non-convex setting. Faster rates of $\widetilde{\mathcal{O}}(\frac{1}{T})$ for adaptive optimization algorithms under the PL condition have been established when $\beta_{2}\to 1$ (He et al., ) or through the use of the AMSGrad trick (Modoranu et al., ). In LDAdam, we apply a uniform version of AMSGrad normalization with arbitrary $\beta_{2}<1$ and derive the rate for compression with any rank $r\geq 1$. Hence, similar to the general non-convex case, we show the best-known convergence rate in the leading term, up to a logarithmic factor. While the last term has higher-order constant dependencies, these are negligible due to the $T^{2}$ damping effect. Thus, the theory suggests that the algorithm’s convergence rate remains comparable to that of the uncompressed version, with only a constant factor affected by the compression rank.
To quantify the slowdown caused by compression, we again consider low-rank compression via SVD decomposition, as discussed earlier. As noted before, $C_{0}=\mathcal{O}(\frac{1}{1-q_{r}})=\mathcal{O}(\frac{d}{r})$, and similarly, from the expression of $C_{1}$, we have $C_{1}=\mathcal{O}(\frac{1}{1-q_{r}})=\mathcal{O}(\frac{d}{r})$. Therefore, the leading term of the rate becomes $\mathcal{O}((\frac{d}{r})^{2}\frac{\log T}{T})$, with a compression slowdown factor of $(\frac{d}{r})^{2}$. The detailed proof is deferred to Appendix E.3.

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

###### Theorem 2 .

## 5 LDAdam in Practice

### 5.1 Training Memory Footprint and Runtime

LDAdam’s memory savings over baseline Adam (Kingma & Ba, ) (see Table 1 for comparison) follow from the low dimensionality of the optimizer states: for gradients of shape $n\times m$ with $n\leq m$, the gradient statistics in a learning subspace of rank $r$ are stored as tensors of shape $r\times m$, and the additional projection matrix is of shape $n\times r$. As a result, the memory footprint of LDAdam’s optimizer states is only a fraction of the size of the model. For an $n\times n$ layer, the compression ratio of LDAdam relative to Adam is of $\frac{3}{2}\frac{r}{n}$. The Table 5 in the Appendix B reports peak memory for fine-tuning and pre-training tasks.

With the error buffer stored in the gradient variable, the error feedback mechanism is memory neutral. However, this comes at the cost of not supporting gradient clipping and per-layer weight updates (Lv et al., ; Zhao et al., ). The latter allows memory savings proportional to the model size, but at the expense of gradient accumulation. Therefore, although applicable to most optimizers, it is not widely adopted, due to the large batch size used for large language model training (Touvron et al., ) and the memory-intensive nature of activations in attention layers.

**Table 1: Optimizer comparison: parameter count during training for a weight layer of shape $n\times m$ with $n\leq m$ (i.e., left projection), training capabilities, and estimates of optimizer states memory footprint in half precision (models architecture is described in the Table 8).**
|  | Adam | LDAdam | GaLore |  |
| --- | --- | --- | --- | --- |
| Token count |  |  |  |  |
| Weights | $nm$ | $nm$ | $nm$ |  |
| Gradients | $nm$ | $nm$ |  |  |
| Optimizer States | $2nm$ | $nr+2rm$ | $nr+2rm$ |  |
| Gradient Clipping | ✓ | ✗ | ✗ |  |
| Gradient Accumulation | ✓ | ✓ | ✗ |  |
| Memory estimates |  |  |  |  |
| RoBERTa-base (r=8) | 0.46 GB | 0.15 GB | 0.15 GB |  |
| Llama 350M (r=256) | 1.37 GB | 0.95 GB | 0.95 GB |  |
| Llama-2 7B (r=32) | 25.1 GB | 1.22 GB | 1.22 GB |  |
| Llama-2 7B (r=512) | 25.1 GB | 4.87 GB | 4.87 GB |  |

#### Runtime.

Compared to Adam, GaLore adds gradient compression; and compared to GaLore, LDAdam adds optimizer states projection-aware update and the generalized error feedback mechanism. Yet, the practical runtime differences are not substantial: concretely, on Llama-350M pre-training, LDAdam has a 10% longer runtime compared to GaLore, and is 15% slower than Adam.
Table 6 in the Appendix B reports runtimes on fine-tuning and pre-training tasks.

### 5.2 Experiments

We empirically evaluate LDAdam for fine-tuning and pre-training large language models against baseline Adam (Kingma & Ba, ) and the memory-efficient GaLore (Zhao et al., ) (Algorithm 2). Similar to most parameter-efficient fine-tuning (PEFT) methods and to GaLore, we apply low-rank compression only to the two-dimensional matrices of the self-attention layers, e.g. wide embedding and output layers are trained using standard Adam optimizer. For a fair comparison, in all experiments we do not apply additional layer-wise learning rate scaling, and apply hyperparameters suggested in the papers introducing the algorithm. To ensure reproducibility, the full list of hyperparameters we used is provided in the Appendix C.

#### Fine-tuning on GLUE.

We evaluate LDAdam for fine-tuning RoBERTa-base model (Liu et al., ) on the General Language Understanding Evaluation (GLUE) benchmark (Wang et al., ). RoBERTa-base is a 125M parameters large language model, built upon the encoder only attention-based BERT architecture (Devlin et al., ), and pre-trained for natural language understanding (NLU). GLUE benchmark includes a wide variety of NLU tasks with training samples of various size. We mimic memory-constrained setups by training for only 3 epochs with a batch size of 16, and used rank r=8 compression for LDAdam and GaLore. For a fair comparison, we tune the learning rate over the set $\{$1e-5, 2e-5, $\ldots$, 5e-5$\}$. Table 2 reports the best average over 3 seeds (Table 7 in Appendix B.3 reports standard deviation), and Figure 2 shows that the error buffer norm is of the same order of magnitude as the gradient norm, and that their ratio is stable throughout training.

**Table 2: Results of fine-tuning RoBERTa-base model on the GLUE benchmark.**
|  | MRPC | STS-B | CoLA | SST-2 | QNLI | QQP | MNLI | Avg |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Training Samples | 3.7k | 7k | 8.5k | 67k | 105k | 364k | 393k |  |
| Adam | 88.97 | 90.08 | 58.57 | 94.34 | 92.81 | 91.38 | 87.85 | 86.28 |
| LDAdam | 88.40 | 90.11 | 59.91 | 95.00 | 92.87 | 91.28 | 87.81 | 86.48 |
| LDAdam no-EF | 88.00 | 89.88 | 56.86 | 95.00 | 92.32 | 89.75 | 86.99 | 85.54 |
| GaLore | 86.19 | 88.97 | 55.12 | 94.15 | 92.01 | 89.86 | 86.80 | 84.73 |

#### Fine-tuning on GSM8K.

We evaluate LDAdam for fine-tuning Llama-2 7B model (Touvron et al., ) on the GSM8K mathematical reasoning dataset (Cobbe et al., ). Llama models are attention-based decoder-only large language models pre-trained for a wide variety of natural language processing tasks, including text generation. The GSM8K dataset contains grade school math word problems and it is used to assess basic mathematical reasoning. We run experiments for ranks r=32 and r=512, on 3 epochs, with a batch size of 32, and tune the learning rate over the set $\{$5e-5, 6e-5, $\ldots$, 9e-5, 1e-4, 2e-4, …, 5e-4$\}$. Table 3 reports the best average accuracy. We also provide peak memory usage read from the GPU, it accounts for the optimizer states memory footprint, but also the fixed memory cost of storing model, gradient and activations.

**Table 3: Results of fine-tuning Llama-2 7B on the GSM8K dataset.**
|  | Adam | LDAdam | LDAdam no-EF | GaLore |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
|  |  | r=32 | r=512 | r=32 | r=512 | r=32 | r=512 |
| Accuracy | 34.72 | 32.53 | 35.86 | 30.70 | 35.78 | 27.07 | 35.18 |
| Peak Memory | 55.34 GB | 32.08 GB | 35.58 GB | 32.08 GB | 35.58 GB | 32.52 GB | 35.58 GB |

Across all fine-tuning experiments, LDAdam significantly outperforms GaLore, and is either on par or better than Adam. When ablating the LDAdam techniques, we observe that LDAdam without error feedback also outperforms GaLore, supporting LDAdam’s continuous subspace adaptation strategy over GaLore’s sequential update approach. The fact that we can match Adam once error feedback is enabled suggests that this technique is key to matching the practical performance of full-parameter training.

#### Pre-training on C4.

We evaluate LDAdam for pre-training Llama models (Touvron et al., ) on the C4 dataset (Raffel et al., ). Due to limited compute resources, we pre-train smaller 130M, 350M and 1.3B parameter variants (Lialin et al., ; Zhang & Sennrich, ; Shazeer, ) of the Llama-2 models. C4 is a clean version of Common Crawl’s web crawl corpus. We adjust the number of training steps according to the Chinchilla scaling law (Hoffmann et al., ), i.e. 20 training tokens per model parameter; and use batch size 512 and sequence length 256, accounting to 131 072 tokens per step, thus using the same settings as GaLore. For a fair comparison, we tune the learning rate over the set $\{$5e-4, 1e-3, 5e-3$\}$, and report the best results for ranks r=8 and r=256 in Table 4. Figure 1 depicts the pre-training dynamics.

**Table 4: Results (validation perplexity $\downarrow$) of pre-training Llama models on the C4 dataset.**
|  | Adam | LDAdam | LDAdam no-EF | GaLore |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
|  |  | r=8 | r=256 | r=8 | r=256 | r=8 | r=256 |
| Llama 130M | 24.64 | 23.82 | 22.65 | 45.17 | 25.32 | 67.78 | 26.04 |
| Llama 350M | 18.08 | 18.37 | 17.30 | - | 19.85 | - | 20.03 |

Furthermore, we scale the experiments up to the Llama 1.3B parameter model, using the same hyperparameters but fixing the learning rate at $5e-4$ due to limited computational resources. Adam achieves a validation perplexity of 14.86
, while LDAdam with rank r=16 achieves 14.09
.

Figure: Figure 1: Pre-training dynamics for Llama 350M (left) and Llama 1.3B (right) on the C4 dataset.
Refer to caption: extracted/5984787/figure/llama350m.png

Pre-training experiments also show that LDAdam has comparable performance to Adam, both in terms of convergence and validation accuracy, and confirm the improvements over GaLore. Further, for very high compression rates (i.e., r=8), LDAdam still converges, but its no-EF variant does not, nor does GaLore. This highlights the need for an error correction mechanism for pre-training tasks, which are known to require models updates of high rank. For lower compression (i.e. r=256), LDAdam
outperforms the no-EF version and GaLore, which is correlated to their consistent projection errors.
LDAdam even appears to slightly outperform Adam; we attribute this small improvement to the regularization effect of compression.

## 6 Discussion

In this work, we proposed a new low-rank learning method with the memory-efficient LDAdam optimizer, analyzed its convergence under standard assumptions, and provided empirical evidence of its ability to match the performance of Adam optimizer at a fraction of its memory cost.

Limitations.
The LDAdam optimizer relies on Adam’s heuristic to estimate gradient statistics, but its development required the design of a specific intermediate optimizer state update rule. Therefore, efficient implementations of Adam are not directly applicable to LDAdam, and extending our work to other optimizers requires further work.
On the experimental side, a natural next step would be to execute large-scale billion-parameter *pre-training* experiments. Doing so is feasible, but requires additional computational infrastructure and technical changes to the algorithm to allow for distributed training.

Impact Statement.
This paper presents work that aims to advance the field of machine learning. We believe that memory-efficient optimization is a step toward democratizing large-scale model training, and thus provides opportunities to foster both the development of new applications and the research in the field. There are many societal concerns about the rapidly growing use of artificial intelligence, we feel that none of them is specifically related to our work.

## Acknowledgements

The authors thank Théo Gaboriaud for his valuable feedback. 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.

## Appendix A Challenges for Estimating the Second-order Statistic of Projected Gradients

Adam Does Not Rely Solely on Linear Operations. Adam follows the proposition from AdaGrad (Duchi et al., ) to use diagonal approximation to estimate: $((1-\beta_{2})\sum_{\tau=0}^{t}\beta_{2}^{t-\tau}g_{\tau}\cdot g_{\tau}^{T})^{
-\frac{1}{2}}$. This leads to $v_{t}^{Adam}=(1-\beta_{2})\sum_{\tau=0}^{t}\beta_{2}^{t-\tau}g_{\tau}\odot g_{\tau}$ and to the matrix inverse square root being replaced by coordinate wise inverse square root operations. Such an approximation saves both memory and computation, but at the cost of leaving the linear framework. Therefore, one cannot expect to draw an autoregressive linear algebra formula to update the second moment estimate.

Adam is Anisotropic. Adaptive optimization learning rate tuning for each parameter is equivalent to a direction-wise rescaling of the training space. It implies that the choice of coordinate system affects the optimization steps. Let $\bar{\mathcal{U}}$ be a fixed orthogonal matrix which acts as rotation in the parameter space and consider the following two minimization problems:

$$ $(P_{I}):\arg\min_{\theta\in\mathbb{R}^{n\times m}}f(\theta)$ (7) $$

$$ $(P_{\bar{\mathcal{U}}}):\arg\min_{\tilde{\theta}\in\mathbb{R}^{n\times m}}f( \bar{\mathcal{U}}\cdot\tilde{\theta})$ (8) $$

These two problems are equivalent up to a change of parametrization. However, applying Adam to each of them will lead to optimization process that are not equivalent: adaptive estimate of the second-order statistics of the gradient leads to $\theta_{t}\neq\bar{\mathcal{U}}\cdot\tilde{\theta_{t}}$ :

$$ $\textstyle\theta_{t+1}=\theta_{t}-\eta_{t}\frac{\sqrt{1-\beta_{2}^{t}}}{1- \beta_{1}^{t}}\frac{(1-\beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}\nabla f( \theta_{t})}{\sqrt{(1-\beta_{2})\sum_{\tau=1}^{t}\beta_{2}^{t-\tau}(\nabla f( \theta_{t}))^{2}}+\epsilon}\approx\theta_{t}-\eta_{t}\frac{\mathbb{E}_{t,\beta _{1}}[\nabla f(\theta_{t})]}{\sqrt{\mathbb{E}_{t,\beta_{2}}[(\nabla f(\theta_{ t}))^{2}]}+\epsilon}$ (9) $$

$$ $\textstyle\tilde{\theta}_{t+1}=\tilde{\theta}_{t}-\eta_{t}\frac{\sqrt{1-\beta_ {2}^{t}}}{1-\beta_{1}^{t}}\frac{(1-\beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t- \tau}\bar{\mathcal{U}}^{\top}\nabla f(\bar{\mathcal{U}}\cdot\tilde{\theta}_{t} )}{\sqrt{(1-\beta_{2})\sum_{\tau=1}^{t}\beta_{2}^{t-\tau}(\bar{\mathcal{U}}^{ \top}\nabla f(\bar{\mathcal{U}}\cdot\tilde{\theta}_{t}))^{2}}+\epsilon}\approx \tilde{\theta}_{t}-\eta_{t}\frac{\mathbb{E}_{t,\beta_{1}}[\bar{\mathcal{U}}^{ \top}\nabla f(\bar{\mathcal{U}}\cdot\tilde{\theta}_{t})]}{\sqrt{\mathbb{E}_{t, \beta_{2}}[(\bar{\mathcal{U}}^{\top}\nabla f(\bar{\mathcal{U}}\cdot\tilde{ \theta}_{t}))^{2}]}+\epsilon}$ (10) $$

Low-rank adaptive optimization follows from $(P_{\mathcal{U}})$ where $\mathcal{U}\in\mathbb{R}^{n\times r}$ is the truncation of the rotation matrix $\bar{\mathcal{U}}$ to its first $r$ rows, and thus yields reparametrization and projection to a lower-dimensional space. However, this leads to $\tilde{\theta}\in\mathbb{R}^{r\times m}$ and the optimization set is restricted to a low-rank subspace. To enable full-parameter training, GaLore’s strategy is to keep the model parameter high-rank while performing low-dimensional updates, leading to:

$$ $\textstyle\theta_{t+1}^{\textsc{GaLore}}=\theta_{t}^{\textsc{GaLore}}-\mathcal {U}\cdot\left(\eta_{t}\frac{\sqrt{1-\beta_{2}^{t}}}{1-\beta_{1}^{t}}\frac{(1- \beta_{1})\sum_{\tau=1}^{t}\beta_{1}^{t-\tau}\mathcal{U}^{\top}\nabla f(\theta _{t}^{\textsc{GaLore}})}{\sqrt{(1-\beta_{2})\sum_{\tau=1}^{t}\beta_{2}^{t-\tau }(\mathcal{U}^{\top}\nabla f(\theta_{t}^{\textsc{GaLore}}))^{2}}+\epsilon} \right)\\ =\theta_{t}^{\textsc{GaLore}}-\mathcal{U}\cdot\left(\eta_{t}\frac{\hat{m}_{t}} {\sqrt{\hat{v}_{t}}+\epsilon}\right)$ (11) $$

With the projection map $\mathcal{U}=\mathcal{U}_{t}$ occasionally updated, GaLore’s low-rank gradient statistic estimate $\hat{m}_{t}$ (resp. $\hat{v}_{t}$) fails to approximate $\mathbb{E}_{t,\beta_{1}}[\mathcal{U}_{t}^{\top}\nabla f(\theta_{t}^{\textsc{
GaLore}})]$ (resp. $\mathbb{E}_{t,\beta_{2}}[(\mathcal{U}_{t}^{\top}\nabla f(\theta_{t}^{\textsc{
GaLore}}))^{2}]$), because gradients projected into different subspaces are accumulated on the same momentum buffers. A simple idea would be to fix a coordinate system for the momentum buffers. However, this would have to be high-dimensional and thus would prevent memory savings. The purpose of LDAdam’s projection-aware update rule is to adapt the optimizer states to the new learning subspace and coordinate system, leading to $\hat{m}_{t-1/2}\approx\mathbb{E}_{t-1,\beta_{1}}[\mathcal{U}_{t}^{\top}\nabla f
(\theta_{t-1}^{\textsc{LDAdam}})]$, and $\hat{v}_{t-1/2}\approx\mathbb{E}_{t-1,\beta_{2}}[(\mathcal{U}_{t}^{\top}\nabla
f
(\theta_{t-1}^{\textsc{LDAdam}}))^{2}]$ and enabling gradient accumulation within low-dimensional momentum buffers. One can thus rewrite LDAdam’s heuristic as performing the following update:

$$ $\textstyle\theta_{t+1}^{\textsc{LDAdam}}\approx\theta_{t}^{\textsc{LDAdam}}- \mathcal{U}_{t}\cdot\left(\eta_{t}\frac{\mathbb{E}_{t,\beta_{1}}[\mathcal{U}_{ t}^{\top}\nabla f(\theta_{t}^{\textsc{LDAdam}})]}{\sqrt{\mathbb{E}_{t,\beta_{2 }}[(\mathcal{U}_{t}^{\top}\nabla f(\theta_{t}^{\textsc{LDAdam}}))^{2}]}+ \epsilon}\right).$ (12) $$

## Appendix B Additional Measurements

### B.1 Peak Memory During Training

Table 5 reports peak memory for fine-tuning and pre-training on a single NVIDIA H100 80GB GPU with micro batch size 1 and without activation checkpointing.

**Table 5: Peak memory for fine-tuning (FT) and pre-training (PT) tasks for micro batch size of 1.**
| Model | Task | Rank | Optimizer |  |
| --- | --- | --- | --- | --- |
|  |  |  | Adam | LDAdam |
| RoBERTa-base | MNLI (FT) | r=8 | 2.39 GB | 1.60 GB |
| Llama 350M | C4 (PT) | r=8 | 4.20 GB | 2.67 GB |
| Llama 350M | C4 (PT) | r=256 | 4.20 GB | 3.01 GB |
| Llama 1.3B | C4 (PT) | r=16 | 15.20 GB | 8.67 GB |
| Llama-2 7B | GSM8K (FT) | r=32 | 57.62 GB | 31.98 GB |
| Llama-2 7B | GSM8K (FT) | r=512 | 57.62 GB | 35.90 GB |

### B.2 Training Runtimes

Table 6 reports runtime for fine-tuning and pre-training on a single (except for Llama 1.3B pre-training were we use four) NVIDIA H100 80GB GPU. We don’t use activation checkpointing, but use a micro batch size of 1 for GSM8K fine-tuning and accumulate gradients over 4 steps in pre-training experiments.

**Table 6: Runtime for fine-tuning (FT) and pre-training (PT) tasks.**
| Model | Task | Rank | Optimizer |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  |  |  | Adam | LDAdam | LDAdam no-EF | GaLore |
| RoBERTa-base | MNLI (FT) | r=8 | 37m | 01h 07m | 01h 01m | 56m |
| Llama 350M | C4 (PT) | r=8 | 21h 59m | 23h 12m | - | - |
| Llama 350M | C4 (PT) | r=256 | 21h 59m | 24h 50m | 24h 44m | 23h 09m |
| Llama 1.3B | C4 (PT) | r=16 | 04d 13h 09m | 05d 04h 24m | - | - |
| Llama-2 7B | GSM8K (FT) | r=32 | 50m | 56m | 55m | 01h 08m |
| Llama-2 7B | GSM8K (FT) | r=512 | 50m | 01h 08m | 01h 07m | 01h 09m |

### B.3 Fine-tuning Standard deviation

**Table 7: Standard deviation over 3 seeds for the results of fine-tuning RoBERTa-base model on the GLUE benchmark.**
|  | MRPC | STS-B | CoLA | SST-2 | QNLI | QQP | MNLI |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Training Samples | 3.7k | 7k | 8.5k | 67k | 105k | 364k | 393k |
| Adam | 1.068 | 0.226 | 2.604 | 0.517 | 0.074 | 0.050 | 0.167 |
| LDAdam | 0.617 | 0.211 | 0.633 | 0.463 | 0.275 | 0.080 | 0.177 |
| LDAdam no-EF | 0.849 | 0.555 | 0.903 | 0.066 | 0.285 | 0.687 | 0.087 |
| GaLore | 0.861 | 0.330 | 2.348 | 0.413 | 0.166 | 0.034 | 0.111 |

### B.4 Error Buffer Norm During Training

Figure: Figure 2: Error buffer norm and gradient norm during the fine-tuning of the RoBERTa-base model on the GLUE benchmark.
Refer to caption: extracted/5984787/figure/glue_norm.png

## Appendix C Materials for Reproducibility

### C.1 Optimizer States Memory Estimates

Table 8 describes the architecture of the trained models architecture and the low-rank structure of the low-dimensional optimizer states. It allows optimizer states token count, thus memory estimates in GB by applying factor $\frac{2}{1024^{3}}$ for half precision training.

**Table 8: Weights and low-dimensional optimizer states shape for used models.**
|  | Weight | Low-rank optimizer states |
| --- | --- | --- |
| RoBERTa-base |  |  |
| Token Embedding | 50265*768+768 | - |
| Positional Embedding | 564*768 | - |
| Attention Head | 12*4*768*768 | 12*4*3*768*r |
| MLP Block | 12*3*3072*768 | 12*3*(2*3072*r+768*r) |
| Normalization | 2*768+12*(9*768+3072) | - |
| Dense Layer | 768*768+768 | - |
| Output Layer | 768*n_label+n_label | - |
| Llama 130M |  |  |
| Embedding Layer | 32000*768 | - |
| Attention Head | 12*4*768*768 | 12*4*3*768*r |
| MLP Block | 12*3*2048*768 | 12*3*(2*2048*r+768*r) |
| Layer Normalization | (12*2+1)*768 | - |
| Output Layer | 768*32000 | - |
| Llama 350M |  |  |
| Embedding Layer | 32000*1024 | - |
| Attention Head | 24*4*1024*1024 | 24*4*3*1024*r - |
| MLP Block | 24*3*2736*1024 | 24*3*(2*2736*r+1024*r) |
| Layer Normalization | (24*2+1)*1024 | - |
| Output Layer | 1024*32000 |  |
| Llama-2 7B |  |  |
| Embedding Layer | 32000*4096 | - |
| Attention Head | 32*4*4096*4096 | 32*4*3*4096*r |
| MLP Block | 32*3*4096*11008 | 32*3*(2*11008*r+4096*r) |
| Layer Normalization | (32*2+1)*4096 | - |
| Output Layer | 4096*32000 | - |

### C.2 Fine-Tuning Hyperparameters

Tables 9 and 10 detail all the hyperparameters we use when fine-tuning respectively RoBERTa-base model on the GLUE benchmark and Llama-2 7B model on the GSM8K dataset.

**Table 9: Hyperparameters used for fine-tuning RoBERTa-base model on the GLUE benchmark.**
|  | Adam | LDAdam | GaLore |
| --- | --- | --- | --- |
| Epochs | 3 |  |  |
| Warm-up | ✗ |  |  |
| Batch Size | 16 |  |  |
| Maximum Length | 128 |  |  |
| Data Type | bfloat32 |  |  |
| Learning Rate | $\{1e-5,2e-5,\ldots,5e-5\}$ |  |  |
| Learning Rate Scheduling | linear to 0% |  |  |
| Decay Rate $\beta_{1}$ | 0.9 | 0.908 | 0.9 |
| Decay Rate $\beta_{2}$ | 0.999 | 0.99 | 0.999 |
| Weight Decay | ✗ | ✗ | ✗ |
| Dropout | ✗ | ✗ | ✗ |
| Gradient Clipping | ✗ | ✗ | ✗ |
| Interpolation Factor $\rho$ | ✗ | 0.908 | ✗ |
| Error Feedback | ✗ | ✓ | ✗ |
| Subspace Frequency | ✗ | 1 | 500 |
| Learning Rate Scaling | ✗ | ✗ | ✗ |

**Table 10: Hyperparameters used for fine-tuning Llama-2 7B model on the GSM8K dataset.**
|  | Adam | LDAdam | GaLore |
| --- | --- | --- | --- |
| Epochs | 3 |  |  |
| Training Steps | 702 |  |  |
| Warm-up Steps | 20 |  |  |
| Batch Size | 32 |  |  |
| Maximum Length | 512 |  |  |
| Data Type | bfloat16 |  |  |
| Learning Rate | $\{5e-5,6e-5,\ldots,9e-5\}$ |  |  |
|  | $\{1e-4,2e-4,\ldots,5e-4\}$ |  |  |
| Warm-up Scheduling | linear from 0% |  |  |
| Learning Rate Scheduling | linear to 0% |  |  |
| Decay Rate $\beta_{1}$ | 0.9 | 0.908 | 0.9 |
| Decay Rate $\beta_{2}$ | 0.999 | 0.99 | 0.999 |
| Weight Decay | ✗ | ✗ | ✗ |
| Dropout | ✗ | ✗ | ✗ |
| Gradient Clipping | 1.0 | ✗ | ✗ |
| Interpolation Factor $\rho$ | ✗ | 0.908 | ✗ |
| Error Feedback | ✗ | ✓ | ✗ |
| Subspace Frequency | ✗ | 1 | 200 |
| Learning Rate Scaling | ✗ | ✗ | ✗ |

### C.3 Pre-Training Hyperparameters

Table 11 details all the hyperparameters we use when pre-training Llama models on the C4 dataset.

**Table 11: Hyperparameters used for pre-training Llama models on the C4 dataset.**
|  | Llama 130M | Llama 350M | Llama 1.3B |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | Adam | LDAdam | GaLore | Adam | LDAdam | GaLore | Adam | LDAdam |
| Training Steps | 20000 | 55000 | 200000 |  |  |  |  |  |
| Warm-up Steps | 2000 | 5500 | 10000 |  |  |  |  |  |
| Maximum Length | 256 | 256 | 256 |  |  |  |  |  |
| Batch Size | 512 | 512 | 512 |  |  |  |  |  |
| Token Batch Size | 131 072 | 131 072 | 131 072 |  |  |  |  |  |
| Total Training Tokens | 2 621 440 000 | 7 208 960 000 | 26 621 440 000 |  |  |  |  |  |
| Data Type | bfloat16 | bfloat16 | bfloat16 |  |  |  |  |  |
| Learning Rate | $\{5e-4,1e-3,5e-3\}$ | $\{5e-4,1e-3,5e-3\}$ | $5e-4$ |  |  |  |  |  |
| Warm-up Scheduling | linear from 0% | linear from 0% | linear from 0% |  |  |  |  |  |
| Learning Rate Scheduling | cosine to 10% | cosine to 10% | cosine to 10% |  |  |  |  |  |
| Decay Rate $\beta_{1}$ | 0.9 | 0.908 | 0.9 | 0.9 | 0.908 | 0.9 | 0.9 | 0.908 |
| Decay Rate $\beta_{2}$ | 0.999 | 0.98 | 0.999 | 0.999 | 0.98 | 0.999 | 0.999 | 0.98 |
| Weight Decay | ✗ | ✗ | ✗ | ✗ | ✗ | ✗ | ✗ | ✗ |
| Dropout | ✗ | ✗ | ✗ | ✗ | ✗ | ✗ | ✗ | ✗ |
| Gradient Clipping | 1.0 | ✗ | ✗ | 1.0 | ✗ | ✗ | 1.0 | ✗ |
| Interpolation Factor $\rho$ | ✗ | 0.908 | ✗ | ✗ | 0.908 | ✗ | ✗ | 0.908 |
| Error Feedback | ✗ | ✓ | ✗ | ✗ | ✓ | ✗ | ✗ | ✓ |
| Subspace Frequency | ✗ | 1 | 200 | ✗ | 1 | 200 | ✗ | 1 |
| Learning Rate Scaling | ✗ | ✗ | ✗ | ✗ | ✗ | ✗ | ✗ | ✗ |

## Appendix D GaLore Algorithm

Algorithm 2 describes GaLore’s (Zhao et al., ) proposal for performing an Adam-type update from a low-rank gradient projection. Note that the learning rate rescaling induced by multiplying by $\alpha\neq 1$ is equivalent to using different learning rates for different layers. This level of hyperparameter tuning is not usual, so for a fair comparison we use a single learning rate for all layers. However, our preliminary experiments suggest that LDAdam would also benefit from layer-wise learning rate tuning. Other than that, in our experiments we use the suggested hyperparameters.

Figure: Algorithm 2 GaLore

## Appendix E Deferred Proofs

Figure: Algorithm 3 LDAdam: Analytical View

### E.1 Key Lemmas

###### Lemma 1 .

With $\Sigma_{T}=T\sigma^{2}+\sum_{t=1}^{T}\mathbb{E}\left[\|\nabla f(\theta_{t})\|^
{2}\right]$, for any $t\geq 1$ the following bounds hold:

$$ $\displaystyle\|B_{t}\|\leq\frac{G}{1-q_{r}},\quad\sum_{t=1}^{T}\mathbb{E}\left [\|B_{t}\|^{2}\right]\leq\frac{1}{(1-q_{r})^{2}}\Sigma_{T}$ (13) $\displaystyle\|\xi_{t}\|\leq\frac{q_{r}G}{(1-\beta_{1})(1-q_{r})},\quad\sum_{t =1}^{T}\mathbb{E}\left[\|\xi_{t}\|^{2}\right]\leq\frac{q_{r}^{2}}{(1-\beta_{1} )^{2}(1-q_{r})^{2}}\Sigma_{T}$ (14) $\displaystyle\left\|\frac{\beta_{1}}{1-\beta_{1}}B_{t}+(1-\beta_{1})\xi_{t+1} \right\|\leq\frac{\beta_{1}+(1-\beta_{1})q_{r}}{(1-\beta_{1})(1-q_{r})}G$ (15) $\displaystyle\sum_{t=1}^{T}\mathbb{E}\left[\left\|\frac{\beta_{1}}{1-\beta_{1} }B_{t-1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]\leq\frac{\beta_{1}+(1-\beta_{ 1})q_{r}^{2}}{(1-\beta_{1})^{2}(1-q_{r})^{2}}\Sigma_{T}.$ (16) $$

###### Proof.

Let us start with the proof of the first bound on $B_{t}$. Denote $\tilde{q}_{r}:=(1-\beta_{1})q_{r}+\beta_{1}\in[q_{r},1)$.

$$ $\displaystyle\|B_{t+1}\|$ $\displaystyle=$ $\displaystyle\left\|\beta_{1}M_{t}+(1-\beta_{1})A_{t+1}\right\|$ $\displaystyle=$ $\displaystyle\left\|\beta_{1}M_{t}+(1-\beta_{1})g_{t+1}+(1-\beta_{1})\xi_{t+1}\right\|$ $\displaystyle=$ $\displaystyle\left\|-(1-\beta_{1})M_{t}+(1-\beta_{1})g_{t+1}+B_{t}\right\|$ $\displaystyle=$ $\displaystyle\left\|-(1-\beta_{1})\mathcal{U}_{t}\mathcal{U}^{\top}_{t}B_{t}+( 1-\beta_{1})g_{t+1}+B_{t}\right\|$ $\displaystyle=$ $\displaystyle\left\|(1-\beta_{1})(B_{t}-\mathcal{U}_{t}\mathcal{U}^{\top}_{t}B _{t})+(1-\beta_{1})g_{t+1}+\beta_{1}B_{t}\right\|$ $\displaystyle\leq$ $\displaystyle(1-\beta_{1})\|B_{t}-\mathcal{U}_{t}\mathcal{U}^{\top}_{t}B_{t}\| +(1-\beta_{1})\|g_{t+1}\|+\beta_{1}\|B_{t}\|$ $\displaystyle\leq$ $\displaystyle\left((1-\beta_{1})q_{r}+\beta_{1}\right)\|B_{t}\|+(1-\beta_{1}) \|g_{t+1}\|$ $\displaystyle=$ $\displaystyle\tilde{q}_{r}\|B_{t}\|+(1-\beta_{1})\|g_{t+1}\|$ $\displaystyle=$ $\displaystyle\tilde{q}_{r}^{t}\|B_{1}\|+(1-\beta_{1})\sum_{\tau=2}^{t+1}\tilde {q}_{r}^{t+1-\tau}\|g_{\tau}\|=(1-\beta_{1})\sum_{\tau=1}^{t+1}\tilde{q}_{r}^{ t+1-\tau}\|g_{\tau}\|.$ $$

Using the bounded gradient assumption, we get

$$ $\|B_{t}\|\leq(1-\beta_{1})G\sum_{\tau=1}^{t}\tilde{q}_{r}^{t-\tau}\leq\frac{(1 -\beta_{1})G}{1-\tilde{q}_{r}}=\frac{G}{1-q_{r}}.$ $$

To derive the bound with expectation, we apply Cauchy-Schwartz inequality and the bounded variance assumption:

$$ $\displaystyle\sum_{t=1}^{T}\mathbb{E}\left[\|B_{t}\|^{2}\right]$ $\displaystyle\leq$ $\displaystyle(1-\beta_{1})^{2}\sum_{t=1}^{T}\mathbb{E}\left[\left(\sum_{\tau=1 }^{t}\tilde{q}_{r}^{t-\tau}\|g_{\tau}\|\right)^{2}\right]$ $\displaystyle\leq$ $\displaystyle(1-\beta_{1})^{2}\sum_{t=1}^{T}\mathbb{E}\left[\left(\sum_{\tau=1 }^{t}\tilde{q}_{r}^{t-\tau}\right)\left(\sum_{\tau=1}^{t}\tilde{q}_{r}^{t-\tau }\|g_{\tau}\|^{2}\right)\right]$ $\displaystyle\leq$ $\displaystyle\frac{(1-\beta_{1})^{2}}{1-\tilde{q}_{r}}\sum_{t=1}^{T}\sum_{\tau =1}^{t}\tilde{q}_{r}^{t-\tau}\mathbb{E}\left[\|g_{\tau}\|^{2}\right]$ $\displaystyle\leq$ $\displaystyle\frac{(1-\beta_{1})^{2}}{(1-\tilde{q}_{r})^{2}}\sum_{t=1}^{T} \mathbb{E}\left[\|g_{t}\|^{2}\right]$ $\displaystyle=$ $\displaystyle\frac{1}{(1-q_{r})^{2}}\sum_{t=1}^{T}\mathbb{E}\left[\|g_{t}- \nabla f(\theta_{t})+\nabla f(\theta_{t})\|^{2}\right]$ $\displaystyle\leq$ $\displaystyle\frac{1}{(1-q_{r})^{2}}\sum_{t=1}^{T}\left(\sigma^{2}+\mathbb{E} \left[\|\nabla f(\theta_{t})\|^{2}\right]\right)$ $\displaystyle\leq$ $\displaystyle\frac{T\sigma^{2}}{(1-q_{r})^{2}}+\frac{1}{(1-q_{r})^{2}}\sum_{t= 1}^{T}\mathbb{E}\left[\|\nabla f(\theta_{t})\|^{2}\right]$ $$

To bound the norm of the error $\|\xi_{t}\|$, notice that

$$ $\|\xi_{t+1}\|=\frac{1}{1-\beta_{1}}\|B_{t}-\mathcal{U}_{t}\mathcal{U}^{\top}_{ t}B_{t}\|\leq\frac{q_{r}}{1-\beta_{1}}\|B_{t}\|\leq\frac{q_{r}G}{(1-\beta_{1}) (1-q_{r})}.$ (17) $$

To get the bound with expectation, we apply previous inequality on $B_{t}$ and get

$$ $\displaystyle\sum_{t=1}^{T}\mathbb{E}\left[\|\xi_{t}\|^{2}\right]$ $\displaystyle=$ $\displaystyle\sum_{t=1}^{T-1}\mathbb{E}\left[\|\xi_{t+1}\|^{2}\right]$ $\displaystyle\leq$ $\displaystyle\sum_{t=1}^{T-1}\frac{q_{r}^{2}}{(1-\beta_{1})^{2}}\mathbb{E} \left[\|B_{t}\|^{2}\right]$ $\displaystyle\leq$ $\displaystyle\frac{q_{r}^{2}}{(1-\beta_{1})^{2}(1-q_{r})^{2}}\left(T\sigma^{2} +\sum_{t=1}^{T}\mathbb{E}\left[\|\nabla f(\theta_{t})\|^{2}\right]\right)$ $$

The fifth bound (15) follows from the triangle inequality and combining the obtained two bounds. For the last bound with expectation, we have

$$ $\displaystyle\sum_{t=1}^{T}\mathbb{E}\left[\left\|\frac{\beta_{1}}{1-\beta_{1} }B_{t-1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]$ $\displaystyle\leq$ $\displaystyle\sum_{t=1}^{T}\left(\frac{\beta_{1}}{(1-\beta_{1})^{2}}\mathbb{E} \left[\|B_{t-1}\|^{2}\right]+(1-\beta_{1})\mathbb{E}\left[\|\xi_{t}\|^{2} \right]\right)$ $\displaystyle\leq$ $\displaystyle\frac{\beta_{1}}{(1-\beta_{1})^{2}}\left(\frac{T\sigma^{2}}{(1-q_ {r})^{2}}+\frac{1}{(1-q_{r})^{2}}\sum_{t=1}^{T}\mathbb{E}\left[\|\nabla f( \theta_{t})\|^{2}\right]\right)$ $\displaystyle\qquad+\frac{q_{r}^{2}}{(1-\beta_{1})(1-q_{r})^{2}}\left(T\sigma^ {2}+\sum_{t=1}^{T}\mathbb{E}\left[\|\nabla f(\theta_{t})\|^{2}\right]\right)$ $\displaystyle=$ $\displaystyle\frac{\beta_{1}+(1-\beta_{1})q_{r}^{2}}{(1-\beta_{1})^{2}(1-q_{r} )^{2}}\left(T\sigma^{2}+\sum_{t=1}^{T}\mathbb{E}\left[\|\nabla f(\theta_{t})\| ^{2}\right]\right).$ $$

∎

###### Lemma 2 .

If $\gamma<1$ and $1-\gamma\leq\frac{1}{2}(1-\beta_{1})(1-q_{r})$, then

$$ $\sum_{t=1}^{T}\gamma^{T-t}\mathbb{E}\left[\left\|\frac{\beta_{1}}{1-\beta_{1}} B_{t-1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]\leq 2C_{2}\left(\frac{\sigma^{ 2}}{1-\gamma}+\sum_{t=1}^{T}\gamma^{T-t}\mathbb{E}\left[\|\nabla f(\theta_{t}) \|^{2}\right]\right),$ (18) $$

where $C_{2}=\frac{\beta_{1}+(1-\beta_{1})q_{r}^{2}}{(1-\beta_{1})^{2}(1-q_{r})^{2}}$.

###### Proof.

From the condition on $\gamma$ and the notation $\tilde{q}_{r}:=(1-\beta_{1})q_{r}+\beta_{1}\in[q_{r},1)$ from the previous proof, we have $1-\gamma\leq\frac{1-\tilde{q}_{r}}{2}$ or equivalently $\gamma-\tilde{q}_{r}\geq\frac{1-\tilde{q}_{r}}{2}$. Using this inequality on $\gamma$ and the previous bound (16) for $\mathbb{E}\left[\|B_{t}\|^{2}\right]$, we have

$$ $\displaystyle\sum_{t=1}^{T}\gamma^{T-t}\mathbb{E}\left[\|B_{t}\|^{2}\right]$ $\displaystyle\leq$ $\displaystyle(1-\beta_{1})^{2}\sum_{t=1}^{T}\gamma^{T-t}\mathbb{E}\left[\left( \sum_{\tau=1}^{t}\tilde{q}_{r}^{t-\tau}\|g_{\tau}\|\right)^{2}\right]$ $\displaystyle\leq$ $\displaystyle(1-\beta_{1})^{2}\sum_{t=1}^{T}\gamma^{T-t}\mathbb{E}\left[\left( \sum_{\tau=1}^{t}\tilde{q}_{r}^{t-\tau}\right)\left(\sum_{\tau=1}^{t}\tilde{q} _{r}^{t-\tau}\|g_{\tau}\|^{2}\right)\right]$ $\displaystyle\leq$ $\displaystyle\frac{(1-\beta_{1})^{2}}{1-\tilde{q}_{r}}\sum_{t=1}^{T}\sum_{\tau =1}^{t}\gamma^{T-t}\tilde{q}_{r}^{t-\tau}\mathbb{E}\left[\|g_{\tau}\|^{2}\right]$ $\displaystyle=$ $\displaystyle\frac{(1-\beta_{1})^{2}}{1-\tilde{q}_{r}}\sum_{\tau=1}^{T}\sum_{t =\tau}^{T}\gamma^{T-t}\tilde{q}_{r}^{t-\tau}\mathbb{E}\left[\|g_{\tau}\|^{2}\right]$ $\displaystyle=$ $\displaystyle\frac{(1-\beta_{1})^{2}}{1-\tilde{q}_{r}}\sum_{\tau=1}^{T}\left( \sum_{t=\tau}^{T}\gamma^{\tau-t}\tilde{q}_{r}^{t-\tau}\right)\gamma^{T-\tau} \mathbb{E}\left[\|g_{\tau}\|^{2}\right]$ $\displaystyle=$ $\displaystyle\frac{(1-\beta_{1})^{2}}{1-\tilde{q}_{r}}\sum_{\tau=1}^{T}\frac{1 -(\tilde{q}_{r}/\gamma)^{T-\tau+1}}{1-\tilde{q}_{r}/\gamma}\gamma^{T-\tau} \mathbb{E}\left[\|g_{\tau}\|^{2}\right]$ $\displaystyle\leq$ $\displaystyle\frac{(1-\beta_{1})^{2}}{(1-\tilde{q}_{r})(1-\tilde{q}_{r}/\gamma )}\sum_{t=1}^{T}\gamma^{T-\tau}\mathbb{E}\left[\|g_{t}\|^{2}\right]$ $\displaystyle\leq$ $\displaystyle\frac{2(1-\beta_{1})^{2}}{(1-\tilde{q}_{r})^{2}}\sum_{t=1}^{T} \gamma^{T-\tau}\mathbb{E}\left[\|g_{t}\|^{2}\right]$ $\displaystyle=$ $\displaystyle\frac{2}{(1-q_{r})^{2}}\sum_{t=1}^{T}\gamma^{T-\tau}\mathbb{E} \left[\|g_{t}-\nabla f(\theta_{t})+\nabla f(\theta_{t})\|^{2}\right]$ $\displaystyle\leq$ $\displaystyle\frac{2}{(1-q_{r})^{2}}\sum_{t=1}^{T}\gamma^{T-\tau}\left(\sigma^ {2}+\mathbb{E}\left[\|\nabla f(\theta_{t})\|^{2}\right]\right)$ $\displaystyle\leq$ $\displaystyle\frac{2}{(1-q_{r})^{2}}\left(\frac{\sigma^{2}}{1-\gamma}+\sum_{t= 1}^{T}\gamma^{T-\tau}\mathbb{E}\left[\|\nabla f(\theta_{t})\|^{2}\right]\right).$ $$

Using the bound (17) for $\|\xi_{t}\|^{2}$, we have

$$ $\displaystyle\sum_{t=1}^{T}\gamma^{T-\tau}\mathbb{E}\left[\|\xi_{t}\|^{2} \right]=\frac{2q_{r}^{2}}{(1-\beta_{1})^{2}(1-q_{r})^{2}}\left(\frac{\sigma^{2 }}{1-\gamma}+\sum_{t=1}^{T}\gamma^{T-\tau}\mathbb{E}\left[\|\nabla f(\theta_{t })\|^{2}\right]\right).$ $$

Combining these two bounds we get

$$ $\displaystyle\sum_{t=1}^{T}\gamma^{T-t}\mathbb{E}\left[\left\|\frac{\beta_{1}} {1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]$ $\displaystyle\qquad\qquad\leq\sum_{t=1}^{T}\gamma^{T-t}\left(\frac{\beta_{1}}{ (1-\beta_{1})^{2}}\mathbb{E}\left[\|B_{t-1}\|^{2}\right]+(1-\beta_{1})\mathbb{ E}\left[\|\xi_{t}\|^{2}\right]\right)$ $\displaystyle\qquad\qquad\leq\frac{\beta_{1}}{(1-\beta_{1})^{2}}\frac{2}{(1-q_ {r})^{2}}\left(\frac{\sigma^{2}}{1-\gamma}+\sum_{t=1}^{T}\gamma^{T-t}\mathbb{E }\left[\|\nabla f(\theta_{t})\|^{2}\right]\right)$ $\displaystyle\qquad\qquad\qquad+\frac{2(1-\beta_{1})q_{r}^{2}}{(1-\beta_{1})^{ 2}(1-q_{r})^{2}}\left(\frac{\sigma^{2}}{1-\gamma}+\sum_{t=1}^{T}\gamma^{T-t} \mathbb{E}\left[\|\nabla f(\theta_{t})\|^{2}\right]\right)$ $\displaystyle\qquad\qquad=\frac{2(\beta_{1}+(1-\beta_{1})q_{r}^{2})}{(1-\beta_ {1})^{2}(1-q_{r})^{2}}\left(\frac{\sigma^{2}}{1-\gamma}+\sum_{t=1}^{T}\gamma^{ T-t}\mathbb{E}\left[\|\nabla f(\theta_{t})\|^{2}\right]\right).$ $$

∎

###### Lemma 3 .

For $\Delta\Gamma_{t}=\Gamma_{t-1}-\Gamma_{t}$ we have

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

###### Proof.

From the definitions of $\Gamma_{t}$ (19) and $\tilde{v}_{t}=\max(v_{t},\|\tilde{v}_{t-1}\|_{\max})$ we have

$$ $\displaystyle\Gamma_{t}$ $\displaystyle=$ $\displaystyle\bar{\mathcal{U}}_{t}\mathrm{Diag}^{-1/2}(\tilde{v}_{t}+\epsilon, \|\tilde{v}_{t}\|_{\min}+\epsilon)\bar{\mathcal{U}}^{\top}_{t}$ $\displaystyle\preceq$ $\displaystyle\bar{\mathcal{U}}_{t}\mathrm{Diag}^{-1/2}(\|\tilde{v}_{t}\|_{\min }+\epsilon)\bar{\mathcal{U}}^{\top}_{t}$ $\displaystyle=$ $\displaystyle\frac{1}{\sqrt{\|\tilde{v}_{t}\|_{\min}+\epsilon}}\bar{\mathcal{U }}_{t}\bar{\mathcal{U}}^{\top}_{t}=\frac{1}{\sqrt{\|\tilde{v}_{t}\|_{\min}+ \epsilon}}I$ $\displaystyle\preceq$ $\displaystyle\frac{1}{\sqrt{\|\tilde{v}_{t-1}\|_{\max}+\epsilon}}I=\frac{1}{ \sqrt{\|\tilde{v}_{t-1}\|_{\max}+\epsilon}}\bar{\mathcal{U}}_{t-1}\bar{ \mathcal{U}}^{\top}_{t-1}$ $\displaystyle=$ $\displaystyle\bar{\mathcal{U}}_{t-1}\mathrm{Diag}^{-1/2}(\|\tilde{v}_{t-1}\|_{ \max}+\epsilon)\bar{\mathcal{U}}^{\top}_{t-1}$ $\displaystyle\preceq$ $\displaystyle\bar{\mathcal{U}}_{t-1}\mathrm{Diag}^{-1/2}(\tilde{v}_{t-1}+ \epsilon,\|\tilde{v}_{t-1}\|_{\min}+\epsilon)\bar{\mathcal{U}}^{\top}_{t-1}= \Gamma_{t-1},$ $$

which implies that $\Delta\Gamma_{t}=\Gamma_{t-1}-\Gamma_{t}$ is positive semidefinite. Hence, $\|\Delta\Gamma_{t}\|=\lambda_{\max}(\Delta\Gamma_{t})\geq 0$. Using the convexity of $\lambda_{\max}$ over symmetric matrices, we get

$$ $\displaystyle\sum_{t=1}^{T}\|\Delta\Gamma_{t}\|$ $\displaystyle=$ $\displaystyle\sum_{t=1}^{T}\lambda_{\max}(\Gamma_{t-1}-\Gamma_{t})$ $\displaystyle\leq$ $\displaystyle\sum_{t=1}^{T}\lambda_{\max}(\Gamma_{t-1})+\lambda_{\max}(-\Gamma _{t})=\sum_{t=1}^{T}\lambda_{\max}(\Gamma_{t-1})-\lambda_{\min}(\Gamma_{t})$ $\displaystyle=$ $\displaystyle\sum_{t=1}^{T}\frac{1}{\sqrt{\|\tilde{v}_{t-1}\|_{\min}+\epsilon} }-\frac{1}{\sqrt{\|\tilde{v}_{t}\|_{\max}+\epsilon}}$ $\displaystyle=$ $\displaystyle\sum_{t=1}^{T}\frac{1}{\sqrt{\|\tilde{v}_{t-1}\|_{\min}+\epsilon} }-\frac{1}{\sqrt{\|\tilde{v}_{t-1}\|_{\max}+\epsilon}}+\frac{1}{\sqrt{\|\tilde {v}_{t-1}\|_{\max}+\epsilon}}-\frac{1}{\sqrt{\|\tilde{v}_{t}\|_{\max}+\epsilon}}$ $\displaystyle\leq$ $\displaystyle\sum_{t=1}^{T}\frac{1}{\sqrt{\|\tilde{v}_{t-1}\|_{\min}+\epsilon} }-\frac{1}{\sqrt{\|\tilde{v}_{t}\|_{\min}+\epsilon}}+\sum_{t=1}^{T}\frac{1}{ \sqrt{\|\tilde{v}_{t-1}\|_{\max}+\epsilon}}-\frac{1}{\sqrt{\|\tilde{v}_{t}\|_{ \max}+\epsilon}}$ $\displaystyle=$ $\displaystyle\frac{1}{\sqrt{\|\tilde{v}_{0}\|_{\min}+\epsilon}}-\frac{1}{\sqrt {\|\tilde{v}_{T}\|_{\min}+\epsilon}}+\frac{1}{\sqrt{\|\tilde{v}_{0}\|_{\max}+ \epsilon}}-\frac{1}{\sqrt{\|\tilde{v}_{T}\|_{\max}+\epsilon}}$ $\displaystyle\leq$ $\displaystyle\frac{2}{\sqrt{\|\tilde{v}_{0}\|_{\min}+\epsilon}}\leq\frac{2}{ \sqrt{\epsilon}}.$ $$

For the second sum of squared norms, notice that for scalars $a\geq b\geq 0$, it holds that

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

Therefore, the above derivation can be repeated without the square roots as follows:

$$ $\displaystyle\sum_{t=1}^{T}\|\Delta\Gamma_{t}\|^{2}$ $\displaystyle=$ $\displaystyle\sum_{t=1}^{T}\left(\lambda_{\max}(\Gamma_{t-1}-\Gamma_{t})\right )^{2}$ $\displaystyle\leq$ $\displaystyle\sum_{t=1}^{T}\left(\lambda_{\max}(\Gamma_{t-1})+\lambda_{\max}(- \Gamma_{t})\right)^{2}=\sum_{t=1}^{T}\left(\lambda_{\max}(\Gamma_{t-1})- \lambda_{\min}(\Gamma_{t})\right)^{2}$ $\displaystyle\leq$ $\displaystyle\sum_{t=1}^{T}\left(\lambda_{\max}(\Gamma_{t-1})\right)^{2}-\left (\lambda_{\min}(\Gamma_{t})\right)^{2}$ $\displaystyle=$ $\displaystyle\sum_{t=1}^{T}\frac{1}{\|\tilde{v}_{t-1}\|_{\min}+\epsilon}-\frac {1}{\|\tilde{v}_{t}\|_{\max}+\epsilon}$ $\displaystyle=$ $\displaystyle\sum_{t=1}^{T}\frac{1}{\|\tilde{v}_{t-1}\|_{\min}+\epsilon}-\frac {1}{\|\tilde{v}_{t-1}\|_{\max}+\epsilon}+\frac{1}{\|\tilde{v}_{t-1}\|_{\max}+ \epsilon}-\frac{1}{\|\tilde{v}_{t}\|_{\max}+\epsilon}$ $\displaystyle\leq$ $\displaystyle\sum_{t=1}^{T}\frac{1}{\|\tilde{v}_{t-1}\|_{\min}+\epsilon}-\frac {1}{\|\tilde{v}_{t}\|_{\min}+\epsilon}+\sum_{t=1}^{T}\frac{1}{\|\tilde{v}_{t-1 }\|_{\max}+\epsilon}-\frac{1}{\|\tilde{v}_{t}\|_{\max}+\epsilon}$ $\displaystyle=$ $\displaystyle\frac{1}{\|\tilde{v}_{0}\|_{\min}+\epsilon}-\frac{1}{\|\tilde{v}_ {T}\|_{\min}+\epsilon}+\frac{1}{\|\tilde{v}_{0}\|_{\max}+\epsilon}-\frac{1}{\| \tilde{v}_{T}\|_{\max}+\epsilon}$ $\displaystyle\leq$ $\displaystyle\frac{2}{\|\tilde{v}_{0}\|_{\min}+\epsilon}\leq\frac{2}{\epsilon},$ $$

which completes the proof.
∎

###### Lemma 4 .

For all iterates $t\geq 1$ the following bound holds

$$ $\|\tilde{v}_{t}\|_{\max}\leq\frac{1+\beta_{2}}{1-\beta_{2}}\frac{(1-\beta_{1}( 1-q_{r}))^{2}}{(1-\beta_{1})^{2}(1-q_{r})^{2}}G^{2}.$ $$

###### Proof.

First, let us bound the $a_{t}$ term using Lemma 1 and Assumption 2.

$$ $\|a_{t}\|=\|\mathcal{U}^{\top}_{t}A_{t}\|\leq\|A_{t}\|\leq\|\xi_{t}\|+\|g_{t} \|\leq\frac{1-(1-q_{r})\beta_{1}}{(1-\beta_{1})(1-q_{r})}G=:CG.$ $$

Next, we bound momentum $m_{t}$:

$$ $\displaystyle\|m_{t}\|$ $\displaystyle\leq$ $\displaystyle\beta_{1}\|m_{t-1/2}\|+(1-\beta_{1})\|a_{t}\|$ $\displaystyle\leq$ $\displaystyle\beta_{1}\|m_{t-1}\|+(1-\beta_{1})CG$ $\displaystyle\leq$ $\displaystyle\beta_{1}^{t}\|m_{0}\|+(1-\beta_{1})CG\sum_{\tau=0}^{t-1}\beta_{1 }^{\tau}=(1-\beta_{1}^{t})CG.$ $$

Next, we bound the intermediate term $\|v_{t-1/2}\|_{1}$. Note that by the triangle inequality we have the following direct bound for it:

$$ $\left\|\left(\mathcal{U}^{\top}_{t}\mathcal{U}_{t-1}\right)^{2}v_{t-1}\right\| _{1}+\left\|\left(\mathcal{U}^{\top}_{t}\mathcal{U}_{t-1}\right)^{2}\left( \frac{m_{t-1}}{1-\beta_{1}^{t-1}}\right)^{2}\right\|_{1}+\left\|\left(\mathcal {U}^{\top}_{t}\mathcal{U}_{t-1}\cdot\frac{m_{t-1}}{1-\beta_{1}^{t-1}}\right)^{ 2}\right\|_{1}$ $$

Now let us bound each term separately. For the first term we have

$$ $\displaystyle\left\|\left(\mathcal{U}^{\top}_{t}\mathcal{U}_{t-1}\right)^{2}v_ {t-1}\right\|_{1}$ $\displaystyle=$ $\displaystyle\sum_{i=1}^{r}\sum_{j=1}^{r}\left\langle\mathcal{U}_{t}[:,i], \mathcal{U}_{t-1}[:,j]\right\rangle^{2}v_{t-1,j}$ $\displaystyle\leq$ $\displaystyle\sum_{j=1}^{r}v_{t-1,j}\cdot\max_{1\leq j\leq r}\sum_{i=1}^{r} \left\langle\mathcal{U}_{t}[:,i],\mathcal{U}_{t-1}[:,j]\right\rangle^{2}$ $\displaystyle=$ $\displaystyle\|v_{t-1}\|_{1}\max_{1\leq j\leq r}\left\|\mathcal{U}^{\top}_{t} \mathcal{U}_{t-1}[:,j]\right\|^{2}$ $\displaystyle\leq$ $\displaystyle\|v_{t-1}\|_{1}\max_{1\leq j\leq r}\left\|\mathcal{U}_{t-1}[:,j] \right\|^{2}=\|v_{t-1}\|_{1},$ $$

since columns $\mathcal{U}_{t}[:,i]$ and $\mathcal{U}_{t-1}[:,j]$ have unit length by construction. Similarly, for the second term we have

$$ $\displaystyle\left\|\left(\mathcal{U}^{\top}_{t}\mathcal{U}_{t-1}\right)^{2} \left(\frac{m_{t-1}}{1-\beta_{1}^{t-1}}\right)^{2}\right\|_{1}$ $\displaystyle=$ $\displaystyle\sum_{i=1}^{r}\sum_{j=1}^{r}\left\langle\mathcal{U}_{t}[:,i], \mathcal{U}_{t-1}[:,j]\right\rangle^{2}\left(\frac{m_{t-1,j}}{1-\beta_{1}^{t-1 }}\right)^{2}$ $\displaystyle\leq$ $\displaystyle\sum_{j=1}^{r}\left(\frac{m_{t-1,j}}{1-\beta_{1}^{t-1}}\right)^{2 }\cdot\max_{1\leq j\leq r}\sum_{i=1}^{r}\left\langle\mathcal{U}_{t}[:,i], \mathcal{U}_{t-1}[:,j]\right\rangle^{2}$ $\displaystyle=$ $\displaystyle\left\|\frac{m_{t-1}}{1-\beta_{1}^{t-1}}\right\|^{2}\max_{1\leq j \leq r}\left\|\mathcal{U}^{\top}_{t}\mathcal{U}_{t-1}[:,j]\right\|^{2}\leq \left\|\frac{m_{t-1}}{1-\beta_{1}^{t-1}}\right\|^{2}\leq C^{2}G^{2}$ $$

Finally, for the third term we have

$$ $\displaystyle\left\|\left(\mathcal{U}^{\top}_{t}\mathcal{U}_{t-1}\cdot\frac{m_ {t-1}}{1-\beta_{1}^{t-1}}\right)^{2}\right\|_{1}=\left\|\mathcal{U}^{\top}_{t} \mathcal{U}_{t-1}\cdot\frac{m_{t-1}}{1-\beta_{1}^{t-1}}\right\|^{2}\leq\left\| \frac{m_{t-1}}{1-\beta_{1}^{t-1}}\right\|^{2}\leq C^{2}G^{2}.$ $$

Combining all three bounds together, we arrive

$$ $\|v_{t-1/2}\|_{1}\leq\|v_{t-1}\|_{1}+2C^{2}G^{2}.$ $$

From this we get the bound for $v_{t}$ using the initialization $v_{0}=0$:

$$ $\displaystyle\|v_{t}\|_{\max}\leq\|v_{t}\|_{1}$ $\displaystyle\leq$ $\displaystyle\beta_{2}\|v_{t-1/2}\|_{1}+(1-\beta_{2})\|a_{t}\|^{2}$ $\displaystyle\leq$ $\displaystyle\beta_{2}(\|v_{t-1}\|_{1}+2C^{2}G^{2})+(1-\beta_{2})C^{2}G^{2}$ $\displaystyle\leq$ $\displaystyle\beta_{2}\|v_{t-1}\|_{1}+(1+\beta_{2})C^{2}G^{2}$ $\displaystyle\leq$ $\displaystyle\beta_{2}^{t}\|v_{0}\|_{1}+(1+\beta_{2})C^{2}G^{2}\sum_{\tau=0}^{ t-1}\beta_{2}^{\tau}\leq\frac{1+\beta_{2}}{1-\beta_{2}}C^{2}G^{2}.$ $$

Hence, using the update rule of $\tilde{v}_{t}$ and initialization $\tilde{v}_{0}=0$, we conclude

$$ $\|\tilde{v}_{t}\|_{\max}\leq\max(\|v_{t}\|_{\max},\|\tilde{v}_{t-1}\|_{\max}) \leq\frac{1+\beta_{2}}{1-\beta_{2}}C^{2}G^{2}=\frac{1+\beta_{2}}{1-\beta_{2}} \frac{(1-(1-q_{r})\beta_{1})^{2}}{(1-\beta_{1})^{2}(1-q_{r})^{2}}G^{2}.$ $$

∎

###### Lemma 1 .

###### Proof.

###### Lemma 2 .

###### Proof.

###### Lemma 3 .

###### Proof.

###### Lemma 4 .

###### Proof.

### E.2 Non-convex Analysis

###### Theorem 3 (Non-convex convergence rate) .

Let Assumptions 1, 2 and 3 hold. Then, choosing step-size $\eta=\min(\eta_{0},\frac{1}{\sqrt{T}})$ with $\eta_{0}=\frac{\epsilon}{4LC_{0}\sqrt{1+C_{2}}}$, Algorithm 3 satisfies

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

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

###### Remark 1 .

Further ignoring absolute constants, the bound becomes

$$ $\frac{C_{0}}{\sqrt{T}}\left(f(\theta_{1})-f^{*}+\frac{L\sigma^{2}}{\epsilon} \right)+\frac{C_{0}}{T}\left(\frac{f(\theta_{1})-f^{*}}{\eta_{0}}+\frac{L^{2}C _{0}C_{2}\sigma^{2}}{\epsilon^{2}}+\frac{C_{1}G^{2}}{\sqrt{\epsilon}}+\frac{C_ {1}^{2}LG^{2}}{\epsilon\sqrt{T}}\right),$ $$

Let us assume that for the low-rank compression of the algorithm we have $1-q_{r}=\mathcal{O}(\frac{r}{d})$ (e.g., SVD option). Then $C_{0}=\mathcal{O}(\frac{d}{r}G),\;C_{1}=\mathcal{O}(\frac{d}{r}),\;C_{2}=
\mathcal{O}(\frac{d^{2}}{r^{2}})$ and $\frac{1}{\eta_{0}}=\mathcal{O}(\frac{d^{2}}{r^{2}}G)$. Plugging this asymptotic expressions into the bound and ignoring other parameters (e.g., $\sigma^{2},L,\epsilon$), we get

$$ $\frac{d}{r}\frac{G}{\sqrt{T}}+\frac{d}{r}\frac{G}{T}\left(\frac{d^{2}G}{r^{2}} +\frac{d^{3}G}{r^{3}}+\frac{dG^{2}}{r}+\frac{d^{2}G^{2}}{r^{2}\sqrt{T}}\right)$ $$

or equivalently

$$ $\frac{d}{r}\frac{G}{\sqrt{T}}+\left(\frac{d}{r}\right)^{4}\frac{G^{2}}{T}+ \left(\frac{d}{r}\right)^{2}\frac{G^{3}}{T}+\left(\frac{d}{r}\right)^{3}\frac{ G^{3}}{T^{3/2}}.$ $$

###### Proof.

Let $\bar{\Gamma}_{t}:=\mathcal{U}_{t}\mathrm{Diag}^{-1/2}(\tilde{v}_{t}+\epsilon)
\mathcal{U}^{\top}_{t}$ be a preconditioning matrix and $M_{t}:=\mathcal{U}_{t}m_{t}$ the exponential moving averages in the full space. With these notations, the update rule of the model becomes $\theta_{t+1}=\theta_{t}-\eta_{t}\bar{\Gamma}_{t}\cdot M_{t}$. As it will be used later, we need to make sure that our preconditioning is positive definite. Due to the structure of $\bar{\Gamma}_{t}$, it is positive semi-definite only. To make it positive definite, notice that $\Gamma_{t}\mathcal{U}_{t}=\bar{\Gamma}_{t}\mathcal{U}_{t}$, where the full-rank preconditioner $\Gamma_{t}$ takes the form

$$ $\Gamma_{t}=\bar{\mathcal{U}}_{t}\mathrm{Diag}^{-1/2}(\tilde{v}_{t}+\epsilon,\| \tilde{v}_{t}\|_{\min}+\epsilon)\bar{\mathcal{U}}^{\top}_{t},$ (19) $$

where $\bar{\mathcal{U}}_{t}\in\mathbb{R}^{d\times d}$ is an orthogonal matrix with the same first $r$ columns as $\mathcal{U}_{t}\in\mathbb{R}^{d\times r}$ and the diagonal matrix in the middle has been extended to meet the sizes of $\bar{\mathcal{U}}_{t}$ by adding the values $\|\tilde{v}_{t}\|_{\min}+\epsilon$ on the diagonal. $\Gamma_{t}\mathcal{U}_{t}=\bar{\Gamma}_{t}\mathcal{U}_{t}$ comes from the observation that $\bar{\mathcal{U}}^{\top}_{t}\mathcal{U}_{t}\in\mathbb{R}^{d\times r}$ is composed of two blocks: upper $r\times r$ block of the identity matrix and $(d-r)\times r$ block of the zero matrix. Hence, the added $(d-r)$ elements on the diagonal do not really affect, so does the last $(d-r)$ columns of $\bar{\mathcal{U}}_{t}$. Therefore, we can write the model update rule as

$$ $\theta_{t+1}=\theta_{t}-\eta_{t}\Gamma_{t}\cdot M_{t},$ (20) $$

with full-rank preconditioning $\Gamma_{t}$. Next, note that

$$ $B_{t}=\beta_{1}M_{t-1}+(1-\beta_{1})g_{t}+(1-\beta_{1})e_{t},$ (21) $$

$$ $m_{t-1/2}=\mathcal{U}^{\top}_{t}M_{t-1}.$ (22) $$

Then, for the low dimensional momentum and the erorr we get

$$ $m_{t}=\beta_{1}m_{t-1/2}+(1-\beta_{1})a_{t}=\mathcal{U}^{\top}_{t}(\beta_{1}M_ {t-1}+(1-\beta_{1})g_{t}+(1-\beta_{1})\xi_{t})=\mathcal{U}^{\top}_{t}B_{t},$ (23) $$

$$ $(1-\beta_{1})\xi_{t+1}=(I-\mathcal{U}_{t}\mathcal{U}^{\top}_{t})B_{t}=B_{t}-M_ {t}.$ (24) $$

Letting $B_{0}=0$, we define virtual iterates $x_{t}$ as follows:

$$ $\displaystyle x_{t+1}$ $\displaystyle=$ $\displaystyle\theta_{t+1}-\eta\Gamma_{t}\left((1-\beta_{1})\xi_{t+1}+\frac{ \beta_{1}}{1-\beta_{1}}B_{t}\right).$ (25) $$

In particular, $x_{1}=\theta_{1}$. Then, we derive the recurrence relation for the new sequence as follows:

$$ $\displaystyle x_{t+1}$ $\displaystyle=$ $\displaystyle\theta_{t+1}-\eta\Gamma_{t}\left((1-\beta_{1})\xi_{t+1}+\frac{ \beta_{1}}{1-\beta_{1}}B_{t}\right)$ $\displaystyle\overset{(\ref{error-update-new})}{=}$ $\displaystyle\theta_{t}-\eta\Gamma_{t}M_{t}-\eta\Gamma_{t}\left(B_{t}-M_{t}+ \frac{\beta_{1}}{1-\beta_{1}}B_{t}\right)$ $\displaystyle=$ $\displaystyle\theta_{t}-\frac{\eta}{1-\beta_{1}}\Gamma_{t}B_{t}$ $\displaystyle\overset{(\ref{bcum-update-new})}{=}$ $\displaystyle\theta_{t}-\frac{\eta}{1-\beta_{1}}\Gamma_{t}\left(\beta_{1}M_{t- 1}+(1-\beta_{1})g_{t}+(1-\beta_{1})\xi_{t}\right)$ $\displaystyle\overset{(\ref{error-update-new})}{=}$ $\displaystyle\theta_{t}-\eta\Gamma_{t}\left(\frac{\beta_{1}}{1-\beta_{1}}(B_{t -1}-(1-\beta_{1})\xi_{t})+\xi_{t}\right)-\eta\Gamma_{t}g_{t}$ $\displaystyle=$ $\displaystyle\theta_{t}-\eta\Gamma_{t}\left(\frac{\beta_{1}}{1-\beta_{1}}B_{t- 1}+(1-\beta_{1})\xi_{t}\right)-\eta\Gamma_{t}g_{t}$ $\displaystyle\overset{(\ref{def:virtual-iter})}{=}$ $\displaystyle x_{t}-\eta\Gamma_{t}g_{t}+\eta\Delta\Gamma_{t}\left(\frac{\beta_ {1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right),$ $$

where $\Delta\Gamma_{t}:=\Gamma_{t-1}-\Gamma_{t}$.

Next we apply smoothness (Assumption 1) 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}.$ $$

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}),\Gamma_{t}g_{t} \right\rangle\right]$ (27) $\displaystyle+\eta\mathbb{E}\left[\left\langle\nabla f(x_{t}),\Delta\Gamma_{t} \left(\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right)\right \rangle\right]$ $\displaystyle+\frac{\eta^{2}L}{2}\mathbb{E}\left[\left\|\Gamma_{t}g_{t}-\Delta \Gamma_{t}\left(\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t} \right)\right\|^{2}\right]$ $\displaystyle=$ $\displaystyle\underbrace{-\eta\mathbb{E}\left[\left\langle\nabla f(\theta_{t}) ,\Gamma_{t}g_{t}\right\rangle\right]}_{I}$ $\displaystyle+\underbrace{\eta\mathbb{E}\left[\left\langle\nabla f(x_{t}), \Delta\Gamma_{t}\left(\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t }\right)\right\rangle\right]}_{II}$ $\displaystyle+\underbrace{\frac{\eta^{2}L}{2}\mathbb{E}\left[\left\|\Gamma_{t} g_{t}-\Delta\Gamma_{t}\left(\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1}) \xi_{t}\right)\right\|^{2}\right]}_{III}$ $\displaystyle+\underbrace{\eta\mathbb{E}\left[\left\langle\nabla f(\theta_{t}) -\nabla f(x_{t}),\Gamma_{t}g_{t}\right\rangle\right]}_{IV},$ $$

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

Bounding term I. Let $\|\Delta\Gamma\|$ be the operator norm (with respect to $\ell_{2}$ norm) of the matrix $\Delta\Gamma$. We have

$$ $\displaystyle I$ $\displaystyle=$ $\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla f(\theta_{t}),\Gamma_{t-1 }g_{t}\right\rangle\right]-\eta\mathbb{E}\left[\left\langle\nabla f(\theta_{t} ),\Delta\Gamma_{t}g_{t}\right\rangle\right]$ (28) $\displaystyle\leq$ $\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla f(\theta_{t}),\Gamma_{t-1 }\nabla f(\theta_{t})\right\rangle\right]+\eta G^{2}\mathbb{E}[\|\Delta\Gamma_ {t}\|].$ $\displaystyle\leq$ $\displaystyle-\eta\lambda_{\min}(\Gamma_{t-1})\mathbb{E}[\|\nabla f(\theta_{t} )\|^{2}]+\eta G^{2}\mathbb{E}[\|\Delta\Gamma_{t}\|]$ $\displaystyle\leq$ $\displaystyle-\frac{\eta}{C_{0}}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}]+\eta G ^{2}\mathbb{E}[\|\Delta\Gamma_{t}\|],$ $$

where we use Assumption 2 and Lemma 4 to bound

$$ $\lambda_{\min}(\Gamma_{t-1})\geq\left(\|\tilde{v}_{t-1}\|_{\max}+\epsilon \right)^{-1/2}\geq\left(\frac{1+\beta_{2}}{1-\beta_{2}}\frac{(1-\beta_{1}(1-q_ {r}))^{2}}{(1-\beta_{1})^{2}(1-q_{r})^{2}}G^{2}+\epsilon\right)^{-1/2}=\frac{1 }{C_{0}}.$ $$

Note that the purpose of making $\Gamma$ matrix positive definite is to have negative term in (28).

Bounding term II. Then we have

$$ $\displaystyle II$ $\displaystyle\leq\eta\mathbb{E}\left[\left\langle\nabla f(\theta_{t}),\Delta \Gamma_{t}\left(\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t} \right)\right\rangle\right]$ $\displaystyle\quad+\eta\mathbb{E}\left[\left\langle\nabla f(x_{t})-\nabla f( \theta_{t}),\Delta\Gamma_{t}\left(\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1- \beta_{1})\xi_{t}\right)\right\rangle\right]$ $\displaystyle\leq\eta\mathbb{E}\left[\|\nabla f(\theta_{t})\|\left\|\Delta \Gamma_{t}\left(\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t} \right)\right\|\right]$ $\displaystyle\quad+\eta^{2}L\mathbb{E}\left[\left\|\Gamma_{t}\left(\frac{\beta _{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right)\right\|\cdot\left\| \Delta\Gamma_{t}\left(\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t }\right)\right\|\right]$ $\displaystyle\overset{(\ref{bound-perturb})}{\leq}\eta C_{1}G^{2}\mathbb{E}[\| \Delta\Gamma_{t}\|]+\frac{\eta^{2}C_{1}^{2}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[ \|\Delta\Gamma_{t}\|],$ (29) $$

where $C_{1}=\frac{\beta_{1}+(1-\beta_{1})q_{r}}{(1-\beta_{1})(1-q_{r})}$ and we used the fact that the largest eigenvalue $\lambda_{\max}(\Gamma_{t})=\|\Gamma_{t}\|=(\|\tilde{v}_{t}\|_{\min}+\epsilon)^
{-1/2}\leq\epsilon^{-1/2}$. The second inequality is because of smoothness of $f(\theta)$, and the last inequality is due to Lemma 1, Assumption 2 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\|\Gamma_{t}g_{t}\right\|^{2} \right]+\eta^{2}L\mathbb{E}\left[\left\|\Delta\Gamma_{t}\left(\frac{\beta_{1}} {1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right)\right\|^{2}\right]$ $\displaystyle\leq\frac{\eta^{2}L}{\epsilon}\mathbb{E}[\|g_{t}-\nabla f(\theta_ {t})+\nabla f(\theta_{t})\|^{2}]+\eta^{2}L\mathbb{E}\left[\left\|\Delta\Gamma_ {t}\left(\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right) \right\|^{2}\right]$ $\displaystyle\leq\frac{\eta^{2}L}{\epsilon}\left(\mathbb{E}[\|\nabla f(\theta_ {t})\|^{2}]+\sigma^{2}\right)+\eta^{2}C_{1}^{2}LG^{2}\mathbb{E}[\|\Delta\Gamma _{t}\|^{2}]$ $\displaystyle\leq\frac{\eta^{2}L}{\epsilon}\mathbb{E}[\|\nabla f(\theta_{t})\| ^{2}]+\frac{\eta^{2}L\sigma^{2}}{\epsilon}+\eta^{2}C_{1}^{2}LG^{2}\mathbb{E}[ \|\Delta\Gamma_{t}\|^{2}],$ (30) $$

where we used Assumption 3 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}),\Gamma_{t-1}g_{t}\right\rangle\right]+\eta\mathbb{E}\left[\left\langle \nabla f(\theta_{t})-\nabla f(x_{t}),\Delta\Gamma_{t}g_{t}\right\rangle\right]$ $\displaystyle\leq\eta\mathbb{E}\left[\left\langle\nabla f(\theta_{t})-\nabla f (x_{t}),\Gamma_{t-1}\nabla f(\theta_{t})\right\rangle\right]+\eta^{2}L\mathbb{ E}\left[\left\|\Gamma_{t}\left(\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{ 1})\xi_{t}\right)\right\|\|\Delta\Gamma_{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}[\|\Delta \Gamma_{t}\|]$ $\displaystyle\overset{(b)}{\leq}\frac{\eta\rho}{2\epsilon}\mathbb{E}[\|\nabla f (\theta_{t})\|^{2}]+\frac{\eta^{3}L^{2}}{2\rho}\mathbb{E}\left[\left\|\Gamma_{ t}\left(\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right)\right \|^{2}\right]+\frac{\eta^{2}C_{1}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|\Delta \Gamma_{t}\|]$ $\displaystyle\leq\frac{\eta\rho}{2\epsilon}\mathbb{E}[\|\nabla f(\theta_{t})\| ^{2}]+\frac{\eta^{3}L^{2}}{2\rho\epsilon}\mathbb{E}\left[\left\|\frac{\beta_{1 }}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]+\frac{\eta^{2} LC_{1}G^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|\Delta\Gamma_{t}\|],$ (31) $$

where (a) is due to Young’s inequality and (b) is based on Assumption 1. Now integrating (28), (29), (30), (31) into (27),

$$ $\displaystyle I$ $\displaystyle\leq$ $\displaystyle-\eta\lambda_{\min}(\Gamma_{t-1})\mathbb{E}[\|\nabla f(\theta_{t} )\|^{2}]+\eta G^{2}\mathbb{E}[\|\Delta\Gamma_{t}\|]$ $\displaystyle II$ $\displaystyle\leq$ $\displaystyle\eta C_{1}G^{2}\mathbb{E}[\|\Delta\Gamma_{t}\|]+\frac{\eta^{2}C_{ 1}^{2}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|\Delta\Gamma_{t}\|]$ $\displaystyle III$ $\displaystyle\leq$ $\displaystyle\frac{\eta^{2}L}{\epsilon}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2} ]+\frac{\eta^{2}L\sigma^{2}}{\epsilon}+\eta^{2}C_{1}^{2}LG^{2}\mathbb{E}[\| \Delta\Gamma_{t}\|^{2}]$ $\displaystyle IV$ $\displaystyle\leq$ $\displaystyle\frac{\eta\rho}{2\epsilon}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2} ]+\frac{\eta^{3}L^{2}}{2\rho\epsilon}\mathbb{E}\left[\left\|\frac{\beta_{1}}{1 -\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]+\frac{\eta^{2}LC_{1 }G^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|\Delta\Gamma_{t}\|],$ $$

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}}{\epsilon}+\frac{\eta^{3}L^{2}}{2\rho\epsilon} \sum_{t=1}^{T}\mathbb{E}\left[\left\|\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1- \beta_{1})\xi_{t}\right\|^{2}\right]$ $\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}[\|\Delta\Gamma_{t}\|]+\eta^{ 2}C_{1}^{2}LG^{2}\sum_{t=1}^{T}\mathbb{E}[\|\Delta\Gamma_{t}\|^{2}]$ $\displaystyle\overset{(\ref{bound-perturb-exp})}{\leq}\left(-\frac{\eta}{C_{0} }+\frac{\eta^{2}L}{\epsilon}+\frac{\eta\rho}{2\epsilon}+\frac{\eta^{3}L^{2}C_{ 2}}{2\rho\epsilon}\right)\sum_{t=1}^{T}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2} ]+\frac{T\eta^{2}L\sigma^{2}}{\epsilon}+\frac{T\eta^{3}L^{2}C_{2}\sigma^{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}[\|\Delta\Gamma_{t}\|]+\eta^{ 2}C_{1}^{2}LG^{2}\sum_{t=1}^{T}\mathbb{E}[\|\Delta\Gamma_{t}\|^{2}],$ $$

where we used (16) of Lemma 1 with constant $C_{2}=\frac{\beta_{1}+(1-\beta_{1})q_{r}^{2}}{(1-\beta_{1})^{2}(1-q_{r})^{2}}$. Choosing $\rho=\frac{\epsilon}{2C_{0}}$ and $\eta\leq\eta_{0}\mathrel{\mathop{:}}=\frac{\epsilon}{4LC_{0}\sqrt{1+C_{2}}}$ and using Lemma 3, we get

$$ $\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}}{\epsilon}+\frac{T\eta^{3}L^{2}C _{0}C_{2}\sigma^{2}}{\epsilon^{2}}$ $\displaystyle\quad+\frac{2\eta(1+C_{1})G^{2}}{\sqrt{\epsilon}}+\frac{2\eta^{2} (1+2C_{1})C_{1}LG^{2}}{\epsilon}.$ $$

Re-arranging terms, we get that

$$ $\displaystyle\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}]$ $\displaystyle\leq$ $\displaystyle 2C_{0}\left(\frac{f(\theta_{1})-f^{*}}{T\eta}+\frac{\eta L\sigma ^{2}}{\epsilon}+\frac{\eta^{2}L^{2}C_{0}C_{2}\sigma^{2}}{\epsilon^{2}}\right)$ $\displaystyle+4C_{0}\left(\frac{(1+C_{1})G^{2}}{T\sqrt{\epsilon}}+\frac{\eta(1 +2C_{1})C_{1}LG^{2}}{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}$. Finally, choosing $\eta=\min(\eta_{0},\frac{1}{\sqrt{T}})$ and considering the two cases, we arrive at the following rate

$$ $\displaystyle\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}]$ $\displaystyle\leq$ $\displaystyle 2C_{0}\left(\max\left(1,\frac{1}{\eta_{0}\sqrt{T}}\right)\frac{f (\theta_{1})-f^{*}}{\sqrt{T}}+\frac{L\sigma^{2}}{\epsilon\sqrt{T}}+\frac{L^{2} C_{0}C_{2}\sigma^{2}}{\epsilon^{2}T}\right)$ $\displaystyle+4C_{0}\left(\frac{(1+C_{1})G^{2}}{\sqrt{\epsilon}T}+\frac{(1+2C_ {1})C_{1}LG^{2}}{\epsilon T^{3/2}}\right)$ $\displaystyle\leq$ $\displaystyle 2C_{0}\left(\frac{f(\theta_{1})-f^{*}}{\sqrt{T}}+\frac{L\sigma^{ 2}}{\epsilon\sqrt{T}}\right)$ $\displaystyle+4C_{0}\left(\frac{f(\theta_{1})-f^{*}}{2\eta_{0}T}+\frac{L^{2}C_ {0}C_{2}\sigma^{2}}{2\epsilon^{2}T}+\frac{(1+C_{1})G^{2}}{\sqrt{\epsilon}T}+ \frac{(1+2C_{1})C_{1}LG^{2}}{\epsilon T^{3/2}}\right),$ $$

which completes the proof of the theorem.
∎

###### Theorem 3 (Non-convex convergence rate) .

###### Remark 1 .

###### Proof.

### E.3 Analysis Under PL Condition

As in the non-convex analysis, here we derive the convergence rate with fixed step-size $\eta$.

###### Theorem 4 .

(Convergence rate under PL)
Let Assumptions 1, 2, 3 and 4 hold. Then, choosing step-size $\eta=\min(\eta_{0},\frac{2C_{0}\log T}{\mu T})$ with $\eta_{0}=\min(\frac{\epsilon}{16LC_{0}},\frac{C_{0}(1-\beta_{1})(1-q_{r})}{2
\mu},\frac{\epsilon^{3/4}}{6L\sqrt{C_{0}C_{2}}})$, Algorithm 3 satisfies

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

###### Proof.

We start from descent lemma

$$ $\displaystyle\mathbb{E}[f(x_{t+1})]-\mathbb{E}[f(x_{t})]$ $\displaystyle=$ $\displaystyle\underbrace{-\eta\mathbb{E}\left[\left\langle\nabla f(x_{t}), \Gamma_{t}g_{t}\right\rangle\right]}_{I^{\prime}}$ (32) $\displaystyle+\underbrace{\eta\mathbb{E}\left[\left\langle\nabla f(x_{t}), \Delta\Gamma_{t}\left(\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1}\xi_{t} )\right)\right\rangle\right]}_{II}$ $\displaystyle+\underbrace{\frac{\eta^{2}L}{2}\mathbb{E}\left[\left\|\Gamma_{t} g_{t}-\Delta\Gamma_{t}\left(\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1} \xi_{t})\right)\right\|^{2}\right]}_{III}.$ $$

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}),\Gamma_{t}g_{t} \right\rangle\right]$ $\displaystyle=$ $\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla f(x_{t}),\Gamma_{t-1}g_{t }\right\rangle\right]-\eta\mathbb{E}\left[\left\langle\nabla f(x_{t}),\Delta \Gamma_{t}g_{t}\right\rangle\right]$ $\displaystyle=$ $\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla f(x_{t}),\Gamma_{t-1}(g_{ t}-\nabla f(x_{t})+\nabla f(x_{t}))\right\rangle\right]-\eta\mathbb{E}\left[ \left\langle\nabla f(x_{t}),\Delta\Gamma_{t}g_{t}\right\rangle\right]$ $\displaystyle=$ $\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla f(x_{t}),\Gamma_{t-1} \nabla f(x_{t})\right\rangle\right]-\eta\mathbb{E}\left[\left\langle\nabla f(x _{t}),\Gamma_{t-1}(g_{t}-\nabla f(x_{t}))\right\rangle\right]$ $\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla f(x_{t}),\Delta\Gamma_{t} g_{t}\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}),\Gamma_{t-1}\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}),\Delta\Gamma_{t}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}),\Gamma_{t-1}\left(\nabla f(\theta_{t})-\nabla f(x_{t})\right)\right \rangle\right]$ $\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla f(\theta_{t}),\Gamma_{t-1 }\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}),\Delta\Gamma_{t}g_{t}\right\rangle\right]$ $\displaystyle-\eta\mathbb{E}\left[\left\langle\nabla f(\theta_{t}),\Delta \Gamma_{t}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}),\Gamma_{t-1 }\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}),\Delta\Gamma_{t}g_{t}\right\rangle\right]+\eta\mathbb{E}\left[\left \langle\nabla f(\theta_{t}),\Delta\Gamma_{t}g_{t}\right\rangle\right].$ $$

To bound the second and third terms above we reuse derivation done in (31) to have

$$ $\displaystyle\eta\mathbb{E}\left[\left\langle\nabla f(\theta_{t}),\Gamma_{t-1} \left(\nabla f(\theta_{t})-\nabla f(x_{t})\right)\right\rangle\right]$ $\displaystyle\overset{(\ref{eq:IV})}{\leq}$ $\displaystyle\frac{\eta\rho}{2\epsilon}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}]$ $\displaystyle+\frac{\eta^{3}L^{2}}{2\rho\epsilon}\mathbb{E}\left[\left\|\frac{ \beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]$ $$

$$ $\displaystyle\mathbb{E}\left[\left\|\nabla f(x_{t})-f(\theta_{t})\right\|^{2}\right]$ $\displaystyle\overset{(\ref{eq:IV})}{\leq}$ $\displaystyle\frac{\eta^{2}L^{2}}{\epsilon}\mathbb{E}\left[\left\|\frac{\beta_ {1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]$ (34) $$

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]$ $\displaystyle+\frac{\eta\rho}{2\epsilon}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2 }]+\frac{\eta^{3}L^{2}}{2\rho\epsilon}\mathbb{E}\left[\left\|\frac{\beta_{1}}{ 1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]$ $\displaystyle+\eta G\mathbb{E}\left[\|\nabla f(x_{t})-\nabla f(\theta_{t})\|\| \Delta\Gamma_{t}\|\right]+\eta G^{2}\mathbb{E}\left[\|\Delta\Gamma_{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]$ $\displaystyle+\frac{\eta\rho}{2\epsilon}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2 }]+\frac{\eta}{\sqrt{\epsilon}}\mathbb{E}\left[\left\|\nabla f(x_{t})-f(\theta _{t})\right\|^{2}\right]+\frac{\eta^{3}L^{2}}{2\rho\epsilon}\mathbb{E}\left[ \left\|\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]$ $\displaystyle+\eta G\frac{\eta LC_{1}G}{\sqrt{\epsilon}}\mathbb{E}\left[\| \Delta\Gamma_{t}\|\right]+\eta G^{2}\mathbb{E}\left[\|\Delta\Gamma_{t}\|\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\rho}{2\epsilon}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2 }]+\frac{\eta}{\sqrt{\epsilon}}\mathbb{E}\left[\left\|\nabla f(x_{t})-f(\theta _{t})\right\|^{2}\right]+\frac{\eta^{3}L^{2}}{2\rho\epsilon}\mathbb{E}\left[ \left\|\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]$ $\displaystyle+\frac{\eta^{2}LC_{1}G^{2}}{\sqrt{\epsilon}}\mathbb{E}\left[\| \Delta\Gamma_{t}\|\right]+\eta G^{2}\mathbb{E}\left[\|\Delta\Gamma_{t}\|\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{\eta\rho}{2\epsilon}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2 }]+\frac{3\eta}{2\sqrt{\epsilon}}\mathbb{E}\left[\left\|\nabla f(x_{t})-f( \theta_{t})\right\|^{2}\right]+\frac{\eta^{3}L^{2}}{2\rho\epsilon}\mathbb{E} \left[\left\|\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right\| ^{2}\right]$ $\displaystyle+\frac{\eta^{2}LC_{1}G^{2}}{\sqrt{\epsilon}}\mathbb{E}\left[\| \Delta\Gamma_{t}\|\right]+\eta G^{2}\mathbb{E}\left[\|\Delta\Gamma_{t}\|\right]$ $\displaystyle\overset{(\ref{aux-bound-2})}{\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\rho}{2\epsilon}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}]$ $\displaystyle+\frac{3\eta}{2\sqrt{\epsilon}}\frac{\eta^{2}L^{2}}{\epsilon} \mathbb{E}\left[\left\|\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{ t}\right\|^{2}\right]+\frac{\eta^{3}L^{2}}{2\rho\epsilon}\mathbb{E}\left[\left \|\frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]$ $\displaystyle+\frac{\eta^{2}LC_{1}G^{2}}{\sqrt{\epsilon}}\mathbb{E}\left[\| \Delta\Gamma_{t}\|\right]+\eta G^{2}\mathbb{E}\left[\|\Delta\Gamma_{t}\|\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\rho}{2\epsilon}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}]$ $\displaystyle+\frac{\eta^{3}L^{2}}{2\epsilon}\left(\frac{3}{\sqrt{\epsilon}}+ \frac{1}{\rho}\right)\mathbb{E}\left[\left\|\frac{\beta_{1}}{1-\beta_{1}}B_{t- 1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]+\frac{\eta^{2}LC_{1}G^{2}}{\sqrt{ \epsilon}}\mathbb{E}\left[\|\Delta\Gamma_{t}\|\right]+\eta G^{2}\mathbb{E} \left[\|\Delta\Gamma_{t}\|\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}[\|\Delta\Gamma_{t}\|]+\frac{\eta^{2}C_{ 1}^{2}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|\Delta\Gamma_{t}\|]$ $\displaystyle III$ $\displaystyle\leq$ $\displaystyle\frac{\eta^{2}L}{\epsilon}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2} ]+\frac{\eta^{2}L\sigma^{2}}{\epsilon}+\eta^{2}C_{1}^{2}LG^{2}\mathbb{E}[\| \Delta\Gamma_{t}\|^{2}]$ $$

into (32), choosing $\rho=\frac{\epsilon}{8C_{0}}$ and using the step-size bound $\eta\leq\frac{\epsilon}{16LC_{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}]+ \frac{\eta\rho}{2\epsilon}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}]$ $\displaystyle+\frac{\eta^{3}L^{2}}{2\epsilon}\left(\frac{3}{\sqrt{\epsilon}}+ \frac{1}{\rho}\right)\mathbb{E}\left[\left\|\frac{\beta_{1}}{1-\beta_{1}}B_{t- 1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]+\frac{\eta^{2}LC_{1}G^{2}}{\sqrt{ \epsilon}}\mathbb{E}\left[\|\Delta\Gamma_{t}\|\right]$ $\displaystyle+\eta C_{1}G^{2}\mathbb{E}[\|\Delta\Gamma_{t}\|]+\frac{\eta^{2}C_ {1}^{2}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|\Delta\Gamma_{t}\|]+\eta G^{2} \mathbb{E}\left[\|\Delta\Gamma_{t}\|\right]$ $\displaystyle+\frac{\eta^{2}L}{\epsilon}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2 }]+\frac{\eta^{2}L\sigma^{2}}{\epsilon}+\eta^{2}C_{1}^{2}LG^{2}\mathbb{E}[\| \Delta\Gamma_{t}\|^{2}]$ $\displaystyle\leq$ $\displaystyle-\frac{\eta}{2C_{0}}\mathbb{E}\left[\left\|\nabla f(x_{t})\right \|^{2}\right]-\frac{\eta}{8C_{0}}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}]+ \frac{\eta^{2}L\sigma^{2}}{\epsilon}$ $\displaystyle+\frac{\eta^{3}L^{2}}{2\epsilon}\left(\frac{3}{\sqrt{\epsilon}}+ \frac{\epsilon}{8C_{0}}\right)\mathbb{E}\left[\left\|\frac{\beta_{1}}{1-\beta_ {1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]$ $\displaystyle+\eta(1+C_{1})G^{2}\mathbb{E}[\|\Delta\Gamma_{t}\|]+\frac{\eta^{2 }(1+C_{1})C_{1}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|\Delta\Gamma_{t}\|]+\eta^{ 2}C_{1}^{2}LG^{2}\mathbb{E}[\|\Delta\Gamma_{t}\|^{2}]$ $\displaystyle\leq$ $\displaystyle-\frac{\eta\mu}{C_{0}}(\mathbb{E}[f(x_{t})]-f^{*})-\frac{\eta}{8C _{0}}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}]+\frac{\eta^{2}L\sigma^{2}}{\epsilon}$ $\displaystyle+\frac{2\eta^{3}L^{2}}{\epsilon^{3/2}}\mathbb{E}\left[\left\| \frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]$ $\displaystyle+\eta(1+C_{1})G^{2}\mathbb{E}[\|\Delta\Gamma_{t}\|]+\frac{\eta^{2 }(1+C_{1})C_{1}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|\Delta\Gamma_{t}\|]+\eta^{ 2}C_{1}^{2}LG^{2}\mathbb{E}[\|\Delta\Gamma_{t}\|^{2}],$ $$

where in the last inequality we applied PL condition from Assumption 4. 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}{8C_{0}}\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}]+\frac{\eta^{2}L \sigma^{2}}{\epsilon}$ $\displaystyle+\frac{2\eta^{3}L^{2}}{\epsilon^{3/2}}\mathbb{E}\left[\left\| \frac{\beta_{1}}{1-\beta_{1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]$ $\displaystyle+\eta(1+C_{1})G^{2}\mathbb{E}[\|\Delta\Gamma_{t}\|]+\frac{\eta^{2 }(1+C_{1})C_{1}LG^{2}}{\sqrt{\epsilon}}\mathbb{E}[\|\Delta\Gamma_{t}\|]+\eta^{ 2}C_{1}^{2}LG^{2}\mathbb{E}[\|\Delta\Gamma_{t}\|^{2}].$ $$

Now we unroll the obtained recursion and invoke Lemma 2 with $\gamma=1-\frac{\eta\mu}{C_{0}}$. From $\eta\leq\frac{\epsilon}{4LC_{0}}\leq\frac{C_{0}}{4\mu}$, we conclude that $\gamma=1-\frac{\eta\mu}{C_{0}}\in(0,1)$. To satisfy the condition on $\gamma$ of Lemma 1, we enforce the bound $\eta\leq\frac{C_{0}}{2\mu}(1-\beta_{1})(1-q_{r})$ on the step size. Therefore,

$$ $\displaystyle\mathbb{E}[f(x_{T+1})]-f^{*}\overset{(\ref{bound-perturb-exp})}{\leq}$ $\displaystyle\left(1-\frac{\eta\mu}{C_{0}}\right)^{T}(f(x_{1})-f^{*})-\frac{ \eta}{8C_{0}}\sum_{t=1}^{T}\left(1-\frac{\eta\mu}{C_{0}}\right)^{T-t}\mathbb{E }[\|\nabla f(\theta_{t})\|^{2}]$ $\displaystyle+\frac{\eta^{2}L\sigma^{2}}{\epsilon}\sum_{t=0}^{T-1}\left(1- \frac{\eta\mu}{C_{0}}\right)^{t}$ $\displaystyle+\frac{2\eta^{3}L^{2}}{\epsilon^{3/2}}\sum_{t=1}^{T}\left(1-\frac {\eta\mu}{C_{0}}\right)^{T-t}\mathbb{E}\left[\left\|\frac{\beta_{1}}{1-\beta_{ 1}}B_{t-1}+(1-\beta_{1})\xi_{t}\right\|^{2}\right]$ $\displaystyle+\eta(1+C_{1})G^{2}\sum_{t=1}^{T}\left(1-\frac{\eta\mu}{C_{0}} \right)^{T-t}\mathbb{E}[\|\Delta\Gamma_{t}\|]$ $\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-t}\mathbb{E}[\|\Delta\Gamma_{t}\|]$ $\displaystyle+\eta^{2}C_{1}^{2}LG^{2}\sum_{t=1}^{T}\left(1-\frac{\eta\mu}{C_{0 }}\right)^{T-t}\mathbb{E}[\|\Delta\Gamma_{t}\|^{2}].$ (35) $$

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

$$ $\sum_{t=0}^{T-1}\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 third sum, we apply Lemma 2 and for the other three sums we bound $1-\frac{\eta\mu}{C_{0}}\leq 1$ and apply the bounds in Lemma 3:

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

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

Plugging all this bounds into (35), cancelling terms involving gradients with step size restriction $\eta\leq\frac{\epsilon^{3/4}}{6L\sqrt{C_{0}C_{2}}}$ 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^{*})- \frac{\eta}{8C_{0}}\sum_{t=1}^{T}\left(1-\frac{\eta\mu}{C_{0}}\right)^{T-t} \mathbb{E}[\|\nabla f(\theta_{t})\|^{2}]$ $\displaystyle+\frac{\eta^{2}L\sigma^{2}}{\epsilon}\frac{C_{0}}{\eta\mu}$ $\displaystyle+\frac{2\eta^{3}L^{2}}{\epsilon^{3/2}}\cdot 2C_{2}\left(\frac{C_{ 0}\sigma^{2}}{\eta\mu}+\sum_{t=1}^{T}\left(1-\frac{\eta\mu}{C_{0}}\right)^{T-t }\mathbb{E}[\|\nabla f(\theta_{t})\|^{2}]\right)$ $\displaystyle+\eta(1+C_{1})G^{2}\frac{2}{\sqrt{\epsilon}}+\frac{\eta^{2}(1+C_{ 1})C_{1}LG^{2}}{\sqrt{\epsilon}}\frac{2}{\sqrt{\epsilon}}+\eta^{2}C_{1}^{2}LG^ {2}\frac{2}{\epsilon}$ $\displaystyle\leq$ $\displaystyle\left(1-\frac{\eta\mu}{C_{0}}\right)^{T}(f(\theta_{1})-f^{*})$ $\displaystyle+\frac{\eta LC_{0}\sigma^{2}}{\mu\epsilon}+\frac{4\eta^{2}L^{2}C_ {0}C_{2}\sigma^{2}}{\mu\epsilon^{3/2}}$ $\displaystyle+\frac{2\eta(1+C_{1})G^{2}}{\sqrt{\epsilon}}+\frac{2\eta^{2}(1+C_ {1})C_{1}LG^{2}}{\epsilon}+\frac{2\eta^{2}C_{1}^{2}LG^{2}}{\epsilon}$ $\displaystyle=$ $\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{2(1+C_{1})G ^{2}}{\sqrt{\epsilon}}\right)+\eta^{2}\left(\frac{4L^{2}C_{0}C_{2}\sigma^{2}}{ \mu\epsilon^{3/2}}+\frac{2(1+2C_{1})C_{1}LG^{2}}{\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{3(1+C_{1})G ^{2}}{\sqrt{\epsilon}}\right)+\eta^{2}\left(\frac{4L^{2}C_{0}C_{2}\sigma^{2}}{ \mu\epsilon^{3/2}}+\frac{5(1+C_{1})C_{1}LG^{2}}{\epsilon}\right).$ $$

Plugging $\eta=\min(\eta_{0},\frac{2C_{0}\log T}{\mu T})$ with $\eta_{0}=\min(\frac{\epsilon}{16LC_{0}},\frac{C_{0}(1-\beta_{1})(1-q_{r})}{2
\mu},\frac{\epsilon^{3/4}}{6L\sqrt{C_{0}C_{2}}})$, we get

$$ $\displaystyle\mathbb{E}[f(\theta_{T+1})]-f^{*}\leq$ $\displaystyle\;\max\left(\frac{1}{T^{2}},\left(1-\frac{\eta_{0}\mu}{C_{0}} \right)^{T}\right)(f(\theta_{1})-f^{*})$ $\displaystyle+\frac{\log T}{T}\frac{2C_{0}}{\mu}\left(\frac{LC_{0}\sigma^{2}}{ \mu\epsilon}+\frac{3(1+C_{1})G^{2}}{\sqrt{\epsilon}}\right)$ $\displaystyle+\frac{\log^{2}T}{T^{2}}\frac{4C_{0}^{2}}{\mu^{2}}\left(\frac{4L^ {2}C_{0}C_{2}\sigma^{2}}{\mu\epsilon^{3/2}}+\frac{5(1+C_{1})C_{1}LG^{2}}{ \epsilon}\right)$ $\displaystyle=$ $\displaystyle\;\frac{\log T}{T}\left(\frac{2LC_{0}^{2}\sigma^{2}}{\mu^{2} \epsilon}+\frac{6C_{0}(1+C_{1})G^{2}}{\mu\sqrt{\epsilon}}\right)+\widetilde{ \mathcal{O}}\left(\frac{G^{4}}{T^{2}}\right),$ $$

which completes the proof.
∎

###### Theorem 4 .

###### Proof.

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: see Weaknesses
soundness: 3
strengths: The paper is overall easy to follow, the idea is interesting, and seems to empirically works great.
summary: The paper proposes a new optimization algorithm to train large (language) models.
Like the recent GaLore algorithm, the proposed algorithm is a variation of ADAM, that performs updates in a lower-dimensional space, to store lower-dimensional quantities to be more memory-efficient.
As I understand (and as written by the authors), the real meat behind this kind of algorithm is the choice of the projection to obtain a meaningful lower-dimensional subspace that is both efficient to compute/store and a good enough approximation of the gradient. The proposed algorithm retains the desirable memory-efficient properties of GaLore, while almost matching the generalization performance of ADAM. Convergence proofs and exhaustive experiments on text benchmarks are provided.
weaknesses: I see that the authors provide some running times and memory metrics in Appendix B, but it seems to me that this part should be more exhaustive in order to prove the proposed algorithm useful.
I know this is hard to monitor for algorithms on GPU, but could authors provide graphs similar to Figure 1, but with running times? and memory usage? with ADAM, GaLore, and LDADAM (proposed).
In other words, could authors provide graphs with perplexity as a function of runtime and perplexity as a function of memory usage? Even if these required graphs will be imperfect for a lot reasons, I think they are crucial to gain insights and potentially prove usefulness.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Overall I think this is a work with interesting observations and methods. I have the following concerns and questions (I'm happy to increase my evaluation if the authors address the issues):

1. To my understanding, GaLore doesn't employ SVD step for every iteration in order to save computation overheads. On the contrary, the proposed method would need a Gram-Schmidt process and several extra steps for updating the momentums and gradients. A major conern is the computation overhead of this proposed method. I udnerstand that the memory consumption is roughly the same as GaLore, yet I'm curious about the time required for training LLMs both in theory and in all the experiment settings using the proposed method.
2. The convergence results in Theorem 1 and 2 are similar to that of full-rank Adagrad/Adam. To my understanding, the low rank parameter $r$ only plays a role in the constant $q_r$? If so, can the author comment on how to choose such $r$ and if the theory has any implication on the choice of $r$. In particular, if $r=m$ and the method reduce back to the original Adam, can we say any matching result with the analysis on the original Adam?
3. This question is related to the previous one. The experiments show that LDAdam is even better than original Adam in some cases (Table 3 with $r=512$, Table 4 with $r=256$). Can the authors expand on the discussion of "regularization effect of compression" or provide some references? It's still not very clear to me right now.
4. Any experiment on even larger LLMs such as Llama 1b or 7b? I'm curious what will happen if the model is large and the training is conducted in parallel on multiple GPUs.
soundness: 3
strengths: The paper is well-written and the contribution is clear. The algorithm is novel with theoretical convergence. The experiments indicates that the new approach is performing not only better than GaLore, but also better than the vanilla Adam for some parameter settings.
summary: The paper proposes a new memory-efficient training algorithm based on Adam via efficient projection onto subspaces. The proposed method enjoys non-asymptotic theoretical rate of convergence and demonstrate good experimental results in training large language models (LLMs).
weaknesses: (Please respond to the "Questions" directly) The computational time of proposed method is not clear; Some discussion on the theoretical results is needed.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: The original GaLore paper supports gradient accumulation, contrary to the authors' claim. I suggest they clarify this point, either correcting the claim or explaining any differences in their interpretation or implementation of GaLore to avoid potential misunderstandings.

Good paper, I suggest acceptance.
soundness: 2
strengths: 1. **Projection-Aware Update Rule**: LDAdam uses a projection-aware update rule to transition between subspaces. This rule allows it to estimate the gradient statistics even after dimensional reduction, adapting efficiently to changes in the subspace basis without losing essential gradient information.

2. **Generalized Error Feedback Mechanism**: To address inaccuracies from low-rank projections, LDAdam introduces a unique error feedback mechanism that accounts for both gradient and optimizer state compression. This mechanism reintroduces projection errors in subsequent iterations, ensuring stability in updates and helping LDAdam match the performance of uncompressed optimizers.

3. **Convergence and Memory Efficiency**: The paper provides a proof of LDAdam’s convergence under standard assumptions. With this design, LDAdam achieves a significantly reduced memory footprint. In practical terms, LDAdam’s optimizer states require only a fraction of the memory needed by Adam, making it feasible to train large language models in resource-constrained environments.
summary: This paper on LDAdam addresses a crucial challenge in training large models: reducing memory consumption of optimizers without sacrificing model quality. The key innovation is the introduction of a lower-dimensional subspace optimization approach, where adaptive optimization occurs within constrained dimensions. The optimizer, though operating in a subspace, ensures exploration across the entire parameter space, thereby maintaining model expressivity while being memory efficient.
weaknesses: Liang et al. [1] provide a convergence analysis for GaLore and related algorithms without relying on a 'stable-rank' assumption. Since LDAdam aligns with this framework, I recommend citing [1] for its convergence insights. Additionally, a comparison of LDAdam’s convergence analysis with [1], specifically on the impact of removing the 'stable-rank' assumption, would help clarify the authors’ theoretical contributions.

[1] Memory-Efficient LLM Training with Online Subspace Descent, Kaizhao Liang, Bo Liu, Lizhang Chen, Qiang Liu.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. How to relax the bounded gradient assumption in the study?
2. How to guarantee a feasible $\mathcal{U}_t$ if Gram-Schmidt process fails?
3. How to validate the theoretical conclusions from both theorems?
soundness: 3
strengths: 1. The investigated topic is quite interesting and critical, particularly when LLMs are popular in every area. A memory-efficient optimizer is required to conduct computationally efficient fine-tuning.
2. The paper is easy to follow and the presentation in this work is clear.
3. This work presents in-depth theoretical analysis for the proposed algorithm and show the explicits the convergence rate.
4. The empirical evidence to validate the proposed algorithm look extensive and thorough.
5. The comparison between the proposed method and baselines is technically convincing.
summary: Thin paper introduces a novel memory-efficient low-dimensional Adam optimizer for training large models. Specifically, the authors propose a new projection-aware update for the optimizer states that allows for transitioning between subspaces. They also perform block power iteration to at each step to efficiently and accurately project the gradient and optimizer states. To reduce the error caused by the low-rank projection, this work leverages a new generalized error feedback mechanism that accounts for both gradient and optimizer state compression. The authors theoretically analyze the convergence rates for different objectives. To validate the proposed algorithm, the authors present extensive experimental results on different datasets with the comparison to different baselines.
weaknesses: 1. Assumption 2 is a quite strong assumption for the gradient, which may limit the applicability of the proposed algorithm. Until now, most existing algorithms have relaxed this assumption in their works such that this assumption imposes some weaknesses to the proposed algorithm.
2. In this work, the authors have mentioned that GaLore is the closest work. Though we have seen the authors have shown in detail how the proposed algorithm differs, it would be great to include a technical discussion for the difference between them in terms of convergence rate. That way, it is clearer to directly show the exact theoretical difference.
3. In the practical view of the algorithm, the authors use Gram-Schmidt process to calculate $\mathcal{U}_t$. However, the process has numerical instability, computational intensity, and even inaccuracies. If this process fails, how can the authors guarantee that there always exists a feasible and reliable $\mathcal{U}_t$ for the subsequent steps?
4. Although the experimental results look promising, how to validate the theoretical conclusions from Theorem 1 and Theorem 2, particularly, the impact of the compression ratio? The authors need to present some ablation studies for the proposed algorithm.


**********************Post-rebuttal Comment***************************
After carefully reading the rebuttal from the authors and taking other reviewers' comments into consideration, I think the authors have addressed my major concerns. The new results also look interesting to me. Thanks for the efforts. I will also raise my score accordingly.

Meta-review: This paper presents LDAdam, a memory-efficient optimizer for training large models. The main idea is to perform adaptive optimization in low-dimensional subspaces while still exploring the full parameter space. The paper’s strengths are in its theoretical guarantees and practical performance. The projection-aware update rule helps to transition between subspaces smoothly, keeping the gradient statistics accurate even with dimensional reduction. The generalized error feedback mechanism helps reduce errors from low-rank projections, which allows LDAdam to perform similarly or better compared to full-rank Adam. The authors give detailed convergence proofs and show strong experiments, including results for models up to 1.3 billion parameters. The results show LDAdam is scalable, stable, and uses much less memory.

Reviewers found the paper clear and the idea interesting. The experiments show that LDAdam saves memory and performs better than baseline optimizers like Adam and GaLore in some cases. The authors answered reviewers’ concerns about computational cost, convergence rates, and compression ratios with more experiments, such as perplexity vs. runtime graphs and throughput vs. rank studies. Some concerns about bounded gradient assumptions and the Gram-Schmidt process were also addressed well. While there are some minor issues, the reviewers agree the paper is a good contribution to memory-efficient optimization methods.

Decision: Accept (Poster)