Lookahead Optimizer/zh

Article
Topic area	Optimization
Prerequisites	Stochastic Gradient Descent, Adam Optimizer

概述

Lookahead 優化器是一種用於深度學習中隨機優化的包裝算法,由 Zhang、Lucas、Hinton 和 Ba 於 2019 年提出。 Lookahead 並不取代 SGD 或 Adam 等現有優化器,而是維護兩組權重:一組快速權重由內部優化器更新若干步,另一組慢速權重周期性地向快速權重靠攏。每次同步之後,快速權重重置為慢速權重,然後過程重複。這種雙循環結構降低了內部優化器軌跡的方差,並傾向於在幾乎不增加成本的情況下改善泛化能力。

Lookahead 被廣泛用作對現有訓練流水線的即插即用增強,因為它保留了內部優化器的行為,只額外引入兩個超參數:內部步數 $ k $ 與慢速權重的插值係數 $ \alpha $。

直覺

隨機梯度方法本質上是有噪聲的:每一步都基於一個小批量計算,因此參數空間中的軌跡會圍繞下降方向振盪而非精確地沿其前進。諸如重球法或自適應學習率等標準的解決方法會對當前的梯度信號進行調整。Lookahead 則在更長的時間尺度上工作:將內部優化器的一連串更新視為一次探索階段,然後只採納所得位移的一部分。

該算法常被描述為口號"k 步前進,1 步後退"。內循環偵察損失曲面的某個區域;外循環則沿着從原慢速點到新快速點的連線保守地移動慢速權重。由於慢速權重只移動位移的一個分數 $ \alpha $,內層軌跡的瞬時振盪會被抑制,而一致的下降方向會逐漸積累。

算法

設 $ \phi $ 表示慢速權重,$ \theta $ 表示快速權重。設 $ A $ 為任意內部優化器(例如帶動量的 SGD 或 Adam),並設 $ L $ 表示損失函數。Lookahead 在一個內循環和一次外部更新之間交替。

1. 初始化 $ \phi_0 $,並令 $ \theta_{0,0} \leftarrow \phi_0 $。
2. 對每一個外部步驟 $ t = 0, 1, 2, \dots $:
   1. 對於 $ i = 1, \dots, k $:抽取一個小批量 $ d_i $,並更新 $ \theta_{t,i} \leftarrow A(\theta_{t,i-1}, \nabla L(\theta_{t,i-1}, d_i)) $。
   2. 更新慢速權重:$ \phi_{t+1} \leftarrow \phi_t + \alpha\,(\theta_{t,k} - \phi_t) $。
   3. 重置快速權重:$ \theta_{t+1,0} \leftarrow \phi_{t+1} $。

慢速權重 $ \phi $ 即訓練結束時返回並在推理時使用的權重。快速權重只是臨時狀態。

更新規則

外循環的更新可以緊湊地寫成線性插值的形式:

$ {\displaystyle \phi_{t+1} = (1 - \alpha)\,\phi_t + \alpha\,\theta_{t,k}.} $

當 $ \alpha = 1 $ 時,慢速權重在每次同步時都跳到快速權重處,Lookahead 退化為只運行內部優化器。當 $ \alpha = 0 $ 時,慢速權重永遠不動。實際取值位於 $ (0, 1] $ 中,常用的默認值為 $ \alpha = 0.5 $。

一種等價的表述是對每 $ k $ 步採樣一次的快速檢查點序列做指數移動平均。不過與標準的移動平均不同,Lookahead 還會在每次外部更新後將快速權重重置為慢速權重,從而將兩條軌跡耦合起來,而不是讓它們各自漂離。

超參數

在內部優化器自身的超參數之外,Lookahead 還額外引入了兩個超參數。

• 同步周期 $ k $:外部更新之間所執行的內部步數。典型取值為 $ 5 $、$ 10 $ 或 $ 20 $。較大的 $ k $ 可以攤銷外部步驟的(微小)開銷,但會讓內部優化器在兩次同步之間漂移更多。
• 慢速權重步長 $ \alpha $:被應用於慢速權重的內部位移的比例。典型值為 $ 0.5 $ 或 $ 0.8 $。較小的 $ \alpha $ 提供更強的方差縮減,但進度也更慢。

原文獻報告在 $ k \in \{5, 10\} $ 與 $ \alpha \in \{0.5, 0.8\} $ 的取值下,該方法在圖像分類與語言建模任務上都表現穩健,表明很少需要精細調參。

方差縮減

在二次損失近似下,Zhang 等人證明:在內部優化器動態固定的情況下,慢速權重的方差滿足

$ {\displaystyle \mathrm{Var}(\phi_{t+1}) = (1-\alpha)^2 \,\mathrm{Var}(\phi_t) + \alpha^2 \,\mathrm{Var}(\theta_{t,k}),} $

因此慢速權重在穩態時的方差是內部優化器方差的 $ \alpha / (2 - \alpha) $。當 $ \alpha = 0.5 $ 時,慢速權重的方差為內部方差的三分之一,這可能轉化為更平坦的極小值與更好的泛化能力。該分析假設內部軌跡不相關且損失局部為二次形式,因此對深度網絡來說只是參考而非保證,但與實踐中觀察到的經驗規律相吻合。

實踐考慮

內存與計算開銷都很有限。Lookahead 需要額外存儲一份參數副本(即慢速權重),並每隔 $ k $ 步進行一次按元素插值,因此對 $ k \geq 5 $ 而言,相較於內部優化器的實際牆鍾開銷通常低於百分之一。

當 Lookahead 與批歸一化或其他帶有運行時統計量的模塊結合使用時,這些層內部維護的緩衝(運行均值與方差)不是優化器的參數,因此不參與插值。多數實現將其繼續附着在模型上,由內循環的前向傳播來更新,這也是推薦的約定。

關於學習率調度,通常保留內部優化器的調度方案而不做修改。Lookahead 本身並不需要預熱,但預熱並無害處,且在 Transformer 訓練中對內部優化器仍然有用。

變體與相關方法

有若干後續工作擴展了 Lookahead 或與其密切相關。

• Stochastic Weight Averaging(SWA)在訓練結束時對權重取平均,而不是在訓練過程中進行。SWA 產生一個平均後的檢查點,而 Lookahead 在整個優化過程中維持一個滑動的插值。
• Ranger 將 Lookahead 與 Adam 的 RAdam 變體結合,在計算機視覺基準上頗受歡迎。
• Polyak 平均對過去迭代進行移動平均,可視為不帶快慢權重重置的、重複 $ \alpha $ 插值的極限形式。
• Reptile 是一種與之具有相同外部更新形式的元學習算法,只是其內循環在一個採樣的任務上訓練,而非一個採樣的小批量。

比較

與原始的 SGD 或 Adam 相比,Lookahead 用少量額外內存換取更低的軌跡方差,並且常常帶來測試準確率上的小幅改善。與 SWA 相比,Lookahead 不需要在訓練末期專門進行平均階段,且可以在任意檢查點上進行評估。與動量相比,這兩個時間尺度是正交的:動量在每一步層面上平滑梯度信號,而 Lookahead 在每 $ k $ 步層面上平滑參數軌跡,二者通常會一起使用。

局限性

Lookahead 的經驗收益較為有限,且在不同任務上並不一致。在調優良好且使用強學習率調度的基線上,改善幅度可能落入運行間波動的噪聲範圍。該方法還引入了兩個超參數,儘管默認值 $ k=5 $ 與 $ \alpha=0.5 $ 通常已經夠用。最後,方差縮減的分析假設損失局部為二次形式,且內部動態近似平穩,但在非凸的深度網絡訓練中,這兩個假設都不嚴格成立,因此理論保證不能直接適用。