A Stable Whitening Optimizer for Efficient Neural Network Training

Thank you for review. It appears the main concern is regarding the wall-clock requirements for the eigendecomposition.

• Implementation of Eigendecomposition

We first note that eigendecomposition is used in all optimizers from the Shampoo family, including SOAP and SPlus. As described in 5.1, we use a distributed method to compute this decomposition for all three methods. Regarding the claim that the additional matrix multipliation is neglible, we agree that this is misleading, and will remove this claim in the revision in favor of a more numerical analysis. See below for an additional analysis we ran investigating wall-clock time.

We note that wall-clock time should be best understood as a rough estimate, as there is an inherent dependency on the hardware used, the specific distributed system, speed of inter-host communication, etc.

• Ablation on Eigendecomposition Frequency

To examine the exact relationship between wall-clock time and eigendecomposition, we ablate the frequency of the matrix eigendecomposition for both Shampoo and SPlus between [1, 5, 10, 25, 100, 250, 500]. Due to computational constraints, we measure validation loss after 1000 gradient steps (vs. 10,000 in the paper) and only on the LLM objective. All other hyperparameters are as discussed in the paper.

Validation Loss after 1K gradient steps:

Steps-to-Adam:

As requested with "Please specify how much more is the significant wall-clock cost for matrix eigen-decomposition with justifications", we make explicit the wall-clock time requrired per-step under various decomposition frequencies:

Average Time Per Update:

Time-to-Adam:

Performance increases as the eigendecomposition frequency increases, but at a diminishing rate. Wallclock time scales roughly linearly with the frequency, as eigendecomposition is the main computational cost. Notably, while there is a sweet spot for the optimal frequency hyperparameter in terms of time-to-Adam (at 250 in this case), Shampoo fails to capitalize on this as training is unstable with an interval >= 100.

We would like to thank you for raising relevant questions and suggestions. We believe the additional ablations have increased the clarity of the results and strengthened the paper. Please let us know if you have any remaining concerns or questions.

Thank you for your detailed review and questions. Please see our response below:

• Ablating the proposed changes

Addressing the request for an ablation of each modification, we ran an additional set of experiments to properly ablate the changes in SPlus, specifically focusing on the interaction between iterate-averaging and instant-sign normalization in the LLM setting. We specifically ran an experiment of SPlus without the use of iterate averaging. We did not ablate shape-aware scaling, as this requires another sweep to find the optimal learning rate, which we are unable to perform due to compute limitations.

To address the request for "I wonder how a Schedule-Free version of SOAP, PSGD and Muon would fare", we additionally compare to SOAP + iterate averaging, as SOAP is the highest-performing baseline method in our experiments. (EMA iterate averaging was more performant than Schedule-Free in our experiments)

Steps-to-Adam:

As shown in the above table, the usage of instant-sign normalization is important to stabilize Shampoo. Both SPlus and SOAP benefit from iterate averaging. SOAP performs slightly better at later checkpoints, but requires keeping an additional set of parameters in memory to track the second moment.

• Wall-clock time and inversion frequency

Thank you for mentioning the skepticism on wall-clock item -- we agree that wall-clock times are hard to standardize and should be seen as a rough estimate -- there is an inherent dependency on the hardware used, the specific distributed system, speed of inter-host communication, etc. The methods with the longest time-per-step in our experiment suite are Shampoo/SOAP/SPlus, especially at frequent inversion times. For all three of these methods, we use the same distributed behavior.

To understand the tradeoff between wall-clock time and performance, we conduct an ablation on the matrix-inversion frequency in Shampoo and SPlus between [1, 5, 10, 25, 100, 250, 500]. Due to computational constraints, we measure validation loss after 1000 gradient steps (vs. 10,000 in the paper) and only on the LLM objective. All other hyperparameters are as discussed in the paper.

Validation Loss after 1K gradient steps:

Steps-to-Adam:

Average Time Per Update:

Time-to-Adam:

Performance increases as the eigendecomposition frequency increases, but at a diminishing rate. Wallclock time scales roughly linearly with the frequency, as eigendecomposition is the main computational cost. Notably, while there is a sweet spot for the optimal frequency hyperparameter in terms of time-to-Adam (at 250 in this case), Shampoo fails to capitalize on this as training is unstable with an interval >= 100.

Thank you for the detailed responses and providing additional ablation studies.

Most of my questions have been addressed, but I do have some remaining ones and follow-up questions: Looking at the ablation study, it appears that instant-normalization has a stabilizing, but almost no acceleration effect as Shampoo with and without instant-norm have a very similar value at a later checkpoint, while the speed-up mainly comes from iterate averaging, which is also evident by the fact that SOAP + iterate-averaging performs very similarly to SPlus if not better. I wonder if iterate-averaging alone already provides some stabilizing effect as well. Have the authors considered running SPlus without instant-normalization?

I have also read your rebuttal to Reviewer Xpwz and noticed that Muon with smaller weight decay is much more competitive (without iterate-averaging). In particular due to smaller per-iteration complexity it converges faster in wall-clock time. I wonder if Muon can achieve even better performance by adding iterate-averaging?

Thank you for your response! Regarding SPlus without instant-sign normalization -- this would essentially be Shampoo with iterate averaging, and as shown in the above table, Shampoo has a tendency to diverge when the eigenbases are cached for long periods. Iterate averaging would not prevent divergence, as the diverged parameters explode in magnitude and are unusable even when averaged.

Regarding Muon, yes, Muon can likely achieve a similar gain in performance from iterate averaging as SPlus/SOAP do. As shown in the table you are referring to, Muon is competitive with SPlus/SOAP in many cases, and has an efficient implementation in terms of wall-clock time -- however, in the above results Muon tends to perform worse when training from later checkpoints, especially on the DiT and ViT cases which have a higher variance batch. (Diffusion has inherently high variance due to the noised objective, and ViT has less signal per-batch due to the single classification target, in comparison to language modelling which calculates a loss at each token). This can be intuitively understood, as Muon approximates the orthogonalization per-batch, whereas SPlus/SOAP/Shampoo can utilize historical information.

We hope these answers your questions. As mentioned in the original paper, the famiIy of 'whitening' optimizers (e.g. Shampoo, SOAP, PSGD, Muon, SPlus) are similar in form. One takeaway from this paper is that in the Shampoo family, SPlus provides a way to roughly match the performance of SOAP while reducing the memory footprint. In all methods, iterate averaging can increase performance at each step (at the cost of additional memory, which we aim to save via the previous technique in SPlus).

If the additional experiments have strengthened your view on the paper, would you consider updating the score accordingly?

Thank you for the clarifications!

I would like to thank the authors for taking the time to clarify my questions and providing detailed explanations. I got a better understanding of SPlus and its performance compared to other optimizers in the family of "whitening" optimizers as a whole now.

I acknowledge the contributions as valuable and suggest acceptance of this work and keep my positive score.

Thank you for your review and detailed questions. Please see the detailed response:

• Ablating iterate averaging $\beta_3$

To answer the question on the sensitivity of the $\beta_3$ parameter, we ran an additional set of ablations, focusing on the LLM setting. Our experimental setup already measures the performance at three stages of training (initialization, 10k checkpoint, 50k checkpoint), and we report values ablating $\beta_3$ for SPlus:

Validation Loss of SPlus:

Steps-to-Adam of SPlus:

As shown in the above table, we can get slightly better performance by tuning the $\beta_3$ parameter per training setting. However, using a default value of 0.999 (as we do in the main paper) allows for nearly the same performance across the board, regardless of the starting checkpoint. The intuitive reason for the effectiveness of iterate averaging, especially under a constant learning rate, can be viewed as reducing the effect of noise during optimization -- if we assume that each gradient steps introduces some independent noise due to linearization/stochastic batches, then iterate averaging enables the noise to cancel out (variance of the mean of i.i.d. random variables decays linearly with the sample size). We will include this discussion in the next revision.

• Memory Usage of SPlus

The additional memory requirement is a fair point and a weakness of the method. However, we note that the core SPlus update saves memory compared to the most competitive baseline (SOAP), which requires an additional set of parameters to track the second moment. In settings where iterative averaging is disabled, SPlus can roughly approximate SOAP while not requiring additional memory. We confirm this via an additional set of experiments, where SPlus is run without iterate averaging:

Validation Loss (no iterate averaging):

Steps-to-Adam (no iterate averaging):

• Notation and figures

Regarding the notation and the axes, thank you for catching these issues, and we have fixed them in the next revision. $g$ in our case represents the gradient with relation to the loss function (i.e. $\nabla L(\theta)$), and $\Delta \theta$ is an arbitrary vector to take the argmin over. $\theta’$ is a set of "live" parameters, which are averaged into $\theta$.

• Why do we use the Steps-to-Adam metric?:

We chose steps-to-Adam as it is an interpretable way to understand the effective speedup of one optimizer to the rest, and many prior works report a speedup over Adam, e.g [1, 2]. Raw validation losses are less interpretable, as the relative difference is harder to compare. That said, in the above tables, we have also provided raw validation losses.

• Regarding the Transformer-only experiments

It is a fair point that our experiments limit only to Transformers, as mentioned in the limitations section. We note that our experimental setup is still larger than many previous works [1, 2] that only consider language modelling -- we additionally consider diffusion modelling and image classification.

We would like to thank you for raising relevant questions and suggestions. We believe the additional ablations have increased the clarity of the results and strengthened the paper. Please let us know if you have any remaining concerns or questions.

[1] Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training. Lui et al, 2023.

[2] SOAP: Improving and Stabilizing Shampoo using Adam. Vyas et al, 2024.

Thank you for the detailed review and analysis. Please find our detailed response below:

• Relation and comparison to Muon

Regarding the relation to Muon, we also agree with your analysis here. The relationship between Shampoo and Muon is studied in [1], and SPlus can be understood as an intermediate point where the eigenvalues are calculated via the sign operator, but the eigenvectors are calculated over the historical average. We will further clarify this relationship in the paper.

Regarding the gap in performance, we found through further analysis that the original Muon reported numbers were using a default weight decay that caused a substantial degradation. In particular, we use weight decay = 0.1*lr for all experiments -- however, this setting is suboptimal for Muon which has an inherently high learning rate. For a more representative comparison, we re-ran the Muon experiments with a weight decay of 0.001 * lr, and report the new numbers:

Steps-to-Adam:

Time-to-Adam:

As shown above, SOAP and SPlus tend to reach the desired validation loss in less gradient steps than Muon. However, Muon can reach this point in less wall-clock time, likely due to its efficient use of Newton-Schulz iteration versus explicit eigendecomposition.

• Ablating the Effect of Iterate Averaging

We ran an additional experiment ablating the effect of iterate averaging (via EMA) on SPlus, focusing specifically on the LLM setting due to compute limitations. We also compare to introducing iterate averaging to SOAP, as SOAP is the most competitive baseline in terms of Steps-to-Adam.

Validation Loss:

Steps-to-Adam:

SPlus without EMA performs similarly to SOAP (and Muon), and all three can be understood as approximating a particular spectral whitening objective (see Section 3, Eq. 4). In our specific setting with a constant learning rate, iterative averaging provides an additional performance gain. The practical benefit of SPlus over SOAP is that SPlus requires one less set of parameters to be kept in memory.

Regarding mathematical typos, thank you for pointing these out, and they have been fixed for the next revision.

We would like to thank you for raising relevant questions and suggestions. We believe the additional ablations have increased the clarity of the results and strengthened the paper. Please let us know if you have any remaining concerns or questions.

[1] Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology. 2024.

import numpy as np

g = np.random.rand(4, 20)
ggt = g @ g.T
print(ggt.shape)
print(ggt)

coeffs = (3, -16/5, 6/5) # Many choices work here.
def newton_schulz_iterator(x):
    a = x @ x.T
    b = coeffs[1] * a + coeffs[2] * a @ a
    return coeffs[0] * x + b @ x
x = ggt
x /= np.linalg.norm(x) + 1e-6  # Singular values must be in [-1, 1]
for _ in range(5):  # Iterate a few times to converge.
    x = newton_schulz_iterator(x)
print('Newton-Schulz orthogonalization\n', x)

# Explicit orthogonalization
u, s, vh = np.linalg.svd(ggt)
print('Explicit orthogonalization\n', u @ vh)

(4, 4)
[[4.52943642 4.34923569 4.21630621 4.21347935]
 [4.34923569 8.74947178 5.75485439 6.18474024]
 [4.21630621 5.75485439 7.15446298 5.41502261]
 [4.21347935 6.18474024 5.41502261 7.6006473 ]]
Newton-Schulz orthogonalization
 [[ 9.96821455e-01 -3.39008106e-04  2.19201580e-03  1.15344933e-03]
 [-3.39008106e-04  9.97126618e-01 -4.52450965e-03  8.55999996e-03]
 [ 2.19201580e-03 -4.52450965e-03  9.95145616e-01  8.73224205e-03]
 [ 1.15344933e-03  8.55999996e-03  8.73224205e-03  9.82407800e-01]]
Explicit orthogonalization
 [[ 1.00000000e+00  2.56983382e-16  1.97278891e-18  1.85961893e-16]
 [ 2.06052557e-16  1.00000000e+00 -5.51062184e-18  1.44590563e-16]
 [ 1.56194126e-16  1.20230648e-16  1.00000000e+00  7.58146768e-17]
 [-1.00514801e-16  1.84793521e-16  2.05541700e-16  1.00000000e+00]]

Thank you for the detailed review and suggestions. Please see our detailed response and additional experiments below:

• How does eigendecompisition frequency affect performance and wallclock time?

This is a great question, and we ran an additional set of experiments to answer it. In the following experiments, we run SPlus and Shampoo over a range of eigendecomposition frequencies between [1, 5, 10, 25, 100, 250, 500]. Due to computational constraints, we measure validation loss after 1000 gradient steps (vs. 10,000 in the paper) and only on the LLM objective. All other hyperparameters are as discussed in the paper.

Validation Loss after 1K gradient steps:

Steps-to-Adam:

Average Time Per Update:

Time-to-Adam:

Performance increases as the eigendecomposition frequency increases, but at a diminishing rate. Wallclock time scales roughly linearly with the frequency, as eigendecomposition is the main computational cost. Notably, while there is a sweet spot for the optimal frequency hyperparameter in terms of time-to-Adam (at 250 in this case), Shampoo fails to capitalize on this as training is unstable with an interval >= 100.

We note that wall-clock time should be best understood as a rough estimate, as there is an inherent dependency on the hardware used, the specific distributed system, speed of inter-host communication, etc.

• Memory Usage of SPlus

The additional memory requirement is a fair point and a weakness of the method. However, we note that the core SPlus update saves memory compared to the most competitive baseline (SOAP), which requires an additional set of parameters to track the second moment. In settings where iterative averaging is disabled, SPlus can roughly approximate SOAP while not requiring additional memory. We confirm this via an additional set of experiments, where SPlus is run without iterate averaging:

Validation Loss (no iterate averaging):

Steps-to-Adam (no iterate averaging):

• Potential methods of improving compute efficiency

As noted in the future work section, we agree that there is a fruitful research direction on methods which can improve the computational efficiency of SPlus and other Shampoo-like methods. We did not conduct specific experiments on algorithmic approximations to the eigendecomposition (e.g. low-rank approximations) as we view them as beyond the scope of this paper. In terms of wall-clock time for billion-parameter LLMs, while we are also curious on this result, we are unable to properly experiment with such large-scale models due to compute limitations. Thank you for the reference to NGPlus! We will update the paper to introduce a citation to this, and describe potential techniques to reduce the computational footprint.

We would like to thank you for raising relevant questions and suggestions. We believe the additional ablations have increased the clarity of the results and strengthened the paper. Please let us know if you have any remaining concerns or questions.