谷歌大腦提出「洗髮水」二階優化算法，Transformer訓練時間減少40%，Jeff Dean點贊

• 2020 年 3 月 31 日
• 筆記

曉查 發自 凹非寺 量子位 報道 | 公眾號 QbitAI

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

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

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

然而出於計算量和存儲成本的考慮，二階優化算法很少用到。

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

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

「洗髮水」算法

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

為何叫「洗髮水算法」？其實是對此類算法的一種幽默稱呼。洗髮水的廣告詞一般是「搓揉、沖洗、重複」，表示簡單重複式的無限循環，最後導致洗髮水用盡（out of bottle）。

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

洗髮水算法需要跟蹤2個預條件算子（Preconditioner）的統計數值Lt和Rt。

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

上述過程像不像一種簡單重複，所以被作者自稱為「洗髮水」。

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

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

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

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

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

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

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

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

使用異步計算

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

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

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

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

最高提速67%

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

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

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

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

洗髮水也可以用在圖像分類任務中。

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

至於在泛化能力上的劣勢，洗髮水算法還有待進一步的改進。

論文地址：

https://arxiv.org/abs/2002.09018

https://arxiv.org/abs/1802.09568

作者系網易新聞·網易號「各有態度」簽約作者

— 完 —