ASGO: Adaptive Structured Gradient Optimization

TL;DR

ASGO targets structured optimization in DNN training by exploiting low-rank gradients and block-diagonal Hessians via a single-side preconditioner update $W_{t+1} = W_t - \eta_t V_t^{-1/2} G_t$, reducing memory and computation relative to Shampoo while achieving superior convergence. The authors provide detailed nonsmooth and smooth convergence analyses under matrix-structured assumptions, show ASGO can leverage low-rank gradients and Hessian diagonality, and connect it to Muon in a momentum-free limit. They propose practical implementations (including NS/PE solvers for $V_t^{-1/2}$ and a Transformer-focused, head-wise design for query/key matrices) and a memory-efficient DASGO variant. Empirically, ASGO matches or surpasses strong baselines like AdamW, Shampoo, and Muon on GPT-2 pretraining and GPT2-Large WikiText-2 finetuning, with competitive wall-clock times and improved stability, supporting its potential for efficient large-scale language-model training.

Abstract

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.

Figures (7)

• Figure 1: Pretraining GPT-2 Train Loss.
• Figure 2: This figure is from zhang2024adam. It depicts the Hessian of different parameter blocks in a small Transformer at the 1% training step. The near-block-diagonal 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.
• Figure 3: This figure is from zhang2024transformerszhang2024adam, calculating the Jensen-Shannon (JS) distance between two eigenvalue densities of all possible block pairs at initialization. It shows that JS distance of blockwise spectra in CNNs is significantly smaller than that in Transformers.
• Figure 4: Train Loss and Test Loss on the NanoGPT and Shakespeare Character Dataset
• Figure 5: Distribution of validation losses after 5 epochs with varying learning rates.

Theorems & Definitions (31)

• Theorem 1: Nonsmooth convergence
• Corollary 2
• Remark 1
• Remark 2
• Theorem 3: Smooth Convergence
• Remark 3
• Remark 4
• Lemma 1: Trace properties
• Lemma 2
• Lemma 3