谷歌大腦提出“洗發水”二階優化算法，Transformer訓練時間減少40%，Jeff Dean點贊
曉查 發自 凹非寺
量子位 報道 | 公眾號 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
- 腦機接口走向現實，11張PPT看懂中國腦機接口產業現狀｜量子位智庫2021-08-10
- 張朝陽開課手推E=mc2，李永樂現場狂做筆記2022-03-11
- 阿里數學競賽可以報名了！獎金增加到400萬元，題目面向大眾公開征集2022-03-14
- 英偉達遭黑客最后通牒：今天必須開源GPU驅動，否則公布1TB機密數據2022-03-05
