全面碾壓AdamW！谷歌新出優(yōu)化器內(nèi)存小、效率高，網(wǎng)友：訓(xùn)練GPT 2果然快

谷歌的 Lion 優(yōu)化器將成為訓(xùn)練大模型或大批量的「福音」。

優(yōu)化器即優(yōu)化算法，在神經(jīng)網(wǎng)絡(luò)訓(xùn)練中起著關(guān)鍵作用。近年來(lái)，研究者引入了大量的手工優(yōu)化器，其中大部分是自適應(yīng)優(yōu)化器。Adam 以及 A****ctor 優(yōu)化器仍然占據(jù)訓(xùn)練神經(jīng)網(wǎng)絡(luò)的主流，尤其在語(yǔ)言、視覺(jué)和多模態(tài)領(lǐng)域更是如此。
除了人工引入優(yōu)化器外，還有一個(gè)方向是程序自動(dòng)發(fā)現(xiàn)優(yōu)化算法。此前有人提出過(guò) L2O（learning to optimize），該方法通過(guò)訓(xùn)練神經(jīng)網(wǎng)絡(luò)來(lái)發(fā)現(xiàn)優(yōu)化器。然而，這些黑盒優(yōu)化器通常是在有限數(shù)量的小任務(wù)上訓(xùn)練而成，很難泛化到大模型。
還有人嘗試了其他方法，他們應(yīng)用強(qiáng)化學(xué)習(xí)或蒙特卡羅采樣來(lái)發(fā)現(xiàn)新的優(yōu)化器。然而，為了簡(jiǎn)化搜索，這些方法通常會(huì)限制搜索空間，從而限制發(fā)現(xiàn)其他優(yōu)化器的可能性。因此，目前方法尚未達(dá)到 SOTA 水平。
最近幾年，值得一提的是 AutoML-Zero，它試圖在評(píng)估任務(wù)時(shí)搜索機(jī)器學(xué)習(xí) pipeline 的每個(gè)組件，對(duì)優(yōu)化器的發(fā)現(xiàn)很有借鑒意義。
在本文中，來(lái)自谷歌、 UCLA 的研究者提出了一種通過(guò)程序搜索發(fā)現(xiàn)深度神經(jīng)網(wǎng)絡(luò)訓(xùn)練的優(yōu)化算法的方法，進(jìn)而發(fā)現(xiàn) Lion（EvoLved Sign Momentum）優(yōu)化器。實(shí)現(xiàn)這一目標(biāo)面臨兩個(gè)挑戰(zhàn)：首先是在無(wú)限稀疏的程序空間中尋找高質(zhì)量的算法，其次是選擇可以從小型任務(wù)泛化到更大、SOTA 任務(wù)的算法。為了應(yīng)對(duì)這些挑戰(zhàn)，該研究采用了一系列技術(shù)，包括具有熱啟動(dòng)和重啟的進(jìn)化搜索、抽象執(zhí)行、funnel 選擇和程序簡(jiǎn)化。

• 論文地址：https://arxiv.org/pdf/2302.06675.pdf
• 項(xiàng)目地址：https://github.com/google/automl/tree/master/lion

與 AdamW 和各種自適應(yīng)優(yōu)化器（這些優(yōu)化器需要同時(shí)保存一階和二階矩）相比，Lion 只需要?jiǎng)恿坎⒗梅?hào)操作來(lái)計(jì)算更新，并且將額外的內(nèi)存占用減半。這在訓(xùn)練大模型或大批量時(shí)很有用。例如，AdamW 需要至少 16 個(gè) TPU V4 來(lái)訓(xùn)練圖像大小為 224、批大小為 4096 的 ViT-B/16，而 Lion 只需要 8 個(gè)。另一個(gè)實(shí)際好處是，由于 Lion 的簡(jiǎn)單性，Lion 在實(shí)驗(yàn)中具有更快的運(yùn)行時(shí)間（steps /sec），通常比 AdamW 和 A****ctor 提速 2-15%，具體取決于任務(wù)、代碼庫(kù)和硬件。
論文一作 Xiangning Chen 表示：我們的符號(hào)程序搜索發(fā)現(xiàn)了一個(gè)只跟蹤動(dòng)量的有效優(yōu)化器 ——Lion。與 Adam 相比，它獲得了 88.3% 的零樣本和 91.1% 的微調(diào) ImageNet 準(zhǔn)確率，以及高達(dá) 5 倍（與 ViT 相比）、2.3 倍（與擴(kuò)散模型相比）和 2 倍（與 LM 想比）訓(xùn)練效率。

Lion 在一系列模型（Transformer、MLP、ResNet、U-Net 和 Hybrid）和任務(wù)（圖像分類、視覺(jué) - 語(yǔ)言對(duì)比學(xué)習(xí)、擴(kuò)散、語(yǔ)言建模和微調(diào)）中展示了出色的性能。值得注意的是，通過(guò)在 BASIC 中用 Lion 替換 A****ctor，該研究在 ImageNet 上實(shí)現(xiàn)了 88.3% 的零樣本和 91.1% 的微調(diào)精度（Pham 等人，2021），分別超過(guò)之前的 SOTA 結(jié)果 2% 和 0.1%。
此外，Lion 將 JFT 上的預(yù)訓(xùn)練計(jì)算量減少了多達(dá) 5 倍，將擴(kuò)散模型的訓(xùn)練效率提高了 2.3 倍，并獲得了更好的 FID 分?jǐn)?shù)，并在語(yǔ)言建模方面提供了類似或更好的性能，節(jié)省了多達(dá) 2 倍的計(jì)算量。
推特用戶 crumb 表示：自己使用谷歌的 Lion 優(yōu)化器訓(xùn)練了一個(gè) 124M 參數(shù) GPT2 模型，發(fā)現(xiàn)達(dá)到與 Adam 相同的損失所需的 step 數(shù)減少了 37.5%。
圖源：https://twitter.com/aicrumb/status/1626053855329898496
算法的符號(hào)發(fā)現(xiàn)
本文使用程序形式的符號(hào)表示具有以下優(yōu)點(diǎn)：（1）它符合算法必須作為程序來(lái)執(zhí)行的事實(shí)；(2) 與神經(jīng)網(wǎng)絡(luò)等參數(shù)化模型相比，程序等符號(hào)表示更易于分析、理解和遷移到新任務(wù)；(3) 程序長(zhǎng)度可用于估計(jì)不同程序的復(fù)雜性，從而更容易選擇更簡(jiǎn)單、通常更具通用性的程序。這項(xiàng)工作側(cè)重于深度神經(jīng)網(wǎng)絡(luò)訓(xùn)練的優(yōu)化器，但該方法通常適用于其他任務(wù)。
下圖中，程序 2 這個(gè)簡(jiǎn)化代碼片段使用與 AdamW 相同的 signature ，以確保發(fā)現(xiàn)的算法具有更小或相等的內(nèi)存占用空間；程序 3 中給出了 AdamW 的示例表示。

該研究采用以下技術(shù)來(lái)解決無(wú)限和稀疏搜索空間帶來(lái)的挑戰(zhàn)。首先應(yīng)用正則化，因?yàn)樗?jiǎn)單，可擴(kuò)展，并在許多 AutoML 搜索任務(wù)中很成功；其次是簡(jiǎn)化程序空間中的冗余；最后，為了降低搜索成本，該研究通過(guò)減少模型大小、訓(xùn)練示例的數(shù)量和距離目標(biāo)任務(wù)的 step 使成本降低。
左：顯示了五次進(jìn)化搜索實(shí)驗(yàn)的均值和標(biāo)準(zhǔn)誤差。右：隨著搜索的進(jìn)行，冗余語(yǔ)句的百分比和緩存命中率都在增加。
 Lion 的推導(dǎo)與分析
研究者表示，優(yōu)化器 Lion 具有簡(jiǎn)潔性、高內(nèi)存效率以及在搜索和元驗(yàn)證方面的強(qiáng)大性能。
推導(dǎo) 
搜索和 funnel 選擇過(guò)程促成了程序 4，它通過(guò)從原始程序 8（附錄）自動(dòng)刪除冗余語(yǔ)句獲得。研究者進(jìn)一步簡(jiǎn)化并得到程序 1 中的最終算法（Lion）。在簡(jiǎn)化過(guò)程中從程序 4 中刪除了幾個(gè)不必要的元素。其中 cosh 函數(shù)被刪除，因?yàn)?m 將在下一次迭代中重新分配（第 3 行）。使用 arcsin 和 clip 的語(yǔ)句也被刪除，因?yàn)檠芯空哂^察到?jīng)]有它們質(zhì)量不會(huì)下降。三個(gè)紅色語(yǔ)句轉(zhuǎn)換為一個(gè)符號(hào)函數(shù)。
盡管程序 4 中同時(shí)使用了 m 和 v，但 v 僅改變了動(dòng)量的更新方式（兩個(gè)具有常數(shù)～0.9 和 ～1.1 的插值函數(shù)等同于一個(gè)具有 ～0.99 的函數(shù)）并且不需要單獨(dú)跟蹤。注意不再需要偏差校正，因?yàn)樗粫?huì)改變方向。 
如下算法 2 顯示了偽代碼。

分析 
符號(hào)更新和正則化。Lion 算法通過(guò)符號(hào)操作在所有維度上產(chǎn)生了具有統(tǒng)一幅度的更新，這在原理上不同于各種自適應(yīng)優(yōu)化器。直觀來(lái)看，符號(hào)操作為更新添加了噪聲，作為了一種正則化形式并有助于泛化。下圖 11（右）展示了一個(gè)證據(jù)。

動(dòng)量跟蹤。與 AdamW 和動(dòng)量 SGD 中常用的 0.9 相比，Lion 中用于跟蹤動(dòng)量的默認(rèn) EMA 因子為 0.99 (β_2)。EMA 因子和插值的這種選擇，使得 Lion 在記住動(dòng)量梯度的 10 倍 history 和在更新中對(duì)當(dāng)前梯度施加更多權(quán)重之間取得平衡。
超參數(shù)和批大小選擇。與 AdamW 和 A****ctor 相比，Lion 更簡(jiǎn)單并且超參數(shù)更少，因?yàn)樗恍枰?? 和因式分解相關(guān)的參數(shù)。Lion 需要更小的學(xué)習(xí)率，進(jìn)而需要更大的解耦權(quán)重衰減才能達(dá)到類似的有效權(quán)重衰減強(qiáng)度 (lr * λ)。
內(nèi)存和運(yùn)行時(shí)優(yōu)勢(shì)。Lion 只保存動(dòng)量，比 AdamW 等流行的自適應(yīng)優(yōu)化器占用的內(nèi)存更小，這在訓(xùn)練大模型和 / 或使用大批量時(shí)非常有用。例如，AdamW 需要至少 16 個(gè) TPU V4 芯片來(lái)訓(xùn)練圖像分辨率為 224 且批大小為 4,096 的 ViT-B/16，而 Lion 只需要 8 個(gè)（均具有 bfloat16 動(dòng)量）。
Lion 評(píng)估結(jié)果
在實(shí)驗(yàn)部分，研究者在各種基準(zhǔn)上對(duì) Lion 進(jìn)行了評(píng)估，主要將它與流行的 AdamW（或者當(dāng)內(nèi)存成為瓶頸時(shí)的 A****ctor）進(jìn)行比較。
圖像分類
研究者在圖像分類任務(wù)上執(zhí)行涵蓋各種數(shù)據(jù)集和架構(gòu)的實(shí)驗(yàn)。除了在 ImageNet 上從頭開始訓(xùn)練之外，他們還在兩個(gè)更大的成熟數(shù)據(jù)集 ImageNet-21K 和 JFT 上進(jìn)行預(yù)訓(xùn)練。圖片大小默認(rèn)為 224。
首先在 ImageNet 上從頭開始訓(xùn)練。研究者訓(xùn)練 ResNet-50 達(dá)到 90 個(gè) epoch，批大小為 1,024，其他模型訓(xùn)練 300 個(gè) epoch，批大小為 4,096。如下表 2 所示，Lion 在各種架構(gòu)上顯著優(yōu)于 AdamW。
其次在 ImageNet-21K 上進(jìn)行預(yù)訓(xùn)練。研究者在 ImageNet-21K 上對(duì) ViT-B/16 和 ViT-L/16 進(jìn)行了 90 個(gè) epoch 的預(yù)訓(xùn)練，批大小為 4,096。下表 2 顯示，即使訓(xùn)練集擴(kuò)大 10 倍，Lion 仍然優(yōu)于 AdamW。

最后在 JFT 上進(jìn)行預(yù)訓(xùn)練。為了突破極限，研究者在 JFT 上進(jìn)行大量實(shí)驗(yàn)。下圖 4 顯示了 JFT-300M 上不同預(yù)訓(xùn)練預(yù)算下三種 ViT 模型（ViT-B/16、ViT-L/16 和 ViT-H/14）的準(zhǔn)確率。Lion 使 ViT-L/16 能夠與 AdamW 在 ImageNet 和 ImageNet V2 上訓(xùn)練的 ViT-H/14 性能相匹配，但預(yù)訓(xùn)練成本降低了 3 倍。

下表 3 顯示了微調(diào)結(jié)果，具有更高的分辨率和 Polyak 平均。研究者使用的 ViT-L/16 與之前由 AdamW 訓(xùn)練的 ViT-H/14 結(jié)果相匹配，同時(shí)參數(shù)減少了 2 倍。在將預(yù)訓(xùn)練數(shù)據(jù)集擴(kuò)展到 JFT-3B 后，Lion 訓(xùn)練的 ViT-g/14 優(yōu)于之前的 ViT-G/14 結(jié)果，參數(shù)減少了 1.8 倍。

視覺(jué)語(yǔ)言對(duì)比學(xué)習(xí)
本節(jié)重點(diǎn)介紹 CLIP 風(fēng)格的視覺(jué)語(yǔ)言對(duì)比訓(xùn)練。研究者沒(méi)有從頭開始學(xué)習(xí)所有參數(shù)，而是使用強(qiáng)大的預(yù)訓(xùn)練模型初始化圖像編碼器。
對(duì)于鎖定圖像文本調(diào)整 (LiT) ，研究者通過(guò)使用相同的 frozen 預(yù)訓(xùn)練 ViT 以對(duì)比方式訓(xùn)練文本編碼器，進(jìn)而在 LiT 上對(duì) Lion 和 AdamW 進(jìn)行比較。下表 4 顯示了 3 種模型尺度上的零樣本圖像分類結(jié)果，Lion 展示出了對(duì) AdamW 的持續(xù)改進(jìn)。

下圖 5（左）展示了 LiT-B/16-B 的示例零樣本學(xué)習(xí)曲線，并在其他兩個(gè)數(shù)據(jù)集上得到了類似的結(jié)果。

擴(kuò)散模型
最近，擴(kuò)散模型在圖像生成方面取得了巨大的成功。鑒于其巨大的潛力，研究者測(cè)試了 Lion 在無(wú)條件圖像合成和多模式文本到圖像生成方面的表現(xiàn)。
對(duì)于 ImageNet 上的圖像合成，研究者利用 2021 年論文《 Diffusion models beat gans on image synthesis 》中引入的改進(jìn) U-Net 架構(gòu)，在 ImageNet 上執(zhí)行 64×64、128×128 和 256×256 圖像生成。如上圖 5（中和右）所示，Lion 可以在 FID 分?jǐn)?shù)上實(shí)現(xiàn)更好的質(zhì)量和更快的收斂速度。
對(duì)于文本到圖像的生成，下圖 6 展示了學(xué)習(xí)曲線。雖然 64 × 64 base 模型沒(méi)有明顯改進(jìn)，但 Lion 在文本條件超分辨率模型上優(yōu)于 AdamW。與 AdamW 相比，Lion 獲得了更高的 CLIP 分?jǐn)?shù)并且具有更小的噪聲 FID 指標(biāo)。

語(yǔ)言建模和微調(diào)
本節(jié)重點(diǎn)介紹語(yǔ)言建模和微調(diào)。在純語(yǔ)言任務(wù)上，研究者發(fā)現(xiàn)調(diào)整 β_1 和 β_2 可以提高 AdamW 和 Lion 的質(zhì)量。
對(duì)于自回歸語(yǔ)言建模，下圖 7 展示了 Wiki-40B 的 token 級(jí)困惑度和 PG-19 的 word 級(jí)困惑度。Lion 始終實(shí)現(xiàn)比 AdamW 更低的驗(yàn)證困惑度。在 Wiki-40B 和 PG-19 上訓(xùn)練中型模型時(shí)，它分別實(shí)現(xiàn)了 1.6 倍和 1.5 倍的加速。當(dāng)模型增加到大尺寸時(shí)，PG-19 進(jìn)一步實(shí)現(xiàn)了 2 倍加速。

對(duì)于掩蔽語(yǔ)言建模，研究者還在 C4 數(shù)據(jù)集上執(zhí)行 BERT 訓(xùn)練，其中 Lion 在驗(yàn)證困惑度方面的表現(xiàn)略好于 AdamW。相關(guān)學(xué)習(xí)曲線如下圖 11（左）所示。

對(duì)于微調(diào)，研究者在 GLUE 基準(zhǔn)上微調(diào) Base（220M）、Large（770M）和最大的 11B T5 模型。下表 6 顯示了 GLUE 開發(fā)集上的結(jié)果，平均而言，Lion 在所有 3  個(gè)模型規(guī)模上都擊敗了 AdamW。

與其他流行優(yōu)化器的比較
該研究還使用四種流行的優(yōu)化器 RAdam、NAdam、AdaBelief 和 AMSGrad 在 ImageNet 上訓(xùn)練 ViT-S/16 和 ViT-B/16（使用 RandAug 和 Mixup）。如下表 7 所示，Lion 仍然是表現(xiàn)最好的。

更多技術(shù)細(xì)節(jié)請(qǐng)參閱原論文。