[论文解读] Shampoo: Preconditioned Stochastic Tensor Optimization

Vineet Gupta, Tomer Koren+1 more|arXiv (Cornell University)|Feb 26, 2018

Stochastic Gradient Optimization Techniques参考文献 21被引用 33

一句话总结

Shampoo 是一种用于深度学习中张量结构参数的新型预条件随机优化算法，通过为每个张量维度维护独立的预条件矩阵，实现了高效的二阶优化。在实践中，其收敛速度显著快于 SGD、Adam 和 AdaGrad，且由于采用了结构化矩阵运算和在 TensorFlow 中的高效实现，每步运行时间与更简单的优化器相当。

ABSTRACT

Preconditioned gradient methods are among the most general and powerful tools in optimization. However, preconditioning requires storing and manipulating prohibitively large matrices. We describe and analyze a new structure-aware preconditioning algorithm, called Shampoo, for stochastic optimization over tensor spaces. Shampoo maintains a set of preconditioning matrices, each of which operates on a single dimension, contracting over the remaining dimensions. We establish convergence guarantees in the stochastic convex setting, the proof of which builds upon matrix trace inequalities. Our experiments with state-of-the-art deep learning models show that Shampoo is capable of converging considerably faster than commonly used optimizers. Although it involves a more complex update rule, Shampoo's runtime per step is comparable to that of simple gradient methods such as SGD, AdaGrad, and Adam.

研究动机与目标

• 解决全矩阵预条件在高维机器学习问题中的可扩展性限制。
• 开发一种实用的预条件方法，利用模型参数中的张量结构（如矩阵和 4D 卷积滤波器）的优势。
• 通过维护针对各维度的预条件矩阵而非完整的 Hessian 近似，实现在深度学习中的高效二阶优化。
• 利用矩阵迹不等式，在随机凸优化设置下建立理论收敛保证。
• 在经验上证明其在收敛速度上优于 Adam 和 SGD 等标准优化器，尽管更新规则更为复杂。

提出的方法

• Shampoo 维护一组预条件矩阵，每个张量维度对应一个，通过累积梯度的二阶统计量在线更新这些矩阵。
• 该算法对每个预条件矩阵执行对称矩阵平方根运算，使用左、右预条件矩阵的几何平均值来缩放梯度。
• 对于矩阵参数 $ W \in \mathbb{R}^{m \times n} $，更新规则为 $ W_{t+1} = W_t - \eta L_t^{-1/4} G_t R_t^{-1/4} $，其中 $ G_t $ 为梯度，$ L_t, R_t $ 为预条件矩阵。
• 预条件矩阵 $ L_t $ 和 $ R_t $ 的更新方式为 $ L_t = L_{t-1} + G_t G_t^\top $ 和 $ R_t = R_{t-1} + G_t^\top G_t $，用于追踪梯度外积的累积。
• 该方法可推广至高阶张量，通过在所有张量模式上应用相同原理，保持优化过程中的张量结构。
• Shampoo 在 TensorFlow 中通过标准张量运算高效实现，仅需极少代码修改，与现有深度学习框架完全兼容。

实验结果

研究问题

• RQ1 能否设计一种可高效扩展至高维、结构化参数空间（如深度学习中的张量）的预条件优化方法？
• RQ2 一种为每个张量维度维护独立矩阵的预条件策略，是否能实现比 Adam 和 SGD 等一阶方法更快的收敛速度？
• RQ3 尽管使用了二阶信息，该方法能否保持较低的每步计算成本？
• RQ4 在随机凸设置下，这种结构化预条件方法能建立怎样的理论收敛保证？
• RQ5 该算法在大规模深度学习模型上，针对不同架构和数据集的实际表现如何？

主要发现

• 在使用 ResNet 和 Inception 架构的 CIFAR-10 和 CIFAR-100 数据集上，Shampoo 的收敛速度显著快于 Adam、SGD 和 AdaGrad，且在更少的训练轮次内达到更低的训练损失。
• 在 LM1B 语言建模基准测试中，Shampoo 的测试困惑度低于 Adam 和 AdaGrad，且在步数维度上收敛更快。
• 尽管更新规则更复杂，涉及矩阵逆平方根运算，Shampoo 的每步运行时间与 SGD 和 Adam 相当，仅在 Tesla K40 GPU 上有轻微性能损失。
• 在某些情况下（如在 CIFAR-100 上使用 ResNet-55），Shampoo 的每步速度甚至快于 Adam，达到每秒 1.249 步，而 Adam 为每秒 1.203 步。
• 采用延迟预条件更新和动量机制可降低平均运行时间而不损失精度，从而实现实际部署。
• 该算法在所有测试模型中均实现了最先进的收敛速度，证明了结构感知预条件在深度学习优化中的有效性。