A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models

Thank you for your time and thoughtful feedback on our manuscript. Below are our detailed responses to the weaknesses and questions you raised:

• Questions:

1. 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.

2. 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.

   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.

   Projection Type	LLaMA-60M	LLaMA-130M
   Gaussian	34.55	25.34
   Haar	34.53	25.06

• Weaknesses:

  As noted above, we believe this concern has already been addressed in our previous response.

We hope these responses address your concerns. Please feel free to reach out with any further questions or feedback.

We appreciate your thorough review and valuable feedback on our manuscript. Below, we provide our detailed responses to the questions you raised.

• Questions:

  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.

  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.

  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.

We hope these responses address your concerns. Please feel free to reach out with any further questions or feedback.

Thank you for taking the time to review our manuscript. We greatly appreciate your valuable feedback. Below, we provide our responses to your comments:

• Weaknesses:

  • 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.

  • We will address the FLOPs-related discussion in the response to the questions section.

• Other Comments or Suggestions:

  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.

• Questions:

  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.

     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.

     We apologize for any confusion and will clarify this point in the next revision.

  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.

  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.

     Method	Perplexity
     RSO-Adam	34.55
     RSO-Shampoo	36.45

  4. In the $k$-th outer iteration, we solve the subproblem $\min_B f(W^k + P^k B),$ 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.

     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 $f(W^{k+1}) \leq f(W^k).$

     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.

  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:

     https://anonymous.4open.science/r/RSO-1CC6/

     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.

     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.

     Rank	64	128	192		Inner Steps	100	200	400	800
     Perplexity	37.60	34.55	33.25		Perplexity	34.54	34.55	34.80	35.30

     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.

We hope these responses adequately address your concerns. If you have further questions or need additional clarifications, we would be happy to provide them.

Thank you for taking the time to review our manuscript. We greatly appreciate your valuable feedback. Below, we provide our responses to your comments:

• Experimental Designs Or Analyses:

  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.

• Essential References Not Discussed:

  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.

  We will incorporate these two works and discuss them in the related work section in the next revision of our manuscript.

• Weaknesses:

  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.

  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.

  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.

  Method	LLaMA-1B (GB)	LLaMA-7B (GB)
  RSO	54.73	68.39
  LoRA	69.93	78.82

• Questions:

  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.

  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.

Thank 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.