選自arXiv
機器之心編譯
參與：思源、一鳴、杜偉
常見的最優化器，如 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 數據集上的準確度測試結果對比。
相信很多讀者都學過最優化方法這類理論課程，我們會發現梯度下降，或者稱之為最速下降法是最簡單的方法，「鋸齒」現象令它在很多領域上都存在問題。我們也會發現各種二階優化、擬二階優化在理論上性質遠遠超過它。然而在深度學習領域，由於數據與模型的規模，我們採用的都是「最速下降法」這個大家庭。
很多時候，我們會想，之前累積的那麼多優秀方法，完全應用不到深度學習嗎？而這篇論文至少告訴我們，計算、內存等各種困難，我們都是有機會克服的，相信那些具有強硬理論支持的最優化方法，最終會在深度學習展現它們的魅力。
閱讀更多 機器之心Pro 的文章
