Fira: Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint?
Authors: Xi Chen
Kaituo Feng
Changsheng Li
Xunhao Lai
Xiangyu Yue
Ye Yuan
Guoren Wang
Venue: NeurIPS 2025 poster ()
Keywords: Large Language Model
Memory Efficient Training
Full-rank Training
Primary area: deep_learning
TL;DR: We propose Fira, a plug-and-play memory-efficient training framework of LLMs to enable full-rank training consistently under the low-rank constraint.

Abstract: Low-rank training has emerged as a promising approach for reducing memory usage in training Large Language Models (LLMs). Previous methods either rely on decomposing weight matrices (e.g., LoRA), or seek to decompose gradient matrices (e.g., GaLore) to ensure reduced memory consumption. However, both of them constrain the training in a low-rank subspace, thus inevitably leading to sub-optimal performance. To resolve this, we propose a new plug-and-play training framework for LLMs called Fira, as the first attempt to consistently preserve the low-rank constraint for memory efficiency, while achieving full-rank training (i.e., training with full-rank gradients of full-rank weights) to avoid inferior outcomes. 
First, we observe an interesting phenomenon during LLM training: the scaling impact of adaptive optimizers (e.g., Adam) on the gradient norm remains similar from low-rank to full-rank training. In light of this, we propose a \textit{norm-based scaling} method, which utilizes the scaling impact of low-rank optimizers as substitutes for that of original full-rank optimizers to achieve this goal.
Moreover, we find that there are potential loss spikes during training. To address this, we further put forward a norm-growth limiter to smooth the gradient.
Extensive experiments on the pre-training and fine-tuning of LLMs show that Fira outperforms both LoRA and GaLore. Notably, for pre-training LLaMA 7B, our Fira uses $8\times$ smaller memory of optimizer states than Galore, yet outperforms it by a large margin.

Fira: Can We Achieve Full-rank Training of LLMs under Low-rank Constraint?
Xi Chen', Kaituo Feng', Changsheng Li Xunhao Lai",
Xiangyu Yue², Ye Yuan', Guoren Wang
Beijing Institute of Technology
School of Intelligence Science and Technology, Peking University
MMLab, The Chinese University of Hong Kong
Hebei Province Key Laboratory of Big Data Science and Intelligent Technology com,
xichen xichen.fy@gmal f.w.mail..a.......m com, kai taittuoffggaaaai lcs@bit. edu. cn, laixunhao@pku. edu. cn, xyyuešie cubheiieecchk... cuhk. eku. yuan- edu cn, wanggrbite126. com com https:/ thub com com/xichen-ffyiiie
Abstract
Low-rank training has emerged as promising approach for reducing memory usage in training Large Language Models (LLMs). Previous methods either rely on decomposing weight matrices (c.g.. LoRA), or seek to decompose gradient matrices (e.g., Gal Lore) to ensure reduced memory consumption. However, both of them constrain the training in a low-rank subspace, thus inevitably leading to sub-optimal performance. To resolve this, we propose a new plug-and-play training framework for LLMs called Fira, as the first attempt to consistently preserve the low-rank constraint for memory efficiency, while achieving full-rank training (i.e., training with full-rank gradients of full-rank weights) to avoid inferior outcomes. First, we observe an interesting phenomenon during LLM training: the scaling impact of adaptive optimizers (e.g., Adam) on the gradient norm remains similar from low-rank to full full-rank training. In light of this, we propose a norm-based scaling method, which utilizes the scaling impact of low-rank optimizers substitutes for that of original full-rank optimizers to achieve this goal. Moreover, we find that low-rank optimizers may lead potential loss spikes during training. To address this, we further put forward a norm- norm-growth limiter to smooth the gradient. Extensive experiments on the pre-training and fine-tuning of LLMs show that Fira outperforms both LoRA and GaLore. Notably, for pre-training LLaMA 7B, our Fira uses 8x smaller memory of optimizer states than Galore, yet outperforms it by a large margin.
Introduction
In recent years, Large Language Models (LLMs) have achieved remarkable advancements in various domains 2 30 10]. While the substantial increase in model size contributes significantly to these advancements, also introduces considerable memory bottlenecks, especially for optimizer states 43 For instance, pre-training LLaMA-7B model from scratch requires at least 58 GB memory, allocated follows: 14GB for loading parameters, 14GB for weight gradients, 28GB for Adam [66 optimizer states, and 2GB for activations 43]. Notably, the optimizer states consume even more memory than the parameters themselves. address this, low-rank training has demonstrated
"Corresponding author
Training the model with single batch size and maximum sequence length of 2048 under BF16 precision.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

W J J( Pun
Pe(P ( [W
Po(r (WW)) +(1 P (WWoWW.
Triiaabee
gacb g
Jaub 9
Frozen
Trainable
Trainable
godre
updato
update
W
sub
W
W
full-rank gradients of low-rank weights
low-rank gradients full-rank weights full-rank gradients of full-rank weights
(a) LoRA
(b) GaLore
(c) Fira (ours)
Figure 1: This analyses three types of memory-efficien approaches at macro level.
its effectiveness reducing the memory usage of the optimizer states by conducting training in low-rank subspace 43 [3]
The current low-rank training methods can be broadly divided into two categories: weight matrix- based and gradient matrix-base low-rank decomposition. For the weight matrix decomposition methods, the most representative one is Low-Rank Adaptation (LoRA) [3], where its basic idea is to use low-rank matrices as decomposed representations of the pre-trained weights during training, as low-rank subspace of the weights. This will inevitably cause the reduction of representation capacity, leading to sub- optimal outcomes [40 33). Although the variant ReLoRA [18] attempts to extend the application of LoRA from fine-tuning to pre-training, by periodically updating high-rank weights with multiple low-rank updates. It still requires full-rank weight training as a warm-up before low-rank training, thus rendering memory efficiency unachievable [43].
shown in Figure [ (a). However, the optimization of LoRA is constrained
For the gradient matrix decomposition based methods, the typical one the gradient low-rank projection (GaLore) proposed recently [43]. In contrast to LoRA, GaLore attempts to reduce the memory usage in optimizer states via decomposing the gradient matrix, as shown in Figure [ (b). While GaLore supports the training of full-rank weights, it leverages only low-rank gradients. restricting them to a low-rank subspace. Consequently, any gradient information outside this subspace is lost, in contrast to training with full-rank gradients. Note that since these methods constrain LLM training to low-rank subspace, this inevitably leads to sub-optimal results compared to full-rank training (i.e., training with full-rank gradients and full-rank weights). This raises the question: Can we achieve full-rank training for L.MSs while consistently maintaining a low-rank constraint?
In light of this, we propose a new memory- efficient training framework for LLMs, called Fira, which, to the best of our knowledge, is the first to achieve full-rank training while consistently maintaining a low-rank con- straint. To achieve this goal, significant challenge is that the low-rank constraint makes it hard to preserve complete optimizer states (e.g., gradient momentum and vari- ance) of full-rank weights in the commonly- used adaptive optimizer (e.g., Adam). As result, the adaptive optimizer fails to correct
Table 1: Cosine Similarity and MSE between low- rank and full-rank scaling factors at the matrix and column levels for pre-training LLaMA models rang- ing from 60M to IB parameters, averaged over 10K steps. Detailed analyses are presented Appendix D
Matrix Level
Columa Level
Size
Cosine Similarity
MSE Cossne Similarity MSE
60M
0.9922
3e-04
0.9273
3e-05
130M
0,9901
2c-04
0.9046
2c-05
350M
0,9893
lc-04
0.9174
1c-05
IB
0.9795
2c-04
0.9229
1c-05
the raw gradients without corresponding op-
timizer states. Without this correction, adaptive optimization algorithms would degrade into simple stochastic gradient descent (SGD), leading to significantly reduced optimization performance [9]
This point is further validated in Section 4. and Section [5.4]
Fortunately, we observe interesting phenomenon during LLM training: at weight matrix level, the scaling factors 5 of low-rank gradients exhibit strong similarity to those of full-rank gradients. As illustrated in Table Cosine Similarity and Mean Squared Error (MSE) between low-rank and full-rank scaling factors cxhibit significant similarity (Cosine Similarity close to 1, while MSE close to 0). Detailed quantitative analyses of this similarity are presented in Appendix D
'The scaling factor de(Re) defined [+())] Rell
where R the norm of the raw gradient, [l)(),))]
the norm of the gradient corrected by the gradient correction function of the optimizer (e.g.. Adam).

---

Based on this observation, we put forward norm-based scaling method that utilizes the scaling factor of a weight matrix in low-rank training replace the corresponding matrix's 's scaling factor in full-rank training. In this way, our scaling factor can also play similar role in correcting the raw gradient, as adaptive optimizers do. Therefore, we can enable full-rank training while preserving the low-rank constraint.
However, we observe potential sudden increases in gradients caused by low-rank optimizers during the early stages of training, which can lead to spikes in training loss, as illustrated in Figure This phenomenon primarily arises because our norm-based scaling method lacks the gradient stability compared to the original full-rank Adam optimizer. Such gradient instability can trigger substantial parameter updates, causing the loss function to reach a much higher value and undermining prior optimization efforts [TT 39]. Despite the use of gradient clipping techniques [25,, this issue may not be adequately resolved, as shown Figure 3 To this end, we propose a norm nom-grrwth limiter, which aims to smooth the gradient by restricting the magnitude of the gradient norm's 's increase. By employing our limiter, we adaptively convert sudden rises gradients into gradual increases, thereby facilitating a smooth update that effectively mitigates the problem of loss spikes.
Our main contributions can be summarized as follows:
1. We propose Fira, a plug-and- and-play memory training framework of LLMs, constitut- ing the first attempt to enable full-rank training consistently under the low-rank constraint. We will release the source code and package our Fira into a Python library for easy use.
2. We design two components in Fira: norm-based scaling strategy that leverages the scaling effects of low-rank optimizers to facilitate full-rank training, and a norm-gre corm-growth limiter to address the issue of loss spikes by limiting the growth of gradient norm.
3. Extensive experiments on the pre-training and fine-tuning of LLMS across various parameter counts (60M, 1 130M, 350M, 1B, 7B) validate the effectiveness of Fira in both pre-training and fine-tuning tasks, surpassing both LoRA and GaLore.
2 Related Work
Low-rank adaptation. Low-Rank Adaptation (LoRA) has been introduced by [13 as an efficient fine-tuning method for LLMs. The core idea of LoRA is to freeze the pre-trained weights and introduce trainable low-rank matrices as decomposed representations of the pre-trained weights.
In this way, the memory usage of training LLMs could be saved. Recently, a variety of methods by extending LoRA have been proposed to further improve the performance [41] 33
35
40 [.
For instance, ReLoRA [X] is proposed to extend the application of LoRA from fine-tuning pre- training. However, still requires full-rank warm-up training before low-rank training, which prevents achieving memory efficiency. It is worth noting that while LoRA- basae methods reduce memory usage by limiting training to a low-rank parameter subspace, they inevitably reduce representation capacity 335
have indicated that LoRA may yield sub-optimal performance since its low-rank constraints in parameters. Inspired by traditional projected gradient descent methods [6 4. GaLore 433 has been proposed recently to mitigate this problem. It enables full-parameter learning under low-rank constraints by projecting the gradient into low-rank subspace, reducing memory usage for optimizer states. However, while GaLore allows memory-efficient full- parameter training, it confines the gradient to a low-rank subspace, discarding the portion outside the subspace and resulting in significant information loss.
Gradient projection. Recent works 40 35
System-Base Based memory-efiee efficient techniques. Many system-based techniques have been developed reduce memory usage in LLM training 27). However, most of these methods achieve memory efficiency by compromising either time or precision. Gradient checkpointing [5] is proposed to reduce memory usage by trading increased computational time for the re-computation of activations. Quantization [9] reduces memory consumption by using lower-bit data types but at the cost of model precision. Memory offloading [38]] 27 reduces GPU memory usage by using non-GPU memory (e.g., CPU) as extension. However, it introduces additional communication overhead, such as CPU-GPU transfer time. It's important to note that our proposed method is complementary to these approaches and can potentially be combined with them to further reduce memory usage.
3

---

3 Preliminaries
3.1 Regular Full-Rank Training
At time step I, we denote the full-rank weight matrix as W. E IR The full-rank gradient can be represented as G2 Vw fWW,) IR=x", where is the objective function. Then the regular full-rank training can be expressed as follows:
W411 = - mG,(
W
(1)
where 17 is the learning rate, and W. is the gradient correction function of the optimizer (for vanilla SGD, (G,) G.). Instead of vanilla SGD, adaptive optimizers (e.g., Adam [16], Adam (2I)) are usually employed to correct the raw gradient for improving the training performance. However, this typically requires additional memory for storing optimizer states used in gradient correction. For
instance, Adam [16] requires storing the optimizer states M and V, which consume 2nn of memory. The gradient correction process is as follows:
Mj = 3 + (1- B1)
(2)
V1 32V (1-32)
(3)
11 33 M₁
(v.
1- 8i V.
where all matrix operations are element-wise. 31 and 82 are Adam's hyper-parameters, and is small constant (e.g., 1 x 10-*) used for numerical stability. Since this regular full-rank training typically consumes large amount of memory for training LLMs, many representative low-rank training methods, e.g.. LoRA [3] and Galore (43, have been proposed reduce memory usage in recent years.
3.2 Low-Rank Adaptation
The basic idea behind LoRA B is to use low-rank matrices decomposed representations of the pre-trained weights during training, in order to reduce memory usage. Formally, LoRA freezes the full-rank weight matrix Wo Rm*n and incorporates two low-rank matrices A, and B, for training
as:
W Wo + B,Ar,
where B, E ER A₂ E R and the rank r < min(m, n). While LoRA reduces memory usage by limiting training low-rank subspace of the weight, it inevitably diminishes the representation capacity of the weight matrix W, 355
3.3 Gradient Low-Rank Projection
In contrast to LoRA, GaLore 4 utilizes a projection matrix P. E Rmxr to project the full-rank
gradient G, E mxx to low-rank gradient RG = P.T.. G, E Rexa (m S n n By doing so, the memory usage of optimizer states could be reduced. The parameter update in GaLore can be formulated as:
W++ W, (Re)
(6)
where the projection matrix Pt can be obtained through singular value decomposition (SVD) of G and can be updated every T step:
G = UVV ~ Jouu, = [21, U2..... Ur]s
(7)
1=1
where is the i-th column vector of the left singular matrix U. By selecting the first columns of matrix U that correspond to the largest singular values, the projection matrix P. effectively captures the most significant directions in the gradient space, leading to faster convergence [43). The optimal switching frequency T usually set to be between 50 to 1000, and the additional computational overhead introduced by SVD is negligible (< 10%), stated in 443. Since Galore restricts the gradient in the low-rank subspace, the gradient information outside this subspace is lost, leading to inferior performance.
FFrr simplicity, we assume S following 43 If
Rr GoQe TXT ,Q: E =

---

Proposed Method
To achieve full-rank training under low-rank constraints, our framework, named Fira, consists of two important components: (i) a norm-based scaling method, enabling full-rank training by leveraging the scaling effects of adaptive optimizers; (ii) a norm-growth limiter, which restricts the growth of the gradient norm to prevent spikes in training loss. Next, we will elaborate on these two components.
4.1 Norm-Based Scaling
The low-rank constraint makes it challenging to record complete optimizer states for correcting raw gradients in full-rank training. Fortunately, we find an interesting phenomenon in L.L training: the scaling factor at the matrix level remains similar from low-rank training to full-rank training. Based on this observation, we propose a norm-based scaling strategy that approximately corrects the raw gradient, similar to adaptive optimizers, thereby enabling full-rank training.
Challenge analysis. Given the difficulty of
8
incorporating trainable low-rank weights into LoRA achieve full-rank weight training [43] we focus on investigating how to achieve full- rank gradient training by extending the gradient projection method, Galore, in this paper. In Ga- Lore, the projection matrix P. E Ruxr projects the full-rank gradient G2 E m2xx of the full- rank weight W, E gmxnn to the low-rank sub- space gradient R2 P. G. E I@"x" Then,
Fira
GaLore
GaLore-add
Training loss
Adam
the gradient outside this subspace can be rep- resented as: (I - P.P. )G, G2 R,R, In other words, the full-rank gradient G. can be divided into two terms: P.R, and (G, P.R.).
400
800
1200
1800
2400
3000
Iterations
Figure 2: Training loss of different methods for pre-training LLaMA 60M on C4 dataset (r dmodel 16/256 and 7'= 200).
In Goloree the optimizer states only store the information of R, instead of G2 to realize the low-rank constraint. The term of (G, P.R,) is directly discarded in Galore due to the lack of corresponding optimizer states for correction in optimizers. This would lead to significant information loss especially when r < domodel, where dmodel = min(m, n)) the full-rank dimension of models (This point can be verified in our experiment section, as illustrated in Figure 5 In Figure [ the validation perplexity of GaLore significantly increases at r = compared to r = 128 when dmodel 256, indicating substantial loss of information and decreased training performance). Intuitively, to capture the information of (G; - P.Re), we can directly add it based on Equ. 6 as follows:
W411 = W. Pebh - n P.Rs).
We denote the update strategy in Eqn. 8 as ard-add. However, as illustrated in Figure 2 GoLere- add exhibits almost improvement compared to updates using Eqn. i Gal.re.. This phenomenon primarily arises because the term of (G, R. doesn't 't have corresponding optimizer states for gradient correction. As result, the optimization of (G, R. uses vanilla SGD, yielding sub- optimal outputs. Besides, in GaLore-add, P,ú, employs the Adam optimizer for training while (G, P.Rc) employs vanilla SGD. This gradient misalignment may also account for the lack of noticeable improvement.
Similarity of scaling factor. To tackle this challenge, propose the concept of the scaling factor, which is defined as follows:
oe (Ro)l
[Rell
where the scaling factor &e represents the magnitude of the correction applied by the adaptive optimizer the gradient norm. Based on the scaling factor ér. we observe an interesting phenomenon during LLM training: the scaling factors at the matrix level exhibit a high degree of similarity between low-rank and full-rank training (column level will introduced later). As shown Table Cosine Similarity and MSE between low-rank and full-rank scaling factors exhibit significant similarity (Cosine Similarity close to 1, while MSE close to 0). Details analyses and additional experiments of this similarity are presented in Appendix D
5

---

Norm-based scaling, Building on the above observation, we propose norm based scaling method that leverages the scaling factor of a weight matrix in low-rank training as substitute for the corresponding factor in full-rank training:
W++1 W. (P. T)0: (R)((R)) - P. P.R.).
(10)
Through Eqn. 10 we can approximately correct (G; - PaR as adaptive optimizers do, so to achieve full-rank training under low-rank constraints.
To further enhance our approach, we introduce a more fine-grain mine-gaiied strategy for computing scaling factors in Eqn. by considering each column of the weight matrix individually:
Φel (Rt) |16b :e
i= 1.1.2....,
(11)
|
where Re... is the i-th column of R. and (e RR); is the i-th scaling factor. As evidenced in Table] scaling factors computed at the column level also demonstrate strong similarity between low-rank and full-rank training, further validating the effectiveness of this fine-grained strategy. Additional theoretical analysis and experimental results on scaling factors are provided in Appendix [
4.2 Norm-Growth Limiter
However, limited by low-rank constraints, we notice that there are still shortcomings in the norm- based scaling approach that need to be improved compared to the full-rank Adam. Specifically, potential sharp increases of the gradient may occur at the early stage of training, leading to spikes of training loss. As shown in Figure 3 Stra-w.o.llimitee (our method without using the proposed norm-growth limiter) experiences spikes in both gradient norm and training loss. In this section, we woongrowthh limiter that transforms abrupt gradient spikes into gradual, smooth increases, and then enhances the norm-based scaling method.
analyze the reasons for this issue and propose
Loss spike analysis. There are two main reasons for the spikes: (i) Switching the projection matrix Pt in gradient projection methods would cause instability during train- ing. As illustrated in figure[] both GaLore and GaLore-ade Lore-add exhibit sig- nificant training loss spikes at inte- ger multiples of T 200 (i.e., the frequency of switching the projec- tion matrix Pt). This instability oc- curs because, when switching pro- jection matrices P,, the optimizer retains states linked to the previous matrix, while the current input gra- dient uses a new projection matrix, leading to significant misalignment. Furthermore, as shown in Figure 2 GaLore-add also exhibits additional training spikes compared to Galore, reinforcing our earlier claim that directly incorporating (G, - P.R.) may introduce instability and hinder training; (ii) Maintaining the original direction of the raw gradient (G, -PRR,) may be insufficient for handling the sharp loss landscapes in LLM training. unlike Adam 339
Loss
Gradient Norm
7.5
Spike
Spike
180
6.0
Spike
Spike
120
4.5-
60
500
1000
1500
2000
500
1000
1500
2000
Fina
Fira-wo.-limlle
Fradggradiett-clipiing
Training loss and gradient norm of three variants of Fira for pre-training LLaMA 60M.
Figure 3:
Due to space constraints, further
analysis is provided in Appendix G
Addressing loss spikes. To address this issue, a straightforward solution is to use gradient clipping techniques (25) to avoid loss spikes. However, clipping based on the absolute norm of gradient matrices fails to account for significant differences between them, leading to sub-optimal results. This point can also be verified in Figure and Table To this end, we propose norm-growth limiter method that constrains the ratio of the current gradient norm to the previous step's 's norm to fixed ratio y when the gradient norm increases:
if ||Sa|| |
S.
then S, + || > S1
(12)
where y is a threshold ensuring that the rate of gradient growth does not exceed this value. S, = (R, P.Ri) is the corrected gradient by applying our based scaling. This approach

---

limits the magnitude of gradient norm increases, converting sudden spikes into gradual rises and thus preventing loss spikes. Moreover, by constraining the relative increase of each gradient matrix's norm, our method is more flexible than the absolute norm clipping. As illustrated in Figure and Figure Fira with our proposed limiter improves the optimization performance without significant spikes.
Discussion with other stabilization methods. variety of approaches have been developed to mitigate loss spikes and enhance training stability, including embedding normalization [[7]. gradient shrink on embedding layers [36]. tensor-wise scaling [8]. While these methods contribute to stabilization, they either primarily concentrate on embedding stabilization (e.g., embedding normalization and gradient shrink), or employ static scaling (e.g.. censor-wise scaling). Compared to our norm- growth limiter, these techniques are insufficient in addressing the adaptive instability that arises from the substantial variations among different weight matrices, thereby inadequately resolving the issue of loss spikes in norm-based scaling. Further detailed analyses and experimental results regarding the norm-growth limiter are discussed Section 5.4 and Appendix G
4.3 Overall Algorithm
We present the overall algorithm of Fira with Adam in Algorithm Tin in Appendix A.1 Our main components, the norm- norm-based scaling method, and the Norm-growthh limiter, are straightforward to implement, requiring only 3 additional lines of code. Moreover, Fira is plug-and play framework that can be easily integrated into the training process without requiring significant modifications. The Alug-and-play Pytorch-like pseudo-code of Fira is provided in Appendix A.2
Table 2: Comparison of different methods. Denote We E R (m S n2, rank r.
Method
GaLore
LoRA
Full-rank
Weights
7731
mn mr + 2ner
mn mr + nr
mn
Optimizer States
mar 2nr +
2mr + 2nr
2mn
Gradients
T7771
mm
ma + 11
mn
Full-Rank Gradients
Full-Rank Weights
X X
Pre-Training Fine-Tuning
It's worth noting that compared to Galore, Fira only introduces one parameter ||S;,1|l for each weight matrix in the optimizer state, which is negligible, as shown in Table Besides, in addition the original hyper-parameters of optimizers and gradient projection methods, Fira only adds one hyper-parameter 7 in the norm-growth limiter, whose choice is not sensitive to the performance of Fira. The hyper-parameter y set 1.01 across all experiments, which consistently yields satisfactory results. Sensitivity analyses of hyper-parameter y to training performance are provided in Appendix
Experiments
In this section, we validate the effectiveness of Fira pre-training and fine-tuning tasks of LLMs. In our experiments, we denote our method using the strategy of Eqn. a0 Fira-matrix, and denote our method additionally using the column-wi strategy of Eqn. TT as Fira.
5.1 Memory-Efficient Pre-training
to conduct the pre-training experiments. We compare Fira with GaLore 43), LoRA [33, ReLoRA 18, and full-rank training baselines. Adam optimizer is used to train all baselines and our method on the C4 dataset in the BF16 format.
Experimental setup. We follow the settings in Galore
4
The settings of these baselines can be found 4 The dataset C4 a colossal, cleaned version of Common Crawl's web crawl corpus, which is widely used in LLM pre-training (26. Following we utilize L.LaMA-based architectures equipped with RMSNorm and SwiGL.U activations
43
37
29 31).
As in 43 our training protocol excludes data repetition and spans sufficiently large
7

---

Table 3: Comparison of different methods for pre-training LLaMA models of various sizes on the C4 dataset. We report validation perplexity (4) with a memory estimate of total parameters, gradients and optimizer states. Results of all baselines are taken from 43 the full-rank dimension of models.
refers to the rank and dmodel is
60M
130M
350M
1B
Full-Rank
34.06 (0.48G)
25.08 (1.0IG)
18.80 (2.74G)
15.56 (10.40G)
Fira
31.06 (0.36G)
22.73 (0.77G)
16.85 (1.90G)
14.31 (6.98G)
GaLore
34.88 (0.36G)
25.36 (0.77G)
18.95 (1.90G)
15.64 (6.98G)
LoRA
34.99 (0.44G)
33.92 (0.99G)
25.58 (2.12G)
19.21 (7.36G)
ReLoRA
37.04 (0.44G)
29.37 (0.99G)
29.08 (2.12G)
18.33 (7.36G)
T r/dmedee
128 / 256
256 / 768
256 1024
512 / 2048
Training Tokens
1.1B
2.2B
6.4B
13.1B
dataset, encompassing a diverse array of model sizes (60M, 130M, 350M, 1B). To guarantee fair comparison, we employ the same learning rate 0.01 as used in GaLore and maintain the same rank for each model size. The detailed settings of pre-training are provided in Appendix B.I A100 80G GPUs to conduct pre-training experiments.
We use
Result analysis. As shown in Table 3 Fira consistently outperforms low-rank training baselines
by a large margin under the same rank constraint, and even surpasses full-rank training. Following 43). we estimate the memory reduction of the optimizer states via the same memory estimation
method introduced in [43]. Detailed memory estimation methods are presented in C.1 Additional
experiments on real memory usage and throughput are presented in C.2 From Table our Fira saves 61.1% % memory usage of the optimizer states when pre- vreitriin the LLaMA 1B architecture compared full-rank training, while Fira achieves better results. Compared to full-rank training, Fira's s superior performance may be attributed to the following reason: the gradient direction in the norm-based scaling method is determined by the current state, rather than by historical gradients in Adam. Therefore, Fira introduces a higher degree of randomness in training, which can enhance the model's ability to escape the local optima, leading to better training performance (44) It's worth noting that previous work 42 45 also reveals this phenomenon. Detailed analysis of this phenomenon is presented in Appendix
5.2 Scaling up to LLaMA 7B Pre-training
120
Fira Rank=64
To validate the scalability of our method, we scale up by pre-training the LLaMA 7B model with the full-rank dimension dmoddl = 4096. We compare Fira with the Galere baseline, which generally achieves the best performance among low-rank training base-
100
Galore Rank=64
Galore Rank=512
Perplexity (1)
80
lines, shown in Table 3 As illustrated in Figure our method demonstrates a significant improvement over GaLore for pre-training LLaMAA 7B, while us- ing an 8x smaller rank (memory of optimizer states). This highlights Fira's effectiveness, suggesting it could be viable solution for large-scale LLM pre- training.
40
20
2K
4K
6K
8K
10K
Training lterations
Figure 4: Pre- tretrrinnin LLaMA 7B with differ- ent methods on the C4 dataset.
5.3 Memory-Efficiene Fine-Tuning
Experimental setup. Following [14, we perform the fine-tuning task to compare Fira with LoRA, GaLore, Flora [1]. ReLoRA, Full-rank training, on the LLaMA-7B model for commonsense reasoning tasks. This task consists of eight sub-tasks, each with its own designated training and testing sets. Following the approach of [14], we combine the training datasets from all eight sub-tasks into a unified training set, while evaluating each sub-task individually using its respective testing dataset. In the fine-tuning task, the rank is set to 32 and the learning rate is set to le-4. The detailed
8

---

Table 4: Accuracy (1) of various fine-tuning methods eight commonsense reasoning datasets with
LLaMA 7B. Results for all baseline methods, except GaLore, are taken from [4
Method
Memory
BoolQ
PIQA
SIQA
HellaSwag WinoGrande ARC-e ARC-c OBQA
Avg
Fira
14.44G
69.4
82.6
78.0
76.8
81.2
82.2
64.4
80.8
76.9
GaLore
14.44G
69.5
82.0
75.1
32.2
18.0
80.7
65.8
78.0
62.7
LoRA
14.53G
68.9
80.7
77.4
78.1
78.8
77.8
61.3
74.8
74.7
ReLoRA
14.53G
68.9
81.2
77.8
46.0
79.4
80.2
64.2
79.6
72.2
Flora
14.44G
50.1
77.5
74.2
53.8
45.5
79
64.6
74.8
64.9
Full-rank
56.00G
64.2
68.1
68.0
42.3
66.5
55.6
43.9
60.0
58.6
settings of fine-tuning are provided in Appendix B.2 experiments.
We adopt RTX 4090 GPUs for fine-tuning
Result analysis. As shown Table our Fira achieves the highest performance on 5 out of datasets, demonstrating better or comparable performance compared to the baseline methods. Notably, GoLore struggles adapt to the HellaSwag and WinoGrande datasets, resulting in a significant decline in scores. In contrast, our Fira adapts to these tasks well and achieves the highest scores on WinoGrande. terms of memory efficiency, our method uses comparable or even less memory than the low-rank training methods LoRA and GaLore. These results illustrate the effectiveness of our method for the fine-tuning LLMs.
5.4 Ablation Study
In this section, conduct ablation study to assess the effectiveness of each component in method. We adopt the same settings in Section 1.. for pre-training the LLaMA 60M model. We design six variants of our method for the ablation study: (1) Fira-w.o. -isalngo-scaalnn: Fira without using the scaling factor cor- rect the gradient (i.e., setting de RR)) to a fixed Fira using the scaling factor at the matrix level instead of the column level. (3) -..miii-iii our Fira without using norm-growth limiter to avoid training loss spikes. (4) fira-graddientelipping our Fira using gradient clipping 25 to avoid loss spikes instead of our proposed norm- woom-growth limiter. (5) Fira-gradient-srinnk our Fira using gradient shrink (36) instead of our proposed norm- wrotghrowhh limiter. (6) fira-fensoo-wiise--cllngg our Fira using tensor-wise scaling [ instead of our proposed norm- corm-growth limiter.
Table Ablation study on the C4 dataset.
Method
Perplexity (L)
-.ra-l.o.-scaling Fira-matrix
37.06
31.52
Fira-w.o.- w....-limmterrr
32,22
aata-grradent-chle
31.22
pira-gradient-shrime
33.98
value of 1). (2) Fira-matrix:
Hira-tensor-wise-ee
33.81
Fira
31.06
Table JJaeseetsseeettt the results. It can be found that Fira outperforms fira-w.o--scalnng thereby demon- strating the effectiveness of our proposed norm-based scaling method for gradient correction. This also suggests that directly incorporating the raw gradient outside the subspace without correction will lead to sub-optimal results. Besides, Fira yields better performance than Fira-matrix, illustrating that more fine-grained consideration of the scaling factor is beneficial. Furthermore, Fira demonstrates improved performance over Fira-w.o.llimite Fira-gradient-clippngg Fira- sira-gradient-shriin and Fira- tensor-wise wise-scaling, indicating the effectiveness of our proposed norm- form-growth limiter addressing the issue of training loss spikes.
5.5
Performance under Varying Ranks
In this section, we illustrate the advantages of our Fira over Galore under a lower rank. We adjust various rank configurations within the set (4, 16, 64, 128) and dmodel 256, and then assess the performance pre-training the LLaMA 60M model on the C4 dataset as outlined Section 5.1 The validation perplexity of Fira and Gal ore after 10K steps across different ranks is depicted in Figure 5 From Figure we can observe that Fira consistently surpasses Gal Lore across all rank configurations. Notably, even when the ranks set very low (4 and 16), Fira still achieves performance comparable

---

Fira
Galore
Full-Rank
Perplexity (
50
(+5
4%)
40
(+22.0%)
30
128
Rank
Figure 5: Validation perplexity of Fira and GaLore for varying ranks when pre-training LLaMA 60M on the C4 dataset with dmodel = 256.
to full-rank training. In contrast, the performance of Grr significantly declines in these cases. These results highlight the superiority of our proposed Fira at lower ranks and its effectiveness in reducing memory usage.
Conclusion
In this paper, present plug-andes aud--ady memory-efficient training framework for LLMs, called Fira, as the first attempt to facilitate full-rank training consistently under low-rank constraints. First, we find notable phenomenon in LLM training: the scaling effect of adaptive optimizers the gradient norm remains similar between low-rank and full-rank training. Building on this observation, we propose a norm-based scaling method that applies the scaling effect of low-rank optimizers place of full-rank optimizers to facilitate full-rank training. This allows us to maintain the low-rank constraint within the optimizer while still benefiting from the advantages of full-rank training for improved performance. Additionally, we observe sudden spikes in gradient values during optimization, which can result in corresponding spikes in the training loss. To mitigate this, we propose norm- womm-rrowth limiter that smooths gradients by regulating the relative increase in gradient norms. Extensive experiments in both pre-training and fine-tuning of LLMS demonstrate the effectiveness of our proposed Fira.
Acknowledgments and Disclosure of Funding
This work was supported by the NSFC under Grants U2441242.
References
[1] Hervé Abdi. The kendall rank correlation coefficient. Encyclopedia of measurement and statistics, 2:508-510 2007.
[2] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXivv2303.08774, 2023.
[3] David Balduzzi, Marcus Frean, Lennox Leary, JP Lewis, Kurt Wan-Duo Ma, and Brian McWilliams. The shattered gradients problem: If resnets are the answer, then what is the question? International conference on machine learning, pages 342-350. PMLR, 2017.
[4] Han Chen, Garvesh Raskutti, and Ming Yuan. Non-convex projected gradient descent for generalized low-rank tensor regression. Journal of Machine Learning Research, 20(5):1-37, 2019.
[5] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv: 1604.06/ 74,00,1644444 2016.
[6] Yudong Chen and Martin Wainwright. Fast low-rank estimation by projected gradient descent: General statistical and algorithmic guarantees. arXiv preprint 1rXiv:1509.0000 2015.
10

---

[7] Israel Cohen, Yiteng Huang, Jingdong Chen, Jacob Benesty, Jacob Benesty, Jingdong Chen, Yiteng Huang, and Israel Cohen. Pearson correlation coefficient. Noise reduction in speech processing, pages 1-4, 2009.
[8] Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. In International Conference on Learning Representations, 2022.
[9] Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized 1lms. Advances in Neural Information Processing Systems, 36, 2024.
[10] Kaituo Feng, Changsheng Li, Xiaolu Zhang, JUN ZHOU, Ye Yuan, and Guoren Wang. Keypoint- based progressive chain-of-thought distillation for lmss. In International Conference on Machine Learning, 2024.
[11] lan Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volume 1. MIT Press, 2016.
[12] Yongchang Hao, Yanshuai Cao, and Lili Mou. Flora: Low-rank adapters are secretly gradient compressors. In Forty-first International Conference on Machine Learning, 2024.
[13] Edward Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022.
[14] Zhiqiang Hu, Lei Wang. Yihuai Lan, Wanyu Xu, Ee-Peng Lim, Lidong Bing, Xing Xu, Soujanya Poria, and Roy Ka-Wei Lee. LIm-adapters: An adapter family for parameter-efficient fine-tuning of large language models. The 2023 Conference on Empirical Methods in Natural Language Processing. 2023.
[15] Dayal Singh Kalra, John Kirchenbauer, Maissam Barkeshli, and Tom Goldstein. When can you get away with low memory adam? arXiv preprint arXiv:2503.00 01833, 2025.
[16] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.6980 2014.
[17] Teven Le Scao, Thomas Wang, Daniel Hesslow, Stas Bekman, M Saiful Bari, Stella Biderman, Hady Elsahar, Niklas Muennighoff, Jason Phang, Ofir Press, al. What language model to train if you have one million gpu hours? In Findings of the Association for Computational Linguistics: EMNLP 2022, pages 765-782, 2022.
[18] Vladislav Lialin, Sherin Muckatira, Namrata Shivagunde, and Anna Rumshisky. Relora: High- rank training through low-rank updates. In The Twelfth International Conference on Learning Representations, 2024.
[19] Liyuan Liu, Haoming Jiang. Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. In International Conference on Learning Representations, 2019.
[20] Xu-Hui Liu, Yali Du, Jun Wang, and Yang Yu. On the optimization landscape of low rank adaptation methods for large language models. In The Thirteenth International Conference on Learning Representations, 2025.
[21] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
[22] Kai Lv, Hang Yan, Qipeng Guo, Haijun Lv, and Xipeng Qiu. Adalomo: Low-memory optimiza- tion with adaptive learning rate. arXiv preprint arXiv:2310.10195. 2023.
[23] Ionut- Modoranu, Mher Safaryan, Grigory Malinovsky, Eldar Kurtié. Thomas Robert, Peter Richtárik, and Dan Alistarh. Microadam: Accurate adaptive optimization with low space overhead and provable convergence. Advances in Neural Information Processing Systems, 37:1-43, 2024.
[24] Robert Nau. Forecasting with moving averages. Fuqua School of Business, Duke University, pages 1-3, 2014.
11

---

[25] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning, pages 1310-1318. Pmlr, 2013.
[26] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1667, 2020.
[27] Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang. Dong Li, and Yuxiong He. {Zero-offload); Democratizing (billion-scale model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21), pages 551-564, 2021.
[28] Philip Sedgwick. Spearman's 's rank correlation coefficient. Bmj, 349, 2014.
[29]
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202, 2020.
[30] Chonghao Sima, Katrin Renz, Kashyap Chitta, Li Chen, Hanxue Zhang, Chengen Xie, Ping Luo, Andreas Geiger, and Hongyang Li. Drivelm: Driving with graph visual question answering. arXiv preprint arXiv: 14150, 2023.
[31] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09286 2023.
[32] Mojtaba Valipour, Mehdi Rezagholizadeh, Ivan Kobyzev, and Ali Ghodsi. Dylora: Parameter efficient tuning of pre-trained models using dynamic search-free low-rank adaptation. arXiv preprint arx.v:22/00005558 2022.
[33] Yeming Wen and Swarat Chaudhuri. Batched low-rank adaptation of foundation models. arXiv preprint arXiv:2312..0660 2023.
[34] Kirk M Wolter and Kirk M Wolter. Introduction to variance estimation, volume 53. Springer, 2007.
[35] Wenhan Xia, Chengwei Qin, and Elad Hazan. Chain of lora: Efficient fine-tuning of language models via residual learning, arXiv preprint 2010:22000..,111,,, 2024.
[36] Aohan Zeng, Xiao Liu, Zhengxiao Du, Zihan Wang, Hanyu Lai, Ming Ding, Zhuoyi Yang. Yifan Xu, Wendi Zheng, Xiao Xia, Weng Lam Tam, Zixuan Ma, Yufei Xue, Jidong Zhai, Wenguang Chen, Zhiyuan Liu, Peng Zhang, Yuxiao Dong, and Jie Tang. GLM-130b: An open bilingual pre-trained model. In The Eleventh International Conference on Learning Representations, 2023.
[37] Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
[38] Haoyang Zhang, Yirui Zhou, Yuqi Xue, Yiqi Liu, and Jian Huang. G10: Enabling an efficient unified gpu memory and storage architecture with smart tensor migrations. In Proceedings of the 56th Annual IEEE/ACM International Symposium on Microarchitecture, pages 395-410, 2023.
[39] Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. In International Conference on Learning Representations, 2020.
[40] Longteng Zhang, Lin Zhang, Shaohuai Shi, Xiaowen Chu, and Bo Li. Lora-fa: Memory--fficient low-rank adaptation for large language models fine-tuning. arXiv preprint 1rXiv:2308.03300 2023.
[41] Qingru Zhang, Minshuo Chen, Alexander Bukharin, Nikos Karampatziakis, Pengcheng He, Yu Cheng, Weizhu Chen, and Tuo Zhao. Adalora: Adaptive budget allocation for parameter- efficient fine-tuning, arXiv preprint 2rX1v:2303.10512, 2023.
12

---

[42] Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi-Qn Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more. arXiv preprint arXiv:2406. 16793, 2024.
[43] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory--fficien Ilm training by gradient low-rank projection. In International Conference Machine Learning. 2024.
[44] Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Chu Hong Hoi, al. Towards theoretically understanding why sgd generalizes better than adam in deep learning. Advances in Neural Information Processing Systems, 33:21285-2112 2020.
[45] Hanqing Zhu, Zhenyu Zhang, Wenyan Cong, Xi Liu, Sem Park, Vikas Chandra, Bo Long, David Z Pan, Zhangyang Wang, and Jinwon Lee. Apollo: Sgd-like memory, adamw-level performance. arXiv preprint arXiv:2412.05270, 2024.
[46] Philip Zmushko, Aleksandr Beznosikov, Martin Takáč, and Samuel Horváth. Frugal: Memory- efficient optimization by reducing state overhead for scalable training. arXiv preprint arXiv:2411.08837,222 2024.
13

---

A Fira Implementation
A.1 Algorithm Pseudocode
Algorithm Fira with Adam
R with S n, rank switching frequency T, hyper-parameter of Galore a, limiter threshold y = 1.01.
Input:
Step size Y/, decay rates (Bi, 82), weight matrices W
Mo. Vo E Rrxa -- 0,0 0 repeat
[Initialize moving 2nd moment and step]
G, _m V (WW,) i mod T = then
(Calculate full-rank gradients of full-rank weights}
U,E, V SVD(G,) Pt Up:r)
(Initialize the projection matrix every T steps)
else
P. 6 Pt-1
[Reuse the previous projection matrix)
end if
Re, St t-
P. Gt, D.P.T
(Divide gradients into two terms by gradient projection
Mr + M- + (1 Rd
(v):(R,): Apply Adam with low-rank gradients R₂
V. 82Vn11 +(1 (1- R2
VI-8 M.
N₂ + I-9 Vi+te
|| ( 1Jll
| 2
|
K
[Norm-Bann -aseddBBedd Scaling
R, I, 1 +e'
Rg
+
R2 8
5: < [k S. k2
:
S. n]]
|
S. S. 7/ m
|S + 7}
(Norm-Growth Limiter)
G 4 a. (P,N, + St)
(Project back and complete full-rank gradients) { UUdatt the weight matrix)
W, W 77 G. t
until convergence criteria met return Wr
A.2 Plug-and-play Framework for Fira
Algorithm 2 Plug-and-play framework for Fira, Pytorch-like.
1: for weight model.parameters() do
2: grad weight.grad
3: sub bub_grad, outer -riigrrrd---- project(grad)
(Gradient projection.
4: sub_adapt adapt(sub_graa)
(Adaptive optimizer, e.g., Adam, RMSProp.
5: outer_Fira = cera(sub_grad, sub_adapt, outer_grad,
(Apply Fira to outer_grad, (full-rank training)
6: veight_update project_ bsauett_bacsuubbbbbbeeea grad) outer_Fira
7: weight.data += weight_update
8: end for
B Details of Experiments
B.1 Detailed Pre-Training Setting
employed This section provides an overview the LLaMA architectures and the 43) T during pre-training. To ensure a fair comparison, adopt the same settings as the hyper-parameters of the LLaMA architectures across various sizes. For all architectures, we utilize maximum sequence length of 256 and batch size of 131K tokens. Furthermore, learning rate warm-up during the initial 10% of training steps and employ cosine annealing for the learning rate schedule, which decreases 10% of the initial learning rate.
implement
For all methods except Fira and Gal Lore, we tune the optimal learning rate from the set (0.01, 0.005, 0.001, 0.0005, 0.0001) across model sizes ranging from 60M to IB, selecting their best validation
14

---

Table 6: Hyper- parper-prraettes of LLaMA architectures for pre- pre-training.
Params
Hidden
Intermediate
Heads
Layers
Steps
Data Amount (Tokens)
60M
512
1376
8
10K
1.3 B
130M
768
2048
12
12
20K
2.6 B
350M
1024
2736
60K
7.8 B
B
2048
5461
24
32
100K
13.1 B
4096
11008
32
150K
19.7
perplexity to report. In contrast, both Fira and GaLore employ the same learning rate 0.01 and subspace change frequency 7 of 200 without tuning. Additionally, the scale factor a is considered fractional learning rate 43 Furthermore, relatively large learning may result in spikes of training loss 43 To address this issue, for models with a size of less than IB, we set a to 0.25, while for models exceeding 1B, we adjust to 0.0625.
B.2 Detailed Fine-tuning Setting
We fine-tune the pre-trained LLaMA-7B model for commonsense reasoning tasks benchmark designed for LLM fine-tuning, which include eight sub-tasks [14]. Table 7 shows the hyper-parameter configurations.
Table Hyper-parameter configurations of fine-tuning LLaMA-7B for Fira.
Hyper-paramette
Setting
Rank
32
a
64
Dropout
0.05
Base optimizer
Adam
LR
le-4
LR Scheduler
Linear
Batch size
16
warm-up Steps
100
Epochs
Where
Q. K. V, Up. Down
C Details of Overhead Comparison
C.1 Memory Estimates
Due to the difficulties associated with directly measuring GPU memory usage for specific com- ponent, we estimate the memory requirements for weight parameters and optimizer states across various methods and model sizes, following GaLore 43 This estimate derived from the number of parameters and optimizer states in BF16 format. In particular, for the memory of parameters, we multiply the total number of parameters by 2; for the memory optimizer states and gradients, we first calculate the total number of them according to Table Saleeland then multiply this total number by 2.
C.2 Additional Experiments on Real Overhead
We conduct additional comparisons regarding real memory usage and throughput of different memory- efficient training methods for both pre-training and fine-tuning, As illustrated in Tables 8 and Fira achieves superior memory efficiency compared to full-rank training without significantly reducing throughput. Although Fira's throughput is slightly lower than that of other memory-efficient methods, delivers exceptional performance. During pre-training, methods like LoRA necessitate maintaining additional higher-rank adapters (c.g. r/dmode! 312/2048) achieve performance comparable
15

---

to full-rank training. However, in practice, maintaining these higher-rank adapters outweighs the benefits of fewer trainable parameters compared to full-rank training. This may be due the fact that higher-rank adapters significantly increase computational demands during pre-training compared to fine-tuning (e.g. r/dmoode 32/2048), thus leading to more memory and less throughput. Furthermore, since full fine-tuning of LLaMA 7B's memory requirements exceeds the A100's 80GB capacity, we utilize DeepSeed's Zero2 technology to mitigate its memory usage.
Table 8: Real memory usage and normalized throughput when pre-training LLaMA 1B on the C4 dataset.
Method
Fira
Galore
Flora
LoRA
ReLoRA
Full-rank
Memory (GB)
54.6
54.6
54.5
59.0
59.0
58.5
Normalized Throughput (%)
94.2
95.9
95.9
67.4
67.4
100
Table 9: Real memory usage and normalized throughput when fine-tuning LLaMA 7B common- sense reasoning datasets.
Method
Fira
Galore
Flora
LoRA
ReLoRA
Full-rank
Memory (GB)
23.4
23.4
23.3
23.7
23.7 232.8
>80
Normalized Throughput (%)
156.1
201.1
210.3
232.8
100
D Additional Quantitative Analysis of Scaling Factor Similarity
In this section, we analyze the similarities of scaling factors between low-rank and full-rank training the matrix and column levels. Specifically, when analyzing similarity, we will form a vector of
[R Mll For the matrix level, we take );
values for different scaling factors (d), d2...... Φk) (Ø:
the different weight matrices of the model as items (i.e., R2(()) E RR*") for the column level, we take
the columns in the weight matrix as items (i.e., RE E RTx1). We use the same low-rank setup as in
Table 3 The models of different training runs are trained from the same random initialization. The only difference between them is the value of the rank hyperparameter (low-rank vs. full-rank).
D.1 Spearman, Kendall, and Pearson Correlation Coeflicients
In this section, we will employ Kendall's Tau correlation coefficient [ and Spearman's rank correlation coefficient (28). Pearson's 's correlation coefficient (7] to evaluate the similarities of scaling factors between low-rank and full-rank training. We conduct our experiment by including all matrices of LLaMA models ranging from 60M to IB. Then, we train these models and assess the similarity of scaling factors averaged over 10,000 steps. Additionally, to evaluate the effectiveness of column- level fine- grained strategy for scaling factors, we perform column-level quantitative similarity analysis. Due to the computational challenges posed by the large number of columns, we randomly sample 100 columns for each weight matrix for analysis. Specifically, in the LLaMA IB model, over 10,000 columns are sampled.
Both Spearman, Kendall, and Pearson correlation coefficients range from -1 to +1. coefficient of 1 signifies a perfect positive correlation, and -1 signifies perfect negative correlation. The p-value helps us determine whether the observed correlation is statistically significant or if it could have occurred by random chance. For instance, a p-value less than 0.05 means there less than 5% probability that the observed correlation happened by chance if there was actually no correlation. Generally, a p-value below 0.05 suggests that significant correlation exists. As shown in Table TO] we can observe the significant similarity of scaling factors between low-rank and full-rank LLM training (all coefficients close to 1, while p-values close to 0). Thus, likely that the observed behavior is an inherent feature of LLM training, manifesting across a broad range of scenarios. This insight provides a robust experimental basis for our proposed norm-baer based scaling in Fira and helps explain its effectiveness.
16

---

Table 10: Spearman, Kendall, and Pearson correlation coefficients (p-values) at both the matrix and column levels for pre-training LLaMA models ranging from 60M 1B parameters, averaged over 10,000 steps.
Matrix Level
Column Level
Size
Spearman
Kendall
Pearson
Spearman
Kendall
Pearson
60M
0.9972 (2e-62) 0.9662 (7e-26) 0.9891 (le-46)
0.9372 (0.0) 0.7942 (0.0) 0.8723 (0.0)
130M
0.9925 (2e-76)
0.9409 (9e-37)
0.9813 (2e-60)
0.8698 (0.0) 0.6830 (0.0) 0.7805 (0.0)
350M
0.9770 (3e-113)
0.8848 (5e-65)
0.9766 (1e-112)
0.9091 (0.0) 0.7400 (0.0) 0.8272 (0.0)
IB
0.9469 (1e-83) 0.8249 (1e-56)
0.9457 (6e-83)
0.8331 (0.0) 0.6513 (0.0) 0.8112 (0.0)
D.2
Similarity Trends.
In this section, we conduct additional experiments on similarity trends of scaling factors through Cosine Similarity and Euclidean Distance using two rank settings: r r/dmodel 16/256 and
T/ /mmdeel 128/256.
0.645
0.15
(i(R)) vs (r(G,)
41(R,) vs ((G,)
(i(R.) 4.IG, - P.R.)
(.(R)) vs A.G - P.R.)
(G(G)) dvIG, P.Rt)
dx(Gi) vs dr(G. P(Rc)
0655
0,500
Eithican daiases
0.00
bodikendom
0..0
0.020
0.010
2000
4000
5003
1003
10000
Derstiee
Figure 6: Fuclidean Distance trends training iterations (r/dmadel 128/250)
Figure 7: Euclidean Distance trends
training iterations (r/ (r/dmodet = 16/250)
1,0000
1.000
0.999
Coliee similett
Ohise iteilaby
0.9994
6097
0.9993
(:(R;) d.GG.)
drd vs dv(G))
0.96
(((R) diGG. (PR,)
Rdd vs diGG, P.Rc)
((G)) va de P.R.)
Gr) vs drlG, P(R.)
03990-
2000
4000
6000
3000
18000
2000
4000
5000
10000
Berationa
Figure 8: Cosine Similarity trends over train- ing iterations (r/ dmoded (28)).
Figure 9: Cosine Similarity trends over train- ing iterations (r /moddl 16/256).
17

---

As shown in Figure 6
8 and the similarity exhibits fluctuations during the initial training phase but achieves relatively steady pattern with high similarity the later iterations. Under lower rank setting r /odeodel 16/256, there is negligible reduction in similarity. Besides, Φi (G,))) and (. R. demonstrate significantly higher similarity owing to their closely aligned dimensions.
E
Theoretical Analysis
E.1
Error Upper Bound for the Approximation of Scaling Factors
In Section 4.1 we use the scaling factors of low-rank gradients approximate that of full-rank gradients. To quantify the effectiveness of this approximation, we will derive its error upper bound theoretically, and verify our analysis experimentally. The error of the approximation s((r) can be written as:
K ((r) = |o] - (Ra)
(13)
where rank r S n, and R, P. E R"*X_ (g1....19H) of low-rank gradients are directly sampled from components (.),,..:9p) of full-rank gradients. Under these conditions, the error can be rewritten as:
To simplify the proof, we consider that r components
Σi= (G)))) 1
Ei 1i (1:)) > of
((r)
(14)
Assumption 1. (Bounded Scaling Factors): we assume that the adaptive optimizer scales each gradient component 9s by the scaling factor that lies within known bounds. Specifically, there exist constants Cmin and Cmax such that for all i:
Cuin I
Cmaxx-
(15)
Ji
This implies:
v (g:))
S
2
(16)
min
g
Theorem 1. (Error Upper Bound for Approximation) Under the assumption that vftg)) are bounded
between constants enin and eunx for all components i, the approximation error k<(r) satisfies:
((+H1) | (1-1(())))
<(r) Aoa
(17)
where Ag² sup supg20 (g) inf Ø (²) (camax emin) defines the squared scaling factor variation, ||||| = Em P. and |,,|| = Σ/=1 9
Proof. We first define the following quantities: Total Gradient Norm 8n Σm G.. Partial Gradient Norm (first r components) S, Σ g Remaining Gradient Norm Sn-r 8n, S, Em L Total Corrected Gradient Norm Sn ΣT= (g.). Partial Adjusted Gradient Norm (first compo- nents) S, ET & (g:), Remaining Adjusted Gradient Norm Sn-p = Sn - S, = Dirr1 1/ (g.).
Then, estimation error s(() can be rewritten using these definitions:
Sn
S
<((r)
(18)
8n
8r
First, we rewrite the estimation error x(r) in terms of Sy, Sp-r, Sy, and 8n-pi
Sr + Sn-r S
(<(r) =
(19)
+ Sn-r Sy
Then, compute the difference in the numerator:
x(r) (S, + Sn-r) sy Sy 8n-r)
(20)
(8, + 8n-r)8 )8r
18

---

Simplify the numerator, thus, the estimation error becomes:
Snn-55 ST
Sp
x(r))
(21)
8r
After that, we factor out Sn-pi
<(r) 8n-r
Sn-r Sr
(22)
8n-
From our bounded assumption, we have:
cuin Sn-r S max 8n-r
min
max
and
(23)
Sr
Therefore, the maximum possible difference between
Sn-r and Se is: Sn-r Sr
Sp-r
Sy
max
max
min
Aga.
(24)
8n-r
8r
In addition, we have:
8p-r
g
Gr
8n
i==1 g
=
||
(25)
Finally, the approximation error x(r) is bounded above by:
(1 (1-
|
<(r) S A
(26)
|G||2
From this theory, we can find that the error upper bound on the approximation of scaling factors is mainly determined by two aspects, and can verify them experimentally:
Variability of Scaling Factor A This term represents the maximum variation in the scaling factors of different gradient components. For further validation, we designed Fira-only- scaling, a variant of Fira. It directly applies the low-rank scaling factors to the full-rank gradients by changing the Eqn. from W W; nP.sh, (R,, no+ (R,) P.R.) to W++1 = W, n Rs R. In this way, we are able to exclude the influence of the original Adam term Pa (Rt) and better analyze the effectiveness of our approximation. As shown in Table [T] Fira-only-scaling (column-level) gains better performance than Fira-only-scalinn (matrix-lewe for its more fine-grained consideration of the scaling factor, which also means smaller maximum variation Age.
Effectiveness of Gradient Sampling (1 (1- |[]]]] (1-1(c/)) ): This term represents the proportion of the gradients norm contributed by the sampled low-rank components from full-rank components. As shown in T22 we conducted ablation experiments Fira-only-scaling- w.o.-vvd, i.c., Fira-only-scaling without SVD in low-rank gradient sampling. As we can SVD is capable of sampling more prominent low low-rank gradients, which leads to a reduction in the upper bound of error and enhanced performance. Similarly, shown in Table [3 employing a higher rank enables the sampling of greater proportion of the gradients norm, resulting in reduced error upper bound and improved performance.
Table 11: Ablation on the level of scaling factors for the variant Fira-only-scaling.
Level
Perplexity (.)
Column
31.68
Matrix
32.05
19

---

Table 12: Ablation on SVD for the variant Fira-only-scaling.
Method
Perplexity (L)
Fira-only-scaling Firz-onlysscaling-wwooss
31.68
32.22
Table 13: Ablation on rank for the variant Fira-only-sealing and full-rank Adam (dmodel = 256).
Rank
16
64
128
256 (Full-rank)
Adam (Full-rank)
Perplexity ()) 35.91
32.90
31.93
31.68
30.84
34.06
E.2
Variance of Scaling Factors.
The variance of adaptive learning rates is significantly elevated during the early stage of training, often necessitating a warm-up to mitigate this variance and stabilize training [[9]. As illustrated in Figure 10 the scaling factor in Fira exhibits similar pattern, characterized by substantial variance during the early stage of training, which also necessitates a warm-up. However, the addition of an extra warm-up hyper-parameter for Fira would be inefficient. Therefore, it is crucial investigate whether the original warm-up would have efficiently mitigated the variance in Fira. the subsequent theoretical analysis,
show that, during the early training phase, the variance of the scaling factor of Fira is less than or equal to that of the adaptive learning rate. This finding suggests that the existing warm-up strategy is sufficient mitigate the variance of Fira, thereby eliminating the need for an additional warm-up hyper-parameter parametrr.
9t Here, the superscript indicates the index of the weight matrix to which the vector belongs, while the subscript j (where ranges from to !) denotes training iterations with each parameter. Following [19] we
Consider independent random vectors (g())) ) where each g = (g)
assume the adaptive learning rate of Adam v0(-)
1-85
and Øj
(6)
((00 for all i
(1-ft) o of
and in the early stage. Additionally, approximate the distribution of the exponential moving average as the distribution of the simple average, p(v())) = p( 12 pl
24 and
then (.) Scale-inv-A )
Theorem 2. (Variance of Scaling Factors) In the early stages of training, if (-)
(i) ~ NN00 for all i,j, then for all p > 4, the scaling factor 6² Σn En v" satisfies Var [d_] I Var we approximate v02 and v the first order, Σr (0
Scale-inv-xX2 ) and Oj
we have Var < Var[v:].
Proof. We express 6² weighted sum:
(27)
[=1
where the weights are defined as:
(0r (1) 2
Wi
En= (2)
(28)
Each wy is neg-tegative random variable satisfying Σn= wi
In the context of adaptive optimization algorithms like Adam, the squared gradients v/ accumulate information from past iterations to adapt the learning rate for each parameter. With 32 = 0.999,
'The assumption of mean-zero normal distribution is valid at the outset of training, as the weights are sampled from normal distributions with mean of zero
20

---

30000
16
Weight matrix ID
Weight matras 1D
rank -5
Weight
10
rank 10
20000
16'
10000
10"
10"
Figure 10: Scaling factor (R,) during the early stage of training (1K iterations of total 10K iterations).
Figure 11: The simulation of variance of the scaling factor Var[ø] across different rank set- tings, The adaptive learning rate vù is equiva- lent when the rank equals
the moving average of the squared gradients places significant weight on historical data, making / dependent mainly past gradients, yielding:
9
(29)
Since v/ primarily depend past gradients 9i (8) gradients gt" we can consider uiद and to be independent random variables. Consequently, express the variance of o²
.... I and w, depend solely the current
Var Var
Σ wh
(30)
Using the law of total variance, we have:
wi...ry Wnn )
Varlo7 = E
Var
w
W1
Wn
Var
E
Σ
w
(31)
1=
Since us are independent of the wi, find:
E
W Wi....y
Wn
B( Σ wi B(wF).
(32)
Var
woe Wi...., Wn
: Var [u.).
(33)
Thus, the variance simplifies to:
Var[ Varl
E
w
(34)
The second term, Var (E [> work? we]] is zero since E(d?
wy]] = E[v?]] is constant.
Let X; -gg )2 where each X, x22 Then, we can express the weights as:
X,
wi Σ Σ X
(35)
Since X1/o² ~ xi. and each w; is the ratio of X; to the sum of all Xj. the vector (wi..., w.) follows Dirichlet distribution with parameters a; z/ 2. where v; 1 is the degrees of freedom of x
21

---

For Dirichlet distribution, the expected value of w/ is given by:
ai 1)
E/w (En= ax) (Ek=1 ak + 1)
(36)
Substituting and Ek=1 Or = yields:
3
4
3
(37)
(2
+1)
n(n++ +2)
n(n + 2)
2
Thus, summing over all gives:
3
E
E[w
(38)
71 + 2
Finally, substituting this result back into the variance expression:
Var[ø = Varlu>
(39)
n 2
Since n 2 1, it follows that:
3 n + 2 S 1,
(40)
which implies:
Var Varl S Var[ub").
(41)
Given p > 4 and (-) ~ Scale-inv- X2 (p: 22) ). the variance of (-) exists [9]
Since vs? and w3 are independent and E E[ws]] = 1:
Eio² =
wath? Elwyl E[d = E[u?] Σ
El E(wyl = =[
(42)
i=1
Thus, have shown that:
Var [ S V,, and F[ = E Fl
(43)
Follow
19 we approximate / 2 and dd to the first order 34
Varlu
Var[
Var[bb]
and Var(ø] P3
(44)
E
4 4E
which implics:
Var[ø] S Var[ø)]
(45)
To further examine our theorem, we conduct simulations calculate the variance of the scaling factor á at ranks within the set {1, 5, 10, 50, 1003. The adaptive learning rate & is equivalent to that of é when the rank equals 1. As shown in Figure 11] the variance decreases as the rank increases. supporting our above theorem Var[ø] < Var[øb] Furthermore, we observe surprisingly large variance during the early stage. which corroborated our initial experiments. Consequently, conclude that our method is efficient without requiring an additional warm-up.
F Potential Reasons for Fira's Enhanced Performance Compared to Full-Rank Optimizer
While our proposed method aims to approximate full-rank training, one might intuitively assume that the performance of full-rank Adam represents the upper bound. However, previous experiments, we noticed that Fira may achieve performance that is comparable to or even better than full-rank training. This raises the question of why Fira outperforms full-rank Adam in certain scenarios. this section, first demonstrate that, despite its goal of approximating full-rank training, Fira is fundamentally
22

---

distinct from full-rank Adam. We then further substantiate this claim through experimental evidence. Finally, we explore potential reasons why Fira can achieve enhanced performance compared to full-rank Adam. Notably, the phenomenon where approximate optimizer correction outperforms the original full-rank optimizer has also been observed in prior work, such as 42 45
consistent with
our findings in Fira.
The approximation of low-rank to full-rank training in Fira is limited to scaling factors. Even in the case of full-rank training, Fira is not equivalent to the original Adam. To illustrate this, we introduce a variant of Fira called sira-only-scaling in Appendix E.1 To isolate the influence of the original Adam term Pa Fira-only- only-scaling directly applies low-rank scaling factors to the full-rank gradients by modifying Eqn. [0 from:
W4+1 Wt- (R.) - noi (R.) P.Re),
(46)
to:
W4+1 = We - (Re))
(47)
As shown in Table 13 higher-rank approximations of scaling factors lead to improved performance, validating the effectiveness of approximation. But can also find that full-rank Fira-only-scaling and Adam are not the same. Sira-only-scaling employs a weight-matrix matrix---xiieeeeeeeeeeee scaling factors Φi (Re) adjust the raw gradients, whereas Adam uses element-wise adaptive learning rates.
As a result, in contrast to full-rank Adam, which is parameter-level adaptive, Fira applies adaptive strategy at matrix-leves (column-level), while maintaining the original gradient direction within matrix (column) (i.e., di (Rt)G.). As a result, the gradient direction which is only determined by the current state can introduce higher degree of randomness in training. Then, the randomness in training can enhance the model's ability to escape the local optima, thus leading to better generalization performance. This insight potentially explains why Fira could match even surpass the full-rank Adam baseline.
To further argue this point, we conduct additional experiments the comparisons of perplexity trends. We compare the perplexity trends of Fira, Fira-only-scaling SGD, and Adam. As illustrated in 12 (a), the performance of vanilla SGD is significantly inferior, highlighting its inadequacy for directly training LLMs. As depicted in 12 (b) and (c), while Adam demonstrates faster convergence during the initial stages, both Fira and Pira-only-scalie
achieve superior performance in the later stages. This may occur because Fira applies an adaptive strategy exclusively at the matrix level while preserving the original gradient direction within each weight matrix. As previously analyzed, Fira may introduce a higher degree of randomness in training and a better ability to escape the local optima, and then enable better generalization performance.
Fira
Fira-only-scaling
15
dam
Adam
10'
102
Fin
$x 1=10
6x
167
4<10' 10
1-10
3x100
3X
1k
10k
25
41c
4%
3k
10K
45
56
R0C
10k
lienaions
Derrations
Hicnatons
(a) Fira V.S. SGD
(b) Fira V.S. Adam
(c) Fira-only-scaling v.s. Adam
Figure 12: Comparisons of perplexity (4) trends for pre-training LLaMA 60M on C4 dataset.
G Additional Analysis of Spikes
Maintaining the direction of the raw gradient without correction might be unable to effectively deal with the steep loss landscapes of LLM training like Adam [39] The steep loss landscapes are likely
to cause abrupt increases in raw gradients. When the raw gradients increase abruptly, the gradients' norm after norm- based scaling may also increase abruptly, illustrated in Figure 3 This arises from
23

---

Fira
Sia-wwo.-liitee Adam
Fira- Stra-grudient-shee -hrink
Fira-tensor--ee tete-tttoor-wseeddee sesoor-wsod-scaliiee
Training loss
Trainin loss
6.0
5.5
5.5
5.0
5,0
4.5
4.0
500
1000
1.00
2000
2500
3000
Iterations
Iterations
Figure 13:
Training loss comparison Adam, Fira, and Fira without limiter.
Figure 14: Training loss comparison of differ- ent gradient stabilization variants of Fira.
the fact that the norm-based scaling method only adjusts the average gradient norm of the gradient the matrix level, failing to make fine-grained adjustments to each parameter, unlike the optimizer Adam. As result, a significant parameter update may occur, undermining previous optimization efforts, i.c. training loss spikes [IIl 39]. As shown in Figure [3] when we directly use Adam to pre-train the LLaMA model, there will be no loss spike. Our corm-growth limiter is mainly aimed at addressing the gradient stability capability that our norm-based scaling method lacks compared to Adam.
For more comprehensive comparisons of our norm- worm-growth limiter, we design two additional gradient stabilization variants solve the loss spike: Gradient Shrink [66 ([Sa) = |S.l f + |S.-1l
- (1 μ)), P.Re is the corrected gradient by applying norm-based scaling. As shown in Figure 14 and Table [4] Fira outperforms other gradient stabilization methods. For further analysis, Gradient Shrink fails to solve the loss spike, while Tensor-Wiss Wise Scaling solves the loss spike but leads sub-optimal results.
and Tensor Wiis Scaling 8 (|S. = |Sr| μ4), where S. = dx R2 (G,
Table 14: Validation perplexity (4) of Fira across different gradient stabilization methods.
Method
Ours
Gradient Shrink
Tensor- Wise Scaling
Gradient Clipping
Without Limiter
Perplexity (L) 31.06
33.98
33.81
31.22
32.22
H
Sensitivity Analyses of
Table 15: Validation perplexity (4) of Fira across different choices of 7.
aC (w.o. limiter)
1.1
1.01
1.001
Perplexity (L)
32.22
32.09
31.06
31.26
31.28
As shown in Table 15 the performance of Fira is not sensitive to the choice of y. Provided that 7 S 1.01, Fira can effectively mitigate spikes in loss, with only marginal decrease in performance when 7 1. As we can see, the setting of 1.01, employed across all experiments in this paper, highly effective. This value neither too small to restrict the normal growth of the gradient nor too large to fail to limit sudden increases in gradient magnitude.
24

---

I Discussion on Concurrent Works
Adam-mini 42 employs block-wise learning rates to replace Adam's element-wise learning rate based the Hessian structure of neural networks. However, it only reduces memory usage for the second-order momentum while keeping the first-order momentum unchanged. In contrast, Fira reduces memory for both first-order and second-order momentum.
SlimAdam 15] uses layer-wise SNR analysis to replace second-moment tensors with their means for memory efficiency. However, like Adam-mini, it only focuses on the second-order momentum and is unable to reduce memory for both first-order and second-order momentum like Fira.
MicroAdam 23 compresses gradient information to save memory with theoretical convergence guarantees, but designed for fine-tuning workloads and cannot be adapted to LLM pre-training. However, Fira works for both LLM pre-training and fine-tuning.
AdaLomo enhances low-memory optimization (LOMO) with adaptive learning rates to improve robustness and convergence. Unlike Fira, which focuses on enabling full-rank training under low-rank constraints, AdaLomo focuses improving LOMO's convergence without involving the low-rank constraints.
FRUGAL 44) addresses gradient information loss outside the subspace in GaLore by directly adding these gradients outside the subspace via SGD or signSGD, without exploring more effective ways to utilize such gradients. In contrast, Fira explores the use of scaling factors to enable more effective updates for gradients outside the subspace. Additionally, FRUGAL provide theoretical convergence guarantees when using SGDM for subspace updates and SGD for updates outside the subspace.
analyzes and compares the loss landscapes of low-rank and full-rank setting under gradient projection of GaLore, demonstrates that low-rank constraints improve the optimization landscape by avoiding sharp minima. This may help explain Fira's superior performance. Although GaLore's ' low-rank constraints enable better optimization landscape, they discard significant gradient information outside the subspace, leading to suboptimal results. In contrast, Fira leverages low-rank constraints for better optimization landscape while effectively capturing gradient information outside the subspace, resulting in superior performance.
GaRare 20
J Additional Experimental Comparison
To further demonstrate the generality of our method, we conduct additional experiments wider range architectures (Gemma-7B, Mistral-7B), datasets (MMLUU), and adaptive optimizer (Adagrad).
Table 16: Comparison results on MMLU tasks. Results of all baselines are taken from
45
Model
Methods
STEM
Social Sciences
Humanities
Other
Average
Full
30.03
37.16
34.08
35.47
34.21
LoRA
26.23
34,94
30.88
36.96
32.18
Gemma-7B
GaLore
25.47
33.21
31.07
33.71
30.95
Fira
29.03
35.27
32.40
36.52
33.26
APOLLO
27.53
36.97
33.99
36.40
33.81
Full
52.40
72.95
55.16
69.05
61.67
LoRA
52.13
72.46
55.05
68.77
61.41
Mistral-7B
GaLore
51.87
72.82
54.94
69.49
61.56
Fira
52.80
72.85
55.07
69.11
61.72
APOLLO
51.63
73.12
54.90
69.58
61.58
Table 17: Validation perplexity ({) for pre-training LLaMA 60M on C4 dataset.
Method
Adagrad
Adagrad
Gal.ore+Adagrad
Perplexity ({)
43.11
103.85
88.87
25

---

As shown in Table 16 and 17 Fira remains competitive with all baselines across different model architectures, datasets, and opiimizers, confirming broad applicability.
To provide a comprehensive comparison against a broader range of baseline methods, we further extend our evaluation under the same experimental settings in Tables and
Table 18: Pre-training LLaMA models on the C4 dataset. Validation perplexity (J) is reported.
60M
130M
350
1B
Fira
31.06
22.73
17.03
14.31
Adam-mini
31.64
28.37
20.72
16.72
SlimAdam
31.22
24.70
18.12
15.42
FRUGAL
34.91
24.88
19.35
15.73
Table 19: Accuracy (f) on eight commonsense reasoning datasets when fine-tuning LLaMA 7B.
Method
BoolQ
PIQA
SIQA HellaSwag WinoGrande ARC-e ARC-c OBQA
Avg
Fira
69.4
82.6
78.0
76.8
81.2
82.2
64.4
80.8
76.9
Adam-mini
69.8
78.3
65.1
38.6
80.3
78.2
64.5
80.4
69.4
SlimAdam
68.3
79.2
77.2
76.9
76.5
78.2
61.1
74.1
73.9
FRUGAL
68.2
78.5
76.1
77.0
77.5
80.1
63.5
76.8
74.7
As shown in Table 18 Fira demonstrates consistent superiority over these baselines when pre- training LLaMA models of different sizes. As shown in Table 19] when fine-tuning LLaMA 7B on commonsense reasoning datasets, Fira achieves the highest performance across most datasets, exhibiting superior average accuracy compared these baselines.
K Robustness of Norm-based Scaling
As discussed in Adam-mini 42 and APOLLO [5]. Adam's 's element-wise scaling contains significant redundancy. Using shared scaling for group of parameters, such column- scaling Fira- only-scaling, is sufficient to replace Adam's 's element wise scaling. In this scenario, by aggregating historical gradient information across multiple parameters, column-wise norm-based scaling of Fira may be more robust than Adam's element-wise scaling when dealing with training data noise.
To empirically assess this robustness, inject Gaussian noise into the gradients during pre-training and measure the resulting performance degradation.
Table 20: Validation perplexity (4) for pre-training LLaMA 60M on C4 dataset.
Method
Without Noise
With Noise (0.00001 -N(00,
Fira-only-scaling
31.68
34.64
Adam
34.06
445.47
As shown in Table 20 Fira-only-scaling exhibits slight performance degradation, while Adam's ' element-wise scaling shows significant performance degradation. This indicates that column-wise scaling is more robust.
L Limitations
This work presents Fira, a novel memory-efficiett training framework that enables full-rank training of Large Language Models (LLMs)) under low-rank constraints, significantly reducing memory usage while maintaining high performance. However, our current research primarily focuses on LLMs. In future work, we aim extend the applicability of Fira to other domains, such as Multimodal Large Language Models (MI.LMs) and diffusion models, to further broaden its impact across diverse machine learning tasks and architectures.
26

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract and introduction clearly reflect the paper's primary contributions
and scope, including the development and evaluation of our proposed method.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
See Appendix
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model sell-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
27

---

Justification:
See Appendix
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The paper offers detailed information required replicate the primary experi- mental results, promoting transparency and reproducibility of the findings.
Guidelines:
The answer NA means that the paper does include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
28

---

Answer: [Yes]
Justification:
The code will be released publicly.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https guwdes/dde/eeee Cooeeeoodeesssseeeeee //nips.com htte:/ubbl///liiiee////eeeueassl ssiniiiiccy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
We detail all training and testing procedures Section 5 and the Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Owing to the high training costs, did not conduct multiple to calculate error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
29

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We detail the information on the computer resources in the Section 5
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: The research in the paper fully complies with the NeurIPS Code of Ethics, with all aspects of our work carefully reviewed to meet the established ethical standards.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification:
The paper has no negative societal impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
30

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper poses such risks, because no data models were released.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
We have appropriately credited and referenced the creators or original owners of assets, such code, data and models, used in our work. And, we have explicitly mentioned and properly respected the licenses and terms of use.
Justification:
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URI.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswithcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
31

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does not include crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does include crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
32

---

16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy buctps://neure wwwparipppsc.coooeeeeeeeee for what should or should be described.
33


Reviews:
[{"review_id": "O4HGOjvsK4", "reviewer": "Reviewer_f4Pe", "summary": "This paper introduces Fira, a novel low-rank training framework that achieves the performance of full-rank training while maintaining the memory efficiency of low-rank methods. Fira restores full-rank training quality while keeping GaLore-style low-rank memory use. Starting from GaLore’s gradient-projection scheme, Fira restores GaLore’s discarded gradient component as a new term when updating, scales it with a norm-based factor that emulates adaptive optimizer corrections without extra state, and keeps training stable via a norm-growth limiter that caps each step’s gradient-norm increase. Experiment results show Fira outperforms full-rank, GaLore, LoRA and ReLoRA on LLaMA pre-training (60 M–1 B) and beats GaLore on LLaMA-7B using 8× smaller rank; fine-tuning on eight reasoning tasks also shows consistent gains.", "questions": "1. The paper emphasizes memory efficiency but provides limited quantitative analysis on computational overhead, especially the cost of additional operations introduced by Fira. Could the authors include a detailed breakdown of computational complexity or measured runtime overhead compared to baseline methods?\n\n2. While comparisons with LoRA, GaLore, and ReLoRA are provided, newer methods involving tensor decompositions or other gradient compression techniques are not discussed. Could the authors elaborate on how Fira compares to other recent low-rank or gradient-efficient methods?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## Strengths\n1. Novel norm-based scaling recovers discarded gradients in GaLore, which improve the performance without extra optimizer state storage.\n2. Simple to implement but effective.\n3. The experiments cover various model sizes (60M to 7B), demonstrate consistent performance gains in both pre-training and fine-tuning across multiple commonsense reasoning tasks, and include thorough and fair ablation studies. And also good result analysis, especially analysis in Appendix F.\n4. Good error-bound analysis in Appendix E.\n5. well-written \n\n## Weaknesses\n1. Although the experimental design is comprehensive, the comparison with more recent training methods is limited.\n\nFor instance:\nLiu, Xu-Hui, et al. \"On the Optimization Landscape of Low Rank Adaptation Methods for Large Language Models.\" The Thirteenth International Conference on Learning Representations. 2025.\n\nHow would this method work with the gradient random projection proposed above? Further discussion and an experiment are needed to draw a convincing conclusion. \n\n\n2. No precise FLOPs since there is additional computing overhead.", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your valuable feedback! All suggested improvements will be developed and carefully incorporated into the final version.\n\n>Weakness1: Additional comparison with GaRare [1].\n\nThanks for your advice. As advised, we compare our method Fira with more recent training approaches, such as GaRare [1], and conduct experiments **integrating the gradient random projection proposed in GaRare with our Fira**. \n\nTable R1: Validation perplexity ($\\downarrow$) for pre-training LLaMA 60M on C4 dataset.\n|Fira + gradient random projection|GaRare|GaLore|Fira|\n|-|-|-|-|\n|33.67|34.45|34.88|31.06|\n\nAs shown in Table R1, Fira with gradient random projection proposed in GaRare outperforms GaRare thereby demonstrating the improvement in Fira remains effective under gradient random projection. Besides, Fira yields better performance than Fira + gradient random projection, indicating that the original SVD-based gradient projection may be more suitable than random projection under the Fira's framework.\n\nIn the final version, we will revise the related work section by adding citations and a discussion of GaRare. Additionally, we will include more experiments and analyses on it to further strengthen the study.\n\n>Weakness2 and Question1: Additional quantitative analysis on computational overhead.\n\n+ Precise FLOPs \n\nThanks for your advice. As advised, we additionally measure the precise FLOPs of GaLore and Fira.\n\nTable R2: Precise FLOPs of GaLore and Fira.\n||Fira|Galore|\n|-|-|-|\n|FLOPs|25.78G|21.47G|\n\n+ Measured runtime overhead\n\nThanks for your advice. As advised, we additionally measure the  real computational time of optmizaion step.\n\nTable R3: Optimizer step time (in seconds) across LLaMA-1B with a sequence length of 1024 on a single A100 GPU.\n|Model Size|GaLore|Fira|\n|-|-|-|\n|1B|1.440s|1.490s|\n\nAs shown in Table R2 and Table R3, compared to GaLore, Fira only slightly increases the computational overhead, which is acceptable.\n\n>Question2: Comparisons with other recent low-rank or gradient-efficient methods.\n\nThanks for your advice. As advised, we additionally elaborate on how Fira compares to other recent low-rank or gradient-efficient methods, e.g., GaRare [1] and FLoRA [2].\n\nTable R4: Pretraining LLaMA models on the C4 dataset. Validation perplexity ($\\downarrow$) is reported.\n|Model Size|60M|130M|350M|1B|\n|-|-|-|-|-|\n|Fira|31.06|22.73|17.03|14.31|\n|GaRare|34.45|25.34|19.35|15.88|\n|FLoRA|36.97|30.22|22.67|20.22|\n\nAs shown in Table R4, Fira achieves better performance compared to these methods across different model sizes.\n\n[1] On the Optimization Landscape of Low Rank Adaptation Methods for Large Language Models, ICLR'2025\\\n[2] Flora: Low-rank adapters are secretly gradient compressors, ICML'2024"}, {"review_id": "91CbrThwnx", "reviewer": "Reviewer_FBSc", "summary": "The paper proposes a new low-memory training technique for LLMs called Fira. This method reduces the memory footprint of Adam by storing optimizer states only for a low-rank approximation of the gradient, while still using the full-rank gradient at each step to improve performance. To achieve this, the authors not only use the low-rank gradient approximation with the standard Adam update (as done in low-rank methods like GaLore), but also utilize the remaining part of the full-rank gradient after applying column-wise or matrix-wise scaling to it based on the same low-rank optimizer states. The paper motivates this scaling approach and proposes a strategy to stabilize the optimization, specifically addressing the spikes in the loss curve that occur during training with this procedure. Empirical results demonstrate that Fira is effective compared to other low-memory methods in both pretraining and fine-tuning of LLMs, and it can even outperform full-rank Adam.", "questions": "All questions and concerns are detailed in the Strengths and Weaknesses section. The main weaknesses 1–3 are the most critical and will have the greatest impact on my final evaluation after the rebuttal.", "limitations": "I believe the limitations section should be more thorough and explicitly mention the limited understanding of the reasons behind the method’s benefits. Additionally, it should acknowledge that the experiments are limited to a single architecture, dataset, and adaptive optimizer.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The paper addresses an important direction: low-memory optimization is crucial from a resource-efficiency perspective, yet the effectiveness of many existing methods remains limited in settings that require substantial training, such as pretraining or fine-tuning on highly dissimilar tasks.\n2. The paper makes a reasonable effort to motivate all proposed modifications and supports them with a broad set of ablation experiments.\n3. The proposed method demonstrates strong performance in both pretraining and fine-tuning settings. It not only outperforms other low-memory methods but also achieves better results than standard Adam in most of the cases.\n4. Overall, while I have some concerns regarding the motivation and generalizability of the results, I find the empirical evidence convincing and believe the proposed method is effective.\n\nMain weaknesses:\n1. Motivation for scaling. Since scaling is a central component of the proposed method, I believe its motivation should be explained more clearly in the paper. Could you clarify how the comparison between scaling factors for the low-rank and full-rank gradients is performed? Are both computed within the same training run, and if so, what training procedure is used? My concern is that observing similar scaling factors along the trajectory of a single optimization procedure (either low-rank or full-rank) does not necessarily imply that the scaling factors are similar across the two procedures. As a result, it is not clear whether the proposed modification is expected to make low-rank training behave more like full-rank training.\n2. Fira vs full rank. Continuing the previous point, it appears that the main accuracy gains of Fira do not result from a better approximation of full-rank training, but rather from generally changing the scaling procedure from element-wise to column-wise. This change is beneficial even in the full-rank setting, as indicated by the results for Fira-only-scaling in Table 13. However, the paper does not provide a clear explanation of why this scaling modification is effective. The discussion in the appendix about escaping local optima is vague and lacks depth. Moreover, column-wise scaling is not always better than element-wise scaling. Fira performs better than Fira-only-scaling, which suggests that element-wise scaling is helpful for the main components of the gradient but may be detrimental for others. Since the choice of scaling strategy is such an important factor in Fira’s performance, it should be analyzed in more detail in the paper.\n3. Related works and baselines. From a practical perspective, my main concern is the limited discussion of other low-memory training techniques and the lack of comparison to them. For example, there is a line of research on low-memory Adam modifications, including Adam-mini (https://arxiv.org/abs/2406.16793), AdaLomo (https://arxiv.org/abs/2310.10195), MicroAdam (https://arxiv.org/abs/2405.15593), and SlimAdam (https://arxiv.org/abs/2503.01843 - this one which is a concurrent work but still relevant). These methods are clearly related to the proposed technique and therefore should be discussed in the related work section and included as baselines in the experiments. Additionally, GaLore-add, which is described as the first step of Fira’s modifications, appears very similar to FRUGAL (https://arxiv.org/abs/2411.07837), so this connection should be clearly stated. Overall, the baselines considered in the paper (LoRA, ReLoRA, GaLore) are quite basic. Including comparisons with current state-of-the-art low-memory training methods would be beneficial to better position the proposed method in context. I do not expect the proposed method to necessarily be state of the art, but such a comparison would provide valuable perspective.\n\nAdditional concerns, comments and questions:\n1. Writing. The writing of the paper could be improved. In particular, the motivation for the proposed scaling technique is not clearly described in the main text. Even after reading the appendix, I still have some questions. Additionally, the comparison between matrix and column scaling is discussed somewhat inconsistently, which could be clarified.\n2. Limited experimental results. All experiments in the paper are conducted using the same architecture (LLaMA) and include only one dataset for pretraining and one for fine-tuning. Additional experiments with different architectures and datasets would strengthen the generality of the claims. Moreover, based on the abstract and introduction, the paper appears to position the proposed modification as general for adaptive optimizers, while the experiments are performed only with Adam. The paper should either narrow its claims or provide results using other adaptive optimizers. Additionally, including full-rank experiments for LLaMA 7B pretraining would be beneficial to evaluate whether Fira consistently outperforms full-rank methods on larger-scale models.\n3. Memory and compute. I am a bit confused about the memory usage comparisons in Tables 3 and 8. Could you please explain why Fira is expected to have a significant memory improvement according to Table 3, but shows almost no real improvement in Table 8? Also, why does Fira provide a substantial decrease in throughput compared to GaLore during fine-tuning but not during pretraining?\n4. The scaling in Fira can be applied at the matrix or column level. Is the same true for the norm growth limiter, or is it always applied at the matrix level? If it is always applied at the matrix level, could you explain why?", "was_revised": true, "final_justification": "**Initial submission.** The paper presented a new low-memory training method with strong performance in both pretraining and fine-tuning settings. However, it had issues with the clarity of the motivation for the method (W1), lacked analysis of the reasons for its effectiveness (W2), and provided only comparisons to basic baselines (W3).\n\n**After discussion.** Most importantly, the authors provided a thorough comparison with stronger baselines (variations of memory-efficient Adam + FLoRa/GaRare) and additional results for other datasets, architectures, and optimizers. The motivation was also mostly clarified, and additional discussion and ablation experiments were provided to explain the method’s effectiveness (these should be further improved for the final version of the paper).\n\n**Final recommendation.** The rebuttal and discussion addressed most of my concerns. While I think the motivation and the explanation of the method’s effectiveness could still be improved, the method itself is clearly effective even when compared to strong baselines. Therefore, I recommend the paper for acceptance and increase my score to 5.", "rebuttal": "Thanks for your valuable feedback! All suggested improvements will be developed and carefully incorporated into the final version.\n\n>Weakness1: Details about motivation for scaling.\n\nSorry for confusing you. Below, we provide clarified details on the comparison between scaling factors for the low-rank and full-rank gradients.\n\n1. Are the scaling factors for the low-rank and full-rank gradients both computed within the same training run?\n\n**No, they are computed within different training runs**. As shown in Table 1 and Table 10 in our manuscript, we conducted multiple experiments under both low-rank and full-rank settings, computing and collecting the scaling factors independently for each experiment. Then, we averaged these scaling factors to assess their similarity between low-rank and full-rank settings.\n\n2. How is the similarity of scaling factors between low-rank and full-rank settings calculated?\n\nFor each training run, we compute column-wise scaling factors separately and combine them into a vector. After multiple experiments, based on two vectors computed under low-rank and full-rank settings, we can calculate their similarity using metrics like cosine similarity.  \n\n>Weakness2: Fira vs full rank.\n\nThanks for your comments. Below, we address the concerns raised from two perspectives. We will include a more detailed analysis in the final version.\n\n1. **Why does column-wise Fira-only-scaling outperform element-wise Adam?**  \n\n- **Robustness to data noise**\n\nAs noted in Adam-mini [1] and APOLLO [2], Adam's element-wise scaling contains significant redundancy. Using a shared scaling for a group of parameters, such as column-wise scaling in Fira-only-scaling, is sufficient to replace Adam's element-wise scaling. In this scenario, by aggregating historical gradient information across multiple parameters, column-wise scaling may be more robust than Adam's element-wise scaling when dealing with training data noise.\n\n- **SGD-like updates improve generalization** \n\nUnlike Adam's element-wise scaling, column-wise Fira-only-scaling preserves original gradient direction within each column, leading to SGD-like updates. In many deep learning tasks, SGD yields better generalization [3, 4]. However, this benefit rarely extends to LLMs, as SGD struggles to train them [5]. As [5] explains, this is because SGD uses a uniform learning rate across all weight blocks, overlooking their inherent differences. However, Fira-only-scaling can take these differences into account by employing adaptive column-wise scaling. This may allow Fira to extend the strong generalization of SGD-like updates to LLMs. A similar perspective is presented in Apollo [2, Section 5.5].\n\n2. **Why does Fira (column-wise + element-wise) outperform Fira-only-scaling?**  \n\n- **Low-rank element-wise scaling benefits main gradient components** \n\nPrior works [6] find that low-rank element-wise scaling may improve the optimization landscape by avoiding sharp minima. In this case, low-rank element-wise scaling may be more effective than corresponding column-wise scaling, and can be helpful for the main components of the gradient. This may explain the improved performance of combination strategy Fira.\n \n- **More precise parameter grouping** \n   \nWhile column-wise scaling reduces redundancy in element-wise Adam, it may not be optimal for all cases. For example, parameters may have diverse training dynamics within a certain column, requiring more fine-grained parameter grouping. Here, element-wise scaling may better handle the more fine-grained training dynamics thus complementing column-wise scaling.\n\n>Weakness3: Related works and baselines.\n\nThanks for your advice. As advised, we will cite and discuss Adam-mini, AdaLomo, MicroAdam, SlimAdam, and FRUGAL in related works and add them as baselines in the experiments.\n\n**Related Work:**\n\nAdam-mini, which changes element-wise Adam to block-wise based on the Hessian structure of neural networks; AdaLomo, which enhances low-memory optimization (LOMO) with adaptive learning rates for better robustness to hyperparameters and convergence; MicroAdam, which compresses gradient information to reduce memory with theoretical convergence guarantees; and SlimAdam, which uses layer-wise SNR analysis to replace second-moment tensors with their means for memory efficiency. \n\nFRUGAL operates almost same as GaLore-add: both directly add the gradients discarded in GaLore, rather than exploring scaling factors for more effective updates of these gradients (as Fira does). But FRUGAL additionally provide theoretical convergence guarantees when using SGDM for subspace updates and SGD for updates outside the subspace. \n\n**Baselines:**\n\nWe conduct additional experiments comparing Fira with Adam-mini, SlimAdam, and FRUGAL. As shown in Table R1, Fira achieves competitive performance against these low-memory training methods.\n\nTable R1: Validation perplexity ($\\downarrow$) for pre-training LLaMA 60M on C4 dataset.\n|Fira|Adam-mini|SlimAdam|FRUGAL|\n|-|-|-|-|\n|**31.06**|31.64|31.22|34.91|\n\nDue to time constraints, we will include all mentioned baselines in the final version.\n\n>Question1: Writing.\n\nSorry for confusing you.\n\n1. Motivation for the proposed scaling technique\n\nAs advised, we add more details in **response to Weakness1** to further explain the motivation for our proposed scaling technique.\n\n2. Inconsistency of the comparison between matrix and column scaling\n\nIn our manuscript, matrix scaling computes and applies scaling factors at matrix level, while column scaling operates at individual columns. In Table 5 of our ablation study, we compare Fira (column scaling) and Fira-matrix (matrix scaling), showing that the finer-grained column scaling performs better. Similarly, in Appendix E, Table 11, confirms that Fira-only-scaling (excluding the Adam term) performs better with column scaling than matrix scaling.\n\n>Question2: Limited experimental results.\n\nThanks for your advice. As advised, we conduct additional experiments on more architectures (Gemma-7B, Mistral-7B), more datasets (MMLU), and more adaptive optimizers (Adagrad) to strengthen the generality of our claims. We will add more experiments in the final version.\n\n1. Architectures and datasets\n\nTable R2. Comparison results on MMLU tasks.\n|Model|Methods|STEM|SocialSciences|Humanities|Other|Average|\n|-|-|-|-|-|-|-|\n|Gemma-7B|Full|30.03|37.16|34.08|35.47|34.21|\n||LoRA|26.23|34.94|30.88|36.96|32.18|\n||GaLore|25.47|33.21|31.07|33.71|30.95|\n||Fira|29.03|35.27|32.40|36.52|33.26|\n|Mistral-7B|Full|52.40|72.95|55.16|69.05|61.67|\n||LoRA|52.13|72.46|55.05|68.77|61.41|\n||GaLore|51.87|72.82|54.94|69.49|61.56|\n||Fira|52.80|72.85|55.07|69.11|61.72|\n\nAs shown in Table R2, Fira remains competitive with the baselines across different architectures and datasets.\n\n2. Adaptive optimizer\n\nTable R3: Validation perplexity ($\\downarrow$) for pre-training LLaMA 60M on C4 dataset.\n|Fira+Adagrad|Adagrad|GaLore+Adagrad|\n|-|-|-|\n|**43.11**|103.85|88.87|\n\nAs shown in Table R3, we conduct experiments on addtional adaptive optimizer Adagrad. Fira's strong performance demonstrates its compatibility with other adaptive optimizer.\n\n3. Including full-rank experiments for LLaMA 7B pretraining.\n\nDue to high computational cost for pretraining LLaMA 7B, we will include it in the final version.\n\n>Question3: Queations about Fira's memory improvement and throughput decrease.\n\nThanks for your comments.\n\n1.Memory improvement\n\nTable 3's memory estimation only includes **parameters, gradients, and optimizer states**. However, Table 8's actual memory estimation includes additional parts, such as **activations, memory occupied by systems (e.g., temporary computation buffer)**, etc. These extra components will reduce the magnitude of the memory improvement in Table 8. Actually, the memory reduction is basically consistent between two tables: Table 3 shows a 3.42GB reduction (10.4GB - 6.98GB), and Table 8 shows a 3.9GB reduction (58.5GB - 54.6GB).\n\n2.Throughput decrease\n\nCompared with GaLore, the decrease in Fira's throughput depends not only on time consumed by additional operations introduced by Fira, but also on **time consumed by the shared operations**. The longer time of shared operations, the smaller the decrease in Fira's throughput.\n\nCompared to fine-tuning, pre-training's shared operations dominate due to a larger batch size (16 for fine-tuning and 512 for pre-training). This explains why the throughput decrease of Fira is relatively small in pre-training.\n\n>Question4: Apply norm growth limiter at the matrix or column level.\n\nThanks for your comments. **Like scaling in Fira, norm growth limiter can also be applied at the matrix or column level.** Below are the experimental results comparing matrix-level and column-level norm growth limiters.\n\nTable R4: Validation perplexity ($\\downarrow$) for pre-training LLaMA 60M on C4 dataset.\n|Norm-growth limiter|Matrix-level|Column-level|\n|-|-|-|\n|PPL|31.06|31.09|\n\nAs shown in Table R4, applying the norm growth limiter at the matrix or column level yields nearly identical results, indicating no significant difference between the two strategies.\n\n>Limitations.\n\nAs advised, we further discuss the reasons behind Fira's benefits in **response to Weakness2**, and include additional experiments on additional architectures, datasets, and adaptive optimizer in **response to Question2**. \n\nWe will further revise the Limitations section in the final version to provide a more comprehensive discussion of Fira's constraints.\n\n[1] Adam-mini: Use Fewer Learning Rates To Gain More, ICLR'2025\\\n[2] APOLLO: SGD-like Memory, AdamW-level Performance, MLSys'2025\\\n[3] Towards Theoretically Understanding Why SGD Generalizes Better Than ADAM in Deep Learning\\\n[4] Improving Generalization Performance by Switching from Adam to SGD, NeurIPS'2020\\\n[5] Why Transformers Need Adam: A Hessian Perspective, NeurIPS'2024\\\n[6] On the Optimization Landscape of Low Rank Adaptation Methods for Large Language Models, ICLR'2025"}, {"review_id": "freYZ1gvYx", "reviewer": "Reviewer_K4iN", "summary": "The authors present Fira, a method to handle low-rank optimizer scaling for full-rank training. The authors observe that the overall magnitude of Adam's update for a layer is very similar whether it's full-rank or low-rank. So Fira replaces Adam's per element update with an update that has the same direction as the full-rank gradient, but the magnitude is derived from the low-rank Adam stats.\n\nFira seems to have two components: a norm-based scaling strategy that leverages the scaling effects of low-rank optimizers to facilitate full-rank training, and a norm-growth limiter to address the issue of loss spikes by limiting the growth of gradient norm.\n\nFira introduces the concept of a scaling factor that represents the magnitude of the correction applied by the adaptive optimizer to the gradient norm. Since at the matrix level, it exhibits a high degree of similarity between low-rank and full-training training, Fira uses it as a norm-based scaling method that leverages the scaling factor of a weight matrix in low-rank.\n\nExtensive experiments in both pre-training and fine-tuning of LLMs demonstrate the effectiveness of Fira.", "questions": "- Does it make sense to compare against system-level baselines such as ZeRO-offload in [23] or quantized optimizers?\n- Is it fair to say that the method reduces the optimizer memory but activation and parameter memory can still dominate the total memory footprint?\n- Downstream impact is somehow limited but it seems enough to demonstrate the applicability of the approach. What limitations do you forecast on vision or SFT or RLHF?\n- How similar is conceptually and its effect the norm-growth limiter to gradient-clipping?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Quality**\n\nStrengths: The methodology looks good: the norm-based scaling rule and the limiter are both derived from concrete empirical observations about Adam's per-matrix norms, and the paper backs them up with ablations showing each part is indispensable (perplexity jumps from 31.06 to 37.06 without scaling, and to 32.22 without the limiter).\n\nExperiments cover four model sizes up to 1 B parameters plus a 7 B scaling test, compare against strong baselines (LoRA, GaLore, ReLoRA, full-rank), and report both memory and accuracy metrics. The algorithm is easy to reproduce (three extra lines of PyTorch) and the authors release code.\n\nWeaknesses: Real application evaluation is limited: pre-training uses only the C4 corpus and runs 10 K steps on the 7 B model, so long-horizon stability is untested.\n\n**Clarity**\n\nStrengths: The paper is easy to follow: Figure 1 contrasts LoRA/GaLore/Fira conceptually, Algorithm 1 gives step-by-step pseudocode, and spike analysis is illustrated with paired loss-and-norm plots. Notation is consistent and an appendix supplies proofs and implementation details.\n\n**Significance**\n\nStrengths: Optimizer-state memory is a real bottleneck for 7 B models and up. Reducing it by about 61% on a 1 B LLaMA while matching or beating full-rank perplexity is practically valuable. Because Fira is optimizer-agnostic, it looks very promising to be adopted widely.\n\nWeaknesses: Memory footprint evaluation seems limited to the optimizer part.\n\n**Originality**\nStrengths: the idea that there’s some sort of norm scaling in Adam and that is preserved in low-rank and full-rank seems to be novel. This is proven empirically in Table 1 via cosine similarity.\n\nWeaknesses: Fira is heavily based on Galore’s approach.", "was_revised": true, "final_justification": "That's great. Thanks", "rebuttal": "Thanks for your valuable feedback! All suggested improvements will be developed and carefully incorporated into the final version.\n\n>Weakness1: Real application evaluation is limited: pre-training uses only the C4 corpus and runs 10 K steps on the 7 B model, so long-horizon stability is untested.\n\n+ **Datasets for pre-training**.\n\nThanks for your advice. Following the experiment setup in GaLore [1], we used the C4 dataset for pre-training. Due to time constraints, we will include pre-training experiments with additional datasets in the final version to further validate our approach.\n\n+ **LLaMA 7B pre-training**.\n\nThanks for your advice. To clarify, pre-training LLaMA 7B for 150K steps requires **significant computing resources**. For example, GaLore utilized  8-node training in parallel with a total of 64 A100 GPUs to complete the LLaMA 7B 150K-step pre-training experiments. If we use 8 A100s GPUs to complete them, it may take more than 2 months.\n\nAt the same time, most of the concurrent work, e.g., Grass [2], FRUGAL [3], were unable to complete the pre-training of LLaMA 7B for 150K steps.\n\nAs advised, we will conduct experiments with longer steps for pre-training LLaMA 7B in our future work to test the long-horizon stability.\n\n>Weakness2: Memory footprint evaluation seems limited to the optimizer part.\n\nThanks for your comments. To clarify, in Tables 3 and 4 in our manuscript, **our memory footprint evaluation is not limited to the optimizer part**, but includes memory footprint evaluation of **total parameters, gradients, and optimizer states**.\n\n>Weakness3: Fira is heavily based on Galore’s approach.\n\nThanks for your comments. While both Fira and GaLore target memory-efficient LLM training, they differ fundamentally: GaLore is restricted in low-rank gradient subspaces, thus completely failing to utilize the gradient information outside the subspace. However, Fira consistently employ full-rank gradients to train full-rank weights, achieving nearly unaffected performance despite reduced rank (PPL($\\downarrow$) 31.06→32.59) compared to GaLore's severe performance degradation (PPL($\\downarrow$) 34.88→56.57).\n\nMoreover, Fira is built upon a profound observation regarding the similarity of scaling factors between full-rank and low-rank training. It represents the first attempt of enabling consistent full-rank training under low-rank constraints. This breakthrough not only improves the memory-performance trade-off but also offers new insights for efficient LLM training.\n\n>Question1: Does it make sense to compare against system-level baselines such as ZeRO-offload in [23] or quantized optimizers?\n\nThanks for your comments. To clarify, our method Fira and system-level baselines (like ZeRO-offload [23] or quantized optimizers) achieve memory efficiency memory efficiency from distinct perspectives. Our method, Fira, achieves memory efficiency **at the algorithm level**, which is **orthogonal to these system-level techniques**. In fact, Fira can be easily combined with them to further boost memory efficiency. For this reason, direct comparisons between them may not be entirely fair. \n\n>Question2: Is it fair to say that the method reduces the optimizer memory but activation and parameter memory can still dominate the total memory footprint?\n\nThanks for your comments. To clarify, in LLM pre-training, especially when the number of parameters is very large, **optimizer memory is often more critical than activation and parameter memory**.\n\nFor example, as mentioned in the Introduction in our manuscript, pretraining a LLaMA 7B model with FP16 using Adam requires 14GB for parameters, 24GB for optimizer states (2x the parameters memory), and only 2GB for activations. In this case, optimizer memory is significantly larger than activation and parameters memory. This effect becomes even more pronounced as model size increases.\n\n>Question3: Downstream impact is somehow limited but it seems enough to demonstrate the applicability of the approach. What limitations do you forecast on vision or SFT or RLHF?\n\nThanks for your comments. As noted, we forecast the limitations when extending Fira to vision. \n\nCompared to LLMs, vision involves diverse modalities (e.g., images, videos, point clouds) and tasks (e.g., generation, detection, tracking). These differences may lead to distinct loss landscapes between LLMs and vision models. Thus, Fira may not work directly for vision. We would need to re-evaluate the phenomena and conclusions observed in LLM training, and adjust the method based on vision-specific characteristics.\n\nWe will explore this point in our future work.\n\n>Question4: How similar is conceptually and its effect the norm-growth limiter to gradient-clipping?\n\nThanks for your comments.\n\nConceptually, the norm-growth limiter and gradient-clipping both address gradient instability but differ in focus. The norm-growth limiter adaptively controls the relative increase magnitude of the gradient norm, accounting for differences across different weight matrices. In contrast, gradient-clipping clip the gradients based on their absolute norm,  setting an upper limit to avoid overly large gradients.\n\nIn terms of effect, as shown in our ablation study (Table 5 in our manuscript), the norm-growth limiter outperforms gradient-clipping.\n\n[1] GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection, ICML'2024\\\n[2] Grass: Compute Efficient Low-Memory LLM Training with Structured Sparse Gradients, EMNLP'2024\\\n[3] FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training, ICML'2025"}, {"review_id": "PYuxeeZPhR", "reviewer": "Reviewer_yQ8K", "summary": "This paper introduces **Fira**, a training framework that enables **full-rank training of LLMs under low-rank memory constraints**. Unlike LoRA and GaLore that restrict training to low-rank subspaces and degrade performance, Fira uses a **norm-based scaling method** to approximate full-rank optimizer behavior using low-rank gradients. It also introduces a **norm-growth limiter** to stabilize training. Experiments show Fira consistently outperforms existing low-rank methods in both pre-training and fine-tuning, with significantly lower memory usage.", "questions": "1. While the empirical similarity between low-rank and full-rank scaling factors is clearly demonstrated, the theoretical justification (e.g., in Appendix E) remains limited and relies on simplifying assumptions. Could the authors provide a more rigorous theoretical explanation or clarify under what conditions such similarity is expected to hold?\n\n2. The norm-based scaling update involves computing the residual term $G_t - P_t R_t$. It appears this term is used during the update step but not explicitly stored. Could the authors confirm whether this computation introduces any additional memory overhead compared to GaLore?\n\n3. In Figure 5, Fira even outperforms full-rank training in some settings. This result appears consistent with findings in [1], which suggest that low-rank constraints can improve the optimization landscape by avoiding sharp minima. Could the authors discuss whether this phenomenon helps explain Fira’s superior performance? If so, it would be helpful to explicitly cite [1] and further analyze this connection.\n\n[1] *On the Optimization Landscape of Low-Rank Adaptation Methods for Large Language Models*, ICLR 2025.", "limitations": "The paper does not explicitly discuss the limitations of the proposed method.\n\n1. The method has only been evaluated on LLaMA architectures, and its generalization to other popular LLM architectures (e.g., GPT, OPT, Mistral) remains unverified.\n2. While the paper demonstrates strong results on LLaMA-7B, larger models have not been tested. It is unclear how Fira performs at even greater scales (e.g., 13B or 65B), where memory and stability challenges may differ.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The paper’s central insight—that the scaling factors of low-rank and full-rank gradients are highly similar—is both surprising and impactful. This observation is well-supported by quantitative evidence (cosine similarity, correlation coefficients), and it forms the basis for a practical and effective norm-based scaling method that approximates adaptive optimizer behavior without storing full-rank optimizer states.\n\n2. A common but under-addressed issue in low-rank training is the occurrence of loss spikes, especially when switching projection subspaces. The proposed norm-growth limiter directly tackles this by adaptively smoothing gradient norm increases. Experiments demonstrate that this mechanism outperforms traditional gradient clipping and other heuristics, contributing to much more stable training.\n\n3. The paper conducts extensive experiments across a wide range of LLaMA models (60M–7B) in both pre-training and fine-tuning settings. Fira consistently outperforms LoRA and GaLore in performance and memory usage, including an 8× reduction in optimizer memory when scaling to LLaMA-7B.\n\n4. Fira is easy to integrate into existing training pipelines, requiring only minimal code changes.\n\nWeaknesses:\n\n1. The theoretical grounding for the norm-based scaling method relies on empirical similarity rather than formal guarantees. While the observed correlation is strong, a deeper theoretical explanation would strengthen the contribution.\n\n2. The comparison scope omits other memory-efficient optimizers, such as Adafactor or 8-bit Adam. Including these baselines could better contextualize Fira’s benefits.\n\n3. Only the LLaMA 7B model is used for large-scale validation. While the results are promising, testing on additional architectures (e.g., GPT, Mistral) would improve generality. That said, the limitation is understandable given the high computational cost.", "was_revised": true, "final_justification": "I believe this is a very good paper, though it doesn’t quite reach the level of having a groundbreaking impact. Therefore, I will maintain my decision to accept it.", "rebuttal": ">Weakness1 & Question1: Theoretical explanation of similarity between low-rank and full-rank scaling factors.\n\nThanks for your advice. As advised, we provide a more rigorous theoretical explanation of similarity between low-rank and full-rank scaling factors. In this theorem, we prove that the expected value of the ratio between full-rank and low-rank scaling factors is a constant, which demonstrates their similarity.\n\nDue to time constraints, we only present the final conclusion along with a brief proof. In the final version, we will further refine the proof details, complete the theoretical proof, and attempt to relax the conditions to derive more general conclusions.\n\n**Theorem.** (Similarity between low-rank and full-rank scaling factors). $G_t \\in \\mathbb{R}^{m \\times n}$ is the full-rank gradient where $m\\le n$, $P_t \\in \\mathbb{R}^{r \\times m}$ is the projection matrix, $R_t=P_tG_t \\in \\mathbb{R}^{r \\times n}$ is the projected low-rank gradient, where $r$ is the rank, and $\\psi$ is the gradient correction function of the optimizer (e.g., Adam). The expected ratio of full-rank scaling factor $s_g=\\frac{||\\psi(G_t)||}{||G_t||}$ to low-rank scaling factors $s_r=\\frac{||\\psi(R_t)||}{||R_t||}$ is a constant $c$:\n\n$$\\mathbb{E}[\\frac{s_g}{s_r}]=\\mathbb{E}[\\frac{||\\psi(G_t)||}{||G_t||} \\cdot \\frac{||R_t||}{||\\psi(R_t)||}]=\\mathbb{E}[\\frac{||R_t||}{||G_t||}]\\cdot\\mathbb{E}[\\frac{||\\psi(G_t)||}{||\\psi(R_t)||}] = c.$$\n\n*Proof:*\n\nAs noted in GaRare [1]、FLoRA [2], the impact of projection matrix $P_t$ (sampled randomly from a normal distribution) on the final results is minimal. For the sake of simplifying the proof, we therefore assume $P_t \\sim \\mathcal{N}(0, \\sigma^2)^{r \\times m}$.\n\nFor the first part $\\frac{||R_t||}{||G_t||}$, it is equal to:\n\n$$\\frac{||R_t||}{||G_t||}=\\sqrt{\\frac{||P_tG_t||^2}{||G_t||^2}}=\\sqrt{\\frac{||G_t^{\\top}P_t^{\\top}P_tG_t||^2}{||G_t||^2}}.$$\n\nFor $(P_t^{\\top}P_t)[i,j]$ where $i\\neq j$，since all entries of $P_t$ are mutually independent:\n\n$$\\mathbb{E}[(P_t^{\\top}P_t)[i,j]]=0.$$\n\nFor the case where $i=j$, $(P_t^{\\top}P_t)[i,i]$ follows a chi-squared distribution, with an expected value of:\n\n$$\\mathbb{E}[(P_t^{\\top}P_t)[i,i]]=\\sum_{k=1}^r P_{k,i}^2=r\\sigma^2.$$\n\nFrom this, we derive:\n$$\\mathbb{E}[P_t^{\\top}P_t]=\\sum_i \\mathbb{E}[P_t^{\\top}P_t[i,i]]=r\\sigma^2\\mathbb{I}_m.$$\n\nThus:\n$$\\mathbb{E}[\\frac{||R_t||}{||G_t||}]=\\mathbb{E}[\\sqrt{\\frac{||G_t^{\\top}P_t^{\\top}P_tG_t||^2}{||G_t||^2}}]=\\sqrt{\\frac{r^2\\sigma^4||G_t||^2}{||G_t||^2}}=r\\sigma^2.$$\n\nFor the second part $\\frac{||\\psi(G_t)||}{||\\psi(R_t)||}$, it is equal to: \n$$\\frac{||\\psi(G_t)||}{||\\psi(R_t)||}=\\frac{\\sqrt{||\\frac{M_t^G}{\\sqrt{V_t^G}}||^2}}{\\sqrt{||\\frac{M_t^R}{\\sqrt{V_t^R}}||^2}}=\\frac{\\sqrt{\\sum^m_{i=1}\\sum^n_{j=1}(\\frac{M_t^G}{\\sqrt{V_t^G}})^2[i,j]}}{\\sqrt{\\sum^r_{i=1}\\sum^n_{j=1}(\\frac{M_t^R}{\\sqrt{V_t^R}})^2[i,j]}},$$\nwhere $M_t$ and $V_t$ are the first and second moments.\n\nDue to the element-wise division between $M_t$ and $V_t$, direct decomposition becomes challenging. To address this, we adopt a conclusion from Adam-mini [3]: the adaptive learning rates associated with $V_t$ in the Adam optimizer contain significant redundancy. Specifically, the $V_t$ values within a parameter block can be replaced by a single, well-chosen $V_t$.\n\nBased on this, we assume that different variances within a matrix can be replaced by the average of all $V_t$ values:\n\n$$\\frac{||\\psi(G_t)||}{||\\psi(R_t)||}=\\sqrt{\\frac{\\frac{\\sum_{i=1}^{r} \\sum_{j=1}^{n}  V^R_t[i,j]}{rn}}{\\frac{\\sum_{i=1}^{m} \\sum_{j=1}^{n}  V^G_t[i,j]}{mn}}\\cdot \\frac{||M_t^G||^2}{||M_t^R||^2}}=\\sqrt{\\frac{m}{r}\\cdot\\frac{\\sum_{i=1}^{r} \\sum_{j=1}^{n}  V^R_t[i,j]}{\\sum_{i=1}^{m} \\sum_{j=1}^{n}  V^G_t[i,j]}\\cdot \\frac{||M_t^G||^2}{||M_t^R||^2}}.$$\n\nAssume that $beta_1$, $beta_2$ are exponential decay rates，for $\\frac{||M_t^G||}{||M_t^R||}$，we have：\n\n$$\\frac{||M_t^G||}{||M_t^R||}=||\\frac{(1-\\beta_1)\\sum_{k=0}^{t-1}\\beta_1^kG_{t-k}}{(1-\\beta_1)\\sum_{k=0}^{t-1}\\beta_1^kR_{t-k}}||.$$\n\nFor each term, result is similar to the derivation of the first part $\\mathbb{E}[\\frac{||R_t||}{||G_t||}]=r\\sigma^2$. Thus:\n\n$$\\mathbb{E}[\\frac{||M_t^G||}{||M_t^R||}]=\\frac{1}{r\\sigma^2}.$$\n\nMoreover:\n\n$$\\frac{\\sum_{i=1}^{r} \\sum_{j=1}^{n}  V^R_t[i,j]}{\\sum_{i=1}^{m} \\sum_{j=1}^{n}  V^G_t[i,j]}=\\frac{\\sum_{i=1}^{r} \\sum_{j=1}^{n} (1-\\beta_2)\\sum_{k=0}^{t-1}\\beta_2^kR_{t-k}[i,j]^2}{\\sum_{i=1}^{m} \\sum_{j=1}^{n} (1-\\beta_2)\\sum_{k=0}^{t-1}\\beta_2^kG_{t-k}[i,j]^2}=\\frac{(1-\\beta_2)\\sum_{k=0}^{t-1}\\beta_2^k||R_{t-k}||^2}{(1-\\beta_2)\\sum_{k=0}^{t-1}\\beta_2^k||G_{t-k}||^2}.$$\n\nBy the same logic, for each term:\n\n$$\\mathbb{E}[\\frac{||R_{t-k}||}{||G_{t-k}||}]=r\\sigma^2.$$\n\nFrom this, we derive:\n$$\\mathbb{E}[\\frac{\\sum_{i=1}^{r} \\sum_{j=1}^{n}  V^R_t[i,j]}{\\sum_{i=1}^{m} \\sum_{j=1}^{n}  V^G_t[i,j]}]=r^2\\sigma^4.$$\n\nThus:\n\n$$\\mathbb{E}[\\frac{||\\psi(G_t)||}{||\\psi(R_t)||}]=\\sqrt{\\frac{m}{r}\\cdot r^2\\sigma^4 \\cdot \\frac{1}{r^2\\sigma^4}}=\\sqrt{\\frac{m}{r}}.$$\n\nGiven that $\\mathbb{E}[\\frac{||R_t||}{||G_t||}]=r\\sigma^2$，we have：\n\n$$\\mathbb{E}[\\frac{s_g}{s_r}]=\\mathbb{E}[\\frac{||\\psi(G_t)||}{||G_t||} \\cdot \\frac{||R_t||}{||\\psi(R_t)||}]=\\mathbb{E}[\\frac{||R_t||}{||G_t||}]\\cdot\\mathbb{E}[\\frac{||\\psi(G_t)||}{||\\psi(R_t)||}]=r\\sigma^2\\sqrt{\\frac{m}{r}}.$$\n\nIt follows that $\\mathbb{E}[\\frac{s_g}{s_r}]$ is a constant.\n\n>Weakness2: Including other memory-efficient optimizers baselines. \n\nThanks for your advice. As advised, we add comparisons with additional memory-efficient optimizers, including Adam-mini [3], SlimAdam [4]. \n\nTable R1: Validation perplexity ($\\downarrow$) for pre-training LLaMA 60M on C4 dataset.\n|Fira|Adam-mini|SlimAdam|\n|-|-|-|\n|31.06|31.64|31.22|\n\nAs shown in Table R1, Fira outperforms these baselines. \n\nWe will include more memory-efficient optimizers baselines in the final version.\n\n>Weakness3: Testing on additional architectures.\n\nThanks for your advice. As advised, we conduct additional experiments on more architectures (Gemma-7B, Mistral-7B) to improve the generality.\n\nTable R2. Comparison results on MMLU tasks.\n|Model|Methods|STEM|SocialSciences|Humanities|Other|Average|\n|-|-|-|-|-|-|-|\n|Gemma-7B|Full|30.03|37.16|34.08|35.47|34.21|\n||LoRA|26.23|34.94|30.88|36.96|32.18|\n||GaLore|25.47|33.21|31.07|33.71|30.95|\n||Fira|29.03|35.27|32.40|36.52|33.26|\n|Mistral-7B|Full|52.40|72.95|55.16|69.05|61.67|\n||LoRA|52.13|72.46|55.05|68.77|61.41|\n||GaLore|51.87|72.82|54.94|69.49|61.56|\n||Fira|52.80|72.85|55.07|69.11|61.72|\n\nResults in Table R2 show that Fira remains competitive compared to baselines across different architectures.\n\nDue to time constraints, further experiments (e.g., pretraining experiments on more architectures) will be included in the final version.\n\n>Question2: Memory overhead of the residual term $G_t - P_tR_t$. \n\nThanks for your comments. The residual term $G_t - P_tR_t$ is computed and used within the same update step and does not need to be stored for other steps. Thus, it introduces no additional memory overhead compared to GaLore.\n\n>Question3: Use findings in [1] to explain Fira’s superior performance.\n\nThanks for your advice. We agree that the **findings in [1] help explain Fira’s superior performance** over full-rank training in some settings, as shown in Figure 5. Specifically, [1] analyzes and compares the loss landscapes of low-rank and full-rank setting，and then demonstrates that low-rank constraints improve the optimization landscape by avoiding sharp minima, which aligns with Fira’s results. \n\nAlthough GaLore's low-rank constraints enable a better optimization landscape, they discard significant gradient information outside the subspace, leading to suboptimal results. In contrast, **Fira leverages low-rank constraints for a better optimization landscape** while **effectively capturing gradient information outside the subspace**, resulting in superior performance.\n\nWe will explicitly cite [1] and **further discuss with this work to clarify its impact on Fira's results in the final version**.\n\n>Limitations:\n\n1. The method has only been evaluated on LLaMA architectures, and its generalization to other popular LLM architectures (e.g., GPT, OPT, Mistral) remains unverified.\n\nSee response to Weakness3.\n\n2. While the paper demonstrates strong results on LLaMA-7B, larger models have not been tested. It is unclear how Fira performs at even greater scales (e.g., 13B or 65B), where memory and stability challenges may differ.\n\nThanks for your comments. Testing Fira's performance at even greater scales (e.g., 13B or 65B) requires significant computational resources. Due to time constraints, we will conduct experiments at greater scales to test our Fira in the final version.\n\n[1] On the Optimization Landscape of Low Rank Adaptation Methods for Large Language Models, ICLR'2025\\\n[2] Flora: Low-rank adapters are secretly gradient compressors, ICML'2024\\\n[3] Adam-mini: Use Fewer Learning Rates To Gain More, ICLR'2025\\\n[4] When Can You Get Away with Low Memory Adam?"}]

Author rebuttal:

Meta-review: The authors tackle the important problem of reducing the memory usage required for training or fine-tuning large language models (LLMs). They propose Fira, a plug-and-play framework for LLM training that maintains low-rank memory efficiency while enabling full-rank training through norm-based scaling. The core idea is to leverage the scaling effects of low-rank optimizers as substitutes for traditional full-rank optimizers. 
 
Experimental results demonstrate that Fira outperforms both LoRA and GaLore in pre-training and fine-tuning, achieving superior performance with lower memory costs. Experiments include challenging reasoning tasks and a comprehensive ablation study. The paper is well-written, presents a novel idea, and offers improvements in a highly challenging and important area of machine learning.
 
All reviewers recognized the significance of the paper and voted for its acceptance. During the rebuttal phase, the authors addressed the reviewers' questions, added comparative baselines, and clarified the theoretical similarities between the scaling factors of low-rank and full-rank methods. Overall, I believe this is an excellent paper that should be published in NeurIPS. I therefore recommend its acceptance.
Decision: Accept (poster)