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

Abstract

This paper studies AdamW-style Shampoo, an effective variant of the classical Shampoo that won the external tuning track of the AlgoPerf neural network training competition. Our analysis unifies one-sided and two-sided preconditioning. When the exponents of the two preconditioners sum to $1/2$, we establish the convergence rate $\frac{1}{K}{\sum }_{k=1}^{K}E\left[||\nabla f(X_k)|{|}_{\ast }\right]\le O(\frac{\sqrt{m+n}C}{K^{1/4}})$, where $K$ represents the number of iterations, $(m,n)$ denotes the dimensions of the matrix-valued parameters, and $C$ matches the constant appearing in the optimal convergence rate of SGD. Theoretically, the nuclear norm and Frobenius norm satisfy $||\nabla f(X)|{|}_F\le ||\nabla f(X)|{|}_{\ast }\le \sqrt{\min \{m,n\}}||\nabla f(X)|{|}_F$, which suggests that our convergence rate is analogous to the optimal $\frac{1}{K}{\sum }_{k=1}^{K}E\left[||\nabla f(X_k)|{|}_F\right]\le O(\frac{C}{K^{1/4}})$ convergence rate of SGD in the ideal case where $||\nabla f(X)|{|}_{\ast }=\hat{I}˜(\sqrt{\min \{m,n\}})||\nabla f(X)|{|}_F$ and $m$ and $n$ are of comparable magnitude. Then, we extend our analysis to settings where the preconditioning exponents do not sum to 1/2, and establish convergence with an explicit but more involved rate.

V3:ICML Camera-Ready. V4 v.s. V3: extend to the more general setting where the exponents of the two preconditioners do not sum to 1/2