Title: ASGO: Adaptive Structured Gradient Optimization
Venue: NeurIPS 2025
Accepted: yes
Keywords: optimization

optimization theory

adaptive gradient method

optimization for deep learning

convex optimization
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.


Peer reviews:
comments: 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.

STRENGTHS AND WEAKNESSES:
Strengths: 
The paper presents clear and well-structured discussions. The convergence analysis and experimental evaluations are thorough and well-executed, particularly the comparisons with existing methods.

Weaknesses: 
1.Although ASGO shows comparable performance to Muon and other baselines, the improvements in final accuracy or perplexity are not consistently significant.
2.The discussion of nonconvex convergence is mostly conjectural, with key claims for ASGO inferred from Muon’s behavior rather than being directly proven.

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$? 
2.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?
3.line 258, Remark 4: In the smooth setting, how the convergence result of Muon compared with your result in Theorem 3?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n### 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.\n\nQUESTIONS:\n- 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).\n\nLIMITATIONS:\nNo, the paper does not cover the cost of the matrix inversion and how it behaves on larger models.", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 3, "CLARITY": 1, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS AND WEAKNESSES:\nOverall, 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\"\n\nQUESTIONS:\nBesides 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\"\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n### 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).\n\nQUESTIONS:\n### 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$.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}