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 2024 poster (2024)
Keywords: adaptive optimization
adam
efficiency
memory efficiency
Primary area: optimization_for_deep_networks
TL;DR: We propose a new memory-efficient version of Adam with strong theoretical and practical convergence, and low space overheads in practice.

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.

MICROADAM: Accurate Adaptive Optimization with Low Space Overhead and Provable Convergence
Tonut- Vlad Modoranu
Mher Safaryan'
Grigory Malinovsky?
Eldar Kurtic
Thomas Robert
Peter Richtárik2 Dan Alistarh"
Institute of Science and Technology Austria (ISTA) 2 King Abdullah University Science and Technology (KAUST)
Abstract
We propose a new variant of the Adam optimizer [Kingma and Ba, 2014] called MICROADAM that specifically minimizes memory overheads, while maintaining theoretical convergence guarantees. We achieve this by compressing the gradient information before is fed into the optimizer state, thereby reducing its memory footprint significantly. We control the resulting compression error via novel instance of the classical error feedback mechanism from distributed optimiza- tion [Seide et al., 2014, Alistarh et al., 2018, Karimireddy et al., 2019] in which the error correction information is itself compressed to allow for practical memory gains. We prove that the resulting approach maintains theoretical convergence guarantees competitive to those of AMSGrad, while providing good practical per- formance. Specifically,
show that MICROADAM can be implemented efficiently on GPUs: on both million-scale (BERT) and billion-scale (LLaMA) models, MI- CROADAM provides practical convergence competitive that of the uncompressed Adam baseline, with lower memory usage and similar running time. Our code is available http://hhw....o//SSt////S///////
Introduction
The Adam [Kingma and Ba, 2014] adaptive optimizer and its variants [Reddi et al., 2019, Loshchilov and Hutter, 2019] has emerged as 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 naive SGD-based optimization, the Adam optimizer states doubles the memory overhead, as it requires storing two additional parame- ters for each variable. For large-scale models, these memory demands pose significant challenge. In turn, this has spurred research into memory-efficien -fficcient 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 metries.
Contributions. In this paper, we address this gap by introducing MICROADAM, an adaptive optimizer which guarantees low memory usage but also ensures provable convergence. We 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 error
"Correspondence conut-vlad.. www.mmammmist..aaccc..
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

due to compression by employing novel variant of error correction, mechanism introduced for distributed optimization [Seide al., 2014]. However, simply using error feedback would not lead to memory savings, since the size of the error correction buffer is comparable to the that of the original optimizer state. Instead, our main algorithmic innovation in showing that the error feedback can itself be compressed in the context of adaptive optimization. This renders the memory overhead of error feedback negligible, while preserving convergence guarantees.
Specifically, on the theoretical side, we provide new analysis showing that, under reasonable assumptions on the loss function being optimized and on the degree of compression, MICROADAM provably guarantees convergence, at asymptotically the same rate as AMSGrad [Zhou et al.. 2024a], i.e. a version of Adam with general convergence guarantees, that fixes fundamental technical issue in the Adam optimizer's 's proof [Reddi et al., 2019]. The key finding is that our approach allows for the overhead of compression to be shifted to the higher-order terms, where it should not impact practical convergence in common cases. This claim holds both for general smooth non-convex functions, and for non-convex functions under the Lojjsi--o (PL) assumption, highlighting trade-off between the degree of compression of the gradients, and that of the enror feedback.
We complement our algorithmic and analytic results with an efficient GPU implementation of MICROADAM, which we validate for fine-tuning language models from the BERT [Devlin et al., 2018], OPT [Zhang et al., 2022] and LLaMA [Touvron et al., 2023] families, with hundreds of millions to billions of parameters. We show that, in practice, gradients can be projected to very high sparsity (99%). while the error correction can be stored convergence. Concretely, our method can significantly improve upon the memory footprint of the extremely popular 8bit Adam [Dettmers et al., 2021] when fine-tuning models such as LLaMA2- 7B/13B [Touvron et al., 2023], at similar or better accuracy. At the same time, MICROADAM provides better accuracy relative to high-compression heuristics such as GaLore [Zhao et al., 2024].
bits per component, without loss of
In summary, we provide a new timization, which has the advantage of providing both theoretical guarantees and good practical convergence, while being scalable to billion-paramene parameter models. MICROADAM could therefore serve as a useful new tool for accurate and memory- optimization of large models.
approach to memory-efficient adaptive op-
2 Related Work
We mainly focus on related work reducing the cost of optimizer states. Dettmers et al. [2021] considers this problem, specifically by performing fine tne-grained quantization of the optimizer states. Their work does not alter the Adam algorithm; instead, it deals with the challenge of accurately compressing the dynamically-changie correspond to roughly halving the memory required by the optimizer states, relative to FP16. In the same vein, AdaFactor [Shazeer and Stern, 2018] and CAME [Luo et al., 2023] reduce memory cost by factorizing the recond-order statistics, while the recent GaLore [Zhao et al., 2024] factorizes the gradients themselves before they enter the optimizer state (but does not use error correction). Importantly, these methods are heuristics: they do not provide theoretical guarantees under standard assumptions,? 2 and in practice require careful tuning to preserve convergence [Luo et al., 2023]. By contrast, our method is theoretically justified, and provides good practical convergence. Earlier work by Anil et al. [2019] provides convergence guarantees for compressed variant of Adagrad [Duchi et al., 2010] called SM3, improving upon earlier work by Spring et al. [2019]. However, it is not clear how to extend their approach to the popular Adam optimizer, and heuristic methods appear to provide superior performance [Luo et al., 2023].
meta-data sequence. As the name suggests, the space savings
Conceptually, our work is related error feedback mechanisms studied in distributed optimization. e.g. [Seide et al., 2014, Alistarh et al., 2018, Karimireddy et al., 2019, Richtárik et al., 202 I]. Specifically, Li al. [2022] proved convergence of AdaGrad-like algorithms in conjunction with error feedback, in distributed environment. Our focus is different: minimizing memory costs in the single-node setting: for this, we show that the error correction buffer can itself be compressed. We provide an analysis for the resulting new algorithm, and efficient CUDA implementations.
More broadly, scaling adaptive or second-order optimizers to large models is very active area. Works such as GGT [Agaraal et al., 2019], Shampoo [Gupta et al., 2018) 8) and M-FAC [Frantar et al.,
'GaLore [Zhao et al., 2024] does state convergence guarantees for variant of the algorithm with fixed projections, but this is under strong "stable rank" assumption, which may not hold in practice.
2

---

2021] provided quadratic-space algorithms that are still feasible to execute for moderate-sized DNNs, but will not scale for aillion-paramete models. uollow-up work such as AdaHessian [Yao al.. 2020], Sophia [L.u et al., 2023], Sketchy [Feinherg et al., 2023] and EFCP [Modoranu et al., 2023], scaled these approaches via additional approximations. Of these, the closest work to ours EFCP, which uses sparsification plus standard error feedback to compress the gradient window employed in the Fisher approximation of the Hessian. However, EFCP does not compress the error accumulator, assumes a different optimization algorithm (Natural Gradient [Amari, 2016)), lacks convergence guarantees, and does not scale to Dallion-paramed models.
3 The MICROADAM Algorithm
Notation. We consider a standard Adam-type algorithm, which we will augment for memory savings. We will use f for the loss function, for the model size, & for the gradient density (sparsity - k), 8, and 9r for the model parameters and gradient at step t respectively, d) for the learning rate, A for the weight decay parameter, m, and v/ for the first and second moment of the gradient, e for the numerical stability constant, 31 and 83 for the momentum coefficients for my and Uy respectively. Furthermore, we use e for the error feedback (EF) vector, b the number of bits for EF quantization, m for the sliding window size, 9 (I, v) for the sliding window of size m l that stores indices I and values V/ selected by the Top-K operator.
Algorithm Description, We provide pseudocode in Algorithm and highlight the parts related to error feedback quantization in blue. The main idea is to compress the gradients via Topk sparsification before they enter the optimizer state, and to correct for the inherent error by applying error feedback e, E R. Instead of storing the optimizer state directly, we maintain a "sliding window" of highly- sparse past gradients and dynamically re-compute the Adam statistics each step based on this window. Yet, this alone does not improve space, as the error buffer partially negates the benefits of gradient compression. Instead, we prove that the error feedback accumulator can itself be compressed via quantization.
In detail, at step i = 1, the error feedback e is completely zero, as initialized in line 2, and thus, at line 5 the accumulator a will only contain the stochastic gradient gr. At line 6, we perform the Top-K compression and only keep the top-1% of values V and their corresponding indices Z1. 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, remove the outliers from the accumulator because they will be transferred to the buffer matrix G. This step is equivalent to a Ta (a) found in theoretical results. After line 7, what is left in is called the error feedback (e.g. the weights which were not chosen by Top-k). At line 8, we compute the statistics ó and A needed for quantization, and 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, 0 (computed by squaring the entries of element-wise) and update the model parameters.
For steps & 2 2, the only change compared to = that error feedback e is not zero anymore. Since the error is compressed,
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, ma/(c+ Vvz) 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 = 10 and suppose that all rows in the indices matrix are disjoint, then the overall density in the update u, will be 90%. The sparsity of We 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, 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 2, + -2,-1+ 1 (1-8)g for the last m sparse gradients as 2n + (1-3) Ei 3t-i, 'g and we also perform the bias correction in the end. Because we compute m, and 8 using the last m sparse gradients in the window, in line we dynamically determine the exponent r for the decay factor 3" based on the current optimization step t, 2th row of the circular buffer G and the window size
3

---

m. The last gradient added to 9 will have r = 0, while the oldest gradient in G will have m 1. In the end, we will add the values 8" V; to the buffer 2 the corresponding indices Z, which is fast operation because we only manipulate 1% of values at 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 incorporated into the optimizer state for Adam, both with and without error feedback Figure shows how EF fixes AdamW W 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 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 MICROADAM, allows for competitive convergence relative to the uncompressed baseline. In Appendix F, we discuss the implications of Error Feedback applied to GaLore.
Algorithm Pseudocode for MICROADAM with quantized EF
Algorithm Adam Statistics, Quantization and Inverse Quantization
1: Input: 31, B2, e, G,T.
1: procedure ADAMSTATS(3,G.CO
2: mo. Up +- 0d, Oa
2:
x- 02 3: for i {1, 2.., min(t, do
81, A <- 0,0
e1 + 0/bb
4:
r T (f (t-18 (1-i-11 1)%mm
3: for f = 21., T do
5:
z[Z.] f 2[[,] + BTV,
4:
St +- (o.)
6: end for
(1-9)²
5:
ar <- Ie + Q-1 dc, At)
7: return 1-8t
6:
I. V. + [aaea)))
8: end procedure
7:
a.I -
procedure Qxx A, b = 4)
8:
8411 A f mim(ma)))) (,), mar(ar)
2: 2
9:
E++1 + Q(a,, Ast1)
3: IQQ -
10:
Gr. t (D)
4: return æQ
11:
me + ADAMSTATS(31
5: end procedure
12:
Ú, + ADAMSTATS(3A g2)
1: procedure Q 6, b)
13:
0.+1 + 02 - Tt etvo
thee
2:
14:
it (i + 1)%m
3: I f IQ u
15: end for
4: return x
5: end procedure
Figure 1: Optimization trajectories of Adam, TopK-Adam and TApm-Adam with EF applied on the
Rosenbrock function far, y) = (1 x²² + 100(y x²)² starting from (ra, yo) = (- 1). Notice the extremely "jagged" profile of TopK-Adaa Adam without EF, and the recovered convergence when EF is added.
Adam
Top-k Adam
Top-k Adam with EF
201
10:
10
30n
30
32
1.2 35 62
23
1.
3.1 Efficient Implementation
A direct implementation (e-g., in Pytorch) of the previous algorithm would not bring significant benefits, and would fact might slow down optimization in terms of wall-clock time, To realize the thcoretical gains, we detail GPU-aware implementation below.

---

Accumulator at. First, we do not use an additional accumulator tensor a; 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 aray of size 2/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 in164 type for the indices matrix Z, assuming that the Top-K operator is applied globally to all the parameters in ag. To avoid this cost, we apply Top-K in blocks of fixed size Ba < 215 -1 1 = 32767 and store block-relative indices in intl6 format (during the development of MICROADAM, PyTorch did not have support for uint16). Applying Top-K per row to o, reshaped to 2D is not only faster, but provides the block-relative indices directly.
Computing Top-K in blocks also allows blocks to dynamically compute the statistics my and éz for Adam, as described in the ADAMSTATS procedure in Algorithm 2. We allocate the maximum possible shared memory for each thread block and store me (first half) and &, (second half) at consecutive locations in the shared memory. Once these statisties are computed, it 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.
allocate and efficiently use CUDA shared memory
Quantization metadata. Our approach also stores two additional vectors and A used for quantiza- tion. Since the quantization block size B, is set to very large value (e.g. 100K), the space required for these two arrays becomes negligible in comparison to the buffer 9 and error feedback e.
Practical memory usage. We note that we apply MICROADAM per layer, and that the size of quantization statistics and A 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- billion- parameter models, where the savings significant.
3.2 Memory footprint analysis for the optimizer states and comparison with other methods
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, and Ut. each of size d, expressed in bytes (B). For concreteness, we report the practical memory usage for the optimizer state for Llama-2 7B model for each optimizer.
Adam W stores states in flout32 format (4 B per component), resulting in a total memory footprint of 4d + 4d = Sd (B), while using bfloat16 would result in 4d (B) memory. We will refer to these memory footprints as Maw32 = Sd (B) = 50.21 (GB) and MAW16 = 4d (B) = 25.10 (GB).
AdamW-8 bit stores states in 8-bit format (1 B per component), both with d components, with memory footprint of M Mwws = d + d = 2d (B) = 12.55 (GB).
MICROADAM DAM stores the error feedback in 4-bit format (0.5 B per component) and the sliding window G that stores the indices I in 16 and V in b format. Both have m x k components, each requiring B per component. In the end, for m 10 and k = d/100, the memory footprint is Maa (m) 0.5d + 4mk (B)= = 0.9d (B)= = 5.65 (GB), that is, half of Adam W-Sbit.
GaLore. Given neural network with L layers, where each layer has weight matrix W; E RAAB (shaped 2D matrix), the model size d Et= A,B, GaLore uses rank-r compression via the SVD composition as W, USVT. where U E A; and we choose the first columns of U as R. E IR A, to project the gradient of W; to the r-dimensional space. As a consequence, the dimension d shrinks to d, E/1 components to be stored in the GPU memory only for the projection matrices R, If we suppose they are stored in 16 (2 B per component), then the entire memory usage for low-rank projection would be 2d,,. 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 6 for memory of rank- 1 layers). In addition, we have to store the states m, and v, for AdamW in bfloat/6 16 format, which adds another 4d, bytes. sum, the total memory footprint of Gal ore is McI.Aw1e(r) = 6d,, + 2e, (B), while for the 8-bit version we get MCLAWs(T) 4d, + 2e, (B). In the end, the practical
i= Air = r Ete Ad, which represents the total number of
5

---

memory usage for Llama-2 7B is McLAW& (256) - 1.36 (GB), McLAWs:102)) 5.43 (GB), McL.wi1(2566) 2.04 (GB) and McL.AWS 1024) 8.15 (GB).
Discussion. Assume our goal is to obtain a lower memory footprint compared to Adam bit. We fix the gradient density to k = d/100 and have determine the number of gradients (window size) m for MICROADAM in order to be competitive with Adam bit.
For this, we have to solve the equation MA (m) Maws for m, resulting in mmax 37.5. Specifically, if we use a gradient history of m < Mmmax gradients in G. MICROADAM will have theoretical memory savings. We will see that, in practice, this history size enough for good practical results. As entries in the window past this range are dampened extremely significantly, their influence is negligible. In Appendix D, provide Python code to compute the theoretical memory usage for the three optimizers for Llama-2 7B.
10 is more than
4
Convergence Guarantees for MICROADAM
In this section, we present our theoretical framework. We begin by introducing and discussing the analytical assumptions used in theory, providing an analytical view of the proposed MICROADAM algorithm, along with two theoretical convergence results.
4.1
Gradient and Error Compression
We now define two classes of compression operators widely used the literature. Assumption 1. The gradient compressor C Rdd I& is g-contractive with 0 S < i.,
[[C(x) - x=| S ||=]], for any x E Rd.
The compression use in Algorithm 1 is the TopK compressor Ta selecting top k coordinates in absolute value. This is known to be contractive with q vI - 1/a. Another popular contractive compressor is the optimal low-rank projection of gradient shaped as a d x d matrix, in which case q = / 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 : Rd + Rd is unbiased and w-bounded with w 2 0, namely, [[[(=)) = J, [1 - x1 S w ||=||: for any E Rd.
One example of w-bounded compressor, a version of which is used in Algorithm 2, is the randomized rounding quantizer we employ, whose properties we provide below.
Lemma Consider Algorithm with randomized rounding. i.e., for a vector x E Rdd with 5 = min; Ij and A = max; F2 let f; := _==ق + E + be the i-th coordinate of the quantized vector i, where & ~ U|0, 1] is the uniform random variable and 2-5 is the quantization level. Then 2-
dd-2 A-5 |||||| B[ x. 22 x 2- VA
for all E ERd
Next, we provide "analytical" view of our method in Algorithm 3. Essentially, use the contractive compressor C for compressing the error corrected gradient information 9: + et, and the unbiased compressor Q to compress the remaining compression error 9. e, C(g. + e,)-
Algorithm 3 MICROADAM: Analytical View
I: Input: parameters B1, B2 E (0, 1). > step-size 17 0, 01 12d, el mo U0 00 0,2
2: for = {1, T do
3: 9t Vef(O:)
Compute unbiased stochastic gradient Add accumulated error e and compress Update and compress the error o Update first-order gradient moment Update second-order gradient moment o Apply AMSGrad normalization o Update the model parameters
4:
Ôt C(g+ + ct)
5:
C++1 1 Q(ei gt - gt
6:
m = 31m4-1 + (1 31 )9:
7:
U, B2U1 1 + (1 )
8:
Or max Ur, 1)
9:
0.+1 0r Vorte
m
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 9t + e, using two compressed vectors: Gr = C(g. + +e) and Q(g + 9t). However, in doing so, we inevitably lose some information about Jr e depending on the degree of compression applied to each term. Thus, the condition (1 + w)g < required by our analysis can be seen as preventing excessive loss of information due to compression.
4.2 Convergence Guarantees for General Smooth Non-convex Functions
Next, state our algorithm's 's convergence guarantees under standard assumptions, stated below: Assumption 3 (Lower bound and smoothness). The loss function f: IRd + RR is lower bounded by some f E R and L-smooth, i.e. [[V - ff('))l S L 1 - for any 0, 0' R
Assumption 4 (Unbiased and bounded stochastic gradient). For all iterates t Z 1, the stochastic gradient St is unbiased and uniformly bounded by constant G 2 0, i.e., Elga] (f()),, l|gell s G.
Assumption 5 (Bounded variance). For all iterates I. 2 1, the variance of the stochastic gradient gr uniformly bounded by some constant o² 2 0, i.e., Elllg: f(0,) 5
Main Result. The above assumptions are standard in the literature, e.g. [Défossez et al., 2022, Li al., 2022, Xie et al., 2023, Zhou et al., 2024a]. Under these conditions, the two compressors satisfy the basic condition (1 w)y < 1, we show:
Theorem 1. (Non-convex convergence rate) Ler Assumptions 1, 2, 3, 4, 5 hold and = (1+w)q < Then, choosing l] min ( ) MICROADAM (Algorithm 3) satisfies
E B|| f(()) 20,
((1))-11 + 140440309) eVT
+ + (++d)
VT
29 with constants Co == 41+ G and C2 == wg(1 + 1-
Discussion. First, notice that the leading term + of the rate is the optimal convergence speed for non-convex stochastic gradient methods [Ghadimi and Lan, 2016]. Furthermore, the obtained convergence rate O + {) 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 Co and C2 affected by compression and appearing in the leading terms be easily estimated once the compressors are fixed. Besides, if we store the full error information without applying Q compressor (i.e., w = 0, qu q), then MICROADAM reduces to the single node Comp--MS method by Li et al. [2022] recovering the same convergence rate. The full proof is provided in the Appendix.
4.3 Convergence Rate for Non-Convex Functions under the PL Condition
Next, we show that we can obtain even stronger bounds when the objective satisfies the PL condition: Assumption 6 (PL-condition). For some > O the loss J satisfies Polyak-Lojasiewizz (PL) inequality [[)] 2u(f(0) f³), for any 0 E Rd.
In this case, show:
Theorem 2. (PL convergence rate) Let Assumptions 2, 3, 4, 5 and 6 hold, and qu < 1. Then, 2Co log T MICROADAM (Algorithm 3) satisfies aT
choosing 7 min TLC
Elf(Or+))) f < 2logT ( LLC HC C6(1+C))(+4088 + Ö (G(())) uyE
(
Jrc
T2
with constant C₁ := I- (1 + C2) + 29
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 PI. condition, which claims to be the first in this direction, focuses on Adam when 32 > 1, achieving a convergence rate of (++) [He et al., 2023]. However, our MICROADAM based on AMSGrad normalization, and constraint on Bg is imposed in the analysis. Therefore, similar to
7

---

the general non 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 T2. Hence, in this case as well, the theory predicts that the convergence rate of the algorithm should be similar to the uncompressed version, modulo a constant that can be controlled using the compression parameters.
5
Experiments
We now validate our optimizer experimentally. We focus on comparing MICROADAM with Adam, Adam-8bit, GaLore and CAME in the context LLM finetuning on different tasks and with SGD, Adam and Adam AdamW-8bit W-8bit in the context of ResNets on ImageNet. Concretely, we test our optimizer in full finetuning (FFT) scenario on LERT-Base/Large [Devlin et al., 2018] and OPT-1.3B [Zhang et al., 2022] 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/MNLL. We first test our integration of MICROADAM HuggingFace Transformers [Wolf ct al., 2020] on moderate-sized language models such as BERT-Base/Len (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 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 the state-of-the-art heuristics Adam-&bit 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 significant accuracy cost for GaLore: across all tasks, drops at least 1% accuracy relative MICROADAM. For BERT- Base we ran GaLore with higher SVD re-computation frequency T' = 20 (10x 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, conclude that MICROADAM can provide better accuracy relative other memory-efficient methods on moderate-sized models, at similar space costs. We show training loss curves in Appendix C.
Table 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).
MICROADAM
GaLore
Model
Metric
Adam
Adam-8b
CAME
(m = 10)
r 256
train loss
0,2651
0,4228 83.53%
0.3402
0,6431 *
0.3908
BASE
accuracy
85.10%
84.61%
7
83.82%
(110M)
memory
2.55 GB
2.70 GB
2.53 GB
2.69 GB
2.53 GB
train loss
0.2509
0.3857
0.2876
0.6658
0.3768*
LARGE (335M)
accuracy
86.17%
84.79%
86.18%
*
84.90%*
memory
5.98 GB
6.64 GB
6.04 GB
6.59 GB
5,85 GB
train loss
0.2122
0.2066
0.2611
0.4959
0.2831
OPT-1.3B
accuracy
88.18%
87.90%
87.81%
83.15%
87.70
(1.3B)
memory
15.28 GB
17.66 GB
15.00 GB
17.13 GB
13.66 GB
Finetuning results for LLaMA2 on GSM-8k. Next, we perform finetuning on Llama-2 7B/13B on GSM-8k, challenging grade-school--eess mathematical reasoning dataset. The baseline model obtains extremely low zero-shot accuracy on this task and therefore fine-tuning is necessary. this setup, we compare MICROADAM with Adam and Adam-8bil in terms of evaluation accuracy and memory usage. In Table we show our results for training epochs, global batch size 32 with micro-batch (per-device) size 1, max sequence length 512 a single GPU, which are the standard
8

---

parameters for this task. We integrated our optimizer with the 11m-foundry repository of MosaicML and tested via In-evaluation -naauattss.-..nesss.... harness.
For the 7B model, out results show that MICROADAM can allow accurate full fine-tuning of a 7B model on this task using single 40GB GPU. Moreover, MICRO preserves accuracy relative to Adam, with lower memory usage than the well-optimized implementation of 8bit Adam W, 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 very similar those in Table 2 and are therefore omitted from Table 3.) Moreover, MICROADAM obtains consistently better accuracy relative Adam-8b, especially on the more challenging ARC-c task.
Table 2: FFT results for Llama-2 7B/13B GSM-8k. 8k.
LLaMA-2 size
Optimizer
Accuracy
State
Total
Runtime
Adam
34.50%
25.1 GB
55.2 GB
Ih 17m
7B
Adam-8b
34.34%
12.55 GB
42.5 GB
18m
MICROADAM (m 10)
34.72%
5.65 GB
37.1 GB
Ih 8m
MICROADAM (m = 20)
35.10%
8.25 GB
39.7 GB
1h 37m
Adam
47.08%
48.42 GB
>80 GB
Ih 20m
13B
Adam-8b
45.19%
24.21 GB
>80 GB
Ih 17m
MICROADAM (m = 10)
44.88 %
10.9 GB
70 GB
1h 38m
Finetuning results for LLaMA2-7B on Open-Platypus. Finally, in Table 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 h: 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 al., 2018], HellaSwag [Zellers al., 2019], MMLU [Hendrycks et al., 20 and Winogrande [Sakaguchi et al., 2019].
Table 3: FFT results on instruction-followie
Ppen-Platypus [Lee et al., 2023] dataset. The results show that MICROADAM fully recovers accuracy relative baseline Adam, and outperforms the 8bit variant, despite using less memory.
Average Accuracy
ARC-c 25-shot
HellaSwag 10-shot
MMLU
Winogrande 5-shot
Optimizer
Memory
5-shot
AdamW
67.17 GB
62.10
52.56
77.38
45.53
72.93
Adam-8b 8bb
53,93 GB
61.84
51.96
77.51
44.11
73.79
MICROADAM
46.63 GB
62.36
53.07
77.46
45.04
73.87
Pre-training results for ResNets on ImageNet. In Table 4 we present our results for pre-training (from scratch, randomly initialized weights) for ResNet- 18/50 on ImageNet (see Figure and Figure 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 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, Adam W-8bit reaches the lowest training loss, followed by AdamW, MICROADAM and SGD. The validation accuracy for AdamW and Adam W-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).
Table 4: Pre-training results for ResNet- 18 and ResNeet-50 on ImageNet.
Model
Metric
SGD
AdamW
MICROADAM
Train Loss
1.416
1.087
1.104
1.218
ResNet- 18
Accuracy
69.96%
69.83%
70.13%
71.6%
State Size
44.59 MB
89.18 MB
22.30 MB
10.03 MB
Train Loss
0,9770
0.5344
0,5158
0.7732
ResNet-50
Accuracy 76.24%
72.05%
72.48%
77.37%
State Size 97.49 MB
194.98 MB
48.75 MB
21.94 MB
Pre-training LLMs. In this section we explain why we do not include L.M pre-training results. Our motivation is twofold. First, MICROADAM is mainly designed for low-memory finetuning, and the experimental section shows that MICROADAM achieved this goal. Surprisingly, updating only 10% of the weights each step yields to significantly better performance compared to SGD for ResNets on ImageNet. Secondly, our experiments on LLM pre- pre-training showed difficulties in achieving the same performance compared to AdamW-Sbit 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 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, well as 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 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 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 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. is important note that while our optimizer can enhance performance and reduce costs, we do not have control over the neural network applications trained with it.
Acknowledgements
The authors thank Razvan Pascanu, Mahdi Nikdan and Soroush Tabesh for their valuable feedback, the IT department from Institute of Science and Technology Austria for the hardware support and Weights and Biases for the infrastructure track all experiments. Mher Safaryan has received funding from the European Union's 's Horizon 2020 research and innovation program under the Marie Aklodowska-Curdd grant agreement No 101034413.
References
N. Agarwal, B. Bullins, X. Chen, F. Hazan, K. Singh, C. Zhang, and Y. Zhang. Efficient full-matrix adaptive regularization. International Conference on Machine Learning, pages 102-110. PMIR, 2019.
10

---

D. Alistarh, T. Hoefler, M. Johansson, N. Konstantinov, S. Khirirat, and C. Renggli. The convergence of sparsified gradient methods. In Advances Neural Information Processing Systems, pages 5973-5983, 2018.
S.-i. Amari. Information geometry and its applications, volume 194. Springer, 2016.
R. Anil, V. Gupta, Koren, and Y. Singer. Memory efficient adaptive optimization. Advances in Neural Information Processing Systems, 32, 2019.
E. Beeching, C. Fourrier, N. Habib, S. Han, N. Lambert, N. Rajani, O. Sanseviero, L. Tunstall, and
T. Wolf. Open Ilm leaderboard. https://huggingfaan httpspacce//tuuggiiggaeeeeeeeeee
11m leaderboard, 2023.
Clark, Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? try arc, the al2 reasoning challenge, 2018.
A. Défossez. Bottou, Bach, and N. Usunier. A simple convergence proof of adam and adagrad. Transactions on Machine Learning Research, 2022. ISSN 2835-8856. URL https: httpge////pperviee..o htt.://wweeettroomrmdddddooooaaa
Dettmers, M. Lewis, S. Shleifer, and L. Zettlemoyer. 8-bit optimizers via block-wise quantization. arXiv preprint arXi::21 2021.
Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. 2018.
J. C. Duchi, E. Hazan, and Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. J. Mach. Learn. Res., 12:2121-2159, 2010.
V. Feinberg, X. Chen, Y. J. Sun, R. Anil, and E. Hazan. Sketchy: Memory-efficient adaptive regularization with frequent directions, 2023.
E. Frantar, E. Kurtic, and D. Alistarh. M-fac: Efficient matrix-free approximations second-order information, 2021.
S. Ghadimi and G. Lan. Accelerated gradient methods for nonconvex nonlinear and stochastic programming. Mathematical Programming, 2016. ISSN 156(1-2):59-99,
V. Gupta, T. Koren, and Y. Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pages 1842-1850. PMLR, 2018.
M. He, Y. Liang. J. Liu, and D. Xu. Convergence of adam for non-convex objectives: Relaxed hyperparameters and non-ergodic case. Journal of Machine Learning Research, 2023.
D. Hendryeks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt. Measuring massive multitask language understanding, 2021.
S. P. Karimireddy, Q. Rebjock, S. U. Stich, and M. Jaggi. Error feedback fixes SignSGD and other gradient compression schemes. In Proceedings of the Thirty-sixth International Conference Machine Learning, pages 3252-3261, 2019.
D. P. Kingma and J. Ba. Adam: A method for stochastic optimization, 2014.
A. N. Lee, C. J. Hunter, and N. Ruiz. Platypus: Quick, cheap, and powerful refinement of Ilms. arXiv preprint arXivv2308.0733 2023.
X. Li, B. Karimi, and P. Li. On distributed adaptive optimization with gradient compression. arXiv preprint arXiv:2205. 05632, 2022.
H. Liu, Z. Li, D. Hall, P. Liang, and T. Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. arXiv preprint arXiv:2305.14344 /4342, 2023.
I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In Proceedings of the Seventh International Conference Learning Representations, 2019,

---

Y. Luo, X. Ren, Z. Zheng, Z. Jiang, X. Jiang, and Y. You. Came: Confidence-guiin guiiedd adaptive memory efficient optimization. arXiv preprint arXiv:2307.02047 2023.
I.-V. Modoranu, A. Kalinov, E. Kurtic, and D. Alistarh. Error feedback can accurately compress preconditioners. arXiv preprint arXiv:2306.06// 2023.
S. J. Reddi, S. Kale, and S. Kumar. On the convergence of adam and beyond. arXiv preprint arXiv:1904.09237, 2019.
Richtárik, I. Sokolov, and I. Fatkhullin. Ef21: A new, simpler, theoretically better, and practically faster error feedback. Advances in Neural Information Processing Systems, 34:4384-4396, 2021.
K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi. WINOGRANDE: an adversarial winograd schema challenge at scale, 2019.
E. Seide, H. Fu, J. Droppo, G. Li, and D. Yu. 1-bit stochastic gradient descent and its application to data-parallel distributed training of speech DNNs. Fifteenth Annual Conference of the International Speech Communication Association, 2014.
N. Shazeer and M. Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596-4604 PMLR, 2018.
R. Spring, A. Kyrillidis, V. Mohan, and A. Shrivastava. Compressing gradient optimizers via count-sketches. In International Conference on Machine Learning, pages 5946-5955, PMLR, 2019.
H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, Rodriguez, A. Joulin, E. Grave, and Lample. Llama: Open and efficient foundation language models, 2023.
T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Fun- towicz, Davison, S. Shleifer, P. von Platen, C. Ma, Jernite, J. Plu, C. Xu, T. L. Scao, Gugger, M. Drame, Q. Lhoest, and A. M. Rush. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Process- ing: System Demonstrations, pages 38-45, Online, Oct. 2020. Association for Computational Linguistics. URL https: //www//wwww.//// aclweb. org/ htt///aathl/////// 2020. memns-demos... 6.
X. Xie, P. Zhou, H. Li, Z. Lin, and S. Yan. Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models, 2023.
Z. Yao, Gholami, Shen, K. Keutzer, and M. W. Mahoney. Adahessian: An adaptive second order optimizer for machine learning. arXiv preprint
2020.
R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi. Hellaswag: Can machine really finish your sentence?, 2019.
S. Zhang, S. Roller, N. Goyal, M. Artetxe, M. Chen, S. Chen, C. Dewan, M. Diab, X. Li, X. V. Lin, T. Mihaylov, M. Ott, S. Shleifer, K. Shuster, D. Simig, P. S. Koura, A. Sridhar, T. Wang, and
L. Zettlemoyer. Opt: Open pre-trained transformer language models, 2022.
J. Zhao, Z. Zhang, B. Chen, Z. Wang, A. Anandkumar, and Y. Tian. Galore: Memory-efficient Ilm training by gradient low-rank projection. arXiv preprint NXKv:22033003007 2024.
D. Zhou, Chen, Y. Cao, Z. Yang, and Q. Gu. On the convergence of adaptive gradient methods for nonconvex optimization. Transactions on Machine Learning Research, 2024a. ISSN 2835-8856. URL
Featured Certification.
Zhou, X. Xie, Z. Lin, and S. Yan. Towards understanding convergence and generalization of adamw. IEEE Transactions on Pattern Analysis and Machine Intelligence, pages 1-8, 2024b. doi: D0.1109/TPAMM2222.333300
12

---

Contents
1
Introduction
2 Related Work
3
The MICROADAM Algorithm
3.1 Efficient Implementation
3.2
Memory footprint analysis for the optimizer states and comparison with other methods
Convergence Guarantees MICROADAM
4.1
Gradient and Error Compression
4.2
Convergence Guarantees for General Smooth Non-convex Functions
4.3
Convergence Rate for Non-Convex Functions under the PI. Condition
5
Experiments
Limitations and Broader Impact
10
A Additional Explanations and Experimental Details
14
Training Settings and Hyper-parameters
14
B.1 GLUE/MNLI
14
B.2 GSM-8k.
15
B.3 ImageNet
15
C Training Graphs
15
Memory footprint for the optimizer state
16
E Deferred Proofs
18
E,1 Intermediate Lemmas
18
E.2 Non-convex Analysis
E.3 Analysis Under PL Condition
26
E,4
Non-convex Analysis with Weight Decay
30
Error Feedback applied to GaLore
34
F.I Behaviour of the Error Feedback Mechanism
34
F.2 Consequences Training
35
13

---

A Additional Explanations and Experimental Details
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 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 subset of layers, such as Q- K-, V-, O-, up- down- and gate- pate-projection layers. Applying GaLore to all layers in the same way
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 result, we omit the GaLore for LLaMa2 and ResNet experiments.
Implicit regularization. In our results so far, we observed MICROADAM having better performance for few LLM finetuning tasks, but especially for the ResNee/IImagNett results, where the difference was statistically significant (around 1 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, 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,, a sparse update would not change the model parameters much dense one. In the ResNet experiments, all optimizers used the same regularization parameter À = le - 4, but the accuracy graph shows that MICROADAM is more regularized than SGD, while the graphs for AdamW and Adam W-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 nigh-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-Sbit in memory--onnstriied settings, and that the empirical results appear to validate our theoretical predictions.
B Training Settings and Hyper-parameters
In this section we provide details about the syper-parametert that we used for each model and dataset. We train all our models in b format, tune the learning rates on grid and report the best accuracy among seeds (7, 42 and 1234) and report the results for the best configuration that converged.
All Adam variants use default parameters B1 = 0.9, 32 = 0.999, c = 10-8 and the regularization parameter > is for finetuning and 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 -1000XXM4-80GB H100-80GB and on RTX 3090 with 24GB RAM in single GPU setup.
B.1 GLUE/MNLI
For GLUE/NNLI, we used the learning rate grid {le- 6, 3c 5c 6. -77-6.2 le 5, 5e 5, 7e 5} for all optimizers and models. Certain optimizers diverge for specific seeds. Next, provide some details about
for each optimizer individually.
MICROADAM. We use = 10 gradients in the sliding window, A = 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 meaningful way).
14

---

Adam and Adam-8bit. All optimizers.
mentioned above apply for these two main baseline
Galore We use rank r = 256 and SVD update interval 7 E [20. 200). In the original Galore paper, the authors tune both learning rate and in experiments keep scale fixed value 1 and augment the learning rate grid with the values {le - 4, 3e - 4, 5e - 4, 7e -3.
CAME. This optimizer has some additional parameters that we keep to default values, such as 3g = 0.9999. Instead of uses é1 = le 30 and 62 = le 16. The authors mention that the learning rate should be much smaller than Adam's 's and because of that we augment the learning rate grid with the values {le 3e 7.5e 7, 7e -7.
B.2 GSM-8k.
5, 6e 5, 7e 5, Se 5, 9e - 5) and reported the model with the best evaluation accuracy. We found that different versions for PyTorch, in-eval-harnen and 1lm-foundry have large variance in the results.
For GSM-8k, we used the learning rate grid {le - -,
MICROADAM. We use similar settings as for GLUE/MNLI above in terms of other hyper- papee--rramttes...
B.3 ImageNet
For ImageNet, we integrate our MICROADAM in the FFCV repository, which is highly tuned for ResNets and SGD. We E 100 epochs, batch size 1024, cosine learning rate schedule with warmup and image resolution 224 x 224 and precision bfloat 16. We started from the initial learning rate 1) 1.024 tuned in the repository which scored highest accuracy for SGD, This learning rate also worked well for MICROADAM, but it didn't It work for AdamW and Adam W-8bit. For these two Adam variants we divided the learning rate by until the models converged.
ResNet-18, For AdamW, the learning rate is 0.016 and for AdamW-8bi is r7 = 0.032.
ResNet-50. For AdamW, the learning rate is 0.008 and for AdamW-8bi is r7 = 0.008.
C Training Graphs
In this section we show training loss curves for BERT-Base, BERT-Large and OPT-1.3b on GLUE/MNLI and Llama-2 7B/13B GSM-8k and MAANet-18/50 on ImageNet.
Figure 2: Training curves for BERT-Base on GLUE/MNLI
Training Loss for BERT-Base
MicroAdam
AdamW-8bit
1.0
GaLore
AdamW
train/loss
0.8
CAME
0.6
0.4
0.2
50
100
150
200
250
300
350
step
15

---

Figure 3: Training curves for BERT-Large GLUE/MNLI
Training Loss for BERT-Large
MicroAdam
AdamW-8bit
1.0
GaLore
AdamW
train/loss
0.8
CAME
0.6
0.4
0.2
50
100
150
200
250
300
350
step
Figure 4: Training curves for OPT-1.3B on GLUE/MNLI
Training Loss for OPT-1.3B
1.2
MicroAdam
AdamW-8bit
1.0
GaLore
AdamW
train/loss
0.8
CAME
0.6
0.4
50
100
200
250
300
350
step
D Memory footprint for the optimizer state
In this section we provide a python script to simulate the memory usage for our optimizer's state for Llama2-7b model. Note that the theoretical memory usage will always be slightly lower than the actual allocated memory on the GPU because PyTorch usually allocates more. To run this script, run the following commands:
im port math
6_738___ 11_411_6616 616 # actual number of parameters for Llama 7 k math. ceil (d / 100)
10
M_AW32
(2 ** 30)
M_W16 = 4 /
(2
** 30)
AW8 2
(2
we 30)
M muA = (0.5 d
+
(2 ** 30)
# B 1o GB
16

---

Figure 5: Training eurves for Llama-2 7B GSM-8k
Training Loss for Llama-2 7B
1.4
MicroAdam
1.2
AdamW-8bit
AdamW
1.0
loss/train/tota
0.8
0.6
0.4
0.2
0.0
100
200
300
400
500
600
700
step
Figure Pre-training for ResNet-18 on ImageNet
Twihing ins Beetien: Maaggeiee
Velidattio arruney Mimagelet
NicrAsoom Adum
0.5
My wmmm
0.2
MicrAddam AdamiAltiit sGD
Adami!
25
150
epoch
epach
Figure 7: Pre-training for ResNet-50 ImageNet
Redhet Samaaeeie
Validation
0.5
NirmAcam
5G3
AdamW auit
Adamm
toge
g.04
WemAdam
0.2
5G0
0.1
Mamm Bait
0.0
Adami
60
100
print (f 2
print (f' 22
print :
print (f' (f [M_muA=:.2f]_
GL sumA 1_423_877 872 # sum sum_Ai from Llama 7B epsilon = 266_240 # sum of sizes for rank -I layers
for bits const in [(8. 4), (16, 6)]:
for rank in [256. 1024]:
dr rank * GL
M. M_GLAW rank (const dr + 2 epsilon) / (2 30) print (f MM GLAW bits] bitss]rank rank rank
222) 2
17

---

E Deferred Proofs
At time step let the uncompressed stochastic gradient be 9: (o)( the error accumulator be and the compressed gradient after the error correction be 0r = Clg + et). The second moment computed by the compressed gradients is denoted as 1 f24 (1 32) and
0 max 6;-1, 2t) is the AMSGrad normalization for the second-orde momentum. Besides the
first-order gradient momentum me used in the algorithm description, we define similar running average sequence my based on the uncompressed gradients !t-
m = m- + (1 31)g. and mi = 3mm + (1 31)9t.
Note that my is used only in the analysis, we do not need store or compute it. By construction we have
m (1 11 31) Σ 3 Tr: m4
(1) Σ sT- gr
Denote by Ee+1 -ez
9x) Q(e: + - ge) (er + 9r 9t) the compression noise from
Q. Due to unbiasedness of the compressor 2 (see Assumption 2), we have E(G, 0t, It. It, e,] = Also, from the update rule of ct+1 get ct+1 ct 9: S.. Moreover, auxiliary sequences.
the following
H1
:+1:=BEE + -1-1811- B1je41 = (1)
8
a4+1-r
er-
r=1
1+1
Z++1 = 8 + 1 = (1-
Σ
8 Gr.
7=1
E.1 Intermediate Lemmas
Lemma Consider Algorithm with randomized rounding, i.e., for vector 12d with ô min; F; and A max; Ti let â == فيمم + E + 5 be the i-th coordinate of the quantized vector it, where & U00, 1] is the uniform random variable and A-8 is the quantization level. Then 2-1
3
B x. 2 x 2
| for all æ E Rdd
Proof. The unbiasedness be verified directly from the definition for each coordinate. Without loss of generality assume that 5 = I1 S 22 I ... < xd-1 S Id = A. By construction of the quantization, we have [ xı í zal = 0 and | I u for the remaining coordinates s Then
(d 2)u²
| - x =
Σ [2 xr/ (d 2)u² =
2x
A + 82
which completes the proof.
Lemma 2. Under Assumptions 1-5.for all iterates and T have
T
||m||| S G, and Σ
Ellmm 12] TO
(())).
Proof.
The first part follows from triangle inequality and the Assumption 4 on bounded stochastic gradient:
2
|mm|l| (1 B1)
Σ 0 gr < 3 Σ 3 g S.
For the second claim, the expected squared norm of average stochastic gradient be bounded by
E
[[]ge]]
! E [llge vf(e.)) + = = +
(I)
18

---

where we use Assumption 5 that 9e is unbiased with bounded variance. Let Ir,j denote the j-th coordinate of 9r- Applying Jensen's inequality for the squared norm, we get
27
E m| = E
(1
31)
31
'
$ (1 -1- 31) Σ
8 T 11 12
1=1
o + (1-11-1 31)
31)
Σ
r=1
Summing over 1. ,,...,? T, we obtain
T
El E| [13] S To + (1- B1)
ΣΣ 3 T - s TO +
BIII (()))l
which completes the proof.
Lemma 3. Let qu = (1+ + w)q < 1. Under Assumptions for all iterates we have
4q G2
1|e.112 - (1 - 2 )
4
2
1 + Tw
1-7 E
Bl|let1 119] S (1 2
(
Proof. We start by using Assumption 1, on compression and Young's inequality to get
||e,11| = =[ + t C(g: + ex)ll
(1 + 22 eell
1 1) :
<1 le 2
|1
p
1 + 2 $
2
||e||
1190 I2, 2
where (2) is derived by choosing p
1-42 and the fact that qu 1. For the first claim
24
recursively apply the obtained inequality and use bounded gradient Assumption 4. For the second claim, initialization = 0 and the obtained recursion imply
24
E[]ee [|3]
Σ
Billg 11
1 L
(1)
4
2
(1
2 2
2
which concludes the lemma.
Lemma 4. Let qw
(1 + w)q Under Assumptions for all iterates we have
2qw
|| < wg
+ 1 o3
G. and | < wq
1 11
G.
1 o
Proof. Using the bounds defining compressors and Lemma 3, we get
IKG.ll [[Q(e, + St - gt) (et + St gt)
wille; + ge- Tell = w|le> + gt Cle: + g0l
S wellet gell
welles |l + wallga
2g
s wa
1 1 o
G.
19

---

For the second claim, recall the definition of 2z and apply triangle inequality:
| S (1 31) Σ 8 1 || S wg 1
2gw
G.
L
Lemma 5, For the moving average sequence &,, it holds that
EBTLE|
47 47
4.
0²
+
4
[(.))
Proof. Let Et, be the j-th coordinate of e₁ and denote
E [[
Applying Jensen's inequality and Lemma 3, get
E E |2] E
(1 31) 3i-1e
T=1
$1 - B1) Σ
4
2
31)
Σ
K
Kr
2
9 7=1
Summing over 1 = I,..., T and using the technique of geometric series summation leads to
T
4T'g
(1-
T
Σ
B11 17|
Kr
(1-43) 1-44)
1=1
r=1
4T g
29
2
K,
(1 2
1=1
4T 4T 2
2
t-r
/(+)))(0))
2 2
(1 2
2
4T 2
4g
T
B
[(0,)1(
(1
2
(1
The desired result is obtained.
Lemma 6. Let q. = (1 + w)q < 1. Under Assumptions 1-5,for all iterates E [T] and coordinates i [d]. the following bound holds
4(1 (
61.5 S
G2
(1 22)
Proof. Lemma and Assumption 4 imply
l 2 |[C(g, + e1
< |[C(g + et) - (g: + et) + (ge + +
< 2(a² + 1)1la + l
$ <4(9² + 1) (6² + (1 2 2 G2
4
4(1 + (1 + 22)22 2
20

---

It's then easy to show by the updating rule of it, there exists j [t] such that 02,, = U/.s. Then
4(1 + (1 + 2.) G²,
= = 32)
Σ
3/ 3 ge. -
(1 3
r=1
which concludes the claim.
Lemma 7. For Dr
Veete we have
vb ite
d Ve
T
d -
Dilli
Σ
D.
Proof. By the update rule, we have ôt-1,8 S o2 for any iterate and coordinate E [d]. Therefore, by the initialization uy- 0, we get
T
T
$ d
Dell
Vo-iut+e
Vo,,te)
00.s +e
VUT, +e
Ve
t=1
[=1 t=1
==1
For the sum of squared la norms, note the fact that for 2 b> 0, it holds that
(a-b)² s b)(a + b) = a² 112
Thus,
T Σ |
T
d
2
T
d
ΣΣ
s
V
Vee, +e
ê-11 +e
te,
c
[=1
1=1
which gives the desired result.
E.2 Non-convex Analysis
Here derive the convergence rate with fixed step-size n). The rate shown in the main part can be obtained by plugging the expression of 1 shown after the proof.
Theorem 3. (Non-convex convergence rate) Let Assumptions 1, 2, 3, 4, 5 hold and qu == (1+w)y < Then choosing any step-size 1) S मíट, MICROADAM (Algorithm 3) satisfies
+ + [())))) 2C, [[()))1'
nLa²
aLC30
L'C
(+CC)0's TVe
x(+20))01109
P'
with constants Co = V (1-05)
4(1+g2 2 G2
+ C
(1 + C2) +
20
C2 == wg(1 +
2
31
1-
1-92
Proof. Similar the proof of Comp-AMS [Li al., 2022], we define two virtual iterates 0 and 2t-
E1
841 01+1
1)
/
0 +
B1 m
+ Zt
Z441 := 0/41
31 V
+
21

---

Then, we derive the recurrence relation for each sequence as follows:
E-1
011= = 8041 - 1)
Vū +
- (1-31) 1-11) 1 Dr + (1-1111
++1 81
02
v0r +
(1-3)) Σ r= 11 (+r + BT+1) + - -3)
0
Vo+t
(1- (1-31) 31 -= 81 [ ter (r)
02
Vol +
(1 -1- 31) Σイ r-
07 Te
ml
2
02
Vi, +E
Vo, Vutte +E
11
V00+e
&
mf
1
2.
= 0/- -
11 +E
-
Vo
+1
(Vi=I++
V0, +e
E
01 +
ml
1
Z,
= 81
Vūr + +1
Eg - -1
01-1t
Ve++e)
V0, +
m/ + Zt
DrEr
i +
where we used the fact that gr + 6++1 gt + et + Çe with quantization noise Go, and eg at initialization. Next, for the Te iterates we have
B1
m + 2
T1+1 = 6//41 - 7
B1 Vēr +
m/ 2,
81
m + Zt
02- 11
" 31 Vor +
+ D,,
Vi= +
31 (m/11 + 2x-1) + (1 - (gt + Gr) + I- (m/_1 + 2x-1) + Bi + Gt)
83
0% 17
v0, +c
+ +nDrEs
B1
m4 + B
St + St
0 17
D,E
B1
Vur +t
Vo, +
St + Ct
31
It
+ 1
De(m/ + 2,-1) + nD..
/0r +
31
Next we apply smoothness (Assumption 3) of the loss function f over the iterates He. From the gradient Lipschitzness we have
L 1 - 12ll..
f(2141) S f((e)) + (Vf(x)), - 21) +
2
22

---

Due to unbiasedness of the compressor Q (see Assumption 2), we have E'Gllgt &c. cz2 0q] = 0.
Taking expectation, we obtain
Vêr +c }
9e + Se
- - [[f(x))]
x -1yL
V fax))
+rylE
<
(Taa) 1- 31
31
De(mb-1 + 2c-1) +
77 E
9t + St
31
22
De(ml/ + Z(-1) 1 D2E,
2
V +
I- B1
-
-1/ -nn
(00))
Vor
gt
(3)
B1
nE
()(1)) 1-B, B1
D,(mL_ + Z,-1) D&E,)
11
1f L
9t + Ce
B1
Dm + Z,-1) - D
2
VUe +
1- 31
111
+ nB
VS((.) vs(x)).
(4)
doc +
In the following, we bound all the four terms highlighted above. Bounding term I. We have
I
V ((.))
8 +
Dt
-n/E
vf(O)),
Votté
1
Vā--1 +
1
) 9t
)
Vf(O.)
7
f(O.),
2
V +
(5)
4(+gaa)' 11 2
G
+E
where we use Assumption 4, Lemma 6 and the fact that la norm no larger than 4 norm.
Bounding term 11. By the definition of &, and Ze, we know that
2q G,
| S (1 (1) Σ - |elll S
oz
29w
|
<
(1 - 31)
Σ
8
--Gll S wa
1
G.
1=1
23

---

Then we have
11 S 1
vI(O.),
31
Di(m/-1 Zz-1) (--1)+ + DiE:)
1- B1
[(V/(o))) VS(()),
31
De(ml/1 Zx-1) DrEs)
1- 11
81 Di Dr Zz-1)
S n/ [[
T-B1 B1
+ LE
I-A
m\ + I-A1
2
B1 1 31
De(m/-1 + Z[-1) + D&E
V +
1f LG D+lld
S G/
(6)
Ve
where C = 1-81 of f()), and the last inequality is due to Lemma 3, Assumption and the property of norms.
8,
(1 + ww [
24 1-08
+
24
The second inequality is because of smoothness
1-
Bounding term III. This term can be bounded as follows:
+ Ce Vū +
31 D + 2(-1) - DEE
III S 1/ LE
+ + LE
1 1 B1
2n L
+ 2072L
Ellgge
vf(0)) + (
B||l<:113
(7)
||pi
(1-)) B1
81
Z1-1 - Ex) =
+ LB
1- 31 m4-1
1-
2n
2172
20 w¹q
29
1 I+ -2
G2
D.11
2n L.
2n L(o² + CZG
$ E [[
+ LG
(8)
where C2 = c1= + 1- and we used Assumption that ga is unbiased with bounded variance o².
Bounding term IV. We have
IV 7/E
(S(0.) V((e)),
gt VoLi+t +e
(9)
Vo-1 +6 It 1
1
7)
(f())) vf(zi),
Vor+c
Vf((i)
1/
v((0)) v)(2))
(10)
V +c
B I-A, m/
+
1-9, Zr-1 + Eg
+ LE
||D
V +
C22
mP
E (0)))) +
"
[((x))]1
E[[
2c
2p
Vt
(8)
nP
[| +
312
T- m61
+
1-
4
Z, + &
27
GG
E
2c
2p
Vū--t +c
Ve
LLL [[)
CFC2 2p
LLC D+)l:
(11)
2c
Vc
where (a) due to Young's inequality and (b) is based on Assumption Now integrating (5), (6), (8), (11) into (4),
24

---

1 S Co =
II S G P+]11]+ C2 B11D.1111
Vc
111 & L =1
+ +GG
) LG
181, C2G2 I
IV & 11P 11
2p + Ve
Blll
2
and taking the telescoping summation over / 1, T, we obtain
E ( - f(x))
$ (-2)
17 1
+
NP
T
V 1(0+))
Lop +GG
TT LTGG2
Co
2e
2pe
t=1
C1
[1+
+
BlIpPih| / CPLG
Σ
E D.1l
Ve
t=1
Setting 7) S 4ico and choosing p = 20% we further arrive at
TTP + CGG
B((())))
Σ B|I
2Co
Tn T
n(L+
I + CC
Ve
where the inequality follows from Lemma 7. Re-arranging terms, we get that
EE [(
nL(o² +CGG9) [[f((1)) - f((T+1)] Tn (1+ GI)GG 71 2C)+
T PL'COOT02
$ 2C
+ 20%
TVe
Te
[(())- + nL C3G Tp LL
S 20
"Tn
m(1 201+ (1+
+ 2Ch
( TVe
Te
where in the last inequality used 21 0 and the lower bound f ((0) for all 0 Rd
25

---

To get the rate mentioned in the main part, choose min ILC'' vT and continue
V f(0))11
[axx {1, 4LCg f(01) f(01)- -f +
Lio +GG9)
L'G
$ 2C
cVT) VT
CVT
eT
(1 (1++G) +
(1+CG)GGGG
+ 20
TVe
CT3/2
f(01) ((01)-
f. -f,
Lio CC2
$ 2Co
VT
cVT
(ALCo f(01) f +
L'
(1 (1+
(1 + 20))
+ 2Co
T
eT
TVe
TT3/2
(f(()) fr
Lio
+
(G(G+d))
2Co
+
VT
c/T
T
where in the second part of the rate we suppressed all the problem and compression dependent constants.
E.3 Analysis Under PL Condition
As in the non-convex analysis, here we derive the convergence rate with fixed step-size The rate shown in the main part can obtained by plugging the expression 7).
Theorem 4. (Convergence rate under PL) Let Assumptions 1, 2, 3, 4, 5 and 6 hold, and qu 1. Then, choosing any step-size 1j S 4LC' MICROADAM (Algorithm 3) satisfies
-[(00+1)1-+- s (1 / (S(0)) f*) +
LCog" LGo + (1+0))
Ve
+n²
(3I2C,C302 24e3/2
+ (1+20))0,
LCR0²
2e
Proof. We start from descent lemma
E[f(#++1)] - f(x2)
< nn ( (aza)
9t + Ge
) n
31 Di(m/_1 + Zc-1) + DrE:)
Vf(m)))
Vor +
1-
T E
+ Ge
B1
12
De(ml 21-1) D&E
2
Vor+ +e
1-31
3
=nnE
(f(a)),
+E
Vf(x))
De(m/ + Z,-1) + DIES)
Vo, +
1-
2
gt +e
11
+ 2
E
De(m, + Z,-1) D,E,
(12)
V0, +c
1 1- 3
111
26

---

We bound part II and part III in the same way as it was done in the non-convex analysis. We now provide bound for part I':
1' = 7/
V ((4)
9t
V0 12
11 ./B (( (VI(E)).
9t
1 (Vi+
n/
V /i--1+ +
V +
Vf(x)); Vf(xz) + +
Vo-1te Va-ite
n/B
nB
Vi(a)). 9t
VoL+e V +E
1
)
-nE
(
- vf((z) -VV(())))
}
n)B
Vi(I)), Vo--tt/
V +
( (Tr(tz)) (ViLte v
n/
We further expand and bound this equation as follows:
1'
Co E V (xx)))
Vi(x) (1000+V/C0
Vo-1+6 (VS(0.) v.((a)))
Voite Vdt-1 te .a)l
V((x)) v((0)) VF(O),
E
Co
Vo--++e +e (V/(0.) - V((z)))
v((m)) v.((.),
v((e). 1
(Vi(0.) - v((ii))
7B
Vo--1+c
mB
V((x)) - vs(00),
gt
Vor+ + Vu.-1 +
((O)). Vor+e Vo-1++ +e 81)
[+x))11 + LE [V]((x)) - 1(0))11)
$
Co
E
Ve
V.S(O)) Vo--+t (Vf(0)) - v((=.)) }
(V)(m)) VS(O.), Dage)) B (V (0.), D.g:))
Next, we use the to bound fxx) - f(())) < L||a, - dell < y.C.@. and the incquality -|[a||2 < =
inequality to bound inner products above, L-smoothness inequality
Ve
27

---

||a - bll for the first term:
I's co E [= J/ [ - f(00))) + G BII (z
+ + [[] (3)+ - 11(0))1|| |. + [1
7 TTA PTC.GGG
S 2 E [V]
2C, |
2C,
V
+n G
E D1 + B [1]1
S LE E [IV 4C
20.
2C,
LPG + L2
+ LLLC
312
Ve
E [IV](E)))]
" B1T
<
2Co
4Co
L11L22 L2
L2
[[Il + T (ID
2/3/2
Ve
Plugging the obtained bound for I' with previously obtained bounds for 1I and TTI
CCFLG P1Ib1
IT S T 11
Ve
III s n'4 =11
L/ + C
into (12) and using the step-size bound 1) S aLCo we get
[[f(+++1] B[f(xt)]
" E [IVS(E))))
" (
2C,
4Co
3n GTG
+ LC LC
+ 263/2
Ve E [[
+ n G
Drlli+ B
VE
+ 2
TT B[]
C3G2) 1
'
1 PTC + C3)
"
[IV
2Co
+ 111 + [II 3T
2/3/2
(((+))) D + + GIGGEIIDDI
+
Ve
< TM (F])(z)) l*) + Fooo L + C)
Co
+ + + (()
E D+lh]]
Ve
LTCF2
+
D1) D.
3/3
28

---

where in the last inequality we applied PL condition from Assumption 6. After some reshufiling of the terms, we obtain the following recursion:
(1 (1- T/ C
(B[f(x))) - 8*) + L2L0² + L + G3 3n L2
([f(+++1) f s
(1-
263/2
+ n1 + CIJC C D111
Ve
+ B
Notice that 17 S ALCo 1 arrive
Co so that the coefficient 1 C. E (0, 1). Unrolling the recursion, we Au
T (BIS(zi)) - (r)
mu -0((1110
[[[(T++)]-S-
Go
Lo
1PLIC + C3)
Σ
T
7/
+
203/2
Co
+ (1 + C1
nH 1
BII
Co
7/p1 )
B[I
Vc
T
+
+
Σ
1
mi Co
D
(13)
For the second sum above we upper bound it by its infinite sum as
Σ
Tp C.
Co 7p
Co
0
For the other three sums we bound
10 Co
S
and apply the bounds in Lemma 7:
m
B D
11]
Blll
d Ve
Co
mu
E
D
B||l
D.
d |. ē
Co
Plugging all this bounds into (13) and noticing that #1 0,, we finally get
T
E[f(xT+1)] - < 1
TH
(f(0)) - fM)
C
+ Co
17 LO
LL + C3)G2 2
2/1/2
mp
m(1 + Gd PlLCG)CLLGR
P
Ve
nM T
-10
Co
LC + (Co(c+C3)GG (1+G)
+1 17
Ve
Me
/3T2C
+ (1+G)LLG CILG2d
+ +1²
2/163/2
29

---

The obtained rate above is with respect to the virtual iterates 24 that we defined for the purposes of analysis. To convert this rate with respect the iterates 0, of the algorithm, we apply L-smoothness to bound the functional difference:
- - - 0.ll T/ (f(z) - f(0.)) = (Vf((0)) 01)1 +
+ LCGGGG
Ve
which implies
m! T
[[f(Or+1)] - -
%)
(1(01)-1)
LCoo²
+ColCi+030 + (1+C C.CaA
+1
Ve
+
Vi
pec
μe
(3T2COPFC (1+G)NGG + CILGd 2/c
+ LCG02
2
and completes the proof.
To get the rate mentioned in the main part of the paper, we plug in the expression 2C, log and collect higher order terms.
min
TLC
pf
1
S4) 4T.
B[](Or+1)] l < max
(1
(1(01) - 1*)
T
Tog 2Co
[ICoo2 +
TCO(CI C2IG² (1+G)09 G Ve
T 1
He
μe
Ve
log" T 4CB
CLL'COC62 (1+G))LGG CLLGd LCGG² +
7
122
24e3/2
2e
2 log'T' LCB 0² + (CI+C33 C2) Co(1+0)(11 + d)c²)
(
+ d)
T
AVE
T2
E,4 Non-convex Analysis with Weight Decay
Algorithm 4 MICROADAMW (MICROADAM with Weight Decay)
1: Input: parameters B1, 82 E (0, 1), >> step-size 0, 01 GRd. e mg vo 0p 0,a
2: for = (1,2,..., 7} do
3: g: Vof(O:)
Compute unbiased stochastic gradient Add accumulated error e, and compress Update and compress the error o Update first-order gradient moment Update second-order gradient moment Update the model parameters with weight decay
4:
St Clg: + Eq)
5:
et+1 Qlei gt - 9t)
6:
m = B1m-11 + (1 31 )9:
7:
Vt B2U1-1 + (1 B2) -2
8:
0.+1 = (1 - 7)2, 7 me
Vve=e
9: end for
Lemma 8. Under Assumptions 1-5, for all iterates of Algorithm 4 we have
E [|m f((.))) < 1 - 91 [1|m-1-- ff(;-=))11]
+ B1
L [110, 0
+ + (((0)) - -
30

---

Proof. We start our proof from
E [[]m= V/((0))] =B
+ Big -
[ (1-
(1 -JV
B [1( -1
(1 -
Using Jensen's inequality we have
E mm V ((() S(1 - -1)1 - -V
+ 3 [1. ff(
Using Young's inequality we have
[( E [[]m= V.(())) - 31)(1 + +b)E -
+ -1-B1) [ (1+5) E| E|V - 1(00)1?
+ 3 (0.)) St 11?
Setting 6 01 we have (1 -11) (1+
- 말 and (1-B1 (1 + 3/ s 3
81
[ (f(())]1] <
1
E ||mm-1 (8(()-))1
2
2
+
E (() - s(0))1 3 (.((O)) 3.11"
31
Combining this bound with L-smoothness we obtain
(1
81
E [1]m ((())) <
E [|mm-11 (S(00-1)1
2
2
+ B1
L 12 0cl1 + 3, ( 3.
Lemma 9, Under Assumptions for all iterates of Algorithm 4 we have
E ((V)
[= < 9E [le: [1 6G +302
Proof. We start from
E [IV]
3e11
= [|[((rr + gt) [[(0))]
=E [[[((er + gt) (e: + gt) + (e+ + gt) (0)))
33E [[[C(e; + g.) - (e, + |
+ 3E [[le: |2| + 3E [ f(0,) - 11.
Using definition of contractive compressor we have
E [[V (0,)
3u11 < 3 [|le: + gt ||
3E [| [| + 3E
[|] ||I f0t) 9.
Using Young's inequality we have
E [[l.)
3.11
< 2
|||e 2 + 2 'E [||a ||
3EE
|||e
||
3E
( 9e [12]
s 9E
[||,,11
+
6G2 +30
31

---

Lemma 10. Under Assumptions 1-5, for all iterates of Algorithm 4 we have
E ||e S (1 4) G2
(1 - (1 w)
Proof. Using definition of contractive compressor we have
E [[lee+1 ||]] EE | (et + 9r Clei + 9c))]
2 (1 + q [lle + 9i C(es + g)ll
Using Young's inequality we have
_ E
E eet1 |2] & (1
+ + w1+ (L+o) [1 a
[II9e1]]]
We need satisfy the following condition:
(1 (1+))
w)q.
It holds for 0 S w. 5 < 1 and a a= [1+wyr 1. Using this parameters we have
1
E [llet2113
$ (1 + w ||| ||| + (1 w
1 (1 (()g))))))))))))
119.11
(1 + 92
$ (1 w [|| +
EE [[]se]]]]
1 (1 w)q
Unrolling this recursion allows us to obtain
w1+ ((1 (1++)) (1 + 22
E [[|e+11] ? ((1 +))g) E ||| ||
1- (1+ w)q G²
(1+ G² t=1
((1+w)q)* |||e
+1 w1+
(1 w) G2,
$ (1 (1 (1+ )
last inequality holds because é1 =
Lemma 11. Under Assumptions 1-5, for all iterates of Algorithm 4 we have
IE
|1 [2] < 4(1 + 92)
4(1
(1
(1 + w
G²
(1- (1+w)g²
Proof.
E
||
[²]
E Ic (g: + ex)]|
E CC (9) + te) (0t + ex) + (a + ec]l
2E ||| (gr + ex) (gr - -=]ill]] 2B [1.. ea
2(1 + Q²) |19 e&
Using Lemma 10 obtain
IF [| [²]
4(1 4411+ + q2) (1+
(1- (1+ (1 w) 6².
(1 w2+
32

---

Lemma 12 (From paper: Zhou et al. [2024b]). Let us consider the update rule:
11
D2+1 = (1 - -1))0+ T Vu+
For brevity, we denote Ty Vue + c. Also we define
Ut == my 20, @9 Uz,
where & denotes element-wise product. Moreover, we also define J (0c) as follows:
(0.)) = f(O+) + Ael10.11..
where x = A Eim (120)² for > 0, Ao = 0 with 0 < q < 1 and ||0elle. = V10c, de @ 01). Also 2 let C1 & 10011100 < C22 then iterates of Algorithm satisfy
((0)) S 7 (0.-1)
+
T |V (0.-1) - m, m,
201
VJ (0.-
T1
||ue-
202
402
Lemma 13 (From paper: Zhou al. [2024b]). Assume that C2 S -1leel 00 S Coo. then we have
(Ut + +) (1e+1 E)P
mo Coo.
|10, c t
E 14, 1 (Vp (0, 1)),
Agc
where 1
02,o0tc tc
Theorem 5. Let Assumptions I to hold. Define D = f(O,) +
"
[[]m= - V With
2013
T 7 91c1
Algorithm 4 satisfies
20g"
(0.1))))
||
2
202 T/T
yo
T
(11
(0
(1 w)
+
22
G² 6G2 + 302
(1- (1)
Proof.
We start from main lemma and lemma for momentum, summing inequalities together we
obtain
F(0.) + VE [[|m,
f(0))l11 s $ F(.-1)) +
me
|V (0(-1) -12
2e1
VJ (0c-1)
1/t
|/ 112
20²
402
31
+
1
16
|||m--1----- (((,--))11)
2
+ v
I
- 0,11 ( 3.
81
Using previous lemmas have f (0.) + VE [[|m+ - (0.))
< 7 (0c-1) +
Stt
(T - m 2
2c1
v (0,
1)
|| 11?
202
40²
31
+
E
[[|m--1
(f(()-1))1]
2
+ V
1 1.2B|(0c11 00,12
11
(1 (1+w)3g²
+ VBi
9
01+
1+
6² + +62 (6002002) +302) (1 - (1+ (1-11+
33

---

Using 8, -Br-1 -17 we_1
we have
2V 2V1221
f(0.) + -VE [||me VI((O)11) (00-1)
91
||ue-112
402
3Ic
V (0 (0._1)
202
+
31
+ 1t
IE
-
2
201
(1 + 92
+ VBi
9
1+
G +G2 + +30²)
(1
ww))
Using V =
R1 21
and v = f(Or) + 2013.
E
[[|m,
[f(e))]]] we have
V S 1
71
4 402L2111
||
402
870
V (0+ (0/-1) l|?
202
1+ (1+
(1 w))
+ U11
- 66 + 302
2c1
- (1+ (1)))
Using ? 5
9101
we have
2L
2cz
1/t 202
vJ (0--1) 2
V < V
(1 11+w7 22
+
7
1+
- + 302
201
(1 (1+))g)²
Summing from = 1 to T we have
VJ (0()))
s
20
C2
9 (1
1+
(1 w1+2)2g²
O2 6G2 +302
nT
ci
(1 (1- (1=
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 ( and the second term is non- oon-iinihing O (31 a a²²²) In
our result, the non-vanishing term is proportional to O (o" G2)
A key difference is that our term is not proportional 81. It is important note that &1 typically takes a value close to in practical applications, meaning its influence on the bound is minimal. Therefore, even though our result does directly involve 31, the impact on the overall bound significantly different.
Moreover, our bound includes additional term proportional to G², 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 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 G². Despite these differences, the practical performance remains largely unaffected, ensuring that the bound remains robust and applicable in a variety of scenarios.
Error Feedback applied to GaLore
F.1 Behaviour of the Error Feedback Mechanism
The Gal oore low-rank updates introduced by [Zhao et al., 2024] enable the compression of optimizer states by performing learning updates on lower-dimensional subspace. In this approach, the
34

---

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 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 result, when the error is passed 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 better understanding, we derive analytical evidence for the described behaviour by induction. Let L be fixed learning subspace and assume that et-1 L Then, gradient passed to the optimizer at step is: CGoLore (ot) = (atj)((tt) = proju (roji(Et-11 + gt) = proju (g:) where error feedback is discarded. Thus, 6 = 0r CGaLore (at) = ct-1 + It proju(g) = ct-1 + proje (gt) E T,- which completes the induction.
Assume now that learning subspace is updated every I' steps, and denote Le the learning subspace at step t. Then, similar induction leads to:
L+J o
e, proji (g.)
Σ
projet1
(9.)
i=l+
i=1
KT-1
о&T = projeer (gKT + exT-1) = pRoj (9KT)
proj o ( j= projey)())
F.2 Consequences on Training
Such behaviour of the error feedback mechanism results in the dominance of the error norm over the gradient norm. Before learning subspace updates, the error is the sum over past gradient components that belong to the orthogonal of the current learning subspaces. Since these components represent descent directions that were not used, they are not expected to compensate each other on average. Consequently, between learning subspace updates, the error norm is expected to grow linearly. Figure provides evidence of such linear growth of the error norm during fine-tuning of RoBERTa- base model GLUE/MNLI task.
implies that known analysis techniques [Alistarh et al., 2018, Karimireddy et al., 2019] of con- vergence 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 eтror be bounded. Given fixed vector, low-rank compression based on its singular value decomposition is a contraction singular value
operator. However, our case, the compression is based on a decomposition and therefore may not be 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 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
35

---

Figure 8: Dynamics of the norm of the compared to of the gradient (of output of the 3rd attention layer) during fine-tuning of RoBERTa- base model on GLUE/MNLI from surrogate GaLore with error feedback optimizer. We used hyperparameters from [Zhao al., 2024], i.e. batch size 16, learning rate 0.00001, projection update gap 200, rank 4 and GaLore scale 4.
nigh frequency dynamics Or-1
Co-a
11
122
1e5
Tel-T
1e5+2T
3ed-3T
Tel+eT 185+ST 3e5-5T Te5+TT Ia5-dT led-or Thaningg step
high frequency
50
23
1c5
465
7et
Talning 5382
convergence by enforcing space exploration. See Figure for examples of how induced orthogonality of successive learning subspaces affects the learning trajectory.
Figure 9: Optimization trajectory for Adam, Arda-mmda and Gal Lore-Adam- for ill-conditioned function (I, cos( x) sin( 4 y) starting from (Io, Vo) (- 1 1) (on first row) and for
Rosenbrock function starting from (Io, yo)
(- 2, 1) (on second row).
Adam
GaLore Adam
GaLore Adam with EF
GaLore Adam
GaLore
32
18-
10
16
1.2
2.2
C6
60
-38
-38
36

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the
paper's contributions and scope?
Answer: [Yes]
Justification: In the abstract and introduction we introduce prior work on memory efficient optimization and claim that work improves the memory usage while preserving the performance. We provide theoretical and experimental justification for our algorithm that reflect the claims in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of our work in Section 6.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
37

---

Answer: [Yes]
Justification: We include brief theoretical justifications for our method in Section 4 and include the complete proofs in the Appendix E.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proof' in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We explain our experiments in Section 5 and provide the complete set of in Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
38

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
We provide a zip file that contains the code for our optimizer.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/ sttee//uuiiee/ooettniiiiiiiiiiii
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ //wpps.ccppubiii wttpccccpppppcpcceeeeeeecc hugpbbiiigguuieeeddSeSuuuuiiiiii
/Cooeeuboeeeeeeeeeee wewibbbmisiioooo ssionnollll mnPoicyy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide information about the training details in Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: We do not provide error bars, but instead explain how report the results in Appendix B. Concretely, we run the same experiment with three different seeds and report the one with best performance.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
39

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: All these information can found in Section 5.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't It make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics stctps//neurips wwwpeerrrps.cpppupppeeeeeeeeeeee httppullt/ttissuiiiiiieeeaeo
Answer: [Yes]
Justification: Yes, the research conducted in this paper conforms, in every respect, with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the broader impact of our work in Section 6.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
40

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 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 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).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
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 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 usage guidelines 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 best faith effort.
12. Licenses for existing assets
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?
Question:
Answer: [Yes]
Justification: We work with open source models that publicly available and cited them properly.
Guidelines:
The answer NA means that the paper does 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 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.
41

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasttssss has curated licenses for some datasets. Their licensing guide can help determine the license of 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 encouraged to reach out the asset's s creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We provide code and instructions how the code in order to reproduce our results.
Guidelines:
The answer NA means that the paper does not release 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.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects/
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
42

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
43


Reviews:
[{"review_id": "pZP5cD7n6O", "reviewer": "Reviewer_2Efz", "summary": "The script proposed a memory-efficient method with a theoretical guarantee.", "questions": "See below.", "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:\n\n**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. \n\n **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. \n\n **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.\n\n\n\n\n\n**Major concern 2: The presentation is quite poor.** For instance:\n\n1. There is no explanation on the principle or idea behind the design of Algorithm 1. \n\n2. 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.\n\n3. The current logic flow of the script is confusing. It would be much easier to read if the authors:\n\n   (i) introduce Algorithm 4 first (generic form of MicroAdam), \n\n   (ii) provide the theoretical guarantee, \n\n   (iii)then introduce Algorithm 1( a  specific form that you chose) , \n\n   (iv) discuss the implementation details,\n\n   (v) then show  the experiments.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 1, "strengths": "The topic on memory-efficient optimizers is important. I did not observe obvious flaws in the theory.", "weaknesses": "See below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the feedback. We address your questions below.\n\n**Limitation 1-1:** Contribution and novelty in MicroAdam\n\n\nWe focus on reducing the memory cost of adaptive optimization, and start from the idea that not all gradient components carry useful information for optimization. Thus, we show that gradients can be sparsified **before used to compute the states in the Adam optimizer**. Intuitively, sparsity leads to significant savings in terms of the size of optimizer states. \nHowever, to ensure convergence both in theory and practice, it is necessary to incorporate an error correction mechanism **error feedback** (EF), made popular in distributed optimization. **Simply using EF does not lead to memory savings**, since the size of the error correction buffer is the same as the model size. At the same time, EF is critical for convergence, as also illustrated in Figure 1. \n\nInstead, **our key algorithmic innovation is in showing that the EF can itself be compressed** in the context of adaptive optimization, leading to significant memory savings while preserving convergence. The new parts of our algorithm, i.e. EF compression and decompression are highlighted **in blue** in Algorithm 1. Our practical contribution is in showing that sparsity in optimizer states can be efficiently leveraged for memory savings, at the scale of billion-parameter models.\n\n**Limitation 1-2:** Pretraining\n\nBecause of limited size for the number of characters in our response, we are kindly asking to read our comment about pretraining in our response to **Weakness 1** of **Reviewer u3Hk**.\n\n**Limitation 1-3:** Novelty in the theory\n\nMotivated by the question of reducing storage cost, we are the first to consider **EF quantization** and to provide a convergence proof for it. \n\nThe key point in the algorithm design, which allows us to obtain practical gains, is the ability to compress the EF accumulator through quantization. In the analytical view of MicroAdam, presented as Algorithm 4, the EF accumulator is denoted by $e_t$, and in line 5 this accumulator is updated and compressed via $e_{t+1} = Q(e_t + g_t - \\tilde{g}_t)$.  As can be seen from examining our proof, quantizing the EF, although very simple algorithmically, significantly complicates the analysis. (Generally, we are not aware of any method that compresses the error accumulator and provides theoretical convergence guarantees in any setup. ) \nMoreover, from the practical side, since the state we maintain are the quantized EF and the sparse gradients, our algorithm must re-compute optimizer states on-the-fly in every iteration based on this compressed state. Implementing this efficiently is a non-trivial challenge.\n\n\n**Major Concern-2:**\n\n**Algorithm 1**:\n\nWe explain Algorithm 1 in detail in Section 3 of our paper. Please check the answer for **Reviewer 8ZAo** where we extend the explanations we already provided in section 3 of our paper to make sure our algorithm is well understood.\n\n**About Figure 1:**\n\nThe motivation for this illustration is to show that some form of EF is **necessary** for good convergence, even in the case of toy instances. Specifically, prior heuristic methods, such as GaLoRE, perform gradient compression but simply **omit EF altogether**, and we wanted to show that this is infeasible in general. We agree with your point that various other heuristics could be applied, and we will clarify the motivation for this example or move it to the Appendix. \n\nIn Figure 1 we show how EF fixes AdamW with Top-K compression. The left plot shows the optimization trajectory of the original Adam optimizer. The center plot shows the convergence of Top-K Adam when we only choose the largest coordinate from the accumulator (equivalent to 50% sparsity, since the problem is two-dimensional). In 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.\n\n**Major Concern-3:** Sections order in the paper\n\nThank you for this suggestion; we will re-order the presentation to introduce the general algorithm, followed by the implementable version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hVaDFrB0JT", "reviewer": "Reviewer_8ZAo", "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. \n\nThe 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.", "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.\n\nIt 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": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The main idea is quite interesting - that most of the values of a gradient are small in magnitude, so need not be used to update the parameters --- only up to 10% of all the parameters are ever updated in any step, and these change only a little across steps. \n\nThe convergence results are a useful validation of the algorithm. The experimental results seem quite strong to me.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the feedback. We address your questions below.\n\n**About Summary:** Memory savings\n\nThe memory savings of MicroAdam are ~20% if we compare the entire memory usage, but the key comparison point should be on the size of optimizer states, as this is the quantity we are minimizing. We restate the memory usages for optimizer states presented in section 3.2 of our paper for Llama-2-7B model:\n- AdamW: 25.10 GB\n- AdamW-8bit: 12.55 GB\n- MicroAdam: **_5.65 GB_**\n\nIn our next revision, we will show the optimizer state in the tables containing results to clearly differentiate from the memory usage that is shared across all optimizers (e.g. model state, such as activations and gradient). As explained in section 3.2 of our paper, the memory footprint of MicroAdam is $0.9d$ bytes in any experiment, where $d$ is the model size, while AdamW and AdamW-8bit have $4d$ and $2d$, respectively. To express the memory savings in percentages, MicroAdam uses only 45% memory compared to AdamW-8bit ($0.9d$ vs $2d$) and 22.5% compared to AdamW ($0.9d$ vs $4d$).\n\n**Weakness 1:** Paper clarity\n\n**Algorithm 1:** We explain Algorithm 1 in detail in Section 3 of our paper. Please check the end of this answer where we extend the explanations to make sure our algorithm is well understood.\n\n**About $k$:** Indeed, $k$ is one of the inputs of the Algorithm 1 and  $\\delta_1$, $\\Delta_1$ should be initialized as scalars. We will fix these in the next revision.\n\n**About error correction:** Let us clarify the main steps of the proposed MicroAdam algorithm and address some of the questions you raised. As you noted, Adam stores two additional values for each parameter: the exponential moving average of the gradients ($m_t$) and the squared gradients ($v_t$). To avoid allocating twice the model size for $m_t$  and $v_t$, we compress the gradients using a sparsity-inducing operator Top-K, storing the most important gradient components at each step, for several past steps (these are stored in the buffer $\\mathcal{G}$). However, we do not ignore the smaller gradients, but we store and accumulate them at 4-bit resolution, adding them to the next gradient to compute the model updates. Intuitively, instead of disregarding small gradient components in each iteration, we accumulate them and apply them in the next step. This ensures that all gradient information is eventually used to update the model parameters. Below we provide a detailed explanation of Algorithm 1.\n\n**Question 1:**\n\nThe idea of our algorithm is to store 99% sparse gradients (e.g. only the largest 1% of values) into the buffer matrix $\\mathcal{G}$. Then, this sparse gradient window is used to “materialize” the Adam statistics $m_t$ and $v_t$. This way, we make sure that we use only $0.9d$ bytes of memory instead of $2d$ as in AdamW-8bit or even $4d$ as in original AdamW. Below we provide a detailed explanation of Algorithm 1.\n\n**Question 2:**\n\nIn all our experiments we report the memory usage for the entire process. We would like to emphasize that what should be compared is the size of optimizer states, since the settings for each experiment is the same and only the optimizer differs (e.g. batch size and model sizes are fixed). In section 3.2 we provide a detailed theoretical memory analysis for the optimizer states of AdamW (25.1GB), AdamW-8bit (12.5GB), **MicroAdam (5.65 GB)**. In the next revision we will clearly make a distinction between the memory usage for the optimizer state and the overall memory usage (which includes, among others, model activations and gradient). Please let us know if we misunderstood your question. We are happy to discuss and clarify everything about our work.\n\n\n**_Explanation for Algorithm 1_**: line $n$ refers to $n^{th}$ line in this algorithm:\n\n**Optimization step $t=1$:**\nIn this case, the error feedback is completely zero as initialized in line 2\n- **Line 4**: the accumulator $a_t$ contains only the stochastic gradient\n- **Line 5**: if we suppose the accumulator $a_t$ is normally distributed with a mean of zero, then this line is equivalent to choosing $k=1\\\\%$ values from the tails (outliers) because we apply Top-K on the absolute values of the accumulator $a_t$. The Top-K operator returns indices of those outliers, which we store in $\\mathcal{I}_t$, as well as the corresponding values $\\mathcal{V}_t$ from $a_t$ **(not from $|a_t|$)**\n- **Line 6**: the outliers are erased from the accumulator $a_t$ because they will be transferred to the buffer matrix $\\mathcal{G}$\n- **Line 7**: compute $\\delta$ and $\\Delta$ (the min and max values from the accumulator $a_t$) which will be used for the quantization at the next step.\n- **Line 8**: quantize the remaining values in the accumulator $a_t$ **without the outliers in the tails** using procedure $Q$ in Alg. 2. **This error will be dequantized and added to the stochastic gradient at the next optimization step $t+1$**. This is the point in our algorithm when we preserve the error instead of discarding it.\n- **Line 9**: add the new set of outliers to the ring buffer $\\mathcal{G}$.\n- **Lines 10 and 11**: compute the first and second order moment statistics $\\hat{m_t}$ and $\\hat{v_t}$ over the ring buffer $\\mathcal{G}$ using Alg. 3.\n- **Line 12**: update the model parameters as usually done in Adam\n\n**Optimization step $t \\geq 2$:**\nThe only change compared to optimization step $t=1$ is that the error feedback $e_t$ is not zero anymore, but will contain the quantized values not selected by Top-K (inliers) and will be added to the gradient. After that, the algorithm works exactly as explained above, starting with line 5.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RvhOAjCOKw", "reviewer": "Reviewer_nVyV", "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.", "questions": "See Weaknesses", "limitations": "See Weaknesses", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- **Theoretical Analysis**: The paper includes comprehensive theoretical analyses of memory consumption and the guarantee of convergence.\n- **Experimental Validation**: The experimental results showcase the potential of MicroAdam in reducing memory consumption, with detailed comparisons to existing methods.", "weaknesses": "- **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.\n\n- **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.\n\n- **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.\n\n- **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.\n\n- **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.\n\n- **Computational Efficiency**: MicroAdam appears to be less computationally efficient. Are there ways to improve its computational efficiency?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the feedback. We address your questions below.\n\n**Weakness 1:** About Figure 1\n\nThe motivation for this illustration is to show that some form of EF is _necessary_ for good convergence, even in the case of toy instances. Specifically, prior heuristic methods, such as GaLoRE, perform gradient compression but simply _omit EF altogether_, and we wanted to show that this is infeasible in general. We agree with your point that various other heuristics could be applied, and we will clarify the motivation for this example or move it to the Appendix. \n\nIn Figure 1 we show how error feedback (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 two-dimensional). 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.\n\n**Weakness 2:** Convergence speed\n\nThe AMSGrad optimizer was introduced as a solution to a fundamental issue in Adam optimizer’s proof, where a state matrix was mistakenly supposed to be positive definite throughout the training (see Section 3 from Reddi et al. 2019 [1]). In Section 5 of [1], the authors provide a one dimensional problem where Adam converges to the worst possible solution. Despite this clear evidence of unsatisfactory solution in simple settings, Adam (and the newer version AdamW [2]) is still the off-the-shelf optimizer when it comes to LLMs finetuning. We chose to build our theoretical analysis on the framework of AMSGrad since this algorithm fixes the above technical issue, and can therefore have a convergence proof (as opposed to Adam). Moreover, CompAMS introduces EF and provides analysis for AMSGrad in a distributed optimization setting. In this context, our work focuses on space optimizations, and introduces compression for EF in the CompAMS framework while preserving convergence.\n\n**Weakness 3:** Memory consumption\n\nWe agree with the example you gave in your comment and we would like to add a few details about our experimental work. We emphasize that the memory usage we reported in the paper was the maximal memory usage, read directly from the GPU. This means that a user must have at least that much memory available on the GPU to perform the experiment. Our memory usage is the overall memory used by the entire program, including model activation and gradient, as well as the batch of data used for forward pass and the optimizer state. In section 3.2 we provide the memory usage only for the optimizer state expressed in bytes ($0.9d$ for MicroAdam, $2d$ for AdamW-8bit and $4d$ for AdamW) and this is the same for any model with $d$ parameters. Concretely, for Llama-2-7b, we have 25.1GB for AdamW, 12.5GB for AdamW-8bit and **5.65 GB for MicroAdam**.\n\nIn practical terms, our experiments consider settings where the memory footprint for orthogonal components (such as activation sizes) is minimized, such as  microbatch size 1 and global batch size 32 for Llama-2-7B. One can have lower memory usage with exactly the same model settings only if the optimizer state is lower and this is what we try to do with MicroAdam.\n\n**Weakness 4:** Top-K Memory Efficiency\n\nThe idea of MicroAdam is to store highly sparse gradients in the buffer $\\mathcal{G}$ and to compute the statistics $m_t$ and $v_t$ on the fly. In practice, we show that very high (at least 99% sparsity) can be induced in the gradients, and this allows us to have a much lower memory usage in practice. It is true that for moderate (50%) sparsity the savings would be minimal or negative, but one of our sources of novelty is in showing that much higher sparsity can be supported in practice, with respect to the information used to generate optimizer states. We can provide more details about our efficient implementation in the discussions after the rebuttal.\n\n**Weakness 5:** Pretraining\n\nBecause of limited size for the number of characters in our response, we are kindly asking to read our comment about pretraining in our response to **Weakness 1** of **Reviewer u3Hk**.\n\n**Weaknesss 6:** Computational Efficiency\n\nTo speed up the main operations at the core of MicroAdam, we already implemented CUDA kernels for quantization and to efficiently compute the statistics $m_t$ and $v_t$. However, we believe there are several ways to further improve computational efficiency. For instance, we can experiment with changing the number of thread blocks in CUDA kernels for the auxiliary operations, such as setting tensors to zero (line 6) and copying the top-k values from the accumulator $a_t$ at indices $\\mathcal{I}_t$ to the matrix $\\mathcal{V}$ that stores the values. At the moment, we use the maximum possible number of CUDA thread blocks to perform the operation, which might be improved in the context of layer-wise preconditioning. We did not prioritize such optimizations during the preparation of the paper because MicroAdam achieved its main goal, that of reducing memory usage by 55% (wrt AdamW-8bit) and 77.5% (wrt AdamW). \nWe skip some details now due to limited sized response, but we are happy to explain more in the discussions phase.\n\nReferences:\n- [1] **On the convergence of Adam and beyond**, Reddi et al., 2019, available at https://arxiv.org/pdf/1904.09237\n- [2] **DECOUPLED WEIGHT DECAY REGULARIZATION**, Loshchilov and Hutter 2017, available at https://arxiv.org/pdf/1711.05101\n- [3] **ON DISTRIBUTED ADAPTIVE OPTIMIZATION WITH GRADIENT COMPRESSION**, Li et al., 2022, ICLR 2022, available at https://arxiv.org/pdf/2205.05632", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MF4qQUa92F", "reviewer": "Reviewer_C8Q3", "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.", "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 are well-addressed", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This work provides theoretical convergence guarantees under commonly used assumptions. Previous work on memory-efficient Adam often relied more on heuristics.\n- Performance on fine-tuning tasks is comparable to uncompressed AdamW, with lower memory cost.", "weaknesses": "- 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. \n- The approach is limited to fine-tuning tasks, and its effectiveness for LLM pre-training or non-language tasks remains unclear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the feedback. We address your questions below.\n\n**Weakness 1:** Complexity of compressors\n\nWe agree that memory savings depend on the choice of compressors, and that efficient implementations are needed. However, we do have the advantage that both gradient sparsification, quantization, and low-rank compression are well-understood from distributed optimization, and one can build on prior work in the area, as well as good support in popular frameworks such as Pytorch. Specifically, the implementation we provide is not very complex, and can be further optimized along the lines we sketched in the response to **Weaknesses 4 and 6** of **Reviewer nVyV**.\n\n**Weakness 2:** Pretraining\n\nWe would like to emphasize that we designed MicroAdam for the finetuning (FT) use case and our research question was “are all gradient entries useful for fine-tuning?”, which is why we test MicroAdam for FT tasks. \n\n\nWe agree with the reviewer regarding MicroAdam’s convergence in Figure 4 and we can explain this behavior by emphasizing the high sparsity in MicroAdam’s update (which, in our experiments, is around 90%). However, our results show that we can achieve at least the same performance with only 45% memory compared to AdamW-8bit ($0.9d$ bytes for MicroAdam compared to $2d$ bites of AdamW-8bit).\n\nWe address your concerns on effectiveness via two additional experiments:\n\n1. To address the training from scratch / non-language tasks question, we trained ResNet-18 and ResNet-50 on ImageNet using MicroAdam and compared it with SGD, AdamW and AdamW-8bit. We provide results in the response PDF. Briefly, MicroAdam manages to obtain close to 72% accuracy, compared to 70% for SGD, AdamW and AdamW-8bit.\n- For ResNet-18, **MicroAdam uses ~10 MB** memory for the model state, compared to ~22 MB for AdamW-8bit ($2.2\\times$ more), ~45 MB for SGD ($4.5\\times$ more) and ~90 MB for AdamW ($9 \\times$ more).\n- For ResNet-50, **MicroAdam uses ~22 MB** memory for the model state compared to ~49 MB for AdamW-8bit ($2.2 \\times$ more), ~97 MB for SGD ($4.4 \\times$ more) and ~195 MB for AdamW ($8.9 \\times$ more).\n\nThis is an indication that MicroAdam is effective for pretraining model for vision tasks.\n\n2. To address the question of scale, in addition to existing experiments, we would like to emphasize the effectiveness of MicroAdam for SFT by training a Llama 2-13B model on the same GSM-8k dataset and comparing it against AdamW-8bit. We provide results in the response PDF. Briefly, our results show that MicroAdam can also recover the accuracy of AdamW-8bit for this task at a lower memory cost: **MicroAdam allows training on a single Nvidia H100 GPU with 80GB RAM**, the total running process requires 70 GB GPU memory (~59 GB for the model and activations and ~11 GB is used for the optimizer state), while AdamW-8bit (24.2 GB for the optimizer state) requires more than one 80GB GPU and need at least 2 GPUs. The need for more than one GPU for AdamW-8bit comes from the size of optimizer state, which is $2.2 \\times$ larger.\n\nFT tasks require adapting a fraction of parameters (for example, LoRA) to be able to learn a downstream task. In contrast, pretraining (PT) requires updating all parameters in the model because there is a lot more knowledge to be embedded into the model compared to FT case and the dataset sizes are also much larger. In the paper we briefly explain that under the settings we choose for MicroAdam ($k=1\\\\%$ and $m=10$ gradients), the parameter update is at least 90% sparse, meaning that MicroAdam updates at most $10\\\\%$ of parameters in each layer at every optimization step.\n\n**Question 1:** Compressors' properties\n\nThe condition $(1+\\omega) q < 1$ is theoretical and parameters $q$ and $\\omega$ are worst case bounds. In the experiments we apply much more compression than the theoretical worst case condition allows us. In general, Top-K, low-rank and quantization are the main compression operators widely studied theoretically and used in practice.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5lo6jPTpyh", "reviewer": "Reviewer_u3Hk", "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", "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": "The limitations are well discussed.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper proposes a new optimizer based on Adam, called MicroAdam, which reduces memory footprints, via sparisification, quantization, and error feedback. \n2. Theoretical convergence guarantees are provided. \n3. Empirical results show good performance on several fine-tuning tasks\n4. Efficient implementation on GPUs is provided.", "weaknesses": "1. The experiments are limitted to finetuning tasks. Note that GaLore also show good performance in pretraining tasks.\n2. The theoretical analysis is actually based on AMSGrad instead of Adam.\n\n---------------\n\nMy concerns are well addressed according to the author's feedback.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the feedback. We address your questions below.\n\n**Weakness 1:** Pretraining\n\nWe would like to emphasize that we designed MicroAdam for the finetuning (FT) use case and our research question was “are all gradient entries useful for fine-tuning?”, which is why we test MicroAdam for FT tasks. \n\nWe agree with the reviewer regarding MicroAdam’s convergence in Figure 4 and we can explain this behavior by emphasizing the high sparsity in MicroAdam’s update (which, in our experiments, is around 90%). However, our results show that we can achieve at least the same performance with only 45% memory compared to AdamW-8bit ($0.9d$ bytes for MicroAdam compared to $2d$ bites of AdamW-8bit).\n\nWe address your concerns on effectiveness via two additional experiments:\n\n1. To address the training from scratch / non-language tasks question, we trained ResNet-18 and ResNet-50 on ImageNet using MicroAdam and compared it with SGD, AdamW and AdamW-8bit. We provide results in the response PDF. Briefly, MicroAdam manages to obtain close to 72% accuracy, compared to 70% for SGD, AdamW and AdamW-8bit.\n- For ResNet-18, **MicroAdam uses ~10 MB** memory for the model state, compared to ~22 MB for AdamW-8bit ($2.2\\times$ more), ~45 MB for SGD ($4.5\\times$ more) and ~90 MB for AdamW ($9 \\times$ more).\n- For ResNet-50, **MicroAdam uses ~22 MB** memory for the model state compared to ~49 MB for AdamW-8bit ($2.2 \\times$ more), ~97 MB for SGD ($4.4 \\times$ more) and ~195 MB for AdamW ($8.9 \\times$ more).\n\nThis is an indication that MicroAdam is effective for pretraining model for vision tasks.\n\n2. To address the question of scale, in addition to existing experiments, we would like to emphasize the effectiveness of MicroAdam for SFT by training a Llama 2-13B model on the same GSM-8k dataset and comparing it against AdamW-8bit. We provide results in the response PDF. Briefly, our results show that MicroAdam can also recover the accuracy of AdamW-8bit for this task at a lower memory cost: **MicroAdam allows training on a single Nvidia H100 GPU with 80GB RAM**, the total running process requires 70 GB GPU memory (~59 GB for the model and activations and ~11 GB is used for the optimizer state), while AdamW-8bit (24.2 GB for the optimizer state) requires more than one 80GB GPU and need at least 2 GPUs. The need for more than one GPU for AdamW-8bit comes from the size of optimizer state, which is $2.2 \\times$ larger.\n\nFT tasks require adapting a fraction of parameters (for example, LoRA) to be able to learn a downstream task. In contrast, pretraining (PT) requires updating all parameters in the model because there is a lot more knowledge to be embedded into the model compared to FT case and the dataset sizes are also much larger. In the paper we briefly explain that under the settings we choose for MicroAdam ($k=1\\\\%$ and $m=10$ gradients), the parameter update is at least 90% sparse, meaning that MicroAdam updates at most $10\\\\%$ of parameters in each layer at every optimization step.\n\n**Weakness 2:** Theoretical analysis based on AMSGrad\n\nThe AMSGrad optimizer was introduced as a solution to a fundamental issue in Adam optimizer’s proof, where a state matrix was mistakenly supposed to be positive definite throughout the training (see Section 3 from Reddi et al. 2019 [1]). In Section 5 of [1], the authors provide a simple, one dimensional problem where Adam converges to the worst possible solution. Despite this clear evidence of unsatisfactory solution in simple settings, Adam (and the newer version AdamW [2]) is still the off-the-shelf optimizer when it comes to LLMs finetuning. We chose to build our theoretical analysis on the framework of AMSGrad since this algorithm fixes the above technical issue, and can therefore have a convergence proof (as opposed to Adam). Moreover, CompAMS introduces error feedback and provides analysis for AMSGrad in a distributed optimization setting. In this context, our work focuses on optimizing space usage, and introduces compression for error feedback in the CompAMS framework.\n\n**Question 1:** Explaining the better performance\n\nThis behavior was also surprising for us. We believe that compression has a regularization effect by updating only 10% of the weights. Adding noise into training dynamics can both increase or decrease the loss or accuracy. It is not the case that lossy compression or noisy training must have higher loss or lower accuracy. For instance, if that were the case, then larger batch size (and implicitly lower noise in the stochastic gradient) should have implied better accuracy, which is not true as far as we are concerned. We are thinking about this in the context of full batch gradient descent, which does not generalize better than the mini-batch version.\n\nReferences:\n- [1] **On the convergence of Adam and beyond**, Reddi et al., 2019, available at https://arxiv.org/pdf/1904.09237\n- [2] **DECOUPLED WEIGHT DECAY REGULARIZATION**, Loshchilov and Hutter 2017, available at https://arxiv.org/pdf/1711.05101\n- [3] **ON DISTRIBUTED ADAPTIVE OPTIMIZATION WITH GRADIENT COMPRESSION**, Li et al., 2022, ICLR 2022, available at https://arxiv.org/pdf/2205.05632", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank all reviewers for the useful feedback! We have provided individual responses to each review, and briefly summarize the main points here:
- To address the concern about Algorithm 1, we provided a detailed explanation of each line in the algorithm at the end of our response to **Reviewer 8ZAo**.
- To address the concern about the absence of pretraining results, we stated our research question which our work successfully answers (that of reducing the memory usage in the finetuning case), having similar motivation as LoRA [1] (e.g. reducing the memory usage for finetuning, but using a completely different approach). In addition, we also provided results for pretraining Computer Vision tasks, where MicroAdam shows ~2% higher validation accuracy on ResNet-18 and ~5% higher than AdamW variants, while still being 1% better than SGD (see the attached PDF).
- To address the concerns about novelty in our contribution, we emphasized our key algorithmic contributions: compressing the error feedback, providing efficient implementation for GPUs and theoretical guarantees as well.
- To resolve the questions about compressor complexity, compression memory efficiency and computational efficiency, we provided detailed explanations.

We hope our responses address the reviewers’ questions, and would be happy to continue the discussion during the rest of the rebuttal period.

References:
- [1] **LORA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS**, Hu et al., 2021, available at https://arxiv.org/pdf/2106.09685

Meta-review: **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.
Decision: Accept (poster)