LAMB Optimizer/zh
| Article | |
|---|---|
| Topic area | Machine Learning |
| Prerequisites | Adam Optimizer, Stochastic Gradient Descent, Gradient Descent |
概述
LAMB 優化器(Layer-wise Adaptive Moments for Batch training,分層自適應矩批量訓練優化器)是一種一階隨機優化算法,旨在使深度神經網絡能夠在非常大的小批量大小下訓練而不損失泛化能力。由 You 等人於 2019 年提出,LAMB 將 Adam 的逐參數自適應矩估計與受 LARS(Layer-wise Adaptive Rate Scaling,分層自適應率縮放)啟發的逐層信任比相結合。該算法因被用於將 BERT 的預訓練時間從單台 TPU pod 上的約三天縮短至 1024 個 TPUv3 芯片上的 76 分鐘而聲名鵲起,同時在 SQuAD 基準上達到或超過了已公布的 F1 分數。[1]
LAMB 處於深度學習中兩個長期關注問題的交匯點:如何利用現代加速器的並行能力(這有利於大批量),以及如何保留人們認為小批量 SGD 所提供的隱式正則化。其核心洞察是:在全局學習率固定的情況下,權重範數差異很大的不同層會經歷差異極大的有效更新幅度,而通過權重範數與更新範數之比的信任比對每層的更新進行重縮放,可以在大批量下恢復穩定的訓練。
動機:大批量訓練
分布式數據並行訓練使有效批量大小與工作節點數量成正比擴展。原則上,將批量大小與學習率同時加倍可以保持梯度下降每輪的軌跡,這一啟發式稱為線性縮放規則。但在實際中,該規則在某個依賴問題的臨界批量大小之外便會失效:訓練要麼發散,要麼在更差的損失上停滯,要麼泛化效果不佳。Goyal 等人證明,對 ImageNet 上的 ResNet-50,使用帶動量、預熱和精細歸一化的 SGD 可以將批量大小線性擴展到 8192,[2] 但同樣的方案對使用 Adam 訓練的 Transformer 模型卻行不通。
由 You 等人較早提出的 LARS 通過引入逐層信任比來處理 SGD 情形,使學習率自適應於每層的權重範數。[3] LARS 將 ResNet-50 的訓練批量大小推進到 32K。然而 LARS 建立在帶動量的 SGD之上,直接將其應用於使用自適應優化器的 Transformer 訓練效果不佳。LAMB 將信任比的思想擴展到基於自適應矩的方法,而後者是 Transformer 預訓練的事實標準選擇。
算法表述
設 $ \theta_t \in \mathbb{R}^d $ 表示第 $ t $ 步的參數向量,按 $ L $ 層劃分為 $ \theta_t = (\theta_t^{(1)}, \dots, \theta_t^{(L)}) $。給定小批量 $ \xi_t $ 上的隨機梯度 $ g_t = \nabla_\theta \ell(\theta_t; \xi_t) $,LAMB 執行如下更新。
首先,它維護梯度及其平方的指數移動平均,這與 Adam 完全一致:
$ {\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 \odot g_t,} $
並通過偏差修正得到 $ \hat m_t = m_t / (1 - \beta_1^t) $ 與 $ \hat v_t = v_t / (1 - \beta_2^t) $。
其次,它構造一個按坐標的 Adam 風格更新,並加入解耦的權重衰減 $ \lambda $:
$ {\displaystyle r_t = \frac{\hat m_t}{\sqrt{\hat v_t} + \epsilon} + \lambda\, \theta_t.} $
第三,也是最關鍵的一步,它通過逐層信任比對每一層 $ i $ 的更新進行重縮放:
$ {\displaystyle \theta_{t+1}^{(i)} = \theta_t^{(i)} - \eta_t \cdot \frac{\phi\!\left(\| \theta_t^{(i)} \|\right)}{\| r_t^{(i)} \|}\, r_t^{(i)},} $
其中 $ \| \cdot \| $ 表示限制在第 $ i $ 層參數上的 L2 範數,$ \eta_t $ 為全局學習率,$ \phi: \mathbb{R}_{\ge 0} \to \mathbb{R}_{>0} $ 是一個縮放函數。在標準實現中取 $ \phi(x) = x $(恆等函數),並可選擇性地裁剪到諸如 $ [\phi_{\min}, \phi_{\max}] $ 之類的範圍,以防止權重範數過小或過大的層產生極端更新。
信任比 $ \| \theta^{(i)} \| / \| r^{(i)} \| $ 確保在期望意義上,相對變化 $ \| \Delta \theta^{(i)} \| / \| \theta^{(i)} \| $ 等於 $ \eta_t $,與該層的絕對尺度無關。Adam 更新相對於當前權重顯得過大的層會被抑制,而過於保守的小更新則會被放大。
偽代碼
input: learning rate eta, betas (b1, b2), epsilon, weight decay lambda
init: theta_0, m_0 = 0, v_0 = 0
for t = 1, 2, ... do
sample mini-batch, compute g_t
m_t = b1 * m_{t-1} + (1 - b1) * g_t
v_t = b2 * v_{t-1} + (1 - b2) * g_t * g_t
m_hat = m_t / (1 - b1**t)
v_hat = v_t / (1 - b2**t)
r_t = m_hat / (sqrt(v_hat) + epsilon) + lambda * theta_{t-1}
for each layer i do
w_norm = ||theta_{t-1}^(i)||
g_norm = ||r_t^(i)||
if w_norm > 0 and g_norm > 0:
trust = phi(w_norm) / g_norm
else:
trust = 1
theta_t^(i) = theta_{t-1}^(i) - eta * trust * r_t^(i)
end for
end for
實踐考量
LAMB 在搭配包含預熱階段的學習率調度時最為有效。最初的 BERT 方案在前數千步採用線性預熱,隨後採用多項式衰減;峰值學習率明顯高於典型的 Adam 範圍,在 Transformer 預訓練中常常達到 $ 10^{-2} $ 或更高,因為信任比已經吸收了更新的絕對幅度。
權重衰減以 AdamW 的方式與梯度解耦;若將衰減折入 $ g_t $,會與二階矩歸一化耦合,從而抹去其大部分收益。偏置和歸一化參數(例如 batch norm 或 layer-norm 的縮放與平移係數)按慣例同時被排除在權重衰減和信任比重縮放之外,因為它們的範數很小,信任比在數值上可能變得不穩定。
數值穩定性還要求防範初始化或剪枝後出現零範數的層。參考實現會在權重範數或更新範數為零時退化為單位信任比。
比較
與 Adam 相比,LAMB 每步僅引入適度的常數倍開銷,主要來自逐層範數計算。在小批量下,LAMB 相對 Adam 的收益可忽略不計,此時 Adam 已能很好收斂;一旦批量大小超過數千樣本,差距才會顯現。與 LARS 相比,LAMB 繼承了 LARS 的逐層自適應,但使用自適應矩估計而非動量,這使其更適合 Transformer 類架構——在這類架構中,各層之間的梯度幅度可能相差多個數量級。
原始論文在標準的光滑性與有界方差假設下給出了收斂性分析,證明對非凸目標可在 $ O(1/\sqrt{T}) $ 速率下收斂到駐點。該分析還表明,LAMB 的更新在某種意義上是使逐層更新範數與權重範數成比例的唯一縮放方式,這為該算法提供了超越純經驗的理論依據。
局限性
LAMB 訓練出的模型並不等同於小批量 Adam 所訓練的模型。小批量與大批量之間的隱式正則化差異並未被完全消除,即使預訓練損失相同,在非常大的批量下下游任務的準確率也可能略有降低。該算法還引入了額外的超參數:信任比的裁剪範圍、按參數組的豁免設置以及 $ \phi $ 的選擇。在實踐中,源自 BERT 方案的默認值能很好地遷移到其他 Transformer 預訓練任務,但在視覺或強化學習任務上需要重新調參。
從經驗上看,LAMB 在 Transformer 預訓練和大規模監督學習中最為成功。在微調階段,由於批量通常較小且預訓練初始化附近的優化曲麵條件良好,它相對於精心調參的 Adam 的優勢較小。對於非常小的模型或表格數據,通常使用更簡單的優化器即可。