Convergence Rate Analysis of the AdamW-Style Shampoo: Unifying One-sided and Two-Sided Preconditioning

1. Problem Statement

The paper addresses the theoretical convergence analysis of the AdamW-Style Shampoo optimizer, a state-of-the-art algorithm that won the AlgoPerf competition for training neural networks. While Shampoo (and its variants such as SOAP and Muon) has demonstrated superior empirical performance over diagonal preconditioning methods (like Adam) through the use of full-matrix preconditioners, theoretical convergence guarantees for nonconvex optimization were either missing or insufficient and limited.

Specifically, the existing literature faced three major gaps:

1. Nonconvexity: Most previous analyses were restricted to convex settings.
2. Scope of Preconditioning: Previous nonconvex analyses (e.g., Xie et al., 2026) were limited to one-sided preconditioning and ignored the more complex and effective two-sided preconditioning used in practical implementations.
3. Practical Components: Existing theories often omitted critical practical features such as momentum and decoupled weight decay (AdamW-style), which are indispensable for training modern large language models (LLMs).

The paper aims to establish a rigorous convergence rate for the general AdamW-Style Shampoo algorithm, unifying both one-sided and two-sided preconditioning under a nonconvex framework.

2. Methodology

The authors analyze Algorithm 1, which represents the AdamW-Style Shampoo optimizer. The update rule includes:

• Exponential Moving Averages (EMA): Tracking the first moments ($M_k$) and second moments ($L_k, R_k$) of the gradients.
• Two-sided Preconditioning: The update step is given by $X_{k+1} = (1-\lambda\eta)X_k - \eta L_{k,\epsilon}^{-\frac{1}{2p}} M_k R_{k,\epsilon}^{-\frac{1}{2q}}$, where $L_k$ and $R_k$ are accumulated outer products of the gradients (approximating $E[GG^T]$ and $E[G^TG]$).
• Unified Exponents: The parameters $p$ and $q$ (where $1/p + 1/q = 1$) allow the algorithm to switch between two-sided ($p, q < \infty$) and one-sided ($p=1, q=\infty$ or vice versa) preconditioning.

Key Technical Innovations:
To handle the complexity of two-sided matrix preconditioning in nonconvex settings, the authors developed new proof techniques:

1. Schatten-p Hölder Inequality: Instead of the standard vector-based Hölder inequalities used in Adam analysis, the authors employ the Schatten-p Hölder inequality to bound the nuclear norm ($\|\cdot\|_*$) of the gradient. This is crucial because preconditioning acts on the singular values of the matrix parameters.
2. Matrix Cauchy-Schwarz Inequality: To bound the spectral norm of the update step ($L^{-1/2p} M R^{-1/2q}$), they apply a matrix-specific Cauchy-Schwarz inequality to handle the interaction between the left and right preconditioners.
3. Spectral Floor Analysis: They introduce a condition (Condition 4) regarding the lower bound of the preconditioner ($\hat{\epsilon}$). They argue that in practice, the effective spectral floor of the preconditioner is significantly larger than the numerical regularization $\epsilon$ (e.g., $10^{-12}$), thereby aligning theoretical bounds with empirical observations.

3. Main Contributions

• Unified Convergence Rate: The paper establishes the first convergence rate for AdamW-Style Shampoo in nonconvex optimization, covering both one-sided and two-sided preconditioning.
• Derivation of the Optimal Rate: They prove that the average expected nuclear norm of the gradient converges at a rate of:
  $$\frac{1}{K} \sum_{k=1}^K \mathbb{E}[\|\nabla f(X_k)\|_*] \leq O\left( \frac{\sqrt{m+n} \cdot C}{K^{1/4}} \right)$$
  where $m, n$ are the matrix dimensions, $K$ is the number of iterations, and $C$ corresponds to the constant in the optimal SGD rate.
• Theoretical Justification for Practical Settings: The analysis explains why setting the numerical regularizer $\epsilon$ to extremely small values (e.g., $10^{-12}$) does not hinder convergence. The theory shows that the effective preconditioner floor is determined by the gradient noise scale rather than the numerical $\epsilon$.
• Norm Relationship: The authors show that the derived nuclear norm rate is analogous to the optimal Frobenius norm rate of SGD ($O(K^{-1/4})$) under the ideal condition that $\|\nabla f(X)\|_* \approx \Theta(\sqrt{\min(m,n)}) \|\nabla f(X)\|_F$.

4. Results

• Theoretical Bounds:
  • The worst-case bound depends on $\sqrt{m+n}$, which is significantly tighter than the $\sqrt{mn}$ (total dimension $d$) dependency found in some diagonal methods or full-matrix lower bounds.
  • The convergence rate matches the known lower bound for nonconvex stochastic optimization ($O(K^{-1/4})$) up to a factor of $\sqrt{m+n}$.
• Empirical Validation (GPT-2 Pretraining):
  • Gradient Norm Ratio: Experiments on GPT-2 pretraining showed that the ratio $\|\nabla f(X)\|_* / \|\nabla f(X)\|_F$ consistently remains close to the theoretical upper bound $\sqrt{\min(m,n)}$. This validates the assumption that the nuclear norm rate effectively corresponds to the Frobenius norm rate in practice.
  • Spectral Floor: The smallest eigenvalues of the preconditioners ($L_k, R_k$) were observed to be orders of magnitude larger than the numerical $\epsilon$ ($10^{-12}$) and comparable to the noise-dependent scale $\hat{\sigma}^2/(m+n)$. This confirms that the "impractical" assumption of a larger spectral floor holds in real LLM training.

5. Significance

• Bridging Theory and Practice: This work provides the first rigorous theoretical foundation for the specific Shampoo variant that won the AlgoPerf competition and explains why it performs so well in practice.
• Advancement over Diagonal Methods: It theoretically validates the superiority of non-diagonal (full-matrix) preconditioning over diagonal methods (like Adam) for matrix-structured parameters and demonstrates that the cross-parameter correlations captured by Shampoo lead to better convergence directions.
• Guide for Hyperparameters: The analysis offers theoretical justification for the choice of hyperparameters (e.g., the relationship between momentum $\theta$ and the second-moment decay factor $\beta$, as well as the role of $\epsilon$) and guides practitioners in fine-tuning these optimizers for large-scale models.
• Foundation for Future Optimizers: By unifying one-sided and two-sided preconditioning and integrating momentum/weight decay, this framework sets a new standard for analyzing advanced, second-order-like optimizers in deep learning.

In summary, the paper successfully unifies the theoretical analysis of Shampoo by proving that it achieves an optimal convergence rate for nonconvex problems while simultaneously capturing the complex matrix structures inherent in modern deep learning models.