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

以下是论文《AdamW 风格 Shampoo 的收敛率分析：统一单侧与双侧预条件》的详细技术总结。

1. 问题陈述

本文针对 AdamW 风格 Shampoo 优化器进行了理论收敛性分析，该算法是 AlgoPerf 神经网络训练竞赛的获胜者，代表了最先进算法。尽管 Shampoo 及其变体（如 SOAP 和 Muon）通过利用全矩阵预条件器，在经验性能上优于对角预条件方法（如 Adam），但其在非凸优化方面的理论收敛保证一直缺失或受限。

具体而言，现有文献面临三个主要空白：

1. 非凸性：大多数先前的分析仅限于凸设置。
2. 预条件范围：先前的非凸分析（例如 Xie 等人，2026）仅限于单侧预条件，忽略了实际实现中更复杂且有效的双侧预条件。
3. 实际组件：现有理论往往忽略了关键的实践特性，如动量和解耦权重衰减（AdamW 风格），而这些对于现代大语言模型（LLM）训练至关重要。

本文旨在为通用的 AdamW 风格 Shampoo 算法建立严格的收敛率，在非凸框架下统一单侧和双侧预条件。

2. 方法论

作者分析了代表 AdamW 风格 Shampoo 优化器的算法 1。其更新规则涉及：

• 指数移动平均（EMA）：跟踪梯度的第一矩（$M_k$）和第二矩（$L_k, R_k$）。
• 双侧预条件：更新步骤为 $X_{k+1} = (1-\lambda\eta)X_k - \eta L_{k,\epsilon}^{-\frac{1}{2p}} M_k R_{k,\epsilon}^{-\frac{1}{2q}}$，其中 $L_k$ 和 $R_k$ 是梯度的累积外积（近似于 $E[GG^T]$ 和 $E[G^TG]$）。
• 统一指数：参数 $p$ 和 $q$（满足 $1/p + 1/q = 1$）允许算法在双侧（$p, q < \infty$）和单侧（$p=1, q=\infty$ 或反之）预条件之间切换。

关键技术创新：
为了处理非凸设置中双侧矩阵预条件的复杂性，作者开发了新颖的证明技术：

1. Schatten-p Hölder 不等式：不同于 Adam 分析中使用的标准基于向量的 Hölder 不等式，作者利用Schatten-p Hölder 不等式来界定梯度的核范数（$\|\cdot\|_*$）。这至关重要，因为预条件器作用于矩阵参数的奇异值。
2. 矩阵 Cauchy-Schwarz 不等式：为了界定更新步骤（$L^{-1/2p} M R^{-1/2q}$）的谱范数，他们采用了特定于矩阵的 Cauchy-Schwarz 不等式，以处理左右预条件器之间的相互作用。
3. 谱底分析：他们引入了关于预条件器下界（$\hat{\epsilon}$）的条件（条件 4）。他们指出，在实践中，预条件器的有效谱底显著大于数值正则化 $\epsilon$（例如 $10^{-12}$），从而使理论界限与经验观察相一致。

3. 主要贡献

• 统一收敛率：本文建立了 AdamW 风格 Shampoo 在非凸优化中的首个收敛率，涵盖了单侧和双侧预条件。
• 最优率推导：他们证明了梯度的平均期望核范数以以下速率收敛：
  $$\frac{1}{K} \sum_{k=1}^K \mathbb{E}[\|\nabla f(X_k)\|_*] \leq O\left( \frac{\sqrt{m+n} \cdot C}{K^{1/4}} \right)$$
  其中 $m, n$ 是矩阵维度，$K$ 是迭代次数，$C$ 与最优 SGD 速率中的常数相匹配。
• 实际设置的理论依据：该分析解释了为何将数值正则化器 $\epsilon$ 设置为极小值（例如 $10^{-12}$）不会阻碍收敛。理论表明，有效预条件器底由梯度噪声尺度决定，而非数值 $\epsilon$。
• 范数关系：作者证明，在理想条件下（即 $\|\nabla f(X)\|_* \approx \Theta(\sqrt{\min(m,n)}) \|\nabla f(X)\|_F$），推导出的核范数速率与 SGD 的最优 Frobenius 范数速率（$O(K^{-1/4})$）是类比对应的。

4. 结果

• 理论界限：
  • 最坏情况界限依赖于 $\sqrt{m+n}$，这比某些对角方法或全矩阵下界中发现的 $\sqrt{mn}$（总维度 $d$）依赖性要紧密得多。
  • 收敛率与非凸随机优化的已知下界（$O(K^{-1/4})$）相匹配，仅相差一个 $\sqrt{m+n}$ 因子。
• 经验验证（GPT-2 预训练）：
  • 梯度范数比率：GPT-2 预训练的实验表明，比率 $\|\nabla f(X)\|_* / \|\nabla f(X)\|_F$ 始终保持在理论上限 $\sqrt{\min(m,n)}$ 附近。这验证了在实际中核范数速率实际上等同于 Frobenius 范数速率的假设。
  • 谱底：观察到预条件器（$L_k, R_k$）的最小特征值比数值 $\epsilon$（$10^{-12}$）大几个数量级，并且与依赖于噪声的尺度 $\hat{\sigma}^2/(m+n)$ 相当。这证实了“较大谱底”这一看似不切实际的假设在现实世界的 LLM 训练中是成立的。

5. 意义

• ** bridging 理论与实践**：这项工作为 AlgoPerf 竞赛获胜的特定 Shampoo 变体提供了首个严格的理论基础，解释了为何它在实践中如此有效。
• 超越对角方法的进步：它在理论上验证了非对角（全矩阵）预条件器对于矩阵结构参数优于对角方法（如 Adam），表明 Shampoo 捕捉到的跨参数相关性带来了更好的收敛方向。
• 超参数指导：该分析为超参数的选择提供了理论依据（例如动量 $\theta$ 与二阶矩衰减 $\beta$ 之间的关系，以及 $\epsilon$ 的作用），指导实践者为大规模模型调整这些优化器。
• 未来优化器的基础：通过统一单侧和双侧预条件并纳入动量/权重衰减，该框架为分析深度学习中的高级二阶类优化器设定了新标准。

总之，本文成功统一了 Shampoo 的理论分析，证明了其在非凸问题上实现了最优收敛率，同时捕捉了现代深度学习模型中固有的复杂矩阵结构。