[论文解读] Scalable Second Order Optimization for Deep Learning

Rohan Anil, Vineet Gupta+3 more|arXiv (Cornell University)|Feb 20, 2020

Stochastic Gradient Optimization Techniques参考文献 49被引用 29

一句话总结

本文提出了一种可扩展的、硬件优化的二阶自适应方法（全矩阵Adagrad的一种变体）的实现，该方法使用分解预条件器以加速深度学习训练。通过利用CPU-加速器流水线、高效的矩阵开方计算以及架构扩展，该方法在大型模型（如Transformer和BERT）上实现了高达47%的运行时间减少，显著优于Adam等一阶方法在收敛速度和训练效率方面的表现。

ABSTRACT

Optimization in machine learning, both theoretical and applied, is presently dominated by first-order gradient methods such as stochastic gradient descent. Second-order optimization methods, that involve second derivatives and/or second order statistics of the data, are far less prevalent despite strong theoretical properties, due to their prohibitive computation, memory and communication costs. In an attempt to bridge this gap between theoretical and practical optimization, we present a scalable implementation of a second-order preconditioned method (concretely, a variant of full-matrix Adagrad), that along with several critical algorithmic and numerical improvements, provides significant convergence and wall-clock time improvements compared to conventional first-order methods on state-of-the-art deep models. Our novel design effectively utilizes the prevalent heterogeneous hardware architecture for training deep models, consisting of a multicore CPU coupled with multiple accelerator units. We demonstrate superior performance compared to state-of-the-art on very large learning tasks such as machine translation with Transformers, language modeling with BERT, click-through rate prediction on Criteo, and image classification on ImageNet with ResNet-50.

研究动机与目标

• 解决由于高内存、计算和通信成本，导致二阶方法在大规模深度学习中不切实际的问题。
• 弥合二阶优化在理论上具有的优势与在大规模实际部署之间的差距。
• 实现在现代异构硬件（CPU + 加速器）上针对最先进模型的高效、可扩展的二阶训练。
• 克服在实现全矩阵预条件化过程中，特别是对于大型层（如嵌入层）所面临的数值和基础设施挑战。
• 在多种大规模任务上，显著提升收敛性和实际运行时间，优于Adam和LAMB等一阶基线方法。

提出的方法

• 设计一种流水线化的优化流水线，将预条件器计算卸载到CPU，而加速器负责前向/反向传播。
• 用一种高效且数值稳定的迭代方法替代昂贵的谱分解（如SVD），以计算矩阵开方（如$L^{-1/4}$、$R^{-1/4}$）。
• 扩展Shampoo算法，使其能够支持大型高秩层（如嵌入矩阵），通过修改预条件器结构和更新规则。
• 采用基于AdaGrad的层内学习率，以稳定训练并允许使用比AdaGrad更大的学习率。
• 通过跨副本分片预条件统计信息，减少TPU核心之间的冗余计算。
• 将预条件梯度计算集成到训练循环中，仅引入极低延迟，利用异步CPU卸载。

实验结果

研究问题

• RQ1 能否使二阶自适应方法在现代深度学习工作负载中实现可扩展性？
• RQ2 如何在不牺牲收敛优势的前提下，降低全矩阵预条件化的计算和内存开销？
• RQ3 在CPU-加速器系统上，实现高效二阶方法所需的架构和算法优化有哪些？
• RQ4 与Adam和LAMB等一阶基线方法相比，二阶方法在多大程度上能减少训练时间和步数？
• RQ5 数值稳定性和可扩展性之间的权衡如何影响大规模二阶优化器的设计？

主要发现

• 在WMT’14英语到法语翻译任务中，Shampoo将标准Transformer的训练时间减少了45%（从12.0小时减少到6.7小时），将更大的Transformer-Big模型的训练时间减少了37%（从47.0小时减少到29.5小时）。
• 在32K批量大小的BERT-Large语言建模任务中，Shampoo在16%更少的训练步数内达到更高的掩码语言模型准确率，并将实际运行时间减少了4%（从3.8小时减少到3.65小时），且未进行超参数调优。
• 在CIFAR-10数据集上使用ResNet-50，Shampoo在143个周期内达到93.45%的准确率，而基线方法需要300个周期，总训练时间减少42%（从1428秒减少到827秒）。
• 在Criteo 1TB点击率预测任务中，Shampoo将训练时间从13分钟减少到8.2分钟，展示了在大规模稀疏模型上的强大效率。
• 即使在预条件器计算尚未完全优化的情况下，该方法在步数和实际运行时间上均实现了显著加速。
• 当前实现中，BERT-Large的每步时间增加了14%，表明通过进一步优化（如跨副本分片）可获得更大的性能增益。