SOAP: Improving and Stabilizing Shampoo using Adam for Language Modeling

Thanks for championing the paper. We are happy to see that the reviewer finds the work clear and easy to follow, and considers the work to have significant contribution to the deep learning optimization community. Specific comments are addressed below:

The authors write SOAP is equivalent to running Adam in a rotated space; however, is the transformation $G' \to Q_L^T G Q_R$ to the eigenbasis necessarily unitary? What is the specific operation in SO(d)?

As pointed out, we project the gradient as $G' \leftarrow Q_L^T G Q_R$. If we represent $g = vec(G)$ and $g' = vec(G')$, where $vec(A)$ represents vectorization of the matrix $A$, then the update can be written as $g' \leftarrow (Q_L \otimes Q_R) g$, where $\otimes$ represents the Kronecker product. $Q_L \otimes Q_R$ represents the orthogonal transformation.

We hope that we have addressed your questions and concerns. If not, we would be happy to answer any additional questions.

We thank the reviewer for their valuable and insightful feedback. We are happy to see that the reviewer finds the connection between Shampoo and Adafactor novel, and our experimental results promising. Specific comments are addressed below:

I am a bit confused by the main contribution of this paper overall; algorithmic insights on the connection between Adafactor and Shampoo are interesting. However, beyond that, the rest seems to be a special case of GaLore.

As detailed in Section 3 (Lines 169-179), SOAP differs from GaLore in two key aspects:

1. The difference in how projection matrices are computed - GaLore uses SVD of the current gradient, while we maintain an exponential moving average of $GG^T$ and $G^T G$.
2. Projecting the momentum on the preconditioned space, every time the preconditoning space is updated. GaLore does not do this.

Appendix D shows these technical differences contribute to the superior performance of SOAP as compared to GaLore.

What happens if you use power iterations to approximate SVD in GaLore; does that still work? It feels like that should be included as a baseline too. [1] Liang, Kaizhao, et al. "Memory-efficient LLM training with online subspace descent." arXiv preprint arXiv:2408.12857 (2024).

Including GaLore with power iteration approximations as a baseline is an interesting idea. But we note that even without any approximations GaLore performs worse than Shampoo in our experiments (Appendix D).

Why does higher frequency lead to worse results? Intuitively, it should provide a better and more up-to-date approximation to the eigenspace.

We apologize for the confusion. The term “preconditioning frequency” refers to the delay between updates. For a frequency of $k$, the preconditioner is updated once every $k$ steps. Higher values of $k$ reduce computational cost but lead to less frequent updates, which can degrade performance.

Please also do some downstream evaluations on checkpoints output by different optimizers, just for the sake of completeness.

Yes, we will add these evaluations in future revisions.

We hope that we have addressed your questions and concerns. If not, we would be happy to answer any additional questions.

We would like to thank the reviewer again for their valuable feedback. We would be grateful if the reviewer would consider updating their score based on our response. If not, we would be happy to answer any additional questions.

We tried preconditioning frequency of 10, 50 and 200 for GaLore (as mentioned in Appendix D) and found* that 200 preconditioning frequency performs the best. This shows that GaLore does not benefit from updating the preconditioning matrices faster. By "performs the best" we refer to final loss given a fixed number of steps (not wall clock time). Quoting appendix D for completeness:

"We swept the cross product over learning rate (3.16e − 4, 1e − 3, 3.16e − 3, 1e − 2), preconditioning frequency (10, 50, 200), both sided and one sided versions. Frequency 200 had the best results matching the observation of Zhao et al. (2024a)."

We hope this answers your question, please let us know if you have any further questions.

*This is likely because GaLore maintains momentum in the rotated (by eigenvectors) space and hence at lower frequencies the momentum estimates possibly become quite bad whenever the eigenvectors change.

"Including GaLore with power iteration approximations as a baseline is an interesting idea. But we note that even without any approximations GaLore performs worse than Shampoo in our experiments (Appendix D)."

I don't think this is a valid argument. Since single step power iteration is not as expensive, you can afford to run it every step to get more up-to-date information of the projected space. Whether it's better or worse is uncertain, but it should be a very obvious baseline.

We think this likely because GaLore maintains momentum in the rotated (by eigenvectors) space (while SOAP maintains it in the original space) and hence at lower frequencies the momentum estimates possibly become quite bad for GaLore whenever the eigenvectors change.

Could you please let us know what the unfalsifiable claims are?

We note that the optimizer code is available at https://anonymous.4open.science/r/SOAP-F93B/soap.py.

We thank the reviewer for their valuable and insightful feedback. We are happy to see that the reviewer finds the paper well-written and novel. Their specific comments are addressed below:

I believe that showing good performance compared to AdamW on the 1B to 7B parameter scale, would strongly support the adoption of SOAP. Also, it would be interesting how the new optimizer SOAP performs in combination with newly emerging Linear Recurrent Models (like Mamba).

We agree that validating SOAP on larger model scales and alternative architectures is important. Due to computational constraints, we focused on 360M and 660M models in this study, but we plan to extend our work to larger scales and diverse model architectures in the future.

Did you observe any influence of SOAP on training instabilities (e.g., grad norm spikes / rise or loss spikes) in language model pretraining?

No, we did not monitor grad norm spikes during training.

It would be helpful for others adopting SOAP if the authors provide learning curves for other common metrics monitored during pre-training (e.g., parameter norms, gradient norms, validation loss, etc.). This would help comparing other training setups to your suggested one.

We are happy to include these learning curves in a future revision to support adoption and facilitate comparisons.

We hope that we have addressed your questions and concerns. If not, we would be happy to answer any additional questions.

We would like to thank the reviewer again for their valuable feedback. We would be grateful if the reviewer would consider updating their score based on our response. If not, we would be happy to answer any additional questions.

We thank the reviewer for their valuable and insightful feedback. We are happy that the reviewer finds the paper well-organized and of significant importance for the LLM community. The specific comments are addressed below:

The evaluation in this paper is limited, as it only tests the proposed SOAP on 360M and 660M language models. It is unclear whether SOAP can scale effectively to larger models, such as 1B or 7B, or how it performs in fine-tuning settings.

We acknowledge that our experiments are limited to 660M-scale models due to computational constraints, as mentioned in the paper (Line 530). However, given the theoretical foundations of SOAP and its promising results at this scale, we are optimistic about its scalability to larger models. Validating SOAP at larger scales remains a priority for future work. We also hope to conduct fine-tuning experiments in the future.

Given that SOAP is intended as a general optimizer, it would be valuable to assess its performance across diverse tasks, including vision and graph tasks, to better understand its generalization.

Currently, our focus is on language modeling, which we believe is an important task in itself. To align with this focus, we can revise our contributions to explicitly state applicability to language modeling setups. We plan to explore SOAP’s performance in other domains as part of future work.

While the paper establishes a theoretical connection between Shampoo and Adafactor that motivates SOAP, it remains unclear why SOAP outperforms both Shampoo and Adam.

As shown in Figure 5 in the Appendix (Line 936), factorized SOAP (i.e., Adafactor in Shampoo’s eigenbasis) performs comparably to SOAP. However, Shampoo only matches the performance of factorized SOAP when using a preconditioning frequency of 1, which incurs significant computational overhead. The main contribution of this work is to decouple the eigenbasis updates from eigenvalue updates, enabling SOAP (or factorized SOAP) to achieve comparable performance even at higher preconditioning frequencies (due to adaptivity in the eigenvalues).

Since SOAP is inspired by the connection between Adafactor and Shampoo, it is essential to include Adafactor as a baseline to fully evaluate SOAP's effectiveness.

As shown in recent work [1], the performance of Adafactor and Adam is comparable to each other. We would be happy to include Adafactor as a baseline if this is critical for the reviewer's evaluation of the work.

[1] - Zhao et al. 2024 - Deconstructing What Makes a Good Optimizer for Language Models

What's the performance of SOAP in other tasks such as vision and graph tasks, and fine-tuning settings?

Evaluating SOAP in these tasks is part of our planned future work.

why can SOAP perform better than Shampoo? Is there any rationale behind this?

SOAP performs better than Shampoo at preconditioning frequency greater than 1. This is because within preconditioning steps, even though the eigenbasis remains fixed, SOAP still updates the eigenvalues, leading to more adaptivity as compared to Shampoo.

What's the performance of SOAP in fine-tuning settings?

Evaluating SOAP for fine-tuning settings is a part of future work.

We hope that we have addressed your questions and concerns. In light of this, we would be grateful if you would consider updating your review of the paper. Alternatively, we would be happy to answer any additional questions that lead you to maintain your current recommendation.

Hi Kyunghun Nam

Thanks a lot for looking at the paper. Yes, these are indeed typos in the Algorithm.

Line 10: It should be $n \times 1$ column vector.

Line 11: It should be $1 \times m$ row vector.

Line 12: It should be AC.

Line 14: It should be $Q_L G_t'' Q_R^T$.

Thanks a lot for catching these typos. We will fix them in the next revision.