Learning Rate Warmup/zh
| Article | |
|---|---|
| Topic area | Optimization |
| Prerequisites | Stochastic Gradient Descent, Learning Rate Schedule, Adam |
概述
學習率熱身是一類學習率調度,其中學習率最初設為較小的值,並在訓練的前幾百至數千次迭代中逐步增加到目標峰值,然後才進入任何標準的衰減階段。該技術於 2010 年代中期非正式地出現在深度學習實踐中,2017 年由 Goyal 等人為大批量 ImageNet 訓練給出了正式方案,並在最初的transformer論文和 BERT 論文中使用之後,成為 Transformer 模型訓練流水線中幾乎通用的組成部分。[1][2]
其動機在於訓練最初的若干次更新具有異常高的風險:參數處於隨機初始化狀態,梯度幅值很大且與真實下降方向相關性較差,而像 Adam 這樣的自適應優化器尚未積累可靠的二階矩估計。在此期間使用較小的學習率可以防止這些早期更新將迭代點推入損失曲面的不良區域,之後便可在訓練的主體階段安全地使用更高的速率。
線性熱身
最常見的形式是線性熱身。給定目標峰值學習率 $ \eta_{\max} $ 以及以迭代(有時為輪次)計的熱身長度 $ T_w $,第 $ t $ 次迭代時的速率為
$ {\displaystyle \eta_t = \eta_{\max} \cdot \min\!\left(1, \frac{t}{T_w}\right).} $
在第 $ T_w $ 步之後,調度會繼續按所規劃的衰減進行:餘弦退火、逆平方根衰減、階梯衰減或保持常數速率。典型的熱身長度介於總訓練量的 1% 到 10% 之間,語言模型預訓練中常見的取值為 500 到 4000 次迭代。
一個近親,常數熱身,在前 $ T_w $ 步使用一個固定的小速率 $ \eta_0 \ll \eta_{\max} $,然後跳躍到 $ \eta_{\max} $。Goyal 等人觀察到,當 $ \eta_0 $ 與 $ \eta_{\max} $ 之間差距較大時,常數熱身較為脆弱,因而推薦漸進式熱身(即上述線性形式)作為大批量訓練更可靠的默認方案。
Transformer 中的逆平方根熱身
最初的transformer論文將線性熱身與隨後的逆平方根衰減結合起來,產生了一種具有單個平滑可微峰值的調度:
$ {\displaystyle \eta_t = d_{\text{model}}^{-1/2} \cdot \min\!\left( t^{-1/2},\, t \cdot T_w^{-3/2} \right),} $
其中 $ d_{\text{model}} $ 是 transformer 模型維度,$ T_w $ 是熱身長度(通常為 4000 步)。最小值的兩條分支在 $ t = T_w $ 處相交,調度從線性增長的熱身階段過渡到逆平方根衰減。這一方案以微小的改動被許多後續的序列到序列系統所繼承,直到線性熱身加餘弦衰減在自回歸 語言模型預訓練中取代了它。
熱身為什麼有效
已有若干互補的解釋,但都尚未完全定論。Liu 等人給出了最具影響力的分析:在 Adam 及其變體中,自適應二階矩估計 $ v_t $ 在最初幾次更新中方差極高,這使得每參數的有效學習率 $ \eta / \sqrt{v_t + \epsilon} $ 不可靠。[3]短暫的熱身階段使 $ v_t $ 在全局步長變大之前有時間趨於穩定。同一篇論文還提出了 RAdam,這是一種 Adam 的變體,包含方差修正項,旨在使顯式熱身變得不必要;在實際中,RAdam 與顯式熱身仍然都在使用。
第二條理由涉及初始化附近的損失曲面。隨機權重產生的激活值和梯度遠離最終的訓練軌跡;早期的大步長可能放大層間的不平衡(例如transformer中注意力子層與前饋子層之間的不平衡),並引發數值不穩定性,如softmax前 logits 的爆炸。熱身可在此過渡期內限制每一步造成的損害。
第三種解釋專門針對大批量訓練,援引梯度噪聲尺度。當批量非常大時,隨機梯度接近真實梯度,因而優化器的行為幾乎是確定性的;此時較高的初始學習率可能導致迭代點持續過沖。Goyal 等人將熱身與線性縮放規則(峰值速率與批量大小成正比)結合起來,把帶動量的SGD擴展到了 ImageNet 上 8192 張圖像的批量,且未損失準確率。
變體與組合
除了線性形式和逆平方根形式外,實踐中還使用其他若干形狀。
- 餘弦熱身:速率從 $ 0 $ 沿半個餘弦曲線上升到 $ \eta_{\max} $,比線性斜坡更為平滑。該形狀隨後由餘弦衰減鏡像,形成單一的鐘形調度。
- 指數熱身:$ \eta_t = \eta_{\max} \cdot (1 - e^{-t / \tau}) $,其中 $ \tau $ 為某一時間常數。在深度學習中較少見,但在強化學習和元學習中有所出現。
- 逐層熱身:在 LAMB 及相關的逐層自適應優化器中,信任比機制產生了一種隱式熱身,其長度因層而異。通常仍會在其之上保留一個顯式的全局熱身。
熱身可以自然地與大多數衰減調度組合。在現代 LLM 預訓練中占主導地位的模式是線性熱身後接 餘弦衰減,衰減到峰值速率的某一比例(通常為 10%)。熱身也常在微調開始時再次應用,理論相同:新初始化的輸出頭以及由框架重置的優化器狀態,都會從短暫的低速率階段中獲益。
實踐考量
原則上熱身會引入兩個超參數($ T_w $ 和起始速率,常為零),但大多數方案會將起始速率固定為零或一個小常數(如 $ 10^{-7} $),僅調整 $ T_w $。經驗性指引:
- 對於 Transformer 預訓練,將 $ T_w $ 設置為總訓練步數的 1% 到 5%。小於 0.5% 的取值經常引發損失尖刺,尤其是在模型寬度較大時。
- 對於有監督微調,典型的熱身步數為 50 到 500 步;最優長度隨數據集規模而變化。
- 擴大批量大小或模型規模時增大 $ T_w $;這兩種調整都會放大熱身所要應對的早期訓練不穩定性。
- 當使用 混合精度或 FSDP / 流水線並行時,更長的熱身有助於容忍損失縮放因子和層歸一化的初始失配。
熱身與 AdamW 中的權重衰減存在相互作用:由於解耦的權重衰減與梯度無關,熱身階段的有效衰減強度並未受到影響,但因為梯度更新較小,朝零方向的相對拉力更大。因此,一些方案會按照熱身階段學習率所用的相同係數來縮放權重衰減,不過這並非通用做法。
局限性與替代方案
熱身在經驗上穩健,但在理論上解釋尚不充分。其有效性取決於尚未完全刻畫的細節:優化器類型、模型架構、運算精度以及初始化方案。RAdam 試圖將熱身階段吸收進優化器本身;Adafactor 與 Lion 同樣在某些情形下減少了對熱身的需求,但並未將其完全消除。架構層面的修正,如 Pre-Norm transformer和精心設計的初始化(Fixup、DeepNorm、ReZero),可降低實現穩定所需的熱身長度,但在大規模訓練中很少能讓省略熱身成為安全選項。
在缺乏單一封閉形式理論的情況下,多個研究團隊達成的共識是:在任何使用自適應優化器、大批量或 transformer 類架構的訓練流水線中,都應保留顯式的熱身階段,並在首次訓練某個新模型類別時,將其長度視為值得做小規模定向掃描的超參數。