Lion : 超越 AdamW 的優化演算法

語言: CN / TW / HK

時間 2023-03-25 19:01:19 Zilliz

主題:

出品人：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，因為噪聲必須適量才能更好地發揮作用。

「其他文章」