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

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?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

The limitations are well discussed.

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?

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

2

Limitations

Limitations are well-addressed

Summary

The paper introduces MicroAdam, a novel optimizer designed to improve memory efficiency while maintaining competitive performance with established optimizers such as Adam. The authors provide theoretical analyses and experimental results to demonstrate the benefits of MicroAdam in various settings.

Strengths

- **Theoretical Analysis**: The paper includes comprehensive theoretical analyses of memory consumption and the guarantee of convergence. - **Experimental Validation**: The experimental results showcase the potential of MicroAdam in reducing memory consumption, with detailed comparisons to existing methods.

Weaknesses

- **Motivation in Figure 1**: Figure 1 does not effectively provide additional motivation for MicroAdam, as it represents a 2-dimensional non-stochastic problem with only choosing k=1 for top-k compression. This toy example can be solved by many heuristic methods or parameter adjustments, which diminishes the unique motivation for MicroAdam. - **Convergence Speed**: The convergence speed of MicroAdam is asymptotically similar to AMSGrad rather than Adam itself. Given that the paper aims to improve Adam's memory footprint, it is unclear why the convergence speed matches that of AMSGrad, which theoretically has no significant speed difference from Adam's with respect to T. - **Memory Consumption Analysis**: Sections 3.2 and C provide theoretical and simulated memory consumption analysis. However, the practical memory consumption during actual LLM training does not necessarily correlate strongly with the theoretical optimizer states' memory usage. For example, a small LLAMA model with sufficiently large batch size and token length can peak at over 70GB of memory usage (mainly due to activation memory), whereas the optimizer states' memory footprint might be under 100MB. Therefore, the paper should report more about real peak memory under appropriate settings rather than theoretical or simulated values, as the differences in optimizer states memory can be mitigated under certain settings and PyTorch's memory management mechanism. - **Top-K Operator Memory Efficiency**: Despite the block-wise operation of the Top-K operator, if the gradients are in bf16 format, the memory overhead for 50% sparsity remains the same (due to the need to store 16-bit indices) or could even be higher when considering other memory overheads in the algorithm. - **Lack of Training from Scratch Results**: A significant issue is the near absence of results for training from scratch, which is crucial for evaluating MicroAdam's effectiveness. For MicroAdam, training from scratch is more important. Table 1 shows that during fine-tuning, the training loss does not strictly correlate with the final metrics. For instance, GaLore has a much worse loss but maintains metrics comparable to those of other methods. This is why LLM fine-tuning metrics are often taken from the best of three epochs, not necessarily the last epoch, even if the training loss is the smallest in the final epoch. Additionally, Figure 4 indicates that MicroAdam's training efficiency is lower than baseline Adam, as MicroAdam's third epoch loss is nearly identical to Adam's second epoch loss. - **Computational Efficiency**: MicroAdam appears to be less computationally efficient. Are there ways to improve its computational efficiency?

Questions

See Weaknesses

Rating

4

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

See Weaknesses

Summary

The paper presents a new optimizer that approximates Adam but has a lower memory footprint. Adam stores for each parameter two additional values --- the exponential moving averages of the gradients and the gradient squared. The current optimizer saves space by using the fact that for most parameters, the gradient is quite small at each step, so could be ignored. The optimizer stores the most important gradients at each step, for several past steps. It also stores a low-resolution version of the gradient, which is added to the next gradient (the reason for this was not clear to me). The exponential moving averages for the gradients and gradients-squared are computed from the most important gradients stored at each step and are used in the Adam update. The paper proves the convergence of the optimizer under some reasonable assumptions. The authors show that the method can be implemented efficiently on a GPU, and the performance is similar to Adam. The authors performed several experiments to show that their optimizer saves about 10-20% memory compared to Adam, while achieving similar accuracy.

Strengths

The main idea is quite interesting - that most of the values of a gradient are small in magnitude, so need not be used to update the parameters --- only up to 10% of all the parameters are ever updated in any step, and these change only a little across steps. The convergence results are a useful validation of the algorithm. The experimental results seem quite strong to me.

Weaknesses

The presentation of the paper was not very clear to me. For example, k is not mentioned as one of the inputs in Algorithm 1. $\delta_1, \Delta_1$ are initialized to vectors in R^d, but then they seem to be scalars. I also did not understand what the error correction was for.

Questions

In your algorithm, why did you not use a_t or the current gradient in steps 9 and 10, since this is already present. It looks like most of the memory is being used for other tasks, so reducing the memory by a factor of 4 resulted in a small overall reduction in memory.

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

N/A

Summary

The script proposed a memory-efficient method with a theoretical guarantee.

Strengths

The topic on memory-efficient optimizers is important. I did not observe obvious flaws in the theory.

Weaknesses

See below.

Questions

See below.

Rating

3

Confidence

4

Soundness

2

Presentation

1

Contribution

1

Limitations

**Major concern 1**: The contribution of the paper is rather weak. After reading the paper, I still did not find what is new in this paper. For instance: **1-1: It is unclear what is new in the compression design of Algorithm 1.** I have a hard time figuring out the contribution of Algorithm 1. Please summarize the significance and novelty of the proposed compression procedure in Algorithm 1. Please discuss how it is different from the existing ones and why do we need this new one. Please discuss the design principles and ideas. **1-2 It is unclear whether the proposed method is truly effective.** There are much much more experiments are needed. Currently, only 7B SFT is provided, which is far from convincing. SFT is not enough to support the effectiveness of a new optimizer. It seems unclear how the method would perform on LLM pre-training or non-language tasks. **1-3: It is unclear what is new in the theory.** Under the assumption of unbiased compression + bounded variance + bounded gradient, this type of Algorithm 4 has already been extensively studied by the distributed community and so. Please highlight what is new and what is nontrivial, if there is any. **Major concern 2: The presentation is quite poor.** For instance: 1. There is no explanation on the principle or idea behind the design of Algorithm 1. 2. It is good to see simple example like Figure 1. But unfortunately, nearly nothing is explained here. We can see EF helps convergence, but we do not understand why. It seems just a toy example without any explanation on the insight. 3. The current logic flow of the script is confusing. It would be much easier to read if the authors: (i) introduce Algorithm 4 first (generic form of MicroAdam), (ii) provide the theoretical guarantee, (iii)then introduce Algorithm 1( a specific form that you chose) , (iv) discuss the implementation details, (v) then show the experiments.

Start discussions with authors, please

Dear Reviewers, Thank you for your valuable contributions to the review process. As we enter the discussion phase (August 7-13), I kindly request your active participation in addressing the authors' rebuttals and engaging in constructive dialogue. Please: - Carefully read the authors' global rebuttal and individual responses to each review. - Respond to specific questions or points raised by the authors, especially those requiring further clarification from you. - Engage in open discussions about the paper's strengths, weaknesses, and potential improvements. - Be prompt in your responses to facilitate a meaningful exchange within the given timeframe. - Maintain objectivity and professionalism in all communications. If you have any concerns or need guidance during this process, please don't hesitate to reach out to me. Your continued engagement is crucial for ensuring a fair and thorough evaluation process. Thank you for your dedication to maintaining the high standards of NeurIPS. Best regards, Area Chair