Breaking the Frozen Subspace: Importance Sampling for Low-Rank Optimization in LLM Pretraining
Authors: Haochen Zhang
Junze Yin
Guanchu Wang
Zirui Liu
Lin Yang
Tianyi Zhang
Anshumali Shrivastava
Vladimir Braverman
Venue: NeurIPS 2025 poster ()
Keywords: low-rank optimization
large language models
Adam
memory efficiency
Primary area: deep_learning
TL;DR: We introduce a new method for selecting subspaces in low-rank optimization for memory-efficient pretraining of large language models (LLMs).

Abstract: Low-rank optimization has emerged as a promising approach to enabling memory-efficient training of large language models (LLMs). Existing low-rank optimization methods typically project gradients onto a low-rank subspace, reducing the memory cost of storing optimizer states. A key challenge in these methods is selecting suitable subspaces to ensure an effective optimization trajectory. Most existing approaches select the dominant subspace to preserve gradient information, as this intuitively provides the best approximation. However, we find that in practice, the dominant subspace stops changing during pretraining, thereby constraining weight updates to similar subspaces. In this paper, we propose importance sampling for low-rank optimization in LLM pretraining with a provable convergence guarantee, which the dominant subspace approach does not have. Empirically, we demonstrate that our method significantly outperforms previous methods in LLM pretraining tasks.

Breaking the Frozen Subspace: Importance Sampling for Low-Rank Optimization in LLM Pretraining
Haochen Zhang', Junze Yin', Guanchu Wang', Zirui Liu', Lin Yang', Tianyi Zhang', Anshumali Shrivastava', Vladimir Braverman
Rice University, University of North Carolina Charlotte, University of Minnesota Twin Cities, UUiiversity of California, Los Angeles, JJohns Hopkins University
Abstract
Low-rank optimization has emerged as a promising approach to enabling memory- efficient training of large language models (LLMs). Existing low-rank optimization methods typically project gradients onto low-rank subspace, reducing the memory cost of storing optimizer states. A key challenge in these methods is selecting suitable subspaces ensure effective optimization trajectory. Most existing approaches select the dominant subspace to preserve gradient information, as this intuitively provides the best approximation. However, find that in practice, the dominant subspace stops changing during pretraining, thereby constraining weight updates to similar subspaces. In this paper, we propose importance sampling for low-rank optimization in LLM pretraining with a provable convergence guarantee, which the dominant subspace approach does have. Empirically, we demonstrate that our method significantly outperforms previous methods in LLM pretraining tasks.
1 Introduction
Large language models (LLMs), pretrained on next-token prediction tasks, achieve human-level text generation capabilities and exhibit zero-shot transferability to various downstream tasks [5] They are also fine-tuned or aligned with human preferences to be expert in downstream tasks [66] [53] Over the past few years, there has been rapid progress in LLM development, characterized by consistent growth in the number of trainable parameters and the scale datasets BI [3] [] The parameter count in
language models has increased from 100 million [56] to over hundred billion [10]. However, despite their enhanced expressiveness, such large models demand extensive GPU memory for pretraining
[22] Thus, critical question arises:
How can we improve the memory efficiency of LLM pretraining?
In LLM pretraining, Adam is commonly used as the optimizer due its superior optimization performance. However, a key limitation of Adam is its memory requirement, as it necessitates storing optimizer states, each consuming as much memory as the model itself. This poses significant challenge, given the substantial memory demands of the model's parameters. To address this issue, researchers have explored low-rank optimization, where gradients are projected onto a low-rank subspace to reduce the memory consumption of optimizer states. These states are then projected back to their original size when updating the weights. For example, GaLore [79] and Q-GaLore
[7] project gradients onto subspaces defined by the leading singular vectors corresponding to the largest singular values, a technique referred to as the dominant subspace. FLora (3 and GoLore
[4]. on the other hand, utilize unbiased random low-rank projections for gradients, employing the Johnson-Lindee Lindenstrauss transform. Grass 51 introduces sparse low-rank projections, which further
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

reduce the gradient memory footprint as well as the computation and communication costs compared dense low-rank projections. Lastly, Fira 77 builds on GaLore by fully leveraging the error in gradient low-rank approximation to achieve improved performance.
These methods are powerful because: 1) the gradients of LLMs during pretraining exhibit an intrinsic low-rank structure, making them well-suited for compression using low-rank ap- proximation, and 2) low-rank approximation can be applied not only to Adam but also to other optimizers that use state information. For in-
Selecting Dominate Subspace
Importance Sampling Subspacee Selection
Adjacent Subspace Overlap
0.8
0.6
0.4
stance, Adafactor [58] employs rank- I factoriza- tion on the second moment in Adam to reduce the memory required for storing the second mo-
0.0
ment. Adam-mini [6] eliminates over 99% of the effective learning rate in the second moment of Adam while achieving performance on par with even better than-Adam Additionally, [1] and [35 propose low- low-precision optimizers with 8-bit and 4-bit optimizer states. Low-rank optimization integrates seamlessly with these Adam variants, further highlighting its impor- tance and underscoring why it deserves signif- icant attention.
10
20 30
40
50
Layer Indices
Figure 1: Adjacent subspace overlap of low- rank optimizer using difference subspace selec- tion methods. Our importance sampling subspace selection can lower the overlap between adjacent subspaces, thus enabling better exploration of more different subspaces in the optimization trajectory.
central question in low-rank optimization is how to maintain the performance of pretrained LLMs while using memory-efficient optimizers, as compared to full-rank optimization. One common paradigm in existing low-rank optimization methods to update weights within the dominant sub- space for certain number of iterations and periodically update this dominant subspace. Nonetheless, the dominant subspaces of gradients in many layers stabilize almost completely after the early stages of pretraining [77]. Consequently, the weight updates during different periods predominantly remain within the same low-rank subspace, resulting in cumulative weight updates that struggle to achieve high rank. This limitation significantly hampers the language modeling capabilities of pretrained LLMs. Thus, it is natural to ask:
Is it possible to overcome the low-rank bottleneck of existing low-rank optimization methods with minimal additional overhead?
In this paper, we provide positive answer this question. We propose a novel method for subspace selection low-rank optimization by introducing an appropriate degree of randomness in the selection process. In summary, the contributions of this study are as follows:
We observe that highly similar adjacent subspaces in existing low-rank optimization methods diminish the diversity of weight updates, degrading the performance of pretrained L.LMs.
To address frozen dominant subspace phenomenon and the low-rank bottleneck of update in existing low-rank optimization methods, we propose an Importance SAmpling method for Low-RAnk optimization (SARA). This method enables low-rank optimizers to explore a broader range of subspaces in the optimization trajectory. Specifically, the low-rank subspace is spanned by singular vectors sampled from m singular vectors for gradient mm Figure illustrates how SARA reduces the overlap between adjacent subspaces during LLM pretraining.
SARA can be integrated with various low-rank optimization methods, such as GaLore and Fira. It is robust to recond-momen factorization and low-precision optimizer state storage. On pretraining tasks for the LLaMA model at different sizes, SARA consistently outperforms dominant subspace selection and reduces the performance gap between low-rank optimizers and full-rank Adam by up to 46.05%.
From theoretical aspect, We prove that SARA achieves comparable convergence rate as GoLore (4 (Theorem [344] proof details are deferred to Appendix A) whereas delivering better empirical results (Section 4 and Appendix F)

---

Roadmap. In Section 2 we present the update rules of GaLore- Adam and Fira-Adam. In Section 3 we describe our methodology for using importance sampling improve the algorithmic design of low rank optimizers- waweeerss-al....
-Ga oore and Fira-ang provide the convergence guarantee of SARA. In Section 4 we present experimental results showing that SARA consistently outperforms dominant subspace selection. In Section 5]
we discuss related work. Finally, in Section 6 we conclude the paper.
Preliminaries
In this section, we present the background required for our theoretical analysis and experiments. In our experiments (Section 4) we apply SARA to two low-rank optimization methods, GaLore and Fira, both of which can be combined with stateful optimizers (e.g., Adam, Adafactor, and Adam-mini).
To ensure clarity, the update rules for Gal -dre-Amaam and Fira-Adam are briefly explained here. For more detailed explanations, please refer to the original papers (79 7 In presenting these methods, we show the update rules for the weights of single layer in the neural network. We assume that the gradient at the f-th iteration is matrix (t)) E Rm*x Without loss of generality, we assume that m and use represent the rank of the low-rank subspace.
Update Rules of GaLore-Adam Galmre-Adam 79 requires storing an orthogonal matrix p(() Rmxr that satisfies (P()) (()) Lr, which is updated periodically. Similar full-rank Adam. GaLoreAAdam also stores the first moment M(2) Re*= and the second moment V(t) E Rxxx for each layer's weights, and updates the weights W(t) as follows: R( = (P(1)) ((t) M =
31 M + (1 ())R()) V (2 + (1 B2) R()) R(), No aP()) arlt)
vVu++" and x00) = (t-1) 77 (0), Here, 31 and 32 are hyperparameters for the online updates of M and V(), as in Adam. The parameter 77 denotes the learning rate, and & is a small positive constant for numerical stability.
Update Rules Fira-Adan Similar to Adam, Fira also needs to store M V and Plo, The difference is that Fira-Adam additionally utilizes the low-rank approximation residual to update wr(O) Let 5()) (I (())))())) and (+) ((t-1) 1) 7 NN 1 ((S())) where (()) represents the low-rank approximation error, ((-) represents a scaling function in Fira (7]. and
N( p(t) M is calculated same GaLore-Adane above.
VV
3
Methodology
In this section, first illustrate the adverse effects of frozen dominant subspace in mini-batch gradients (Section B.I). To address this issue, we then propose SARA for low-rank optimization (Section 3.2 Finally, present convergence analysis of low-rank optimization using SARA (Section 3.3
3.1 Frozen Dominant Subspace of Mini-batch Gradient
77 observes that the cosine similarity between adjacent dominant subspaces approaches 1.0 in some layers after certain stage of L.LM pretraining. indicating that the dominant subspace of the gradient almost stops evolving. We observe similar phenomenon in our experiment as well. Figure 2 shows the average result of dominant subspace overlap different layers across all blocks different itera- tions. We notice that dominant subspace overlaps are low in all layers at the carly stage of pretraining, but they increase drastically as pretraining progresses, eventually becoming stable at different levels. Among all layers, gate_proj overlap between adjacent subspaces is harmful for low-rank optimization. Considering an extreme case, when the overlap reaches 1.0, the low-rank optimizer can only change the weights within fixed low-rank subspace. However, when the low-rank subspace shifts significantly over time, the overall weight update-formed by summing updates from various low-rank subspaces- can overcome the constraints of the low-rank bottleneck. For readability, we refer to this phenomenon as the frozen dominant subspace.
_proj and jp_proj exhibit the highest subspace overlaps. Intuitively, a high
3

---

Algorithm Low-rank Optimization with SARA
1: Input: The I-th layer weight 2i Rm for all / E [N]. Learning rate scale factor decay rates $1.52, rank r, subspace change frequency 2,, small constant for numerical stability &.
2: Initialize: for all / [N]] 1, .M
R
3: for t do
4:
forl=1-1 N do
Compute the mini-batch gradient: G( E mm
5:
p( - SARA(G ( .7)
6:
P see Algorithm input parameters see Section
8 (V((-1)) M M(-11) at P(o) of B2,,
7:
(1) t GALORE- (AAMMS)ADAM(S) or FIRM-S
8:
9:
end for
10:
end for
Return ,(T) (a] 12
(T)
3.2 SARA: Importance SAmpling for Low-RAnk Optimization
Algorithm SARA: Importance sampling subspace selection for low-rank optimization
Input: The mini-batch gradient at the iteration t, G(t) E Rmx*,, where [N]] denotes the layer. Subspace change frequency E Z..
2: if I. mod
then
3: U1") .1 Vi <- (VD(G(())
<(t)
I (AMPLE([m] num num = r, weight = )(())
4:
5:
I (TRT(I)
p(() + U 1:I I
6:
7: else
p(t) +- p(t-1)
8:
D. Reuse the previous projector
9: end if
10: Return P!O)
To overcome the problem of the frozen dominant subspace problem, we propose SARA to construct low- rank subspace. Low-rank optimiza- tion with SARA is given Algo- that SARA
q_proj
V_praj
___o
up_proj
k_proj
proj
down pro
Dominate subspace Overiap
0.7
rithm [ It can be does not change the overall structure of the original low-rank optimization algorithm but is plug-and-pla sub- stitute for dominant subspace selec- tion. Algorithm 2 gives the procedure of SARA. 4ine/4dennter the weighted sampling without replacement. More precisely, each of the m left singu- lar vectors is equipped with a weight w; E (0, 1) proportional sponding singular value S..
0.6
0.5
0.4
10000
20000
30000
Update Steps
corre-
Figure 2: The average mini-batch gradient dominant sub- space overlap different linear layers over 8 blocks in the LLaMA-60M model during pretraining. We measure the overlap between adjacent subspaces every 200 iterations.
S
Wi
Σラ=1
T/
S
For an index set sample = (I1.... ....,,,)). Ir), the sampling probability can be written as
P( ... I,) (h. ..4)
II k=1
Wix
1 Wi, Wig-1

---

Line | sorts the sampled indices in ascending order so that the newly updated subspace basis vectors can align with optimizer states well. Line construutss the orthogonal basis of the new subspace.
By using weighted sampling without replacement, we make adjacent subspaces more distinct and prevent the optimization trajectory from being trapped in overly similar subspaces during training. Another advantage of SARA is that it introduces negligible additional overhead example, computing an SVD on a 2048 x 2048 matrix takes 0.34 seconds, while sampling adds only 0.0005 seconds on average.
3.3 Provable Convergence of SARA
24 points out that selecting the dominant subspace in low-rank optimization, as done in Goreore, does not always guarantee convergence to the optimal solution. Although GoLore ensures provable convergence, it does not significantly close the performance gap between GaLore-Adan and full-rank Adam in pretraining tasks, as reported in [24]. In this section, we show that SARA achieves provable convergence, representing a key advantage over GaLore and comparable to GoLore. The empirical results of SARA are shown in Section and Appendix F representing a key advantage over GaLore and GoLore.
We treat an LLM as a neural network with layers, and each layer has a weight matrix, i.e., for all I E [N]. E R=xx Without loss of generality, we assume that S nn. In practice most L.LMs do not have biases for attention blocks and MI.P blocks, and low-rank optimization is only applied to the weight matrix, but not to biases. Therefore, this abstraction is reasonable. Mathematically, our objective function is f Rmxm x Rm X .... X Ruxnn + R. For all (#1,.... Iiy, .NN),y (y1...., ., "....... ................................ YN) E dom(f), we denote (if(x) and (if as E mRx,, respectively. Below, we use the following two assumptions similar
af 3e
E R and of Tyt
to
24
Assumption 3.1 (L-smoothness). Let f Ruxm x Rmxna x ... x Ruxnn + IR be our objective function. Let T. >> 0. For all 1 E [N.. we let (I) (.i,..- .......... ,i,,.... ................................ ,NNN, y (y),---: -..--..... UN
dom(f) be any arbitrary N-tuples satisfying that if E [N] \ (1), then H, = yi- We assume f is L-smooth that satisfies:
Vif(2) - Vif(') L -|
Assumption 3.2 (Bounded, Centered, and Independent Mini-batcl Gradient Noise). Let (z) E Rmxm: be the gradient of our objective function for the 1-th layer at the 1-th iteration, where E Z,. Let al E RRxn be the mini-batch gradient which is the noisy version of V=f() For all I [N], we assume there exists least upper bound of E R for || (1 namely
|a -Vi(z(()))) /|= S of
and
E
[ = ))
Furthermore, we define o² = Σi of.
To analyze the convergence of SARA, we characterize the error introduced by projecting gradients onto the sampled low-rank subspaces. Specifically, we bound the discrepancy between the original gradient and its projection under the importance sampling scheme of SARA. This projection error plays a central role in the convergence analysis, as it quantifies how well the sampled subspace preserves gradient information. Because of the page limit, we defer its proof to Appendix A
Lemma 3.3 (Error of SARA's Projection, see Lemma for proof). Let T be the update period of SARA, and r be the rank of low-rank subspace in SARA. For all i E [m], [N], k E N. let p, (0) (i) denote the probability that the i-th basis vector is selected for the 1-th layer at time
kt, and
define 5/0) == min;g|m)
[m] Pi (t), 5 == minje[m] 5/") Let p(() E TRoor denote the orthonormal
(t)
projection matrix and let V (t)( mmxm be the gradient matrix of the 1-th layer at time t. Then,
the following inequality holds:
E [|(1 (1- Vif(a(0))1 (1 6)- 1E
/ F
5

---

With the projection error bounded in Lemma 3.3 we are now equipped to analyze the convergence behavior of SARA-base low-rank optimization. The theoretical result below demonstrates that SARA achieves provable convergence at rate comparable to prior work, while the experimental results in the following section show improved empirical performance.
Theorem 3.4 (Convergence complexity of Low-rank MSGD with SARA, see Corollary A.6 for proof).
By Assumption 3. 3.2
i 2 2 + 128/ (36) + (128a)² (11)))//(////)))) for A = f(x((0) inf (x), we
choose 31
1
-
3631
64
and ll
[AL
SOL2 309
80 30
16r
3/1
low-rank MSGD-SARA with momentum re-projection converges as
T EEB E [[I
T-1
LA 82.0T
LA
=
83.5T
Comparing with [24), we adopt the same hyperparameters used in their study in Theorem 3.4 When examining the convergence rate, we that the primary distinction lies in our convergence rate depends on 5 (Theorem 3.4,, whereas GoLore depends on & = (Theorem 3,5).
Theorem 3.5 (Convergence of MSGD with GoLore, Corollary 3 of (24)). Under Assumption 3.1 3.2
let every notation be defined as in Theorem 3.4 and using the same hyperparameterss 31, T, T, Let Then, AISGD-GoLore with momentum re-projection converges
5
m
1
Σ E
T-1
LA
LAG2 83.T
+
T
625T
Because SARA adopts importance sampling, we have < 5. Thus, the convergence rate of MSGD- SARA is slower than MSGD-GoLore up to constant factor. Compared to MSGD-GaLore (using dominant subspace), which does not have a provable convergence guarantee, SARA has the advantage in the theoretical convergence rate.
4 Experimental Results
we describe our experimental setup. In Section SARA when combined with various low-rank Adam optimizers. In Section 4.3 we show that SARA
In Section 4.1
4.2
we evaluate the efficacy of
promotes subspace exploration and enables higher-rank updates. Scction 4.4 we further evaluate
SARA additional baselines and datasets.
4.1 Experiment Setting
In this section, describe our experimental setup. We present the C4 dataset [57] architecture, and hyperparameters.
Pre-training on C4 Dataset. C4 [57] short for Colossal Clean Crawled Corpus, is large-scale, open-source text dataset widely used practice for pretraining transformer models such as BERT
(55). T5 688, and GPT models. C4 is also commonly used in the memory-efficien optimization community to evaluate the performance of memory-efficient optimizers
[23 79 77 24). our
experiments, pretrain LLaMA models of different sizes on the C4 dataset without data repetition, using a sufficiently large amount of data 225
Architecture and Hyperparameters We evaluate the performance of different optimizers on LLaMA models with 60 million, 130 million. 350 million, and 1.1 billion parameters, using the same architecture as in (79]. For full-rank Adam, we use 31 0.9, ,32 = 0.999, and a learning rate of 0.001, except for the LLaMA -0M model, where the learning rate is set to 0.0025. More detailed hyperparameters for our re- i--mlmmmmnttioonn are provided in Appendix B All experiments are conducted using one GPU node with Nvidia A40 GPUs.
4.2 Efficacy of SARA with Low-Rank Adam Optimizers
In this section, we evaluate the efficacy of SARA with different low-rank Adam optimizers across multiple model sizes.

---

Table 1: Validation perplexity (PPL) LLaMA models pretrained on the C4 dataset with 60M, 130M, and 350M parameters, comparing various low-rank optimizers with and without SARA. SARA consistently reduces the PPI. gap relative to full-rank Adam, demonstrating its effectiveness across different optimizer variants and model scales.
60M
130M
350M
Full-Rank Adam
27.71
23.27
18.21
MiLLor-SARA-Adee
30.47
24.21
19.16
GaLore-Adan
31.50
24.88
19.68
PPL gap reduction
27.17%
41.61 %
35.37%
Fira-SARA-Adam
28.12
22.22
17.25
Fira-Adam
28.42
22.37
17.35
PPL gap reduction
42.25%
GaLore- SAw.SAAARAAAaaaae
30.06
24.09
18.88
GaLore-Adaaacto
31.13
24.79
19.45
PPL gap reduction
31.28%
46.05%
45.96%
GaLore-SARA- La-ARA-Ada--man Adam-mini
31.66
24.87
19.41
-almore-Admm mini
32.08
25.46
19.89
PPL. gap reduction
9.61%
26.94%
28.57%
(8bit)
30.55
24.67
18.16
GaLore-Adam (8bit)
31.62
25.35
18.63
PPL gap reduction
27.36%
32.69%
T/ dmddel
128/256
256/768
256/1024
Tokens
1.5B
2.2B
6B
Efficacy of SARA with different low-rank Adam optimizers First, evaluate the efficacy of SARA when combined with various low-rank Adam optimizers. Table [ shows that SARA consistently outperforms dominant subspace selection. In cases where full-rank Adam achieves the lowest PPL, we also report the percentage reduction in the PPL gap achieved by SARA compared the dominant subspace baseline. As shown in Table T SARA reduces the PPL gap by up 46.05%. %. scenarios where full-rank Adam does not achieve the lowest PPL., SARA still improves performance leading singular vector selection. SARA proves effective not only with low-rank Adam variants such as GaLore-Adane and Fira-Adam, but also with optimizers that approximate second moments, e.g., Adaffccoo and GaLLoreAAdam-me Results with the 8-bit optimizer further highlight the robustness of SARA under low-precision optimizer state storage.
Scale Up to Llama-1.1B We also evaluate the efficacy of SARA on the pretraining of LLaMA-I.IB. Due to limited computational resources, we conduct experiments using only GaLore-Adam, As shown in Table SARA remains effective on LLaMA-1.1B.
Table PPL on ILaMA-1.IB pretrained with full-rank Adam, GaLore-Adam, and GaLore-SARA- Adam on the C4 dataset. Despite the larger model size, SARA continues to outperform dominant subspace selection, confirming scalability and robustness.
Full
GaLore GaLLore-SRA--- SARA-Adam GaLore-Adam
I.IB
15.90
15.36
15.47
r dmadel Tokens
512/2048
512/2048
512/2048
13.4B
13.4B
13.4B
4.3 SARA Encourages Subspace Exploration and Higher-Rank Updates
In this section, we empirically show that SARA encourages subspace exploration and enables higher-rank updates.
7

---

modellayers. wwwp4.pppppppp.p..mm _ppppr
loydel.layers...moo wttp:r.pmmmpuppppp.
Adjacent Subspace Overtap
0.7
Anchor Subspace Overlap
0.7
0.6
0.6
Ga Adam GaLore-Ade Adam
GaLore-SARAAAdam
0.5
0.5
GaLore-Adam
0.4
0.4
0.3
0.3
2500
3000 Update Steps
3500
4000
2500
3000 Update Steps
3500
4000
Figure 3: a). The left figure shows the overlap between adjacent subspaces in GaLore-Adam and GaLore-SARA-Adaee during pretraining on the LLaMA-60M model between 2200-th and 4000-th iteration. b). The right figure takes the low-rank subspace the 2000-th iteration as the anchor subspace, and shows the overlap between subspaces in later iterations and the anchor subspace.
SARA encourages subspace exploration
provides interesting observation that the similar- ity between adjacent subspaces in some layers gradually becomes very high during pretraining, we We adopt a metric to measure overlap between
observe a similar phenomenon shown in Figure 2 two subspaces from [21]. Given two orthonormal matrices U, V E Ruxr, we have
= -VV-VV
the overlap between subspaces spanned by U and V are defined
overlap (U, V) Σ V..lé.
where Vi, denotes the i-th column of V. We adopt the above metric to show that the observation in (77 is not because of using cosine similarity as the measure, but the frozen subspace phenomenon also exists when using other metrics to measure subspace overlap (or subspace similarity).
An interesting fact is that the overlap between adjacent subspaces in MaLorr-SARA-Adaa is shown in B (a). In Figure 3 (b), we chose the sub- space the 2000-th iteration as the anchor sub- space and examined the overlap between sub- spaces from later iterations, specifically between the 2200-th and 4000-th iterations. We observe that the overlap between the anchor subspace is
Average result across all layers
much lower than GaLore-Adam,
Normalized Singular Values
1.0
Gal.ore-SARAAAde
Figure
-aLarr-Addmm
Full-rank Adam
0.6
0.4
0.2
and later subspaces of woa.-oo--RR--AAA lower than that of GaLore-Adam. This indi- cates that SARA encourages the optimization trajectory to explore more different subspaces compared to using the dominant subspace.
100
200
300
400
500
Singular Value Indices
Figure The average result of normalized singular values of the weight difference between the 28k- step checkpoint and 30k-step checkpoint across all layers in LLaMA- -60 during pretraining.
SARA Enables Higher-rank Update Fig- ure 4 shows that the update produced by SARA the weight matrix exhibits more evenly dis- tributed singular values compared to the update
using the dominant subspace. This suggests that SARA helps overcome the low-rank bottleneck asso- ciated with the dominant subspace approach. We observe both this higher-rank update and improved subspace exploration occurring simultaneously. We believe these two phenomena are correlated. One possible explanation is that better exploration of diverse subspaces leads higher-rank update.

---

4.4 Additional Baselines and Datasets
In this section, we present additional baselines (GoLore 224 and online PCA 422 and dataset
(SlimPajama).
More Baselines for Pretraining on C4 To provide a more comprehensive evaluation of SARA, conduct extensive pretraining experiments Beyond the comparisons shown in our main results, we include two particularly relevant baselines: GoLore 24 and the online PCA approach from 42 These baselines were selected because they also competitive alternatives to our method. The results, presented in Table reveal several important insights. First, we observe that Gol Ard- performs reasonably well, achieving validation perplexity of 1661 and 24.01 for the 60M and 1 130M parameter models respectively. However, SARA consistently outperforms GoLore by significant margins (1.14 and 0.93 perplexity points for the two model sizes), demonstrating the effectiveness of our approach in learning more efficient low-rank representations. The comparison with 42's ss online PCA method is particularly illuminating. Their approach, while computationally efficient, shows substantially higher perplexity (33.69 and 30.62) compared to both GoLore and SARA. This is mainly due to the unstable training loss curve during the online PCA method pretraining.
the C4 dataset with additional baseline methods.
Table 3: Validation perplexity comparison on the C4 dataset for LLaMA models with 60M and 130M parameters using additional baselines.
C4 Dataset (validation perplexity)
60M (1,5B tokens)
130M (3B tokens)
GoLore-Adan
31.61
24.01
42
with Adam
33.69
30.62
MaLLor-SARA-Adee
30.47
23.08
Full rank Adam
27.71
22.19
Pretraining Results on More Datasets To demonstrate the generalizability of SARA beyond the C4 dataset, we conduct additional pretraining experiments on the SlimPajama dataset, carefully filtered and deduplicated subset of the Pile corpus. The results Table 4 confirm that our method's advantages not dataset- -ataset-speciffc. Several interesting patterns emerge from the SlimPajama results. First, the performance gaps between methods are slightly smaller on SlimPajama compared to C4. For the 130M parameter model, SARA improves upon standard GaLore-Adam by 0.57 perplexity points (25.23 vs 25.80), while trailing full-rank Adam by only 0.23 points. This suggests that our method may be particularly effective higher-quality, deduplicated datasets like SlimPajama. We also note that the absolute perplexity values are slightly lower on SlimPajama than on C4 for comparable model sizes, which may reflect the dataset's more careful curation. This makes the strong performance of SARA even more noteworthy, as it demonstrates effectiveness across different difficulty levels.
Table 4: Validation perplexity the SlimPajama dataset for LLaMA models with 60M and 130M parameters.
SlimPajama (validation perplexity)
60M
130M
Full rank Adam
27.79
25.00
GaLore-Adam
31.76
25.80
GaLore-SARA-Adam
30.79
25.23
Additional experiments on high-rank updates, anchor similarity, and adjacent overlap can found in Appendix F.I
F22
and F3 respectively.
5 Related Work
Memory Efficient Parametrization. LoRA 26 can be seen as memory-efficien parametrization of weights L.MMs and widely used in fine-tuning, LoRA's bottleneck lies its low-rank structure, which impedes its expressiveness. COLA (67, DeRA [8I), and PLoRA [47] propose to

---

increase the rank and improve the performance of LoRA. ReLoRA 40 and SLTrain [22 extend LoRA to pre-training tasks by merging and resetting adapters, and adopting low-rank plus sparse parameterization, respectively. MoRA 22 alleviates the shortcoming of the low-rank disadvantage of LoRA by sharing the same trainable parameters to achieve a higher-rank update. Additionally, 45 analyzes the sparsity-based parameter-efficient fine-tuning (SPEFT) for L.Mss, which is an alternative method of LoRA. [54] designs a novel method for memory-efficiet fine-tuning for LLMs. It has been shown that it can outperform LoRA and full-parameter training in many cases. Similarly, 28 proposes another novel fine-tuning method called Half Fine-Tuning (HFT), which can mitigate "catastrophic forgetting" in LLMs during sequential training and instruction tuning. Finally, [36] also proposes a novel memory-efficieee fine-tuning method by strategically selecting layers to update based on outlier statistics. Our paper also considers the memory and convergence behavior, but the difference is that we mainly focus on the LLM pretraining instead of fine-tuning.
Memory Efficient Optimizer. One way to achieve memory-efficien optimization is by using memory-efficien optimizers, which primarily aim reduce the memory cost of optimizer states in Adam [33]. A series of works [58] 76 46 80) factorize the second moment in Adam. Quantizing optimizer states and storing them in low-precision formats has also proven successful [3 12].
Another line of work focuses on gradient compression methods. GaLore [79] and Q-GaLore 77 use
SVD to apply dense low-rank projections gradients. FLora 23 and GoLore [24] adopt random projection, while Grass [1] employs sparse low-rank projection to gradients.
Subspace Learning. Existing studies provide sophisticated analyses of various subspace learning algorithms [1] 34 29). 21 claims that gradient descent primarily occurs in the dominant subspace, which is spanned by the top eigenvectors of the Hessian. In contrast, [60] argues that, due noise in SGD, the alignment between the gradient and the dominant subspace is spurious, and learning does occur in the dominant subspace but rather in its orthogonal complement, i.e., the bulk subspace. Intuitively, our findings align with those of (60), suggesting that selecting basis vectors based on specific sampling probabilities can enhance the performance of I LLMs during pre-training.
6 Conclusion
In this paper, we propose SARA for low-rank optimization in LLM pretraining. The motivation is to find an effective subspace selection method to overcome the low-rank bottleneck caused by the frozen dominant subspace in low-rank optimization. SARA samples singular vectors of mini-batch gradients with probabilities proportional to their singular values, this enables optimization trajectory to explore more different subspaces. Theoretically, in Theorem 3.4 we show that GaLore GaLore-SARA.
MSGD achieves comparison convergence rate as GoLore-MSGD, which is O
LA
LAg2
J
d
Empirically, we find that SARA improves the language modeling capability pretrained models compared to using the dominant subspace, as verified by experiments involving SARA and dominant subspace selection with multiple low-rank optimizers.
Acknowledgment
Lin F. Yang is supported in part by NSF Grant 2221871 and an Amazon Faculty Award. Vladimir Braverman supported in part by NSF Grant CNS 2528780.
References
[1] Marah Abdin, Jyoti Aneja, Hany Awadalla, Ahmed Awadallah, Ammar Ahmad Awan, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Jianmin Bao, Harkirat Behl, et al. Phi-3 technical report: highly capable language model locally on your phone. arXiv preprint arXiv:2404. 14214, 2024.
[2] Nikhil Abhyankar, Vivek Gupta, Dan Roth, and Chandan K. Reddy. H-star: Lim-driven hybrid sql-text adaptive reasoning tables, 2025.
10

---

[3] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, 1lge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08877, 2023.
[4] Maryam Aliakbarpour, Vladimir Braverman, Junze Yin, and Haochen Zhang. Support basis: Fast attention beyond bounded entries. arXiv preprint arXiv:2510.01643, 2025.
[5] Tom B Brown. Language models are few-shot learners. arXiv preprint arXiv:2005. 2020.
[6] Defu Cao, Furong Jia, Sercan O Arik, Tomas Pfister, Yixiang Zheng, Wen Ye, and Yan Liu. Tempo: Prompt-based generative pre-trained transformer for time series forecasting, arXiv preprint MXXiv:2310.04900 2023.
[7] Xi Chen, Kaituo Feng, Changsheng Li, Xunhao Lai, Xiangyu Yue, Ye Yuan, and Guoren Wang. Fira: Can we achieve full-rank training of Ilms under low-rank constraint? arXiv preprint arXiv:241000162,, 2024.
[8] Yifang Chen, Jiayan Huo, Xiaoyu Li, Yingyu Liang, Zhenmei Shi, and Zhao Song. Fast gradient computation for rope attention in almost linear time. arXiv preprins arXiv:2412.113366 2024.
[9] Yifang Chen, Xiaoyu Li, Yingyu Liang, Zhenmei Shi, and Zhao Song. The computational limits of state-space models and mamba via the lens of circuit complexity. arXiv preprint arXiv:2412.06148, 2024.
[10] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-
113, 2023.
[11] Romain Cosson, Ali Jadbabaie, Anuran Makur, Amirhossein Reisizadeh, and Devavrat Shah. Low-rank gradient descent. IEEE Open Journal of Control Systems, 2023.
[12] Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. arXiv preprint arXiv:2 10.02861, 2021.
[13] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783. 2024.
[14] Logan Engstrom, Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Firdaus Janoos, Larry Rudolph, and Aleksander Madry. Implementation matters in deep rl: A case study on ppo and trpo. In International conference on learning representations, 2019.
[15] Yeqi Gao, Zhao Song, Weixin Wang, and Junze Yin. A fast optimization view: Reformulat- ing single layer attention in LLM based tensor and SVM trick, and solving it in matrix multiplication time. In The 41st Conference on Uncertainty in Artificial Intelligence, 2025.
[16] Yeqi Gao. Zhao Song, and Shenghao Xie. In-context learning for attention scheme: from single softmax regression to multiple softmax regression via a tensor trick. arXiv preprint arXiv:2307.002419, 2023.
[17] Yeqi Gao, Zhao Song, and Junze Yin. Gradientcoin: peer-to-peer decentralized large language models. arXiv preprint arXiv:2308. 10502, 2023.
[18] Yeqi Gao, Zhao Song, and Junze Yin. An iterative algorithm for rescaled hyperbolic functions regression. In International Conference on Artificial Intelligence and Statistics, 2025.
[19] Yuzhou Gu. Zhao Song, and Junze Yin. Binary hypothesis testing for softmax models and leverage score models. In Forty-second International Conference on Machine Learning, 2025.
[20] Yuzhou Gu, Zhao Song, Junze Yin, and Lichen Zhang. Low rank matrix completion via robust alternating minimization in nearly linear time. In The Twelfth International Conference on Learning Representations, 2024.

---

[21] Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in tiny subspace. arXiv preprint arXiv: 1812.04754, 2018.
[22] Andi Han, Jiaxiang Li, Wei Huang, Mingyi Hong, Akiko Takeda, Pratik Jawanpuria, and Bamdev Mishra. Sltrain: a sparse plus low-rank approach for parameter and memory efficient pretraining. arXiv preprint arXiv:2406.022 022214,,,, 2024.
[23] Yongchang Hao, Yanshuai Cao, and Lili Mou. Flora: Low-rank adapters are secretly gradient compressors. arXiv preprint arXiv:2402.00 03093,, 2024.
[24] Yutong He, Pengrui Li, Yipeng Hu, Chuyan Chen, and Kun Yuan. Subspace optimization for large language models with convergence guarantees. arXiv preprint 1.XXi:22410.112888 2024.
[25] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, al. An empirical analysis of compute-optimal large language model training. Advances in Neural Information Processing Systems, 35:30016-30030, 2022.
[26] Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint wxxiv:2206.00966 2021.
[27] Xiang Huang, Sitao Cheng, Shanshan Huang, Jiayu Shen, Yong Xu, Chaoyun Zhang, and Yuzhong Qu Queryagent: reliable and efficient reasoning framework with environmental feedback-based self-correction, 2024.
[28] Tingfeng Hui, Zhenyu Zhang, Shuohuan Wang. Weiran Xu, Yu Sun, and Hua Wu. HFT: Half fine-tuning for large language models. In Proceedings of the 63rd Annual Meeting of the Association for Compulational Linguistics (Volume 1: Long Pupers), 2025.
[29] Ali Jadbabaie, Anuran Makur, and Amirhossein Reisizadeh. Adaptive low-rank gradient descent. In 2023 62nd IEEE Conference on Decision and Control (CDC), pages -310-3.0. IEEE, 2023.
[30] Deyi Ji, Lanyun Zhu, Siqi Gao, Peng Xu, Hongtao Lu, Jieping Ye, and Feng Zhao. Tree-of-table: Unleashing the power of IIms for enhanced large-scale table understanding, 2024.
[31] Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint 1rX.v:23100068825 2023.
[32] Ting Jiang, Shaohan Huang, Shengyue Luo, Zihan Zhang, Haizhen Huang, Furu Wei, Weiwei Deng, Feng Sun, Qi Zhang. Deqing Wang, et al. Mora: High-rank updating for parameter- efficient fine-tuning. arXiv preprint arXiv:2405. 12130, 2024.
[33] Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint
141 2014.
[34] David Kozak, Stephen Becker, Alireza Doostan, and Luis Tenorio. Stochastic subspace descent. arXiv preprint
145, 2019.
[35] Bingrui Li, Jianfei Chen, and Jun Zhu. Memory efficient optimizers with 4-bit states. Advances in Neural Information Processing Systems, 36, 2024.
[36] Pengxiang Li, Lu Yin, Xiaowei Gao, and Shiwei Liu. Outlier-weighed layerwise sampling for IIm fine-tuning. In Findings of the Association for Computational Linguistics: ACL 2025, pages 19460-19473, 2025.
[37] Yunfan Li. Yiran Wang, Yu Cheng, and Lin Yang. Low-switching policy gradient with explo- ration via online sensitivity sampling. International Conference on Machine Learning, pages 19995-20034. PMLR, 2023.
[38] Yunfan Li and Lin Yang. On the model-misspeciiication in reinforcement learning. In In- ternational Conference on Artificial Intelligence and Statistics, pages 2764-2772, PMLR, 2024.
12

---

[39] Zhihang Li, Zhao Song, Zifan Wang, and Junze Yin. Local convergence of approximate newton method for two layer nonlinear regression. arXiv preprint arXiv:2311.55390, 2023.
[40] 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, 2023.
[41] Jiehao Liang, Zhao Song, Zhaozhuo Xu, Junze Yin, and Danyang Zhuo Dynamic maintenance of kernel density estimation data structure: from practice to theory. In Proceedings of the Forty-First Conference on Uncertainty in Artificial Intelligence, 2025.
[42] Kaizhao Liang, Bo Liu, Lizhang Chen, and qiang liu. Memory-efficient L.M training with online subspace descent. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[43] Yingyu Liang. Heshan Liu, Zhenmei Shi, Zhao Song, and Junze Yin. Conv-basis: A new paradigm for efficient attention inference and gradient computation in transformers. arXiv preprint arXiv:2405.05211, 2024.
[44] Junyan Liu, Yunfan Li, Ruosong Wang, and Lin Yang. Uniform last-iterate guarantee for bandits and reinforcement learning. In The eighth Annual Conference on Neural Information Processing Systems, 2024.
[45] Xinxin Liu, Aaron Thomas, Cheng Zhang, Jianyi Cheng, Yiren Zhao, and Xitong Gao. Refining salience-aware sparse fine-tuning strategies for language models. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), 2025.
[46] Yang Luo, Xiaozhe Ren, Zangwei Zheng, Zhuo Jiang, Xin Jiang, and Yang You. Came: Confidence- guided adaptive memory efficient optimization. arXiv preprint WxXiv2237.000047 2023.
[47] Xiangdi Meng, Damai Dai, Weiyao Luo, Zhe Yang, Shaoxiang Wu, Xiaochen Wang, Peiyi Wang, Qingxiu Dong, Liang Chen, and Zhifang Sui. Periodiclora: Breaking the low-rank bottleneck in lora optimization. arXiv preprint arXiv:2402.16141 2024.
[48] Amirreza Neshaei Moghaddam and Bahman Gharesifard. Gradient flow approximations in temporal difference learning. IEEE Control Systems Letters, 2024.
[49] Amirreza Neshaei Moghaddam, Alex Olshevsky, and Bahman Gharesifard. Sample complexity of linear quadratic regulator without initial stability. arXiv preprint arXiv:2502. 14210, 2025.
[50] Amirreza Neshaei Moghaddam, Alex Olshevsky, and Bahman Gharesifard. Sample complexity of the linear quadratic regulator: A reinforcement learning lens. Journal of Machine Learning Research, 26(151): 1-50, 2025.
[51] Aashiq Muhamed, Oscar Li, David Woodruff, Mona Diab, and Virginia Smith. Grass: Com- pute efficient low-memory 11m training with structured sparse gradients. arXiv preprint arXiv:2406. 1.....17700 17660, 2024.
[52] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, al. Efficient large-scale language model training on gpu clusters using megatron-Im. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-15, 2021.
[53] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida. Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances neural information processing systems, 15-27730-27744, 2022.
[54] Rui Pan, Xiang Liu, Shizhe Diao, Renjie Pi, Jipeng Zhang, Chi Han, and Tong Zhang. Lisa: Lay- erwise importance sampling for memory-efficient large language model fine-tuning. Advances in Neural Information Processing Systems, 37:57018-57044 2024.
13

---

[55] Jacob Portes, Alexander Trott, Sam Havens, Daniel King, Abhinav Venigalla, Moin Nadeem, Nikhil Sardana, Daya Khudia, and Jonathan Frankle. Mosaicbert: A bidirectional encoder optimized for fast pretraining. Advances in Neural Information Processing Systems, 36:3106 106- 3130, 2023.
[56] Alec Radford, Karthik Narasimhan, Tim Salimans, llya Sutskever, et al. Improving language understanding by generative pre- vre-training. OpenAI, 2018.
[57] 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):166) 2020.
[58] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596-4604. PMLR, 2018.
[59] Ritwik Sinha, Zhao Song, and Tianyi Zhou. A mathematical abstraction for balancing the trade- off between creativity and reality in large language models. arXiv preprint arXiv:2306.0222 2023.
[60] Minhak Song, Kwangjun Ahn, and Chulhee Yun. Does sgd really happen in tiny subspaces? arXiv preprint arXiv:2405. 16002, 2024.
[61] Zhao Song, Chongxi Wang, Guangyi Xu, and Junze Yin. The expressibility of polynomial based attention scheme. In Proceedings of the 31st ACM SIGKDD Conference Knowledge Discovery and Data Mining V.2, 2025.
[62] Zhao Song, Weixin Wang, and Junze Yin. A unified scheme of resnet and softmax. arXiv preprint TrXiv:2309.13488 2023.
[63] Zhao Song, Mingquan Ye, Junze Yin, and Lichen Zhang. Efficient alternating minimization with applications to weighted low rank approximation. In The Thirteenth International Conference on Learning Representations, 2025.
[64] Zhao Song, Junze Yin, and Lichee Zhang. Solving attention kernel regression problem via pre-conditioner. In International Conference on Artificial Intelligence and Statistics, pages 408-216. PMLR, 2024.
[65] Songyuan Sui, Hongyi Liu, Serena Liu, Li Li, Soo-Hyun Choi, Rui Chen, and Xia Hu. Chain- of-query: Unleashing the power of Ilms in sql-aided table understanding via multi-agent collaboration, 2025.
[66] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov. Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint 1i0Xi:::377009288,. 2023.
[67] Wenhan Xia, Chengwei Qin, and Elad Hazan. Chain of lora: Efficient fine-tuning of language models via residual learning, arXiv preprint arXiv:2401.0m 2400:24001.0111,, 2024.
[68] I. Xue. mt5: A massively multilingual pre-trained text-to-text transformer. arXiv preprint arXiv:2010. / 1934, 2020.
[69] Wen Ye, Jinbo Liu, Defu Cao. Wei Yang, and Yan Liu. When Ilm meets time series: Can Ilms perform multi-step time series reasoning and inference. arXiv preprint arXiv:2509.01822, 2025.
[70] Anton Zakharenkov and Ilya Makarov. Deep reinforcement learning with dqn vs. ppo in vizdoom. In 2021 IEEE 21st international symposium on computational intelligence and informatics (CINTI), pages 000131 1011110--336...1 IEEE, 2021.
[71] Haochen Zhang, Xi Chen, and Lin F Yang. Adaptive liquidity provision in uniswap v3 with deep reinforcement learning. arXiv preprint arXiv:2309. 10129, 2023.
[72] Haochen Zhang, Zhiyun Peng, Junjie Tang, Ming Dong, Ke Wang, and Wenyuan Li. A multi-layer extreme learning machine refined by sparrow search algorithm and weighted mean filter for short-term multi-step wind speed forecasting. Sustainable Energy Technologies and Assessments, 50: 101698, 2022.
14

---

[73] Haochen Zhang, Tianyi Zhang, Junze Yin, Oren Gal. Anshumali Shrivastava, and Vladimir Braverman. CoVE: Compressed vocabulary expansion makes better LLM-based recommender systems. In Findings of the Association for Computational Linguisties: ACI. 2025, 2025.
[74] Junzi Zhang, Jongho Kim, Brendan O' O'Donoghue, and Stephen Boyd. Sample efficient reinforce- ment learning with reinforce. Proceedings of the AAAI conference on artificial intelligence, 2021.
[75] Mumin Zhang, Yuzhi Wang, Haochen Zhang, Zhiyun Peng, and Junjie Tang. A novel and robust wind speed prediction method based on spatial features of wind farm cluster. Mathematics, 11(3):499, 2023.
[76] Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more. arXiv preprint arXivv:4066. 16793, 2024.
[77] Zhenyu Zhang. Ajay Jaiswal, Lu Yin, Shiwei Liu, Jiawei Zhao, Yuandong Tian, and Zhangyang Wang, Q-galore: Quantized galore with int4 projection and layer-adaptive low-rank gradients. arXiv preprint arXiv:2407.08296 2024.
[78] Zhi Zhang, Chris Chow, Yasi Zhang. Yanchao Sun, Haochen Zhang, Eric Hanchen Jiang, Han Liu, Furong Huang, Yuchen Cui, and Osear Hernan Madrid Padilla. Statistical guarantees for lifelong reinforcement learning using pac-bayesian theory. arXiv preprint 1xXx:24411000000 2024.
[79] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficieet Ilm training by gradient low-rank projection. arXiv preprint arXiv:2403.03507, 2024.
[80] Pengxiang Zhao, Ping Li, Yingjie Gu, Yi Zheng, Stephan Ludger Kölker, Zhefeng Wang, and Xiaoming Yuan. Adapprox: Adaptive approximation in adam optimization via randomized low-rank matrices. arXiv preprint arXiv:2403. 240..2458.. 2024.
[81] Bojia Zi, Xianbiao Qi, Lingzhi Wang, Jianan Wang, Kam-Fai Wong, and Lei Zhang. Delta- lora: Fine-tuning high-rank parameters with the delta of low-rank matrices. arXiv preprint (tti::2309.02411,, 2023.
15

---

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: We have explicitly presented our main claims in our abstract and introduction (see Section
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 are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have detailedly discussed the limitations of our paper in Section D
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 well-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]
16

---

Justification:
The assumptions we have can be found at Assumption 3.1 and Assumption 3.2
The proof of Theorem 3.4 can be found in Appendix A found in Appendix A
The proof of Lemma 3.3 can be
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 in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or 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 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: We have provided the pseudocode for our algorithms-SAke (Algorithm 2 as well as details the LLaMA models, datasets, and evaluation metrics. The hyper- parameters is included in Appendix B The GitHub Repo will be open to public after the review session.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a 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 a specific model and empirical evaluation, it may be necessary to either make it 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 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, other means that appropriate to the research performed.
also be provided via detailed
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 new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.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 or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe 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 (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
17

---

reproducible, they 5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [No]
Justification: The datasets we used are public and can be accessed by everyone. Also, we will release code after the review session.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for 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
/nips. wwwpccplccgg.uiieooooooococpcppe / sesiissinnn ssiniioiscoyycy for more details.
The authors should provide instructions on data access and preparation, including how 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 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 in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide the training and test details, hyperparameters, experiment settings, etc (see Section [4.]]and and Appendix [B. Once the review session is done, will provide the code of our experiments.
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:
We do not have 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 emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
18

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run 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).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It 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 show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain 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 provide this information in Section 4
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 as 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 or 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 https: tttttps:///euriiii hww.npeu...../. httpsscccpuubllp cc/puulii/Bthhicchiiddee
Answer: [Yes]
Justification: All authors have reviewed the NeurIPS Code of Ethics and confirm that the research presented in this paper fully complies with
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
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 (c.g.. if there 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: [Yes]
Justification:
We have discussed the positive and negative societal impact in Section E
Guidelines:
The answer NA means that there is no societal impact of the work performed.
19

---

If the authors answer NA or No, they should explain why their work has no societal impact or 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.
The conference expects that many papers will be foundational research and not tied to 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 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 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 are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to 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:
Our paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 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 the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We cited the previous papers that produced the code package or dataset.
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 URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
20

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
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 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:
We do 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/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing 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 details about compensation (if any)?
Answer: [NA]
Justification:
We do not involve 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 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 to the NeurIPS Code of Ethics, 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 institution) were obtained?
Answer: [NA]
Justification:
We involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending 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.
21

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
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, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [Yes]
Justification:
We only use LLM to check our grammar.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy bottts://neurre http://arrissccccoomeeerrrrr//// cc//oonferencee//2225/LL. for what should or should be described.
22

---

Appendix
A Proof Details for SARA
Lemma A.1 (Gradient connections, Lemma 2 in (4)). Let (z)(z) denote the gradient with respect to the t-th layer parameters at point I, and suppose x0)) denotes the parameters at iteration t. Then for any integers t 2 0 and > O, we have
Then, for all L, O, we have
2
(e]((((t)) (r- Σ Ve((x(tttt)) -
Ves(x((tr)) -ee())())
Lemma A.2 (Error of SARA's Projection). Let T be the update period of SARA, and be the rank of low-rank subspace in SARA. For all i E [mm] E [N], k E N, let pi the i-th basis vector is selected for the 1-th layer time Kt, and define
(t) (i) denote the probability that
min Pi (0) EE
(t), (i),
8= min te(m),420
Let p(") E &Rxxr denote the orthonormal projection matrix and let V vz((ttt) E Ruxx be the gradient matrix of the 1-th layer at time I. Then, the following inequality holds:
E P\( VIS
(1-
Proof. We analyze the expected projection residual:
||(: P(() (P())) Vis (((()())) 1
E
Eppe ||( [ p| (p(0)) Vis(ar
/ [
Evis(ato)
E ir (Bpeo
p((0) (P!)))) (
Let { be fixed orthonormal basis for IR. and define the indicator variable 1 (1) denote whether U, is selected. Then,
P(00) 1 (P)(0(
Σa (1 ELDDU,,
Therefore,
ELOGNU,,U M(((011) ]
E [1[1 p(0
Vis(E(0))
= E
tr
Σ 1=1
(1
E
(1))
E Vif(o()))))
<(1 min E
E [if(ot>))
(1 o(t) E Vif(
<1 E
Lemma A.3 (Momentum Contraction of SARA). For SARA with momentum re-projection, for all I [N], have the inequalities below hold.
23

---

Part 1 (t = 0).
E
Vi5(E ||
- (1 (23)
| [1
-8)00
+ f9o8.
F
Part 2 (1 = KT, E No.
E
(z)()!())) [1]]]
1
1
8
Br
E
-
4
2(1 - 6)
Σ EE
_-1
Vis(a(
(
5(1
B1)
Vif(x((0) Vis(o((-1))
1 rr
E
318
x-2
Vef(ollrtrtt) -vi((alhtt) ||21 p. B..0..
+ -
- 1)(1 5) Σ
E
Part 3 (t = KT + r, S 1).
E
M
(t)
vif(2 (1) ||
1 1 1 5 81) E
6)
| (1
2
[
6
B E
Vis
[1=]
5(1 31) B
Vi[(o!0)- Vis(
2
318
10r31
Vif(o(rr+0)
V [ - -BO
+
Proof. Proof of Part 1.
Suppose 0. By definition of the momentum estimator M, 3,P, (P(0)) pose the error as:
( G we decom-
E
M
Vif(=!0)>|
E P. P.: (0 Vif(x!o))
E (13
(P - Vil( | F
The first term is bounded using the assumption that G(0) is an unbiased estimator with variance of:
E 1...
V.f(x -V ) 1|- s S0907.
For the second term, we apply Lemma A.2
E [|( ((
Virla $ (1-(1:) ())6) |1
Combining both bounds:
E
| Vi < (1--(23, (231 39)6) E
|115 [r]
+ B.f
Proof of Part 2.
Suppose f = kr. At the projection step, we have the update rule:
M P..' (P(")) (11 (1-33)00 1.),) + +,,,(())
Then:
M Vi 122] = | (11 (1)))))(( + 3 - Vif(a!1)) (1)
E
24

---

I--
p! (I-P! (P(()))) Vir(a(o>>1
(2)
(()() (1 B1)aa +3,000 -
( Vif(a(0))/
) 1
The second equality is because of the Pythagorean Theorem.
Applying Lemma A.2 2 to the second term:
P/" (P()) ( (1- 5) ||
E
The first term is bounded as below:
P 1 (P(T)TT (1-8)) - -
V )|²
E
(1 M (1-1) B (i)(a((1) //2]
< E
(1L (M) Vif(x!0)) + ((
vif(x((()) [2] | Vif(=(t))
E
Fl
|M vi)(= 1/
S (1 31)
+ +
|M770-11 Vif((
[|3]
S (1 31)
+ BIot.
where the first inequality because of (P/)) 1, the second inequality is because of the independent noise of mini-batch gradient, the third inequality is because of the bounded noise assumption.
To bound EE
M
1((-1)
apply Young's inequality:
E
M((-1)
Vis(c(()))
[
8,31
(
(z 1[2] (1+
4 8,3,
E ViS(o(t))
-V(((())
1
E
F
So far, we can have the upper bound for the first term in Eq. 4 By applying Eq. (5). Eq. Lemma A.I which naturally holds in our problem setting, we have shown Part 2.
4 and
Proof of Part 3.
Suppose kt +r. 1r.1srsT-- 1). this case, p(() = p((T)) reused. The update becomes:
M (1- (1-8,) 3, (P())
Using the standard decomposition:
M ( (2()) = (1 (1- (1) + (.(P)) (P(")T-TT)))
+ (p(") vif(x())).
By unbiasedness, we have
[2] M Vif(x(0))
F
25

---

(1 (1- -I Vif(z!1)) + 1f
EE
||
+ B
Recall that by (P0)) = 1 and the noise assumption, we have
|| PP
ves(a()))| s 8800.
I
By applying Jensen's 's inequality, we have
E M ( - vif(a()))
||
S = (1- B1)B
M M - vi((=(()) 2
1/2
(1()) - + 8fot.
+ B
We now bound the projection error using Lemma A.2
E (P))) - -)V
(P (P("0) - 1) (xx
1
1
E
4
11 Vif(2t0) - (if(tc( 1
+
1
+
E
38
E
V V.f(c((r))
4
4
-
- Vif(z((r)) 122
+ 1 +
8)
E
The first equality is because of Young's inequality, the second inequality because of Lemma A.2
This concludes the bound. The final contraction inequality is then followed by applying this bound and collecting terms.
For Part 3 of Lemma
we get exactly the same bound for our SARA compared with their GoLore.
Though our Momentum Contraction result is little worse than the one in [4 we can still get the same result for Momentum Error Bound, as shown in Lemma A.4
Lemma A.4 (Momentum Error Bound of MSGD with SARA). Define
0²
Σ of [E(N
Then we have
Kr
M - v -V(0)) |
E
1=
5(1 - 3(1
55(1 r)
+
--
2
Kr
Σ
E
| (0+11) - 2/00
5/4)88/
8/4)8
(1 8/4)3
1==
1 1-//2
2
Kr-2
2
+ KTB100
+
E
(())))))))
T- 6/4
(1-8/4)T84
F
8/4
Proof. First we apply summation to Part 3 of Lemma
A.3
as follows:
(k+1)7-1
(k 1)r-1
Σ
E
M
1 1 (1
8
)
Σ
E
-vi(o((
4
1=kr+1
+1
26

---

(k++))-11
1 <(1-
<
B1
Σ
E
1=
(1++)7-1
5(1 - -11
Vel(a(0)) - vis( (2)
2
E
318
1=
1033 1-1
Vi - Vif(z []: 8
Σ
E
i=1
+ 0f 1)
(k+1)r-1
(1 (1-2) 2
0 31
Σ
1=
(x+1)--2
5(1 31)
- 1)
- P.
+
310
+ 19 -T)
(6)
Then add Eq. (6) and Part 2 of Lemma
A.3 together, we have
(k+1)r-1
(A 1)--1
Σ
M (z)(
/
1 1
5 )81
Σ
E
-
|
t=
(k+1)T-1
S
(1
31 + 2(1 -
Σ
E
|
F
t=kr
(x+1)T-2
5(1 (1) 5BIT(T -1) 1)
(T
1)(1 1)(1-0)) 00
Σ
EE
Vis(alltt))
(1)( /
318
+ 83O8T
Then applying summation over & from 0 to K and summation over all 7 [N] gives us the desired result.
So far, we have the comparable result of upper bound of momentum error, in the next step, apply the same proof procedure as in (24) give us the convergence of low-rank MSGD with SARA. Theorem A.S. Under Assumptions 1-3, if hyperparameters
64
36
38
331 16T L
0 < 81 11
3
0 0< min
4L
8022'
22 V
MSGD-SARA with momentum re-projection converges as
1
Kr-1
Σ 0-
E
013
<
16A
02
KT
kn
35
for any K 2 1, where A (x) - inf, f(x).
Corollary A.6 (Convergence complexity of Low-rank MSGD with SARA), Under Assumption 3.1
3.2 if TT= 2 2 128/(30) + ((880)/(9VVLL)) and we choose
-1
03/202T
64
801. 36,87
80r212
16r
31 1
==
n= 4L
+
LA
3531
35
3
low-rank MSGD-SARA -SRRA with momentum re-projection converges as
T Σ
[1
=0
823T
L.A
+
83.3T
T.A2
where A = f()())) inf fz)
27

---

B Experiment Implementation
To enable re-implementation. we provide hyperparameters herein. For fair comparison, we adopt the same hyperparameters for dominant subspace selection and SARA. Hyperparameters for GaLore
with dominant subspace selection and SARA are shown in Table S Hyperparameters for Fira with
dominant subspace selection and SARA are shown in Table 6
Table Hyperparameters for experiments with GaLore
Name
Values
Batch Size
512
Maximum Sequence Length
512
Warmup Steps
1000 for 60m, 2000 for 130m, 6000 for 350m, 10000 for 1.1B
Rank
128 for 60m, 256 for 130m, 256 for 350m, 512 for 1.1B
Weight Decay
0
Learning Rate
0.01
Scheduler
Cosine
Optimizer Specific Parameters
Adam: |31 0.9, f82 0.999
Adafactor: 31 0.9, (t( -
Adam-mini: 0.9, 32 0.95
Table 6: Hyperparameters for experiments with Fira
Name
Values
batch Size
512
Maximum sequence length
512
Warmup Steps rank
1000 for 60m, 2000 for 130m, 6000 for 350m
128 for 60m, 256 for 130m, 256 for 350m
Weight Decay
Learning Rate
0.01
Scheduler
Cosine
optimizer specific parameters
Adam: B1 0.9, B2 0.999
C More Related Work
Large language models
Many works study the LLM from other aspects. For example, low rank
approximation 220 63 also be applied to improve the computational complexity of (masked)
attention approximation
43
8].
Gaussian kernel density estimation is closely related to attention op-
timization, and similar our work, [4]] uses importance sampling study the dynamic maintenance
of KDE data structures. [6 59 62 15
64 39 analyze the attention regression problems. [ study the computational limits of Mamba. 61 investigates the expressibility of polynomial attention.
18
[17] applies the sketching technique to develop the decentralized large language model. [ studies the attention optimization without requiring strict assumptions.
(73) focuses on efficiently aligning LLMs with recommendation tasks. [19] studies the binary hypothesis testing for softmax models.
Reinforcement Learning In reinforcement learning (RL) 37 38
48 50 49), an agent
learns to make sequential decisions by interacting with environment to maximize cumulative reward. RL algorithms, especially policy gradient methods (c.g.. REINForCE, PPO, TRPO) [74] 70, 14], often rely stochastic gradient descent (SGD) Adam for optimization. Our low-rank optimization techniques for Adam, which could, in theory, be applied to RI. training to make policy optimization more memory-efficient.
Broader Application of Large Language Models The application of large transformers with long-context capabilities is not limited to language processing; broader domains such as meries are also of practical importance. Another line of work applies LLM for
prediction [75 72 69
28

---

structured data understanding [27 30 2 65 Extending and validating our method to preserve the
accuracy of LLMs these tasks is, therefore, worthwhile direction for future work.
D Limitations
Our work does have any noteworthy limitations, as we directly address the key limitation of dominant subspace selection in prior work without requiring extra assumptions, In our paper, we provide a theoretical analysis of the convergence rate, though still relies on the assumptions of T-smoothness (Assumption 3.1 and bounded, centered mini-batch gradient noise (Assumption 3.2
However, note that these assumptions standard in the optimization literature.
E
Societal Impact
Regarding the positive societal impact, by reducing memory requirements for training, SARA method enables more organizations-inelneee those with limited compute budgets-to train or fine-tune competitive L.Mss.
To the best our knowledge, we do not anticipate any negative socictal impacts.
F More experimental results
F.1 High-Rank Updates
Now, we provide more experimental results for high-rank updates.
29

---

mod.ilemddde Eaparco nis dewn, HIGHT
m113 middddddd bypar mE GMA, WWN1T1
middil medeln byrre duni weight
da In SAAN-ANN AONE
@n LY SARN SOVE
Nomaaaee Steple Nelvea
@n on SARA san
Wonatoe
0--
Menatoee
6A
220
320
302
200
300
502
200
502
Updice Stoga
Updira Stopa
Updics Stopa
(a) Normalized singular values of coppdoww_prooo
(b) Normalized singular values of pep-gott_prooo
(c) Normalized singular values of mlp.up_proy
ms model dusel arm weight
module Layars d sef arm awiiht
modakummdde Jaya self aer pro weiight
1.0
1.0
3.0
Full
2.-
2.4
302
220
320 Stago
400
502
1.00
300 U2dite Stago
220
400
502
1.00
200
300
400
502
Updite Staoo
(d) Normalized singular values of self_attn.k _etttt.___ppp.. _proj.
(e) Normalized singular values of certfffttttt_tteeeeet projj.
(f) Normalized singular values of self_attn.g aatn.q_proj.
2.8
GuzaroAAiu Fu Kari Assin
02-
(g) Normalized singular values
sclf_saaeeeeee aadm.._ppooj.
Figure 5: Normalized singular values of the weight difference between the 28k-step checkpoint and 30h-step checkpoint in different layers of Block of LLaMA-60M model during pretraining.
30

---

mnd mow.ilemmdde Enparcl wis dewn, MIGHE
m113 middddddddllli bypar mE GMA, WWNTTT
middis byer Le weight
da In SAAN-ANN SONE
Sn Lm SARN SOVE
Namaaaee Sieple Nelvea
6n on SARA san
1.8
2.-
Wonstoe
Menatoee
220
320
302
200
300
502
200
502
Updise Stoga
Updira Stopa
Updics Stopa
(a) Normalized singular values of mip.down _ooj__roj.
(b) Normalized singular values of papp-gtte_poooo-
(c) Normalized singular values of mip.up_pree
msciue model Lust ar weighht
madilemade Eayara sel arm any awight
modikummdde Jaya sel aer pro weiight
Wwwakoed couurriauee
1.0
1.0
3.0
Ful
2.-
2.4
302
220
320 Stago
400
502
1.00
220
300
400
502
1.00
200
300
400
502
U2dite Stago
Updite Stago
(d) Normalized singular values of self_attn.k _etttt.___ppp _proj.
(c) Normalized singular values of cerifffttttt_poee proj.
(f) Normalized singular values of self_attn.gp watt.q_ppoo.....
GuzaroAAiu Tu Ma-i Assi
0.-
(g) Normalized singular values
sclf_saaeeeeee aadm.u_ppooj.
Figure 6: Normalized singular values of the weight difference between the 28k-step checkpoint and 3tep checkpoint in different layers of Block of LLaMA-60M model during pretraining.
31

---

mnd mod.ilemmddl laparc2 wis dewn, MAGHE
m113 modd bypar22 mE GMA, WWNTTT
medeinmeseeebee byeee 2,m weight
da In SAAN-AN SONE
SN LY SARN SOVE
Namaaaec Steple Nilvea
6n on SARAAiin san
1.8
24
Wonstoe
Menatoee:
220
320
302
200
300
502
200
502
Updice Stoga
Updira Stops
Updics Stopa
(a) Normalized singular values of coppdoww_prooo
(b) Normalized singular values of pellggote_poo...
(c) Normalized singular values of mlp.up_proy
msciue model 2.se! arm weightt
madile.mode Layar sef arm anl awight
modikemmdde Jaya sel arr pro wwiight
1.0
1.0
1.0
Ful
Wwwakoed
2.-
2.4
302
220
320 Stago
400
502
1.00
300 U2dite Stago
220
400
502
1.00
200
300
400
502
Up2ite Stago
(d) Normalized singular values of self_attn.k _etttt.___ppp _ppoj.
(e) Normalized singular values of celiffftttttppoee proj.
(f) Normalized singular values of self_attn.gp watt.q_proo.....
GuzaroAAiu nu mar--nt Assi
02-
(g) Normalized singular values
sclf_saaeeeeee Llw.aatn___eeee p_poj.
Figure 7: Normalized singular values of the weight difference between the 28k-step checkpoint and 3te-step checkpoint in different layers of Block of LLaMA-60M model during pretraining.
32

---

mnd mod.ilemmddd Eaparc) nis dewn, MAGHT
mn13 modd bypar 2. GMA, WWNTTT
midci byere An weight
da In SAAN-ANN SON
On Le SARN SOVE
Namaaaoc Steple Nilvea
6n om SARAAiin san
Menstree:
220
320
302
200
300
502
200
502
Updise Stoga
Updira Stops
Updits Stops
(a) Normalized singular values of coppdoww_prooo
(b) Normalized singular values of papp-gtte_poooo-
(c) Normalized singular values of mlp.up_prop
ms model 2use! arm weightt
madile Layara 2. sef arm any awight
modikummdde Jaya zel arr pro weiight
1.0
1.0
3.0
Guiare Fdl
Wwwakoed
2.-
2.4
302
220
320 Stago
400
502
1.00
300 U2dite Stago
220
400
502
1.00
200
300
400
502
Up2ite Stago
(d) Normalized singular values of self_attn.k _etttt.___ppp.. _proj.
(e) Normalized singular values of cerifffttttt_poee proj.
(f) Normalized singular values of self_attn.gp watt.q_proo.....
2.8
GuzaroAAiu nu Ka-i Assi
02-
(g) Normalized singular values
sels_atin.. L___a_nivv_ppp p_poj.
Figure 8: Normalized singular values of the weight difference between the 28k-step checkpoint and 3te-step checkpoint in different layers of Block of LLaMA-60M model during pretraining.
33

---

mod.ilemmdde Enparct nis dewn, MAGHT
mn13 moddd bypart mE GMA, WWNTTT
midci www.:hieeebbbbe.c......o 4.m p- weight
da In SARN-AAN SON
On LY SARN SOVE
Nomaaaae Nalvea
Gal ve DS kaln
Menstoee Siopsa
Wenatow
24-
0.4
220
320
302
200
300
502
200
502
Updise Stoga
Updira Stopa
Updics Stops
(a) Normalized singular values of copddoww_prooo
(b) Normalized singular values of pap-gote_prooo
(c) Normalized singular values of mlp.up_proy
msciue model dutt ar weightt
madile.mode Layar sel arm anl awight
modikumodee Iay.n sel aer pro weiight
1.0
1.0
1.0
Ful
Wwwakoed
2.-
2.4
302
220
320 Stago
400
502
1.00
220
300
400
502
1.00
200
300
400
502
U2dite Stago
Up2ite Stago
(d) Normalized singular values of self_attn.k _etttt.___ppp _proj.
(c) Normalized singular values of fatff__ttt____tet
(f) Normalized singular values of self_attn.gp watt.q_ppoo.....
2.8
GuzaroAAiu nu mar- Assi
02-
(g) Normalized singular values
sclf_saaeeeeee aadn.u_projj.
Figure 9: Normalized singular values of the weight difference between the 28k-step checkpoint and 3te-stepp checkpoint in different layers of Block of LLaMA-60M model during pretraining.
34

---

mnd mod.ilemmddl laparst nis dewn, MIGHE
mn13 model bypart mE GMA, WWNTT1
medeln byrre S.ng weight
da In SARN-ANN SONE
Mamaaaee Sieple Nilvea
6n on SARs san
Menatoee
220
320
302
200
300
502
200
502
Updise Stoga
Updira Stopa
Updics Stops
(a) Normalized singular values of mip.down _ooj__roj.
(b) Normalized singular values of papp-gtte_poooo.
(c) Normalized singular values of mip.up_pree
msciue model S.se! ar weightt
madulemmde sel arm anl awight
modikumede Jaya St aer pro weiight
1.0
1.0
3.0
Guiare Ful
Wwwakoed
2.-
2.4
302
220
320 Stago
400
502
1.00
220
300
400
502
1.00
200
300
400
502
U2dite Stago
Up2ite Stago
(d) Normalized singular values of self_attn.k _etttt.___ppp _proj.
(c) Normalized singular values of cetifffttttt_tteeeeet proj.
(f) Normalized singular values of self_attn.gp watt.q_ppoo.....
GuzaroAAiu nu Wa-i Assi
02-
(g) Normalized singular values self_atin.v v_proj.
Figure 10: Normalized singular values of the weight difference between the 28k-step checkpoint and 30k-step checkpoint in different layers of Block of LLaMA-60M model during pretraining.
35

---

mod.ilemoddl Eaparci mis dewn, HIGHT
mn13 model bypar mE GMA, WWNSTT
midcis byrre G weight
da In SAAN-AN SONI
ON LY SARN SOVE
Namaaaoc Steple Nilvea
6n om SARs san
1.8
24
Wonatoe
Menstree
24-
220
320
302
200
300
502
200
502
Updice Stoga
Updira Stops
Updics Stops
(a) Normalized singular values of coppdoww_prooo
(b) Normalized singular values of papp-gtte_poooo.
(c) Normalized singular values of mlp.up_prop
msciue model Casl! ar weightt
madilemode Layars C sel arm awiihh
modikummde aer pro weiight
1.0
1.0
1.0
Ful
2.-
2.4
302
220
320 Stago
400
502
1.00
300 U2dite Stago
220
400
502
1.00
200
300
400
502
Up2ite Stago
(d) Normalized singular values of self_attn.k _etttt.___ppp.. _proj.
(e) Normalized singular values of celfffftttttppoe proj.
(f) Normalized singular values of self_attn.gp aatt.q_proj.
2.8
GuzaroAeiu Tu Ma-i Assi
0.2
(g) Normalized singular values self_atin.v v_proj.
Figure 11: Normalized singular values of the weight difference between the 28k-step checkpoint and 30k-stepp checkpoint in different layers of Block 6 of LLaMA-60M model during pretraining.
36

---

mnd mow.ilemoddi layarc? wis dewn, MigHt
mn13 mwdd bypart ME GMA, WWNTTT
midkis mocel.meeee.ooo byere 7.m weight
da In SAAN-AN
ON LY SARN SOVE
Nomeaaae Siepla Nalvea
6n on SARA san
Menstoee Strppa
Menatoe
24
6.4
12
220
320
302
200
300
502
200
502
Updite Stoga
Updira Stops
Updics Stopa
(a) Normalized singular values of mlp.down _ooj__rooo..............
(b) Normalized singular values of fesp-gote_prooj
(c) Normalized singular values of mlp.up_pree
ms model Tise! ar weightt
module Layare sel arm any wwight
modakemmdee Jayar self aer pro weiight
1.0
1.0
1.0
Guiare Ful
2.-
2.4
302
220
320
400
502
1.00
220
300
400
502
1.00
200
300
400
502
Stago
Usdite Stago
Up2ite Stago
(d) Normalized singular values of self_attn.k _etttt.___ppp.. _proj.
(c) Normalized singular values of ceatfffttttt_toee projj.
(f) Normalized singular values of self_attn.gp aatt.q_proj.
2.8
GuzaroAAeu nu rari Assi
02-
(g) Normalized singular values self_atin.. p_proj.
Figure 12: Normalized singular values of the weight difference between the 28k-step checkpoint and 30k-step checkpoint in different layers of Block of LLaMA-60M model during pretraining.
F.2 Anchor Similarity
Now, provide more experimental results for anchor similarity.
37

---

medel syjer Satt A.ttrkppn
midet byree arth 0,
mose Dartr a,ans
Botter Deaecae.eeeet
0.50
0.6
Komter www.eetaaaaeeee..
0.45
0.5
Gelors-Azam
0.40
Gelere-Azam
2500
3000
3500
4000
2500
3000
3500
4000
2500
3000
2500
4000
Upsto Strge
Upsine Strpe
Updotse Singe
(a) Anchor subspace overlap of(b) Anchor subspace overlap of(c) Anchor subspace overlap of attn.o attn.o_proj.
attn.k __roj.
proj..
model duttruc ano)
bayem rilp down pmm
model Jayant relp gats pnd
0.6
Andhr Sbbees
0.40
Colem SM66-At
0.5
Golome 5MM Aih AALmP-Asam
Anchor Subupess
Gol Siil SAW Gol JoLsP--ssm
0.4
0.35
4000
25:00
4,000
2500
Uadit Staca
Updace Stops
Updacce Staps
(d) Anchor subspace overlap of(e) Anchor subspace overlap of(f) Anchor subspace overlap of mlp.doum_proo proj.
attn. prpjj.
mip.gate_pproj
0.8
Modeir S.boper
0.6
Galom SAA---Ah
0.4
3500
4,000
(g) Anchor subspace overlap of chppuppppooj.
Figure 13: The low-rank subspace of different layers in Block 0 the 2000-th iteration is taken as the anchor subspace. The figure shows the overlap between subspaces the corresponding layer in Block O in later iterations and the anchor subspace.
38

---

medel bayer Lam 5.5
mider Jatr o,arn
msont 3.811 0,
0.50
0.45
Calars-Aaam
Gelors-Azam
0.45
Gelore-Azam
Beder
4ocher
0.40
0.40
2500
3000
2500
4000
2500
3000
3500
4000
2500
3000
2500
4000
Ujstee Stres
Upstoe Strppe
Upplae Sreps
(a) Anchor subspace overlap of(bb Anchor subspace overlap of(c) Anchor subspace overlap of attn.o attn.o_proj.
attn.k attn.k_proj.
attn.q_proj.
model byeer La
tayem. ripp down pmm
model Jayant. mep gats pnd
Anchoo Subopess Crette
0.7
0.50
Andia Subipess
Calem 5466-At0 Armm Arraaam
0.5
dolame SMM Aihn Jolmm-asaa
0.6
Gol 5AA SAW GoL CoLsm-Asss
0.5
0.3
0.4
4000
25:00
4,000
2500
Uadit Staca
Updace Stops
Updace Staps
(d) Anchor subspace overlap of(e) Anchor subspace overlap of(f) Anchor subspace overlap of mlp.doum_proo)
attn.v proj.
mip.gate_ppooj
0.7
0.6
Galom So0---
0.5
3500
(g) Anchor subspace overlap of mlp.uppprae proppppro.....
Figure 14: The low-rank subspace of different layers in Block the 2000-th iteration is taken as the anchor subspace. The figure shows the overlap between subspaces the corresponding layer in Block in later iterations and the anchor subspace.
39

---

medel syeer A.prn
mide byree o
mosetbyrrt s,ans
0.7
0.55
0.6
0.50
Azam
0.5
Gelors-Azam
Gelere-Azam
Dotter
Kender
0.4
2500
3000
3500
4000
2500
3000
3500
4000
2500
3000
2500
4000
Upoton Strge
Upsine Strpe
Updltse Singe
(a) Anchor subspace overlap of(b) Anchor subspace overlap of(c) Anchor subspace overlap of attn.k __roj.
attn.o_proj.
attn.q_proj.
1ayem rilp down pmm
model Jayant relp gats pnd
0.7
Anchoo Subopess
0.6
0.50
Calem 5M66-A10 Anmm
dalare 5M6 Aih JoLmm-Ataa
Gol SAA AAW Ar--ASI
0.4
0.5
0.3
4000
25:00
4,000
2500
Upditt
Updace Staps
Updacee Staps
(d) Anchor subspace overlap of(e) Anchor subspace overlap of(f) Anchor subspace overlap of attn.v v_poj.
mlp.doum_proo)
mip.gate_pproj
Galom Sot------
25.00
3500
4,000
(g) Anchor subspace overlap of chppuppppooj
Figure 15: The low-rank subspace of different layers in Block 2 the 2000-th iteration is taken as the anchor subspace. The figure shows the overlap between subspaces the corresponding layer in Block 2 in later iterations and the anchor subspace.
40

---

medel syeer A.pr)
mide Boree aattoo,
mosetbyrre Sart s,ars
0.7
Aocher Sucezace Drattsp
0.5
0.45
Azam
Gelors-Azam
Gelere-Azam
detter
0.4
2500
3000
3500
4000
2500
3000
3500
4000
2500
3000
2500
4000
Upotoe strge
Upsine Strpp
Updlts Singe
(a) Anchor subspace overlap of(b) Anchor subspace overlap of(c) Anchor subspace overlap of attn.o p_roj.
attn.k __roj.
attn.q_proj.
model bywer pmm)
Layen. rilp down
model Jayant melp gats pnd
0.6
0.7
Ancher Subqges
0.5
Calume 5466-A10
0.5
Golome 5M6 Aih AAL-P-ata
Anchoo Subupess
0.6
Gol SiAl AAW AP----
0.4
0.5
0.4
0.3
0.4
4000
25:00
4,000
2500
Upditt
Updace Stops
Updacce Staps
(d) Anchor subspace overlap of(e) Anchor subspace overlap of(f) Anchor subspace overlap of attn.v proj.
mlp.doum_pror,
mip.gate_pproj
0.7
0.6
Galom SoAA-ich
0.4
3500
(g) Anchor subspace overlap of mlp.up_pre proppppoo...
Figure 16: The low-rank subspace of different layers in Block 3 the 2000-th iteration is taken as the anchor subspace. The figure shows the overlap between subspaces the corresponding layer in Block 3 in later iterations and the anchor subspace.
41

---

Ride syjer ANRY A.pr)
mide Dorey 3rth o
mese 4 a,ans
Aocter Sucezare Drattp
0.6
0.6
0.5
0.45
Azaa
0.5
Gelors-Azam
Gelere-Azam
Gotter
0.40
0.4
0.35
2500
3000
3500
4000
2500
3000
3500
4000
2500
3000
2500
4000
Upoton Strge
Upsine Strpe
Updots Singe
(a) Anchor subspace overlap of(b) Anchor subspace overlap of(c) Anchor subspace overlap of attn.o attn.o_proj.
attn.k __roj.
attn.q_proj.
model bypere amt
bayent rrlp down pmm
model Jayart, melp gats pnd
kndto Cuette
0.6
0.5
Anchor Subrpes
0.6
0.5
Colume 5466-A40 Armm Ar--aam
dolore 5M6 Aih GaL wr--aam
Gol SAAl SAW SP---A
dentre
0.3
4000
25.00
4,000
2500
Upditt
Updace Stops
Updacee Staps
(d) Anchor subspace overlap of(e) Anchor subspace overlap of(f) Anchor subspace overlap of attn. p_pjj.
mlp.doum_pprop
mip.gate_pproj
0.7
0.6
Galom 508--60h
3500
4,000
(g) Anchor subspace overlap of mlp.uppprae propp.prj...
Figure 17: The low-rank subspace of different layers in Block 4 the 2000-th iteration is taken as the anchor subspace. The figure shows the overlap between subspaces the corresponding layer in Block 4 in later iterations and the anchor subspace.
42

---

mede syeer Sant A.pr)
mide Duree o
moseebyrre S.rtr anans
0.6
0.55
Aocher Suctaace Orattp
0.50
0.6
0.5
Adam
Gelors-Azam
Gelere-Azam
0.40
0.4
0.4
0.35
2500
3000
3500
4000
2500
3000
3500
4000
2500
3000
2500
4000
Upoton Strge
Upsire Strpe
Updlts Singe
(a) Anchor subspace overlap of(b) Anchor subspace overlap of(c) Anchor subspace overlap of attn.k __roj.
attn.o attn.o_proj.
attn.q_proj.
model www.ryyyertttton. am)
medel byyert
model Jayant mep gate pnd
Ancher Subupes Ceette
0.6
Anctto Subogacs Cverte
0.7
0.6
0.5
Salume 5066 aor--a Jsan
dalare 5MM Aihn Golmm-attam
Gol Siil AAW Ao----
3500
3500
Upditt
Updita
Updacee Staps
(d) Anchor subspace overlap of(e) Anchor subspace overlap of(f) Anchor subspace overlap of attn.v p_poj.
mlp.doum_proo) proj.
mip.gate_prooj
0.7
0.8
Galom 5086--0h
3500
4,000
(g) Anchor subspace overlap of mlp.up_ppe proj_
Figure 18: The low-rank subspace of different layers in Block 5 the 2000-th iteration is taken as the anchor subspace. The figure shows the overlap between subspaces the corresponding layer in Block 5 in later iterations and the anchor subspace.
43

---

mede syeer SART A.pr)
mide Doree 2.tth o.atto..
mosel Deree GLETh RLA
0.5
0.6
0.5
Azam
0.5
Gelors-Azam
Gelore-Azam
Ander
0.4
0.4
0.4
2500
3000
3500
4000
2500
3000
3500
4000
2500
3000
3500
4000
Upote Strge
Upsine Strpe
Upsine Strps
(a) Anchor subspace overlap of(b) Anchor subspace overlap of(c) Anchor subspace overlap of attn.o attn.o_proj.
attn.k __roj.
proj.
bryers attro
medel byyert
model Jayartu melp gats pnd
Annte Gabvgar Certee
0.6
0.5
0.7
0.5
Anctoo Subuggss
Golome 5M6 Aih Jolmm-ataa
0.6
Gol SiAl AAW MALSP-ASSI
25.00
3500
4,000
3500
Upditse Staps
Updita
Updacce Staps
(d) Anchor subspace overlap of(e) Anchor subspace overlap of(f) Anchor subspace overlap of attn.v v_poj.
mlp.doum_proo) proj.
mip.gate_prooj
Galom S0A6-
0.5
0.4
0.3
25000
3500
4,000
(g) Anchor subspace overlap of mlp.uppprae propppprj...
Figure 19: The low-rank subspace of different layers in Block 6 the 2000-th iteration is taken as the anchor subspace. The figure shows the overlap between subspaces the corresponding layer in Block 6 in later iterations and the anchor subspace.
44

---

medel syeer 7. A.prs)
mide Byree arth
mosel byree 7.am RL
0.5
0.6
0.5
Azam
0.5
Gelors-Azam
Gelore-Azam
Andeer
0.4
2500
3000
3500
4000
2500
3000
3500
4000
2500
3000
3500
4000
Upoton Strge
Upsine Strpe
Upsiie Strps
(a) Anchor subspace overlap of(b) Anchor subspace overlap of(c) Anchor subspace overlap of attn.k __roj.
attn.o_proj.
proj.
Bryert aattr.aar
medel byyert?
model Jayant, relp gats pnd
Aocte Sabvear Certee
0.40
0.5
Galare 5MA Aih Colmm-Atta
Gol Siil AAW Gol SoLSP--sss
0.35
0.4-
0.4
0.3
25.00
4,000
3500
Upditse Staps
Updita
Updacee Staps
(d) Anchor subspace overlap of(e) Anchor subspace overlap of(f) Anchor subspace overlap of attn.v u_pjj.
mlp.doum_proo)
mip.gate_ppojj
0.7
0.6
0.5
Galom S0A---A--
0.4
0.3
3500
(g) Anchor subspace overlap of mlp.up_pee proj_
Figure 20: The low-rank subspace of different layers in Block 7 the 2000-th iteration is taken as the anchor subspace. The figure shows the overlap between subspaces the corresponding layer in Block 7 in later iterations and the anchor subspace.
45

---

E3 Adjacent Overlap
medktayss mpp.own
moddd Fayos.0 mp gete
mase//ayess Surii 201
0.6
0.50
0.60
Culare GAFA-Ader GALAm Adam
Gelore GAPA A0LT Galore-Ass
0.45
Gelate-5AR
0.35
Kebcart
0.40
Golome-sssm
0.30
0.3
0.30
2500
3000
3500
4000
2500
3000
3500
4000
2500
3000
3500
4000
Upditt
Updare Strpe
Updats Strpe
(a) Adjacent subspace overlap ol(b) Adjacent subspace overlap of(c) Adjacent subspace overlap of malp.gate proj
valp.down__ee _roj
ralp.up nnrp.up_projj
attia peot
0.425
0.50
Adtcom: Dve
C-lore SATA-A
0.375
CALaT--AAA-AA Almm
Galare-AAp-Atun Culore-Aaum
Celare-kzam
0.60
0.350
0.40
0.35
(d) Adjacent subspace overlap of(e) Adjacent subspace overlap of(() Adjacent subspace overlap eelf_attn.k_pro
fel__attn.q_pp
model www.raayraddd aull
0.50
0.60
0.35
2500
3000
3500
4000
(g) Adjacent subspace overlap of falf___ttt____e__
Figure 21: The overlap between adjacent subspaces of optimization trajectory of different layers in Block 0 in GaLoor-Adam and BaLLoreSARR-AAaa during pretraining on the LLaMA-60M model between 2200-th and 4000-th iteration
46

---

mede snyertl rip COWR,
mode bayan. Aip Ear PMd
mos0i Pyeer Lumb
0.8
0.5
Adtacurt Subesace
0.6
Cularr-Azam Adam
Gelors-Azam
Gelore-Azam
0.4
2500
3000
3500
4000
2500
3000
3500
4000
2500
3000
3500
4000
Hprine Stepe
Upsire Strpp
Upsiie Strps
(a) Adjacent subspace overlap of(b) Adjacent subspace overlap of(c) Adjacent subspace overlap of mip.gate__pin
hnlp.down_ppvee
mip.epp_pre
mode Layenst aronn groj
model Jayensl proj
model laya anl atti
0.5
0.6
0.6
Galem SM66-At
Galare SMA Aih Jolsm--stam
Gol Siil AAW AP--A
0.4
0.4
4,000
2500
Updits Staps
Updit
Updsce Staps
(d) Adjacent subspace overlap of(e) Adjacent subspace overlap of(f) Adjacent subspace overlap of selfnotttto_pp
fet__attnk__poo
0.6-
0.5
3500
4.000
(g) Adjacent subspace overlap of _a_t__tttt__p_oo
Figure 22: The overlap between adjacent subspaces of optimization trajectory of different layers in Block GaLore-Adam and GaLore-SARR SARA-Adam during pretraining on the LLaMA-600 -60 model between 2200-th and 4000-th iteration
47

---

mede Jnyen.2 rip CONR
made byyan Arp gat PM
mosoi pyeee Jumbb
0.6
0.5
0.5
Papove Guberarr
Adticurt Adticutt Subecacc
0.6
0.6
Calaer-Azam Adam
Gelore-Azam
Gelore-Azam
0.3-
0.4
2500
3000
3500
4000
2500
3,00
3500
4000
2500
3000
3500
4000
Hprine Stepe
Uppire Strpp
Upsiin Strps
(a) Adjacent subspace overlap of(b) Adjacent subspace overlap of(c) Adjacent subspace overlap of mlp.gate_ppi
hell.downmppvg
mip.up_proo
model Syarn prnj
Jayan22 proj
model Jayan attr-a peo
Aigernn Dubicso Courto
Adjecens Coeetto
Adgecan subrcaes Couetto
0.6
0.5
Calume 5M66 Arbm Aar---a Asam
Galare Set---------- Aih AoLsP--saa
0.5
Gol SiAl AAW vir- Aw--A
0.5
4000
35.00
4,000
Upditt
Updita
Updatt
(d) Adjacent subspace overlap of(e) Adjacent subspace overlap of(f) Adjacent subspace overlap of self_attn.k _attn.k_proj felf_attt.o_pres
0.5
adgeret
30.00
4000
(g) Adjacent subspace overlap of sel Lattn.e_pror
Figure 23: The overlap between adjacent subspaces of optimization trajectory of different layers in Block in GaLore-Adam and GaLore-SARR SARA-Adam during pretraining on the LLaMA-60N -60 model between 2200-th and 4000-th iteration
48

---

mede payert hip COWK,
model bayan. Arl gat PMM
moshi pyeet Aumbbaa
0.5
Papeoee Guberarn
Adticurt Scbssata
0.6
Calarr-Azam Adam
0.5
Gelore-Azam
0.5
Gelore-Azam
0.3
0.4
0.4
2500
3000
3500
4000
2500
3000
3500
4000
2500
3000
3500
4000
Hprine Stepe
Uppire Strpp
Upsiin Strps
(a) Adjacent subspace overlap of(b) Adjacent subspace overlap of(c) Adjacent subspace overlap of mip.gate__pin
help.down_pppe
mip.up_proo
model byyrr pmo
Jayonc3 proj
model Jayant, and attr- peo
Adgecan subizare Couatto
Adprcee Subrracc Coeetto
Adgeaa Subizies
0.5
Calem 5466-A10 Armm
Golavr 5AM Athi GAL-P-AA130 Ar--aamm
Gol Siil AAW www- Aw--AsI
0.4
35.00
Upditt
Updit
Updalo
(d) Adjacent subspace overlap of(e) Adjacent subspace overlap of(f) Adjacent subspace overlap of self helf_attn.o_pro
setf_attn.q
Adgeent Sucsppce Overts
0.5-
0.5
30.00
(g) Adjacent subspace overlap of sel Lattn.e_pror
Figure 24: The overlap between adjacent subspaces of optimization trajectory of different layers in Block in GaLore-Adam and GaLore-SARR SARA-Adam during pretraining on the LLaMA-60N -60 model between 2200-th and 4000-th iteration
49

---

mede Inyert rii COWR, 5T1
mode byyam. Prl Ear PMd
mosol pyaet tmrb
0.7
0.5
0.6
0.4-
Gelors-Azam
0.3
Gelore-Azam
0.4
0.3
0.3
2500
3000
3500
4000
2500
3000
3500
4000
2500
3000
3500
4000
Hprine Stepe
Uppire Strpp
Upsiin Strps
(a) Adjacent subspace overlap of(b) Adjacent subspace overlap of(c) Adjacent subspace overlap of valp.dowm_ppe proj
mlp.gate_ppi
mippup__roo
model byyrr prnj
layantd wwwo proj
model Jayart, attr- peo
0.55
Adgecan Subocses Coeatto
0.55
0.50
0.50
Salum 5466-A10 Anmm araa
Galare 5M6 Aih GoLir-araa
Gol Siil AAW Ar--s-
0.45
0.4
35000
Upditt
Updita
Updaos
(d) Adjacent subspace overlap of(e) Adjacent subspace overlap of(f) Adjacent subspace overlap of fet__attnk__pa aelf_aatt.oppr
self_attn.q_pre
Adgeent Suzupace
0.5
0.4
25.00
30.00
3500
4000
(g) Adjacent subspace overlap of sel Lattn.e_pror
Figure 25: The overlap between adjacent subspaces of optimization trajectory of different layers in Block in GaLore-Adam and GaLore-SARR SARA-Adam during pretraining on the LLaMA-600 -60 model between 2200-th and 4000-th iteration
50

---

mode by len 5 Sean,or
made Dayan.s Ail Eat PMd
mossi syaet Su
0.6
Adtvort Schoacn
Adam
0.5
Gelors-Azam
Gelore-Azam
0.4
0.4
0.3
2500
3000
3500
4000
2500
3000
3500
4000
2500
3000
3500
4000
Upoton Strge
Upsire Strpp
Upsine Strps
(a) Adjacent subspace overlap of(b) Adjacent subspace overlap of(c) Adjacent subspace overlap of dowdddowppprr
mlp.gate_ppi
mippupp_rro
model Syyens proj
Iayanss attr proj
model Jayan attr- peo
0.55
0.5
Calume 5466-A10 Arom cor--amm
Golare 5MM Aihn Galir--aaa Ara-
0.45
Gol Siil AAW Gol Ao----
0.4-
4000
35.00
4,000
Upditu
Updita
Updaos
(d) Adjacent subspace overlap of(e) Adjacent subspace overlap of(f) Adjacent subspace overlap of felf_attnkk_ppt [elf_oatn.o_pae
Adpeent Susppec
0.5
0.4-
30.00
4000
(g) Adjacent subspace overlap of _attn_t_tppprro
Figure 26: The overlap between adjacent subspaces of optimization trajectory of different layers in Block in GaLore-Adam and GaLore-SARR SARA-Adam during pretraining on the LLaMA-600 -60 model between 2200-th and 4000-th iteration
51

---

model by fet & ceAn,
mode bay.n. Arl gar AM
mosol Pyeer Gorbb
0.50
0.6
0.5
Adam
Gelors-Azam
Gelore-Azam
0.35
0.4
0.4
0.30
2500
3000
3500
4000
2500
3000
3500
4000
2500
3000
3500
4000
Upoton Strge
Upsire Strpp
Uppiin Strps
(a) Adjacent subspace overlap of(b) Adjacent subspace overlap of(c) Adjacent subspace overlap of mip.upp_pre
holp.down_ppree
milp.gate_prt
model Syyrr proj
mode Jayans attia proj
model Jayartu peri
0.50
0.6
Colume SM66-A10 Arim vrr-aam
0.5
Galare 5M66 Ath GaLir-araa
0.40
Gol Siil SAW Gol AALSP-ASSI
4000
25.00
4,000
3000
35.00
4,000
Upditt
Updace Staps
Updaos
(d) Adjacent subspace overlap of(e) Adjacent subspace overlap of(f) Adjacent subspace overlap of felf_aatnkk_ppr self_attn.o_prr
0.55
0.40
30.00
(g) Adjacent subspace overlap of sel Lattn.t_pror
Figure 27: The overlap between adjacent subspaces of optimization trajectory of different layers in Block in GaLore-Adam and GaLore-SARR SARA-Adam during pretraining on the LLaMA-600 -60 model between 2200-th and 4000-th iteration
52

---

mode by.en. rip soan,ar
model Layan. Ar gar PM
mosol pyeer 7.mbb
0.6
0.40
0.5
0.5
Azaa
Gelors-Azam
Gelore-Azam
0.35
0.30
2500
3000
3500
4000
2500
3000
3500
4000
2500
3000
3500
4000
Upotoe Strge
Upsire Strpp
Upsiie Strps
(a) Adjacent subspace overlap of(b) Adjacent subspace overlap of(c) Adjacent subspace overlap of mip.upp_rro
hellp.dwwn_prg
mlppgate__pis
model Syyerc
proj
mode Jayans attra proj
model Jayant, peo
0.8
Dubossse
Colum 5466-A/0 Armm Aiam
Golare 5AM Aihi JoLYP-Atta
Gol SM6 GAW GoL CoLsP--sss
0.35
4000
25:00
4,000
35.00 4,000
Upditt
Updace Staps
Updals
(d) Adjacent subspace overlap of(e) Adjacent subspace overlap of(f) Adjacent subspace overlap of falf_aatnkk_ppr aelf__attto_ppe
0.50
0.45
pdprent
0.40
30.00
4000
(g) Adjacent subspace overlap of _att__t_tt_ap_oo
Figure 28: The overlap between adjacent subspaces of optimization trajectory of different layers in Block in GaLore-Adam and GaLore-SARR SARA-Adam during pretraining on the LLaMA-600 -60 model between 2200-th and 4000-th iteration
53


Reviews:
[{"review_id": "qfqhrYg4zX", "reviewer": "Reviewer_ZP7Q", "summary": "This paper investigates the issue of \"frozen\" dominant subspaces in low-rank optimization during LLM pretraining, demonstrating that conventional dominant subspace selection leads to highly overlapping and restrictive update directions. To address this, the authors introduce SARA, i.e., importance sampling for low rank optimization, which randomly samples singular vectors weighted by their importance to construct more diverse update subspaces. Theoretical convergence guarantees are presented and SARA is shown empirically to outperform prior low-rank optimization approaches across several LLM sizes and optimizers, demonstrating reduced perplexity gaps versus full-rank baselines.", "questions": "Please see the weaknesses. In general:\n\n1.Beside validation perplexity, more metrics are needed to show the model performance.\n\n2.More analysis of hyperparameter sensitivity is needed.\n\n3.It's better to scale up the model size to prove the generalization ability of the method.", "limitations": "Yes.", "rating": 3, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n1.The paper provides a clear and well-motivated analysis of the limitations of dominant subspace selection in existing low-rank optimization methods, supported by empirical observations (see Figure 2) that show increasing and later-stable overlap of dominant subspaces during pretraining.\n\n2.SARA is a simple yet principled modification that introduces importance sampling into the subspace selection process, which reduces overlap between adjacent subspaces (evident in Figure 3) and encourages more effective exploration without adding computational overhead.\n\n3.The methodology is well-integrated with existing low-rank optimizers (e.g., GaLore and Fira), and the plug-and-play nature is likely to make adoption easy within the community.\n\nWeaknesses\n\n1.The evaluation focuses mainly on perplexity during pretraining. While this is the standard criterion for language modeling, it is not evident how SARA’s purported improvements translate into downstream performance (such as finetuning, alignment, or other application-level metrics). The practical significance of reduced PPL gap, especially at scale, is not elaborated on beyond pretraining.\n\n2.While the method is said to be plug-and-play and robust, there is little discussion or systematic ablation regarding the impact of SARA’s hyperparameters (e.g., the subspace refresh frequency τ, the sampled rank r, the total number of candidate singular vectors m) on convergence and final perplexity. This leaves uncertainty on how sensitive SARA is to these settings and may make reproducibility and deployment less straightforward.\n\n3.While the experiments include the model ranging from 60M to 1.1B, there is still a gap between this and the current mainstream LLM size, such as 7B or 32B.  Considering the problem of training cost, if the loss/perplexity changing chart and table under less training tokens can be provided, the paper can also be made more solid.", "was_revised": false, "final_justification": "", "rebuttal": "We express our deepest gratitude to the reviewer for the time and effort in reviewing our work. It is very encouraging for us to see your evaluation of our work as 1). clear and well-motivated, 2). simple and has principled modification, and 3). well-integrated with existing low-rank optimizers. Below, we want to address your concerns about our work.\n\n**Weakness and Question 1**: Besides validation perplexity, more metrics are needed to show the model's performance.\n\n**Response to Weakness and Question 1**: Thank you for your insightful comments. We have conducted additional experiments using multiple downstream tasks from the GLUE benchmark to further evaluate the effectiveness of our method. The table below compares fine-tuning performance on GLUE tasks using pre-trained RoBERTa-Base with different optimizers and LoRA configurations. Results for LoRA and GaLore-Adam are from Zhao et al. [1]. GaLore-SARA-Adam uses identical hyperparameters to GaLore-Adam for the fairness of comparison.\n\n| | CoLA | STS-B | MRPC | RTE | SST2 | MNLI | QNLI | QQP | Average |\n|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|\n|LoRA (rank 4) | 61.38 | 90.57 | 91.07 | 78.70 | 92.89 | 86.82 | 92.18 | 91.29 | 85.61 |\n|GaLore-Adam (rank 4) | 60.35 | 90.73 | 92.25 | 79.42 | 94.04 | 87.00 | 92.24 | 91.06 | 85.89 |\n|GaLore-SARA-Adam ( rank 4) | 61.34 | 90.81 | 91.17 | 80.50 | 94.15 | 87.48 | 92.73 | 91.27 | 86.18|\n|LoRA (rank 8) | 62.83 | 90.80| 91.90 | 79.06 | 93.46 | 86.94 | 92.25 | 91.22 | 85.93|\n|GaLore-Adam (rank 8) | 60.06 | 90.82 | 92.01 | 79.78 | 94.38 | 87.17 | 92.20 | 91.11 | 85.94|\n|GaLore-SARA-Adam (rank 8) | 61.65 | 90.72 | 92.90 | 78.70 | 94.15 | 87.52 | 92.91 | 91.51 | 86.25|\n\n\n\n**Weakness and Question 2**: More analysis of hyperparameter sensitivity is needed.\n\n**Response to Weakness and Question 2**: We thank you for pointing this out. We have performed additional sensitivity analysis on the subspace refresh frequency $\\tau$ and the sampled rank $r$ by training a LLaMA-130M model on the C4 dataset with sequence length 256. The sensitivity of performance to subspace refresh frequency $\\tau$ is presented in the table below, with an optimal setting around $\\tau=200$. Additionally, when $\\tau$ varies from 100 to 500, our algorithm's performance remains stable:\n\n\n| $\\tau$ | 100 | 200 | 300 | 400 | 500 |\n|-----|-----|-----|-----|-----|-----|\n|Evaluation PPL | 24.78 | 24.46 | 24.51 | 24.78 | 24.82|\n\nThe impact of the sampled rank $r$ on perplexity is demonstrated in the table below, confirming our intuition that a higher rank consistently yields lower evaluation perplexity:\n\n| rank $r$ | 100 | 200 | 300 | 400 |\n|-----|-----|-----|-----|-----|\n|Evaluation PPL | 27.70 | 25.23 | 24.04 | 23.29 |\n\n\n**Weakness and Question 3**: It’s better to scale up the model size to prove the generalization ability of the method.\n\n**Response to Weakness and Question 3**: Thank you very much for your insightful comments. We currently report experimental results on models with 60M, 130M, 350M, and 1.1B parameters (Tables 1 and 2), demonstrating that SARA scales effectively across a range of model sizes. Due to limitations in time and computational resources, it is very challenging to present results on larger models at this stage. Even conducting experiments on the 1.1B model takes more than 10 days by using all of our GPUs. Nevertheless, we would like to emphasize that our theoretical convergence guarantee (Theorem 3.4) is independent of model size and thus supports general scalability. \n\nAlthough we lack the time and resources to run very large-scale experiments, we have fully utilized the rebuttal period to present the following experimental results that are computationally efficient and feasible within this limited timeframe. \n\n$\\bullet$ **Hyperparameter sensitivity**: Demonstrated stable performance across critical hyperparameters (see our **Response to Weakness and Question 2** to **Reviewer ZP7Q**).\n\n$\\bullet$ **Additional metrics beyond validation perplexity**: Confirmed enhanced downstream task performance (see our **Response to Weakness and Question 1** to **Reviewer ZP7Q**).\n\n$\\bullet$ **Comparison with uniform sampling**: Highlighted advantages of our sampling strategy (see our **Response to Weakness and Question 2** to **Reviewer yFMY**).\n\n$\\bullet$ **Multiple random seeds**: Verified the robustness and reliability of our method across different runs (see our **Response to Weakness and Question 2** to **Reviewer yFMY**).\n\nOnce again, we are deeply grateful for your insightful comments. We have addressed all identified weaknesses and questions, and we respectfully and genuinely hope you will consider raising your score.\n\n[1] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. \"Galore: Memory-efficient llm training by gradient low-rank projection.\" ICML'24."}, {"review_id": "LPme1MxbZG", "reviewer": "Reviewer_CjAj", "summary": "The paper tackles a limitation of recent low-rank optimizers such as GaLore and Fira: after a few thousand iterations the dominant gradient subspace stabilises, so subsequent updates accumulate in almost the same low-rank slice and stall expressiveness.\n\nTo break this \"frozen subspace\", the authors introduce SARA—Importance SAmpling for Low-RAnk optimisation. Every $\\tau$ steps they compute the SVD of the minibatch gradient and sample r left singular vectors without replacement in proportion to their singular values (Alg. 2). The resulting projector replaces the usual top-r basis and can be plugged into any gradient-projection optimiser.\n\nKey claims are:\n\n* Theory – SARA's stochastic projector yields an $O(L\\Delta/(\\delta^{2.5}T)+\\sqrt{L\\Delta}\\sigma/(\\delta^{3.5}T))$ convergence rate (Theorem 3.4), matching the expected $O(1/\\sqrt{T})$ for SGD momentum in non-convex smooth objective. In comparison, there is no convergence guarantees for the vanilla GaLore method.\n\n* Practice – On LLaMA-60 M/130 M/350 M/1.1 B pre-training over C4 and SlimPajama, SARA paired with GaLore or Fira achieves slightly better or on-par performance when compared with Galore in terms of perplexity.\n\n* Dynamics – SARA lowers adjacent-subspace overlap (Figure 1, 3) and produces higher-rank weight updates (Figure 4), supporting the motivation.", "questions": "* What is the performance of SARA in larger-sized models, such as Llama-3.1-8B?\n\n* Do the improvements in subspace diversity lead to any performance gains for downstream tasks, such as math, commonsense reasoning, or instruction-following?\n\n* I am wondering why the similarity between adjacent projection matrices is large, given the diversity in sampled batches, and the noise is expected to be a dominant term in the gradient in the small batch training setting?", "limitations": "Although the theoretical analysis is conducted under SGD with momentum, the experimental results are based on the Adam optimizer. Given the challenges in establishing convergence guarantees for Adam, it would be helpful to explicitly acknowledge this discrepancy in the limitations section (e.g., Appendix C).", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strength\n\n1. Applying sampling to subspace selection in low-rank projection methods is novel, and compelling to me.\n\n2. The theoretical analysis attains the convergence rate expected for momentum-based SGD on smooth objectives, and addresses the non-convergence issue of Galore.\n\n3. The paper is clearly written and straightforward to follow.\n\n4. The experimental results are comprehensive and convincingly support the authors' claims.\n\n### Weakness\n\n1. **Scalability**: Adding experiments on larger language models—for instance, Llama-3.1-8B—would strengthen the empirical evidence for the method’s scalability.\n\n2. **Theory-practice Gap**: While the theorem is proven for SGD momentum, the experiments use Adam for optimization. I understand that it is difficult to prove convergence guarantees for Adam given its non-projection-free property, but it would be beneficial to mention this in the limitations (Appendix C).\n\n3. **Minor presentation issues**:\n\n* Some missing citations [1–4].\n\n* Algorithm 1, line 7: The formatting could be refined for readability.\n\n* Tables 1–4: Sorting the results in descending order would make the comparisons clearer.\n\n* Line 469: The statement should involve inequalities rather than equalities.\n\n### References\n\n[1]: Liu, Xinxin, et al. \"Refining Salience-Aware Sparse Fine-Tuning Strategies for Language Models.\" arXiv preprint arXiv:2412.13488 (2024).\n\n[2]: Pan, Rui, et al. \"LISA: layerwise importance sampling for memory-efficient large language model fine-tuning.\" NeurIPS 2024.\n\n[3]: Hui, Tingfeng, et al. \"Hft: Half fine-tuning for large language models.\" arXiv preprint arXiv:2404.18466 (2024).\n\n[4]: Li, Pengxiang, et al. \"Owlore: Outlier-weighed layerwise sampled low-rank projection for memory-efficient llm fine-tuning.\" arXiv preprint arXiv:2405.18380 (2024).", "was_revised": true, "final_justification": "The rebuttal addresses most of my concerns, and I would like to maintain my current rating.", "rebuttal": "We express our deepest gratitude to the reviewer for the time and effort in reviewing our work. It is very encouraging for us to see your evaluation of our work as 1). novel and compelling, 2). clearly written, and 3). containing sufficient theoretical and experimental results. Below, we want to address your concerns about our paper.\n\n**Weakness and Question 1**: Adding experiments on larger language models—for instance, Llama-3.1-8B—would strengthen the empirical evidence for the method’s scalability. \n\n**Response to Weakness and Question 1**: Thank you very much for your insightful comments. We currently report experimental results on models with 60M, 130M, 350M, and 1.1B parameters (Tables 1 and 2), demonstrating that SARA scales effectively across a range of model sizes. Due to limitations in time and computational resources, it is very challenging to present results on larger models at this stage. Even conducting experiments on the 1.1B model takes more than 10 days by using all of our GPUs. Nevertheless, we would like to emphasize that our theoretical convergence guarantee (Theorem 3.4) is independent of model size and thus supports general scalability. \n\nAlthough we lack the time and resources to run very large-scale experiments, we have fully utilized the rebuttal period to present the following experimental results that are computationally efficient and feasible within this limited timeframe. \n\n$\\bullet$ **Hyperparameter sensitivity**: Demonstrated stable performance across critical hyperparameters (see our **Response to Weakness and Question 2** to **Reviewer ZP7Q**).\n\n$\\bullet$ **Additional metrics beyond validation perplexity**: Confirmed enhanced downstream task performance (see our **Response to Weakness and Question 1** to **Reviewer ZP7Q**).\n\n$\\bullet$ **Comparison with uniform sampling**: Highlighted advantages of our sampling strategy (see our **Response to Weakness and Question 2** to **Reviewer yFMY**).\n\n$\\bullet$ **Multiple random seeds**: Verified the robustness and reliability of our method across different runs (see our **Response to Weakness and Question 2** to **Reviewer yFMY**).\n\n**Weakness 2**: While the theorem is proven for SGD momentum, the experiments use Adam for optimization. I understand that it is difficult to prove convergence guarantees for Adam given its non-projection-free property, but it would be beneficial to mention this in the limitations (Appendix C).\n\n**Response to Weakness 2**: We thank you for your insightful comments. We use Adam since it is a widespread adoption in large-scale LLM pretraining for its empirical performance. We completely agree with your insightful comments that proving convergence guarantees for Adam, especially under low-rank projection with importance sampling, remains challenging due to its complex dynamics and lack of a projection-free formulation. We have included this discussion in the limitations part of our paper. We thank you again for pointing this out!\n\n**Weakness 3**: Some missing citations [1–4]. Algorithm 1, line 7: The formatting could be refined for readability. Tables 1–4: Sorting the results in descending order would make the comparisons clearer. Line 469: The statement should involve inequalities rather than equalities.\n\n\n**Response to Weakness 3**: Thank you very much for pointing out these wonderful papers [1–4]. We have carefully read these recent works, and we believe they are highly related to our work. [1] analyzes the sparsity-based parameter-efficient fine-tuning (SPEFT) for LLMs, which is an alternative method of LoRA. [2] designs a novel method for memory-efficient fine-tuning for LLMs. It has been shown that it can outperform LoRA and full-parameter training in many cases. Similarly, [3] proposes another novel fine-tuning method called Half Fine-Tuning (HFT), which can mitigate \"catastrophic forgetting\" in LLMs during sequential training and instruction tuning. Finally, [4] also proposes a novel memory-efficient fine-tuning method by strategically selecting layers to update based on outlier statistics. Our paper also considers the memory and convergence behavior, but the difference is that we focus on the LLM pretraining instead of fine-tuning. We have included the above discussion in the revised version of our paper.\n\nRegarding Algorithm 1 line 7, we have updated it as the following two lines:\n\\begin{align*}\n    \\mathcal{S} &\\gets \\\\{V_l^{(t - 1)}, M_l^{(t - 1)}, x_l^{(t)}, P_l^{(t)}, G_l^{(t)}, \\beta_1, \\beta_2, \\xi, \\eta, \\alpha\\\\}. ~~~~~~~~~\\text{Input parameters}\\\\\\\\\n    x_l^{(t)} &\\gets \\text{GaLore-Adam}(\\mathcal{S}) \\text{ or } \\text{Fira-Adam}(\\mathcal{S}).\n\\end{align*}\nThank you for your insightful comments. \n\nAlso, we have fixed our Tables 1–4 to make \"the results in descending order\" to make them clearer.\n\nFinally, regarding Line 469, we first want to sincerely thank you for carefully checking the details of our work, even though this is in the appendix. In Line 469, we have:\n\\begin{align*}\n\\mathbb{E} \\left[ \\left\\\\| \\tilde{M}_l^{(0)} - \\nabla_l f(x^{(0)}) \\right\\\\|_F^2 \\right] \n&= \\mathbb{E} \\left[ \\left\\\\| \\beta_1 P_l^{(0)} (P_l^{(0)})^\\top \\left( G_l^{(0)} - \\nabla_l f(x^{(0)}) \\right) \\right\\\\|_F^2 \\right] \\\\\\\\\n&\\quad + \\mathbb{E} \\left[ \\left\\\\| \\left( \\beta_1 P_l^{(0)} (P_l^{(0)})^\\top - I \\right) \\nabla_l f(x^{(0)}) \\right\\\\|_F^2 \\right].\n\\end{align*}\nThis is because in general $\\mathbb{E}[\\\\|A + B\\\\|_F^2] = \\mathbb{E}[\\\\|A\\\\|_F^2] + \\mathbb{E}[\\\\|B\\\\|_F^2] + 2\\mathbb{E}[\\langle A, B \\rangle]$, but here, we have $\\mathbb{E}[\\langle A, B \\rangle] = 0$, where $A = \\beta_1 P_l^{(0)} (P_l^{(0)})^\\top \\left( G_l^{(0)} - \\nabla_l f(x^{(0)}) \\right)$ and $B = \\left( \\beta_1 P_l^{(0)} (P_l^{(0)})^\\top - I \\right) \\nabla_l f(x^{(0)})$. The reason is that $A$ has ``zero-mean'' since we have $G_l^{(0)}$ is unbiased, as our assumption.\n\n**Question 2**: Do the improvements in subspace diversity lead to any performance gains for downstream tasks, such as math, commonsense reasoning, or instruction-following?\n\n**Response to Question 2**: Thank you for your insightful comments. We have conducted additional experiments using multiple downstream tasks from the GLUE benchmark to further evaluate the effectiveness of our method. The table below compares fine-tuning performance on GLUE tasks using pre-trained RoBERTa-Base with different optimizers and LoRA configurations. Results for LoRA and GaLore-Adam are from Zhao et al. [7]. GaLore-SARA-Adam uses identical hyperparameters to GaLore-Adam for the fairness of comparison.\n\n| | CoLA | STS-B | MRPC | RTE | SST2 | MNLI | QNLI | QQP | Average |\n|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|\n|LoRA (rank 4) | 61.38 | 90.57 | 91.07 | 78.70 | 92.89 | 86.82 | 92.18 | 91.29 | 85.61 |\n|GaLore-Adam (rank 4) | 60.35 | 90.73 | 92.25 | 79.42 | 94.04 | 87.00 | 92.24 | 91.06 | 85.89 |\n|GaLore-SARA-Adam ( rank 4) | 61.34 | 90.81 | 91.17 | 80.50 | 94.15 | 87.48 | 92.73 | 91.27 | 86.18|\n|LoRA (rank 8) | 62.83 | 90.80| 91.90 | 79.06 | 93.46 | 86.94 | 92.25 | 91.22 | 85.93|\n|GaLore-Adam (rank 8) | 60.06 | 90.82 | 92.01 | 79.78 | 94.38 | 87.17 | 92.20 | 91.11 | 85.94|\n|GaLore-SARA-Adam (rank 8) | 61.65 | 90.72 | 92.90 | 78.70 | 94.15 | 87.52 | 92.91 | 91.51 | 86.25|\n\n**Question 3**: I am wondering why the similarity between adjacent projection matrices is large, given the diversity in sampled batches, and the noise is expected to be a dominant term in the gradient in the small batch training setting?\n\n**Response to Question 3**: The randomness in mini-batch gradients can indeed introduce some variation in the dominant subspaces across steps. However, this typically does not have a significant impact on the similarity between adjacent projection matrices, as demonstrated empirically in [5]. In the problem setting of [6], the subspace spanned by the top-$k$ singular vectors tends to stabilize early in training and remains nearly fixed, even though individual singular vectors may vary. The gradient dynamically adapts to reside within this top subspace, resulting in a consistent projection structure across steps [6]. \n\nWe sincerely thank you for raising this important point. We have incorporated this clarification into the final version of our paper, as we believe it strengthens the motivation behind our work.\n\nOnce again, we are deeply grateful for your insightful comments. We have addressed all identified weaknesses and questions, and we respectfully and genuinely hope you will consider raising your score.\n\n \n[1] Xinxin Liu, Aaron Thomas, Cheng Zhang, Jianyi Cheng, Yiren Zhao, and Xitong Gao. \"Refining Salience-Aware Sparse Fine-Tuning Strategies for Language Models.\" ACL'25.\n\n[2] Rui Pan, Xiang Liu, Shizhe Diao, Renjie Pi, Jipeng Zhang, Chi Han, and Tong Zhang. \"Lisa: Layerwise importance sampling for memory-efficient large language model fine-tuning.\" NeurIPS'24.\n\n[3] Tingfeng Hui, Zhenyu Zhang, Shuohuan Wang, Weiran Xu, Yu Sun, and Hua Wu. \"Hft: Half fine-tuning for large language models.\" Preprint'24.\n\n[4] Pengxiang Li, Lu Yin, Xiaowei Gao, and Shiwei Liu. \"Owlore: Outlier-weighed layerwise sampled low-rank projection for memory-efficient llm fine-tuning.\" Preprint'24.\n\n[5] Zhenyu Zhang, Ajay Jaiswal, Lu Yin, Shiwei Liu, Jiawei Zhao, Yuandong Tian, and Zhangyang Wang. \"Q-galore: Quantized galore with int4 projection and layer-adaptive low-rank gradients.\" Preprint'24.\n\n[6] Guy Gur-Ari, Daniel A. Roberts, and Ethan Dyer. \"Gradient descent happens in a tiny subspace.\" Preprint'18.\n\n[7] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. \"Galore: Memory-efficient llm training by gradient low-rank projection.\" ICML'24."}, {"review_id": "Cx4U8qaaD4", "reviewer": "Reviewer_yFMY", "summary": "This paper addresses a key issue in low-rank optimizers for LLM training: the update directions tend to collapse into the same subspace over time, limiting learning. The authors propose a simple method to fix this while keeping the efficiency benefits of low-rank updates.\n\n**Main contributions**\n\n- Propose **SARA**, a sampling method that selects singular vectors based on their importance to keep update subspaces diverse.  \n- Prove that SARA maintains the convergence rate of standard low-rank methods while preserving memory and compute efficiency.  \n- Show that SARA can be easily added to existing optimizers like GaLore, Fira, and Adafactor without major code changes.  \n- Demonstrate improved results on LLaMA models, reducing the gap between low-rank and full-rank Adam on C4 and SlimPajama.  \n- Provide evidence that SARA improves subspace diversity using overlap and rank-based metrics.", "questions": "1. **Why sample by singular values?**  \n   Why are the sampling probabilities proportional to singular values, for instance, why not proportional to exp(.) of singular values to resemble a softmax function? Have you tried uniform sampling? An ablation would clarify this design choice.\n\n> **I understand that running additional experiments might not be feasible, I'd be happy to increase my score if authors can provide some discussion and maybe run some simple experiments with uniform sampling.**\n\n2. **Aggregation of subspace metrics**  \n   Section 4.3 shows overlap for only one parameter group. Please aggregate results over all layers and groups (e.g., average or any reasonable summarization metric for overlap per layer and per training step) to give a complete picture.\n\n> **I am willing to increase my score if authors can aggregate the results for subspace metrics.**\n\n3. **How often is SVD computed?**  \n   How frequently do you compute the SVD and update the sampling probabilities? I could not find details about the hyperparameter $ \\tau $ in the paper. I assume it matches the baseline methods, but please clarify and report it in Appendix B. Also, a rough estimate of the added runtime or FLOP overhead would be helpful.\n\n4. **Is sampling stable across runs?**  \n   Does the added randomness affect training stability or convergence? Showing results across multiple seeds would help. Also, does random sampling affect the optimal value for hyperparameter $\\tau$?\n\n5. **Typo in Figure 4?**  \n   The x-axis in Figure 4 (and similar plots in the appendix) is labeled \"update steps,\" but it appears to show singular value indices. Please clarify or correct the label.", "limitations": "Add a short paragraph acknowledging technical limits: extra SVD cost, possible instability from random sampling, etc.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "## Strengths\n\n- **Clear motivation.** The paper addresses a real issue in low-rank training: the update directions collapse into the same subspace over time, limiting learning. This is a well-motivated and practically important problem.  \n- **Simple and useful method.** The proposed SARA algorithm is easy to implement, requires no extra memory or compute, and can be applied to many existing optimizers.  \n- **Theoretical support.** While not central to the paper’s impact, the convergence analysis adds technical credibility.  \n- **Strong experiments.** The paper includes extensive experiments on LLaMA models, showing consistent improvements and closing much of the gap with full-rank Adam.  \n\n## Weaknesses\n\n- **Unorganized experiments.** Section 4.3 visualizes subspace metrics for just one parameter group, with the rest of the figures relegated to the appendix. The paper should aggregate and discuss results across all layers and groups in the main text (for example, plotting the average subspace overlap per layer and the average overlap over training steps for each parameter group) to give a clearer, more comprehensive view of the method’s behavior. \n- **No comparison with alternative sampling methods.** The paper does not benchmark SARA against simpler schemes such as uniform, top-k, or random-k sampling, leaving it unclear how much of the gain is due to importance sampling versus just adding noise or diversity.\n- **Writing could be improved.** Some parts of the paper are under-explained, especially the sampling procedure. The design choices, such as sampling probabilities being proportional to singular values, are not well motivated or justified.\n- **No analysis of randomness.** It is unclear whether the randomness introduced by sampling could hurt performance or stability in some cases. A discussion or ablation would help. For instance, does it affect the distributed training/parallelization?", "was_revised": true, "final_justification": "The rebuttal addressed most of my actionable concerns: it added uniform and top-r baselines over three seeds showing consistent PPL gains for importance sampling, clarified SVD refresh ($\\tau=200$) with negligible overhead, and reported multi-seed (3 seeds) stability and a $\\tau$ sweep (best $\\approx 200$). They also averaged subspace overlap across layers in the rebuttal and indicated that figure labeling/organization will be fixed in the revision.\n\nRemaining issues: there is still no comparison to a softmax sampling variant, and the rationale for using probabilities proportional to singular values (rather than other importance-sampling schemes) needs a clearer justification. In addition, the layer-wise aggregation and the new baselines should be included in the main paper with transparent reporting (means, standard deviations or 95% CIs, number of seeds, and training tokens). With these caveats, the rebuttal strengthens the work enough to warrant a borderline accept (4).", "rebuttal": "We express our deepest gratitude to the reviewer for the time and effort in reviewing our work. It is very encouraging for us to see your evaluation of our work as 1). having clear motivation, 2). simple and useful, and 3). containing theoretical support and strong experiments. Below, we want to address your concerns about our work.\n\n**Question 1 and Weakness 3**: Why are the sampling probabilities proportional to singular values, for instance, why not proportional to exp(.) of singular values to resemble a softmax function? Have you tried uniform sampling? An ablation would clarify this design choice.\n\n\n**Response to Question 1 and Weakness 3**: The frozen dominant subspace phenomenon was recently formalized in [1, 2], and to the best of our knowledge, our paper is the first to propose a theoretically grounded and empirically validated method to mitigate it. Specifically, while prior methods like GaLore and $Q$-GaLore project gradients onto the dominant subspace (defined by leading singular vectors), we show empirically (e.g., Figure 3, Table 1) that such a strategy can lead to stagnation in subspace evolution during pretraining, ultimately limiting the expressiveness of weight updates.\n\nTo tackle this, we introduce an importance sampling mechanism inspired by classical sketching techniques in randomized linear algebra. Given a big data matrix $A \\in \\mathbb{R}^{n \\times d}$, sketching aims to reduce dimensionality by sampling a small matrix $SA$, where $S \\in \\mathbb{R}^{s \\times n}$ and $s \\ll n$. A principled way to construct $S$ is by leveraging leverage scores, which quantify each row’s contribution to the row space of $A$. Rather than selecting the top-$k$ leverage rows (which could lead to overfitting to a narrow subspace), sketching theory recommends sampling **proportionally** to leverage scores, which better preserves the full subspace structure [3].\n\nWe draw a direct analogy between this and our method: we treat the singular values (quantifying each singular vector’s contribution) of the gradient as analogous to leverage scores, and sample left singular vectors without replacement, **proportional** to their singular values. This avoids always selecting the same top-$k$ directions, like in GaLore and $Q$-GaLore, and instead promotes diversity across training steps. Our importance sampling method is not a heuristic adjustment made purely for empirical performance. Rather, it is a theoretically motivated approach rooted in randomized linear algebra and directly addresses the frozen subspace problem. While LoRA-related work focuses on augmenting low-rank structure, our work focuses on dynamically diversifying it during optimization, which is a crucial distinction. \n\nMoreover, SVD is a well-established and widely used technique in machine learning. Many methods—such as principal component analysis (PCA) [4] and attention approximation techniques [5]—rely on projecting data onto the top-$k$ left singular vectors of a data matrix to reduce dimensionality. Our randomized importance sampling strategy offers a novel perspective: instead of deterministically selecting the top directions, probabilistic selection may enhance subspace diversity and performance. We believe this insight could inspire future work in improving other applications that currently rely on fixed top-$k$ subspace projections and in theoretically and/or empirically studying the similarities and differences between our importance sampling and leverage score sampling.\n\n**Weakness and Question 2**: No comparison with alternative sampling methods. The paper does not benchmark SARA against simpler schemes such as uniform, top-$k$, or random-$k$ sampling, leaving it unclear how much of the gain is due to importance sampling versus just adding noise or diversity.\n\n**Response to Weakness and Question 2**: We run experiments with uniform sampling, our proposed importance sampling, and top-$r$ selection (the subspace selection used in GaLore) using three different seeds on the LLaMA-130M model with a sequence length of 256. The models are trained for 2 billion tokens. The results are provided in the table below. We can see that importance sampling outperforms both uniform sampling and top-$r$ selection.\n\n|   | Seed 1 | Seed 2 | Seed 3 | mean|\n|-----|-----|-----|-----|-----|\n|Top-$r$ | 24.92 | 24.95 | 24.99 | 24.96|\n|Uniform Sampling | 24.69 | 24.68 | 24.69 | 24.68|\n|Importance Sampling | 24.46 | 24.47 | 24.54 | 24.49|\n\n\n\nWe average the subspace metrics across all layers, and the resulting averages are presented in the table below. We observe that the overlap between the anchor subspace and subsequent subspaces is lower for GaLore-SARA-Adam compared to GaLore-Adam. This indicates that SARA encourages the optimization trajectory to explore a more diverse set of subspaces relative to using the dominant subspace. The averaged anchor subspace overlap metrics for GaLore-Adam and GaLore-SARA-Adam across all layers are summarized as follows:\n\n|Training iterations | 2200 | 2400 | 2600 | 2800 | 3000 | 3200 | 3400 | 3600 | 3800 | 4000|\n|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|\n|GaLore-Adam | 0.6328 | 0.6328| 0.6328| 0.6289| 0.6289| 0.6211| 0.6211| 0.6172| 0.6133 | 0.6094|\n|GaLore-SARA-Adam | 0.3516| 0.3535| 0.3496| 0.3496| 0.3477| 0.3477| 0.3477| 0.3438| 0.3418| 0.3398|\n\n\n**Question 3**: How frequently do you compute the SVD and update the sampling probabilities? I could not find details about the hyperparameter $\\tau$ in the paper. I assume it matches the baseline methods, but please clarify and report it in Appendix B. Also, a rough estimate of the added runtime or FLOP overhead would be helpful.\n\n**Response to Question 3**: SVD is computed every $\\tau$ iterations, where $\\tau$ is the subspace update period of our algorithm. In our experiments, we set $\\tau$ to the same value as in GaLore’s original paper, which is 200. We have included it in the updated version of our paper. \n\nThe additional overhead introduced by sampling is minimal. To address your concern, we conducted an additional experiment. In the experiment, we measure the wall-clock time of SVD computation and sampling. We calculate the average wall-clock time of ten runs, the result for SVD is 0.34 seconds for a 2048 by 2048 matrix, and the result for sampling is 0.0005 seconds. The empirical result shows that the addition computational cost brought by sampling is negligible.\n\n\n**Weakness 4 and Question 4**: Does the added randomness affect training stability or convergence? Showing results across multiple seeds would help. Also, does random sampling affect the optimal value for the hyperparameter $\\tau$?\n\n**Response to Weakness 4 and Question 4**: Sampling is stable across runs. In our **Response to Weakness and Question 2**, we show the results of GaLore-SARA-Adam across three seeds, demonstrating its stability. We also test whether random sampling affects the optimal subspace update frequency, i.e., $\\tau$. Specifically, we train the LLaMA-130M model with different $\\tau$ values on the C4 dataset with a sequence length of 256. The results are shown in the table below. We find that the optimal $\\tau$ remains 200. This indicates that random sampling does not affect the optimal value of $\\tau$, which is an important hyperparameter.\n\n| $\\tau$ | 100 | 200 | 300 | 400 | 500 |\n|-----|-----|-----|-----|-----|-----|\n|Evaluation PPL | 24.78 | 24.46 | 24.51 | 24.78 | 24.82|\n\nAlso, in our experiment, we did not find that sampling has any negative effect on distributed training.\n\n**Weakness 1 and Question 5**: Organization and typos. Section 4.3 visualizes subspace metrics for just one parameter group, with the rest of the figures relegated to the appendix. The paper should aggregate and discuss results across all layers and groups in the main text (for example, plotting the average subspace overlap per layer and the average overlap over training steps for each parameter group) to give a clearer, more comprehensive view of the method’s behavior.\nThe $x$-axis in Figure 4 (and similar plots in the appendix) is labeled \"update steps,\" but it appears to show singular value indices. Please clarify or correct the label.\n\n\n**Response to Weakness 1 and Question 5**: We sincerely appreciate your valuable comments on the organization and typos in our work. We completely agree with your valuable comments, and we have fixed these issues in the revised version of our paper.\n\n\nOnce again, we are deeply grateful for your insightful comments. We have addressed all identified weaknesses and questions, and we respectfully and genuinely hope you will consider raising your score.\n\n\n[1] Zhenyu Zhang, Ajay Jaiswal, Lu Yin, Shiwei Liu, Jiawei Zhao, Yuandong Tian, and Zhangyang Wang. \"Q-galore: Quantized galore with int4 projection and layer-adaptive low-rank gradients.\" Preprint'24.\n\n[2] Guy Gur-Ari, Daniel A. Roberts, and Ethan Dyer. \"Gradient descent happens in a tiny subspace.\" Preprint'18.\n\n[3] Michael W. Mahoney \"Randomized algorithms for matrices and data.\" Foundations and Trends in Machine Learning'11.\n\n[4] Jianqing Fan, Qiang Sun, Wen-Xin Zhou, and Ziwei Zhu. \"Principal component analysis for big data.\" Preprint'18.\n\n[5] Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. \"Linformer: Self-attention with linear complexity.\" Preprint'20."}, {"review_id": "BgBaMJTJZv", "reviewer": "Reviewer_TBg6", "summary": "This paper identifies and addresses the \"frozen subspace\" problem in low-rank optimization methods used for LLM pretraining. The authors observe that methods like GaLore, which rely on projecting gradients onto a dominant subspace, often suffer from this subspace becoming static early in training, thus limiting the model's learning capacity. To counteract this, the paper proposes SARA (Importance Sampling for Low-Rank optimization), which replaces the deterministic selection of top singular vectors with a stochastic, importance-sampling-based approach. By sampling basis vectors with probabilities proportional to their singular values, SARA encourages the exploration of more diverse subspaces throughout training. The work is supported by a theoretical convergence guarantee and empirical results on LLaMA models, demonstrating that SARA consistently improves performance over dominant subspace baselines and narrows the gap with full-rank Adam.", "questions": "Please refer to the Weaknesses.", "limitations": "Yes", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "This paper presents a contribution to memory-efficient optimization, with several key strengths but also noteworthy weaknesses that temper its overall impact.\n\n**Strengths**\n\n1. Clear Motivation and Problem Diagnosis: The identification and empirical illustration of the \"frozen dominant subspace\" phenomenon provide a compelling rationale for the proposed work.\n\n2. Simple and Practical Method: As a plug-and-play module that modifies only the subspace selection step, SARA is highly practical and can be easily integrated into existing low-rank optimization frameworks like GaLore and Fira.\n\n3. Theoretical Backing: Providing a formal convergence guarantee (Theorem 3.4) is a significant strength. It adds a layer of theoretical rigor that is absent in the popular dominant subspace approach, bolstering the credibility of the method.\n\n4. Good Clarity and Presentation: The paper is well-written and easy to follow. The problem, method, and results are presented logically and supported by clear figures and algorithms.\n\n**Weaknesses**\n\n1. Incremental Novelty: The core idea of analyzing and modifying low-rank subspaces to improve performance is an established paradigm, particularly in work on improving LoRA. This positions the contribution as an effective but arguably incremental advancement rather than a fundamentally new approach.\n\n2. Limited Experimental Scope and Scale: The empirical evaluation, while covering multiple optimizer variants, lacks the scale and breadth expected of top-tier work in this area. \n\ni) The experiments do not extend to larger, more commonly benchmarked models like a LLaMA-7B, and larger datasets, which would be a more convincing test of scalability.\n\nii) The evaluation is confined to perplexity. It critically lacks evaluation on a suite of downstream fine-tuning tasks (e.g. Avg Acc), which is the standard for demonstrating the quality and generalizability of a pretrained model.\n\niii) There is a crucial absence of any ablation or sensitivity analysis on the rank r. (a key hyperparameter)\n\n3. Marginal Performance Gains: The practical significance of the improvements is questionable in some cases. For instance, on the 130M model, SARA's perplexity gain over the Fira-Adam baseline is minimal (22.22 vs. 22.37). Maybe the SOTA baseline has implicitly solved the subspace problem.\n\n4. \"No Overhead\" Claim: The paper claims SARA \"does not bring extra overhead\". However, as the method introduces an explicit weighted sampling step after the SVD. This claim requires empirical validation with wall-clock time and memory usage comparisons to prove that the additional computational cost is indeed negligible.\n\n5. The paper's own theory suggests SARA's convergence rate is slightly worse than GoLore's by a constant factor , yet SARA consistently outperforms GoLore in experiments. This gap is not adequately discussed, leaving open questions about the factors driving the empirical success.", "was_revised": true, "final_justification": "I keep my score unchanged. The authors have addressed some of my concerns, but others remain. E.g., the experimental setting is not reasonable, and I have concerns on the scalability of the method. It is crucial to test models on a large model, at least 7B.", "rebuttal": "We express our deepest gratitude to the reviewer for the time and effort in reviewing our work. It is very encouraging for us to see your evaluation of our work as 1). having clear motivation, 2). simple and practical, 3). theoretically rigorous, and 4). having clear presentation. Below, we want to address your concern about our work.\n\n**Weakness 1**: The idea of analyzing low-rank subspaces to improve performance is an established paradigm, particularly in work on improving LoRA. This positions the contribution as an effective but arguably incremental advancement.\n\n**Response to Weakness 1**: We appreciate your observation that analyzing and modifying low-rank subspaces is an established paradigm, especially in the context of LoRA. We agree that this is a widely explored area. However, we would like to clarify that the main novelty of our work is to address a newly identified challenge within this paradigm—namely, the frozen dominant subspace phenomenon.\n\nThis phenomenon was recently formalized in [1, 2], and to the best of our knowledge, our paper is the first to propose a theoretically grounded and empirically validated method to mitigate it. Specifically, while prior methods like GaLore and $Q$-GaLore project gradients onto the dominant subspace (defined by leading singular vectors), we show empirically (e.g., Figure 3, Table 1) that such a strategy can lead to stagnation in subspace evolution during pretraining, ultimately limiting the expressiveness of weight updates.\n\nTo tackle this, we introduce an importance sampling mechanism inspired by classical sketching techniques in randomized linear algebra. Given a big data matrix $A \\in \\mathbb{R}^{n \\times d}$, sketching aims to reduce dimensionality by sampling a small matrix $SA$, where $S \\in \\mathbb{R}^{s \\times n}$ and $s \\ll n$. A principled way to construct $S$ is by leveraging leverage scores, which quantify each row’s contribution to the row space of $A$. Rather than selecting the top-$k$ leverage rows (which could lead to overfitting to a narrow subspace), sketching theory recommends sampling **proportionally** to leverage scores, which better preserves the full subspace structure [3].\n\nWe draw a direct analogy between this and our method: we treat the singular values (quantifying each singular vector’s contribution) of the gradient as analogous to leverage scores, and sample left singular vectors without replacement, **proportional** to their singular values. This avoids always selecting the same top-$k$ directions, like in GaLore and $Q$-GaLore, and instead promotes diversity across training steps. Our importance sampling method is not a heuristic adjustment made purely for empirical performance. Rather, it is a theoretically motivated approach rooted in randomized linear algebra and directly addresses the frozen subspace problem. While LoRA-related work focuses on augmenting low-rank structure, our work focuses on dynamically diversifying it during optimization, which is a crucial distinction. As far as we know, no prior work explores importance-sampled low-rank subspace selection in this way, particularly within the context of LLM pretraining.\n\nMoreover, SVD is a well-established and widely used technique in machine learning. Many methods—such as principal component analysis (PCA) [4] and attention approximation techniques [5]—rely on projecting data onto the top-$k$ left singular vectors of a data matrix to reduce dimensionality. Our randomized importance sampling strategy offers a novel perspective: instead of deterministically selecting the top directions, probabilistic selection may enhance subspace diversity and performance. We believe this insight could inspire future work in improving other applications that currently rely on fixed top-$k$ subspace projections and in theoretically and/or empirically studying the similarities and differences between our importance sampling and leverage score sampling.\n\n**Weakness 2**: \n\n**2a**. The experiments do not extend to larger models. \n\n**2b**. The evaluation is confined to perplexity. It critically lacks evaluation on a suite of downstream fine-tuning tasks, which is the standard for demonstrating the quality and generalizability of a pretrained model. \n\n**2c**. There is a crucial absence of any ablation or sensitivity analysis on the rank $r$. (a key hyperparameter).\n\n**Response to Weakness 2**: Thank you very much for your insightful comments. We currently report experimental results on models with 60M, 130M, 350M, and 1.1B parameters (Tables 1 and 2), demonstrating that SARA scales effectively across a range of model sizes. Due to limitations in time and computational resources, it is very challenging to present results on larger models at this stage. Even conducting experiments on the 1.1B model takes more than 10 days by using all of our GPUs. Nevertheless, we would like to emphasize that our theoretical convergence guarantee (Theorem 3.4) is independent of model size and thus supports general scalability.\n\nAlthough we lack the time and resources to run very large-scale experiments, we have fully utilized the rebuttal period to present the following experimental results that are computationally efficient and feasible within this limited timeframe. The following results can address your concerns in **2b** and **2c**.\n\n$\\bullet$ **Hyperparameter sensitivity**: Demonstrated stable performance across critical hyperparameters (see our **Response to Weakness and Question 2** to **Reviewer ZP7Q**).\n\n$\\bullet$ **Additional metrics beyond validation perplexity**: Confirmed enhanced downstream task performance (see our **Response to Weakness and Question 1** to **Reviewer ZP7Q**).\n\n$\\bullet$ **Comparison with uniform sampling**: Highlighted advantages of our sampling strategy (see our **Response to Weakness and Question 2** to **Reviewer yFMY**).\n\n$\\bullet$ **Multiple random seeds**: Verified the robustness and reliability of our method across different runs (see our **Response to Weakness and Question 2** to **Reviewer yFMY**).\n\n**Weakness 3**: The practical significance of the improvements is questionable in some cases. On the 130M model, SARA's perplexity gain over the Fira-Adam baseline is minimal (22.22 vs. 22.37). \n\n**Response to Weakness 3**: We thank you very much for pointing this out. We want to highlight that the key measure is not the raw perplexity improvement, but the reduction in the performance gap between the low-rank optimizer and full-rank Adam. In the 130M case, SARA reduces this gap by 42.25\\%, which is substantial given the tight margins typically observed in large-scale LLM pretraining. Also, SARA consistently improves performance across a wide range of model sizes (60M to 1.1B) and optimizer variants (Fira, GaLore, Adafactor, Adam-mini, 8-bit), as shown in Table 1 and Table 2. This consistency indicates that the improvements are not isolated noise but the result of meaningful enhancements to subspace exploration during optimization. \n\nAdditionally, we only pretrain LLaMA models with a certain number of tokens, so we can only observe relatively small performance gains in our paper. In fact, the performance gain of our method compared to GaLore increases as the language models are trained with more tokens. Therefore, once the language models are fully trained on the pretraining task, we expect to see much larger performance gains.\n\n**Weakness 4**: \"The paper claims SARA \"does not bring extra overhead\". This claim requires empirical validation with wall-clock time and memory usage comparisons.\n\n\n**Response to Weakness 4**: Thank you for pointing this out. Therefore, we conducted an additional experiment to address your concern. In the experiment, we measure the wall-clock time of SVD computation and sampling. We calculate the average wall-clock time of ten runs; the result for SVD is 0.34 seconds for a 2048 by 2048 matrix, and the result for sampling is 0.0005 seconds. The empirical result shows that the additional computational cost brought by sampling is negligible.\n\n**Weakness 5**: The paper's own theory suggests SARA's convergence rate is slightly worse than GoLore's by a constant factor, yet SARA consistently outperforms GoLore in experiments. \n\n**Response to Weakness 5**: We thank you very much for pointing out this important point. SARA picks singular vectors proportional to their singular values: this can not only ensure that we pick useful information (high singular values) but also address frozen dominant subspace phenomenon. GoLore, on the other hand, projects gradients onto a random low-rank subspace using unbiased random projections. This random projection cannot guarantee that they can pick so useful information compared to our SARA. Additionally, our theoretical guarantee considers the worst case. In practice, \"the gradient lies\npredominantly in the smaller, top subspace\" [2], and our importance sampling concentrates on informative directions, leading to a better effective convergence behavior than the worst-case suggests. We have included the above discussion to our paper. \n\nOnce again, we are deeply grateful for your insightful comments. We have addressed all identified weaknesses and questions, and we respectfully and genuinely hope you will consider raising your score.\n\n[1] Zhenyu Zhang, Ajay Jaiswal, Lu Yin, Shiwei Liu, Jiawei Zhao, Yuandong Tian, and Zhangyang Wang. \"Q-galore: Quantized galore with int4 projection and layer-adaptive low-rank gradients.\" Preprint'24.\n\n[2] Guy Gur-Ari, Daniel A. Roberts, and Ethan Dyer. \"Gradient descent happens in a tiny subspace.\" Preprint'18.\n\n[3] Michael W. Mahoney \"Randomized algorithms for matrices and data.\" Foundations and Trends in Machine Learning'11.\n\n[4] Jianqing Fan, Qiang Sun, Wen-Xin Zhou, and Ziwei Zhu. \"Principal component analysis for big data.\" Preprint'18.\n\n[5] Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. \"Linformer: Self-attention with linear complexity.\" Preprint'20."}]

Author rebuttal:

Meta-review: The paper gets a borderline rating. For the reviewers who gave negative rating, one didn't participate in the discussion actively. After taking a look at the authors' response, I think most of the reviewers' concern have been addressed reasonably. The AC agrees with Reviewer TBg6 that more experiments on larger models would make the paper stronger, but the AC also thinks that the paper already reaches good quality for publication. The AC read through the manuscript, all reviews, the discussion, and the rebuttal. The authors are highly encouraged to improve the paper quality according to the reviewers' feedback in the camera-ready version. The AC decided to accept this submission.
Decision: Accept (poster)