ASGO: Adaptive Structured Gradient Optimization
Authors: Kang An
Yuxing Liu
Rui Pan
Yi Ren
Shiqian Ma
Donald Goldfarb
Tong Zhang
Venue: NeurIPS 2025 poster ()
Keywords: optimization
optimization theory
adaptive gradient method
optimization for deep learning
convex optimization
Primary area: optimization
TL;DR: We develop an optimizer ASGO that can provably exploit the low-rank gradients and block-wise diagonal Hessians in training.

Abstract: Training deep neural networks (DNNs) is a structured optimization problem, because the parameters are naturally represented by matrices and tensors rather than simple vectors. Under this structural representation, it has been widely observed that gradients are low-rank and Hessians are approximately block-wise diagonal. These structured properties are crucial for designing efficient optimization algorithms but may not be utilized by current popular optimizers like Adam. In this paper, we present a novel optimization algorithm ASGO that capitalizes on these properties by employing a preconditioner that is adaptively updated using structured gradients. By fine-grained theoretical analysis, ASGO is proven to achieve superior convergence rates compared to existing structured gradient methods. Based on the convergence theory, we further demonstrate that ASGO can benefit from the low-rank and block-wise diagonal properties. We also discuss practical modifications of ASGO and empirically verify the effectiveness of the algorithm on language model tasks.

ASGO: Adaptive Structured Gradient Optimization
Kang An ; Yuxing Liu²; Rui Pan?, Yi Shiqian Ma', Donald Goldfarb', Tong Zhang²
Rice University Universityy of Illinois Urbana- Urbana-Champaige Champaign $ Meta Platforms, Inc. *Columbia University
Orang. an, chiqion. mxltrice, oda, Kyuxingo, wwipt tozhang) anniiniii............m.o eds, yiren con, goldfardecolubia.cce
Abstract
Training deep neural networks is structured optimization problem, because the parameters are naturally represented by matrices and tensors rather than by vectors. Under this structural representation, it has been widely observed that gradients are low-rank and Hessians are approximately block diagonal. These structured properties are crucial for designing efficient optimization algorithms, but are not utilized by many current popular optimizers like Adam. In this paper, we present novel optimization algorithm ASGO that capitalizes these properties by em- ploying a preconditioner that is adaptively updated using structured gradients. By a fine-grained theoretical analysis, ASGO is proven achieve superior convergence rates compared to existing structured gradient methods. Based on this convergence theory, we further demonstrate that ASGO can benefit from low-rank gradients and block diagonal Hessians. We also discuss practical modifications of ASGO and empirically verify ASGO's effectiveness on language model tasks. Code is available at https: //g thub. com/ com/inffi stasstySttrrs/AAAG
Introduction
Numerical optimization algorithms, especially those that can efficiently train large foundation mod- els [Devlin et al., 2018, Brown et al., 2020, Touvron et al., 2023a,b, Ouyang al., 2022], play an important role in the modern machine learning field. Among them, adaptive gradient methods like AdaGrad [Duchi et al., 2011] and Adam [Kingma and Ba, 2014] are popular choices, gaining huge success in training state-of-the-art models in many tasks. These algorithms typically apply a diagonal matrix preconditioner the gradient It to update the deep neural network (DNN) parameters we:
W2+1 w/ A,A, 9t, where w, E RR 9t E RR and A, Rdxd is diagonal matrix.
This coordinate- wise step size design has been theoretically verified to be effective as can exploit the sparsity of the gradient vectors 9r [Duchi et al., 201 1]. Also, when the Hessian is well-approximated by a diagonal matrix whose diagonal entries have very different scales, adaptive gradient methods have been proven to be beneficial [Liu et al., 2024, Jiang et al., 2024a, Xie et al., 2024]. While these results seem to be convincing, common DNNs do not necessarily have sparse gradients or Hessians that are well-approximated by ill-conditioned diagonal matrices. Instead, if we take the matrix structure of gradients in DNNs into account, it has been widely observed that these structured gradients usually low-rank [Zhao et al., 2021, Yang al., 2023, Cosson al., 2023], and the Hessians are well-approximated by block diagonal matrices [Collobert, 2004, Zhang et al., 2024a,b]. Since adaptive gradient methods like Adam, treat the parameters vectors and ignore the matrix structure of gradients, they are generally unable to exploit these structured properties. This leads us ask the following question: How can we effectively use the matrix structure of gradients to exploit their low-rank and block diagonal properties?
Equal Contribution. Ordering for the first two authors determined by a coin flip.
39th Conference on Neural Information Processing Systems (NeurIPS 2025).

---

One possible answer is provided by Shampoo [Gupta et al., 2018]:
W4+1 = W - - R, where Ws, G2 IR m and Lt E Rmxm R, E Ru are full matrices. The main motivation for such a design is that if we apply vectorization to the update, the Shampoo preconditioner is a single matrix that is the Kronecker product of L, and R which approximates the full-matrix preconditioner for AdaGrad [Duchi et al., 201 1]. However, the theoretical convergence of Shampoo is worse than that for AdaGrad or even SGD when the dimension is large. Also, Shampoo needs more memory and much heavier computation than adaptive gradient methods because it requires two preconditioners, making it less suitable for training large-scale DNNs.
In this paper, we provide an answer to the aforementioned question by proposing ASGO (Adaptive Structured Gradient Optimization), which significantly improves the convergence guarantees of Shampoo while requiring less memory and computation. In light of the analysis that demonstrates the benefits of adaptive gradient methods [Duchi et al., 2011, Liu et al., 2024, Jiang et al., 2024a, Xie et al., 2020], we use appropriate assumptions to enable more fine-grained convergence analysis, that yields superior convergence results and an explanation of how ASGO can benefit from the low-rank and block diagonal properties of given problem. We also discuss the connection between ASGO and Muon [Jordan et al., 2024], a structured gradient method based upon the steepest descent algorithm with respect to the spectral norm, to conjecture relation between ASGO and Muon analogous to the relation between AdaGrad and SignSGD [Bernstein et al., 2018, Kunstner et al., efficient design for query-key attention parameters in transformer
2023]. Furthermore, we develop models and examine its empirical performance on pretraining transformer model tasks. Our main contributions are summarized as follows.
We propose the structured gradient based algorithm ASGO, theoretically analyze its convergence, and show that it converges faster than full-matrix AdaGrad and Shampoo.
We further demonstrate that ASGO can effectively exploit the low-rankness of gradients as well as the approximate block diagonal property of Hessians that is typically observed in training DNNS, indicating good theoretical properties under realistic settings.
We develop a practical implementation of ASGO that is efficient while enhancing ASGO's s effec- tiveness. We further empirically validate the effectiveness and efficiency of this implementation on language model tasks, demonstrating the algorithm's great potential real applications.
2 Related Work
Adaptive Gradient Methods. Adaptive gradient methods that use diagonal preconditioners to speedup convergence are extremely popular for solving many real-world optimization problems. To the best of our knowledge, the first method of this kind for machine learning, AdaGrad [Duchi et al., 2011, Streeter and McMahan, 2010], was developed based on rigorous theory that showed the benefits of using this kind of preconditioner. Adam [Kingma and Ba, 2014, Loshchilov and Hutter, 2017] modified AdaGrad and has become the default choice for training large foundation models. In theory, has been proven that adaptive gradient methods can benefit from sparse gradients and approximately ill-conditioned Hessians [Duchi et al., 2011, I, Liu et al., 2024, Jiang et al., 2024a, Xie et al., 2024]. It is worth noting that the original AdaGrad paper [Duchi al., 2011] also proposed a version of AdaGrad that uses a full-matrix preconditioner instead of the diagonal one, which is believed to perform even better. However, this full-matrix AdaGrad method suffers from large memory costs for storing the preconditioner, while having no better convergence guarantee compared to diagonal AdaGrad and SGD under the assumption of objective function convexity as described in Section 4 below.
Optimization with Matrix Structure. In the standard optimization literature (i.c., excluding areas such as conic optimization), is common to consider variables as vectors. However, recently, optimization methods for machine learning that consider variables as matrices have been rapidly gaining attention. Adafactor [Shazeer and Stern, 2018], LAMB [You al., 2019], and Adam- mini [Zhang et al., 2024b] consider the matrix or layer structure to help reduce the memory cost of Adam and enable more efficient training. Shampoo [Gupta et al., 2018] and KFAC [Martens and Grosse, 2015] are two pioneering works that approximate, respectively, the full-matrix preconditioner of AdaGrad and the true Fisher matrix (FM), using Kronecker products of smaller matrices, making the memory cost more affordable. Unfortunately, the analysis in [Gupta et al., 2018] shows that

---

the rate of convergence for Shampoo is no better than that for the full-mat mull-matrix AdaGrad. We note that there is another method, TNT [Ren and Goldfarb, 2021] that is very closely related to Shampoo. TNT was developed as natural gradient method, approximating the true FM by the covariance of block-wise sampling-basee gradients assuming that they are Vensor-Normally distributed. The main difference between TNT and Shampoo, is that TNT uses true FMs and inverses of their Kronecker factors, whereas Shampoo uses empirical FMs and the 1/4 power of their factors. More discussion on related recent and concurrent work is presented in Appendix A.
3 Our ASGO Algorithm
3.1 Notation and problem setting
Throughout this paper, we use capital letters such as W to represent matrices and [W]. to denote the (i. j-th entry of W. For an arbitrary matrix W E Rmxx, we denote
opp as the spectral norm matrix W, i.e., its largest singular value;
[[WI,, as the trace norm of a matrix W, i.e., the summation of its singular values, which is well-known the dual norm of the spectral norm:
|| Ir as the Frobenius norm of W, which also equals ir (WTW) W). where ir (-) is the trace;
|[ a (WT LW), where L. E mm is real symmetric positive definite matrix.
For symmetric square matrices, A, B E Rmxm A = B denotes that A is positive semidefinite and A < B denotes that B - A is positive definite. > and 2 are defined accordingly. We study the following stochastic optimization problem:
min WERm x S(W) - Eg [S(W,)))
where we only have access to a stochastic gradient oracle (V(: W.
3.2 ASGO (Algorithm 1)
We propose ASGO (Adaptive Structured Gradient Optimization) in Algorithm 1. which is an algorithm with single-side preconditioner. Compared to full-matrix AdaGrad, ASGO preserves and utilizes the matrix structure of We and G,, avoiding the huge memory cost for storing its preconditioner. ASGO's preconditioner consists of a single matrix compared to the two matrices used by Shampoo, leading to the following main update rule:
W++1 W₁- - V,,"" Gr, where W, E _ and V3 E p is a full matrix.
ASGO only needs to store one preconditioner matrix and compute its matrix square root and inverse on each iteration, versus two matrices for Shampoo. On the other hand, ASGO's preconditioner may not be as good an approximation to the full-matrix AdaGrad preconditioner or the empirical Fisher matrix [Gupta al., 2018, Morwani et al., 2024] Shampoo's 's. However, as we shall see in the following sections, this design can exploit the low-rankness of gradients and the block diagonal nature of Hessians achieve better convergence rates.
4 Nonsmooth Theory
Although DNNs are generally nonconvex globally, convexity may apply locally in some regions. Thus, convex analysis can be helpful for understanding the behavior of DNN training algorithms.
Assumption (Convexity). f(-) is convex and W. one of its minimizers.
Theorem 1 (Nonsmooth convergence). Under Assumption 1, for Algorithm I with n v) Dop it holds that
b11
eDi
-(W-) 1
Σ
G.G.
Dop
D
T
where Dop maxosisT-1 I|W, -Wollop and DF maxo |W, -WelF
3

---

Algorithm ASGO (Adaptive Structured Gradient Optimization)
1: Input: Wo E m=xx, schedule {me} batch size M E N, and the number of iterations
<> 0, (e should be small, similar to the for Adam or AdaGrad.)
2: Initialize V-1 = m
3: for ( 0 to 7 do
4: Sample mini-bater B, with [B,] = M uniformly
5: G₁ M Eeer, Vw ((W:g)
6: Vi = V
7:
A, = V + elm
8: W,+1 W n,
9: end for
Corollary 2. If we also assume an upper bound for each stochastic sub-grutd sub-ggradint such that
E. E.G.G.]] Q², where Q E mm is positive definite matrix, Theorem also implies
(1211. Dop
<Dr. DopT)
Σ [LF(WOI-F(M))) SO
Remark 1. Note that treat Dop constant here, but it may increase the iteration number increases. This can be addressed, for example, by invoking projection onto a bounded convex set W with respect the norm I1-I1x, in each iteration, in AdaGrad [Duchi et al., 2011]. Dop would then be bounded by the spectral norm of W. However, since this projection is rarely used in training DNNs, we follow Gupta al. [2018] and omit it in Algorithm 1. We refer interested readers to Appendix H for a discussion of how one can incorporate the projection operation and obtain Dop and Dy independent of More importantly, this theoretical bound depends the trace norm of gradients and the spectral norm of weights, showing that the algorithm can make use of the low-rank property of gradients.
One can easily check that this convergence rate for convex nonsmooth problems O(1/VT), the same as SGD [Zinkevich, 2003] (see below) and AdaGrad [Duchi et al., 2011].
Comparison with SGD, The convergence rate for SGD under the assumptions of Corollary is:
T-1 SGD:
BiS(W))) s(w.) < oo DF 1Qll (
Σ 1=0
T
where Dr and || the Frobenius upper bounds for the weights and gradients, respectively. By comparing this bound with Corollary 2, we have
||Ql| < ||QII_ < Vra [[Q]pp where no the rank of Q. Thus when the G, are low-rank, or have very imbalanced singular values, ||Q],, can be close to |QQ|
Dr/VrD Dop < Dr, where rn max rank of. W.). Thus, when the W. W. relatively high-rank or have lots of singular values of a similar scale, Dop can be much smaller than Dp.
Therefore, ASGO should work well when G, are low-rank and W W. are relatively high-rank.
Intuition about ASGO in practice. We argue that Ge is low-rank and W, W, is often relatively high-rank in many practical tasks, revealing the potential of ASGO in applications. As we have discussed in Section 2, gradients are commonly low-rank in DNNs, verified by Zhao al. [2021], Yang et al. [2023], Cosson et al. [2023]. Meanwhile, given the huge success of LoRA [Hu et al., 2022] in fine-tuning foundation models, which employs low-rank total updates of W. there may seem to be a conflict to assume that W, - W. has a high rank. However, it has been observed that Wo - W. should relatively high-rank to obtain a better result, at least pretraining and some complex fine-tuning tasks for large foundation models [Lialin et al., 2023, Jiang et al., 2024b, Huang al., 2025]. Further exploration of the connection between the rank of weight updates and the performance of algorithms is an interesting topic for future research.

---

Comparison with Full-Matrix AdaGrad and Shampoo. Shampoo and the full-matrix AdaGrad achieve the following convergence rates under the same settings as Theorem 1.
m
n
T-1
Full-Matrix AdaGrad:
Dr
Σ Gells
j=1
#1
tr E ))
[-1 #
Shampoo:
O
VrDp . tr Σ
G.G.
We check that Theorem 1 indicates convergence speed that is at least Dp/ Dop times faster than full-matrix AdaGrad and vraDp/DD Dop times faster than Shampoo; (see Appendix E for proofs). This also provides theoretical evidence that single-side preconditioning may be better at exploiting low- rankness of gradients, yielding a faster convergence speed compared to Shampoo-like preconditioning. Remark 2. Proofs of all of the results in this section given in Appendix E. We keep the matrix structure of Wt and Gt throughout our analysis, in contrast to standard analyses of the convergence of AdaGrad and Shampoo, which are based on vectorizations of W, and Gt. This is important for proving that ASGO can exploit the structured properties. We note that for simplicity in Algorithm I above, the preconditioner operates left-side of Gt. However, in practical implementation of ASGO, using the right-sided preconditioner is preferable unless the number of columns of G. is much greater than the number of rows. See the discussion toward the end of Section 5 and in Section 7.1.2 and Algorithm 2 in Appendix B.
Smooth Theory
It is also important to study the performance of Algorithm in smooth settings, as many real training tasks have been widely observed to be smooth, at least locally. Also, only smooth settings can we properly describe the importance of batch size.
Assumption 2 (Smoothness). f is 1-smooth with respect to 11-122- where L E mm is symmetric positive definite matrix and for any X E RRxx
X 2 LXT
If = [.1,.... [x1, Nn), where each x, E Rmm and we vectorize X, i.e., x vec(X) ".....,"""" "T, and form the block diagonal matrix L =diag[[,.... T.), we obtain that ir LX) Ei x2 Lx, = Lx. This means that Assumption 2 is equivalent to the existence of sym- metric matrix L mmmm L > O such that for any w E Rm -L (²f,(w) x L, where fo(w)= S(W), W [Wi...., Wn]] E Rmxx and w vec(W) [w] ...... w T Rm Hence, is closely related to the block-wise diagonal structure of the Hessian as observed by many re- searchers with each block corresponding to a column of W; (see Figure 3 in Appendix A). Also note that Assumption implics the standard smoothness assumption with respect to the Frobenius norm by fo (w) S ||Z|lop- This bisck-wise diagonal smoothness is extension of the standard smoothness, which is related to but different from the diagonal anisotropic smoothness employed for analyzing sign-based and adaptive gradient methods [Bernstein et al., 2018, Liu et al., 2024].
Assumption 3 (Variance). Let N, A V E) - V((W)) Ruxn be the stochastic gradient noise. We assume that [[N.] and there exists a symmetric positive definite matrix V such that
E N.N. = 2
One can check that Assumption 3 implies the standard variance bound E < ||V The assumption shares some similarity with the coordinate-wist variance bounds in Bernstein et al. [2018), Crawshaw et al. [2022], Liu et al. [2024], in the sense that it allows more fine-graines analysis. This matrix-form variance upper bound may better describe the real case since it takes the structure of the noise into account, which is relevant to matrix rank and other structural properties.
Theorem 3 (Smooth Convergence). Under Assumptions and for Algorithm I with m 17 Dop and a batch size of M, it holds that
4D2 op
||||. T
+ 2/2Don IV
2cD
Σ B[((W)]] - S(W.) <
+
VMT
Dop T
5

---

where Dop maNosIsT-1 IW: W. and DE maxogigT-1 [W; W-lpp
As discussed in Section 4. Dop and Dp can be bounded if we add a projection step in the update. The convergence rate specified in Theorem 3 is similar to the rate in Corollary 2 if the batch size M is small, when the O(L/VMT) /MT) term dominates the rate, and thus shares the same properties as we discussed in Section 4. This means that ASGO can benefit if the stochastic gradient noise V is generally low-rank. Furthermore, when the batch size M is large so that the O[I /T) term contributes significantly to the bound, Theorem 3 implies more, which we now discuss.
Comparison with SGD, The convergence rate of SGD [Garrigos and Gower, 2023] is:
SGD: 1 T-1
Σ B[f(Wa)] f(W.) so
D ITllop
DE IV
VMT
where Dr is defined in Theorem 3. Since the comparison between the ((1/vMM) term is generally consistent with the discussion in Section 4, we further compare the O(1/I)) term here to provide some more intuition. Specifically, we have
|| = |b|l, & TL 11Lllopp where Tz is the rank of L. Thus when L is low-rank, or has very imbalanced singular values, ||]||, can be close to |'l
DF/yrD < Dop I Dp, where rn max rank of (W, W.). Thus when W. W. are of relatively high-rank or have lots of singular values of a similar scale, Dop can be much smaller than Dp.
Therefore, can see that in general, ASGO should work well when the Hessian can be well approximated by block diagonal matrix with low-rank L for each block and W, - W, are of relatively high-rank. Note that Hessians have been found to be low-rank in DNNs, especially after some steps of training [Sagun et al., 2016, 2017, Wu et al., 2020].
Block diagonal structure of DNNs. It has been widely observed that the Hessian of MLPs and other DNNs, are approximately block-wise diagonal [Collobert, 2004, Zhang et al., 2024a,b, Bahamou al., 2022], shown in Figure 2 in Appendix A. This block-wise diagonal structure naturally arises from the structure of MLPs, where all the parameters associated with a particular neuron (i.e., all elements of row of W) are more closely related and whose corresponding sub-block of the Hessian matrix is denser and has elements that are larger in absolute value than a set of parameters that do not share such an association. Intuitively, the elements of the rows of W seem to be more likely to be closely related than those of the columns of W. Based on Theorem 3, our algorithm should perform well in this block Hessian setting, showing great potential in real applications.
Intuition on single-side preconditioners. Since ASGO only uses a single-side preconditioner, is important to determine which side they should be applied. As demonstrated in Figure 2, block in the Hessian commonly corresponds to all input weights into a neuron in the next layer, i.e., as a row of the weight matrix W. Therefore, based on the block-wise smoothness convergence results, we should put the preconditioner A; on the right-hand side of Gt, i.c., computing A, = v/l/2, where = V + G,T G, E Ra*", In this way, we can better exploit the superior adaptability of block diagonal Hessians with correct block partitions (each row of W Algorithm 2 for more details of our implementation. We also provide empirical results that show that right-sided outperforms left-sided preconditioning in Section 7.1.2.
block). One may also refer to
Remark 3. The proof of Theorem 3 is given in Appendix F. The analysis is similar in its general outlnee to the smooth analysis for AdaGrad [Levy al., 2018, Liu et al., 2024], but it is more complex because of the involvement of matrix operations in ASGO.
Further Discussions on ASGO
Connection with Muon. Muon [Jordan et al., 2024] can be interpreted as a standard steepest descent algorithm utilizing the spectral norm [Bernstein and Newhouse, 2024b] with momentum. we ignore the incorporated momentum, Muon computes
1 W|W Wallep 1
W e++ argmin
(G,, W W.) + 2mp-
(2)
W W

---

preserving the matrix structure of the problem and naturally exploiting the structural properties because of the involvement of spectral norm in the steepest descent framework. This aligns with ASGO's exploitation of the structured properties of the gradient. Moreover, if we ignore the mo- mentum in both the gradient and the preconditioner, we can see that ASGO is equivalent to Muon. Interestingly, this equivalence between ASGO and Muon in both theoretical basis and algorithm design is analogous to that between diagonal AdaGrad and SignSGD [Bernstein et al., 2018, Kun- stner et al., 2023]. In this sense, we may interpret ASGO and Muon AdaGrad and SignSGD for structured gradients, respectively. Note that Shampoo also admits such relation with Muon in algorithmic design. However, discussed in Section 4, Shampoo has worse convergence rate and thus fails to benefit from the structured properties as does ASGO. To some extent, this means that ASGO may be more appropriate approach than Shampoo such an analog to diagonal AdaGrad.
From this intuition, we may also conjecture what the nonconvex convergence rate of ASGO is based on what it is for Muon. Hence, we prove in Appendix G that this rate for Muon is:
min IVS(W S O (IEI.. (r(W)) (1I... ( 0SIST-1
T
where we assume that f inf f(W) > -o0. Since diagonal AdaGrad and SignSGD have the same convergence rate in nonconvex settings up to logarithmic factors [Bernstein et al., 2018, Sun et al., 2023, Liu et al., 2024], we expect that ASGO has a nonconvex convergence rate comparable to (3) obtained by Muon up to logarithmic factors. It is an interesting future topic to prove this conjecture and further explore the nonconvex behavior of ASGO theoretically. Also, we note that as an intuitively smoother version of Muon, ASGO has good convergence properties in nonsmooth settings as shown in Section 4, where Muon, like SignSGD [Bernstein et al., 2018], may fail to converge [Karimireddy et al., 2019].
Remark 4. In smooth settings, nonconvex convergence analysis of Muon has been presented in Li and Hong [2025]. These convergence results are established under more general setling, involving gradient noise and momentum following the analysis in Cutkosky and Mehta [2020]. However, if we only look at the deterministic case, their result is worse than (3) because of the explicit dependence on the dimension n. The key here is that the standard smoothness condition with respect to the Frobenius norm is not a good fit for analyzing structured gradient al algorithms like Muon or ASGO. Using Assumption 2, we obtain better convergence results for Muon in (3).
Practical Implementations of ASGO. We also provide Algorithm 2, a practical implementation of ASGO, together with a memory-efficiet diagonal variant of ASGO, named DASGO in Algorithm 4. DASGO is basically light-weigh optimizer with the preconditioner A, of ASGO being diagonalized, which is efficient in both memory and computations. To make ASGO effective and efficient in practice, critical aspect is the efficient computation of the preconditioner's 's inverse square root, V. V/ Standard methods such as Singular Value Decomposition (SVD) are prohibitively expensive for large-scale applications. To ensure our optimizer is practical, we adopt the Newton-Schultz (NS) iterative method as summarized in Algorithm 3, which substantially reduces the wall-clock time per step. The applicability of this approach stems from the deep connection between the matrix inverse square root and the matrix sign function. As established in Thm 5.2 [Higham, 2008b], computing A A1/ is equivalent to extracting a sub-block from the matrix sign function of an augmented matrix:
msign ( 0 ) = 1-1/2
0
1
This equivalence implies that the computational complexity of our inverse square root step is on the same order as the matrix operation in the Muon optimizer. It is worth noting, however, that the NS iteration in ASGO is applied to this 2 x 2 block matrix, leading to a slightly heavier per-iteration cost compared to Muon. Finally, we highlight that our framework is modular. Recent advances [Amsel et al., 2025, Grishina et al., 2025] compute the matrix sign function more efficiently than classical Newton-Schults for instance, by optimizing the iteration's hyperparameters. This modularity allows such methods to be readily incorporated. In our empirical implementation, we validated this by experimenting with both the classical Newton-Schultz and PolarExpress [Amsel et al., 2025] methods. Please refer to Appendix B for more details of other modifications in our practical implementation.
A proof of this equivalence can be found in Appendix G.
7

---

7 Empirical Results
We empirically evaluated the effectiveness of ASGO (Algorithm 2) and DASGO (Algorithm 4) on pretraining and finetuning tasks for Large Language models (LLMs). We compared our methods against the established optimizers, AdamW [Kingma and Ba, 2014, Loshchilov and Hutter, 2017], Shampoo [Gupta et al., 2018], and Muon [Jordan et al., 2024]. Several important implementational details should be noted for fair comparison. First, since Muon is designed to operate exclusively on matrix parameters, we followed Jordan et al. [2024] and applied AdamW update rules to all ID parameters within the Muon optimizer to ensure that it can handle the complete model. All experiments were conducted using NVIDIA V100s SMX2 and NVIDIA GH200 GPUs. Specifically, for the larger-scale pretraining of GPT2, we utilized a configuration of four GH200 GPUs, while other experiments were performed on single V100 GPU.
7.1 Pretraining GPT2
To further evaluate the efficacy of ASGO and DASGO, we extended our investigation to larger-scale pretraining tasks. We adopted the configuration of the GPT2 model as described by [Karpathy, 2022], comprised of 12 Transformer blocks, each with hidden dimension of 768. All 1D parameters like bias and layer normaliztion and the embedding layer were trained by AdamW. There are 12 projection layers with dimensions Shampoo needs to store two preconditioner matrices of size 768x x 768 and size 2304 x 230. In contrast, ASGO requires storing only one preconditioner matrix of size 768 768. This substantial difference highlights ASGO's advantages in memory requirements compared with Shampoo. We trained the model for 2400 steps using a batch size of 64, 4 H100 GPUs, sequence length of 512, and gradient accumulation steps. This setup corresponds to a total token budget of: 512 x 64 x 4 x 8 x 2400 & 2.5 Billion tokens. This budget represents approximately 20 tokens per parameter
in the Transformer blocks. For each of these layers alone,
suggested by the Chinchilla scaling laws [Hoffmann et al., 2022]. We used both the PolarExpress (PE) and Newton- Schultz (NS) algorithms to compute the inverse square root of in ASGO. The parameters that we used in these algorithms can be found in Section C. To ensure a fair comparison, we carefully tuned the learning rates and 82 values (where applicable) for all optimizers. (Tuning details are provided in Section C.) Furthermore, we employed learning rate schedule consisting of 240 linear warm-up steps followed by a cosine decay, which is common practice for training LLMs. Figure and Table 1 present comparison of the training and validation loss achieved by ASGO and DASGO against Shampoo, Muon and AdamW.
Table 1: Pretraining GPT2 Train loss and Validation Loss
Optimizer
Training Loss
Validation Loss
A
3.442
3.455
Muon
3.332
3.342
Shampoo
3.405
3.417
DASGO (Ours)
3.497
3.509
(Ours)
3.335
3.347
ASGO-PolarExpress (Ours)
3.332
3.342
We now present the performance evaluation of ASGO against several baselines on the GPT-2 pretraining task. The final training and validation losses after 2400 steps are summarized in Table 1, and the full training loss dynamics are depicted in Figure 1. Although the results obtained using ASGO-PE and ASGO-NS are similar, ASGO-PE performs slightly better than ASGO-NS. These results clearly demonstrate that ASGO achieves state-of-the-art performance, matching the strongest baseline while significantly outperforming AdamW and Shampoo. Focusing on the final validation loss in Table 1, ASGO--PlarExpres, achieves a validation loss of 3.342. This result is on par with the Muon opti- mizer (3.342). More importantly, ASGO's performance is substantially better than both the standard AdamW (3.455) and Shampoo (3.417).
Figure 1: Pretraining GPT-2 Train Loss.
8

---

The training curves in Figure/ I provide deeper insights into the training dynamics. ASGO and Muon exhibit nearly identical training trajectories; both converge faster than AdamW. In contrast, while Shampoo displays loss spikes during the training process. 'Tii phenomenon suggests potential numerical instability even with preconditioner update frequency I, a behavior not observed ASGO.
These findings strongly support the efficacy of ASGO. It not only matches the final performance and rapid convergence of the highly-tuned Muon optimizer but also remedies the training instabilities observed in Shampoo.
However, DASGO (3.509) underperformed AdamW, Muon, Shampoo and ASGO in the GPT2 pretraining task. This performance gap was also observed in the pretraining of the smaller Na model. A primary reason for the difference between DASGO and ASGO likely stems from its diagonal preconditioning. By only retaining and using only the diagonal elements of the GG", DASGO essentially disregards the inter-dependencies between different parameter gradients from neurons within a layer, moving away from a true matrix- adaptive approach towards a per- parameter scaling akin to vector-based methods. This highlights the value of preserving non-diagonal elements in the preconditioner matrices, particularly for capturing parameter interactions in attention- based architectures. Nevertheless, DASGO's significantly reduced memory footprint makes it compelling option for resource-limited settings where computational efficiency is paramount.
7.1.1 Wall time comparison for pretraining GPT2
We now compare the Wall-Clock Times (WCTs) per training step for Adam W, Muon and ASGO, using the exact experimental setup from the GPT-2 pretraining task (Section 7.1).
Table 2: Wall-clock time per training step, difference compared with AdamW update.
Optimizer
Single Train Step WCT
Time Difference
AdamW
5,451 sec
sec
Muon
5.519 sec
0.068 sec
ASGO (Ours)
5.947 sec
0,468 sec
For this analysis, we used 10 iterations of the NS algorithm for ASGO instead of iterations used for Muon's 's update, which are the settings required to achieve the optimal convergence performance reported previously. The measured WCT for a single training step (including forward, backward, gradient accumulation passes, and optimizer steps) is shown in Table 2. Although computing the inverse square root of Vg, in ASGO's optimizer step is more expensive than in Adam W, the forward and backward passes remain the dominant computational bottleneck. As shown in Table 2, full training step with ASGO costs only 0.468 seconds more than the A update. This overhead (8.6% increase over Ad is acceptable, and consequently, the total end-to-end training time is only marginally affected.
These tables show that although AdamW is computationally cheapest at the optimizer level, the total training times for all three methods are very similar. This demonstrates that the total training time may not be main concern for ASGO. Note that ASGO achieves the smallest training loss while AdamW achieves the largest loss.
7.1.2 Ablation Study on the preconditioning side:
In Section 5, we theoretically motivated our choice of applying the single-sided preconditioner to the right side of the gradient (or its momentum). based on block-wise smoothness convergence results. To empirically validate this crucial design choice, conducted ablation study comparing the performance of ASGO with left-sided versus right-sided preconditioning. The experimental setup is identical to that used for the GPT-2 pretraining in Section 1. As presented in Table 3, the results demonstrate that right-sided preconditioning consistently and significantly outperforms its left-sided counterpart across all tested learning rates (15e 4, le 3, 5e 3, le 21).
This finding empirically confirms our theoretical intuition and highlights the clearly asymmetric impact of the preconditioning side. This observed asymmetry also suggests potential limitation of Shampoo-like double-sided preconditioners. Beyond their higher memory and computational costs, their symmetric approximation (derived from a Kronecker product) may fail to capture the

---

true underlying structured curvature. This could lead to misplacing the preconditioning effect on the less critical side, resulting in less effective preconditioning We hypothesize that this structural mis- approximation could be contributing factor to the sub-optimal performance and training instabilities (e.g., the loss spike in Figure 1) observed Shampoo.
Table 3:
Performance comparison between precondition side
Learning Rate
Validation Loss of left side
Validation Loss of right side
5e-4
3.582
3,569
le-3
3,468
3,447
5e-3
3.399
3.362
le-2
3.412
3.386
7.2 Finetuning GPT2-Large WikiText-2
To further assess the performance of our proposed optimizers, we conducted fine-tuning experiments on the GPT2-Large(774NN) model [Radford et al., 2019] using the WikiTexx-2 dataset. We explored two distinct fine-tuning objectives:
First, we fine-tuned GPT2-Large using the standard Causal Language Modeling (CLM) loss, which is the conventional approach for autoregressive language models.
Second, we also fine-tuned the same GPT2-Large model on WikiText-2 but employed the Fill-in- the-Middle (FIM) training objective, following the setting in [Bavarian et al., 2022]. The FIM objective modifies the training process to enable the model to learn to infill text by rearranging document spans.
To account for statistical variability, we conducted five experimental runs with different random seeds under the same hyperparameter settings. Table presents the average perplexity results after fine-tuning for 2 epochs. For both objectives, we employed cosine decay learning rate scheduler. Hyperparameter search details and 95% confidence intervals provided in Table 8 in Appendix C3. Under the CLM objective, ASGO (13.88 perplexity) and DASGO (13.84 perplexity) achieved lower perplexity than both Muon (13.91 perplexity) and AdamW (14.01 perplex- ity). This suggests that both ASGO and its memory-efficieee variant, DASGO, fective for traditional LLM fine-tuning. When the model was fine-tuned using the FIM objective, we observed similar trend, with ASGO (15.66 perplexity) and DASGO (15.45 perplexity) outperform- ing both Muon (15.93 perplexity) and AdamW (17.46 perplexity). Across both fine-tuning scenarios, ASGO and DASGO demonstrated competitive or superior performance compared these baseline optimizers.
Table 4: Finetuning Large Perplexity
AdamW
Muon
ASGO
DASGO
CLM
14.01
13.91
13.88
13.84
FIM
17.46
15.93
15.66
15.45
be ef-
8 Conclusions
In this paper, we proposed novel algorithm ASGO, which achieves significantly better convergence rates compared to full-matrix AdaGrad and Shampoo. Based on the theory, we demonstrated that ASGO can benefit from low-rank gradients and block-wise diagonal Hessians, which are widely observed structural properties of DNNs. We further proposed some practical modifications to ASGO, and verified its effectiveness empirically. Currently, ASGO still has two major limitations: 1) compu- tationally intensive compared to Adam because of the matrix operation; 2) it is not straightforward to extend the algorithm to apply to tensors. We plan to look into these issues future work.
10

---

Acknowledgment
Shiqian Ma was partially supported by ONR grant N00014-24-1-2705, NSF grants CCF-2311275 and ECCS-232659) This work was also partially supported by NSF grant No. 2416897 and ONR grant No. NO001425 42,12218, and used both the DeltaAl advanced computing and data resource, supported by NSF(award OAC 2320345) and the State of Illinois, and the Delta advanced computing and data resource, supported by NSF (award OAC 2005572) and the State of Illinois.
References
Noah Amsel, David Persson, Christopher Musco, and Robert M Gower. The polar express: Optimal matrix sign methods and their application to the muon algorithm. arXiv preprint arXiv:2505.1693 2025.
Tsuyoshi Ando, Chi-Kwong Li, and Roy Mathias. Geometric means. Linear algebra and its applications, 385:305-334, 2004.
Achraf Bahamou, Donald Goldfarb, and Yi Ren. A mini- biokkblock fisher method for deep neural networks, 2022. URL https: h org/ abs abss/202..e 04124.
Mohammad Bavarian, Heewoo Jun, Nikolas Tezak, John Schulman, Christine McLeavey, Jerry Tworek, and Mark Chen. Efficient training of language models to fill in the middle, 2022. URL atttp://arxiin www.aaawww..iioaaaaaaaaaaaaea hwtgg/bs//22/ 2207 14255.
James Bergstra and Yoshua Bengio. Random search for hyper-parameter optimization. Journal of Machine Learning Research, 13(10):811335, 2012. URL http://jmlm org http/apppss/v//
bergstra12a. html.
arXiv preprint Jeremy Bernstein and Laker Newhouse. Modular duality in deep learning. arXiv:2410.21265, 2024a.
Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology. arXiv preprint arXiv:2409.20325, 2024b.
Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Animashree Anandkumar. signsgd: Compressed optimisation for non-convex problems. In International Conference Machine Learning, pages 560-569. PMLR, 2018.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877-1901, Curran Associates, Inc., 2020. URL https: /ttppe///proceeddnggs neurips. cttccc//aaaa /0200// 4
David Carlson, Volkan Cevher, and Lawrence Carin. Stochastic spectral descent for restricted boltzmann machines. In Artificial intelligence and statistics, pages 111-119. PMLR, 2015a.
David Carlson, Ya-Ping Hsieh, Edo Collins, Lawrence Carin, and Volkan Cevher. Stochastic spectral descent for discrete graphical models. IEEE Journal of Selected Topics in Signal Processing. 10 (2):296-311, 2015b.
Lesi Chen, Chengchang Liu, and Jingzhao Zhang. Second-order min-max optimization with lazy hessians. arXiv preprint 2xxiv22410.09568, 2024.
Dami Choi, Christopher J Shallue, Zachary Nado, Jaehoon Lee, Chris J Maddison, and George E Dahl, On empirical comparisons of optimizers for deep learning. arXiv preprint arXiv: 10.0444,, 2019.
Ronan Collobert. Large scale machine learning. Technical report, Université de Paris VI, 2004.

---

Romain Cosson, Ali Jadbabaie, Anuran Makur, Amirhossein Reisizadeh, and Devavrat Shah. Low- rank gradient descent. IEEE Open Journal of Control Systems, 2:380-395, 2023.
Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang. Robustness to unbounded smoothness of generalized signsgd. Advances in neural information processing systems, 35:9955-9996 9968,, 2022.
Ashok Cutkosky and Harsh Mchta. Momentum improves normalized sgd. In International conference on machine learning, pages 2260-2268. PMLR, 2020.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: weto:088000000 2018.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning reseurch, 12(7), 2011.
Runa Eschenhagen, Aaron Defazio, Tsung-Hsien Lee, Richard E. Turner, and Hao-Jun Michael Shi. Purifying shampoo: Investigating shampoo's heuristics by decomposing its preconditioner, 2025. URL https: /arxxvv.vvv/v/////////a///aaaav/ org/ abs/2506. 03595.
Vladimir Feinberg, Xinyi Chen, Y Jennifer Sun, Rohan Anil, and Elad Hazan. Sketchy: Memory- efficient adaptive regularization with frequent directions. Advances in Neural Information Process- ing Systems, 36:7591 1-75924, 2023.
Guillaume Garrigos and Robert M Gower. Handbook of convergence theorems for (stochastic) gradient methods. arXiv preprint arXiv:2301.11233, 2023.
Ekaterina Grishina, Matvey Smirnov, and Maxim Rakhuba. Accelerating newton-schulz iteration for orthogonalization via chebyshev-type polynomials. arXiv preprint arXiv:2506. 10935, 2025,
Vineet Gupta, Tomer Koren, and Yoram Singer. A unified approach to adaptive regularization in online and stochastic optimization. arXiv preprint arXiv:/706.06550 2017.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimiza- tion. In International Conference on Machine Learning, pages 1842-1850, PMLR, 2018.
Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in tiny subspace. arXiv preprint arXiv:/8/2.0477: 2018.
Elad Hazan, Amit Agarwal, and Satyen Kale. Logarithmic regret algorithms for online convex optimization. Machine Learning, 69:169-192, 2007.
Nicholas J. Higham. Functions of Matrices: Theory and Computation. Society for Industrial and
Applied Mathematics, United States, 2008a. ISBN 9780898717778. ISBN 1778-0-898871--4-7.
Nicholas J. Higham. Functions of Matrices: Theory and Computation. Society for Industrial and Applied Mathematics, Philadelphia, PA, USA, 2008b. ISBN 978-0-898716-46--7, doi: 10.1 117/1 97818987777877.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 2022.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, al. Lora: Low-rank adaptation of large language models. ICLR, 1(2):3, 2022.
Qiushi Huang, Tom Ko, Zhan Zhuang, Lilian Tang, and Yu Zhang. Hira: Parameter-efficient hadamard high-rank adaptation for large language models. In The Thirteenth International Conference Learning Representations, 2025.
Ruichen Jiang, Devyani Maladkar, and Aryan Mokhtari. Convergence analysis of adaptive gradient methods under refined smoothness and noise assumptions. arXiv preprint 240.22406604592. 2024a.
12

---

Ting Jiang, Shaohan Huang, Shengyue Luo. Zihan Zhang, Haizhen Huang. Furu Wei, Weiwei Deng. Feng Sun, Qi Zhang, Deqing Wang, al. Mora: High-rank updating for parameter-efficien fine-tuning. arXiv preprint arXiv:2405. 12130, 2024b.
Keller Jordan, Yuchen Jin, Vlado Boza, You Jiacheng, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https:
//kellerjordan github. io/iio/pooe http//sss///eeeo haaaa//mm////..//
Sai Praneeth Karimireddy, Quentin Rebjock, Sebastian Stich, and Martin Jaggi. Error feedback fixes signsgd and other gradient compression schemes. In International Conference on Machine Learning, pages 3252-3261. PMLR, 2019.
Andrej Karpathy. NanoGPT. https: htwph///..cc.aaaaaatttaaaaooo.. 2022
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint 0rX6v.1412.6980, 2014.
Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt, Noise not main factor behind the gap between sgd and adam on transformers, but sign descent might be. arXiv preprint arXiv:2304.13900 2023.
Tim Large, Yang Liu, Minyoung Huh, Hyojin Bahng, Phillip Isola, and Jeremy Bernstein. Scalable optimization in the modular norm. arXiv preprint arXiv:2405. 14813, 2024.
Kfir Y Levy, Alp Yurtsever, and Volkan Cevher. Online adaptive methods, universality and accelera- tion. Advances in neural information processing systems, 31, 2018.
Jiaxiang Li and Mingyi Hong. A note on the convergence of muon and further. arXiv preprint arXiv:2502. 02900, 2025.
Vladislav Lialin, Namrata Shivagunde, Sherin Muckatira, and Anna Rumshisky. Relora: High-rank training through low-rank updates. arXiv preprint arXiv:2307.05665, 2023.
Wu Lin, Felix Dangel, Runa Eschenhagen, Juhan Bae, Richard E. Turner, and Alireza Makhzani. Can remove the square-root in adaptive gradient methods? a second-order perspective, 2024. URL tistps://arxiie org/ abaaaba///222 03496.
Jingyuan Liu, Jianlin Su, Xingeheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for Ilm training. arXiv preprint ar arXiv22502.1688. 2025a.
Liming Liu, Zhenghao Xu, Zixuan Zhang, Hao Kang, Zichong Li, Chen Liang, Weizhu Chen, and Tuo Zhao. Cosmos: A hybrid adaptive optimizer for memory-efficee training of llms. arXiv preprint arXiv:2502.17410 2025b.
Yuxing Liu, Rui Pan, and Tong Zhang. Adagrad under anisotropic smoothness. arXiv preprint arXiv:2406.155 15244, 2024.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv: 2017.
Karl Löwner. Über monotone matrixfunktionen. Mathematische Zeitschrift, 38(1):177-216, 1934.
James Martens and Roger Grosse.
Optimizing neural networks with kronecker-factored approximate curvature. In International conference on machine learning, pages 2408-2417. PMLR, 2015.
Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham Kakade, and Lucas Janson. A new perspective on shampoo's 's preconditioner. arXiv preprint arXiv:2406. /7748, 2024.
Yurii Nesterov et al. Lectures convex optimization, volume 137. Springer, 2018.
Son Nguyen, Bo Liu, Lizhang Chen, and Qiang Liu. Improving adaptive moment optimization via preconditioner diagonalization. arXiv preprint arXiv:2502.07488; 2025.
13

---

Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang. Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35: 27730-27744, 2022.
Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falks Falls, and Volkan Cevher. Training deep learning models with norm-constrained Imos. arXiv preprint arXiv:2502.07529, 2025.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8)99 2019.
Yi Ren and Donald Goldfarb. Tensor normal training for deep learning models, 2021. URL org/ awwaaasss/o// 02925..
Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint arXiv:1611.07470 2016.
Levent Sagum, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over-parametrized neural networks. arXiv preprint 1rXiv:1706.0445. 2017.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost.
In International Conference on Machine Learning pages 4596-4604. PMLR, 2018.
Hao-Jun Michael Shi, Tsung-Hsien Lee, Shintaro Iwasaki, Jose Gallego-Posada, Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, and Michael Rabbat. A distributed data-parallel pytorch im- plementation of the distributed shampoo optimizer for training neural networks at-scale. arXiv preprint arXiv:2309. 209443309 2023.
Matthew Streeter and H Brendan McMahan. Less regret via online conditioning. arXiv preprint arXi::1002.4862, 2010.
Tao Sun, Qingsong Wang, Dongsheng Li, and Bao Wang. Momentum ensures convergence of signsgd under weaker assumptions. In International Conference on Machine Learning, pages 33077-33099, PMLR, 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. /3971, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023b.
Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam. arXiv preprini 13X11 2024.
Hongyi Wang, Scott Sievert, Shengchao Liu, Zachary Charles, Dimitris Papailiopoulos, and Stephen Wright. Atomo: coammunication-effiicee learning via atomic sparsification. Advances in neural information processing systems, 31, 2018.
Yikai Wu, Xingyu Zhu, Chenwei Wu, Annie Wang, and Rong Ge. Dissecting hessian: Understanding common structure of hessian in neural networks. arXiv preprint arXiv:2010.04261, 2020.
Shuo Xie, Mohamad Amin Mohamadi, and Zhiyuan Li. Adam exploits f -geometyy of loss landscape via coordinate- wise adaptivity. arXiv preprint arXiv:2410.08199. 2024.
Shuo Xie, Tianhao Wang. Sashank Reddi, Sanjiv Kumar, and Zhiyuan Li. Structured preconditioners in adaptive optimization: A unified analysis. arXiv preprint 1rXiv:2503.10533 2025.
Yuege Xie, Xiaoxia Wu, and Rachel Ward. Linear convergence of adaptive stochastic gradient descent. In International conference on artificial intelligence and statistics, pages 1475-1485.
PMLR, 2020.
14

---

Greg Yang, James B Simon, and Jeremy Bernstein. A spectral condition for feature learning. arXiv preprint wrXiv:2310.178111 2023.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint ar04: 00962, 2019.
Huizhuo Yuan, Yifeng Liu, Shuang Wu, Xun Zhou, and Quanquan Gu. Mars: Unleashing the power of variance reduction for training large models. arXiv preprint arXiv:2441. 2024.
Yushun Zhang, Congliang Chen, Tian Ding. Ziniu Li, Ruoyu Sun, and Zhi-Quan Luo. Why trans- formers need adam: hessian perspective. arXiv preprint arXiv:2402.16788, 2024a.
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Diederik P Kingma, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more. arXiv preprini arXiv:2406. 16793, 2024b.
Jiawei Zhao, Florian Schäfer, and Anima Anandkumar. Zero initialization: Initializing neural networks with only zeros and ones. arXiv preprint arXiv:2110. 12661, 2021.
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficiet IIm training by gradient low-rank projection. arXiv preprini arXiv.2403.03507, 2024.
Hanqing Zhu, Zhenyu Zhang, Wenyan Cong, Xi Liu, Sem Park, Vikas Chandra, Bo Long, David Z Pan, Zhangyang Wang, and Jinwon Lee. Apollo: Sgd-like memory, adamw-level performance. Proceedinggs of Machine Learning and Systems, 7, 2025.
Martin Zinkevich. Online convex programming and generalized infinitesimal gradient ascent. In Proceedings of the 20th international conference on machine learning (icml-03), pages 928-936, 2003.
NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We claim our main contribution both in abstract and summerized in the end of introduction part.
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 the paper and important assumptions and limitations. A No or
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 the work performed by the authors?
Answer:
[Yes]
Justification:
We discuss the limitations of the proposed algorithm in the conclusion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
15

---

The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, 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 on the scope of the claims made, c.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should 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 is low or images are taken in low lighting. Or speech-to-texl 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 fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, 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 of assumptions and a complete (and correct) proof?
Answer:
[Yes]
Justification:
We provide proofs and assumptions for all the theorems.
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 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
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We clearly clarify all details including experiment setting and hyperparameter tuning method in Appendix Section C which make the experiments reproducible.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
16

---

If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a 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 to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., . a large language model), then there should either be way to access this model for reproducing the results or a 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 tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g.. registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We plan to make the code publicly available in GitHub repository upon publication. All datasets used are publicly available and cited, with setup details in Appendix C.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / //ipps.c.. wwwpcc/ppc//puuecceeeecceeeeoooe
CCoaae//dddSSbbbiaaa wwtuuuimmsssooii wssssssnools. insiinnnoccy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
17

---

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: All training and test details, including model setting, hyperparameter selection, and optimizer configurations, are specified in Appendix C.
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 appreciate the results and make sense of them.
The full details can 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 or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: To account for statistical variability, set the fixed random seed report mean training performance among difference random seeds to verify our proposed algorithm performance.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars 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).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but 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 to show in tables or figures symmetric error bars that would yield results that are out of range (e.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 on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We report the computer resources we use Section 7.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
18

---

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 or 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 NeurIPS Code of Ethics https:/ h/tppp/////uuis.../// cctppclccppeeeeececcccccccccccet wttp//hhhhhhheeeeeeeeeeeeeeeeeee
Answer: [Yes]
Justification: Our research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper proposes a foundational optimization algorithm aimed at improving the training efficiency and effectiveness of deep neural networks. This work does not introduce direct applications or systems with immediate societal consequences beyond the general implications of advancing ML research.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact 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), fairesss 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 tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 to 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 unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
19

---

This paper proposes a foundational optimization algorithm aimed at improving the training efficiency and effectiveness of deep neural networks.
Justification:
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 to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
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: All existing assets, such as datasets, baseline model architectures, and opti- mizers are properly credited via citations in the paper. Key open-source projects utilized include Karpathy's NanoGPT (Karpathy [2022], MIT License) and the specific version of Muon code we adapted from wtt////lleeooooraa//iiiioic
(Jordan al. [2024], MIT License). For finetuning GPT-2 (Radford et al. [2019].. we utilized example scripts from the Hugging Face Transformers library (which Apache 2.0 licensed). The datasets are standard benchmarks with established public licenses, and we have respected their terms of use.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms 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-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's ' creators.
13. New assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: An anonymized version of our code and experimental scripts is provided the supplementary material. This includes initial documentation (c.g., README with setup instructions and guidance for running key experiments), which will be further detailed for the full public release.
Guidelines:
The answer NA means that the paper does not release new assets.
20

---

Researchers should communicate the details of the dataset/codeemmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset 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 involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethies, 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
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 an equivalent approval/review based the requirements of your country or institution) were obtained?
Question:
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 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 we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: Core method development in this research does not involve LLMs as any important, original, non-standard components. We only used an LLM for assistance with writing, editing, or formatting purposes, and this usage does not impact the core methodology, scientific rigor, or originality of the research.
Guidelines:
21

---

The answer NA means that the core method development in this research does not involve LLMs as any important. original, non non-standard components.
Please refer to our LLM policy (https : /ttp///ee.ppp/eeeee cc/Conferences//220 /LLM) for what should or should not be described.
22

---

A Additional Related Work
Optimization with Matrix Structure. Much recent work has focused on improving full-matrix AdaGrad and Shampoo. Feinberg et al. [2023] uses sketching-bane based approach to approximate the full-matrix AdaGrad preconditioner with lower memory cost. Morwani al. [2024] provides theoretical intuition and empirical evidence to claim that Shampoo should use the - -1/2 power in its preconditioners to better approximate the Empirical Fisher (EF) matrix. Vyas et al. [2024] demonstrates that Shampoo is like doing Adafactor in the eigenspace of gradients and proposes a novel algorithm, SOAP, that performs Adam this eigenspace. SOAP is observed to achieve better performance than Adam and Shampoo, but suffers from high computation load per iteration. Galore [Zhao et al., 2024] and Apollo [Zhu et al., 2025] share a similar algorithmic design with SOAP, with extra focus on lowering memory costs. Muon [Jordan et al., 2024] follows this line of work, derived from the steepest descent (in the spectral norm) framework, which it shows to be scalable and effective in training large foundation models [Liu al., 2025a]. This spectral descent method has its origin in restricted Boltzmann machines and discrete graphical models [Carlson et al., 2015a, Large et al. [2024], Bernstein and Newhouse [2024a] propose the modular approach that has also been applied to improve Muon. Following Muon, Pethick et al. [2025] proposes the linear minimization oracle (LMO) framework and derives Scion based on it, showing good convergence hyperparameter transfer properties. More recently, Nguyen et al. [2025] proposes AdaDiag, which may be viewed as SOAP doing SVD without gradient accumulation. Liu et al. [2025b] proposes COSMOS, combination of SOAP and Muon, that trades off between performance and computational efficiency.
Rank of Gradients and Weight Updates. It has been widely observed that gradients are naturally low-rank in DNNs, even when a large batch size is employed [Gur-Ari et al., 2018, Zhao et al., 2021, Yang et al., 2023, Cosson al., 2023]. This property has been widely utilized for computation and memory efficiency in training [Wang ct al., 2018, Cosson et al., 2023, Zhao et al., 2024]. On the other hand, the rank of the total weight update AW WT - Wo. depends lot on the training method. If we use LoRA [Hu et al., 2022], AW is determined to be low-rank. However, in pretraining or even many complex fine-tuning tasks, LoRA's performance is much worse than methods that produce high-rank weight updates like full- hull-parameter training. which is conjectured to be due to the weight update rank [Lialin et al., 2023, Jiang et al., 2024b, Huang et al., 2025].
03
85
150
297
(a) Hcssian of Query (4 heads)
(b) Hessian of Key (4 heads)
(c) Hessian of Value (4 heads)
Figure 2: This figure from Zhang al. [2024b]. It depicts the Hessian of different parameter blocks in a small Transformer at the 1% training step. The bear-block-diagona structure maintains throughout training. But different parameter blocks have different numbers of small dense matrices, where Query and Key correspond to the number of heads.
Block-wise Diagonal Hessian. It has been observed that the Hessian of neural network tends to be block-wise diagonal with each block corresponding to neuron both in experiments and theory for small MI.Ps [Collobert, 2004]. Zhang et al. [2024a.b] recently numerically verified this property in small transformers, and as illustrated in Figure 3. further empirically showed that transformers may exhibit heterogeneity between blocks, while CNNs may not.
Concurrent Work. When we were finishing writing this paper, we noticed that the paper [Xie et al., 2025], which had just appeared on arXiv, proposes and studies an algorithm, referred to as One-Sided Shampoo, that is identical to ASGO. Although the theoretical techniques and convergence results proved in [Xie et al., 2025] and here are very similar general, there are still some notable differences between the two works. First, the motivations are different. In Xie et al. [2025], the
23

---

(a) VGG16
(b) VIT base
(c) GPT2
Figure 3: This figure is from Zhang et al. [2024a,b], calculating the Jensen-Shannon (JS) distance
between two eigenvalue densities of all possible block pairs at initialization. shows that JS distance of blockwise spectra in CNNS is significantly smaller than that in Transformers.
authors develop their method from the unified preconditioning method framework AdaptReg [Gupta al., 2017] and mainly highlight its superior convergence results compared to Shampoo and full- matrix AdaGrad. In our paper, we focus more theoretically discussing how ASGO can utilize the structured properties of optimization problems including low-rank gradients and block-wise diagonal Hessians, to highlight the potential of ASGO as practical algorithm for training deep learning models. Second, our empirical results provide evidence that ASGO can perform well in practical tasks, whereas Xie et al. [2025] focuses on convex settings and examines One-Sided Shampoo only on linear regression. Moreover, our implementation includes specialized designs for Transformer architectures (particularly query/key) improve performance deep learning training tasks. As direct byproduct of our main algorithm, we also developed lightweight diagonal version named DASGO to trade off memory consumption and performance. Third, Xie et al. [2025] focuses primarily on developing a general proof framework applicable to multiple optimizers including One-Sided Shampoo, whereas work specifically examines the theoretical and practical benefits of ASGO. To conclude, both works contribute to better understanding of ASGO/Onne Shampoo.
B More Discussions on ASGO
practical implementation of ASGO. We present practical implementation of ASGO in Algorithm 2, with the input W as module in the network. Our implementation incorporates several common modifications to enhance computational efficiency and stability.
Algorithm 2 A Practical Implementation of ASGO (for layer W E mx such that m n)
1: Input: Wo E mm Ir schedule {mt momentums 31, B2 [0, 1), batch size M E -, e E R (E
should be small, similar to the c for Adam.)
2: Initialize M-1 = 0 E m V-1 = 0 E Rax
D If m n, then we set V-1 = 0 mxm
3: for = 0 to T do
4: Sample mini-batch B, with |B, = M uniformly
5: G. = M7 ΣeES, WwJ
6: Mr = M1 + B.
7:
Ve 32V2-1 + (1- G2 G.
D If <,, V, =00=-/VV.. (1- Run Alg. 3 to compute Vn, square root inverse D If < n, W411 We 0.2/mmaA!" || dielr
8:
AjaV SprrtivversellwtooSScuukee
M.
9: W41 W. U) MM lip
end for
(i) Exponential moving averages used to update the preconditioner and incorporate momentum, following established practice, which results in Adam's improvements across various tasks over AdaGrad.
(ii) As demonstrated in Section 6, we use an efficient numerical algorithm (Algorithm 3) compute the matrix square root for ASGO, which is mathematically equivalent to applying the Newton-Schul
iteration to the 2 x block matrix )
A
The parameters a, b, c can be chosen as established
0
24

---

Algorithm 3 SqrttveerseNewtooSSShuul
1: Input: X E _=** parameters a, e (can be chosen in ways like Jordan et al. [2024], Amsel et al. [2025], Grishina et al. [2025]), step number K,
2: Obtain normalized input & X with ax = | +
ax'
3: Initialize identity matrix Z t- In E R
4: for k
K do
5: A < ZY
6: B bA cA2
7:
Y aY + YB
8:
2+ aZ + BZ
9:
end for
10:
return Vax
Return the approximation of X-1, and vax Y an approximation of X
in Jordan et al. [2024], Amsel et al. [2025], Grishina et al. [2025]. In each of the Newton-Schulz iteration, we perform 4 matrix multiplications, while the Muon Newton-Schull iteration requires 3.
(iii) We adaptively select the preconditioning side. Instead of fixing the preconditioner to one side, we dynamically choose to precondition on the side corresponding to the smaller dimension of the gradient matrix G, in Rmx" Specifically, if m < n. the preconditioner is formed from G,G, and applied from the left; otherwise (if m > n), it is formed from G? G2 and applied from the right. Based on the intuition interpreted in Section 5, we also apply the preconditioner from the right when m ni.
(iv) We incorporate a similar update alignment to ASGO, as presented in line 9 of Algorithm 2. Specifically, we normalize the update of ASGO for each matrix weight W by the root mean square norm (RMS norm) and multiplies 0.2 to match the RMS norm of Adam update, as suggested by Liu al. [2025a]. In this way, the training becomes more stable and easier tune.
(v) There are more possibilities in doing practical implementation for ASGO. For example, an alternative approach to improve the computational efficiency is doing matrix operations in lazy way, i... the operations only happen every several iterations, which is widely employed in previous algorithms [Shi et al., 2023, Vyas al., 2024, Zhao et al., 2024, Chen et al., 2024]. In this way, we can further lower the computational cost ASGO, but more efforts on preventing performance degradation will be needed. We may also consider incorporating the variance reduction techniques as in Yuan et al. [2024], Cutkosky and Mehta [2020].
Special Design for Transformers. Furthermore, we introduce a specialized adaptation of ASGO for query and key matrices in attention layers. Recent work by Zhang et al. [2024b] has demonstrated that the Hessian structure of query and key layers differs significantly from conventional MLP layers. As illustrated in Figure 2, the number of dense blocks corresponds to the number of attention heads rather than output neurons. This observation aligns with the forward computation of multi-head attention, where attention scores are computed independently across different subspaces, suggesting that query and key parameters could be optimized in head-wise manner. To leverage this insight, we reshape query and key parameters from matrices in R" xhd to dineesiimesional tensors Rh xxxxxd and apply our optimization algorithm independently to each head's 's subspace. This restructuring reduces both memory consumption and computational complexity by decreasing the matrix size from ( into Ohha²). For instance, in NanoGPT [Karpathy, 2022], this adjustment reduces the preconditioning size of a single query/key parameter from approximately 106 elements into 105. The empirical performance of this modification is evaluated in NanoGPT pretrain task in Section C.I.
A Diagonal Variant of ASGO. Drawing inspiration from Duchi al. [2011] we also implement a variant of ASGO using diagonal matrices, which we denote as DASGO, presented in Algorithm 4. DASGO can be viewed as lightweight version of ASGO, which eliminates the need to compute the inverse square root of full matrices, and reduces memory requirements to a level comparable with Adam-mini [Zhang et al., 2024b]. It also makes the choice of which side of M, to precondition unimportant in terms of computational effort, in contrast to ASGO. For DASGO, we choose to apply the diagonal preconditioner on the right side, aligning with the neuron architecture in DNNs. However, since DASGO only employs diagonal preconditioner, it fails to recover the superior
25

---

theoretical properties of ASGO under block-wise diagonal Hessian settings We further empirically
examine this tradeoff in Section
Algorithm 4 Implementation DASGO (Diagonal Adaptive Structured Gradient Optimization)
1: Input: Wo R= é E R. Ir schedule {me), momentums $1, 32 E [0,1), and batch size M E
2: Initialize M-1 = 0 R 8-1 0 R
3: for t = 0 to T - 1 do
4: Sample mini-batch B, with |B, = M uniformly
5: G₁ LEEs W ((W
6: M, = -81 + 11 31
7: Ur = B2 + (1 B2) diag(G; G₂)
8: W4+1 W. Mr diag(er + +e)
9: end for
C Details of Empirical Experiments
C.1 Pretraining NanoGPT
As initial experiment, compared the performance of ASGO, DASGO, Muon, Shampoo, and AdamW for training NanoGPT, which consists of Transformer layers, attention heads, and an embedding dimension of 384, the Shakespeare character-level dataset with sequence length of 256 tokens. For all algorithms, we trained the model for 20 epochs, where each epoch contained 128 steps, using a batch size 128 and the OneCycle learning rate (Ir) schedule. In contrast to our experiments on GPT2, where we trained ID and Embedding layer parameters as done in AdamW, we trained all the parameters independently in ASGO.
Before discussing hyperparameter (HP) tuming, we note that recent research [Shi et al., 2023, Morwani et al., 2024, Lin et al., 2024] has suggested treating the inverse order (IO) in Shampoo's preconditioner as a tunable HP rather than using the standard value of -1/4. Notably, Morwani et al. [2024] demonstrated that Shampoo with 10= -1 1 provides superior approximation of the full-matrix AdaGrad optimizer. Consequently, we treated the 10 as a Shampoo HP in our pretraining experiments. Additionally, Shampoo's performance is highly dependent on the initialization of its preconditioner matrices, as this is crucial for accurately approximating the Empirical Fisher Information Matrix [Morwani al., 2024]. To address this initialization sensitivity, we implemented a preconditioner warmup phase specifically for Shampoo. Following the methodology in Ren and Goldfarb [2021]. J. we dedicated the first epoch exclusively to accumulating statisties for the preconditioner matrices without updating model parameters. This approach yields more robust estimation of the preconditioner, which significantly improves the stability of Shampoo during subsequent training iterations. Figure depicts the training loss and validation loss for training NanoGPT with Shampoo, DASGO, Muon, ASGO and Adam-W.
Examining both test and training loss curves, ASGO consistently outperforms Shampoo despite requiring only half the memory consumption and computational effort, which highlights ASGO's practical advantages for training language models. Furthermore, ASGO and Muon achieve the lowest final training and test losses, outperforming all other methods. This consistent performance between ASGO and Muon aligns well with our discussions in Section 6. Moreover, The lightweight DASGO optimizer achieves competitive results against AdamW in both training and test loss metries. DASGO demonstrates particularly strong performance during the initial training phase (first 200 training steps, which is about the first two epochs), but eventually exhibits performance gap with ASGO. We discuss this performance gap in Section 7.1.
Hyperparameter Tuning For consistency across all optimization methods, we employ the OneCy- cleLR learning rate schedule, which has been shown to provide stable convergence properties in deep learning tasks. To ensure optimal performance for each optimizer, we conducted a comprehensive HP search using random search [Bergstra and Bengio, 2012, Choi et al., 2019] within the following predefined ranges: initial Ir: [10-5, 10-1]; B1: [0.7, 0.99]; :3: [0.7, 0.99]; and selected the warmup factor from the set [0.1, 0.15,,151100 0.2, 0.25, 0.3). For the optimizers Shampoo and ASGO, we selected
26

---

Mun
(a) Train
(b) Test Loss
Figure 4: Train Loss and Test Loss on the NanoGPT and Shakespeare Character Dataset
the update frequency from the set [5, 10, 15] and for Shampoo, selected its IO as either - J. HPs were selected based on validation loss after 20 epochs of training. For each optimizer, performed a random search for 16 hours to find the optimal HP combination. These best HP choices for training NanoGPT are given in Table
Table Hyperparameter selection for NanoGPT experiment
Optimizer learning rate
B1
B2
warmup factor
update frequency
inverse order
Muon
0.00349
0.9881
N/A
0.3
N/A
N/A
AdamW
0.00450
0.9332
0.9528
0.2
N/A
N/A
DASGO
0.060
0.9584
0.9435
0.2
N/A
N/A
Shampoo 0.00593
0,9402
0,9760
0.3
15
ASGO
0.01470
0.9541
0.8487
0.3
15
N/À
Ablation Study: Impact of Special Query/Key Design In Appendix B, we highlighted the compu- tational benefits of specialized processing for Query and Key matrices in Transformer architectures. To validate these theoretical advantages empirically, we conducted an ablation study comparing the performance and training stability of ASGO and Muon pretraining NanoGPT, both with and without the specialized Query-Kep Key processing. We maintained the optimal HP configuration established in Table 5 for all parameters except the learning rate. For each algorithm variant (with and without specialized processing), we randomly sampled learning rates from a log uniform distribution ranging from 10-5 to 10-1 Each configuration after which we recorded the validation loss. This approach allowed us to evaluate both performance and robustness to learning rate selection. Figure presents the distribution of validation losses after epochs for both algorithms. The violin plots* reveal several key insights: (1) For ASGO (Figure 5a), the benefits of specialized Query-Key processing are substantial. The specialized implementation demonstrates narrower, lower distribution of validation losses (centered around 2.5), indicating greater stability across different learning rates. In contrast, the vanilla implementation shows tong-tailed distribution extending beyond loss values of 6.0, with several outliers representing training instability certain learning rates. (2) For Muon (Figure 5b), both implementations demonstrate nearly identical performance distributions, with their median validation losses and distribution shapes showing minimal differences. However, the specialized implementation achieves this comparable performance while reducing memory requirements and computational complexity. This represents clear efficiency advantage the specialized Query-Key processing effectively provides computational savings with no performance penalty.
trained for 5 epochs
the NanoGPT model,
C.2 Pretraining GPT2
Experimental Setup: The model configuration consists of 12 Transformer layers, 12 attention heads, and embedding dimension of 768. The total number of training steps was 2400, with batch
VViolin plot outlines depict empirical probability density; i.e., the width of the shaded area represents the proportion of the data located there. Box plots within violin plot display the median and inter-quartile range.
27

---

AS60
Muan special
(a) ASGO
(b) Muon
Figure Distribution of validation losses after 5 epochs with varying learning rates.
size per GPU of 64. The model was trained using Distributed Data Parallel (DDP) on 4 NVIDIA GH200 GPUs, employing gradient accumulation over steps. To ensure fair comparison across all optimization algorithms, a consistent learning rate schedule was utilized: a linear warmup for the first 200 steps, followed by cosine annealing decay to final learning rate of 10 for the remainder of the training process. Training was conducted on the OpenWebTex dataset with sequence length of 512 tokens. For Shampoo optimizer, we used the Adam update grafting mentioned in [Shi et al., 2023, Eschenhagen et al., 2025]. For Muon, we used the Adam update RMS norm alignment in [Liu et al., 2025a]. We ran both the Tewton-Schulle and the PolarExpress algorithms for 10 steps to compute the inverse square root of V in the ASGO optimizer using the following parameters:
Newton-Schultz: (a, bg.bec c) (2, 11.5
PolarExpress:
coeff's list
(22822208844432233998811998111 1,.0://9999998777777700007713333
6$946008544299662999999111118422 0.51119113944999711177
11991111007773388110099922916111 0.37680408948554483333
(18875001880853447)) 4.330044997,,,30000003399999, 037700017700))00000044444444
(1.875.-1.25 0.375),
(1.875, -1.25, 0.375). (1.875, -1.25, 0.375)]
Hyperaarameter Tuning: To ensure optimal performance for each optimizer, conducted
hyperparameter search using grid search. The search space for learning rate (Ir) and 2nd momentum 32 (where applicable for the optimizer) included the following values:
Learning rate (Ir): [1x10- 1x10 5x 1x103
10 11111x11x10000001111111,xx111111
1,110=,,xx00-,,,
:2: (0.7,008,009000955009999
Optimal hyperparameters were selected based on the validation loss achieved after 1000 training steps. The best hyperparameter choices for pretraining GPT-2 are presented in Table 6.
Table 6: Optimal hyperparameter selection for pretraining GPT-2.
Optimizer Learning Rate
B2
Weight Decay
Damping e
Update Freq. (r)
Muon
0.01
N/A
N/A
N/A
AdamW
0.005
0,9
0.99
0.1
1 1 10
N/A
Shampoo
0.005
0,9
0.99
1 10-10
DASGO
0.01
0.9
0.9
0.1
1 10
N/A
ASGO
0.01
0.9
0.8
0.1
1 10
I
28

---

DA2G3 EacoPanrtarees
E/G00
4920
Adar/i
4.0
20
Shap
Smpp
(a) Train Loss
(b) Test Loss
Figure 6: Pretraining Performance GPT2
C.3 Finetune GPT2-Large
We fine-tuned the GPT-2 Large model on the WikiText-2 dataset. We utilized the GPT-2 Large model, which comprises approximately 774 million parameters. This model is transformer-basee language model pretrained on English text using Causal Language Modeling (CLM) objective, initially developed by OpenAI and was accessed via the Hugging Face Hub. The fine-tuning was performed on the WikiText-2 dataset, a standard benchmark for evaluating language model performance. The dataset was loaded directly using standard library functions. We investigated two distinct fine-tuning objectives:
Experimental Setup:
Causal Language Modeling (CLM): The conventional autoregressive language modeling task.
Fill-in-the-Middle (FIM): Adopting the methodology from Bavarian al. (2022). This objective trains the model to infill masked text spans within a document.
All models were fine-tuned for a total of epochs. For the learning rate, we employed a cosine annealing decay schedule over the course of these epochs, with the learning rate decaying by the end of training. No warmup phase was used for the fine-tuning learning rate schedule. The fine-tuning process used batch size of 16 and a sequence length of 128 tokens. For each optimizer evaluated (AdamW: W. Muon, ASGO, and DASGO), all hyperparameters, with the exception of the learning rate, were maintained at the same values used in our GPT-2 pretraining experiments detailed in Table 6).
Hyperparameter Tuning: For the fine-tuning experiments WikiTeext-2, hyperparameter tuning focused exclusively on identifying the optimal learning rate for each optimizer under both learning rate candidates: [1 x 10-6,3 x 10-5, x 10-4,3 5x x 10-4, x 10-$, 10-3). The optimal learning rate for each optimizer and fine-tuning objective combination was determined by selecting the learning rate that yielded the lowest validation perplexity on the WikiText-22 validation set after the full epochs fine-tuning Table shows the optimal learning rate.
the CLM and FIM objectives. We performed a grid search over a predefined
Table 7: Optimal Learning rate from Finetuning GPTg-Large
AdamW
Muon
ASGO
DASGO
CLM
x 10-5
5 10-4
1x 10-3
1x 10-3
FIM
1x 10
5x 10
1x 10
1 10-3
Table 8 presents the average perplexity results obtained after fine-tuning the GPT2-Large (774M) model for epochs on the WikiText-2 dataset. The indicates the 95% confidence intervals mean perplexity value over these runs, starting from different random seeds.
Damping parameter e sensitivity study: In this section, we discuss the practical differences between ASGO and Muon. Theoretically, as discussed in Section 6, ASGO can be viewed as an adaptive extension of Muon. Specifically, ASGO's update rule degenerates to that of Muon in the absence of momentum and adaptive scaling (i.c., when B1 0 and 32 0).

---

Table 8: Finetuning GPT2-Large Perplexity
AdamW
Muon
ASGO
DASGO
CLM FIM
14.010 + 0.016
13.912 + 0.003
13.879 + 0.010
13.841 + 0.019
17.458 + 2.254
15.933 + 0.102
15.661 0.097
15.451 0.072
Moreover, if we modify the formulas for V22 and W4+1 in ASGO (Algorithm 2) to:
if m < n then V2 = B2 - + (1 (2) M,M,"
else V, = 82V,-1 + (1 32) M, M,
(4)
if m < n then W = W, - 111 (Ve + +e)) M
else W2+1 W. (V, + c
and set B2 ASGO becomes identical to Muon, even with momentum. However, in practical train- ing implementations, a subtle yet crucial difference exists between ASGO and Muon. This distinction arises because ASGO applies additional preconditioner with damping the momentum term, step not present in the standard Muon update.
To precisely study these differences, background optimizer to train a GPT2-model on the NanoGPT dataset for 1000 steps. At each step, computed and compared the Muon and ASGO update directions AW in the embedding, query, key, value, and MLP layers. Specifically, we computed:
set up an experiment using the Muon optimizer as the
Standard Muon Update Direction (SVD-based): A AW = UM UMV where UM and VM the left and right singular vectors obtained from the first-order momentum M.
Modified ASGO Update Direction (SVD-based, 0): AWASGO
SVD,c=0 = V Mr, where
V, is computed via standard Singular Value Decomposition (SVD).
Modified ASGO Update Direction (SVD-based, 10- AWASOO
SVD,c=10- (V
10-81y- M., where (V, + 10- is computed via standard SVD.
Modified ASGO Update Direction (Coupled Newton, 10- AWASOO
CN,==10-x (V-
10 10-8 where (V, + 10 100 for 50 iteration steps. Higham [2008a], Shi al. [2023]
-1)- à is computed using the Coupled Newton algorithm
Modified ASGO Update Direction (Newton-Schulz, 10-8y; AWASGO
(V.
10 10-$1) 10-8 Mr. where (Ve + 10 10- computed using the Newton-Schus algorithm for 50 iteration steps.* Higham [2008a], Bernstein and Newhouse [2024b], Jordan et al. [2024]
We then computed the cosine similarity between each of the four modified ASGO update directions and the standard Muon update direction A AWMM and plotted these similarities in Figure 7.
Unlike AdamW, ASGO exhibits sensitivity to the choice of the damping parameter. As depicted in Figure 7(a), when damping term is used = 0), ASGO's update direction, though exhibiting some instability, largely reconstructs Muon's update direction, with cosine similarities generally exceeding 0.8, particularly for the MLP and Embedding layers, as training progresses. However, the introduction of a small damping hyperparameter (e = 10-$) significantly alters these similarities. Figure 7(b) shows that even with precise SVD computation, the similarities for layers like Value and Key only reach approximately 0.7 gradually. This suggests that the small value can introduce and accumulate errors during training, potentially leading to degradation optimizer performance, behavior not typically observed in other common adaptive methods like AdamW: W. This phenomenon highlights the critical importance of small singular values in the preconditioning process, as applying damping term can disproportionately affect them, thereby altering the matrix multiplication and
"The Newton-Sehhe -chroln-SShhzzaa algorithm is derived from Algorithm 6.35 on Page 153 in Higham [2008a]. We utilized the quintic version with parameters (a, b, c) = (2, -1.5, 0.5) as described in Jordan et al. [2024].
*CosineSimilarityy Tx(47 B)
30

---

04
Opiny
Qptry
(a) CosineSimilarity (AW SVD
A WrASO SVD,e=0)
(b) CosincSimilarity Maon A W
SVD,c=10-5
06
02
(c) CosineSimiilarity( W Mion
W ASGO
(d) CosineSimillritty (A W Muuo
A ASAGO
cN-=-10-8)
SVD NS,c=
Figure 7: Comparison of Cosine Similarities (CSs) between the Muon update direction (A W and various modified ASGO update directions different layers (MLP, Value, Key, Query, Embedding) during 1000 steps of GPT2 pretraining. Subfigures (a) and (b) plot CSs for SVD-based ASGO update directions with = 0 and e 10 8, respectively. Subligure (c) and (d) plot CSs for Coupled Newton-basee and Sewton-Schull-basedd, respectively, ASGO update directions with 10
amplifying errors. To circumvent this sensitivity, we opted to use an value of 0 in our large-scale GPT2 pretraining experiments described in Section 7.1. In this setting, directly compute the SVD of the preconditioning matrix and its inverse square root. The results presented in Figure 6 indeed demonstrate that ASGO achieves performance highly similar to Muon.
In Figure 7(c) and (d), we explore the use of more computationally efficient methods, Coupled Newton (CN) and Newton-Schulz (NS), respectively, to approximately compute V; 1 as alternatives SVD. The results show that, aside from some numerical stability issues observed in the Embedding layer for both CN and NS at certain steps-whichh reasonable given that the sparse nature of embedding layers often requiring larger damping parameters- the cosine similarities for the remaining parameters are highly comparable those obtained with SVD. This demonstrates the feasibility employing efficient matrix algorithms to replace SVD, thereby enhancing ASGO's computational efficiency. This direction represents one of our promising avenues for future research.
D Auxiliary Lemmas for the Proof
Lemma 1 (Trace properties). For arbitrary matrices A E mx B Ruxm X,Y YE have the following basic properties for the trace:
E mxx we
1. tr (X) tr Tx
2. (AB) Ir (BA);
3. if X symmetric positive semidefinite, (X) ||]]], 0;
4. if X, Y and X symmetric, (XY) 2
The following lemma notes the operator monotonicity of the power functions, which is a classic result [Löwner, 1934, Ando et al., 2004, Gupta et al., 2018].
Lemma The function f > a with a E [0, 1] is operator- iec i 3 5 B, holds that A < Ba.
31

---

Lemma 3. For symmetric positive definite matrices X, Y E R R=x" it holds that
ir (=(x+1)) tr X + +Y
Proof. It holds that
tr (X =tr + (X+ (x
=Ir (X+ + ir (x++))=
2tr X + tr {
where the last inequality is based on Lemma 2 such that (X + Y) > X and the fact that A- = B-! if A - B for symmetric positive definite matrices A, B. Further based on Lemma 1, can finish the proof.
Lemma 4, For symmetric positive semidefinite matrix X E Rmxm it holds that
tr x
Proof. The inequality is equivalent to that
m
m
VA((X)
where (X) denotes the i-th largest eigenvalue of X (the same singular values for real symmetric positive semidefinite matrices). Firstly, we have the fact that the eigenvalues {{){))} majorize the diagonal entries {[X] mh i.e. for all
Σ xX)
1X and Σ xXX
1
Also, we know g [ Ex xf is Schur-concave operator. Then obtain the equality based on the Schur-Horn theorem, which implies that for a Schur-concave operator if sequence (x:)) ym majorizes Hy then ((x,)) S g (
Lemma 5. ||-11, is norm, i.e., it satisfies the basic properties of norms. Its dual norm is I1-1lr---
Proof. Denote X = [ri,..., [a]] E Rmxx where each d. E Rx1mm Then we have
L
tr (XTLX) LX) = Σ x. LE:
F1
L
L
T
which is norm for the space mx whose dual norm is ||-11r, This concludes the proof.
Lemma 6. Assume a non-negative sequence { and a positive sequence { with S
Em 83, it holds that
n
Σ 1=1
イ S
(5)
S
8j
j=1
The inequality holds as equality if and only if for all i
1, n and j 1, 1,...,70
Tj
8; $j
Proof. A proof can be found in Lemma G.5 of the ICLR version of Liu et al. [2024].
32

---

E Proof of Nonsmooth Convergence of ASGO
Proof of Theorem 1. Denote AW, - W W. where W. is an optimum of Problem (1). From the algorithm update, we have
Ar. AAW-- = AW, A,W; - nu (G) AW; + AW, G.) + Gr AT G. Then, by taking trace of the above equality and rearranging, we can obtain that (AW, A,W,, AW AW (AW net (C) A, G.).
2n, (G) AW,)
By convexity, we know
E tr (G, AW.) (WI) 2 B[f(Wa) f(W.). Then combining these and taking summation over t and taking 1): 17, we have
2
E [[((WW]] f(W.)
T-
Ir (AWTTAAW, AWT, A.W.H.WWHH)
+ +
Σ
Lr (G
1
+ -r (AW. AW)) L |tr (AWAAAAA--AWW))
tr (AW, (AW,AAWW, AAW, - AWT.A--AW
"
tr (G. AG
+ n
<-tr (AW, A 1)
-E 1
tr (AW, (A. (Ai-AAm -A--))Ai)) + +n
Σ tr t=
(G) A,
(GT
where we note A = emm- Thus the first term on the RHS of (6) can be bounded by 6 as
1 tr (AW,JA = = 1la - DD
(7)
We then deal with the second and third terms separately. For the second term, we have
(AW, (A.
A-1))W,) -A-
tr (AW,AW, (A₂ At-1))
(AW,AWT, M -
JAW,AW, 1lop IAA+ - 1A--Ac-1ll.
lwwll (A.-A,11 S -AT- (A-1))
(8)
Note that the first and last equalities are based on the properties of the trace, (see Lemma 1). The first inequality is based on the duality of the I-Illop and II-I1, norms. The third equality relies on the
positive semidefiniteness of A, - A for all based on Lemma 2, since A - A G.G.T > 0.
For the third term of (6), have
tr (G. A, G.) =tr (G.G, A, tr ((A) - 1
=tr (((A, A-11) 2A, (Ai - Am-1) )2 A AcA A71) (9)
-2tr (A, - -A-1) - tr ((A+ - - + ((Ar-1A4 A
22t (A, - Ar-1),
33

---

where the last inequality follows from the fact that
tr ((A - -A-1))) = tr -(cA--AA-1)A4 A--))4 $] [Ad - AA-1)A } 20
and
tr ((Ag-1A, - A tr (At-1) (AcA4-1A,7)) =
which based the positive definiteness of Ar and the properties of the trace (Lemma 1). Therefore, by substituting (7), (8), and (9) into (6), we have
T-1
B(((W.)) - S(W.)
t=
T-1
tr (AW AW,AW))
LE
Σ tr (AW (A, (A-A-AA-)AW)))
+ +n
Σ tr (G) A, A,) A=
1
6
D
Dop
+
+
2n
E ft (AT-1-AA - A-1
17
17
$) }{
D
tr Σ EGG. G.G.
+
2n
E
1
Taking
Dop completes the proof.
Based Theorem and Lemma 8, can also prove Corollary
Proof of Corollary 2. Based on the results in Theorem if we additionally have
E G.G. 3 Q?
then using Lemma can obtain that
Σ G.G.
(10
SE
||Q1I.
T-1
G.G.T
Q-1
T-1
Σ
G.G.T
Σ GGTT-1)
E
|1Q11, tr
Q111
(E
QQ11.
|Q11. VT 11
where the first equality is based on Lemma and the second inequality is based on the fact that g(x) = vz is concave for 2 0. This concludes the proof.
We also provide the following proof for the comparison between Theorem and the convergence rates of Shampoo and full-matrix AdaGrad.
Proof of Comparison with Shampoo and full-matri AdoGrad. We list the convergence rates here.
m Dr
T-1
Full-Matrix AdaGrad: O
Σ
G.R.R
j=1=1
t=
34

---

[-1
G.G. )
ir (
1-1
1)
Shampoo:
VraDD..n .tr ir
Σ
Gi
1)
ASGO:
O Dop tr
EGG!
0
Dop
Σ
G
1=0
t=
The last inequality for ASGO is based on Lemma AdaGrad:
We first compare ASGO and full-matrix
))
( Σ 1=0
T Σ t=0
T
ur
[GG
G.G.
Σ
G
Σ
G.
>ΣΣ 0=1
G.
:=1 t=0
where the first inequality is based on Lemma 4 and the second inequality is simply a fact that for any vector we have | 2 ||>]]] Thus we prove that the proven convergence rate of ASGO is at least Dp/ Dop times faster than full-matrix AdaGrad.
Then we compare ASGO and Shampoo. We have
* )
T-2
-GG
G.T.
Σ E G.G.
tr
1=0
1=
(2-1 Σ
Str
Σ G
G.G
1=0
1pp
=tr
Σ
G.G.
Σ
G.G.
1=0
t=0
where the inequality is based on the fact that I1-llop and ||-11- are dual norms. Then we have
Σ
G.G.T
Σ I...G.
= Σ G. Gellon
up
Σ G.. G. )
T-
$ tr Σ
E...
G.
s tr
where the first inequality is based on that Il is a norm and the second inequality is based on the fact that Lr (X)= 2 || llop for symmetric positive semidefinite X, and the third inequality based on
the fact that (tr (X)) s tr (xx) because if we denote og as the j-th singular value of X, we have
Σ0,
(tr (X))+
03
Thus we can conclude that
))
1) ((
tr (EEaai)
tr Σ GGT
St
Therefore, the proven rate of ASGO is at least vraDe/ Dop times faster than Shampoo.
35

---

F Proof of Smooth Convergence of ASGO
The starting point of the smooth analysis is Theorem For notation simplicity, we denote N₂ -
G. - VS(W)) and V - (W(,) in this section.
Lemma 7 (Separate Gradient and Noise). Under the same settings as Theorem it holds that
B[s(W))- S(W.) Dop E
Dop
T-1
+ DDi
VAVA
R
Σ
N
T
T
Dup
T
Proof. Based on Theorem 1, we have
Σ ElS(W,) S(W.) D E
EDI
Σ G.G.
DapT
Dop E
+ EDG DOpT
2V + +2N,NT
T
Dop E
DOP T
DD Dop
2 2 VAVA
+
E
2 Σ
N.N
+
1=0
t=
Here the first inequality comes directly from Theorem 1. The second inequality is based the fact that for all A, B E IR mi' we have
2AAT 2BB - B) B) =AA T + BB A B--BAA
= =(A - (A - B) T = 0.
The last inequality is based on Lemma 3.
The following lemma is key technical lemma for proving the smooth convergence results, which is a generalization to the matrix case of Lemma 6.
Lemma 8 (An upper bound on I1-11.), For symmetric positive definite matrix A E R and matrix G E gmxn, it holds that
||[| S
(10)
Proof. We first consider the case A diag(Ai, Am] is diagonal matrix. We have
||]|_ =tr ((GGT)
Σ v GGT
[GG
153
Aj
||AIl, ir (CGTA-1)= VIAII Ir (GTA-
where the first inequality is based on Lemma 4 and the second inequality is based on Lemma 6.
Then we prove that this holds for general symmetric positive definite matrix A. Let the singular value decomposition be A U.UT. Denote & = UTG. G. Since Σ E mm is diagonal matrix, it holds that
|G]]_ - ||||_ s / (GTE-IC) = UI
||A||, tr (GT
which concludes the proof.
36

---

We also use the following lemma to indicate the reduction of variance by batch size M. Lemma (Variance reduction by batch size). Under Assumption we have
E N.N,
V2
M
where we denote N; - G. VW( for Algorithm 1.
Proof. For notation simplicity, we denote ((W & VF(WE:E) - Vf(W)) and thus
N₂ Σ N(W.:))
EEB
Then it holds that
Σ [ )
N
E
Σ N(W
M
SEB
CEB
EE (N(WGE)N(WW:
M
SEB
Σ ([N(::))N(W:::]
M
1
Σ
V2
V²
M
M
ŞEB
where the last equality based on that (W(;;; are mutually independent and I and the last inequality is based on Assumption 3.
Then based on these lemmas, we can prove the smooth results.
Proof of Theorem 3. We separately deal with the bias and variance terms, which refer to the first two terms on the RHS of Lemma For the bias part, by plugging in the smoothness matrix into Lemma 8, we can obtain
1 1-1 VINV
(L
E
VAVS
(10) S
||-|I.
Σ
E
|1411 tr
VIVITLL
|ZI. Σ E [IVAI
|/|. Σ 2(((W)) S(W.))
where the equalities are based Lemma 1 and the second inequality is based on the fact that g(z) y is concave for x 2 0. The last inequality is based on the properties of smoothness [Nesterov al., 2018] and the fact that Il-Il,,- is the dual norm of I|-I1,--
Then, the variance part, let first assume M for simplicity, then we have
$
1
(0)
N.NT)
E
Σ
NONT
SE
V-1
Σ N.NT
VII. tr Σ N.NT
E
37

---

VIl. Σ tr [E[NNN]] V-1)
T-1
VII. Σ 20
I VTIIVIL,
where the equality is based on Lemma 1. and the second inequality based on the fact that g(x) /x is concave for 2 0. Then plugging these in Lemma have
$
Dop E T
DOP T
+ DD
[[((W)) s(W.)]
2 VAVS
+
E
2 2
N.N!
T t=
Dop
1=0
oD S T
o2 [
De DopT
||
B[s(W,) f(W.)] +
VT
E
Thus we denote x =
V 7 ETE B[f(W,) - f(W.)], we can write the inequality as
x² s bæ +
where
2D op
oo IVI. +
cD Dop
b
c=
VT
vT
Then as x 2 we can solve this simple quadratic inequality obtain that
x² $ b+ V2 4c)
2 - 26² + 2c
4D2 op B(f(W,) f(W.))
|
2v2D op
IV
2cD3
T 0=
T
VT
Dop
which concludes the proof for M Then, based on Lemma 9 to incorporate batch size M > we finish the proof.
G Proof of Section 6
Muon [Jordan al., 2024] is essentially the following algorithm:
B, = μB,-1 + G.
U,, Sr, V. Compact SVD(B.)
(11)
W411 W, UUV.
where G, is the gradient obtained W. and U. E Rm V₁ Rr the matrices of the r, left and right singular vectors corresponding to the non-zero singular values B, that are the diagonal components of the diagonal matrix S, E x with being the rank of B,. Note that U,S.V, referred to as the "compact" or "Reduced" SVD of B₁. Muon implements this algorithm using Newton-Schui matrix iterations to approximately compute U.V. instead of directly employing SVD to improve computational efficiency. The relationship between (11), with B, = G.. and (2) interpreted in Bernstein and Newhouse [2024b] (see Story II). In the following discussion, refer Muon as the one described in (11). We first show that ASGO and Muon, without gradient accumulation and momentum, equivalent.
Proposition 1. If in ASGO (Algorithm version) we set B1 32 = 0, O, and T = 1, and in Muon (11) we set μ 0, ASGO is equivalent Muon.
Proof. Let U.S.V. be the compact SVD of G.. For ASGO, the update is
W41 =W; (TG - G.
38

---

WW - ([
=W: -
=W: nUS
W. M.V.
Clearly, this is the same as Muon.
We now prove a nonconvex convergence result for Muon the deterministic case with μ = 0.
Theorem 4 (Nonconvex convergence of Muon). We consider Muon presented (11) with p = In the deterministic case, i.e. no gradient noise, if we assume that there exists lower bound f such that f(W) 2 f for all W and Assumption 2, by taking Th = 1J 2()(W))- it holds that
LI.
1
IIII. (f(Wo) f*))
VV((WIl.
T
21
Proof. Since we assume the deterministic setting, we have G2 V(W,). Also, since we assume 1 =-: O, we have U,, V2 are the left and right singular vectors of G, such that G, U,S.V,T. Then based on the smoothness assumption, it holds that
f(W(1) Sf(W)) (Vf(W)), W4+1 W.) - tr ((W411 Wo) L(W Wa)
"i 2
=((W) net (G. + tr ((U.V,T) " I((UVTT)
Th
=f(W)) th (V.S,U, U.V.T) + tr (V.U, LUVV))
2
=f(W)) tetr (V.S,V,)) + nh
tr (U. LU.VVVV)
2
=f(W) net (S.V Vi) + n
tr (LU;U?)
2
<f(W)) ||Grl, + T 2
|||..
where the last inequality is because of the fact that UUT 2 and Lemma Then by summing up over and rearrangement, we can obtain that
ne ||all S f(Wo) f(Wr))
Σ
1 ||]|.
2
t=
Then we take 7 1) to obtain that
[V(WWII| ((W)) 1"
|
n/T
Setting
2(f(Wo)=1" |...T
finishes the proof.
H Incorporating Projection in ASGO
To achieve fully theoretically rigorous convergence bounds for ASGO, i.e., with Dop and De strict constants as increases, one cell-established way to incorporate a projection operation in the Algorithm [Hazan et al., 2007, Duchi et al., 2011]. Consider closed convex set W C gm with an optimum W, inside and diameters D Dp such that for any W1, W2 W, it holds that
W1 W2llop < D.p and
| - Wallp S Dp.
39

---

For Algorithm 1, we desire to search solution inside W and change the update step line 8 to
Z++1 W, 7, and Wi+1 II (Ze+1).
where the anisotropic projection operator x [2007], Duchi et al. [2011], is defined as
R W, as also incorporated in Hazan al.
IIK, (X) aig || Yll
Yew
for any X TR which enjoys the non-expansiveness of a projection operator:
AW 41 W.ll.. Ir (AW < tr (AZ44,,,Z44) = -| 12
(12)
with AW & W441 W, and AZ Z - W, since W, e W. Substituting this non- expansiveness property in the proof of nonsmooth convergence, more specifically, in (6), we can obtain that
2
B[((W))- - f(W.)
E
tr (AW, AAW, AZA AZZZAAZZ
+ nE
(12)
$
E
tr (AWTAAM) - AWT A.A.H.WHH))
+ +
tr [GT AT
E
tr (AW, AAW, - AWTA AWTAA))WW))
+ tr (AW, A.
E (rW7Ar--AWW
17
17
+
Σ tr (G. AT G.)
(AW,
E 7
tr (AW, (A, AA-)) + +E Σ (GT
which yields exactly the same result (6). Since the rest of the proofs of both nonsmooth and smooth theory are based on (6), we can achieve the same convergence results as Theorems and 3, while we have the upper bound on maxogesT IW, W. Wl-loop < Dop and maxogisT-1 ||W; - W-llr S
Dr. where the diameters Dop and DE depend on the geometry of the convex set W and remain constants when T increases.
40


Reviews:
[{"review_id": "cla8RtpyGV", "reviewer": "Reviewer_Vxtn", "summary": "The paper designed an optimization algorithm when the decision variable has a matrix structure. The algorithm has some connections with AdaGrad. Compared with Shampoo, it does not need to maintain two preconditioning matrices. Theoretically, the authors provided the convergence guarantees for the convex nonsmooth setting and the convex smooth setting.", "questions": "### Major Question:\n\n1. How should $\\epsilon$ be chosen? Is it introduced to ensure the invertibility of $\\Lambda$?\n\n2. On Page 6, it is mentioned that the preconditioner should be on the right side of $G_t$. Why is it placed on the left side in Algorithm 1?\n\n3. What does the sentence \"we expect that ASGO has a nonconvex convergence rate comparable to (3) obtained by Muon up to logarithmic factors\" mean? Are the authors aiming to establish a similar rate in the nonconvex regime?\n\n4. Is there any comparison between ASGO and other algorithms in terms of wall-clock time?\n\n\n### Minor Suggestions:\n\n1. At the top of Page 5, I suggest including $T$ in the convergence rates for clarity.\n\n2. When describing the algorithm in Section 3.2, I recommend defining $V$.", "limitations": "N/A", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths:\n\n1. The authors provide a convergence guarantee for the algorithm. Furthermore, they characterize scenarios where the algorithm benefits, such as when the gradient is of low rank or when there is a block-wise diagonal Hessian structure. The notation introduced in Assumption 2 aids in this characterization.\n\n2. The experimental results appear promising.\n\n\n\n### Weaknesses:\n\n1. The bounds in Theorems 1 and 3 include a term $D_{op} = \\max_t ||W_t - W_*||$, which may grow with $T$. Consequently, it is not entirely rigorous to claim convergence, particularly in the absence of a bounded domain.\n\n2. The paper compares ASGO with SGD. However, I am uncertain whether this is a fair comparison, as SGD can converge without requiring a bounded domain, and its complexity can depend only on $||W_0 - W_*||^F$.  This makes it challenging to directly compare with $D_{op}$.\n\nReference.\nOrabona, Francesco. \"Last iterate of SGD converges (even in unbounded domains).\" Blogpost on http://parameterfree. com, available at https://parameterfree. com/2020/08/07/last-iterate-of-sgd-converges-even-in-unbounded-domains (2020).", "was_revised": true, "final_justification": "I raised my score because the authors provided a detailed explanation during the rebuttal on how to incorporate projections, which enhances the soundness of the theory.", "rebuttal": "Thank you very much for your careful reading and insightful opinions and questions, which we now address. \n\n**Weakness 1: $D_{op}$ may grow with $T$.**\n\nThanks for the meaningful question. It is indeed true that $D_{op}$ could increase when $T$ increases in the settings of Theorems 1 and 3, but as we also mentioned in Remark 1 (after Corollary 2), we can address this by invoking a projection onto a bounded convex set, where the projection should be with respect to $|| \\cdot ||_{\\Lambda_t}$. This technique is well-established in standard optimization literature [1,2]. Here, since the projection operation is rarely used in practical tasks like training DNNs, we follow [3] and omit it in the algorithm and theorem. See the discussion after Theorem 7 in [3] for a similar opinion.\n\n**Weakness 2: The comparison between ASGO and SGD may be unfair**\n\n1. We agree with the reviewer that the comparison is not fully rigorous, but we included it to provide some intuition about when ASGO may have a good convergence guarantee and how ASGO can utilize structural properties (low-rank gradients for the nonsmooth part), and we are not claiming that ASGO is a better algorithm than SGD. Note that a similar comparison was also considered in, for example, the original AdaGrad paper [1], also for explaining the potential benefits of AdaGrad. \n\n2. Also, it is possible to prove rigorously the convergence of ASGO in an unbounded domain, following [4] for AdaGrad. However, since this analysis introduces novel assumptions and can be much more complex, it goes beyond the scope of this paper, and we leave it as an interesting future direction of research.\n\n**Major question 1: Choosing $\\epsilon$ and the purpose**\n\nYour question is insightful, and your intuition is exactly right. The parameter $\\epsilon$ is indeed a standard damping term, introduced in Algorithm 1 primarily to ensure \nstrict positive definiteness and thus the invertibility of the preconditioner matrix $V_t$. This is a common practice in adaptive optimization methods to guarantee numerical stability.\n\nHowever, in our specific implementation, the matrix $V_t$ is constructed using an exponential moving average (EMA), which provides inherent smoothing and accumulation of second-order information. As a result, the matrix $V_t$ becomes progressively better-conditioned during training, making the risk of singularity negligible after the initial iterations.\n\nAfter careful tuning, we found that choosing the damping parameter as small as possible always gave the best performance. Due to the stability provided by employing an EMA, we found we could set $\\epsilon=0$ in our experiments without encountering any numerical issues. We compute the inverse square root $V_t^{-1 / 2}$ directly and robustly using Singular Value Decomposition (SVD). Setting $\\epsilon = 0$ removes a tunable hyperparameter and allows the preconditioner to adapt more precisely to the loss geometry without any artificial damping.\n\n**Major question 2: Why is $\\Lambda_t$ placed on the left side in Algorithm 1?**\n\nWe place $\\Lambda_t$ on the left side in Algorithm 1 to make the algorithm easier to understand, since this form is closer to the standard AdaGrad algorithm, where the parameters are treated as vectors and a diagonal preconditioner is used. Also, we point out that no matter on which side $\\Lambda_t$ is placed, our theoretical analysis still holds. This is because, based on the algorithmic design, optimizing $f: \\mathbb{R}^{m\\times n} \\to \\mathbb{R}$, with $\\Lambda_t$ on the right side is equivalent to optimizing $g: \\mathbb{R}^{n\\times m} \\to \\mathbb{R}$, where $g(W^\\top) \\equiv f(W)$, with $\\Lambda_t$ on the left side. We thank the reviewer for the question and will add more explanation on this point in our revision. \n\n**Major question 3: Nonconvex analysis of ASGO**\n\nThe answer to your insightful question is \"Yes\". By the quoted sentence, we mean that we expect to obtain a nonconvex convergence result for ASGO that is comparable to the Muon convergence rate that we established, since ASGO and Muon share some important algorithmic similarities. We think this would be an interesting topic for future research.\n\n\n**Major question 4: Wall-clock time analysis**\n\nWe strongly agree that a fair comparison requires a thorough analysis of runtime, and we appreciate the opportunity to elaborate on the complexity of ASGO and provide detailed wall-clock time (WCT) measurements.\n\nTo compare the WCT for ASGO with that for other algorithms, we conducted a new pre-training experiment for the 125M parameter GPT-2 model on the OpenWebText dataset with a significantly larger token budget. We trained the model for 1200 steps using a batch size of 64, 4 A6000 GPUs, a sequence length\nof 512, and 8 gradient accumulation steps. This setup corresponds to a total token budget of: $512 \\times 64 \\times 4 \\times 8 \\times 1200 \\approx \\mathbf{1 . 2 6}$ Billion tokens. This budget represents approximately 10 tokens per parameter, which attains the limit of our computational resources.\n\nAs we mentioned regarding question 1, our initial version of ASGO uses SVD to compute the matrix inverse square root. While mathematically exact, this approach carries a significant computational burden. For the new experiments presented in this rebuttal, we have adopted the more efficient Newton-Schultz (NS) iterative method, which substantially reduces the wall-clock time of the optimizer step and makes the algorithm more practical. The following runtime analysis is based on this improved NS-based implementation.\n\nAs shown in Thm 5.2 [6], computing $A^{-1 / 2}$ can be achieved by computing the sign function of a block matrix \n    $$\n    \\operatorname{sign}\\left(\\left[\\begin{array}{cc}\n    0 & A \\\\\\\\\n    I & 0\n    \\end{array}\\right]\\right)=\\left[\\begin{array}{cc}\n    0 & A^{1 / 2} \\\\\\\\\n    A^{-1 / 2} & 0\n    \\end{array}\\right].\n    $$\n    This implies that computing the inverse square root in ASGO shares the same order of complexity as computing the matrix sign function in Muon. However, since the NS iteration in our case is applied to a matrix twice the dimension, the per-iteration cost is approximately double that of Muon.\n\n**Empirical Wall-Clock Time Analysis.** To provide a concrete analysis, we measured the WCT under the exact settings of our 1.26B token experiment.\n| **Optimizer** | **Single Train Step WCT** | **Time Difference** |\n| :--- | :---: | :---: |\n| AdamW | 5.451 sec | 0 sec |\n| Muon | 5.519 sec | 0.068 sec |\n| **ASGO (Ours)** | 5.947 sec | 0.468 sec |\n\nTable 1: Wall-clock time per training step, difference compared with AdamW update.\n\n- Single-Step Optimizer Cost: For a fair comparison of convergence, we use 10 iterations of the NS algorithm for ASGO instead of 5 iterations used for Muon's update. The measured WCT for a single train step (including forward, backward, gradient accumulation passes, and optimizer step) is shown in Table 1.\n\nAlthough computing the inverse square root of $V_t$, in ASGO's optimizer step is more expensive than in AdamW, the forward/backward pass is the dominant bottleneck. In our setup, a full training step in the ASGO update only costs 0.468s more than the AdamW update. This is acceptable in the training process. As a result, the total end-to-end training time is only marginally affected. Table 2 shows the total time to complete the 1200-step training run.\n\n| **Metric** | **AdamW** | **Muon** | **ASGO (Ours)** |\n| :--- | :---: | :---: | :---: |\n| WCT | 127.76 min | 129.01 min | 133.50 min |\n| Train Loss | 3.73 | 3.58 | 3.53 |\n\nTable 2: The total Wall-clock Time (WCT) and final training loss for training 1200 steps.\n\nThese tables show that although AdamW is computationally cheapest at the optimizer level, the total training times for all three methods are very similar. This demonstrates that the total training time may not be a main concern for ASGO.\n    Note that ASGO achieves the smallest training loss while AdamW achieves the largest loss. \n\n\n\n**Minor Suggestions**\n\nWe thank the reviewer for carefully reading our submission and suggesting ways to improve its clarity, which we will follow.\n\n**References**\n\n[1] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research 12.7 (2011).\n\n[2] Elad Hazan, Amit Agarwal, and Satyen Kale. Logarithmic regret algorithms for online convex optimization. Machine Learning 69.2 (2007): 169-192.\n\n[3] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. International Conference on Machine Learning. PMLR, 2018.\n\n[4] Amit Attia and Tomer Koren. Sgd with adagrad stepsizes: Full adaptivity with high probability to\nunknown parameters, unbounded gradients and affine variance. International Conference on Machine Learning. PMLR, 2023.\n\n[5] Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks\n\n[6] Nicholas J. Higham. Functions of Matrices: Theory and Computation."}, {"review_id": "O2K9RL3Uir", "reviewer": "Reviewer_WuZT", "summary": "In light of the recent resurgence of \"matrix\" layer-wise preconditioners in deep learning optimization such as Shampoo and Muon, this paper considers understanding the principled benefit of these methods via regret / convergence bounds. In particular, the authors present an algorithm denoted Adaptive Structured Gradient Optimization (ASGO); compared to Shampoo and Muon which can be viewed as preconditioning the gradient matrix in weight space from both sides, ASGO proposes only storing a one-sided preconditioner. Though seemingly a more restrictive optimizer class, the authors show that from the view of convex or smooth matrix-variable optimization, the convergence upper bound of ASGO may exhibit more favorable dependence on problem parameters than Shampoo or Muon, or even full-matrix Adagrad. These hints from theory are corroborated in experiments, where ASGO is demonstrated to consistently match the performance of other state-of-the-art methods on various deep learning tasks, indicating that it may suffice to perform salient one-sided preconditioning of the gradient in practice.", "questions": "Besides the points laid out under Weaknesses, I have the following auxiliary questions:\n\n1. A critical aspect of deep learning optimization is ensuring \"scale-invariance\" of the hidden activations across width/depth, such as ensured by $\\mu$P [1,2] or spectral scaling [3] and developments thereafter. How should ASGO be scaled layer-wise to ensure stability of hyperparameters across architecture sizes?  \n\n2. Regarding ASGO's compatibility with $d > 2$ dimensional tensor weights, what prevents handling higher-dimensional tensors as in, for example, Muon?\n\n[1] Yang and Hu, \"Feature Learning in Infinite-Width Neural Networks\"\n\n[2] Yang et al, \"Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer\"\n\n[3] Yang et al, \"A Spectral Condition for Feature Learning\"", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Overall, this paper presents an interesting perspective to gradient preconditioning for deep learning optimization. In particular, the observation that a \"weaker\" preconditioner class might actually yield better results is nominally surprising and potentially interesting for the community. Furthermore, strong empirical results and implementation details are provided, which will be helpful for the community. Therefore, I recommend acceptance of the paper.\n\n**Strengths**\n\nTo motivate where a smaller preconditioner can actually benefit convergence, convergence rate upper bounds are provided for convex non-smooth and smooth settings where the optimization variable is matrix-shaped, demonstrating the performance bounds on ASGO can exhibit better dependence on certain instance-dependent quantities than Shampoo or full-matrix Adagrad. This serves as evidence that there may not be a monotonic relationship between the memory/algorithmic complexity of a preconditioner and practical performance.\n\nTo complement the theoretical observations, the paper provides solid experiments on standard GPT2 pretraining and finetuning benchmarks, showing that its performance can outperform standard AdamW baselines, while nearly matching Muon. Empirically, it is also noted that ASGO may be generally more robust to learning rate selection than Muon.\n\n**Weaknesses**\n\nMy main concern regarding the theory of in the paper are as follows:\n\n1: Conclusions are drawn by comparing performance upper bounds. For the purpose of motivating empirical performance this may be fine, but a rigorous comparison should provide performance lower bounds on, e.g. full-matrix Adagrad on classes of instances, where it may then be shown that ASGO's upper bounds rigorously outpace other methods. In general, walking through some problem instances (even if simple) can greatly improve the claim that a simpler preconditioner can rigorously outperform larger ones.\n\n2: The theoretical set-up considers a single matrix optimization variable. However, the non-triviality of establishing feature learning in neural network theory comes from the compositionality of layers. A simple example could be linear regression with two-layer linear networks $f(x) = W_2 W_1 x$. I am not confident the analysis in this paper can predict or describe the unique optimization geometry arising from matrix-valued weights interacting via composition. This also now adds the dimension of comparing **layer-wise** versus potentially cross-layer preconditioning (see e.g. KFAC variants [1] and PSGD [2]).\n\n**Minor Comments**\n\nRegarding the typical rank-structure / evolution in neural network weights, this paper may be insightful for the authors [3].\n\nThe statement of the Muon bound (3) seems wrong and differs from the more reasonable statement in Appendix G.\n\n\n[1] Martens and Grosse, \"Optimizing Neural Networks with Kronecker-factored Approximate Curvature\"\n\n[2] Li, \"Preconditioned Stochastic Gradient Descent\"\n\n[3] Martin and Mahoney, \"Implicit Self-Regularization in Deep Neural Networks: Evidence from Random Matrix Theory and Implications for Learning\"", "was_revised": true, "final_justification": "This paper presents an interesting new member of the \"Factored Adagrad\" family. Interestingly, it is argued that a \"weaker\" class of preconditioners --- in this case one-sided preconditioning versus, e.g., Shampoo's two-sided one --- can actually lead to performance **improvement**. Some theory via optimization over a matrix-shaped variable is derived to justify this phenomenon by demonstrating different preconditioning schemes lead to different manifestations of problem-parameters, where certain preconditioners can outperform other \"stronger\" ones if these problem-parameters line up in the right way.\n\nThe theory admittedly does not capture a full story in relation to deep learning, e.g. restricted to convex optimization over a single-block matrix variable, comparisons of performance **upper bounds** etc, which most deep learning-oriented methods tend to lack. However, more importantly, the authors support the efficacy of their optimizer with some LLM pretraining experiments, showing performance on-par with the state-of-the-art in Muon, with similarly general algorithmic set-up. The main difference is that ASGO maintains an Adagrad-like moment accumulator/preconditioner, while Muon performs orthonormalization per-step, which gives Muon the edge in computational efficiency via more specific iterative methods tailored to orthonormalization rather than generic matrix inverse sqrts. However, as understood by the deep learning optimization community, for large networks such as in LLMs, the FLOPs introduced by layer-wise preconditioning (within reason) is quite small proportional to a backwards pass.\n\nPutting these contributions together, I think various communities will find use in these results, and recommend acceptance (5).", "rebuttal": "We would like to extend our sincere appreciation to the reviewer for all the constructive questions and suggestions and recognition of our paper. The reviewer's remarks provide many important insights for future research directions relevant to ASGO. Many thanks!\n\n**Weakness 1: Conclusions are drawn by comparing performance upper bounds.**\n\n\nThis is an important question. We agree with the reviewer that comparing upper bounds is not a rigorous way to compare the algorithms, and here we compare the upper bounds mainly to give some intuition on how ASGO can be potentially beneficial from a convergence guarantee viewpoint.\nWe are definitely interested in finding lower bounds for the class of full-matrix AdaGrad or Shampoo, and some important scenarios where one-sided preconditioners can largely outperform more complex ones.\nHence, this\nis an interesting future topic and may inspire even more ways to improve ASGO.\n\n**Weakness 2: The theoretical set-up considers a single matrix optimization variable.**\n\nThanks a lot for this constructive and meaningful comment! Considering the layer-wise structure of the problem is definitely a very interesting topic. It is indeed the case that the analysis in the paper may not be trivially generalized to the setting where a lot of weight matrices are mixed together, as in neural networks. We consider this a potential future research topic and feel excited to work more on this question. For now, we think that the\nrecent paper [1] may provide some viewpoints and intuitions into how to prove the properties with a layer-wise structure, where the authors establish the convergence rates under multi-layer settings. However, [1] still doesn't take the interactions between matrices into consideration, which can make the problem much more complex. It would be really interesting to further explore this.\n\n\n**Minor comments 1: The paper [2] for rank-structure / evolution in neural network weights.**\n\nWe thank the reviewer for bringing this interesting and relevant paper to our attention.\nThe structured properties of matrices in DNNs discussed in the paper, especially the weight rank dynamics during the training process, are really insightful. This provides some significant evidence on what the weight rank should be in the training process, which is also a critical component of the ASGO convergence rate, as presented in Theorems 1 and 3. We will incorporate a discussion on this point in future revisions. Thanks again for introducing this insightful paper!\n\n\n\n**Minor comments 2: The Muon bound (3) differs from the statement in Appendix G** \n\nThanks for the careful reading and the question. We kindly note that we use different forms of the convergence rate in Equation (3) and Theorem 4 in Appendix G, and we can derive (3) based on Theorem 4:\n$$\n\\frac{1}{T} \\sum_{t=0}^{T-1} || \\nabla f(W_t) || \\leq \\sqrt{\\frac{|| L || (f(W_0) - f^*)}{2T}} \\quad \\Rightarrow \\quad \\min_{0\\leq t \\leq T-1} || \\nabla f(W_t) ||^2 \\leq || L || (f(W_0) - f^{\\*}) / (2T).\n$$\nHere, the norms are the trace norm $ || \\cdot ||_{\\*} $.\nTherefore, Equation (3) and Theorem 4 actually indicate the same convergence rate. We apologize for causing confusion because of the use of different forms and will make them consistent in future revisions.\n\n**Question 1: The use of $\\mu$P.**\n\nThis is definitely an important and interesting point for the scaling of ASGO. Intuitively, it is natural for ASGO to incorporate $\\mu$P based on [3], just as it is for Muon. This is because we can expect the spectral norm of the weight update to have an approximately $ \\mathcal{O}(1) $ scale (for the practical version (Algorithm 2)):\n\n$$\n|| \\Delta W || = || \\Lambda_t^{-1} M_t || = || (\\sum_{s=0}^{t} \\beta^{t-s} (1-\\beta) G_t G_t^\\top)^{-1} M_t || \\approx \\mathcal{O}( || (G_t G_t^\\top)^{-\\frac{1}{2}} G_t || ) = \\mathcal{O}(1).\n$$\nHere the norms are the spectral norm $ ||\\cdot||_{op} $.\nTherefore, to implement $\\mu$P for ASGO, we only need to use an adjusted learning rate $\\tilde{\\eta}_t = \\sqrt{n/m} \\eta_t $ to fulfill the requirement $|| \\Delta W || = \\mathcal{O}(\\sqrt{n/m})$, and set the initialization accordingly. This formulation is straightforward and can be important for scaling up ASGO, which we consider an important future research direction.\n\n**Question 2: ASGO for higher-dimensional tensors.**\n\nWe thank the reviewer for this insightful question. ASGO can indeed handle higher-dimensional tensors just like Muon, which basically flattens the tensor into 2-dimensional matrices. We are also considering other possible solutions that may utilize the structure of higher-dimensional tensors better, which could be an interesting future\nresearch topic.\n\n\n**References**\n\n[1] Riabinin, Artem, et al. \"Gluon: Making Muon \\& Scion Great Again! (Bridging Theory and Practice of LMO-based Optimizers for LLMs).\" arXiv preprint arXiv:2505.13416 (2025).\n\n[2] Charles H. Martin and Michael W. Mahoney. \"Implicit self-regularization in deep neural networks: Evidence from random matrix theory and implications for learning.\" Journal of Machine Learning Research 22.165 (2021): 1-73.\n\n[3] Greg Yang, James B. Simon, and Jeremy Bernstein. \"A spectral condition for feature learning.\" arXiv preprint arXiv:2310.17813 (2023)."}, {"review_id": "FoXZJsjtfV", "reviewer": "Reviewer_uDaJ", "summary": "This paper proposes ASGO (Adaptive Structured Gradient Optimization), an optimization algorithm that exploits structured properties of neural network gradients by using a single preconditioner matrix, requiring less memory than existing structured methods like Shampoo. The authors provide theoretical analysis and demonstrate empirical performance on language modeling tasks.", "questions": "- How does the computational cost of computing $V_t^{-1}$ scale with dimension m?\n- How does this compare to the computational cost of AdamW or Muon in practice on scale? \n\nThe paper covers an interesting problem but it contains multiple false claims and provides a weak experimental evaluation. In my opinion, an accept would require the correction of all weaknesses I listed above and an evaluation either on a broader benchmark or on larger language modes (or a strongly limited scope of the paper but this reduces the benefit). Also it should include a runtime analysis and a comparison with other optimizers given the same compute budget (for the runtime but also for the HPO).", "limitations": "No, the paper does not cover the cost of the matrix inversion and how it behaves on larger models.", "rating": 2, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 1, "clarity": 1, "significance": 2, "originality": 2, "strengths_and_weaknesses": "### Strengths\n\n- ASGO requires only one preconditioner matrix compared to Shampoo's two, reducing memory requirements by approximately half while achieving better theoretical convergence rates.\n- The paper clearly explains why single-side preconditioning may be preferable.\n\n\n### Weaknesses\n\n- The authors claim that \"adaptive gradient methods like AdaGrad and Adam\" “apply a diagonal matrix preconditioner to the gradient\" (line 19). This is wrong, also Adam does not use matrix inversion $Λt^(-1)$ - it uses element-wise division by the square root of the second moment estimates. This is a fundamental misrepresentation of how Adam actually operates.\n- The authors state \"adaptive gradient methods like Adam, treat the parameters as vectors and ignore the matrix structure of gradients.\" (line 30) This is incorrect. Adam performs element-wise operations on parameters regardless of their shape (vector, matrix, or tensor) - it doesn't reshape matrices into vectors.\n- The authors introduce a new optimizer for deep neural networks, “like Adam”, but evaluate only on small language model tasks. In my opinion one could introduce a general optimizer but have to evaluate it on common benchmarks like AlgoPerf or DeepOBS, or one introduces a specialized optimizer for small language models but communicate this limitation clearly.\n- In section 3.2, where the algorithm is introduced, it is unclear what $V_t$ is.\n- The claim “DNNs are generally nonconvex globally, convexity may apply locally in some regions” (line 121) is not backed.\n- In line 236, “Moun computes ….”, I can’t see how this represents the Muon optimizer. \n- The paper refers to Algorithm 2 and to Algorithm 3 (line 265, 267, 271) which are not in the paper.\n-  The authors perform in 7.2 a training on a ~125M parameter model on a token budget of 1000 steps x 32bs x 8 grad accum x 4 GPUS x 512 seq length = 525M tokens which is way below Chinchilla optimal. If one proposes an optimizer for language model pretraining, I would like to see experiments on at least 1B and on at least 8x Chinchilla optimal token budgets (at least for smaller models).\n- The optimizer contains a matrix inversion, this is any time consuming operation, however the authors do not discuss or provide measurements on the runtime of the optimizer. A fair comparison would include a comparison on the same runtime (compute budget) against other optimizers. \n- In contrast to the \"YES\" in the Paper Checklist 5, the authors do not provide code to replicate the experiments or test the optimizer. This is crucial for a newly introduced optimizer.", "was_revised": true, "final_justification": "The authors addressed many of my points in the review but one major criticism is unsolved. Either one could introduce a general optimizer but have to evaluate it on common benchmarks like AlgoPerf or DeepOBS, or one introduces a specialized optimizer for small language models but communicate this limitation clearly. But introducing a general optimizer and viewing the proper evaluation as future work doesn't work in my opinion. So far, we only know that the general optimizer works on 125M LLM sub-epoch training. \nAlso, the lack of source code to test the performance and runtime is in this case not acceptable. Therefore, I keep my score.", "rebuttal": "We thank the reviewer for the constructive and detailed feedback.\n\n**On Weaknesses 1 \\& 2: Diagonal preconditioner**\n\nWe agree that Adam and AdaGrad are implemented using element-wise operations. However, our claim is that this element-wise division is mathematically equivalent to preconditioning with a diagonal matrix. This is a standard and widely accepted viewpoint in the optimization literature, and Adam is described in this fashion.\nLet's consider the update rule for Adam:\n$$\n\\theta_{t+1}=\\theta_t-\\eta \\cdot m_t /\\left(\\sqrt{v_t}+\\epsilon\\right),\n$$\nwhere $m_t$ ($v_t$) is the first-order (second-order) momentum and the division is elementwise. This operation is identical to:\n$$\n\\theta_{t+1}=\\theta_t-\\eta \\cdot P_t^{-1} m_t\n$$\nwhere $P_t = \\operatorname{diag}\\left(\\sqrt{v_t}+\\epsilon\\right)$ is a diagonal preconditioner. The inverse of this diagonal matrix, $P_t^{-1}$, is simply a diagonal matrix with elements $1 /\\left(\\sqrt{v_{t, i}}+\\epsilon\\right)$, and \npre-multiplying the momentum\ngradient $m_t$ results in the element-wise division. This interpretation is discussed in foundational texts, e.g., [1, Chapter 5.4]. The original AdaGrad paper also introduced both full-matrix and diagonal versions [2].\n\nThis directly relates to the reviewer's second point. Since Adam performs element-wise operations, its preconditioning model is conceptually equivalent to flattening all parameters (e.g., matrices $W \\in \\mathbb{R}^{d_1 \\times d_2}$ ) into a single long vector and then applying the diagonal preconditioner discussed above. This approach implies that each parameter is scaled independently, a decoupling that renders the preconditioner agnostic to the parameters' inherent tensor structure. \nThis is why we say Adam generally ignores the structure of parameters.\n\nThis exact limitation of Adam is the primary motivation for our work. ASGO is designed to employ preconditioners that explicitly model the rich interactions within a parameter tensor (e.g., the row-and-column structure of a matrix), which a diagonal preconditioned method like Adam cannot. \n\n**On Weakness 3: Evaluation on benchmarks**\n\nThe primary objective of our experiments was to provide a clear and controlled empirical validation of our theoretical analysis, specifically demonstrating why and how a structure-aware optimizer can outperform a structure-agnostic one like Adam. For this purpose, small- to medium-scale language models provide an ideal and interpretable testbed.\nThis focused approach also aligns with our available computational resources, allowing us to conduct a thorough and rigorous investigation. Our experimental setting is similar in spirit to that of other recently proposed optimizers; for instance, Muon [3] was also initially evaluated on similarly-scaled models (e.g., NanoGPT) to establish its fundamental properties.\n\nFurthermore, to demonstrate the versatility of ASGO, we extended our evaluation beyond pre-training to include fine-tuning experiments, showcasing its potential in both training from scratch and adaptation scenarios.\n\nWhile the central contribution of this paper is to establish the theoretical motivation and empirical viability of structure-aware preconditioning, we concur with the reviewer that broader evaluation is a crucial next step. We appreciate the suggestions of benchmarks like AlgoPerf and DeepOBS, and we will prioritize them in future work to rigorously assess the scalability and generalizability of ASGO.\n\n**On weakness 4: Definition of $V_t$**\n\n$V_t$ is defined in Line 6 of Algorithm 1, (cited on the first line of Sect. 3.2, but which the Latex editor put on the next page in Sect. 4). We will add the definition where $V_t$ first appears.\n\n**On weakness 5: Local convexity** \n\nOur claim is based on the widely studied property that while DNN loss landscapes are globally non-convex, they often possess favorable local geometric properties. This is supported by a growing body of theoretical work. For instance, [6] introduces a class of functions satisfying local strong convexity [6, Assumption 1] and shows that deep ReLU networks are a practical realization of this class [6, Section 4.2]. Similarly, [7] establishes the restricted strong convexity of DNNs within a sequence of ball-shaped regions around the weight iterates. In our revision, we will add a discussion of this context and include these supporting citations.\n\n**On weakness 6: Muon updates**\n\nThe statement on Line 236 and the accompanying Equation (2) describe the momentum-free update step of the Muon optimizer. This step corresponds to a steepest descent update under the spectral norm, a definition that is detailed in [4, Proposition 5]. In the revised manuscript, we will restructure this section to make the statement unambiguous.\n\n**On weakness 7: Algorithm 2 and 3**\n\nAlgorithm 2 is in Appendix B page 15, Algorithm 3 is in Appendix C page 17. In the revised manuscript, we will ensure that every reference to these algorithms is accompanied by an explicit note directing the reader to the correct appendix.\n\n**On weakness 8: Larger training size**\n\nWe agree that evaluating optimizers on a larger, more representative token budget is essential for pre-training language models. Consequently, we have conducted a new pre-training experiment for the 125M parameter GPT-2 model on the OpenWebText dataset with a significantly larger token budget. \n\nWe trained the model for 1200 steps using a batch size of 64, 4 A6000 GPUs, a sequence length of 512, and 8 gradient accumulation steps, yielding a total of 1.26B tokens. This budget represents approximately 10 tokens per parameter, which attains the limit of our computational resources.\n\nAlthough we list a fair Hyperparameter Parameter Optimization (HPO) in Appendix C.2, we performed a new HPO for the learning rate of all three optimizers under this new setting:\n\nAdamW: 1r in $\\{5 e^{-4},1 e^{-3},5 e^{-3},1 e^{-2}\\}$\n\nMuon \\& ASGO: 1r in $\\{1 e^{-3},5 e^{-3},1 e^{-2},5 e^{-2}, 1 e^{-1}\\}$\n\n| Optimizer | **315M** | **630M** | **945M** | **1.26B** |\n| :--- | :---: | :---: | :---: | :---: |\n| AdamW | 4.71 | 4.05 | 3.79 | 3.73 |\n| Muon | 4.17 | 3.77 | 3.62 | 3.58 |\n| **ASGO** | 4.11 | 3.74 | 3.58 | 3.53 |\n\nTable 1: Training loss at different checkpoints for the 125M GPT-2 model on OpenWebText.\n\nAs Table 1 clearly indicates, the structure-aware preconditioning methods (Muon and our ASGO) consistently achieve a lower final training loss compared to the vector-based AdamW optimizer when trained on a larger token budget. This new experiment provides strong empirical evidence for our paper's central thesis: that leveraging the matrix structure of parameters is beneficial for training DNNs. We will replace our previous results with this new, more rigorous experiment in the revised manuscript. We plan to conduct more experiments and include these in future revisions. \n\n**On weakness 9: Runtime comparison** \n\nFirst, our practical implementation for all experiments follows Algorithm 2, which directly computes $\\left(V_t+\\epsilon I_n\\right)^{-1 / 2}$ instead of the matrix inverse. This is mathematically identical to the formulation in Algorithm 1.\n\nSecond, our initial version of ASGO relied on SVD to compute $\\left(V_t+\\epsilon I_n\\right)^{-1 / 2}$. This approach carries a significant computational burden. For the new experiments presented in this rebuttal, we have adopted the Newton-Schultz (NS) iterative method, which substantially reduces the wall-clock time (WCT) of the optimizer step and makes the algorithm more practical. The following runtime analysis is based on this improved NS-based implementation.\n\nAs shown in [5, Thm. 5.2], computing $A^{-1 / 2}$ can be achieved by computing the matrix of a block matrix \n$$\n\\operatorname{sign}\\left(\\left[\\begin{array}{cc}\n0 & A \\\\\\\\\nI & 0\n\\end{array}\\right]\\right)=\\left[\\begin{array}{cc}\n0 & A^{1 / 2} \\\\\\\\\nA^{-1 / 2} & 0\n\\end{array}\\right],\n$$\nwhich is the same operation as computed in Muon. This implies that computing the square root inverse in ASGO shares the same order of complexity as computing the matrix sign function in Muon. \n\nTo provide a concrete analysis, we measured the wall-clock time under the exact settings of our 1.26 B token experiment.\n| **Optimizer** | **Single Train Step WCT** | \n| :--- | :---: | \n| AdamW | 5.451 sec | \n| Muon | 5.519 sec | \n| **ASGO (Ours)** | 5.947 sec | \n\nTable 2: Wall-clock time per training step.\n\nWhile ASGO's optimizer step is more expensive than AdamW, the forward/backward pass is the dominant bottleneck. The table 3 shows total end-to-end training time is only marginally affected. \n\n| **Metric** | **AdamW** | **Muon** | **ASGO (Ours)** |\n| :--- | :---: | :---: | :---: |\n| WCT | 127.76 min | 129.01 min | 133.50 min |\n| Train Loss | 3.73 | 3.58 | 3.53 |\n\nTable 3: The total Wall-clock Time (WCT) and final training loss for training 1200 steps.\n\nWe think this additional experiment also addresses the two questions of the reviewer.\n\n**On weakness 10: Public code**\n    \nWe apologize for the confusion caused by the checklist. As stated in our justification for Checklist 5, we will make the complete source code for our experiments publicly available in a GitHub repository upon acceptance of the paper.\n\nWe believe that we have addressed the weaknesses and questions mentioned in the review, and hope that it convinces the reviewer to change the negative view of our paper.\n\n**References**\n\n[1] R.-y. Sun. Optimization for Deep Learning: An Overview\n\n[2] J. Duchi et al. Adaptive Subgradient Methods for Online Learning and Stochastic Optimization\n\n[3] K. Jordan et al. Muon: An optimizer for hidden layers in neural networks\n\n[4] J. Bernstein te al. Old Optimizer, New Norm: An Anthology\n\n[5] N. J. Higham. Functions of Matrices: Theory and Computation\n\n[6] F. Liao et al. Provable Accelerated Convergence of Nesterov’s Momentum for Deep ReLU Neural Networks\n\n[7] A. Banerjee et al. Restricted strong convexity of deep learning models with smooth activations"}, {"review_id": "CBlkA3FpLk", "reviewer": "Reviewer_hvYf", "summary": "The authors propose ASGO (Adaptive Structured Gradient Optimization), a new structured gradient method that exploits the matrix structure of gradients, including their low-rank and block-wise diagonal properties. ASGO improves the convergence guarantees of the existing algorithm Shampoo and is compared with the existing algorithm Muon. Experiments on language model tasks show that ASGO performs competitively, matching or outperforming methods such as AdaGrad, Adam, Muon, and Shampoo.", "questions": "1.line 113, in the update equation: would it be clearer to use  $\\Lambda_t$  instead of $V_t$, since Algorithm 1 (line 8) uses $\\Lambda_t$, which is a slight modification of $V_t$? \n2.line 239: you claim an equivalence between ASGO and Muon (without momentum) and draw an analogy to AdaGrad and SignSGD. Could the authors provide a more concrete explanation for this equivalence?\n3.line 258, Remark 4: In the smooth setting, how the convergence result of Muon compared with your result in Theorem 3?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths: \nThe paper presents clear and well-structured discussions. The convergence analysis and experimental evaluations are thorough and well-executed, particularly the comparisons with existing methods.\n\nWeaknesses: \n1.Although ASGO shows comparable performance to Muon and other baselines, the improvements in final accuracy or perplexity are not consistently significant.\n2.The discussion of nonconvex convergence is mostly conjectural, with key claims for ASGO inferred from Muon’s behavior rather than being directly proven.", "was_revised": true, "final_justification": "Thanks for the authors' responses. According to other reviewers' comments and the authors' responses, I still remain the original score for this paper.", "rebuttal": "We sincerely appreciate the reviewer's recognition of our paper. Many thanks! Below, we address the weakness and question parts of the review.\n\n**Weakness 1: The experimental results are not consistently significant.**\n\nWe thank the reviewer for their insightful comment. You are right to point out the significance of the empirical results. \nThe central contribution of this paper is to establish the theoretical motivation and empirical viability of structure-aware preconditioning. Specifically, we want to demonstrate why and how a structure-aware optimizer can outperform a structure-agnostic one like Adam. Based on this motivation, ASGO is designed to overcome this by employing preconditioners that explicitly model the rich interactions within a parameter tensor (e.g., the row-and-column structure of a matrix), which a structure-agnostic optimizer (like AdamW) cannot.\n\nTo provide stronger empirical validation for this core thesis, we have conducted a new, more rigorous pre-training experiment for a 125M parameter GPT-2 model on OpenWebText, increasing the token budget to 1.26 billion tokens more than double our original run. This budget represents approximately 10 tokens per parameter, which attains the limit of our computational resources. We adopted the Newton-Schultz(NS) algorithm [1, Thm. 5.2] for 10 steps instead of Singular Value Decomposition (SVD) to compute the inverse square root of preconditioning. which substantially reduces the wall-clock time of the optimizer step. \n\n| **Optimizer** | **315M** | **630M** | **945M** | **1.26B** |\n| :--- | :---: | :---: | :---: | :---: |\n| AdamW | 4.71 | 4.05 | 3.79 | 3.73 |\n| Muon | 4.17 | 3.77 | 3.62 | 3.58 |\n| **ASGO (Ours)** | 4.11 | 3.74 | 3.58 | 3.53 |\n\nTable 1: Training loss at different checkpoints for the 125M GPT-2 model on OpenWebText. The results are from the extended 1.26B token training run.\n\nAfter a fair hyperparameter search, the results (Table 1) show that structure-aware optimizers (our ASGO and Muon) achieve a consistently lower final training loss than the structure-agnostic AdamW. This new result directly validates our paper's central argument: that leveraging parameter structure leads to a better-trained neural network.\n\nFurthermore, we analyzed the practical wall-clock time by using the efficient Newton-Schultz (NS) method (10 steps) to compute the preconditioner's inverse square root. While the optimizer step for ASGO is computationally more intensive than AdamW, the total training time is remarkably similar across all methods, as shown in Table 2. This demonstrates that the per-step overhead of ASGO is not a bottleneck for overall training throughput.\n\n| **Metric** | **Adam** | **Muon** | **ASGO (Ours)** |\n| :-------------- | :----------: | :----------: | :-------------: |\n| WCT             | 127.76 min   | 129.01 min   | 133.50 min      |\n| Train Loss      | 3.73         | 3.58         | 3.53            |\n\nTable 2: The total Wall-clock Time(WCT) and final train loss for training 1200 steps\n\nWe believe that by leading with our theoretical motivation and supporting it with these new, stronger empirical results, we have addressed the reviewer's concern. We will update the manuscript with this new analysis. We plan to conduct more experiments\nand include these in future revisions.\n\n**Weakness 2: The discussion of nonconvex convergence is mostly conjectural.**\n\nThe nonconvex convergence of ASGO is indeed conjectured from Muon's convergence guarantee, based on the close relation between the algorithms. We think it is an interesting future research direction, and we will actively work on this. Thanks a lot for raising this point!\n\n**Question 1: Use $\\Lambda_t$ instead of $V_t$ in line 113.**\n\nWe thank the reviewer for the careful reading and the helpful suggestion. We will modify the formulation in line 113 and add a more detailed definition in future revisions accordingly.\n\n**Question 2: More discussions on the relation between Muon and ASGO.**\n\n\nThank you for this excellent question, which allows us to elaborate on a key intuition behind our work. First, we want to give some intuition on the relationship between Muon and SignSGD. As analyzed in [2], both Muon (without momentum) and SignSGD (without momentum) can be interpreted in the steepest descent framework (Equation (2) in the paper), but with respect to different matrix norms. This provides a clear theoretical connection. Meanwhile, both methods can be seen as performing whitening on the first-order momentum. SignSGD performs element-wise whitening while Muon performs spectral whitening.\n\nThen we kindly lead you to the connection between ASGO and Muon. In Section 6, we note that ASGO without accumulating $GG^\\top$ in the preconditioner is equivalent to Muon without momentum, and provide a corresponding proof in Appendix G. To be more concrete, this equivalence comes from the equation $(G_tG_t^\\top)^{-1/2} G_t = U_t V_t^\\top$, where $U_t, V_t$ are the left and right singular vectors in the compact singular value decomposition of $G_t$, i.e., $G_t = U_t \\Sigma_t V_t^\\top$. Intuitively, ASGO uses a smoother preconditioning than Muon, since it updates by $(\\sum_{s=1}^{t} G_sG_s^\\top)^{-1/2} G_t$ instead of directly using $(G_tG_t^\\top)^{-1/2}$ as the preconditioner. \nThis is similar to the %relation\nrelationship\nbetween AdaGrad and SignSGD, since if we remove the accumulation in the preconditioner of AdaGrad and the momentum of SignSGD, we can also obtain equivalence between the two algorithms:\n$$\n\\text{AdaGrad:} \\quad w_{t+1} = w_{t} - \\eta_t (g_t^2)^{-1/2} g_t \n$$\n$$\n\\text{SignSGD:} \\quad w_{t+1} = w_t - \\eta_t \\text{sign}(g_t) \n$$\nwhere $(g_t^2)^{-1/2}$ here denotes coordinate-wise operations. Based on this similarity between the relations, we say the connection between ASGO and Muon is analogous to the connection between AdaGrad and SignSGD.\nThanks for this insightful question, and we will also add the discussion context to the paper to make it clearer.\n\n**Question 3: How does the convergence of Muon compare with Theorem 3.**\n\nThanks for this important question. Though the results are both obtained under smooth settings, we kindly note that the two results are generally incomparable, since Theorem 3 indicates the loss convergence of ASGO under convex settings, while Equation (3) introduces the convergence rate of Muon to a stationary point under nonconvex settings. We consider the nonconvex convergence of ASGO as an interesting topic for future work, and expect that it should be on par with the convergence rate in Equation (3) for Muon.\n\n**References**\n\n[1] Nicholas J. Higham. Functions of Matrices: Theory and Computation.\n\n[2] Jeremy Bernstein, Laker Newhouse. Old Optimizer, New Norm: An Anthology"}]

Author rebuttal:

Meta-review: This paper proposes ASGO, a new adaptive approach for stochastic gradient optimization. ASGO differs from prior work Shampoo and K-FAC in that it uses a single matrix conditioner to condition the gradient updates of matrix-valued parameters. The authors provided a convergence theory of this in convex problems, and shown that the convergence bounds of this approach might be even better than Shampoo-style preconditioning and comparable with Muon but with better theoretical guarantees.

Reviewers raised the questions about empirical timing and the authors provided a version in the rebuttal that utilizes Newton-Schultz iterations to compute the matrix inverse square-root. After rebuttal, 3 reviewers are on the positive side, with ratings of 1 Accept and 2 borderline accepts. One reviewer still stands on the negative side. The main remaining issues are that the algorithm is only tested on a single LLM benchmark, and that the code was not available. 

The AC believes that the availability of the code should not be used to hold against authors. And also that the paper's contribution is of theoretical nature hence can enjoy some leniency on the results. Thus, AC recommends acceptance of this paper.
Decision: Accept (poster)