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

Authors: ionut-vlad  modoranu
mher  safaryan
grigory  malinovsky
eldar  kurtic
thomas  robert
peter  richtárik
dan  alistarh
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 10:51:22

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

Main paper: 
 introduction: 
 The Adam [Kingma and Ba, 2014] adaptive optimizer and its variants [Reddi et al., 2019, Loshchilov and Hutter, 2019] has emerged as a dominant choice for training deep neural networks (DNNs), especially in the case of large language models (LLMs) with billions of parameters. Yet, its versatility comes with the drawback of substantial memory overheads: relative to naive SGD-based optimization, the Adam optimizer states doubles the memory overhead, as it requires storing two additional parameters for each variable. For large-scale models, these memory demands pose a significant challenge. In turn, this has spurred research into memory-efficient adaptive optimizers, such as AdaFactor [Shazeer and Stern, 2018] , 8-bit Adam [Dettmers et al., 2021] , or the very recent GaLore [Zhao et al., 2024] low-rank projection approach. Despite their popularity and practical utility, the above methods lack rigorous convergence guarantees, and often trade off memory reductions with decreased convergence in practice. This raises the question of whether it is possible to design adaptive optimizers that are not only memory-efficient, but also maintain strong theoretical and practical performance metrics.In this paper, we address this gap by introducing MICROADAM, an adaptive optimizer which guarantees low memory usage but also ensures provable convergence. We develop our approach to improve the performance of finetuning LLMs and mainly focus on the research question "are all gradient entries important for optimization?" To answer this question, we start from the idea that we can allow the (lossy) sparse projection of gradient information before it enters the optimizer states; crucially, different from prior work, we ensure convergence by correcting for the inherent errorIn detail, at step t = 1, the error feedback e 1 is completely zero, as initialized in line 2, and thus, at line 5 the accumulator a 1 will only contain the stochastic gradient g 1 . At line 6, we perform the Top-K compression and only keep the top-1% of values V 1 and their corresponding indices I 1 . The compression is equivalent to choosing the top-1% values in the left and right tails (outliers) due to the absolute value we apply on top of accumulator a. At line 7, we remove the outliers from the accumulator because they will be transferred to the buffer matrix G. This step is equivalent to e ← a -T k (a) found in theoretical results. After line 7, what is left in a is called the error feedback (e.g. the weights which were not chosen by Top-k). At line 8, we compute the statistics δ and ∆ needed for quantization, and at line 9, we effectively quantize the accumulator (e.g. error feedback after line 7). At line 10 we update the buffer G, in lines 11, 12 and 13 we compute the statistics m, v (computed by squaring the entries of G element-wise) and update the model parameters.For steps t ≥ 2, the only change compared to t = 1 is that error feedback e is not zero anymore. Since the error is compressed, we need to decompress it and add it to the gradient. This process happens at line 5 and it is the point where we feed back the error: the accumulator will store the gradient whose direction is corrected by the error (e.g. the cumulative history of weights not chosen by Top-k at the previous steps).Properties and Limitations. We would like to point out that the resulting update u t = m t /(ϵ + √ v t )will always be highly sparse when the window size m is small. For illustration, if we use density k = d/100 (e.g. 1% density equivalent to 99% sparsity) with m = 10 and suppose that all rows in the indices matrix I are disjoint, then the overall density in the update u t will be 90%. The sparsity of u t increases if rows in I have common values. MICROADAM yields good results for LLM finetuning and pre-training computer vision models, as the experimental section shows. However, we noticed the update u t of MICROADAM is too sparse to be able to provide good enough updates for LLM pre-training. We believe this happens because the attention layers must receive dense updates to be able to learn the correlations between words.Dynamic Statistics. In ADAMSTATS procedure in Algorithm 2 we implement the unrolled recursion of momentum z t ← βz t-1 +(1-β)g t for the last m sparse gradients as z t ← (1-β) t i=t-m β t-i g i and we also perform the bias correction in the end. Because we compute mt and vt using the last m sparse gradients in the window, in line 4 we dynamically determine the exponent r for the decay factor β r based on the current optimization step t, i th row of the circular buffer G and the window size m. The last gradient added to G will have r = 0, while the oldest gradient in G will have r = m -1.In the end, we will add the values β r V i to the buffer z at the corresponding indices I i , which is a fast operation because we only manipulate 1% of values at a time. We discuss the efficient CUDA implementation in the Appendix.Algorithm Intuition. To gain intuition, we illustrate the impact of compressing gradients via TopK before they are incorporated into the optimizer state for Adam, both with and without error feedback (EF). Figure 1 shows how EF fixes AdamW with Top-K compression. The plot on the left shows the optimization trajectory of the original Adam optimizer. The center plot illustrates the convergence of Top-K Adam when we only choose the largest coordinate from the accumulator (equivalent to 50% sparsity since the problem is 2D). In the end, on the right side we show that adding EF to Top-K Adam recovers the same optimization trajectory as the original Adam optimizer. Extrapolating to higher dimensional problems, our MICROADAM approach helps recover the trajectory of the original Adam optimizer, while using less memory. The results clearly show that EF is essential for fast convergence. Besides, TopK with EF, which is a surrogate of MICROADAM, allows for competitive convergence relative to the uncompressed baseline. In Appendix F, we discuss the implications of Error Feedback applied to GaLore.Algorithm 1 Pseudocode for MICROADAM with quantized EF 1: Input: for i ∈ {1, 2, ..., min(t, m)} do 4: r ← (t -i -1)%m 5:β 1 , β 2 , ϵ, G, T, d, k 2: m 0 , v 0 ← 0 d , 0 d δ 1 , ∆ 1 ← 0, 0 e 1 ← 0 4bg t ← ∇ θ f (θ t ) 5: a t ← g t + Q -1 (e t ,z[I i ] ← z[I i ] + β r V i 6:end for 7:return (1-β)z 1-β t 8: end procedure 1: procedure Q(x, δ, ∆, b = 4) 2: u ← ∆-δ 2 b -1 3: x Q ← ⌊ x-δ u + 1 2 ⌋ 4: return x Q 5: end procedure 1: procedure Q -1 (x Q , δ, ∆, b) 2: u ← ∆-δ 2 b -1 3: x ← x Q • u + δ 4:return x 5: end procedure f (x, y) = (1 -x) 2 + 100(y -x 2 ) 2 starting from (x 0 , y 0 ) = (-12 , 1). Notice the extremely "jagged" profile of TopK-Adam without EF, and the recovered convergence when EF is added. 
efficient implementation: 
 A direct implementation (e.g., in Pytorch) of the previous algorithm would not bring significant benefits, and would in fact might slow down optimization in terms of wall-clock time. To realize the theoretical gains, we detail a GPU-aware implementation below.Accumulator a t . First, we do not use an additional accumulator tensor a t ; instead, we use a CUDA kernel to dequantize the error buffer, and store the result in the grad attribute of the model parameters. This allows us to accumulate the error feedback into the gradients, without allocating a full or half precision d-dimensional array. Each component of the EF has 4 bits and the entire EF is stored in an array of size d/2 of uint8 values.Top-K. Since we run on LLMs with billions of parameters, naive storage of the sparse indices would require using an int64 type for the indices matrix I, assuming that the Top-K operator is applied globally to all the parameters in a t . To avoid this cost, we apply Top-K in blocks of fixed size B d < 2 15 -1 = 32767 and store block-relative indices in int16 format (during the development of MICROADAM, PyTorch did not have support for uint16). Applying Top-K per row to a t reshaped to 2D is not only faster, but provides the block-relative indices directly.Computing Top-K in blocks also allows us to allocate and efficiently use CUDA shared memory blocks to dynamically compute the statistics mt and vt for Adam, as described in the ADAMSTATS procedure in Algorithm 2. We allocate the maximum possible shared memory for each thread block and store mt (first half) and vt (second half) at consecutive locations in the shared memory. Once these statistics are computed, it is easy to update the model parameters. Note that the block-relative indices returned by Top-K will be directly used as indices in the shared memory array of CUDA thread blocks to retrieve values from I and V.Quantization metadata. Our approach also stores two additional vectors δ and ∆ used for quantization. Since the quantization block size B q is set to a very large value (e.g. 100K), the space required for these two arrays becomes negligible in comparison to the buffer G and error feedback e.Practical memory usage. We note that we apply MICROADAM per layer, and that the size of quantization statistics δ and ∆ are allocated based on the layer size. Having many such small tensors may result in slightly sub-optimal memory allocation from Pytorch. This is why our reported memory usage can be higher than the theoretical usage for small models, in the 100M parameter range; these effects disappear for billion-parameter models, where the savings are significant.We now compare the theoretical memory footprint of MICROADAM with AdamW [Loshchilov and Hutter, 2019] , AdamW-8 bits [Dettmers et al., 2021] , and GaLore [Zhao et al., 2024] , focusing on memory usage of the optimizer states m t and v t , each of size d, expressed in bytes (B). For concreteness, we report the practical memory usage for the optimizer state for a Llama-2 7B model for each optimizer.• AdamW stores states in float32 format (4 B per component), resulting in a total memory footprint of 4d + 4d = 8d (B), while using bfloat16 would result in 4d (B) memory. We will refer to these memory footprints as • GaLore. Given a neural network with L layers, where each layer has weight matrixEQUATIONW i ∈ R Ai×Bi (shaped as a 2D matrix), the model size d = L i=1 A i B i .GaLore uses a rank-r compression via the SVD composition as W i = U SV T , where U ∈ R Ai×Ai and we choose the first r columns of U as R i ∈ R Ai×r to project the gradient of W i to the r-dimensional space. As a consequence, the dimension d shrinks to d r = L i=1 A i r = r L i=1 A i , which represents the total number of components to be stored in the GPU memory only for the projection matrices R i . If we suppose they are stored in bfloat16 (2 B per component), then the entire memory usage for low-rank projection would be 2d r . Note that some layers in the model are rank-1 and they do not need compression, but will still have associated states in Adam, which means they must be tan into consideration when computing the theoretical memory (we will use ϵ 1 for memory of rank-1 layers).In addition, we have to store the states m t and v t for AdamW in bfloat16 format, which adds another 4d r bytes. In sum, the total memory footprint of GaLore is M GLAW 16 (r) = 6d r + 2ϵ 1 (B), while for the 8-bit version we get M GLAW 8 (r) = 4d r + 2ϵ 1 (B). In the end, the practical memory usage for Llama-2 7B is M GLAW 8 (256) = 1.36 (GB), M GLAW 8 (1024) = 5.43 (GB), M GLAW 16 (256) = 2.04 (GB) and M GLAW 8 (1024) = 8.15 (GB).Discussion. Assume our goal is to obtain a lower memory footprint compared to AdamW-8 bit. We fix the gradient density to k = d/100 and we have to determine the number of gradients (window size) m for MICROADAM in order to be competitive with AdamW-8 bit.For this, we have to solve the equation M µA (m) = M AW 8 for m, resulting in m max = 37.5. Specifically, if we use a gradient history of m < m max gradients in G, MICROADAM will have theoretical memory savings. We will see that, in practice, this history size m = 10 is more than enough for good practical results. As entries in the window past this range are dampened extremely significantly, their influence is negligible. In Appendix D, we provide Python code to compute the theoretical memory usage for the three optimizers for Llama-2 7B. 
convergence guarantees for microadam: 
 In this section, we present our theoretical framework. We begin by introducing and discussing the analytical assumptions used in our theory, providing an analytical view of the proposed MICROADAM algorithm, along with two theoretical convergence results.We now define two classes of compression operators widely used in the literature. Assumption 1. The gradient compressor C : R d → R d is q-contractive with 0 ≤ q < 1, i.e., ∥C(x) -x∥ ≤ q ∥x∥ , for any x ∈ R d .The compression we use in Algorithm 1 is the TopK compressor T k selecting top k coordinates in absolute value. This is known to be contractive with q = 1 -k /d. Another popular contractive compressor is the optimal low-rank projection of gradient shaped as a d × d matrix, in which case q = 1 -R/d where R is the projection rank.The second class of compressors, which we use for the error feedback, requires unbiasedness and relaxes the constant in the uniform bound. Assumption 2. The error compressor Q : R d → R d is unbiased and ω-bounded with ω ≥ 0, namely,E[Q(x)] = x, ∥Q(x) -x∥ ≤ ω ∥x∥ , for any x ∈ R d .One example of ω-bounded compressor, a version of which is used in Algorithm 2, is the randomized rounding quantizer we employ, whose properties we provide below. Lemma 1. Consider Algorithm 2 with randomized rounding, i.e., for a vector x ∈ R d with δ = min i x i and ∆ = max i x i , let xi := ⌊ xi-δ u + ξ⌋u + δ be the i-th coordinate of the quantized vector x, where ξ ∼ U[0, 1] is the uniform random variable and u = ∆-δ 2 b -1 is the quantization level. ThenE[x] = x, ∥x -x∥ ≤ √ d-2 2 b -1 ∆-δ √ ∆ 2 +δ 2 ∥x∥, for all x ∈ R d .Next, we provide an "analytical" view of our method in Algorithm 3. Essentially, we use the contractive compressor C for compressing the error corrected gradient information g t + e t , and the unbiased compressor Q to compress the remaining compression error g t + e t -C(g t + e t ). Algorithm 3 MICROADAM: Analytical View 1: Input: parameters β 1 , β 2 ∈ (0, 1), ϵ > 0, step-size η > 0, θ 1 ∈ R d , e 1 = m 0 = v 0 = v0 = 0 d 2: for t = {1, 2, ..., T } do 3: g t = ∇ θ f (θ t ) ⋄ Computem t = β 1 m t-1 + (1 -β 1 )g t ⋄ Update first-order gradient moment 7: v t = β 2 v t-1 + (1 -β 2 )g 2 t ⋄θ t+1 = θ t -η mt √ vt+ϵ⋄ Update the model parameters 10: end for It is clear from this description that our objective with these two compressors, C and Q, is to approximate the dense gradient information g t + e t using two compressed vectors: gt = C(g t + e t ) and Q(g t + e t -gt ). However, in doing so, we inevitably lose some information about g t + e t depending on the degree of compression applied to each term. Thus, the condition (1 + ω)q < 1 required by our analysis can be seen as preventing excessive loss of information due to compression.Next, we state our algorithm's convergence guarantees under standard assumptions, stated below: Assumption 3 (Lower bound and smoothness). The loss function f : R d → R is lower bounded by some f * ∈ R and L-smooth, i.e., ∥∇f (θ) -∇f (θ ′ )∥ ≤ L ∥θ -θ ′ ∥, for any θ, θ ′ ∈ R d . Assumption 4 (Unbiased and bounded stochastic gradient). For all iterates t ≥ 1, the stochastic gradient g t is unbiased and uniformly bounded by a constant G ≥ 0, i.e., E[g t ] = ∇f (θ t ), ∥g t ∥ ≤ G. Assumption 5 (Bounded variance). For all iterates t ≥ 1, the variance of the stochastic gradient g t is uniformly bounded by some constantσ 2 ≥ 0, i.e., E[∥g t -∇f (θ t )∥ 2 ] ≤ σ 2 .Main Result. The above assumptions are standard in the literature, e.g. [Défossez et al., 2022 , Li et al., 2022 , Xie et al., 2023 , Zhou et al., 2024a] . Under these conditions, if the two compressors satisfy the basic condition (1 + ω)q < 1, we show: Theorem 1. (Non-convex convergence rate) Let Assumptions 1, 2, 3, 4, 5 hold and q ω := (1+ω)q < 1.Then, choosing η = min{ ϵ 4LC0 , 1 √ T }, MICROADAM (Algorithm 3) satisfies 1 T T t=1 E[∥∇f (θ t )∥ 2 ] ≤ 2C 0 f (θ1)-f * √ T + L(σ 2 +C 2 2 G 2 ) ϵ √ T + O G 3 (G+d) T with constants C 0 := 4(1+q 2 ω ) 3 (1-q 2 ω ) 2 G 2 + ϵ and C 2 := ωq(1 + 2qω 1-q 2 ω ).Discussion. First, notice that the leading term 1 √ T of the rate is the optimal convergence speed for non-convex stochastic gradient methods [Ghadimi and Lan, 2016] . Furthermore, the obtained convergence rate O( 1√ T + d T )asymptotically matches the rate of uncompressed AMSGrad in the stochastic non-convex setup [Zhou et al., 2024a] . Hence, the added compression framework of the MICROADAM together with error feedback mechanism can slow down the convergence speed only up to some constants including the dimension. Evidently, the additional constants C 0 and C 2 affected by compression and appearing in the leading terms can be easily estimated once the compressors are fixed. Besides, if we store the full error information without applying Q compressor (i.e., ω = 0, q ω = q), then MICROADAM reduces to the single-node Comp-AMS method by Li et al. [2022] recovering the same convergence rate. The full proof is provided in the Appendix.Next, we show that we can obtain even stronger bounds when the objective satisfies the PL condition: Assumption 6 (PL-condition). For some µ > 0 the loss f satisfies Polyak-Lojasiewicz (PL) inequality∥∇f (θ)∥ 2 ≥ 2µ(f (θ) -f * ), for any θ ∈ R d .In this case, we can show: Theorem 2. (PL convergence rate) Let Assumptions 1, 2, 3, 4, 5 and 6 hold, and q ω < 1. Then, choosing η = min{ ϵ 4LC0 , 2C0 log T µT }, MICROADAM (Algorithm 3) satisfiesE[f (θ T +1 )] -f * ≤ 2 log T T LC 2 0 µ σ 2 +(C1+C 2 2 )G 2 µϵ + C0(1+C1)(1+d)G 2 µ √ ϵ + O G 4 (G+d) T 2 with constant C 1 := β1 1-β1 (1 + C 2 ) + 2qω 1-q 2 ω .Discussion. In contrast to the general non-convex setup, the study of non-convex analysis under the PL condition for AMSGrad or Adam-type methods is much less extensive. The only work we found analyzing the PL condition, which claims to be the first in this direction, focuses on Adam when β 2 → 1, achieving a convergence rate of O( 1 T ) [He et al., 2023] . However, our MICROADAM is based on AMSGrad normalization, and no constraint on β 2 is imposed in the analysis. Therefore, similar to the general non-convex case, we are able to achieve the best-known convergence rate in the leading term, up to a logarithmic factor. The third, higher-order term has higher constant dependencies, but they should be negligible as the term is dampened by T 2 . Hence, in this case as well, the theory predicts that the convergence rate of the algorithm should be similar to the uncompressed version, modulo a constant that can be controlled using the compression parameters. 
experiments: 
 We now validate our optimizer experimentally. We focus on comparing MICROADAM with Adam, Adam-8bit, GaLore and CAME in the context of LLM finetuning on different tasks and with SGD, Adam and AdamW-8bit in the context of ResNets on ImageNet. Concretely, we test our optimizer in full finetuning (FFT) scenario on BERT-Base/Large [Devlin et al., 2018] and OPT-1.3B [Zhang et al., 2022] on GLUE/MNLI and Llama2-7B/13B [Touvron et al., 2023] on the GSM8k math reasoning dataset and on the Open-Platypus instruction tuning dataset, as well as pre-training ResNet models on ImageNet. We provide full details regarding training settings hyper-parameters in Appendix B.Finetuning results on GLUE/MNLI. We first test our integration of MICROADAM in HuggingFace Transformers [Wolf et al., 2020] on moderate-sized language models such as BERT-Base/Large (110M and 335M parameters) and OPT-1.3B, comparing with Adam, Adam-8bit, CAME and GaLore. The results are shown in Table 1 . Certain optimizers, notably CAME and GaLore, had numerical stability issues across runs; for a fair comparison, we report the numbers for the run with maximum accuracy. We emphasize that all methods were tuned using the same protocol.The results show that MICROADAM achieves comparable memory usage to the state-of-the-art heuristics Adam-8bit and GaLore, while being surprisingly lower than CAME on all tasks. The memory savings for GaLore are more visible when the model size increases, which follows our analysis of theoretical memory usage. However, we see that these gains come at a significant accuracy cost for GaLore: across all tasks, it drops at least 1% accuracy relative to MICROADAM. For BERT-Base we ran GaLore with a higher SVD re-computation frequency T = 20 (10× lower) and the results did not improve, but its running time was much higher. Relative to 8bit Adam, MICROADAM uses essentially the same memory, but achieves slightly better accuracy.From these results, we conclude that MICROADAM can provide better accuracy relative to other memory-efficient methods on moderate-sized models, at similar space costs. We show training loss curves in Appendix C.Table 1 : Finetuning results on GLUE/MNLI. We report the entire memory usage read from the GPU during training, that includes the optimizer state, activations and gradients. The asterisk flags the runs for which one or two seeds did not converge (we report the run with maximum performance). Finetuning results for LLaMA2 on GSM-8k. Next, we perform finetuning on Llama-2 7B/13B on GSM-8k, a challenging grade-school-level mathematical reasoning dataset. The baseline model obtains extremely low zero-shot accuracy on this task and therefore fine-tuning is necessary. In this setup, we compare MICROADAM with Adam and Adam-8bit in terms of evaluation accuracy and memory usage. In Table 2 we show our results for 3 training epochs, global batch size 32 with micro-batch (per-device) size 1, max sequence length 512 on a single GPU, which are the standard parameters for this task. We integrated our optimizer with the llm-foundry repository of MosaicML and tested via lm-evaluation-harness.For the 7B model, out results show that MICROADAM can allow accurate full fine-tuning of a 7B model on this task using a single 40GB GPU. Moreover, MICROADAM preserves accuracy relative to Adam, with lower memory usage than the well-optimized implementation of 8bit AdamW, and marginally lower running time for the shorter gradient window m = 10. Increasing the window size m to 20 gradients leads to slightly better accuracy, at the cost of higher runtime and space, but still in the 40GB limit. Running GaLore in this setup was infeasible since using SVD decomposition for all layers in the model was too slow. Preliminary experiments (with high runtimes) did not yield competitive accuracy. We show training loss curves in Appendix C.The results show that MICROADAM allows for full accuracy recovery on this task as well relative to Adam, despite using 50% less memory. (The memory usage and runtime are very similar to those in Table 2 and are therefore omitted from Table 3 .) Moreover, MICROADAM obtains consistently better accuracy relative to Adam-8b, especially on the more challenging ARC-c task. Finetuning results for LLaMA2-7B on Open-Platypus. Finally, in Table 3 we present FFT results with various optimizers on the popular instruction-tuning Open-Platypus dataset [Lee et al., 2023] .To ensure fair comparisons, we perform the same grid search for each optimizer to find the best performing learning-rate, while keeping all other hyperparameters at their default values. We use m = 10 gradients for the sliding window and gradient density k = 1%. Evaluations are conducted following the standard few-shot setup of the Open LLM Leaderboard [Beeching et al., 2023] on the following datasets: ARC-c [Clark et al., 2018] , HellaSwag [Zellers et al., 2019] , MMLU [Hendrycks et al., 2021] , and Winogrande [Sakaguchi et al., 2019] .Table 3 : FFT results on instruction-following Open-Platypus [Lee et al., 2023] dataset. The results show that MICROADAM fully recovers accuracy relative to baseline Adam, and outperforms the 8bit variant, despite using less memory. 7 in Appendix C). We compare our MICROADAM with SGD, Adam and AdamW and report the training loss, validation accuracy and only the optimizer state memory because the total memory usage is not an issue for ResNets on ImageNet (here we focus on the results to emphasize the pre-training performance). For ResNet-18, AdamW reaches the lowest training loss, followed by AdamW-8bit, MICROADAM and in the end MICROADAM. However, despite slightly larger loss, MICROADAM yields the best result among all optimizers, with 2% more than the highly tuned SGD, while having the lowest memory footprint for the optimizer states. For ResNet-50, AdamW-8bit reaches the lowest training loss, followed by AdamW, MICROADAM and SGD. The validation accuracy for AdamW and AdamW-8bit is surprisingly small compared to SGD and MICROADAM. As it is widely known in the community, Adam variants have lower performance than SGD for Computer Vision tasks and MICROADAM fixes this issue (see the Discussion section for an intuitive explanation for this phenomenon). of the weights at each step yields to significantly better performance compared to SGD for ResNets on ImageNet. Secondly, our experiments on LLM pre-training showed difficulties in achieving the same performance compared to AdamW-8bit. Our explanation is that projection matrices from the attention layers must receive dense updates to learn the correlations between words. In contrast to the convolutional filters for CV models, the weights in attention are much larger and try to capture global correlations (features) between words, while the convolutional filters are smaller and capture local features.Discussion. In Appendix A we provide information about the optimization set, intuitive explanations for the implicit regularization effect of MICROADAM, as well as an overview of our results. 
limitations and broader impact: 
 The MICROADAM algorithm we propose is designed and tested with fine-tuning workloads in mind, where the user aims to minimize the memory cost of optimizing over a powerful pre-trained model. Additional work is needed to adapt our approach to the case of LLM pre-training, which presents a different set of challenges, both in terms of implementation and optimization trajectory. We plan to undertake this study in future work as the current implementation works for ResNets.Another limitation we aim to address in future work is that we have only focused on sparsity as a form of gradient projection. However, our theoretical analysis also applies to low-rank projection of gradients. We believe that our practical implementation can be extended to this case as well, although providing a general, accurate, and efficient implementation will require non-trivial efforts.Our work introduces a new, accurate, and memory-efficient optimizer for fine-tuning LLMs. The major positive impact of our approach is its ability to maintain performance while reducing memory requirements, thereby lowering the cost of running experiments due to the reduced hardware expenses.It is important to note that while our optimizer can enhance performance and reduce costs, we do not have control over the neural network applications trained with it. Optimization set. The parameters included in the optimization set usually vary depending on the model and optimizer type. For GLUE, we do not include the embeddings in the optimization set for any of the optimizers because our experiments showed no significant difference when optimizing the embeddings. Moreover, it is more fair for GaLore which would have an increased memory usage due to the projection matrix for the embeddings. For LLaMa2 and ResNet models, we include all layers in the optimization set, regardless of their type. This means that MICROADAM updates at most 10% of the weights in each layer. In the original work, GaLore was applied only to a subset of layers, such as Q-, K-, V-, O-, up-, down-and gate-projection layers. Applying GaLore to all layers in the same way as we do with MICROADAM would result in much larger memory usage because of the projection matrices for the embeddings. Moreover, computing SVD for the embedding layer would be infeasible. As a result, we omit the GaLore for LLaMa2 and ResNet experiments.Implicit regularization. In our results so far, we observed MICROADAM having better performance for a few LLM finetuning tasks, but especially for the ResNet/ImageNet results, where the difference was statistically significant (around 1%). Our intuition for this behavior mainly comes from the accuracy curves in Figure 6 and Figure 7 . The trajectories of MICROADAM and SGD are typical for regularized training. We hypothesize that MICROADAM has an implicit regularization mechanism because the model update u t is 90% sparse, which leads to only updating 10% of the model parameters at each step in each layer. In contrast to a 100% dense update u t , a sparse update would not change the model parameters as much as a dense one. In the ResNet experiments, all optimizers used the same regularization parameter λ = 1e -4, but the accuracy graph shows that MICROADAM is more regularized than SGD, while the graphs for AdamW and AdamW-8bit look like the regularization does not have any effect.Discussion. In summary, the experimental results have shown that MICROADAM can recover the state-of-the-art accuracy of the the uncompressed Adam baseline, while providing significant memory gains and matching wall-clock speed on billion-parameter models. Specifically, our approach matches and outperforms Adam-8b and CAME both in terms of memory use and in terms of final accuracy. Relative to the high-compression GaLore method, MICROADAM provides consistently higher accuracy, as well as more stable practical convergence. We conclude that MICROADAM should be a good alternative to Adam-8bit in memory-constrained settings, and that the empirical results appear to validate our theoretical predictions.In this section we provide details about the hyper-parameters that we used for each model and dataset. We train all our models in bfloat16 format, tune the learning rates on a grid and report the best accuracy among 3 seeds (7, 42 and 1234) and report the results for the best configuration that converged.All Adam variants use default parameters β 1 = 0.9, β 2 = 0.999, ϵ = 10 -8 and the regularization parameter λ is 0 for finetuning and 3e -4 for ImageNet pre-training. MICROADAM uses a window size of m = 10 gradients with k = 1% density (equivalent to 99% sparsity and quantization bucket size is set to 64 for the error feedback.For GaLore we use rank r = 256 and the SVD update interval is set to T = 200, as suggested by the original paper. We run our experiments on NVidia GPUs A100-SXM4-80GB, H100-80GB and on RTX 3090 with 24GB RAM in single GPU setup.For GLUE/MNLI, we used the learning rate grid {1e -6, 3e -6, 5e -6, 7e -6, 1e -5, 3e -5, 5e -5, 7e -5} for all optimizers and models. Certain optimizers diverge for specific seeds. Next, we provide some details about hyper-parameters for each optimizer individually.MICROADAM. We use m = 10 gradients in the sliding window, k = 1% density (e.g. 99% sparsity) and quantization bucket size 64 (we also tried 100 000, but this didn't affect performance or memory usage in a meaningful way).Adam and Adam-8bit. All hyper-parameters mentioned above apply for these two main baseline optimizers.GaLore. We use rank r = 256 and SVD update interval T ∈ {20, 200}. In the original GaLore paper, the authors tune both learning rate and in our experiments we keep scale fixed to value 1 and augment the learning rate grid with the values {1e -4, 3e -4, 5e -4, 7e -4}.CAME. This optimizer has some additional parameters that we keep to default values, such as β 3 = 0.9999. Instead of ϵ, it uses ϵ 1 = 1e -30 and ϵ 2 = 1e -16. The authors mention that the learning rate should be much smaller than Adam's and because of that we augment the learning rate grid with the values {1e -7, 3e -7, 5e -7, 7e -7}.B.2 GSM-8k.For GSM-8k, we used the learning rate grid {1e -5, 2e -5, 3e -5, 4e -5, 5e -5, 6e -5, 7e -5, 8e -5, 9e -5} and reported the model with the best evaluation accuracy. We found that different versions for PyTorch, lm-eval-harness and llm-foundry have large variance in the results.MICROADAM. We use similar settings as for GLUE/MNLI above in terms of other hyper-parameters.For ImageNet, we integrate our MICROADAM in the FFCV repository, which is highly tuned for ResNets and SGD. We use E = 100 epochs, batch size 1024, cosine learning rate schedule with warmup and image resolution 224 × 224 and precision bfloat16. We started from the initial learning rate η = 1.024 tuned in the repository which scored highest accuracy for SGD. This learning rate also worked well for MICROADAM, but it didn't work for AdamW and AdamW-8bit. For these two Adam variants we divided the learning rate by 2 until the models converged.ResNet-18. For AdamW, the learning rate is η = 0.016 and for AdamW-8bit is η = 0.032.ResNet-50. For AdamW, the learning rate is η = 0.008 and for AdamW-8bit is η = 0.008.In this section we show training loss curves for BERT-Base, BERT-Large and OPT-1.3b on GLUE/MNLI and Llama-2 7B/13B on GSM-8k and ResNet-18/50 on ImageNet. In this section we provide a python script to simulate the memory usage for our optimizer's state for Llama2-7b model. Note that the theoretical memory usage will always be slightly lower than the actual allocated memory on the GPU because PyTorch usually allocates more. To run this script, run the following commands: p r i n t ( f ' {M_AW32= : . 2 f } GB ' ) p r i n t ( f ' {M_AW16= : . 2 f } GB ' ) p r i n t ( f ' {M_AW8= : . 2 f } GB ' ) p r i n t ( f ' {M_muA= : . 2 f } GB ' ) GL_sumA = 1 _423_872 # sum_Ai fAt time step t, let the uncompressed stochastic gradient be g t = ∇ θ f (θ t ), the error accumulator be e t , and the compressed gradient after the error correction be gt = C(g t + e t ). The second moment computed by the compressed gradients is denoted asv t = β 2 v t-1 + (1 -β 2 )g 2 t ,and vt = max{v t-1 , v t } is the AMSGrad normalization for the second-order momentum. Besides the first-order gradient momentum m t used in the algorithm description, we define similar running average sequence m ′ t based on the uncompressed gradients g t .m t = β 1 m t-1 + (1 -β 1 )g t and m ′ t = β 1 m ′ t-1 + (1 -β 1 )g t , Note that m ′t is used only in the analysis, we do not need to store or compute it. By construction we have m t = (1 -β 1 ) t τ =1 β t-τ 1 gτ , m ′ t = (1 -β 1 ) t τ =1 β t-E t+1 := β 1 E t + (1 -β 1 )e t+1 = (1 -β 1 ) t+1 τ =1 β t+1-τ 1 e τ . Z t+1 := β 1 Z t + (1 -β 1 )ζ t+1 = (1 -β 1 ) t+1 τ =1 β t+1-τ 1 ζ τ .Lemma 1. Consider Algorithm 2 with randomized rounding, i.e., for a vector x ∈ R d with δ = min i x i and ∆ = max i x i , let xi := ⌊ xi-δ u + ξ⌋u + δ be the i-th coordinate of the quantized vector x, where ξ ∼ U[0, 1] is the uniform random variable and u = ∆-δ 2 b -1 is the quantization level. ThenE[x] = x, ∥x -x∥ ≤ √ d-2 2 b -1 ∆-δ √ ∆ 2 +δ 2 ∥x∥, for all x ∈ R d .Proof. The unbiasedness can be verified directly from the definition for each coordinate. Without loss of generality assume that δ= x 1 ≤ x 2 ≤ • • • ≤ x d-1 ≤ x d = ∆. By construction of the quantization, we have |x 1 -x 1 | = |x d -x d | = 0 and |x i -x i | ≤ u for the remaining coordinates 2 ≤ i ≤ d -1. Then ∥x -x∥ 2 = d i=1 |x i -x i | 2 ≤ (d -2)u 2 ≤ (d -2)u 2 ∆ 2 + δ 2 ∥x∥ 2 ,which completes the proof.Lemma 2. Under Assumptions 1-5, for all iterates t and T we have∥m ′ t ∥ ≤ G, and T t=1 E[∥m ′ t ∥ 2 ] ≤ T σ 2 + T t=1 E[∥∇f (θ t )∥ 2 ].Proof. The first part follows from triangle inequality and the Assumption 4 on bounded stochastic gradient:∥m ′ t ∥ = (1 -β 1 ) t τ =1 β t-τ 1 g τ ≤ (1 -β 1 ) t τ =1 β t-τ 1 ∥g τ ∥ ≤ G.For the second claim, the expected squared norm of average stochastic gradient can be bounded byEQUATIONwhere we use Assumption 5 that g t is unbiased with bounded variance. Let g t,j denote the j-th coordinate of g t . Applying Jensen's inequality for the squared norm, we getE[∥m ′ t ∥ 2 ] = E   (1 -β 1 ) t τ =1 β t-τ 1 g τ 2   ≤ (1 -β 1 ) t τ =1 β t-τ 1 E[∥g τ ∥ 2 ] ≤ σ 2 + (1 -β 1 ) t τ =1 β t-τ 1 E[∥∇f (θ τ )∥ 2 ],Summing over t = 1, . . . , T , we obtainT t=1 E[∥m ′ t ∥ 2 ] ≤ T σ 2 + (1 -β 1 ) T t=1 t τ =1 β t-τ 1 E[∥∇f (θ τ )∥ 2 ] ≤ T σ 2 + T t=1 E[∥∇f (θ t )∥ 2 ],which completes the proof.Lemma 3. Let q ω = (1 + ω)q < 1. Under Assumptions 1-5, for all iterates t we have∥e t ∥ 2 ≤ 4q 2 ω (1 -q 2 ω ) 2 G 2 , E[∥e t+1 ∥ 2 ] ≤ 4q 2 ω (1 -q 2 ω ) 2 σ 2 + 2q 2 ω 1 -q 2 ω t τ =1 1 + q 2 ω 2 t-τ E[∥∇f (θ τ )∥ 2 ].Proof. We start by using Assumption 1, 2 on compression and Young's inequality to getEQUATIONwhere ( 2) is derived by choosing ρ = 1-q 2 ω 2q 2 ω and the fact that q ω < 1. For the first claim we recursively apply the obtained inequality and use bounded gradient Assumption 4. For the second claim, initialization e 1 = 0 and the obtained recursion implyE[∥e t+1 ∥ 2 ] ≤ 2q 2 ω 1 -q 2 ω t τ =1 1 + q 2 ω 2 t-τ E[∥g τ ∥ 2 ] (1) ≤ 4q 2 ω (1 -q 2 ω ) 2 σ 2 + 2q 2 ω 1 -q 2 ω t τ =1 1 + q 2 ω 2 t-τ E[∥∇f (θ τ )∥ 2 ],which concludes the lemma.Lemma 4. Let q ω = (1 + ω)q < 1. Under Assumptions 1-5, for all iterates t we have∥ζ t ∥ ≤ ωq 1 + 2q ω 1 -q 2 ω G, and ∥Z t ∥ ≤ ωq 1 + 2q ω 1 -q 2 ω G.Proof. Using the bounds defining compressors and Lemma 3, we get∥ζ t ∥ = ∥Q(e t + g t -gt ) -(e t + g t -gt )∥ ≤ ω∥e t + g t -gt ∥ = ω∥e t + g t -C(e t + g t )∥ ≤ ωq∥e t + g t ∥ ≤ ωq∥e t ∥ + ωq∥g t ∥ ≤ ωq 1 + 2q ω 1 -q 2 ω G.For the second claim, recall the definition of Z t and apply triangle inequality:∥Z t ∥ ≤ (1 -β 1 ) t τ =1 β t-τ ∥ζ τ ∥ ≤ ωq 1 + 2q ω 1 -q 2 ω G.Lemma 5. For the moving average error sequence E t , it holds thatT t=1 E[∥E t ∥ 2 ] ≤ 4T q 2 ω (1 -q 2 ω ) 2 σ 2 + 4q 2 ω (1 -q 2 ω ) 2 T t=1 E[∥∇f (θ t )∥ 2 ].Proof. Let e t,j be the j-th coordinate of e t and denoteK t := t τ =1 1+q 2 ω 2 t-τ E[∥∇f (θ τ )∥ 2 ].Applying Jensen's inequality and Lemma 3, we getE[∥E t ∥ 2 ] = E   (1 -β 1 ) t τ =1 β t-τ 1 e τ 2   ≤ (1 -β 1 ) t τ =1 β t-τ 1 E[∥e τ ∥ 2 ] ≤ 4q 2 ω (1 -q 2 ω ) 2 σ 2 + 2q 2 ω (1 -β 1 ) (1 -q 2 ω ) t τ =1 β t-τ 1 K τ ,Summing over t = 1, . . . , T and using the technique of geometric series summation leads toT t=1 E[∥E t ∥ 2 ] ≤ 4T q 2 ω (1 -q 2 ω ) 2 σ 2 + 2q 2 ω (1 -β 1 ) (1 -q 2 ω ) T t=1 t τ =1 β t-τ 1 K τ ≤ 4T q 2 ω (1 -q 2 ω ) 2 σ 2 + 2q 2 ω (1 -q 2 ω ) T t=1 K t = 4T q 2 ω (1 -q 2 ω ) 2 σ 2 + 2q 2 ω (1 -q 2 ω ) T t=1 t τ =1 1 + q 2 ω 2 t-τ E[∥∇f (θ τ )∥ 2 ] ≤ 4T q 2 ω (1 -q 2 ω ) 2 σ 2 + 4q 2 ω (1 -q 2 ω ) 2 T t=1 E[∥∇f (θ t )∥ 2 ],The desired result is obtained.Lemma 6. Let q ω = (1 + ω)q < 1. Under Assumptions 1-5, for all iterates t ∈ [T ] and coordinates i ∈ [d], the following bound holdsvt,i ≤ 4(1 + q 2 ω ) 3 (1 -q 2 ω ) 2 G 2 .Proof. Lemma 3 and Assumption 4 imply∥g t ∥ 2 = ∥C(g t + e t )∥ 2 ≤ ∥C(g t + e t ) -(g t + e t ) + (g t + e t )∥ 2 ≤ 2(q 2 + 1)∥g t + e t ∥ 2 ≤ 4(q 2 + 1) G 2 + 4q 2 ω (1 -q 2 ω ) 2 G 2 = 4(1 + q 2 )(1 + q 2 ω ) 2 (1 -q 2 ω ) 2 G 2 .It's then easy to show by the updating rule of vt , there exists a j ∈ [t] such that vt,i = v j,i . Thenvt,i = (1 -β 2 ) j τ =1 β j-τ 2 g2 τ,i ≤ 4(1 + q 2 )(1 + q 2 ω ) 2 (1 -q 2 ω ) 2 G 2 ,which concludes the claim.Lemma 7.For D t := 1 √ vt-1+ϵ -1 √ vt+ϵ we have T t=1 ∥D t ∥ 1 ≤ d √ ϵ , T t=1 ∥D t ∥ 2 ≤ d ϵ .Proof. By the update rule, we have vt-1,i ≤ vt,i for any iterate t and coordinate i ∈ [d] . Therefore, by the initialization v0 = 0, we getT t=1 ∥D t ∥ 1 = T t=1 d i=1 1 vt-1,i + ϵ - 1 vt,i + ϵ = d i=1 1 v0,i + ϵ - 1 vT,i + ϵ ≤ d √ ϵ .For the sum of squared l 2 norms, note the fact that for a ≥ b > 0, it holds that(a -b) 2 ≤ (a -b)(a + b) = a 2 -b 2 .Thus,T t=1 ∥D t ∥ 2 = T t=1 d i=1 1 vt-1,i + ϵ - 1 vt,i + ϵ 2 ≤ T t=1 d i=1 1 vt-1,i + ϵ - 1 vt,i + ϵ ≤ d ϵ ,which gives the desired result.Here we derive the convergence rate with fixed step-size η. The rate shown in the main part can be obtained by plugging the expression of η shown after the proof.Theorem 3. (Non-convex convergence rate) Let Assumptions 1, 2, 3, 4, 5 hold and q ω := (1+ω)q < 1. Then, choosing any step-size η ≤ ϵ 4LC0 , MICROADAM (Algorithm 3) satisfies1 T T t=1 E[∥∇f (θ t )∥ 2 ] ≤ 2C 0 f (θ1)-f * T η + ηLσ 2 ϵ + ηLC 2 2 G 2 ϵ + η 2 L 2 C0C 2 1 G 2 ϵ 2 + (1+C1)G 2 d T √ ϵ + η(1+2C1)C1LG 2 d T ϵ , with constants C 0 := 4(1+q 2 ω ) 3 (1-q 2 ω ) 2 G 2 + ϵ, C 1 := β1 1-β1 (1 + C 2 ) + 2qω 1-q 2 ω , C 2 := ωq(1 + 2qω 1-q 2 ω ).Proof. Similar to the proof of Comp-AMS [Li et al., 2022] , we define two virtual iterates θ ′ t and x t .θ ′ t+1 := θ t+1 -η E t+1 √ vt + ϵ x t+1 := θ ′ t+1 -η β 1 1 -β 1 m ′ t + Z t √ vt + ϵ .Then, we derive the recurrence relation for each sequence as follows:θ ′ t+1 = θ t+1 -η E t+1 √ vt + ϵ = θ t -η (1 -β 1 ) t τ =1 β t-τ 1 gτ + (1 -β 1 ) t+1 τ =1 β t+1-τ 1 e τ √ vt + ϵ = θ t -η (1 -β 1 ) t τ =1 β t-τ 1 (g τ + e τ +1 ) + (1 -β)β t 1 e 1 √ vt + ϵ = θ t -η (1 -β 1 ) t τ =1 β t-τ 1 (g τ + e τ + ζ τ ) √ vt + ϵ = θ t -η (1 -β 1 ) t τ =1 β t-τ 1 e τ √ vt + ϵ -η m ′ t √ vt + ϵ -η Z t √ vt + ϵ = θ t -η E t vt-1 + ϵ -η m ′ t √ vt + ϵ + η 1 vt-1 + ϵ - 1 √ vt + ϵ E t -η Z t √ vt + ϵ = θ ′ t -η m ′ t √ vt + ϵ + η 1 vt-1 + ϵ - 1 √ vt + ϵ E t -η Z t √ vt + ϵ = θ ′ t -η m ′ t + Z t √ vt + ϵ + ηD t E t ,where we used the fact that gt + e t+1 = g t + e t + ζ t with quantization noise ζ t , and e 0 = 0 at initialization. Next, for the x t iterates we havex t+1 = θ ′ t+1 -η β 1 1 -β 1 m ′ t + Z t √ vt + ϵ = θ ′ t -η m ′ t + Z t √ vt + ϵ -η β 1 1 -β 1 m ′ t + Z t √ vt + ϵ + ηD t E t = θ ′ t -η β 1 (m ′ t-1 + Z t-1 ) + (1 -β 1 )(g t + ζ t ) + β 2 1 1-β1 (m ′ t-1 + Z t-1 ) + β 1 (g t + ζ t ) √ vt + ϵ + ηD t E t = θ ′ t -η β 1 1 -β 1 m ′ t-1 + Z t-1 √ vt + ϵ -η g t + ζ t √ vt + ϵ + ηD t E t = x t -η g t + ζ t √ vt + ϵ + η β 1 1 -β 1 D t (m ′ t-1 + Z t-1 ) + ηD t E t .Next we apply smoothness (Assumption 3) of the loss function f over the iterates x t . From the gradient Lipschitzness we havef (x t+1 ) ≤ f (x t ) + ⟨∇f (x t ), x t+1 -x t ⟩ + L 2 ∥x t+1 -x t ∥ 2 .Due to unbiasedness of the compressor Q (see Assumption 2), we have E[ζ t |g t , gt , e t , vt ] = 0.Taking expectation, we obtainEQUATIONIn the following, we bound all the four terms highlighted above.Bounding term I. We haveEQUATIONwhere we use Assumption 4, Lemma 6 and the fact that l 2 norm is no larger than l 1 norm.Bounding term II. By the definition of E t and Z t , we know that∥E t ∥ ≤ (1 -β 1 ) t τ =1 β t-τ 1 ∥e t ∥ ≤ 2q ω 1 -q 2 ω G, ∥Z t ∥ ≤ (1 -β 1 ) t τ =1 β t-τ 1 ∥ζ t ∥ ≤ ωq 1 + 2q ω 1 -q 2 ω G.Then we haveEQUATIONwhereC 1 = β1 1-β1 1 + ωq 1 + 2qω 1-q 2 ω + 2qω 1-q 2 ω. The second inequality is because of smoothness of f (θ), and the last inequality is due to Lemma 3, Assumption 4 and the property of norms.Bounding term III. This term can be bounded as follows:EQUATIONwhere C 2 = ωq(1 + 2q 1-q 2 ) and we used Assumption 5 that g t is unbiased with bounded variance σ 2 . Bounding term IV. We haveEQUATIONwhere (a) is due to Young's inequality and (b) is based on Assumption 3. Now integrating ( 5), ( 6), ( 8), ( 11) into (4),I ≤ - η C 0 E[∥∇f (θ t )∥ 2 ] + ηG 2 E[∥D t ∥ 1 ] II ≤ ηC 1 G 2 E[∥D t ∥ 1 ] + η 2 C 2 1 LG 2 √ ϵ E[∥D t ∥ 1 ] III ≤ η 2 L ϵ E[∥∇f (θ t )∥ 2 ] + η 2 L(σ 2 + C 2 2 G 2 ) ϵ + η 2 C 2 1 LG 2 E[∥D t ∥ 2 ] IV ≤ ηρ 2ϵ E[∥∇f (θ t )∥ 2 ] + η 3 L 2 2ρ C 2 1 G 2 ϵ + η 2 LC 1 G 2 √ ϵ E[∥D t ∥ 1 ],and taking the telescoping summation over t = 1, . . . , T , we obtainE[f (x T +1 ) -f (x 1 )] ≤ - η C 0 + η 2 L ϵ + ηρ 2ϵ T t=1 E[∥∇f (θ t )∥ 2 ] + T η 2 L(σ 2 + C 2 2 G 2 ) ϵ + T η 3 L 2 C 2 1 G 2 2ρϵ + η(1 + C 1 )G 2 + η 2 (1 + C 1 )C 1 LG 2 √ ϵ T t=1 E[∥D t ∥ 1 ] + η 2 C 2 1 LG 2 T t=1 E[∥D t ∥ 2 ].Setting η ≤ ϵ 4LC0 and choosing ρ = ϵ 2C0 , we further arrive atE[f (x T +1 ) -f (x 1 )] ≤ - η 2C 0 T t=1 E[∥∇f (θ t )∥ 2 ] + T η 2 L(σ 2 + C 2 2 G 2 ) ϵ + T η 3 L 2 C 0 C 2 1 G 2 ϵ 2 + η(1 + C 1 )G 2 d √ ϵ + η 2 (1 + 2C 1 )C 1 LG 2 d ϵ .where the inequality follows from Lemma 7. Re-arranging terms, we get that1 T T t=1 E[∥∇f (θ t )∥ 2 ] ≤ 2C 0 E[f (x 1 ) -f (x T +1 )] T η + ηL(σ 2 + C 2 2 G 2 ) ϵ + η 2 L 2 C 0 C 2 1 G 2 ϵ 2 + 2C 0 (1 + C 1 )G 2 d T √ ϵ + η(1 + 2C 1 )C 1 LG 2 d T ϵ ≤ 2C 0 f (θ 1 ) -f * T η + ηL(σ 2 + C 2 2 G 2 ) ϵ + η 2 L 2 C 0 C 2 1 G 2 ϵ 2 + 2C 0 (1 + C 1 )G 2 d T √ ϵ + η(1 + 2C 1 )C 1 LG 2 d T ϵ ,where in the last inequality we used x 1 = θ 1 and the lower boundf * ≤ f (θ) for all θ ∈ R d .To get the rate mentioned in the main part, choose η = min{ ϵ 4LC0 , 1 √ T } and continue1 T T t=1 E[∥∇f (θ t )∥ 2 ] ≤ 2C 0 max 1, 4LC 0 ϵ √ T f (θ 1 ) -f * √ T + L(σ 2 + C 2 2 G 2 ) ϵ √ T + L 2 C 0 C 2 1 G 2 ϵ 2 T + 2C 0 (1 + C 1 )G 2 d T √ ϵ + (1 + 2C 1 )C 1 LG 2 d ϵT 3/2 ≤ 2C 0 f (θ 1 ) -f * √ T + L(σ 2 + C 2 2 G 2 ) ϵ √ T + 2C 0 4LC 0 ϵ f (θ 1 ) -f * T + L 2 C 0 C 2 1 G 2 ϵ 2 T + (1 + C 1 )G 2 d T √ ϵ + (1 + 2C 1 )C 1 LG 2 d ϵT 3/2 = 2C 0 f (θ 1 ) -f * √ T + L(σ 2 + C 2 2 G 2 ) ϵ √ T + O G 3 (G + d) T ,where in the second part of the rate we suppressed all the problem and compression dependent constants.As in the non-convex analysis, here we derive the convergence rate with fixed step-size η. The rate shown in the main part can be obtained by plugging the expression of η.Theorem 4. (Convergence rate under PL) Let Assumptions 1, 2, 3, 4, 5 and 6 hold, and q ω < 1.Then, choosing any step-size η ≤ ϵ 4LC0 , MICROADAM (Algorithm 3) satisfiesE[f (θ T +1 )] -f * ≤ 1 -ηµ C0 T (f (θ 1 ) -f * ) + η LC0σ 2 +LC0(C1+C 2 2 )G 2 µϵ + (1+C1)G 2 d+C1G 2 √ ϵ + η 2 3L 2 C0C 2 1 G 2 2µϵ 3/2 + (1+2C1)C1LG 2 d ϵ + LC 2 1 G 2 2ϵ.Proof. We start from descent lemmaEQUATIONWe 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 ′ :I ′ = -ηE ∇f (x t ), g t √ vt + ϵ = -ηE ∇f (x t ), g t vt-1 + ϵ -ηE ∇f (x t ), g t 1 √ vt + ϵ - 1 vt-1 + ϵ = -ηE ∇f (x t ), g t -∇f (x t ) + ∇f (x t ) vt-1 + ϵ -ηE ∇f (x t ), g t 1 √ vt + ϵ - 1 vt-1 + ϵ = -ηE ∇f (x t ), ∇f (x t ) vt-1 + ϵ -ηE ∇f (x t ), g t -∇f (x t ) vt-1 + ϵ -ηE ∇f (x t ), g t 1 √ vt + ϵ - 1 vt-1 + ϵ .We further expand and bound this equation as follows:I ′ ≤ - η C 0 E ∥∇f (x t )∥ 2 -ηE ∇f (x t ) -∇f (θ t ) + ∇f (θ t ), 1 vt-1 + ϵ (∇f (θ t ) -∇f (x t )) -ηE ∇f (x t ) -∇f (θ t ) + ∇f (θ t ), 1 √ vt + ϵ - 1 vt-1 + ϵ g t = - η C 0 E ∥∇f (x t )∥ 2 -ηE ∇f (x t ) -∇f (θ t ), 1 vt-1 + ϵ (∇f (θ t ) -∇f (x t )) -ηE ∇f (θ t ), 1 vt-1 + ϵ (∇f (θ t ) -∇f (x t )) -ηE ∇f (x t ) -∇f (θ t ), 1 √ vt + ϵ - 1 vt-1 + ϵ g t -ηE ∇f (θ t ), 1 √ vt + ϵ - 1 vt-1 + ϵ g t ≤ - η C 0 E ∥∇f (x t )∥ 2 + η √ ϵ E ∥∇f (x t ) -f (θ t )∥ 2 -ηE ∇f (θ t ), 1 vt-1 + ϵ (∇f (θ t ) -∇f (x t )) + ηE [⟨∇f (x t ) -∇f (θ t ), D t g t ⟩] + ηE [⟨∇f (θ t ), D t g t ⟩] .Next, we use the Cauchy-Schwartz inequality to bound inner products above, L-smoothness inequality to bound ∥∇f (x t ) -∇f (θ t )∥ ≤ L∥x t -θ t ∥ ≤ ηLC1G √ ϵ , and the inequality -∥a∥ 2 ≤ -1 2 ∥b∥ 2 + ∥a -b∥ 2 for the first term:I ′ ≤ - η C 0 E ∥∇f (x t )∥ 2 + η √ ϵ E ∥∇f (x t ) -f (θ t )∥ 2 + ηG √ ϵ E [∥∇f (θ t ) -∇f (x t )∥] + ηGE [∥∇f (x t ) -∇f (θ t )∥∥D t ∥] + ηG 2 E [∥D t ∥] ≤ - η 2C 0 E ∥∇f (x t )∥ 2 - η 2C 0 E ∥∇f (x t )∥ 2 + η √ ϵ η 2 L 2 C 2 1 G 2 ϵ + η 2 LC 1 G 2 ϵ + ηG ηLC 1 G √ ϵ E [∥D t ∥ 1 ] + ηG 2 E [∥D t ∥ 1 ] ≤ - η 2C 0 E ∥∇f (x t )∥ 2 - η 4C 0 E[∥∇f (θ t )∥ 2 ] + η 2C 0 E[∥∇f (x t ) -∇f (θ t )∥ 2 ] + η 3 L 2 C 2 1 G 2 ϵ 3/2 + η 2 LC 1 G 2 ϵ + η 2 LC 1 G 2 √ ϵ E [∥D t ∥ 1 ] + ηG 2 E [∥D t ∥ 1 ] ≤ - η 2C 0 E ∥∇f (x t )∥ 2 - η 4C 0 E[∥∇f (θ t )∥ 2 ] + 3η 3 L 2 C 2 1 G 2 2ϵ 3/2 + η 2 LC 1 G 2 ϵ + η 2 LC 1 G 2 √ ϵ E [∥D t ∥ 1 ] + ηG 2 E [∥D t ∥ 1 ] .Plugging the obtained bound for I ′ with previously obtained bounds for II andIII II ≤ ηC 1 G 2 E[∥D t ∥ 1 ] + η 2 C 2 1 LG 2 √ ϵ E[∥D t ∥ 1 ] III ≤ η 2 L ϵ E[∥∇f (θ t )∥ 2 ] + η 2 L(σ 2 + C 2 2 G 2 ) ϵ + η 2 C 2 1 LG 2 E[∥D t ∥ 2 ]into (12) and using the step-size bound η ≤ ϵ 4LC0 we getE[f (x t+1 )] -E[f (x t )] ≤ - η 2C 0 E ∥∇f (x t )∥ 2 - η 4C 0 E[∥∇f (θ t )∥ 2 ] + 3η 3 L 2 C 2 1 G 2 2ϵ 3/2 + η 2 LC 1 G 2 ϵ + η 2 LC 1 G 2 √ ϵ E [∥D t ∥ 1 ] + ηC 1 G 2 E[∥D t ∥ 1 ] + η 2 C 2 1 LG 2 √ ϵ E[∥D t ∥ 1 ] + ηG 2 E [∥D t ∥ 1 ] + η 2 L ϵ E[∥∇f (θ t )∥ 2 ] + η 2 L(σ 2 + C 2 2 G 2 ) ϵ + η 2 C 2 1 LG 2 E[∥D t ∥ 2 ] ≤ - η 2C 0 E ∥∇f (x t )∥ 2 + η 2 Lσ 2 ϵ + η 2 L(C 1 + C 2 2 )G 2 ϵ + η(1 + C 1 )G 2 E[∥D t ∥ 1 ] + 3η 3 L 2 C 2 1 G 2 2ϵ 3/2 + η 2 (1 + C 1 )C 1 LG 2 √ ϵ E[∥D t ∥ 1 ] + η 2 C 2 1 LG 2 E[∥D t ∥ 2 ] ≤ - ηµ C 0 (E[f (x t )] -f * ) + η 2 Lσ 2 ϵ + η 2 L(C 1 + C 2 2 )G 2 ϵ + η(1 + C 1 )G 2 E[∥D t ∥ 1 ] + η 2 (1 + C 1 )C 1 LG 2 √ ϵ E[∥D t ∥ 1 ] + η 2 C 2 1 LG 2 E[∥D t ∥ 2 ] + 3η 3 L 2 C 2 1 G 2 2ϵ 3/2, where in the last inequality we applied PL condition from Assumption 6. After some reshuffling of the terms, we obtain the following recursion:E[f (x t+1 )] -f * ≤ 1 - ηµ C 0 (E[f (x t )] -f * ) + η 2 Lσ 2 ϵ + η 2 L(C 1 + C 2 2 )G 2 ϵ + 3η 3 L 2 C 2 1 G 2 2ϵ 3/2 + η(1 + C 1 )G 2 E[∥D t ∥ 1 ] + η 2 (1 + C 1 )C 1 LG 2 √ ϵ E[∥D t ∥ 1 ] + η 2 C 2 1 LG 2 E[∥D t ∥ 2 ].Notice that η ≤ ϵ 4LC0 ≤ C0 4µ , so that the coefficient 1 -ηµ C0 ∈ (0, 1). Unrolling the recursion, we arriveEQUATIONFor the second sum above we upper bound it by its infinite sum asT t=1 1 - ηµ C 0 t ≤ ∞ t=0 1 - ηµ C 0 t = C 0 ηµ .For the other three sums we bound 1 -ηµ C0 ≤ 1 and apply the bounds in Lemma 7:T t=1 1 - ηµ C 0 t E[∥D t ∥ 1 ] ≤ T t=1 E[∥D t ∥ 1 ] ≤ d √ ϵ , T t=1 1 - ηµ C 0 t E[∥D t ∥ 2 ] ≤ T t=1 E[∥D t ∥ 2 ] ≤ d ϵ .Plugging all this bounds into (13) and noticing that x 1 = θ 1 , we finally getE[f (x T +1 )] -f * ≤ 1 - ηµ C 0 T (f (θ 1 ) -f * ) + C 0 ηµ η 2 Lσ 2 ϵ + η 2 L(C 1 + C 2 2 )G 2 ϵ + 3η 3 L 2 C 2 1 G 2 2ϵ 3/2 + η(1 + C 1 )G 2 d √ ϵ + η 2 (1 + C 1 )C 1 LG 2 d ϵ + η 2 C 2 1 LG 2 d ϵ ≤ 1 - ηµ C 0 T (f (θ 1 ) -f * ) + η LC 0 σ 2 µϵ + LC 0 (C 1 + C 2 2 )G 2 µϵ + (1 + C 1 )G 2 d √ ϵ + η 2 3L 2 C 0 C 2 1 G 2 2µϵ 3/2 + (1 + C 1 )C 1 LG 2 d ϵ + C 2 1 LG 2 d ϵ .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 θ t of the algorithm, we apply L-smoothness to bound the functional difference:|f (x t ) -f (θ t )| ≤ |⟨∇f (θ t ), x t -θ t )⟩| + L 2 ∥x t -θ t ∥ 2 ≤ ηC 1 G 2 √ ϵ + η 2 LC 2 1 G 2 2ϵ ,which impliesE[f (θ T +1 )] -f * ≤ 1 - ηµ C 0 T (f (θ 1 ) -f * ) + η LC 0 σ 2 µϵ + LC 0 (C 1 + C 2 2 )G 2 µϵ + (1 + C 1 )G 2 d √ ϵ + C 1 G 2 √ ϵ + η 2 3L 2 C 0 C 2 1 G 2 2µϵ 3/2 + (1 + C 1 )C 1 LG 2 d ϵ + C 2 1 LG 2 d ϵ + LC 2 1 G 2 2ϵand completes the proof.To get the rate mentioned in the main part of the paper, we plug in the expression η = min{ ϵ 4LC0 , 2C0 log T µT } and collect higher order terms.E[f (θ T +1 )] -f * ≤ max 1 T 2 , 1 - ϵµ 4L T (f (θ 1 ) -f * ) + log T T 2C 0 µ LC 0 σ 2 µϵ + LC 0 (C 1 + C 2 2 )G 2 µϵ + (1 + C 1 )G 2 d √ ϵ + C 1 G 2 √ ϵ + log 2 T T 2 4C 2 0 µ 2 3L 2 C 0 C 2 1 G 2 2µϵ 3/2 + (1 + C 1 )C 1 LG 2 d ϵ + C 2 1 LG 2 d ϵ + LC 2 1 G 2 2ϵ ≤ 2 log T T LC 2 0 µ σ 2 + (C 1 + C 2 2 )G 2 µϵ + C 0 (1 + C 1 )(1 + d)G 2 µ √ ϵ + O G 4 (G + d) T 2 .Algorithm 4 MICROADAMW (MICROADAM with Weight Decay) 1: Input: parameters β 1 , β 2 ∈ (0, 1), ϵ > 0, step-size η > 0, θ 1 ∈ R d , e 1 = m 0 = v 0 = v0 = 0 d 2: for t = {1, 2, ..., T } do 3: g t = ∇ θ f (θ t ) ⋄ Computem t = β 1 m t-1 + (1 -β 1 )g t ⋄ Update first-order gradient moment 7: v t = β 2 v t-1 + (1 -β 2 )g 2 t ⋄Update second-order gradient moment 8:θ t+1 = (1 -η t λ)θ t -η mt √ vt+ϵ⋄ Update the model parameters with weight decay 9: end for Lemma 8. Under Assumptions 1-5, for all iterates of Algorithm 4 we haveE ∥m t -∇f (θ t )∥ 2 ≤ 1 - β 1 2 E ∥m t-1 -∇f (θ t-1 )∥ 2 + 2 β 1 L 2 E ∥θ t -θ t-1 ∥ 2 + β 1 E∥∇f (θ t ) -gt ∥ 2 .Proof. We start our proof fromE ∥m t -∇f (θ t )∥ 2 =E ∥(1 -β 1 )m t-1 + β 1 gt -∇f (θ t )∥ 2 =E [∥(1 -β 1 )m t-1 + (1 -β 1 )∇f (θ t-1 ) -(1 -β 1 )∇f (θ t-1 ) + β 1 gt -∇f (θ t )∥ 2 =E [∥(1 -β 1 )m t-1 + (1 -β 1 )∇f (θ t-1 ) -(1 -β 1 )∇f (θ t ) -(1 -β 1 )∇f (θ t-1 ) + β 1 gt -β 1 ∇f (θ t )∥ 2 .Using Jensen's inequality we haveE ∥m t -∇f (θ t )∥ 2 ≤(1 -β 1 )E ∥m t-1 -∇f (θ t-1 ) + ∇f (θ t-1 ) -∇f (θ t )∥ 2 + β 1 E ∥g t -∇f (θ t )∥ 2 .Using Young's inequality we haveE ∥m t -∇f (θ t )∥ 2 ≤(1 -β 1 )(1 + b)E ∥m t-1 -∇f (θ t-1 )∥ 2 + (1 -β 1 ) 1 + 1 b E∥∇f (θ t-1 ) -f (θ t )∥ 2 + β 1 E ∥∇f (θ t ) -gt ∥ 2 . Setting b = β1 2 we have (1 -β 1 ) 1 + β1 2 ≤ 1 -β1 2 and (1 -β 1 ) 1 + 2 β1 ≤ 2 β1 : E ∥m t -∇f (θ t )∥ 2 ≤ 1 - β 1 2 E ∥m t-1 -∇f (θ t-1 )∥ 2 + 2 β 1 E∥∇f (θ t-1 ) -f (θ t )∥ 2 + β 1 E ∥∇f (θ t ) -gt ∥ 2 .Combining this bound with L-smoothness we obtainE ∥m t -∇f (θ t )∥ 2 ≤ 1 - β 1 2 E ∥m t-1 -∇f (θ t-1 )∥ 2 + 2 β 1 L 2 E∥θ t-1 -θ t ∥ 2 + β 1 E ∥∇f (θ t ) -gt ∥ 2 .Lemma 9. Under Assumptions 1-5, for all iterates of Algorithm 4 we haveE ∥∇f (θ t ) -gt ∥ 2 ≤ 9E ∥e t ∥ 2 + 6G 2 + 3σ 2 .Proof. We start fromE ∥∇f (θ t ) -gt ∥ 2 =E ∥C(e t + g t ) -∇f (θ t )∥ 2 =E ∥C(e t + g t ) -(e t + g t ) + (e t + g t ) -∇f (θ t )∥ 2 ≤3E ∥C(e t + g t ) -(e t + g t )∥ 2 + 3E ∥e t ∥ 2 + 3E ∥∇f (θ t ) -g t ∥ 2 .Using definition of contractive compressor we haveE ∥∇f (θ t ) -gt ∥ 2 ≤ 3q 2 E ∥e t + g t ∥ 2 + 3E ∥e t ∥ 2 + 3E ∥∇f (θ t ) -g t ∥ 2 .Using Young's inequality we haveE ∥∇f (θ t ) -gt ∥ 2 ≤ 6q 2 E ∥e t ∥ 2 + 6q 2 E ∥g t ∥ 2 + 3E ∥e t ∥ 2 + 3E ∥∇f (θ t ) -g t ∥ 2 ≤ 9E ∥e t ∥ 2 + 6G 2 + 3σ 2 .Lemma 10. Under Assumptions 1-5, for all iterates of Algorithm 4 we haveE ∥e t ∥ 2 ≤ (1 + ω) 2 q 2 (1 -(1 + ω)q) 2 G 2Proof. Using definition of contractive compressor we haveE ∥e t+1 ∥ 2 = E ∥Q (e t + g t -C(e t + g t ))∥ 2 ≤ (1 + ω) 2 q 2 E ∥e t + g t -C(e t + g t )∥ 2 .Using Young's inequality we haveE ∥e t+1 ∥ 2 ≤ (1 + ω) 2 q 2 (1 + a) E ∥e t ∥ 2 + (1 + ω) 2 q 2 1 + 1 a E ∥g t ∥ 2 .We need to satisfy the following condition:(1 + ω) 2 q 2 (1 + a) ≤ (1 + ω)q.It holds for 0 ≤ ω, 0 ≤ q < 1, (1 + ω)q < 1 and a = 1 (1+ω)q -1. Using this parameters we haveE ∥e t+1 ∥ 2 ≤ (1 + ω)qE ∥e t ∥ 2 + (1 + ω) 2 q 2 1 1 -(1 + ω)q ∥g t ∥ 2 ≤ (1 + ω)qE ∥e t ∥ 2 + (1 + ω) 2 q 2 1 -(1 + ω)q E ∥g t ∥ 2 .Unrolling this recursion allows us to obtainE ∥e t+1 ∥ 2 ≤ ((1 + ω)q) t E ∥e 1 ∥ 2 + T t=1 ((1 + ω)q) t (1 + ω) 2 q 2 1 -(1 + ω)q G 2 ≤ ((1 + ω)q) t E ∥e 1 ∥ 2 + 1 1 -(1 + ω)q (1 + ω) 2 q 2 1 -(1 + ω)q G 2 ≤ (1 + ω) 2 q 2 (1 -(1 + ω)q) 2 G 2 ,last inequality holds because e 1 = 0.Lemma 11. Under Assumptions 1-5, for all iterates of Algorithm 4 we haveE ∥g t ∥ 2 ≤ 4(1 + q 2 ) 1 + (1 + ω) 2 q 2 (1 -(1 + ω)q) 2 G 2 .Proof.E ∥g t ∥ 2 = E ∥C (g t + e t )∥ 2 = E ∥C (g t + e t ) -(g t + e t ) + (g t + e t )∥ 2 ≤ 2E ∥C (g t + e t ) -(g t + e t )∥ 2 + 2E ∥g t + e t ∥ 2 ≤ 2(1 + q 2 )E ∥g t + e t ∥ 2 .Using Lemma 10 we obtainE ∥g t ∥ 2 ≤ 4(1 + q 2 ) 1 + (1 + ω) 2 q 2 (1 -(1 + ω)q) 2 G 2 .Lemma 12 (From paper: Zhou et al. [2024b] ). Let us consider the update rule:θ t+1 = (1 -ηλ)θ t -η t m t √ v t + ε .For brevity, we denote v t = √ v t + ε. Also we defineu t := m t + λθ t ⊗ v t ,where ⊗ denotes element-wise product. Moreover, we also define f (θ t ) as follows:f (θ t ) = f (θ t ) + λ t ∥θ t ∥ 2 vt where λ t = λ 2 t i=11-q 2 i for t > 0, λ 0 = 0 with 0 < q < 1 and∥θ t ∥ vt = ⟨θ t , vt ⊗ θ t ⟩. Also let c 1 ≤ ∥ v t ∥ ∞ < c 2 , then iterates of Algorithm 4 satisfy f (θ t ) ≤ f (θ t-1 ) + η t 2c 1 ∥∇f (θ t-1 ) -m t-1 ∥ 2 - η t 2c 2 ∇ f (θ t-1 ) 2 - η t 4c 2 ∥u t-1 ∥ 2 .Lemma 13 (From paper: Zhou et al. [2024b] ). Assume that c s,∞ ≤ ∥g t ∥ ∞ ≤ c ∞ , then we have∥m t ∥ ∞ ≤ c ∞ , ∥v t + ϵ∥ ∞ c 2 ∞ + ϵ, (v t + ϵ) p (v t+1 + ϵ) p ∞ ∈ [1 -µ, 1 + µ] (∀p ∈ (0, 1)) , where µ = β2c 2 ∞ c 2 s,∞ +ϵ . Theorem 5. Let Assumptions 1 to 5 hold. Define Ψ t = f (θ t ) + ηt 2c1β1 E ∥m t -∇f (θ t )∥ 2 . With η t = η ≤ β1c1 2L c1 2c2 , Algorithm 4 satisfies 1 T T t=1 ∇ f (θ t-1 ) 2 ≤ 2c 2 ηT Ψ 0 + c 2 c 1 9 1 + (1 + ω) 2 q 2 (1 -(1 + ω)q) 2 G 2 + 6G 2 + 3σ 2 .Proof. We start from main lemma and lemma for momentum, summing inequalities together we obtainf (θ t ) + V E ∥m t -∇f (θ t )∥ 2 ≤ f (θ t-1 ) + η t 2c 1 ∥∇f (θ t-1 ) -m t-1 ∥ 2 - η t 2c 2 ∇ f (θ t-1 ) 2 - η t 4c 2 ∥u t-1 ∥ 2 + V 1 - β 1 2 E ∥m t-1 -∇f (θ t-1 )∥ 2 + V 2 β 1 L 2 E∥θ t-1 -θ t ∥ 2 + V β 1 E ∥∇f (θ t ) -gt ∥ 2 .Using previous lemmas we havef (θ t ) + V E ∥m t -∇f (θ t )∥ 2 ≤ f (θ t-1 ) + η t 2c 1 ∥∇f (θ t-1 ) -m t-1 ∥ 2 - η t 2c 2 ∇ f (θ t-1 ) 2 - η t 4c 2 ∥u t-1 ∥ 2 + V 1 - β 1 2 E ∥m t-1 -∇f (θ t-1 )∥ 2 + V 2 β 1 L 2 E∥θ t-1 -θ t ∥ 2 + V β 1 9 1 + (1 + ω) 2 q 2 (1 -(1 + ω)q) 2 G 2 + 6G 2 + 3σ 2 . Using θ t -θ t-1 = -η t ut-1 vt-1 we have f (θ t ) + V E ∥m t -∇f (θ t )∥ 2 ≤ f (θ t-1 ) - η t 4c 2 - 2V L 2 η 2 t β 1 c 2 1 ∥u t-1 ∥ 2 - η t 2c 2 ∇ f (θ t-1 ) 2 + V 1 - β 1 2 + η t 2c 1 E ∥m t-1 -∇f (θ t-1 )∥ 2 + V β 1 9 1 + (1 + ω) 2 q 2 (1 -(1 + ω)q) 2 G 2 + 6G 2 + 3σ 2 . Using V = ηt 2c1β1 and Ψ t = f (θ t ) + ηt 2c1β1 E ∥m t -∇f (θ t )∥ 2 we have Ψ t ≤ Ψ t-1 - η t 4c 2 1 - 4c 2 L 2 η 2 t β 2 1 c 3 1 ∥u t-1 ∥ 2 - η t 2c 2 ∇ f (θ t-1 ) 2 + η t 2c 1 9 1 + (1 + ω) 2 q 2 (1 -(1 + ω)q) 2 G 2 + 6G 2 + 3σ 2 . Using η t = η ≤ β1c1 2L c1 2c2 we have Ψ t ≤ Ψ t-1 - η t 2c 2 ∇ f (θ t-1 ) 2 + η 2c 1 9 1 + (1 + ω) 2 q 2 (1 -(1 + ω)q) 2 G 2 + 6G 2 + 3σ 2 .Summing from t = 1 to T we have1 T T t=1 ∇ f (θ t-1 ) 2 ≤ 2c 2 ηT Ψ 0 + c 2 c 1 9 1 + (1 + ω) 2 q 2 (1 -(1 + ω)q) 2 G 2 + 6G 2 + 3σ 2 .Discussion. This result is similar to the one from Zhou et al. [2024b] in the non-convex case, where the decay rate for the first term is O 1 T and the second term is a non-vanishing O β 1 c2 c1 σ 2 . In our result, the non-vanishing term is proportional to O c2 c1 σ 2 + G 2 . A key difference is that our term is not proportional to β 1 . It is important to note that β 1 typically takes a value close to 1 in practical applications, meaning its influence on the bound is minimal. Therefore, even though our result does not directly involve β 1 , the impact on the overall bound is not significantly different. Moreover, our bound includes an additional term proportional to G 2 , which represents the gradient norm squared. This makes the bound slightly worse compared to the result in Zhou et al. [2024b] . However, this degradation is only by a constant factor, which means that while the theoretical bound may be worse, the practical implications are often negligible.In summary, our result aligns closely with previous findings, with differences primarily in the constant factors and the presence of G 2 . Despite these differences, the practical performance remains largely unaffected, ensuring that the bound remains robust and applicable in a variety of scenarios.The GaLore low-rank updates introduced by [Zhao et al., 2024] enable the compression of optimizer states by performing learning updates on a lower-dimensional subspace. In this approach, the optimizer receives gradients projected on a defined learning subspace. Theoretical convergence guarantees are provided under a "stable rank" assumption, where learning subspace is fixed during training. However, in practice, convergence is attained by occasionally updating the learning subspace and allowing full space learning to better align with the gradient trajectory during training.Here, it is useful to draw an analogy with the TopK method, as the occasional updates of the learning subspace resembles working with a fixed mask for many steps. Using a fixed mask would result in discarding the same coordinates of the gradient at each step. Similarly, in the case of low-rank updates, components orthogonal to the same learning subspace are discarded at each step.The systematic nature of the information discarded by compression carries significant implications for error feedback behavior. Over multiple steps, the error accumulates gradient components belonging to the orthogonal space of the same learning subspace. Consequently, by linearity, the error itself resides in the orthogonal space of this learning subspace. As a result, when the error is passed to the accumulator, its projection onto the learning space is effectively disregarded until it is potentially utilized at the specific step when the learning subspace is updated. Therefore, the behavior of error feedback in the case of low-rank updates is non-standard: it accumulates gradient components over numerous steps before unloading them all at once. For a better understanding, we derive analytical evidence for the described behaviour by induction. Let L be fixed learning subspace and assume that e t-1 ∈ L ⊥ . Then, gradient passed to the optimizer at step t is: C GaLore (a t ) = proj L (a t ) = proj L (e t-1 + g t ) = proj L (g t ) where error feedback is discarded. Thus, e t = a t -C GaLore (a t ) = e t-1 + g t -proj L (g t ) = e t-1 + proj L ⊥ (g t ) ∈ L ⊥ which completes the induction.Assume now that learning subspace is updated every T steps, and denote L t the learning subspace at step t. Then, a similar induction leads to:e t = t i=1 t • j=i proj L ⊥ j (g i ) = t i=1 ⌊ t T ⌋ • j=⌊ i T ⌋ proj L ⊥ jT (g i ) a kT = proj L kT (g kT + e kT -1 ) = proj L kT (g kT ) + kT -1 i=1 proj L kT • ( t • j=i proj L ⊥ j )(g i )Such behaviour of the error feedback mechanism results in the dominance of the error norm over the gradient norm. Before learning subspace updates, the error is the sum over past gradient components that belong to the orthogonal of the current learning subspaces. Since these components represent descent directions that were not used, they are not expected to compensate each other on average. Consequently, between learning subspace updates, the error norm is expected to grow linearly.Figure 8 provides evidence of such linear growth of the error norm during fine-tuning of RoBERTabase model on GLUE/MNLI task.It implies that known analysis techniques [Alistarh et al., 2018 , Karimireddy et al., 2019] of convergence for the error feedback mechanism do not apply to GaLore. Indeed, such proofs rely on the assumption that the compression operator is contractive, as it allows the error to be bounded. Given a fixed vector, low-rank compression based on its singular value decomposition is a contraction operator. However, in our case, the compression is based on a previously-computed singular value decomposition and therefore may not be a contraction operator for newly computed gradients. The extreme case being when the gradient is orthogonal to the learning subspace, in which case the compression operator returns the null vector. Figure 8 shows that during training the error norm is not on the same order of magnitude of the gradient norm.The dominance of the error over the gradient also has effects on space exploration, as the learning subspaces are computed from the singular value decomposition of the accumulator (i.e. the sum of the gradient and the error). Since the main components of the accumulator belong to the orthogonal of current learning subspaces, successive learning subspaces will tend to be orthogonal to each other. This allows errors to be effectively passed to the optimizer, but all at once which can introduce irregularities in the learning trajectory. However, it also implies that learning is performed on a learning subspace that is suboptimal in terms of the direction of the gradient, but this may help convergence by enforcing space exploration. See Figure 9 for examples of how induced orthogonality of successive learning subspaces affects the learning trajectory.Figure 9 : Optimization trajectory for Adam, GaLore-Adam and GaLore-Adam-EF for ill-conditioned function f (x, y) = cos( 5π 4 x) + sin( 7π 4 y) starting from (x 0 , y 0 ) = (-1 4 , 1 4 ) (on first row) and for Rosenbrock function starting from (x 0 , y 0 ) = (-1 2 , 1) (on second row).• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [No] Justification: We do not release any models.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification:We work with open source models that are publicly available and we cited them properly.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [Yes] Justification: We provide code and instructions on how to run the code in order to reproduce our results. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.GaLore[Zhao et al., 2024] does state convergence guarantees for a variant of the algorithm with fixed projections, but this is under a strong "stable rank" assumption, which may not hold in practice. 
 
 Appendix:  


Decision: **Summary:**

MicroAdam is a new optimizer that aims to reduce memory usage compared to Adam while maintaining performance. It uses gradient sparsification, quantization, and error feedback to compress optimizer states. Theoretical convergence guarantees are provided. Experiments show competitive performance on language model fine-tuning tasks with reduced memory usage.

**Strengths:**

- Novel approach to reducing optimizer memory footprint
- Provides theoretical convergence guarantees
- Shows good empirical performance on fine-tuning tasks
- Efficient GPU implementation

**Weaknesses:**

- Initially limited to fine-tuning experiments (though pretraining results were added in rebuttal)
- Theoretical analysis based on AMSGrad rather than Adam
- Complex implementation details for compressors
- Unclear why performance sometimes exceeds uncompressed Adam

**Key Discussion Points:**

- *Pretraining results:* Authors added ImageNet pretraining experiments for ResNet models in rebuttal to address this limitation
- *Theoretical analysis:* Authors clarified choice of AMSGrad as basis due to known issues with Adam's convergence proof
- *Implementation complexity:* Authors acknowledged this but noted existing work/tools can be leveraged
- *Memory savings:* Clarified that main savings are in optimizer state size, which is significantly reduced
- *Performance exceeding Adam:* Authors hypothesized this may be due to regularization effects of compression
- *Applicability beyond fine-tuning:* Additional experiments helped demonstrate broader applicability

Overall, reviewers generally found the work to be technically sound with good theoretical and empirical contributions, though some concerns remained about novelty and broader applicability. The authors' responses and additional experiments addressed many of the initial concerns. Most reviewers maintained or increased their scores to weak accept or higher after the rebuttal.

Reviews: review: paper_summary: paper_summary: The script proposed a memory-efficient method with a theoretical guarantee.
strength_weakness: strength_weakness: See below.
questions: questions: See below.
limitations: limitations: **Major concern 1**: The contribution of the paper is rather weak. After reading the paper,  I still did not find what is new in this paper.  For instance:

**1-1: It is unclear what is new in the compression design of Algorithm 1.** I have a hard time figuring out the contribution of Algorithm 1.  Please summarize the significance and novelty of the proposed compression procedure in Algorithm 1. Please discuss how it is different from the existing ones and why do we need this new one. Please discuss the design principles and ideas. 

 **1-2 It is unclear whether the proposed method is truly effective.**  There are much much more experiments are needed. Currently, only 7B SFT is provided, which is far from convincing. SFT is not enough to support the effectiveness of a new optimizer.  It seems unclear how the method would perform on LLM pre-training or non-language tasks. 

 **1-3:  It is unclear what is new in the theory.**  Under the assumption of unbiased compression +  bounded variance + bounded gradient, this type of Algorithm 4 has already been extensively studied by the distributed community and so. Please highlight what is new and what is nontrivial, if there is any.





**Major concern 2: The presentation is quite poor.** For instance:

1. There is no explanation on the principle or idea behind the design of Algorithm 1. 

2. It is good to see simple example like Figure 1. But unfortunately, nearly nothing is explained here. We can see EF helps  convergence, but we do not understand why. It seems just a toy example without any explanation on the insight.

3. The current logic flow of the script is confusing. It would be much easier to read if the authors:

   (i) introduce Algorithm 4 first (generic form of MicroAdam), 

   (ii) provide the theoretical guarantee, 

   (iii)then introduce Algorithm 1( a  specific form that you chose) , 

   (iv) discuss the implementation details,

   (v) then show  the experiments.
paper_summary: paper_summary: The paper presents a new optimizer that approximates Adam but has a lower memory footprint. Adam stores for each parameter two additional values --- the exponential moving averages of the gradients and the gradient squared. The current optimizer saves space by using the fact that for most parameters, the gradient is quite small at each step, so could be ignored. The optimizer stores the most important gradients at each step, for several past steps. It also stores a low-resolution version of the gradient, which is added to the next gradient (the reason for this was not clear to me). The exponential moving averages for the gradients and gradients-squared are computed from the most important gradients stored at each step and are used in the Adam update. 

The paper proves the convergence of the optimizer under some reasonable assumptions. The authors show that the method can be implemented efficiently on a GPU, and the performance is similar to Adam. The authors performed several experiments to show that their optimizer saves about 10-20% memory compared to Adam, while achieving similar accuracy.
strength_weakness: strength_weakness: The presentation of the paper was not very clear to me. For example, k is not mentioned as one of the inputs in Algorithm 1. $\delta_1, \Delta_1$ are initialized to vectors in R^d, but then they seem to be scalars.  I also did not understand what the error correction was for.
questions: questions: In your algorithm, why did you not use a_t or the current gradient in steps 9 and 10, since this is already present.

It looks like most of the memory is being used for other tasks, so reducing the memory by a factor of 4 resulted in a small overall reduction in memory.
limitations: limitations: N/A
paper_summary: paper_summary: The paper introduces MicroAdam, a novel optimizer designed to improve memory efficiency while maintaining competitive performance with established optimizers such as Adam. The authors provide theoretical analyses and experimental results to demonstrate the benefits of MicroAdam in various settings.
strength_weakness: strength_weakness: - **Motivation in Figure 1**: Figure 1 does not effectively provide additional motivation for MicroAdam, as it represents a 2-dimensional non-stochastic problem with only choosing k=1 for top-k compression. This toy example can be solved by many heuristic methods or parameter adjustments, which diminishes the unique motivation for MicroAdam.

- **Convergence Speed**: The convergence speed of MicroAdam is asymptotically similar to AMSGrad rather than Adam itself. Given that the paper aims to improve Adam's memory footprint, it is unclear why the convergence speed matches that of AMSGrad, which theoretically has no significant speed difference from Adam's with respect to T.

- **Memory Consumption Analysis**: Sections 3.2 and C provide theoretical and simulated memory consumption analysis. However, the practical memory consumption during actual LLM training does not necessarily correlate strongly with the theoretical optimizer states' memory usage. For example, a small LLAMA model with sufficiently large batch size and token length can peak at over 70GB of memory usage (mainly due to activation memory), whereas the optimizer states' memory footprint might be under 100MB. Therefore, the paper should report more about real peak memory under appropriate settings rather than theoretical or simulated values, as the differences in optimizer states memory can be mitigated under certain settings and PyTorch's memory management mechanism.

- **Top-K Operator Memory Efficiency**: Despite the block-wise operation of the Top-K operator, if the gradients are in bf16 format, the memory overhead for 50% sparsity remains the same (due to the need to store 16-bit indices) or could even be higher when considering other memory overheads in the algorithm.

- **Lack of Training from Scratch Results**: A significant issue is the near absence of results for training from scratch, which is crucial for evaluating MicroAdam's effectiveness. For MicroAdam, training from scratch is more important. Table 1 shows that during fine-tuning, the training loss does not strictly correlate with the final metrics. For instance, GaLore has a much worse loss but maintains metrics comparable to those of other methods. This is why LLM fine-tuning metrics are often taken from the best of three epochs, not necessarily the last epoch, even if the training loss is the smallest in the final epoch. Additionally, Figure 4 indicates that MicroAdam's training efficiency is lower than baseline Adam, as MicroAdam's third epoch loss is nearly identical to Adam's second epoch loss.

- **Computational Efficiency**: MicroAdam appears to be less computationally efficient. Are there ways to improve its computational efficiency?
questions: questions: See Weaknesses
limitations: limitations: See Weaknesses
paper_summary: paper_summary: This paper proposes a memory-efficient Adam-based optimizer called MicroAdam. The key idea is to compress gradients using Top-K sparsification before passing them to the optimizer state, along with an error feedback vector that is itself compressed via quantization. For general smooth non-convex (and PL) functions, MicroAdam’s convergence rates are provided with constants depending on the choice of compressors. Experiments on fine-tuning tasks with BERT and LLaMA suggest that MicroAdam performs competitively with AdamW while using less memory.
strength_weakness: strength_weakness: - Both theoretical and practical performance, as well as memory savings, depend heavily on the choice of compressors. These compressors can be complex to implement and require specific techniques to avoid memory overhead on GPUs. The difficulty and complex details of implementation might hinder impact as it is not easy to incorporate in practice. 
- The approach is limited to fine-tuning tasks, and its effectiveness for LLM pre-training or non-language tasks remains unclear.
questions: questions: The condition $(1+\omega)q \leq 1$ is necessary for theoretical convergence. Do the compressors used in the experiment section satisfy this requirement? Generally, how do the authors suggest finding appropriate compressors for unfamiliar tasks?
limitations: limitations: Limitations are well-addressed
paper_summary: paper_summary: This paper proposes a new optimizer based on Adam, called MicroAdam, which reduces memory footprints, via sparisification, quantization, and error feedback. Theoretical convergence guarantees are provided. Empirical results show good performance on several fine-tuning tasks
strength_weakness: strength_weakness: 1. The experiments are limitted to finetuning tasks. Note that GaLore also show good performance in pretraining tasks.
2. The theoretical analysis is actually based on AMSGrad instead of Adam.

---------------

My concerns are well addressed according to the author's feedback.
questions: questions: 1. In many cases, MicroAdam shows even better loss or accuracy than the original Adam, which doesn't really make sense since the compression methods should be lossy which typically incurs worse loss or accuracy. Is there any explanation of this phenomenon?
limitations: limitations: The limitations are well discussed.

Comments: title: Response to authors
Author Response
I will keep my score
Start discussions with authors, please
comment: Thank you for addressing my concerns and conducting additional vision pretraining experiments. The accuracy of AdamW seems a bit lower than usual. I don't think one can conclude that MicroAdam achieves better accuracy than AdamW in this case without further fine-tuning. However, I believe it sufficiently demonstrates MicroAdam's behavior in pretraining, so I will raise my score.
Dear Reviewer,

Thank you for your response. We provide some brief clarifications below: 

> I still think it is rather weak to validate the efficacy of a new optimizer by SFT on Llama. 

Unfortunately the reviewer may have completely missed our _pre-training results_, presented in the rebuttal. Specifically, in the PDF attached to the main response, we have presented pre-training results for ResNet18 and ResNet50 on ImageNet, trained _from scratch_ showing that our optimizer outperforms Adam and even SGD in this setting, while using _less than half of the memory_ relative to Adam-8bit, and less than 1/4 relative to SGD! 

Our SFT results on Llama are on 7B and 13B parameter models (the latter is at the end of the rebuttal PDF), where our method presents significant practical advantages--specifically, it allows the user to use a single GPU for SFT, rather than several, while providing a solid convergence theory. 

We unfortunately do not have the GPU resources currently to perform pre-training at this scale. With all due respect, we do not believe that being able to perform billion-parameter-scale pre-training experiments should be a requirement for acceptance at NeurIPS. 

> The whole script is written in the style of proposing a new optimizer for generic training procedures (including pre-training), not just for SFT.

We respectfully disagree on this point: as the reviewer can see by examining our Limitations section (lines 327 to 331 of the submission), the lack of large-scale pre-training experiments is clearly stated as a limitation of our paper, and extending MicroAdam to large-scale pre-training is cited as our main direction of future work. We believe MicroAdam can be extended to large-scale (LLM) pre-training, but doing so would require significant application-specific effort and major computational resources for validation. 

We believe we have established the viability of our approach in the submission, both via strong theoretical guarantees and via medium-scale SFT experiments. The rebuttal shows that our approach is also valid for vision pre-training experiments. We would be happy to amend the submission to further clarify the fact that we do not currently aim to do large-scale pre-training, which is left for future work. 

Respectfully,\
The authors
Thanks a lot for the response. 

I still think it is rather weak to validate the efficacy of a new optimizer by  SFT on Llama.  Further, I kindly disagree with authors comment that "our research question was "are all gradient entries useful for fine-tuning?” " This is not shown in the script. The whole script is written in the style of proposing a new optimizer for generic training procedures (including pre-training), not just for SFT. 

I will keep my score.
I thank the authors for the feedback. It seems that most of my concerns are addressed and I will raise the score.
Dear Reviewers,

Thank you for your valuable contributions to the review process. As we enter the discussion phase (August 7-13), I kindly request your active participation in addressing the authors' rebuttals and engaging in constructive dialogue.

Please:

- Carefully read the authors' global rebuttal and individual responses to each review.

- Respond to specific questions or points raised by the authors, especially those requiring further clarification from you.

- Engage in open discussions about the paper's strengths, weaknesses, and potential improvements.

- Be prompt in your responses to facilitate a meaningful exchange within the given timeframe.

- Maintain objectivity and professionalism in all communications.

If you have any concerns or need guidance during this process, please don't hesitate to reach out to me.

Your continued engagement is crucial for ensuring a fair and thorough evaluation process. 
Thank you for your dedication to maintaining the high standards of NeurIPS.

Best regards,

Area Chair