Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Huan Li, Yiming Dong, Zhouchen Lin

TL;DR

This work develops a unified convergence analysis for AdamW-style Shampoo that handles both one-sided and two-sided matrix preconditioning in nonconvex optimization. By leveraging Schatten-$p$ Hölder inequalities and a matrix Cauchy–Schwarz bound, the authors derive an explicit nonasymptotic rate: $\frac{1}{K}\sum_{k=1}^K \mathbb{E}[\|\nabla f(\mathbf{X}_k)\|_*] \leq \mathcal{O}\left(\frac{\sqrt{m+n}\,C}{K^{1/4}}\right)$, with $C$ depending on gradient noise and smoothness, and show this aligns with SGD bounds under the ideal relationship between the Frobenius and nuclear norms. The analysis unifies two distinct preconditioning regimes, provides concrete parameter regimes (including $\hat{\varepsilon}$, $\varepsilon$, momentum, and decoupled weight decay) to guarantee convergence, and discusses practical implications for large-scale neural networks. The results offer theoretical grounding for Shampoo-style optimizers in nonconvex settings and connect their performance to matrix dimension through the nuclear-norm measure.

Abstract

This paper studies the AdamW-style Shampoo optimizer, an effective implementation of classical Shampoo that notably won the external tuning track of the AlgoPerf neural network training algorithm competition. Our analysis unifies one-sided and two-sided preconditioning and establishes the convergence rate $\frac{1}{K}\sum_{k=1}^K E\left[\|\nabla f(X_k)\|_*\right]\leq O(\frac{\sqrt{m+n}C}{K^{1/4}})$ measured by nuclear norm, where $K$ represents the iteration number, $(m,n)$ denotes the size of matrix parameters, and $C$ matches the constant in the optimal convergence rate of SGD. Theoretically, we have $\|\nabla f(X)\|_F\leq \|\nabla f(X)\|_*\leq \sqrt{m+n}\|\nabla f(X)\|_F$, supporting that our convergence rate can be considered to be analogous to the optimal $\frac{1}{K}\sum_{k=1}^KE\left[\|\nabla f(X_k)\|_F\right]\leq O(\frac{C}{K^{1/4}})$ convergence rate of SGD in the ideal case of $\|\nabla f(X)\|_*= Θ(\sqrt{m+n})\|\nabla f(X)\|_F$.

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

TL;DR

This work develops a unified convergence analysis for AdamW-style Shampoo that handles both one-sided and two-sided matrix preconditioning in nonconvex optimization. By leveraging Schatten- Hölder inequalities and a matrix Cauchy–Schwarz bound, the authors derive an explicit nonasymptotic rate: , with depending on gradient noise and smoothness, and show this aligns with SGD bounds under the ideal relationship between the Frobenius and nuclear norms. The analysis unifies two distinct preconditioning regimes, provides concrete parameter regimes (including , , momentum, and decoupled weight decay) to guarantee convergence, and discusses practical implications for large-scale neural networks. The results offer theoretical grounding for Shampoo-style optimizers in nonconvex settings and connect their performance to matrix dimension through the nuclear-norm measure.

Abstract

This paper studies the AdamW-style Shampoo optimizer, an effective implementation of classical Shampoo that notably won the external tuning track of the AlgoPerf neural network training algorithm competition. Our analysis unifies one-sided and two-sided preconditioning and establishes the convergence rate measured by nuclear norm, where represents the iteration number, denotes the size of matrix parameters, and matches the constant in the optimal convergence rate of SGD. Theoretically, we have , supporting that our convergence rate can be considered to be analogous to the optimal convergence rate of SGD in the ideal case of .
Paper Structure (12 sections, 14 theorems, 99 equations, 1 figure, 1 algorithm)

This paper contains 12 sections, 14 theorems, 99 equations, 1 figure, 1 algorithm.

Table of Contents

  1. Introduction
  2. Contributions
  3. Problem Settings, Notations, and Assumptions
  4. Convergence Rate of the AdamW-Style Shampoo
  5. Proof of the Theorem
  6. Bounding the Nuclear Norm of Gradient by Schatten-$p$ Holder Inequality
  7. Bounding the Spectral Norm of Update by Matrix Cauchy-Schwarz Inequality
  8. Proof Sketch of Theorem \ref{['main-theorem']}
  9. Conclusion
  10. Basic Properties in Matrix Analysis
  11. Proof of Theorem \ref{['main-theorem']}
  12. Supporting Lemmas

Key Result

Theorem 2.1

Suppose that Assumptions 1-3 and hold for some $\hat{\varepsilon}\geq\varepsilon$. Let $\hat{\sigma}^2=\max\left\{\sigma^2,\frac{L\left(f(\mathbf{X}_1)-f^*\right)}{K\gamma^2}\right\}$ with any $\gamma\in(0,1]$, $\frac{1}{p}+\frac{1}{q}=1$, $1-\theta=\sqrt{\frac{L\left(f(\mathbf{X}_1)-f^*\right)}{K\hat{\sigma}^2}}$, $\theta\leq\beta In the worst case, when $\hat{\varepsilon}=\varepsilon$, we have

Figures (1)

  • Figure 1: Illustrations of $\frac{1}{k}\sum_{t=1}^k\|\nabla f(\mathbf{X}_t)\|_F$ (left) and $\|\mathbf{X}_k-\mathbf{X}^*\|_F$ (right) over steps on the toy example (\ref{['toy-example']}).

Theorems & Definitions (23)

  • Theorem 2.1
  • Lemma 2.2
  • Definition 3.1
  • Lemma 3.2
  • Lemma 3.3
  • proof
  • Lemma 3.4
  • proof
  • Lemma 3.5
  • Lemma 3.6
  • ...and 13 more