A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
Authors: Yiming Chen
Yuan Zhang
Yin Liu
Kun Yuan
Zaiwen Wen
Venue: ICML 2025 poster (2025)
Keywords: randomized subspace optimization
large language model training
stochastic optimization
Primary area: optimization->stochastic
TL;DR: We propose a randomized subspace optimization framework for LLM training that reduces both activation and optimizer memory costs while ensuring convergence and achieving performance on par with GaLore and Adam.

Abstract: The memory challenges associated with training Large Language Models (LLMs) have become a critical concern, particularly when using the Adam optimizer. To address this issue, numerous memory-efficient techniques have been proposed, with GaLore standing out as a notable example designed to reduce the memory footprint of optimizer states. However, these approaches do not alleviate the memory burden imposed by activations, rendering them unsuitable for scenarios involving long context sequences or large mini-batches. Moreover, their convergence properties are still not well-understood in the literature. In this work, we introduce a Randomized Subspace Optimization framework for pre-training and fine-tuning LLMs. Our approach decomposes the high-dimensional training problem into a series of lower-dimensional subproblems. At each iteration, a random subspace is selected, and the parameters within that subspace are optimized. This structured reduction in dimensionality allows our method to simultaneously reduce memory usage for both activations and optimizer states. We establish comprehensive convergence guarantees and derive rates for various scenarios, accommodating different optimization strategies to solve the subproblems. Extensive experiments validate the superior memory and communication efficiency of our method, achieving performance comparable to GaLore and Adam.

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
Yiming Chen 1
Yuan Zhang 2
Yin Liu
Kun Yuan
Zaiwen Wen
Abstract
1. Introduction
Large Language Models (LLMs) have achieved remarkable success across various domains (Achiam al., 2023; Brown, 2020; Dubey et al., 2024), primarily driven by the increasing scale of datasets and model parameters. The Adam opti- mizer (Kingma, 2014; Loshchilov & Hutter, 2019) is widely recognized as the default choice for training these models, owing to operation efficiency and robust performance.
The memory challenges associated with training Large Language Models (LLLMs) have become a critical concern, particularly when using the Adam optimizer. To address this issue, numer- memory-efficient techniques have been pro- posed, with GaLore standing out as notable ex- ample designed to reduce the memory footprint of optimizer states. However, these approaches do not alleviate the memory burden imposed by acti- vations, rendering them unsuitable for scenarios involving long context sequences or large mini- batches. Moreover, their convergence properties are still not well-understood in the literature. In this work. we introduce a Randomized Subspace Optimization framework for pre-training and fine- tuning LLMs. Our approach decomposes the high-ding -igh-dimensional training problem into series of lower-dimensional subproblems. At each it- eration, random subspace is selected, and the parameters within that subspace are optimized. This structured reduction in dimensionality allows our method to simultaneously reduce memory usage for both activations and optimizer states. We establish comprehensive convergence guar- antees and derive rates for various scenarios, ac- commodating different optimization strategies solve the subproblems. Extensive experiments validate the superior memory and communication efficiency of our method, achieving performance comparable to GaLore and Adam.
However, as the scale of LLMs continues to grow, the associ- ated memory demands have emerged significant bottle- neck. This challenge stems from the need to store optimizer states, such as first-order and second-order moments, along- side the activations required for gradient computations. For instance, training a LLaMA-7B model necessitates 28GB of memory to store optimizer states in FP16 precision (Zhao et al., 2024a), while a GPT-3 model with 175B parameters requires extraordinary 1.4TB memory in FP32 precision. Additionally, in scenarios involving long sequence lengths or large mini-batches, activation memory dominates as the primary constraint (Zhang et al., 2024b). These substan- tial memory requirements necessitate either deploying addi- tional GPUs or reducing batch sizes. However, increasing the number of GPUs introduces additional communication overhead, potentially limiting training scalability (Malladi et al., 2023), while smaller batch sizes prolong training time due to reduced throughput.
Memory-efficient training algorithms. Significant efforts have been made to address the memory overhead in LLMs training. One line of research focuses on parameter-effiient methods, such as Low-Rank Adaptation (LoRA) and its vari- ants (Hu et al., 2022; Lialin al., 2023; Xia al., 2024), which constrain trainable parameters to low-rank subspaces for each weight matrix. Similarly, sparsity-based techniques (Thangarasa et al., 2023) reduce memory usage by training only a subset weights. These strategies decrease the num- ber of trainable parameters, thereby reducing the memory requirements for storing gradients and optimizer states. An- other research direction aims to achieve memory savings through the compression of optimizer states. For instance, Gal Lore and its variants (Chen et al., 2024h; Hao et al., 2024; He et al., 2024; Zhao et al., 2024a) project gradients onto low-rank subspaces, leveraging the compressed gradients
'Equal contribution 'Bejjing International Center for Mathemat- ical Research, Peking University. Beijing, China Center for Data Science, Peking University. Beijing, China "Center for Machine Learning Research, Peking University, Beijing, China. Correspon- dence to: Kun Yuan <kunyuan @ pku.edu.cn>>
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
to compute the first- and second-order moments, which sig- nificantly reduces their memory footprint. Alternatively, Adam-mini (Zhang et al., 2024a) uses block-wise second- order moments for learning rate adjustments to reduce mem- ory redundancy. A recent study, Apollo (Zhu et al., 2024), reinterprets Adam as an adaptive learning rate algorithm applied to the gradient. Instead of the coordinate-we wise ap- proach used in Adam, it employs a column-wise adaptive learning rate, thereby effectively reducing the memory over- head associated with optimizer states.
a Randomized Subspace Optimization (RSO) frame- work for LLM training, which decomposes the original training problem into series of lower-dimensional subproblems. This decomposition simultaneously re- duces the memory required for optimizer states and activations, effectively addressing Limitation L1. Fur- thermore, the framework can reduce communication overhead in distributed training scenarios.
C2. Theoretical convergence guarantees. We provide a comprehensive convergence analysis for the RSO framework. The established guarantees and rates apply across various scenarios. These include subproblems solved using zeroth-order, first-order, second-order algorithms, well optimization methods like gra- dient descent, momentum gradient descent, adaptive gradient descent, and their stochastic variants. This addresses Limitation L.2. Notably, we present refined convergence guarantees for scenarios where subprob- lems are solved using the Adam optimizer.
Limitations in existing approaches. Despite the progress in memory-efficient algorithms for training LLMs, two crit- ical limitations persist in the aforementioned approaches:
LI. Inability to reduce activations. While the aforemen- tioned approaches effectively reduce memory associ- ated with optimizer states, they fail address the mem- ory burden posed by activations. This limitation stems from their reliance on computing full-rank gradients, which necessitates storing the complete activations. As a result, these methods are unsuitable for scenarios in- volving long context sequences or large
C3. Improved experimental performances. We conduct extensive experiments to evaluate the proposed RSO framework. The experimental results demonstrate that our approach significantly enhances memory efficiency compared to state-of-the-arr methods, such as GaLore and LoRA. Additionally, our method achieves faster training speeds by reducing communication overhead, outperforming both GaLore and Adam while maintain- ing comparable performance levels. These findings highlight the practical values of our approach.
L2. Insufficient convergence guarantees. While the afore- mentioned approaches demonstrate strong empirical performance, their theoretical convergence properties remain less understood. For instance, GaLore (Zhao et al., 2024a) provides convergence analysis only for fixed projection matrices, rather than for the period- ically updated projection matrices used in practical implementations. This lack of comprehensive theoret- ical guarantees raises concerns about whether these methods reliably converge to the desired solution and the rates at which such convergence occurs.
2. Related Works
Parameter-efficient methods. A promising approach to memory-efficiet training involves parameter-efficient meth- ods, which reduce the number of trainable parameters and consequently lower the memory required for storing op- timizer states. For example, (Hu et al., 2022) propose Low-Rank Adaptation (LoRA), which restricts trainable parameters a low-rank subspace for each weight matrix. Similarly, (Thangarasa et al., 2023) incorporate sparsity by training only subset of weights. While these meth- ods effectively reduce memory consumption, the reduction in trainable parameters can sometimes lead to suboptimal model performance (Biderman et al., 2024). To address this limitation, recent advancements suggest using multiple LoRA updates to enable high-rank weight updates (Lialin et al., 2023; Xia et al., 2024). However, in pre-training set- tings, this approach still relies on a full-rank weight training phase as a warm-up before transitioning to low-rank training (Lialin et al., 2023), thereby limiting its memory efficiency.
Main results and contributions. In this work, we pro- pose a method that concurrently reduces memory consump- tion for both the optimizer states and activations. The cen- tral idea behind our approach is to decompose the original nigh-ddmeensioal training problem into series of lower- dimensional subproblems. Specifically. at each iteration, we randomly select subspace and optimize the parameters within this subspace. After completing the optimization in one subspace, we switch to different subspace and con- tinue the process. Since each subproblem operates in a lower-dimensional space, it requires smaller gradients and optimizer states. As we will demonstrate, the reduced di- mensionality of the subproblems also leads to a significant reduction in the memory required for storing activations. Furthermore, the smaller scale of the subproblems results in reduced communication overhead when training across multiple workers. Our main contributions are as follows:
Optimizer-effiiennt methods. An alternative approach to memory savings focuses on compressing optimizer states
C1. Subspace method for LLM training. We introduce
2

---

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
while maintaining the number of trainable parameters. Ga- Lore (Zhao et al., 2024a) achieves this by compressing the gradient matrix through projection onto a subspace and leveraging the compressed gradient to compute first- and second-order moments. This projection reduces the gradient size and is typically derived via the Singular Value Decom- position (SVD) of the true gradient (Zhao et al., 2024a). To mitigate the computational cost of SVD, alternative methods have been proposed, such using random matrices (Hao et al., 2024; He et al., 2024) or generating the projection matrix through online Principal Component Analysis (PCA) (Liang et al., 2024). Fira (Chen et al., 2024a) and dDA (Robert et al., 2024) employ an error- error-feedback mechanism. The former combines the true gradient with the Gal.ore update to improve performance, while the latter explicitly accounts for both gradient and optimizer state compression. adaptive
System-based methods. Several system-level techniques have been proposed improve memory efficiency. Acti- vation checkpointing (Chen et al., 2016) reduces memory usage by recomputing activations on demand rather than storing them throughout the entire iteration, though this comes at the cost of increased computational complexity. Quantization (Dettmers et al., 2023) lowers memory con- sumption by using lower-bit bit data representations, but this may introduce a trade-off between memory efficiency and training precision. Additionally, methods such as those intro- duced by (Ren et al., 2021; Zhang et al., 2023a) reduce GPU memory usage by offloading data to non-GPU resources, which can lead to additional communication overhead.
3. Preliminaries
Apollo (Zhu et al., 2024) interprets Adam learning rate algorithm and uses compressed optimizer states directly as scaling factors for the true gradient. Additionally, Adafactor (Shazeer& & Stern, 2018) discards the first-order moment and approximates the second-order moment with two low-rank matrices, while Adam-mini (Zhang et al., 2024a) proposes that block-wise second-order moments are sufficient for adjusting learning rates. (Das, 2024) integrates the GaLore method with natural gradient optimizer to enhance performance. BAdam (Luo et al., 2024) and Block- LLM (Ramesh et al., 2024) incorporate block coordinate descent strategies into LLM training, restricting the number of parameters optimized in each epoch to reduce the mem- ory overhead associated with optimizer states. Meanwhile, (Wen et al., 2025) applies wavelet transforms compress gradients beyond the low-rank structures.
This section introduces the optimization framework for LLM pre-training and fine-tuning, followed by a review of several memory- methods.
3.1. LLM Optimization
When addressing the pre-training or fine-tuning of LLMs, the problem can be formulated as follows:
min f(W) = E5 (F), W
(I)
where W (We) represents the set of trainable pa- rameters with total dimension of d. Here, We E R denotes the weight matrix for the f-th layer, and L is the total number of layers. The function F(W;g) is the loss function, which depends on the random variable & represent- ing individual data samples.
Activation-efficcent methods. Although the aforemen- tioned methods effectively reduce memory consumption for optimizer states, they do not address the memory costs asso- ciated with activations. To reduce activations, zeroth-order (ZO) algorithms have been introduced in LLM training (Mal- ladi et al., 2023). These methods can be further improved through variance reduction techniques (Gautam et al., 2024), while (Zhao et al., 2024b) utilizes ZO approaches to approx- imate a natural gradient algorithm. Moreover, (Chen et al., 2024b) proposes novel ZO framework to enhance perfor- mance. Unlike oidst- (FO) methods, 70 algorithms ap- proximate gradients by finite differences in function values, eliminating the need for explicit gradient computation. This approach bypasses backpropagation and activation storage, significantly reducing memory demands. However, due to their slower convergence rates (Berahas et al., 2022; Duchi et al., 2015; Nesterov & Spokoiny, 2017), ZO methods are primarily suitable for fine-tuning applications. Similarly, FO methods can achieve activation savings by layer-wise training (Lai et al., 2024), but their use also predominantly targets fine-tuning phases.
To address the optimization problem defined in (1), com- monly used approaches include SGD (Bottou, 2010), Mo- mentum SGD (Sutskever et al., 2013), and Adam (Kingma, 2014). The iterative update rule for Adam is as follows:
M = 31 - + (1 - (1) VF(W :5",
(2a)
V = 32 vr-1 + (1 (1))))) (VF(W' :'))
(2b)
M' M2/(1 -1)(), 31) v,=V./11- 8)),
(2c)
W = W - a M /V/(VV/ V +c).
(2d)
Here, M and V represent the first-order and recond-order moments, respectively, and e >>> 0 is small constant.
3.2. Memory Consumption in LLM Training
The key memory components involved in the training pro- cess include four primary elements: model parameters, opti- mizer states, gradients, and activations. The model compo- nent stores parameters required for training. In the case of the Adam optimizer, the optimizer states are represented by the first and second moment estimates, denoted as M and
3

---

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
This leads notable memory savings compared to the Adam optimizer. However, shown in Figure the op- timizer state contributes significantly to memory costs pri- marily when using a small batch size. Conversely, with larger batch sizes-more practical in many scenarios-the memory efficiency advantages of GaLore diminish, as acti- vation memory becomes the dominant component of overall memory consumption.
LLaMA-1B
Model
OptStates
Gradient
2
O 49
Activations
4. Randomized Subspace Optimization
batchsize 1
batchsize = 64
In this section, we present the randomized subspace op- timization (RSO) method, tailored explicitly for the pre- training and fine-tuning of LLMs.
Figure 1. Memory components involved in training the LLaMA- IB model using the Adam optimizer under varying batch sizes. The reported values indicate memory usage in GB.
4.1. Algorithm Framework
V. The gradient corresponds to the memory cost associated with E) With the Adam optimizer, both the opti- mizer state and the gradient are determined by the number of trainable parameters, see recursions (2b)-
As previously discussed, the memory overhead in LLM training primarily stems from the large scale of the models. In other words, the primary source of memory consumption arises from the high dimensionality of the LLM training problem (1). This observation motivates us decompose the original problem into a series of lower-dimensional sub- problems. By partitioning the problem into smaller com- ponents, we can effectively reduce memory usage, as each subproblem requires less memory to process.
Another sigmificant memory cost arises from the activations, which represent the intermediate values computed during forward propagation. Unlike the optimizer state and gradi- ents, the memory for activations depends on multiple factors, including model size, batch size, and sequence length.
Figure I illustrates the memory consumption during the training of the LLaMA- 1B model. For small batch sizes, the optimizer state constitutes substantial portion of the memory usage. In contrast, for large batch sizes, activations dominate and account for nearly the entire memory cost.
Similar to the random coordinate descent (Wright, 2015), which optimizes the objective function one coordinate at time, address problem (1) incrementally, subspace by subspace. The proposed update rules are as follows:
D arg B W W
f((* + B)
1 21
B
J (4a)
3.3. Memory--fficient Method
B*
(4b)
As previously diseussed, the optimizer state imposes sub- stantial memory overhead. To address this challenge, Ga- Lore (Zhang et al., 2023b) introduces a projection technique that generates a compressed representation of the optimizer state, eliminating the need to store its full version. Conse- quently, the update rule of GaLore as follows:
Here, f(Wk + + = E2[F(W* + P=) in which pà = [P denotes the subspace projection matrices. randomly selected matrix with {Be} consist vari-
Each Ph E Rexr
Te & m.. The parameters ables with significantly smaller dimensions compared to W. Specifically, in the b-th layer, B& has dimensions ra nyg whereas We has dimensions m. ng. A proximal term BI == E. Br 3-2 is introduced to (4a) to ensure conver- gence, with coefficient r/* to regulate its influence.
M B1 -M -M- (1-33)) PT PTVF(W': ;t),,
(3a)
V = B2 ( (PVF
(3b)
M, (1/(11ft)) V.= V.((1
(3c)
In the k-th iteration, a subspace projection matrix pk is ran- domly selected, and the subproblem (4a) solved. This process approximately minimizes the objective function within the chosen subspace. Upon solving the subproblem, the current parameters are updated, and new subspace projection matrix, pk+1, is selected for the subsequent it- eration. When addressing the subproblem in (4a), standard optimizers such as GD, SGD, momentum SGD or Adam can be employed. Notably, obtaining an exact solution (4a) is not required; an inexact solution suffices for the proposed
W = W a ....MM///VV.
(3d)
Here, P represents the projection matrix, which maps the gradient matrix onto dimer-dimeniional subspace. Specifi- cally, GaLore selects P as the top left singular vectors of the gradient matrix, capturing its most important components.
Since the projected gradient P P PFFF(W g lies within a low-dimensional subspace, the associated optimizer states M and V in GaLore are also substantially reduced in size.

---

3: Solve subproblem (4a) and obtain the approximate solution 8 using given optimizer such as Adam. A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
approach. The RSO algorithm is presented in Algorithm 1.
Memory
Algorithm
Optimizer States
Activations
Algorithm 1 Randomized Subspace Optimization
Input: Initialization W
RSO
24nr
8bsn + 4bsr + 2bs²
Output: Solution W
Gal.ore
24nr
15bsm + 2hs²
LoRA
48nr
15bsn + 26s²
1: for == - do
Adam
24n²
15bsn + 2bs²
2: Sample P according given distribution.
Table 1. Memory analysis of different algorithms in terms of op- timizer states and activations for one typical transformer block. Here, b, and n represent the sequence length, batch size, and embedding dimension, respectively. The intermediate dimension of the feed-forward network is assumed to be 4n.
4: Update the weights by W W + B*
5: end for
4.2. Memory Efficiency
We now demonstrate that the proposed RSO approach offers superior memory efficiency. Unlike other memory-efficient methods (e.g., GaLore, Adam-min mini, Apollo, etc.) that pri- marily focus reducing the memory usage of optimizer states, the RSO method additionally achieves substantial savings in gradient and activation memory requirements.
In the backward--roppagatoon process, Adam and RSO com- putes the weight gradient as follows:
dy Yi
Oy
(Adam)
(7)
ow. oZe
Memory for optimizer states. When solving (4a), the reduced dimensionality of the subproblem significantly de- creases the memory requirements for optimizer states. For instance, the memory required for both the first-order and second-order moment estimates in each subproblem is rene parameters per é-th layer, which is substantially lower than the meny memory overhead in the standard Adam optimizer.
(RSO):
dy
(YePo)T
0y
(8)
aB:
0%.
Adam requires storing the activation Y, E Rexmm in (7) to compute gradients with respect to We. In contrast, RSO only needs store Y,P, E R to compute gradients with respect Be. Since Tę < mg, this approach achieves significant memory savings. As result, the RSO method substantially reduces the memory overhead associated with activations in layers of the form (6).
Memory for gradients. Specifically, for the subproblems in (4a), it is sufficient to compute the gradient with respect to B, i.e., VB F(Wk + pk B; E), rather than calculating the full-dimensional gradient V Ww F (WW;; {) with respect to the original weight matrix W, as outlined in the GaLore recursion in (3a)-(3b). This results in considerable memory savings associated with the gradient computation.
We analyze the memory overhead of the proposed RSO method for a typical transformer block. Table summarizes the results, comparing memory usage for optimizer states and activations across various algorithms. Details of this memory analysis is presented in Appendix A.
Memory for activations. The RSO method not only re- duces memory usage for gradients but also significantly minimizes the memory required store activations. For example, consider neural network where the f-th layer is defined as follows:
4.3. Communication Efficiency
As previously mentioned, the RSO algorithm solves smaller subproblems at each iteration, resulting in gradients with reduced dimensionality compared to methods like Adam and GaLore, which rely on full-dimensional gradients. This reduction in gradient size enables RSO to achieve improved communication efficiency.
(Adam): Z Ye- We,
y = L(Zz). (RSO): Ze Y (We PrBB), y L(Z)).
(5)
(6)
Expression (5) represents the forward process of Adam, where the f-th layer is associated with the weight matrix We, while (6) corresponds the RSO method associated with the weight matrix Bg. Here, Y: Xm denotes the output of the previous layer (i.e., the activation), and Z; serves as the input to the next layer. The function L(-), encompassing all subsequent layers and the loss function, depends only on Ze and not on Yg. Thus, once Ze is com- puted, Ye is no longer required for calculating the loss y.
Specifically, in data-parallel framework such as Dis- tributed Data Parallel (DDP), the model is replicated across multiple devices, with each device computing gradients its local data batch. These gradients are then aggre- gated across devices, necessitating gradient communication. By operating with lower-dimensional gradients, the RSO method effectively reduces communication overhead com- pared to existing approaches.
5

---

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
Subproblem Solver
Subproblem Complexity
Total Complexity
Zero-Order (ZO) Methods
Stochastic ZO Method (Shamir, 2013)
nerr))
0( p-1 nero)²
First-Order (FO) Methods
GD
O
Accelerated GD (Nesterov et al., 2018)
(logg=
SGD (Bottou et al., 2018)
O
Momentum SGD (Yuan et al., 2016)
(
0((3)
Adam-family (Guo et al., 2024)
O(e-))
0(1-)
Second-Order (SO) Methods
Newton's 's method (Boyd & Vandenberghe, 2004)
Olog((oge))))
Stochastic Quasi-Newton method (Byrd et al., 2016)
O(e
O
Table The sample complexities of the RSO method with various subproblem solvers. For the ZO solver, refers to the number of stochastic function value evaluations; for the FO solver, it refers to the number of detemmnistic/stochastti gradient computations; and for the SO solver, it refers to the number of determiinisicc/toohhsiii Hessian or estimated Hessian computations. Ó(-) hides logarithm terms.
5. Convergence Analysis
5.3, Pt should be drawn from a Haar distribution or con- structed as a random coordinate matrix (see (Kozak et al., 2023), Examples I and 2 for further details).
In this section, we present the convergence guarantees for the RSO method. To account for the use of various opti- mizers in solving the subproblem (4a), we assume that, at each iteration k; the chosen optimizer produces an expected e-inexact solution. Such an expected e-inexact solution is defined below:
The following theorem establish the convergence rate of the RSO algorithm. Detailed proofs are provided in Appendix B.
Theorem 5.5. Under Assumptions 5.2 and 5.3, let each subproblem in (4b) be solved starting from the initial point Bd to an expected e-inexact solution B& with suitable choice of n². The sequence (W*), generated by the RSO method satisfies the following bound:
Definition 5.1 (Expected e-inexact solution). A solution B* is said to be an expected e-inexact solution it satisfies:
Elg" (B*) g" (BB)) S e,
(9)
where
= (B)
+ p* B) + E B12 and B is
the optimal solution define as B == arg min 9" (B).
1
K-1
E B
18f.D. +18C, K
(10)
K
When of is properly chosen, it can be guaranteed that (BB is a strongly convex function hence B is unique.
k=
where Ao = s(W9) f" and : maxt (ms//me
To establish convergence guarantees for the RSO algorithm,
Sample complexity with different optimizers. When all subproblems are solved to expected e-inexact solutions, the RSO method achieves an e-stationary point within O(c-1) iterations. As each iteration requires solving the subprob- lem (4a), the total sample complexity of the RSO method depends on the solver employed for this subproblem. For instance, since gradient descent solves (4a) in Ologg")) inner iterations, the RSO method using gradient descent attains a total sample complexity of O( e -1). Table 2 summarizes the sample complexities for the RSO method when equipped with various solvers, including zeroth-order, first-order, and second-order scenarios with optimizers such as gradient descent, momentum gradient descent, adaptive gradient descent, and their stochastic variants.
we require the following assumptions:
Assumption 5.2. The objective function S(W) is L-smooth, i.e., it holds for any W and 2 that
[VS(W*)- VS(W3) & L Wi W
where [W]] = VE Well for any W (W Assumption 5.3. The random matrix P = (Pe) is sampled from a distribution such that P7 Pt = (me/ro)Ire and E[P.PT]] = Imz for each E.
Remark 5.4. In practice, when m. >> Te, sampling each Pt from normal distribution (0, 1) yields an approxima- tion P7 Pt z5 (me This approach provides computa- tional efficiency. However, to rigorously satisfy Assumption
Comparable complexity with vanilla Adam. It observed
6

---

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
Algorithm
60M
130M
350M
1B
Adam*
34.06 (0.22G)
25.08 (0.50G)
18.80 (1.37G)
15.56 (4.99G)
GaLore*
34.88 (0.14G)
25.36 (0.27G)
18.95 (0.49G)
15.64 (1.46G)
LoRA*
34.99 (0.16G)
33.92 (0.35G)
25.58 (0.69G)
19.21 (2.27G)
ReLoRA*
37.04 (0.16G)
29.37 (0.35G)
29.08 (0.69G)
18.33 (2.27G)
RSO
34.55(0.1444 14G)
25.34 (0.27G)
18.86 (0.49G)
15.86 (1.46G)
r r/dmodel
128 / 256
256 / 768
256 / 1024
512 / 2048
Training Tokens (B)
1.1
2.2
6.4
13.1
Table Comparison of validation perplexity and estimated memory usage for optimizer states across different algorithms during the pre-training of LLaMA models of various sizes on the C4 dataset. The optimizer states are stored in BF16 format. Results marked with * are sourced from (Zhao al., 2024a).
CoLA
STS-B
MRPC
RTE
SST2
MNLI
QNLI
QQP
Avg
Adam
62.24
90.92
91.30
79.42
94.57
87.18
92.33
92.28
86.28
GaLore (rank=4)
60.35
90.73
92.25
79.42
94.04
87.00
92.24
91.06
85.89
LoRA (rank=4)
61.38
90.57
91.07
78.70
92.89
86.82
92.18
91.29
85.61
RSO (rank=4)
62.47
90.62
92.25
78.70
94.84
86.67
92.29
90.94
86.10
GaLore (rank=8)
60.06
90.82
92.01
79.78
94.38
87.17
92.20
91.11
85,94
LoRA (rank=8)
61.83
90.80
91.90
79.06
93.46
86.94
92.25
91.22
85.93
RSO (rank=8)
64.62
90.71
93.56
79.42
95.18
86.96
92.44
91.26
86.77
Table 4. Evaluation of various fine-tuning methods the GLUE benchmark using the pre-trained RoBERTa Base model. The average score across all tasks is provided.
in Table 2 that RSO with Adam to solve subproblem has sample complexity O), which is on the same order as vanilla Adam (Kingma, 2014) without subspace projection.
using the C4 dataset, a large-scale, cleaned version of Com- mon Crawl's web corpus, which is primarily intended for pre-training language models and word representations (Raf- fel al., 2020). We compare our method against LoRA (Hu et al., 2022), GaLore (Zhao al., 2024a), ReLoRA (Lialin et al., 2023), and Adam (Kingma, 2014) base- line methods. We adopt the same configurations as those reported in (Zhao et al., 2024a), and the detailed settings for pre-training are provided in Appendix D.
6. Experiments
In this section, we present numerical experiments to evaluate the effectiveness of our RSO method. We assess its perfor- mance on both pre-training and fine-tuning tasks across models of varying scales. Additionally, we compare the memory usage and time cost of our RSO method with exist- ing approaches highlight its advantages in memory and communication efficiency. In all tests, the RSO method uses the Adam optimizer with fixed number of steps to solve each subproblem. The random projection matrices P= { are independently sampled from a normal distribution (O,, 1/re) for each layer. We set all re to the same value, which we define as the "rank" of our method to facilitate consistent comparison with other approaches that explicitly specify rank parameter.
Main results. As shown in Table 3, under the same rank constraints, RSO outperforms other memory-efficient meth- ods in most cases. We also report the estimated memory overhead associated with optimizer states. From Table 3, we observe that for LLLMA-350M RSO achieves nearly the same performance as Adam while reducing the memory re- quired for optimizer states by 64.2%. For LLaMA-1B, this reduction increases to 70.7%. The results for LLaMA-7B are provided in Appendix C.I.
6.2. Fine-tuning with RSO
6.1. Pre-training with RSO
Experimental setup. We extend the application of our RSO algorithm to fine-tuning tasks. Specifically, we fine- tune pre-trained RoBERTa models (Liu al., 2019) on the GLUE benchmark (Wang al., 2019), which encompasses diverse range of tasks, including question answering, sen-
Experimental setup. We evaluate the performance of our RSO method on LLaMA models with sizes ranging from 60M to 7B parameters. The experiments are conducted

---

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
LLMMA-35MM
LLaMA-18
LLaMA-78
80
Adam RSD
Adbm 85.0 Gal.ore
R5O
70
Momory (GB)
50
Memory [6
60-
40
Memory
30
20
20
10
0
64
512
256
Rank
Rank
Rank
Figure 2. Comparison of peak memory usage (in GB) per device for RSO and Gal during L.LMMA training with varying ranks. All hyperparameters, except rank, are consistent with (Zhao et al., 2024a). Adam's 's memory usage is reported for LLAMA-350M and LLaMA-ID IB but excluded for LLaMA-7B due to out-of out-of-memory memory (OOM) error.
LL (Seconds)
LLaMA-7B (Seconds)
Method
Seq 64
Seq 128
Seq 256
Seq 64
Seq 128
Seq 256
RSO
0.94
1.70
3.29
2.40
2.94
4.60
GaLore
1.12
1.84
3.35
7.86
8.26
9.12
Adam
1.11
1.81
3.32
7.84
8.23
OOM
Table 5. Comparison of iteration time (in seconds) for different methods in LLaMA training across various sequence lengths. All hyperparameters, except sequence length, follow (Zhao al., 2024a). IJaMA-IB on 4x A800 GPUs, while I.I.MA-7B 8x A800 GPUs. SVD decomposition time in GaLore is excluded. Additionally, for LLaMA-7B with a sequence length of 256, the Adam optimizer encounters an out-of- of-mmorry (OOM) error.
timent analysis, and semantic textual similarity. Detailed settings can be found in Appendix D.2.
GaLore widens as the rank decreases. This is because, as indicated Table I, RSO method further reduces mem- ory consumption for activations with lower ranks, whereas GaLore does not benefit in this regard. For LLaMA-350M and LLaMA-IB, GaLore's memory usage is observed to be comparable to that of Adam, as activation memory domi- nates in these cases. However, RSO still achieves superior memory efficiency due to its reduced activation cost.
Main results. As shown in Table 4, our RSO method sur- passes other memory-effice -ffiiieetf approaches and delivers per- formance comparable to Adam across most datasets in the GLUE benchmark. Notably, RSO significantly outperforms Adam on the CoLA and MRPC datasets when the rank is set to 8. Additional fine-tuning experiments on LLaMA and OPT models are provided in Appendix C.2.
RSO method requires less time per iteration. Table presents comparison of the time required for one itera- tion across different methods when training LLaMA models. As shown, the RSO method requires significantly less time compared to GaLore or Adam due to its improved commu- nication efficiency, achieved by reducing the dimensionality of gradients. For example, when training LLaMA-7B with sequence length of 64, the time required by RSO is only one-third that of GaLore or Adam. Notably, while GaLore involves SVD decomposition (which is excluded from this measurement), RSO demonstrates even greater efficiency in training time.
6.3. Memory and Communication Efficiency
To evaluate the memory and communication efficiency of our proposed method, we measure the peak memory usage and iteration time during the pre-training of LLaMA models of various sizes.
RSO method requires less memory overhead. Figure 2 illustrates the actual memory usage during the training of LLaMA models. As shown, the RSO method incurs significantly lower memory overhead compared to GaLore and Adam. This reduction attributed to RSO's ability save memory for activations and use low-dimensional gradients. For instance, in the case of LLaMA-7B with rank of 64, the RSO method achieves over a 40% reduction in memory overhead compared to GaLore.
As shown in Table 5, the difference iteration time between RSO and other approaches becomes more pronounced as model size increases or sequence length decreases. This phenomenon can be attributed the communication over- head, which primarily stems from the synchronization of gradients across devices. Such overhead is more closely tied to model size while being less affected by sequence length.
Additionally, in Figure 2, the memory gap between RSO and
8

---

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
In contrast, the computational overhead is highly sensitive sub-
References
sequence length. Therefore, RSO exhibits stantial advantage when the sequence length is considerably smaller than the model size, as communication overhead constitutes a larger fraction of the total iteration time under these conditions.
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L.,, Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. GPT-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
Berahas, A. S., Cao, L.,, Choromanski, K., and Scheinberg, K. A theoretical and empirical comparison of gradient approximations in derivative- optimization. Foun- dations of Computational Mathematics, 22(2):507-560, 2022.
7. Conclusion
We propose Randomized Subspace Optimization (RSO) method, aiming for Large Language Model (LLM) pre- training and fine-tuning, By decomposing the origanal train- ing problem into small subproblems, our method achieves both memory and communication efficiency, while reach same level performance compared with GaLore and Adam.
Biderman, D., Portes, J., Ortiz, J. G., Paul, M., Greengard, P., Jennings, C., King, D., Havens, S., Chiley, V., Frankle, J., et al. Lora learns less and forgets less. arXiv preprint arXhv:2405.0961 2024.
We outline two directions for future work on the RSO method. First, further reduction of memory overhead for activations could be explored. While part of the activation memory has already been reduced, the remaining portion might be further optimized through alternative strategies for partitioning the original problem. Second, it is worth investigating the performance of various methods for solv- ing the subproblems. Given the low-dimensional nature of the subproblems, exploring the application of second-order methods could be particularly promising.
Bottou, L. Large-scale machine learning with stochastic gradient descent. In Proceedings of COMPSTAT 2010: 19th International Conference on Computational Statis- tics, Paris France, August 22-27, 2010 Keynote, Invited and Contributed Papers, pp. 177-186. Springer, 2010.
Bottou, L.,, Curtis, F. E., and Nocedal, Optimization methods for large-scale machine learning. SIAM review, 60(2):223-311, 2018.
Impact Statement
Boyd, and Vandenberghe, L. Convex optimization. Cam- bridge university press, 2004.
This study focuses on improving the memory efficiency of large language model (LLM) training. By enabling LLM training devices with lower memory capacity, ap- proach helps reduce resource consumption and carbon emis- sions associated with LLM training. Furthermore, this en- hancement may make LLM training more accessible to smaller institutions and research teams operating under con- strained budgets.
Brown, T. B. Language models are few-shot learners. arXiv preprint arXiv:2005. 2020.
Byrd, R. H., Hansen, S. L., Nocedal, J., and Singer, Y. A stochastic quasi-newton method for large-scale optimiza- tion. SIAM Journal on Optimization, 26(2):1008-110 2016.
Acknowledgments
Chen, T., Xu. B., Zhang, C., and Guestrin, C. Training deep nets with sublinear memory cost. arXiv preprint arXiv: 1604.6614474 2016.
The computational resources were supported by the Center for Intelligent Computing and Song-Shan Lake HPC Center (SSL-HPC) in Great Bay University, Dongguan, China. This work was supported part by National Key Research and Development Program of China under the grant numbers 2024YFA1012902 and 2024YFA1012900 and the National Natural Science Foundation of China under the grant num- bers 1233 1010, 12288100 101, 12401408 and W2441021. Kun Yuan is also supported by AI for Science Institute, Beijing, China, and National Engineering Labratory for Big Data Analytics and Applications. We also thank the anonymous reviewers for their valuable feedback.
Chen, X., Feng, K., Li, C., Lai, X., Yue, X., Yuan, Y., and Wang. G. Fira: Can we achieve full-rank train- ing of llms under low-rank constraint? arXiv preprint arXiv:2410.01e6 2024a.
Chen, Y., Zhang, Y., Cao, L., Yuan, K., and Wen, Z. Enhanc- ing zeroth- order fine-tuning for language models with low-rank structures. arXiv preprint urXiv:2410.07600 2024b.
Das, A. Natural galore: Accelerating galore for memory- efficient IIm training and fine-tuning. arXiv preprint arXiv:2410. 16029, 2024.

---

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
Dettmers, T., Pagnoni, A., Holtzman, A., and Zettlemoyer, L. Qlora: Efficient finetuning of quantized L.Mss. Ad- vances in neural information processing systems, 36: 10088-10115, 2023.
Liang, K., Liu, B., Chen, L., and Liu, Q. Memory-efficient
LLM training with online subspace descent. arXiv preprint arXiv:2408. 12857, 2024.
Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoye,, L., and Stoyanov, V. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: /907. 1692, 2019.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The llama herd of models. arXiv preprint arXiv:2407.21788 2024.
Loshchilov, I. and Hutter, F. Decoupled weight decay reg- ularization. In International Conference on Learning Representations, 2019.
Duchi, J. Jordan, M. I., Wainwright, M. J., and Wibisono, A. Optimal rates for zero-order convex optimization: The power of two function evaluations. IEEE Transactions on Information Theory, 61(5):7788-288 2806, 2015.
Luo, Q. Yu, H., and Li, X. Badam: A memory efficient full parameter optimization method for large language models. Advances in Neural Information Processing Systems, 37: 2024.
Gautam, T., Park, Y., Zhou, H., Raman, P., and Ha, W. Variance-reducced zeroth-order methods for fine-tuning language models. In International Conference Ma- chine Learning, pp. 15180-15208. PMLR, 2024.
Malladi, S., Gao, T., Nichani, E., Damian, A., Lee, J. D., Chen, D., and Arora, Fine-tuning language models with just forward passes. Advances in Neural Information Processing Systems, 36:53038-53000 2023.
Guo, Z. Xu, Y., Yin, W., Jin, R., and Yang, T. Unified convergence analysis for adaptive optimization with mov- ing average estimator. arXiv preprint arXiv:2104.1488 2024.
Nesterov, Y. and Spokoiny, V. Random gradient-free mini- mization of convex functions. Foundations of Compulu- tional Mathematics, 17(2)5277-566, 2017.
Hao, Y., Cao, Y., and Mou, L. Flora: Low-rank adapters are secretly gradient compressors. In International Con- ference on Machine Learning, pp. 17554-17577 (1551..PII. PMLR, 2024.
Nesterov, Y. et al. Lectures on convex optimization, volume 137. Springer, 2018
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21 (140): 2020.
He, Y., P., Y., Chen, C., and Yuan, K. Subspace optimization for large language models with convergence guarantees. arXiv preprint arXiv:2410. 11289, 2024.
Hu, E. J., Wallis, P., Allen-Zhu, Z., Li, Y., Wang. S., Wang. L., Chen, W., et al. Lora: Low-rank adaptation of large language models. In International Conference on Learn- ing Representations, 2022.
Ramesh, A. V., Ganapathiraman, V., Laradji, I. H., and Schmidt, M. Blockllm: Memory-efficient adaptation of Ilms by selecting and optimizing the right coordinate blocks. arXiv preprint arXiv:2406. 1726,, 2024.
Ren, J., Rajbhandari, S., Aminabadi, R. Y., Ruwase, O., Yang, S., Zhang, M., Li, D., and He, Y. Zero-offload: Democratizing billion-scale model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21), pp. 551-564, 2021.
Kingma, D. P. Adam: A method for stochastic optimization. arXiv preprint 411.::142..6980, 2014.
Kozak, D., Molinari, C., Rosasco, I.., Tenorio, L., and Villa, S. Zeroth-order optimization with orthogonal random directions. Mathematical Programming, 199(1):1179- 1219, 2023.
Robert, T., Safaryan, M., Modoranu, I.-V., and Alistarh, D. LDadam: Adaptive optimization from -iweennnniooal gradient statistics. arXiv preprint arXiv:2410. 16103, 2024.
Lai, X., Tian, Z., Chen, Y., Li, Y., Yuan, Y., Liu, S., and Jia, J. Lisa: Reasoning segmentation via large language model. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition. pp. 9579- 9589, 2024.
Shamir, O. On the complexity of bandit and derivative- free stochastic convex optimization. In Conference on learning theory, pp. 3-24. PMLR, 2013.
Lialin, V., Muckatira, S., Shivagunde, N., and Rumshisky, A. Relora: High-rank training through low-rank updates. In The Twelfth International Conference on Learning Representations, 2023.
Shazeer, N. and Stern, M. Adafactor: Adaptive learning rates with sublinear memory cost. In International Con- ference on Machine Learning. pp. 4596-4604, PMLR, 2018.
10

---

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
Sutskever, I., Martens, J., Dahl, G., and Hinton, G. On the importance of initialization and momentum in deep learn- ing, In International conference machine learning, pp. 1139-1147, PMLR, 2013.
Zhao, J., Zhang, Z., Chen, B., Wang. Z., Anandkumar, A., and Tian, Y. Galore: Memory-efficient LLM training by gradient low-rank projection. In International Conference on Machine Learning, pp. 61 61121-6114 61143. PMLR, 2024a.
Thangarasa, V., Gupta, A., Marshall, W., Li, T., Leong, K., DeCoste, D., Lie, S., and Saxena, S. SPDF: Sparse pre- training and dense fine-tuning for large language models. In Uncertainty in Artificial Intelligence, pp. 2134-2146.
Zhao, Y., Dang, S., Ye, H., Dai, G., Qian, Y., and Tsang, 1. W. Second-order fine-tuning without pain for LLMs: Hessian informed zeroth-order optimizer. arXiv preprint W4Xiv:240..155 15773, 2024b.
PMLR, 2023.
Zhu, H., Zhang, Z., Cong, W., Liu, X., Park, S., Chandra, V., Long, B., Pan, D. Z., Wang, Z., and Lee, J. Apollo: SGD- like memory, adamw-level performance. arXiv preprint arXiv:4412.05270, 2024.
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Wang, A., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. R. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In International Conference Learning Representations, 2019.
Wen, Z., Luo, P., Wang, Deng. X., Zou, J., Yuan, K., Sun, T., and Li, D. Breaking memory limits: Gradi- ent wavelet transform enhances LLMs training. arXiv preprint arXiv:2501.072:3 2025.
Wright, S. J. Coordinate descent algorithms. Mathematical Programming, 3-1(1)33-34 2015.
Xia, W., Qin, C., and Hazan, E. Chain of lora: Efficient fine- tuning of language models via residual learning. arXiv preprint 3xXiv:2401.04155 2024.
Yuan, K., Ying, B., and Sayed, A. H. On the influence of momentum acceleration on online learning. Journal of Machine Learning Research, 17(192):1-66, 2016.
Zhang, H., Zhou, Y., Xue, Y., Liu, Y., and Huang, J. G10: Enabling an efficient unified gpu memory and storage architecture with smart tensor migrations. Proceedings of the 56th Annual IEEE/ACM International Symposium Microarchitecture, pp. 395-410, 2023a.
Zhang, Y., Chen, C., Li, Z., Ding, T., Wu, C., Ye, Y., Luo, Z.-Q., and Sun, R. Adam-mini: Use fewer learning rates to gain more. arXiv preprint arXiv:2406.1 (tf.:2406..6793,, 2024a.
Zhang, Y., Li, P., Hong, J., Li, J., Zhang, Y., Zheng, W., Chen, P.-Y., Lee, J. D., Yin, W., Hong, M., et al. Revis- iting zeroth-order optimization for memory-effiiint Ilm fine-tuning: A benchmark. In International Conference on Machine Learning, pp. 59173-59190. PMLR, 2024b.
Zhang, Z., Liu, B., and Shao, J. Fine-tuning happens in tiny subspaces: Exploring intrinsic aask-specific subspaces of pre-trained language models. The 6l Annual Meeting of the Association For Computational Linguistics, 2023b.

---

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
A. Memory Complexity Analysis
In this section, we analyze the memory overhead our proposed RSO algorithm for one typical transformer block.
A.1. Transformer Structure
Transformers (Vaswani, 2017) have become a foundational component of LLMss Here, we focus on the forward and backward propagation processes within a single transformer block using our RSO algorithm.
Forward Propagation. Consider the input R= to a transformer block, where is the sequence length and n the emhedding dimension. The attention mechanism within the transformer block performs the following linear operations:
Q = XW + P. = X + Bx),), V X(W, + P.B.,
(11)
where Wg. Wk. W, E R"** are the original weight matrices, Pq. Pk. E R"x* are the projection matrices, and Bg, Bg, Bp E R2*x are the low-rank weight matrices used in the RSO method for each subproblem. These intermediate values are then combined as follows:
Ag Vn
A, = QKT A
A A,V, A, = An + PoBo),
(12)
where represents the softmax activation function, and W. E Ruxn is the output projection matrix. Next, the feed-forward network consists of two fully-connected layers, which are computed as:
Ż AW + PiB.), Zi =Z=), Z2 (W P2122)
(13)
where W1 12x4n and W2 E 1 are the weights of the ferd-forward layers. We assume that the intermediate dimension of the feed-forward network is four times the embedding dimension. Similarly, Pi Rgar P2 R are the projection matrices, and Bi Brxdn B2 E Rr are the low-rank trainable parameters for RSO. The function represents the activation function.
For the Adam and GaLore methods, there are no and B matrices, as they directly work with the original weight matrices. In the case of the LoRA method, each projection matrix replaced by a trainable parameter A.
Backward Propagation. To calculate the gradients of all the weight matrices, the backward propagation begins with the partial gradient of the loss function F with respect the output of this block, denoted DZ2 D2, represent the derivative of F with respect to any matrix. Note that in RSO algorithm, compute the gradient with respect to B, the low-rank trainable parameters, instead of the original weight matrix W. The gradients for the weights in the feed-forward network are computed as follows:
OF Here, we D to
DB, (2, Dz2, Dž DZ2,,
D22 (Z,) DB, (A,P) D2, DA, (W1+ (14)
For the attention mechanism, the gradients for the corresponding matrices are calculated as:
DB, = (A) DA, = D DA.(W. + P.B.),, DA, = DA, T
(15)
To compute the gradients for the matrices Q, K, V. the following equations are used:
1 _o Å K, DK = DA. 2
]
DV = DJD,,, DQ = DA, の de
Q.
(16)
vn
The gradients for the low-rank weight matrices By, Be, B, are computed as follows:
DB, = (XP.)TDV, DB, (XP,) DQ, DB, = (XP.)
(17)
Finally, to ensure that the backward propagation process can continue, the derivative with respect the input X must also be calculated. This is given by:
DX DQ(W, + P,B,) + DK(W, P.Bc) DV(W, P.B.)
(18)
12

---

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
When using the Adam GaLore algorithms, the derivatives must be computed with respect to the original weight matrix W instead of the low-rank matrix B. As result, all occurrences of DB need to be replaced with DW. For example, the derivatives with respect to Wa, q, Wa, and W, are computed as follows:
DW, = x DV, DW, = = DW₂ = DK.
A.2. Memory for Optimizer States Analysis
For Adam algorithm, the trainable parameters include Wg. Wk. W., Wo, W1, W2. It is straightforward to compute the total number of parameters 12n Consequently, the optimizer states, considering both the first-order and second-order moments, require 24n² storage.
For RSO method, the trainable weights for each subproblem are By, Bg, BB, Bo, B1, B2, with total of 12nr parameters per subproblem, leading to an optimizer state storage requirement of 24nr. LoRA trains an additional matrix A (corresponding to the matrix used above), resulting in twice the optimizer state memory required compared to RSO. Gore projects each gradient matrix from "xx to grx" resulting in the same optimizer state memory requirement of 24nr.
A.3. Memory for Activations Analysis
From the backward propagation process, it evident that the activations generated during forward propagation are required. Specifically, in (1I), the matrices X, Q, K, V need to be stored, resulting in the following memory requirement: M1 = 4sn.
However, in our RSO algorithm, X is only needed to compute DBg DB2 DB., where only the projections XP, XPp, q X PR are required. By setting Pg Pk P. P, only need store XP, reducing the memory requirement to M1 = 3sn + sr.
Additionally, in (12), the matrices A,, A,, An, A. must be stored, requiring M2 = 2s² + 2sn. It is worth noting that A_ yn does not need to be stored, A, can be used to recover it due to the properties of the softmax function. For the RSO algorithm, storing An and A, is unnecessary, as P, and AgPi suffice. Consequently, the memory requirement is reduced to M2 2s + 2sr.
For the feed-forward network in (13), the matrices 2,, Z1, Zz need to be stored, resulting a memory requirement of M3 9sn. the RSO algorithm, Z1 can be replaced with Z₁ P2, reducing the memory requirement to M3 = 5sn + 8r.
Combining these results, the total memory cost for activations in the RSO algorithm is
Musal = 8sn + 2s² + 4sr,
compared to the memory cost in Adam or GaLore:
Mosal 15sn + 2s².
As the LoRA method trains parameters A (corresponding to P in method), it requires the activations as Adam, resulting in the same memory overhead as Adam or Gal.ore.
B. Convergence Analysis
In this section, we present the convergence analysis of the RSO algorithm and provide a detailed proof of Theorem 5.5. Under Assumption 5.3, the following properties hold and can be straightforwarddy derived:
tr(w Pf S /max(mm//rNN/111
where W = [We) denotes any family of matrices with We E IR For simplicity, we will not explicitly reference these properties when they are used.
Lemma B.1. Under Assumptions 5.2 and 5.3. (B) is G -)-strongly convex and (7 + L--smooth, with 0 < < 1/2 and L max{me//re]
13

---

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
Proof. We denote the gradient of /with respect to the parameters in the ê-th layer by VeJ(W). Let hB) s(W*+ pk B). It can be shown that ht is L--Lipschitz smooth, as follows:
Of ( + P!) d
0f
Pp* BP3 )
(
W
d
Σ
(((P)] - pk B²)
I -v
Vef(W" + P*B") - Vef(W* + +P*
S max
rt
(m Tz.
IV +P'BB VS(W* B3)
max
< L max
me
||P B' p B² < B' B²1
Te
As htt is Í--Lipschizz smooth, we can conclude that g²(B) = h²* + have the inequality
1 2xy²
B1| is (,)= + L) smooth. Furthermore, we
i
B
- アリ
B² B²
Based on this inequality, with the definition of g²² we have
1
(V -B" i
BB² 2 g 21
12
1 +
B 2
2Tp
i
= (B) + vha (Bi) (B1) + 17
Br. B2 -B:)
+
| -B21
27/h
2
(1 L
= (Vg (Bi) B2 - Bi) +
[B| -B2l3
21y
2
which shows that g² (B) is G+ - L)-stronngl convex when 0 < < 1/2.
Theorem B.2 (Theorem 5.5). Under Assumptions 5.2 and 5.3, let the subproblem (4b) is solved from the initial point
Bd =0
0 to an expected e-inexact solution B& with T/² 2L the sequence {W*}, generated by the RSO method satisfies
K-1
I3 18 - + 18Lc.
(19)
K
K
Proof. For simplicity, denote = # L. As " (B) is μ strongly convex, V B, we have
get B < g B - -
Let B- = O and using the definition of h& we can obtain a descent condition as
< (B')
1
Taking the expectation with respect to the initial condition W and random matrix P by the tower rule, it can be derived
Elg (0) - " 11
2
14

---

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
Telescoping the above inequality from k- to K we have
K
K
K-1 Σ k=
Bla
Elg
(B*)
tc.
2
R=
k=
the terms in the RHS of the above inequality can be canceled with each other and resulting in the following inequality:
Notice that, by the update rule of (4b), 9 ++ (0) (W f(W*
P
B²) -
2
|B
2n
K-1 Σ
K-2
IP S (0) - Elg k (B 1-
Σ
El
B+1
B
117
27t
k=
2 0
Dividing K on both sides and using the fact - (BK-1)
f + p -1 B 1) ス f*, we can derive
f(W -1
K-1 K
K-2
El | v g 8
K + c.
2
Next, we need to establish the connection between B || and the final stationary measure [// As (B) (L/1/* + --mmmohhh, AB it holds that
2 ( m
(
IIB B+1 -
With B = 0, it holds
1 K
-2
f(W)) f K
+
(20)
K
2
1p
k=0
Furthermore, notice that (g) = (pt) fixed W have
T T (WW and the random matrix Pk is sampled independently to W for any
E Ep
V * [0>1]3] = Σ B (VVVV
(ve((W*)))
Σ Epp (.,
Σ p
Σ [I
Taking expectation with the respect the randomness in W we can claim E|| (0)) ][] Inserting it can be written
this result back to (20) with 1/"
27.
K
18L(](WO)) 1*) (")
[[]
K + 18Lc.
K 0=
Thus complete the proof.
C. More Experimental Results
C.1. Pre-training on LLaMA-7B Model
Table 6 compares the performance of our RSO method with GaLore and Adam on the LLaMA-7B model, where evaluations are conducted for 50K steps due to limited computational resources. We also report the memory overhead and total training time for each method. As shown in the table, RSO exhibits performance comparable to that GaLore and Adam. Notably, RSO requires less than half the training time of the other methods.
15

---

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
Method
Memory (GB)
Training Time (h)
Perplexity
Adam
78.92
216
15.43
GaLore
75.33
134
15.59
RSO
54.81
64
15.99
Table Comparison of various pre-training methods for the LLaMA-7B model on the C4 dataset. Perplexity is reported at 50K steps. The training is conducted on 8x A800 GPUs. The actual memory cost per device and the total training time also reported. RSO and Galore are configured with a batch size of 16, while Adam uses batch size of 8.
C.2. Fine-tuning on LLaMA and OPT Models
Table 7 compares RSO with other fine-tuning methods on LLaMA and OPT models across two datasets. As shown in the table, RSO outperforms other memory-efficient methods terms of accuracy on most tasks.
WinoGrande
Copa
Model
Adam
LoRA
GaLore
RSO
Adam
LoRA
GaLore
RSO
LLaMA-7B
64.4
70,9
70.9
71.0
84.0
84.0
85.0
86.0
LLaMA-13B
73.3
76.6
74.6
74.7
90.0
92.0
92.0
92.0
1.3-B
60.4
57.3
58.3
58.9
76.0
73.0
72.0
74.0
OPT-6.7B
62.2
64.7
66.8
69.2
78.0
80.0
80.0
82.0
Table Comparison of various methods for fine-tuning LLaMA and OPT models on the WinoGrande and COPA datasets. The test accuracy for each method is reported.
Params
Hidden
Intermediate
Heads
Layers
Steps
Data amount
60M
512
1376
8
8
10K
1.3 B
130M
768
2048
12
12
20K
2.6 B
350M
1024
2736
16
24
60K
7.8 B
1 B
2048
5461
24
32
100K
13.1 B
7 B
4096
11008
32
32
150K
19.7 B
Table Hyperparameter configurations LLaMA models different scales, along with corresponding number training steps. Due to limited computational resources, only the first 50K steps are completed for LLaMA-7B.
D. Experimental Details
D.I. Pre-training Experimental Setup
For the pre-training of LLaMA models across all scales, adopt configuration consistent with that used in (Zhao et al., 2024a). The main model hyperparameters and training steps for each method are summarized in Table 8. Specifically, in the RSO method, the Adam optimizer used to perform multiple steps for solving each subproblem, resulting in an outer-inner iteration structure: the outer iterations correspond to subproblem updates, while the inner iterations represent Adam steps. For fair comparison, we report the total number of inner iterations as number RSO steps. experiments, the maximum sequence length
to 256, and the total training batch size is fixed at 512, corresponding approximately 131K tokens per batch. A linear warm-up of the learning rate is applied over the first 10% of training steps, followed cosine annealing schedule that decays the learning rate to 10% of its initial value.
For the RSO method, the learning rate is selected from the set (0.05, 0.02, 0.013. Consistent with the configuration used in GaLore, a learning rate scaling factor is applied to the weights all multi-head attention and feed-forward layers in the model. The number of Adam steps used to solve each subproblem is set to either 200 or 500 depending on the specific experiment.
16

---

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
Task
MNLI
SST-2
MRPC
CoLA
QNLI
QQP
RTE
STS-B
Batch Size
16
16
16
32
16
16
16
Epochs
30
30
30
30
30
30
30
30
Learning Rate (Rank 4)
1E-05
3E-05
3E-05
3E-05
1B-05
1E-05
1E-05
1E-05
Learning Rate (Rank= = 8)
1E-05
2E-05
2E-05
IB-05
IE-05
2B-05
2E-05
3E-05
Scaling Factor
{8, 16, 32}
Steps per Subproblem
{ [300, 500}
Max Sequence Length
512
Table Hyperparameter settings for fine-tuning the RoBERTa-Base model the GLUE benchmark using the RSO method with different rank configurations.
D.2. Fine-tuning Experimental Setup
Fine-tuning on the GLUE Benchmark. To fine-tume the pre-trained RoBERTa- Base model on the GLUE benchmark, we train for 30 epochs using batch size of 16 across all tasks, except for CoLA, which uses batch size of 32. Consistent with the GaLore setting, a learning rate scaling factor is applied to the weights of all multi-head attention and feed-forward layers. Detailed hyperparameter configurations are listed in Table 9.
Fine-tuning on the WinoGrande and COPA Datasets. For fine-tuning LLaMA and OPT models on the WinoGrande and COPA datasets, randomly sample 1,000 training examples, 500 validation examples, and 1,000 test examples from each dataset. All experiments are run for 1,000 training steps. The corresponding hyperparameter settings are summarized in Table 10.
Experiment
Hyperparameters
Values
Batch Size
16
FT
Learning Rate
(IE-07, IE-06, 1E-05)
Weight Decay
0
Batch Size
16
Learning Rate
(1B-07, 1E-06, 5E-06)
LoRA
Rank
8
Weight Decay
0
Batch Size
16
Learning Rate
(1E-07, 1E-06, SE-06)
GaLore
Rank
8
SVD Update Interval
(300, 500}
Scaling Factor
{4, 8} 0
Weight Decay
Batch Size
16
Learning Rate
(1E-07, 1E-06, 5E-06)
Rank
8
RSO
Steps per Subproblem
(300, 500}
Scaling Factor
(4, 8}
Weight Decay
0
Table 10. Hyperparameter settings for fine-tuning LLaMA and OPT models the WinoGrande and COPA datasets.
17


Reviews:
[{"review_id": "N0fKUTdORd", "reviewer": "Reviewer_3u9E", "summary": "The paper introduces a memory-efficient optimization method named RSO, which optimizes subspace variables rather than the original weight. Different from the conventional subspace method such as Galore, which projects the original weight's gradient into subspace, they propose to directly optimizes the subspace variable. Hence, they can reduce the activation memory by a large margin. Such a method also avoids heavy computation of SVD as in Galore. Experiments on pre-training Llama shows that RSO achieves similar convergence speed as Galore with reduced memory and time cost.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our manuscript. We greatly appreciate your valuable feedback. Below, we provide our responses to your comments:\n- Experimental Designs Or Analyses:\n\n   In Table 5, we set the rank to one-fourth of the embedding dimension for both the LLaMA-1B and LLaMA-7B models, consistent with the perplexity comparison experiments and the setting used in the GaLore paper. Each matrix $P$ is constructed as a random Gaussian matrix in our experiments. We will include these implementation details in the next revision of the manuscript.\n\n- Essential References Not Discussed:\n\n   Thank you for the suggestion. These two works introduce block coordinate descent methods into LLM training. BAdam [1] divides the model parameters into several blocks and updates them sequentially, with one block optimized per epoch using Adam. BlockLLM [2] partitions the parameters by layers and updates those with larger gradient norms in each outer loop. It also incorporates a pruning mechanism to further reduce the number of parameters being updated. By limiting the number of parameters optimized in each epoch, both methods effectively reduce the memory required for storing optimizer states, as well as the computational cost associated with backpropagation.\n\n   We will incorporate these two works and discuss them in the related work section in the next revision of our manuscript.\n\n- Weaknesses:\n\n   In mixed-precision training, activations are stored in FP16, while model parameters and optimizer states are stored in FP32. Compared to LoRA, RSO incurs slightly lower memory usage for optimizer states, as LoRA trains two matrices per layer instead of one. On the other hand, RSO introduces additional memory overhead due to storing the full model in FP32. However, as illustrated in Figure 1 of the manuscript, when training large models with practical batch sizes, activation memory dominates the overall memory consumption, accounting for nearly all of the total usage. As a result, even with the extra memory required for model storage, the activation efficiency of RSO still leads to overall memory savings compared to LoRA.\n\n   Moreover, although LoRA avoids storing the full model in FP32, it significantly reduces the number of trainable parameters, which may lead to performance degradation in pre-training tasks.\n\n   We compare the actual memory consumption of RSO and LoRA during the pre-training of LLaMA-1B and LLaMA-7B under mixed-precision training. To avoid out-of-memory (OOM) errors, we set the rank-to-embedding-dimension ratio to $1/4$ for LLaMA-1B and $1/8$ for LLaMA-7B. As shown in the following table, the results confirm the memory efficiency of our RSO method even in the mixed-precision setting.\n\n   | Method | LLaMA-1B (GB) | LLaMA-7B (GB) |\n   |--------|---------------|---------------|\n   | RSO    | 54.73         | 68.39         |\n   | LoRA   | 69.93         | 78.82         |\n\n- Questions:\n\n   1. In our experiments, we use the Adam optimizer to solve each subproblem, and the matrix $B$ is initialized as a zero matrix. This ensures that the objective value of the next subproblem is initialized at the function value evaluated at the solution of the previous subproblem, which helps maintain a more stable training trajectory.\n\n   2. Each matrix $P$ is constructed as a random Gaussian matrix in our experiments. In the memory analysis, we assume a shared $P$ for the query, key, and value weight matrices to achieve optimal memory efficiency. Without sharing $P$, one would need to store $XP_q$, $XP_k$, and $XP_v$ separately instead of a single $XP$, resulting in an additional memory cost of $2sr$, which remains relatively minor.\n\nThank you again for your valuable feedback, and we hope our responses address your concerns. If you have further questions, we would be happy to provide additional clarification.", "claims_and_evidence": "Yes. The memory cost and convergence are carefully analyzed in this paper.", "methods_and_evaluation_criteria": "Yes, the C4 pre-training is a standard benchmark for the considered setting.", "theoretical_claims": "Yes. I went through the key proof steps of the main theorem.", "experimental_designs_or_analyses": "Yes. The experimental design is fair for comparing the memory/time cost of RSO and the baseline methods. However, the RSO rank for Table 5 experiments is missing. The strategy of constructing P matrices are not explicitly discussed in the paper.", "relation_to_broader_scientific_literature": "The key contribution of reducing the activation memory cost is of general interest for LLM research community as well as practitioners.", "essential_references_not_discussed": "For the literature review of optimizer-efficient methods, it is suggested to discuss block coordinate descent methods, e.g. [1, 2], which share similar time efficiency feature as RSO by reducing gradient computational cost.\n\n[1] BAdam: A Memory Efficient Full Parameter Optimization Method for Large Language Models\n\n[2] BlockLLM: Memory-Efficient Adaptation of LLMs by Selecting and Optimizing the Right Coordinate Blocks", "other_strengths_and_weaknesses": "**Strengths**\n\n* The proposed method reduces the memory cost of storing activation through algorithmic-level design, which is novel and may inspire development of new activation-efficient methods. Given that the activation memory is usually the main bottleneck for long-sequence training, this contribution is critical.\n\n* Since the algorithm only computes the gradient of the subspace variable, the gradient computation is much cheaper than full backward. Hence, the time efficiency of RSO is much better than the Adam and other baseline approaches.\n\n* The algorithm exhibits competitive performance in C4 pre-training and GLUE benchmark.\n\n* The sample complexity result is established for RSO under strongly convex and smoothness assumption.\n\n**Weaknesses**\n\n* The method seems not very memory-efficient under mixed precision training setting compared to LoRA, whose master fp32 copy, gradient, and optimizer states are all low-rank matrices.  RSO cannot avoid storing the full master copy as the update will be merged into the weight after (approximately) solving each sub-problem.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "* What is the initialization of B matrix? The initialization may have huge influence on the convergence.\n\n* The construction of P matrix is not explicitly discussed. In memory analysis, P is shared for query, key, and value. Is there any reason for doing this or it is mainly used for memory efficieny?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TFrBB6oUmu", "reviewer": "Reviewer_FKkJ", "summary": "The paper introduces a new optimization method designed to reduce communication costs in distributed training of large language models (LLMs). The method proposed by the authors is based on breaking up the problem into smaller low-rank subproblems that we can optimize using arbitrary optimizers, by introducing a new update equation to the weights. The paper proposes optimizing over a low rank matrix passed through a random projection, which reduces the memory requirement for the activations and implicitly the gradients. The authors also provide guarantees for the number of steps required to find an approximate stationary point under mild assumptions, as well as an empirical evaluation of their method on models at various scales and across datasets, and a comparison with existing methods such as GaLore and LoRA.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our manuscript. We greatly appreciate your valuable feedback. Below, we provide our responses to your comments:\n\n- Weaknesses:\n\n   - In the RSO method, a series of subproblems need to be solved. However, this does not introduce additional complexity in tuning hyperparameters. In our experiments, for each task, we employ the Adam optimizer with the same hyperparameters across all subproblems, and apply a unified learning rate schedule rather than assigning separate schedules to individual subproblems.\n\n   - We will address the FLOPs-related discussion in the response to the questions section.\n\n- Other Comments or Suggestions:\n\n   Thank you for the suggestion. As discussed in Sections 4.3 and 6.3, RSO can also improve iteration time by reducing communication overhead when training across multiple devices, as shown in Table 5. We will move this discussion earlier in the manuscript and place greater emphasis on this point in the next revision.\n\n- Questions:\n\n   1. In our experiments, we use the Adam optimizer to perform multiple steps to solve each subproblem in RSO, resulting in an outer–inner iteration structure: outer iterations correspond to different subproblems, while inner iterations correspond to Adam steps. For fair comparison, we count the total number of inner iterations as the number of RSO steps; that is, solving one subproblem with $\\tau$ Adam steps is counted as $\\tau$ steps. Under this definition, the per-step computational cost of RSO is comparable to that of Adam and lower than GaLore, which incurs additional overhead due to costly SVD operations. Since all methods are evaluated under the same number of steps, we believe the computational complexity is well controlled and the comparison is fair.\n\n      In addition, RSO significantly reduces communication overhead (see Section 4.3). Table 5 compares per-step iteration time (for RSO, per inner iteration), further supporting our claim.\n\n      We apologize for any confusion and will clarify this point in the next revision.\n\n   2. The choice of $\\eta_k$ is discussed in Appendix B; see Lemma B.1 and Theorem B.2 for more details, where we restrict the range of $\\eta_k$ to facilitate the convergence analysis. In our experiments, we choose a relatively large constant for $\\eta_k$ (e.g., 10) to accelerate convergence in practice. The results show that RSO still converges well under this setting.\n\n   3. We conduct additional experiments on LLaMA-60M, where RSO is evaluated using the Shampoo optimizer for solving each subproblem. The results are presented in the following table.\n\n      |Method|Perplexity|\n      |-|-|\n      | RSO-Adam   |34.55|\n      | RSO-Shampoo|36.45|\n\n   4. In the $k$-th outer iteration, we solve the subproblem\n   $$\\min_B f(W^k + P^k B),$$\n   using a given optimizer to obtain an approximate solution $B^k$, rather than performing a single gradient step, where $P^k$ is a random projection matrix.\n\n      Since choosing $B = 0$ recovers $W^k$, it follows that $f(W^k + P^k B^k) \\leq f(W^k)$. Updating the weight as $W^{k+1} = W^k + P^k B^k$ thus ensures\n      $$f(W^{k+1}) \\leq f(W^k).$$\n\n      This demonstrates that solving each subproblem in a new random subspace does not compromise the progress made in previous iterations. On the contrary, each update builds upon the last, despite the randomness in the projection. As a result, errors do not accumulate across outer iterations; instead, the optimization process consistently advances in a descent direction in expectation.\n\n   5. We track both training and evaluation loss during the training of LLaMA-60M and LLaMA-130M, under configurations using 200 and 400 Adam steps per subproblem. Results are available at:\n\n      https://anonymous.4open.science/r/RSO-1CC6/\n\n      As shown, the evaluation loss decreases consistently. While the training loss may briefly fluctuate at the start of each subproblem, it also follows a downward trend overall.\n\n      We further perform an ablation study on LLaMA-60M by varying the projection rank and the number of inner Adam steps per subproblem. Results are summarized below. Increasing the rank generally improves performance, though at the cost of higher memory usage. To balance performance and efficiency, we choose rank 128—matching GaLore—for fair comparison. For the number of inner steps, moderate values (e.g., 200) perform well, whereas excessively large values may degrade performance.\n\n      |Rank|64|128|192||Inner Steps|100|200|400|800|\n      |-|-|-|-|-|-|-|-|-|-|\n      |Perplexity|37.60|34.55|33.25||Perplexity|34.54|34.55|34.80|35.30|\n\n      Finally, we adopt a unified learning rate schedule across the entire training process, rather than resetting it for each subproblem. This schedule decays adaptively, requiring no manual tuning between subproblems.\n\nWe hope these responses adequately address your concerns. If you have further questions or need additional clarifications, we would be happy to provide them.", "claims_and_evidence": "The paper provides sufficient evidence that their proposed method can achieve improvements in optimization speed and memory reduction. Concretely, while the performance of RSO is similar to that of GaLore - as shown in Table 3 - for pretraining, or marginally better compared to GaLore - as shown in Table 4 - for finetuning, GaLore does not target reducing communication costs as it still requires all-gathering the whole gradient across devices. RSO on the other hand, in Table 5, shows that it can achieve much faster iteration time than Adam or GaLore during optimization.", "methods_and_evaluation_criteria": "The authors sweep over hyperparameters across all baseline methods as shown in Appendix D. Due to computational constraints the authors provide training runs up to 1B, but they run fine-tuning on several datasets and larger scales than pre-training.", "theoretical_claims": "While I have not checked the proofs of their theorems regarding guarantees, I do have several questions regarding the method and the claims that I will post in the Questions for Authors section.", "experimental_designs_or_analyses": "See methods and evaluation criteria.", "relation_to_broader_scientific_literature": "The method is in general related to other low rank training works such as: [1], [2], [3], [4].\n\n[1] uses a similar low rank approximation of the weight updates as RSO in order to minimize memory cost. [2] uses only the top subspace of the gradients in the optimization procedure, which is also a memory reduction technique. [3] uses a similar idea for the momentum term. [4] does not base its compression on rank, but on the top frequencies of the momentum term. \n\n[1]: Hu, Edward J., et al. \"Lora: Low-rank adaptation of large language models.\" ICLR 1.2 (2022): 3.\n[2]: Zhao, Jiawei, et al. \"Galore: Memory-efficient llm training by gradient low-rank projection.\" arXiv preprint arXiv:2403.03507 (2024).\n[3]: Vyas, Nikhil, Depen Morwani, and Sham M. Kakade. \"Adamem: Memory efficient momentum for adafactor.\" 2nd Workshop on Advancing Neural Network Training: Computational Efficiency, Scalability, and Resource Optimization (WANT@ ICML 2024). 2024.\n[4]: Peng, Bowen, Jeffrey Quesnelle, and Diederik P. Kingma. \"Decoupled Momentum Optimization.\" arXiv preprint arXiv:2411.19870 (2024).", "essential_references_not_discussed": "None, see above.", "other_strengths_and_weaknesses": "Strengths:\n* The paper proposes a method to reduce memory usage an in turn communication cost in distributed settings\nRSO achieves faster iteration time as it does not need to communicate the full gradient at each step, without sacrificing performance (when compared to the existing works such as GaLore)\n* The authors perform a thorough empirical evaluation and show convergence guarantees for their method\n\nWeaknesses:\n* The method adds another layer of complexity by solving an internal suboptimization problem, which leads to additional hyperparameters to be tuned such as the optimizer choice for this problem\n* The paper does not provide a FLOPs controlled experiment, see questions", "other_comments_or_suggestions": "In general I think the paper could benefit more from a rewriting. Currently, the most important finding regarding RSO which is that even though it achieves similar performance to previous works in the literature while reducing drastically the time per iteration appears very low in the paper - Table 5. Moving this section up would greatly benefit this work.", "questions_for_authors": "* When controlling for compute, how does RSO stand when compared to the other methods? Concretely, it would seemingly appear that because this method requires optimizing a subproblem for each of the K steps, this would incur a lot more FLOPs when compared to Adam/Galore etc. Could the authors provide an experiment in this direction?\n* The authors mention that “when \\eta_k is properly chosen” and “with suitable choice of \\eta_k”, but it is not clear to me what this choice of \\eta_k should be and how it should be chosen in practice. Could the authors be more specific here? \n* If it is not a computational constraint, could the authors also try optimizing the subproblems with a nondiagonal preconditioner such as SOAP/Shampoo? As this requires quite a few extra runs, I will not base my decision on this experiment, but it is more of a scientific curiosity\n* Why does optimizing in a random subspace given by a random rotation matrix not accumulate errors over time? More concretely, why is it the case that at each time step k we can optimize via a random projection matrix and this does not break the progress we have made so far in optimization in the previous steps?\n* Could the authors provide plots of their training losses? How stable is the proposed method to hyperparameter choices? Does the learning rate change over time for the subproblem optimization?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "M43kgYnNbz", "reviewer": "Reviewer_75Nn", "summary": "The paper introduces a Randomized Subspace Optimization (RSO) framework for LLM training, breaking the problem into lower-dimensional subproblems to reduce memory and communication overhead. It also offers comprehensive convergence guarantees for various optimizers, with refined results for Adam.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your thorough review and valuable feedback on our manuscript. Below, we provide our detailed responses to the questions you raised.\n\n- Questions:\n\n   In our convergence analysis, we allow the use of different optimizers to solve each subproblem. In our experiments, we employ the Adam optimizer with a fixed number of iterations to solve each subproblem.\n\n   In our implementation of RSO, we adopt a LoRA-style design. For a linear transformation with a frozen weight matrix $W$, we introduce a trainable low-rank matrix $B$ and a projection matrix $P$. Both $W$ and $P$ are treated as fixed parameters with `requires_grad = False`, while $B$ is defined as a learnable parameter. During each forward pass, we compute $XW$ and $XPB$ separately and sum the results to obtain the final output. Since the optimizer (e.g., Adam) updates only $B$, gradients are neither computed nor stored for $W$ or $P$. As a result, PyTorch automatically avoids storing intermediate activations associated with their gradients, thereby improving memory efficiency without requiring manual intervention.\n\n   After each subproblem is solved, the update is merged into the model by setting $W \\leftarrow W + P B$, followed by resampling a new random projection matrix $P$ for the next subproblem. \n\nWe hope these responses address your concerns. Please feel free to reach out with any further questions or feedback.", "claims_and_evidence": "Yes, it's great.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "The paper highlights that prior methods, such as GaLore, lack comprehensive convergence guarantees—often only analyzing fixed projection matrices. In contrast, it provides a complete convergence analysis for the RSO framework across various subproblem solvers, including zeroth-, first-, and second-order methods, and optimizer variants like gradient descent and Adam.", "experimental_designs_or_analyses": "The proposed RSO framework significantly improves memory efficiency and speeds up training by reducing communication overhead compared to state-of-the-art methods such as GaLore, LoRA, and Adam, while maintaining competitive performance. These results underscore the practical value of the approach.", "relation_to_broader_scientific_literature": "Their work extends current literature on deep learning optimization by decomposing high-dimensional training into manageable subproblems using random subspaces. This approach not only reduces memory and communication overhead but also aligns with recent theoretical advances by providing comprehensive convergence guarantees and competitive empirical performance relative to methods like GaLore and Adam.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "The theoretical analysis and experimental validations presented in the paper are robust and well-supported.", "other_comments_or_suggestions": "See in next part.", "questions_for_authors": "[1] How do you address the subproblem in (4a)? Could you provide the detailed code implementation of how to use Adam to optimize the parameter B (how to make B get the gradient)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "A6tkQiVjYb", "reviewer": "Reviewer_RgSv", "summary": "The paper highlights a critical limitation in existing methods: while prior approaches have focused on reducing the memory burden of optimizer states, they have largely overlooked the substantial memory consumption imposed by activations—especially in scenarios with long context sequences or large mini-batches. To address this challenge, the authors propose a novel randomized subspace optimization framework that decomposes the high-dimensional training problem into a series of lower-dimensional subproblems. This innovative approach significantly reduces memory requirements for both activations and optimizer states while maintaining competitive performance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and thoughtful feedback on our manuscript. Below are our detailed responses to the weaknesses and questions you raised:\n\n- Questions:\n\n1. We would like to clarify that, as stated in our manuscript, $\\eta_k$ is required to satisfy $\\eta_k \\le 1/\\hat{L}$ to guarantee theoretical convergence. This condition is explicitly imposed to facilitate the convergence analysis. In our experiments, we set $\\eta_k$ to a relatively large constant (e.g., 10) in order to accelerate convergence in practice. The results demonstrate that RSO still converges reliably under this setting.\n\n2. In all of our experiments, we actually use Gaussian distributions instead of Haar distributions to achieve better computational efficiency. We apologize for any confusion caused by the previous version of the manuscript, and we will explicitly clarify this choice in the next revision.\n\n   As discussed in Remark 5.4, when the rank $r$ is much smaller than the embedding dimensions $m$ or $n$, Gaussian projections closely approximate the behavior of Haar projections. For completeness, we additionally provide experimental results using Haar-distributed projection matrices for training LLaMA-60M and LLaMA-130M, as shown in the following table.\n\n   | Projection Type | LLaMA-60M  | LLaMA-130M  |\n   |------------------|-------------|---------------|\n   | Gaussian         | 34.55       | 25.34         |\n   | Haar             | 34.53       | 25.06         |\n\n- Weaknesses:\n\n    As noted above, we believe this concern has already been addressed in our previous response.\n\nWe hope these responses address your concerns. Please feel free to reach out with any further questions or feedback.", "claims_and_evidence": "The main claim of this paper is that RSO significantly reduces memory usage for optimizer states, gradients, and activations. To support this, the authors conduct a comprehensive theoretical analysis of memory consumption in Table 1, as well as empirical study in Tables 3 and Figure 2.", "methods_and_evaluation_criteria": "The proposed method uses a randomized subspace projection to transform the full-scale optimization problem into a series of lower-dimensional subproblems, which is well-aligned with the problem of reducing training overhead in LLMs.", "theoretical_claims": "Yes, I have checked the proof of memory complexity for activations and optimizer states in Appendix A and convergence analysis in Appendix B.", "experimental_designs_or_analyses": "Yes. The experiments are designed to validate the memory and communication efficiency of RSO as well as its training performance. The main experiments include: 1) pretraining on LLaMA with RSO, showing that RSO could reduce memory consumption. 2) fine-tuning RoBERTa on GLUE benchmark, showing that RSO can achieve competitive performance.", "relation_to_broader_scientific_literature": "The paper is well-situated in the context of existing research on memory-efficient LLM training. It builds upon prior works such as GaLore, LoRA by addressing a key gap: the memory overhead associated with activations.", "essential_references_not_discussed": "I believe that this paper have cited most related works.", "other_strengths_and_weaknesses": "Strengths:\n\n* The paper offers a robust combination of theoretical convergence guarantees and practical memory efficiency analyses.\n\n* Extensive experiments on both pre-training and fine-tuning verify the method’s effectiveness.\n\nWeaknesses:\n\n* More detailed ablation studies on the choice of the distribution of the random subspaces could offer deeper insights into the method’s sensitivity and robustness.", "other_comments_or_suggestions": "No.", "questions_for_authors": "1. For each subproblem (4a), each $\\eta^k$ shoud be chosen such that $\\eta^k \\le 1/ \\hat{L}$. How did you set this value in experiments?\n\n2. The convergence analysis relies on specific properties of the random projection matrices. How robust is RSO when these matrices are sampled from simpler distributions (e.g., Gaussian) rather than the theoretically Haar distribution?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: In this paper, the authors propose a memory-efficient optimizer using random subspace optimization (RSO). Specifically, the method decomposes the original optimization problem into sub-problems via random projection matrices, and update parameters based on these projections. The authors also provide a theoretical convergence analysis for RSO. Importantly, this approach not only achieves memory efficiency through reduced optimizer states but also decreases activation memory cost due to the low-rank projections.

Reviewers all acknowledge the paper's contributions and practical significance. Furthermore, the authors' rebuttal has satisfactorily addressed most concerns raised by reviewers. Thus, this paper is recommended for acceptance.
Decision: Accept (poster)