解析AdamW版Shampoo的收敛速度,统一单边与双边预条件方法。
Convergence Rate Analysis of the AdamW-style Shampoo: Unifying One-Sided and Two-Sided Preconditioning
- 提出统一单边与双边预条件的理论框架,揭示其收敛机制。
- 在指数和为1/2时,实现接近SGD最优的1/4次方收敛率。
- 适用于矩阵参数优化,尤其适合深度学习中的复杂模型训练。
本文研究了获胜算法竞赛外部调优赛道的AdamW风格Shampoo,这是一种经典Shampoo的有效变体。分析统一了单边与双边预条件方法。当两个预条件算子的指数之和为1/2时,建立了收敛速率:\frac{1}{K}\sum_{k=1}^KE\left[||\nabla f(X_k)||_*\right]\leq O(\frac{\sqrt{m+n}C}{K^{1/4}}),其中K为迭代次数,(m,n)表示矩阵参数维度,C与SGD最优收敛率中的常数一致。理论上,核范数与弗罗贝尼乌斯范数满足||\nabla f(X)||_F\leq ||\nabla f(X)||_*\leq \sqrt{\min\{m,n\}}||\nabla f(X)||_F,表明该速率在理想情况下(||\nabla f(X)||_*= Θ(\sqrt{\min\{m,n\}})||\nabla f(X)||_F 且 m、n量级相当)相当于SGD的最优\frac{1}{K}\sum_{k=1}^KE\left[||\nabla f(X_k)||_F\right]\leq O(\frac{C}{K^{1/4}})。进一步将分析扩展至指数和不为1/2的情形,推导出显式但更复杂的收敛速率。
原文摘要 · Abstract (English)
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}^KE\left[||\nabla f(X_k)||_*\right]\leq 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\leq ||\nabla f(X)||_*\leq \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}^KE\left[||\nabla f(X_k)||_F\right]\leq O(\frac{C}{K^{1/4}})$ convergence rate of SGD in the ideal case where $||\nabla f(X)||_*= Θ(\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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。