SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training
Authors: Chao Ma
Wenbo Gong
Meyer Scetbon
Edward Meeds
Venue: ICML 2025 poster (2025)
Keywords: Efficient LLM optimization
Stateless Optimizers
Primary area: deep_learning->algorithms
TL;DR: Training LLMs without storing optimizer states

Abstract: Adaptive optimizers such as Adam (Kingma & Ba, 2015) have been central to the success of large language models. However, they often require maintaining optimizer states throughout training, which can result in memory requirements several times greater than the model footprint. This overhead imposes constraints on scalability and computational efficiency. Stochastic Gradient Descent (SGD), in contrast, is a stateless optimizer, as it does not track state variables during training. Consequently, it achieves optimal memory efficiency. However, its capability in LLM training is limited (Zhao et al., 2024b)  In this work, we show that pre-processing SGD using normalization and whitening in a stateless manner can achieve similar performance as Adam for LLM training, while maintaining the same memory footprint of SGD. Specifically, we show that normalization stabilizes gradient distributions, and whitening counteracts the local curvature of the loss landscape. This results in SWAN (SGD with Whitening And Normalization), a stochastic optimizer that eliminates the need to store any optimizer states. Empirically, SWAN achieves ~50% reduction on total end-to-end memory compared to Adam. Under the memory-efficienct LLaMA training benchmark of (Zhao et al., 2024a), SWAN reaches the same evaluation perplexity using half as many tokens for 350M and 1.3B model.

SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training
Chao Ma" Wenbo Gong"
Meyer Scetbon 1 Edward Meeds'
Abstract
2018; Pagliardini al., 2024; Liu et al., 2023; Zhao et al.,
2024a), have been central to the success of training large language models (LLMs) (Brown et al., 2020; Touvron et al., 2023b; Dubey et al., 2024; Bi et al., 2024; Bai et al., 2023; Zhang et al., 2022). However, most adaptive optimizers for LLMs are stateful, meaning they require tracking and maintaining internal states. While achieving remarkable empirical success, these states introduce significant memory overhead. For instance, Adam (Kingma & Ba, 2015) the de facto optimizer for LLM training involves the track- ing of exponential moving averages (EMAs), effectively doubling memory requirements. AdEMAMix (Pagliardini et al., 2024), an extension of Adam that achieves significant convergence speed boost, requires storing even more states, tripling the memory requirements. This overhead can be significant especially in distributed settings, where the opti- mizer states could consume a significant amount of the GPU memory (Dubey et al., 2024; Korthikanti et al., 2023). On the other hand, while stochastic gradient descent (SGD) is optimal in terms of memory efficiency (i.e., it is stateless), its capability to train LLMs is limited (Zhao et al., 2024b; Zhang et al., 2020; Kunstner al., 2023; 2024). Therefore, natural question arises:
Adaptive optimizers such as Adam (Kingma & Ba, 2015) have been central to the success of large language models. However, they often require maintaining optimizer states throughout training. which can result in memory requirements several times greater than the model footprint. This over- head imposes constraints on scalability and com- putational efficiency. Stochastic Gradient Descent (SGD), in contrast, is stateless optimizer, it does not track state variables during training. Con- sequently, it achieves optimal memory efficiency. However, its capability in LLM training is lim- ited (Zhao et al., 2024b). In this work, we show that pre-processing SGD using normalization and whitening in a stateless manner can achieve simi- performance as Adam for LLM training, while maintaining the same memory footprint of SGD, Specifically, we show that normalization stabi- lizes gradient distributions, and whitening coun- teracts the local curvature of the loss landscape. This results SWAN (SGD with Whitening And Normalization), a stochastic optimizer that elimi- nates the need store any optimizer states. Em- pirically, SWAN achieves 50% reduction on total end-to-end memory compared to Adam. Un- der the memory-efficienet LLaMA training bench- mark of (Zhao et al., 2024a), SWAN reaches the same evaluation perplexity using half as many tokens for 350M and 1.3B model.
Can LLMs be trained efficiently with stateless optimizers?
There a growing body of research that has contributed to answering this question positively. by developing novel optimizers that reduce the memory requirements associated with tracking internal state variables, while achieving simi- lar even speedup boost performance compared to Adam." For instance, some methods rely solely on tracking the first moment of gradients (Xu et al., 2024a; Jordan et al., 2024), while others introduce an additional one-dimensional track- ing variable on top of first moments (Zhang et al., 2024b; Zhao et al., 2024c). Alternatively. approaches focusing exclusively pre-conditioner tracking have also been pro- posed (Pooladzandi & Li, 2024; Li, 2017). Another line of work focuses on using low-rank approximations to store the first and second moments, thereby reducing the memory cost associated with tracking optimizer states (Lialin et al., 2023; Hao et al., 2024; Zhao et al., 2024a).
1. Introduction
Adaptive optimizers, such as Adam and its variants (Kingma & Ba, 2015; Loshchilov & Hutter, 2019; Shazeer & Stern,
'Equal contribution 'Microsoft Research. Correspondence to: Chao Ma <chao.ma microsoft.comm: Wenbo Gong <wenbo- gong microsoff..co...
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
In this work, we address this question by proposing to sim- 'See Appendix A for detailed related work review.

---

SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training
ply pre-process the instantaneous stochastic gradient in stateless manner. The result is SWAN (SGD with Whitening And Normalization), a novel stochastic optimizer that elimi- nates all internal optimizer states and empirically achieves comparable or even better performance compared to Adam on several LLM pre-training tasks. Our optimizer consists of combining two known operators to pre-process the raw gradients: GradNorm and GradWhitening
an example of stateful optimizer, which means it accumu- lates and stores internal states throughout training. Consider loss function Lw : X + RR parameterized by weight ma- trices W E Rm*x, and denote x+() mini-batch of inputs provided at the 1-th training step that is sampled from data distribution (z). Let G be the stochastic gradient of Lw (i.e., a random variable induced by sampling x x!!)). Then, Adam can be broken down into the following steps:
GradNorm ap- plies row-wise standardization on the gradient matrix, while GradWhitening orthogonalizes the normalized gra- dient matrix. On a theoretical model of transformer (Tian et al., 2023), we show that these operators help stabilizing the stochasticity of gradient distributions during training, and mitigating the local geometry of the loss landscape, respectively. When applied in tandem, they eliminate the need track state variables, thereby recovering the memory footprint of SGD, while matching the performance of Adam. To summarize, our contributions are:
m = 3/ (1- (1-1) 1) + (1 (1)
(00)= = 32p(1-11 + +(1 B2) G
m
W(+1)= w(0- V
where m(()) and (t)) are EMAS of the first and second mo- ments of the gradients; and is global step size. Track- ing and storing these two EMA estimates triples the total memory consumption required to train a LLM model. The LLaMA 405B model, for example, requires an additional 1.6TB of memory for Adam compared to SGD.
SWAN, practical, stateless, adaptive optimizer, pro- vides the following properties:
1. Memory efficiency: SWAN achieves the memory footprint of SGD, that is se 50% reduction on total memory, and = 100% reduction optimizer states compared to Adam.
Desired properties of Adam. There is a rich literature on understanding adaptive methods' inner workings and unreasonable effectiveness. Notably, we outline Adam's key properties that relate this work, that have been proposed and supported by studies in the literature:
2. Sample efficiency: with SWAN, we obtain the first evidence that stateless training of LLMs can be possible. Across several LLM scales, SWAN con- sistently achieves the same or better performance than Adam. On the memory-efficient LLaMA train- ing benchmark (Zhao et al., 2024a), SWAN even improves upon the token-efficiency of Adam (2x speedup).
Gradient whitening 3 it known that the inverse second moment - of Adam performs gradient whitening by approximating the square root inverse of the diago- nal of Fisher information matrix (Kingma & Ba, 2015; Hwang. 2024), as opposed the inverse in the standard natural gradient descent. This step biases the optimiza- tion trajectories towards well-conditioned regions (Jiang et al., 2024) and provides a better approximation to the geodesic flow when compared with the natural gradient update (Yang & Laaksonen, 2008).
Vour4e
3. Robustness to ill-conditioned problems and hyperpa- rameters: on L.MM pre-training, SWAN consistently and efficiently converges without learning rate warm- up (Section 5. Section 5.2) and extensive tuning.
Theoretical consistency with LLM dynamnies. We show that (1) GradNorm can stabilize the heterogeneous covariance of LLM gradients, leveraging the redundan- cies in transformer gradient flows (Theorem in Ap- pendix C); and (2) GradWhiteming can be derived as a non-diagonal second-order update under a spe- cific structural assumption of the Hessian (Section 4.3). Additionally, we highlight that in the quadratic case, GradWhitening leads to convergence rates that are ro- bust to the condition number of the local curvature (The- orem 2, Appendix C.
Gradient smoothing: the EMA operations in Adam natu- rally reduce the influence of mini-batch noise (Cutkosky & Mehta, 2020; Crawshaw et al., 2022);
Grudient invariance: recent work suggest the perfor- mance gap between SGD and Adam might lie in Adam's sign-descent-like nature (Bernstein et al., 2018; Craw- shaw et al., 2022; Chen et al., 2023). For example, Adam is robust to the rescaling of gradient diagonals (Kingma & Ba, 201 5); and is invariant to any sign-preserving scalings (under 31 B2 0) (Bernstein et al., 2018).
2. Preliminaries
"Whitening" is an loadel term. In this paper, we primar- ily refer to natural gradient descent with exponent - (Yang & Laaksonen, 2008), as opposed to the work of (Amari, 1998) (which is also referred to as whitening in some context.)
Adam (Kingma& & Ba, 2015) is the current standard choice for adaptive optimizers on LLM pre-training tasks. Adam is
2

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Table 1: Memory consumption breakdown of optimizers, and the corresponding learning rate expressiveness. We assume all optimizers are applied on 2D tensor of m by n (m < n). The optimizer state is defined as all intermediate variables maintained over time; expressiveness is the number of adaptive learning rates that a optimizer can express per tensor.
Adam
SGD
SGD-Sal
Apollo (rank r)
Apollo (mini)
GaLore (rank r) SWAN
Optimizer States |
2mn
0
mu
2nr + + (mr)
2n + 2 + (m)
2nr + mr
0
Model weights 4
mn
mn
mn
mn
mn
mn
Expressiveness f
mn
blockwise
7
m7
Test PPL, 350M Llama
Test PPL, 1B Llama
26
Adam
26
Adam
Adam Swan!
Galore
SNNAN
Swan
Mamory Consumption (Gad)
5GB momenitam)
24
SNAVN LONO
Test Perplexity
2.3 X speed
2.4 X speed up
20
18
10
14.
166
Tokens Seen (Billions)
Tokens Seen (Billions)
Model Size
(a) 350 M LLaMA model Pretraining
(b) 1.3 B LLaMA model Pretraining
(C) End-to-end memory footprint
Figure I: SWAN performance preview on memory-efficient LLaMA pretraining benchmark (Zhao et al., 2024a). (a) and (b): On both 350M and 1.3B models, SWAN achieves > 2X speed-up vs Adam in terms of tokens seen. (c): memory footprint. We directly measure end-to-end memory under full-model training, with batch size sequence. SWAN achieves near-SGD optimizer memory reduction (== 50% reduction total memory, and 100% reduction optimizer states).
For more comprehensive discussion of these properties, please refer to appendix B.
operations that exploit rich information contained in the gra- dient matrix. To this end, we compose two well-known ma- trix operators, namely normalization and whitening. When applied in tandem, they achieve similar desirable properties of adaptive optimizers, without the need to store historical gradient moments. The result is SWAN (SGD with Whiten- ing And Normalization), new stateless optimizer.
Adam as SGD pre-processing. Adam can be understood as computing history-dependent pre-processing of the Veun4e to achieve
gradients ((G G()) ...,, G
in!" (9)
the desired properties described above. A key observation is that all of these properties are achieved through element- wise operations, where each element of the gradient matrix independently pre-processed and re-scaled. This approach does not take into account the interactions and structures between diffferent variables, and we hypothesize that this is the reason why additional history information is necessary to bridge this gap, ultimately leading to the requirement for EMA states. We believe that designing stateless adaptive optimizers is possible if we can achieve similar properties by applying mevrix-level operations that pre-process the instantaneous stochastic gradients of SGD (G)) >> (()).
3.1. SWAN Update Rules
In SWAN (Algorithm 1), the raw SGD gradient G' is pro- cessed by the following operations:
(G(t) +- GradNorm (ra)Norm(G(()) AW <- GradWhitening (Q())
(SWAN)
The weight is then updated by W W(() - A W The GradNorm operator (Equation (1)) denotes the normal- ization of the gradient matrix row-wise (Figure 2); and the GradWhitening operator denotes the whitening of the gra- dients (Figure 2).
3. The SWAN Optimizer: Preprocessing SGD with Normalization and Whitening
As discussed Section 2, 4 the key to designing stateless adaptive optimizers lies in the incorporation of matrix-level
On the GradNorm Step. Consider the gradient matrix G E Rmxn with rows and columns corresponding to input and output dimensions, respectively, of some block of model
mr comes from the storage random projection tensors.
3

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Normalize output
GradNorm
Gradwhitening
View each col as data
Gradient matrix
Normalized Gradient matrix
Whitened Gradient matrix
Figure 2: Illustration of GradNorm and GradWhitening operators. In GradNorm operator, we perform normalization across the output dimensions (columns), using statistics computed row-wise. In GradWhitening operator (illustration adapted from Huang et al. (2019)), we treat each column of the gradient matrix G as a separate data sample. Then, GradWhitening can be seen squeezing the gradient data along all eigen directions, such that the covariance matrix is the identity matrix.
parameters. Let 1 S i S represent the input indices and 1 $ j S n represent the output indices. Instead of per- forming element-wise EMA to stabilize and normalize the noisy gradients, as done in Adam, we propose to standardize across the output dimensions at each time step l:
where & denotes Kronecker product. This leads to the following whitening step:
GradWhitening((G) (GG y
(2)
where the exponent - 1 stands for matrix inverse square root.
G
(GGT)-& à G is simply the orthogonalization of G, i.e. the closest orthogonal matrix to G (w.r.t the Frobenius norm). The derivation of (2) as structured Hessian approximation is discussed in Section 4.3. Finally, Equation (2) can also be interpreted as minimizing the condition number of G.
GradNorm G) == $1
(1)
where :- - (G.,) is m-dimensional the root mean square across dimension mm-dimensional column vector); 1, a v--dimensional column vector of ones.
Similar to GradNorm, GradWhitening (Equation (2)) as matrix operation has been widely used as forward- pass operator in the form of decorrelated batch normal- ization (Huang et al., 2018; Li et al., 2018); and it has also shown initial success in processing backward gradi- ents (Tuddenham et al., 2022; Jordan et al., 2024; Gupta et al., 2018). As shown in Figure 2, by treating each column of G as i.i.d. vector-valued data samples G (g1 ... Py. (..). GradWhiteming can be seen as effec- tively stretching/squeezingg this data matrix along the eigen- vectors to whiten its covariance matrix. This forces the gradients to traverse all eigen- dieenddirestions at the same rate.
GradNorm is the forward pass operator RMS-Norm (Zhang & Sennrich, 2019) (without learnable parameters) applied on backward gradients. Similar operations have been consid- ered in the literature to process gradients (You et al., 2019). GradNorm allows the optimizer to be invariant under matrix- wise and row-wise rescaling. In Section 4 we will show that this is the key to stabilizing the L.LM gradient distribution. Meanwhile, compared the sign operation (sign(Gr (G)) and matrix-wisee normalization [ used in signed descent, GradNorm preserves richer information of the gradient scal- ing while offering invariance properties.
On the GradWhiten teniingttning Step. As discussed in Section 2, Adam relies on a second moment estimate to perform element-wise gradient whitening. More formally, the second moment of Adam estimate a diagonal approximation of the Fisher information matrix (FIM) via (Hwang, 2024):
3.2. Practical Considerations
Rescaling Updates. The operator GradWhiten maps the normalized gradient (()) onto the closest orthogonal matrix, and as such might drastically change its effective learning rate. In practice, we propose the following re- scaling before updating the weights:
E ((Vec(G) Vec Rs Diag(VVe[[E(
where Vec(.) denotes the vectorized operation, Diag(-) de- notes the operation that produces a diagonal matrix given a vector, and @ denotes the Hadamard product. Hence, whitening using the inverse square root of this diagonal FIM is equivalent element-wise rescaling with VE(GOG)
( ||
VInAAW W
(3)
|AWO|
JAW
This helps rescale the norm of the whitened gradient back to the norm of G that is || = vmn. After we proposed Equation (3), follow-up work of (Gong et al., 2025) showed that Equation (3) can naturally derived
Here, we consider non-diagonal approximation:
E (eeec(G)Vec(G) Inxn 8 GG

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
GPU-friendly estimation. This is given by:
Algorithm 1 SWAN Optimizer
Yk+1 = 2 Y Z, Yn))
Input: weight matrix W E R with m S n. Step size 7). Number of GradWhitening iteration K (default 10).
(4)
ZR+1 = 1/331 Z5 xx)
where Yo GG Zo = I. At convergence, Z (GGT)-1/2, hence GradWhitening(G: = ZG (Algo- rithm 2). See Appendix for inplementation details. The use of the same N-S procedure for Shampoo variants has been discussed in (Mei al., 2023) and (Jackson, 2023). Another N-S procedure was applied in (Bernstein & New- house, 2024b), and optimized in (Jordan et al., 2024). How- ever, estimating (GG D with NS requires ((m") (as- suming m < n) complexity, making its scalability for larger models, especially under certain frontier settings unclear (Essential AI, 2025). Here, we propose a heuristic scheme that has Omm²) complexity:
Initialize step 5- 0
repeat
G( E Rmax x- V W (W(()))))))))))))))))))))()
G( <- GradNorm ()))
A AW(O < GradWhitening (( K))
(optional) AW()) mmn A
W W(t-) A (1-1)))
++++1
until convergence criteria met return (
Algorithm 2 GradWhitening Operator
Yi+11 1 Y Dia - Zz
2
(5)
Input: Gm with m S n. Number of iterations K. Step size 3. Boolean diag indicating if use diagonal substitution scheme.
Z 1k+1 = (31 Dia Y Diag(Z
where Diag(.) returns a diagonal matrix that has the same diagonal elements as the input matrix. Basically, whenever we encounter matrix multiplication in NS iterations, we re- place one of them by its diagonal approximation. We refer to this as the NS with diagonal substitution (NSDS) scheme. In Appendix I, we empirically demonstrate that NSDS per- forms well in minimizing the matrix condition number of gradients. As shown in Figure 6, on the gradient distribu- tion induced by Llama model training, the performance of NSDS (when combined with GradNorm) is comparable to the standard NS scheme in reducing the condition number.
Initialize +- GGT .z I+1
for = I.....Kau do if diag then
Y +- BYDiag331 ZDiag(Y))
Z + 83I Yiag(Z)))Diaan)
else
Y + (3I ZY).
Z + 3(1 ZY)Z
end if
end for return ZG
We found that NSDS does have the side effect of signifi- cantly slower early stage convergence speed. However, it of- fer stronger tail convergence in return. For L.M pretraining NSDS helps SWAN achieve similar or better final validation loss performance than the original NS scheme (Section 5.1), enabling Adam-level training throughput without the help of sharding NS computation (Section 5.3).
Algorithm 3 GradNorm Operator
Input: mm return
Figure 3: SWAN Optimizer.
4. Why SWAN: A LLM Dynamics Perspective
As a first analysis, we consider SWAN from a learning dy- namics perspective, specifically the dynamies of an LLM based upon simplified transformer block. It is this analysis that led to the design choices for SWAN. Theoretical in- sights of this section is verified empirically in Appendix J.I.
from a unifying framework for LLM optimization based on Fisher information. Later, the same scaling (noticing ||A(( || = vm) ) has been used in (Liu et al., 2025) to apply whitening operations larger models.
4.1. Setup
Fast GradWhiteming Exact GradWhitening is expen- sive, as it involves solving the matrix square-root inverse. We can instead directly apply the Newton-Schulz (NS) vari- ant of decorrelated batch normalization (Song et al., 2022; Li al., 2018; Huang et al., 2019), which allows a more
We consider the simplified transformer block (STB) archi- tecture proposed in (Tian et al., 2023).
Definition 1 (Simplified Transformer Block), Given the input x E MMxx query q. context embed- ding Uc RdxMe, and the query embedding
5

---

SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training
us E Rdx1, the STB computes the output as h o (WT (Uc (exp(zg) o x) + un)), where Mc is the con- text length. The attention logits Za E (RMxx1 are given by zql- = u W Q W k with Wa. WK E Raxd being weights for queries and keys, respectively, W gdxx is the MLP project-up weights, and & is a nonlinear function.
Approximating Hessian with Kronecker factorization is not new and has been extensively studied in the litera- ture (Martens & Grosse, 2015; George et al., 2018; Gao et al., 2023; 2021; Koroko et al., 2022; Eschenhagen et al., 2024; Gupta et al., 2018). This specific structure is useful in our context because it aligns with the statistical property of STB, as shown later in Proposition 1.
Given a STB, we consider a loss function Lw,z(xx()) where (+)) is a mini-batch of inputs provided at the t-th training step sampled from data distribution Paata (z). Stan- dard mini-batch learning dynamics is then given by
By leveraging Assumption 1, we can now effectively es- timate H by only using one single gradient matrix sam- ple G = [g),..., In]] E Rmxxn Recall that the Fisher information formulation of Hessian is defined as H = E| Vec G) Vec where Vec(-) denotes the vectorized operation. Under assumption 1, we can estimate H = Ixx X H by H R 1 E 9:9. GG In fact, in the follow-up work of (Gong et al., 2025) was formally showed that GGT is the solution of
W iLw.z, R,, (2)(1) OW
(() OLw.z, (2) Ōzg
29
In this case, both (+ and 2 are viewed as random vari- ables induced by random mini-batch w0), For example, for each row i, W [] can be re-written as W(O [t.=] E[VV &
n.
n
( [t, .]. where (i.
[i, |] is zero mean random
min Lax @ M - E Vec( Vec T M
(6)
variable with covariance CovlW [[,,],
Hence, GradWhitening can be seen as applying a second order update under our structural assumption:
4.2. GradNorm: stabilizing gradient distributions of STB
Below we show that, based on the dynamics of the STB, GradNorm stabilizes W-
GradWhitening(G)) v //maxmm, G)G
Theorem 1 (GradNorm stabilizes gradient distributions across time for the STB). Consider the STB (Definition 1). Assuming we inherit the assumptions in Theorem 1 of Tian et al. (2023), as described in Appendix C. Then consider U W, the composition of the MLP project-up matrix and the embedding matrix as a whole. Then, its standardized
Noticing | whitening of eq. (3). In the following Proposition, we show that the assumption of GradWhitening aligns with the equilibrium Hessian structure in the STB regime.
/m, we recovered the rescaled
Proposition 1 (Shared structures in the block-diagonal of Hessians at transformer equilibrium). Consider a STB (1), trained with full- gradient descent. Next, assume we inherit all the assumptions from Theorem I of (Tian et al., 2023). Then, as t + co, we have the following shared Hessian structure along the diagonal blocks:
stochastic gradients (1) Gugw == GradNorm( dLw ((a satisfy:
U
Cov Cor(Guwwl..j(u) = Cov(
For all t1, t2. and i. In other words, the covariance struc- ture of G is identical across all time steps t, achieving distributional stability across time. The same relation- ship also holds for the gradient of attention score G ==
H
(7)
Ens Hok
Eno
HoW
0Lw (m)
For all 1 S s, S H(W) ((WW)k,,'k+ sk, "waedherr
d. 1
S k, S n, Where
GradNorm(
BC.
Bzg
Theorem suggests that GradÑorm implicitly aligns with the dynamics of transformer architectures and removes the time-heterogeeneity in gradient covariance structures.
Proposition shows that, under a simplified setting of the transformer, the Hessian will also converge to an equilib- rium solution where the Mc x Mc blocks over the diagonal direction of Hessian shares an identical structure, which sup- ports the assumption of GradWhitening We verify this ex- perimentally (Appendix, Figure 5). Furthermore, Appendix Theorem 2 shows that GradWhi vergence more robust to local curvature conditions and out- performs both SGD and Adam in ill-conditioned scenarios.
4.3. GradWhitening: exploiting the tied block-diagonal structures of transformer Hessian
makes SGD's con-
Here show that GradWhitening is equivalent non- diagonal second-order method under a specific Kronecker factorization assumption of the Hessian/FIM.
Assumption (Assumption of GradWhitening). At time t, the local Hessian H of the loss has shared diagonaa structure H I Iaxn @ H, where H E Rmxmm m <, n.
'Using the natural gradicnt update rule with exponent = (Gong et al., 2025; Yang & Laaksonen, 2008) Veo (W(1)))) Vec(W") or -F -F(Veo(G) Vec --1// (e))))
6

---

SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training
5. Experiments
Results As shown in Table and Figure 1, all SWAN variants achieve strong performance, requiring the lowest memory consumptions comparable to SGD. Across all mod- els, SWAN-O surpasses the performance of the Adam and (muon-like) baseline in terms of validation perplexities. SWAN further delivers a com- parable performance to Apollo series; and finally, SWAN with NSDS scheme delivers the strongest, SOTA-level per- formance under this setup. Notably, on the 350M and 1.3B models, all SWAN variants reaches at least 2x speedup (in steps or tokens) relative to Adam (Figure and Figure 9 (a)). Finally, we observe that in general, SWAN has a slower convergence speed in the early stage than SWANT. How- ever, offers strong tail convergence in return and surpasses SWANT in terms of final performance.
5.1. SWAN Performance on LLM Pre-training Tasks
Setup We evaluate SWAN on memory-efficient Llama (Touvron et al., 2023a) pre-training tasks, using the stan- dardized settings of (Zhao al., 2024a), which has been adopted by many recent works (Zhao et al., 2024a; Chen et al., 2024; Zhao et al., 2024a; Zhu et al., 2024). We con- sider models with 60M, 130M, 350M, and 1.3B parameters, all trained on the C4 dataset (Raffel et al., 2020) using an effective batch size of 130K tokens. Following the setup of (Zhao et al., 2024a), SWAN is applied to all linear modules in both attention and MLP blocks. Training uses 6F16 by default unless specified, see Appendix K. The other evalua- tion settings follows Zhao et al. (2024a).
5.2. Ablation of SWAN on LLM pretraining
SWAN optimizer We consider three variants: SWAN-O, which aims show-case the robustness and effectiveness of our method out-of-the-box, with almost no tuning. It uses naive NS-iteration for whitening, disabled learning rate warmup, and use similar learning rates optimized for Adam. SWAN the vanilla version of our method, in which we enabled learning rate warmup, and allowed the use of optimized learning rates that largely differ from Adam. Fi- nally, SWAN the strongest and most efficient variant that employs the proposed NSDS scheme for fast whitening (sec- tion 3.2). For all variants, we adopt a lazy-tuning approach (hyperparameters are set without extensive search) to reduce the possibility of unfair performance distortion. Notably, for SWAN#, we share the same hyperparameters across all model sizes. Full details can be found in Appendix K.
We take SWAN-O from the SWAN series as an example and conduct the following ablation studies.
Effect of GradNorm and GradWhitening on Perfor- mance We consider six ablation settings: (1) SWAN-O (2) SWAN-0 (GradNorm only). (3) SWAN-O StradWhitening only), (4) Adam, (5) Adam (momentum only), and (6) Adam (second moment only). As shown in Figure 4 (a), both GradNorm and GradWhitening contribute SWAN's fi- nal performance. Removing either results in performance degradation. Similarly, Adam also requires first and second moment estimates for optimal performance.
Does SWAN work only because it inereases effective learning rate via GradNorm? To rule out this possibility, we remove the GradNorm operator from SWAN-O and run a learning rate sweep. Starting with the default Ir of SWAN-O, we apply multipliers from 1 to 10 In Figure 4 (b), the blue line shows the final PPL for SWAN w/o GradNorm at different learning rates multipliers, while the dashed red line represents the performance of full SWAN-O at the default learning rate. Although raising learning rate can improve the performance of SWAN-O without GradNorm, the gap to full SWAN-0 remains large. This indicates that GradNorm's 's gra- dient noise stabilization is essential and cannot be replaced simply by increasing the learning rate.
Baselines We compare SWAN with Adam (Kingma & Ba, 2015) and other memory-efficied optimizers. All base- line results are directly quoted from corresponding papers as they all share the same standard setup. The baselines include Adam (Kingma& & Ba, 2015), which is standard choice for training large models; and Galore (Zhao et al., 2024a), a memory-efficient Adam variant with low-rank gradient projections. We also consider Low-Rank (Ka- malakara et al., 2022), a low-rank factorization approach (W = BA); and LoRA (Hu et al., 2021), which applies the LoRA method for pre-training in (Zhao al., 2024a). Additionally, we include ReLoRA (Lialin et al., 2023), a full-rank extension of LoRA with parameter merging, which applies Newton- Schulz GradWhitening on top of momentum instead of GradÑorr; this is similar to Muon (Jordan et al., 2024) with Nesterov acceleration turned off. Finally, we compare with Apollo-mini and Apollo (Zhu et al., 2024). Full details can be found in Appendix K.
How does warm-up affect the performance? Section 5. shows that SWAN can train with no warm-up phase even under a relatively large learning rate (0.001). Here, we compare Adam and SWAN with and without warm-ups. As seen in Figure 4 (c), SWAN without the warm-up phase gives better performance, and it still outperforms Adam under Adam's 's own warm-up schedule. On the other hand, Adam's 's performance decreases drastically without a proper warm-up. These suggest that SWAN is more robust to warm-
and Momentum + www--Gaaahieennne
*Code will be released at w.t.uu..com/icccoo msr_optin
7

---

SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training
Table 2: Comparison with Adam and its memory- -efooreefficient low-rank variants on pre-training various sizes of LLaMA models on C4 dataset. Validation perplexity is reported, along with memory estimate of the total of parameters and optimizer states based on BF16 format. Baseline results are directly taken from the official numbers reported in Zhao al. (2024a); Zhu et al. (2024), as they shares exactly the same setup. Note that for Adam we report both the official results from Zhao et al. (2024a) and our reproduced result.
60M
130M
350M
1.3 B
Adam
33.02 (0.32G)
24.44 (0.75G)
19.24 (2.05G)
16.44 (7.48G)
Adam (cited)
34.06 (0.32G)
25.08 (0.75G)
18.80 (2.05G)
15.56 (7.48G)
SWAN
30.59 (0.23G)
22.61 (0.43G)
16.63 (0.93G)
13.56 (2.98G)
SWANT
30.00 (0.23G)
22.83 (0.43G)
17.14 (0.93G)
14.42 (2.98G)
SWAN-O
32.28 (0.23G)
24.13 (0.43G)
18.22 (0.93G)
15.13 (2.98G)
Momeltum-GGadWWitennin
31.6 (0.27G)
24.59 (0.59G)
19.30 (L.49G)
16.08 (5.23G)
Galore
34.88 (0.26G)
25.36 (0.57G)
18.95 (1.29G)
15.64 (4.43G)
Apollo-mini
31.93 (0.23G)
23.53 (0.43G)
17.18 (0.93G)
14.17 (2.98G)
Apollo
31.55 (0.26G)
22.94 (0.57G)
16.85 (1.29G)
14.20 (4.43G)
Low-Rank
78.18 (0.26G)
45.51 (0.57G)
37.41 (1.29G)
142.53 (4.43G)
LoRA
34.99 (0.36G)
33.92 (0.80G)
25.58 (1.76G)
19.21 (6.17G)
ReLoRA
37.04 (0.36G)
29.37 (0.80G)
29.08 (1.76G)
18.33 (6. 17G)
SWAN speed up vs Adam
1.52 X
1.6 X
> 2.3 X
> 2.4 X
r of low-rank methods Training Steps
128
256
256
512
10K
20K
60K
100K
Ablation 1, Test PPL, 130M Llama
Ablation 2, LR vs Eval PPL
Ablation 3, Test PPL, 130M Llama
SWAN (GredNorm only)
140
SWAN (w/o GredNorm)
SWAN (w/ warmup)
SWAN (GradWhitening only)
SWAN
Adam (wlo warmup)
Adam (momentum only)
(wi warmup) SWAN (wio warmup)
Perplextty
107
Adam (second moment only)
Adam
Final PPL
100
perplexity
10'
SWWN
80
60-
40
20
0.0
1.0
15
2.0
2.5
10' Multiplier Base
30
0.0
1.0
1.5
2.0
2.5
Takens Seen (Billions)
Takens Seen (Billions)
(a)
(b)
(c)
Figure 4: Ablation studies on 130M model. (a) On the contribution of each components in SWAN and Adam. (b) On removing GradNorm and compensate with larger learning rates. (c) On the effect of learning rate warm-ups.
up schedule and can train effectively with or without it.
5.3. Memory Efficiency and Throughput Analysis
Memory Footprint We compare SWAN-0, Adam, and Galore on a single A A100 GPU. Unlike (Zhao et al., 2024a; Zhu et al., 2024), which report layer-wise training memory usage, we measure total end-to-end memory consumption under full-model training using batch-size of 1 for LLaMA with 1.3B, 7B, and 13B parameters. As shown in Figure I (c), SWAN's memory usage is on par with SGD, providing nearly a 50% reduction in total memory. This underlines the benefit of the stateless design.
Is the speed-up Multiplicative or Additive? A key ques- tion in assessing speedup factors is whether the improve- ment over Adam is multiplicative or additive. A multiplica- tive speedup implies that the optimizer's relative advantage remains proportionally consistent over time, while addi- tive speedup suggests a less desired constant step advantage. In Appendix J.2 and Figure 9, we show that for smaller models (60M and 130M), the empirical speedup trajectories indicate a primarily additive speedup. For larger models (350M and 1.3B), they suggest a multiplicative speedup.
Effective Throughput We assess throughput when train- ing a 1.3B LLama model on 8 A100 GPUs with a batch size of 130K. All gradient processing are done BF16. We
8

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Table 3: Raw and effective throughput analysis, under dif- ferent model parallelization (MP) settings.
et al., 2025) introducing general framework called multi- normalized gradient descent (MNGD).. MNGD considers set of K 2 1 norm functions (.1.....9 .9K)) and defines the update direction AW as the solution to:
Method
Raw / eff. throughput (IB)
Adam
117872/ / 117872 (tokens/s)
SWANT w/o MP
58600 / 117200 (tokens/s)
SWAN w/o MP
107808/ / 258739 (tokens/s)
arg max s.t. Vie [1, g.(2) =1. (8)
SWANT w/ MP
114160 / 228320 (tokens/s)
SWAN w/ MP
115872 / 278092 (tokens/s)
where G is the current gradient (can be raw gradient or momentum). In other words, MNGD chooses an update direction that is simultaneously normalized under all К norms (each gz(2) is set to 1). (Scetbon et al., 2025) showed that solving this problem leads to an iterative normalization strategy one can loop over = to K for several passes, normalizing the gradient with respect to each gi in turn.
use two metrics: raw throughput: number of tokens pro- cessed per second. Effeettvee throughput: raw throughput ad- justed by SWAN's token efficiency relative to Adam. These metrics evaluate the impact of different GradWhiteming schemes on training speed, and also account for the fact that some optimizers make more effective use of train- ing tokens. As shown in Table 3, SWAN with naive NS GradWhiten tening requires model parallelization (where NS of different tensors are distributed to different GPUs) to achieve competitive throughput. Without model paralleliza- tion, the throughput of SWANT is 50 % lower than Adam. With the proposed NSDS scheme, SWANT achieves an raw throughput comparable to Adam, without any need for model parallelization. Consequently, SWAN# exhibits 2 x higher effective throughput than Adam.
Notably, SWAN fits into this framework as a special case. It can be viewed with K = 2 norms: 9i is the row-wise & norm and g2 is the spectral norm. This generalized viewpoint offers greater flexibility for designing new optimizers. As demonstrated in (Scetbon et al., 2025), one may design more elegant and effi- cient optimizers than SWAN (and other Shampoo variants). For example, instead of using a single spectral norm that requires expensive NS iterations, one could enforce multiple norms that have simpler projection operations.
single-iteration approximation of MNGD
5.4. Verifying the Theoretical insights of Section 4
We empirically validate the theoretical benefits of GradNorm and GradWhitening. As detailed in Ap- pendix J.1, we demonstrate that GradNorm effectively stabi- lizes stochastic gradient distributions during SGD training of a scaled-down LLaMA model on the C4 dataset, evidenced by significantly reduced KL divergence fluctuations com- pared to standard training. Additionally, GradWhitening significantly enhances optimization performance across high-dimensional and ill-conditioned classic optimization functions, consistently outperforming both SGD and Adam. These findings confirm that GradNorm promotes gradient stability and GradWhitening addresses local curvature challenges, enabling faster and more reliable performance.
7. Conclusion, Limitations and Future Work
We introduced SWAN, a stateless optimizer for LLM train- ing that preprocesses raw gradients via a chain of two operators. GradNorm and GradWhitening. These design choices were intended to leverage structural redundancies in transformer training dynamics. Through theoretical anal- ysis and empirical evidence, we demonstrated that SWAN maintains an SGD-level memory footprint while achieving performance on par with Adam. Notably, on the memory- efficient LLaMA training benchmark of Zhao al. (2024a), when training 350M and 1B models, SWAN attained 2x speedup in tokens processed compared Adam, marking a new of-he-of-the-art on this benchmark.
6. Beyond SWAN: the General Principle of Constructing Chain of Gradient Processors
Our findings provide first proof-of-concey that the state- less approach has the potential to serve as a practical, ef- ficient alternative to optimizers that require tracking inter- nal states. However, we acknowledge a gap between our setup and industry-scale LLM training (scale, context length, batch size, architecture, dataset, etc). These differences pose challenges for generalizing our insights, and we aim to ad- dress them in future work. Furthermore, future research may explore alternative design choices for stateless optimiz- ers and extend SWAN to other complex training regimes beyond standard LLM pre-training.
Section implied our thought process when designing SWAN: the operators GradNorm and GradWhitening were chosen to exploit the structures of transformer dynamics. However, the idea of preprocessing gradients via a chain of matrix operators is more fundamental. This raises a key question: given an arbitrary model, what gradient operators can we choose, so that the resulting optimizer is valid? Af- ter our ICML submission, the follow-up work of (Scethon

---

SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which feel must be specifically highlighted here.
10

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
References
Wentao Zhang, Yichao Zhang, Chenggang Zhao, Yao Zhao, Shangyan Zhou, Shunfeng Zhou, Qihao Zhu, and Yuheng Zou. Deepseek Ilm: Scaling open-source lan- guage models with longtermism. ArXiv, 0bs/2401.029544 2024. URL https: /apii semant i cscholars org/ http//CouaDD:266618000
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251-776 1998.
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning, 2021. URL https://arxiv.. abs/2002.050 09018.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Nee- lakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In NeurIPS, 2020.
Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hin- ton. Layer normalization. ArXiv, 2016. URL https::/apt www.sseimmiiiiiieeeeee...e org/CorpusID: :2336111 8236317.
Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang. Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang. et al. Qwen technical report. arXiv preprint WrXiv:230.166609 2023.
Jeremy Bernstein and Laker Newhouse. Modular duality in deep learning, 2024a. URL https://arxiv... org/ abs/ 241 10.10000 212655
David E Carlson, Edo Collins, Ya-Ping Hsich, Lawrence Carin, and Volkan Cevher. Preconditioned spectral de- scent for deep learning. Advances in neural information processing systems, 28, 2015.
Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology. arXiv preprint arXiv:2409.20325, 2024b.
Xi Chen, Kaituo Feng, Changsheng Li, Xunhao Lai, Xi- angyu Yue, Ye Yuan, and Guoren Wang. Fira: Can we achieve full-rank training of Ilms under low-rank con- straint? arXiv preprini arXiv:2410.01625 2024.
Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology, 2024c. URL htt.s://///// rxiv.
/org/abs/20 ha//bbbb///99 2409 20325..
Xiangning Chen, Chen Liang, Da Huang. Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, and Quoc V. Le. Symbolic discovery of optimization algorithms. In NeurlPS, 2023.
Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzade- nesheli, and Animashree Anandkumar. signsgd: Com- pressed optimisation for non-convex problems. In Inter- national Conference on Machine Learning, pp. 560-569. PMLR, 2018.
Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang. Robustness to unbounded smoothness of generalized signsgd. Advances in neural information processing systems, 35:9955- 2022.
DeepSeek-A1 AI Xiao Bi, Deli Chen, Guanting Chen, Shan- huang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong. Qiushi Du, Zhe Fu, Huazuo Gao, Kaige Gao, Wenjun Gao, Ruiqi Ge, Kang Guan, Daya Guo, Jianzhong Guo, Guangbo Hao, Zhewen Hao, Ying He, Wen-Hui Hu, Panpan Huang, Erhang Li, Guowei Li, Jiashi Li, Yao Li, Y. K. Li, Wenfeng Liang, Fangyun Lin, Aixin Liu, Bo Liu (Benjamin Liu), Wen Liu, Xi- aodong Liu, Xin Liu, Yiyuan Liu, Haoyu Lu, Shang- hao Lu, Fuli Luo, Shirong Ma, Xiaotao Nie, Tian Pei, Yishi Piao, Junjie Qiu, Hui Qu, Tongzheng Ren, Zehui Ren, Chong Ruan, Zhangli Sha. Zhihong Shao, Jun-Mei Song, Xuecheng Su, Jingxiang Sun, Yaofeng Sun, Min Tang, Bing-Li Wang, Peiyi Wang, Shiyu Wang, Yao- hui Wang, Yongji Wang, Tong Wu, Yu WI, Xin Xie, Zhenda Xie, Ziwei Xie, Yi Xiong, Hanwei Xu, Ronald X Xu, Yanhong Xu, Dejian Yang, Yu mei You, Shuiping Yu, Xin yuan Yu, Bo Zhang, Haowei Zhang, Lecong Zhang, Liyue Zhang, Mingchuan Zhang, Minghu Zhang,
Ashok Cutkosky and Harsh Mehta. Momentum improves normalized sgd. In International conference machine learning. pp. 2260-2268. PMLR, 2020.
André Belotto Da Silva and Maxime Gazeau. A general sys- tem of differential equations to model first-order adaptive algorithms. Journal of Machine Learning Research, 21 (129):1-42, 2020.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Ab- hishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurélien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Rozière, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Nayak, Chloe Bi, Chris Marra, Chris McConnell, Chris- tian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Di- nan, Eric Michael Smith, Filip Radenovic. Frank Zhang. Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Ander- son, Graeme Nail, Grégoire Mialon, Guan Pang, Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, llyan Zarov, Imanol Arricta Ibarra, Is- abel M. Kloumann, Ishan Misra, Ivan Evtimov, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der Linde, Jen- nifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jian- feng Chi, Jianyu Huang. Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Va- suden Alwala, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, and Kevin Stone. The llama herd of models. CoRR, abs/2407.22 21783, 2024.
Wenbo Gong, Meyer Scetbon, Chao Ma, and Edward Meeds. Towards efficient optimizer design for Ilm via structured fisher approximation with a low-rank extension. arXiv preprint arXiv:2502.07752, 2025.
Vineet Gupta, Tomer Koren, and Yoram Singer. Sham- poo: Preconditioned stochastic tensor optimization, 2018. URL https: /ttpps://arxiii org/ abs / 18000 69568.
Yongchang Hao, Yanshuai Cao, and Lili Mou. Flora: Low-rank adapters are secretly gradient compres- sors. ArXiv, abs/2402.03293, 2024. URL https:
267412117.
Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen- Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language mod- els. arXiv preprint arXiv:2 106.09685, 2021.
Lei Huang, Dawei Yang. Bo Lang, and Jia Deng. Decorre- lated batch normalization. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 791-800, 2018.
Runa Eschenhagen, Alexander Immer, Richard Turner, Frank Schneider, and Philipp Hennig. Kronecker-factoree approximate curvature for modern neural network archi- tectures. Advances in Neural Information Processing Systems, 36, 2024.
Lei Huang, Yi Zhou, Fan Zhu, Li Liu, and Ling Shao. erative normalization: Beyond standardization towards efficient whitening. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 4874-4883, 2019.
Essential AL. Layer sharding for large-scale training with muon. https::///// httpps://wwwwwe http://wwwwssntaallaaa. infra, 2025.
Dongseong Hwang. Fadam: Adam is a natural gradient optimizer using diagonal empirical fisher information. arXiv preprint arXiv:2405. 1280, 2024.
Jingwen Fu, Bohan Wang, Huishuai Zhang, Zhizheng Zhang, Wei Chen, and Nanning Zheng. When and why momentum accelerates sgd: An empirical study. arXiv preprint arXiv:2306.099 09000, 2023.
Jacob Jackson. An isometric stochastic optimizer. arXiv preprint arXiv:2307..12979 2023.
Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv preprint arXiv: 2017.
Kaixin Gao, Xiaolei Liu, Zhenghai Huang, Min Wang, Zi- dong Wang, Dachuan Xu, and Fan Yu. A trace-restricted Kronecker-factoree approximation natural gradient. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 35, pp. 7519-7527, 2021.
Samy Jelassi and Yuanzhi Li. Towards understanding how momentum improves generalization in deep learning. In International Conference on Machine Learning, pp. 9965 10040. PMLR, 2022.
Kaixin Gao, Hai Huang, Xiaolei Liu, Min Wang. Shuangling Wang, Zidong Wang, Dachuan Xu, and Fan Yu. Figenvalue-corrected natural gradient based on new approximation. Asia-Pacific Journal of Operational Research, 40(01):2340005 2023.
Kaiqi Jiang, Dhruv Malik, and Yuanzhi Li. How does adap- tive optimization impact local neural network geometry? Advances in Neural Information Processing Systems, 36, 2024.
Thomas George, César Laurent, Xavier Bouthillier, Nico- las Ballas, and Pascal Vincent. Fast approximate natural gradient descent in a kronecker factored eigenbasis. Ad- vances in Neural Information Processing Systems, 31, 2018.
Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cecista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https: /eeleereerrr jordan. github. co/postes ptppps/msuoott.. s/munn/.
12

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Siddhartha Rao Kamalakara, Acyr Locatelli, Bharat Venkitesh, Jimmy Ba, Yarin Gal, and Aidan N Gomez. Exploring low rank training of deep neural networks. arXiv preprint 2tt.::::2009113000 2022.
Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for Ilm training. 2025.
arXiv preprint
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poster), 2015.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In ICLR (Poster). OpenReview.net, 2019.
Abdoulaye Koroko, Ani Anciaux-Sedrakinn Ibtihel Ben Gharbia, Valérie Garès, Mounir Haddou, and Quang Huy Tran. Efficient approximations of the fisher matrix in neural networks using kronecker product singular value decomposition. arXiv preprint arXiv:2201.10288 2022.
James Martens and Roger Grosse. Optimizing neural net- works with kronecker-factoree approximate curvature. In International conference on machine learning, pp. 2408- 2417. PMLR, 2015.
James Martens, Jimmy Ba, and Matt Johnson. Kronecker- factored curvature approximations for recurrent neural networks. In International Conference on Learning Rep- resentations, 2018.
Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation re- computation in large transformer models. Proceedings of Machine Learning and Systems, 5:341-353, 2023.
Jonathan Mei, Alexander Moreno, and Luke Walters. Krada- grad: kronecker opproximation-domnntteo gradient pre- conditioned stochastic optimization. In Uncertainty in Artificial Intelligence, pp. 1412-1422. PMLR, 2023.
Frederik Kunstner, Jacques Chen, Jonathan Wilder Laving- ton, and Mark Schmidt. Noise is not the main factor be- hind the gap between sgd and adam on transformers, but sign descent might be. arXiv preprint arXiv:2304.133: 13960,, 2023.
Igor Molybog, Peter Albert, Moya Chen, Zachary Vitoo, David Esiobu, Naman Goyal, Punit Singh Koura, Sharan Narang, Andrew Poulton, Ruan Silva, et al. A theory on adam instability in large-scale machine learning. arXiv preprint arXiv:2304.09871, 2023.
Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, and Alberto Bietti. Heavy-tailed class imbal- ance and why adam outperforms gradient descent on lan- guage models. arXiv preprint S4ti::2402.1444/4 2024.
Yurii Nesterov. Introductory lectures on convex optimiza- tion: A basic course, volume 87. Springer Science Business Media, 2013.
Peihua Li, Jiangtao Xie, Qilong Wang, and Zilin Gao. To- wards faster training of global covariance pooling net- works by iterative matrix square root normalization. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 947-955, 2018.
Matteo Pagliardini. Pierre Ablin, and David Grangier. The ademamix optimizer: Better, faster, older. arXiv preprint 0rXiv:2409.03137, 2024.
Xi-Lin Li. Preconditioned stochastic gradient descent. IEEE
Abel Peirson, Ehsan Amid, Yatong Chen, Vladimir Fein- berg, Manfred K Warmuth, and Rohan Anil. Fishy: Lay- erwise fisher approximation for higher-order neural net- work optimization. In Has it Trained Yet? NeurIPS 2022 Workshop, 2022.
Iransactions on neural networks and learning systems, 29 (5):1454-1466, 2017.
Vladislav Lialin, Namrata Shivagunde, Sherin Muck- atira, and Anna Rumshisky. Relora: High- rank training through low-rank updates. In In- ternational Conference on Learning Representations, 2023. URL /tstps::/api... semant icschoolr. org/
Omead Pooladzandi and Xi-Lin Li. Curvature-informel sgd via general purpose lie-group preconditioners. arXiv preprint arXiv:2440.00455. 2024.
259836974.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Jour- nal of Machine Learning Research, 21(140):1-67, 2020. URL http: /ttp:://mmr.../ jmlr.mmrrr//.. wwwll.org//aaae wttppppppers/////22000777 html.
Wu Lin, Felix Dangel, Runa Eschenhagen, Juhan Bac, Richard E Turner, and Alireza Makhzani. Can we remove the square-root in adaptive gradient methods? a second- order perspective. arXiv preprint arXiv:2402.03446, 2024.
Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. CoRR, abs/2305.14342. 2023.
Meyer Scetbon, Chao Ma, Wenbo Gong, and Edward Meeds. Gradient nultinnormalizztion for stateless and scalable Ilm training. arXiv preprint 2rXiv:2502.06742, 2025.
13

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In ICML., volume 80 of Proceedings of Machine Learning Research, pp. 4603 4611. PMLR, 2018.
Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409. 11211, 2024.
Hao-Jun Michael Shi, Tsung-Hsien Lee, Shintaro Iwasaki, Jose Gallego-Posada, Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, and Michael Rabbat. A distributed data-parallel pytorch implementation of the distributed shampoo optimizer for training neural networks at-scale. arXiv preprint (tt:2339.06499 2023.
Sike Wang, Pan Zhou, Jia Li, and Hua Huang. 4-bit shampoo for memory-efficient network training. arXiv preprint wrXiv:2405.18144 2024.
Minghao Xu, Lichuan Xiang, Xu Cai, and Hongkai Wen. No more adam: Learning rate scaling at initialization is all you need, 2024a. URL https:/ petttp:://ariive /ttp//axxvv.rr///////.///.// abs /2122 11766...
Yue Song, Nicu Sebe, and Wei Wang. Fast differentiable matrix square root and inverse square root. IEEE Trans- actions on Pattern Analysis and Machine Intelligence, 45 (6):7367-7380, 2022.
Minghao Xu, Lichuan Xiang, Xu Cai, and Hongkai Wen. No more adam: Learning rate scaling at initialization all you need. arXiv preprint arXiv:2412.1 2024b.
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du Joma: Demystifying multilayer trans- formers via joint dynamics of mip and attention. arXiv preprint arXiv:2310.00533 2023.
Zhirong Yang and Jorma Laaksonen. Principal whitened gradient for information geometry. Neural Networks, 21 (2-3):222-240, 2008
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, San- jiy Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in min- utes. arXiv preprint arXiv:1904. 00962, 2019.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Mar- tinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, al. Llama: Open and efficient foundation language models. arXiv preprint ar arXiv:2302.13977 2023a.
Biao Zhang and Rico Sennrich. Root mean square layer nor- malization. Advances in Neural Information Processing Systems, 32, 2019.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale. Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marein Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Ko- renev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yun- ing Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Sub- ramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Tay- lor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic. Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023b. URL https:t////....... http//axvvvv/v org/ abs/ 2307 09288.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Su- vrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33:15383-1553 113933, 2020.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre- trained transformer language models. arXiv preprint arXiv:2205.01000 2022.
Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi thi-Quan Luo. Why transformers need adam: A hessian perspective. arXiv preprint arXiv:2402. 16788, 2024a.
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chen- wei Wu, Yinyu Ye, Zhi- Luo, and Ruoyu Sun. Adam- mini: Use fewer learning rates to gain more. arXiv preprint arXiv:2406. 16793, 2024b.
Mark Tuddenham, Adam Prügel-
-ennett, and Jonathan Hare. Orthogonalising gradients to speed up neural net- work optimisation. arXiv preprint arXiv:2202.07052, 2022.
Jiawei Zhao, Zhenyu (Allen) Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficien IIm training by gra- dient low-rank projection. ArXiv, abs/2403.00307,
14

---

SWAN: with Normalization and Whitening Enables Statcless LLM Training
2024a. URL / htt. TTg/CoorpuSIDD:26255355
Rosie Zhao, Depen Morwani, David Brandfinbrenner, Nikhil Vyas, and Sham Kakade. Deconstructing what makes good optimizer for language models. arXiv preprint arXiv:2407.07972 2024b.
Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham Kakade. Deconstructing what makes good optimizer for language models, 2024c. URL https: http:://axxvvvvooe org/ 0bs/2407, 07972.
Shen- Yi Zhao, Chang- Wei Shi, Yin-Peng Xie, and Wu-Jun Li. Stochastic normalized gradient descent with momen- tum for large-batch training. Science China Information Sciences, 61(11)2 2024d.
Hanqing Zhu, Zhenyu Zhang, Wenyan Cong, Xi Liu, Sem Park, Vikas Chandra, Long, David Z Pan, Zhangyang Wang, and Jinwon Lee. Apollo: Sgd-like memory, adamw-level performance. arXiv preprint arXiv:2412.05270. 2024.
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jin- wen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. arXiv preprint arXiv: 10955 2018.
15

---

SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training
A. Related works and extended discussions
Towards Stateless Optimizers for LLM Training. Adaptive optimizers generally rely on tracking internal state variables to perform weight updates, which can substantially increase memory consumption when training large models. Several recent works have successfully managed to reduce the memory requirements associated with storing additional state variables for training LLMs. Muon (Jordan et al., 2024; Bernstein & Newhouse, 2024b), a newly proposed optimizer, has demonstrated strong acceleration and memory saving for LLM training by simplifying thampoo-like optimizers (Gupta et al., 2018; 88; Anil et al., 2021; Shi et al., 2023; Wang ct al., 2024; Vyas et al., 2024; Peirson et al., 2022; Lin et al., 2024) and requiring only the tracking of a first-moment estimate. SGD-Sal (Xu et al., 2024b) only stores the first moment with a learning rate estimated at the beginning of training. Sign-based methods such as (Chen et al., 2023) have also demonstrated success on training transformer- based models by only tracking first moments. There are also several works that aim to enhance the memory efficiency of Adam by reducing the memory cost associated with second moments. Adam-mini (Zhang et al., 2024b) significantly reduces memory usage by storing only scalar values for each parameter block, while Adalayer (Zhao et al., 2024c) retains only the scalar average of the second moment for each layer. Alternatively, PSGD (Pooladzandi & Li, 2024; Li, 2017) focuses on exclusively tracking pre-conditioner, eliminating the need to track a first moment estimate. Finally, the concurrent work of (Zhu et al., 2024) proposed an approximate gradient scaling method, that tracks the channel-wise or tensor-wii wise surrogate learning rates, further improving the memory efficiency. However, all the aforementioned optimizers still require the storage of state variables. In contrast, SWAN completely eliminates the need to store internal states for both the first and second moments by employing combination of GradNorm and GradWhi steps, which is discussed next.
Pre-processing Gradients. Gradient pre-processing is a common technique used to enhance performance of optimizers.
Various pre-processing procedures have been proposed in the literature, such as signed gradient (Bernstein et al., 2018; Crawshaw et al., 2022; Chen et al., 2023; Kunstner et al., 2023), gradient clipping (Zhang et al., 2020), normalization (Zhang al., 2020; You et al., 2019). and whitening (Yang & Laaksonen, 2008; Kingma & Ba, 2015; Hwang, 2024; Jordan et al., 2024; Bernstein & Newhouse, 2024c;a; Carlson et al., 2015). In this work, we particularly focus on normalization and whitening. We apply normalization row-wise on gradient matrices, similar to (You et al., 2019); together with gradient whitening under a specific structural assumption of the Fisher Information (Hwang, 2024; Martens et al., 201 8), recovering the orthogonalization step used in (Jordan et al., 2024; Tuddenham et al., 2022). Our key result is that composing normalization and whitening on stochastic gradients is sufficient to enable the efficient training of LLMs in a completely stateless manner. Compared to Lamb (You et al., 2019), our normalization operation is applied on raw gradients instead of Adam states. Compared to Muon (Jordan et al., 2024), SWAN removes first-moment tracking and instead uses normalization, making the optimizer fully statless. Moreover, one of our proposed efficient NS scheme for computing square root inverse requires O(m²) (assuming m n) instead of O(m²) of the standard Newton-Schie Schulzz scheme. We show empirically that removing any one of these two pre-processing steps from SWAN results in significant performance degradation (Section 5.2).
The Role of Momentum There are a growing body of work that showed the benefits of applying momentum for optimization (Jelassi & Li, 2022; Cutkosky& & Mehta, 2020; Fu et al., 2023; Crawshaw et al., 2022; Zhang et al., 2020; Zhao al., 2024d). While our work aims at reducing all optimizer states including momentum, we emphasize that our results does not contradict the effectiveness of momentum; in fact we found that enabling momentum with the suitable hyperparameter can indeed further improve the result. However, our core contribution is to push the boundaries and demonstrate that: it is possible to train LLMs matching the performance of Adam using a completely stateless optimizer. Whether enabling additional states for SWAN can further improve the performanceis orthogonal to our contribution. This provides more options under certain practical scenarios. For example, due to its memory efficiency, the stateless design of SWAN without momentum may allow training with 2X larger batch sizes compared with using momentum; depending on the users's s application, the benefits of being able to run 2X larger batch size might out-weigh the benefits of momentum.
Comparisons to Apollo optimizer (Zhu et al., 2024) A closely related concurrent work is Apollo (Zhu et al., 2024). which proposes to re-scales the raw gradient of SGD to achieve SGD-like memory. Given a 2D gradient matrix of shape mn by n, Apollo propose to perform structured n- channel scaling (i.e., column-wise scaling), where the scaling factor is estimated via heuristic approaches similar to Fira (Chen et al., 2024). Compared to Apollo, SWAN enjoys the following advantages. 1), using the perspective of Apollo, SWAN can also be viewed as gradient rescaling G GredVaicentan(((dddoooC..) The key difference here is, Apollo can only perform channel n rescaling, and requires Or((n + +)) memory to maintain low-rank optimizer states to achieve channel n resealing. Apollo-mini further reduces the memory cost to (n): however it can only
G
16

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
achieves channel m rescaling with zero memory cost, offers richer representation power with much less memory cost. 2), In SWAN, the gradient preprocessing operators are chosen based on key theoretical insights of LLM learning dyanmics discussed in Section 4; while the Apollo design is mainly based on heuristics. 3), in practice. SWAN has stronger performance in terms of LLM pretraining. In the same setup, SWAN? achieves 1.46X and 1.36X speedup versus Apollo-mini on 350M and 1.3B Llama models, respectively.
cope with channel-1 (global) scaling. On the contrary, the SWAN rescaling G- GradBiiteenttggGrraaerroCCC
G
Comparison to Muon Apart from al algorithmic details, the key differences between our work and Muon (Bernstein & Newhouse, 2024b) is twofold. First, compared with Muon (Jordan et al., 2024), our core contribution is to push the boundaries and demonstrate that: is possible to train LLMs matching the performance of Adam using a completely stateless optimizer. Whether our method can consistently outperform all existing non-Adam optimizers (such as shampoo/Muon/SOOAP is orthogonal to this contribution. Our work thus highlights viable extreme stateless alternative for memory-consttrinee settings, and this distinction is central to our contribution. Second, the chain of thought behind the design choice of SWAN (exploiting structural redundancies of LLM dynamics to design composition of gradient processors) leads to fundamentally different research paths. From our view, the follow up work of (Gong et al., 2025) shows that the whitening operation used in both Muon and SWAN is a direct result of structured Fisher information matrix (FIM) approximation. As demonstrated in (Gong et al., 2025), this leads to potentially richer design choices that cannot be easily motivated within the normed steepest descent theory of Muon. Interestingly, (Gong et al., 2025) also shows that the structured FIM view leads to one additional coefficient Muon (v 1/max(m, n)) rescaling) that were recently re-discovered in (Liu et al., 2025) (for both easier hyperparameter transfer from Adam and stabilizing training). Finally, SWAN can be seen as fixed point solution of a iterative projection scheme now-normalizationn - whitening + row-normalization - whitening + ... This leads to a new theoretical framework that directly generalizes (Bernstein & Newhouse, 2024b). This is demonstrated in the other follow up work of (Scethon et al., 2025), in which the authors shows that the SWAN-motivated formulation opens up more design spaces for optimizers. Specifically, (Scetbon et al., 2025) demonstrated that more elegant and efficient matrix-based optimizers can be easily derived. This is also briefly discussed in Section 6.
Iterative methods for efficient whitening Computing Gradwhitening exactly can be expensive, as it involves solving the matrix square-root inverse. In the literature, iterative solvers for inverse square root has been explored. One of the most used methods is the Newton-Schulz method (Song et al., 2022; Li et al., 2018; Huang et al., 2019), which matrix- multiplication rich and allows more GPU-friendly estimation. In the context of optimization, the same N-S procedure has been proposed in (Mei et al., 2023) and (Jackson, 2023) for preconditioned optimizer. An alternative N-S procedure has also been re-discovered in (Bernstein & Newhouse, 2024b), and optimized in (Jordan et al., 2024). Although our inital implementation was inspired by (Song et al., 2022; Li et al., 2018; Huang et al., 2019), in principle any estimators can also be used. Furthermore, we proposed a efficient heuristic scheme, focusing on condition number minimization as a proxy for whitening. The proposed scheme requires O(m²) to estimate square root inverse instead of O(m²) of the original N-S scheme, and offers competitive performance in the context of gradient preprocessing.
Low-rank methods. Low-rank optimization techniques have been explored in the context of large language model (LLM) training as means to reduce memory consumption. These methods focus on applying low-rank approximations to model weights, gradients, and/or optimizer state variables. A seminal work in this domain is LORA (Hu et al., 2021) to fine-tune pre-trained models using additional low-rank weight matrices each layer, thereby significantly reducing memory usage to update the weights. More recently, methods such as ReLoRA (Lialin et al., 2023), FLORA (Hao et al., 2024), and Galore (Zhao et al., 2024a) have advanced low-rank optimization techniques for memory memory-efficient LLM pre-training. These approaches leverage low-rank gradient projections to enable full-rank learning, thereby achieving memory savings without compromising model capacity. Notably, they have achieved substantial reductions in the memory consumption of optimizer states, with only minimal impact on model performance. While these approaches effectively reduce the memory footprint of LLM training, they still necessitate storing internal states, resulting in higher memory consumption compared SWAN.
B. Desired properties of adaptive optimizers
There is a rich literature on understanding adaptive methods' inner workings and unreasonable effectiveness. Using Adam as an example, we first summarize from the literature below the key desired properties of stateful adaptive optimizers that contribute to their empirical success: gradient smoothing, gradient invariance, and gradient whitening. Then we discuss how these understandings will leads to the design of stateless adaptive optimizers.
17

---

SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training
Gradient Smoothing. Under the stochastic optimization setting, mini-batch sampling introduces heterogeneous distribu- tion shift on the gradient distribution: G F[G + +() where g(*) is time-heterogeneous noise induced by mini-batch sampling, While e(t) helps SGD escapes local optima (Jastrzębski et al., 2017; Zhu et al., 2018), the covariate shift of e(t) over time also present challenges to learning as the model needs to adjust and compensate for this shift, especially under the emergence of heavy tailed gradient distributions (Zhang et al., 2020)* ?. Following this viewpoint, it has been proven that momentum reduces the influence of noises for SGD (Cutkosky & Mehta, 2020; Crawshaw et al., 2022). Therefore we hypothesis that the first moment estimate (t) of Adam also effectively stabilizes gradient distribution and reduces effect of e(t), This smoothing stabilizes the variance caused by noisy stochastic gradients across time.
Gradient Invariance. More recently it has also been identified (Kunstner et al., 2023: 2024) that the major factor contributing to the performance gap between SGD and Adam might lie in Adam's Sign-descent-lik nature (Bernstei et al., 2018; Crawshaw et al., 2022; Chen et al., 2023). Intuitively, Adam without bias correction under 31 0 and [2 0 is equivalent to signed gradient descent (AW sign Indeed, the performance of Adam can be closely reproduced (Kunstner et al., 2023; Crawshaw et al., 2022) or even surpassed (Chen et al., 2023) by variants of signed descent with momentum. Apart from sign-based methods, evidence on performance boost using gradient hiiping/normalization was also discussed in the context of understanding Adam (Zhang et al., 2020). Therefore, we hypothesize that one of the key properties of Adam is that offers invariance over certain transformations on gradients. Particularly, the original Adam is invariant diagonal rescaling of the gradients (Kingma Ba, 2015); the signed gradient method is invariant to any scaling that preserves the sign of gradients; and the clipped SGD variant is invariant extreme gradient magnitude spikes.
Gradient Whitening. Finally, argue that the empirical success of adaptive methods also lies in that they model the curvature by first-order information. This is realized by the second moment estimate vit), which approximates the diagonal of the Fisher information matrix (Kingma& & Ba, 2015; Hwang, 2024); helping to counteract local curvatures of the problem. Specifically, Adam computes a trailing estimation of the diagonal coefficients of the Fisher matrix F E Eg by tracking F = diag( F) = diag[E[g"]], where g = vec G) is the vectorized gradient. Interestingly, instead of preconditioning the first moment as F -e'vve ve) Adam uses whitening-like preconditioned update F vec(m)) suggesting an element-wise approximate whitening of the gradient. It has been shown that such element-wise whitening leads to diagonal approximation to inverse Hessian F- x ) (Molybog et al., 2023). Recent empirical studies show that Adam biases optimization trajectories towards regions where the condition number of Hessian is low (Jiang et al., 2024). Therefore, we hypothesize that Adam approximately whitens the gradients element-wise, leading to well-conditioned regions.
C. Additional Analysis
C.1. Analyzing the Grad Pt. II: Robustness Against Local Curvature
In this section, we present main results regarding the convergence rate of the GradWhitening method, understand its implications, and compare it with the lower bounds of GD and Adam.
First, for simplicity, we focus on the following quadratic problem:
1 TT (w))-Hww-- (T((CT)) 2
C(W)-
where We E RRxn is the parameter matrix, H E Rmxm is a positive definite matrix, and C E Rmxx is a constant matrix.
For simplicity and without loss of generality, we assume C This is because minimizing L(W)- = 1/4 ]T((W HWW) Cr(C W) is equivalent to minimizing L(W)= [ (W- W
T H(W W*)], where W HH H''C. By defining Z = W - W the problem reduces to minimizing ((Z)= TT( HZ).
Remark Most results in this note can be easily extended to any loss function that are either i) strongly convex; ii) has twice differentiable functions and Lipschitz continuous Hessian, by considering their the second order approximation around W
Next, to understand the effect of GradWhitening we will examine the gradient flow dynamics induced by GradWhitening.
Such shift cannot be removed by forward covariate-shift reduction architectures such Layer Norm, as it is only invariant to global scaling and re-centering, such as wi (W(") + 71' for some scalar and incoming vector shift y (Ba et al., 2016).
18

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Consider the GradWhiteming--mdiiee gradient descent:
AW -/rraWhhitening((( (G()
(10)
its exact convergence rate is given by the result as below:
Theorem 2 (Contraction factor of bradWhiteningg Consider the quadratic loss function Equation (9). Assume the initialization distribution of w assigns zero probability to any set of zero Lebesgue measure in R Let our update rule be:
Wit+1)
w(O
whitened - -/rraWhhiteii ((9))
whitened
where the learning rate is Then, with probability 1, we have:
The optimal dynamic learning rate to achieve the fastest convergence is given by
HWW()) (t)
whitened
(11)
Tr|H]
where | (t)
whitened 1 denotes the Schatten p-norm with 1 (i.e,, sum of singular values).
Under 0f the contraction factor
factor of loss function at t is given by:
L(W
whitened
(t++)
L.
=1 =1-
HW (t)
whitened
(12)
C(W
whitened
(t)
Tri(W
(t)
whitened
HW (t)
whitened [r
Furthermore, we additionally enforce W Vinxn (R), iee, initialized as element in Steifel manifold. Then have
L(W (Wimleee)) L C(W9)
(13)
C-
That is, GradWhitening solves the optimization problem (Equation (9)) with 1 step iteration.
Theorem has the following key implications.
Convergence rate is condition number agnosite Unlike the convergence rates of GD and Adam presented in Zhang et al. (2024a), as well as Theorem 3 and Corollary Appendix, the optimal convergence rate (12) of GradWhitening no
(HW() 11
longer explicitly depends on the condition number x of H. In fact, consider a lower bound Tr((W olionee
votemed
Whiened Tr(H] 2
J TTWW"
Tr/#/ chicned!
Hw Aitered H)) since trace of H appear both in the nominator and denominator, we expect that to be more robust
Tr(CW ) to ill-conditioned problems. For example, consider the specific initialization W(()
whitened = cT, it is straightforward to show L K, which is completely disentangled from the condition
(
(
waismalll lf
Tr( (t)
that
IV
vaiemall
(0((W-()) ) H
vanensal
[T(H]
Tr((W:0)
vanamaa)
HW
vineraal Tr|H]
[[HWO)
number. Hence
streeee
would not shrink as k co. See Proposition 2 for less extreme situations.
starred
HW
TH
whitened initialized the Stiefel manifold, then wradWhitening reaches superlinear convergence rate (= Newton's 's method), while being cheaper. In fact, it is straightforward to verify that GradWhiteming reaches optimal solution with step update. This implies GradWhitening is theoretically the optimal optimization algorithm if W is initialized in the Stiefel manifold.
Superlinear convergence with Stiefel manifold initialization Theorem suggests that if (t)
Estimation and interpretation of optimal learning rate Compared the optimal dynamic learning rate of gradient G descent G .G the optimal learning rate rn(*** of GradWhitening is much easier to compute. Tri
vitamal can be
Tr
seen as balancing the average gradient magnitude against the average curvature. A higher trace of gradient (HW whitoned) (strong gradients) relative to H (steep curvature) suggests larger learning rate, promoting faster updates. Conversely, a higher trace of would imply smaller learning rate to ensure stable convergence in regions with high curvature.
19

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Next, show that the convergence speed of GradWhitening update is indeed robust the condition number of local curvature.
Proposition (Robustness of GradWhitening update convergence rate against the condition number of local Hessian), Consider the quantity:
T whitoned!
Tr[(W
(t) whitened)
y HW
[r|H]
whitemed
Assume: i), w
whitened + W"; and ii) the norm of H is bounded. Then, there exist some finite positive constant c, such that
Q>
This holds even if + +oc, where is the condition number of H.
Remark Proposition 2 does not that the bound cannot be arbitrarily small, but rather: a small condition number does not necessarily lead to a diminished bound. This distinguishes our method from standard SGD, where such a decrease would be expected.
Below, we provide comparison between GradWhitening modified gradient descent and Adam. We only consider non- Stiefel initialization for GradWhiteningg since with non-Stiefel initialization GradWhi tening is optimal according to Theorem 2. Our results below shows that, for poor conditioned problems GradWhitening with properly chosen single global learning rate always outperforms Adam even with optimally tuned sub-group learning rates, in terms of convergence speed.
with single vs Adam with tuned group Ir). Consider the optimization problem Equation (9). Assume H is block-diagonal, Le., H diag H1, H2...., H1), where each H, Imm.xm; positive definite matrix for 1 = 1.2....,L. L, and Σ my = m. Assuming for GradWhitening we use one global learning rate for all parameters; and for Adam, we use the optimally chosen group learning rate mp and initial condition wo for each block H
Proposition (GradWhiteninn
Assume either i) certain regularity conditions are met (see proof in Appendix), or ü), H is poorly-conditooned (its condition number is large enough). Then: regardless of its initialization, GradWhitening with properly chosen learning rate will still have strictly better convergence speed (i.e., smaller contraction factor) across all blocks E [L] than Adam (B) = 0, = 1) under optimal group-wise learming rates and initial condition.
As pointed out by (Zhang et al., 2024a) and (Da Silva Gazeau, 2020), Adam with 32 < will have issues with convergence, which will not be completely removed even with Ir decay. Therefore, we will not discuss the case of B3 to avoid the complication.
Remark
C.2. Numerical Verification of Proposition 1
Given a STB, we consider the following standard full-batch learning dynamics (Tian al., 2023). Define the conditional expectation iF [+ == El:1g = m). Consider the dynamics of the weight matrix W and the attention logits Zy. if we train the model with a batch of inputs that always end up with query g[i] m. The weight update for W and 2g are given by the following noisy updates:
W = Eg=m (G O [
+ m = Eggm
(c00)
ab
T
g
(14)
(t)
Where (()) = (Uc (exp(z(()) (1) 2 2z) u.)) (h( o' ((W()) (() is the derivative of the current activation,
G
Vaco C is the gradient of the loss function with respect to the hidden activation (1), and (t) pco = Σn
(00 (h
(1)
(0)
(1)
is the sum of the gradients with respect to the attention logits. Here, w is the k-th column of W(O) g (g backpropagated gradient sent to node k at sample i.
be the
Then, we numerically solving the STB ODE with n = 12, Mc = 10 in Equation (14). During all training steps, we analytically track the evolution Hessian of rmW. Results are shown in Figure As predicted by Proposition I, we see very similar structures across the diagonal blocks of the Hessian.
20

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Normalized Diagonal Blocks
H.2
H.3
H.4
forametors
Portmmotors
forameteons
Porometoors
0.5
Peraimeters
Parameters H. 1.11
Parameters H,12
H.9
H
Figure 5: Normalized Hessian Blocks of size Mc x Mc along the diagonal direction of the Hessian, obtained from numerically solving the STB ODE (with n= 12, Mc = 10) (1) given by the full-batch dynamics (i.e., removing noise in Equation (14)). During all training steps, we analytically track the evolution of Hessian. As predicted by Proposition 1, we see very similar structures across the diagonal blocks of the Hessian.
D. Proof of Theorem
Proof. We first consider the noiseless, full batch dynamics. Define V Mexn V U. W. Then following Theorem
2 in (Tian et al., 2023), each column of V satisfies the following differential equation:
[ = (V? /2 +
E, (gh
(15)
The corresponding dynamics of attention score is given by:
Zn
v2 [d)]
(16)
2
Without loss of generality, in this proof we only consider C = 0.
Now, following the argument of Lemma B.6 of (Zhao al., 2024a), we reparameterize the dynamics row-wise. For this, consider instead
u
U2
V
1 Mc
21

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Then, equation 15 becomes:
ú = (exp(u,") 1|ps
(17)
where E x is given by [ull := 19,,,, TzI. Therefore, it is clear that u, always move along the direction of μed due to the stationary back-propagated gradient assumption. Hence, úi; = (z(t)) for some scalar dynamics ou(t).
Next, consider the mini-batch version the dynamics. this case, the packpropagated gradient term [po]; = E 5
corrupted by some i.idd mini-batch noise &. The noisy row-wise dynamics now becomes:
' at) + E,)
(18)
Therefore, after row-wise standardization, the new dynamics becomes
a;(1)( + - E, Mij + I, Eb) à; (t(t)( E,(ni + Eig 1, " E, Mij - 1 E, {1))
(1i + 41 - ( E, Mij + 56) Eo) C, + Eij - :/ E, Mig - - E,
Therefore, the normalized noisy gradient û; no longer depend on the time-variant component (t)). Hence, we have proved:
Cor[Gugw(i.:] Cov ((a) for all &1, 12, and i.
The corresponding result for G(t) == GradNorm (T,.),(TTx!) can be trivially derived due to Equation (16).
E. Proof of Theorem 2
Proof. We first show that VZ(W(()) = HW (and hence V VL(W(t) whitemed) with o oo) are non-zero with probability I under Assumption of the theorem. Given VL(W(")) HW the set of matrices W such that (r)) = forms hyperplane in the space of d d matrices. Specifically, it is defined by the linear equation: TH(:: = 0. Since H is positive definite, at least one entry of H is non-zero. Thus, the hyperplane Tr(HW(()) = 0 has zero Lebesgue measure in the space of d x d matrices. Given that (' is sampled from a continuous distribution, the probability that (r(HW(()) 0 is zero. Therefore, V.((())) j 0 (and hence VL(W(()
whitened) with 1 o 2o) with probability 1.
Next, we define the cost-to-go as:
((W(9)- - -- 1 Tr (W(C))
and the per-step improvement is (since L = 0 under = ==0,)
- -
1 Tr Tr (Wr))) T HW 2
Tr (wi++1)) (
2
Substituting the update rule W(++1) w(() nGGaddhittnnngg Gwhitened.)) w(()) - -U we get:
((W()) L( = Tr
( (aw(O)THWiCH - Tr
Expanding the right-hand side, we have
Tr [UV [H(VT)) - 2
- 9
((W()) - ((W(++)) = n/T (wi()) T HHUV]]
22

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Now, noticing that G HW UEVT, have:
HW Tr (HW(o)TTUT) (T Tr UVU] Tr [VEV]] Tr
Since V is orthogonal, V = I, and is diagonal, we obtain:
Tr (W()) [WUV Tr(2)) |HW (t)
whitened 1I
Similarly:
Tr ((UVT)"H(UV))) -VU''uuVV"] [VAV]]]] Tr(H), where A is the eigenvalue matrix of H. Given those intermediate results, have:
((W(+1)) L- C(W(O)- L"
- - ((W((1))
((W()) - L
H (1) whitened |1 TT()) Tr (W(O))TWW))
Noticing that this quadratic function and the second order coefficient is positive, is straightforward to verify via the quadratic formula that the optimal learning rate is given by
HW
whitened 11 Tr(H)
ni
Under which the optimal contraction factor is given by
((W(t+1) (t)
C
HW (t) whitened
whitened
=1
L(W whitened
Tr
(w
whitened)
HW (4)
whitened
Tr(H)
Finally, if we additionally enforce (0)) (R), i... we can parameterize w O where O is orthogonal, then it is trivial verify that GradWhitening reaches the optimal solution with 1-step update. To this, consider the GradWhitening update:
[HO||
(1)
whitened = O - "rrdWWahhttttnn HO) = 0 GradWhitening
Tr(H)
[HO noticing that Tr( L. and GradWhitening(HO) P(QA)TTPPOQTTT O qQ O. Hence:
W (L)
whitened = -7 = -r"GradWhitenning/ (HO) = O - 0 = 0 = W
Hence, the proof is complete.
F. Proof of Proposition
Proof. Since (e) * W the square of the trace of the gradient Tr HW (0)
whitened must exceed some positive constant
whitened
Cc. that is, Tr HW (f)
12 Cc.
whitcned!
On the other hand, because:
1. The quadratic loss term Tr (W)))) whitened)
HWW
whitened upper-bounded on RaxR, and
23

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
2. Tr[( (r((Wot)) whitened
HW whitenedd
# (due to ((()
whitened o W
we have that there exists positive number 0 < Cc such that
< Tr(( whitened
HW
whitened
< Cc.
Finally, since the norm of H is upper bounded, its trace must also be upper bounded by some constant CH. Therefore, putting everything together, we have:
Q CcCh
This inequality holds even as the condition number R - +00.
G. Proof of Proposition 3
To prove Proposition 3, we first generalize existing work on the convergence rate lower bound (via contraction factor) of gradient descent and Adam (we only consider B2 1) under the same setting:
Theorem 3 (Contraction factor lower bound for gradient descent, generalized based on Zhang et (2024a)). Consider the optimization problem in Equation (9). Let Wt be the output of GD after t steps. Then, for any step size v, there exists an initial condition such that the following lower bound the contraction rate holds:
L² (1- (C(WED) L*),
L(WL GD
where L- = -(W"). Furthermore, under optimal 17 Ju +mm- the bound becomes tight regardless of the settings of H. where dn and Dm are the largest and smallest eigen values of respectively.
Proof. The proposition in Zhang et al. (2024a) has shown that the lower bound holds for diagonal positive definite Hessian H. To show that the lower bound holds for a general positive definite Hessian H we will reformulate the problem to align with the setup in diagonal case (Prposition of Zhang et al. (2024a)).
First, for any positive definite Hessian H, we can perform an eigen decomposition USUT, where U is an orthogonal matrix and is a diagonal matrix containing the eigenvalues of H. Define change of variables Z U W. Then, the optimization problem becomes
1 Tr(Z SZ),
C(Z)
2
which reduces the problem to the case of diagonal H with condition number A1. where a and Am are the largest and smallest eigenvalues of H, respectively.
Thus, by applying Proposition of Zhang et al. (2024a) to this transformed problem, we conclude that there exists initial point such that the lower bound on the contraction rate
(E(W GD) -2*) (
((w(tt)) - L² -
L G
K
holds for the transformed variables and, equivalently, for the original variables W since the condition number is preserved under orthogonal transformations.
Therefore, the lower bound for gradient descent applies any general positive definite Hessian H provided the condition number k remains unchanged.
Finally, under the optimal step size A the bound becomes tight regardless of the settings of H. This is achieved by selecting to minimize the contraction factor, aligning with well-known results regarding the optimal convergence rate of gradient descent on quadratic objectives (Nesterov, 2013).
This completes the proof of Theorem
24

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Corollary (Lower bound Adam (32 1)). Consider the optimization problem in Equation (9) Assume the weight
initialization 8 assigned zero probability to any set of zero Lebesgue measure in Rm Let W be the parameter
Adam
after L iterations of Adam with hyperparameters 31 = and 32 1. Then, for any step size t, the following lower bound on the contraction rate holds:
L L(W(++) Adam
2 --- 22 R' (WW +1) (
(E(W Adam) -C-), (E()))am)-
where k2(W)) is the W -epeendent condition number of the preconditioned Hessian diag( H-)H- and L L(W*). Proof.
The update rule Adam with 31 = 0 and B2 is given by Zhang et (2024a):
w (+++)) Adam
w(t) (+)
77 diag (HW Hw Adam
Adam
This can be interpreted gradient descent with preconditioned Hessian matrix diag H I"HH. By applying Theorem 3, we conclude that the contraction rate for Adam under these settings satisfies the lower bound:
((W Adam
L(
-- (1
x+11
(E(W Adam ---).
where K is the condition number of the Hessian matrix H. Therefore, the proof is complete.
Next, We extend our Theorem block-diagonal Hessian case to prepare for discussions on group learning rates when comparing Adam.
Corollary (Upper Bound Convergence Rate of SWAN). Consider the same quadratic loss function L(W) [T(W HW) with H being block-diagonal. That is, H diag(H, H2...., H1), where each Hi Rmuxx is positive definite matrix for | = 1,2,..... L, and El= mi = m and El= 14 =n Assume the initialization distribution of O()) assignes zero probability to any zero measure set Rmx* Let w(l) whitened be the parameter matrix after iterations of the SWANoptimizer defined in Theorem 2, with learning rate r7. Then, under the conditions that 8:
H.
( 12 - TH(H) Tr (wW()) (.,((")
2d1.m
O,
A1,1
A1
where Aι,1 and Me, are the largest and smallest singular value of Hi, respectively; then there exists a proper learning rate t such that: with probability 1, the loss satisfies:
C(W(t+1)
2 Ki + 1
21)
(C(W whitened) -C").
whitened)
L.
max [E[L]
1
(19)
where K₁ is the condition number of H,
Proof. Applying the arguments in the proof of Theorem 2 each block I, we have (for simplicity, will drop the subscript "whitened"" when there is no confusion):
L(w((+)) L- =1 ((W(") ((W()+1) ((W(")
((W())-2 C.
n H, (01 TTT((
1T (w/) (W(") (H,"
'Nole that according to Proposition 2, this always achievable when H. is poorly-conditioned
Atdd small enough). The lower Ara
Tx( (t))
interval above converges to zero, and one can simply pick e.g., = minie())
25

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
is straightforward to verify via the quadratic formula that if one chooses satisfying:
H. W 11 11 H H Tr(H.) (r( TH, H.W, (9)
Ar.
<1
Tr(H,)
2A H 11 + | |H l1 Tr(H)) TT((W)
A1,
Tr(H.)
then we have:
Tr (w w
K] +1
H 11 - A TH(,)
2
Rearranging, we obtain:
((W((+1) L ((W()) Lr
2
Since H is block diagonal, the updates for each block are independent. Summing the loss over all blocks, we obtain:
I.
C ((W whitened) - -- Σ (c(w(())) --
2 ks +
(C(W("))
Taking the maximum contraction factor across all blocks:
7 ((W(")) - -2-)
((W(()+)) L whitened)
L" max [E[L]
2
Σ
Ki +
1=1
2
(C(W (C(W(hheeed) whitemed) -c*).
max lE(r]
1
Kit
+
Thus, the overall contraction factor for the SWANoptimizer is:
PSWAN
max [E[L]
1
Ki + 1
Finally, we are ready prove Proposition
Proposition (GradWhi GGradWhitening with single vs Adam with tuned group Ir). Consider the optimization problem Equation (9). Assume H is block-diagonal, ie.,
Hy), where each H, pm, positive definite matrix for 1 = 1,2...., L, and Σl= my = m. Assuming for GradWhitening we use one global learning rate for all parameters; and for Adam, we use the optimally chosen group learning rate 1n and initial condition wp for each block H4 Assume either i) certain regularity conditions are met (see proof Appendix), or ii), H is poorly-conditioned (its condition number is large enough). Then: regardless of its initialization, GradWhitening with properly chosen learning rate will still have strictly better convergence speed (i.e., smaller contraction factor) across all blocks [L] than Adam (B1 = 0, 1) under optimal group-wise learning rates and initial condition.
Proof. For simplicity, we will drop the subscript "whitened" when there is no confusion. Let Ki ((W denote the w( (0)
dependent condition number of the 1-th block preconditioned Hessian diag( H, w HH, Let A (W(")) and
1 (W())) be the smallest and largest eigenvalues of diag( H, W HH respectively. Then,
26

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Case 1: Under the conditions that:
1. Existence of roots: vI E [L].
mm, (W (w(o)) (10)
H. W (9) I - Tr(::) Tr (ww(")H,,("" (H, N At (
> 0,
and
2. Overlap condition:
( + Tr(H) Tr(H.) Tr )T,"))
TH:W(")
2m (w)
(W, At,m (W
min [E[4]
Hr(
T((H, W(()2 -T(H)) (W,W")
2A (W
AH, A (W,))
> max [
Tr(H;)
Then, there exists global learning rate y, such that for all E [Z].
L(W whitened): Li
((w((t+))
2
L L(W (t+1) Adam ) Li
[(W((+1))
<1
C( C(W (t)
whitened))
LT
(0)
+1
C( Adam
LJ
At,m: (W)
Case 2: If Hiin poorly-conditioned, i.e.,
O, then Proposition 2 asserts that the following term
(W
((() TT( - ;r Tr ((W))) (T) [E[T]
2A (W + (W 0,
Tr(H;)
Tr(H, w THHH))
and one can simply choose, for example, 1) minre(r)
Under this choice of 17. still have
L( ((w(++1)
C( L(w(t+1)
whitened ): Li
Adam ) Li
<
L( (W(()
whitened ) Li
K((W 00) +
((W(t)
Adam ): Li
for all I [L].
H. Proof of Proposition
Proof. First, define V E Mexx as = VUUUV W, and consider the Hessian with respect to instead of W. Notice that although the loss function C is unknown, its first-order derivatives are known. Specifically, they are given by:
ac = in = Eg=m [ghy 5z e² Σ, of.
Beik
Therefore, the second-order derivatives, i.c., the Hessian matrix H(V), are:
H(V)..P' = 8 2
E Egmm (gn, mij e E, v8.
Ooyk where 6f is the Kronecker delta, which is 1 if = l' and 0 otherwise.
Based on Lemma B.6 in Zhao et al. (2024a), as f + oc, there exists an index subset O₁ C fl. Mc} such that:
up*k k > Utk,
iy >> Dik, 4I E o I E Oi, Vk.
27

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Consequently,
H(V) > H(V)RR,E,,
J. =1 O1, 1,1' & Ot, Wk
After normalization, as t + 0o, we have
H(V)uuK JuNe En H(
[1, if 1' O
1000
0,
otherwise.
Reverting back to the W space, we have
H(W) (Ix @ (V) & Uc)
where & denotes the Kronecker product and I is the identity matrix of appropriate dimensions.
Therefore, for all 1 S. s, o' < d and S k, k' S n, we obtain
H(W) skow
H(W) sk s' E.,o HWW),,
as - 00.
E.,' ((W) o'k
This holds for all 1 S 8, S d and < k.Ms
Remark As pointed out by one of the reviewers, at equilibrium, Or contains only single index. We acknowledge that during early training before convergence- mulliple dominating indices may appear (as observed in numerical experiment in Figure 5), and similar structures are still present across the normalized diagonal blocks.
I. Acceleration of Newton-Schulz iteration via diagonal substitution
L1. Algorithnms
Computing GradWhitening exactly be expensive, it involves solving the matrix square-root inverse. One option is to directly apply the Newton-Schulz variant of decorrelated batch normalization (Song et al., 2022; Li et al., 2018; Huang et al., 2019), which allows a more GPU-friendly estimation. This is given by (Song et al., 2022; Lii et al., 2018):
Y 14+1 = 1 Y Zz vn)
Z (3I Z₁ Y
where Y, = GG Zo = I. At convergence, GradWhitening(G) ZG (Algorithm 2). However, estimating
(GG T) ) with NS requires (mm)) (assuming m < n) complexity. Here, we propose heuristic scheme that has (mm²)
complexity to estimate square-root inverse:
kt+ Y Diag 3I Diag (Y+))) Z } (31 - Diag Z₁ Y Diag(aZ (Z.)
where Diag(-) returns a diagonal matrix that has the same diagonal elements as the input matrix. Basically, whenever encounter matrix multiplication in NS iterations, replace one of them by its diagonal approximation. We refer to this as the NS with diagonal substitution (NSDS) scheme.
Note that in the above standard presentation we have fixed that both NS and NSDS uses coefficients = 0.5 for Y updates and Z updates, respectively. In practice we may further tune these coefficients to compensate for short number of iterations (usually under 10).
1.2. Experiment: LLM Gradient Condition Number Reduction
Setup In this synthetic experiment, we assessed the effectiveness two whitening methods, Newton-Schult and the proposed Newton-Schulz with diagonal substitution (NSDS) scheme, on gradient matrices obtained from LLM training. The exact GradWhitenie teninng operator results in matrices with optimal condition number (= 1); therefore, we hereby investigate
28

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
the matrix condition numbers of the processed gradients obtained from different methods. Specifically,, both methods use 5 NS iterations with NS step size optimized. We train a 130M L.Lama model following the architecture setting of Section 5.1 on randomly generated sequences for 1000 steps, and take the MLP weights of middle layer (we take the fifth layer without loss of generality) and use different methods to whiten the gradient matrices. We consider three methods: standard NS; NSDS; and SWAN with NSDS (that is, composing GradNorm with SDDS-GradWhitening condition number reduction ratio of different method was calculated for both whitening methods (the higher the better). Note that for all methods, the gradients have been pre-normalized by its norm.
At each training step, the
Results Results are shown in Figure 6 (a). We notice that NSDS alone (orange curve) not sufficient to reach a good condition number reduction ratio. However, when combined with GradNorm (i.e., SWAN with NSDS, the green curve), its performance started to catch up and even outperform the standard NS method after 500 training steps. This show the effectiveness of the proposed scheme. One potential caveat that spot is that the condition number produced by SWAN with NSDS is more noisy than the standard NS iteration; which might lead to improvements that will be addressed in future work.
The significance of GradNorm and NSDS The results above highlight the importance of GradNorm. A potential question is whether GradNorm, when followed by GradWhiten is merely a no-op that rescales the initial location of the NS iteration for better convergence. In fact, for all methods considered Figure 6, the gradients have been pre-normalized by their (global) norm before being fed into each method. This re-scaling applied to all methods provides negative answer to the question. To clarify further, we estimate the cosine similarities between the processed gradients produced by different pairs of methods. The results, shown in Figure 6 (b), reveal the following: As the training iterations increase, the cosine similarity score between SWAN-NSDS and NS (denoted as cos (SWAN-SSDS, NS)) monotonically decreases to small values, indicating a wear-orthogonal relationship. Both (SWAN, suggesting that both GradNorm and NSDS contribute to the orthogonality of cos(c (SWAN-NSDS, NS). This demonstrates that the changes introduced by both GradWhitening and NSDS-Gr GradWhitening are significant, rotating the update by a relatively large magnitude. This might also explain the observation Section 5.1 that SWAN with NSDS behaves differently from other variants, showing slower early convergence but stronger long-term convergence.
and cos( SWAAN, SWAN-NSDS) decrease over time,
Condition Number Reduction over Training Iterations
Cosine Similarity Between Preprocessing Approaches
SWAN NS2S NS
SWWAN vs AS
SWAN Vs SWAN-N5DS
SWAN NSDS NSDS
Condition Redaction
Cogine Simiarty
0.6
0.4
0.2
NS only
NSDS only
SWAN (NSDS)
0.0
0.0
200
600
800
1000
200
400
1000
Treining Iteration
Triniin iteration
(a)
(b)
Figure 6: Comparison between different whitening schemes. (a) Performance comparison by condition number reduction ratio (higher the better). (b) Cosine similarities between the whitened matrix produced by different preprocessing schemes, respectively.
I.3. Ablation: effect of NSDS iterations
We examine the impact of the number of iterations of our Newton Schulz with Diagonal Substitution (NSDS) scheme. With SWANt, we compare the test PPL performance on 130M model. Results are shown in Table 4. As we can see, the improvement brought by additional NSDS iteration marginal; hence in this paper, we only apply 2 iterations of NSDS.
29

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
KL Divergence to Initial Step (model.layers.2.mlgg wppppro..weihtt
KL Divergence to Initial Step modell.ayerr..2.see. wtw.lleess..s...atttn..... (lr..roo.llgg
Without GradNorm
Without GradNorm
With GradNorm
With GradNorm
(Diiergence
Divergence
3.43
43000
.101
.Ee Training Steps
A46
128
1286
Training Steps
(a)
(h)
Figure 7: KL divergence comparison of gradient distributions against initial gradient distribution across training Steps. We the projection weights in attention and MLP modules of the second layer example. The plots compare standard training with GradNorm- Gradloom-augmenttee training. Lower KL divergence values indicate greater stability in gradient distributions.
Table 4: Effect of NSDS iterations on test PPL. Results are obtain using FP32 precision.
NSDS iterations
Test PPL
(LLM loss diverge)
22.63
22.62
10
22.61
1.4. Ablation:
Precision
We compare the performance of different precisions of GradWhitening with NSDS scheme. As shown in Table 5, on 130M model (2 step NSDS iterations) ablation we show that BF16 can be used without major performance degrade compared to FP32.
Table 5: BF16 vs FP 32 NSDS
NSDS Precision
Test PPL on 130M
BF16
22.61
FP32
22.63
J. Additional Experiments
J.1. Empirical verification of theoretical insights from Section 4 regarding GradNorm and GradWhitening
J.1.1. DOES GradNorm STABILIZE GRADIENT DISTRIBUTIONS OF SGD?
To examine whether GradNorm stabilizes the distribution of the stochastic gradients suggested by Theorem 1, we conduct controlled experiments using scaled-down LLaMA- based model (about 10 million parameters) (Lialin et al., 2023), trained on the C4 dataset. Our goal is to measure how GradNorm affects the distribution of stochastic gradients over multiple training steps. Specifically, we employ a small-scale I.I L.aMM-bssed model with approximately 10 million parameters (Lialin et al. 2023). Training is conducted on the C4 dataset (Raffel et al., 2020).
Baselines
We compare:
Standard training: This uses an SGD optimizer with learning rate of 5 x 10 and linear learning rate scheduler, including 10% warm-up of total training steps (10,000 steps).
30

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Comparison of Optimizers Quadratic Problem
Comparison of Optimizers on Quadratic Problem
10c
Gradient Descent
Gradient Descent
10"
10'
Adam (B,=0,9, 8x=0.999)
Adam (B,=0.9, Bv=0.999)
(Log Scale)
10a
GradWhitening (Orthogonal Init)
Loss (Log Scale)
10a
GradWhitening (Orthogonal Init)
10"
GradWhitening (Non-Orthogonal Init)
10
Newton's Method
Newton's 's Method
10-2
10
10
10
10-4
10-6
10
10
10-10
10-16
800
Iterations
Iterations
(a) Quadratic optimization, well-conditioned
(b) Quadratic optimization, ill-conditioned
Comparison on Rastrigin Function (2500 dim)
Comparison on Rastrigin Function (2500 dim)
101
10a
(Log Scale)
102
(Log Scale)
10a
107
Gradient Descent
Gradlent Descent (Orthogona Init)
Adam (Bæ=0.9, 9x=0,
10-2
GradWhitening (Orthogonal Init)
10-2
Adam (8s=0.9, f2=00999,, Orthogonal Init)
GradWhittning (Orthogonal Init)
10-1
GradWhitening (Non-Orthnn Orn-Orrhhaanal Init)
10
Newton's 's Method (Orthogonal Init)
Newton's Method
10-4
10-6
10a1
10-a
10-10
10-10
200
400
600
800
1000
200
400
600
800
1000
Iterations
Iterations
(c) Rastrigin function optimization
(d) Effect of initialization
Figure 8: Comparison of convergence rate of different methods on quadratic and non-convex optimization problems. (a): -i0-ddimnssional quadratic optimization with well-conditineed H. (b): 2500-dimensional quadratic optimization with ill-conditioned (c): 2500-dimensional Rastrigin function optimization. (d): 2500-dimensional Rastrigin function optimization, but foreing all methods to use the same orthogonal initial location.
GradNorm-process training: This applies GradNorm to pre-process the stochastic gradient before the parameter update. All other settings match the standard training baseline.
Methodology We measure gradient statistics in the presence of mini- batch noise. At step t = O, we sample 16 additional mini-batches (batch size 64 each) and compute the mean and standard deviation of the corresponding raw or GradÑorm gradients in each batch, and obtain the approximated initial gradient distribution. After each training step of baseline methods, we perform the same procedure and calculate the Kullback-Leibler divergence between the resulting gradient distributions and the initial gradient distributions. This process tracks how the gradient distribution changes over time.
Results Figure shows the KL divergence of gradient distributions for standard and GradNorm augmeette training. relative to the corresponding initial approximated distributions. Apart from early spikes, GradNorm reduces fluctuations in the gradient distribution throughout training.
1.1 DOES GradWhitening COUNTERACTS LOCAL CURVATURE AND PROVIDE FAST CONVERGENCE ON ILL-CONDITIONED PROBLEMS?
This subsection evaluates the optimization performance of gradient descent when combined with GradWhitening We use three classic problem settings:
High-dimensional quadratic optimization, quadratic problem the form in Equation (9), where W E 5550x50
31

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Ill-conditioned quadratic optimization. Same setup as above, but with deliberately chosen ill-conditioned H.
Non-convex optimization with multiple local optima. We use the multivariate Rastrigin function:
f(W) = m2 2 A Wr - A E cos(2
where W is an m x m matrix and m = 50. This function has 10T" possible local optima.
Baselines We compare five methods all three problems: gradient descent (GD) with the theoretical optimal learning rate in Theorem 3, Adam with 31 0.9, B2 0.999 and hand-tuned learning rate, Newton's method with tuned learning rate, and two GradWhi tening-based variants (with and without orthogonal initialization). This is to verify, under orthogonal initialization, GradWhite tening--roceseed GD behaves similarly to Newton's 's method, as discussed in Section 4.3 and Theorem 2. All methods share the same initialization, except for the orthogonal aradWhitening variant, which projects the initial parameters onto an orthogonal matrix.
Results From Figure (a)-(c) we summarizes the following outcomes:
Quadratic problems (Figure 8 (a) and (b)). GradWhitening with orthogonal initialization and Newton's 's method converge optimum in one step, aligning with the theoretical predictions in Section 4.3 and Theorem 2.
Well- vs. ill-conditioned cases (Figure 8 (a) and (b)). In the well conditonit setting (a), standard GD outperforms both Adam and GradWhitening (non-orthogon initialization). In the ill-conditioned setting (b), GradWhitening (non-orthogonal initialization) outperforms GD by a large margin, while GD experiences slow convergence.
Comparison with Adam (Figure 8 (a)-()), In all three settings, GradWhitening with non-orthogonal initialization consistently outperforms Adam, consistent with Proposition 3.
Rastrigin function (Figure (c)). this non-convex problem, GradWhitening performs comparably to Newton's method, with or without orthogonal initialization.
Effect of initialization (Figure 8 (d)) Furthermore, force all methods share the same orthogonalized initialization. As shown by the result, GradWhiteminn GD with orthogonal initialization still consistently outperforms all baselines, confirming that this initialization is only beneficial to GradWhitening GD among all basclines.
J.2. Is the speed-up Multiplicative or Additive?
A key question in assessing speedup factors is whether the improvement over Adam is multiplicative or additive. A multiplicative speedup implies that the optimizer's relative advantage remains proportionally consistent over time, while an additive speedup suggests less desired constant step advantage. To investigate this, We focus on the SWAN-0 variant as an example, and explore these scenarios using two plots, the speed-up ratio comparison and the perplexity comparison (Figure 9), across different model sizes.
Speedup Ratio Definition We define the speedup ratio R(P) for a given perplexity (PPL. threshold P as the ratio of the number of training steps Adam requires to reach specific evaluation perplexity (PPL) to the number of required steps for SWAN:
R(P) = SAdam ( (P)
(20)
SSWAN (P)*
where Sddam (P) and SswAN(P) are the training steps required by Adam and SWAN, respectively, to reach perplexity P. To test whether the speedup is additive, we compute a counterfactual additive curve by assuming SWAN gains a fixed step advantage over Adam in early training (approximately the first 10%-20% of total steps):
32

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Speedup Factor Comparison Multiple Models
PPL Acceleration Comparison for Multiple Models
60M Actual
60M Counterfactual Additive
130M
130M Counterfactual Additive
Speedup Factor (Ours/Adam)
350M Actual
350M Counterfactual Additive
1.3B Actual
1.38 Counterfactual Additive
Ours PPL
60M Actual
60M Counterfactual Additive
130M Actual
130M Counterfactua Additive
350M Actual
350M Counterfactual Additive
1.38 Actual
1.38 Counterfactual Additive
Training Steps
(a) SWAN speedup factor vs training steps
(b) Ours PPL. vs Adam PPI. scatter plots
Figure 9: Comparative analysis of SWAN and Adam optimizers: speedup ratios and perplexity metrics across various model sizes. (a) shows how SWAN reduces the number of training steps needed to achieve the same evaluation perplexity Adam for models ranging from 60M to 1.3B parameters. A speedup ratio greater than indicates that SWAN reaches target PPL values faster than Adam. (b) presents direct comparison of perplexity scores between SWAN and Adam. In both plots, we also provide counterfactual additive curves (dashed lines), which displays hypothetical baselines that has a constant step advantages over Adam. Together, these plots highlight the nature of SWAN's speedup over Adam across different model scales.
Counterfactual Additive Curve Estimation To test whether the speedup is additive, we compute counterfactual additive curve by assuming SWAN gains fixed step advantage A over Adam in early training (approximately the first 10%-20% of total steps):
A
SS (P:) SSWAN
(21)
where N is the number of PPL thresholds considered. We then use A to define the counterfactual additive speedup ratio:
SAdam (P)
Radiiniv (P)=
(22)
SAdam (P) A
and the counterfactual additive perplexity estimate:
PPLaddiive (S) PPLAdam (S + A).
(23)
This represents the expected perplexity of SWAN if it only consistently outperforms Adam by A steps.
Results We focus on the SWAN-0 variant as example. Figure 9 compares SWAN-O's actual performance against the counterfactual additive curves. If its actual curves exceed these additive estimates, it indicates tendency towards multiplicative speedup, instead of the additive advantage. We summarize the observations for model sizes of 60M, 130M, 350M, and 1.3B parameters:
For smaller models (60M and 130M), the actual speedup trajectories align closely with the additive baseline, indicating primarily additive speedup.
For larger models (350M and 1.3B), the actual curves rise noticeably above the additive estimates, suggesting a multiplicative speedup. This indicates that SWAN yields increasing efficiency gains as model size grows.
33

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
J.3. On Adam hyperparameters sweep
In our main results presented in Section 5, the hyperparameters of Adam (81, 82, etc) follows the setups of (Zhao et al., 2024a). Moreover, weight decay is not used. Below, we further fine-tune those hyperparameters and comapare with method. We conducted an additional parameter sweep for Ir. 31, 82, weight decay, and specifically for the 1B model scale, and obtained the following optimal values:
Learning rate: 0.0007
Betas: (0.9, 0.95)
c: le-8
Weight decay: 0.1 (we observed that weight decay of 0.1 outperforms no weight decay, though we did not perform an exhaustive search)
We further trained the IB model 20B tokens using Adam with these optimal settings and our method. Apart from hyperparameters, the general training setup still follows (Zhao et al., 2024a). The following table summarizes the test loss (lower is better) at various training steps:
Table 6: Comparison of the test loss obtained during training when training IB LLaMA with SWANF v.s. IB LLaMA under optimally tuned Adam.
Method
40K
80K
120K
150K
Adam
2.880
2.728
2.659
2.651
SWAN!
2.810
2.661
2.574
2.556
For our method, we still use the default parameter setting described in Appendix K. Our method still outperforms Adam by a large margin despite we spent much more compute to sweep the Adam hyperparameters.
K. Implementation details
General setup We describe the implementation setups for SWAN used in LLM pre-training tasks. To enable more straightforward and comparable analysis, we simply replicate the setting of (Zhao et al., 2024a), under exactly the same model configs and optimizrr hyperparameter configs, whenever possible. This includes the same model architecture, tokenizer, batch size, context length, learning rate scheduler, learning rates, subspace scaling, etc.
Precision All baselines uses BF16 for model weights, gradients, and optimizer states storage. For all SWAN variants, we use BF16 for model weights and gradients. For the GradWhitening step of SWAN-0 and SWANT we use FP32 to whiten the BF16 gradients and then convert it back to BF16. We found that this helps to improve trainining stability and performance. However, for the NSDS scheme of SWAN; we oberve that FP32 does not offer performance boost over BF16 (Appendix I.4), therefore stick with BF16.
Learning rate scheduling we use exactly the same scheduler as in (Zhao et al., 2024a), with the exception of SWAN-0, which does not require any learning rate warmup. Therefore, for SWAN-0, directly start with maximum learning rate, and enter the learning rate decay phase, using the same decay parameters (Zhao al., 2024a).
Reproducing baseline results Most baseline results are cited from respective papers (Zhao al., 2024a: Zhu al., 2024)
share the exact same setup. We also tried reproduce their results using the same opensourced code, and generally obtain slightly worse results for Galore, Apollo and Adam for larger models (350M and IB, see Table 7). Therefore in the main paper we compare with both the official results and our reproduced results. For the reproduced results of Adam, we specify the details below as was not disclosed in (Zhao et al., 2024a). We use same learning rate tuning procedure as suggested by (Zhao et al., 2024a) (i.e,, performing grid search over (0.01, 0.005, 0.001, 0.0005, 0.0001)). We found that the optimal learning rates for Adam is 0.001. The only exception is that for model of size 1.3B: as we already know that
34

---

SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training
Table 7: Reproduced results.
130M
350M
1.3 B
Adam (reproduced)
24.44 (0.75G)
19.24 (2.05G)
16.44 (7.48G)
Apollo-mini (reproduced)
23.97 (0.43G)
17.60 (0.93G)
14.37 (2.98G)
Galore (reproduced)
24.67 (0.57G)
19.74 (1.29G)
15.89 (4.43G)
r of low-rank methods Training Steps
256
256
512
20K
60K
100K
larger model requires smaller learning rates, conduct a learning search for Adam over smaller but more fine- fine-grained grid of {0.001, 0.0007, 0.0005, 0.0003, 0.0001 }. As a result, the optimal learning rate found for Adam on 1.3B is 0.0007. For other hyperparamters, we follow (Zhao et al., 2024a) where we use 31 = 0.9, B2 = 0.99, and no weight decay. Finally, one baseline that does not exist in the literature is the Momentum + GradWhitening (Muon-like optimizer without Nestrov acceleration) baseline; and we report our own results. We start from the default learning rates used by Muon (Jordan et al., 2024) and tuned them over a grid of 0.01, 0.02, 0.03, 0.04, 0.05.
SWAN Settings and hyperparameters Since SWAN utilizes matrix-level operations on gradients, it can only be applied to 2D parameters. Therefore, in our experiments, we only apply SWAN on all linear projection weights in transformer blocks. Similar Galore (Zhao et al., 2024a), the rest of the non-linear parameters still uses Adam as the default choice. Therefore, we follow the learning rate setup of Galore, where we fix some global learning rate across all model sizes and all modules. Then, for the linear projection modules where SWAN is applied, we simply apply a scaling factor on top of the global learning rate. For all SWAN variants, we adopt a lazy-tuning approach (hyperparameters are set without extensive search), as detailed below. This helps to reduce the possibility of unfair performance distortion due to excessive tuning.
SWAN-0 uses naive NS-iteration for whitening, disabled learning rate warmup, and use similar learning rates optimized for Adam. We fix the global learning rate to be the same as Adam, and fix 1. 'hhe only exception is the 1.3 B case. This is because we observe that the optimal learning rate of Adam under 1.3B becomes smaller than 0,001, hence we also reduce the learning rate on SWAN, where we used a = 0.3, resulting an effective learning rate of 0.0003. To summarize the hyperparameter of SWAN-0 is set to be similar to Adam, without any tuning. This is to demonstrate the robustness of SWAN series optimizers and their capability to work out-of-the-bor as replacement for Adam.
SWANT, is the vanilla version of our method, in which we enabled learning rate warmup, and allowed the use of optimized learning rates that largely differ from Adam. We notice that SWAN allows larger learning rates than Adam. We use global learning rate of 0.02, as well as the scaling factor a 0.05. This is selected by simply searching the learning rate over a constraint grid (0.01, 0.02, 0.05). and then setting a 0.05 such that the effective learning rate scaled back to 0.001. There is no guarantee that this heuristic rule is optimal; but we found that this usually does not make run fail (e.g., with loss divergence).
Finally, SWAN' the most efficient version of SWAN that employs the proposed NSDS scheme for fast whitening (section 3.2). Similar to SWAN', we use the same global learning rate of 0.02, as well as the scaling factor 0.05 across all model sizes. We suspect with more careful tuning, its performance can be significantly improved: however, this is out of the scope of the paper.
The configurations of GradWhitening is discussed next.
Implementation of GradWhitening For GradWhitening, before N-S iteration, we further normalize its input matrix by its Frobenius norm. This is applied in all our ablation studies as well, regardless of whether the gradient processed by GradNorm. For our proposed SSwton-Schulz with Diagonal Substitution (NSDS) used in SWAN!, we only run it for steps across all model sizes. We found that using a step size / * 0.5 in the GradWhi tening operator (Algorithm 2) can improve its convergence. We set B = 0.4. We found that NSDS is generally robust to 3 as long as it is not too large; and our specific choice of parameters usually already gives satisfactory performance in LLM pretraining. For the naive N-S iteration used in SWAN-O and SWANT, we run 10 steps which is usually sufficient. We set 3 = 0.8. The naive NS is run in FP32 precision, on top of 6F16 gradients and weights; while the NSDS is run in BF16
35

---

SWAN: SGD with Normalization and Whitening Enables Statcless LLM Training
Computational Overhead. Below, discuss the computational overhead of Naive Newton-Schulz. In practice, only the naive Newton-Schulz for 10 iterations, which corresponds Σ 50 matrix multiplications. These matrix multiplications are in general GPU friendly, hence for the task of training LLMs, the batch size is the more dominant factor for compute. For example QWn 14B (Bai et al., 2023) has 4M batch size model dimension Amodel = 5120, DeepSeek (Bi et al., 2024) 67B has 6M batch size vs dmodel = 8192, and LLaam (Dubey et al., 2024) 405B has 4-16M batch size vs dmwdee 16384. To estimate the computational overhead, assuming the N-S iteration involves approximately 50 x model FLOPs. contrast, the primary training cost scales with the batch size and is proportional batch_size x dooda FLOPs. In those examples, the estimated computational overhead of Newton- Schoolz is typically below 7%. A similar estimation has been given in (Jackson, 2023) (S 3%), well in (Jordan al., 2024) (s 1%).
However, note that whether the above analysis hold for large scale, distributed LLM training is still unclear. The N-S iteration is after all a (m))) complexity operation and might need to scale as model size increases. As argued by (Essential AI, 2025), for MoE models the ratio of NS flops vs model flops can be much worse than dense models. This exactly the motivations that drives us to explore alternative schemes.
36


Reviews:
[{"review_id": "uH30gfMY3N", "reviewer": "Reviewer_G4T7", "summary": "The paper proposed a \"stateless\" optimizer using gradient-normalization and gradient-whitening. The proposed method saves half memory over Adam and reaches 2x speedup. The idea is interesting and the writing is clear.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**We thank the reviewer for their careful reading and for the positive comments regarding the clarity of our writing and the significance of our contribution. We address each question below:**\n\n**Q1:**  \nRegarding the sufficiency of 12B tokens for 1B models, we have conducted additional experiments training a 1B model for 20B tokens. For the Adam baseline, we performed an extensive parameter sweep and found that the optimal parameters were **lr = 0.0007** and **betas = (0.9, 0.95)**. Using these settings, we compared the performance of Adam and SWAN (using the SWAN$^\\ddagger$ setting with the same learning rate as Adam). The table below summarizes the training loss at various steps:\n\n| Training Steps | Adam Loss | SWAN Loss |\n|----------------|-----------|-----------|\n| 20K            | 3.054     | 2.989     |\n| 40K            | 2.880     | 2.810     |\n| 60K            | 2.792     | 2.726     |\n| 80K            | 2.728     | 2.661     |\n| 100K           | 2.681     | 2.609     |\n| 120K           | 2.659     | 2.574     |\n| 150K           | 2.651     | 2.556     |\n\nAt the end of training, SWAN achieved a roughly **1.8× speedup** compared to Adam while reaching lower validation loss values. These detailed comparisons confirm that the advantage of SWAN persists even when training with more tokens and under a longer training schedule. (In fact, in this run, the extra steps Adam must take to reach SWAN's performance keep growing.)\n\n**Q2:**  \nIn Figure 1 (c), the term “SGD” refers to SGD without momentum. We acknowledge that including comparisons with other memory-efficient methods (e.g., Adam-mini, Muon) in the bar plot would provide a clearer picture. In our revision, we will update the figure to include these baselines. Notably, our method achieves a near memory-optimal footprint—comparable to vanilla SGD (i.e., without momentum)—which is one of the key strengths of SWAN.\n\n**Q3:**  \nIndeed, in some settings, this additional momentum before SWAN operations could provide further acceleration. However, there's a trade-off given hardware constraints. For example, consider the task of training a 350M model with 1024 context length under mixed precision on a node of 8XA100 40G GPUs. Including momentum forces the users to use larger gradient accumulation. In contrast, in this scenario, the stateless design of SWAN without momentum allows training with 2X larger batch sizes compared with using momentum. We will clarify these trade-offs in our revision.\n\n\n**Q4:**  \nWe appreciate the reviewer’s reference to Zhang et al. (2024) and acknowledge its relevance in motivating the need for stateless methods. While we have already cited this paper in our submission, we will expand our discussion in the revision to provide a deeper analysis of how the Hessian perspective discussed in [1] further motivates the design of new stateless optimizers like SWAN.\n\n**Once again, we thank the reviewer for their constructive feedback and insightful questions, which will help us improve the manuscript in the final version.**", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "yes", "experimental_designs_or_analyses": "yes", "relation_to_broader_scientific_literature": "no", "essential_references_not_discussed": "see below", "other_strengths_and_weaknesses": "see below", "other_comments_or_suggestions": "see blow", "questions_for_authors": "Q1: In Figure 1,  12B tokens are not sufficient for 1B models. Does the advantage of SWAN maintain if we train more tokens?\n\nQ2: In Figure 1 (c), does SGD refers to SGD with momentum or without momentum? Further, to make a clearer comparison with existing works, please also include other memory-efficient methods (e.g., Adam-mini, Muon) in the bar plot of Figure 1 (c).\n\nQ3: In Algorithm 1 SWAN, what if keep track of a 1st-order momentum M and apply the  GradNorm() and GradWhitening() to M, instead of G? Does it bring extra acceleration?\n\nQ4: [1]  studied why SGD performs poorly on Transformers.  Please discuss [1] as a motivation to re-design new stateless methods.\n\n[1] Zhang, Y., Chen, C., Ding, T., Li, Z., Sun, R., & Luo, Z. (2024). Why transformers need adam: A hessian perspective. Advances in Neural Information Processing Systems, 37, 131786-131823.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "msVbEsl30I", "reviewer": "Reviewer_A9LJ", "summary": "The paper proposes SWAN, an optimizer which is completely stateless. They claim that SWAN outperforms existing optimizers while also using lesser memory (since it is stateless). They support these claim by doing LLM pretraining experiments. \n\nSWAN is similar to another previously proposed optimizer Muon the following changes: 1. They remove momentum, 2. They add column normalization and 3. In one of the versions of SWAN they propose a simplified version of Newton-Schultz iterations which are needed for matrix whitening.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**We thank the reviewer for their detailed and constructive feedback. We address the main points raised below:**\n\n1. **\"The claim that SWAN outperforms Muon is problematic\"**  \n   - Our core contribution is to push the boundaries and demonstrate that: **it is possible to train LLMs matching the performance of Adam using a completely stateless optimizer**.  Our baseline “Momentum-GradWhitening” is included mainly for completeness, and whether SWAN can consistently outperform other existing non-Adam optimizers (such as shampoo/Muon/SOAP) is orthogonal to our contribution. \n\n    - While the reviewer notes that it is questionable whether SWAN outperforms Muon, it is important to emphasize that Muon is not a stateless optimizer and is primarily designed for wall-clock time acceleration. Our work thus highlights a viable extreme stateless alternative for memory-constrained settings, and this distinction is central to our contribution. We will clarify this distinction in the revision. \n\n    - Finally, we would like to note that compared to Muon, the most relevant baseline should in fact be the concurrent work of Apollo [2] (which is a low-rank/rank-1 optimizer) that also claims to match the performance of Adam. In our experiments, we have demonstrated that SWAN consistently outperforms Apollo.\n\n2. **Experiments with Larger Batch Sizes:**  \n   We have conducted additional experiments training a 130M model with a 1M batch size over 2B tokens. For the Adam baseline, we performed a learning rate sweep and found the optimal parameters to be **lr = 0.00075** and **betas = (0.9, 0.95)**. The results are summarized in the table below:\n\n   | Training Steps | Adam Val Loss | SWAN (SWAN$^\\dagger$) Val Loss |\n   |----------------|---------------|-------------------------------|\n   | 500            | 4.1796        | 4.0755                        |\n   | 1.5K           | 3.483         | 3.477                         |\n   | 2.5K           | 3.398         | 3.370                         |\n\n   These results demonstrate that SWAN remains on par with or slightly better than Adam in terms of validation loss, even when training with a significantly larger batch size. This again confirms our core finding of \"it is possible to train LLMs matching the performance of Adam using a completely stateless optimizer\"\n\n3. **\"For Muon optimizers, was Adam used for first and last layer?\"**  \n   Yes, following standard practice (which predates both our work and Muon), Adam is used for the first and last layers in Muon. SWAN also employs Adam for these layers, aligning with practices used in other baselines such as Galore and Apollo.\n\n4. **\"Could the authors share their codebase?\"**  \n   We are fully committed to open-sourcing our codebase. We are actively working on a robust open-source release that will include detailed configurations as well as more new optimizers that were not included in our submissions\n\n5. **Regarding hyperparameter settings of baselines:**  \n   As stated in our paper, we strictly follow the experimental setup of Zhao et al. (2024a) [1]. Our experiment code is based on their implementation, and the configurations—including weight decay and other hyperparameters for all baselines—are directly taken using the config files that can be found in their open-source repository. We will further clarify these details in the revised version of the paper.\n\n**Once again, we thank the reviewer for their insightful comments and suggestions, which will help us further improve the manuscript.**\n\n[1] Zhao, Jiawei, et al. \"Galore: Memory-efficient llm training by gradient low-rank projection.\" arXiv preprint arXiv:2403.03507 (2024).\n\n[2] Zhu, Hanqing, et al. \"Apollo: Sgd-like memory, adamw-level performance.\" arXiv preprint arXiv:2412.05270 (2024).", "claims_and_evidence": "I think the claim that SWAN outperforms Muon is problematic since their comparison is done at a small batch size (130K) while it is well know from prior works that the benefit of momentum emerges at larger batch sizes.", "methods_and_evaluation_criteria": "Yes, except for the small batch size.", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "NA", "relation_to_broader_scientific_literature": "This paper adds to many new recent optimizers trying to improve on Adam in both speed and memory requirements.", "essential_references_not_discussed": "Most references have been discussed.", "other_strengths_and_weaknesses": "The authors do not describe how they set hyperparameters such as weight decay.", "other_comments_or_suggestions": "I think the paper would benefit from focusing on one main benefit and showing evidence for it so that no doubt remains. For example, the aforementioned use of small batch sizes. I would also recommend that the authors release the codebase so that it is easy to see details such as weight decay (which is not specified in the paper).", "questions_for_authors": "1. For Muon optimizers, was Adam used for first and last layer? (as in recommended in the referenced blogpost on muon).\n2. Could the authors do experiments with 1 million batch size (with tune LRs). Any >100M sized model should suffice.\n3. Could the authors share their codebase?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TYWMlWWkjg", "reviewer": "Reviewer_urpW", "summary": "This paper introduces a new stateless optimizer SWAN (SGD with Whitening And Normalization) with the same performance as the Adam optimizer for LLM training. The author analyses that SGD with GradNorm and GradWhitening applied in tandem can minimize the condition number, stabilize gradient distributions across transformer block and coverage more robust to the local curvature condition. This paper evaluates the SWAN on LLM pre-trained tasks, where SWAN outperformed other optimizers, and analyses the effect and efficiency of SWAN.\n\n\n\n## update after rebuttal:\n\nI have read the response by the authors and my concerns are mostly addressed. I keep my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**We thank the reviewer for their insightful feedback and for acknowledging many of the strengths of our work. We address the specific points raised below:**\n\n**1, Iteration Count and Motivation for Whitening:**  \nRegarding the number of iterations in the Newton–Schulz procedure, our ablation experiments in Appendix B show that the performance does not peak at 2 iterations; rather, increasing the iteration count (e.g., to 5) continues to improve the approximation accuracy of the whitening operator. We chose 2 iterations in our main experiments as a trade-off between computational cost and performance, while the improvement with additional iterations is clearly demonstrated in our supplementary results. \n\nFor motivation for whitening, please refer to our response below.\n\n**2, Evidence for the claim that Adam needs historical information because it ignores non-diagonal interaction:**  \nIn Section 5.3 of our paper, we provide theoretical motivation for why Adam requires historical information—and why our method does not. Below, we rephrase the analysis from the natural gradient descent/Fisher information perspective. First, recall that the Fisher information matrix (FIM) is defined as $F = \\mathbb{E}[gg^T]$ where $G$ is the gradient matrix and $g = \\operatorname{vec}(G)$ is the flattened gradient (hence $F$ is $mn$ by $mn$.). Adam approximates the FIM using two key approximations:\n- It uses exponential moving averages (EMAs) over time to estimate the expectation $\\mathbb{E}[\\cdot]$.\n- It further adopts a diagonal approximation to $F$, thereby ignoring interactions between different parameters.\n\nIn contrast, our gradient whitening operation is based on a block (identical) diagonal structural assumption:\n$$\n\\tilde{F} = I \\otimes M,\n$$\nwhere $\\otimes$ is the Kronecker product,  $I$ is the $n$ by $n$ identity matrix and $M$ is a $m$ by $m$ matrix represents the FIM for each identical block diagonals. We can find the optimal $M$ by solving the optimization problem\n$$ \\min_{\\tilde{F} = I \\otimes M}  ||\\tilde{F} - F||^2_{Frobenius} $$\nthe optimal solution for each block is given by the unbiased estimate\n$$ M \\approxeq \\frac{1}{n} \\sum_{j=1}^n g_j g_j^T = GG^T $$\nwith the summation taken over the columns of $G$ (each column corresponds to a \"diagonal block\" in the $F$). Then, we can use the optimal $\\tilde{F}$ to perform natural gradient descent where the update is given by:\n$$\\Delta  \\operatorname{vec}(W) = \\tilde{F}^{-1/2} g$$\nsimplifying to matrix form, this is equivalent to\n$$\\Delta  W \\propto  M^{-1/2} G$$, which is exactly (up to scaling) $\\text{GradWhitening}(G) = (GG^T)^{-1/2}G$.\n\nIn plain language, we have shown that: because $\\tilde{F}$ assumes each diagonal block is identical, one can estimate the full FIM using spatial information (i.e., average over blocks using $\\frac{1}{n} \\sum_{j=1}^n g_j g_j^T$ instead of over time) rather than relying on temporal averaging. This key insight allows the whitening operation to be derived as the solution to the FIM approximation problem under our block diagonal assumption. This analysis shows that our structural assumption permits bypassing the need for historical averaging.\n\n**3, Clarification on Terminology for the Whitening Operation:**  \nWe follow standard terminology from the literature (e.g., in Decorrelated Batch Normalization by Huang et al., 2018) where the operation$(GG^T)^{-1/2}G$\nis commonly referred to as “whitening”. Another reason is that performing natural gradient descent using the inverse square root of FIM is also referred to as whitening in the literature. (Adam can also be seen as a diagonal approximation of the inverse square root of FIM). We will add further clarification in the revised manuscript to ensure this point is unambiguous.\n\n**4, Generalization Beyond the LLMs:**  \nAlthough our experiments focus on LLM pretraining with the C4 dataset, we expect that with minimal changes, similar benefits could be achieved in vision or multimodal tasks. We plan to explore these extensions in future work.\n\n**5, Additional Revisions:**  \nWe appreciate the reviewer’s detailed suggestions regarding the clarity of the appendix (including the presentation of assumptions from Theorem 1 of Tian et al., 2023), the labeling in Figure 2, and consistency in punctuation after equations. We will carefully revise these sections to improve readability and ensure consistency throughout the manuscript.\n\n**Once again, we thank the reviewer for their constructive comments, which will help us improve the quality and clarity of our paper in the final revision.**", "claims_and_evidence": "The claims are overall supported by theoretical or experimental evidence. The evidences are clear and convincing.\nHowever, the motivation of applying normalization and whitening onto gradient is not very clear. Especially the best performance is usually obtained with around 2 iterations when using Newton’s iteration to obtain an approximated whitened/orthogonalized representation. What is the results when using more iterations? E.g. 5 iterations.", "methods_and_evaluation_criteria": "The methods and evaluation criteria are reasonable.", "theoretical_claims": "I do not find remarkable errors in theoretical claims. The claims are supported by analysis and previous studies. The theoretical analysis in this paper is somewhat solid, providing both practical considerations and dynamic analysis. One main concern is that the assumption 1 is too strong, and I donot think assumption 1 hold in practice.", "experimental_designs_or_analyses": "The experiments are overall comprehensive. The results show that SWAN outperformed other optimizers. The authors conduct ablation experiments to figure out the effect of GradNorm and GradWhitening and why and how they help optimization. The author also analyses memory efficiency and throughput of SWAN.", "relation_to_broader_scientific_literature": "The authors provide a new optimizer combining with normalization and whitening technology, which is effective and efficient in memory according to the experiments.", "essential_references_not_discussed": "I think this paper should give credit to the normalized gradient method in training DNNs, e.g., the paper [1]\n\n[1]Block-normalized Gradient Method: an Empirical Study for Training Deep Nerual Network. preprint arXiv:1707.04822", "other_strengths_and_weaknesses": "The paper is well organized.\nHowever, the colored column in Figure 2 is confusing. It would be better if more detailed explanation is marked in this figure.", "other_comments_or_suggestions": "Typo problem: The punctuation after the equations are not unified.", "questions_for_authors": "1.\tIn Section2, the author hypothesize the additional history information in Adam is because the approach does not take into account the interactions and structures between different variables. Could the authors provide more evidence?  \n2.\tAccording to my understanding, the gradient matrix G refers to the stochastic gradient of a particular weight matrix. The elements among columns of G are all variables. Therefore, how to understand whitening the gradient data in GradWhitening? May be the term “orthogonal the gradient vector” better?  \n3.\tThe experiments only involve C4 dataset. Could the advantage of SWAN preserve when it transfers to other tasks, for example vision or Multimodal tasks? (Even though this paper addresses for LLM training, it seems the method is general without considering what type data is)", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "NA"}}, {"review_id": "FikGjlQE96", "reviewer": "Reviewer_KRz9", "summary": "This paper combines two classic technique normalization and whitening to improve SGD and achieves better performance than Adam and other optimizers while saving memory cost by not saving optimizer state. Theoretical insights are also provided to explain the effect of each technique.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**We thank the reviewer for their thorough and constructive feedback. We are grateful that the reviewer acknowledged the novelty, writing, empirical performance, and the overall significance of our work. Below, we address each point raised:**\n\n\n---\n\n**1. On Combining GradNorm and GradWhitening:**  \n- This is an excellent question.  Indeed, we found there is a theoretical explanation behind this, which in fact gave rise a class of more general algorithms than SWAN.  However, this insight is non-trivial and, due to space limitations, we plan to present a more comprehensive analysis in follow-up work. Below we briefly describe the high level idea:\n\n- To start with, given a gradient matrix $G$ with shape $m$ by $n$, GradNorm and GradWhitening can be interpreted as projection operators under specific norm constraints. The extended version of SWAN can be thus viewed as an iterative process:\n$$\nG \\rightarrow \\text{GradNorm}(G) \\rightarrow \\text{GradWhitening}(\\text{GradNorm}(G)) \\rightarrow \\text{GradNorm}(\\text{GradWhitening}(\\text{GradNorm}(G))) \\rightarrow \\ldots\n$$\nuntil a fixed point is reached. In other words, it corresponds to performing steepest descent under multiple (non-Euclidean) norm constraints, as opposed to the standard SGD update which is steepest descent under a single Euclidean norm constraint.\n-  In theory, one could choose an arbitrary collection of norms—provided they satisfy certain theoretical properties—to guide the update. When \\(G\\) is nearly a square matrix, even a single iteration (as implemented in SWAN) is nearly sufficient. \n- In revision, we will present the main results as extended discussion.\n\n---\n\n**2. Order of GradNorm and GradWhitening (Reviewer Question Q1):**  \nFrom a fixed-point perspective above, the final result is invariant to the order as long as the iterative process converges to a multi-norm fixed point. Hence, the key is not the order but achieving convergence under the imposed norm constraints.\n\n---\n\n**3. Normalizing Along Different Axes (Reviewer Question Q2):**  \nWe indeed experimented with normalizing along alternative axes. We found that normalizing along the row-wise direction (as in our current formulation) indeed yields superior performance. This is consistent with Theorem 1, which suggests that the primary source of gradient noise exhibits a row-wise scaling structure. We will include additional discussion and experimental evidence of this in the revised manuscript.\n\n---\n\n**4. Regarding the Proof of Proposition 1:**  \nThe reviewer asks whether the subset $O_l$ can have multiple elements. In fact, at equilibrium, $O_l$ contains only a single index. We acknowledge that during early training—before convergence—multiple dominating indices may appear (as observed in numerical experiment in Figure 8), and similar structures are still present across the normalized diagonal blocks. We will clarify this point in the revision to make it explicit that our theoretical results assume convergence, at which point $O_l$ becomes a singleton.\n\n---\n\n**5. Regarding the Proof of Proposition 2:**  \nOur argument in Proposition 2 is not that the bound cannot be arbitrarily small, but rather that a small condition number does not necessarily lead to a diminished bound. This distinguishes our method from standard SGD, where such a decrease would be expected. We will revise our presentation to better articulate that our bound remains meaningful even when the gradient matrix is well-conditioned.\n\n---\n\n**6. Typos and Presentation of Assumptions:**  \nWe thank the reviewer for pointing out typos and suggesting improvements in the presentation of our assumptions (e.g., in Theorem 1 and throughout the appendix). We will carefully revise the manuscript to correct these issues and to list all relevant assumptions explicitly.\n\n---\n\n**Once again, we appreciate the reviewer's insightful comments, which help us improve both the theoretical exposition and empirical presentation of our work.**", "claims_and_evidence": "The superior performance of SWAN is supported by the results in table 1. The authors also conduct ablation study to show that both modifications are necessary. I am not convinced by the theoretical analysis in section 4. See the theoretical part and weakness part for details.", "methods_and_evaluation_criteria": "The proposed method combines two classic technique that are widely used by previous work. It is interesting that both techniques together provide a much better performance than each of them alone. The selected benchmark and models of different sizes are reasonable.", "theoretical_claims": "I checked the proof of theorem 2, proposition 2 and proposition 1. I feel the results are either trivial or their proofs lack detail for me to confirm correctness. \n\n1. In the proof of proposition 2, can you explain the existence of C_G and C_L? If we just consider $H=diag(h_i)$ and $W_{whitened}^{(t)}=diag(w_i)$, then $Q=\\frac{(\\sum h_i w_i)^2}{(\\sum h_i w_i^2)(\\sum h_i)}$ can be arbitrarily small even with fixed $h_i$.  \n2. In the proof of proposition 1, can the subset $O_l$ have multiple elements? Is there any conclusion for the relationship between $H(V)\\_{lk,lk’}$ and $H(V)_{l’k,l’k’}$ for both $l,l’ \\in O_l$? Moreover, the convergence result of $H(V)$ after normalization can’t hold when the size of $O_l$ is larger than 1.", "experimental_designs_or_analyses": "The experiment design is clear, with three different version of SWAN to make a fair comparison with previous methods. I have a question for figure 6. When you say you compute mean of GradNorm gradient, do you compute the mean of GradNorm(G) or the original G but obtained from checkpoints trained with GradNorm update rule? If it is the mean of GradNorm(G), then the results sound less interesting because the distribution indeed should change less when you have a hard constraint on the scale of random variables.", "relation_to_broader_scientific_literature": "This paper follows previous work on designing a more efficient optimizer for training LLMs.", "essential_references_not_discussed": "The paper has included recent progress on designing new optimizers.", "other_strengths_and_weaknesses": "Strength\n1. The paper proposed an efficient algorithm for GradWhitening. It is interesting that it converges very fast. \n\nWeakness\n1. Even though the authors try to give justification for the methods from the theoretical prospective, they only prove some results for each technique separately. It is unclear why combining them together can improve the results so much. Also each theoretical result will break when you add another operation. So I feel all the theoretical results are not so meaningful.", "other_comments_or_suggestions": "1. There is a typo in the definition of GradNorm (right part of line 146). You seem to miss a square term when defining $s$. \n2. Another typo in theorem 1 when you define the standardized stochastic gradient. \n3. It will be good to list all the assumptions rather than just saying inheriting from other paper.", "questions_for_authors": "1. Have you considered switch the order of gradnorm and gradwhitening? \n2. Have you tried normalizing along another axis when doing gradnorm? I am curious whether we can get a similar analysis (theorem 1) for normalizing along the other axis. If not, why is one better than the other?\n3. Do you have any insight on why combining the two techniques can improve performance over each of them own?\n4. Question 1 in theoretical part. \n5. Question 2 in theoretical part.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This introduces SWAN, a completely stateless optimizer for LLM pre‐training that combines a row‐wise GradNorm step with a matrix‐level GradWhitening step. By preprocessing raw SGD gradients in this way, SWAN matches or surpasses Adam’s token‐efficiency while cutting end‑to‑end memory by ≈ 50 % and eliminating all optimizer states. SWAN represents a significant and practical step toward memory‐efficient LLM training, with strong empirical results and solid theoretical grounding. All the reviewers recommend (weak) accept. The authors should incorporate the discussions and additional experimental results presented during rebuttal into the next revision.
Decision: Accept (poster)