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

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 frac1Ksumk=1KEleft[||nabla f(Xk)||_*right]≤ O(frac√m+nCK1/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≤ ||nabla f(X)||_*≤ √minm,n||nabla f(X)||F, which suggests that our convergence rate is analogous to the optimal frac1Ksumk=1KEleft[||nabla f(Xk)||Fright]≤ O(fracCK1/4) convergence rate of SGD in the ideal case where ||nabla f(X)||_*= Theta(√minm,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.