本文經AI新媒體量子位（公眾號ID:QbitAI）授權轉載，轉載請聯(lián)系出處。

機器學習的優(yōu)化步驟，目前都是一階方法主導。

無論是SGD還是Adam，此類優(yōu)化算法在都是計算損失函數(shù)的一階導數(shù)——梯度，然后按照某種規(guī)定的方式讓權重隨梯度下滑方向迭代。

其實二階梯度會有更好的特性，因為它是計算梯度的導數(shù)，能夠更快地找到最合適的下降方向和速度。

然而出于計算量和存儲成本的考慮，二階優(yōu)化算法很少用到。

最近，谷歌大腦提出了一種新的二階預處理方法，帶來很大改進，優(yōu)于SGD、Adam和AdaGrad等一階算法，縮短了神經網(wǎng)絡的訓練時間。

它在Transformer訓練任務中比任何一階方法都快得多，而且能達到相同甚至更高的精度。連Jeff Dean也不禁在Twitter上點贊。

“洗發(fā)水”算法

這篇文章是對之前一種二階方法洗發(fā)水算法（Shampoo algorithm）做的實用化改進。

為何叫“洗發(fā)水算法”？其實是對此類算法的一種幽默稱呼。洗發(fā)水的廣告詞一般是“搓揉、沖洗、重復”，表示簡單重復式的無限循環(huán)，最后導致洗發(fā)水用盡（out of bottle）。

而這種算法用于機器學習優(yōu)化，最早來自于本文通訊作者Yoram Singer在2018年被ICML收錄的一篇文章Shampoo: Preconditioned Stochastic Tensor Optimization。

洗發(fā)水算法需要跟蹤2個預條件算子（Preconditioner）的統(tǒng)計數(shù)值Lt和Rt。

然后計算這2個預條件算子的四次根再求逆。將這兩個矩陣分別左乘和右乘梯度向量，迭代出t+1步的梯度再由以下公式得出：

上述過程像不像一種簡單重復，所以被作者自稱為“洗發(fā)水”。

2018年的那篇論文更側重于理論解釋，然而就是如此簡單的“洗頭”步驟實際應用起來也會面臨諸多困難。

這一步中最大的計算量來自于Lt-1/4和Rt-1/4。計算這個兩個數(shù)需要用到代價高昂的奇異值分解。

實際上，四次逆根不僅可以用SVD方法算出，也可以用舒爾-牛頓法（Schur-Newton algorithm）算出，而且隨著矩陣維度的增大，后者節(jié)約的時間越來越可觀。

舒爾-牛頓法可以在普通CPU上計算，不必消耗GPU、TPU這類神經網(wǎng)絡加速器的計算資源。

但即使是這樣，計算矩陣根的逆仍然相當耗時。如果不解決這個問題，訓練速度就不可能提高。

所以作者使用了異步計算的方法，并使用了TensorFlow中的Lingvo來對訓練循環(huán)進行改進。

CPU負責收集和處理訓練數(shù)據(jù)以及輔助活動，例如檢查點和訓練狀態(tài)摘要。而在GPU、TPU等加速器運行訓練循環(huán)時通常處于空閑或低利用率狀態(tài)，并自動提供雙精度計算。

這使它們成為計算預條件算子的理想選擇，而不會增加訓練消耗的資源。

使用異步計算

他們在每一步中都計算所有張量的預條件算子，但是預處理后的梯度卻是每N步計算一次，并交由CPU處理。

這期間，GPU或TPU依然在計算，過去的預條件算子在訓練過程中會一直使用，直到獲得更新后的預訓練算子為止。

計算過程像流水線一樣，并且異步運行而不會阻塞訓練循環(huán)。結果是，洗發(fā)水算法中最難計算的步驟幾乎沒有增加總的訓練時間。

僅有這些還不夠，作者對洗發(fā)水算法又做了幾點改進，使它可以適應大型模型的訓練。包括解耦步長大小和方向、預處理大型張量還有將大型張量劃分成多個塊。

最高提速67%

在WMT’14英語到法語翻譯的Transformer訓練任務中，該算法實現(xiàn)了1.67倍的加速，將時間減少了40%。

洗發(fā)水算法在和Adam或AdaGrad精度相同的情況下，只需后兩者實現(xiàn)了約一半的相同的精度AdaGrad或亞當許多步驟，而且對學習率的寬容度比AdaGrad高。

之前異步計算中的N是一個可調參數(shù)，決定了訓練的計算量，N越大，計算量越小。當然N也會對結果造成影響。我們需要在訓練過程的性能和結果的質量之間做出權衡。

實驗表明，這種方法可以承受多達1200個步驟的延遲，而不會造成任何明顯的質量損失。

洗發(fā)水也可以用在圖像分類任務中。

作者還在ImageNet-2012數(shù)據(jù)集上訓練了ResNet-50模型，結果比帶動量的SGD收斂更快，但是訓練損失與SGD相近，但是在測試集上的效果不如后者。

至于在泛化能力上的劣勢，洗發(fā)水算法還有待進一步的改進。

論文地址：

https://arxiv.org/abs/2002.09018

https://arxiv.org/abs/1802.09568