Lion Optimizer/zh
| Article | |
|---|---|
| Topic area | optimization |
| Prerequisites | Stochastic gradient descent, Momentum (optimization), Adam optimizer |
概述
Lion 優化器(全稱 EvoLved Sign Momentum,演化式符號動量)是一種用於訓練神經網絡的一階隨機優化算法,由 Chen 等人於 2023 年提出。[1] 與大多數廣泛使用、基於理論原理手工設計的優化器不同,Lion 是通過對候選優化器程序空間進行符號程序搜索而發現的。其更新規則在概念上比 Adam 更簡單,但在大規模 深度學習 工作負載上的準確率通常能夠與之相當或更優,同時所占用的優化器狀態內存大約只有一半。
Lion 的定義性特徵是參數更新方向是插值動量項的逐元素符號。這使得更新的每個坐標都具有相同的大小(由學習率縮放的固定步長),與 Adam 不同——後者的逐坐標自適應縮放由梯度平方的指數移動平均決定。Lion 已被應用於多種大規模 語言模型 和 視覺訓練流水線,並成為自動發現優化器這一研究方向的標杆。
背景與動機
到 2020 年代初,AdamW 已成為訓練 Transformer 的事實標準優化器,儘管存在眾所周知的局限:每個參數需要存儲兩個矩緩衝區(一階矩和二階矩),使優化器內存達到參數本身的兩倍,並且其超參數對模型規模十分敏感。一長串研究工作試圖通過手工設計來取代 Adam,經驗結果褒貶不一:LAMB、Adafactor、Sophia、Shampoo 等各自解決了特定問題,但都未能撼動 Adam 的默認地位。
Lion 源於一個不同的前提:與其設計一個優化器,不如去搜索一個優化器。作者將優化器形式化為一個對梯度和狀態進行操作的短程序,定義了一個由原始操作(加法、乘法、符號、指數移動平均、裁剪等)構成的搜索空間,並使用 演化式程序搜索 配合 元學習風格的代理訓練任務來評估候選項。Lion 是在對更大的發現候選項進行大量剪枝後湧現出的、最簡潔的高性能程序。
算法
設 $ g_t $ 為損失函數在第 $ t $ 步關於參數 $ \theta_t $ 的梯度,設 $ m_t $ 為初始化為零的單個動量緩衝區,設 $ \eta $ 為學習率。Lion 有兩個插值係數 $ \beta_1, \beta_2 \in [0,1) $ 和一個權重衰減係數 $ \lambda \ge 0 $。第 $ t $ 步的更新為:
$ {\displaystyle c_t = \beta_1 \, m_{t-1} + (1 - \beta_1) \, g_t} $
$ {\displaystyle \theta_t = \theta_{t-1} - \eta \, \big( \operatorname{sign}(c_t) + \lambda \, \theta_{t-1} \big)} $
$ {\displaystyle m_t = \beta_2 \, m_{t-1} + (1 - \beta_2) \, g_t} $
第一行通過在先前的動量與當前梯度之間進行插值,得到一個更新方向 $ c_t $。第二行僅使用 $ c_t $ 的逐元素符號來執行參數更新,並附加一個採用 AdamW 風格的解耦權重衰減。第三行則使用另一個不同的插值係數來更新持久動量緩衝區 $ m_t $。關鍵之處在於,用於下一步方向的緩衝區($ m_t $)與用於當前方向的緩衝區(涉及 $ m_{t-1} $ 的混合)由不同的 $ \beta $ 值控制。
Chen 等人報告的默認超參數為 $ \beta_1 = 0.9 $、$ \beta_2 = 0.99 $,學習率大致比 AdamW 的取值小一個數量級,而權重衰減則大致大一個數量級。
直觀理解
三點觀察有助於解釋 Lion 為何有效。
首先,每次更新在每個坐標上都具有單位幅度。由於更新為 $ \eta \cdot \operatorname{sign}(c_t) $,任何一個坐標的步長都不會超過 $ \eta $。這種隱式裁剪類似於梯度裁剪,但它是在逐坐標層面、且在每一步都進行,從而對 語言模型 訓練中常見的異常梯度提供魯棒性。
其次,雙 $ \beta $ 結構將方向平滑與緩衝區累積解耦。當 $ \beta_1 < \beta_2 $ 時,方向 $ c_t $ 對當前梯度的反應比持久緩衝區 $ m_t $ 更加靈敏。這使 Lion 內建了一種類似 Nesterov 的前瞻特性:在時刻 $ t $ 所走的一步比單獨由動量緩衝區給出的步驟更多地反映 $ g_t $。
第三,取符號會丟棄梯度的幅度信息,而 Adam 是通過其二階矩來利用這一信息的。Lion 在丟棄幅度信息的情況下仍能良好訓練這一經驗事實表明,對許多大型神經網絡而言,方向比精確的幅度縮放更重要,前提是學習率和權重衰減得到相互兼容的調優。
內存與計算
Lion 每個參數僅存儲一個動量緩衝區 $ m_t $;而 Adam 和 AdamW 則存儲兩個(一階矩和二階矩)。對於一個擁有 $ P $ 個參數、以 bfloat16 或 32 位精度訓練的模型而言,這將優化器狀態內存減少了一半:從 $ 2P $ 個值降到 $ P $ 個值。對於參數規模達到數百億乃至數千億的模型,在加速器內存預算固定的情況下,這種節省對分布式訓練而言意義重大。
Lion 每步的計算也略微更便宜:僅需一次符號運算和兩次插值,無需對二階矩估計取平方根並進行除法。在實踐中,每步的時間主要由前向傳播和反向傳播主導,因此優化器算術更便宜帶來的實際加速幅度不大,但確實非零。
超參數調優
Lion 的超參數並不能直接從 Adam 沿用過來。作者經驗觀察到、並在獨立復現中得到佐證的推薦調整方式如下:
- 學習率:大致比同一模型下適用的 AdamW 取值小 $ 3\times $ 到 $ 10\times $。由於 Lion 的更新在每個坐標上具有單位幅度,有效步長更可預測,但在條件良好的坐標上其絕對值小於 Adam 的自適應縮放。
- 權重衰減:大致比 AdamW 取值大 $ 3\times $ 到 $ 10\times $,以在學習率較小的同時,保持有效的正則化強度 $ \eta \lambda $。
- 批量大小:在某些情形下,Lion 比 Adam 更能從更大的批量中受益,可能是因為基於符號的更新每步方差更高,在更大批量上能更乾淨地求平均。
- $ \beta_1, \beta_2 $:默認值 $ (0.9, 0.99) $ 較為穩健;原論文發現進行調優可獲得些許提升,但建議從默認值開始。
經驗性能
在原始論文中,Lion 在 ImageNet 上的圖像分類、圖文對比學習(CLIP 風格)、擴散模型訓練以及最高達數十億參數的語言模型預訓練上進行了評估。在所有這些設置下,Lion 在驗證準確率或困惑度上與 AdamW 持平或略有超越,且通常具有可比的實際訓練時間。
在較小模型上的獨立復現總體證實 Lion 具有競爭力,但並非在所有情形下都更優。其結果取決於模型架構、批量大小以及調優投入。在某些微調場景下,尤其是數據集非常小的情形,Lion 對學習率調度的選擇比 AdamW 更敏感。
變體與後續工作
自 2023 年以來,出現了多種變體和分析:
- Tiger 採用了類似的動量符號思想,但僅使用單個 $ \beta $,據報道在略微犧牲準確率的代價下進一步提升了內存效率。
- Lion-K 將逐元素符號替換為更一般的非線性,並通過理論分析將 Lion 與正則化鏡像下降的解釋聯繫起來。
- 研究者已經發展出 Lion 的連續時間和收斂性分析,將該算法視為帶符號非線性的特定常微分方程的離散時間近似,為所觀察到的魯棒性提供了部分理論依據。
這些後續工作普遍認同,Lion 的行為是由符號非線性與動量的指數移動平均之間的相互作用所支配,而非單獨由其中任一組件決定。
與相關優化器的比較
- 與 signSGD 對比:signSGD 直接對原始梯度取符號,而不使用動量緩衝區。Lion 可視為採用了特定雙 $ \beta $ 動量方案的 signSGD;該動量對 Lion 的經驗性能至關重要。
- 與 Adam / AdamW 對比:Lion 使用一個緩衝區而非兩個,並丟棄了幅度信息。Adam 通過二階矩進行逐坐標尺度自適應,而 Lion 則不會。
- 與帶動量的 SGD 對比:帶動量的 SGD 直接將原始動量向量用作方向;Lion 使用其符號,從而對每個坐標的更新幅度設了上限。
- 與 Adafactor 對比:Adafactor 通過對二階矩進行因子分解來降低優化器內存;Lion 則完全消除了二階矩。
局限性
Lion 並非 Adam 的通用替代品。已報告的局限性包括:
- 超參數敏感性:由於 Lion 的更新基於符號,在學習率、權重衰減和預熱上的細微改動可能比 Adam 更劇烈地改變其行為。
- 小批量場景:基於符號的更新每步方差更大;在批量非常小時,Adam 的逐坐標自適應縮放可起到穩定作用,而 Lion 缺乏這種機制。
- 稀疏梯度:在內在具有稀疏梯度結構的模型中(部分推薦系統、嵌入表),取符號的操作可能放大那些很少被更新的坐標上的噪聲。
- 理論保證有限:儘管已有後續分析,Lion 在凸收斂理論方面的深度仍不及 SGD 或 Adam 變體。
歷史
Lion 見於 Google 的 Chen 等人在 2023 年發表的論文《Symbolic Discovery of Optimization Algorithms》。該優化器是更廣泛的符號搜索方法學的代表性成果,作者認為這一方法可以推廣到機器學習訓練流水線的其他組件。在論文發布後的數月內,PyTorch、JAX 和 Optax 上就出現了開源實現,該算法被納入多個大規模訓練棧,用於語言模型和圖像生成的預訓練等任務。