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

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[\mathrm{\parallel }\mathrm{\nabla }f(X_k){\mathrm{\parallel }}_{\ast }]\le O(\frac{\sqrt{m+n}C}{K^{1\mathrm{/}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 $\mathrm{\parallel }\mathrm{\nabla }f(X){\mathrm{\parallel }}_F\le \mathrm{\parallel }\mathrm{\nabla }f(X){\mathrm{\parallel }}_{\ast }\le \sqrt{m+n}\mathrm{\parallel }\mathrm{\nabla }f(X){\mathrm{\parallel }}_F$, supporting that our convergence rate can be considered to be analogous to the optimal $\frac{1}{K}{\sum }_{k=1}^{K}E[\mathrm{\parallel }\mathrm{\nabla }f(X_k){\mathrm{\parallel }}_F]\le O(\frac{C}{K^{1\mathrm{/}4}})$ convergence rate of SGD in the ideal case of $\mathrm{\parallel }\mathrm{\nabla }f(X){\mathrm{\parallel }}_{\ast }=\mathrm{\Theta }(\sqrt{m+n})\mathrm{\parallel }\mathrm{\nabla }f(X){\mathrm{\parallel }}_F$ and balanced $m$ and $n$.