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

TL;DR

This paper analyzes the convergence rate of the AdamW-style Shampoo optimizer, unifying preconditioning methods and establishing a rate comparable to SGD, with theoretical support for its effectiveness in neural network training.

Contribution

It provides a unified convergence analysis for AdamW-style Shampoo, connecting one-sided and two-sided preconditioning, and relates its rate to that of SGD.

Findings

Convergence rate of O((sqrt(m+n)C)/K^{1/4}) measured by nuclear norm.

Theoretical bounds relate nuclear norm and Frobenius norm of gradients.

Supports the optimizer's effectiveness with convergence comparable to SGD.

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)\|_*= \Theta(\sqrt{m+n})\|\nabla f(X)\|_F$.