選自arXiv

機(jī)器之心編譯

常見的最優(yōu)化器，如 Adam、AdaGrad、SGD+Momentum 等，都是一階的。但是二階梯度的收斂速度相比它們就快了太多。近日，谷歌研究者聯(lián)合普林斯頓大學(xué)等，提出了真正應(yīng)用的二階梯度最優(yōu)化器 Shampoo，讓這個(gè)理論上頗有前景的設(shè)想變?yōu)楝F(xiàn)實(shí)。

目前，無論是從理論還是應(yīng)用層面來說，機(jī)器學(xué)習(xí)中的優(yōu)化都是以隨機(jī)梯度下降等一階梯度方法為主。囊括二階梯度和/或二階數(shù)據(jù)統(tǒng)計(jì)的二階優(yōu)化方法雖然理論基礎(chǔ)更強(qiáng)，但受限于計(jì)算量、內(nèi)存和通信花銷等因素，二階梯度優(yōu)化方法的普及度不高。

可能你用各種框架搭建過各種神經(jīng)網(wǎng)絡(luò)，同時(shí)也嘗試過調(diào)用 Adam、AdaGrad、SGD+Momentum 等形形色色的最優(yōu)化器。但是你會發(fā)現(xiàn)，它們都采用一階梯度，二階梯度似乎仍然遙遙無期。

然而在谷歌大腦與普林斯頓大學(xué)等研究者的努力下，二階梯度優(yōu)化終于在實(shí)戰(zhàn)大模型上展現(xiàn)出獨(dú)特的優(yōu)勢。

研究者表示，為了縮短理論和實(shí)際優(yōu)化效果之間的差距，該論文提出了一種二階優(yōu)化的概念性驗(yàn)證，并通過一系列重要的算法與數(shù)值計(jì)算提升，證明它在實(shí)際深度模型中能有非常大的提升。

論文地址：https://arxiv.org/abs/2002.09018

具體而言，在訓(xùn)練深度模型過程中，二階梯度優(yōu)化 Shampoo 能高效利用由多核 CPU 和多加速器單元組成的異構(gòu)硬件架構(gòu)。并且在大規(guī)模機(jī)器翻譯、圖像識別等領(lǐng)域?qū)崿F(xiàn)了非常優(yōu)越的性能，要比現(xiàn)有的頂尖一階梯度下降方法還要好。

我們可以先看看它在 WMT 英-法翻譯數(shù)據(jù)集上的效果，它采用的是標(biāo)準(zhǔn)的 Transformer。除了這一個(gè)實(shí)驗(yàn)，研究者還測試了 Big Transformer 以及 ImageNet 上的 ResNet，這些實(shí)驗(yàn)結(jié)果都展示在后文中。

WMT 14 英法翻譯數(shù)據(jù)集上的 Transformer，二階梯度算法 Shampoo 在迭代數(shù)上快了 1.95 倍，且就算要計(jì)算二階梯度，每一次迭代也只慢了 16%，總體上來說節(jié)省了 40% 的執(zhí)行時(shí)間。

從上圖可以清楚地看到，如果 Adam 和 Shampoo 要訓(xùn)練到相同的準(zhǔn)確度，Adam 需要迭代 30 萬次，而 Shampoo 大概只需要迭代 11 萬到 12 萬次就差不多了。二階梯度，果然收斂速度快了太多。

二階梯度是什么

想象一下，如果我們希望找到「谷底」，那么沿著坡度一步一步往下走就行了。對于機(jī)器學(xué)習(xí)來說，「谷底」就是最優(yōu)解，一步一步就是迭代過程。之前，我們采用一階梯度，也就是坡度的陡和緩來確定步子要邁多大。而當(dāng)坡度是有變化的，即逐漸變陡或變緩，根據(jù)當(dāng)前坡度來確定步子大小就有一些問題。

之前我們可以慢慢多走幾步，就能根據(jù)坡度的變化直接調(diào)整。現(xiàn)在如果能用二階梯度，相當(dāng)于梯度的梯度，那么也就知道坡度變化的趨勢，因此一步就能走到位。所以二階梯度本質(zhì)上比一階梯度多出一些信息，模型收斂也就會更快。

對于凸優(yōu)化，二階梯度一次就能找到最優(yōu)解，而一階梯度找到的方向必定垂直于當(dāng)前點(diǎn)的等高線，因此出現(xiàn)這種「鋸齒」現(xiàn)象。

二階梯度策略無疑是數(shù)學(xué)優(yōu)化中性能最有保障的算法之一。在這種算法中，我們使用預(yù)條件算子（preconditioner）矩陣轉(zhuǎn)換梯度，然后應(yīng)用到每個(gè)迭代步驟上。一般來說，這包括計(jì)算/近似估計(jì)二階導(dǎo)矩陣，如 Hessian。另一方面，AdaGrad 和其他相關(guān)的算法主要針對隨機(jī)優(yōu)化，使用二階梯度的方差矩陣來構(gòu)建預(yù)條件算子。

雖然二階方法在收斂性上比一階好很多，但是其計(jì)算量限制了實(shí)際的應(yīng)用。因?yàn)樵诿看翁荻雀碌臅r(shí)候，這種算法需要平方級別的存儲和立方級別的計(jì)算時(shí)間。因此，這些方法在現(xiàn)在的機(jī)器學(xué)習(xí)優(yōu)化方法中并不常見。

現(xiàn)代優(yōu)化策略面對的最大的一個(gè)挑戰(zhàn)是在理論和實(shí)際優(yōu)化方法中搭建一個(gè)橋梁，使得二階優(yōu)化方法能夠更合理地被應(yīng)用和部署。

在這篇論文中，研究者真正提出了一種二階梯度改進(jìn)方法，它與 Adam 等算法一樣是適應(yīng)性梯度，但它能利用二階梯度信息加速收斂，甚至在大型機(jī)器翻譯模型中收斂快了一半。

二階梯度，超越 Adam 的最優(yōu)化

為了完成這項(xiàng)工作，研究者仔細(xì)思考了二階優(yōu)化存在的問題與困難，并改進(jìn)了名為 Shampoo 的二階梯度方法。研究者表示，二階優(yōu)化最無解的是目前深度學(xué)習(xí)庫對一階梯度已經(jīng)有了大量優(yōu)化，它們對計(jì)算量與內(nèi)存的要求都不高。然而，對于 Shampoo，每一次迭代的密集計(jì)算，都對應(yīng)用大模型產(chǎn)生了不可逾越的阻攔。

和一階梯度優(yōu)化方法相比，Shampoo 要走向?qū)嵺`，還有如下三大挑戰(zhàn)。

1. 算法上的挑戰(zhàn)

現(xiàn)代機(jī)器學(xué)習(xí)架構(gòu)通常使用很大的嵌入層，維度可能多達(dá)百萬級別。Shampoo 需要對每個(gè)維度計(jì)算一個(gè)預(yù)條件子，但是不管是計(jì)算還是存儲，百萬次的百萬維度矩陣都幾乎是不可處理的。因此從算法上，我們首先就要設(shè)計(jì)一種新方法來解決這類問題。

2. 計(jì)算上的挑戰(zhàn)

Shampoo 的權(quán)重更新式大概是如下這樣的，其中 L 和 R 都是矩陣，它們需要求逆與求根，在計(jì)算上會顯得非常復(fù)雜，這也會拖慢整個(gè)迭代的速度。

之前矩陣求逆與根是可以使用 SVD 計(jì)算的，但是它們太慢了。因此可以考慮如 Schur-Newton 等一些算法，可以將逆 P 次根問題轉(zhuǎn)換為一系列矩陣-向量和矩陣-矩陣的積，所以對于優(yōu)化很有幫助。

圖 1：在不同維度上，計(jì)算某矩陣逆 P 次根的基準(zhǔn)對比。Schur-Newton 迭代方法能夠在 CPU 上高效運(yùn)行，而且相比 SVD 有很大的冗余用于提升。

3. 基礎(chǔ)設(shè)施挑戰(zhàn)

神經(jīng)網(wǎng)絡(luò)加速器通常是定制的，用來讓機(jī)器學(xué)習(xí)程序運(yùn)行地更快，開銷更少。加速器設(shè)計(jì)傾向于低精度（8bit/16bit），能夠滿足現(xiàn)有的基準(zhǔn)。研究者的方法需要雙精度運(yùn)算，因此已有的加速器甚至都不會啟動。

此外，TensorFlow 等深度學(xué)習(xí)庫提供的最優(yōu)化器 API 適應(yīng)于隨機(jī)、一階梯度下降那種模式。而二階優(yōu)化器需要與訓(xùn)練循環(huán)做交互，因此從實(shí)現(xiàn)上需要對框架底層做出修正。

雖然難，但還是能攻破

三大挑戰(zhàn)使實(shí)現(xiàn)二階梯度優(yōu)化器異常復(fù)雜，研究者針對不同的問題提出了一系列優(yōu)化組件、優(yōu)化算法。最終搞定的分布式 Shampoo 在 CPU、TPU、GPU 等硬件上高效運(yùn)行，這也是我們第一次看到二階優(yōu)化器在大模型上能 Work 的新研究。

研究者們首先分析了標(biāo)準(zhǔn)的數(shù)據(jù)并行方法。每個(gè)加速器的核都會在一個(gè)分批數(shù)據(jù)上進(jìn)行前向和反向傳播。然后算法會對數(shù)據(jù)批進(jìn)行梯度聚合，并使用 all-reduction 的方法獲得平均梯度。聚合的梯度被用來進(jìn)行權(quán)重更新。前向和反向傳播在所有核中并行計(jì)算。

為什么使用 All-reduction？這增加了一個(gè)屏障，使得所有核同步聚合批的梯度，并進(jìn)行權(quán)重更新。圖 3 中，研究者評價(jià)了 Transformer 模型每一步的計(jì)算開銷。

圖 3：使用 Diagonal AdaGrad 優(yōu)化器的 Transformer 模型的每迭代步延遲時(shí)間為 134 毫秒，其中（1）前向傳播 57 毫秒；后向傳播 71 毫秒；all reduction：4 毫秒；權(quán)重更新：2 毫秒。

本文分布式系統(tǒng)實(shí)現(xiàn)的整體設(shè)計(jì)如下圖 4 的時(shí)間軸所示，具體分布式架構(gòu)可查閱原論文。

圖 4：本文優(yōu)化算法設(shè)計(jì)的時(shí)間軸。在每一步上計(jì)算所有張量的預(yù)調(diào)節(jié)器統(tǒng)計(jì)數(shù)據(jù)。預(yù)調(diào)節(jié)器只在每 N 步上計(jì)算，并且計(jì)算會分配給訓(xùn)練系統(tǒng)上可用的所有 CPU 核心。運(yùn)算進(jìn)行流水線處理，這樣就實(shí)現(xiàn)了開銷均攤。

可實(shí)戰(zhàn)的二階梯度優(yōu)化

研究者在包含 3630 萬個(gè)句對的 WMT 14 英法標(biāo)準(zhǔn)機(jī)器翻譯數(shù)據(jù)集上驗(yàn)證了分布式系統(tǒng)實(shí)現(xiàn)方法的有效性。此外，他們在實(shí)驗(yàn)中使用了當(dāng)前 SOTA Transformer 架構(gòu)，該架構(gòu)包含 9330 萬個(gè)參數(shù)和 6 層的編碼器-解碼器結(jié)構(gòu)。

實(shí)驗(yàn)在 32 核谷歌 Cloud TPU v3 Pod 上運(yùn)行，結(jié)果如下圖 6 所示，本文提出的 Shampoo 算法只需一半迭代數(shù)就能實(shí)現(xiàn)與 AdaGrad 和 Adam 相同的準(zhǔn)確度。

圖 6：WMT 14 英法翻譯數(shù)據(jù)集上的 Transformer 模型，Shampoo 二階梯度算法的收斂速度在迭代數(shù)上快了 1.95 倍，且就算要計(jì)算二階梯度，每一次迭代也只慢了 16%，總體上來說節(jié)省了 40% 的執(zhí)行時(shí)間。

研究者還利用一個(gè)大型 Transformer 模型進(jìn)行實(shí)驗(yàn)，該模型包含 3.754 億個(gè)參數(shù)和 6 層的編碼器-解碼器結(jié)構(gòu)。實(shí)驗(yàn)結(jié)果如下圖 12 所示，端到端的執(zhí)行時(shí)間實(shí)現(xiàn)了提升。

圖 12：WMT 14 英法翻譯數(shù)據(jù)集上的 Transformer-Big 模型，Shampoo 二階梯度算法的收斂速度在迭代數(shù)上快了 2 倍，且就算要計(jì)算二階梯度，每一次迭代也只慢了 40%，總體上來說節(jié)省了 30% 的執(zhí)行時(shí)間。

最后，研究者在 ImageNet-2012 數(shù)據(jù)集上訓(xùn)練了 ResNet-50 模型，并對使用 SGD+Momentum 的 SOTA 基準(zhǔn)方法、本文提出的 Shampoo 二階梯度算法以及 Adagrad 方法的測試結(jié)果進(jìn)行對比，結(jié)果如下圖 14 和表 2 所示。

研究者發(fā)現(xiàn)，Shampoo 二階梯度算法雖然未能在測試損失或準(zhǔn)確度方面實(shí)現(xiàn)任何改進(jìn)，但與調(diào)整好的 SGD+Momentum 基準(zhǔn)方法相比，該方法能夠更快地減少訓(xùn)練損失。

圖 14：在 ImageNet-2012 數(shù)據(jù)集上訓(xùn)練 ResNet-50 模型時(shí)的訓(xùn)練（圖左）和測試（圖右）交叉熵變化圖。

表 2：當(dāng)訓(xùn)練 ResNet-50 模型的 batch size=4096 時(shí)，三種方法在 ImageNet-2012 數(shù)據(jù)集上的準(zhǔn)確度測試結(jié)果對比。

相信很多讀者都學(xué)過最優(yōu)化方法這類理論課程，我們會發(fā)現(xiàn)梯度下降，或者稱之為最速下降法是最簡單的方法，「鋸齒」現(xiàn)象令它在很多領(lǐng)域上都存在問題。我們也會發(fā)現(xiàn)各種二階優(yōu)化、擬二階優(yōu)化在理論上性質(zhì)遠(yuǎn)遠(yuǎn)超過它。然而在深度學(xué)習(xí)領(lǐng)域，由于數(shù)據(jù)與模型的規(guī)模，我們采用的都是「最速下降法」這個(gè)大家庭。

很多時(shí)候，我們會想，之前累積的那么多優(yōu)秀方法，完全應(yīng)用不到深度學(xué)習(xí)嗎？而這篇論文至少告訴我們，計(jì)算、內(nèi)存等各種困難，我們都是有機(jī)會克服的，相信那些具有強(qiáng)硬理論支持的最優(yōu)化方法，最終會在深度學(xué)習(xí)展現(xiàn)它們的魅力。

推薦閱讀 中文版開源！ 一份來自亞馬遜工程師寫的 Google 面試指南，太火了10個(gè)必會的 PyCharm 技巧 為了追到小姐姐，我用 Python 制作了一個(gè)機(jī)器人 青出于藍(lán)而勝于藍(lán)，這是一款脫胎于Jupyter Notebook的新型編程環(huán)境 【中文教程】簡單粗暴入門TensorFlow 2.0 | 北大學(xué)霸出品

總結(jié)

以上是生活随笔為你收集整理的二阶梯度优化新崛起，超越 Adam，Transformer 只需一半迭代量的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。