AdamW/zh

Article
Topic area	optimization
Prerequisites	Adam, Stochastic gradient descent, Weight decay

概述

AdamW 是一種用於訓練神經網絡的隨機優化算法，它將權重衰減與 Adam 的基於梯度的更新解耦。該算法由 Loshchilov 和 Hutter 於 2017 年提出，修正了自適應優化器中長期存在的一個實現缺陷：在標準 Adam 中，向損失函數添加 L2 懲罰項並不會產生真正的權重衰減，因為該懲罰項會被逐參數的自適應學習率重新縮放。AdamW 通過在自適應更新之後對參數施加一個獨立的、固定比率的收縮，恢復了 Hanson and Pratt 最初的表述。這一改動只需要寥寥幾行代碼，卻能持續改善泛化能力，並已成為基於 Transformer 的模型的默認優化器，包括 BERT、GPT 風格的語言模型以及 Vision Transformer。

動機：在 Adam 中 L2 正則化並非權重衰減

對於普通的 SGD，在損失函數中添加 L2 懲罰項 $ \tfrac{\lambda}{2}\|\theta\|^2 $，在數學上等價於在每一步將參數乘以 $ (1 - \eta\lambda) $，其中 $ \eta $ 是學習率。這兩種表述——L2 正則化與權重衰減——是一致的。

這種等價性在自適應方法中不成立。Adam 會按照對二階矩 $ \hat{v}_t $ 的估計來縮放每一個梯度分量，因此附加到梯度上的 L2 項 $ \lambda\theta $ 在被應用之前會被除以 $ \sqrt{\hat{v}_t}+\epsilon $。歷史梯度較大（即條件較好的方向）的參數所受到的正則化反而比梯度較小的參數更少，這與權重衰減本應起到的作用恰好相反。Loshchilov 和 Hutter 證明，這種耦合導致 Adam 在圖像分類基準上的泛化表現不及帶動量的 SGD，而將權重衰減解耦能夠彌合大部分差距。

算法

設 $ \theta_t $ 表示第 $ t $ 步的參數向量，$ g_t = \nabla_\theta f_t(\theta_{t-1}) $ 是在第 $ t $ 個小批量上損失函數的隨機梯度，$ \eta_t $ 是（可能按調度變化的）學習率。AdamW 以衰減率 $ \beta_1, \beta_2 \in [0,1) $ 維護梯度及梯度平方的指數滑動平均：

$ {\displaystyle m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t} $

$ {\displaystyle v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2} $

經過偏差修正後的估計為

$ {\displaystyle \hat{m}_t = \frac{m_t}{1-\beta_1^t}, \qquad \hat{v}_t = \frac{v_t}{1-\beta_2^t}.} $

隨後參數的更新公式為

$ {\displaystyle \theta_t = \theta_{t-1} - \eta_t \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda\,\theta_{t-1} \right).} $

關鍵之處在於 $ \lambda\,\theta_{t-1} $ 這一項被施加在自適應分母之外。相比之下，原始的 Adam-加-L2 更新會把 $ \lambda\theta_{t-1} $ 摺疊進 $ g_t $，得到 $ \theta_t = \theta_{t-1} - \eta_t (\hat{m}_t + \lambda \theta_{t-1}\cdot\text{scaling})/(\sqrt{\hat{v}_t}+\epsilon) $，此時權重衰減項與梯度共享同一個逐參數的自適應縮放因子。

在大多數實現中，默認超參數為 $ \beta_1 = 0.9 $、$ \beta_2 = 0.999 $、$ \epsilon = 10^{-8} $，其中 $ \lambda $ 在 Transformer 預訓練時通常取 $ [10^{-2}, 10^{-1}] $，在微調時通常取 $ [10^{-4}, 10^{-2}] $。

解耦權重衰減在實踐中的表現

解耦帶來了兩個實際後果。首先，最優權重衰減係數 $ \lambda $ 現在在很大程度上獨立於學習率 $ \eta $，從而簡化了超參數調優——在原始 Adam 中，改變學習率實際上也會改變正則化強度，迫使開發者必須聯合掃參。其次，AdamW 的最優 $ \lambda $ 通常比在 Adam 中可用的 L2 係數大一到兩個數量級，因為自適應縮放不再對其進行衰減。

一個常見的細節問題是：在使用學習率調度時是否應將 $ \lambda $ 也乘以 $ \eta_t $。原論文將更新寫作 $ \theta_t = \theta_{t-1} - \eta_t \hat{m}_t/(\sqrt{\hat{v}_t}+\epsilon) - \eta_t \lambda \theta_{t-1} $，即權重衰減也按調度進行縮放。某些實現（尤其是 PyTorch 的一個早期版本）則直接施加 $ \lambda\theta_{t-1} $ 而不乘以 $ \eta_t $；如今這被普遍視為一個 bug，當前的 PyTorch、JAX 與 TensorFlow 實現都遵循原論文的約定。

變體與擴展

若干優化器對 AdamW 進行了擴展或修改：

• Lion（EvoLved Sign Momentum，Chen 等人 2023）——用符號算子代替二階矩估計，保留解耦的權重衰減；顯存占用約為 AdamW 的一半。
• 帶 梯度裁剪的 AdamW ——全局範數或逐層裁剪是大語言模型預訓練的標準做法，用於抑制損失尖峰。
• LAMB（Layer-wise Adaptive Moments）——在 AdamW 之上加入逐層歸一化以適配超大批大小（32k+），曾用於刷新 BERT 預訓練耗時紀錄。
• AdaFactor ——通過對二階矩矩陣進行因式分解來節省顯存；支持解耦衰減。
• Adan 與 Sophia ——受二階方法啟發的優化器，均沿用了解耦衰減的設計。

FP16/bfloat16 訓練時代引入了額外的超參數：在實踐中，混合精度下 $ \epsilon $ 經常被提高到 $ 10^{-6} $ 或 $ 10^{-5} $，以避免 $ \sqrt{\hat{v}_t}+\epsilon $ 的下溢。

與 Adam 和 SGD 的比較

經驗上，AdamW 彌合了最初促使從業者在視覺任務中偏好帶動量 SGD 的泛化差距。在 ImageNet 上，調優良好的 AdamW 訓練的 ResNet-50 所達到的準確率與 SGD+動量相差不到 0.1–0.3 個百分點，而樸素的 Adam-加-L2 通常要落後 1–2 個百分點。對於 Transformer 而言，AdamW 幾乎是通用選擇：注意力子層與前饋網絡子層之間的梯度動態範圍很大，需要逐參數的自適應縮放來處理；解耦衰減則避免了嵌入層和層歸一化參數被欠正則化。

一條實用的經驗法則：如果你從零開始訓練模型，並且架構中含有 LayerNorm 或 RMSNorm，請使用 AdamW。如果你是在對預訓練模型進行微調，請配合較小的 $ \lambda $ 使用 AdamW，並可能採用較小的 $ \beta_2 $（例如 0.95），以避免陳舊的二階矩估計沖淡預訓練權重。

局限性

AdamW 繼承了 Adam 的顯存開銷：每個參數需要額外存儲兩個張量（$ m_t $、$ v_t $），優化器狀態相比樸素 SGD 增加到三倍。對於參數量達數十億的模型而言，這是主要開銷，因而催生了分片變體（如 ZeRO）和將優化器狀態量化的 8-bit AdamW。

解耦衰減並非萬能。它假設 $ \lambda $ 遵循一個固定的調度；循環或熱重啟調度（Loshchilov 與 Hutter 的 SGDR）與二階矩緩衝之間會產生非平凡的相互作用，目前的最佳實踐仍是先在前幾百到幾千步內對學習率進行預熱，然後再施加完整的權重衰減。在數據量較少的場景中，AdamW 對 $ \beta_2 $ 仍然敏感，因為此時二階矩估計噪聲較大；在強化學習和持續學習中，$ \beta_2 = 0.95 $ 或 $ 0.98 $ 等取值十分常見。

最後，L2 正則化與權重衰減之間的等價性，在 AdamW 中不成立，在 Adam 中同樣不成立——它們如今是兩種不同的正則化器；在報告 「weight decay」 時若不說明它指的是 AdamW 中的 $ \lambda $ 還是某種L2 損失項，就會成為再現性錯誤的常見來源。