本文经AI新媒体量子位授权转载（公众号ID：QbitAI），转载请联系出处转载。机器学习的优化步骤目前以一阶方法为主。无论是SGD还是Adam，这类优化算法都是计算损失函数的一阶导数——梯度，然后让权重按照某种规定的方式沿着梯度的方向迭代。其实二阶梯度会有更好的特性，因为它是计算梯度的导数，可以更快的找到最适合的方向和下降速度。但是，出于计算和存储成本的考虑，很少使用二阶优化算法。最近，GoogleBrain提出了一种新的二阶预处理方法，带来了很大的改进，优于SGD、Adam、AdaGrad等一阶算法，缩短了神经网络的训练时间。它在Transformer训练任务上比任何一阶方法都快得多，并且达到相同甚至更高的精度。就连JeffDean也忍不住在Twitter上点赞。文章《Shampoo》算法是对之前二阶法Shampoo算法的实用改进。为什么叫“洗发水算法”？事实上，它是此类算法的一个幽默名称。洗发水的广告语一般是“rub,rinse,repeat”，意思是简单重复的无限循环，最终导致使用洗发水（瓶外）。使用该算法进行机器学习优化最早来自本文通讯作者YoramSinger的一篇名为Shampoo:PreconditionedStochasticTensorOptimization的文章，2018年被ICML收录。shampoo算法需要跟踪统计值两个预处理器（Preconditioner）的Lt和Rt。然后计算这两个预处理算子的四次根，然后求逆。将梯度向量左右相乘两个矩阵，迭代第t+1步的梯度，然后由下面的公式得到：上面的过程是不是看起来像是简单的重复，所以作者自称“洗发水”。2018年的论文更侧重于理论解释，但即使是这么简单的“洗头”步骤，在实际应用中也会面临很多困难。这一步计算量最大的是Lt-1/4和Rt-1/4。计算这两个数字需要昂贵的奇异值分解。事实上，四次逆根不仅可以用SVD方法计算，还可以用Schur-Newton算法计算，而且随着矩阵维数的增加，后者越来越节省时间。Schur-Newton方法可以在普通CPU上进行计算，无需消耗GPU、TPU等神经网络加速器的计算资源。但即便如此，计算矩阵根的倒数仍然非常耗时。不解决这个问题，就不可能提高训练速度。于是作者采用了异步计算的方式，在TensorFlow中使用了Lingvo来改进训练循环。CPU负责收集和处理训练数据以及检查点和训练状态摘要等辅助活动。然而，GPU和TPU等加速器在运行训练循环时通常处于空闲或低利用率状态，并自动提供双精度计算。这使它们成为计算预处理器的理想选择，而不会增加训练消耗的资源。他们使用异步计算在每一步计算所有张量的预处理算子，但是每N步计算预处理的梯度并交给CPU处理。在此期间，GPU或TPU仍在计算，训练过程中会使用之前的预处理算子，直到得到更新的预训练算子。计算流水线化并异步运行，不会阻塞训练循环。事实证明，Shampoo算法中计算难度最大的步骤对总训练时间的影响很小。光有这些还不够，作者对洗发水算法进行了多次改进，使其能够适应大型模型的训练。包括解耦步长和方向，预处理大张量，以及将大张量分成块。最高加速比为67%。在WMT’14英法翻译的Transformer训练任务中，该算法实现了1.67倍的加速，缩短了40%的时间。在与Adam或AdaGrad精度相同的情况下，Shampoo算法只需要后两者就可以达到AdaGrad一半左右的精度或Adam的很多步，对学习率的容忍度比AdaGrad高。前面异步计算中的N是一个可调参数，决定了训练的计算量。N越大，计算量越小。当然N也会影响结果。我们需要在训练过程的性??能和结果的质量之间做出权衡。实验表明，这种方法可以容忍多达1200步的延迟，而不会造成任何明显的质量损失。Shampoo也可以用于图像分类任务。作者还在ImageNet-2012数据集上训练了ResNet-50模型，结果比带动量的SGD快，但是训练损失和SGD差不多，但是在测试集上效果不如后者。至于泛化能力的不足，洗发水算法还有待进一步改进。论文地址：https://arxiv.org/abs/2002.09018https://arxiv.org/abs/1802.09568
