PaperlayerView exact revision on arXiv

Source-linked AI summary

ASGO: Adaptive Structured Gradient Optimization

Kang An, Yuxing Liu, Rui Pan, Yi Ren, Shiqian Ma, Donald Goldfarb, Tong Zhang

arXiv:2503.20762v3cs.LGmath.OC

TL;DR

Deep neural-network optimizers often overlook low-rank gradients and approximately block-diagonal Hessians. ASGO uses an adaptively updated structured preconditioner and achieves better convergence guarantees than full-matrix AdaGrad and Shampoo, with empirical validation on language-model tasks.

  • Problem

    Popular optimizers often overlook the low-rank gradients and approximately block-diagonal Hessians observed in deep neural networks.

  • Method

    ASGO adaptively updates a single-matrix preconditioner using structured gradients while preserving matrix structure and reducing memory versus full-matrix methods.

  • Results

    ASGO achieves faster convergence than full-matrix AdaGrad and Shampoo and is empirically effective and efficient on language-model tasks.

  • Takeaways & Limitations

    ASGO provides theoretical and empirical evidence that structured optimization can exploit low-rank gradients and approximately block-diagonal Hessians in neural-network training.

  • Takeaways & Limitations

    DASGO underperformed AdamW, Muon, Shampoo, and ASGO on GPT2 and NanoGPT pretraining, likely because diagonal preconditioning discards parameter-gradient interactions.

Abstract

from arXiv · show

Training deep neural networks is a 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 a novel optimization algorithm ASGO that capitalizes on these properties by employing a preconditioner that is adaptively updated using structured gradients. By a fine-grained theoretical analysis, ASGO is proven to 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://github.com/infinity-stars/ASGO.

1 Introduction

The introduction argues that common diagonal adaptive optimizers do not fully match the structured properties of DNN optimization, while Shampoo incurs unfavorable convergence, memory, and computation costs. It presents ASGO as a structured-gradient method designed to improve Shampoo’s convergence guarantees and exploit low-rank gradients and approximately block-diagonal Hessians.

  • Motivation: Adaptive gradient methods use diagonal preconditioners and can exploit sparse gradients or diagonally scaled Hessians, but common DNNs do not necessarily have sparse gradients.The passage motivates moving beyond coordinate-wise preconditioning because common DNN gradients may not be sparse.
  • Limitations of Shampoo: Shampoo uses structured matrix preconditioners, but its convergence is worse than AdaGrad or even SGD in high dimensions and it requires more memory and computation.Its heavier cost arises from maintaining two preconditioners.
  • ASGO: ASGO is proposed as a structured-gradient algorithm that improves Shampoo’s convergence guarantees while requiring less memory and computation.The paper positions ASGO as an answer to the limitations of Shampoo.
  • Theoretical contributions: The analysis shows that ASGO converges faster than full-matrix AdaGrad and Shampoo and can exploit low-rank gradients and approximately block-diagonal Hessians.These properties are described as typically observed during DNN training and theoretically useful under realistic settings.
  • Practical implementation: The paper also develops a practical ASGO implementation intended to improve efficiency on large-scale neural-network training tasks.The supplied contribution list introduces this implementation but truncates its further description.

2 Related Work

Prior work includes adaptive gradient methods with diagonal preconditioners and increasingly popular matrix-structured optimizers for machine learning. Related structured methods include Adafactor, LAMB, Adam-mini, Shampoo, and TNT, with Shampoo’s convergence rate matching full-matrix AdaGrad.

  • Adaptive Gradient Methods: Adaptive gradient methods use diagonal preconditioners to accelerate convergence and are widely used in real-world optimization.AdaGrad was an early machine-learning method of this kind, developed with rigorous convergence theory.
  • Optimization with Matrix Structure: Matrix-structured optimization methods for machine learning have gained increasing attention as alternatives to vector-based formulations.Adafactor, LAMB, and Adam-mini consider matrix or layer structure.
  • Optimization with Matrix Structure: Shampoo’s convergence rate is no better than that of full-matrix AdaGrad.This result follows the analysis cited in the passage.
  • Optimization with Matrix Structure: TNT is closely related to Shampoo but uses true factor matrices and inverses of their Kronecker factors in its natural-gradient construction.TNT approximates the true Fisher matrix using covariance from block-wise sampling-based gradients under a Tensor-Normal assumption.

3 Our ASGO Algorithm

ASGO is an adaptive structured-gradient optimizer with a single-side preconditioner that preserves matrix structure while reducing preconditioner storage relative to full-matrix AdaGrad and Shampoo. Its design trades approximation quality for the ability to exploit low-rank gradients and block-diagonal Hessians.

  • Algorithm design: ASGO uses a single-side preconditioner that preserves the matrix structure of W_t and G_t, avoiding full-matrix AdaGrad’s huge preconditioner memory cost.The algorithm is introduced as Adaptive Structured Gradient Optimization with a matrix-structured update.
  • Algorithm design: ASGO stores one preconditioner matrix, whereas Shampoo uses two matrices.ASGO computes the preconditioner’s matrix square root and inverse on each iteration.
  • Algorithm design: ASGO’s preconditioner may approximate the full-matrix AdaGrad preconditioner or empirical Fisher less accurately than Shampoo’s, while exploiting low-rank gradients and block-diagonal Hessians.The tradeoff arises from using a single matrix rather than Shampoo’s two-matrix preconditioner.

4 Nonsmooth Theory

Under convexity, ASGO attains an O(1/√T) convergence rate for nonsmooth problems, matching SGD and AdaGrad while exploiting matrix structure. Its advantage is strongest when gradients are low-rank and weight differences are relatively high-rank, yielding faster theoretical rates than full-matrix AdaGrad and Shampoo.

  • Nonsmooth convergence: Under convexity, ASGO achieves an O(1/√T) convergence rate for convex nonsmooth problems, matching SGD and AdaGrad.The result follows from Theorem 1 and is summarized as the same rate as SGD and AdaGrad.
  • Practical limitation: The analysis treats Dop as constant, although it may increase with T; projection onto a bounded convex set can keep it bounded but is omitted because such projection is rarely used in DNN training.The omission follows the practical convention cited in the paper, while the bounded-set projection is discussed as a remedy.
  • Rank-dependent advantage: ASGO benefits from low-rank gradients because ∥Q∥F ≤ ∥Q∥∗ ≤ √rG∥Q∥F, while high-rank weight differences can make Dop much smaller than DF.The bounds imply ∥Q∥∗ can approach ∥Q∥F for low-rank or spectrally imbalanced gradients, and Dop can shrink when weight differences are relatively high-rank.
  • Rank-dependent advantage: Consequently, ASGO should work well when Gt are low-rank and Wt − W∗ are relatively high-rank.The paper argues that these structural conditions are common in practical deep-learning tasks.
  • Practical limitation: Although LoRA uses low-rank total updates, the paper notes that W0 − W∗ can be relatively high-rank for pretraining and some complex foundation-model fine-tuning tasks.This observation addresses the apparent tension between low-rank adaptation methods and the theory’s favorable high-rank weight-difference condition.
  • Comparison with structured optimizers: Theorem 1 gives ASGO a convergence speed at least DF/Dop times that of full-matrix AdaGrad and √rGDF/Dop times that of Shampoo.This comparison provides theoretical evidence that single-side preconditioning can exploit gradient low-rankness more effectively than Shampoo-like preconditioning.

5 Smooth Theory

The smooth theory models curvature with block-wise matrix smoothness and shows that ASGO’s convergence benefits from low-rank stochastic noise and favorable Hessian structure. It further identifies conditions under which ASGO performs well and motivates right-sided preconditioning for neuron-wise Hessian blocks.

  • Block-wise smoothness: Assumption 2 requires f to be 1-smooth with respect to a symmetric positive definite matrix L.This is equivalent to a block-wise Hessian bound −L ⪯ ∇²f_v(w) ⪯ L and implies standard Frobenius-norm smoothness.
  • Smooth convergence: Theorem 3 establishes ASGO’s smooth convergence rate under Assumptions 1, 2, and 3.The rate has terms involving batch size M, stochastic-noise quantity V, and the nuclear norm ∥L∥_*; the supplied passage does not include the complete bound.
  • Smooth convergence: When M is small, ASGO shares the Section 4 rate properties and benefits when stochastic gradient noise V is generally low-rank.In this regime, the O(1/(MT)) term dominates the convergence bound.
  • Structural implications: ASGO should work well when the Hessian is approximately block diagonal with low-rank L in each block and W_t − W_* is relatively high-rank.The supporting inequalities are ∥L∥_op ≤ ∥L∥_* ≤ r_L∥L∥_op and D_F/√r_D ≤ D_op ≤ D_F.
  • Preconditioner placement: Because Hessian blocks commonly correspond to neuron-input weights arranged as rows of W, ASGO applies its single-side preconditioner on the right.The paper also reports that right-sided preconditioning outperforms left-sided preconditioning empirically in Section 7.1.2.

6 Further Discussions on ASGO

ASGO is theoretically and algorithmically connected to Muon, becoming equivalent when momentum is ignored, while its nonconvex convergence rate remains a conjecture. Practical implementations include memory-efficient DASGO and Newton–Schultz-based inverse-square-root computation, with complexity comparable to Muon but slightly higher per iteration.

  • Connection with Muon: Ignoring momentum in both methods’ gradient and preconditioner makes ASGO equivalent to Muon, reflecting shared exploitation of matrix structure and structured gradients.Muon uses spectral-norm steepest descent with momentum, while ASGO’s equivalence holds without momentum.
  • Nonconvex convergence: ASGO is expected to have a nonconvex convergence rate comparable to Muon’s up to logarithmic factors, but proving this remains future work.The expectation assumes f* = inf f(W) > −∞ and draws on matching rates for diagonal AdaGrad and SignSGD up to logarithmic factors.
  • Practical Implementations of ASGO: Algorithm 2 provides a practical ASGO implementation, while DASGO diagonalizes the preconditioner for lower memory and computational costs.DASGO is described as a lightweight optimizer based on a diagonalized Λ_t.
  • Practical Implementations of ASGO: Newton–Schultz iteration replaces prohibitively expensive SVD for computing the preconditioner’s inverse square root, substantially reducing wall-clock time per step.The approach relies on the connection between matrix inverse square roots and matrix sign functions.
  • Practical Implementations of ASGO: ASGO’s inverse-square-root computation has complexity on the same order as Muon’s matrix operation, though its 2 × 2 block-matrix iteration is slightly more expensive per step.The framework is modular and can incorporate more efficient matrix-sign-function methods.

7 Empirical Results

ASGO matches Muon on GPT-2 pretraining while converging faster than AdamW and avoiding Shampoo’s loss spikes, whereas DASGO underperforms because diagonal preconditioning discards gradient interdependencies. On GPT2-Large fine-tuning, both ASGO and DASGO outperform Muon and AdamW under CLM and FIM objectives.

  • GPT-2 pretraining: ASGO and Muon follow nearly identical training trajectories and converge faster than AdamW, while Shampoo exhibits loss spikes suggesting numerical instability.The instability occurs even with preconditioner update frequency 1 and is not observed in ASGO.
  • GPT-2 pretraining: DASGO records 3.509 and underperforms AdamW, Muon, Shampoo, and ASGO on GPT2 pretraining, likely because diagonal preconditioning discards interdependencies between parameter gradients.The same performance gap also appears in smaller NanoGPT pretraining.
  • Wall-clock time: Total training times for AdamW, Muon, and ASGO are very similar despite AdamW being cheapest at the optimizer level, while ASGO achieves the smallest training loss.ASGO’s optimizer step has additional inverse-square-root computation, but this does not make total training time a main concern in the reported comparison.
  • Preconditioning-side ablation: The ablation study confirms that preconditioning side has an asymmetric impact, supporting right-sided preconditioning and suggesting limitations in Shampoo-like double-sided preconditioners.The paper attributes possible suboptimal performance and training instabilities in Shampoo to structurally misapproximating curvature and placing preconditioning on the less critical side.
  • GPT2-Large fine-tuning: For GPT2-Large fine-tuning, ASGO and DASGO outperform Muon and AdamW under both CLM and FIM objectives.Under CLM, ASGO achieves 13.88 perplexity, DASGO 13.84 perplexity, Muon 13.91 perplexity, and AdamW 14.01 perplexity; under FIM, ASGO achieves 15.66 perplexity, DASGO 15.45 perplexity, Muon 15.93 perplexity, and AdamW 17.46 perplexity.

8 Conclusions

The paper concludes that ASGO improves convergence rates over full-matrix AdaGrad and Shampoo while exploiting low-rank gradients and block-wise diagonal Hessians. Practical modifications and empirical results further support its effectiveness, though the passage notes remaining limitations.

  • ASGO achieves significantly better convergence rates than full-matrix AdaGrad and Shampoo.
  • ASGO benefits from low-rank gradients and block-wise diagonal Hessians, structural properties widely observed in deep neural networks.
  • The authors propose practical modifications to ASGO and verify its effectiveness empirically.

NeurIPS Paper Checklist

The paper’s checklist responses state that it presents assumptions and proofs, discusses limitations, and documents experimental procedures, code/data access, statistical variability, and compute resources for reproducibility.

  • Claims and limitations: The authors state that the paper discusses the proposed algorithm’s limitations in the conclusion.They also state that the abstract and introduction present and summarize the main contribution.
  • Theory assumptions and proofs: The authors state that assumptions and proofs are provided for all theorems.This addresses the checklist requirement for complete theoretical-result documentation.
  • Experimental reproducibility: Experiment settings and hyperparameter-tuning details are specified in Appendix C to make the experiments reproducible.The paper also states that the code will be publicly available and that all datasets are publicly available and cited.
  • Experimental setting/details: All training and test details, including model settings, hyperparameter selection, and optimizer configurations, are specified in Appendix C.The checklist response marks the experimental-setting question “Yes.”
  • Experiment statistical significance: The paper reports fixed random seeds or mean training performance across different random seeds to account for statistical variability.The response marks the statistical-significance question “Yes.”
  • Experiments compute resources: The authors state that they report the computer resources used for the experiments in Section 7.The checklist response marks the compute-resources question “Yes.”

A Additional Related Work … C.2 Pretraining GPT2

The paper situates ASGO within structured optimization research, describes practical and architecture-specific variants, and reports language-model experiments spanning NanoGPT and GPT-2.

  • A Additional Related Work: Prior work motivates ASGO through low-rank gradients, block-wise diagonal Hessians, and matrix-structured optimizer developments.Gradients are widely observed to be low-rank, while neural-network Hessians tend to be block-wise diagonal; recent work also improves full-matrix AdaGrad and Shampoo.
  • A Additional Related Work: ASGO is compared with concurrent One-Sided Shampoo work, which is described as algorithmically identical but differently motivated and evaluated.The paper emphasizes its focus on structured optimization properties and practical deep-learning tasks, whereas the concurrent work focuses on convex settings and linear regression.
  • B More Discussions on ASGO: The practical ASGO implementation uses momentum-based exponential moving averages, Newton-Schulz matrix inverse-square-root computation, adaptive preconditioning sides, and update alignment.The preconditioning side is selected according to the smaller gradient-matrix dimension, and updates are RMS-normalized and scaled by 0.2 to match Adam updates.
  • B More Discussions on ASGO: ASGO includes specialized processing for Transformer query and key matrices because their Hessian block structure differs from conventional MLP layers.The dense-block count for query and key matrices corresponds to attention heads rather than output neurons.
  • B More Discussions on ASGO: DASGO provides a lightweight diagonal ASGO variant that avoids full-matrix inverse-square-root computation and reduces memory requirements to a level comparable with Adam-mini.The diagonal variant also makes the choice of preconditioning side unimportant in terms of computation.
  • C.1 Pretraining NanoGPT: In NanoGPT pretraining, ASGO consistently outperforms Shampoo while using half its memory and computational effort, and ASGO and Muon achieve the lowest final losses.The experiment compares ASGO, DASGO, Muon, Shampoo, and AdamW on the Shakespeare character-level dataset.
  • C.1 Pretraining NanoGPT: NanoGPT ablations evaluate specialized query-key processing by comparing ASGO and Muon with and without that design under matched optimal hyperparameters.Learning-rate robustness is assessed by sampling rates from 10^-5 to 10^-1 and recording validation loss after 5 epochs.
  • C.2 Pretraining GPT2: GPT-2 pretraining uses a 12-layer, 12-head, 768-dimensional Transformer trained for 2400 steps on 4 NVIDIA GH200 GPUs, with grid-search tuning of learning rate and β2.Optimal hyperparameters are selected using validation loss after 1000 training steps and reported in Table 6.

C.3 Finetune GPT2-Large · D Auxiliary Lemmas for the Proof

The paper evaluates ASGO and related optimizers by fine-tuning GPT2-Large on WikiText-2 and studying update-direction sensitivity to damping and matrix approximations. It also collects auxiliary trace, operator-monotonicity, spectral, norm, and sequence inequalities used in the proofs.

  • C.3 Finetune GPT2-Large: GPT2-Large fine-tuning uses the 774M-parameter model on WikiText-2 under causal language modeling and fill-in-the-middle objectives.All models are trained for 2 epochs with batch size 16 and sequence length 128; learning rates follow cosine annealing without warmup.
  • C.3 Finetune GPT2-Large: Learning-rate tuning searches six candidates separately for each optimizer and objective, selecting the value with the lowest validation metric.The candidates are [1×10−6, 5×10−5, 1×10−4, 5×10−4, 1×10−3, 5×10−3].
  • C.3 Finetune GPT2-Large: Table 8 reports average perplexity after 2-epoch GPT2-Large fine-tuning, with ± denoting 95% confidence intervals across five random-seed runs.The provided passage identifies the evaluation setup but does not include the table’s numerical perplexity values.
  • C.3 Finetune GPT2-Large: ASGO theoretically degenerates to Muon when β1 = 0 and β2 = 0, and becomes identical to Muon with momentum when β2 = 0 under modified formulas.In practical implementations, ASGO additionally applies ϵ damping as a preconditioner to the momentum term.
  • C.3 Finetune GPT2-Large: Without damping, ASGO approximately reconstructs Muon’s update direction, with cosine similarities generally exceeding 0.8, especially in MLP and Embedding layers.The undamped direction shows some instability, while ϵ = 10−8 significantly changes the similarities.
  • C.3 Finetune GPT2-Large: Coupled Newton and Newton-Schulz approximate ASGO’s matrix computation, producing highly comparable cosine similarities to SVD outside some Embedding-layer stability issues.The passage attributes the Embedding behavior to sparse embeddings potentially requiring larger damping parameters.

E Proof of Nonsmooth Convergence of ASGO

The section proves ASGO’s nonsmooth convergence by bounding the terms in its potential inequality, then compares its proven rate with full-matrix AdaGrad and Shampoo. The comparisons show speedups of DF/Dop and √rGDF/Dop, respectively.

  • Nonsmooth convergence proof: ASGO’s convergence proof bounds the terms in its main inequality using trace identities, norm duality, and positive semidefiniteness.The proof separately handles the second and third terms before substituting the resulting bounds into the main inequality.
  • Comparison with Shampoo and full-matrix AdaGrad: ASGO’s proven convergence rate is at least DF/Dop times faster than full-matrix AdaGrad.The comparison uses Lemma 4 and the inequality ∥x∥2 ≥ ∥x∥1.
  • Comparison with Shampoo and full-matrix AdaGrad: ASGO’s proven rate is at least √rGDF/Dop times faster than Shampoo.The derivation uses duality of ∥·∥op and ∥·∥∗ together with trace and singular-value inequalities for symmetric positive semidefinite matrices.

F Proof of Smooth Convergence of ASGO

The smooth-convergence proof develops separate bounds for gradient and noise, establishes a matrix norm bound and batch-size variance reduction, and uses them to prove Theorem 3 by treating bias and variance separately.

  • Technical lemmas: Smooth analysis begins with Theorem 1 and separates the gradient from the noise term through Lemma 7.The section denotes the noise by N_t ≜ G_t − ∇f(W_t) and the gradient by ∇f_t ≜ ∇f(W_t).
  • Technical lemmas: Lemma 8 provides an upper bound on ∥·∥_* for a symmetric positive definite matrix Λ and matrix G.The proof first handles diagonal Λ and then extends the result to general Λ using singular value decomposition.
  • Technical lemmas: Lemma 9 establishes variance reduction by batch size M under Assumption 3.Its proof uses mutual independence, zero-mean noise, and the variance assumption.
  • Proof of Theorem 3: Theorem 3 is proved by separately bounding the bias and variance terms appearing on the right-hand side of Lemma 7.The bias bound applies the smoothness matrix L in Lemma 8, while the variance argument first treats M = 1 and then incorporates M > 1 using Lemma 9.

G Proof of Section 6 · H Incorporating Projection in ASGO

Section 6 proves that ASGO reduces to Muon under specific parameter settings and establishes a deterministic nonconvex convergence result for Muon. Projection extends ASGO’s convergence guarantees while keeping operator- and Frobenius-norm distances uniformly bounded by constants determined by the feasible set.

  • G Proof of Section 6: Muon uses singular vectors of the gradient and approximates their product with Newton–Schulz iterations instead of directly computing an SVD.This implementation is intended to improve computational efficiency.
  • G Proof of Section 6: ASGO is equivalent to Muon without gradient accumulation or momentum when β1 = β2 = 0, ϵ = 0, τ = 1, and µ = 0.The equivalence is stated for the Algorithm 2 version of ASGO and Muon in equation (11).
  • G Proof of Section 6: The paper then proves a deterministic nonconvex convergence theorem for Muon with µ = 0.The theorem assumes a global lower bound f(W) ≥ f ∗ and Assumption 2.
  • G Proof of Section 6: In the deterministic setting, Gt equals ∇f(Wt), and with µ = 0 the singular-vector factors reconstruct Gt through its compact SVD.The proof uses these identities together with the smoothness assumption and a matrix inequality.
  • G Proof of Section 6: The resulting proof sums the smoothness-based inequality over iterations and rearranges it to complete the convergence argument.The displayed conclusion includes the factor 2(f(W0)−f ∗) ∥L∥∗T.
  • H Incorporating Projection in ASGO: To obtain rigorous convergence bounds with constants that remain fixed as T increases, the paper incorporates projection onto a closed convex set W containing an optimum W∗.The set is characterized by operator- and Frobenius-norm diameters ˜Dop and ˜DF.
  • H Incorporating Projection in ASGO: Projection preserves the key proof inequality, yielding the same nonsmooth and smooth convergence results as Theorems 1 and 3.The projection map Λt maps Rm×n into W, and its non-expansiveness is substituted into the proof.
  • H Incorporating Projection in ASGO: The projected iterates satisfy max0≤t≤T −1 ∥Wt −W∗∥op ≤˜Dop and max0≤t≤T −1 ∥Wt −W∗∥F ≤˜DF, with diameters remaining constant as T increases.The constants depend on the geometry of the convex set W.
Loading 2503.20762v3…