A Stable Whitening Optimizer for Efficient Neural Network Training
Authors: Kevin Frans
Sergey Levine
Pieter Abbeel
Venue: NeurIPS 2025 poster ()
Keywords: optimization
preconditioning
whitening
second-order
shampoo
Primary area: deep_learning

Abstract: In this work, we take an experimentally grounded look at neural network optimization. Building on the Shampoo family of algorithms, we identify and alleviate three key issues, resulting in the proposed SPlus method. First, we find that naive Shampoo is prone to divergence when matrix-inverses are cached for long periods. We introduce an alternate bounded update combining a historical eigenbasis with instantaneous normalization, resulting in across-the-board stability and significantly lower computational requirements. Second, we adapt a shape-aware scaling to enable learning rate transfer across network width. Third, we find that high learning rates result in large parameter noise, and propose a simple iterate-averaging scheme which unblocks faster learning. To properly confirm these findings, we introduce a pointed Transformer training benchmark, considering three objectives (language modelling, image classification, and diffusion modelling) across different stages of training. On average, SPlus is able to reach the validation performance of Adam within 44-58% of the gradient steps and 62-83% of the wallclock time.

A Stable Whitening Optimizer for Efficient Neural Network Training
Kevin Frans UC Berkeley kvfrans@berkeley.co edu
Sergey Levine UC Berkeley
Pieter Abbeel UC Berkeley
Abstract
In this work, we take an experimentally grounded look at neural network optimiza- tion. Building the Shampoo family of algorithms, we identify and alleviate three key issues, resulting in the proposed SPlus method. First, we find that naive Shampoo is prone to divergence when We introduce an alternate bounded update combining historical eigenbasis with instantaneous normalization, resulting in across-the-boiie stability and significantly lower computational requirements. Second, we adapt a shape-aware scaling to enable learning rate transfer across network width. Third, find that high learning rates result in large parameter noise, and propose a simple iterate-averaging scheme which unblocks faster learning. To properly confirm these findings, we introduce a pointed Transformer training benchmark, considering three objectives (language modelling, image classification, and diffusion modelling) across different stages of training. On average, SPlus is able to reach the validation performance of Adam within 58% of the gradient steps and 62 83% of the wallclock time.
are cached for long periods.
Introduction
The backbone of modern deep learning stochastic gradient descent method which, while particularly effective in practice, is time-consuming by its iterative nature. As large neural networks are scaled up billions of parameters [1, 8, 38], and are trained on datasets of similar scale, becomes increasingly important that each gradient update is efficient in making learning progress.
A core optimization strategy involves adapting updates second-orde curvature [3, 24], allowing for faster learning progress in certain directions while preventing divergence others. However, practical challenges arise with regards large neural network training. For example, it is computationally intractable to calculate the Hessian of any reasonably-sized network [25, 33], and even further intractable invert this matrix. Numerical stability is also concern, and complex methods often require additional hyperparameter tuning and regularization [2, 32).
In this work, take cautiously empirical view neural network optimization. We carefully design an evaluation suite that is well-aligned with common uses today. We consider the optimization of standard Transformer model [52] on three distinct training settings autoregressive language modelling [42], diffusion modelling [16, 37], and image classification [10]. We further consider the performance of optimizers various stages of training (e.g. starting from checkpoints near the start, middle, and end), to avoid biasing towards early-stage performance gains. To our knowledge, this is currently the widest-scope comparison of adaptive optimizers on Transformer training-
From findings in this setting, we develop method we refer to as SPlus. SPlus builds on the Shampoo family of algorithms [15], which can broadly be
approximating whitening metric from historical gradients, and performing steepest descent along this metric. We identify and address three issues with the naive Shampoo update. First, we find that Shampoo often diverges at high learning
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

rates or under less-frequent inversion rates. To fix this issue. we introduce an alternative bounded update combining the historical eigenbasis with instantaneous normalization, resulting in across- the- board stability even with significantly lower- seeer-frqueecy inversions. Second, we find that Shampoo update magnitudes are not properly scaled with relation to network width. We adopt methodology derived for SGD and Adam [55] to our setting, enabling easy hyperparameter tuning via learning rate transfer across network widths. Third, we find that higher learning rates in whitening- optimizers result in significant parameter noise. We propose a simple iterate-averaging scheme to alleviate this issue, unblocking fast learning with a much lower degradation in performance.
In experiments, SPlus achieves the fastest loss decrease among wide suite of previous optimizers. SPlus matches Adam within 44 58% of the gradient steps, and within 62 83% of the wallclock time. SPlus does not introduce any critical hyperparameters to tune, and does not require additional forward/backward passes. We hope that this effectiveness and simplicity enables the community to easily adopt SPlus to their existing training settings.
We provide the full code replicate experiments at github. github.com/kv con/kvraas//s//lusss The repo also contains single-file implementations of SPlus in JAX and Pytorch, along with basic reccomendations for usage.
Related Work
Second-order optimization of neural networks. Second-order methods, which we broadly define as methods which further modify the first-order gradient, can be largely categorized among two axes the second-order metric used, and the way in which the metric is approximated and applied. A common metric choice is the Hessian, which in Newton's 's method can be directly applied by multiplying the inverse Hessian and the gradient. Neural-network specific methods have approximated the Hessian via conjugate- conjugate-gradient methods [33], iterative fitting [25, 41], or diagonal approximation [24, 27]. An alternate metric is the Fisher, motivated by natural gradient methods [3, 49J. Closest related to our work is the use of the "empirical" Fisher metric [22], i.e. an uncentered covariance matrix of gradients. The most well-tested of these are the diagonal Adam and its variants [11, 20, 28, 46).
Kronecker factorized metries. Our work utilizes Kronecker factorization of the empirical Fisher, building on the techniques introduced in K-FAC [32], and more directly Shampoo [15]. Shampoo has been extended to utilize distributed inversion and exponential averaging [47], and/or perform inversion via Newton iteration [4]. Closest related to our work is SOAP [53], which similarly utilizes an cigen-decompositton of the factor matrices, however, SOAP maintains additional set of second-moments per parameter, whereas in our work this is not neccessary.
Non- EaelEuclid gradient descent. We build on previous works which have studied gradient descent over non-Euclidean distance metrics. In particular, such a framing may manifest not as linear transformations of the gradient, but also methods which explicitly parameterize learning rate [23, 55], or instantaneously transform the gradient via sign descent [5]. Newton-Shule orthogonalization of gradients [18, 29] can also be motivated as maximum spectral descent.
Iterate averaging. Whitening-basee based optimizers have poor convergence properties due to the uniform size of each update, and instead typically rely on learning-rate decay. An alternative strategy is to average fast-moving snapshots of parameters [17, 35, 39, 40, 44], which has been shown effective in approximating learning-rate -ratingg-rtteaae decay [9] and is often used in image generation models [19, 57].
Background and Preliminaries
Gradient descent methods can be seen as following the steepest descent direction under some metric. Naive gradient descent implicitly assumes a Euclidean distance metric over parameters, in which case the update vector is simply the scaled gradient g = VeL(e, .r) itself:
u arg min A8
TTA (a/2)110013
ag.
(1)
Improvement
Distance Penalty
However, is often helpful to impose other metrics. For example, certain parameters may be more sensitive to second-order changes, and thus should be assigned larger penalty. We can generally express distance using Riemannian metric tensor M, which a symmetric positive-definite matrix of shape IR (ddi())) Under M, the distance an update can be expressed as the matrix product:
A 2
AOT M00.
(2)

---

No Whitening (Gradient Descent)
Elementwise Whitening (Adam, Sign Descent)
Full Whitening (Shampoo, SPlus)
Whitening normalizes gradients to have uniform magnitude along each axis of descent. This decouples the updates from gradient magnitude. Elementwise whitening imposes independent axis per dimension, whereas full whitening uses the axes that maximally explain gradient covariance.
Figure I:
Gradient descent can now be performed using M the distance metric. The solution then becomes:
= arg min A
AO
+ +1/22)00" MAO
M-
(3)
Improvement
Distance Penalty
which, in the case of Euclidean distance represented by the identity metric, reduces Equation (1).
Whitening metric. Empirically, the whitening metric [56] has proven to be reliable choice for neural network optimization. The whitening metric is the square-root of the uncentered covariance:
1/2 Eg [g9"] 1/2
Eg [e)
(4)
and is guaranteed to be positive definite. The whitening metric can be calculated via empirical gradi- ents and does not require additional forward/backward passes. Notably, Adam [20] performs descent on a diagonal approximation of the whitening metric, assuming each parameter as independent.
What about the Hessian the Fisher? While potentially effective, these metrics often expensive to compute as they require additional information outside the standard gradient. We refer to detailed discussions in [21, 31, 34, 49], as well brief overview in Appendix A.4.
Approximations for neural network learning. To make storing and inverting the metric amenable for large neural networks, it is common to assume per-layer blockwise approximation [15, 32). In this way, the full whitening metric can be represented as set of smaller block matrices, one per layer, which can each be independently inverted.
To further reduce memory and computation, each block can be further approximated by Kronecker product of two smaller matrices:
011 B
B12B
01m B
M Mmn. A mmm
B
0m1 B
Qm2B
Omm B
Kronecker products have a useful property that the inverse (at any power) of Kronecker product equivalent to the Kronecker product of the inverse factors:
if M = A B, then M- = (A-112 B-122).
(6)
Additionally, multiplication by Kronecker product can be performed without explicitly forming the full product matrix. Consider the flattened gradient and update vectors g. u E Rimm and their corresponding matrix forms G, U pmxx The following operations are identical:
(A 1)/g
U A-122/BB--1
(7)
In Shampoo [15] (Algorithm 1), the above techniques are utilized to derive an efficient update. The factor matrices of the whitening metric can be directly calculated from the matrix matrix-shaped gradients:
M = E Elgg 1-1/2
~ (L R) where = E[GGT)-1/4
R= E
(8)
after which the update can be caleulated as:
u = M-
U-I-MAR--A
(9)
In practice, matrix inversion is slow, so the above inversion is performed only every N steps, and the results are cached until recomputation. This caching can result in unstable training (Figure 2), and one of the issues will discuss in the following sections.

---

Shampoo
SPlus
Cache 5uration=500
11
10
Cache Duration=100
Cache Duration=25
10
Cache Duration=10
Cache Duration=5
Loss
9
8
200
400
600
800
1000
200
400
600
800
1000
Gradient Steps
Gradient Steps
Figure 2: Shampoo is prone to divergence, but SPlus remains stable under the same settings. Plotted above are loss curves on language modelling, sweeping over learning rate between (0.0001, 0.000215, 0.000464, 0.001) and cache duration between (5, 10, 25, 100, 500). SPlus is sig- nificantly more robust to hyperparameters than Shampoo. This robustness is crucial for improving practical training speed in our setting. Shampoo diverges when caching for 100 gradient steps while SPlus remains stable, enabling a faster wall-clock performance than Adam.
4 SPlus: A Stable Whitening Optimizer
Our main contribution is SPlus, an efficient optimizer which builds upon Shampoo to stabilize training and reduce overall gradient-step and wall-clock time. In developing SPlus, we take fundamentally empirical and experimental approach we identify three core shortcomings of the naive Shampoo method, examine their causes, and propose a series of nuanced improvements to alleviate these issues. In aggregate, these changes lead to a significant improvement in reliability and training speed.
4.1 To reliably prevent divergence, utilize instant-sign normalization
As shown in Figure 2, find that naive Shampoo is prone divergence. When examining range of learning rates and matrix-inversion frequencies, we find that Shampoo diverges > 50% of the trials in our setting. Notably, Shampoo regularly diverges when matrix-inversion is cached for over 25 gradient steps.
We hypothesize that the interaction between cached matrix- inverses and incoming gradients is the cause of frequent divergence. To provide intuition on this behavior, we can rewrite the square-root matrix inverse in terms of its sigen-decompossitinn!
U Shampoo (QLA) 1/A G (QQA Q
(10)
where eigenvectors Q are orthonormal, and eigenvalues A are diagonal.
In the above decomposition, the eigenvectors can be understood as basis directions that maximally explain the covariance between gradients. Each eigenvalue represents the historical squared mag- nitude of gradients along each basis. The Shampoo update normalizes incoming gradients by their respective historical magnitudes along each basis.
The risk is when incoming gradients align with tail-end basis direction (which has small historical magnitude), which case the update can diverge. This risk is especially prominent when the cached matrix-inverse is stale, as incoming gradients may no longer align with the historical distribution.
To alleviate this risk, instead propose normalization scheme that does not rely historical magnitudes at all. Sign-normalization has been previously studied as the equivalence of Adam without a running average, accomplishing a similar normalizing behavior [5, 6]. We therefore opt to ignore historical magnitudes, and instead perform normalization instantancously via the 'sign'
"Eigen-decompee decomposition is commonly used under the hood for calculating symmetric matrix inverses, as matrix powers share the same eigenbasis, and the diagonal term easily raised any power.

---

Adam
SPlus (no updale scaling)
SPlus
Widh=
10
Width=
9.0
Wisth=
Wdth-16
Wiath-16
8.5
Wiath-16
Width=32
Wiath=
Wiath=
Wdth-64
-
8.0
Width-64
Wddh=118
Width=122
7.5
Width= 128
Validation
256
Wiath= 256
Wiath=
Wdth=512
51dth=512
7.0
Width=512
7-
6.5
6-
6.0
5.5
51
0.0001
0.001
0.01
0.0001
0. 0.001
0.01
0.1
1.0
10.0
Learning Rate
Learning Rate
Learning
Figure 3: Optimal learning rates for SPlus transfer across network widths. This is achieved by normalizing per-layer update magnitudes by constant shape-dependent factor. Notably, this learning rate transfer does not hold by default for Adam or Shampoo.
function. The eigenbasis remains as the historical eigenhasis. We refer to this update as instant-sign normalization:
U Q,sigg(QF
(11)
Instant-sign normalization has a hard bound preventing divergence. As Qr/Qn are orthonormal, and the inner sign-man matrix contains only 1 or 1, the resulting update will satisfy:
I spentral S |rall
|0 S max( Vn). von and
(12)
Additionally, instant-sign normalization provides a more fine grained elementwise normalization than naively using the Kronecker- saapckkrr-pppoximateedee inverse factors. As motivating example, under an identity eigenbasis, the Shampoo update would be:
U ~= GA -1/4 1/4
++ u == & Az)-1/4 g
(13)
where notably, the diagonal component of (Az, @ AR) is not fully expressive due to being constructed out of a Kronecker product. This notion is studied in SOAP [53], who note that Shampoo (with 1/2 power) is equivalent to rank-1 Adam approximation method alleviates this issue with an additional elementwise normalization matrix for each parameter. In contrast, instant-sign normalization does not require additional parameters in memory.
a rotated eigenbasis. Their proposed
As shown in Figure 2 (left), the instant-sign normalization of SPlus eliminates divergence across the board. Crucially, SPlus allows for the matrix-inversion to be cached for significantly longer intervals without collapse. Prior works on Shampoo perform recomputation every 10 steps [15, 47]. and our empirical findings support that this frequency needed to prevent divergence. contrast, SPlus remains stable even when results are cached for over 100 steps. Utilizing this stability, SPlus can be run at a speed which outperforms Adam in reaching an equivalent validation loss not only in gradient steps, but also in wall-clock time (Figure 2, right).
4.2 To standardize learning rate across network widths, use symmetric shape-aware scaling
Learning rate is often the first hyperparameter to tune due to its outsized impact on performance. Recent works have shown that for SGD and Adam, it is possible to naturally parameterize updates such that the optimal learning rate remains constant even as network width is adjusted [54, 55]. As shown in Figure 3, neither Shampoo nor the instant-sign update above display the correct learning-rate transfer across widths. this section, we derive simple adjustment to introduce learning-rate transfer to our setting as well.
We start by defining a desired property, following [55] - after an update, the expected magnitude of change in individual intermediate activations should be invariant, regardless of network width:
for any intermediate activation vector a::
T Σ (Ar;)==001
(14)
5

---

Lanauge Modelling (LLM)
Image Classification (VIT)
Diffusion Modelling (DIT)
4.00
0.800
SPlus
3.75
3.50
3.25
Valltation
2.75
2.50
0.780
2.25
2.00
0.775
10
10a²
10a
10-1
10a
10a
10-2
10
10c
Leaming rate
Learing rate
Learning rate
Figure 4: Iterate averaging enables the use of higher learning rates without degradation. Training with higher learning rate creates tradeoff between faster learning progress and increased parameter noise. By averaging previous iterates, parameter noise is lessened, and we can surpass the tradeoff to reveal a stronger optimal learning rate.
To achieve the above property, the norm of the update to each dense layer of shape mx must be properly scaled. In SPlus, Equation (12) states that the Frobenius norm of raw instant-sign update is /nm. Thus, when considering network with c-times larger width, one should divide the update by factor of We implement this scaling without a reference width by introducing per-layer scaling factor 2/(m + n):
U (ii.sii((.,,RROO
2/ +m+n)
(15)
As shown in Figure 3, the ahape-aware scaling factor enables a natural parametrization where learning rates transfer between network width. This property is especially useful for tuning, and enables more robust default learning rate for the optimizer.
Our scaling factor is different than the "spectral" scaling of 1/m introduced in [54], and this is by design. We found that spectral scaling was harmful in the MLP block. Consider (256, 1024) layer and (1024, 256) layer. Under spectral scaling, the first layer would have 4x larger per-parameter learning rate than the second. contrast, our method is symmetric and assigns the same learning rate both layers (as Adam does) while properly normalizing when both widths increased. We made this design choice purely on empirical findings - our average scaling outperformed the spectral scaling in all cases, see Figure 7 and provide brief discussion in Appendix A.3.
4.3 To reduce parameter noise, make use of iterate averaging
Stochastic gradient descent methods inherently follow noisy descent directions. This noise can be broadly categorized as 1) noise from using stochastically sub-sampled batch rather than the full dataset [34, 48], and 2) noise from linearizing the loss and taking a discrete step.
The second noise source is especially prominent in whitening- optimizers. In naive gradient descent, the magnitude of updates will decrease as the magnitude of gradients decreases, providing natural annealing. However, whitening-buat -baee optimizers instead utilize normalized updates which aim to take uniformly-sized updates. This behavior manifests trajectories where parameters "overshoot" their ideal value, and oscillate back and forth.
While learning rate decay (and lower learning rates in general) address these issues, they result in tradeoff of slower learning progress. At low learning rates, learning progress is slow, and loss magnitude remains large. At high learning rates, loss is again large, but for different reason the presence of large noise in the parameters. Learning rate tuning can be used to locate a tradeoff' between these two factors, however, is there is way get the best of both worlds?
We find that a more effective method of controlling the noise-progres tradeoff is via iterate averaging [39, 40, 44]. Specifically, of slow parameters are calculated via an exponential moving average of the live parameters:
of live parameters updated with a large learning rate. A second set
(16) In this way, learning progress can remain fast, yet the effect of gradient noise is diminished, as discussed in [17] and [35]. We note that parameter averaging is common technique in machine
8' - L(
+ (1 - 3)0 + 33'.

---

Algorithm 1 Shampoo
Algorithm 2 SPlus (changes in red)
for each layer gradient G do
for each layer gradient G do
G= VoL(e, x) where G R=
G = VoL(('), 2) where G E R=*x
++ (1 B2) + 3GG
L t- (1 -32 + SGG
R 11-1111))) -,) + SGTG
R 11-11) 30)++BGGGG G
G G (1 11--8)- + 3G
G (1-68 (1-31 -B1) + 3G
if mod N then
if n mod N then
I.-1/4 +- matpow L, -1/4)
Qr, Az < eigh(I.
it-1/4 + marpow R. 11/4) end if
AR + eigh(
end if
U- L - 1/4
- st QQR) 2/(m n)
k- + aU
0' - - + al
end for
0 + 8a)0 + 30
end for
learning and has been effective in a range of nomaann-spcciiic methods, e.g. image generation [19, 57], reinforcement learning [12, 51], and representation learning [14].
Figure 4 highlights the benefits of simple iterate averaging. Across the board, evaluating at the exponentially averaged parameters achieves lower validation loss. Note that series of live parameters in the averaged and non-averaged cases are equivalent. The averaged parameters more closely reveal the "true" learning progress of utilizing a higher learning rate, which is otherwise obscured by parameter noise causing an increase in validation loss.
5 How does SPlus compare to prior optimizers?
We now present a thorough evaluation of SPlus alongside previous optimizers. Intentionally, we focus specifically on the Transformer architecture [52], as it has been adapted as the backbone for most large-scale neural networks today, regardless of domain or modality [10, 30, 37, 50]. Transformers are a general architecture. This flexibility means that we must be careful in evaluating their training, to avoid overfitting on a specific domain, e.g. only language modelling. To our knowledge, our setting currently represents the widest-scope evaluation of optimizers on Transformer training.
To demonstrate robustness across settings, we examine neural networks trained on three different objectives and datasets. First, we examine an autoregressive language model (LLM), trained on the OpenWebTest [13] dataset with a sequence length of 256. Second, we examine a latent diffusion model (DiT) [37], trained via flow-matching [26] to generate Imagenet images encoded via a pretrained variational auto-encoder [43]. Third, we examine an image classification network (ViT) [10], trained to classify raw-pixel Imagenet images. All three settings are adapted directly from prior work, and utilize the same Transformer backbone.
The specific Transformer architecture is adapted from GPT-2 [8]. Layer normalization terms are applied pre-attention and pre-MI... We remove bias terms from the network. Fach objective also includes different input/output heads token embedding and logit predictor for language modelling, a patch embedder and patch output for diffusion modelling, and patch embedding and class predictor for image classification. We use a momentum of 0.9 when applicable, a linear warmup of 200 steps followed by constant schedule, and a weight decay of 0.1. We train in bfloat16. See the provided code for further details.
To thoroughly compare between optimizer types, we consider the performance across different stages of training. Concretely, we construct base checkpoints by with Adam, and saving check- points at fixed intervals (initialization, ten thousand, and fifty thousand steps). We then evalu- ate each optimizer on the three checkpoints, training for ten thousand additional gradient steps. Learning rate is swept independently for each optimizer type, along a resolution of 10 e.g. (0.0001, 0.000215, 0.000464, 0.001, Final performance is reported as validation loss after this procedure, measured on a fixed held-out validation set. The same random seed and data order are used in each run.
As loss scales vary per objective, we focus on steps-to-Adam as the main metric. We record the fraction of gradient steps and fraction of wallclock time required to match the performance of Adam
7

---

5
Adam
Shampoo
SOAP
SPlus
10000
20000
50000
3.8
3.6
3.1
3.4
3.2
3.0
3.0
2000
4000 6000
8000
10000
10000 12000 14000 16000 18000 20000
50000 52000 54000 56000 58000 60000
Figure 5: Optimizers are evaluated over 10k gradient steps, starting from three distinct check- points per objective. We design this setting test robustness across objectives and across stages of training. As shown above for the L.MM case, SPlus consistently reaches the same validation performance Adam within smaller fraction of gradient steps (dotted line).
Method
Steps-To-Ade
STA (LLM)
STA (ViT)
STA (DiT)
Time-to-Adam
Naive SGD
10.0
>
10.0
10.0
> 10.0
> 10.0
Adam
1.0
1.0
1.0
1.0
1.0
Sch.Free Adam
0.679
0.674
0.698
0.664
0.654
Sophia
1.0
1.0
n/a
n/a
> 1.0
Shampoo
0.699
0,699
Diverge
Diverge
2.426
SOAP
0,575
0.683
0.567
0.477
0.807
PSGD
0.652
0.705
0.615
0.636
0.940
Muon
0.832
>1.0
0.920
0.877
0.934
SPlus (ours)
0.439
0.419
0.504
0.396
0.617
Average values over LLM, ViT, and DIT.
Only considering non-divergent settings.
Figure 6: SPlus outperforms prior methods in both gradient steps and wallelock time, and matches Adam performance within 44 58% of the gradient steps. Learning rates swept independently for each method. We examine the training a 160M-parameter Transformer with a batch size of 1024, and sequence/patch length of 256. Results averaged starting from three base checkpoints, and the full results are in Table 3 of the Appendix.
on the task, measured via validation loss. Wall-clock results machine-specific and should be as a rough estimate; we run all experiments on the same set of 32 TPUv3 pods, a typical run takes half a day. Results are reported utilizing the best-performing learning rate for each optimizer.
We measure performance against the following optimizers which broadly span the literature:
Naive SGD, which does not modify the gradient except for scaling by global learning rate.
Adam [20]. the main baseline, which keeps track of an elementwise uncentered variance, then scales the gradient elementwise.
Schedule-Free Adam [9], which replaces traditional momentum with set of live and slow
parameters. Gradients are evaluated at a linear interpolation of the two.
Sophia [27]. which computes an elementwise estimate of the Hessian, then scales the gradient followed by a clipping step. Sophia requires auxiliary backwards passes for the Hessian calculation, and caches the results for 10 steps.
Shampoo [15], deseribed in the background section. We do not learning -rat grafting. Matrix inversion is performed every 10 steps (and otherwise diverges). We additionally consider Shampoo with 1/2 power, following [36, 47).
SOAP [53], variant on Shampoo which also tracks elementwise uncentered variance, akin to running Adam in the eigenbasis Shampoo. Matrix inversion is performed every 100 steps.

---

For frequency
1
10
25
100
250
500
Steps-to-Adane SPlus
0.78
0.71
0.75
0.76
0.77
0.77
0.81
Steps-to-Adam, Shampoo
0.86
0.80
0.80
0.79
Div.
Div.
Div.
Time per Update, SPlus
7.39
1.79
1.09
0.70
0,48
0.44
0,43
Time per Update, Shampoo
7.41
1.78
1.09
0.67
0.42
0,41
0.40
Time per Update, Adam
0.34
Time-to-Adam, SPlus
16.21
3.34
2.05
1.20
0.74
0.65
0.67
Time-to-Adam, Shampoo
18.11
3.79
2.18
1.21
Div.
Div.
Div
Table 1: Figendecomposition frequency trades off between performance and wallclock time. Each trial run for 1000 gradient steps on the LLM-Init objective. Note that specific wallclock times are machine-depender (we train on a v3-32 TPU) and may vary between systems.
Method
LLM-Init
LLM-10K
LLM-50K
Memory Usage
Adam
1.0
1.0
1.0
3mn
S.Free Adam
1.0
0.532
0.49
3mn
Adam + Iterate Averaging
0.677
0.52
0.81
4mn
Shampoo
Div.
Div.
0.699
2mn + 2(m² + n²)
SPlus (No Averaging)
0.729
0.662
0.685
2mn + 2(m2 + n²)
SPlus
0.487
0.422
0.348
3mn + 2(m2 + n²2)
SOAP
0.712
0.66
0.677
3mm + 2(m2² + n²)
SOAP Iterate Averaging
0.513
0.416
0.317
4vnn + 2(m² + n²)
Table 2: Ablating the usage of iterate averaging vs. the underlying gradient transformation.
PSGD [25], which keeps track an inverse whitening matrix calculated via iterated gradient descent rather than explicit matrix inversion. We the Kron version which is known to perform the best.
Muon [18], which performs an orthogonalization procedure on each gradient via Newton- Schulz iteration every update, without historical information. As Muon utilizes higher inherent learning rate, we use a weight decay of 0.001.
As displayed in Figure 6, SPlus is able outperform prior methods across the board both gradient steps and wallclock time. We find that a well-tuned Adam is a hard baseline to beat. For example, we were unable to match the performance Adam with Sophia (a similar finding was reported in [58]) or with Muon for LLM training. Shampoo was especially unstable (as discussed in earlier sections), and we find that when divergence does not occur, Shampoo training curves roughly match those of SOAP. as also reported in [53]. In terms of wall-clock performance, strong contender Schedule-Fres Adam, which does not perform any matrix-band based computation and only utilizes elementwise operators. For the main table, we did not heavily tune the matrix-inversion frequency for SPlus or SOAP, and assign it default value of 100 Table shows the effect of ablating this frequency.
5.1 What implementation details matter for SPlus?
We only apply the SPlus update on two-dimensional dense layers, which in the case of Transformer, composes majority of the backbone (the exception being LayerNorm scale parameters). We also do not apply SPlus to the domain-specific input and output layers - e.g. the token embedding, the classification head, and the convolutional patch layers. For these nonstandard parameters, simply set the update as the sign of the momentum values. Additionally, for nonstandard parameters where the dependee scaling term of Equation (15) is undefined, we fixed constant scaling (0.001 in our experiments). We found that this constant is not sensitive to even 10x or 0.1x pertubations, and does not need to be tuned.
We perform the above experiments over a pod of TPUv3 machines, parameters are distributed in fully- data parallel [59] setup. The SPlus update is distributed among devices. Specifically,

---

the per-step computations are calculated independently on each device as usual. However, for the eigendecomposition which occurs every N steps, we instead broadcast the L and R matrices evenly among devices. In parallel, each device then performs the eigendecomposition for its assigned matrices, then re-broadcasts the results. In this way, the most expensive step of the SPlus update can be reduced by a factor of (1 (1/num devices).
To ablate the specific changes proposed in SPlus, we show in Table 2 an ablation on the interaction between iterate-averaging and instant-sign normalization, focusing specifically on the LLM setting. Iterate averaging is consistently helpful regardless of the underlying gradient transformation. This trend may be result of our usage of a constant learning rate; as discussed in [9, 17], iterate averaging can be seen as a stand-in for learning rate decay. The downside the memory requirement for an additional set of parameters. We note that SOAP + Iterate Averaging is competetive and can outperform SPlus. One way to contrast these methods is that the instant-sign normalization used in SPlus is memory-efficient approximation of the internal Adam optimizer used in SOAP, which otherwise requires an additional set of parameters in memory.
6 Discussion and Conclusion
In this work, we present SPlus, a stable whitening optimizer for neural network training. Through fundamentally experimental approach, we introduce three key changes to improve scalability. First, direct multiplication by the square-root inverse is replaced by instant-sign normalization, which dramatically improves stability, Second, updates are correctly scaled for learning rate to transfer among network widths. Third, iterate averaging is applied to the live parameters, which reduces parameter noise and enables using a larger learning rate.
Empirically, we show that SPlus can achieve the same validation performance as Adam with 44- 58% of the gradient steps, and in ~ 62-83% the wall-clock time. Over a range of training objectives and checkpoints, SPlus achieves superior performance in comparison to previous optimizers.
Code. We provide full open-source code to replicate these experiments at github. wwt.s/oooo//o.sss//oo//o
Limitations. In comparison to Adam, SPlus requires storing 3nm + 2(n2 + m²) parameters per dense layer-three instances for the live parameters, slow parameters, and momentum, along with two Kronecker factors for the gradient covariances plus the cached eigenvectors. For square matrix, this is roughly 60% more memory than Adam. Furthermore, SPlus requires nontrivial wall-clock time for matrix cigendecomposstion. We note that SPlus (without iterate-averaging) uses the same amount of compute Shampoo, and less memory than SOAP. The training settings in this work only consider Transformer architectures. While this is an intentional choice, as neural networks today have largely converged on the Transformer backbone, it remains unanswered how performance would vary on non-Transformer architectures. We utilize a constant learning rate as is common in diffusion model training [37], but is less standard in language model training.
Future directions. We believe that SPlus, as well as our scientific setup as a whole, opens up directions in scalable neural network optimizers. We are curious the results of applying SPlus to large-scale training at the partmmeer scale. By providing a thorough evaluation setup, we hope to lower the bar of experimenting with new strategies, including extensions to SPlus. For example, alternate factorizations or low-rank approximations could reduce the computational cost of whitening, and strategy may exist in-between whitening and Hessian-based conditioning. Such ideas should be evaluated in a reproducible way, following the methodology developed here. On a practical level, we hope the efficiency and simplicity of SPlus allows the community to easily plug-am and--aa into their desired training objectives, and train neural networks faster as whole.
Acknowledgments
This work was supported part by an NSF Fellowship for KF, under grant No. DGE 2146752. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the NSF. PA holds concurrent appointments as a Professor at UC Berkeley and as an Amazon Scholar. This paper describes work performed at UC Berkeley and is not associated with Amazon. We thank Google TPU Research Cloud (TRC) for granting access to TPUs for research.
10

---

References
[1] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08874, 2023.
[2] Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Disentangling adaptive gradient methods from learning rates. arXiv preprint arXiv:2002./1800 2020.
[3] Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2): 251-276, 1998.
[4] Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018, 2020.
[5] Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology. arXiv preprint 2:Xi::2409200323 2024.
[6] Jeremy Bernstein, Yu-Xiang Wang. Kamyar Azizzadeneshelii and Animashree Anandkumar. signsgd: Compressed optimisation for non-convex problems. In International Conference on Machine Learning, pp. 560-569. PMLR, 2018.
[7] Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM review, 60(2):223 311,,,31111100111 2018.
[8] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[9] Aaron Defazio, Xingyu Yang, Ahmed Khaled, Konstantin Mishchenko. Harsh Mehta, and Ashok Cutkosky. The road less scheduled. Advances in Neural Information Processing Systems, 37:9974-10007, 2024.
[10] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint
arXiv:2010.11929 2020.
[11] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
[12] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International conference on machine learning. pp. 1587-1596. PMLR, 2018.
[13] Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus. http://Skylioon007, github. hoo/peenneeTexttCrrpss. 2019.
[14] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems, 22221221-2222 2020.
[15] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo:
Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pp. 1842-1850. PMLR, 2018
[16] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6855 6851, 2020.
[17] Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv:1803.05407.020 2018. 2018.
11

---

[18] K Jordan, Y Jin, V Boza, Y Jiacheng, F Cecista, L Newhouse, and J Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024b. URL https://kellerjordnnn github. io/posts/muon.
[19] Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyzing and improving the training dynamics of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 24174-24184, 2024.
[20] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprini arXiv:1412.6980, 2014.
[21] Mikalai Korbit, Adeyemi D Adeoye, Alberto Bemporad, and Mario Zanon. Exact nauss-newton optimization for training deep neural networks. arXiv preprint arXiv:2405.14402, 2024.
[22] Frederik Kunstner, Philipp Hennig, and Lukas Balles. Limitations of the empirical fisher approximation for natural gradient descent. Advances in nearul information processing systems, 32, 2019.
[23] Tim Large, Yang Liu, Minyoung Hub, Hyojin Bahng, Phillip Isola, and Jeremy Bernstein. Scalable optimization in the modular norm. arXiv preprint arX1v:240S.14813, 2024.
[24] Yann LeCun, Léon Bottou, Genevieve B Orr, and Klaus-Robert Müller. Efficient backprop. In Neural networks: Tricks of the trade, pp. 9-50. Springer, 2002.
[25] Xi-Lin Li. Preconditioned stochastic gradient descent. IEEE transactions on neural networks and learning systems, 29(5):1454-1455 2017
[26] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
[27] Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. arXiv preprint 2rXiv:2305.14344 2023.
[28] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 17X1v:1711.05101, 2017.
[29] Chao Ma, Wenbo Gong, Meyer Scetbon, and Edward Meeds. Swan: Preprocessing sgd enables adam-level performance on 11m training with significant memory reduction. arXiv preprint arX1v:2412.13118, 2024.
[30] Xin Ma, Yaohui Wang, Gengyun Jia, Xinyuan Chen, Ziwei Liu, Yuan-Fang Li, Cunjian Chen, and Yu Qiao. Latte: Latent diffusion transformer for video generation. arXiv preprint arXiv:2401.0300 2024.
[31] James Martens, New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21(146):1-76, 2020.
[32] James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored ap- proximate curvature. In International conference on machine learning, pp. 2408 2417. PMLR, 2015.
[33] James Martens et al. Deep learning via hessian-free optimization. In Icml, volume 27, pp. 735-748 742, 2010.
[34] Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team, An empirical model of batch training. arXiv preprint arXiv: 181200181..06622, 2018.
[35] Daniel Morales- Morales-Brotons, Thijs Vogels, and Hadrien Hendrikx. Exponential moving average of weights in deep learning: Dynamics and benefits. arXiv preprint arXiv:2411.18704 2024.
[36] Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham Kakade, and Lucas Janson. new perspective shampoo's preconditioner. arXiv preprint 177748 2024.
12

---

[37] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 4195-4205, 2023.
[38] Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.01952, 2023.
[39] Boris T Polyak. New stochastic approximation type procedures. Automat. i Telemekh, 7(98-107):
1990.
[40] Boris T Polyak and Anatoli B Juditsky. Acceleration of stochastic approximation by averaging. SIAM journal control and optimization, 30(4: 1992.
[41] Omead Pooladzandi and Xi-Lin Li. Curvature-informed sgd via general purpose lie-group preconditioners. arXiv preprint arXiv:2402.04553, 2024.
[42] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog. 188):9 2019.
[43] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10684-10695, 2022.
[44] David Ruppert. Efficient estimations from a slowly convergent robbins-monro process. Techni- cal report, Cornell University Operations Research and Industrial Engineering, 1988.
[45] Nicol N Schraudolph. Fast curvature matrix-vector products for second-order gradient descent. Neural compulation, 14(7):1723-1738, 2002.
[46] Noam Shazeer and Mitchell Stem. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596-4604. PMLR, 2018.
[47] Hao-Jun Michael Shi. Tsung- Tsung-Hsien Lee, Shintaro Iwasaki, Jose Gallego-Posan Posada, Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, and Michael Rabbat. A distributed data- aattaparallel pytorch implementation of the distributed shampoo optimizer for training neural networks at-scale. arXiv preprint arXiv:23309 06497, 2023.
[48] Samuel L Smith, Pieter-Jan Kindermans, Chris Ying, and Quoc V Le. Don't 't decay the learning rate, increase the batch size. arXiv preprint 1711:::1111004899 2017.
[49] Jascha Sohl-Dickstein. The natural gradient by analogy signal whitening, and recipes and tricks for its use. arXiv preprint /r05 1828, 2012.
[50] Octo Model Team, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees. Sudeep Dasari, Joey Hejna, Tobias Kreiman, Charles Xu, et al. Octo: An open-source generalist robot policy. arXiv preprini arXiv: 12213, 2024.
[51] Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q-learnin. In Proceedings of the AAAI conference on artificial intelligence, volume 30, 2016.
[52] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[53] Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409. 1321, 2024.
[54] Greg Yang. Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi. Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203. 03466, 2022.
[55] Greg Yang, James B Simon, and Jeremy Bernstein. spectral condition for feature learning. arXiv preprint arXin:2310.17813, 2023.
13

---

[56]
Zhirong Yang and Jorma Laaksonen. Principal whitened gradient for information geometry. Neural Networks, 21(2-3):232-4000 2008.
[57] Yasin Yaz, Chuan-Sheng Foo, Stefan Winkler, Kim-Hui Yap, Georgios Piliouras, Vijay Chan- drasekhar, et al. The unusual effectiveness of averaging in gan training. In International Conference on Learning Representations, 2018.
[58] Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham Kakade. Decon- structing what makes good optimizer for language models, arXiv preprint arXiv:2407.07972, 2024.
[59] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Out, Sam Shleifer, al. Pytorch fidp:: experiences on scaling fully sharded data parallel. arXiv preprint 12X7 2023.
14

---

A Appendix supplemental material
A.1 Pseudocode of SPlus
We provide here snippet of the core components of SPlus, implemented in JAX. For full imple- mentation, check out the repo at github.com/kvrranssppp
class SPlusState (NamedTuple):
ema: chex.Are
momentum: chex. chex.Array
sides: chex. Array
q- q_sides: chex. Array
step: int
ema_rate: float
def laalus_ge_eevl___ pa..eallaarmmmstttte)
ema_hat jax. tree_may mapee_mappaa (lambda e: / (1 state. ema_rate ** state. stete, state .ma))
return ema_hat
def splus(
learning_rate: base ScalaarrrSchehule,
b1: float 0.9,
b2: float 0.999,
ema_rate: float 0.999,
eps: float = 1e-30,
inverse_every: int 100, nonstandard_constants float 0.001, weight biight_eeca:: float = 1e-2,
mask: Optional [Union [Any, Callable [base. Params], Anyll] None, max_dim: int 10000,
def (para_n(()
momentum otu. tree_ zeros_ like (params)
ema otu. wwwettrreeesss... _eerosslikee ikke (params)
def sides_des sides_decomp(p))
if len (en(p.shape) 2:
return [jnp. zeros d)) if d < max_dim else None for d in p. shape]
return None
sides jax. tree (es___pppide__decoom decomp, params)
def qs ss_decomp(p):
if lpn shape) == 2:
return [jnp.ep eye (d) if d max dim else None for d in shape]
s_sides jax tatetrreeappppeee (qs decomp, params)
step 0 return SPlusState momentum, sides, q_sides, step, ema_rate)
def apdate_sides(e s): if len(g. shape)) == 2:
return
b2 * s[0] + (1 62) 0 g.T T if 8 [[0] is not None else None, b2 s[1] (1 b2)
g.T g if s[1] None else None,
]
else:
return None
def rot (p, q):
if len (p. shape) 2:
q[0] T @ if q[0] is not None else p
15

---

P q[1] if q[1] is not None else return P
def unrot(p, q):
if len (p. shape) == 2:
= q[0] p if [00] is not None else
p q[1].T if q[1] is not None else p
return P
jjax.jit def eit_eiivecs(s);
if s is None: return None
jmp. eigh eps jnp.eye shappe [0]))
return q
def (odate_inverseesides))
q_sides jax cree_map wat.taaappppeeeieeeeeeeeeee eigvecs, sides)
return q_sides
def pddate_fngrraee state, params): step state. step
Rotate to eigenbasis, take sign, unrotate.
momentum jax. tree tateerreepppppee map.mmppmmmmaaa m, g: m (1 b1) g, state momenttu, grads) momentum_rot = jax trretteeeaeeeee taeeeereeeooo,.....eaattt momentum, state, q_sides)
updates_rot jax toxetree_mapy (lambda m: sign momentum_rot)
updates jax tree_map wwt. (umrot, updates_rots state. q_sides)
sides jax. tree_ _aaeppmppppp (update_sidess grads, state. sides)
ema jax. tree_map (lambda g: ema_rate
+ (1 ema_rate) g, state ema, params)
Every iinerree every" steps, update the inverse eigendecomposition.
do io_irverse (step % severse_every == 0) (step == 1) q_sides jax con (do inverse, update_ interse, lambda state sidis,, sides)
return updates, SPlusState (ema, momentum, sides, q_sides, step, state ema_rate)
def shape_ (updates, state, params):
def shape_scale (path, u)
path_str '/'njoin(( [p. key for p in path]) if len (u. shape) shapee == and u. shape[o] max__in dim and u. shspape[1] max _im
scale (1 (u. shape u. shape[t])//
else:
scale nonstandard_constan
return scale
return jax wtw.tree.il....r.eaapp.....
htt.aaeree....p.i.p..aahhhpp..e. updates), None
splus_main base. GraadeentTraansoommtiiinniii, update_fn) splus_scaling = base. GraddeanTTenneoormattooo (1ammaa None, shape_scaling)
return combine. chaine
splus_main,
transform add_ decaaed_weigeee (weige (weightdeecayyyeeeeeeeeeeeeeeaae mask),
transform scale_ haale_by_learning rate (learning_rate)
splus_scaling
16

---

A.2 Full results of optimizer comparisons
We present in Table 3 and Table 4 an extended table of results, comparing optimizer performance under each specific objective and starting checkpoint. As described in Section 5, all experiments are conducted using 160M parameter Transformer model. In each plot, we compare the amount of gradient steps and/or wallclock time required to match the performance of Adam. All optimizers are trained for 10k gradient steps, and the learning rate is tuned independently. "Div." indicates that training diverges under any non-trivial learning rate. "> 1.0" indicates that at the 10k step mark, the method is unable outperform Adam.
Method
LLM-Init
LLM-10K
LLM-50K
ViT-Init
VIT-10K
VIT-50K
DiT-Init
DiT-10K
DIT-50K
Naive SGD
10.0
> 10.0
> 10.0
> 10.0
10.0
10.0
> 10.0
> 10.0
10.0
Adam
1.0
1.0
1.0
1.0
S.Free Adam
0.532
0.49
0.629
0.467
1.0
0.507
0.487
Sophia
1.0
1..
>.0
n/a
n/a
n/a
n/a
n/a
n/a
Shampoo
Div.
Div.
0,699
1.0
1.0
Div.
Div.
Div.
Div.
SOAP
0.712
0.66
0.677
0.574
0.57
0.557
0.486
0.459
0.488
PSGD
0.895
0.628
0.594
0.535
0.458
0.852
0.768
0.412
0.728
Muon
>1.
0.978
0.783
0.878
0.833
SPlus
0.487
0.422
0.348
0.586
0.475
0.452
0.459
0.359
0.371
Table 3: Full results comparing steps-to-Adam.
Method
LLM-Init
LLM-10K
LLM-50K
ViT-Init
VYT-10K
VIT-50K
DiT-Init
DÝT-10K
DŁT-50K
Naive SGD
10.0
10.0
10.0
10.0
10.0
10.0
10.0
10.0
10.0
Adam
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
S.Free Adam
1.0
0.48
0.44
1.0
0.593
0.441
1.0
0.475
0.459
Sophia
> 1.0
> >1.0
>1.0
n/a
n/a
n/a
n/a
n/a
n/a
Shampoo
Div.
Div.
2.426
1.0
Div.
Div.
Div.
Div.
SOAP
0.951
0.864
0.886
0.844
0.811
0.78
0.734
0.676
0.72
PSGD
1.08
0.836
0.808
0.854
1.18
1.11
0.73
1.064
Muon
1.0
1.0
1.0
0.996
0.79
1.0
0.915
0.881
0.827
SPlus
0.651
0.547
0.447
0.832
0.674
0.628
0.707
0.523
0.545
Table 4: Full results comparing wallclocktto-Adam.
A.3
Spectral scaling symmetric scaling
Checkpoint 0
Checkpoint 10,000
Checkpoint 50,000
Validation Loss
3.10
3.04
3.05
3.02
3.00
Leaming Rate
Spectral Scaling; 1/m
Symmetric Scaling (ours); 2/(m+n)
Figure 7: While both valid strategies that enable learning transfer across width, we find that symmetric scaling leads to a better final performance versus spectral scaling. Dotted line shows the best-performing SPlus run without scaling (i.e. without learning-rate transfer properties).
In Section 4.2, we mention difference between the SPlus symmetric scaling factor 2/(m + n) versus the "spectral" scaling [55] which argues for 1/m. The spectral scaling is the correct factor such that regardless of any dense layer input/outpu ratio, the scale of activation updates remains constant. However, we find experimentally that this property harmful for transformer training. One hypothesis is that for non-square dense layers, gradients are effectively low-rank and/or certain eigenbases do not align with incoming activation vectors.
As shown in Figure 7, the symmetric scaling that we opt for maintains the original performance, outperforming spectral scaling while allowing for learning rate transfer. The rationale behind symmetric scaling is that for the core Transformer backbones specifically the MLP block which often consists of two layers of (m, 4m) and (4m, m) size learning follows the same trajectory as no scaling factor was used at all, i.c. single global learning rate is applied to all parameters.
17

---

A.4 A discussion on the Hessian, Fisher, and Empirical Fisher.
In this section, we provide a brief discussion on various distance metrics used in optimization methods. Recall that generalized gradient descent follows the steepest direction of improvement, where distance can be defined in terms of metric matrix M:
arg min A
A0 Improvement
+1/22400 MAO
M M
(17)
Distance Penalty
where M is also referred to as preconditioner.
Hessian and Newton's method. A particularly suitable choice for M is the Hessian matrix, which is a matrix of second-order derivatives:
H ==EEDD (O)(0,
(18)
When the Hessian is used preconditioner, we arrive at Newton's 's method. Newton's 's method can be seen utilizing a quadratic approximation of the loss function rather than a linear one, where the penalty for taking large steps is defined entirely by the second-order order effects of that step the loss. For this reason, Newton's 's method is sometimes proposed as a way to avoid tuning a learning rate, and in fact for purely quadratic loss functions, Newton's method can find the global optimum in single iteration.
Intuitively, one would desire that the Hessian is positive definite, such that the second-order term always results in a positive distance penalty - however, this is generally only true for convex loss functions. When the Hessian has negative eigenvalues, Newton's 's method can step in non-descending directions, or even diverge completely.
matrix is an approximation to the Hessian using only first-order terms. For simplicity, let's assume the loss function is the equan- error of a single output. Denoting the network output as f(0, 27), we can expand the Hessian for single 2, pair:
Gauss-Newton matrix. The
(8L(0,2) (ef(0, x Vof +
y)
(19)
Gauss-Newton term
Dropped second-order term
The Gauss-Newtor approximation is often desired as it is strictly positive semi-definite, avoiding negative distance issues that the full Hessian has. Additionally, the Gauss-Newton term simple to calculate as it only requires first-order gradients. The Gauss-Newton can be generalized to non-MSE losses [7, 21, 45] by introducing PSD matrix A between the two gradient terms:
G E [Vo A (of(0,
(20)
Fisher information matrix and natural gradient descent. For neural networks defining probability distributions, we can use a metric that is particular to the distribution itself:
F =
Ve log Po ( V log (o(l)))?)
(21)
which known as the Fisher information matrix. The Fisher does not take the loss function into account. It is only affected by the shape of the probability distribution itself, as defined by the current neural network. When descent is performed using the Fisher as a preconditioner, it is often referred to as natural gradient descent [3, 49].
Natural gradient descent has the nice property that it is invariant to the parameterization of the network to first order, optimization will follow the same trajectory regardless of neural networks' internal structure. This property also true for Newton's 's method under affine transformations of parameters.
Note the expectation in Equation (21), which notes that must be sampled from the current distribu- tion. This means that the Fisher cannot be calculated by taking the loss over samples from the dataset, and must instead use sampled outputs.
The Fisher can also be interpreted the Hessian of a particular loss function, namely the expectation of log- likelihoooo under sampled outputs:
L Egw (x) [log po (w(x:)
(22)
where notably the second-order terms in the form of Equation (19) evaluate to zero
18

---

Empirical Fisher. An approximation often used in practice is to calculate a Fisher-like like matrix, but over dataset labels:
F P Vo log pa (v log (og
(23)
This should not be confused with the true Fisher, as studied in [34]. In fact, the empirical Fisher is generalized Gauss-Newton matrix
closer in nature to the Gauss-Newton matrix, and is equivalent where the inner PSD matrix is constructed as A = V ((0,
The empirical Fisher can be seen an (uncentered) covariance of gradients. The Fisher itself is actually centered covariance, since the expectation of ongllikelhood gradients under sampled outputs zero. In practice, we found that centering the empirical Fisher held no practical difference.
Whitening metric. The whitening metric [56] is the matrix square-root of the empirical Fisher:
W = Ez.yr log po(y Ve log (v|x²²² 1/2
(24)
This metric is widely used in neural network training, for example, Adam is diagonal estimate of the whitening metric, and Shampoo is Kronecker approximation of this term. The name whitening refers to property that when projected onto the whitening metric basis, the resulting preconditioned gradients have identity covariance:
Coo V, V) =
where
w L(0,z).
(25)
The whitening metric has the same eigenbasis the empirical Fisher, they symmetric matrix powers of each other.
19

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
We describe the empirical setting, then show the according results in the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We include limitations section the final page.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
20

---

Justification:
We do not provide theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibiliyy
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We have put forth best-faith effort to provide reproducible training details and methodology in section 5. Nuanced specifics can be found within the provided code.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
21

---

Answer: [Yes]
Justification:
We only utilize open-source datasets in this work, and provide access to our code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: //n /nips. ctsips.ccpuubic/guiii wttiiicgggde/////eee Sodeeee/ wwwpomisssiooooooooooooooooo onlolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We have put forth best-faith effort provide reproducible training details and methodology section Nuanced specifics can be found within the provided code.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Due computational constraints, we are only able to run one trial each for each optimizer, under specific objective, checkpoint, and learning rate. That said, we control for random seed and data order, and heuristically the training runs do not large independent variance.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
22

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We describe the machines used to run the experiments, and give an estimate of the time taken.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: //te::////eeeee http//nurrrrrreeeeeeeee http.///pc...i//.icccici......e/
Answer: [Yes]
Justification:
We follow the ethics guidelines.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: We do not find the need to discuss specific societal impacts of this work, as purely relates computational problem.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
23

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We propose a computational method, and do not release any new data.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite the two spercsource datasets that we utilize.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswi wwwwweettoe...... com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
24

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not release assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moodl as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
We do involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We do not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
25


Reviews:
[{"review_id": "UE3TV213pF", "reviewer": "Reviewer_mSRL", "summary": "This paper introduces SPlus, a novel optimizer designed to stabilize and accelerate training for large neural networks, particularly Transformers. It builds upon the Shampoo family of whitening-based optimizers, which estimate a preconditioning matrix via Kronecker-factored approximations of gradient covariances.\n\nThe authors identify three critical weaknesses in existing Shampoo-like optimizers: \n- Instability from stale matrix inverses. \n- Learning rate scaling issues across network widths. \n- Parameter noise at high learning rates.\nTo address these, SPlus introduces:\n- Instant-sign normalization in place of magnitude-based scaling, ensuring bounded updates and eliminating the divergence observed in Shampoo.\n- Shape-aware scaling to ensure that optimal learning rates can transfer across different model widths.\n- Iterate averaging to mitigate parameter noise while maintaining fast learning dynamics.\n\nExtensive experiments on Transformer architectures (language modeling, image classification, diffusion modeling) show that SPlus matches the validation loss of Adam in ~44% fewer gradient steps and in ~62% of the wall-clock time.", "questions": "One paper I think is very related with the proposed method is NGPlus [1], which is a second-order method and saves computational cost compared with KFAC and Shampoo. \n\n[1] Yang, Minghan, et al. \"An efficient Fisher matrix approximation method for large-scale neural network optimization.\" IEEE Transactions on Pattern Analysis and Machine Intelligence 45.5 (2022): 5391-5403.\n\nDid you experiment with low-rank approximations for L and R? For example, truncating eigen decompositions might reduce both memory and compute costs. Do you have empirical results or intuition on how such approximations might affect stability or convergence speed?\n\nHave you measured wall-clock speed on significantly larger models (e.g. billion-parameter LLMs)? The eigen decompositions could become disproportionately costly.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\nThe paper convincingly demonstrates how cached matrix inverses can destabilize training in whitening-based optimizers. The explanation of eigen-decomposition and sign-normalization provides strong intuition for SPlus’ improvements.\n\n\nAdapting insights from width-scaling theory to whitening optimizers is a significant practical contribution. This should reduce the trial-and-error process for hyperparameter tuning as models scale.\n\nEvaluations are thorough, covering different tasks and various training stages. SPlus consistently outperforms other methods, including Shampoo, SOAP, PSGD, and Sophia.\n\n\nWeaknesses\n\nAll experiments focus on Transformer-based models. It is unclear whether the same improvements would hold for other architectures such as CNNs, RNNs, or graph neural networks. \n\nSPlus increases memory usage significantly (∼60% more than Adam). For very large models (e.g. billion-parameter LLMs), this could become a practical barrier. While the authors mention low-rank factorization as future work, no experiments or ablations are presented to quantify potential savings.\n\nSPlus performs eigendecomposition every N steps (defaulting to 100). The paper does not fully explore how sensitive performance is to this frequency, especially in larger models where eigen decompositions can become bottlenecks.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review and suggestions. Please see our detailed response and additional experiments below:\n\n* **How does eigendecompisition frequency affect performance and wallclock time?**\n\nThis 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. \n\n*Validation Loss after 1K gradient steps:*\n|Frequency | 1 | 5 | 10 | 25 | 100 | 250 | 500 |\n| -- | -- | -- | -- | -- | -- | -- | -- |\n|SPlus | **3.316** | 3.332 | 3.332 | 3.335 | 3.336 | 3.336 | 3.337 |\n|Shampoo | 3.332 | 3.342 | 3.340 | 3.340 | Div. | Div. | Div. |\n| Adam | 3.36 |\n\n*Steps-to-Adam*:\n|Frequency | 1 | 5 | 10 | 25 | 100 | 250 | 500 |\n| -- | -- | -- | -- | -- | -- | -- | -- |\n|SPlus | 0.78 | **0.71** | 0.75 | 0.76 | 0.77 | 0.77 | 0.81|\n|Shampoo | 0.86 |   0.80  |   0.80 |    0.79 | Div. | Div. | Div. |\n\n*Average Time Per Update*:\n|Frequency | 1 | 5 | 10 | 25 | 100 | 250 | 500 |\n| -- | -- | -- | -- | -- | -- | -- | -- |\n|SPlus | 7.39 | 1.79 | 1.09 | 0.70 | 0.48 | 0.44 | 0.43 |\n|Shampoo | 7.41 | 1.78 | 1.09 | 0.67 | 0.42 | 0.41 | 0.40 |\n|Adam | 0.34 |\n\n*Time-to-Adam*:\n|Frequency | 1 | 5 | 10 | 25 | 100 | 250 | 500 |\n| -- | -- | -- | -- | -- | -- | -- | -- |\n|SPlus | 16.21 | 3.34 | 2.05 | 1.20 | 0.74 | **0.65** | 0.67 |\n|Shampoo | 18.11 | 3.79 | 2.18 |1.21 | Div. | Div. | Div. |\n\nPerformance 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.\n\nWe 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.\n\n* **Memory Usage of SPlus**\n\nThe 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:\n\n*Validation Loss (no iterate averaging)*:\n| Method | LLM-Init | LLM-10K | LLM-50K | Approx Memory Usage| \n| -- | -- | -- | -- | -- |\n| Adam  | 3.132  | 3.114 | 3.012 | 3nm |\n| SPlus No EMA  | 3.087  | 3.057 | 3.010 | 2nm + 2(n^2 + m^2) |\n| SOAP | 3.085 | 3.072 | 2.995 | 3nm + 2(n^2 + m^2) |\n\n*Steps-to-Adam (no iterate averaging)*:\n| Method | LLM-Init | LLM-10K | LLM-50K | Approx Memory Usage| \n| -- | -- | -- | -- | -- |\n| Adam  | 1.0  | 1.0 | 1.0 | 3nm |\n| SPlus No EMA  | 0.729  | 0.662 | 0.685 | 2nm + 2(n^2 + m^2) |\n| SOAP | 0.712 | 0.660 | 0.677 | 3nm + 2(n^2 + m^2) |\n\n* **Potential methods of improving compute efficiency**\n\nAs 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.\n\nWe 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."}, {"review_id": "CgA4kw9vv0", "reviewer": "Reviewer_oPR8", "summary": "This paper proposes a new optimizer termed SPlus.  The algorithm is positioned as an improvement over shampoo.  SPlus maintains exponential moving averages of the left and right squared gradients, and then normalizes the new gradient by conducting sign-normalization in the shampoo eigenbasis, which is re-computed every so often.  The authors say that SPlus allows for these eigenvalue computations to be performed less regularly than in shampoo, which yields wall clock gains.  The main justification is empirical effectiveness.  SPlus also includes two other components: a scaling adjustment which allows the learning rate to be transferred across widths (as in Mu-P), and iterate averaging, which averages out the parameter noise/oscillations and is claimed to allow larger learning rates.", "questions": "If I understand things correctly, if we put aside iterate averaging and shape-aware scaling, then SPlus is very similar to Muon.  In particular, SPlus with $\\beta_2 = 0$ is almost exactly Muon (or rather, an idealized version of Muon which exactly computes the matrix sign rather than using the Newton-Schultz iteration), with the only remaining differences being: whether L/R is computed from $G$ (SPlus) or $\\overline{G}$ (Muon); and (2) whether momentum is parameterized as $\\overline{G} \\leftarrow \\(1 - \\beta_1) \\overline{G} + \\beta_1 G$ (SPlus) vs $\\overline{G} \\leftarrow \\beta_1 \\overline{G} + G$ (Muon), which are the same up to a rescaling of the learning rate by $1 / (1 - \\beta_1)$.  I am trying to understand which of the differences between SPlus and Muon are responsible for SPlus's superior performance that was observed in experiments.  I have the following questions:\n\n - If you remove iterate averaging and shape-aware scaling from SPlus, or alternatively if you add them to Muon, how do the two algorithms compare?  I am especially thinking about iterate averaging.\n- How does SPlus with $\\beta_2 = 0$ compare to SPlus with $\\beta_2 = 0.9$ (i.e the value used in your experiments)?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n - the experiments (all on transformers) show that SPlus beats not only Adam but also several newer methods including Muon, SOAP, and PSGD.\n\nWeaknesses:\n - I have questions about the comparison to Muon  -- see below\n - There are a number of mathematical typos: \n   - equation 1, which is the first equation of the paper, is wrong; the RHS should be $-(1/2\\alpha) g$, not $\\alpha g$\n   - equation 3 is also wrong; the RHS should be $- \\frac{1}{2} M^{-1} g$\n   - equation 11, which is the main equation describing the proposed algorithm, is wrong; the transpose should be on the second $Q_L$ and not the first one\n   - I believe the SPlus algorithm box is missing $2$ subscripts for some of the $\\beta_2$'s.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review and analysis. Please find our detailed response below:\n\n* **Relation and comparison to Muon**\n\nRegarding 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.\n\nRegarding 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:\n\n*Steps-to-Adam:*\n| Method | LLM-Init | LLM-10K | LLM-50K | ViT-Init | ViT-10K | ViT-50K | DiT-Init | DiT-10K | DiT-50K |\n|---|---|---|---|---|---|---|---|---|---|\nSOAP |  0.712 | 0.66 | 0.677 | 0.574 | 0.57 | 0.557 |0.486 | 0.459 | 0.488 |\nMuon (new) |0.636|0.68|0.91|0.567|0.652|1.0+|0.556|0.507|0.507|\nSPlus |  0.487 | 0.422 | 0.348 | 0.586 | 0.475 | 0.452 | 0.459 | 0.359 | 0.371\n\n*Time-to-Adam:*\n| Method | LLM-Init | LLM-10K | LLM-50K | ViT-Init | ViT-10K | ViT-50K | DiT-Init | DiT-10K | DiT-50K |\n|---|---|---|---|---|---|---|---|---|---|\nSOAP |  0.951 | 0.864 | 0.886 | 0.844 | 0.811 | 0.78 | 0.734 | 0.676 | 0.72 |\nMuon (new) |0.611|0.66|0.89|0.618|0.698|1.0+|0.57|0.515|0.59|\nSPlus |  0.651 | 0.547 | 0.447 | 0.832 | 0.674 | 0.628 | 0.707 | 0.523 | 0.545 |\n\nAs 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. \n\n* **Ablating the Effect of Iterate Averaging**\n\nWe 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.\n\n*Validation Loss*:\n| Method | LLM-Init | LLM-10K | LLM-50K |\n| -- | -- | -- | -- | \n| Adam  | 3.132  | 3.114 | 3.012 | \n| SPlus No EMA  | 3.087  | 3.057 | 3.010 | \n| SOAP | 3.085 | 3.072 | 2.995 |\n| Muon | 3.061 | 3.061 | 3.007 |\n| SPlus  | 3.003  | 2.978 | 2.936 | \n| SOAP + EMA | 2.998| 2.993 | 2.931 |\n\n*Steps-to-Adam:*\n| Method | LLM-Init | LLM-10K | LLM-50K |\n| -- | -- | -- | -- | \n| Adam  |  1.0 | 1.0 | 1.0 | \n| SPlus No EMA  | 0.729  | 0.662 | 0.685 | \n| SOAP | 0.712 | 0.660 | 0.677 |\n| Muon | 0.636| 0.68| 0.91 |\n| SPlus  | 0.487 | 0.422 | 0.348 | \n| SOAP + EMA | 0.513 | 0.416 | 0.317 |\n\nSPlus 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.\n\nRegarding mathematical typos, thank you for pointing these out, and they have been fixed for the next revision.\n\nWe 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.\n\n\n[1] Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology. 2024."}, {"review_id": "zMyKQZqoue", "reviewer": "Reviewer_Xpwz", "summary": "This paper introduces SPlus, a stable and efficient optimizer that addresses three key limitations of the Shampoo algorithm: susceptibility to divergence, shape sensitivity across network widths, and elevated gradient noise under large learning rates. To mitigate these issues, the authors propose three targeted modifications. Extensive evaluations on Transformer architectures across three representative tasks demonstrate that SPlus yields substantial improvements in both convergence speed and wall-clock efficiency.", "questions": "1. The effect of the averaging coefficient $\\beta_3$ in parameter averaging is unclear. It would be helpful to provide empirical sensitivity analysis or practical guidance on choosing $\\beta_3$ .\n2. Can a fixed choice of the hyperparameter $\\beta_3$  maintain an appropriate update magnitude across different training phases? Specifically, does the parameter averaging mechanism risk inducing insufficient updates in the early stages, or causing instability in later training due to oscillatory behavior?\n3. Could additional theoretical or intuitive justification be provided for why parameter averaging helps mitigate optimization noise?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n\n1.The paper proposes targeted solutions to common optimization challenges such as instability and sensitivity to hyper-parameters of Shampoo algorithm, making approximately second-order optimizer potentially applicable to train deep neural networks. \n\n2.The experiments cover multiple training stages and diverse tasks (LLM, ViT, DiT), with comparisons against a strong set of baselines. SPlus shows clear improvements in both convergence and wall-clock efficiency.\n\n3.The paper is well-written and easy to follow. Source code is provided, supporting strong reproducibility.\n\n \n\n**Weaknesses**\n\n**Limited Theoretical Depth:** The proposed improvements are primarily empirical in nature, targeting issues such as instability and gradient noise during optimization. While effective, techniques like instant-sign normalization and parameter averaging build on existing methods, and the paper lacks deeper theoretical analysis—e.g., why parameter averaging effectively mitigates noise in this setting.\n\n**Increased Memory Overhead:** As noted in the limitations, the method incurs additional memory cost compared to Adam, due to the need to maintain historical parameters for averaging, which may limit its applicability in resource-constrained environments. It is better for this paper provide the comparison of wall-clock time cost. \n\n**Limited Architectural and Scale Generalization:** The generalization of SPlus to larger-scale models and non-Transformer architectures remains untested. The steps-to-Adam metric is not sufficient to evaluate the performance. It is better to provide the final results (e.g., Accuracy on ImageNet classification and others for training LLMs) when using different optimizer with equal/sufficient training steps.  Besides, It is not clear   how the introduced hyper-parameters $\\beta_3$ affects the performance. I think this paper should provide the ablation study for hyper-parameters $\\beta_3$ of SPlus. \n\n\n\nOther minors:\n\n(1) The notation is not self-contained, e.g, What is the notation $g, \\Delta \\theta$ meaning? This paper should provide the illustration for self-containing.  Similar problems also happens in Eqn. 17, what is $\\theta^{'}$? It seems has something wrong in the left formulation  of Eqn.17\n\n (2)The axes in Figures 3 and 5 lack clear labels.", "was_revised": true, "final_justification": "The proposed method is effecitve and the writing is good, and my concerns are well addressed in rebuttal. I keep positive with \"Borderline accept\"", "rebuttal": "Thank you for your review and detailed questions. Please see the detailed response:\n\n* **Ablating iterate averaging $\\beta_3$**\n\nTo 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:\n\n*Validation Loss of SPlus*:\n| Starting Checkpoint | $\\beta_3 = 0$ | $\\beta_3 = 0.9$ | $\\beta_3 = 0.99$ | $\\beta_3 = 0.995$ | **$\\beta_3 = 0.999$** | $\\beta_3 = 0.9999$ | \n| -- | -- | -- | -- | -- | -- | -- |\n| LLM at Init  | 3.085 | 3.064 | 3.028 | 3.017 | **3.003** | 3.077 |\n| LLM at 10K  | 3.062 | 3.047 | 3.012 | 3.002 | **2.988** | 3.303 |\n| LLM at 50K | 3.010 | 2.995 | 2.964 | 2.954 | **2.936** | 3.035 |\n\n*Steps-to-Adam of SPlus*:\n| Starting Checkpoint | $\\beta_3 = 0$ | $\\beta_3 = 0.9$ | $\\beta_3 = 0.99$ | $\\beta_3 = 0.995$ | **$\\beta_3 = 0.999$** | $\\beta_3 = 0.9999$ | \n| -- | -- | -- | -- | -- | -- | -- |\n| LLM at Init  | 0.73 | 0.616 | 0.479 | **0.461** | 0.486 | 0.771 |\n| LLM at 10K  | 0.667 | 0.556 | 0.409 | **0.384** | 0.419 | Div. |\n| LLM at 50K | 0.998 | 0.792 | 0.495 | 0.425 | **0.383** | Div. |\n\nAs 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.\n\n\n* **Memory Usage of SPlus**\n\nThe 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:\n\n*Validation Loss (no iterate averaging)*:\n| Method | LLM-Init | LLM-10K | LLM-50K | Approx Memory Usage| \n| -- | -- | -- | -- | -- |\n| Adam  | 3.132  | 3.114 | 3.012 | 3nm |\n| SPlus No EMA  | 3.087  | 3.057 | 3.010 | 2nm + 2(n^2 + m^2) |\n| SOAP | 3.085 | 3.072 | 2.995 | 3nm + 2(n^2 + m^2) |\n\n*Steps-to-Adam (no iterate averaging)*:\n| Method | LLM-Init | LLM-10K | LLM-50K | Approx Memory Usage| \n| -- | -- | -- | -- | -- |\n| Adam  | 1.0  | 1.0 | 1.0 | 3nm |\n| SPlus No EMA  | 0.729  | 0.662 | 0.685 | 2nm + 2(n^2 + m^2) |\n| SOAP | 0.712 | 0.660 | 0.677 | 3nm + 2(n^2 + m^2) |\n\n* **Notation and figures**\n\nRegarding 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$.\n\n* **Why do we use the Steps-to-Adam metric?**:\n\nWe 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.\n\n* **Regarding the Transformer-only experiments**\n\nIt 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.\n\nWe 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.\n\n[1] Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training. Lui et al, 2023.\n\n[2] SOAP: Improving and Stabilizing Shampoo using Adam. Vyas et al, 2024."}, {"review_id": "JIC9VzxKm9", "reviewer": "Reviewer_YNFF", "summary": "This paper takes an empirical approach to study neural network optimization based on the Shampoo family of algorithms. \nThe authors identify three key issues with the Shampoo method and address them by proposing a method called SPlus. \n\nThe three key issues are namely that: 1) Shampoo diverges when matrix-inverses are cached for too long, 2) Shampoo updates are not properly scaled with the network width and 3) Shampoo is more sensitive to high learning rates, which are addressed by 1) performing a so-called \"instant-sign normalization\", 2) by introducing a scaling factor of $1/m$, where $m$ are the number of rows in the layer gradient $G \\in \\mathbb{R}^{m \\times n}$, and finally by 3) using a iterate-averaging scheme, higher learning rates can be used. \n\nThe proposed method SPlus is compared against other preconditioned optimizers, including Adam, Sophia, Shampoo, SOAP, and Muon on Transformer training on language modeling, image classification and diffusion modeling, which highlight the proposed method over standard Adam.", "questions": "- I wonder whether the hypothesis of the authors for the divergence of Shampoo when the matrix-inversion is cached too long, can be empirically tested? Perhaps by checking how well the gradient and matrix-inverses are aligned throughout training?\n\n- Can the authors also provide ablation studies for the final comparison in Steps-To-Adam and Time-to-Adam for the different modifications that they proposed? (i.e. only using the instant-normalization, only using shape-aware scaling, only using iterate-averaging, as well as \ninstant-normalization + iterate-averaging and instant-normalization + shape-aware scaling?) It would be sufficient for me to see it for a single task, e.g. only on the LLM or ViT.", "limitations": "yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 1, "strengths_and_weaknesses": "Strength:\n- The paper is clearly written and easy to follow. \n- The improvement through the suggested modifications are illustrated well and together they are able to outperform Adam (and many other optimizers) both in terms of iterations and clock-time across all tested settings. \n- The authors provide Code, details about the experiments as well as Pseudocode of SPlus in the appendix. \n\nWeaknesses:\n- The proposed modifications are not necessarily novel per se. Instant-sign normalization is essentially Signum [1] in Shampoo's eigenbasis, which is reminiscent of SOAP being a rank-1 Adam approximation in a rotated eigenbasis. Also iterate-averaging has already been quite known [2]. As such the I am not sure about the novelty of this work. (The authors also refer to previous work, this is not my critique point.)\n- I think the comparison of SPlus to Adam is overselling the proposed method a bit. If one checks Figure 6, then many of recently proposed methods are quite competitive to SPlus. For instance Schedule-Free Adam is only 6% slower than SPlus. Given that Schedule-Free Adam improves considerably on Adam, I wonder how a Schedule-Free version of SOAP, PSGD and Muon would fare. \n- In general I am always a bit skeptical about wall-clock time runs, as it also highly depends on how each of the methods are implemented in Code, which have the risk to obscure insights. For instance the authors describe how the updates are computed in SPlus in lines 270-276 and how this is optimized for SPlus (which is a great thing to mention), but it is less clear how other methods are implemented and used. \n- The ablation studies clearly show the effect and benefit of each modification to Shampoo. \n\n[1] Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Animashree Anandkumar. 320 signsgd: Compressed optimisation for non-convex problems. In International Conference on 321 Machine Learning, pp. 560–569. PMLR, 2018\n[2] Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon 348 Wilson. Averaging weights leads to wider optima and better generalization. arXiv preprint 349 arXiv:1803.05407, 2018.", "was_revised": true, "final_justification": "This paper identifies three key limitations in Shampoo and proposes a new optimizer SPlus, which addresses all of them by introducing three modifications, namely 1) instant-normalization, 2) a width-dependent scaling factor and 3) employing iterate-averaging. \n\nThe discussion phase was more focused on understanding the effect of point 1) and 3) in SPlus through further ablation studies, and its comparison to other methods from the family of \"whitening optimizers\", including Muon, SOAP and Shampoo. \n\nI suggest acceptance by giving a score of 4 because the proposed method SPlus does address important limitations in Shampoo, further expanding the family of \"whitening optimizers\". The reason for not giving a higher score is because the proposed modifications are not that novel and mainly taken from other optimization methods (e.g. instant-sign normalization resembles Signum and the benefit of iterate-averaging was already known for SGD.)", "rebuttal": "Thank you for your detailed review and questions. Please see our response below:\n\n* **Ablating the proposed changes**\n\nAddressing 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. \n\nTo 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)\n\n*Steps-to-Adam:*\n| Method | LLM-Init | LLM-10K | LLM-50K |\n| -- | -- | -- | -- | \n| Adam  |  1.0 | 1.0 | 1.0 | \n| Shampoo  | Div.  | Div. | 0.699 | \n| Shampoo + Instant-Sign (i.e. SPlus without EMA) | 0.729  | 0.662 | 0.685 |\n| Shampoo + Instant-Sign + Iterate Averaging (SPlus) | 0.487 | 0.422 | 0.348 |\n| SOAP | 0.712 | 0.660 | 0.677 |\n| SOAP + Iterate Averaging | 0.513 | 0.416 | 0.317 |\n\nAs 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.\n\n* **Wall-clock time and inversion frequency**\n\nThank 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.\n\nTo 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. \n\n*Validation Loss after 1K gradient steps:*\n|Frequency | 1 | 5 | 10 | 25 | 100 | 250 | 500 |\n| -- | -- | -- | -- | -- | -- | -- | -- |\n|SPlus | **3.316** | 3.332 | 3.332 | 3.335 | 3.336 | 3.336 | 3.337 |\n|Shampoo | 3.332 | 3.342 | 3.340 | 3.340 | Div. | Div. | Div. |\n| Adam | 3.36 |\n\n*Steps-to-Adam*:\n|Frequency | 1 | 5 | 10 | 25 | 100 | 250 | 500 |\n| -- | -- | -- | -- | -- | -- | -- | -- |\n|SPlus | 0.78 | **0.71** | 0.75 | 0.76 | 0.77 | 0.77 | 0.81|\n|Shampoo | 0.86 |   0.80  |   0.80 |    0.79 | Div. | Div. | Div. |\n\n*Average Time Per Update*:\n|Frequency | 1 | 5 | 10 | 25 | 100 | 250 | 500 |\n| -- | -- | -- | -- | -- | -- | -- | -- |\n|SPlus | 7.39 | 1.79 | 1.09 | 0.70 | 0.48 | 0.44 | 0.43 |\n|Shampoo | 7.41 | 1.78 | 1.09 | 0.67 | 0.42 | 0.41 | 0.40 |\n|Adam | 0.34 |\n\n*Time-to-Adam*:\n|Frequency | 1 | 5 | 10 | 25 | 100 | 250 | 500 |\n| -- | -- | -- | -- | -- | -- | -- | -- |\n|SPlus | 16.21 | 3.34 | 2.05 | 1.20 | 0.74 | **0.65** | 0.67 |\n|Shampoo | 18.11 | 3.79 | 2.18 |1.21 | Div. | Div. | Div. |\n\nPerformance 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."}, {"review_id": "5ZQ0QM2Ge3", "reviewer": "Reviewer_usQx", "summary": "This paper proposes a novel whitening optimizer for neural network training that improves the convergence of learning. The major contributions include: 1. stability enhancement and computational saving for long-term matrix calculations; 2. enabling learning rate transfer across network width; 3. iterate-average learning scheme for faster learning. Numerical studies support the effectiveness of the proposed optimizer with signficant improvements on learning efficiency.", "questions": "1. Why is the additional wall-clock time required for SPlus negligible? (Page 9) If this is merely a numerical observation, please clarify and provide as much detail as possible for theoretical justifications.\n\n2. Please specify how much more is the significant wall-clock cost for matrix eigen-decomposition with justifications. (Page 9) Monetary/space/time costs and their tradeoffs should be made clearer instead of only mentioning the time efficiency.", "limitations": "No. The tradeoffs among monetary/space/time costs should be better justified and compared with existing approaches for better supporting the advantages of the learning scheme and clearer statement on its disadvantages.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n1. The paper is well written.\n\n2. Using bounded update with simple iterative-averaging scheme is novel and enhances the network stability.\n\n3. Numerical results clearly show the superiority of the proposed optimizer over previous ones.\n\nWeaknesses\n\nThe computational cost of the proposed optimizer should be stated and justified more clearly.", "was_revised": true, "final_justification": "The complexity issue was addressed.", "rebuttal": "Thank you for review. It appears the main concern is regarding the wall-clock requirements for the eigendecomposition.\n\n* **Implementation of Eigendecomposition**\n\nWe 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.\n\nWe 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.\n\n* **Ablation on Eigendecomposition Frequency**\n\nTo 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. \n\n*Validation Loss after 1K gradient steps:*\n|Frequency | 1 | 5 | 10 | 25 | 100 | 250 | 500 |\n| -- | -- | -- | -- | -- | -- | -- | -- |\n|SPlus | **3.316** | 3.332 | 3.332 | 3.335 | 3.336 | 3.336 | 3.337 |\n|Shampoo | 3.332 | 3.342 | 3.340 | 3.340 | Div. | Div. | Div. |\n| Adam | 3.36 |\n\n*Steps-to-Adam*:\n|Frequency | 1 | 5 | 10 | 25 | 100 | 250 | 500 |\n| -- | -- | -- | -- | -- | -- | -- | -- |\n|SPlus | 0.78 | **0.71** | 0.75 | 0.76 | 0.77 | 0.77 | 0.81|\n|Shampoo | 0.86 |   0.80  |   0.80 |    0.79 | Div. | Div. | Div. |\n\nAs 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:\n\n*Average Time Per Update*:\n|Frequency | 1 | 5 | 10 | 25 | 100 | 250 | 500 |\n| -- | -- | -- | -- | -- | -- | -- | -- |\n|SPlus | 7.39 | 1.79 | 1.09 | 0.70 | 0.48 | 0.44 | 0.43 |\n|Shampoo | 7.41 | 1.78 | 1.09 | 0.67 | 0.42 | 0.41 | 0.40 |\n|Adam | 0.34 |\n\n*Time-to-Adam*:\n|Frequency | 1 | 5 | 10 | 25 | 100 | 250 | 500 |\n| -- | -- | -- | -- | -- | -- | -- | -- |\n|SPlus | 16.21 | 3.34 | 2.05 | 1.20 | 0.74 | **0.65** | 0.67 |\n|Shampoo | 18.11 | 3.79 | 2.18 |1.21 | Div. | Div. | Div. |\n\nPerformance 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.\n\nWe 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."}]

Author rebuttal:

Meta-review: The paper introduces SPlus, an optimizer designed to improve the stability and acceleration of training large neural networks and especially transformers. It addresses three key weaknesses of existing Shampoo-like optimizers: instability from stale matrix inverses, learning rate scaling issues across network widths, and parameter noise at high learning rates. SPlus proposes instant-sign normalization for bounded updates, shape-aware scaling for optimal learning rate transfer, and iterate averaging to mitigate parameter noise. Through extensive experiments on transformer architectures, SPlus demonstrates significant improvements in convergence speed and wall-clock time over Adam.

A positive point of the paper is how it demonstrates how cached matrix inverses destabilize training in whitening-based optimizers and SPlus' improvements through eigen-decomposition and sign-normalization. Adapting width-scaling theory insights to whitening optimizer is a significant practical contribution that should reduce hyperparameter tuning effort. Thorough evaluations across different tasks and training stages show SPlus' consistent outperformance of various methods (including Shampoo, SOAP, PSGD, and Sophia). The paper is clearly written, easy to follow, and additionally provides source code.

Despite its strengths, several concerns were raised. 
* There is limited architectural generalization, as all experiments focus on transformers.  (And while low-rank factorization is mentioned as future work, no experiments quantify potential savings)
* The increased memory usage of SPlus (∼60% more than Adam), which could be a practical barrier.
* There is doubt over the quality of the muon baseline (but the authors have/will address this in revision)
* The lack of exploration into the sensitivity of performance to the frequency of eigendecomposition, which can become a bottleneck in larger models. 
* There is a lack deeper theoretical analysis for some of the proposed modifications, such as why parameter averaging effectively mitigates noise.

The reviewers explicitly found some minor corrections, which should be addressed, and the authors have promised a list of changes. 

(Accept)
Decision: Accept (poster)