LDAdam: Adaptive Optimization from Low-Dimensional Gradient Statistics

Authors: Thomas Robert, Mher Safaryan, Ionut-Vlad Modoranu, Dan Alistarh
Primary area: optimization
Keywords: adaptive optimization, memory efficiency, low-rank learning, low-rank compression, convergence rates

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

Decision: accept (poster)
Reviewer scores: 6, 8, 6, 8

Reviews and discussion:
<thread>
creation date:
1730250596889

thread_type:
official_review

<thread_object>
creator: reviewer_tz3y

<content>
summary:
Thin paper introduces a novel memory-efficient low-dimensional Adam optimizer for training large models. Specifically, the authors propose a new projection-aware update for the optimizer states that allows for transitioning between subspaces. They also perform block power iteration to at each step to efficiently and accurately project the gradient and optimizer states. To reduce the error caused by the low-rank projection, this work leverages a new generalized error feedback mechanism that accounts for both gradient and optimizer state compression. The authors theoretically analyze the convergence rates for different objectives. To validate the proposed algorithm, the authors present extensive experimental results on different datasets with the comparison to different baselines.

strengths:
1. The investigated topic is quite interesting and critical, particularly when LLMs are popular in every area. A memory-efficient optimizer is required to conduct computationally efficient fine-tuning.
2. The paper is easy to follow and the presentation in this work is clear.
3. This work presents in-depth theoretical analysis for the proposed algorithm and show the explicits the convergence rate.
4. The empirical evidence to validate the proposed algorithm look extensive and thorough.
5. The comparison between the proposed method and baselines is technically convincing.

weaknesses:
1. Assumption 2 is a quite strong assumption for the gradient, which may limit the applicability of the proposed algorithm. Until now, most existing algorithms have relaxed this assumption in their works such that this assumption imposes some weaknesses to the proposed algorithm.
2. In this work, the authors have mentioned that GaLore is the closest work. Though we have seen the authors have shown in detail how the proposed algorithm differs, it would be great to include a technical discussion for the difference between them in terms of convergence rate. That way, it is clearer to directly show the exact theoretical difference.
3. In the practical view of the algorithm, the authors use Gram-Schmidt process to calculate $\mathcal{U}_t$. However, the process has numerical instability, computational intensity, and even inaccuracies. If this process fails, how can the authors guarantee that there always exists a feasible and reliable $\mathcal{U}_t$ for the subsequent steps?
4. Although the experimental results look promising, how to validate the theoretical conclusions from Theorem 1 and Theorem 2, particularly, the impact of the compression ratio? The authors need to present some ablation studies for the proposed algorithm.


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

questions:
1. How to relax the bounded gradient assumption in the study?
2. How to guarantee a feasible $\mathcal{U}_t$ if Gram-Schmidt process fails?
3. How to validate the theoretical conclusions from both theorems?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730584061575

thread_type:
official_review

<thread_object>
creator: reviewer_2ksc

<content>
summary:
This paper on LDAdam addresses a crucial challenge in training large models: reducing memory consumption of optimizers without sacrificing model quality. The key innovation is the introduction of a lower-dimensional subspace optimization approach, where adaptive optimization occurs within constrained dimensions. The optimizer, though operating in a subspace, ensures exploration across the entire parameter space, thereby maintaining model expressivity while being memory efficient.

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

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

3. **Convergence and Memory Efficiency**: The paper provides a proof of LDAdam’s convergence under standard assumptions. With this design, LDAdam achieves a significantly reduced memory footprint. In practical terms, LDAdam’s optimizer states require only a fraction of the memory needed by Adam, making it feasible to train large language models in resource-constrained environments.

weaknesses:
Liang et al. [1] provide a convergence analysis for GaLore and related algorithms without relying on a 'stable-rank' assumption. Since LDAdam aligns with this framework, I recommend citing [1] for its convergence insights. Additionally, a comparison of LDAdam’s convergence analysis with [1], specifically on the impact of removing the 'stable-rank' assumption, would help clarify the authors’ theoretical contributions.

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

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

Good paper, I suggest acceptance.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730602582694

thread_type:
official_review

<thread_object>
creator: reviewer_haqr

<content>
summary:
The paper proposes a new memory-efficient training algorithm based on Adam via efficient projection onto subspaces. The proposed method enjoys non-asymptotic theoretical rate of convergence and demonstrate good experimental results in training large language models (LLMs).

strengths:
The paper is well-written and the contribution is clear. The algorithm is novel with theoretical convergence. The experiments indicates that the new approach is performing not only better than GaLore, but also better than the vanilla Adam for some parameter settings.

weaknesses:
(Please respond to the "Questions" directly) The computational time of proposed method is not clear; Some discussion on the theoretical results is needed.

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

1. To my understanding, GaLore doesn't employ SVD step for every iteration in order to save computation overheads. On the contrary, the proposed method would need a Gram-Schmidt process and several extra steps for updating the momentums and gradients. A major conern is the computation overhead of this proposed method. I udnerstand that the memory consumption is roughly the same as GaLore, yet I'm curious about the time required for training LLMs both in theory and in all the experiment settings using the proposed method.
2. The convergence results in Theorem 1 and 2 are similar to that of full-rank Adagrad/Adam. To my understanding, the low rank parameter $r$ only plays a role in the constant $q_r$? If so, can the author comment on how to choose such $r$ and if the theory has any implication on the choice of $r$. In particular, if $r=m$ and the method reduce back to the original Adam, can we say any matching result with the analysis on the original Adam?
3. This question is related to the previous one. The experiments show that LDAdam is even better than original Adam in some cases (Table 3 with $r=512$, Table 4 with $r=256$). Can the authors expand on the discussion of "regularization effect of compression" or provide some references? It's still not very clear to me right now.
4. Any experiment on even larger LLMs such as Llama 1b or 7b? I'm curious what will happen if the model is large and the training is conducted in parallel on multiple GPUs.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730802846652

thread_type:
official_review

<thread_object>
creator: reviewer_hwuw

<content>
summary:
The paper proposes a new optimization algorithm to train large (language) models.
Like the recent GaLore algorithm, the proposed algorithm is a variation of ADAM, that performs updates in a lower-dimensional space, to store lower-dimensional quantities to be more memory-efficient.
As I understand (and as written by the authors), the real meat behind this kind of algorithm is the choice of the projection to obtain a meaningful lower-dimensional subspace that is both efficient to compute/store and a good enough approximation of the gradient. The proposed algorithm retains the desirable memory-efficient properties of GaLore, while almost matching the generalization performance of ADAM. Convergence proofs and exhaustive experiments on text benchmarks are provided.

strengths:
The paper is overall easy to follow, the idea is interesting, and seems to empirically works great.

weaknesses:
I see that the authors provide some running times and memory metrics in Appendix B, but it seems to me that this part should be more exhaustive in order to prove the proposed algorithm useful.
I know this is hard to monitor for algorithms on GPU, but could authors provide graphs similar to Figure 1, but with running times? and memory usage? with ADAM, GaLore, and LDADAM (proposed).
In other words, could authors provide graphs with perplexity as a function of runtime and perplexity as a function of memory usage? Even if these required graphs will be imperfect for a lot reasons, I think they are crucial to gain insights and potentially prove usefulness.

questions:
see Weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1732121675520

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Authors' General Response for the Rebuttal

comment:
We would like to thank the reviewers for their useful and diverse feedback! We provide individual answers to all questions, as well as additional experimental results, as asked for by the reviewers. Specifically, we outline the main questions addressed by our rebuttal and included in our PDF revision:

- To address questions by Reviewers, we have added pretraining experiments with up to 1.3B parameter models, and expanded our algorithm to the multi-GPU setting. The results show the same trends as in the original submission, where LDAdam has slightly better convergence than Adam, with a much lower memory overhead. They also confirm the scalability and stability of our approach. 

- We have added an in-depth discussion of computational cost, with detailed asymptotics for LDAdam and competing approaches. This shows that LDAdam provides the same or better memory cost relative to prior work on memory-efficient optimization, while at the same time providing better convergence guarantees. We have also provided results in convergence-vs-wall-clock-time format as well as ablations between throughput (tokens/second) and rank, outlining the smooth performance trade-offs that LDAdam’s implementation allows. 

- We have also added several detailed answers, such as clarifying that the bounded gradient assumption is standard in the analysis of adaptive optimization methods, discussing the superiority of our analysis approach relative to that of GaLore.

- We discussed the influence of the rank parameter r on computation and memory overhead, particularly in the context of practical considerations for selecting its value. As a sanity check, we reconfirmed that using the maximum possible rank allows our theory to match the results of the original Adam/AMSGrad. Additionally, we provided a detailed technical comparison between our theoretical framework and results and those of GaLore and approaches based on continuous-time ODEs. Finally, we justified our use of the bounded gradient assumption, which is a standard practice in the adaptive optimization literature, as there is no strictly better alternative.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732550461694

thread_type:
official_comment

<thread_object>
creator: area_chair_nm2u

<content>
title:
Important: Please Review Rebuttals and Update Reviews as Needed

comment:
Dear Reviewers,

Thank you for your hard work and dedication to providing thoughtful reviews for this year’s ICLR submissions. Your efforts play a vital role in maintaining the conference’s high standards and fostering meaningful discussions in the community.

As we are close to the end of the discussion phase, I kindly urge you to read the authors’ responses and reevaluate your reviews carefully, especially if they address your concerns. If you find that the authors have resolved the issues or clarified misunderstandings, please consider adjusting your comments and scores accordingly. This ensures fairness and gives each submission the opportunity to be judged on its merits.

Your continued commitment is greatly appreciated—thank you for contributing to the success of ICLR!
</content>
</thread_object>

</thread>


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

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