FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training
Authors: Philip Zmushko
Aleksandr Beznosikov
Martin Takáč
Samuel Horváth
Venue: ICML 2025 poster (2025)
Keywords: Memory-efficient training
Optimization
Full-Rank Update
Large Language Models
LLM
Pre-training
Fine-tuning
Primary area: optimization->large_scale_parallel_and_distributed
TL;DR: We create a memory-efficient optimization framework that performs full-rank updates by combining advanced methods like Adam with state-free methods like signSGD.

Abstract: With the increase in the number of parameters in large language models, the training process increasingly demands larger volumes of GPU memory. A significant portion of this memory is typically consumed by the optimizer state. To overcome this challenge, recent approaches such as low-rank adaptation (LoRA), low-rank gradient projection (GaLore), and blockwise optimization (BAdam) have been proposed. 
However, in all these algorithms, the *effective rank of the weight updates remains low-rank*, which can lead to a substantial loss of information from the gradient. This loss can be critically important, especially during the pre-training stage. In this paper, we introduce FRUGAL (**F**ull-**R**ank **U**pdates with **G**r**A**dient sp**L**itting), a new memory-efficient optimization framework. FRUGAL leverages gradient splitting to perform low-dimensional updates using advanced algorithms (such as Adam), while updates along the remaining directions are executed via state-free methods like SGD or signSGD. Our framework can be integrated with various low-rank update selection techniques, including GaLore and BAdam. We provide theoretical convergence guarantees for our framework when using SGDM for low-dimensional updates and SGD for state-free updates. Additionally, our method consistently outperforms concurrent approaches, achieving state-of-the-art results in pre-training and fine-tuning tasks while balancing memory efficiency and performance metrics.

FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training
Philip Zmushko 1
Aleksandr Beznosikov 2341 Martin Takáč
Samuel Horváth 5
Abstract
1. Introduction
In recent years, Large Language Models (LLMs) such as GPT (OpenAI, 2023) and LLaMA-3 (Dubey et al., 2024) have demonstrated remarkable performance across vari- ous disciplines (Brown, 2020; Yang et al., 2024; Romera- Paredes et al., 2024). However, a critical factor in achieving these results is the size of these models (Hoffimann et al., 2022). Increasing the number of parameters leads to higher computational and memory costs. For example, an bil- lion parameter LLaMA-3 model in 16-bit format requires 32GB just for parameters and gradients. Using the standard Adam optimizer (Kingma, 2014) adds another 32GB for m and statistics. Moreover, achieving high-quality results often requires 32-bit precision for weights and optimizer states (Zamirai et al., 2020), pushing memory requirements beyond even high-end GPUs like the A100-80GB.
With the increase in the number of parameters in large language models, the training process in- creasingly demands larger volumes of GPU mem- ory. A significant portion of this memory is typ- ically consumed by the optimizer state. To over- come this challenge, recent approaches such as low-rank adaptation (LoRA), low-rank gradient projection (oaL,re), and blockwise optimization (BAdam) have been proposed. However, in all these algorithms, the effective rank of the weight updates remains low-rank, which can lead to a substantial loss of information from the gradi- ent. This loss be critically important, espe- cially during the pre-training stage. In this paper. we introduce FRUGAL (uull-Rank Updates with GrAdient spLitting), a new memory-efficient op- timization framework. FRUGAL leverages gra- dient splitting to perform low-dimensional up- dates using advanced algorithms (such Adam), while updates along the remaining directions are executed via state-free methods like SGD or signSGD, Our framework can be integrated with various low-rank update selection techniques, in- cluding GaLore and BAdam. We provide theo- retical convergence guarantees for our framework when using SGDM for low-dimensional updates and SGD for state-free updates. Additionally, our method consistently outperforms concurrent approaches, achieving state-of-ther the-r--rrtttt results in pre-training and fine-tuning tasks while balancing memory efficiency and performance metrics.
Numerous research projects have been aimed at reducing these significant costs. These approaches include engineer- ing solutions like gradient checkpointing (Chen et al., 2016) and memory offloading (Rajbhandari et al., 2020), which do not change the training trajectory. There are also methods that adjust the training algorithm by decreasing the number of trainable parameters (Frankle & Carbin, 2018; Horváth et al., 2024) or their bit precision (Wortsman al., 2023), as well optimizer statistics (Dettmers et al., 2021; Shazeer & Stern, 2018; Zhang et al., 2024c).
Parameter- Fine-Tuning (PEFT) methods, such as LoRA (Hu et al., 2021) and Dora (Liu al., 2024b) reduce memory costs by training a relatively small number of pa- rameters compared to the size of the original model, while the remaining modules are frozen. This approach has proven effective for the task of efficient fine-tuning of pre-trained models. However, PEFT methods have fundamental limi- tation: parameter updates always lie in a low-dimensional subspace L, which prevents the use these methods for pre-training (Lialin et al., 2023) and may restrict their capa- bilities fine-tuning (Zhang et al., 2024a).
'Yandex, Russia "Moscow Institute of Physics and Technology, Russia "Ivannikov Institute for System Programming RAS, Russia "kookkovo Institute of Science and Technology, Russia Mohamed bin Zayed University Artificial Intelligence, UAE Correspon- dence to: Philip Zmushko <zmushko,,be ph. @ggmii. com.>
Recent works, such as GaLore (Zhao et al., 2024a). ReLoRA (Lialin et al., 2023) and BAdam (Luo et al., 2024) offer a solution to this problem. These methods enable higher-dimensional full-parameter learning by periodically
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

FRUGAL: Memory-Efficien Optimization by Reducing State Overhead for Scalable Training
Algorithm FRUGAL (State-Full, State-Free)
SGD, SignSGD
AdamW
Input:
model so with p parameter sets {0; RR Re=1 loss L, gradient projectors {Pr,))e= number of steps
FRUGAL
Update Dimension
1: for = K do
get data batch (xr y)
3: compute f [(ff(x)), y) {Forward)
4: for = 30,
at from Backward do
5: Iiull t (gis(gi), {Project Grad]
giree - 9: Ph.j (grmn,)) {Residual}
7:
5a, [P
se,] {Project state)
ReLoRA, GaLore, BAdam
Optimizer State Memory
8: Mfull, + State-Ful www..pppatteeeeeeeeeepppppaaaaet frill sa.)
9: Ufee. + State-Free wpdat Irre.
10: 0g 4- 0; + Pr- (wan,,) + Ufree,
Memory Usage
11: end for
12:
end for
Figure 1. FRUGAL reduces memory usage by splitting gradient updates into low-dimensional updates with advanced optimizers (c.g., AdamW) and using state-free methods (c.g., SignSGD and SGD) for the rest.
free optimizer, and show that FRUGAL matches the best-known convergence rate in many scenarios.
To verify the practical applicability of FRUGAL, we con- duct extensive experiments in popular real-world scenar- ios In these experiments, we pre-train LLaMA-like mod- els (up to IB parameters) the Colossal Clean Crawled Corpus (C4) dataset (Raffel et al., 2020) and fine-tune RoBERTa (Liu, 2019) on the GLUE benchmark (Wang, 2018). The results show that our method significantly outperforms previous memory algorithms while using less memory budget.
changing the optimizable low-rank subspace L. How- ever, even though these methods result in overall param- eter changes that are high-dimensional, the updates in each step remain low-dimensional. The dimensionality of the frozen subspace dim M dim L significantly exceeds dim L. The remaining information contained in the gradi- ent is not utilized for parameter updates. Nevertheless, this information can still be leveraged to train the model.
We present the FRUGAL framework, designed to bridge this gap. Our approach stems from a crucial observation: al- though memory constraints prevent using optimizers with auxiliary optimizer state such Adam (Kingma, 2014)
We demonstrate that only the Output layer in transformer- like models requires advanced optimizers like Adam, while other modules (including RMSNorms and Embed- dings) can use simpler methods like signSGD without significant performance loss. This opens up new possibil- ities for memory-efficient training and provides crucial insights into Transformers learning dynamics.
in the remaining subspace M, one still can update M using state-free optimization algorithms like Stochastic Gra- dient Descent (SGD) or signSGD (Bernstein et al., 2018). This solution allows for high-dimensions updates, which provides additional opportunities to explore the parameter space and improves convergence. We will further refer to the subspaces T. and M according the types of optimizers used for their updates- state-full and state-free.
2. Related work
Memory-eflicient full-parameter learning. Recent re- search has focused on reducing the memory footprint of LLM by decreasing the size of the optimizer states while maintaining their performance. Low-rank adaptation meth- ods, such LoRA (Hu et al., 2021), inject trainable rank decomposition matrices into linear layers, reducing memory requirements by optimizing only a few learnable adapters. ReLora (Lialin al., 2023) builds upon this by merging low- rank adaptations into the main model weights during train- ing, increasing the total rank of the update. BAdam (Luo et al., 2024) leverages Block Coordinate Descent for full- parameter training by switching active blocks during fine-
Contributions. We summarize the main contributions of our work as follows:
We present a new memory-efficien optimization frame- work that combines the use of advanced optimization algorithms for the state-full subspace with state-free algo- rithms for the complementary subspace. The framework supports various types of state-full optimizers, state-free optimizers, and different methods for projecting the gradi- ent onto the state-full full subspace.
We provide theoretical convergence guarantees for our framework. the proof, we consider the case with SGDM as the frlle-full optimizer and SGD as the state-
'The https:///nnnee 4open wwciiinnee//RGGA.L....
code is available
2

---

FRUGAL: Memory-Effiiient Optimization by Reducing State Overhead for Scalable Training
tuning. MicroAdam (Modoranu et al., 2024) compresses gradient information before feeding it into the optimizer state, significantly reducing the memory footprint while enabling full parameter learning with eтror feedback mecha- nisms. Gal oree (Zhao et al., 2024a) maintains full parameter learning by projecting gradients onto a low-rank subspace using SVD decomposition, storing optimizer states in this reduced space. However, while these methods effectively reduce memory overhead, they all perform low-rank up- dates at each iteration. In contrast, our approach utilizes all available gradient information to perform full-dimensional updates at each optimizer step, offering novel perspective on memory-efficienn optimization for LLM.
Table I. Comparison of different projection and state-free subspace optimization strategies on pre- pre-training LLaMA-130M on C4 with AdamW the state-full algorithm.
Projection
Optimizes state- free subspace
Validation perplexity +
type
4k
40k
200k
SVD
No
39.75
24.38
21.11
Random
No
42.31
23,55
20.01
Random
Yes
37.26
21.53
18.64
SVD
Yes
33.96
21.01
18.35
RandK
Yes
36.38
21.25
18.63
Blockwise
Yes
37.20
21.42
18.60
AdamW
33.95
20.56
18.13
However, we note that there also exist several concurrent works- Fira (Chen et al., 2024a), LDAdam (Robert al., 2024), and Adamem (Vyas et al., 2024) that also adopt full-dimensional updates for similar goals. See Appendix B for a comparison and detailed discussion.
induced errors, Kunstner et al. (2023) demonstrated that this superiority persists even in full-batch training. They proposed a new hypothesis suggesting that Adam's key suc- cess factor is related to its similarity to signSGD (Balles & Hennig, 2018; Balles et al., 2020), and both Kunstner et al. (2023) and Zhao et al. (2024b) showed that the signed descent with momentum reduces the performance gap with Adam. In contrast, to the best of our knowledge, we are the first to train the majority of language model parameters using signSGD without momentum, achieving minimal loss in quality. This approach further demonstrates the effec- tiveness of sign-based methods for LLM training, paving the way for more efficient and scalable optimization strategies.
Other memory-efficient optimization. Several other meth- ods have been proposed to reduce the memory footprint of optimizers. AdaFactor (Shazeer & Stern, 2018) attempts to mimic Adam's behavior while reducing memory usage through factorization of the variance matrix Adam- mini (Zhang et al., 2024c) further reduces memory by stor- ing only one value per block. Dettmers et al. (2021) and Li et al. (2024) decrease memory footprint by quantizing optimizer states to lower-precision representations. Lv et al. (2023) proposed to reduce memory by fusing the backward operation with the optimizer update. Notably, these ap- proaches are orthogonal to our method FRUGAL and can be combined with it for further memory efficiency.
3. Empirical Analysis and Motivation
3.1. The importance of exploring the entire space during the training process
Block Coordinate Descent. Block Coordinate Descent (BCD) is a well-established optimization method with a rich history in mathematical optimization (Ortega & Rheinboldt, 2000; Tseng, 2001; Richtárik & Takáë, 2014; 2015a;b). In recent years, a specific instance of BCD, known as layer- wise learning, has been applied to deep learning. Notable examples include Luo et al. (2024); Pan et al. (2024), which leverage this approach for LLM fine-tuning, To the best of our knowledge, our work presents the first theoretical anal- ysis of an extended BCD framework (Section 5) where the remaining layers are also updated with a different algorithm. This novel approach extends traditional BCD techniques, opening new avenues for full model optimization.
In recent work, Zhao et al. (2024a) proposed GaLore, an optimization method based on projecting the gradient matrix G of each Linear layer² onto a low-dimensional subspace. To obtain the projection matrix P. they use the SVD decom- position of G., which is recomputed with frequency T. The vectors or rows of G are projected onto the first left or right singular vectors, respectively. This approach has theo- retical foundations: the first r singular vectors correspond to the first r singular values and, therefore, should better utilize information from the spectrum of G.
Given the computational burden of SVD decomposition, natural question arises about the possibility of employing random semi- ceni-orthogonal projection matrix R as an alter- native to projecting onto the first r singular columns with P. Surprisingly, while the SVD decomposition provides better initial performance, the random projection proves superiority in long-term training, yielding significant im- provements. As illustration, we took the preitrain
Sign-based methods for training language models. Since its introduction, Adam has become the de facto primary opti- mization algorithm, demonstrating superior practical results compared to SGD-based algorithms across various deep learning tasks. This difference is particularly noticeable when training Transformers on language tasks. While Zhang et al. (2020) hypothesized that Adam outperforms SGD in this setup due to the heavy-tailed distribution of sampling-
Since Linear layers contain most parameters and require most memory, we primarly focus on them. "See Section 6.1 for detailed description and discussion.
3

---

FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training
1000 and 2000
1000 and 10000
1000 and 99000
Random
Frequency
20
13-
20-
10-
10
Principal cosines
Principal cosines
Principal cosines
Principal cosines
Figure 2. Histograms of principal angle cosines. The first three are taken between P. and Py from different iterations t and t'. P is obtained from the truncated SVD decomposition of the gradient G of the Key projection from the 5th layer. The last histogram is taken between two random semi-orthogonal projections R and R'' for comparison comparison.
of a 130M model with LLaMA-like architecture on the C4 dataset. The results are presented in the first part of Table 1, where we compare SVD and Random projections.
updating parameters is not readily obvious. Using algo- rithms like Adam (Kingma, 2014) is not an option due to the memory overhead they introduce, which exactly what we aim to avoid. An alternative approach is to use state- free optimizers such SGD signSGD (Bernstein et al., 2018). Unfortunately, SGD has been shown to be ineffective for training transformer models, as shown in Zhang et al. (2020); Pan & Li (2023).
To investigate this phenomenon, we pre-trained the LLaMA- 60M model and collected gradients G2 from different it- erations /. for examination. We evaluated the similarity of the projection matrices by calculating the principal angles between the projections P, from different steps. Similarly to the observations in Q-Galore (Zhang et al., 2024d), we found that these projections show minimal change during training; see Figure 2 for details.
Nevertheless, a recent study Zhao et al. (2024b) suggests a promising methodology: while SGDM generally does not work well with transformers, using SGDM for the majority of parameters and Adam for a selected subset can lead to effective training. This raises the question: Could hybrid approach using SGD or signSGD instead of SGDM be vi- able? If the key subset of parameters is handled by advanced algorithms, can the other parameters be trained effectively with state-free optimizers?
Here, we take the projection matrix of k-proj from 5-th layer and plot histograms of the cosine of the principal angles between pairs P. and Pe from different iterations. For comparison, we also include the random projections on the right. As can be seen, the distributions of cosines differ significantly for P. and for R., While R2 feature no angles with cosines higher than 0.9, the top 57 cosines for P. surpass 0.9, even for gradients 1000 steps apart.
To address this question, we conducted an experiment on ALLMA-130M, where utilized the Adam (Kingma, 2014) for state-full parameters and signSGD (Bernstein et al., 2018) for state-free parameters'. Once again we used Ran- dom projection and highlighted the result in the second part of Table I. Full-rank updates significantly enhance perfor- mance, approaching the efficiency of the memory memory-intensive intensive Adam optimizer. These findings underscore the potential of state-free algorithms for updating substantial portion of the parameter space, paving the way for efficient and scal- able optimization methods that deliver high performance without the significant memory costs traditionally associ- ated with state-of-the thee--rr-- optimizers.
This leads to the conclusion that although the SVD decom- position generally better captures the information contained in G., the original GaLore algorithm updates the weights only in a small subspace. We hypothesize that training with random projections yields superior results due the more extensive investigation of the optimizable space during the training process This finding indicates that to achieve better convergence, it is important to find optimization algorithmss that explore the entire space during the training process.
3.2. Advantage of the Fank-Rank Updates
The insight from Section 3.1 suggests that the training of language models performs significantly better when the entire parameter space is explored during the training pro- cess. Given the importance of updating parameters in all directions, this poses the question: Is it optimal to use low- rank updates, as employed by methods such as GaLore, ReLoRA, and BAdam? The effective rank of low-rank up- dates is significantly smaller than the full dimensionality of the parameter space, inevitably leading to a loss of valuable information contained in the gradient.
4. Full-Rank Updates with GrAdient spLitting
General framework. The setup outlined in the conclusion of Section 3.2 results in general framework for memory- efficient optimization. It operates as follows: the entire space is partitioned into state-full and state-free subspaces. The state-full subspace is updated using an advanced al- gorithm, while the state-free subspace updated using a state-free method. After a certain number of steps, the state-
*See detailed description of the setup in Appendix A.I
However, the method to leverage the full-rank gradient for
4

---

FRUGAL: Memory-Effiiient Optimization by Reducing State Overhead for Scalable Training
full subspace is changed to better explore the optimization space. A formal description is presented in Algorithm 1.
Algorithm 2 FRUGAL (SGDM, SGD)
Input: momentum weight 3 E [0, 1), initialization æ² E R and m² = 0, step sizes {at > 0) momentum set
We note that this framework allows for variation not only in the state-full optimizer but also in the choice of projection and state-free optimizer. However, determining the optimal state-free optimizer and the projection method onto the state- full subspace is not readily apparent. In this section, we strive find the optimal configuration.
Ju C [d] for k- = 1,2.....
1: for k = 1,2....00 do
2: g + vfcx(=*)
m
k-1
if Jk. otherwise;
3: my (1 - B)0J +
State-free optimizer. We conducted preliminary exper- iment using different state-free algorithms to choose be- tween SGD and signSGD (Bernstein et al., 2018). Table 10 shows that signSGD outperforms SGD, leading us to favor signSGD. We attribute this performance to the similarities between signSGD and Adam (Kingma, 2014), as noted in Balles& & Hennig (2018); Balles al. (2020); Kunstner et al. (2023). Additionally, signSGD produces updates of similar magnitude to those generated by Adam, which simplifies the calibration of the learning rate for state-free parameters.
m gj
if j E Jk.
4: 0f
otherwise;
5: x x+1
x²
a ahaa
6: end for
when the projection changes from one round to another. However, our preliminary experiments with RandK selec- tion showed that resetting states performs comparably to projection. Therefore, we could replace this projection with state resetting when the projection changes, which also aligns with blockwise subspace selection. However, either resetting or projecting states is important since want pro- jected gradients and optimizer states to reside in the same space. For instance, GaLore ignores this step, which Ieads to degraded performance when projections are updated fre- quently; see Appendix D and Section 6.4 for details.
Projection type. When selecting a projection method, it crucial to strike balance between quality and memory efficiency. When using SVD decomposition for projection matrices, as in GaLore (Zhao al., 2024a), the method bet- ter preserves the information embedded in the gradient but requires additional memory for storing projection matrices and computational resources for performing the SVD. To reduce computational demands, one could employ random coordinate projection denoted as RandK, but this requires additional memory or recomputation'. A more structured alternative is to select not random entries but entire random columns or rows. The most aggressive approach follows the method from BAdam, wherein an entire block is cho- sen as the state-full subspace. The performance results obtained with all these variants are presented in the second part of Table 1. SVD slightly outperforms both RandK and Block projections, demonstrating comparable performance. Nonetheless, a downside is the increased compute and mem- ory demand from SVD. Therefore, we opt for the blockwise selection, as it is the most memory-efficient- requiring only the storage of active block indices.
5. Theoretical Results
For the theoretical analysis, we consider the case where the State-Free optimizer is SGD and the State-Full optimizer SGD with momentum (SGDM). For the projection, use coordinate-win projection. This special case of FRUGAL is provided in Algorithm 2. We minimize the objective
minzeR= (f(x) == Ex+ (z))
(1)
where we access of via a stochastic oracle that takes x as input and returns (f(+()),, V[{+
5.1. Notation and Preliminaries
In experiments in Section 6, we use a specific variant with AdamW as the Stale-Full optimizer and signSGD as the State-Free optimizer. We primarily employ blockwise pro- jection but switch to column-wise projection when the num- ber of parameters in any single block exceeds memory bud- get, as detailed in Section 7. In addition, PyPorch-like pseudocode of our framework is presented in Appendix G.
We use || || for the vector fn-norm, and (:.)) stands for the dot product. Let g denote the full gradient of ʃ at x* i.c., g == Vf(=*) g denote the stochastic gradient g² v fax( for random sample ç* and f == min f(x). We use subscript j to denote the j-th coordinate. We call a function I-smooth if it is continuously differentiable and its gradient is Lipschitz continuous:
For Line 7, state projection, in Algorithm 1, we note that if the projection does not change, i.e., Pk,: = P then Thus, we only need to project states
( < -| Bl|-
Pa. (t.((Pp1..ss))
Assumption 5.1. We make the following assumptions, which are standard in non-convex stochastic optimization; see (Liu et al., 2020).
'See Appendix C for discussion on the memory requirements for different projection methods.
5

---

FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training
Table Comparison of validation perplexity and memory estimation for various optimization methods across I. L.MMMA model scales trained on C4. We also indicate the additional memory overhead introduced by the optimization algorithm. The values are calculated assuming that each float value occupies bytes (float32). P denotes the proportion the Linear layer parameters the state-full subspace. Note that Embeddings, RMSNorms, and Oulput layer are always trained with Ad
60M
130M
350M
IB
AdamW
22.73 (0.43G)
18.13 (1.00G)
14.43 (2.74G)
12.02 (9.98G)
GaLore, P 0.25
25.68 (0.30G)
21.11 (0.54G)
16.88 (1.10G)
13.69 (3.41G)
BAdam, P 0.25
24.86 (0.29G)
20.34 (0.52G)
16.41 (1.05G)
13.75 (3.23G)
FRUGAL, p 0.25
23.59 (0.29G)
18.60 (0.52G)
14.79 (1.05G)
12.32 (3.23G)
FRUGAL, = 0.0
24.06 (0.24G)
18.90 (0.37G)
15.03 (0.49G)
12.63 (0.98G)
Training tokens
20B
20B
24B
30B
Number of iterations
200k
200k
240k
300k
1. Smoothness: The objective fxx) in equation 1 is L- smooth (Equation (2)).
Table 3. Perplexity and memory consumption (weights, gradients and optimizer states) of different size LLaMA models pre-trained on C4 100k iterations (10B tokens) using AdamW with pure 6 of mixed precision.
2. Unbiasedness: At each iteration k, ge satisfies Be-[ō"]= =g g".
Model size
Format
Memory
Perplexity
3. Independent samples: The random samples {c* are independent.
175M
Mixed Precision
2.0GB
17.43
350M
Pure bf16
2.IGB
17.75
4. Bounded variance: The variance of
53 with respect
350M
Mixed Precision
4.2GB
15.16
to & satisfies Varg= (Ø)) = Fxc+ [10 - 9 [I2] S o3 for
1.3B
Pure b
7.7GB
16.51
some of > O. We denote o² = Σラ= of.
same rate if p ( B) or p = O (3). Finally, in the worst case (e.g.. Jk is deterministic and 0 < |J&] d), we have P = 0 and Pl = 1. Thus, the rate becomes O (1/ka + Laao" +Lad"/1-3 //-////1)) 9), which is worse by factor of 1/1-9. However, this is expected since the bias from momentum is not outweighed by the variance reduction effect, as only the coordinates with momentum enjoy reduced variance; see Lemmas E.2 and E.3 in the appendix for details.
Finally, we define the probability that index j E Jn selected, conditioned on the prior iteration ko- 1, as P = Prk--lj Jzl. Other useful quantities are plaax = max (n3) and p p == min (1/}.
5.2. Convergence of Algorithm 2
Below, we present the main convergence theorem.
6. Pre-training experiments
Theorem 5.2. Let Assumption 5.1 hold and a = a S
B KT- Then, the iterates of Algorithm satisfy
In this section, we evaluate the performance of FRUGAL on the language models pre-training.
1 Σ [11lg1111= o( (((x))-
ka
6.1. Comparison to existig memory-efficient algorithms
Lao? (1+ Pinax (1 pru )3
To begin, we compare our framework with existing memory- efficient methods across four sizes of LLaMA-based archi- tectures: 60M, 130M 350M, and IB.
(1 - B)
where phinn = 1 E Piii and phnnx = (axiE(M) (a fPiaxt }.
Setup. The core setup for pre-training is taken from Zhao et al. (2024a). We utilize LLaMA -base (Touvron et al., 2023a) model architectures and train them on the Colossal Clean Crawled Corpus (C4) dataset (Raffel et al., 2020). The C4 dataset is intended for pre- pre-training, making this setup good approximation of real-world applications. A detailed description of the setup can be found in Appendix A.I.
The proof is deferred to Appendix E. Let us analyze the obtained result. Firstly, if J4 [d] or J/+ = 0, Algorithm 2 becomes SGDM and SGD, respectively. In this case, we = 0 for SGD. There- fore, the resulting rate is O (1/Ka + Loo²), which recovers the best-known rate for both SGD and SGDM under these assumptions (Liu et al., 2020). Furthermore, if at each step each coordinate is sampled independently with probability p, we have P fl = p. Therefore, we recover the
have Ph. = 1 for SGDM and P
However, we made several critical modifications compared to Zhao et al. (2024a) to align the experimental setup with
6

---

FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training
Table 4. Perplexity of 1.L.MA-330M models pre-trained on C4 for 100k iterations (10B tokens). The leftmost column indicates the modules moved to the state-free set and trained using signSGD. The results show that Output layer, unlike Embeddings and RM- SNorms, are exceptionally responsive to the choice of optimization algorithm from AdamW to signSGD.
Table 5. Pre-training I.a.MA 3B on C4 dataset for 300K steps.
Validation perplexity for different iterations is reported.
Method
100k
200k
300k
AdamW
14.2
12.25
10.93
FRUGAL, P 0.25
14.33
12.42
11.07
FRUGAL, p 0.0
14.78
12.76
11.35
State-free modules
Perplexity +
Linear (FRUGAL p p= 0.0 from Table 2)
20.02
Precision training for its illustrative value in understand- ing each method's potential. However, for completeness, we also conducted experiments in pure 6 format, detailed in our ablation study Section 6.4.
Linear, RMSNorms
20.07
Linear, Embeddings
20.48
Linear, Embeddings, RMSNorms
20.55
Linear, Output layer
34.66
Baselines. We use the following methods as baselines:
practical training scenarios. Below, we discuss each modifi- cation and provide a detailed rationale for these decisions.
Full-rank Training. Training using memory-inefficient AdamW (Loshchilov, 2017). Weights, gradients, and statistics are stored and computed for all parameters. This serves as an upper bound for model performance.
Training Duration. The training approach in Zhao al. (2024a) aligns with the empirical rule from seal- ing laws (Hoffmann et al., 2022), which suggests using approximately 20 times the size of the model in tokens for training. However, this number of tokens is far from achieving convergence. In practice, models are typically trained for significantly longer periods (Touvron et al., 2023b; Zhang et al., 2024b). One reason for this discrep- ancy is that the original scaling laws do not account for the inference of the model after training (Sardana & Frankle, 2023). For our experiments, we chose 200k steps for the 60M and 130M models, 240k for the 350M model, and 300k for the IB and 3B models.
GaLore. Zhao et al. (2024a) proposed GaLore, a memory- efficient optimization algorithm that uses a low-rank pro- jection of gradient matrices G. Every T steps, the current gradient matrix G; is used to compute the projection matrix via SVD decomposition. The gradient is then projected onto the low-rank space, where the optimization step is performed. Subsequently, the resulting low-rank update is projected back into the full-rank space and added to the weights W.
BAdam. Luo et al. (2024) proposed a block coordinate descent (BCD)-type optimization method termed BAdam. The parameters are divided into blocks, which are then updated one by one using AdamW. The optimized block is changed every T steps. Although this method was initially proposed only for fine-tuning, it is the closest method to our F FRUGAL. Unlike BAdam, in our algorithm, state-free blocks are not frozen but are updated using signSGD.
Mixed Precision. Pure 16-bit training has been shown to potentially compromise model convergence and ac- curacy (Zamirai et al., 2020). This degradation occurs because formats such f or bfloat16, used to store master weights, lack the numerical precision needed for accurate and fine-gree fineggrained weight updates. Consequently, mixed precision training has become a more common ap- proach for training language models (Le Scao et al., 2023; Almazrouei et al., 2023). Moreover, even when training with fp8, the master weights are typically stored in fp32 format (Liu et al., 2024a).
Other Algorithms. Among other relevant methods, ReLoRA (Lialin et al., 2023), MicroAdam (Modoranu et al., 2024), Fira (Chen al., 2024a), LDAdam (Robert et al., 2024), and Adamem (Vyas et al., 2024) can also be highlighted. However, we did not include them for comparison here for the following reasons: 1. ReLoRA: This method was evaluated in (Zhao et al., 2024a), where significantly underperformed compared GaLore. 2. MicroAdam: Its current implementation only supports bfloat16 16 master weights, whereas our main experiments conducted with mixed precision. 3. Fira, LDAdam, and Adamem: These methods are concurrent works that were published during the final stages of our work. Accord- ingly, the majority of our experiments were completed before we became aware of these recent developments.
Our experimental results strongly support the importance of precision choice: in Table 3 we show that adopting pure b training led to such sigmificant performance degradation that doubling the model size failed to com- pensate for it, effectively negating any memory benefits from reduced precision storage. While training in pure 16-bit format is also possible, stochastic rounding (Gupta et al., 2015; Zamirai et al., 2020) is often employed to mitigate the aforementioned issue. Given that the goal of this research is to identify the optimal optimization algorithm, we deemed it more appropriate to compare optimizers in a transparent and stable setup that does not require auxiliary tricks. Hence, we primarily used Mixed
Main results. The results of our experiments are presented in Table 2, which includes both validation perplexity and
7

---

FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training
Table Evaluating PRUGAI for memory-efficient fine-tuning RoBERTa-Base on GLUE benchmark. Results represent the mean and
standard deviation across 3 independent runs. Upper 1 is better.
Method
Modules Rank
CoLA
STS-B
MRPC
RTE
SST2
MNLI
QNLI
QQP
Avg
Full-parameter
63.6
91.2
90.2
78.7
94.8
87.6
92.8
91.9
86.4
LoRA
QV
8
63.8
90.9
89.1
79.2=111
94.8
87.6
93.1
90.6
86.1
GaLore
All
8
|60.02a 90.8
79.0
79.7
94.9
87.6
93.3
91.1 2.1
85.8
GaLore
QV
8
56.11.
90.8
88.1
74.7=19
94.3
86.61.1
92.61.1
89.4
84.1
FRUGAL
QV
8
64.5 91.1
89.2
82.4
94.8
87.41.
92.8
91.4
86.7
FRUGAL
None
64.8 0
91.1
89.1
81.62.6
94.9 87.31.1 92.8
91.31.1
86.6
memory footprint estimations for each method. We com- pared all memory-efficient methods under the same memory budget with a density p = 0.25. Here, P refers to the propor- tion of Linear layer parameters belonging to the state-full subspace. Similarly to GaLore, non-Linear modules (Em- beddings, RMSNorms, Output layer) are optimized with AdamW. See Appendix A.I for details.
6.3. LLaMA 3B training
To demonstrate the practical viability of method for large-scale applications, we evaluated F FRUGAL against AdamW on the pre-training of the LLaMA 3B model. Due to computational constraints, we conducted a single training run of 300k steps using a cosine learning rate scheduler with 10% warmup steps. We used learning rate of Se-4, weight decay of 0.1, and gradient clipping of 1.0, with other hyperparameters consistent with Appendix A.1. The re- sults in Table 5 confirm that FRUGAL successfully scales to billion-paree parameter models without performance degradation, making it viable option for industrial-scat applications.
We conducted grid search to determine the optimal learn- ing rate for Adam W. which we then applied to FRUGAL and BAdam (Luo et al., 2024). For GaLore (Zhao et al., 2024a), we found that using this same learning rate produced better results than the originally suggested rate. This discrepancy might be attributed to our experiments involving a signifi- cantly larger number of training steps than those for which GaLore's 's original learning rate was optimized.
6.4. Ablation study
We also conducted additional experiments to verify the ro- bustness of our framework to various hyperparameters.
Table 2 demonstrates that FRUGAL significantly outper- forms memory-efficient baselines across all model sizes with the same memory budget, coming close to the perfor- mance of AdamW.
First, we began by evaluating different model architectures. Experiments with GPT-2 124M (Radford et al., 2019) in Ta- ble 12 show that FRUGAL maintains its strong advantage over memory-efficiend baselines, albeit with somewhat wider gap to AdamW. Second, an ablation study on the state-full subspace update frequency T in Table 14 shows that the performance keeps improving up to T' = 200. We note that, unlike in Zhao et al. (2024a), the perplexity does not decrease significantly even when reducing the update frequency to T 10 (~ 0.2 drop vs. 4. drop for Ga- Lore). A detailed explanation for this result can be found in Appendix D. After that, when using other schedulers, the performance gap between FRUGAL and baselines remains consistent, as shown in Tables 15 and 16. Table 8 shows that the same holds for B3 = 0.95- - another popular value for the second moment decay parameter in AdamW-likk meth- ods. FRUGAL also provides improvement over baselines for other state-full optimizers, as can be seen in experiments with Lion (Chen et al., 2024b) presented in Table 11. Then, the results of the training in pure 6 are presented in Table 9, demonstrating consistency with our main exper- iments in Table 2, i.e., FRUGAL significantly outperforms the baselines across these variations. We also conducted experiments to show how perplexity changes with varying p. and the results are presented in Table 17. Finally, we
6.2. Zero-density training
Table 2 also reveals a surprising result: FRUGAL with p= 0.0 outperforms both GaLore and BAdam, even when these competing methods use a higher density of p = 0.25. Essentially, for FRUGAL with p 0.0, the parameters are divided into two parts - a state-full part consisting of the Embeddings, RMSNorms, and Output layer, and a fate-free part consisting of all other parameters. This division remains fixed throughout the training. We conducted additional ex- periments to determine the maximum subset of parameters that can be trained with state-free optimizer without signifi- cant quality degradation. We systematically moved different combinations of the Embeddings, R.MNNomss and Output layer from the state-full to the state-free set and observed the results during the training of LLaMA-130M, Table 4 reveals that the Output layer demonstrates dramatically higher sensitivity, with changes to its optimizer resulting in severe performance degradation. This finding aligns with results from Zhao et al. (2024b), where the authors demonstrated that most parameters can be trained using SGDM, but the Output layer require training with AdamW.
8

---

FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training
Table 7. Accuracy 1 comparison of memory- efficciett methods on LLaMA 3.1-8B across commonsense reasoning tasks.
Method
Rank BoolQ
PIQA SIQA HellaSwag WinoGrande ARC-e ARC-c OBQA
Avg
LoRA
32
73.39
88,19
79.48
95.13
83.50
90.24
79.52
84.00
84.18
GaLore
32
74.11
88,36
78.97
94.88
83.82
91.08
80.2
85.80
84.65
FRUGAL
32
74.25
88.47
79.12
95.15
86.11
91.37
79.86
84.80
88.89
FRUGAL
71.44
88.47
80.40
94.74
86.35
90.95
81.14
86.00
18.94
conducted an experiment compare different strategies for selecting state-full blocks during training. The results in Ta- ble 13 show that there is no significant difference between random and structured block selection.
the choice of optimizer, which can be seen in Table 19 where the model's performance significantly deteriorates when using signSGD for classification head optimization.
7.2. Fine-tuning LLaMA on commonsense reasoning
These experimental results validate that our framework's superiority is resilient to hyperparameter variations.
In addition to our experiments with RoBERTa, we con- ducted experiments on LLM fine-tuning to evaluate our framework's 's performance in this practically important area. To assess this capability, we chose L.LaMA 3. 3-1 (Dubey et al., 2024) and commonsense reasoning benchmark, as this domain represents a fundamental capability requiring both factual knowledge and logical inference. This benchmark includes 8 subtasks each containing its training and test splits. Following the setup from Hu et al. (2023) we train the model on a single combined Commonsensel 70K dataset (Hu et al., 2023). We refer readers to the original paper for detailed descriptions of these tasks and the con- struction of the Commonsense 70K dataset.
7. Fine-tuning experiments
7.1. Fine-tuning RoBERTa on GLUE
We evaluated the performance of our framework in memory- efficient fine-tuning using the GLUE benchmark (Wang, 2018), a widely-used collection of tasks for evaluating lan- guage models. Following the approach from Zhao et al. (2024a), we fine-tuned RoBERTa-base (Liu, 2019) using LoRA (Hu et al., 2021) and GaLore as baselines for com- parison. We adhered the setup described in LoRA, where low-rank updates of rank 8 were applied only to the Q and V matrices. See detailed description in Appendix A.2.
Following the experimental protocol from Hu al. (2023), apply memory efffcieeeficceent methods to the same parameter subsets: the Q K, V. Up, and Down projection matrices. We used the same hyperparameter configuration as in the original work, except for the learning rate, which we varied across [5e-6, 1e-5, 2e-5, 5e-5, Ic-4, 2c-4] for each algorithm to ensure optimal performance. The results in Table show that FRUGAL again slightly outperforms both LoRA and GaLore in the average accuracy across all 8 tasks. Remark- ably, our method attains this advantage even with p O (effectively signSGD), requiring zero memory for optimizer state storage while delivering better accuracy.
For this experiment we opted for columnwise selection of active parameters. This transition from blockwise to colum- nwise selection was necessary to maintain comparable mem- ory usage across methods, as the number of trainable pa- rameters in LoRA with rank 8 is approximately 2.5 times fewer than the number of parameters in any RoBERTa ma- trix. For the same reason, we did not include comparisons with BAdam (Luo et al., 2024) in this setup.
The results are presented in Table 6. Since the LoRA setup adds trainable adapters only to the Q and V matrices, while the GaL code uses all modules as projectable parame- ters, we conducted experiments in both setups. The results demonstrate that FRUCAL significantly outperforms GaLore and shows comparable results to LoRA.
8. Conclusion
In this work, we introduce a new memory-efficient opti- mization framework, FRUGAL, Within this framework, the optimization space is divided into two subspaces: the first is updated using a state-full algorithm such as Adam, while the second is updated using state-free algorithm such as signSGD. We prove theoretical convergence guarantees for our framework with SGDM serving as the state-full algo- rithm and SGD as the state-free algorithm. In experiments involving pre-training and fine-tuning of language models, FRUGAL outperforms other approaches.
As in Section 1. we conducted additional experiments with FRUGAL using p = 0.0. In this setup, only the classification head is trained using AdamW, W. while the embedding parame- ters remain frozen, and the remaining parameters are trained using signSGD, The results demonstrate that this training approach barely compromises performance compared to FRUGAL with rank 8, and still outperforms GaLore.
Similar to our findings in Section 1. we observe that the classification head parameters are particularly sensitive to
9

---

FRUGAL: Memory-Effiiien Optimization by Reducing State Overhead for Scalable Training
Acknowledgements
Chen, X., Feng, K., Li, C., Lai, X., Yue, X., Yuan, Y., and Wang. G. Fira: Can achieve full-rank train- ing of Ilms under low-rank constraint? arXiv preprint arXiv:441001623, 2024a.
The work on the final version was conducted at Moscow Institute of Physics and Technology and was supported by grant for research center in the field of artificial intelligence, provided by the Ministry of Economic Development of the Russian Federation (agreement No. 139-15-2025-013, dated June 20, 2025, subsidy identifier 000000C3 BANNS333992242000022 The work was partially conducted while Philip Zmushko visited Mohamed bin Zayed University of Artificial Intelli- gence (MBZUAI).
Chen, X., Liang, C.,, Huang. D., Real, E.,, Wang, K., Pham, H., Dong, X., Luong, T., Hsieh, Lu, Y., et al. Sym- bolic discovery of optimization algorithms. Advances in neural information processing systems, 36, 2024b.
Dettmers, T., Lewis, M., Shleifer, S., and Zettlemoyer, L. 8- bit optimizers via alock-wise quantization. arXiv preprint arXiv:2110.02880 2021.
Impact Statement
Our work makes large-scale model training more accessible to the broader research community by reducing memory requirements. This reduction in computational demands not only decreases the financial cost of training but also potentially lowers the environmental impact. The proposed methods enable researchers with limited resources partic- ipate in large-scale machine learning research, potentially accelerating progress in the field through wider community involvement.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The llama 3 herd of models. arXiv preprint 2rX7::2407.21788 2024.
Frankle, and Carbin, M. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprin! FrXXiv:803.03666 2018.
Gupta, S., Agrawal, A., Gopalakrishnan, K., and Narayanan, P. Deep learning with limited numerical precision. In International conference on machine learning, pp. 1737- 1746. PMLR, 2015.
References
Almazrouei, E., Alobeidli, H., Alshamsi, A., Cappelli, A., Cojocaru, R., Debbah, M., Goffinet, É.,, Hesslow, D.,, Lau- nay, J., Malartic, Q., et al. The falcon series of open lan- guage models. arXiv preprint arXiv:2311. 2023.
Hao, Y., Cao. Y., and Mou, L. Flora: Low-rank adapters are secretly gradient compressors. arXiv preprint arXiv:2402.03293, 2024.
Balles, L. and Hennig, P. Dissecting adam: The sign, magni- tude and variance of stochastic gradients. In International Conference on Machine Learning, pp. 404-413. PMLR, 2018.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 0...2223...55.. 2022.
Balles, L., Pedregosa, F., and Roux, N. L. The geometry of sign gradient descent. arXiv preprint arXiv:2002.08056, 2020.
Horváth, S., Laskaridis, S., Rajput, S., and Wang, Mae- stro: Uncovering low-rank structures via trainable decom- position. In Forty-first International Conference Ma- chine Learning, 2024. URL https::/ settps://openreviee net/ forum? d=
Bernstein, J., Wang, Y.-X., Azizzadenesheli, K., and Anand- kumar, A. signSGD: Compressed optimisation for non- convex problems. In Dy, and Krause, A. (eds.), Pro- ceedings of the 35th International Conference Ma- chine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 560-569. PMLR, 10-15 Jul 2018. URL ht..ttpp:/// http::::///ooceeeinie http///pooeeen....ee/ mlr www.l..ppres/veee. press/v80/ bernstein88e laa.inn8a..... html.
7b deddd=-jjymiii yambg4xx.
Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., and Chen, W. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.0966 2021.
Hu, Z., Wang, L., Lan, Y., Xu, W., Lim, E.PP., Bing, L., Xu, X., Poria, S., and Lee, R. K.W. LIm-adapters: An adapter family for parameter-efficiet fine-tuning of large lan- guage models. arXiv preprint arXiv:2304.01933 2023.
Brown, T. B. Language models are few-shot learners. arXiv preprint arXiv:2005.14166 2020.
Chen, T., Xu, B., Zhang, C., and Guestrin, C. Training deep nets with sublinear memory cost. arXiv preprint arXiv: /604.0000.000 06/444 2016.
Kingma, D. Adam: method for stochastic optimization. arXiv preprint arXiv: /412.6:411.69888 2014.
10

---

FRUGAL; Memory-Efficient Optimization by Reducing State Overhead for Scalable Training
Kunstner, F., Chen, J., Lavington, J. W., and Schmidt, M. Noise is not the main factor behind the gap between sgd and adam transformers, but sign descent might be. arXiv preprint arXiv:2304.13966 2023.
Pan, R., Liu, X., Diao, S., Pi, R., Zhang, J., Han, C., and Zhang, T. Lisa: Layerwise importance sampling for memory-efficient large language model fine-tuning, arXiv preprint arXiv:2403.17911 2024.
Le Scao, T., Fan, A., Akiki, C., Pavlick, E., Ilié, S., Hesslow, D., Castagné, R., Luccioni, A. S., Yvon, F., Gallé, M., et al. Bloom: A 176b-parameter open-access multilingual language model. 2023.
Pan, Y. and Li, Y. Toward understanding why adam con- verges faster than sgd for transformers. arXiv preprint arXiv:2306 00204, 2023.
Radford, A.,
Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., al. Language models unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Li, B., Chen, J., and Zhu, J. Memory efficient optimiz- with 4-bit states. Advances in Neural Information Processing Systems, 36, 2024.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S.,
Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21 (1 (140):1-67) 2020.
Lialin, V., Muckatira, S., Shivagunde, N., and Rumshisky, A. Relora: High-rank training through low-rank updates. In The Twelfth International Conference on Learning Representations. 2023.
Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Per- formance Computing, Networking, Storage and Analysis, pp. 1-16. IEEE, 2020.
Liu, A., Feng, B., Xue, B., Wang, B., Wu, B., Lu, C.,, Zhao, C., Deng, C., Zhang, C., Ruan, C., et al. Deepseek- v3 technical report. arXiv preprint arXiv:2412. 2024a.
Liu, S.-Y., Wang, C.-Y.. Yin, H., Molchanov, P., Wang. Y.C. F, Cheng, K.-T., and Chen, M.-H. Dora: weight-decomposed low-rank adaptation. arXiv preprint arXiv:2402.09333, 2024b.
Richtárik, P. and Takáč, M. Iteration complexity of random- ized block-coordinate descent methods for minimizing a composite function. Mathematical Programming, 144(1): 1-38, 2014.
Liu, Y. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: /907. 1692, 2019.
Richtárik, P. and Takáč, M. On optimal probabilities in stochastic coordinate descent methods. Optimization Let- ters, 2015, pp. 1-11, 2015a.
Liu, Y., Gao, Y., and Yin, W. An improved analysis of stochastic gradient descent with momentum. Advances in Neural Information Processing Systems, 33:18261- 18271, 2020.
Richtárik, P. and Takáč, M. Parallel coordinate descent methods for big data optimization. Mathematical Pro- gramming. Series A, pp. 1-52, 2015b.
Loshchilov, I. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05100 2017.
Robert, T., Safaryan, M., Modoranu, I.-V., and Alistarh, D. Ldadam: Adaptive optimization from low-dimensional gradient statisties. arXiv preprint arXiv:2410.16103. 2024.
Luo, Q., Yu, H., and Li, X. Badam: A memory efficient full parameter training method for large language models. arXiv preprint arXiv:2404.0288
2024.
Komera-Paredes. B., Barekatain, M., Novikov, A., Balog, M., Kumar, M. P., Dupont, E., Ruiz, F. J., Ellenberg, J. S., Wang, P., Fawzi, O., et al. Mathematical discoveries from program search with large language models. Nature, 625
Lv, K., Yang, Y., Liu, T., Gao, Q., Guo, Q., and Qiu, X. Full parameter fine-tur tinettunnng for large language models with lim- ited resources. arXiv preprint arXiv:2306.00976 09782, 2023.
(7995):468-455, 2024.
Modoranu, I.-V., Safaryan, M., Malinovsky, G., Kurtie, E., Robert, T., Richtarik, P., and Alistarh, D. Microadam: Ac- curate adaptive optimization with low space overhead and provable convergence. arXiv preprint 2rXiv:2405.1559) 2024.
Sardana, N. and Frankle, J. Beyond chinchilla-optimal: Accounting for inference in language model scaling laws. arXiv preprint arXiv:2401.00445 2023.
Shazeer, N. Glu variants improve transformer. arXiv preprint arXiv:2002.05202,, 2020.
OpenAL. Gpt-4 technical report. arXiv:2303.08877, 2023.
arXiv preprint
Shazeer, N. and Stern, M. Adafactor: Adaptive learning rates with sublinear memory cost. International Con- ference on Machine Learning. pp. 2018.
Ortega, J. M. and Rheinboldt, W. C. Iterative solution of nonlinear equations several variables. STAM, 2000.
PMLR,

---

FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint arXiv:2302.13397. 2023a.
Zhang, Y., Chen, C., Li, Z., Ding, T., Wu, C., Ye, Y., Luo, Z.-Q., and Sun, R. Adam-mini: Use fewer learning rates to gain more. arXiv preprint arXiv:2400 16793, 2024c.
Zhang, Z., Jaiswal, A., Yin, L.,, Liu, S., Zhao, J., Tian, Y., and Wang. Z. Q-galore: Quantized galore with int4 projection and layer-adaptive low-rank gradients. arXiv preprint www.iv:2407.02222. 2024d.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaci, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., al. Llama 2: Open foundation and fine- tuned chat models. arXiv preprint w00292230..09288, 2023b.
Zhao, J., Zhang, 7.,, Chen, B., Wang, Z.,, Anandkumar, A., and Tian, Y. Galore: Memory-efficient Ilm train- ing by gradient low-rank projection. arXiv preprint 0rXiv:2403.03507 2024a.
Tseng, P. Convergence of block coordinate descent method for nondifferentiable minimization. Journal of optimiza- tion theory and applications, 109:475 494, 2001.
Zhao, R., Morwani, D., Brandfonbrener, D., Vyas, N., and Kakade, Deconstructing what makes good optimizer for language models. arXiv preprint 2024b.
Vyas, N., Morwani, D., and Kakade, S. M. Adamem: Memory efficient momentum for adafactor. In 2nd Work- shop on Advancing Neural Network Training: Computa- lional Efficiency, Scalability, and Resource Optimization (WANT@ ICMI. 2024), 2024.
Wang, A. Glue: A multi-task benchmark and analysis plat- form for natural language understanding. arXiv preprint arXiv:1804.07460 074611, 2018.
Wortsman, M., Dettmers, T., Zettlemoyer, L.., Morcos, A., Farhadi, A., and Schmidt, L. Stable and low-precision training for large-scale vision-language models. Ad- vances in Neural Information Processing Systems, 36: 10271-10298, 2023.
Yang, J., Jin, H., Tang, R., Han, X., Feng, Q., Jiang, H., Zhong, S., Yin, B., and Hu, Harnessing the power of 1lms in practice: A survey chatgpt and beyond. ACM Transactions Knowledge Discovery from Data, 18(6):
1-32, 2024.
Zamirai, P., Zhang, J., Aberger, C. R., and De Sa, C. Revis- iting b training. 2020.
Zhang, B. and Sennrich, R. Root mean square layer nor- malization. Advances in Neural Information Processing Systems, 32, 2019.
Zhang, B., Liu, Z., Cherry, C., and Firat, O. When scal- ing meets 1lm finetuning: The effect of data, model and finetuning method. arXiv preprint arXiv:2402. 2024a.
Zhang, J., Karimireddy, S. P., Veit, A., Kim, S., Reddi, S., Kumar, S., and Sra, S. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 08:1533-15393, 2020.
Zhang, P., Zeng G., Wang, T., and Lu, W. Tinyllama: An upen-source small language model. arXiv preprint arXiv:2401.00 023385 2024b.
12

---

FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training
A. Experimental setups
This section describes the main setups used in the experiments and presents additional experiments.
To begin, we introduce the hyperparameter density p. This hyperparameter represents the fraction of the total space in Linear layers that is updated with stateful optimizer. For GaLore, this parameter is equal to p -r/h, where r is the projection rank, and h is the hidden size of the model. For the RandK projection, this parameter can be expressed as means sparsity. For BAdam and FRUGAL with the blockwise update, this parameter denotes the ratio of the number of active blocks Oblock to the total number of blocks that is, Oolock/P /p. When using FRUGAL with the column-wise update, as in Section 7, p is equal to the ratio of the number of active columns Gcolumn to their total number h, i.e, p = asoluma/ / h.
where
A.I. Pre-training setup
We adopt a LLaMA- based architecture with RMSNorm (Zhang & Sennrich, 2019) and SwiGLU (Shazeer, 2020) activations on the C4 dataset. Following Zhao et al. (2024a), we trained using a batch size of 512 sequences, sequence length of 256, weight decay of 0, and no gradient clipping. We used T5 tokenizer, since it also was trained on C4 with dictionary size equal to 32k. The update frequency T' is set to 200.
Since, unlike GaLore, we consider not only matrix projections, decided to generalize the concept of rank Instead, we use density p. which represents the proportion of Linear layer parameters in the state-full subspace. Thus, for SVD-like projection as in Galeore, the density equals point out that similarly to Zhao et al. (2024a), we keep Embeddings, RMSNorms, and Output layer in the state-full subspace throughout the training and don't 't reset the optimizer state for them.
r/h, where h denotes the hidden dimension of the model. We also should
We used standard Adam hyperparameters: 31 = 0.9, B3 = 0.999, = le 8.
8. For all methods except GaLore,, we selected the learning rate equal to the optimal learning rate for Adam, which we determined through a grid search among set equal to that for the
values [lc 3e 4, le- 3, 3e 3). FRUGAL's learning rate for the state-free optimizer state-full optimizer for simplicity and ease of tuning. For a fair comparison with Gal Gore (Zhao et al., 2024a), we conducted experiments with two learning rate values: 1) the one specified by the authors in the original paper and 2) the optimal learning rate for Adam, as used for other methods. We did this because the learning rate in the original paper could have been optimized for a different number of iterations.
To match the learning rate changes in the first steps of our training with Zhao et al. (2024a), we used a cosine learning rate schedule with restarts, with a warmup of 10% of the steps in a cycle length, and decay of the final learning rate down repeated the
to 10% of the peak learning rate. To verify that our results are not sensitive to the choice of scheduler, experiments for LLaMA-130M with other schedulers. The results for constant with warm-up and cosine (one cycle) with warm-up schedulers can be found in Tables 15 and 16.
For pre-training GPT-2 124M (Radford et al., 2019) we followed the setup described above except for the tokenizer. We utilized the GPT-2 original tokenizer, with 50257 vocabulary size. The results are presented Table 12.
Table 8. Comparison of validation perplexity and memory estimation for various optimization methods across LLaMA model scales trained on C4 with B₂ = 0.95. We also indicate the additional memory overhead introduced by the optimization algorithm. The values are calculated assuming that cach float value occupies 4 bytes (float32). P denotes the proportion of the Linear layer parameters in the state-full subspace. Note that Embeddings, RMSNorms, and Output layer are always trained with AdamW.
60M
130M
350M
AdamW
23.51 (0.43G)
18.32 (1.00G)
14.57 (2.74G)
GaLore, p == 0.25
26.66 (0.30G)
21.03 (0.54G)
16.79 (1.10G)
BAdam, p 0.25
25.40 (0.29G)
20.17 (0.52G)
16.54 (1.05G)
FRUGAL, p = 0.25
24.07 (0.29G)
18.79 (0.52G)
14.96 (1.0SG)
FRUGAL, P p= 0.0
24.54 (0.24G)
19.11 (0.37G)
15.20 (0.49G)
Training tokens
20B
20B
24B
Number of iterations
200k
200k
240k
13

---

FRUGAL: Memory-Effiiient Optimization by Reducing State Overhead for Scalable Training
Table Perplexity of I.I 1aMA-1300M models pre-trained C4 using pure bíloat16 format both for model weights and optimizer statistics.
Table 10. Perplexity of I.L.MA-33MM models pre-trained on C44 for 200k steps with different state-free optimizers for FRUGAL..
Method
100k iterations
State-free optimizer
Validation perplexity
Method
Adam
21.88
GaLore, P 0.25
24.19
Adam
18.13
BAdam, 0.25
25.03
FRUGAL, / = 0.25
signSGD
18.60
FRUGAL, P = 0.25
23.17
FRUGAL, P = 0.25
SGD
19.11
FRUGAL, = 0.0
22.64
Table 12. Validation perplexity of GPT-2 124M model pre-trained C4 for 200k steps with various optimization methods.
Table 11. Perplexity LLaMA-13000 130M models pre-trained on C4 with Lion state-full optimizer for 200k steps.
Method
Validation perplexity
Method
200k
Adam
21.94
Adam
13
GaLore, = 0.25
25.84
Lion
18.55
BAdam, P 0.25
25.43
GaLore (+ Lion), p 0.25
21.65
FRUGAT, P 0.25
23.23
FRUCAL (+ Lion), p 0.25
18.89
FRUGAL, P 0.0
25.04
Table 14. Perplexity of LLaMA-130M models pre-trained on C4 for 200k iterations (20B tokens) using FRUGAL with p = 0.25 and different update frequency T.
Table 13. Perplexity of LLaMA-130M models pre-trained on C4 for 200k iterations using FRUGAL with p = 1/3 and different Block update strategy, taken from Luo et al. (2024).
Update frequency T
Perplexity
10
18.82
Method
Perplexity
20
18.73
Random 18.50
50
18.69
Ascending 18.54
100
18.65
Descending
18.50
200
18.60
500
18.60
1000
18.61
Table 15. Perplexity LLLMA- 130M models pre-trained on using constant scheduler with warmupp up at various training itera- tions.
Table 16. Perplexity of LLaMA- 130M models pre- trained on C4 using cosine scheduler with warm-up at various training iterations.
Method
100k
200k
Method
100k
200k
Adam
19.38
17.95
Adam
19.51
18.51
GaLore, == 0.25
22.30
20.60
GaLore, p= 0.25
22.63
21.03
BAdam, == 0.25
22.35
20.07
BAdam, P 0.25
22.31
20.66
FRUGAL, P 0.25
19.62
18.16
FRUGAL, P = 0.25
19.97
18.85
FRUGAL, P 0.0
19.83
18.34
FRUGAL, P = 0.0
20.33
19.14
A.2. RoBERTa fine-tuning setup
The batch size and learning rate values used for FRUGAL in the experiments from Table 6 are presented in Table 18. In all experiments, we set the learning rate for the state-free optimizer to 1/10 of the learning rate of the state-full optimizer. Other hyperparameters, such as scheduler, number of epochs, maximum sequence length, and warmup ratio, were taken from Hu al. (2021).
14

---

FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training
Table 17. Perplexity of 1.MMA-300 models pre-trained C4 for 200k iterations (20B tokens) using PRUGAL with different density p.
FRUGAL
1.0 (Adam)
0.5
0.33
0.25
0.125
0.0625
0.0 18.90
signSgd
Perplexity
18.13
18.40
18.50
18.63
18.71
18.80
33.22
Table 18. Hyperparameters of fine-tuning RoBERTa- base for FRUGAL.
MNLI
SST-2
MRPC
CoLA
QNLI
QQP
RTE
STS-B
Batch Size
128
128
16
256
256
128
32
16
State-full Learning Rate
5E-05
5E-05
2E-04
5E-04
1E-04
5E-05
2E-04
1E-04
State-free Ir multiplier
0.1
Rank/Density
r= 8 / r
r= (p=0)
We also present a comparison between fine-tuning using FRUGAL with = 0.0 and full fine-tuning using signSGD. Essentially, the only difference is that the second case, the classification head is updated with signSGD instead of Adam. The results in Table 19 show that the classification head extremely sensitive to the optimizer type, and switching the optimizer significantly drops the accuracy,
Table 19. Results of fine-tuning RoBERTa-Bare on several tasks from GLUE, The left column indicates which modules were trained using the state-full optimizer Adam. The remaining modules, except for the frozen Embedding layer, were trained using the state-free signSGD.
Method
SST2 QNLI
QQP
Classification head (corresponds to the FRUGAL, with p= 0.0) 94.9 92.8
91.34.1
None (corresponds to the fine-tuning using signSGD)
89.7
81.6
74.3
B. Comparison with Concurrent Methods
In this section, we present comparison with several concurrent works 1155-nnnn namely, Fira (Chen et al., 2024a), IDAdam (Robert et al., 2024), and Adamem (Vyas et al., 2024) which we discovered only during the final stages of our project. Due to this timing, we did not have the opportunity to conduct comprehensive experimental comparisons with them as we did with other basclines in Table 2.
B.I. Algorithmic comparison
We begin by comparing these methods with FRUGAL from an algorithmic perspective.
Adamem. Similarly to framework, Vyas et al. (2024) divides the gradient into two parts: the first part is a projection of the gradient onto the top SVD subspace, while the second part is the residual outside this subspace. The first part is used to update momentum, which is then fed into Adafactor's preconditioner Shazeer & Stern (2018), while the second part is fed directly into a one-sided Adafactor preconditioner Vyas et al. (2024). The outputs of these preconditioners are then used for the final parameter update. Although the core idea of splitting the gradient into two orthogonal subspaces analogous to Algorithm 1, Adamem implements only one possible variant for choosing subspace and for updates of each component. Thus, Adamem represents special case of FRUGAL: with SVD-based projection, Adafactor with momentum as the state-full optimizer, and one-sided Adafactor the state-free optimizer.
Fira. Chen et al. (2024a) also apply the idea decomposing the full-rank gradient into a low-rank subspace gradient and residual gradient. Similarly to GaLore, the low-rank part is obtained through SVD and passed through an AdamW step; however, unlike GaLore, the residual part is not discarded, but is also used in the update in an SGD-like format without using additional optimizer state. An important feature that Fira introduces is norm- bosm- scaling: this technique adaptively scales the per-column learning rate for the residual part by
(R.)] where th: denotes the AdamW update rule. In addition, for [R,
15

---

FRUGAL: Memory-Effiiien Optimization by Reducing State Overhead for Scalable Training
Table 20. Comparison of validation perplexity for AdamW, FRUGA and AdaMeM across ILaMM model scales trained C4.
60M
130M
350M
AdamW
22.73
18.13
14.43
AdaMeM, 0.25
23.81
18.99
15.10
FRUGAL, p 0.25
23.59
18.60
14.79
FRUGAL, 0.0
24.06
18.90
15.03
Training tokens
20B
20B
24B
Number of iterations
200k
200k
240k
enhanced training stability, Fira proposes replacing standard gradient clipping with norm-growth limiter, which transforms abrupt gradient spikes into gradual, smooth increases.
Thus, the key difference between Fira and FRUGAL is that the state-free update sophisticatedly utilizes information from the state-full update. However, Fira has several limitations. First, Fira relies on computing SVD for projection onto the low-rank subspace, which introduces additional memory and computational overhead, with the computational burden becoming more pronounced as model size increases (see Section 4 and Appendix C. Morcover, it is not readily obvious how to propose an alternative to SVD to eliminate this overhead, because, for example, block projection from BAdam (Luo et al., 2024) cannot work with norm- bean-baseddde scaling (Chen et al., 2024a) by construction. Second, following GaLore, Fira continues using the old optimizer state after updating the projection, which should be suboptimal (see Appendix D). Furthermore, according to the original paper Chen et al. (2024a), the reason for instability standard gradient clipping may be precisely the moments of low-rank subspace transitions.
We also would like note separately that neither Fira (Chen et al., 2024a) nor Adamem (Vyas al., 2024) provide theoretical convergence proofs. In contrast, we provide proof that recovers the best known convergence rate under conventional assumptions.
LDAdam. Notably, unlike Fira and Adamem, Robert et al. (2024) identify the addition of momentum and gradient from different subspaces as mathematically incorrect (similarly to Appendix D). To resolve this inconsistency, Robert et al. (2024), propose reproject the previous optimizer state (similarly to Algorithm 1). Interestingly, LDAdam differs from FRUGAL, Adamem, and Fira in that each individual step remains low-rank, but unlike GaLore and BAdam, the remaining information is not discarded but preserved in an error feedback buffer, thus allowing LDAdam to have convergence guarantees and demonstrate high performance.
However, LDAdam also has a drawback. The method requires updating the projection at every step; and although Robert al. (2024) proposed replacing the expensive SVD decomposition with significantly cheaper block power iteration, the necessity of performing this operation every step causes a slowdown of 2 15%.
B.2. Experimental comparison
We also present preliminary experimental comparison of FRUGAL with Adamem (Vyas et al., 2024), Fira (Chen et al., 2024a), and LDAdam (Robert et al., 2024).
Adamem. We compared FRUGAL with our reimplementation of Adamem (Vyas et al., 2024) on pre-training LLaMA model of 3 sizes - 60M, 130M, and 350M - in the same setup as in Section 6.1 and Table 2. Unique Adamem hyperparameters were taken from Vyas et al. (2024). The results presented in Table 20 show that while Adamem performs significantly better than GaLore due to utilizing information from the gradient residual, it still falls slightly short of FRUGAL. However, at this point, we are not ready to claim whether this is consequence of less favorable choice of state-full and state-free update rules or potentially suboptimal hyperparameter settings for experimental setup.
Fira and LDAdam For experiments with Fira (Chen et al., 2024a) and LDAdam (Robert et al., 2024), we deviated from the setup in Appendix and Table 2. This modification was necessary since our main setup following GaLore (Zhao et al., 2024a) does not use gradient clipping, while its counterpart the norm-growth limiter is critically important for Fira's performance. Additionally, we observed that Fira significantly benefits from using weight decay, so for the experiments
16

---

FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training
Table 21. Comparison of validation perplexity for AdamW, PRUGAI, Fira and I. I.Addam across L. LLMM model scales trained C4. We
also report approximate slowdown comparing to AdamW (taken from original papers).
Approximate Slowdown
130M
350M
A
17.52
13.81
Fira, P 0.25
10%
17.58
13.74
LDAdam, p 0.25
15%
17.54
13.54
FRUGAL, p = 0.25
0%
17.58
13.99
in Table 21 we enabled both gradient clipping (norm- woorm-growth limiter for Fira) and weight decay. The results show that all three methods achieve performance very close to AdamW. Surprisingly, LDAdam even outperforms the full-rank A baseline for the 350M model (which, however, may be a result of insufficient hyperparameter grid search). Fira performs slightly better than FRUGAL.
However, we would like to note that these improvements are valid in terms of the number of iterations, and not in terms of wall clock time, where both Fira and LDAdam introduce noticeable time overhead. In the same table, we report the approximate wall-clock slowdown of the methods compared plain AdamW. We provide approximate values taken from the original papers because the actual slowdown can vary significantly depending on the setup (number of workers in training, GPU type, etc.). One can observe that the slowdown due to the need to perform additional heavyweight operations during the optimizer step can negate the advantages gained from a more carefully designed update rule. Thus, practitioners should choose the algorithm for their specific use case.
C. Memory estimation
In this section, we will examine memory requirements for different projection types using the LLaMA-like architecture as an example and show that RandK, column-wise, and blockwise projections result in approximately the same amount of additional memory for a given density value p Appendix A. In contrast, the semi-orthogonal projection matrix (GaLore-like requires slightly larger value in this setup. Recall that we follow the setup from Zhao et al. (2024a), where Embeddings, RMSNorms, and Output layer remain in the state-full subspace throughout the training, so the projection does not interact with them, and they give the same memory overhead for all projection methods.
Let the number of parameters in the remaining projectable parameters be P. Then, training using Adam gives additional overhead of 2P float values for storing m and for each parameter. Now, let's consider blockwise and column- gaise projections and suppose we want to achieve density p. For blockwise, we take round L.) layers, where L is the total number of transformer layers, and for column-wist we take round(p k) columns for each matrix of size n X k. Since the memory required to store block or column indices is negligible compared to other costs, we find that the total size of the optimizer state when using Adam as a state-full optimizer will be 2p- P. with an adjustment for rounding.
In the case of RandK projection, we have the same 2p- P float values M and V in the optimizer state. However, we must also know the current indices corresponding to these values. On the other hand, it is widely known that if one needs to save set of random values, they don't 't need to store all these values- it's sufficient to store only the seed from which they were generated. Thus, for RandK, the total memory also equals 2p- P.
If we recalculate this considering a specific LLaeA-like architecture, each layer consists of 7 matrices: 4 matrices of size x h (Query, Key, Value, Output) and matrices of size h x h, (Gate, Down, Up), where h is the hidden size of the model, and h the FFN hidden size. In the LLaMA architecture, it's typically:
2
has = 4h-
3 h
3
Then, the amount of memory for RandK projection (and consequently for all others mentioned above) is:
8
2 (4- (ph?) + 3- (p h ((-)))))) = 2- (4. ph2 + 3- (7.. 122)) 24p h?
for each layer on average (2 corresponds to the number of matrices M and V).
17

---

FRUGAL: Memory-Effiiien Optimization by Reducing State Overhead for Scalable Training
In the case of GaLore-like semi-orthogonal projection matrix, the situation is follows. We have projections onto low-rank subspace of rank r, where round(p h). Then, for Query, Key, Value, and Output projections, need to store P. M,, V E g and for Gate, Down and Up projections either P E g M, V E R or P E R M, V E Rixe, Since the second option requires less memory, it is used by default in (Zhao et al., 2024a) and, therefore, in FRUGAL, too. Then, the total memory requirements are:
8 3irh =6phh.
4. (3 rh) + (2.r. hr. hyy) = 12rh + 6rh + 3rh (12 6
To sum up. RandK, column-wise and blockwise projection requires 2pP additional memory, while semi-orthogonal projection (GaLore-like) requires 26
24 2pP = 13 12 2pP additional memory.
Let's recall that in addition this, SVD requires additional computation, which can take up 10% as the model size increases (Zhao al., 2024a). Therefore, for our method, we settled on blockwise projection.
D. Optimizer state management
In this section, we would like to propose some modifications to the GaLore algorithm. These modifications are also used in our framework as SVD projection.
Specifically, we want to consider the projection of the state when changing the active subspace. In GaLore (Zhao et al., 2024a), when updating the projection, the optimizer states M and V do not change. This results in new projected gradients and old M and V being in different subspaces. This implementation has little effect the result with large values of update frequency T, the values of M and V from the previous subspace decay exponentially quickly. However, more frequent changes T significantly affect the result. We hypothesize that this is why in Zhao et al. (2024a) the model quality degraded so significantly when T was decreased, while as seen in Table 14, FRUGAL experiences much less degradation.
There are two different ways to overcome this obstacle: either project the state back to full-rank space or reset the state before a new round. However, the first option may be challenging in the case of arbitrary projection. Specifically, while it's possible to project momentum back to full-rank space (see Alg. 2 in Hao et al. (2024)), the same cannot be easily done with variance because its values depend quadratically on the projection matrix. However, the projection of variance will also be trivial if the set of basis vectors for the projection is fixed, which is true, for example, for coordinate projection with RandK.
To demonstrate the effectiveness of this improvement, we provide a toy example. We consider quadratic minimization problem [ W IR 10x100 For optimization, we use GaLore-like SGDM and GaLore-like SGDM with Momentum state projection. This projection is similar to Alg. from (Hao et al., 2024), except we additionally normalize the new momentum by the ratio of norms before and after re-projection to preserve momentum mass. We use ranks of and 6, and an update frequency T 10 and plot mean and standard deviation across 5 independent runs. The results are presented in Figure 3. As can be seen, the variant with state projection converges much faster.
18

---

FRUGAL: Memory-Efficiien Optimization by Reducing State Overhead for Scalable Training
GaLore-like SGDM, rank
10a
GaLore-like SGDM with Momentum re-projection, rank 3
GaLore-like SGDM, rank
CaLore-like SCDM with Momentum re-projection, rank
Loss
10
Iteration
Figure 3. Toy example of solving quadratic minimization problem with GaLore-like SGDM with and without re-projection of optimizer state, Algorithm with re-projection converges much faster.
19

---

FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training
E. Convergence Theory
Firstly, we provide ommited definition of L-smooth function.
Definition E.1. We say that Rd R is smooth with L 2 0, it differentiable and satisfies
I. lly =112 y
f(0) S f() + (V](x). x) +
Below, we provide an equivalent formulation of Algorithm 2 that enables us to use the proof of the similar structure to SGDM momentum analyis of (Liu et al., 2020).
Algorithm 3 FRUGAL (SGDM, SGD) Equivalent Algorithm 2 for constant step size
Input: momentum weight B E [0, 1), initialization x² R2d and m O, step sizes {ar ;- >> 0J k= momentum set
JR C [d] for k: = 1,2... for k == 1.2...... do
Compute stochastic gradient g fx*h
m if E JR. otherwise
Update momentum vector r (1 B) + 3
4: Update iterate k
+- x am
//2/2
Bxt
5:
k+1
1-3 k 3
I-3
if $ Jkt1.
otherwise
6: end for
Next, we present several key ingredients the proof. Firstly, we can express the momentum term mf as
k
m = (1 - B)
Σ gh-igh gj.
(3)
t=
where t5 = ma & Jol, i.e., the last time when the momentum buffer was released. We denote
k
m = (1 - B)
g
(4)
1=1
Using this notation, we proceed with two lemmas, one showing variance reduction effect momentum, the other boundess of momentum bias. Lemma E.2. Under Assumption 5.1, the update vector vitk in Algorithm satisfies
E
m
13]
1 1 0²
Proof.
Since n. = =(1 - (1-8)) Ding >) Bh- we have
E m m m²11
Σ E
111
jEld]
< (1 - 3)2
E jEld]
3
(0)
Moreover, since s!.52.... ch are independent random variables (item 3 of Assumption 5.1), we can use conditional expectation to show that [(0) 9²" - = for é # i2. Therefore,
20

---

FRUGAL: Memory-Efficieni Optimization by Reducing State Overhead for Scalable Training
k
E |l -
1a04-mmm1' m²]]²]
5 3)
E
o/1l² g2
SEId
1=
1 3
Σ
E
(1
12(0- 03
1 + 3
sE(d
1- 3
1 1 3 2 1 + 3
Σ
oj
1+ 3
DE[4
Lemma E.3. Under Assumption 5.1, the update vector rh in Algorithm 3 further satisfies
m
E
Σ
(1
9
S
pm
E
Σ a ||
- -all==-1~=-- -2
(1 Bkc,
JE
where KÚ = k - 14 + 1, and
Æx 1hhh-i kk i 3
B
(5)
Proof.
Let Pra-1lj E JRl = r and piaxx = max (15}. (1 Then.
27
my
1 3
k
E
Σ (1- - رk
9
=E
Σ(1 (1- ggg
3 3h - 935)
(1 Bk,)
1 1 Bkg
JE
JEJN
i=
Σ h
= (1 - =E
95), 3 100
(t) - r;,))
JEJU 8,1=1)
82k- 1-1 (21l0f
-55121 + 1 111 03112
$ B)
J. 8,1=1
(1 - 3)
82k-i-1
1 E1119
2
JEJN 1=1
1
+ (1 -3
92k-
プタ E11oo - oř11
JEJU 1=1
gh- i(1 3kg)
(1 B
10 o8112
1- 3
JEJU 1=1
k
$ (1 T)T
gh 110 13
JEJU 1=1
$ (1 - fu
E
gh -1 11g 1²1
1=1
where we applied Cauchy-Schwarr to the first inequality.
21

---

FRUGAL: Memory-Effiiien Optimization by Reducing State Overhead for Scalable Training
By applying triangle inequality and the smoothness of f (item I in Assumption 5.1), we further have
E
Σ
)²
m
9j
(1
3)pmax IE
Σ Bk (k
> 1g g'1?
(1
342
JEJx
A
E
3)p 12 Σ
-0 gh- - 1) || x/12
Therefore, by defining d2., = (1--3)13 |/ Ei= Bh-i(k:
i), we get
m
E
Σ (1-
(1 Bk: 9/
p E
Σ 0kklll** - -2'1
(6)
JEJN
Furthermore, dh, be calculated
+ L2 )
3
de = L3gh
1-
4-1+
(7)
1-
1 -
Notice that
B
d aad = 1²
1-1+
(8)
1 B
Combining this with equation 6, arrive at
m
E
Σ JEJU
(1
(1 Bhg
gj
k
max
E
Σ 0 1
i-1
x²
where
B k 1 it
αk gh2
-3
From Lemma E.3, we know that the distance of the non-stochastic momentum from g² bounded by the weighted sum of past successive iterate differences. Furthermore, the coefficients aR decays exponentially in B.
Therefore, the following Lyapunov function
L (²=(f((A)-
Σ c cilla*+1 1
-
(9)
for some positive o that we specify later. As is common for convergence theory of SGDM to analyze an auxiliary sequence * defined as
h: = 1,
2)
(10)
1/2
k- h: 2 2,
1-3
which behaves more like SGD iterate, although the stochastic gradient 3 is not taken at * Lemma E.4. Let k be iterates of Algorithm then * defined in equation 10 satisfies
k+1 2 -og
22

---

FRUGAL: Memory-Effiiient Optimization by Reducing State Overhead for Scalable Training
Proof. We have to consider two different cases. Firstly, if k or j & Jk. then
1/1/2
3xy J am, - 3x - -(1 -(1- B)ay
a(1 - B]0} -ag;
2) 2/41 2j
1 3
1-
1 B
Secondly, if 2 Js. then
x++1
2/
(I)
1-1/2 - 23
//1//
1- 3 (af 2
1 B
1 1-1/22 - -zg))
(1)
1 1-8 (nf az 2)
B
(-amm) 1 1- (-arm) 1)
B
(-amy + +asm 1) -agj.
Before procceding with the main convergence theory, we require one more proposition that shows descent in objective value. Proposition E.5. Take Assumption 5.1. Then, for z defined equation 10, we have
[[ S <BJ((*))} (-a+
1 1
1
La +
B
B[1la*119]
1-
2
(11)
2(1 2) + /) Lazo² +
12
Ta² E
m2
Σ (1 -
Bk,) (1)) - 3 -
1 - B
(1
JEJ&
Proof. The smoothness of f yields
Ex [f( k S f(2h) { + Be+l(( 241 -1 +
2 Be 111 111
(12)
La²
+ (e -agh)] Ec 11g
where we have applied Lemma E,4 in the second step.
For the inner product term, we can take full expectation E = EE....E to get
- - -ag
which follows from the fact that * is determined by the previous - random samples Ç) <2 (k.....k-,, which is independent of k and Ex [g*] = gf.
So, we can bound
-agh)) =
aElll |||
a 2 2
- -"E1
1 2po
Elllo -
where po > 0 can be any positive constant (to be determined later).
Combining equation 12 and the last incquality, we arrive at
El * []] E[f( ) + a
Po 2 -011
2
+ (a 1 - -)E
x 1
La²
2po
23

---

FRUGAL: Memory- Eefiiry-Effiiiett Optimization by Reducing State Overhead for Scalable Training By construction, o -- for j Jg. 0 otherwise. Consequently,
=
3 L2 2
3
Y'
Σ
1mm
-
112
JEJ
(13)
Laa2
+
(a 2pu
-
a
E
111g
+
B118*11
2
Let K = k - -17-11 + 1. Then, from Lemma E.2 we know that
E
m k
1
S
2E
Σ
1
my
112
+ 21E
|mma k-1 ||
JEJ
JE
JEJ
3
Σ
oJ
+
Σ
||mf-111 112
+ 3
JEJI
JEJL
mi h
(14)
E
Σ
m
112
E
Σ
B(k-1))32
(1 11()
jEJx
JEJ
m k-
27
S 21
Σ (1- (1- 1202
9
+ 21
Σ
||
(1- ((x-))))
JE
jEJu
[ S o² + [|l1s*111
Putting these into equation 13, we arrive at
E[f(z*t)) E[f(=*]]+
(aa+
a
2po
1
+
2a POL2
(_ 3
(1-3)
})
)² La²
+ 2
E Bllgg Bllla*131
a POL?
9
1
8
La?
02
1- 3
1+
2
B
m k-1
2a 20
E
Σ
(1
g(k-1)g
9,
1 3
3)b
JEJx
Notice that if j E Jk, then (k- 1), = kg- - 1. Therefore,
m k
27
3m k- 1 + +111B))g
9
E
9
gk:)
(1 (1-34hs)
(11--t1)))
k --1
3²B
m
9
3Kss
(1
g(h 1)g
Substituting the above into the last inequality produces
1
B
)2 + La²
B[f(+++]]] [] +
+
20
2
2po
3
2
+ 43 1i2
2
3
La²
0²
3
+
B
2
(15)
1
m x
+ 3 10
E
Σ
(1
3h )2
g/
1-
(1 - gkg
JEJ.
24

---

FRUGAL: Memory-Effiient Optimization by Reducing State Overhead for Scalable Training
1-2 Finally, Po
2La gives
E E[f(z*+|))
{{f(( (-a+
1
La + 1
La
Blll=*111
1-3
2
8²
+ 1
La²
m
La +
1 E
Σ
(1 - B
(1- Bh,) 3
2(1 8)
JEJx
E.1. Convergence of Algorithm
Firstly, by combining results from prior section, we can bound our Lyapunov function L defined in equation 9.
Proposition E.6. Let Assumption 5.1 hold and a 2/2 P v8+82 in Algorithm 3. Let {c) in equation 9 be defined
1-8
by
1 4a2
(1-B33
R+
8+.82 (1-M)7
L²
L
C+11 --
[Aeqa² + 1
La²
3 ) Bi+
1 1-3"
for all i 2 1.
Then, Q > 0ffor all i 1, and
-a + 3-8+ 32
ELL 1b1
La² + 4c10²) B[]lo*113]
2(1 - 3)
(16)
B
Lazo" + LLa'0 + 20qa²0²
+
2(1+ 3)
Proof. Recall that & is defined as
k-1
=== f(2h) - fr +
Σ
ci cilla - -
Therefore, by equation 15 we know that
E h+1 1
(-a+ + 1 1-3
1 32
La
- ([1l
1
2
(G+1 c;) c:)BBllbAtt1- - x 112
B[ll 11
(17)
1=1
2(1 + B) +22 211 La Laz02
82
Ta
(1-38t)) - g/:
m (1 gk.
E
9
1 1-3
jeJu
To bound the ej B[||a++1 _*1 2 term, we need the following inequalities, which are obtained similarly as equation 14.
1 - B 0² +
Bllim 111
21
m |1
1 3
Blllm*1|3 < 21E
Σ (1 8h5
m
g k
+ 21E
119* |1² [13]
(18)
(1
8kg
JEJ
||||~\19 0² + 1111g~~113.
25

---

FRUGAL: Memory-Effiiient Optimization by Reducing State Overhead for Scalable Training
Let Pra-ili J&] DJ and pmiin minzetd (. Then, e Blll|a*t1
can be bounded
c BBlll ==11 2 ||
Σ
11 12
JEJL
JeJu
||| + (1 Pmni |1
( 3
s qaa²
+ -1- - *
Pmin
o² + 510111 9 |12]
1 + 3
m
+ 4
Σ
(1 - 3b
(1 - 3k3)
JEdz
Combine this with equation 17, obtain
E E[LA+I= L*]
(-a+ + 1-3 La² + 1 La² + 5c1a 2 [|]
32
82
1 C1
( 1- B
plusa) )) LOPO
+
2
2(1 + 3)
2
L
1 B
k-1
(c+1 c))E[l1=** htl 11
(19)
i=1
27
+ (401a²+
La² +
E
s4=
m
9j
1 1- B/
(1 3h:
jEJ
In the rest of the proof, let show that the sum of the last two terms in equation 19 is non-positive.
First of all, by Lemma E.3 we know that
m
E
Σ (1 JEJE
-
9
Pmax
k
x² 1/
(1 - Bhg)
1=1
where
3
0k. 1.2gh
k: i
-
Or equivalently,
E
Σ (1
Bh:T²
m²
9/ |
E
ak -illllllt+1 20-111
(1 - gks)
jEJu
where
B
-2
i+
1 1 3
Therefore, to make the sum of the last terms of equation 19 be non-positive, we need have
C++1 S [Acia² + 1 1-
La²
12
(1+ 1-
(
B
for all 2 1. To satisfy this inequality, we choose
C+1 = c: (Accaa +
Laz
2
128'pt pinxx
( B 3
1-3
26

---

FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training
for all 2 which implies that
Lad
1-1
3 1 B
G C1
Aega² +
1.2
Pmax
i
1
1=1
To have > 0 for all 2 1, we can set C₁ as
Laz
2
3 1
e1 =
[ccaa²
1
Pmax
8 i
1 -
B
where, Pmaxx max (P Since
i
1
3(1 31)
1 3 533
8
have Ein i5
8 (1-3y7 and
La²
B
[Acia²
L 1 max (1 - 9)2
61 =
1
which implics that
22 Ph (1 (1-9)3
84.32
C1
1 4a 9+0
(1-3y Planx 12
(20)
Notice that <
1-8 2/21 v Pinnss v /3
ensures CJ > 0.
Therefore,
E -6 -a + 3 3- 2,²
La² + +5eja²) Elllo*11" [12]
2(1 3)
g²
2 11-B
pluin)) ))
La" 0² + Laa'd 2
2(1 + 3)
1
By telescoping equation 16, we obtain the convergence bound our proposed algorithm under nonconvex settings,
Theorem E.7. Let Assumption 5.1 hold and & a S
1-3 Then, the iterates of Algorithm satisfy
L(4-
(((o)) ( + Lao? (1+ Piaax (1 --Ftinh )8 )) k =
)))
(21)
ka
(1 - B)
where phi + Σ Di Pinn and
pla
[aa ( (pimax)
Proof. From equation 16 we know that
BlL 14 -R, ntS
(22)
27

---

FRUGAL: Memory-Effiiiee Optimization by Reducing State Overhead for Scalable Training
where
R1 + 3 8-84 8
La 4cqa²
2(1 3)
8
1
LLa tqaa (2 2 21-3 +
R2 2(1 3)
La
+ 2
11
Pmiia
We further define
8 P₂ 2(1 + ()
FLa²d² cia20" (
2 2-1 + Piiin
La + 2
1+
where p = ト
P
Telescoping equation 22 yields
R1 Σ
B
Bilgg 11
Σ R
and therefore
I,a
R₂
1
+
(23)
k
kR1
Ri
In the rest of the proof, we will appropriately bound R and R2.
First, let us show that R 2 2 and < min
1-9 L(1-B+3*)
1-9
2VILV Plana v B+00
From equation 20 we know that
a 2 L P max (1-9)3
844
C1
1 4a2 8+ L Pmaax
(1-3y²
1-3 Since S 2//11 Phans v 18433
we have
40. B+ 32
1 k
Pmax S
(1- By²
2
Thus,
Ci S a L3 Pmax (1- B)3
B2
L
S
8(1 - B)
Therefore, in order to ensure R₁ 2 & it suffices to have
2
3 B2
aL
La +
2(1 - 3) < 2
2(1 33)
which is equivalent to our condition a S
1-8 E14-8-3'
For R2, we can upperbound ej using our condition a S L(4-84844 Thus,
++
I 1 BT.
C1 S 2 2 Pmax
(1 313 B)3 2(1 B)
28

---

FRUGAL: Memory-Effiiien Optimization by Reducing State Overhead for Scalable Training
Therefore,
g
1
( 1-8 3
R2
2(1 La
+ z + taaa0²
2 1
+
P
2
12
2 Pmax o2a
1
La
+
+La² o (1
Pmiia
2(1 + 3)
2
(1+3)
1
232 + 8pmax
p max (1 -mmmn) )
La²
+
2(1 + ()
2
8(1 8(1- - 3)
By putting them all together, we obtain
B|llg'|1 2(f(x)) + (2 + 8p
p (1 Fmin)
Lao
ka
2(1 + B)
8(1 - 9)
(f ka
[((x)) - + Lao? (1+
1max (1 --Ftt Fminn3
(1 3)
29

---

FRUGAL: Memory-Effiiient Optimization by Reducing State Overhead for Scalable Training
F. Limitations
We would also like to acknowledge the limitations of this work. Due to computational constraints, we were unable to conduct experiments on pre-training 7B+ LLMs, which crucial for understanding the potential of our approach when scaling, Furthermore, our experiments are limited to training language models, although memory-cfficiett optimization could also be beneficial for training diffusion models. Finally, there may be better method for selecting the next state-ful full subspace during the training. We leave the exploration of more sophisticated selection strategies for future work.
G. Simplified algorithms pseudocode
In this section we present the simplified pseudocode of FRUGAL. In Algorithm one can find optimizer steps both for FRUGAL with SVD projection (GaLore-like (Zhao al., 2024a)) and Block projection (BAdam-like (Luo et al., 2024)).
30

---

FRUGAL: Memory-Efficien Optimization by Reducing State Overhead for Scalable Training
Algorithm FRUGAL step pseudocode, PyTorch-like
1:
def svd.or ord.or.rae or.rkkk step (self):
2: for param in self, params:
3: grad param.grad.
4: paranistate self. [param]
5:
update projector necessary
6:
if self.step $ self.update..ag
0:
7:
param state(""=ojectoo"] selÉ. proj (grad)
8:
projector = (arramssaate("proeeeee
9:
# obtain state-full grad and state-free grad
10:
grad full = projector.proj down (grad)
11:
grad.free gradfull - srojectorrproo.up (grad.full)
reset state-full optimizer state if necessary if self. self.step $ self.update.gap
0:
14:
param.state ("eppee[ee___e [vgx torch. zeros like (ull)
15:
param.state ["""p_avg.sq" = torch. zeros (grad.full)
16:
4 state-full subspace update
17:
self. +=
18:
update.full self.staate.full.stee (grad.full, paran.state)
19:
update.full projector.prot (update.ful1)
state-free subspace update
21:
update_free self.state_free (grad_free)
22:
# perform resulting update
23:
update update full update free add
24:
25:
def block_step (self):
change state-full and state-free blocks if necessary if self.step self.update.gap 0:
28:
29:
indices.full self indices (indices.full) for idx, param in enumerate (aal).)
30:
31:
grad = param.grad
32:
param_state = self. state [param]
33:
if idx in indices.full:
reset state-full optimizer state
35:
param_state ["exp_avg"] torch. zeros_like (grad)
36:
param state ("exp avg sq"] torch. zeros like (grad)
37:
param_state ("fllttteeffllllsssssssss = True
38:
else:
39:
# free state-full optimizer state to save memory param_statte..leazr(()
40:
param_state ("full False
perform updates
43:
for param in self.param
44:
grad param.grad
45:
param state = (param)
46:
# choose the optimizer depending on the block type if param.state (ft.staateefulllssuuppaee"""""
47:
update stat full. (grad, param_state) else:
50:
update self. Salf.state_ffi stat__ (grad)
51:
# perform resulting update param.add (update)
52:
31

---

FRUGAL: Memory-Efficiiee Optimization by Reducing State Overhead for Scalable Training
Algorithm 5 Examples of state-full and state-free steps for Algorithm 4
I: def state.ful (self, grad, paramistate):
2: exp_avg = pa raw.state
3: exp.avg.sq p arraamstatte"eex.avg...a "=xp_avg_sq"]
4: step self.ster
5: betal, beta2 self.betas
6:
1appa=1..--btall
7: http:/gsg__mml.ttttteaaaooooood.
grad, value=1. 0 0-beta2)
8: denom = exp avg (q.
9: step.size = self.1r.full
10:
if self.correet.bias:
bias.correctionl 1.0 betal ** step
correctire 1.0 beta2 step
step.size ffulll borreccorettionl
14:
biss.correciinn-ssrtt path.sqri (bias.correction22
15:
danom.div.bbia. (bias cowrection sgrt)
16:
add (sel..eps)
17:
apdate_full = exp.avg denom * sstep_size
18:
return update.full
20:
def state_free.signnee sereerfreess.nnssdssseeeeee (self, grad)
21:
update free -self. 1r free sign ()
22:
return update_free
32


Reviews:
[{"review_id": "IAqcrwbGei", "reviewer": "Reviewer_HZzU", "summary": "The paper introduces FRUGAL, a memory-efficient optimization framework designed for scalable training of large language models (LLMs). The key idea behind FRUGAL is gradient splitting, which enables a mix of stateful optimizers (e.g., AdamW) for a low-dimensional subspace and state-free optimizers (e.g., signSGD, SGD) for the remaining directions. This allows full-rank parameter updates while keeping memory overhead low.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and address their questions below.\n\nThe tables and proof that we will refer to using an apostrophe (e.g., Table 1') can be found at the anonymous link https://anonymous.4open.science/r/frugal-618F/rebuttal.pdf.\n\n>Computational Overhead: How does FRUGAL compare in terms of computational efficiency, compared to GaLore and BAdam?\n\nWe 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 Table 7'. 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`.\n\nThe 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.\n\n>Hyperparameter Sensitivity: most of experiments set rho=0.25. How does it get selected?\n\nThe value $\\rho=0.25$ was chosen to match the number of trainable parameters in the GaLore baseline, where the rank of the projection in most experiments was $1/4$ of the hidden size. We appreciate the reviewer's question and will add this explanation to the text.\n\n>How about its sensitivity?\n\nWe conducted an ablation study to verify the robustness of our algorithm to density $\\rho$. This experiment is described in Section 6.4, lines 420-422, and in Table 15. The results indicate that perplexity increases gradually as we transition from $\\rho=1.0$ (which essentially coincides with AdamW) to $\\rho=0.0$.", "claims_and_evidence": "1. FRUGAL enables full-rank updates with lower memory overhead than existing methods. Supported by experimental results showing that FRUGAL achieves performance close to full AdamW training while using significantly less memory.\n\n2. State-free optimizers (e.g., signSGD) are effective for certain LLM components. The authors provide evidence that embeddings, RMSNorms, and all but the Logits layer can be trained with signSGD with minimal accuracy loss.\n\n3. FRUGAL achieves state-of-the-art memory-efficient training for both pre-training and fine-tuning. Results show that FRUGAL outperforms GaLore and BAdam in pre-training while achieving comparable fine-tuning performance to LoRA with lower memory costs.\n\n4. FRUGAL maintains convergence rates comparable to standard optimizers, supported by theoretical proofs", "methods_and_evaluation_criteria": "FRUGAL splits gradient updates into two subspaces:\n1. Stateful subspace (L) – Updated using AdamW or another optimizer that maintains state.\n2. State-free subspace (M) – Updated using signSGD or SGD, eliminating the need for momentum and variance buffers.\n\nAlso it supports various subspace selection strategies, including:\n1. SVD-based projections (like GaLore).\n2. Random projections (for computational efficiency).\n3. Block-wise updates (similar to BAdam).", "theoretical_claims": "The paper provides theoretical guarantees for FRUGAL’s convergence. The proofs are generally correct.", "experimental_designs_or_analyses": "The experimental evaluation includes:\n1. Pre-training performance (LLaMA, C4 dataset)\n2. Fine-tuning performance (RoBERTa, GLUE benchmark)\n\nFindings:\n1. FRUGAL achieves similar performance to AdamW at lower memory cost.\n2. It outperforms GaLore and BAdam on pre-training tasks.\n3. SignSGD works well for embeddings and normalization layers but degrades performance for the Logits layer.", "relation_to_broader_scientific_literature": "memory-efficient training approaches (LoRA, GaLore, BAdam)", "essential_references_not_discussed": "Not I aware of", "other_strengths_and_weaknesses": "None", "other_comments_or_suggestions": "None", "questions_for_authors": "Computational Overhead: How does FRUGAL compare in terms of computational efficiency, compared to GaLore and BAdam?\n\nHyperparameter Sensitivity: most of experiments set rho=0.25. How does it get selected? How about its sensitivity?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dvYYWgyUbJ", "reviewer": "Reviewer_wtSQ", "summary": "The paper proposes a memory efficient way of combining existing stateful (like Adam) and stateless (like signSGD) optimizers, by running stateful optimizers in a low dimensional space and stateless optimizers in the complementary space. They provide results on pretraining as well as finetuning setup and show that the method outperforms other low memory counterparts such as LORA and BAdam, while being comparable to Adam in some setups.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their detailed comments. We appreciate their commendation of our extensive experimental evaluation and the ablation study we conducted. We address their concerns and questions below.\n\nThe tables and proof that we will refer to using an apostrophe (e.g., Table 1') can be found at the anonymous link https://anonymous.4open.science/r/frugal-618F/rebuttal.pdf.\n\n>1. Would it be possible to rerun some of the pretraining experiments with beta2 = 0.95?\n\nWe conducted additional experiments with a $\\beta_2=0.95$ for models with sizes 60M, 130M, and 350M. The results are presented at Table 1'.\n\nThe results are consistent with the observations from the main paper- FRUGAL performs significantly closer to the full-rank baseline than GaLore and BAdam.\n\nWe would like to point out that, when selecting the value for $\\beta_2$ for the main experiments in the paper, we followed the setup from our primary baseline, GaLore. We agree that $\\beta_2=0.95$ is a standard choice for training LLMs. \n\n>3. How does the proposed method differ from AdaMem?\n\nWe thank the reviewer for bringing this paper to our attention. This algorithm indeed shares similarities with our framework, as it also uses residual for ensuring a full rank update. We will certainly include a discussion of this paper in the camera-ready version.\n\nWe would also like to emphasize that FRUGAL is a more general method compared to AdaMeM. Specifically, Algorithm 1 allows for choosing among various types of state-free optimizers, state-full optimizers, and projections, making it significantly more flexible than AdaMeM. Furthermore, AdaMeM can be considered a special case of our proposed FRUGAL framework, with Adafactor with momentum as the state-full optimizer, one-sided Adafactor as the state-free optimizer, and SVD-based projection.\n\n>2. Would it be possible to include AdaMem as a baseline? \n\nWe reimplemented AdaMeM (see GaLore-based implementation at the https://anonymous.4open.science/r/frugal-618F/adamem.py) and conducted pre-training experiments on models up to 350M in size. Following the original work, we swept the learning rate over the range [1e-4, 3e-4, 1e-3, 3e-3] and the delta parameter over [0.5, 1.0]. The results in Table 2' indicate that AdaMeM slightly underperforms compared to the FRUGAL. \n\n>4. Can the gains of the given method be theoretically shown in a similar setting as Theorem 4.1 of AdaMem? This is to provide a setting where the convergence rate with momentum is better, and thus, it could be understood how the method interpolates between the two rates.\n\nWe agree with the reviewer that having an analysis for such a setup would help provide better intuition about the capabilities of the proposed method, and the paper would definitely benefit from it. However, such an analysis is a non-trivial task, and our attempts during the process of creating a paper showed that it is not readily obvious how to achieve such a result.\n\nRegarding Theorem 4.1 from [1], we believe there was a mistake in its proof, and the theorem is *incorrect*. See the description of the error in Proof 1'.\n\nWhile slightly modifying the formulation—specifically, replacing top-k with bottom-k—could yield desirable results, we believe they would not be useful or illustrative. This is because applying accelerated methods to the bottom-k eigenspace is somewhat counterintuitive and contradicts methods used in practice (including AdaMeM). Furthermore, we believe that considering a quadratic setup with SVD is an oversimplified setup as computing SVD is harder than solving the problem and yields a closed-form solution, so no optimization is needed beyond SVD. \n\nWe would also like to emphasize that the existing analysis presented in Section 5 is conducted under assumptions that closely match real-world conditions, achieves optimal convergence rates, and thus serves as a significant theoretical grounding for our method.\n\nWe would be glad to continue the discussion and address any follow-up questions the reviewer may have. For now we believe that we addressed all the reviewer's concerns and questions, none of which is a serious issue with our approach. Therefore, we would kindly request the reviewer to reconsider their score.", "claims_and_evidence": "The paper provides **mostly** adequate evidence for the claims made about the superiority of the FRUGAL optimizer over existing low memory optimizers. I have mentioned mostly as the hyperparameter used for Adam as mentioned in the Appendix has beta 2 of 0.999. Although this has been the default beta2 value in vision literature, in recent LLM works such as Zhao, the optimal beta2 is close to 0.95.", "methods_and_evaluation_criteria": "Yes methods and evaluation criteria make sense.", "theoretical_claims": "No I did not verify the correctness of the theoretical claims.", "experimental_designs_or_analyses": "I looked at the hyperparameter sweeps for most of the experiments. My main concern is the use of beta2 = 0.999 for Adam, which is not standard in language model setups.", "relation_to_broader_scientific_literature": "The work is well placed within the literature of memory efficient optimizers.", "essential_references_not_discussed": "One of the crucial references not discussed is AdaMem. It also has the similar idea of decomposing the current update into a ‘top’ direction where momentum is maintained and a bottom direction where only preconditioning happens without any momentum. This is an essential baseline to be compared with.\n\nVyas et al. 2024 - AdaMeM: Memory Efficient Momentum for Adafactor", "other_strengths_and_weaknesses": "One of the main strength of the paper is the extensive evaluation of the FRUGAL optimizer on various pretraining and fine tuning settings, and also various ablations showing the effect of various choices on the performance.\n\nThe main weakness is the missing reference to AdaMem and a comparison to the optimizer.", "other_comments_or_suggestions": "The theory provided within AdaMem work clearly shows an improvement based on how momentum is maintained, while in this case, the theory is in a setting where the rates for SGD and SGDM are the same. It would be better to consider a setting where the rates with momentum differ and show how the given optimizer interpolates between the two rates.", "questions_for_authors": "1. Would it be possible to rerun some of the pretraining experiments with beta2 = 0.95? (even at scale of 130m works)\n\n2. Would it be possible to include AdaMem as a baseline?\n\n3. How does the proposed method differ from AdaMem?\n\n4. Can the gains of the given method be theoretically shown in a similar setting as Theorem 4.1 of AdaMem? This is to provide a setting where the convergence rate with momentum is better, and thus it could be understood how the method interpolates between the two rates.\n\n--------------------------------\n\nUpdated the score based on rebuttal.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0YKKedAUbn", "reviewer": "Reviewer_Hho9", "summary": "This paper focuses on memory-efficient training by using different optimizers on different subspaces of the gradient, and uses different methods for projecting the gradient onto the state-full subspace. They extend AdaLayer to use signSGD instead of SGDM.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's comprehensive feedback. We are glad that they appreciated the theoretical convergence guarantees and strong experimental results. We also answer their questions below. The tables that we will refer to using an apostrophe (e.g., Table 1') can be found at the anonymous link https://anonymous.4open.science/r/frugal-618F/rebuttal.pdf.\n\n>The paper combines different lineages of memory-efficient optimizers …  but this makes it difficult to see where the efficiency is actually coming from.\n\nWe agree that understanding how different components of the overall framework affect the final results is critically important. Below, we address each of the enumerated points:\n\n>1) low-rank projection-based, 2) block coordinate descent,\n\nWithin our framework, we consider two main options for selecting the state-full subspace: 1. a low-rank projection for each trainable matrix (denoted as SVD in the paper), and 2. a block coordinate descent-like approach where some subset of matrices is fully within the state-full subspace while all others are in the state-free subspace (denoted as Blockwise). Experimental results comparing these approaches are presented in Table 1. Also, see the detailed discussion of memory and computational requirements in Section 4 and Appendix B.\n\n>3) sign-based\n\nAs state-free optimizers, we only considered signSGD and SGD. The discussion and experimental results can be found in Section 4 and Table 8.\n\n>4) state-free/full hybrid\n\nThe incorporation of state-free subspace optimization is the main contribution of our framework. The advantages of this approach are demonstrated, for example, in Tables 1 and 2.\n\nTo summarize the results, the most significant impact on the metrics is due to the addition of the state-full/state-free hybrid, followed by the use of sign-based optimization for the state-free subspace. The type of projection has the least effect. \n\n>Adding \"AdamW, ρ=0.25\" to Table 2 would help quantify the contribution to memory-efficiency of each method.\n\nWe assume that by \"AdamW, $\\rho=0.25$\" the reviewer means that only $\\rho=0.25$ of the Linear layer parameters are unfrozen and trained using AdamW. We would like to note that this setup is very similar to our baseline BAdam, where the set of unfrozen parameters changes every $T$ steps. For a more complete picture, we also conducted experiments with a setup where the active parameters do not change throughout the training process. As active parameters, we experimented with selecting the first and last $\\rho=0.25$ layers. The results are presented in Table 3'.\n\n>I am particularly interested in the extension of AdaLayer's state-free optimizer from SGDM to signSGD. It would be nice to see an ablation for just this change. I would also like to see experiments for 0<ρ<0.25.\n\nSince the original version of AdaLayer (Algorithm 1 from [1]) does not use SGDM, we assume that the reviewer was referring to its version in Section 3.2 [1]. \nAs requested, we conducted pretraining experiments on models up to 350M, replacing SGDM with signSGD in this version of the algorithm and also using FRUGAL with AdaLayer (see implementation here https://anonymous.4open.science/r/frugal-618F/adalayer.py) as the state-full optimizer with different values of $\\rho$. The final results, which also include the results from Table 9 of the original paper, are presented in Table 4'. As can be seen, the variant with AdaLayer demonstrates a very similar trend to the original FRUGAL, albeit with slightly worse values.\n\n>The essential reference [Zhao et al. 2024b] is cited, but is not emphasized properly.\n\nWhile the AdaLayer shares some similarities with FRUGAL with ρ=0.0, we consider our framework more general since *it allows configurations at other values of $\\rho$ that have no equivalent in AdaLayer•. Therefore, we believe our work is still closer to our primary baselines, such as GaLore.\n\n>It would also be nice to have some comments about how the memory-efficiency improves with scale (number of parameters).\n\nAn exact estimate of the memory savings from using FRUGAL compared to AdamW is provided in Table 5'.\n\n> It is unclear how the authors choose which linear layers to optimize with the state-full optimizer. Depending on this choice, perhaps \"FRUGAL, ρ=0.25\" could yield even better results?\n\nWe compared several methods of alternating between state-full subspaces in Section 6.4 and Table 11. The optimal selection of stateful subspaces is a complex task, and we leave it for future work.\n\n>State-free optimizers are more sensitive to hyperparameters, so adapting them in a certain sub-space may inherit this weakness. Have the authors encountered such problems?\n\nWe did not observe significant sensitivity in our experiments. For example, varying $\\beta_2$, as presented in Table 1', affects FRUGAL in approximately the same way as it does AdamW.\n\n[1] Zhao et al., Deconstructing what makes a good optimizer for language models, 2024.", "claims_and_evidence": "The proposed method outperforms existing methods such as GaLore and BAdam when training Llama-60M, 130M, 350M, and 1B on the C4 dataset.", "methods_and_evaluation_criteria": "The authors perform extensive experiments across various model scales and hyperparameter settings, providing a thorough analysis of optimizer performance and stability. The benchmark datasets are standard and I did not find anything missing in the evaluation criteria.", "theoretical_claims": "This work presents the first theoretical analysis of an extended block coordinate descent framework where the remaining layers are also updated with a different algorithm. I did not find any errors in the proofs.", "experimental_designs_or_analyses": "The paper combines different lineages of memory-efficient optimizers 1) low-rank projection-based, 2) block coordinate descent, 3) sign-based, and 4) state-free/full hybrid, but this makes it difficult to see where the efficiency is actually coming from. Adding \"AdamW, ρ=0.25\" and \"AdaLayer [Zhao et al. 2024b] ρ=0\" to Table 2 would help quantify the contribution to memory-efficiency of each method. It would also be nice to have some comments about how the memory-efficiency improves with scale (number of parameters).", "relation_to_broader_scientific_literature": "Although the paper puts a lot of emphasis on optimizers such as GaLore, ReLoRA, and BAdam, the key contribution of this paper is more related to AdaLayer [Zhao et al. 2024b], but with signSGD instead of SGDM. If a majority of layers can be optimized with state-free optimizers, the efficiency of the state-full part is less significant.", "essential_references_not_discussed": "The essential reference [Zhao et al. 2024b] is cited, but is not emphasized properly.", "other_strengths_and_weaknesses": "In the paper, it says \"ρ denotes the proportion of the Linear layer parameters in the state-full subspace\", but it is unclear how the authors choose which linear layers to optimize with the state-full optimizer. Depending on this choice, perhaps \"FRUGAL, ρ=0.25\" could yield even better results?", "other_comments_or_suggestions": "I am particularly interested in the extension of AdaLayer's state-free optimizer from SGDM to signSGD. It would be nice to see an ablation for just this change. I would also like to see experiments for 0<ρ<0.25 to see whether it yields similar accuracy while reducing the memory consumption.", "questions_for_authors": "State-free optimizers are more sensitive to hyperparameters, so adapting them in a certain sub-space may inherit this weakness. Have the authors encountered such problems?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pqtRc3VIG6", "reviewer": "Reviewer_3FWu", "summary": "This paper presents a novel approach to reduce memory overhead during LLM training by dividing the model parameters into two distinct groups. One group is optimized using Adam-family optimizers, which maintain optimizer states, while the other group is trained with state-free optimization methods such as SGD and sign-SGD. This design aims to minimize memory consumption while maintaining acceptable performance degradation.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed feedback and respond to their concerns and questions hereafter.\n\nThe tables that we will refer to using an apostrophe (e.g., Table 1') can be found at the anonymous link https://anonymous.4open.science/r/frugal-618F/rebuttal.pdf.\n\n>Recent advancements, such as the Lion optimizer, exhibit a highly similar approach that closely aligns with this claim.\n\nWe kindly disagree with the reviewer that the existence of the Lion work diminishes the contribution of our research. While some prior works [1,2] have indeed used sign-based approaches, it is important to note that all these methods **incorporated momentum**. To the best of our knowledge, we are the first to use signSGD **without momentum** to successfully train the majority of language model parameters.\n\n>Algorithm 2 seems problematic in its treatment of parameters that are not part of the set $J_k$. Specifically, the algorithm updates $m_j^k$ for these parameters, which equivalently assumes their gradients are zero.\n\nWe thank the reviewer for pointing out this minor inaccuracy with the update of $m$ in line 3. We intended the following update formula:  \n\n$\\tilde{m}_j^{k} \\leftarrow \\begin{cases}  (1-\\beta)\\tilde{g}_j^{k} + \\beta\\tilde{m}_j^{k-1} & \\text{if } j \\in J_k; \\\\ \n0 & \\text{otherwise;} \n\\end{cases}$\n\nThis ensures that when $j$ re-enters $J_k$, the value $m_j$ is reset to 0 (this is assumed in Equation 3).\n\n>If the set $J_k$ is sampled randomly, Algorithm 2 would quite resemble standard SGD. Furthermore, as $\\beta$ approaches 1, $m_j^k$ converges to zero, preventing the algorithm from achieving convergence. \n\nWe kindly disagree with the reviewer on this issue. Since $j\\notin J_k$ are updated without using $\\beta$, and for $j\\in J_k$, the algorithm implements standard SGDM, which works well with $\\beta$ values close to 1, we do not see why this should prevent the convergence of our algorithm. Moreover, this statement contradicts Theorem 5.2, which we consider to be correct.\n\n>The purpose and functionality of Line 6 in Algorithm 1 are unclear. This step fails to effectively distinguish between state-full and state-free parameters, diminishing the clarity of the proposed method's design.\n\nBy $P^{-1}$ in Line 6, we mean the right inverse of $P$. Thus, $P^{-1}(P(g))$ represents the projection of $g$ onto the low-rank subspace. We will add this clarification in the final version of the paper.\n\n>Theoretical findings, particularly Theorem 5.2, suggest that combining SGD with SGDM results in a consistently worse upper bound.\n\nWe note that this is the case only in deterministic scenario, where the main issue is that the bias from momentum affects all the coordinates. At the same time, the variance reduction effect is only present in momentum coordinates. This is not improvable in the worst case, but in the average (best) case, the benefit of momentum can be more prevalent compared to (sign)SGD, which explains our numerical findings.\n\n>The experimental results provided for fine-tuning are insufficient. Notably, the fine-tuning experiments on RoBERTa using the GLUE benchmark are somewhat outdated in the context of modern LLM benchmarks.\n\nAt the reviewer's request, we conducted additional fine-tuning experiments. Specifically, we evaluated the effectiveness of various fine-tuning algorithms on LLaMA 3.1-8B in Commonsense Reasoning.\n\nFollowing the setup in [3], we fine-tuned the model on the Commonsense170K dataset [3] and evaluated accuracy across 8 datasets (see the full list in [3], Section 3.1). We used the same hyperparameters as in [3], and varied the learning rate among [5e-6, 1e-5, 2e-5, 5e-5, 1e-4, 2e-4] for all methods.\n\nThe results, presented in Table 6', show that FRUGAL slightly outperforms both LoRA and GaLore. Notably, it achieves this even with $\\rho=0.0$ (using 0 memory for optimizer state).\n\n>Since SVD can be viewed as a special case of block-wise projection (i.e., full-rank projection without dimensionality reduction), it is unexpected that the block-wise approach yields inferior results. \n\nWe believe there has been a slight misunderstanding here regarding the SVD and Blockwise methods of selecting the state-full subspace. By SVD, we meant the approach of applying a low-rank projection to **each** trainable matrix. By block-wise, we refer to a block coordinate descent-like strategy where a subset of matrices resides  in the state-full subspace, while **the others entirely remain in the state-free subspace**.\n\nTo sum up, we believe that we addressed all the reviewer's concerns and questions, none of which is a serious issue with our approach. Therefore, we would kindly request the reviewer to reconsider their score.\n\n[1] Chen et al., Symbolic discovery of optimization algorithms, NeurIPS 2023.\n\n[2] Zhao et al., Deconstructing what makes a good optimizer for language models, ICLR 2025.\n\n[3] Hu et al., Llm-adapters: An adapter family for parameter-efficient fine-tuning of large language models, EMNLP 2023.", "claims_and_evidence": "The claim that this work is \"the first to train the majority of language model parameters using signSGD without momentum\" appears to be overstated. Recent advancements, such as the Lion optimizer, exhibit a highly similar approach that closely aligns with this claim. Moreover, with carefully chosen hyperparameters, the claim (using signSGD for LLM) may effectively reduce to a special case of Lion, further diminishing the novelty of the stated contribution.", "methods_and_evaluation_criteria": "- Algorithm 2 seems problematic in its treatment of parameters that are not part of the set $J_k$. Specifically, the algorithm updates \n$m_j^k$  for these parameters, which equivalently assumes their gradients are zero. If the set $J_k$  is sampled randomly, Algorithm 2 would quite resemble standard SGD. Furthermore, as $\\beta$ approaches 1, $m_j^k$ converges to zero, preventing the algorithm from achieving convergence. This issue raises concerns about the proposed method's stability and effectiveness.\n\n- The purpose and functionality of Line 6 in Algorithm 1 are unclear. This step fails to effectively distinguish between state-full and state-free parameters, diminishing the clarity of the proposed method's design.", "theoretical_claims": "- Theoretical findings, particularly Theorem 5.2, suggest that combining SGD with SGDM results in a consistently worse upper bound. This outcome appears to undermine the theoretical soundness of the proposed FRUGAL method. A more comprehensive discussion or additional theoretical justification is necessary to clarify this apparent limitation.", "experimental_designs_or_analyses": "The experimental results provided for fine-tuning are insufficient. Notably, the fine-tuning experiments on RoBERTa using the GLUE benchmark are somewhat outdated in the context of modern LLM benchmarks. Consequently, the conclusions drawn from these experiments may have limited generalizability to contemporary LLM training and fine-tuning settings.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "Not Found.", "other_strengths_and_weaknesses": "See above", "other_comments_or_suggestions": "See above", "questions_for_authors": "An intriguing observation arises from Table 1, where the block-wise results underperform compared to the SVD-based method. Since SVD can be viewed as a special case of block-wise projection (i.e., full-rank projection without dimensionality reduction), it is unexpected that the block-wise approach yields inferior results. A deeper investigation into this discrepancy would enhance the paper's insights.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The authors propose a hybrid optimization method aimed at improving memory efficiency without compromising performance. The core idea is intuitive and somewhat straightforward, with similar approaches having been explored in the literature. It is recommended that the authors cite them in relevant works, such as Fira [1]. Therefore, the novelty of the methodology appears incremental.

Nevertheless, the reviewers recognize the value and contributions of this work, particularly noting the strength of the empirical evaluation and comprehensive ablation studies provided. Additionally, exploring the method's performance in scenarios involving longer context lengths and over-training setups could further highlight the robustness and practical utility of the proposed approach.

[1] Chen, Xi, et al. "Fira: Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint?." arXiv preprint arXiv:2410.01623 (2024).
Decision: Accept (poster)