2020/02/27 15:01

Rohan Anil等作者思源、一鸣、杜伟编译

二阶梯度优化新崛起，超越 Adam，Transformer 只需一半迭代量

常见的最优化器，如 Adam、AdaGrad、SGD+Momentum 等，都是一阶的。但是二阶梯度的收敛速度相比它们就快了太多。近日，谷歌研究者联合普林斯顿大学等，提出了真正应用的二阶梯度最优化器 Shampoo，让这个理论上颇有前景的设想变为现实。

目前，无论是从理论还是应用层面来说，机器学习中的优化都是以随机梯度下降等一阶梯度方法为主。囊括二阶梯度和/或二阶数据统计的二阶优化方法虽然理论基础更强，但受限于计算量、内存和通信花销等因素，二阶梯度优化方法的普及度不高。

可能你用各种框架搭建过各种神经网络，同时也尝试过调用 Adam、AdaGrad、SGD+Momentum 等形形色色的最优化器。但是你会发现，它们都采用一阶梯度，二阶梯度似乎仍然遥遥无期。

然而在谷歌大脑与普林斯顿大学等研究者的努力下，二阶梯度优化终于在实战大模型上展现出独特的优势。

研究者表示，为了缩短理论和实际优化效果之间的差距，该论文提出了一种二阶优化的概念性验证，并通过一系列重要的算法与数值计算提升，证明它在实际深度模型中能有非常大的提升。

论文地址：https://arxiv.org/abs/2002.09018

具体而言，在训练深度模型过程中，二阶梯度优化 Shampoo 能高效利用由多核 CPU 和多加速器单元组成的异构硬件架构。并且在大规模机器翻译、图像识别等领域实现了非常优越的性能，要比现有的顶尖一阶梯度下降方法还要好。

我们可以先看看它在 WMT 英-法翻译数据集上的效果，它采用的是标准的 Transformer。除了这一个实验，研究者还测试了 Big Transformer 以及 ImageNet 上的 ResNet，这些实验结果都展示在后文中。

WMT'14 英法翻译数据集上的 Transformer，二阶梯度算法 Shampoo 在迭代数上快了 1.95 倍，且就算要计算二阶梯度，每一次迭代也只慢了 16%，总体上来说节省了 40% 的执行时间。

从上图可以清楚地看到，如果 Adam 和 Shampoo 要训练到相同的准确度，Adam 需要迭代 30 万次，而 Shampoo 大概只需要迭代 11 万到 12 万次就差不多了。二阶梯度，果然收敛速度快了太多。

二阶梯度是什么

想象一下，如果我们希望找到「谷底」，那么沿着坡度一步一步往下走就行了。对于机器学习来说，「谷底」就是最优解，一步一步就是迭代过程。之前，我们采用一阶梯度，也就是坡度的陡和缓来确定步子要迈多大。而当坡度是有变化的，即逐渐变陡或变缓，根据当前坡度来确定步子大小就有一些问题。

之前我们可以慢慢多走几步，就能根据坡度的变化直接调整。现在如果能用二阶梯度，相当于梯度的梯度，那么也就知道坡度变化的趋势，因此一步就能走到位。所以二阶梯度本质上比一阶梯度多出一些信息，模型收敛也就会更快。

对于凸优化，二阶梯度一次就能找到最优解，而一阶梯度找到的方向必定垂直于当前点的等高线，因此出现这种「锯齿」现象。

二阶梯度策略无疑是数学优化中性能最有保障的算法之一。在这种算法中，我们使用预条件算子（preconditioner）矩阵转换梯度，然后应用到每个迭代步骤上。一般来说，这包括计算/近似估计二阶导矩阵，如 Hessian。另一方面，AdaGrad 和其他相关的算法主要针对随机优化，使用二阶梯度的方差矩阵来构建预条件算子。

虽然二阶方法在收敛性上比一阶好很多，但是其计算量限制了实际的应用。因为在每次梯度更新的时候，这种算法需要平方级别的存储和立方级别的计算时间。因此，这些方法在现在的机器学习优化方法中并不常见。

现代优化策略面对的最大的一个挑战是在理论和实际优化方法中搭建一个桥梁，使得二阶优化方法能够更合理地被应用和部署。

在这篇论文中，研究者真正提出了一种二阶梯度改进方法，它与 Adam 等算法一样是适应性梯度，但它能利用二阶梯度信息加速收敛，甚至在大型机器翻译模型中收敛快了一半。

二阶梯度，超越 Adam 的最优化

为了完成这项工作，研究者仔细思考了二阶优化存在的问题与困难，并改进了名为 Shampoo 的二阶梯度方法。研究者表示，二阶优化最无解的是目前深度学习库对一阶梯度已经有了大量优化，它们对计算量与内存的要求都不高。然而，对于 Shampoo，每一次迭代的密集计算，都对应用大模型产生了不可逾越的阻拦。

和一阶梯度优化方法相比，Shampoo 要走向实践，还有如下三大挑战。

1. 算法上的挑战

现代机器学习架构通常使用很大的嵌入层，维度可能多达百万级别。Shampoo 需要对每个维度计算一个预条件子，但是不管是计算还是存储，百万次的百万维度矩阵都几乎是不可处理的。因此从算法上，我们首先就要设计一种新方法来解决这类问题。

2. 计算上的挑战

Shampoo 的权重更新式大概是如下这样的，其中 L 和 R 都是矩阵，它们需要求逆与求根，在计算上会显得非常复杂，这也会拖慢整个迭代的速度。

之前矩阵求逆与根是可以使用 SVD 计算的，但是它们太慢了。因此可以考虑如 Schur-Newton 等一些算法，可以将逆 P 次根问题转换为一系列矩阵-向量和矩阵-矩阵的积，所以对于优化很有帮助。

图 1：在不同维度上，计算某矩阵逆 P 次根的基准对比。Schur-Newton 迭代方法能够在 CPU 上高效运行，而且相比 SVD 有很大的冗余用于提升。

3. 基础设施挑战

神经网络加速器通常是定制的，用来让机器学习程序运行地更快，开销更少。加速器设计倾向于低精度（8bit/16bit），能够满足现有的基准。研究者的方法需要双精度运算，因此已有的加速器甚至都不会启动。

此外，TensorFlow 等深度学习库提供的最优化器 API 适应于随机、一阶梯度下降那种模式。而二阶优化器需要与训练循环做交互，因此从实现上需要对框架底层做出修正。

虽然难，但还是能攻破

三大挑战使实现二阶梯度优化器异常复杂，研究者针对不同的问题提出了一系列优化组件、优化算法。最终搞定的分布式 Shampoo 在 CPU、TPU、GPU 等硬件上高效运行，这也是我们第一次看到二阶优化器在大模型上能 Work 的新研究。

研究者们首先分析了标准的数据并行方法。每个加速器的核都会在一个分批数据上进行前向和反向传播。然后算法会对数据批进行梯度聚合，并使用 all-reduction 的方法获得平均梯度。聚合的梯度被用来进行权重更新。前向和反向传播在所有核中并行计算。

为什么使用 All-reduction？这增加了一个屏障，使得所有核同步聚合批的梯度，并进行权重更新。图 3 中，研究者评价了 Transformer 模型每一步的计算开销。

图 3：使用 Diagonal AdaGrad 优化器的 Transformer 模型的每迭代步延迟时间为 134 毫秒，其中（1）前向传播 57 毫秒；后向传播 71 毫秒；all reduction：4 毫秒；权重更新：2 毫秒。

本文分布式系统实现的整体设计如下图 4 的时间轴所示，具体分布式架构可查阅原论文。

图 4：本文优化算法设计的时间轴。在每一步上计算所有张量的预调节器统计数据。预调节器只在每 N 步上计算，并且计算会分配给训练系统上可用的所有 CPU 核心。运算进行流水线处理，这样就实现了开销均摊。

可实战的二阶梯度优化

研究者在包含 3630 万个句对的 WMT'14 英法标准机器翻译数据集上验证了分布式系统实现方法的有效性。此外，他们在实验中使用了当前 SOTA Transformer 架构，该架构包含 9330 万个参数和 6 层的编码器-解码器结构。

实验在 32 核谷歌 Cloud TPU v3 Pod 上运行，结果如下图 6 所示，本文提出的 Shampoo 算法只需一半迭代数就能实现与 AdaGrad 和 Adam 相同的准确度。

图 6：WMT'14 英法翻译数据集上的 Transformer 模型，Shampoo 二阶梯度算法的收敛速度在迭代数上快了 1.95 倍，且就算要计算二阶梯度，每一次迭代也只慢了 16%，总体上来说节省了 40% 的执行时间。

研究者还利用一个大型 Transformer 模型进行实验，该模型包含 3.754 亿个参数和 6 层的编码器-解码器结构。实验结果如下图 12 所示，端到端的执行时间实现了提升。

图 12：WMT'14 英法翻译数据集上的 Transformer-Big 模型，Shampoo 二阶梯度算法的收敛速度在迭代数上快了 2 倍，且就算要计算二阶梯度，每一次迭代也只慢了 40%，总体上来说节省了 30% 的执行时间。

最后，研究者在 ImageNet-2012 数据集上训练了 ResNet-50 模型，并对使用 SGD+Momentum 的 SOTA 基准方法、本文提出的 Shampoo 二阶梯度算法以及 Adagrad 方法的测试结果进行对比，结果如下图 14 和表 2 所示。

研究者发现，Shampoo 二阶梯度算法虽然未能在测试损失或准确度方面实现任何改进，但与调整好的 SGD+Momentum 基准方法相比，该方法能够更快地减少训练损失。

图 14：在 ImageNet-2012 数据集上训练 ResNet-50 模型时的训练（图左）和测试（图右）交叉熵变化图。

表 2：当训练 ResNet-50 模型的 batch size=4096 时，三种方法在 ImageNet-2012 数据集上的准确度测试结果对比。

相信很多读者都学过最优化方法这类理论课程，我们会发现梯度下降，或者称之为最速下降法是最简单的方法，「锯齿」现象令它在很多领域上都存在问题。我们也会发现各种二阶优化、拟二阶优化在理论上性质远远超过它。然而在深度学习领域，由于数据与模型的规模，我们采用的都是「最速下降法」这个大家庭。

很多时候，我们会想，之前累积的那么多优秀方法，完全应用不到深度学习吗？而这篇论文至少告诉我们，计算、内存等各种困难，我们都是有机会克服的，相信那些具有强硬理论支持的最优化方法，最终会在深度学习展现它们的魅力。

理论二阶优化Adam优化器谷歌普林斯顿大学Transformer

相关数据

深度学习技术

深度学习（deep learning）是机器学习的分支，是一种试图使用包含复杂结构或由多重非线性变换构成的多个处理层对数据进行高层抽象的算法。深度学习是机器学习中一种基于对数据进行表征学习的算法，至今已有数种深度学习框架，如卷积神经网络和深度置信网络和递归神经网络等已被应用在计算机视觉、语音识别、自然语言处理、音频识别与生物信息学等领域并获取了极好的效果。

来源：LeCun, Y., Bengio, Y., & Hinton, G. (2015). Deep learning. nature, 521(7553), 436.

随机优化技术

随机优化（SO）方法是生成和使用随机变量的优化方法。对于随机问题，随机变量出现在优化问题本身的表述中，其涉及随机目标函数或随机约束。

来源：维基百科

权重技术

线性模型中特征的系数，或深度网络中的边。训练线性模型的目标是确定每个特征的理想权重。如果权重为 0，则相应的特征对模型来说没有任何贡献。

来源：Google AI Glossary

交叉熵技术

交叉熵（Cross Entropy）是Loss函数的一种（也称为损失函数或代价函数），用于描述模型预测值与真实值的差距大小

来源：维基百科

机器学习技术

机器学习是人工智能的一个分支，是一门多领域交叉学科，涉及概率论、统计学、逼近论、凸分析、计算复杂性理论等多门学科。机器学习理论主要是设计和分析一些让计算机可以自动“学习”的算法。因为学习算法中涉及了大量的统计学理论，机器学习与推断统计学联系尤为密切，也被称为统计学习理论。算法设计方面，机器学习理论关注可以实现的，行之有效的学习算法。

来源：Mitchell, T. (1997). Machine Learning. McGraw Hill.

迭代技术

模型的权重在训练期间的一次更新。迭代包含计算参数在单个批量数据上的梯度损失。

来源：Google ML glossary

基准技术

一种简单的模型或启发法，用作比较模型效果时的参考点。基准有助于模型开发者针对特定问题量化最低预期效果。

来源：Google ML Glossary

参数技术

在数学和统计学裡，参数（英语：parameter）是使用通用变量来建立函数和变量之间关系（当这种关系很难用方程来阐述时）的一个数量。

来源：维基百科

收敛技术

在数学，计算机科学和逻辑学中，收敛指的是不同的变换序列在有限的时间内达到一个结论（变换终止），并且得出的结论是独立于达到它的路径（他们是融合的）。通俗来说，收敛通常是指在训练期间达到的一种状态，即经过一定次数的迭代之后，训练损失和验证损失在每次迭代中的变化都非常小或根本没有变化。也就是说，如果采用当前数据进行额外的训练将无法改进模型，模型即达到收敛状态。在深度学习中，损失值有时会在最终下降之前的多次迭代中保持不变或几乎保持不变，暂时形成收敛的假象。

来源：Wikipedia Google ML glossary

凸优化技术

凸优化，或叫做凸最优化，凸最小化，是数学最优化的一个子领域，研究定义于凸集中的凸函数最小化的问题。凸优化在某种意义上说较一般情形的数学最优化问题要简单，譬如在凸优化中局部最优值必定是全局最优值。凸函数的凸性使得凸分析中的有力工具在最优化问题中得以应用，如次导数等。凸优化应用于很多学科领域，诸如自动控制系统，信号处理，通讯和网络，电子电路设计，数据分析和建模，统计学（最优化设计），以及金融。在近来运算能力提高和最优化理论发展的背景下，一般的凸优化已经接近简单的线性规划一样直捷易行。许多最优化问题都可以转化成凸优化（凸最小化）问题，例如求凹函数f最大值的问题就等同于求凸函数 -f最小值的问题。

来源：维基百科

TensorFlow技术

TensorFlow是一个开源软件库，用于各种感知和语言理解任务的机器学习。目前被50个团队用于研究和生产许多Google商业产品，如语音识别、Gmail、Google 相册和搜索，其中许多产品曾使用过其前任软件DistBelief。

来源：维基百科

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数，这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在维空间内，有个分量的一种量，其中每个分量都是坐标的函数，而在坐标变换时，这些分量也依照某些规则作线性变换。称为该张量的秩或阶（与矩阵的秩和阶均无关系）。在数学里，张量是一种几何实体，或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达，记作标量的数组，但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中，表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了，它们都是二阶张量，对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

来源：维基百科

机器翻译技术

机器翻译（MT）是利用机器的力量「自动将一种自然语言（源语言）的文本翻译成另一种语言（目标语言）」。机器翻译方法通常可分成三大类：基于规则的机器翻译（RBMT）、统计机器翻译（SMT）和神经机器翻译（NMT）。

来源：机器之心

神经网络技术

（人工）神经网络是一种起源于 20 世纪 50 年代的监督式机器学习模型，那时候研究者构想了「感知器（perceptron）」的想法。这一领域的研究者通常被称为「联结主义者（Connectionist）」，因为这种模型模拟了人脑的功能。神经网络模型通常是通过反向传播算法应用梯度下降训练的。目前神经网络有两大主要类型，它们都是前馈神经网络：卷积神经网络（CNN）和循环神经网络（RNN），其中 RNN 又包含长短期记忆（LSTM）、门控循环单元（GRU）等等。深度学习是一种主要应用于神经网络帮助其取得更好结果的技术。尽管神经网络主要用于监督学习，但也有一些为无监督学习设计的变体，比如自动编码器和生成对抗网络（GAN）。

来源：机器之心

梯度下降技术

梯度下降是用于查找函数最小值的一阶迭代优化算法。要使用梯度下降找到函数的局部最小值，可以采用与当前点的函数梯度（或近似梯度）的负值成比例的步骤。如果采取的步骤与梯度的正值成比例，则接近该函数的局部最大值，被称为梯度上升。

来源：Vapnik V. N. (2000). The Nature of Statistical Learning Theory. Information Science and Statistics. Springer-Verlag.Wikipedia

随机梯度下降技术

梯度下降（Gradient Descent）是遵循成本函数的梯度来最小化一个函数的过程。这个过程涉及到对成本形式以及其衍生形式的认知，使得我们可以从已知的给定点朝既定方向移动。比如向下朝最小值移动。在机器学习中，我们可以利用随机梯度下降的方法来最小化训练模型中的误差，即每次迭代时完成一次评估和更新。这种优化算法的工作原理是模型每看到一个训练实例，就对其作出预测，并重复迭代该过程到一定的次数。这个流程可以用于找出能导致训练数据最小误差的模型的系数。

来源：机器之心

最速下降法技术

在数学中，最速下降法或固定相法或鞍点法是拉普拉斯逼近积分方法的扩展，使复平面中的轮廓积分变形以通过静止点（鞍点）附近，大致在最速下降或固定相的方向。

来源：维基百科

优化器技术

优化器基类提供了计算梯度loss的方法，并可以将梯度应用于变量。优化器里包含了实现了经典的优化算法，如梯度下降和Adagrad。优化器是提供了一个可以使用各种优化算法的接口，可以让用户直接调用一些经典的优化算法，如梯度下降法等等。优化器（optimizers）类的基类。这个类定义了在训练模型的时候添加一个操作的API。用户基本上不会直接使用这个类，但是你会用到他的子类比如GradientDescentOptimizer, AdagradOptimizer, MomentumOptimizer（tensorflow下的优化器包）等等这些算法。

来源：维基百科

二阶优化方法技术

最优化方法，是指解决最优化问题的方法。所谓最优化问题，指在某些约束条件下，决定某些可选择的变量应该取何值，使所选定的目标函数达到最优的问题。即运用最新科技手段和处理方法，使系统达到总体最优，从而为系统提出设计、施工、管理、运行的最优方案。

来源：百度百科