全面碾壓AdamW！谷歌新出優化器記憶體小、效率高，網友：訓練GPT 2果然快

優化器即優化演算法，在神經網路訓練中扮演關鍵角色。近年來，研究者引入了大量的手工優化器，其中大部分是自適應優化器。 Adam 以及 Adafactor 優化器仍然佔據訓練神經網路的主流，尤其在語言、視覺和多模態領域更是如此。

除了人工引入優化器外，還有一個方向是程式自動發現最佳化演算法。先前有人提出 L2O（learning to optimize），透過訓練神經網路來發現優化器。然而，這些黑盒優化器通常是在有限數量的小任務上訓練而成，很難泛化到大模型。

還有人嘗試了其他方法，他們應用強化學習或蒙特卡羅取樣來發現新的最佳化器。然而，為了簡化搜索，這些方法通常會限制搜索空間，從而限制發現其他優化器的可能性。因此，目前方法尚未達到 SOTA 水準。

最近幾年，值得一提的是 AutoML-Zero，它試圖在評估任務時搜尋機器學習 pipeline 的每個元件，對優化器的發現很有借鑒意義。

在本文中，來自Google、 UCLA 的研究者提出了一種透過程式搜尋發現深度神經網路訓練的最佳化演算法的方法，進而發現Lion（EvoL「ign Mo#me

• ##n##tum）優化器
• 。要實現這一目標面臨兩個挑戰：首先是在無限稀疏的程式空間中尋找高品質的演算法，其次是選擇可以從小型任務泛化到更大、SOTA 任務的演算法。為了應對這些挑戰，該研究採用了一系列技術，包括具有熱啟動和重啟的進化搜尋、抽象執行、funnel 選擇和程序簡化。 論文網址：https://arxiv.org/pdf/2302.06675.pdf

專案網址：https://github.com/google/automl/tree/master/lion

與AdamW 和各種自適應優化器（這些優化器需要同時保存一階和二階矩）相比，Lion 只需要動量並利用符號操作來計算更新，並且將額外的記憶體佔用減半。這在訓練大模型或大批量時很有用。例如，AdamW 需要至少 16 個 TPU V4 來訓練影像大小為 224、批次大小為 4096 的 ViT-B/16，而 Lion 只需要 8 個。另一個實際好處是，由於 Lion 的簡單性，Lion 在實驗中具有更快的運行時間（steps /sec），通常比 AdamW 和 Adafactor 提速 2-15%，具體取決於任務、程式碼庫和硬體。

論文一作 Xiangning Chen 表示：我們的符號程式搜尋發現了一個只追蹤動量的有效優化器 ——Lion。與Adam 相比，它獲得了88.3% 的零樣本和91.1% 的微調ImageNet 準確率，以及高達5 倍（與ViT 相比）、2.3 倍（與擴散模型相比）和2 倍（與LM 想比）訓練效率。

############Lion 在一系列模型（Transformer、MLP、ResNet、U-Net 和Hybrid）和任務（圖像分類、視覺- 語言對比學習、擴散、語言建模和微調）中展示了出色的性能。值得注意的是，透過在BASIC 中以Lion 取代Adafactor，該研究在ImageNet 上實現了88.3% 的零樣本和91.1% 的微調精度（Pham 等人，2021），分別超過先前的SOTA 結果2% 和0.1 %。 ############此外，Lion 將JFT 上的預訓練計算量減少了多達5 倍，將擴散模型的訓練效率提高了2.3 倍，並獲得了更好的FID 分數，並在語言建模方面提供了類似或更好的效能，節省了多達2 倍的運算量。 ######

推特用戶 crumb 表示：自己使用Google的 Lion 優化器訓練了一個 124M 參數 GPT2 模型，發現達到與 Adam 相同的損失所需的 step 數減少了 37.5%。

圖來源：https://twitter.com/aicrumb/status/1626053855329898496

演算法的符號發現

本文使用程式形式的符號表示具有以下優點：（1）它符合演算法必須作為程式來執行的事實；(2) 與神經網路等參數化模型相比，程式等符號表示更易於分析、理解和遷移到新任務；(3) 程式長度可用於估計不同程式的複雜性，從而更容易選擇更簡單、通常更具通用性的程序。 這項工作著重於深度神經網路訓練的最佳化器，但該方法通常適用於其他任務。

下圖中，程式2 這個簡化程式碼片段使用與AdamW 相同的signature ，以確保發現的演算法具有較小或相等的記憶體佔用空間；程式3 中給出了AdamW 的範例表示。

該研究採用以下技術來解決無限和稀疏搜尋空間帶來的挑戰。首先應用正則化，因為它簡單，可擴展，並在許多AutoML 搜索任務中很成功；其次是簡化程序空間中的冗餘；最後，為了降低搜索成本，該研究通過減少模型大小、訓練示例的數量和距離目標任務的step 使成本降低。

左：顯示了五次演化搜尋實驗的平均值和標準誤差。右：隨著搜尋的進行，冗餘語句的百分比和快取命中率都在增加。

 Lion 的推導與分析

研究者表示，優化器Lion 具有簡潔性、高記憶體效率以及在搜尋和元驗證方面的強大性能。

推導 

##搜尋和funnel 選擇過程促成了程式4，它透過從原始程式8（附錄）自動刪除冗餘語句獲得。研究者進一步簡化並得到程式 1 中的最終演算法（Lion）。在簡化過程中從程式 4 中刪除了幾個不必要的元素。其中 cosh 函數被刪除，因為 m 將在下一次迭代中重新分配（第 3 行）。使用 arcsin 和 clip 的語句也被刪除，因為研究者觀察到沒有它們品質不會下降。三個紅色語句轉換為一個符號函數。

儘管程式4 中同時使用了m 和v，但v 僅改變了動量的更新方式（兩個具有常數∼0.9 和∼1.1 的插值函數等同於一個具有∼ 0.99 的函數）並且不需要單獨追蹤。注意不再需要偏差校正，因為它不會改變方向。 

#如下演算法 2 顯示了偽代碼。

#分析 

符號更新與正規化。 Lion 演算法透過符號運算在所有維度上產生了具有統一幅度的更新，這在原理上不同於各種自適應最佳化器。 直觀來看，符號操作為更新添加了噪聲，作為了一種正則化形式並有助於泛化。下圖 11（右）展示了一個證據。

動量追蹤。與 AdamW 和動量 SGD 中常用的 0.9 相比，Lion 中用於追蹤動量的預設 EMA 因子為 0.99 (β_2)。 EMA 因子和插值的這種選擇，使得 Lion 在記住動量梯度的 10 倍 history 和在更新中對當前梯度施加更多權重之間取得平衡。

超參數和批次大小選擇。與 AdamW 和 Adafactor 相比，Lion 更簡單且超參數更少，因為它不需要 ϵ 和因式分解相關的參數。 Lion 需要更小的學習速率，進而需要更大的解耦權重衰減才能達到類似的有效權重衰減強度 (lr * λ)。

記憶體和運行時優勢。 Lion 只保存動量，比 AdamW 等流行的自適應優化器佔用的記憶體更小，這在訓練大模型和 / 或使用大批量時非常有用。例如，AdamW 需要至少 16 個 TPU V4 晶片來訓練影像解析度為 224 且批次大小為 4,096 的 ViT-B/16，而 Lion 只需要 8 個（兩者皆具有 bfloat16 動量）。

Lion 評估結果

在實驗部分，研究者在各種基準上對Lion 進行了評估，主要將它與流行的AdamW（或當記憶體成為瓶頸時的Adafactor）進行比較。

圖像分類

#研究者在影像分類任務上執行涵蓋各種資料集和架構的實驗。除了在 ImageNet 上從頭開始訓練之外，他們還在兩個更大的成熟資料集 ImageNet-21K 和 JFT 上進行預訓練。圖片大小預設為 224。

首先在 ImageNet 上從頭開始訓練。研究者訓練 ResNet-50 達到 90 個 epoch，批次大小為 1,024，其他模型訓練 300 個 epoch，批次大小為 4,096。如下表 2 所示，Lion 在各種架構上顯著優於 AdamW。

其次在 ImageNet-21K 上進行預訓練。研究者在 ImageNet-21K 上對 ViT-B/16 和 ViT-L/16 進行了 90 個 epoch 的預訓練，批次大小為 4,096。下表 2 顯示，即使訓練集擴大 10 倍，Lion 仍優於 AdamW。

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

#下表3 顯示了微調結果，具有更高的解析度和Polyak 平均。研究者使用的 ViT-L/16 與先前由 AdamW 訓練的 ViT-H/14 結果相匹配，同時參數減少了 2 倍。在將預訓練資料集擴展到 JFT-3B 後，Lion 訓練的 ViT-g/14 優於先前的 ViT-G/14 結果，參數減少了 1.8 倍。

視覺語言對比學習

#本節重點介紹 CLIP 風格的視覺語言對比訓練。研究者沒有從頭開始學習所有參數，而是使用強大的預訓練模型初始化影像編碼器。

對於鎖定圖像文字調整 (LiT) ，研究者透過使用相同的 frozen 預訓練 ViT 以對比方式訓練文字編碼器，進而在 LiT 上對 Lion 和 AdamW 進行比較。下表 4 顯示了 3 種模型尺度上的零樣本影像分類結果，Lion 顯示了對 AdamW 的持續改進。

下圖5（左）展示了LiT-B/16-B 的範例零樣本學習曲線，並在其他兩個數據集上得到了類似的結果。

#擴散模型

#最近，擴散模型在影像生成方面取得了巨大的成功。鑑於其巨大的潛力，研究者測試了 Lion 在無條件圖像合成和多模式文字到圖像生成方面的表現。

對於ImageNet 上的圖像合成，研究者利用2021 年論文《 Diffusion models beat gans on image synthesis 》中引入的改進U-Net 架構，在ImageNet 上執行64×64 、128×128 和256×256 影像生成。如上圖 5（中和右）所示，Lion 可以在 FID 分數上實現更好的品質和更快的收斂速度。

對於文字到圖像的生成，下圖 6 展示了學習曲線。雖然 64 × 64 base 模型沒有明顯改進，但 Lion 在文字條件超解析度模型上優於 AdamW。與 AdamW 相比，Lion 獲得了更高的 CLIP 分數並且具有更小的噪音 FID 指標。

#語言建模與微調

#本節重點介紹語言建模和微調。在純語言任務上，研究者發現調整 β_1 和 β_2 可以提高 AdamW 和 Lion 的品質。

對於自回歸語言建模，下圖 7 展示了 Wiki-40B 的 token 級困惑度和 PG-19 的 word 級困惑度。 Lion 始終實作比 AdamW 更低的驗證困惑度。在 Wiki-40B 和 PG-19 上訓練中型模型時，它分別實現了 1.6 倍和 1.5 倍的加速。當模型增加到大尺寸時，PG-19 進一步實現了 2 倍加速。

#對於掩蔽語言建模，研究者也會在C4 資料集上執行BERT 訓練，其中Lion 在驗證困惑度方面的表現略優於AdamW。相關學習曲線如下圖 11（左）所示。

#對於微調，研究者在GLUE 基準上微調Base（220M）、Large（ 770M）和最大的11B T5 模型。下表 6 顯示了 GLUE 開發集上的結果，平均而言，Lion 在所有 3  個模型規模上都擊敗了 AdamW。

#與其他流行最佳化器的比較

#研究也使用四種流行的優化器RAdam、NAdam、AdaBelief 和AMSGrad 在ImageNet 上訓練ViT-S/16 和ViT-B/16（使用RandAug 和Mixup）。如下表 7 所示，Lion 仍然是表現最好的。

# 更多技術細節請參閱原始論文。

以上是全面碾壓AdamW！谷歌新出優化器記憶體小、效率高，網友：訓練GPT 2果然快的詳細內容。更多資訊請關注PHP中文網其他相關文章！