FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training

We appreciate the reviewer's comprehensive feedback. We are glad that they appreciated the theoretical convergence guarantees, strong experimental results, as well as the provided code implementation.

The empirical and theoretical support behind FRUGAL is not solid enough

We would like to emphasize that our method consistently beats all baselines, often by a large margin, as demonstrated in Tables 2, 4-7. These results validate the practical effectiveness of FRUGAL across diverse settings.

On the theoretical side, our framework provides the strongest guarantees we are aware of, especially given its ability to handle arbitrary projections. One should note that unlike GaLore, FRUGAL does not require projections to remain fixed for theory (see Theorem 3.8 in GaLore paper).

FRUGAL's stateful optimizer is basically either Galore or BAdam.

We kindly disagree with the reviewer on this point. FRUGAL does not recover GaLore as a special case, even when the state-free optimizer is discarded. This distinction arises from the reprojection step (Line 7, Algorithm 1), which is critical to enabling FRUGAL to work effectively with arbitrary projections, a capability that sets it apart from GaLore. We provide both discussion and experimental validation of this issue in lines 307-314 and Appendix C.

This evidence alone is not convincing enough for an assured generalization to other non-Llama architectures.

We have conducted additional experiments on pre-training GPT-2 124M to further strengthen our findings. See details in Table 18. Similarly to experiments with LLaMA, we found that FRUGAL significantly outperforms GaLore and BAdam.

We believe that the current results already provide strong support for both the theoretical and empirical contributions of our work. However, we kindly request the reviewer to specify any additional comparisons, results, or insights that would help address their concerns. We are open to elaborating further or including additional experiments to ensure clarity and completeness.

An ablation study on projector type versus stateful optimization density $\rho$ is definitely needed.

To address the reviewer's request and further demonstrate the effectiveness of our method, we conducted an additional ablation study. The experiments follow the setup of Table 1 but explore different density $\rho$ values: 0.0625, 0.125, 0.333, and 0.5, while in Table 1 we use $\rho = 0.25$. The results, presented in Table 15, align with our findings from Table 1. Specifically, training with random projection significantly outperforms SVD projection when training without optimizing the state-free subspace. When state-free subspace optimization is employed, SVD projections marginally outperform their Blockwise counterparts.

We believe these empirical results provide sufficient additional evidence to support our conclusions drawn from Figure 2 and Table 1. We remain open to the reviewer's suggestions regarding additional experiments that could further strengthen this argument.

The presentation of the Algorithm needs to be clearer. It is hard to understand the exact algorithm (which is actually simple) in the first time of reading Algorithm 1 and Section 3.

We appreciate the reviewer's feedback. Indeed, Algorithm 1 is written from a maximally general perspective to cover all possible types of projections. To facilitate understanding of the final algorithm variant used in our experiments, we have provided PyTorch-style pseudocode in Algorithms 4, 5.

To sum up, we believe that we addressed all the reviewer's concerns, none of which is a serious issue with our approach. Furthermore, all the mentioned strengths highlight impactful and timely contributions of our work. Therefore, we would kindly request the reviewer to reconsider their score.

Dear reviewer yRQJ,

The discussion period is about to end and we would kindly like to hear from you regarding our rebuttal.

Thanks for the response.

I previously overlooked Line 7, Algorithm 1 and it is a weakness worthwhile of discussion for Galore. I would suggest add a paragraph of discussion to the main paper besides putting it on Appendix C.

Thanks for the additional experiment on GPT2 and Llama-3B, and the ablation on density $rho$.

I still have concerns on the intuition on most parameters in LLM can be optimized by stateless optimizers part. Even though the authors provide some ablations, it is better to give an empirical insight beyond the result number (accuracy/perplexity/etc.) itself to strengthen this argument. The insight means a deeper analysis on any underlying variables or empirical facts that make most parameters in LLM can be optimized effectively by stateless optimizers (something at least can have a causal relationship). I do not see such analysis from this paper so I will still retain my original score.

We would like to thank the reviewer for their detailed comments. We are delighted that they commended the novelty of our extensive experiments and the theoretical and practical justifications of our method.

The GLUE benchmarks is little bit outdated

Thank you for this feedback. We chose GLUE for experiments because we followed the setup from GaLore (see lines 511-512).

However, we agree that additional experiments on more recent than GLUE fine-tuning benchmarks would provide an additional argument supporting the superiority of our method. We will try to add an additional fine-tuning experiment before the end of the Discussion period.

We would also like to emphasize again that our paper already includes an extensive set of experiments confirming the superiority of our method over baselines, both in pre-training and fine-tuning. Additionally, for this rebuttal, we have already conducted additional experiments on pre-training GPT-2 124M (see details in Table 18), where our method also significantly outperformed the baselines, which once again demonstrates the stability of our method across different setups.

Is there any explanations about which part of the parameters can be directly optimized with SGD type optimizer with other requires adam and why?

We appreciate the reviewer raising this question. As another reviewer had a similar concern, we have included our response in the general response section.

For $\rho=0$ in Table 2, is it equal to fully optimized with SGD?

We thank the reviewer for their question. Since a similar question was raised by another reviewer, we have addressed this in the general response.

The concepts of state-full and state-free subspace in line 80/82 is hard to understand, it's better to formally define these two concepts.

We are grateful for the reviewer's feedback. As it mirrors a question from another reviewer, we have included our response in the general response section.

this sequence make it a little bit confusing whether the "Low-rank Random" in Table 1 is training of entire random projection or first SVD and later random.

We are very grateful to the reviewer for pointing out this ambiguity. We have reformulated this sentence, see lines 203-205. If the reviewer has any remaining concerns about the clarity, we would be happy to address them.

it's better to define the meaning of K in the inputs of Algorithm 1, as well as $s$.

We appreciate the reviewer's feedback. We have modified the algorithm description to include $K$'s definition in the Input section and added a brief note clarifying that $s$ denotes the optimizer state. While space constraints prevent us from adding more detailed explanations, we welcome any specific suggestions from the reviewer on how to make the Algorithm 1 more formally precise and clear.

To sum up, we believe that we addressed all the reviewer's concerns, none of which is a serious issue with our approach. Furthermore, all the mentioned strengths highlight impactful and timely contributions of our work. Therefore, we would kindly request the reviewer to reconsider their score.

We thank the reviewer for their detailed review. We are glad that they liked our paper's presentation and our method's practical impact.

The paper would benefit from a discussion of the method's limitations and potential failure modes.

We thank the reviewer for this suggestion. To provide a more comprehensive understanding of potential trade-offs associated with our method, we have added a Limitations section to Appendix H.

Given the importance of "state-full" and "state-free" in the proposed method, the paper should offer clearer definitions of these terms.

We thank the reviewer for this feedback. Since this overlaps with another reviewer's inquiry, we have addressed it in the general response.

Beyond memory efficiency, how does the method’s running time compare to that of other baseline approaches?

As requested by the reviewer, we have measured the running time of all methods used in the paper. We present the average computational time of the optimizer step for different sizes of LLaMA models in Appendix F, Table 19. The measurements for memory-efficient methods were made with density $\rho=0.25$ and update gap $T$ equal to $200$. We report the average time over $200$ steps (to capture precisely one step with the state-full subspace update). Measurements were conducted on a single A100-80G GPU using PyTorch 2.4.1. We note that these experiments were conducted without using torch.compile.

The results show that memory-efficient methods requiring gradient projection within each Linear layer matrix (GaLore, RandK) stand out negatively. GaLore requires more time than RandK due to SVD decomposition. As model size increases, blockwise-projection methods even start outperforming Adam, despite being implemented through a for-loop over all parameters, while PyTorch uses an efficient Adam implementation by stacking updates into a single shared tensor (flag foreach=True) to better utilize the parallelization capabilities of modern GPUs. This occurs because Adam's update step requires significantly more operations than the state-free step in FRUGAL. Therefore, approximately 75% of updates in FRUGAL's for-loop require significantly fewer computations and, consequently, less time.

To sum up, we believe that we addressed all the reviewer's concerns, none of which is a serious issue with our approach. Furthermore, the reviewer highlighted the practicality of our approach. Therefore, we would kindly request the reviewer to reconsider their score.

Including more experiments comparing the method with various stateful and stateless optimizers would enhance the paper.

Following the reviewer's request, we conducted additional experiments. We explored two variations: 1) replacing AdamW with Lion as the state-full optimizer, and 2) substituting signSGD with SGD as the state-free optimizer.

The results are presented in Tables 16 and 17. As observed, the results for Lion are comparable to those obtained with AdamW - the additional optimization of the state-free subspace significantly improves performance. Training with SGD as the state-free optimizer also demonstrates satisfactory results. However, we would like to note that unlike signSGD, hyperparameter tuning for SGD training is considerably more challenging. This is because, unlike signSGD, whose update magnitudes approximately equal to those of popular Adam-like algorithms, the magnitude of updates (essentially, gradients) in SGD differs substantially, necessitating learning rates that deviate significantly from those used with the state-full optimizer. Furthermore, successful training with SGD absolutely requires gradient clipping, while the absence of such clipping is not a critical impediment for signSGD.

Testing models with larger sizes (e.g., 3B and 7B) could further demonstrate the generalizability of the proposed method.

We agree with the reviewer on this matter. The paper already includes preliminary results for training LLaMA 7b in Appendix D, Table 14. While these results don't allow for a fair comparison with standard Adam training (because 1. the GaLore baseline uses 8-bit Adam, which apparently significantly degrades the final result, and 2. the training was conducted in pure bfloat-16 format), we believe, as stated in lines 991-992, that these results demonstrate FRUGAL's potential for scaling.

Furthermore, to more thoroughly examine our method's potential with increasing model sizes, we are working on experiments comparing Adam and FRUGAL training on a 3B model. Since this experiment requires substantial computational resources, obtaining results takes considerable time. However, we plan to share preliminary results within a week before the end of the Discussion period and include the final results in the camera-ready version of the paper.

Please clarify the reasons for selecting the specific optimizers in the theoretical section.

We appreciate the reviewer’s comment on the selection of specific optimizers in the theoretical section. The primary reason for choosing SGD with momentum (SGDM) and SGD in our theoretical analysis is the lack of guaranteed convergence for Adam in certain scenarios, as demonstrated by Reddi et al. (2018) [4]. Therefore, we discard normalization in our analysis, leading to a combination of SGDM and SGD.

Additionally, we note that methods such as GaLore focus solely on gradient descent (according to Theorem 3.8 in GaLore paper, analysis is limited to constant projection of updates). In contrast, FRUGAL's framework is designed to accommodate arbitrary projections, significantly broadening its applicability and demonstrating its theoretical superiority over existing methods. This flexibility underpins the superiority of our algorithmic choices.

The paper doesn’t explain why FRUGAL with $\rho = 0$ achieves optimal performance in certain scenarios.

W thank the reviewer for raising this point. As it mirrors a question from another reviewer, we have included our detailed response in the general response.

Incorporating diverse architectures could strengthen the argument for FRUGAL's superior characteristics

We have conducted additional experiments on pre-training GPT-2 124M to further strengthen our findings. See details in Table 18. Similarly to experiments with LLaMA, we found that FRUGAL significantly outperforms GaLore and BAdam.

To sum up, we believe that we addressed all the reviewers' concerns, none of which were serious issues, mainly focusing on issues with presentation. Furthermore, all the mentioned strengths highlight impactful and timely contributions of our work. Therefore, we would kindly request the reviewer to reconsider their score.

[1] Liu et. al., An improved analysis of stochastic gradient descent with momentum. NeurIPS 2020.

[2] Hu et. al., Lora: Low-rank adaptation of large language models. ICLR 2022.

[3] Zhao et. el., Deconstructing what makes a good optimizer for language models. arXiv:2407.07972, 2024.

[4] Reddi et al. On the convergence of adam and beyond. arXiv preprint arXiv:1904.09237, 2019.

We are grateful to the reviewer for their thorough review. We appreciate that they commended the novelty of our idea, the theoretical and practical justifications of our method, as well as our results in analyzing the training dynamics of transformer models.

The paper's structure would greatly benefit from a clearer organization. Currently, some analysis and experimental results appear within the Methods section, which disrupts the logical flow and makes it challenging for readers to follow the methodology.

We appreciate the reviewer's feedback. To enhance the paper's readability and clarity, we have reorganized the Method section. Specifically, we have divided it into two sections: 1. Empirical Analysis and Motivation, which includes subsections ''The importance of exploring the entire space during the training process'' and ''Advantage of the Full-Rank Updates''. 2. FRUGAL: Full-Rank Updates with GrAdient spLitting, which includes the subsection ''FRUGAL as Optimization Framework''. We have also restructured the text within this section: to achieve a more coherent narrative flow, we moved the introduction of the general framework with reference to Algorithm 1 to the beginning of the section.

However, we believe that the presence of experiments later in this section remains appropriate. As FRUGAL is positioned as a general framework that supports arbitrary types of projections, State-Full and State-Free optimizers, these experiments are crucial for justifying our specific implementation choices. While future research may propose improved hyperparameter settings and components that could be incorporated into FRUGAL, we needed to identify and validate the optimal configuration from existing options for our main experimental evaluation. Therefore, maintaining these experiments within this section provides the necessary justification for our implementation choices and serves as a foundation for potential future improvements to the framework.

If the reviewer has additional suggestions regarding the structure of this section, we welcome their input and would be happy to incorporate them further.

Several notations (e.g., g~) are introduced without proper definitions, which assumes too much prior knowledge from readers. Additionally, concepts like smoothness and unbiasedness are only vaguely referenced and would benefit from clearer definitions.

We thank the reviewer for highlighting these opportunities for improvement. In response, we have incorporated definitions for $\tilde g$ and random sample $\zeta$, as well as a definition for the L-smooth function.

However, we have maintained the current formulation of Unbiasedness as presented in Assumption 1, as it is both precise and consistent with standard definitions in the literature (e.g., [1]).

We welcome any further suggestions from the reviewer regarding potential improvements to the theory section.

Including a full-parameter fine-tuning baseline in Table 4 would provide a valuable benchmark

We thank the reviewer for this suggestion. We have added the full-parameter fine-tuning baseline to Table 4. We would like to note that, following [2], rather than conducting these experiments ourselves, we have adopted the results from prior works.

Definitions for Full-Rank SVD/Random and Low-Rank SVD/Random are scattered across Table 1 and lack clear differentiation.

To improve the readability of Table 1, we have modified the column headers. Instead of a single 'Method' column, we have added two columns: 'Projection type' and 'Optimizes state-free subspace', which better reflect the distinctions between the algorithms compared in this table. Additionally, for consistency with the text, we have added clarifications in lines 206-208, 250-251, and 296-298.

If the reviewer believes this presentation can be further improved, we welcome their feedback.

Lastly, there are deviations from the primary algorithm, such as using column-wise projection instead of block-wise projection.

We appreciate the reviewer's attention to this discrepancy. To ensure complete coherence throughout the paper, we have revised the text at the end of this subsection accordingly. Please refer to lines 302-306.

We welcome any additional suggestions from the reviewer regarding potential improvements to this section of the paper.

We thank the reviewer for the response. We address their concern in detail below.

A key aspect of the proposed method is its strategy for determining which layers or parameters should be optimized using a stateful optimizer versus a state-free one. While Table 3 presents some results, it lacks an analysis of the selection of linear layers.

We respectfully disagree with the statement that this is a key aspect of our work. While our paper presents multiple contributions, the mentioned aspect is not central to our findings. The primary contributions are the reprojection of states and the updating of all network parameters. The reviewer's comment seems to address the lack of explanation regarding why we initially considered only Linear layers as candidates for the state-free subspace in all experiments except those presented in Table 3 and Table 13. To address this, we followed the setup from GaLore, where the authors also projected gradients onto a low-rank subspace solely for Linear layers. This approach is natural, given our primary objective: minimizing memory consumption. As model sizes increase, Linear layers in Transformers account for a growing proportion of the total parameters. For example, in LLaMA 3B, Linear layers constitute approximately 94% of all model parameters. Consequently, the memory demands of other layer types are comparatively negligible, which justifies focusing on Linear layers in this context. We have included this clarification in the paper on line 194 and footnote 2. We welcome additional feedback if the reviewer has further questions or if we have misunderstood their concern.

How do you propose bridging the gap between theoretical proofs and practical implementation in the context of optimizer selection?

We would like to reiterate that our theory provides the strongest justification compared to related work. As we argue, our theory is the closest to what is actually implemented when compared to related work. However, as discussed, the difference arises from fundamental challenges in analyzing adaptive algorithms such as Adam.

Additionally, our theory cannot be used to determine which parameters should be selected for a state-free optimizer versus a stateful optimizer. However, this limitation is not specific to our theory, as SGD is already the worst-case optimal algorithm. Consequently, we cannot achieve better results when considering combinations of SGDM and SGD. Nevertheless, our results are tight, as we recover the best-known outcomes for both SGD and SGDM under the analyzed setting.

In Table 15, SVD consistently outperforms other selection methods, albeit with potentially significant computational costs. This seems inconsistent with the paper’s claim that random projection offers superior performance. Clarification on this discrepancy is necessary.

There is no inconsistency here. We never claim that random projection offers superior performance in general. On the contrary, we assert that Random projection outperforms SVD when training without optimizing state-free subspace (see lines 200-208). Furthermore, in Table 1 and lines 296-301, we explicitly state that "SVD outperforms both RandK and Block projections" when training with optimizing state-free subspace. In the same paragraph, we explain our ultimate choice of Blockwise projection over SVD based on reduced computational and memory costs while maintaining comparable (just slightly worse) performance.

The results in Table 15 are fully aligned with our previous findings in Table 1 and our statements in lines 200-208 and lines 296-301. While SVD significantly underperforms compared to Random projection when training without optimizing state-free subspace (which further supports our hypothesis from Section 3.1), SVD shows slightly better results than Blockwise projection when training with optimizing state-free subspace.

We note that we compared explicitly against Blockwise projection rather than Random projection because they showed similar results in Table 1, leading us to focus on Blockwise projection in subsequent experiments as part of FRUGAL. We included these entries for completeness and to maintain consistency with Table 1's style.

We thank the reviewer for their valuable feedback. We are pleased that they appreciated the novelty of our idea, our empirical results, and extensive ablation studies.

Line 249 introduces state-free and stateful parameters but could provide more explicit explanation on the selection criteria. Are parameters randomly selected to each category? In that case the assumption is all the parameters are equally important for that iteration. The work could benefit from more detailed study on how to choose the parameters for state free updates.

We thank the reviewer for their insightful comment on the selection criteria for state-free and stateful parameters. Our primary contribution emphasizes updating all parameters, which is crucial for achieving state-of-the-art performance, as detailed in Sections 3.1 and 3.2. This approach contrasts with methods such as GaLore, which only updates a subset of parameters during each iteration.

In our work, we extensively explore various parameter selection strategies, as highlighted in Table 1 and the discussion in Sections 3.1 and 3.2. Among these, we consider both random parameter selection approaches (RandK, Random projection) and deterministic ones (SVD decomposition, blockwise projection with descending/ascending block changes). While a comprehensive study on the importance of individual parameters may be valuable, it is beyond the scope of this work. Such an analysis would require developing additional adaptive metrics for parameter importance assessment, and we leave this investigation for future work.

Nonetheless, we believe that our investigations in Sections 3-4 and Table 1, which examine 4 different types of state-free subspace selection, cover a sufficient range of possible projections. Furthermore, in Table 3, we conduct an additional study regarding the relative importance of different parameters. If the reviewer has specific suggestions regarding any other particular algorithms for state-free subspace selection, we would be happy to consider them and incorporate them into our paper.

The purpose of the density parameter (ρ) is not thoroughly explained, especially in relation to zero-density training. Please clarify whether zero-density training implies all parameters are state-free (i.e., trained exclusively with SGD). The selection of ρ is not mentioned in the algorithm as well.

Thank you for this question. Given that another reviewer raised a similar point, we have provided our answer in the general response.

But this paper seems to suggest the opposite. Do you see a space where these two ideas coexist? For example, low rank for certain tasks vs full rank for other tasks?

We thank the reviewer for this excellent question. Indeed, GaLore theoretically proves that the gradient is low-rank, and we confirmed this finding in our initial experiments - for a 60M model with hidden size of 512, the gradients stable ranks varied between 1 and 10, depending on the layer.

However, we believe our work does not contradict the low-rank nature of gradients. Our results show that although the gradient's Frobenius norm is predominantly concentrated in the leading singular values, utilizing the remaining part of the gradient is crucial for successful training. We propose the following potential explanation for this phenomenon: gradient magnitudes are generally proportional to weight magnitudes. However, even weights with small values can be critical for the overall LLM performance, as the neural network function is non-linear, meaning small weight changes can lead to significant output variations. Furthermore, while individual updates may be low-rank, the resulting weight matrix after training typically becomes high-rank, as demonstrated in Chen et al., 2021 (Figure 2).

Chen, Patrick, et al. "Drone: Data-aware low-rank compression for large nlp models." Advances in neural information processing systems 34 (2021): 29321-29334.

Introduce abbreviations for better readability. For example SGD as Stochastic Gradient Descent.

We are grateful to the reviewer for this comment. Indeed, we introduce SGD in the abstract without expanding the abbreviation. However, this was done to avoid overloading the abstract. When SGD is first mentioned in the main text of the paper, we provide its full form. We also believe that specifically in the case of SGD, using this abbreviation in the abstract without expansion is appropriate because Stochastic Gradient Descent is widely known and considered general knowledge.

If the reviewer can point out other instances in the text where we use abbreviations without proper expansion, we would be grateful for such feedback.

We thank the reviewer for acknowledging our response. However, as we have provided detailed explanations addressing each concern, we would appreciate clarification on what specific issues remain that justify maintaining the original score. We would be happy to further engage in discussion if there are some unresolved concerns.

General Inquiries

Also, multiple reviewers raised similar questions, so we addressed them in the general response.

Reviewers Cacw and f3bb asked to provide a formal definitions of state-full and state-free.

To enhance the readability and clarity of our paper, we have added more formal definitions of state-full and state-free subspaces and optimizers in the Introduction (see paragraphs in lines 077-090). If the reviewers believe these definitions could be further improved, we are open to suggestions.

Reviewers YtQB and f3bb were curious about the explanation of which certain parts of the parameters can be optimized with a state-free optimizer and why.

We thank the reviewers for raising this interesting question. In Table 3, Table 13, and lines 471-481 and 526-532, we provide detailed experimental investigation and discussion of which model modules are sensitive to the choice between Adam and other less advanced optimization algorithms.

As we write in lines 480-481, our results show the remarkable sensitivity of Logits to optimizer choice, which is aligned with findings from [1]. The authors of that paper hypothesized that this sensitivity is related to different columns of the Logits matrix corresponding to different tokens that appear with dramatically different frequencies in texts, and it is precisely here that the adaptive per-parameter learning rate enabled by Adam becomes crucial. However, we want to specifically note that this explanation remains a hypothesis and requires separate investigation to fully understand this phenomenon (for instance, similar token-related logic applies to the Embeddings matrix, yet it doesn't exhibit comparable sensitivity). Such investigation is an interesting direction for future research but lies outside the scope of our current work.

We would also like to emphasize that, to the best of our knowledge, our paper is the first to demonstrate that almost all modules in the model can be optimized using state-free optimizers. We believe this contribution already represents a significant advancement for further research into this phenomenon.

Reviewers 2Mhk and f3bb raised questions about the definition of $\rho$, particularly about the meaning of zero-density training with $\rho=0$, and how it differs from training all parameters using a state-free optimizer.

We would like to note that this parameter was introduced in lines 456-457 (previous revision) with reference to Appendix A.1, where it was described in lines 774-776. We also emphasize that training with $\rho=0$ is not equivalent to training the entire model with a state-free optimizer but rather indicates that only all Linear layers are trained using the state-free optimizer, while the remaining parameters are trained using standard Adam (this setup is similar to GaLore). This is also explicitly noted in the caption of Table 2 and in the paragraph dedicated to zero-density training in lines 471-481. Moreover, lines 471-481 and Table 3 are devoted to understanding why training the entire model with a state-free optimizer yields such poor quality, while training with $\rho=0$ does not degrade performance significantly.

To make the description of the density parameter $\rho$ clearer and more comprehensible, we have made several clarifications to the paper. Specifically, we expanded the introduction of parameter $\rho$ in lines 460-463, added its description in the caption of Table 2, included its mention for context in line 473, and extended the explanation of how $\rho$ relates to rank $r$ from GaLore in Appendix A.1 line 774-778.

We welcome any questions from the reviewers regarding the value of this parameter or suggestions about its description in the paper.

[1] Zhao et. el., Deconstructing what makes a good optimizer for language models. arXiv:2407.07972, 2024.