Title: MicroAdam: Accurate Adaptive Optimization with Low Space Overhead and Provable Convergence
Venue: NeurIPS 2024
Accepted: yes
Keywords: adaptive optimization

adam

efficiency

memory efficiency
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

Abstract

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

1 Introduction

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

Contributions. In this paper, we address this gap by introducing MICROADAM, an adaptive optimizer which guarantees low memory usage but also ensures provable convergence. We 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

due to compression by employing a novel variant of error correction , a mechanism introduced for distributed optimization [Seide et al., 2014]. However, simply using error feedback would not lead to memory savings, since the size of the error correction buffer is comparable to the that of the original optimizer state. Instead, our main algorithmic innovation is in showing that the error feedback can itself be compressed in the context of adaptive optimization. This renders the memory overhead of error feedback negligible, while preserving convergence guarantees.

Specifically, on the theoretical side, we provide a new analysis showing that, under reasonable assumptions on the loss function being optimized and on the degree of compression, MICROADAM provably guarantees convergence, at asymptotically the same rate as AMSGrad [Zhou et al., 2024a], i.e. a version of Adam with general convergence guarantees, that fixes a fundamental technical issue in the Adam optimizer'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 Polyak-Lojasiewicz (PL) assumption, highlighting a trade-off between the degree of compression of the gradients, and that of the error feedback.

We complement our algorithmic and analytic results with an efficient GPU implementation of MICROADAM, which we validate for fine-tuning language models from the BERT [Devlin et al., 2018], OPT [Zhang et al., 2022] and LLaMA [Touvron et al., 2023] families, with hundreds of millions to billions of parameters. We show that, in practice, gradients can be projected to very high sparsity (99%), while the error correction can be stored at 4 bits per component, without loss of convergence. Concretely, our method can significantly improve upon the memory footprint of the extremely popular 8bit Adam [Dettmers et al., 2021] when fine-tuning models such as LLaMA27B/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].

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

2 Related Work

We mainly focus on related work reducing the cost of optimizer states. Dettmers et al. [2021] considers this problem, specifically by performing fine-grained quantization of the optimizer states. Their work does not alter the Adam algorithm; instead, it deals with the challenge of accurately compressing the dynamically-changing meta-data sequence. As the name suggests, the space savings correspond to roughly halving the memory required by the optimizer states, relative to FP16. In the same vein, AdaFactor [Shazeer and Stern, 2018] and CAME [Luo et al., 2023] reduce memory cost by factorizing the second-order statistics, while the recent GaLore [Zhao et al., 2024] factorizes the gradients themselves before they enter the optimizer state (but does not use error correction). Importantly, these methods are heuristics : they do not provide theoretical guarantees under standard assumptions, 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 a compressed variant of Adagrad [Duchi et al., 2010] called SM3, improving upon earlier work by Spring et al. [2019]. However, it is not clear how to extend their approach to the popular Adam optimizer, and heuristic methods appear to provide superior performance [Luo et al., 2023].

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

More broadly, scaling adaptive or second-order optimizers to large models is a very active area. Works such as GGT [Agarwal et al., 2019], Shampoo [Gupta et al., 2018] and M-FAC [Frantar et al.,

2021] provided quadratic-space algorithms that are still feasible to execute for moderate-sized DNNs, but will not scale for billion-parameter models. Follow-up work such as AdaHessian [Yao et al., 2020], Sophia [Liu et al., 2023], Sketchy [Feinberg et al., 2023] and EFCP [Modoranu et al., 2023], scaled these approaches via additional approximations. Of these, the closest work to ours is EFCP, which uses sparsification plus standard error feedback to compress the gradient window employed in the Fisher approximation of the Hessian. However, EFCP does not compress the error accumulator, assumes a different optimization algorithm (Natural Gradient [Amari, 2016]), lacks convergence guarantees, and does not scale to billion-parameter models.

3 The MICROADAM Algorithm

Notation. We consider a standard Adam-type algorithm, which we will augment for memory savings. We will use f for the loss function, d for the model size, k for the gradient density (sparsity d -k ), θ t and g t for the model parameters and gradient at step t respectively, η t for the learning rate, λ for the weight decay parameter, m t and v t for the first and second moment of the gradient, ϵ for the numerical stability constant, β 1 and β 2 for the momentum coefficients for m t and v t respectively. Furthermore, we use e t for the error feedback (EF) vector, b the number of bits for EF quantization, m for the sliding window size, G = ( I , V ) for the sliding window of size m × k that stores indices I and values V selected by the Top-K operator.

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

In detail, at step t = 1 , the error feedback e 1 is completely zero, as initialized in line 2, and thus, at line 5 the accumulator a 1 will only contain the stochastic gradient g 1 . At line 6, we perform the Top-K compression and only keep the top1% of values V 1 and their corresponding indices I 1 . The compression is equivalent to choosing the top1% values in the left and right tails (outliers) due to the absolute value we apply on top of accumulator a . At line 7, we remove the outliers from the accumulator because they will be transferred to the buffer matrix G . This step is equivalent to e ← a -T k ( a ) found in theoretical results. After line 7, what is left in a is called the error feedback (e.g. the weights which were not chosen by Top-k). At line 8, we compute the statistics δ and ∆ needed for quantization, and at line 9, we effectively quantize the accumulator (e.g. error feedback after line 7). At line 10 we update the buffer G , in lines 11, 12 and 13 we compute the statistics ˆ m, ˆ v (computed by squaring the entries of G element-wise) and update the model parameters.

For steps t ≥ 2 , the only change compared to t = 1 is that error feedback e is not zero anymore. Since the error is compressed, we need to decompress it and add it to the gradient. This process happens at line 5 and it is the point where we feed back the error: the accumulator will store the gradient whose direction is corrected by the error (e.g. the cumulative history of weights not chosen by Top-k at the previous steps).

Properties and Limitations. We would like to point out that the resulting update u t = m t / ( ϵ + √ v t ) will always be highly sparse when the window size m is small. For illustration, if we use density k = d/ 100 (e.g. 1% density equivalent to 99% sparsity) with m = 10 and suppose that all rows in the indices matrix I are disjoint, then the overall density in the update u t will be 90% . The sparsity of u t increases if rows in I have common values. MICROADAM yields good results for LLM finetuning and pre-training computer vision models, as the experimental section shows. However, we noticed the update u t of MICROADAM is too sparse to be able to provide good enough updates for LLM pre-training. We believe this happens because the attention layers must receive dense updates to be able to learn the correlations between words.

Dynamic Statistics. In ADAMSTATS procedure in Algorithm 2 we implement the unrolled recursion of momentum z t ← βz t -1 +(1 -β ) g t for the last m sparse gradients as z t ← (1 -β ) ∑ t i = t -m β t -i g i and we also perform the bias correction in the end. Because we compute ˆ m t and ˆ v t using the last m sparse gradients in the window, in line 4 we dynamically determine the exponent r for the decay factor β r based on the current optimization step t , i th row of the circular buffer G and the window size

m . The last gradient added to G will have r = 0 , while the oldest gradient in G will have r = m -1 . In the end, we will add the values β r V i to the buffer z at the corresponding indices I i , which is a fast operation because we only manipulate 1% of values at a time. We discuss the efficient CUDA implementation in the Appendix.

Algorithm Intuition. To gain intuition, we illustrate the impact of compressing gradients via TopK before they are incorporated into the optimizer state for Adam, both with and without error feedback (EF). Figure 1 shows how EF fixes AdamW with Top-K compression. The plot on the left shows the optimization trajectory of the original Adam optimizer. The center plot illustrates the convergence of Top-K Adam when we only choose the largest coordinate from the accumulator (equivalent to 50% sparsity since the problem is 2D). In the end, on the right side we show that adding EF to Top-K Adam recovers the same optimization trajectory as the original Adam optimizer. Extrapolating to higher dimensional problems, our MICROADAM approach helps recover the trajectory of the original Adam optimizer, while using less memory. The results clearly show that EF is essential for fast convergence. Besides, TopK with EF, which is a surrogate of MICROADAM, allows for competitive convergence relative to the uncompressed baseline. In Appendix F, we discuss the implications of Error Feedback applied to GaLore.

Algorithm 1 Pseudocode for MICROADAM with quantized EF

Algorithm 2 Adam Statistics, Quantization and Inverse Quantization

Adam

Top-kAdam

Top-kAdamwithEF

1.0

102

1.0

102

1.0

102

100

100

100

0.5

0.5

0.5

102

102

102

10-4

10-4

10-4

0.0

0.0

0.0

1.0

0.5

0.0

0.5

1.0

1.0

0.5

0.0

0.5

1.0

1.0

0.5

0.0

0.5

1.0

3.1 Efficient Implementation

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

Accumulator a t . First, we do not use an additional accumulator tensor a t ; instead, we use a CUDA kernel to dequantize the error buffer, and store the result in the grad attribute of the model parameters. This allows us to accumulate the error feedback into the gradients, without allocating a full or half precision d -dimensional array. Each component of the EF has 4 bits and the entire EF is stored in an array of size d/ 2 of uint8 values.

Top-K. Since we run on LLMs with billions of parameters, naive storage of the sparse indices would require using an int64 type for the indices matrix I , assuming that the Top-K operator is applied globally to all the parameters in a t . To avoid this cost, we apply Top-K in blocks of fixed size B d < 2 15 -1 = 32767 and store block-relative indices in int16 format (during the development of MICROADAM, PyTorch did not have support for uint16 ). Applying Top-K per row to a t reshaped to 2D is not only faster, but provides the block-relative indices directly.

Computing Top-K in blocks also allows us to allocate and efficiently use CUDA shared memory blocks to dynamically compute the statistics ˆ m t and ˆ v t for Adam, as described in the ADAMSTATS procedure in Algorithm 2. We allocate the maximum possible shared memory for each thread block and store ˆ m t (first half) and ˆ v t (second half) at consecutive locations in the shared memory. Once these statistics are computed, it is easy to update the model parameters. Note that the block-relative indices returned by Top-K will be directly used as indices in the shared memory array of CUDA thread blocks to retrieve values from I and V .

Quantization metadata. Our approach also stores two additional vectors δ and ∆ used for quantization. Since the quantization block size B q is set to a very large value (e.g. 100 K ), the space required for these two arrays becomes negligible in comparison to the buffer G and error feedback e .

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

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 t and v t , each of size d , expressed in bytes (B). For concreteness, we report the practical memory usage for the optimizer state for a Llama-2 7B model for each optimizer.

AdamW stores states in float32 format (4 B per component), resulting in a total memory footprint of 4 d +4 d = 8 d (B), while using bfloat16 would result in 4 d (B) memory. We will refer to these memory footprints as M AW 32 = 8 d (B) = 50 . 21 (GB) and M AW 16 = 4 d (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 AW 8 = d + d = 2 d (B) = 12 . 55 (GB).

MICROADAM stores the error feedback e in 4-bit format (0.5 B per component) and the sliding window G that stores the indices I in int16 and V in bfloat16 format. Both have m × k components, each requiring 2 B per component. In the end, for m = 10 and k = d/ 100 , the memory footprint is M µA ( m ) = 0 . 5 d +4 mk (B) = 0 . 9 d (B) = 5 . 65 (GB), that is, half of AdamW-8bit.

GaLore. Given a neural network with L layers, where each layer has weight matrix W i ∈ R A i × B i (shaped as a 2D matrix), the model size d = ∑ L i =1 A i B i . GaLore uses a rankr compression via the SVD composition as W i = USV T , where U ∈ R A i × A i and we choose the first r columns of U as R i ∈ R A i × r to project the gradient of W i to the r -dimensional space. As a consequence, the dimension d shrinks to d r = ∑ L i =1 A i r = r ∑ L i =1 A i , which represents the total number of components to be stored in the GPU memory only for the projection matrices R i . If we suppose they are stored in bfloat16 (2 B per component), then the entire memory usage for low-rank projection would be 2 d r . Note that some layers in the model are rank1 and they do not need compression, but will still have associated states in Adam, which means they must be tan into consideration when computing the theoretical memory (we will use ϵ 1 for memory of rank1 layers). In addition, we have to store the states m t and v t for AdamW in bfloat16 format, which adds another 4 d r bytes. In sum, the total memory footprint of GaLore is M GLAW 16 ( r ) = 6 d r +2 ϵ 1 (B), while for the 8-bit version we get M GLAW 8 ( r ) = 4 d r + 2 ϵ 1 (B). In the end, the practical

memory usage for Llama-2 7B is M GLAW 8 (256) = 1 . 36 (GB), M GLAW 8 (1024) = 5 . 43 (GB), M GLAW 16 (256) = 2 . 04 (GB) and M GLAW 8 (1024) = 8 . 15 (GB).

Discussion. Assume our goal is to obtain a lower memory footprint compared to AdamW-8 bit. We fix the gradient density to k = d/ 100 and we have to determine the number of gradients (window size) m for MICROADAM in order to be competitive with AdamW-8 bit.

For this, we have to solve the equation M µA ( m ) = M AW 8 for m , resulting in m max = 37 . 5 . Specifically, if we use a gradient history of m < m max gradients in G , MICROADAM will have theoretical memory savings. We will see that, in practice, this history size m = 10 is more than enough for good practical results. As entries in the window past this range are dampened extremely significantly, their influence is negligible. In Appendix D, we provide Python code to compute the theoretical memory usage for the three optimizers for Llama-2 7B.

4 Convergence Guarantees for MICROADAM

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

4.1 Gradient and Error Compression

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

Assumption 1. The gradient compressor C : R d → R d is q -contractive with 0 ≤ q < 1 , i.e.,

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

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

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

Lemma 1. Consider Algorithm 2 with randomized rounding, i.e., for a vector x ∈ R d with δ = min i x i and ∆ = max i x i , let ˆ x i := ⌊ x i -δ u + ξ ⌋ u + δ be the i -th coordinate of the quantized vector ˆ x , where ξ ∼ U [0 , 1] is the uniform random variable and u = ∆ -δ 2 b -1 is the quantization level. Then

Next, we provide an 'analytical' view of our method in Algorithm 3. Essentially, we use the contractive compressor C for compressing the error corrected gradient information g t + e t , and the unbiased compressor Q to compress the remaining compression error g t + e t -C ( g t + e t ) .

Algorithm 3 MICROADAM: Analytical View

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

4.2 Convergence Guarantees for General Smooth Non-convex Functions

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

Assumption 3 (Lower bound and smoothness) . The loss function f : R d → R is lower bounded by some f ∗ ∈ R and L -smooth, i.e., ∥∇ f ( θ ) -∇ f ( θ ′ ) ∥ ≤ L ∥ θ -θ ′ ∥ , for any θ, θ ′ ∈ R d .

Assumption 4 (Unbiased and bounded stochastic gradient) . For all iterates t ≥ 1 , the stochastic gradient g t is unbiased and uniformly bounded by a constant G ≥ 0 , i.e., E [ g t ] = ∇ f ( θ t ) , ∥ g t ∥ ≤ G.

Assumption 5 (Bounded variance) . For all iterates t ≥ 1 , the variance of the stochastic gradient g t is uniformly bounded by some constant σ 2 ≥ 0 , i.e., E [ ∥ g t -∇ f ( θ t ) ∥ 2 ] ≤ σ 2 .

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

Theorem 1. (Non-convex convergence rate) Let Assumptions 1, 2, 3, 4, 5 hold and q ω := (1+ ω ) q < 1 . Then, choosing η = min { ϵ 4 LC 0 , 1 √ T } , MICROADAM (Algorithm 3) satisfies

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

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 µ > 0 the loss f satisfies Polyak-Lojasiewicz (PL) inequality

In this case, we can show:

Theorem 2. (PL convergence rate) Let Assumptions 1, 2, 3, 4, 5 and 6 hold, and q ω < 1 . Then, choosing η = min { ϵ 4 LC 0 , 2 C 0 log T µT } , MICROADAM (Algorithm 3) satisfies

with constant C 1 := β 1 1 -β 1 (1 + C 2 ) + 2 q ω 1 -q 2 ω .

Discussion. In contrast to the general non-convex setup, the study of non-convex analysis under the PL condition for AMSGrad or Adam-type methods is much less extensive. The only work we found analyzing the PL condition, which claims to be the first in this direction, focuses on Adam when β 2 → 1 , achieving a convergence rate of O ( 1 T ) [He et al., 2023]. However, our MICROADAM is based on AMSGrad normalization, and no constraint on β 2 is imposed in the analysis. Therefore, similar to

the general non-convex case, we are able to achieve the best-known convergence rate in the leading term, up to a logarithmic factor. The third, higher-order term has higher constant dependencies, but they should be negligible as the term is dampened by T 2 . Hence, in this case as well, the theory predicts that the convergence rate of the algorithm should be similar to the uncompressed version, modulo a constant that can be controlled using the compression parameters.

5 Experiments

We now validate our optimizer experimentally. We focus on comparing MICROADAM with Adam, Adam-8bit, GaLore and CAME in the context of LLM finetuning on different tasks and with SGD, Adam and AdamW-8bit in the context of ResNets on ImageNet. Concretely, we test our optimizer in full finetuning (FFT) scenario on BERT-Base/Large [Devlin et al., 2018] and OPT-1.3B [Zhang et al., 2022] on GLUE/MNLI and Llama2-7B/13B [Touvron et al., 2023] on the GSM8k math reasoning dataset and on the Open-Platypus instruction tuning dataset, as well as pre-training ResNet models on ImageNet. We provide full details regarding training settings hyper-parameters in Appendix B.

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

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

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

Finetuning results for LLaMA2 on GSM-8k. Next, we perform finetuning on Llama-2 7B/13B on GSM-8k, a challenging grade-school-level mathematical reasoning dataset. The baseline model obtains extremely low zero-shot accuracy on this task and therefore fine-tuning is necessary. In this setup, we compare MICROADAM with Adam and Adam-8bit in terms of evaluation accuracy and memory usage. In Table 2 we show our results for 3 training epochs, global batch size 32 with micro-batch (per-device) size 1 , max sequence length 512 on a single GPU, which are the standard

parameters for this task. We integrated our optimizer with the llm-foundry repository of MosaicML and tested via lm-evaluation-harness .

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

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

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

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 6 and Figure 7 in Appendix C). We compare our MICROADAM with SGD, Adam and AdamW and report the training loss, validation accuracy and only the optimizer state memory because the total memory usage is not an issue for ResNets on ImageNet (here we focus on the results to emphasize the pre-training performance). For ResNet-18, AdamW reaches the lowest training loss, followed by AdamW-8bit, MICROADAM and in the end MICROADAM. However, despite slightly larger loss, MICROADAM yields the best result among all optimizers, with 2% more than the highly tuned SGD, while having the lowest memory footprint for the optimizer states. For ResNet-50, AdamW-8bit reaches the lowest training loss, followed by AdamW, MICROADAM and SGD. The validation accuracy for AdamW and AdamW-8bit is surprisingly small compared to SGD and MICROADAM. As it is widely known

in the community, Adam variants have lower performance than SGD for Computer Vision tasks and MICROADAM fixes this issue (see the Discussion section for an intuitive explanation for this phenomenon).

Pre-training LLMs. In this section we explain why we do not include LLM 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 at each step yields to significantly better performance compared to SGD for ResNets on ImageNet. Secondly, our experiments on LLM pre-training showed difficulties in achieving the same performance compared to AdamW-8bit. Our explanation is that projection matrices from the attention layers must receive dense updates to learn the correlations between words. In contrast to the convolutional filters for CV models, the weights in attention are much larger and try to capture global correlations (features) between words, while the convolutional filters are smaller and capture local features.

Discussion. In Appendix A we provide information about the optimization set, intuitive explanations for the implicit regularization effect of MICROADAM, as well as an overview of our results.

6 Limitations and Broader Impact

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

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

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

References

Agarwal, B. Bullins, X. Chen, E. Hazan, K. Singh, C. Zhang, and Y. Zhang. Efficient full-matrix adaptive regularization. In International Conference on Machine Learning , pages 102-110. PMLR, 2019.

Alistarh, T. Hoefler, M. Johansson, N. Konstantinov, S. Khirirat, and C. Renggli. The convergence of sparsified gradient methods. In Advances in Neural Information Processing Systems , pages 5973-5983, 2018.

S.-i. Amari. Information geometry and its applications , volume 194. Springer, 2016.

Anil, V. Gupta, T. Koren, and Y. Singer. Memory efficient adaptive optimization. Advances in Neural Information Processing Systems , 32, 2019.

Beeching, C. Fourrier, N. Habib, S. Han, N. Lambert, N. Rajani, O. Sanseviero, L. Tunstall, and T. Wolf. Open llm leaderboard. https://huggingface.co/spaces/HuggingFaceH4/open_ llm_leaderboard , 2023.

Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018.

Défossez, L. Bottou, F. Bach, and N. Usunier. A simple convergence proof of adam and adagrad. Transactions on Machine Learning Research , 2022. ISSN 2835-8856. URL https://openreview.net/forum?id=ZPQhzTSWA7 .

Dettmers, M. Lewis, S. Shleifer, and L. Zettlemoyer. 8-bit optimizers via block-wise quantization. arXiv preprint arXiv:2110.02861 , 2021.

Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 2018.

C. Duchi, E. Hazan, and Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. J. Mach. Learn. Res. , 12:2121-2159, 2010.

Feinberg, X. Chen, Y. J. Sun, R. Anil, and E. Hazan. Sketchy: Memory-efficient adaptive regularization with frequent directions, 2023.

Frantar, E. Kurtic, and D. Alistarh. M-fac: Efficient matrix-free approximations of second-order information, 2021.

Ghadimi and G. Lan. Accelerated gradient methods for nonconvex nonlinear and stochastic programming. Mathematical Programming , 2016. ISSN 156(1-2):59-99.

Gupta, T. Koren, and Y. Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning , pages 1842-1850. PMLR, 2018.

He, Y. Liang, J. Liu, and D. Xu. Convergence of adam for non-convex objectives: Relaxed hyperparameters and non-ergodic case. Journal of Machine Learning Research , 2023.

Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt. Measuring massive multitask language understanding, 2021.

P. Karimireddy, Q. Rebjock, S. U. Stich, and M. Jaggi. Error feedback fixes SignSGD and other gradient compression schemes. In Proceedings of the Thirty-sixth International Conference on Machine Learning , pages 3252-3261, 2019.

P. Kingma and J. Ba. Adam: A method for stochastic optimization, 2014.

N. Lee, C. J. Hunter, and N. Ruiz. Platypus: Quick, cheap, and powerful refinement of llms. arXiv preprint arXiv:2308.07317 , 2023.

Li, B. Karimi, and P. Li. On distributed adaptive optimization with gradient compression. arXiv preprint arXiv:2205.05632 , 2022.

Liu, Z. Li, D. Hall, P. Liang, and T. Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. arXiv preprint arXiv:2305.14342 , 2023.

Loshchilov and F. Hutter. Decoupled weight decay regularization. In Proceedings of the Seventh International Conference on Learning Representations , 2019.

Luo, X. Ren, Z. Zheng, Z. Jiang, X. Jiang, and Y. You. Came: Confidence-guided adaptive memory efficient optimization. arXiv preprint arXiv:2307.02047 , 2023.

I.-V. Modoranu, A. Kalinov, E. Kurtic, and D. Alistarh. Error feedback can accurately compress preconditioners. arXiv preprint arXiv:2306.06098 , 2023.

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.

Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi. WINOGRANDE: an adversarial winograd schema challenge at scale, 2019.

Seide, H. Fu, J. Droppo, G. Li, and D. Yu. 1-bit stochastic gradient descent and its application to data-parallel distributed training of speech DNNs. In Fifteenth Annual Conference of the International Speech Communication Association , 2014.

Shazeer and M. Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning , pages 4596-4604. PMLR, 2018.

Spring, A. Kyrillidis, V. Mohan, and A. Shrivastava. Compressing gradient optimizers via count-sketches. In International Conference on Machine Learning , pages 5946-5955. PMLR, 2019.

Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample. Llama: Open and efficient foundation language models, 2023.

Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, J. Davison, S. Shleifer, P. von Platen, C. Ma, Y . Jernite, J. Plu, C. Xu, T. L. Scao, S. Gugger, M. Drame, Q. Lhoest, and A. M. Rush. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations , pages 38-45, Online, Oct. 2020. Association for Computational Linguistics. URL https://www.aclweb.org/anthology/2020.emnlp-demos.6 .

Xie, P. Zhou, H. Li, Z. Lin, and S. Yan. Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models, 2023.

Yao, A. Gholami, S. Shen, K. Keutzer, and M. W. Mahoney. Adahessian: An adaptive second order optimizer for machine learning. arXiv preprint arXiv:2006.00719 , 2020.

Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi. Hellaswag: Can a machine really finish your sentence?, 2019.

Zhang, S. Roller, N. Goyal, M. Artetxe, M. Chen, S. Chen, C. Dewan, M. Diab, X. Li, X. V. Lin, T. Mihaylov, M. Ott, S. Shleifer, K. Shuster, D. Simig, P. S. Koura, A. Sridhar, T. Wang, and L. Zettlemoyer. Opt: Open pre-trained transformer language models, 2022.

Zhao, Z. Zhang, B. Chen, Z. Wang, A. Anandkumar, and Y. Tian. Galore: Memory-efficient llm training by gradient low-rank projection. arXiv preprint arXiv:2403.03507 , 2024.

Zhou, J. Chen, Y. Cao, Z. Yang, and Q. Gu. On the convergence of adaptive gradient methods for nonconvex optimization. Transactions on Machine Learning Research , 2024a. ISSN 2835-8856. URL https://openreview.net/forum?id=Gh0cxhbz3c . Featured Certification.

Zhou, X. Xie, Z. Lin, and S. Yan. Towards understanding convergence and generalization of adamw. IEEE Transactions on Pattern Analysis and Machine Intelligence , pages 1-8, 2024b. doi: 10.1109/TPAMI.2024.3382294.

Contents

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 in the optimization set for any of the optimizers because our experiments showed no significant difference when optimizing the embeddings. Moreover, it is more fair for GaLore which would have an increased memory usage due to the projection matrix for the embeddings. For LLaMa2 and ResNet models, we include all layers in the optimization set, regardless of their type. This means that MICROADAM updates at most 10% of the weights in each layer. In the original work, GaLore was applied only to a subset of layers, such as Q-, K-, V-, O-, up-, down- and gate-projection layers. Applying GaLore to all layers in the same way as we do with MICROADAM would result in much larger memory usage because of the projection matrices for the embeddings. Moreover, computing SVD for the embedding layer would be infeasible. As a result, we omit the GaLore for LLaMa2 and ResNet experiments.

Implicit regularization. In our results so far, we observed MICROADAM having better performance for a few LLM finetuning tasks, but especially for the ResNet/ImageNet results, where the difference was statistically significant (around 1%). Our intuition for this behavior mainly comes from the accuracy curves in Figure 6 and Figure 7. The trajectories of MICROADAM and SGD are typical for regularized training. We hypothesize that MICROADAM has an implicit regularization mechanism because the model update u t is 90% sparse, which leads to only updating 10% of the model parameters at each step in each layer. In contrast to a 100% dense update u t , a sparse update would not change the model parameters as much as a dense one. In the ResNet experiments, all optimizers used the same regularization parameter λ = 1 e -4 , but the accuracy graph shows that MICROADAM is more regularized than SGD, while the graphs for AdamW and AdamW-8bit look like the regularization does not have any effect.

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

B Training Settings and Hyper-parameters

In this section we provide details about the hyper-parameters that we used for each model and dataset. We train all our models in bfloat16 format, tune the learning rates on a grid and report the best accuracy among 3 seeds (7, 42 and 1234) and report the results for the best configuration that converged.

All Adam variants use default parameters β 1 = 0 . 9 , β 2 = 0 . 999 , ϵ = 10 -8 and the regularization parameter λ is 0 for finetuning and 3 e -4 for ImageNet pre-training. MICROADAM uses a window size of m = 10 gradients with k = 1% density (equivalent to 99% sparsity and quantization bucket size is set to 64 for the error feedback.

For GaLore we use rank r = 256 and the SVD update interval is set to T = 200 , as suggested by the original paper. We run our experiments on NVidia GPUs A100-SXM4-80GB, H100-80GB and on RTX 3090 with 24GB RAM in single GPU setup.

B.1 GLUE/MNLI

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

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

Adam and Adam-8bit. All hyper-parameters mentioned above apply for these two main baseline optimizers.

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

CAME. This optimizer has some additional parameters that we keep to default values, such as β 3 = 0 . 9999 . Instead of ϵ , it uses ϵ 1 = 1 e -30 and ϵ 2 = 1 e -16 . The authors mention that the learning rate should be much smaller than Adam's and because of that we augment the learning rate grid with the values { 1 e -7 , 3 e -7 , 5 e -7 , 7 e -7 } .

B.2 GSM-8k.

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

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

B.3 ImageNet

For ImageNet, we integrate our MICROADAM in the FFCV repository, which is highly tuned for ResNets and SGD. We use E = 100 epochs, batch size 1024, cosine learning rate schedule with warmup and image resolution 224 × 224 and precision bfloat16 . We started from the initial learning rate η = 1 . 024 tuned in the repository which scored highest accuracy for SGD. This learning rate also worked well for MICROADAM, but it didn't work for AdamW and AdamW-8bit. For these two Adam variants we divided the learning rate by 2 until the models converged.

ResNet-18. For AdamW, the learning rate is η = 0 . 016 and for AdamW-8bit is η = 0 . 032 .

ResNet-50. For AdamW, the learning rate is η = 0 . 008 and for AdamW-8bit is η = 0 . 008 .

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 on GSM-8k and ResNet-18/50 on ImageNet.

Training Loss for BERT-Base

0

50

100

150

200

250

300

350

step

0.2

0.4

0.6

0.8

1.0

train/loss

MicroAdam

AdamW-8bit

GaLore

AdamW

CAME

Training Loss for BERT-Large

0

50

100

150

200

250

300

350

step

0.2

0.4

0.6

0.8

1.0

1.2

train/loss

MicroAdam

AdamW-8bit

GaLore

AdamW

CAME

Training Loss for OPT-1.3B

0

50

100

150

200

250

300

350

step

0.2

0.4

0.6

0.8

1.0

1.2

train/loss

MicroAdam

AdamW-8bit

GaLore

AdamW

CAME

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:

i mport math

Training Loss for Llama-2 7B

0

100

200

300

400

500

600

700

step

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

loss/train/total

MicroAdam

AdamW-8bit

AdamW

0

20

40

60

80

100

1

2

3

4

5

6

train loss

Training loss for ResNet-18/ImageNet

MicroAdam

AdamW-8bit

SGD

AdamW

0

20

40

60

80

100

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

top 1

Validation accuracy for ResNet-18/ImageNet

MicroAdam

AdamW-8bit

SGD

AdamW

0

20

40

60

80

100

1

2

3

4

5

6

train loss

Training loss for ResNet-50/ImageNet

MicroAdam

SGD

AdamW-8bit

AdamW

0

20

40

60

80

100

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

top 1

Validation accuracy for ResNet-50/ImageNet

MicroAdam

SGD

AdamW-8bit

AdamW

E Deferred Proofs

At time step t , let the uncompressed stochastic gradient be g t = ˜ ∇ θ f ( θ t ) , the error accumulator be e t , and the compressed gradient after the error correction be ˜ g t = C ( g t + e t ) . The second moment computed by the compressed gradients is denoted as v t = β 2 v t -1 + (1 -β 2 )˜ g 2 t , and ˆ v t = max { ˆ v t -1 , v t } is the AMSGrad normalization for the second-order momentum. Besides the first-order gradient momentum m t used in the algorithm description, we define similar running average sequence m ′ t based on the uncompressed gradients g t .

Note that m ′ t is used only in the analysis, we do not need to store or compute it. By construction we have

Denote by ζ t = e t +1 -( e t + g t -˜ g t ) = Q ( e t + g t -˜ g t ) -( e t + g t -˜ g t ) the compression noise from Q . Due to unbiasedness of the compressor Q (see Assumption 2), we have E [ ζ t | θ t , g t , ˜ g t , e t ] = 0 . Also, from the update rule of e t +1 we get e t +1 = e t + g t -˜ g t + ζ t . Moreover, we use the following auxiliary sequences,

E.1 Intermediate Lemmas

Lemma 1. Consider Algorithm 2 with randomized rounding, i.e., for a vector x ∈ R d with δ = min i x i and ∆ = max i x i , let ˆ x i := ⌊ x i -δ u + ξ ⌋ u + δ be the i -th coordinate of the quantized vector ˆ x , where ξ ∼ U [0 , 1] is the uniform random variable and u = ∆ -δ 2 b -1 is the quantization level. Then

Proof. The unbiasedness can be verified directly from the definition for each coordinate. Without loss of generality assume that δ = x 1 ≤ x 2 ≤ · · · ≤ x d -1 ≤ x d = ∆ . By construction of the quantization, we have | ˆ x 1 -x 1 | = | ˆ x d -x d | = 0 and | ˆ x i -x i | ≤ u for the remaining coordinates 2 ≤ i ≤ d -1 . Then

which completes the proof.

Lemma 2. Under Assumptions 1-5, for all iterates t and T we have

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

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

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

Summing over t = 1 , . . . , T , we obtain

Lemma 3. Let q ω = (1 + ω ) q < 1 . Under Assumptions 1-5, for all iterates t we have

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

where (2) is derived by choosing ρ = 1 -q 2 ω 2 q 2 ω and the fact that q ω < 1 . For the first claim we recursively apply the obtained inequality and use bounded gradient Assumption 4. For the second claim, initialization e 1 = 0 and the obtained recursion imply

which concludes the lemma.

Lemma 4. Let q ω = (1 + ω ) q < 1 . Under Assumptions 1-5, for all iterates t we have

Proof. Using the bounds defining compressors and Lemma 3, we get

For the second claim, recall the definition of Z t and apply triangle inequality:

Lemma 5. For the moving average error sequence E t , it holds that

Proof. Let e t,j be the j -th coordinate of e t and denote

Applying Jensen's inequality and Lemma 3, we get

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

The desired result is obtained.

Lemma 6. Let q ω = (1 + ω ) q < 1 . Under Assumptions 1-5, for all iterates t ∈ [ T ] and coordinates i ∈ [ d ] , the following bound holds

Proof. Lemma 3 and Assumption 4 imply

It's then easy to show by the updating rule of ˆ v t , there exists a j ∈ [ t ] such that ˆ v t,i = v j,i . Then

which concludes the claim.

Lemma 7. For D t := 1 √ ˆ v t -1 + ϵ -1 √ ˆ v t + ϵ we have

Proof. By the update rule, we have ˆ v t -1 ,i ≤ ˆ v t,i for any iterate t and coordinate i ∈ [ d ] . Therefore, by the initialization ˆ v 0 = 0 , we get

For the sum of squared l 2 norms, note the fact that for a ≥ b > 0 , it holds that

Thus,

which gives the desired result.

E.2 Non-convex Analysis

Here we derive the convergence rate with fixed step-size η . The rate shown in the main part can be obtained by plugging the expression of η shown after the proof.

Theorem 3. (Non-convex convergence rate) Let Assumptions 1, 2, 3, 4, 5 hold and q ω := (1+ ω ) q < 1 . Then, choosing any step-size η ≤ ϵ 4 LC 0 , MICROADAM (Algorithm 3) satisfies

with constants

C

:=

√

4(1+

(1

0

-

q

q

2

ω

2

ω

)

)

3

2

1

1

β

1

-

β

1

2

2

q

-

1

ω

q

2

ω

2

2

1

q

-

ω

q

2

ω

Proof. Similar to the proof of Comp-AMS [Li et al., 2022], we define two virtual iterates θ ′ t and x t .

2

G

+

ϵ, C

:=

(1 +

C

) +

, C

:=

ωq

(1 +

)

.

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

where we used the fact that ˜ g t + e t +1 = g t + e t + ζ t with quantization noise ζ t , and e 0 = 0 at initialization. Next, for the x t iterates we have

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

Due to unbiasedness of the compressor Q (see Assumption 2), we have E [ ζ t | g t , ˜ g t , e t , ˆ v t ] = 0 . Taking expectation, we obtain

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

Bounding term I. We have

where we use Assumption 4, Lemma 6 and the fact that l 2 norm is no larger than l 1 norm. Bounding term II. By the definition of E t and Z t , we know that

Then we have

where C 1 = β 1 1 -β 1 ( 1 + ωq ( 1 + 2 q ω 1 -q 2 ω )) + 2 q ω 1 -q 2 ω . The second inequality is because of smoothness of f ( θ ) , and the last inequality is due to Lemma 3, Assumption 4 and the property of norms.

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

where C 2 = ωq (1 + 2 q 1 -q 2 ) and we used Assumption 5 that g t is unbiased with bounded variance σ 2 . Bounding term IV. We have

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

and taking the telescoping summation over t = 1 , . . . , T , we obtain

Setting η ≤ ϵ 4 LC 0 and choosing ρ = ϵ 2 C 0 , we further arrive at

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

where in the last inequality we used x 1 = θ 1 and the lower bound f ∗ ≤ f ( θ ) for all θ ∈ R d .

To get the rate mentioned in the main part, choose η = min { ϵ 4 LC 0 , 1 √ T } and continue

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 be obtained by plugging the expression of η .

Theorem 4. (Convergence rate under PL) Let Assumptions 1, 2, 3, 4, 5 and 6 hold, and q ω < 1 . Then, choosing any step-size η ≤ ϵ 4 LC 0 , MICROADAM (Algorithm 3) satisfies

Proof. We start from descent lemma

We bound part II and part III in the same way as it was done in the non-convex analysis. We now provide a bound for part I ′ :

We further expand and bound this equation as follows:

Next, we use the Cauchy-Schwartz inequality to bound inner products above, L -smoothness inequality to bound ∥∇ f ( x t ) - ∇ f ( θ t ) ∥ ≤ L ∥ x t -θ t ∥ ≤ ηLC 1 G √ ϵ , and the inequality -∥ a ∥ 2 ≤ -1 2 ∥ b ∥ 2 +

∥ a -b ∥ 2 for the first term:

Plugging the obtained bound for I ′ with previously obtained bounds for II and III

into (12) and using the step-size bound η ≤ ϵ 4 LC 0 we get

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

Notice that η ≤ ϵ 4 LC 0 ≤ C 0 4 µ , so that the coefficient 1 -ηµ C 0 ∈ (0 , 1) . Unrolling the recursion, we arrive

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

For the other three sums we bound 1 -ηµ C 0 ≤ 1 and apply the bounds in Lemma 7:

Plugging all this bounds into (13) and noticing that x 1 = θ 1 , we finally get

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

which implies

and completes the proof.

To get the rate mentioned in the main part of the paper, we plug in the expression η = min { ϵ 4 LC 0 , 2 C 0 log T µT } and collect higher order terms.

E.4 Non-convex Analysis with Weight Decay

Lemma 8. Under Assumptions 1-5, for all iterates of Algorithm 4 we have

Proof. We start our proof from

Using Jensen's inequality we have

Using Young's inequality we have

Combining this bound with L -smoothness we obtain

Lemma 9. Under Assumptions 1-5, for all iterates of Algorithm 4 we have

Proof. We start from

Using definition of contractive compressor we have

Using Young's inequality we have

Lemma 10. Under Assumptions 1-5, for all iterates of Algorithm 4 we have

Proof. Using definition of contractive compressor we have

Using Young's inequality we have

We need to satisfy the following condition:

It holds for 0 ≤ ω, 0 ≤ q < 1 , (1 + ω ) q < 1 and a = 1 (1+ ω ) q -1 . Using this parameters we have

Unrolling this recursion allows us to obtain

last inequality holds because e 1 = 0 .

Lemma 11. Under Assumptions 1-5, for all iterates of Algorithm 4 we have

Proof.

Using Lemma 10 we obtain

Lemma 12 (From paper: Zhou et al. [2024b]) . Let us consider the update rule:

For brevity, we denote ̂ v t = √ v t + ε . Also we define

where ⊗ denotes element-wise product. Moreover, we also define ˜ f ( θ t ) as follows:

where λ t = λ 2 ∑ t i =1 ( 1 -q 2 ) i for t > 0 , λ 0 = 0 with 0 < q < 1 and ∥ θ t ∥ ˆ v t = √ ⟨ θ t , ˆ v t ⊗ θ t ⟩ . Also let c 1 ≤ ∥ ̂ v t ∥ ∞ < c 2 , then iterates of Algorithm 4 satisfy

Lemma 13 (From paper: Zhou et al. [2024b]) . Assume that c s, ∞ ≤ ∥ ˜ g t ∥ ∞ ≤ c ∞ , then we have

where µ = β 2 c 2 ∞ c 2 s, ∞ + ϵ

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

Using previous lemmas we have

Discussion. This result is similar to the one from Zhou et al. [2024b] in the non-convex case, where the decay rate for the first term is O ( 1 T ) and the second term is a non-vanishing O ( β 1 c 2 c 1 σ 2 ) . In our result, the non-vanishing term is proportional to O ( c 2 c 1 ( σ 2 + G 2 ) ) .

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

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

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

F Error Feedback applied to GaLore

F.1 Behaviour of the Error Feedback Mechanism

The GaLore low-rank updates introduced by [Zhao et al., 2024] enable the compression of optimizer states by performing learning updates on a lower-dimensional subspace. In this approach, the

optimizer receives gradients projected on a defined learning subspace. Theoretical convergence guarantees are provided under a 'stable rank' assumption, where learning subspace is fixed during training. However, in practice, convergence is attained by occasionally updating the learning subspace and allowing full space learning to better align with the gradient trajectory during training.

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

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

For a better understanding, we derive analytical evidence for the described behaviour by induction. Let L be fixed learning subspace and assume that e t -1 ∈ L ⊥ . Then, gradient passed to the optimizer at step t is: C GaLore ( a t ) = proj L ( a t ) = proj L ( e t -1 + g t ) = proj L ( g t ) where error feedback is discarded. Thus, e t = a t - C GaLore ( a t ) = e t -1 + g t -proj L ( g t ) = e t -1 + proj L ⊥ ( g t ) ∈ L ⊥ which completes the induction.

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

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 8 provides evidence of such linear growth of the error norm during fine-tuning of RoBERTabase model on GLUE/MNLI task.

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

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

175

high frequency dynamics of||et-1l|

I| et-1 II

150

= 124

[-

122

125

=

120

100

1e5

1e5+T1e5+2T 1e5+3T 1e5+4T 1e5+5T 1e5+6T 1e5+7T 1e5+8T 1e5+9T

Training step

75

highfrequency dynamicsof Ilgt ll

50

25

二

6

4

2

0

0

1e5

1e5+T1e5+2T 1e5+3T 1e5+4T 1e5+5T 1e5+6T 1e5+7T 1e5+8T 1e5+9T

1e5

4e5

7e5

Training step

Training step

convergence by enforcing space exploration. See Figure 9 for examples of how induced orthogonality of successive learning subspaces affects the learning trajectory.

Adam

GaLoreAdam

0.5

0.5 -

0.5 -

1

1

0

0

0

0.0

0.0-

0.0-

1

-1

2

0.5

1.5

1.0

0.5

0.0

0.5

0.5

1.5

1.0

0.5

0.0

0.5

0.5

1.5

1.0

0.5

0.0

0.5

Adam

GaLoreAdam

1.0 -

102

1.0

102

1.0 -

102

100

100

100

0.5 -

0.5

0.5

10-2

10-2

10-2

10-4

10-4

10-4

0.0-

0.0-

0.0-

-1.0

0.5

0.0

0.5

1.0

1.0

0.5

00

0.5

1.0

1.0

0.5

0:0

0.5

1.0

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 we claim that our 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 in 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 it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of 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 no 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 a separate "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 on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 of 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 set of assumptions and a complete (and correct) proof?

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 proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We explain our experiments in Section 5 and provide the complete set of hyper-parameters 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 a 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 a novel architecture, describing the architecture fully might suffice, or if the contribution is 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 to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a 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 be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case 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 to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not 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 NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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 a 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 as much information as possible in supplemental material (appended to the paper) is 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, hyperparameters, 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 of the paper to a 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, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: We do not provide error bars, but instead explain how we 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, confidence intervals, or statistical significance tests, at 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 run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to 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 or 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 a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: All these information can be 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 as 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 or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Yes, the research conducted in this paper conforms, in every respect, with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal 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 or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [No]

Justification: We do not release any models.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We work with open source models that are publicly available and we cited them properly.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We provide code and instructions on how to run the code in order to reproduce our results.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing 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 as 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or 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 not involve crowdsourcing nor research with human subjects.

Depending on 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.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: 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

STRENGTHS:
1. This paper proposes a new optimizer based on Adam, called MicroAdam, which reduces memory footprints, via sparisification, quantization, and error feedback. 
2. Theoretical convergence guarantees are provided. 
3. Empirical results show good performance on several fine-tuning tasks
4. 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.
2. The theoretical analysis is actually based on AMSGrad instead of Adam.

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

My concerns are well addressed according to the author's feedback.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-10

comments: 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.

STRENGTHS:
- This work provides theoretical convergence guarantees under commonly used assumptions. Previous work on memory-efficient Adam often relied more on heuristics.
- 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. 
- The approach is limited to fine-tuning tasks, and its effectiveness for LLM pre-training or non-language tasks remains unclear.

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
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- **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.\n\nWEAKNESSES:\n- **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?\n\nQUESTIONS:\nSee Weaknesses\n\nLIMITATIONS:\nSee Weaknesses", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nIn 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.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe script proposed a memory-efficient method with a theoretical guarantee.\n\nSTRENGTHS:\nThe topic on memory-efficient optimizers is important. I did not observe obvious flaws in the theory.\n\nWEAKNESSES:\nSee below.\n\nQUESTIONS:\nSee below.\n\nLIMITATIONS:\n**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.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-15"}