AdamW訓練動態下威布爾權重尺度參數的演化
該研究基於雙參數威布爾框架分析Transformer權重分佈，揭示了AdamW訓練中權重尺度參數λ的增長、過沖和鬆弛行為。通過AdamW更新的三力分解（對齊力、注入力、衰減力），發現對齊力主導上升階段，貢獻88-94%的力預算。飽和時對齊力與衰減力接近平衡。研究還提出了樣條位移方法，從稀疏檢查點恢復對齊力，準確率達92-94%。峯值λ隨訓練數據相干性變化，表明數據依賴的權重尺度增長成分。
近日，一篇題為《Weibull Weight-Scale Parameter Evolution under AdamW Training Dynamics》的論文在arXiv上發表，深入研究了Transformer在AdamW優化器訓練下權重分佈演化的動力學機制。該研究基於雙參數威布爾框架，揭示了權重尺度參數λ在訓練過程中的增長、過沖和鬆弛現象，併為之提供了理論解釋。
論文作者、來自Tiexin Ding的研究團隊從AdamW更新規則出發，推導出平方權重範數的主導階三力分解模型。這一模型包含三個關鍵分量：對齊力，衡量權重與自適應更新方向之間的相關性；注入力，源自自適應步長的大小；以及衰減力，來自解耦權重衰減。通過在自己的Pythia-70M模型上進行的實驗，研究者發現對齊力是λ上升階段的主導因素，在四個隨機種子下貢獻了88%至94%的絕對力預算，並且對超權重的移除保持穩健。
隨着訓練接近飽和，對齊力與衰減力逐漸達到平衡，這解釋了λ從增長到鬆弛的轉變。這些力的動態直接控制了λ(t)所依賴的平方範數分量。剩餘的均方根到威布爾重建偏移是可測量的，並分解為橋接和積分分量，在密集採樣區域約佔5-6%。
為了將分析擴展到無法獲取優化器矩的實際模型，研究者提出了一種樣條位移方法，該方法能夠從稀疏檢查點恢復對齊力，準確率約為92-94%，幾乎是樸素兩點基線的兩倍。此外，實驗中觀察到λ(t)的峯值隨訓練數據相干性而變化，這表明了權重尺度增長存在數據依賴的成分，作者計劃在後續控制實驗中進行進一步研究。
該研究的代碼和數據已在GitHub上開源，為機器學習社區提供了診斷和理解Transformer訓練動力學的有力工具。