隨機梯度下降的頻率偏差與Adam如何解決它
現代語言模型因token分佈極不均勻面臨最佳化挑戰:常見token的引數頻繁接收梯度更新,而罕見token的引數可能數百步無有效訊號。標準SGD以相同學習率更新所有引數,導致罕見token權重難以學習。Adam透過方差歸一化自動為罕見token提供更大的有效學習率,使其快速收斂。在控制實驗中,使用頻率跨四個數量級的六個token,SGD對最罕見token的權重僅學到0.15(真實值為1.0),而Adam使其接近1.0。
現代語言模型訓練資料中的token分佈極不均勻,常見詞如“the”幾乎出現在每個句子中,而罕見但重要的詞如“thalweg”則極少出現。這種不平衡帶來了一個隱藏的最佳化挑戰:與常見token相關的引數持續接收梯度更新,而與罕見token相關的引數可能數百甚至數千步都收不到有意義的訊號。在標準隨機梯度下降(SGD)中,所有引數使用相同的學習率,導致常見token的權重快速收斂,而罕見token的權重往往停留在隨機初始化附近。
Adam最佳化器透過自適應學習率解決了這一問題。雖然Adam常被描述為帶動量的SGD,但其最關鍵的特性是方差歸一化:它單獨跟蹤每個引數的歷史梯度統計,並根據可靠梯度資訊的觀測頻率自動調整更新大小。更新頻率低的引數會獲得成比例更大的有效學習率,使罕見特徵比在普通SGD下學習得更快。
為了具體展示這一行為,研究者設計了一個受控的NumPy實驗,使用包含六個token的詞彙表,這些token的頻率跨越四個數量級——從幾乎每個批次都出現的token到僅0.1%批次出現的token。他們對同一線性模型分別使用SGD和Adam訓練,保持所有目標權重相同(真實權重均為1.0),透過比較最終引數值、非零梯度次數以及Adam的有效學習率,直接觀察自適應最佳化如何補償頻率不平衡。
實驗設定了高度簡化的訓練環境,僅聚焦於token頻率因素。詞彙表包含從極常見詞“the”到極罕見詞“thalweg”的六個token,出現機率跨越四個數量級。每個token被賦予相同的重要性——真實權重均為1.0,從而去除語義複雜性,專注於引數接收梯度更新的頻率。每個訓練樣本表示為稀疏二元向量,表示批次中出現的token,目標值為啟用token權重之和加少量噪聲。梯度僅針對批次中的token計算,因此罕見token自然比常見token接收的更新少得多。
訓練3000步後,結果清晰顯示出SGD與Adam的差異。對於常見token,兩種最佳化器都成功學得正確權重,因為這些token幾乎出現在每個批次中。但對於罕見token,SGD表現糟糕:“xenobiotic”僅在約15%的步驟中接收到梯度,權重停留在0.53左右(真實值為1.0),而最罕見的“thalweg”僅3.4%的步驟有梯度,權重僅約0.15。相比之下,Adam儘管接收到同樣稀疏的梯度訊號,卻使兩個罕見token的權重都接近正確值。
Adam之所以成功,是因為它的方差歸一化機制。Adam維護每個引數的動量估計(平均值)和方差估計(平方梯度的平均值),並在更新前進行偏差校正。引數v_hat跟蹤梯度平方的指數移動平均,即歷史梯度幅度的近似值。在罕見token的案例中,大多數時間梯度為零,因此v_hat很小,導致有效學習率(lr/√(v_hat))非常大,從而在接收梯度時給予大幅更新。隨著訓練推進,v_hat逐漸增加,有效學習率下降,避免過沖。實驗測量了每個token在最後100步的平均有效學習率,發現引數頻率與有效學習率之間呈反向關係:頻率越低的token,Adam分配的有效學習率越高,近乎完美地補償了梯度稀疏性。
這項實驗直觀地證明了Adam為何能有效處理長尾token分佈。它並非僅僅增加所有引數的學習率,而是智慧地將大更新分配給最需要它們的引數。這一特性對現代語言模型至關重要,因為這類模型必須從頻繁出現的語法功能詞到極少出現的專業術語中學習所有型別的token。