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