ASGO: Adaptive Structured Gradient Optimization

We sincerely appreciate the reviewer's recognition of our paper. Many thanks! Below, we address the weakness and question parts of the review.

Weakness 1: The experimental results are not consistently significant.

We thank the reviewer for their insightful comment. You are right to point out the significance of the empirical results. The 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.

To 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.

Table 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.

After 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.

Furthermore, 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.

Table 2: The total Wall-clock Time(WCT) and final train loss for training 1200 steps

We 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 and include these in future revisions.

Weakness 2: The discussion of nonconvex convergence is mostly conjectural.

The 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!

Question 1: Use $\Lambda_t$ instead of $V_t$ in line 113.

We 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.

Question 2: More discussions on the relation between Muon and ASGO.

Thank 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.

Then 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. This is similar to the %relation relationship between 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:

where $(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. Thanks for this insightful question, and we will also add the discussion context to the paper to make it clearer.

Question 3: How does the convergence of Muon compare with Theorem 3.

Thanks 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.

References

[1] Nicholas J. Higham. Functions of Matrices: Theory and Computation.

[2] Jeremy Bernstein, Laker Newhouse. Old Optimizer, New Norm: An Anthology

We thank the reviewer for the constructive and detailed feedback.

On Weaknesses 1 & 2: Diagonal preconditioner

We 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. Let's consider the update rule for Adam:

where $m_t$ ($v_t$) is the first-order (second-order) momentum and the division is elementwise. This operation is identical to:

where $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 pre-multiplying the momentum gradient $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].

This 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. This is why we say Adam generally ignores the structure of parameters.

This 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.

On Weakness 3: Evaluation on benchmarks

The 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. This 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.

Furthermore, 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.

While 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.

On weakness 4: Definition of $V_t$

$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.

On weakness 5: Local convexity

Our 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.

On weakness 6: Muon updates

The 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.

On weakness 7: Algorithm 2 and 3

Algorithm 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.

On weakness 8: Larger training size

We 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.

We 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.

Although 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:

AdamW: 1r in $\{5 e^{-4},1 e^{-3},5 e^{-3},1 e^{-2}\}$

Muon & ASGO: 1r in $\{1 e^{-3},5 e^{-3},1 e^{-2},5 e^{-2}, 1 e^{-1}\}$

Table 1: Training loss at different checkpoints for the 125M GPT-2 model on OpenWebText.

As 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.

On weakness 9: Runtime comparison

First, 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.

Second, 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.

As shown in [5, Thm. 5.2], computing $A^{-1 / 2}$ can be achieved by computing the matrix of a block matrix

which 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.

To provide a concrete analysis, we measured the wall-clock time under the exact settings of our 1.26 B token experiment.

Table 2: Wall-clock time per training step.

While 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.

Table 3: The total Wall-clock Time (WCT) and final training loss for training 1200 steps.

We think this additional experiment also addresses the two questions of the reviewer.

On weakness 10: Public code

We 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.

We 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.

References

[1] R.-y. Sun. Optimization for Deep Learning: An Overview

[2] J. Duchi et al. Adaptive Subgradient Methods for Online Learning and Stochastic Optimization

[3] K. Jordan et al. Muon: An optimizer for hidden layers in neural networks

[4] J. Bernstein te al. Old Optimizer, New Norm: An Anthology

[5] N. J. Higham. Functions of Matrices: Theory and Computation

[6] F. Liao et al. Provable Accelerated Convergence of Nesterov’s Momentum for Deep ReLU Neural Networks

[7] A. Banerjee et al. Restricted strong convexity of deep learning models with smooth activations

Thank you very much for your careful reading and insightful opinions and questions, which we now address.

Weakness 1: $D_{op}$ may grow with $T$.

Thanks 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.

Weakness 2: The comparison between ASGO and SGD may be unfair

1. 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.

2. 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.

Major question 1: Choosing $\epsilon$ and the purpose

Your 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 strict 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.

However, 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.

After 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.

Major question 2: Why is $\Lambda_t$ placed on the left side in Algorithm 1?

We 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.

Major question 3: Nonconvex analysis of ASGO

The 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.

Major question 4: Wall-clock time analysis

We 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.

To 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 of 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.

As 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.

As shown in Thm 5.2 [6], computing $A^{-1 / 2}$ can be achieved by computing the sign function of a block matrix $\operatorname{sign}\left(\left[\begin{array}{cc} 0 & A \\\\ I & 0 \end{array}\right]\right)=\left[\begin{array}{cc} 0 & A^{1 / 2} \\\\ A^{-1 / 2} & 0 \end{array}\right].$ 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.

Empirical Wall-Clock Time Analysis. To provide a concrete analysis, we measured the WCT under the exact settings of our 1.26B token experiment.

Table 1: Wall-clock time per training step, difference compared with AdamW update.

• 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.

Although 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.

Table 2: The total Wall-clock Time (WCT) and final training loss for training 1200 steps.

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 a main concern for ASGO. Note that ASGO achieves the smallest training loss while AdamW achieves the largest loss.

Minor Suggestions

We thank the reviewer for carefully reading our submission and suggesting ways to improve its clarity, which we will follow.

References

[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).

[2] Elad Hazan, Amit Agarwal, and Satyen Kale. Logarithmic regret algorithms for online convex optimization. Machine Learning 69.2 (2007): 169-192.

[3] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. International Conference on Machine Learning. PMLR, 2018.

[4] Amit Attia and Tomer Koren. Sgd with adagrad stepsizes: Full adaptivity with high probability to unknown parameters, unbounded gradients and affine variance. International Conference on Machine Learning. PMLR, 2023.

[5] Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks

[6] Nicholas J. Higham. Functions of Matrices: Theory and Computation.

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!

Weakness 1: Conclusions are drawn by comparing performance upper bounds.

This 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. We 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. Hence, this is an interesting future topic and may inspire even more ways to improve ASGO.

Weakness 2: The theoretical set-up considers a single matrix optimization variable.

Thanks 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 recent 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.

Minor comments 1: The paper [2] for rank-structure / evolution in neural network weights.

We thank the reviewer for bringing this interesting and relevant paper to our attention. The 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!

Minor comments 2: The Muon bound (3) differs from the statement in Appendix G

Thanks 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:

Here, the norms are the trace norm $|| \cdot ||_{\*}$. Therefore, 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.

Question 1: The use of $\mu$P.

This 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)):

Here the norms are the spectral norm $||\cdot||_{op}$. Therefore, 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.

Question 2: ASGO for higher-dimensional tensors.

We 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 research topic.

References

[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).

[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.

[3] Greg Yang, James B. Simon, and Jeremy Bernstein. "A spectral condition for feature learning." arXiv preprint arXiv:2310.17813 (2023).

I thank the authors for their clarifications. I maintain my positive evaluation.