Lion : 超越 AdamW 的優化算法

出品人:Towhee 技術團隊

優化算法,即優化器,在訓練神經網絡中起着基礎作用。

近年來引入了大量手工優化器,其中大部分是自適應優化器。然而,具有解耦權重衰減的 Adam,也稱爲 AdamW,和具有因數二次矩的 Adafactor,仍然是訓練大多數深度神經網絡的事實上的標準優化器,尤其是最近最先進的語言、視覺和多模態模型。

另一個方向是自動發現這樣的優化算法。學習優化 (L2O) 方法建議通過訓練參數化模型(例如神經網絡)來發現優化器以輸出更新。然而,那些通常在有限數量的小任務上訓練的黑盒優化器很難推廣到最先進的設置,在這些設置中,更大的模型需要更多的訓練步驟來訓練。

另一類方法應用強化學習或蒙特卡洛採樣來發現新的優化器,其中搜索空間由由預定義操作數和運算符組成的樹定義。爲了使搜索易於管理,他們通常通過使用固定操作數和限制樹的大小來限制搜索空間,從而限制發現的可能性。因此,所發現的算法尚未達到最先進的水平。

Google 和 UCLA 提出 Lion (EvoLved Sign Momentum) 優化算法,以優越的性能和良好的效果超越了經典的優化算法。Lion 通過程序搜索來發現優化算法,並且用其改進深度神經網絡訓練。它能夠解決在無限、稀疏空間的搜索挑戰。

除了更具效率以外,Lion 還具有強大的泛化能力,跨越了體系結構、數據集和任務的限制。在圖像分類方面,Lion 在 ImageNet 上將 ViT 的準確性提高了高達 2%,並在 JFT 上節省了高達 5 倍的預訓練計算。

在視覺語言對比學習方面,Lion 在 ImageNet 上實現了 88.3% 的 zero-shot 和 91.1% 的微調準確度,分別超過了之前的最佳結果 2% 和 0.1%。在擴散模型上,Lion 的表現優於 Adam,實現了更好的 FID 得分並將訓練計算減少了高達 2.3 倍。對於自迴歸、掩碼語言建模和微調,Lion 也表現出與 Adam 相當或更好的性能。

Lion 利用了高效的搜索技術來探索無限稀疏的程序空間。爲了彌合代理和目標任務之間的廣義間隙,它還引入了程序選擇和簡化策略。與自適應優化器不同,Lion 對於通過符號操作計算的每個參數都採用了相同的幅度的更新。它的性能增益隨着訓練批次大小的增加而增加。由於符號函數產生的更新的範數更大,它還需要較小的學習率。

AdamW vs. Lion optimizer.

上圖爲 AdamW 與 Lion 優化算法的對比。

可以看到,與 AdamW 和各種自適應優化器需要同時保存一階和二階矩相比,Lion 只需要動量,將額外的內存佔用減半。這在訓練大型模型和/或大批量時很有用。

例如,AdamW 需要至少 16 個 TPU V4 芯片來訓練圖像大小爲 224、批量大小爲 4,096 的 ViT-B/16,而 Lion 只需要八個。

另一個實際好處是,由於 Lion 的簡單性,Lion 在我們的實驗中具有更快的運行時間(步數/秒),通常比 AdamW 和 Adafactor 提速 2-15%,具體取決於任務、代碼庫和硬件。

因此,Lion 是一種表現出色的優化器,比AdamW 表現更強大。Lion 具有更快速和更節省顯存的特點,有望成爲未來主流優化器之一。如果你在訓練時,感覺保存的訓練過程中的狀態太大,導致性能緩慢,不妨試試用 Lion。

自從 Adam 被提出以來,由於其快速收斂的特性,已成爲許多模型的默認優化器。但是,一些學者擔心這種現象可能導致所有的模型改進都朝着有利於 Adam 的方向發展。因此,發現比 Adam 更簡單、更有效的優化器是很重要的。Lion 能夠獲得更好的泛化性能,這可能是因爲它引入了額外的噪聲,使得模型進入了更平坦但未必更小的損失區域。雖然 Lion 有更強的魯棒性,但它並不是完美的,比如在小的 batch_size 時,Lion 的表現不如 AdamW,因爲噪聲必須適量才能更好地發揮作用。

相關資料

代碼地址:https://github.com/google/automl/tree/master/lion 論文鏈接:Symbolic Discovery of Optimization Algorithms

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章