Research map / Optimization / Optimizers
ASGO: Adaptive Structured Gradient Optimization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes ASGO, a single-side (one-sided) full-matrix adaptive structured gradient optimizer that preconditions with an accumulated G G^T statistic, with convergence theory and GPT-2 pretraining/fine-tuning experiments against AdamW, Shampoo and Muon.
Problem
Adaptive optimizers like Adam ignore the matrix structure of DNN parameters, while Shampoo-style two-sided preconditioners are memory-heavy and have weaker convergence guarantees than full-matrix AdaGrad or SGD.
Approach
Use a single-sided full-matrix preconditioner V_t^{1/2} accumulated from structured gradients (G_t G_t^T), computed via a Newton-Schulz iteration on a 2x2 block matrix; theory shows it exploits low-rank gradients and block-diagonal Hessians, and it is shown to be the adaptive analogue of Muon.
Key results
- On GPT-2 pretraining (2.5B tokens), ASGO-PolarExpress reaches validation loss 3.342, matching Muon (3.342) and beating AdamW (3.455) and Shampoo (3.417).
- On GPT-2 Large CLM fine-tuning on WikiText-2, ASGO achieves perplexity 13.88 vs AdamW 14.01 and Muon 13.91.
- Right-sided preconditioning gives lower validation loss than left-sided at all four tested learning rates (e.g., 3.362 vs 3.399 at lr 5e-3).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use ASGO-PolarExpress instead of AdamW | GPT-2 pretraining · GPT-2 pretraining | validation loss: 3.455 -> 3.342 | Helped |
| Use ASGO-PolarExpress instead of Shampoo | GPT-2 pretraining · GPT-2 pretraining | validation loss: 3.417 -> 3.342 | Helped |
| Use ASGO-PolarExpress instead of Muon | GPT-2 pretraining · GPT-2 pretraining | validation loss: 3.342 -> 3.342 | No clear effect |
| Use diagonal preconditioning (DASGO) instead of full-matrix ASGO (ablation) | GPT-2 pretraining · GPT-2 pretraining | validation loss: 3.342 -> 3.509 | Hurt |
| Use PolarExpress instead of Newton-Schulz for the inverse square root (ablation) | ASGO on GPT-2 · GPT-2 pretraining | validation loss: 3.347 -> 3.342 | No clear effect |
| Apply the single-sided preconditioner on the right rather than the left (ablation) | ASGO on GPT-2 · GPT-2 pretraining, learning rate 5e-4 | validation loss: 3.582 -> 3.569 | Helped |
| Apply the single-sided preconditioner on the right rather than the left (ablation) | ASGO on GPT-2 · GPT-2 pretraining, learning rate 1e-3 | validation loss: 3.468 -> 3.447 | Helped |
| Apply the single-sided preconditioner on the right rather than the left (ablation) | ASGO on GPT-2 · GPT-2 pretraining, learning rate 5e-3 | validation loss: 3.399 -> 3.362 | Helped |
| Apply the single-sided preconditioner on the right rather than the left (ablation) | ASGO on GPT-2 · GPT-2 pretraining, learning rate 1e-2 | validation loss: 3.412 -> 3.386 | Helped |
| Use ASGO instead of AdamW | GPT2-Large fine-tuning · WikiText-2 CLM | perplexity: 14.01 -> 13.88 | Helped |
| Use DASGO instead of AdamW (ablation) | GPT2-Large fine-tuning · WikiText-2 FIM | perplexity: 17.46 -> 15.45 | Helped |
| Add specialized Query-Key processing (ablation) | ASGO on NanoGPT · Shakespeare character dataset, validation loss after 5 epochs | validation loss | Helped |
Surprising findings
- DASGO (diagonal variant) underperforms AdamW on GPT-2 pretraining (3.509 vs 3.455 validation loss), suggesting off-diagonal preconditioner terms matter.
- ASGO and Muon show nearly identical training trajectories, supporting the claimed equivalence in the deterministic case.
Limitations
- Per-iteration cost is higher than Adam due to matrix operations (Newton-Schulz on a 2x2 block matrix, 10 iterations in the wall-time test).
- Extension to tensors is not straightforward.
- Nonconvex convergence rate is only conjectured, not proven.
Experimental setup
Datasets: OpenWebText Shakespeare character dataset WikiText-2
Benchmarks: OpenWebText validation loss WikiText-2 perplexity Shakespeare validation loss
Metrics: training loss validation loss perplexity wall-clock time per step
Baselines: AdamW Shampoo Muon DASGO
Builds on: AdaGrad Shampoo Muon Adam
Closest papers by meaning
- MARS: Unleashing the Power of Variance Reduction for Training Large Models (ICML 2025 · similarity 0.86)
- A Stable Whitening Optimizer for Efficient Neural Network Training (NeurIPS 2025 · similarity 0.86)
- AdaFisher: Adaptive Second Order Optimization via Fisher Information (ICLR 2025 · similarity 0.86)
- DeepZero: Scaling Up Zeroth-Order Optimization for Deep Model Training (ICLR 2024 · similarity 0.86)
- Adam-mini: Use Fewer Learning Rates To Gain More (ICLR 2025 · similarity 0.85)
- Adam Reduces a Unique Form of Sharpness: Theoretical Insights Near the Minimizer Manifold (NeurIPS 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.