Gradient Multi-Normalization for Efficient LLM Training
Authors: Meyer Scetbon
Chao Ma
Wenbo Gong
Edward Meeds
Venue: NeurIPS 2025 poster ()
Keywords: LLM
Optimizer
Gradient
Normalization
Primary area: deep_learning
TL;DR: We introduce a new design principle for LLM matrix optimizers - gradient multi-normalization, unifying previous work, and enabling faster and memory-efficient training of LLMs.

Abstract: Training large language models (LLMs) commonly relies on adaptive optimizers such as Adam (Kingma & Ba 2015), which accelerate convergence through moment estimates but incur substantial memory overhead. Recent stateless approaches such as SWAN (Ma et al., 2024) have shown that appropriate preprocessing of instantaneous gradient matrices can match the performance of adaptive methods without storing optimizer states. Building on this insight, we introduce \emph{gradient multi-normalization}, a principled framework for designing stateless optimizers that normalize gradients with respect to multiple norms simultaneously. Whereas standard first-order methods can be viewed as gradient normalization under a single norm (Bernstein & Newhouse, 2024), our formulation generalizes this perspective to a multi-norm setting. We derive an efficient alternating scheme that enforces these normalization constraints and show that our procedure can produce, up to an arbitrary precision, a fixed-point of the problem. This unifies and extends prior stateless optimizers, showing that SWAN arises as a specific instance with particular norm choices. Leveraging this principle, we develop SinkGD, a lightweight matrix optimizer that retains the memory footprint of SGD while substantially reducing computation relative to whitening-based methods. On the memory-efficient LLaMA training benchmark (Zhao et al., 2024), SinkGD achieves state-of-the-art performance, reaching the same evaluation perplexity as Adam using only 40\% of the training tokens.

Gradient Multi-Normalization for Efficient
LLM Training
Meyer Scetbon" Microsoft Research
Chao Ma" Microsoft Research
Wenbo Gong". Microsoft Research
Ted Meeds Microsoft Research
Abstract
Training large language models (LLMs) commonly relies on adaptive optimizers such as Adam (Kingma & Ba 2015). which accelerate convergence through mo- ment estimates but incur substantial memory overhead. Recent stateless approaches such SWAN (Ma et al. 2024) have shown that appropriate preprocessing of instantaneous gradient matrices can match the performance of adaptive methods without storing optimizer states. Building this insight, introduce gradient mulli-normalization, a principled framework for designing stateless optimizers that normalize gradients with respect multiple norms simultaneously. Whereas standard first-order methods can be viewed as gradient normalization under a single norm [Brmsttin & Newhouse) 2024), our formulation generalizes this perspective to multi-norm setting. We derive an efficient alternating scheme that enforces these normalization constraints and show that our procedure can produce, up an arbitrary precision, a fixed-point of the problem. This unifies and extends prior stateless optimizers, showing that SWAN arises as a specific instance with partic- ular norm choices. Leveraging this principle, we develop SinkGD, lightweight matrix optimizer that retains the memory footprint of SGD (w/o momentum) while substantially reducing computation relative to whitening-based methods. On the memory-efficien LLaMA training benchmark (Zhao et al., (004a)), SinkGD achieves state-of-the-ard performance, reaching the same evaluation perplexity Adam using only 40% of the training tokens.
Introduction
The training of Large Language Models (LLMs) relies heavily on adaptive optimization algorithms, such as Adam Kingma & Ba (2015), which dynamically adjust learning rates for each parameter based on past gradient information, leading faster convergence and improved stability. However, these optimizers introduce substantial memory overhead due to the storage of internal states, typically moment estimates of gradients, a challenge that becomes particularly pronounced in distributed training settings where memory constraints and communication overhead are critical concerns Rajb handari et al. (2020): Korthikanti al.) (2023); Dubey et (024)). contrast, simpler first-order optimization methods such as Stochastic Gradient Descent (SGD) require significantly less memory but fail to adequately train LLMss Zhho et al. (0024b) Zhang et al. 2020);; Kunstner et al. (2023) 2024). As result, there is ongoing need for developing new optimization strategies that resolves the memory efficiency v.s. training performance dilemma for large-scale models training.
Recent research has made significant strides in improving memory efficiency of optimization. A key focus on low rank compression of optimizer states (Hu 2021 Lialin et 2023) Zhao
et al. 2024a Hao et al. 2024 Xu et al.) 2024a Zhang et al. 2024 Gong et al. 2025 Chen et al.
2024a| Zhu et al.) 2024). More recently, a new approach that directly removes certain optimizer
states started to emerge, and has shown stronger results compared with direct state compression (Ma
"quual contribution. This work was done when Meyer Sectbon was affiliated with Microsoft Research.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

350M Model
1.3B Model
Fifective Through put Comparison
20-
Adam
Adam StnkGD
SinkGD
1406
StaGG odtedtivet
Apolla- mini
Apalla-mini
120K
SWAN
PPL
SWAN
TAM229
100
2. speal
2 spooduu
15
<0k
13X spx
gnk
206
20CC
63K
20K
GAK
70X
9ox
Stops
Steps
(a) 350M LLaMA model
(b) 1.3B LLaMA model
(c) Training Throughputs
Figure 1: SinkCD performance preview on memory-efficient LLaMA pretraining benchmark
Zhao et al. (2024a). (a) and (b): Comparison of the test perplexities obtained by Adam Kingma & Ba (2015); Zhao et al. 2024a), SWAN (Ma et al, 2024). , Apollo (Zhu et al., 2024), and our proposed IB LLaMA pretraining task with C4 dataset. All loss curves of Adam and Apollo-mini are reproduced from the corresponding opensource codes. We also compare with their official results in Table On both 350M and 1 1.3B LLama architectures, SinkGD achieves > 2x speed-up vs Adam in terms of tokens seen; and 1.3 to 1.5 X speed-up vs SWAN and Apollo. (c): Training throughput analysis on training 1.3 B model on 8 x A A100, under constant batch size 130K tokens. We present two metrics: raw throughput, measured by number of training tokens consumed per second; and effective throughput, which is throughput adjusted by the token efficiency of optimizer relative to Adam. SinkGD has a raw throughput that is marginally higher than Adam, while improving the effective throughput by 3x.
SinkGD (Algorithm
et al., 2024; Jordan et al. 2024 Chen et al. 2024c). Among these advancements, Ma et al. (2024) propose to replace all moving average states with sequence of stateless matrix operations acting on the raw gradient tensors. This results in SWAN, an optimizer that achieved the same memory footprint as SGD (w/o momentum) while delivering comparable or even better performances than Adam. Collectively, they demonstrate that memory efficiency and loss throughput are not mutually exclusive. The goal of this paper is to understand the design principle of those empirical methods, and further improve their efficiency.
Contributions. Motivated by the insights of SWAN (Ma et al. 2024). we introduce gradient multi- normalization. a framework for designing efficient optimizers based on novel mullii-nrmaaliztioon scheme. Unlike standard first-order methods that can be interpreted as gradient normalization according to a single norm (Bernstein & Newhouse) 2024). our approach aims at normalizing gradients according to multiple norms. We demonstrate that gradient matrix processing-bai based techniques proposed in (Ma et al.] (024) can be understood as a particular instance of our approach with carefully chosen norms and shorter iterative steps. Given this new perspective, we instantiate a more efficient and scalable stateless optimizer (SinkGD) that achieves Adam-level computational cost, while having the same memory footprint as SGD. Moreover, it achieves on par or even outperforms Adam in LLM pretraining tasks, as well as various existing memory-efficient baselines under LLaMA architecture. Our contributions are summarized below:
We propose a novel family of first-order methods, called Multi-Normalized Gradient Descent (MNGD), that aims at normalizing gradients according to multiple norms. Our framework general- izes the steepest descent viewpoint of Bernstein & Newhouse (024) that recasts popular first-order optimizers as normalization of gradients under a single norm.
We then propose a simple alternating scheme in order to effectively compute the multi- normalization of gradients. and show that our algorithm can provide a fixed-point solution up to an arbitrary precision, ensuring the normalization of its output with respect to the norms considered.
We leverage our framework and design scalability of SWAN. Our algorithm, namely SinkGD (Algorithm) 4. alternatively performs row- wise and column-wise normalization according to the Fuclidean geometry. We show that SinkGD exactly recovers the square-root iterates of Sinkhorn algorithm (Sinkhorn 1964).
lightweight, stateless optimizer that improves the
Finally, we evaluate our Sinkhorn-baan based stateless optimizer SinkCD by training LlaMA models various scales, from 60m to 1.3B. Results show that SinkGD manages to be on par or

---

outperforms the Adam optimizer, as well as other memory-efficien baselines. On the memory- efficient LLaMA training benchmark Zhao al. (2024a), SinkGD achieves a 3x speedup over Adam with significantly reduced memory requirements.
2 Background
2.1 From Adam to Stateless Optimizers
Adam Optimizer. Adam (Kingma & Ba 2015) relies on accumulating internal states throughout training in order to improve the convergence. More formally, given a loss function (0, O x X + C(0, .z) E I. where θ Ie is the set of learnable parameters and X' is the set where the data resides, Adam aims at minimizing + Eg=z, = (C(0, x)) where Pa the distribution data on X. To achieve this, Adam computes every step > 1 stochastic gradient associated with mini-batch of input data (*)), and performs the following updates (denoting V. V VoL(O,, xx.)))) x0))):
(1 V
St = B25--1 + (1 V2)
m
S₁
m.
M
8t
B2t1 0, T) vs. +
3i
33
where @ is the Hadamard product, 7) are global step-sizes, and 31, B2 0 are the weights of the exponential moving averages (EMAs) for the first and second moments respectively. During training, Adam optimizer stores two additional states (m,, effectively tripling the memory required train the model compared to a simple stochastic gradient descent (SGD) scheme.
SWAN: Stateless Optimizer. Ma et al. 2024) propose to move away from the paradigm of keeping track of internal states, and propose SWAN, a stateless optimizer that only pre-processes the stochastic gradients before updating the parameters. More precisely, they propose to update the learnable weight matrices involved the model using two matrix operators. Given a weight matrix W E R gmmx n with m & n, at time 2 1, the SWAN update (denoting V. Vw LWW, (x) ))):
V. V.V.
V. Vn(V.V, 1 W++ W. V
(1)
where for matrix W E R Q(W) Diag(1 Wi.:lla.... || m.: (l21) is the diagonal matrix of size where the diagonal coefficients are the E--norm the rows of W. To compute (V ( ) the authors leverage the Newton-Schule algorithm (Song et al., 2022 Li et al., 2018 Huang et al. 2019) instead of computing the SVD. While this approach does not require storing any additional states, it still suffers from a computational burden due to the Omm2(n + m)) computation of (V.27)-1//V, which may limit its practical usage for certain scenarios.
2.2 Steepest Descent as Gradient Normalization
Bernstein & Newhouse 2024) interpret several gradient descent schemes steepest descent methods under specific norms. More formally, they propose to minimize local quadratic model of the loss ((()x(()) 0g w.r.t to a given norm || Il, that
2 == C(O,,
(V., 2) ||
where A > O are the sharpness parameters and V. = voL(0,,x(1)) is the current stochastic gradient. As shown in (Bernstein & Newhouse 2024 finding a minimizer of 211-11 can be equivalently formulated solving:
V X
arg max (V,,2)
(2)
ZER2 ||=|==
where ||~|||- := supzER4; | (2, 2) is the dual norm of ||z|| Several popular gradient-descent schemes can be recovered using the above approach. For example, when the εn-norm is used, one recovers standard gradient descent, while the fxx. leads to signed gradient descent (Carlson et al.) 2015). However, this framework considers only a single norm for pre-processing the raw gradient V,. In the following, we extend this approach to incorporate multiple norms for gradient pre-processing. enabling the design of efficient and stateless optimizers for LLM training.
3

---

Multi-Normalized Gradient Descent
Notations. For a vector IRd, we call its normalized projection w.r.t a given norm || solution to the following optimization problem:
-II, the
PIH (a) arg max 2)
(3)
=|||==
is matrix norm. We also extend the definition of this notation if R is matrix and ||
3.1 Gradient Multi-Normalization
Let us now consider a finite family of K 2 1 norms ........9 ( In order pre-process the gradient V jointly according to these norms, we propose to consider the following optimization problem:
arg max V 2) s.t. V [[1, KI], g=(=) =
Assuming the constraint set is non-empty, the existence of maximum is guaranteed. However, this problem is NP-hard and non-convex due to the constraints, making hard to solve efficiently for the general case arbitrary norms.
Algorithm MultiNorm V, L,)
Algorithm Multi-Normmalizd GD (MIGD)
the stochastic gradient VoL(0,, x()), the norms g == (g1..... 9k), and 2 the number of iterations.
Input:
the number of updates, (Me)sseTT the global step-sizes, L the loss to minimize. L 2 1 the number of itera- tions for the nultii-ormaalizatinn, and == (01,... ........ (..k...9k]) the norms.
Input: T
Initialize æ = VgL(e,, ( for L 1 to L do
for = 1 to do 2 Pg. (2)
for = 1 to T do
V. VoL(O,, (0) with 2(() Ps
arg max (2, 2) (()===
V. MultiNorm Ves L, g) in Alg.
end for end for Return x
0++1 + 0, 1/e V end for Return x
Remark 3.1. Observe that when K the problem 4 recovers exactly the single normalization step used in Bernstein & Newhouse (2024).
Remark 3.2. The convex relaxation of defined as
arg max V, z) s.t. V e [[I, g<(=) <
in fact equivalent to the single normalization case discussed Section 2.2
where the norm
considered is ||a||| : max (z))). Thus, solving (5) is equivalent computing the projection
sE||L,K
(V). (). In Appendix C we provide general approach to compute it using the so-called Chambolle-
Pock algorithm Chambolle & Pock (201T).
While solving ( exactly might not practically feasible general, we propose a simple alternating projection scheme, presented in Algorithm I Notably, our method assumes that the projections P can be efficiently computed for all i [[t, K]. Fortunately, when the gi 's correspond to l, -ormms with p E ||1, +ool], or Schatten p-norms for matrices, closed-form solutions for these projections exist. See Appendix C for more details.
an Instance of MultiNorm. SWAN Ma et al. 2024 applies two specific prepproeessing steps to the raw gradients order to update the weight matrices. In fact, each of these pre-prov processing steps can be seen as normalized projections with respect to a specific norm. More precisely, for W E mm and m n, let us define
SWAN:
max
|Wa
91 (W): = NEl 1mml]
and 92(W) |
Vn

---

where for E [1, +o0, || following equalities:
ll.p is the Schatten p-norm of W. Simple derivations leads to the
Psi (W) V))(((W))-V)
(W) V
Therefore applying single iteration 1) of Algorithm 1 with norms 91 and d2 defined above on the raw gradient V, exactly leads to the SWAN update (Eq. )
3.2 On the Convergence of MultiNorm
We aim now providing some theoretical guarantees on the convergence of MultiNorm (Algo- rithm T]. More precisely, following the SWAN implementation Ma et al. (2024), we focus on the specific case where K = 2 and the normalized projections associated with the norms 91 and O2 have constant normo. More formally, we consider the following assumption.
Assumption 3.3. Let g be a norm on RR. We say that it satisfies the assumption if for all r e Rd. (1)l2=) where > 0 is an arbitrary positive constant independent of x and 1122 represents the Euclidean norm.
Remark 3.4. Observe that both norms in SWAN satisfies Assumption 3.3 and their normalized projections have the same fz-norm, as for any W E Rmxx with S n, we have ||Pg, (W)ll2 [P (W)ll2 Vnm.
This assumption enables to obtain useful properties on Pg as we show in the following Lemma: Lemma 3.5. Let g norm satisfying Assumption
Then
Pg Pg - Ps
for all x E R.A, (P)(()) = [P, = c², where g" is the dual norm associated with g-
Let now introduce some additional notation clearly state result. Let zo Rdd and let us define for n 2 0:
22n+1 Pan (22u), 22n+2 Pgz (I2m+11
(6)
which is exactly the sequence generated by Algorithm| 1 when K and (oL(01,x(t)) Let us now show our main theoretical result, presented in the following Theorem.
Theorem 3.6. Let gi and g2 two norms on Rd satisfying Assumption normalized projections have the same E2 norm. Let also (In ( be defined as in (6) and let us define the set of fixed-poin as:
3.3 and such that their
F: (x: Pa (P) x}
Then by denoting d(x, F) == min | - 2112 we have
sEF
d(In, F)
0
This Theorem states that if MultiNorm runs for sufficient amount of time, then the returned point can be arbitrarily close to fixed-point solution. While we cannot guarantee that solves (4). we can assert that our algorithm converges to a fixed-point solution with arbitrary precision, and by-product produces solution x normalized w.r. both norms 91, 92 (up to an arbitrary precision).
Remark 3.7. Note that in Theorem 3.6 we we assume that the normalized projections associated to g1 and 92 have the same norms. However, given two norms 91 and 92 satisfying Assumption 3.3 i.c. such that for all x::
PP g. (z)ll2 = c1 |P y2 (x)
for some C1, C2 O. and given target value >> O, normalized projections have the same &g norm equal to a. More formally, by denoting ÿi = 6s
always rescale the norms such that their
9r and
f2 g2, obtain that
(x) 2 = P (2) = a.
Remark 3.8. It is worth noting that, for squared matrices (m n), single iteration (L 1) of MultiNorm using the norms considered in Ma et al. (2024). immediately converges to fixed-point- precisely recovering SWAN.
5

---

3.3 MNGD: New Family of Stateless Optimizers.
We now introduce our family of optimizers: Multii-ooralizeed Gradient Descents (MNGDs) (Algo- rithm 2). The key distinction from the framework proposed in Bernstein & Newhouse (2024) is that MNGDs normalize the gradient with respect to multiple norms using the MultiNorm step. whereas in Bernstein & Newhouse (2024), the gradient is normalized using single norm. In the following, we focus on the MNGD scheme with specific choice of norms, for which we can efficiently compute the gradient nultinnormalization step. This enables the application of stateless optimizers to LLMs.
4 Sinkhorn: a Multi-Vormalization Procedure
As in SWAN Ma et al. (2024), we propose to normalize the weight matrices according multiple norms, We still leverage the row-wise fz-norm to pre- pre-process raw gradients, however, rather than using the spectral norm, we propose to consider instead a relaxed form of this constraint and use the column-wis wise nn-mm.. More formally, consider the two following norms on matrices of size mxx,
max +E[[L.m]]
[Walll2
max
[Walll
]E[LL.n]
9 (W)
92 (W)
Vn
Vm
which leads to the following two normalized projections:
Pg. (W) = ~ (s.
VmWR(W)-
where R(W)= == Diag( W.,1l2:..... ([.,lll)) E R is the diagonal matrix of size with the f--norm of the columns W as diagonal coefficients. For such choice of norms, the MultiNorm reduces to a simple procedure as presented in Algorithm 3
Remark 4.1. For such choice of norms, we obtain IP 21 (W) = P aa (W)ll2 vnm for any
W E Rmxx, That is, both norms satisfy Assumption 3.3
and their E2 norms are equal to /nm.
For completeness we include the MNGD scheme (Algorithm that replaces the MultiNorm step with SR-Sinkhorn (Algorithm 3
Algorithm 3 SR- ER--Sinhhorn V
Algorithm 4 Sinkhorn GD (SinkGD)
Input: the stochastic gradient V Vw (wW, and L 2 the number of iterations.
Input: T 2 1 the number of updates, (NoosesT the global step-sizes, L the loss to minimize, and T, 2 the number of iterations for the SR-Sinkhorn procedure. for f = 1 to 7 do
Initialize X = V L(W,, (1)) E Rmx".
for f = 1 to L do
Q(X) = Diag(lIX X1.: 12:12 12.12........ (
V. VgL (0)) with (1() Pz
X + (nQ((x--))))) 1
R(X) = Diag( [X 1 2 m2
V. - SR- (anSiinhoorn Ve, L) as in Alg.
X /mx VmXR(())
0.41 + 0g V.
end for Return X
end for Return xr
The Sinkhorn Algorithm. Before explicitly showing the link between Algorithm 3 and the Sinkhorn algorithm, let us first recall the Sinkhorn theorem [Sinkhorn (1944) and the Sinkhorn algorithm Sinkhom & Knopp (1967). Given a positive coordinate- wise matrix A RRT*" there exists unique matrix P E R of the form P QAR with Q and R positive coordinate-wise and diagonal matrices of size m and n respectively. such that P1,, = m1m and p 1m m1,, To find P, one can use the Sinkhorn algorithm that initializes Po == A and computes for l: 2 0:
P
Pk+1 1/ 1/agg///DDiPP__111 11/-1/21m)--
Equivalently, these updates on P can be directly expressed as updates on the diagonal coefficients of Q Diag(u) and R = Diag(v) with u E T and v E RT. By initializing uo = Im an 1m the above updates can be reformulated as follows:
1m
1
Uk+1
Ave
Uk+1
A uk+1

---

where / denote the coordinate-wei -witi division. Franklin & Lorenz 1989) show the linear convergence
of Sinkhorn's iterations. More formally, they show that (ug, 2/) converges some (u* e*) such that p Diag u² )ADiag((") satisties P1, n1m and P 1m mln, and:
dx( (uk, u²) ((MA)2)). and (uk, v²) E (({(A)²*)
where dz is the Hilbert projective metric De La Harpe 1993) and X(A) < 1 is a contraction factor associated with the matrix A.
Links between Sinkhorn and Algorithm 3
Algorithm 3 can be seen as simple reparameterizatioo
of the updates presented in 7 More precisely, given a gradient V E IR=x and denoting == V
we obtain that the iterations of Algorithm 3 exactly compute:
1/2 x+4
Avk
1m
"K+1
1/2
m
A ATux11
1
(8)
where the square-root applied coordinate-wise and returns after L iterations X Diag(u, Diag(u//")VDiee Diagg(?; 1/22 Therefore the linear convergence of Algorithm follows directly from the convergence rate of Sinkhorn, and Algorithm Aagorithm3can be thought as applying the square-root Sinkhorn algorithm, thus the name SR-Sinhkorn, Note also that at convergence (L + +00) we obtain X e Rmxa which is fixed-pe point of both normalized projections, that is Pe. (X*) = Po, (X*) X
from which we deduce that |X vn and ||X", l22 vm as demonstrated in Theorem 3.6
On the Importance of the Scaling. Now that have shown the convergence SR-Sinkhorn, let us explain in more detail the scaling considered for both the row-wise and column-wise normalizations. First recall that both norm gi and 92 satisfy Assumption B.3 and that the norm of their normalized projections is equal to Vnvn. The reason for this specific choice of scaling ( (ynm) is due to the global seep-size in Algorithm 4 In our proposed MNGD, we did not prescribe how to select T/-- In practice, we aim to leverage the same global step-sizes as those used in Adam KKigmma & Ba (2015) for training LLMs, and therefore we need to globally rescale the (pre-processed) gradient accordingly. To achieve that, observe that when EMAs are turned-off, Adam corresponds to simple signed gradient descent, and therefore the Frobenius norm of the pre-processed gradient is simply ynm. Thus, when normalizing either the rows or the columns, we only need to rescale the normalized gradient accordingly.
Computational Efficiency of SinkGD over SWAN. Compared to SWAN Ma et al. (024)) the proposed approach, SinkGD, is more efficient only requires ((u)) numerical operations. In contrast, SWAN, even when implemented with Newton-Schule matrix multiplications, which have a time complexity of O(m² (m + n)). In the next section, we will demonstrate the practical effectiveness of MNGD with SR-S Sinninnhorn, that is SinkGD. This approach manages to be on par with, and even outperforms, memory- efficient baselines for pretraining the family of LLaMA models up to IB scale.
still requires performing matrix-
5 Experimental Results
In this section, we evaluate the empirical performance of applying SinkGD optimizer to LLM pretraining tasks. All experiments were performed on NVIDIA A100 GPUs.
5.1 LiaMA Pre-training Tasks
Setup. We evaluate SinkGD on the memory-efficient LLaMA training benchmark proposed by Zhao et al. (2024a). This benchmark uses LLaMA-bau architecture (Touvron et al. (003)) with RMSNorm and SwiGLU activations (Zhang & Sennrich 2019 Gao al. 2023). We consider models with 60M, 130M, 350M, and 1.3B parameters, all trained on the C4 dataset Raffel et al. (2020) using an effective token batch size of 130K tokens (total batch size 512, context length 256). Specifically, for both 130M and 350M, we use 128 batch size with 4 accumulations, For 60M and 1B, we uses 256 batch with 2 accumulation, and 32 per-device batch size with 2 accunulation and 8xA 100s, respectively. Following the setup of /Zhao et al. (0024)): Zhu et al. (2024), SD is applied to all linear modules both attention and MI. blocks with T. = 5 iterations for the SR-Sinkhorn procedure. For all other modules, that are the embedding layer, the RMSnorm layers, and the last
7

---

Table 1: Comparison with Adam and memory-efficien baselines pre-training LLaMA models with C4 dataset. Test PPL reported, along with memory estimate of the total parameters and optimizer states in BF16 format. The PPLs reported for all competitive methods are taken from Zhao et al. (2024a): (huu et al. (2024) and Ma et al. (2024). We report both the reported Adam results from Zhao et al. (2024a) and our reproduced result. Note that the memory estimations from/Zhau al. (2024a); (Zhu et al. (2024) did not consider the fact. that Adam optimizer was used for emhedding layers. This corrected in our estimates.
Methods
60M
130M
350M
1.3B
PPL
MEM
PPL
MEM
PPL
MEM
PPL
MEM
Adam (reproduced)
33.94
0.32G
25.03
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
Galore
34.88
0.26G
25.36
0.57G
18.95
1.29G
15.64
4.43G
31.06
0.26G
22.73
0.57G
17.03
1.29G
14.31
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
SWAN
32.28
0.23G
24.13
0.43G
18.22
0.93G
15.13
2.98G
SinkGD
30.99
0.23G
22.75
0.43G
16.51
0.93G
13.51
2.98G
SinkGD speed up v.s. Adam (reproduced)
1.60
1.56 X
2.42
2.79
SinkGD speed up v.s. Adam (cited)
1.66
1.73 X
2.10
2.17
Total Training Steps
10K
20K
60K
100K
output layer, Adam optimizer Kingma & Ba (015) is used. We use exactly the same cosine learning rate scheduler as in Zhao et al. (2024a). w Where 10% of total training steps is used for warm-up. Note that, as Zhao et al. (2024a): Zhu et al. (2024), we use a group- learning rate for our optimizer. The effective learning rate used for linear modules in the transformer blocks is of the form oo where " is global learning rate provided by the scheduler and a fixed hyperparameter that we set to = 0.05. For Adam, we use 7/: the learning rate.
Baselines. We consider the following memory-eeficiienn optimizers baselines: Adam (Kingma & Ba 2015); Galore Zhao et al.] (2024a); Fira Chen et al.] (2024b), Apollo and Apollo-mini Zhu et al. (2024), and SWAN Ma al. (2024). For all methods, training uses BF16 precision for weights, gradients and optimizer states by default, except for SWAN that uses FP32 preci- sion to pre-process the gradient Ma et al. (2024). We also perform a grid search learning rate for Adam over {0.01, 0.005, 0.001, 0.0005, 0.0001], except for IB model which we search over {0.001, 0.0007, 0.0005, 0.0003, 0.0001 We do not perform any weight decay for all optimizers.
Table 2: Comparison of the test perplexities obtained during training when training 1B LLaMA with SinkGD v.s. 7B LLaMA using different baselines. For Apollo, Apollo-mini, 8-bit Adam and Galore,
we cite the number from Zhu et al. (2024).
Method
Mem.
40K
80K
120K
150K
8-bit Adam (7B)
26G
18.09
15.47
14.83
14.61
8-bit Galore (7B)
18G
17.94
15.39
14.95
14.65
Apollo (7B)
15,03G
17.55
14.39
13.23
13.02
Apollo-mini (7B)
13.53G
18.03
14.60
13.32
13.09
SinkGD (IB)
2.98G
16.44
14.27
13.17
12.97
Performance evaluation and memory efficiency analysis. The results presented in Table T demonstrate the effectiveness of SinkGD in terms of both computational efficiency and model performance. Notably, SinkGD achieves competitive performance while maintaining the lowest estimated memory consumption, comparable to that of SGD. Across all evaluated models, our method performs on par with or even surpasses Adam and other memory- test perplexity. In particular, SinkGD outperforms all other baselines in this experimental setup for the 350M and 1.3B model variants. Additionally, we quantify the computational efficiency of SinkGD by measuring the speed-up up relative to Adam. This is determined by computing the ratio of the total training steps of Adam to the number of steps needed for SinkGD reach the same final test perplexity. Note also that the reported memory consumption values in Table [ account for three components: (1) memory allocated for model parameters, (2) optimizer-related costs for linear modules within transformer blocks, and (3) Adam's 's memory footprint for remaining parameters.
baselines in terms of
8

---

Table 3: Raw and effective throughput analysis.
Method
Raw / eff. throughput
Adam
53047/ / 53047 (tokens/s)
SinkGD
57982 / 161769 (tokens/s)
Comparative analysis of 1B and 7B LLaMA training. To further evaluate the efficacy of our proposed optimizer, we replicate the experimental setup of [Zhu et al. (2024), but instead train a IB-parameter LLaMA model using SinkGD and compare its performance against a 7B 7B-parameter LLaMA model trained with Apollo, Apollo-mini, 8-bit Adam, and 8-hit Galore. As shown in Table 2 the 1B model trained with SinkGD achieves comparable test perplexities those of the 7B model trained with Apollo after 150K optimization steps, while incurring significantly lower costs. Notably, training the 7B LLaMA model with Apollo requires 15 days on an 8xA 100 GPU setup complete 150K steps, whereas our approach achieves a similar loss in 3.3 days. The reported memory estimates correspond to the total memory cost detailed in the previous paragraph.
5.2 Ablation Study
Throughput analysis. We also assess throughput when training 1.3--parameter model on 8xA 100 GPUs. We use two metrics: (1) the raw throughput which the number of tokens processed per second, and (2) the effective throughput defined as the total training token used by Adam divided by the time (in seconds) used by SinkGD to reach the same test perplexities. These metrics evaluate the impact of the multi-normalization step on training speed, and also account for the fact that some optimizers make more effective use of training tokens. As shown in Table B SinkGD achieves competitive raw throughput compared Adam, suggesting the nulti-normaiization step does not require expensive computations. Furthermore. SinkGD exhibits a 3 x higher effective throughput than Adam, indicating significantly faster wall-clock time convergence.
On the effect of the number of iterations. In this experiment, we measure the effect of applying different iterations of our proposed MultiNorm (Algorithm IT scheme in the specific case of the SWAN and SinkGD methods. More specifically, we train a 130M LLaMA model on C4 datasets and compare the test perplexities obtained after 10K steps. We observe a small but consistent improvement when using L, = 5 iterations, we decide to use this number of iterations in our benchmark evaluation, as reported in table
Table 4: Comparison of the test PPLs obtained during training at 10K steps when training 130M LLaMA model with either SWAN SinkGD using different number of iterations in MultiNorm.
Method
PPL
SWAN (L. =)
26.79
SWAN (I. =5) 5)
26.56
SinkGD (L (L=1)
26.21
SinkGD (L=
26.13
Related Work
Gradient Normalization. Gradient normalization has emerged key technique in optimization, complementing its well-established role in forward-pass operations such as Layer Normalization (LayerNorm) (Ba al. 2016). LARS and LAMB (You al. 2017 2019) employ global normal- ization to raw gradients and Adam's 's layer-wise updates, respectively, improving convergence and mitigating gradient pathologies large-batch training. Apollo (Zhu al. 2024) introduces channel- wise scaling approach, while SWAN (Ma et al. (2024) replaces Adam's 's first-moment estimate with normalized gradients stabilize gradient distributions. Theoretical analyses further underscore the importance of gradient normalization. HHazan et al. (2015) study its convergence properties in SGD, while Cutkosky& & Mehta (2020) demonstrate that incorporating momentum enhances convergence without requiring large batches. |ernstein & Newhouse (2024) interpret normalization in certain optimizers as a form of steepest descent under specific norm, with SignSGD (Bernstein et al. 2018), or standard gradient descent, serving as examples of gradient normalization. More recently, the concurrent work of Vyass et al.) considered iterative processing of gradient matrices from the perspective of whitening.

---

Memory Efficient Optimizers. Optimizers for large-scale training can reduce memory consump- tion primarily through two approaches: (1) low-rank approximation and (2) elimination of internal state dependencies. Low-rank optimizers project gradients onto a reduced subspace, allowing internal state updates within this subspace. ReLoRA (Lialin et al. (023) periodically merges LoRA (Hu) et al. 202T) weights to restore full-rank representations. FLORA (Hao et al. (024) employs random Gaussian projections, whereas GaLore (Zhao et al., 2024a) utilizes singular value decomposition (SVD) for structured projections, further improved by Fira (Chen et al.
2024a) via a compensation state representations. An al-
term. Apollo (Zhu et al. 2024) minimizes memory overhead using rank-
ternative approach eliminates the need for internal states altogeeher. SWAN (Ma et al. (024) removes Adam's ' first and second moments through gradient normalization and whitening. Adam-mini (Zhang ct al. 2024) reduces memory by leveraging block- twice second moment estimation. SGD-Sal -Sa (Xu 2024b) obviates Adam's second moment by precomputing learning rate scaling. Sign-hased optimization (Chen et al. (002cc enables large-scale training using only first-moment updates. Muon 2024), a simplification of Shampoo (Gupta et al. (2088). accelerates large model training whitened hrst-moment updates, further demonstrating the viability of reduced-memory optimizers. 2025) unified memory-efficien optimizers under the perspective of structured Fisher information approximation, and proposed low rank approximation with error corrections reduce memory and accelerate computation.
et al.
(Jordan et al.
(Gong et al.
Alternating Projection. Many iterative fixed-point algorithms employ alternating updates to en- force constraints or refine estimates. A classical example is the Von Neumann algorithm Von Neumann (1950), which alternates projections onto affine subspaces and converges their intersection. The Sinkhorn algorithm Sinkhorn & Knopp ( 1967) similarly alternates row and column normalizations, which can be seen as Bregman projections Benamou al. (2015) onto affine spaces, to approximate entropy-regularized optimal transport. While effective in Hilbert spaces, these algorithms do not generalize arbitrary convex sets. Dykstra's 's algorithm Dykstra (1983) extends these methods by in- troducing correction terms, ensuring convergence to the exact projection. More generally, alternating projection methods have been extended through Pierra's 's product space reformulation |Pieraa (1984), as well as modern techniques like ADMM Boyd et al. (2011) and block-coordinatee methods Tibshi rani (0177) in large-scale optimization. Despite these theoretical advances, extending alternating projection methods to non-convex settings remains significant challenge. Recent progress includes manifold- projection methods[Lewi & Malick (2008), and proximal alternating techniques
al. (2014), which aim to improve convergence in non-convex problems, yet comprehensive theory for convergence remains an open question.
7 Conclusion.
In this work, introduced gradient multi-normalizaton, a novel optimizer design framework. Our approach formalizes gradient matrix processing as the normalization of stochastic gradients with respect to multiple norms, and we propose an alternating optimization procedure achieve this normalization efficiently. We establish that our arbitrary precision, a fixed point of the optimization problem, thereby ensuring that the gradient is ap- propriately scaled according to both norms. Using this principle, we developed SinkGD, lightweight matrix optimizer with reduced computational cost while retaining the same memory footprint as SGD, Experiments on pretraining L.LMMA models with up to 1B parameters demonstrate strong speedup over Adam with significantly reduced memory requirements. Open questions for future research include: 1, understanding the mathematical mechanisms of how gradient multi-normalization enables the removal of EMA states, and when such approaches will fail; 2, the understanding the theoretical properties of multiinnomallzztionn based optimization; and 3, extending the applicability of gradient multi-normalization to other training tasks beyond LLM pre-training, as well as its scalability to frontier model size training regimes.
scheme can approximate, to
10

---

References
Ba. J., Kiros, J. R., and Hinton, G. E. Layer normalization. ArXiv, abs/1607.06450, 2016. URL https://api. httpcc//rpp...mmmrrcccooo.... htt.ac/o......oooooooooouuaaaa:.
Benamou, J.-D., Carlier, G., Cuturi, M., Nenna, I.., and Peyré, G. Iterative bregman projections for
regularized transportation problems. SIAM Journal on Scientific Computing, 37(2)A1111-11138 2015.
Bernstein, and Newhouse, L. Old optimizer, new norm: anthology. arXiv:2409.20325 2024.
arXiv preprint
Bernstein, J., Wang. Y.-X., Azizzadenesheli, K., and Anandkumar, A. signsgd: Compressed optimisa-
tion for non-convex problems. In International Conference on Machine Learning, pp. 560-569 PMLR, 2018.
Bolte, J., Sabach, S., and Teboulle, M. Proximal alternating linearized minimization for nonconvex and nonsmooth problems. Mathematical Programming, 146(1):459-494, 2014.
Boyd, S., Parikh, N., Chu, E., Peleato, B., Eckstein, J., et al. Distributed optimization and statistical learning via the alternating direction method of multipliers. Foundations and Trends@ in Machine learning, 3(1):1-12,, 2011.
Carlson, D., Cevher, V., and Carin, I.. Stochastic spectral descent for restricted boltzmann machines. In Artificial Intelligence and Statistics, pp. 111-119. PMLR. 2015.
Chambolle, A. and Pock, T. A first-order primal-dual algorithm for convex problems with applications to imaging. Journal of mathematical imaging and vision, 10:1200145, 2 2011.
Chen, X., Feng, K., Li, C., Lai, X., Yue, X., Yuan, Y., and Wang. G. Fira: Can we achieve full-rank training of Ilms under low-rank constraint? arXiv preprint arXiv:2410001623, 2024a.
Chen, X., Feng, K., Li, C., Lai, X., Yue, X., Yuan, Y., and Wang, G. Fira: Can we achieve full- rank training of Ilms under low-rank constraint?, 2024b. URL https://arxii..aatt///// org ab1/2410 01623]
Chen, X., Liang, C., Huang, D., Real, E., Wang, K., Pham, H., Dong, X., Luong, T., Hsieh, C.-J., Lu, Y., et al. Symbolic discovery of optimization algorithms. Advances in neural information processing systems, 36, 2024c.
Cutkosky, A. and Mehta, H. Momentum improves normalized sgd. In International conference on machine learning, pp. 2260-222022222 1266-2268. PMLR, 2020.
De La Harpe, P. On hilbert's metric for simplices. Geometric group theory, 1:97-119, 1993.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A.,
Yang, A., Fan, A., Goyal, A., Hartshorn, A., Yang, A., Mitra, A., Sravankumar, A., Korenev, A.,
Hinsvark, A., Rao, A., Zhang, A., Rodriguez, A., Gregerson, A., Spataru, A., Rozière, B., Biron, B., Tang, B., Chern, B., Caucheteux, C., Nayak, C., Bi, C., Marra, C., McConnell, C., Keller, C., Touret, C., Wu, C., Wong, C., Ferrer, C. C., Nikolaidis, C., Allonsius, D., Song, D., Pintz, D., Livshits, D., Esiobu, D., Choudhary, D., Mahajan, D.. Garcia-Olano, D.. Perino, D., Hupkes, D.. Lakomkin, F.,, AlBadawy, F.,, Lobanova, E., Dinan, E.,, Smith, F. M., Radenovic, F., Zhang, F,, Synnaeve, G., Lee, G., Anderson, G. I.., Nail, G., Mialon, G., Pang. G., Cucurell, G., Nguyen, H., Korevaar, H., Xu, H., Touvron, H., 7arov, I., Ibarra, A., Kloumann, M., Misra, I., Evtimov, I., Copet, J., Lee, J., Geffert, J., Vranes, J., Park, J., Mahadeokar, J., Shah, J., der Linde, J., Billock, J., Hong, J., Lee, J., Fu, J., Chi, J., Huang, J., Liu, J., Wang J., Yu, J., Bitton, J., Spisak, J., Park, J., Rocca, J., Johnstun, J., Saxe, J., Jia, J., Alwala, K. V., Upasani, K., Plawiak, K., Li, K., Heafield, K., and Stone, K. The lama 3 herd of models. CoRR, abs/2407.21783. 2024.
Dykstra, R. L. An algorithm for restricted least squares regression. Journal of the American Statistical Association, 78(384):837-844 1983.
Franklin, J. and Lorenz, J. On the scaling multidimensional matrices. Linear Algebra and its applications, 11::7177735, 1989.

---

Gao, K., Huang, Z.-H., Liu, X., Wang, M., Wang. S., Wang, Z., Xu, D., and Yu, F. Eigenvalue- corrected natural gradient based on a new approximation. Asia-Paciifa Journal of Operational Research, 40(01) (0001):334005, 2023.
Gong, W., Scetbon, M., Ma, C., and Meeds, E. Towards efficient optimizer design for Ilm via structured fisher approximation with a low-rank extension. arXiv preprint arXiv:2502.07752, 2025.
Gupta, V., Koren, T., and Singer, Y. Shampoo: Preconditioned stochastic tensor optimization, 2018. URL
org/ abs/ 095./ 09568
Hao, Y. Cao, Y., and Mou, L. Flora: Low-rank adapters are secretly gradient compressors. ArXiv, Abs/2402.03223, 2024. htt.ttttpps//ap...eeentttt wttp://chccoolooooooororraaaaaaa
Hazan, E., Levy, K., and Shalev-Shwartz, Beyond convexity: Stochastic quasi-convex optimization. Advances in neural information processing systems, 28, 2015.
Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang. L.,, and Chen, W. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106. 100.::2000600685, 2021.
Huang, L.. Zhou, Y., Zhu, F., Liu, L., and Shao, L. Iterative normalization: Beyond standardization towards efficient whitening. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 1874-888,, 2019.
Jordan, K., Jin, Y., Boza, V., You, Cecista, F., Newhouse, L., and Bersstein, Muon: An optimizer for hidden layers in neural networks, 2024. URL URL.ttts:/klllororrddaa wew.eerr.rrnnn......ub.b.. github. iot.pots/ muon/
Kingma, D. and Ba, J. Adam: A method for stochastic optimization. In ICLR (Poster), 2015.
Korthikanti, V. A., Casper, J., Lym, S., McAfee, L., Andersch, M., Shoeybi, M., and Catanzaro, B. Reducing activation recomputation in large transformer models. Proceedings of Machine Learning and Systems, 5:341-353 2023.
Kunstner, F., Chen, J., Lavington, J. W. and Schmidt, M. Noise not the main factor behind the gap between sgd and adam on transformers, but sign descent might be. arXiv preprint arXiv:2304. 13960, 2023.
Kunstner, F., Yadav, R., Milligan, A., Schmidt, M., and Bietti, Heavy-tailed class imbalance and why adam outperforms gradient descent on language models. arXiv preprint arXiv:2402.1942 2024.
Lewis, A. S. and Malick, Alternating projections on manifolds. Mathematics of Operations Research, 33(1):216-23,, 2008.
Li, P., Xie, J., Wang, Q., and Gao, Z. Towards faster training of global covariance pooling networks by iterative matrix square root normalization. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 947-955, 2018
I Lialin, V., Shivagunde, N., Muckatira, S., and Rumshisky, A. Relora: High-rank training through low-rank updates. In International Conference on Learning Representations, 2023. URL hhtpps Saapiissmanticscco....o www.oolrrrrrrro/ooooooooo :D:
Ma, C., Gong, W., Scetbon, M., and Meeds, E. Swan: Sgd with normalization and whitening enables stateless Ilm training. 2024. URL http:/ttt/////aatte http://aaaxxv.oorraa///a/aaa
org/ awwp///bs//412... 13148
Pierra, G. Decomposition through formalization product space. Mathematical Programming. 28: 96-115, 1984.
Raffel, C., Shazeer, N., Roberts, A., Lee, K. Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020.
- html
Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-16. IEEE, 2020.
12

---

Rockafellar, R. T. Conjugate duality and optimization. SIAM, 1974.
Sinkhorn, R. A relationship between arbitrary positive matrices and doubly stochastic matrices. The annals of mathematical statistics, 35(2):876-879 1964.
Sinkhorn, R. and Knopp. P. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics, 21(2):343-343 1967.
Song, Y., Sehe, N., and Wang, W. Fast differentiable matrix square root and inverse square root. IEFE Transactions on Pattern Analysis and Machine Intelligence, 45(6):737773800 2022.
Tibshirani, R. J. Dykstra's algorithm, admm, and coordinate descent: Connections, insights, and extensions. Advances Neural Information Processing Systems, 30, 2017.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F,, et al. Llama: Open and efficient foundation language models. arXiv preprint ar arXiv:2302. 13971, 2023.
Von Neumann, J. Functional operators: Measures and integrals, volume 1. Princeton University Press, 1950.
Vyas, N., Zhao, R., Morwani, D., Kwun, M., and Kakade, S. Improving soap using iterative whitening and muon.
Watson, G. A. Characterization of the subdifferential of some matrix norms. Linear Algebra Appl, 170(1):33-44 1992.
Xu, M., Xiang, L., Cai, X., and Wen, H. No more adam: Learning rate scaling at initialization all you need, 2024a. URL http:://axxx org abs/ 2412. 11768
Xu, M., Xiang, L.,, Cai, X., and Wen, H. No more adam: Learning rate scaling at initialization is all you need. arXiv preprint 211v:2412.11766, 2024b.
You, Y., Gitman, I., and Ginsburg, B. Large batch training of convolutional networks. arXiv preprint arXiv:1708.00 038888 2017.
You, Y., Li, J., Reddi, S., Hseu, J., Kumar, S., Bhojanapalli, S., Song. X., Demmel, J., Keutzer, K., and Hsich, C.-J. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv: /904. 00962, 2019.
Zhang, and Sennrich, R. Root mean square layer normalization. Advances Neural Information Processing Systems, 32, 2019.
Zhang. J., Karimireddy, S. P., Veit, A., Kim, S., Reddi, S., Kumar, S., and Sra, S. Why adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33: 15383-15393, 2020.
Zhang, Y., Chen, C., Li, Z., Ding, T., Wu, C., Ye, Y., Luo, Z.-Q., and Sun, R. Adam-mini: Use fewer learning rates to gain more. arXiv preprint arXiv:2406. 16793, 2024.
Zhao, J., Zhang, 7. A., Chen, B., Wang, Z.,, Anandkumar, A., and Tian, Y. Galore: Memory- efficient Ilm training by gradient low-rank projection. ArXiv, abs/2403.035077 2024a. URL https://apt www.stiiscooollaaat..e.ee org/ CorpusID: (0orpus1D:228253599 268253596
Zhao, R., Morwani, D., Brandfonbrener, D., Vyas, N., and Kakade, S. Deconstructing what makes good optimizer for language models. arXiy preprint arXiv:2407.0799 2024b.
Zhu, H., Zhang, Z., Cong, W., Liu, X., Park, S., Chandra, V., Long, B., Pan, D. Z., Wang, Z., and Lee, Apollo: Sgd-like memory, adamw-level performance. arXiv preprint arXiv:2412.055777 2024.
13

---

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

---

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

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [NA]
Justification: our main results used the exact standardized benchmark setup introduced in Zhao al. (2024a), where research code/dataset/confggs habe been open sourced. Methods introduced in this paper will be integrated into an opensource repo
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed run to reproduce the results. See the NeurIPS code and data submission guidelines (https www.eeeeeeeeemeeeeeeeeeeeeeeeeee
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: Appendix we have discussed experimental settings and hyperparameters in detail.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: note that error bars are uncommon for LLM pre- pre-training experiments due to its computational demanding nature. However, the gap between method and baselines are non-trivial. For example, the validation PPL difference hetween ours and Adam on 1.3B model is around 2, which usually considered very significant.
Guidelines:
The answer NA means that the paper does not include experiments.
16

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (c.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: we have provided details on compute types (A100).
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: research presented in the paper is purely methodological and were conducted with only theoretical derivations and standardized benchmark experiments.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
17

---

Justification: we focus on fundamental research on optimization and does not have negative social impacts such as unintended uses, security, privacy issues. In the long run potential social impact would be making the LLM training more economic and sustainable for the society.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: no models/datasets where released.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: we mainly use the LLaMA architecture and C4 dataset, both were correctly credited.
Guidelines:
The answer NA means that the paper does not existing assets.
18

---

The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- re-packaged, both the original license and the license of the derived asset (if has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: our paper concerns fundamental research in methodology and does not release dataset/code/model.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
19

---

Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question:
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: the core method development in this research does not involve LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our L.LM policy for what should or should not be described.
20

---

Appendix
A Implementation details
General setup We describe the implementation setups for the LLM pre-training tasks. enable a more straightforward and comparable analysis, simply replicate the setting of Zhao et al. (2024a) under exactly the same model configs and optimizer 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 BFI6 for model weights, gradients, and optimizer states storage. For SWAN and SWANI, we follow the original paper and use FP32 in there whitening step.
Learning rate scheduling we use exactly the same scheduler as in Zhao al. (2024a) for all methods.
Hyperparameters Since SinkGD utilizes matrix-level operations on gradients, it can only be applied to 2D parameters. Therefore, in our experiments, we only apply SinkGD on all linear projection weights in transformer blocks. Similar to Galore (Zhao et al., 2024a). the rest of the non-linear parameters still uses Adam as the default choice. Therefore, setup of Galore, where we fix some global learning rate across all model sizes and all modules. Then, for the linear projection modules where SinkGD is applied, we simply apply scaling factor top of the global learning rate. For all SinkGD, we adopt a lazy-tuning approach (hyperparameters set without extensive search), detailed below. This helps reduce the possibility of unfair performance distortion due to excessive tuning.
follow the learning rate
Adam For Adam we use same learning rate tuning procedure as suggested by Zhao al.] (2024a) and Ma al. (2024) (i.e., performing grid search over {0.01, 0.005, 0.001, 0.0005, 0.0001f). 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 larger model requires smaller learning rates, we conduct learning search for Adam over a smaller but more fine-grained grid of {0.001, 0.0007, 0.0005, 0. 0.0003, 0.0001]. As result, the optimal learning rate found for Adam on 1.3B is 0.0007. Other hyperparameters of Adam (31, 32, etc) follows the implementation of (Zhao et al., 2024a), that is, we use 31 0.9, 82 0.999. We also perform further ablation on extensive grid-search over all hyperparameters, see Appendix D.I
SWANI, is the tuned version of SWAN presented in Ma al. (2024). The original results of SWAN from Ma et al. (2024) assumes no learning rate warm-up up and no learning rate tuning, in order to demonstrate the robustness of the method. This setting is more challenging than the setting of the usual literature (Zhao et al.) 2024a Zhu et al. 2024). Hence, for fair comparison relax those constraints and matches the setting of Galore and Apollo: now allow learning rate warm-up (set to the same as Adam and Apollo), as well as larger learning rates for SWAN. This improved version of SWAN is denoted by SWANI. We use a global learning rate of 0.02, as well as the scaling factor a = 0.05. This is selected by simply searching the learning rate
constraint grid (0.01, 0.02, 0.05). and then setting 0.05 such that the effective learning rate is scaled back to 0,001, Finally, for other
a = hyperparameters, we follow Ma et al. (2024).
Finally, SinkGD, we use the same global learning rate of 0.02, as well as the scaling factor a 0.05 which are the same as SWANi, 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. For SR-S: Sinkhor V, L.) operation used in SinkGD, we simply use steps.
21

---

B
Proofs
B.1 Proof of Lemma 3.5
Proof. Let us assume that [ and so for any I. Then we have that:
[1I|, o P (a)ll2 9 (P (z)) = o P sup ا (z))
where the first inequality follows from Cauchy-Schward and the second inequality follows from the definition of P, Now recall by definition, that g (x)) S 1, and therefore we can select (, in the right inequality which gives:
PP, o (2 (z) 2 (P)(z))
|P (z) 2
However because IP, o (2( |1p)(1)]2 = c, we obtain that
g' = (
and by optimality, we also deduce that P, o P, (,(x) = Pa(x).
B.2 Proof of Thoerem 3.6
Proof. First observe that thanks to Lemma 3.5 we have for any n 2 1:
| 112 = g (x2n-1) = 92 (22n) = c²
(9)
where of and 92 the dual norms of Φ1 and 92 respectively. We also have that for g2(22n) S 1, 91 (22n+1) <
(10)
by definition of the normalized projections. We even have
(2(22m) 91 (T2m+1) =1
by optimality of the normalized projections. Let assume now that n 2 is even, then we have that:
(In+1, )n)) (P (In), Fnd (=n)
2 (2, xn) V 2 E B, (02)
where B,, (0.) the unit ball centered O₂ associated to the norm øi and the inequality follows from the definition of P In particular by taking = In- Pa, (In-2) E B,, (0z), we obtain that:
(on+1, an)
( 1 1, 3n
A similar proof can be conducted when is odd using the definition of P.,. Therefore the sequence ((#n+1, 1, zn))m,) m22 is increasing and bounded it converges to certain constant > 0. From this result we directly deduces that:
(of (#2a))n21 is monotonic increasing and converges towards
(0t (x2n+ ))) is monotonic increasing and converges towards
Because (120+1).20 and (22n)n20 bounded, we can extract common subsequence (#20(n)+1)221 and (I20(n))221 that converge to some cluster points æ1 and respectively.
Now by continuity of the dual norms and of the inner product obtain that:
lim z (42600+1) g (x1)
n000
lim of (I26()) gi (Iz) n
lim (I26(m)+1: to(n) = (21,
40000
22

---

However observe that these three sequences are subsequences of (f Fat2))22 which converges towards r, therefore we obtain that:
r r= (=))))==)((1)) = (11, x2)
Additionally, remark that
92 (I2d(n)+1) (P (I2
(11)
Let us now show that 220(n)+1 psi (I20(n))
Psi (122). Indeed we have that:
n--00
(Pa. (122(()): x2p(m)) = 9i (T20(n))
yi (222
n--20
Then, because (Por (122(m)) 220 is bounded, we can extract subsequence that converges towards such that g1 (z) I 1, from which follows that:
(2, = (Po1 (22)
then by optimality of Po. (22) over the unit ball induced by 91. deduce that Pa. (22)) This is true for all converging sub-sequences of (Po. (#24(1)))n20 therefore we have that Pai (#20(n)) = > Pgi (I2), and by unicity of the limit, it follows that n-
31 = Ps. (x2)
Now from the equality (x1) = (x1 x2). and given the fact that 92(x2) (as for all g2(222(n) & which is obtained from (10), we deduce that
T2 = Pzz (x1)
thanks to the optimality of Poz Now observe now that:
92 (m1) = ( (x22) = (Poz Poi (x2), Poi (12)) (x2, P (x2))
where the equality follows from the fact that:
Pg2 Psi (x2) (2)(21) 12
and the two equalities follows the previous results obtained. Therefore we obtain that
g2 (z1) = 95 =
where the last equality follows from Thus, obtain that
r == 2 = 02 (x1) = (411 .2) & ||01||2|1=211
but from ( |21 ||>2||2 c. from which follows that 1 122 x, and Ps, (x) Poz (2)(2) As by-product, also obtain that (Zn, In+1)
c², and therefore | Zn
2c2 2(,,, (n+1)
From the above proof, also conclude that if is a cluster point of (In)n20, then there exists eb such that (Z 2220 converges towards y that satisfies: Pa, (v) = Poz (γ) = y. Indeed this follows simply from the fact that we can extract a subsequence of (Irr()) 222 which has all indices that are either even odd.
Let us now show that
91 (In)
1, and (zx)) 1.
Indeed for a convergent subsequence, if the subsequence has infinitely many odd indices the result trivial from the fact 91 (I2n+1) 1. Now if the indices are even, we obtain that 91 (Z2o(n))
9i (z:), however has to be fixed-point so 9i (x) 91 (P pn (z:) 1. This hold for any subsequences, therefore we have 91 (Ix) - 1. Similarly, we can apply the same reasoning for g2 (r,,).
n-
Let us now show the following Lemma.
23

---

Lemma B.1. Let 91 and 92 two norms satisfying the same assumption as in Theorem 3.6 that is for
all 2, IP gt (z)ll2 = |P (z)ll2 with e O. Then by denoting Sg the unit sphere associated
norm y, we have:
Sg gg n F where F {x : Pen (x)) P (2) = x)
Proof. Indeed F C Sn n Sn
n Sel follows directly from the definition of PaL Par and from Now let E Sn n gg S8dt Observe that
Assumption 3.3
e 2 (2, Pi(2))
sup (2,9) gi (a)
where the inequality follows from the assumption Ps. and from the definition of 2. Then as 9n (z) = L, we deduce that:
2 2 (z, P1(2)) |=|2
from which folows that Poi (2) = 2. Similarly we deduce that P P3.(2) = 2, and thus we have Sg. Sg n F.
O. Additionally, from we have -+ O since all these spaces are d(En, Set,) = O, therefore we have that d(x, Sgt n Sg n ( n-oo closed, and the result follows from Lemma B.1
Now observe that d(In, Sg.) x-*
0, and d(zn, Sgy) n-x
C On the Convex Relaxation of Problem
(4
Given K norms, (gi...., ........... in this section interested in solving:
arg max (V, s.t. V e [|1, K]), g((2) s
(12)
which as stated in the main paper is equivalent to solve
arg max V V.2)
s.t. |||||| <
where
||||| == max g:(z). xE|[L,
(13)
Note that this constrained optimization problem exactly finding the subdifferential of the dual norm ||- ||. To see this, let recall the following Lemma with its proof Watson (992)..
Lemma C.1. The subdifferential of a norm || || at x is given by
Ø|l [1(x) = {p IR. Illll s (0, x) = ||-||)
where the dual norm is defined as
||||| max (2,
||
isi
Proof. We can show this result by double inclusion. Let us define the subdifferential of a norm as
a|| [xxx) {p: || 2 |||||| + (p.v x) Vy}
and let denote our set of interest
V(x) = {pe { |p||- < 1, (p. x) = ||=||}
Let E (xx). Then we have
| + (p.v x)) = {p,
< ||p||- |
< ||v||
24

---

where the first equality comes from the definition of the first inequality comes from Holder, and the last one obtained by definition of p. So we deduce that V(x) C 0l| ((x)). Let us now take
p E O|l (a)), that is such that for all y || 2 ||-||| + (p. x.. Then we have for all y that:
(p, x) - || 2 (p. y) 11yl
sup(p. 9) 1
IPll"
where -| ( is the Fenchel Leggndre transform of the norm || ||. From Lemma C.3
we deduce that
(p. σ)
|| 2 1s, (p)
where B1 is the unit ball associated with the dual norm || As the left-hand side is finite, deduce that Bj. Then we deduce that
| (p, x) 2 ||||||
where the left inequality is obtained by applying Holder the inner-product, from which follows the result.
For simple norms, such as the n--mr with 1 < p +00, obtaining an element of 0|| (V- can be done in closed-form. For Ep norms, recall the their dual norm the l norms with q the dual exponent respectively. The following Lemma provide analytic formulas for such norms.
Lemma C.2. For x E IRd, let us define the ep- ---orm as
1/p
d Σ
|0lle =
Case 1: 1 p co. Let us define the dual exponent q by
+
p q
Then the subdifferential of ||olp
(1 10--2200 201 221..... ..... [x p-- P-22xd
x + 0,
0||l|2
||=1|/--
(gE Rdd 119lla S 11}.
x
Case 2: p For the mmmmrmm the subdifferential R2' is given by
0| {9 E : sign(z:))
Here, sign(r;) is +1 +1 if,;
0, -1 if x; O, and can be any value in [-1, if 0.
Case 3: p = oc. For the 80 -norm, let M ||~||o and let us define
S(x) {g g 0 if [.a| M, g: sign(x.) else, |lg 12 1}
Then by denoting for any set A IR.d, conv( A) the convex hull of the A, we have:
x 0 a|| = i{ E R² 11g/1 S $1),
(co)
x=0 0.
However, for general norms, there are not known closed-form solutions of their associated subdiffer- entials. In particular, if the norm is defined as in (13), even when the y, 's are simple norms (i.e. norms for which we can compute the subdifferential their dual norms), then no closed-form solution can be obtained in general.
25

---

C.1 A Dual Perspective
In this section, we propose an algorithmic approach solve the convex relaxation the problem introduced in 4 More formally, given family of simple norms (g: and some positive constants (s)) we consider the following problem:
max (V, d0) s.t. g(d0) S 1
(14)
doERd
where
(:(
g(x) = max [E[IIK]] &i
which is also a norm. For such problems, as long as V o O, then the solutions lies in the level set {d0: g(d0) 1}. Even if the subdifferentials of (the dual norm of) each di can be derived in closed form, there is not known closed-form for the subdifferential of (the dual norm of) To solve (14), we propose consider a coordinate gradient descent on the dual. simple application of the Fenchel duality Rockatellar (1974) leads to the following equivalent optimization problem:
inf A1BoAK
Edgl(A) s.t. VL(0)
Σ A
(15)
where of is the dual norm of 9i and so for all i E [[I, Kl). from which primal solution can be recovered by simply finding yi s.t. ); > and such that (Ae, ui) e.g, (1s) under the condition that ge(1s) Es, which is equivalent to solve:
Vi := E, arg max (2, Ax)
2: (:(=
Proof. Let (Bi(e:)) the ball associated with the norm (g.) with radius (e:) respectively. Let us also denote for any set A c RR the indicator function as
0 x A
1x(x)
+o0 otherwise
In the following we denote f(x) = (x, V). Then 14 optimization problem:
can be reformulated as the following
inf (d0) +
(d)
db
which can be again reparameterized (up to the sign) as
K
inf z= WIElI. f(x) +
1=1
Now the Lagrangian associated with this problem is:
F((A,) (
f(x) - (z. Σ A1) (B)) Ac)
And taking the infimum of the Lagrangian w.r.t the primal variables leads to the following optimization problem:
K
inf f(x
(x.
A1)
+
inf 1 (..)))())) + (Vi. As)
yi
26

---

Now observe that
inf f(z) - (x, Σ x1) =
sup T
Σ x1) f(x)
m
i=1
K
where f is the Fenchel-LLegedrre transform of f. Similarly, we have:
inf 1 (c. + (ys. Ae) sup -x3) (b)
yi
1 (((e)) (->a)
Finally the dual of the problem is:
K
K ( (-xc)
sup Ax
Now recall that f(x) = (x, V), therefore we have that
f*(x) (z)
Also, we have that where of is the dual norm of gi, from which it follows the final dual formulation:
1 (e(es) (x) Eig]
K
inf Σ eagl (A:) s.t. VL(0))
A
Finally, Slater condition are verified, thus strong duality holds, and the KKT conditions gives the following primal-dual conditions:
VC(0) ΣK A A E (18,(c)((i)) Vi
= yi Vi
Now according to Lemma
C.4
we have that
(0) g. <E
())()) =
g() >E, (p: (0.z) e.o! (p)} if gx)) E2
from which follows that can recover primal solution by simply finding y: X; and such that (Ne, Uz) = E.g. (y:) under the condition that (v:) = E,, which is equivalent to solve:
vi := E arg max (z, 2: 9((=)51
To solve the dual problem introduced in 15 precisely, we can reformulate the problem as an unconstrained optimization one by considering:
apply coordinate gradient descent the X, More
(VC(0)-
K
K
inf coe Ag
Σn 2
A
+
49! (A.)
2
27

---

Algorithm 5 Primal-Dual Algorithm to solve (T6)
Input: 3x _2 E1. ck 0, m 72 O, s.t. 7 1.
Initialize X == ====000 ===0000 0a. 0d.
for to L do
Adld + A (en) + m/2 (u Bk))
A prox mgcng! () - 1/22)
2A Aoid
end for Return X
Starting with 2 =...= 1 Od we propose apply the following updates time 2 0 and so for all & E [2, K[]:
(++1)
arg min c1 (s ->x) (xg)((A))
(16)
A
where 0 = V (()
Σ (0) In order solve (6),, we leverage the so-called Chantolle-Pock
algorithm Chambolle Pock (2011). Let us denote hi (A) == Eig A) and ha(A) = s..g Then we can write
inf (A) + h4 (A) inf ha (A) + sup(z, A) - (2(
=infsup(z,,>) - (1())) + (x(
where hi is the Fenchel-Legendre transform of h given by hi(x) = (xx, (0) + 18, (e.) where B3 (ct) is the ball induced by the norm 91 of radius ET. We now ready to present the Chambolle- Pock algorithm for our setting as presented Algorithm This algorithm requires to have access to the projection operation w.r.t the norm 9i and the proximal operator w.r.t the norm d that is, requires to have access to:
projs (c1) (z) arg min || æ:|/2
(=)
|| 1 + (2)
prox Xg. (z) = arg min
Computing proximal and projection operators of norms and their duals can also be done using the Moreau decomposition property which states that:
prox, (x) + prox (x) = x
in particular if f == || is norm, we have:
prox 1n (x) + proj (1) (x)
x
where B_ is the unit ball of the dual norm of ||- ||. Finally, the full coordinate gradient scheme is presented in Algorithm 6 which returns a solution of the primal problem defined in (12).
Lemma C.3. Let ||-II I| be norm on T with dual norm ||=||- = maI (2, x), then the Fenchel- Legendre transform of is the indicator function of the unit ball induced by its dual norm. More formally, we have
0
| S
sup (2, - -||||| = ZERd
+oo otherwise
28

---

Algorithm 6 Coordinate Gradient Descent to solve (T5
Input: the gradient V ( and ......, > 0 Initialize X2 for /. = 1 to 7 do
Aσ- = Og-
for
K do
$ V [C)
Σ (1) kyth
(t+1)
arg ming h1 (2) + he (2) with Alg.
end for end for
Find k such that dE > 0
Return x= == 8h arg max (z, Ax)
g=)
Proof. Using the fact that
||>||| sup=:||=||=<1 (2, 2),, have:
sup
(z, σ)
||=||
max x2,
sup (y1
zER4
=ER*
yel vll.
sup inf
(2, 4)
=eldd villvll-ss
inf sup (z, y)
s:llvll. xER
(0 if y y=
inf slllel--s
+00 otherwise
which gives the desired result. Note that the third equality follows from Sion's minimax theorem.
Lemma C.4. Let || a norm on Ted and > Then we have:
{0) [ <e
(L)()) = 0 if |=|| >
{p: (p,z) = el[pll.) if||=]|| =
where || ||+ is the dual norm of I|. and B(s) is the bull of radius w.r. the norm || ||-
Proof. Recall that the definition of the subdifferential is:
(z) {p: (B) 2 1B(e) (x) (p.y x) Vy}
If ||a||]| < e, then we have that p must satisfy for all E B(g):
(p. - x) S 0
By taking 7 sufficiently small we can therefore choose + P. E B(e) which leads to
>||p|12 <
which is only true for 0 as y can be selected to be negative positive. Now if || > E, then the subdifferential is clearly empty. Finally, let us consider the case where ||>||| = E. We deduce that:
(p,x) Z (p. y) (V(
and so for all y. Therefore we obtain that
(p, x) sup(p, y) (v)
=||pp]]
But we also have that:
e|lp|]- = |||||||><|] 2 (p, x))
from which follows that (p,x) =||p]l+ which conclude the proof.
29

---

D Further ablations
D.1 On Adam hyperparameters sweep
In our main results presented in Section 5 the hyperparameters of Adam (31, 52, etc) follows the setups of (Zhao et 2024a). Moreover, weight decay is used. Below, further fine-tune those hyperparameters and comapare with our method. We conducted an additional parameter sweep for Ir, 31. $2, weight decay, and specifically for the 1B model scale, and obtained the following optimal values:
Learning rate: 0.0007
Betas: (0.9, 0.95)
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 on 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 5: Comparison of the test loss obtained during training when training 1B LLaMA with SinkGD
1B LLaMA under optimally tuned Adam.
Method
40K
80K
120K
150K
Adam
2.880
2.728
2.659
2.651
SinkGD
2.799
2.658
2.578
2.561
For our method, we still use the default parameter setting described in Appendix A Our method still outperforms Adam by a large margin despite spent much more compute to sweep the Adam hyperparameters.
30


Reviews:
[{"review_id": "vUFifCGk90", "reviewer": "Reviewer_yet7", "summary": "This paper proposes an optimization framework which processes the gradient/momentum by applying multiple normalizations. Under this framework, they develop a practical algorithm named SinkGD, which iteratively performs row normalization and column normalization on the momentum without introducing heavy computation overhead. The algorithm is shown to be convergent based on the theory of Sinkhorn algorithm. Numerical experiments show that SinkGD achieves significantly faster convergence than AdamW (2.4-2.8 $\\times$) on the task of pretraining LLMs in C4 dataset.", "questions": "* Can authors provide other norm examples except for the Euclidean norm that satisfies Assumption 3.3?\n* In Table 3, the raw throughput of SinkGD is higher than AdamW, which is confusing to me. AdamW only needs to perform element-wise operations twice, while SinkGD requires iteratively normalizing the gradient for multiple steps.", "limitations": "Yes.", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n* The perspective of steepest descent under multiple norm measurements is novel and unifies existing algorithms such as Muon and SWAN.\n* SinkGD avoids storing second moment term as in AdamW, which is memory-efficient. The iterative normalization operation seems cheap based on the throughput in Table 3.\n* SinkGD attains superior convergence under the LLM pretraining tasks and outperforms AdamW by a large margin, which partially justifies the effectiveness of the multi-normalization step. The algorithm can be of general interest given its memory efficiency and fast convergence.\n\n**Weaknesses**\n* The experiments focus on pretraining stage. To broaden its applicability, I suggest to conduct additional experiments of SinkGD under finetuning stage, where the memory budget is usually more limited.\n* The experiment results are based on pure BF16 precision. It is suggested to provide the numerical results under mixed-precision training scheme, which is often used in LLM's pretraining.", "was_revised": true, "final_justification": "I appreciate the author's response, which resolves my concerns on SinkGD's generalization on finetuning tasks and mixed-precision training paradigm. The stateless feature of SinkGD is appealing and should be of broad interest of the community. I decide to raise my confidence to 5 and support for acceptance with score 5.", "rebuttal": "We thank the reviewer for the constructive feedback and for recognising the novelty of our multi-norm perspective, the memory efficiency of SinkGD, and its strong empirical performance in LLM pretraining.\n\nThe concerns raised are mostly due to scope clarification and implementation details. We address them below.\n\n---\n\n- **SinkGD results beyond pretraining.**\n\n  Our focus is on scalable stateless training for LLMs, where pretraining from scratch provides a clean and controlled setting to isolate the optimizer behaviour. We also acknowledge the importance of fine-tuning, nevertheless, they introduces more confounding factors compared to pretraining (e.g., model quality, task-specific dynamics, optimizers used to pretrain the model) that obscure optimizer effects. That said, we include preliminary fine-tuning results on GLUE datasets (comparing SinkGD to Galore, rank=8):\n\n  | Dataset | QNLI | MRPC | RTE | STS-B |\n  |---------|------|------|-----|--------|\n  | Galore (rank=8) | 92.20 | 92.01 | 79.78 | 90.82 |\n  | SinkGD          | 92.05 | 92.38 | 79.78 | 91.19 |\n\n  SinkGD performs comparably while maintaining lower memory usage. \n\n---\n\n- **BF16 was chosen for consistency with prior work.**\n\n  BF16 is the standard precision used in Galore, Apollo, and SWAN. It reflects realistic memory constraints in large-scale training. To address the reviewer’s concern, we ran additional experiments in mixed-precision (FP32) on a 350M model (batch size: 250K tokens, 30K steps):\n\n  | Steps | BF16 SinkGD Val PPL | Mixed-Precision SinkGD Val PPL |\n  |-------|----------------------|-------------------------------|\n  | 10K   | 20.34                | 20.64                         |\n  | 20K   | 17.10                | 17.38                         |\n  | 30K   | 16.40                | 15.66                         |\n\nWe observed that under mixed precision, the convergence is indeed stronger; however, BF16 SinkGD offers stronger early-mid stage convergence. (A side note: this observation might imply a hybrid approach for efficient training.)\n\n---\n\n- **Assumption 3.3 beyond Euclidean norm (l2 norm)** \n> Can authors provide other norm examples except for the Euclidean norm that satisfies Assumption 3.3?\n\n  Sure. For example, the common choices for optimization, such as $l_\\infty$ norm (used in Adam), and spectral norm (matrix case, used in SWAN). We will clarify this in revision.\n\n---\n\n- **SinkGD achieves higher throughput due to statelessness and parallelism.**\n\n  While SinkGD performs iterative normalization, it avoids maintaining and updating first/second moment states (as in AdamW), as well as subsquent element-wise operations. Additionally, our implementation shards the Sinkhorn computation across devices. Each local rank processes a subset of rows/columns, enabling efficient parallelism. This explains the higher raw throughput observed in Table 3.\n\n---\n\nWe appreciate the reviewer’s suggestions and will incorporate the clarifications and additional results in the revised version."}, {"review_id": "plb6IlHHmh", "reviewer": "Reviewer_Xept", "summary": "The paper tackles the high-memory cost of adaptive optimizers such as Adam in LLM training and proposes a stateless alternative that remains competitive in speed and final accuracy. It does so by framing gradient preprocessing as a multi-normalization problem, proving that simple alternating projections can enforce several norm constraints simultaneously, and instantiating this idea in a lightweight optimizer dubbed SinkGD. Experiments on 350 M- and 1.3 B-parameter LLaMA variants show that SinkGD matches or surpasses Adam’s perplexity while cutting optimizer-state memory to (roughly) SGD levels and delivering 2-3× token-efficiency gains.", "questions": "The experiments focus on LLaMA language model pre-training (with model sizes from 60M up to 1.3B parameters). It would help if the authors discuss the broader applicability of their optimizer. For example, can the proposed stateless optimizer scale to even larger models (tens or hundreds of billions of parameters) without issues? Would it be effective for other training scenarios like fine-tuning, different architectures, or large-batch training beyond the C4 dataset?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\n1.Theoretical Foundation: The paper provides a principled framework (MNGD) and proves key properties of the proposed optimizer. In particular, it shows that the alternating multi-normalization procedure can converge to a fixed-point solution with arbitrary precision. This theoretical result, along with connections to known algorithms, lends credibility to the method’s soundness.\n\n2.Memory and Efficiency Gains: By eliminating optimizer state, SinkGD matches the memory footprint of plain SGD. The method significantly reduces computation compared to SWAN by removing the expensive whitening step. This yields Adam-level step time complexity instead of the higher complexity introduced by SWAN’s whitening, making the approach more scalable without sacrificing accuracy.\n\n**Weakness:**\n\n1.Limitations in Theorem: While the authors prove convergence to a fixed point, they acknowledge that the algorithm is not guaranteed to *optimally* solve the original optimization problem in a theoretical sense. The convergence analysis assumes certain conditions and does not provide guarantees for non-convex objectives beyond reaching a fixed point. This leaves a gap in understanding whether the procedure always leads to the best solution or just a stable one. \n\n2.Scope of Empirical Validation: The experiments, though thorough on LLaMA pre-training, are limited to language model pretraining tasks up to 1.3B parameters. The method’s efficacy on *larger* models or other domains (e.g. fine-tuning, different architectures) remains unproven. \n\n3.Not Enough Contribution: The main algorithm SinkGD is essentially an “improved SWAN” that relaxes SWAN’s constraints for efficiency. Thus,  the work is building on an existing concept (stateless gradient normalization) more than creating a wholly new paradigm.", "was_revised": true, "final_justification": "The topic of stateless optimization for LLM training is very important topic, and this paper provides technique progress either in providing insights or an improving method. Even though the effectiveness of the proposed method should be further validated (e.g., larger models with longer training), I feel this paper can be accepted by NeurIPS.", "rebuttal": "We thank the reviewer for the thoughtful review and for highlighting the theoretical foundation of our framework, the memory and efficiency gains of SinkGD, and its strong empirical performance in LLM pretraining.\n\nWe address the concerns below.\n\n---\n\n- **On the theoretical limitations of convergence guarantees.**\n\n  We acknowledge that our current analysis only guarantees convergence to a fixed point and does not establish global optimality of the multi-norm problem (Eq.4 of the paper). While we do not yet have a complete theoretical explanation, our empirical results suggest that a fixed point solution provide sufficiently adaptive regularization to enable effective training. We believe this opens a promising direction for future theoretical work.\n\n---\n\n- **On broader applicability and scaling beyond 1.3B.**\n\n  Our goal is to enable scalable stateless training. Prior to our work, even demonstrating fast and stable stateless pretraining at the 1.3B scale had not been established. We believe showing that a simple method like SinkGD can match or outperform Adam (with similar throughput) in this setting is already a meaningful contribution. \n\n  Furthermore, in Table 2, we show that a 1.3B model trained with SinkGD matches the performance of a 7B model trained with memory-efficient or quantized Adam variants from the literature. This suggests that SinkGD is competitive even when compared to significantly larger models trained with more complex optimizers.\n\n  Regarding other domains: we include preliminary fine-tuning results on GLUE datasets (comparing SinkGD to Galore, rank=8):\n\n  | Dataset | QNLI | MRPC | RTE | STS-B |\n  |---------|------|------|-----|--------|\n  | Galore (rank=8) | 92.20 | 92.01 | 79.78 | 90.82 |\n  | SinkGD          | 92.05 | 92.38 | 79.78 | 91.19 |\n\n  SinkGD performs comparably while maintaining lower memory usage. We will include these results in the revision and plan to explore larger-scale and cross-domain applications in future work.\n\n---\n\n- **On the novelty of SinkGD beyond SWAN.**\n\n  While SinkGD builds on the idea of stateless gradient normalization, our contribution is not merely an incremental improvement over SWAN. The SWAN paper motivates its design from a specific observation of transformer dynamics in a toy setting. Their work mainly says: the first order dynamics of transformer implies you should use row-normalization; and its second order dynamics implies you should use orthogonalization (ie normed GD under spectral norm). However, SWAN does not discuss why these two operations should be used *jointly*.  In contrast, we address this issue and generalize the core principles behind SWAN into a unified framework of multi-normalization, which recovers SWAN as a special case and opens up a broader design space for stateless optimizers. This new development is not covered/implied in the SWAN paper, we believe this generalization is conceptually significant and will motivate further development of efficient optimizers for large-scale models. \n\n---\n\nWe appreciate the reviewer’s suggestions and will incorporate the clarifications and additional results in the revised version."}, {"review_id": "MLXIfJAchY", "reviewer": "Reviewer_YGTY", "summary": "This paper propose a new framework for stateless optimizers named as Gradient Multi-Normalization. Previous approaches such as SWAN could be interpreted as specializations of Gradient Multi-Normalization. The authors also provide theoretical convergence guarantee for iterative multi-normalization. Based on their framework, the authors propose SinkGD, a new memory-efficient optimizer. SinkGD achieves on-par or better performance compared with previous methods on LLM pretraining tasks, with improved memory costs and convergence speed.", "questions": "- What is the rationale for retaining Adam for the embedding and output layers? How much additional overhead will it bring? What would happen if SinkGD is used for all layers?\n\n- What is the intuition behind normalizing over multiple norms in MNGD?\n\n- While stateless optimization reduces memory costs compared with Adam-based methods, how much additional computational cost does the iterative sub-process in Algorithm 3 incur?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Clarity:\nThe paper is well-written, with clear presentation of theoretical derivations and detailed experimental evaluations.\n\n### Quality:\nThe theoretical insights behind the gradient multi-normalization framework are well-explained and supported with rigorous proofs.  The experiments are comprehensive and demonstrates the superiority of proposed method.\n\n### Significance\n\n#### Strengths\nIn general, I think this paper makes a solid improvement in stateless optimizer design. The MNGD framework provides a new, insightful perspective on stateless optimizers, unifying and generalizing previous work like SWAN. The author demonstrates the strong empirical performance of their optimizer in memory efficiency and training speed-up.\n\n#### Weaknesses\n- The experiments are only conducted on the C4 dataset. Generalization to other tasks such as (parameter-efficient) finetuning is underexplored.\n- The authors did not mention how hyper-parameters in baseline methods (such as rank $r$ in Galore/Fira) are tuned. This could potentially lead to unfair comparisons in Table 1.\n\n### Originality\n\nResults obtained in this paper are new and original.\n\nOther comments:\n- Line 208: typo \"SR-Sinhkorn\" \n- The supplementary material appears to be wrongly uploaded appendix of another paper", "was_revised": true, "final_justification": "The added experiments on finetuning performance have addressed my concerns.", "rebuttal": "We thank the reviewer for the positive assessment and for highlighting the clarity of our presentation, the theoretical rigor of our framework, and the empirical strength of SinkGD.\n\nWe address the concerns below.\n\n---\n\n- **On generalization beyond C4 and pretraining.**\n\n  Our focus is indeed on scalable stateless for pretraining. Pretraining from scratch provides a clean and controlled setting to isolate the optimizer behavior. We also acknowledge the importance of fine-tuning, nevertheless, they introduces more confounding factors compared to pretraining (e.g., model quality, task-specific dynamics, optimizers used to pretrain the model) that masks optimizer effects. That said, we include preliminary fine-tuning results on GLUE datasets (comparing SinkGD to Galore, rank=8):\n\n  | Dataset | QNLI | MRPC | RTE | STS-B |\n  |---------|------|------|-----|--------|\n  | Galore (rank=8) | 92.20 | 92.01 | 79.78 | 90.82 |\n  | SinkGD          | 92.05 | 92.38 | 79.78 | 91.19 |\n\n  SinkGD performs comparably while maintaining lower memory usage. We will include these results in the revision and plan to explore larger-scale and cross-domain applications in future work.\n\n---\n\n- **On hyperparameter tuning of baselines.**\n\n  For Galore and Apollo, we directly use the optimal configurations reported in their respective papers. This is rigorous for the following reasons:\n  1. We use the exact same dataset, model architecture, and training setup.\n  2. We use the same training and evaluation codebases (Apollo’s repo inherits from Galore).\n  3. Both papers report that their methods are relatively insensitive to hyperparameter choices such as learning rate.\n  4. For SinkGD, we reuse the global learning rate and local scaling from Apollo without additional tuning.\n  5. Regarding the rank parameter: our method is stateless, which corresponds to *rank zero* in Galore/Apollo. Thus, the comparison is conservative and favors the baselines (as both Galore and Apollo uses $r>1$).\n\n---\n\n- **On the use of Adam for embedding and output layers.**\n\n  We follow the Galore and Apollo setup, which also retain Adam for these layers. After submission, we found that SinkGD can be applied to all layers. However, for embedding and LM head layers, we observed that reintroducing momentum is important to match performance of SinkGD for Attention/MLP + Adam for embedding. We will include this observation in the revision.\n\n---\n\n- **On the intuition behind multi-norm normalization.**\n\n  Two perspectives:\n  1. From prior work: SWAN, the first stateless optimizer, uses multiple norms motivated by transformer dynamics. Our framework generalizes this into a principled multi-norm formulation.\n  2. From optimization theory: the choice of norm affects convergence. For example, recent work (Xie et al., 2024) shows that Adam exploits $\\ell_\\infty$ geometry to reduce the empirical smoothness constant. This suggests that norm choice should reflect the loss landscape. Our approach can be viewed as an *ensemble of norms*, providing regularization and robustness across diverse geometries.\n\n---\n\n- **On the computational cost of Algorithm 3.**\n\n  Both SinkGD and Adam have the same complexity ($\\mathcal{O}(m^2)$ for square matrices). In practice, especially under distributed training, the cost of Sinkhorn iterations is negligible compared to forward/backward passes. For example, 5 iterations of Sinkhorn involve roughly $5d^2$ flops, while training cost scales with batch size $\\times d^2$. Since batch sizes are large in practice, the relative overhead of SinkGD diminishes at scale.\n\n---\n\nWe also appreciate the reviewer’s suggestions and will incorporate the clarifications and additional results in the revised version.\n\n\n---\n\n**Reference**\n\n1, Xie, Shuo, Mohamad Amin Mohamadi, and Zhiyuan Li. \"Adam Exploits $\\ell_\\infty $-geometry of Loss Landscape via Coordinate-wise Adaptivity.\" arXiv preprint arXiv:2410.08198 (2024).)"}, {"review_id": "GArm4X31Ry", "reviewer": "Reviewer_JbUr", "summary": "This paper proposes a new family of stateless optimizers, termed Multi-Normalized Gradient Descent (MNGD), to improve the memory efficiency of large language model (LLM) training. The method generalizes prior work such as SWAN by normalizing stochastic gradients with respect to multiple matrix norms through a proposed multi-step alternating projection algorithm. The authors further present a specific instance, Sinkhorn GD (SinkGD), which uses the Sinkhorn algorithm to implement the multi-normalization efficiently. Experimental results on LLaMA pretraining tasks (up to 1B parameters) show that SinkGD achieves comparable or better perplexity than Adam while significantly reducing memory consumption and training time.", "questions": "1.Broader Benchmark Evaluation: The paper only provides experiments on LLaMA pretraining with the C4 dataset. Have the authors considered evaluating SinkGD on downstream tasks such as language modeling benchmarks (e.g., WikiText, OpenBookQA, etc.) or real-world applications to demonstrate generalization?\n2.Baseline Coverage: While comparisons to Adam, SWAN, and a few others are presented, important recent optimizers (e.g., Sophia, Lion variants, AdaFactor) are missing. Could the authors justify these omissions or provide further experiments?\n3.Ablation Scope: The only ablation study concerns the number of Sinkhorn iterations. Could the authors provide more analysis on the role of norm choices, projection quality, or convergence stability across different architectures and model sizes?\n4.Scaling and Practicality: While the paper claims that SinkGD is more memory efficient and computationally attractive, the actual implementation details (e.g., compatibility with distributed training frameworks or wall-clock speed) are limited. Can the authors clarify whether SinkGD integrates smoothly into existing LLM training pipelines at scale?", "limitations": "While the paper proposes a promising stateless optimizer with theoretical grounding, several limitations remain. \n1. The current experiments are confined to LLaMA pretraining on the C4 dataset. Without evaluation on downstream tasks, transfer benchmarks, or domain-shifted settings, it remains unclear how well the method generalizes beyond its pretraining context.\n2. The ablation analysis is limited to varying the number of Sinkhorn iterations. Other key aspects—such as the choice of norm, stability across training regimes, or sensitivity to optimizer hyperparameters—are not investigated, leaving gaps in understanding the method’s behavior.\n3.  The empirical section omits comparisons with several competitive recent optimizers, including Sophia, Lion variants, and other stateless or memory-efficient methods. This weakens the strength of the performance claims made in the paper.\n4. While the formulation is well-integrated and practically motivated, the core components—such as multi-norm projections and the use of Sinkhorn iterations—are largely drawn from existing techniques. As such, the contribution may be interpreted more as an effective synthesis rather than a fundamentally novel optimization approach.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths：\nThe paper presents a novel family of stateless optimizers, Multi-Normalized Gradient Descents (MNGDs), which generalizes prior work like SWAN by normalizing gradients with respect to multiple norms. The proposed SinkGD variant leverages the Sinkhorn algorithm for efficient and scalable normalization, and is theoretically justified through convergence analysis. Empirical results demonstrate improvements in memory efficiency and throughput in large language model (LLM) pretraining tasks, suggesting practical utility.\n\nWeaknesses:\n1.Limited experimental scope: All experiments are restricted to LLaMA pretraining on C4. There is no downstream task or broader benchmark evaluation to verify robustness or generality.\n2.Weak ablation analysis: The only ablation varies the number of Sinkhorn iterations. Other key factors, such as norm selection, convergence stability, or sensitivity to optimizer hyperparameters, are not analyzed.\n3.Insufficient baselines: Several strong recent optimizers (e.g., Sophia, Lion variants, or other stateless optimizers) are missing from the comparison, which limits the strength of the empirical claims.\n4.Incremental novelty: While the proposed formulation is cohesive and practically motivated, it mainly combines existing components (e.g., Sinkhorn projection). The contribution may be viewed more as a practical synthesis than a fundamentally new optimization paradigm.", "was_revised": true, "final_justification": "All concerns have been resolved.", "rebuttal": "We thank the reviewer for the detailed feedback and for acknowledging the theoretical grounding, practical motivation, and empirical performance of SinkGD. We address the concerns below.\n\n\n---\n\n- **On ablation analysis.**\n\n  (1) Hyperparameter sensitivity: SinkGD has very few hyperparameters. The number of Sinkhorn iterations is ablated in the paper. The learning rate is directly transferred from Apollo [2] and used across all model sizes without tuning. This demonstrates robustness to hyperparameter choices.\n\n  (2) Norm choice: we emphasize that SinkGD *is* defined by row- and column-wise norm normalization. All other common matrix norms (e.g., spectral norm, $l_\\infty$ norm) should correspond to other baselines (SWAN, AdamW) already included in our comparisons. Thus, norm choice is implicitly ablated through baseline comparisons.\n\n  (3) Stability: SinkGD is stateless and does not accumulate second moments. This avoids instability issues observed in Adam due to large update norms from moment ratios [1]. SinkGD also avoids instability from low-rank approximations, which are observed in [4] and [6]. Moreover In Table 2, our 1 B model trained with SinkGD matches the performance of 7B models trained with Galore/Apollo/Adam across training stages, showing no signs of instability. In contrast, instability of Adam under the same setup is reported in [2].\n\n---\n\n- **On baseline coverage.**\n\n  We acknowledge that including more baselines is beneficial. Due to time constraints during the rebuttal, we were unable to complete benchmarking of Lion/Sophia under our setup. However, our paper targets memory-efficient LLM pretraining, and we follow the baseline setup of Galore [3], Apollo [2], SWAN [5], Fira [4], and GWT [6], none of which compare with Lion/Sophia as they operate under different memory regimes.\n\n  Additionally, our paper compares with 8 strong baselines, including SWAN, Apollo, and Fira, which are more recent than Lion and Sophia. We believe this provides a fair and comprehensive evaluation within the scope of memory-efficient stateless training.\n\n---\n\n- **On generalization beyond C4 pretraining.**\n\n  Our focus is on scalable stateless training. Prior to our work, even demonstrating fast and stable stateless pretraining at the scale considered in this paper had not been established. We believe showing that a simple method (SinkGD) derived from MNGD can match or outperform Adam (with similar computational cost) in this setting is already a meaningful contribution.\n\n  Moreover, our setup closely aligns with recent work on memory-efficient optimization such as Galore [3], Apollo [2], and SWAN [5], which also focus primarily on pretraining, where memory constraints and efficiency are critical.\n\n\n---\n\n- **On novelty and contribution.**\n\n  We respectfully disagree with the assessment that our contribution is merely a synthesis of existing components. If we look at the literature\n  - signSGD [7] (and arguably, reduced Adam) uses the sign operator, a known mathematical primitive.\n  - SWAN [5] and Muon use whitening/orthogonalization operators from prior literature such as computer vision architectures.\n  - Galore [3] uses SVD/subspace projections, well-known in optimization.\n\n  However, these optimizers are all considered significant contributions. This tells us that we must carefully distinguish the *projection function* itself used in updates, vs *the underlying optimization principle that derived and propose to use those projections*, that leads to practical gains that are not predicted by prior works. \n\n  In our case:\n  1. To the best of our knowledge, no prior work has proposed using Sinkhorn iterations as a gradient projector; and no prior work has made sufficient hint that could trivially predict the practical gain of such operations.\n  2. Our work is the first to propose performing steepest descent under an ensemble of multiple norms. This meta-viewpoint enables the synthesis of new optimizers beyond SinkGD and opens up a rich design space. While SWAN touches on this from the perspective of transformer dynamics, it does not explicitly uncover or formalize this principle.\n\n---\n\n**References**  \n[1] Molybog et al., \"A theory on Adam instability in large-scale machine learning\", arXiv:2304.09871 (2023)  \n[2] Zhu et al., \"Apollo: SGD-like memory, AdamW-level performance\", arXiv:2412.05270 (2024)  \n[3] Zhao et al., \"Galore: Memory-efficient LLM training by gradient low-rank projection\", arXiv:2403.03507 (2024)  \n[4] Chen et al., \"Fira: Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint?\", arXiv:2410.01623 (2024)  \n[5] Ma et al., \"SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training\", arXiv:2412.13148 (2024)  \n[6] Wen et al., \"Breaking memory limits: Gradient wavelet transform enhances LLMs training\", arXiv:2501.07237 (2025)  \n[7] Bernstein et al., \"signSGD: Compressed optimization for non-convex problems\", ICML (2018)"}]

Author rebuttal:

Meta-review: The paper introduces Gradient Multi-Normalization (MNGD), a principled framework for designing stateless optimizers by normalizing gradients across multiple norms, generalizing SWAN and enabling SinkGD, a highly efficient optimizer for large-scale LLM pretraining.

The main contributions include a novel multi-normalization scheme, theoretical guarantees for convergence to a fixed point, and the development of SinkGD, which achieves ~3× faster training than Adam while retaining comparable or superior performance, particularly for LLaMA models up to 1.3B parameters. The strengths of the paper lie in its strong theoretical foundation, a unified perspective linking SWAN as a special case, and comprehensive empirical validation across multiple scales with significant computational and memory benefits. 

Additionally, the work addresses an urgent challenge in training LLMs efficiently and shows a clear practical impact. The primary weakness is limited evaluation beyond LLaMA pretraining, leaving its generalizability to other architectures or modalities less explored; however, this does not diminish the technical depth and significance of the presented contributions. 

During the rebuttal period, reviewers questioned scalability, implementation practicality, and comparisons against recent memory-efficient optimizers, but the authors addressed these thoroughly with additional experimental clarifications and theoretical insights, which strengthened confidence in the results. 

Overall, the paper offers a well-justified, theoretically sound, and practically impactful contribution.
Decision: Accept (poster)