Subspace Optimization for Large Language Models with Convergence Guarantees

We thank the reviewer for acknowledging our theoretical findings and for the valuable comments. All questions have been clarified as best as we can, and we are glad to address any further comments or questions. We include all new experiments in this anonymous link.

Weakness 1. GaLore is mainly developed for (pre-)training, current experiments only considering fine-tuning are not enough.

We would like to respectfully remind the reviewer that we have already included a pre-training result in Figure 4, where the algorithms are applied to pre-train a LLaMA-60M model on the C4 dataset for 30K iterations. To further address this concern, we have conducted additional experiments on pre-training LLaMA-350M models, with results provided in the anonymous link (Figure 1). These results show that transitioning from GaLore to GoLore leads to a faster decrease in the loss curve, further demonstrating the effectiveness of GoLore in pre-training tasks.

Weakness 2. The improvement in the fine-tuning task is minor.

We respectfully disagree with the reviewer’s assessment. Notably, in the GLUE benchmark, our method achieves an average score of 86.03, compared to GaLore’s 85.77. This represents a threefold reduction in the performance gap between GaLore and AdamW (86.15). We believe this result highlights the effectiveness of our method, demonstrating that a simple modification can significantly close the performance gap.

Weakness 3. The proposed algorithm still relies on GaLore in the initial stage in the practical implementation.

Our hybrid strategy, which relies on GaLore in the initial stage, aligns with our theoretical insights. The limitation of GaLore arises when anisotropic gradient noise starts to dominate the true gradient, which typically occurs in the later stages of training. In the early stage, when true gradients are dominant, GaLore’s greedy strategy remains effective, and no modification is necessary.

Empirically, applying GoLore from the very beginning (i.e., GoLore@100%) results in performance similar to that of GaLore, likely due to a slower initial phase. Therefore, we believe the hybrid strategy is preferable, as it ensures both a fast initial convergence and higher accuracy in the later stages.

Question 1. It would be better to include important parameters and discussion on these parameters in the main text.

The primary hyperparameter introduced beyond GaLore is the switching point, where we transition to GoLore. We kindly refer the reviewer to our rebuttal to Reviewer iLWQ (Weakness 1 and 2) for a detailed discussion, as well as to the anonymous link (Figure 2 and 3) for ablation results.

Question 2. It seems that the analysis didn't consider momentum and Adam cases. It would be better to clarify this.

We kindly remind the reviewer that GoLore's convergence analysis is based on momentum SGD, which explicitly accounts for momentum. Additionally, in our non-convergence analysis, Theorem 4 holds for any optimizer $\rho$, including Adam, AdamW, and momentum SGD.

However, we acknowledge that whether GoLore converges under Adam remains an open question. To clarify this, we will revise the conclusion section as follows:

Original: A limitation of this paper is that recent GaLore variants, such as Fira, are not readily covered by our analysis framework.

Revised: A limitation of this paper is that our convergence analysis framework has not readily covered the use of the Adam optimizer and recent GaLore variants such as Fira.

Question 3. In your counter example and proof, can we reduce the stationary gap via reducing the step size? Is it an issue faced by all stochastic algorithms?

We cannot reduce the gap by decreasing step size. As demonstrated in Theorem 4, the non-convergence result holds for any optimizer $\rho$, regardless of the algorithm framework or hyperparameters, including the step size. Therefore, the counterexample illustrates an issue that is inherent to all stochastic algorithms.

Question 4. If your claim that SVD decomposition biases the gradient holds, why does the deterministic GaLore converge (Theorem 5)?

It may be more accurate to say that GaLore's greedy SVD projection biases the stochastic gradient, including the gradient noise. Intuitively, the bias introduced by gradient noise will not vanish as the gradient noise itself will not vanish.

On the other hand, the bias in the true gradient can be bounded, as shown in Lemma 1, and this bias vanishes as the true gradient approaches zero. This is why GaLore can still converge in the deterministic case, despite the bias in the gradient.

We hope these responses can clarify the reviewer's questions and are more than happy to address any further comments or questions.

We thank the reviewer for acknowledging our theoretical results and appreciate the efforts made to check our convergence proofs. All questions have been clarified as best as we can, and we are glad to address any further comments or questions. We include all new experiments in this anonymous link.

Weakness 1. The dependence of the assumptions may limit broader applicability. In a noisy setup, a random projection may add more uncertainty in the final claims.

We respectfully disagree with the reviewer’s concern that GoLore's broader applicability is limited by the dependence on assumptions, particularly regarding the noise scale. While introducing random projection may add more uncertainty, we argue that such randomness is, in fact, essential to mitigate the impact of gradient noise in noisy settings.

To see it, we first consider the deterministic projection used by GaLore. Let $G=G_0+E$ represent the stochastic gradient, where $G_0$ is the true gradient and $E$ is the gradient noise. In noisy settings where $\\|E\\|\gg\\|G_0\\|$, GaLore deterministically select the projection matrix $P$ using the information of $G\approx E$. The noisier $E$ is, the noisier $P$ becomes. Moreover, GaLore's Top-K selection can degrade the unbiasedness of $E$, potentially leading to non-convergence.

In contrast, GoLore randomly samples $P$ from the same distribution, regardless of how noisy $E$ is. This approach shields $P$ from being influenced by the dominant gradient noise, ensuring its stability regardless of the noise level in $E$. Furthermore, GoLore's random projection maintains the unbiasedness in subspace selection, as shown by $\mathbb{E}[PP^\top]=\frac{r}{m}\cdot I$ in Lemma 5.

Additionally, the assumptions regarding gradient noise in our analysis are standard in stochastic optimization and are widely accepted in convergence studies. Therefore, we believe that our theoretical results are not unduly restricted by these assumptions and that random projection enhances the algorithm’s robustness in noisy settings.

Question 1. Can a sampling based row rank approximation be used instead of random projection of GoLore or singular value decomposition of GaLore?

Thank you for the interesting question. To explore this, we have additionally evaluated the performance of a sampling-based method as follows:

Let $G\in\mathbb{R}^{m\times n}\ (m\le n)$ be the stochastic gradient matrix. Similar to GaLore, we first compute SVD of $G$, obtaining $G=U\Sigma V^\top$. However, instead of selecting the first $r$ columns of $U$ as GaLore does, we sample $r$ columns with probabilities proportional to the corresponding singular values $\sigma_i$'s. As shown in Figure 1 in our anonymous link, the performance of this sampling-based strategy is similar to that of GaLore's, outperformed by our GoLore method. This is potentially because such an importance sampling strategy cannot lead to unbiased projection matrices, and still suffer from the bias induced by gradient noise.

We thank the reviewer again for the careful comments and valuable suggestions. We hope these responses can clarify the reviewer's questions and are more than happy to address any further comments or questions.

We thank the reviewer for acknowledging our theoretical results and experimental designs, as well as the detailed comments and suggestions. All questions have been clarified as best as we can, and we are glad to address any further comments or questions. We include all new experiments in this anonymous link.

• Weakness 1. The paper does not provide a guideline on how to choose the switching point.

  We appreciate the reviewer’s valuable feedback. In our experiments, we manually selected the switching points from \\{20\\%,30\\%,40\\%,50\\%\\} based on intuition and prior experience. Specifically, we choose an earlier switching point if we expect the algorithm to converge more quickly to the solution's neighborhood and a later one if we anticipate slower convergence. To provide a broader guideline, we offer the following insights:

  • If access to the true gradient is available, a reasonable switching point would be when gradient noise begins to dominate the true gradient.

  • If only stochastic gradients are available, and assuming the gradient noise has a roughly constant scale, this dominance can be estimated by monitoring whether the norm of the stochastic gradients falls below a certain threshold. This threshold serves as a hyperparameter that depends on both the training task and the batch size used in the algorithm.

  • The optimal switching point is task-dependent. Empirically, it is recommended to switch when the rate of decrease in the loss curve starts to slow down.

• Weakness 2. The ablation study on the subspace update frequency is missing.

  We appreciate the reviewer’s suggestion. We did not initially include an ablation study on the subspace update frequency $\tau$ because the original GaLore paper (Zhao et al., 2024) states that values of $\tau$ between 50 and 1000 have minimal impact on performance. Based on this, we set $\tau=500$ to:

  • Reduce computational overhead, and
  • Ensure a sufficient number of GoLore projection steps to enhance performance.

  However, in response to the reviewer's concern, we have provided additional ablation results in Figure 3 in the anonymous link. These results confirm that varying $\tau$ within the range of 50 to 1000 has negligible effect on performance.

• Suggestion 1. It may be informative to plot the norm of gradient bias of both Galore and Golore.

  We appreciate the reviewer’s valuable suggestion. In response, Figures 4 and 5 in the anonymous link have additionally plotted the gradient bias norms in our experiments, demonstrating that GoLore achieves a smaller gradient bias.

We thank the reviewer again for the careful comments. We hope these responses can clarify the reviewer's questions and are more than happy to address any further comments or questions.