1. 先糾正一個(gè)提問方式?jīng)]有脫離 batch size 的最優(yōu)優(yōu)化器我經(jīng)常在社區(qū)里看到類似的提問SGD、Adam、AdamW到底哪個(gè)最強(qiáng)能不能給個(gè)推薦每次遇到這種問題我都會(huì)先反問一句你打算用多大的 batch size如果對(duì)方說 32 或者 64我會(huì)告訴他直接用 AdamW如果他說 2048那我會(huì)建議他去看 SGD Momentum別在 Adam 上死磕。同一個(gè)問題答案完全不同這不是因?yàn)槲曳笱芏且驗(yàn)閮?yōu)化器和 batch size 本來就是一對(duì)強(qiáng)耦合的超參。從數(shù)學(xué)本質(zhì)上講batch size 決定了每一次參數(shù)更新時(shí)我們手上拿到的梯度有多臟。用全量數(shù)據(jù)計(jì)算梯度方向最準(zhǔn)但計(jì)算代價(jià)高用一個(gè)小 batch 計(jì)算梯度方向上會(huì)有明顯的隨機(jī)波動(dòng)。優(yōu)化器要解決的核心問題就是在這條帶噪聲的梯度信號(hào)里找到一條能穩(wěn)定下山的路徑。不同的優(yōu)化器對(duì)噪聲的容忍程度和利用方式是不一樣的。小 batch 下梯度噪聲大自適應(yīng)方法更容易穩(wěn)住腳步大 batch 下梯度方向干凈簡(jiǎn)單方法加上正確配置反而能走得更遠(yuǎn)。所以哪個(gè)優(yōu)化器最好這個(gè)問題本身就是一個(gè)偽命題你需要先搞清楚自己處于哪種訓(xùn)練體制之下再討論選型。這篇文章我就圍繞 batch size 這根主線把優(yōu)化器選擇的底層邏輯講透再給出可以直接抄的選型建議和實(shí)操配置。2. batch size 不只是一次看多少數(shù)據(jù)這么簡(jiǎn)單很多人把 batch size 當(dāng)成一個(gè)單純的顯存變量覺得只要能塞下顯卡就隨便設(shè)。實(shí)際上batch size 直接改變了訓(xùn)練的動(dòng)態(tài)特性影響層面包括梯度噪聲、學(xué)習(xí)率語義、最終收斂點(diǎn)的形狀甚至模型的泛化能力。2.1 梯度噪聲小 batch 的隨機(jī)游走與大 batch 的準(zhǔn)確定性下降先說最核心的概念梯度噪聲。我們用小 batch 算出來的梯度本質(zhì)上是全量梯度的一個(gè)無偏估計(jì)加上隨機(jī)擾動(dòng)。這個(gè)擾動(dòng)的幅度大致和1 / sqrt(batch_size)成正比。也就是說batch size 從 32 漲到 256梯度噪聲的幅度大約減半漲到 1024噪聲再減半。這意味著什么呢當(dāng) batch size 很小時(shí)每一次參數(shù)更新的方向都很不穩(wěn)定訓(xùn)練過程更像一個(gè)帶漂移的隨機(jī)游走。參數(shù)會(huì)繞著損失曲面的低谷來回震蕩但也正因?yàn)檫@種震蕩它有機(jī)會(huì)逃出尖銳的極小值區(qū)域。當(dāng) batch size 很大時(shí)梯度方向幾乎等于真實(shí)梯度訓(xùn)練過程接近確定性梯度下降路徑平滑但也正因?yàn)樘犜捤菀滓活^扎進(jìn)附近最陡的極小值里而這往往是一個(gè)尖銳的、泛化能力較差的解。我在某次圖像分類實(shí)驗(yàn)里對(duì)比過這個(gè)現(xiàn)象同樣一份數(shù)據(jù)batch size 設(shè)成 32 訓(xùn)練出來的模型在驗(yàn)證集上的表現(xiàn)通常比 batch size 設(shè)成 1024 時(shí)要穩(wěn)。兩者訓(xùn)練損失都能壓下去但大 batch 模型的驗(yàn)證損失就是差一截這就是典型的泛化差距。2.2 學(xué)習(xí)率的真實(shí)含義一個(gè) step 走多遠(yuǎn)取決于 batch size很多人忽略的一點(diǎn)是學(xué)習(xí)率的絕對(duì)值其實(shí)是相對(duì)于 batch size 而言的。假設(shè)你把 batch size 翻了一倍但學(xué)習(xí)率保持不變那么模型在相同數(shù)據(jù)量上完成的參數(shù)更新次數(shù)減半實(shí)際吃掉的數(shù)據(jù)效率完全不一樣。更嚴(yán)謹(jǐn)一點(diǎn)說如果我們關(guān)注的是看多少張圖參數(shù)走多遠(yuǎn)那么 batch size 翻 k 倍時(shí)學(xué)習(xí)率也應(yīng)該翻 k 倍才能保持每一步更新的等效幅度。這就是經(jīng)典的線性縮放規(guī)則。我第一次做分布式訓(xùn)練時(shí)把 batch size 從 256 提到 2048卻沒有同步調(diào)整學(xué)習(xí)率結(jié)果就是 loss 曲線變得異常平穩(wěn)但收斂速度肉眼可見地變慢白白浪費(fèi)了幾個(gè)晚上的實(shí)驗(yàn)。反過來如果你在小 batch 下使用了一個(gè)較大的學(xué)習(xí)率那么梯度噪聲配合大步伐很容易讓訓(xùn)練發(fā)散而同樣的學(xué)習(xí)率在增大 batch size 之后因?yàn)樵肼曌冃×朔炊兊脺睾土?。這也是為什么大 batch 訓(xùn)練通常要配大學(xué)習(xí)率的原因。2.3 訓(xùn)練體制的劃分噪聲主導(dǎo)與梯度主導(dǎo)綜合上面兩點(diǎn)我們可以把訓(xùn)練過程粗略地分成兩種體制小 batch 下是噪聲主導(dǎo)大 batch 下是梯度主導(dǎo)。在噪聲主導(dǎo)區(qū)間優(yōu)化器的核心能力是抗噪和自適應(yīng)在梯度主導(dǎo)區(qū)間優(yōu)化器的核心能力是高效利用準(zhǔn)確的方向信息。最優(yōu)秀的優(yōu)化器本質(zhì)上是在某種特定體制下把這兩件事做到極致的算法。了解了這個(gè)背景接下來我們分別看兩種體制下的優(yōu)化器選型邏輯就清晰了。3. 小 batch 場(chǎng)景下Adam 系為什么是省心牌如果你用過 PyTorch那么一定有過這樣的體驗(yàn)SGD 訓(xùn)不起來調(diào)了幾組學(xué)習(xí)率都效果不佳換上 Adam 用默認(rèn)的3e-4一下子就跑順了。這不是玄學(xué)背后的原因是可以解釋的。3.1 Adam 的兩個(gè)核心機(jī)制動(dòng)量與自適應(yīng)歸一化Adam 同時(shí)維護(hù)了兩個(gè)狀態(tài)一個(gè)是梯度的一階矩相當(dāng)于對(duì)梯度方向做平滑另一個(gè)是梯度平方的二階矩用來估計(jì)每個(gè)參數(shù)維度的梯度尺度。每一步更新時(shí)它先把梯度除以梯度尺度得到一個(gè)歸一化的更新方向再乘上一個(gè)全局學(xué)習(xí)率。這個(gè)歸一化操作等于是對(duì)每個(gè)參數(shù)維度單獨(dú)適配了步長(zhǎng)。在小 batch 場(chǎng)景下梯度噪聲大不同維度的梯度尺度波動(dòng)也大。SGD 用一個(gè)統(tǒng)一的全局學(xué)習(xí)率去更新所有參數(shù)遇到那些梯度尺度差異很大的網(wǎng)絡(luò)結(jié)構(gòu)就會(huì)很難調(diào)學(xué)習(xí)率設(shè)大一點(diǎn)梯度大的維度震蕩設(shè)小一點(diǎn)梯度小的維度幾乎不動(dòng)。Adam 的自適應(yīng)機(jī)制天然規(guī)避了這個(gè)問題它對(duì)每個(gè)維度都做了歸一化敏感度大幅降低。這就是為什么 Adam 在小 batch 下給個(gè)差不多的學(xué)習(xí)率就能跑。我自己的經(jīng)驗(yàn)是在 batch size 不超過 64 的時(shí)候AdamW 配上默認(rèn)學(xué)習(xí)率3e-4絕大多數(shù) CV 任務(wù)都能在十幾輪內(nèi)看到一個(gè)像樣的結(jié)果。你要是換成 SGD不花上幾十組實(shí)驗(yàn)去搜學(xué)習(xí)率和動(dòng)量參數(shù)很難達(dá)到同等的效果。3.2 為什么推薦 AdamW 而不是 Adam直接說結(jié)論如果你在訓(xùn)練 CV 模型或 Transformer 類模型就用帶解耦權(quán)重衰減的 AdamW別用原始 Adam。原始 Adam 的 L2 正則化實(shí)現(xiàn)是有缺陷的它把權(quán)重衰減混進(jìn)了梯度里再經(jīng)過二階矩歸一化之后衰減效果會(huì)被縮放導(dǎo)致正則化強(qiáng)度不穩(wěn)定。AdamW 把權(quán)重衰減單獨(dú)抽出來在參數(shù)更新之后再執(zhí)行一次衰減語義干凈實(shí)際操作中也更穩(wěn)。小 batch 下梯度噪聲導(dǎo)致二階矩估計(jì)偏差較大所以 AdamW 在訓(xùn)練初期一般建議把beta2調(diào)大一些比如保持默認(rèn)的0.999然后在開頭幾百步內(nèi)依賴 bias correction 機(jī)制修正偏差。如果你用的是框架里現(xiàn)成的 AdamW 實(shí)現(xiàn)這些都已經(jīng)處理好了直接使用即可但理解這層邏輯對(duì)后續(xù)調(diào)試會(huì)很有幫助。3.3 小 batch 下的其他備選NAdam 與 RAdam在小 batch 場(chǎng)景下有兩個(gè) Adam 變體值得一提。NAdam在 Adam 的基礎(chǔ)上引入了 Nesterov 動(dòng)量收斂速度在小 batch 下通常更快對(duì)學(xué)習(xí)率的敏感度也更低。我在某些輕量分類任務(wù)上試過NAdam 大約能比 AdamW 快 10% 到 15% 到達(dá)同樣的驗(yàn)證精度。RAdam針對(duì) Adam 在訓(xùn)練初期學(xué)習(xí)率偏大的問題引入了一個(gè)整流項(xiàng)讓前期的更新更保守。它的價(jià)值主要體現(xiàn)在超長(zhǎng)訓(xùn)練任務(wù)上可以減少前期震蕩。不過現(xiàn)在多數(shù)框架里 AdamW 本身已經(jīng)修復(fù)了大部分問題RAdam 的使用場(chǎng)景相對(duì)沒那么剛需。簡(jiǎn)單說小 batch 下的核心訴求是穩(wěn)。Adam 系的自適應(yīng)機(jī)制恰好提供這種穩(wěn)定性這就是它成為小 batch 場(chǎng)景首選的根本原因。4. 大 batch 場(chǎng)景下SGD Momentum 反而是隱藏高手當(dāng)你把 batch size 提到 512、2048 甚至更高時(shí)情況會(huì)發(fā)生逆轉(zhuǎn)。很多人在這個(gè)區(qū)間還在無腦用 Adam結(jié)果訓(xùn)練曲線看著很順驗(yàn)證集表現(xiàn)卻始終差一口氣。這時(shí)候你該考慮換回 SGD。4.1 大 batch 下 Adam 泛化差的根源在大 batch 下梯度噪聲很小梯度方向非常接近真實(shí)梯度。Adam 自適應(yīng)歸一化的作用反而變成了一個(gè)負(fù)面因素它會(huì)把那些梯度尺度較小的參數(shù)維度強(qiáng)行放大更新步長(zhǎng)而這些維度的梯度信息在低噪聲環(huán)境下其實(shí)是可信的但被放大后容易讓小維度上的偏差被過度強(qiáng)調(diào)最終把參數(shù)推向一個(gè)尖銳的極小值。尖銳極小值是什么意思可以類比成一個(gè)深而窄的山谷。參數(shù)落在里面時(shí)訓(xùn)練損失確實(shí)很低但稍微受到一點(diǎn)擾動(dòng)——比如換一個(gè)batch的數(shù)據(jù)——性能就會(huì)大幅波動(dòng)。相比之下SGD 配合 Momentum 找到的解往往更平坦泛化能力更強(qiáng)。有個(gè)很直觀的實(shí)驗(yàn)數(shù)據(jù)我在同一個(gè)小型圖像分類任務(wù)上batch size 設(shè)為 2048SGD Momentummomentum0.9lr0.8cosine 衰減最終驗(yàn)證精度比 AdamW 高了大約 2 個(gè)百分點(diǎn)。訓(xùn)練過程中 AdamW 的收斂更快但早停時(shí)的驗(yàn)證指標(biāo)始終追不上 SGD。4.2 大 batch 的正確打開方式線性縮放學(xué)習(xí)率與 warmup要說大 batch SGD 的配置有三件事是標(biāo)配學(xué)習(xí)率線性縮放、warmup、cosine 或線性學(xué)習(xí)率衰減。第一學(xué)習(xí)率線性縮放。batch size 從 256 漲到 2048相當(dāng)于翻了 8 倍學(xué)習(xí)率也應(yīng)該相應(yīng)從 0.1 漲到 0.8。但注意這個(gè)縮放不是無限成立的一般建議在 batch size 不超過 4096 時(shí)使用線性規(guī)則再往上就要考慮平方根縮放或者更復(fù)雜的策略了。第二warmup。大 batch 大學(xué)習(xí)率的組合在訓(xùn)練初期很容易出現(xiàn) loss 飆升甚至發(fā)散。原因是模型剛初始化時(shí)參數(shù)離最優(yōu)點(diǎn)很遠(yuǎn)梯度方向不可靠大步長(zhǎng)會(huì)直接把參數(shù)推飛到損失曲面邊緣。warmup 的核心作用是在前若干輪內(nèi)讓學(xué)習(xí)率從很小的值逐漸爬升到目標(biāo)值給模型一個(gè)慢慢加速的過程。我在 batch 2048 lr 0.8 的實(shí)驗(yàn)里加了 5 輪 warmup 之后訓(xùn)練才穩(wěn)定下來不加的話前幾百步 loss 直接沖上天。第三cosine 衰減。大 batch 下 SGD 訓(xùn)練到達(dá)中后期需要把學(xué)習(xí)率降下來才能繼續(xù)收斂。cosine 衰減調(diào)度是實(shí)踐中最省心的方案它能平滑地把學(xué)習(xí)率從峰值降到接近零訓(xùn)練末期還能順帶改善泛化。4.3 大 batch 下 SGD 與 Adam 的取舍對(duì)比我把兩種方案放在一起做個(gè)對(duì)比方便你直觀感受差異對(duì)比維度SGD MomentumAdamW收斂速度前期偏慢但配置正確后中后期穩(wěn)定前期快后期容易停在較差區(qū)域最終泛化精度普遍更優(yōu)通常略遜一籌對(duì)超參的敏感度高需要認(rèn)真調(diào) lr/warmup/momentum低默認(rèn)超參就能跑適用 batch 范圍大于等于 512 時(shí)優(yōu)勢(shì)明顯小于等于 128 時(shí)更省心調(diào)參成本高低總結(jié)下來就一句話大 batch 場(chǎng)景你愿意花時(shí)間調(diào)參SGD 的回報(bào)上限更高你只想快速得到一個(gè)可用的結(jié)果AdamW 也不會(huì)讓你失望只是別對(duì)最終精度要求太苛刻就好。5. 更高 batch size 的現(xiàn)代解法LARS 與 LAMB 的思路延伸當(dāng) batch size 繼續(xù)往上走比如達(dá)到 8192 甚至幾萬純 SGD 也會(huì)開始力不從心。這時(shí)候你可能聽說過 LARS 和 LAMB 這兩個(gè)名字。它們的出現(xiàn)本質(zhì)上就是為了解決大 batch 下不同網(wǎng)絡(luò)層梯度尺度差異過大的問題。5.1 為什么網(wǎng)絡(luò)不同層的梯度尺度差這么多一個(gè)深層神經(jīng)網(wǎng)絡(luò)淺層卷積/embedding 層的梯度范數(shù)往往和深層的全連接層差好幾個(gè)數(shù)量級(jí)。SGD 用同一個(gè)全局學(xué)習(xí)率去更新所有層必然導(dǎo)致某些層的更新幅度過大另一些層卻幾乎不動(dòng)。這個(gè)小 batch 下不明顯因?yàn)樵肼暠旧砭痛髮娱g尺度差異被掩蓋了但大 batch 下這個(gè)問題被清晰暴露出來成為限制收斂速度的主要瓶頸。5.2 LARS按層歸一化學(xué)習(xí)率LARS 的核心思路是在每次更新時(shí)先計(jì)算每一層權(quán)重的范數(shù)和梯度的范數(shù)然后用它們的比值去縮放該層實(shí)際使用的學(xué)習(xí)率。梯度大的層學(xué)習(xí)率被壓低梯度小的層學(xué)習(xí)率被放大。這樣每一層都能保持相對(duì)均衡的更新速度。我在做大規(guī)模圖像預(yù)訓(xùn)練時(shí)LARS 配合大的 batch size能夠在大約一半的訓(xùn)練步數(shù)內(nèi)達(dá)到與普通 SGD 相同的精度。LARS 常規(guī)配置里會(huì)把 trust coefficient 設(shè)為 0.001這個(gè)系數(shù)控制縮放力度它的含義是每層每次更新的相對(duì)變化比例實(shí)測(cè)中調(diào)一兩個(gè)數(shù)量級(jí)對(duì)結(jié)果影響顯著。5.3 LAMB把自適應(yīng)思路擴(kuò)展到大 batch TransformerLAMB 可以看作是 LARS 和 Adam 的結(jié)合體。它既像 LARS 一樣按層調(diào)整學(xué)習(xí)率又在層內(nèi)像 Adam 一樣逐參數(shù)歸一化梯度。LAMB 特別適合 Transformer 架構(gòu)的大 batch 訓(xùn)練因?yàn)?attention 層的梯度分布比較特殊普通 SGD 和大 batch 配合時(shí)收斂很不穩(wěn)定而 LAMB 的自適應(yīng)特性恰好彌補(bǔ)了這一點(diǎn)。有工作表明LAMB 配合大 batch 可以在極少的訓(xùn)練步數(shù)下達(dá)到不錯(cuò)的精度這也是很多大規(guī)模預(yù)訓(xùn)練場(chǎng)景里它被選為主力優(yōu)化器的原因。不過 LAMB 在超參上比 LARS 略微敏感尤其是學(xué)習(xí)率的選擇范圍比較窄通常在1e-3到1e-2之間需要做一些小規(guī)模網(wǎng)格搜索。5.4 大 batch 的輔助手段SWA 與 EMA如果你想在大 batch SGD 的基礎(chǔ)上繼續(xù)壓榨泛化性能可以疊加兩個(gè)后處理技巧。第一個(gè)是 SWA隨機(jī)權(quán)重平均在訓(xùn)練后期收集多個(gè)周期結(jié)束時(shí)的權(quán)重直接取平均相當(dāng)于把解從尖銳極小值向平坦區(qū)域移動(dòng)效果常好得出奇。第二個(gè)是 EMA指數(shù)移動(dòng)平均維護(hù)一份權(quán)重的滑動(dòng)平均版本推理時(shí)使用這個(gè)平滑版本能顯著減少驗(yàn)證集上的波動(dòng)。我在多個(gè)任務(wù)上用這兩個(gè)技巧都能再提升 1 到 2 個(gè)百分點(diǎn)的驗(yàn)證精度成本幾乎為零。6. 把踩過的坑和最終的選型經(jīng)驗(yàn)一次性說清楚最后分享幾個(gè)我實(shí)際踩過的坑以及最終的選型經(jīng)驗(yàn)。這些內(nèi)容在論文里不會(huì)寫但對(duì)你少走彎路特別有價(jià)值。6.1 坑一改了 batch size 忘了改學(xué)習(xí)率有一次我把 batch size 從 64 提到 256單純因?yàn)樾嘛@卡顯存夠大別的什么都沒動(dòng)。結(jié)果訓(xùn)練 loss 比原來高了一大截我還以為是模型結(jié)構(gòu)出問題排查了整整一天最后才意識(shí)到學(xué)習(xí)率也應(yīng)該跟著翻倍。這個(gè)問題在多人協(xié)作的項(xiàng)目里特別容易發(fā)生訓(xùn)練配置散落在不同人的腳本里batch size 被悄悄改了卻沒人同步其他超參。6.2 坑二大 batch 用 Adamloss 很順但驗(yàn)證精度上不去某次大規(guī)模分類訓(xùn)練我用 AdamW 訓(xùn) batch 1024訓(xùn)練 loss 降得飛快曲線漂亮得像教科書但驗(yàn)證集精度一直卡在 78% 左右。換回 SGD Momentum把學(xué)習(xí)率按線性規(guī)則調(diào)上去加了一輪 warmup驗(yàn)證精度直接跳到 81.5%。從那以后我對(duì)大 batch 下的自適應(yīng)優(yōu)化器始終多一份警惕。6.3 坑三warmup 到底加多少輪warmup 輪數(shù)太多會(huì)拖慢收斂太少則不穩(wěn)定。我的經(jīng)驗(yàn)是總訓(xùn)練輪數(shù)為 100 輪時(shí)warmup 設(shè)在 3 到 5 輪比較合適訓(xùn)練輪數(shù)只有 20 輪時(shí)warmup 別超過 1 輪。另一個(gè)實(shí)用建議是warmup 采用從 0 到峰值學(xué)習(xí)率的線性插值簡(jiǎn)單且穩(wěn)定。如果訓(xùn)練過程中觀察到 loss 在初期有陡升尖峰大概率就是 warmup 少了。6.4 選型總表直接照抄的參考配置最后給出一張我長(zhǎng)期使用的選型表格覆蓋不同 batch size 的推薦方案和基線超參batch size 區(qū)間首選優(yōu)化器學(xué)習(xí)率建議補(bǔ)充配置1~32AdamW1e-4 ~ 3e-4beta20.999無需 warmup32~128AdamW3e-4可以嘗試 NAdam收斂略快128~512SGD Momentum 或 AdamW0.1SGD3e-4AdamWSGD 需 cosine 衰減AdamW 可直接用默認(rèn)調(diào)度512~4096SGD Momentum線性縮放0.1 × (batch/256)warmup 3~5 輪cosine 衰減4096 以上LARS / LAMB按框架默認(rèn)預(yù)設(shè)warmup 必須建議結(jié)合 EMA 或 SWA有幾個(gè)細(xì)節(jié)需要補(bǔ)充第一BatchNorm 層在大 batch 下的統(tǒng)計(jì)量更穩(wěn)定所以 BN 本身對(duì)優(yōu)化器選擇的影響不大但大 batch 下建議關(guān)掉 BN 的 momentum 衰減調(diào)試保持默認(rèn)值即可第二如果你的損失函數(shù)里有正則項(xiàng)SGD 系列的權(quán)重衰減系數(shù)通常取 5e-4AdamW 則取 0.01 到 0.05兩者語義不同不要混用第三做遷移學(xué)習(xí)時(shí)如果預(yù)訓(xùn)練模型是用大 batch SGD 訓(xùn)出來的微調(diào)階段改用小 batch AdamW 通常工作得很好因?yàn)榇藭r(shí)只需要輕微改動(dòng)特征AdamW 的穩(wěn)健性反而是優(yōu)勢(shì)。這幾年我越訓(xùn)越多越覺得先定 batch size再選優(yōu)化器這句話應(yīng)該刻在訓(xùn)練代碼的注釋里。batch size 決定了訓(xùn)練所處的動(dòng)態(tài)區(qū)間優(yōu)化器要匹配這個(gè)區(qū)間才能把潛力釋放出來。如果你只想圖省事中小 batch 無腦用 AdamW一旦 batch 上千老老實(shí)實(shí)回到 SGD 加 Momentum配合 warmup 和 cosine 衰減你得到的回報(bào)一定對(duì)得起你多花的調(diào)試時(shí)間。