缩放权重衰减并实现更好训练

中文摘要 AI 辅助

缩放定律的发现促使在越来越多的数据上训练神经网络，通常采用恒定解耦权重衰减，使网络权重在训练过程中稳步收缩。受罗宾斯 - 蒙罗条件启发，我们提出按峰值学习率的分数\(\eta/\eta_{\max}\)缩放权重衰减。证明了这种缩放权重衰减为随机梯度下降和非欧几里得谱优化器Muon保留了相应无正则化方法的渐近平稳性保证，避免了恒定解耦权重衰减引入的额外渐近偏差。通过稳态分析解释了标准权重衰减下权重范数随训练稳步收缩，而缩放权重衰减下趋于大致恒定值的原因。应用于专家混合模型训练时，带缩放权重衰减的Muon（Muon - SW）始终优于具有相同超参数的Muon，在我们最大规模下，从7200万到9.3亿参数、每个活跃参数约600个token的模型训练中，达到相同验证损失的速度快30%。如果这种趋势持续，该方法有望大幅加速前沿模型的预训练，且只需几行代码实现。

英文摘要

The discovery of scaling laws has motivated training neural networks on ever increasing quantities of data. This is typically done with a constant decoupled weight decay which causes the network weights to shrink steadily over the course of training. Taking inspiration from the Robbins--Monro conditions, we propose to scale weight decay by the fraction of the peak learning rate $η/η_{\max}$. We prove that this scaled weight decay preserves the asymptotic stationarity guarantees of the corresponding unregularized methods for both stochastic gradient descent and the non-Euclidean spectral optimizer Muon, thereby avoiding the additional asymptotic bias introduced by constant decoupled weight decay. This retains the stability benefits of weight decay without changing the asymptotic optimization target. Using a steady-state analysis, we explain why under standard weight decay the weight norm shrinks steadily as training proceeds, whereas under scaled weight decay it settles to a roughly constant value. When applied to the training of mixture-of-experts models, Muon with scaled weight decay (Muon-SW) consistently outpaces Muon with identical hyperparameters, reaching the same validation loss $\mathbf{30\%}$ faster at our largest scale across models from $72 - 930$ million parameters trained at $\sim 600$ tokens per active parameter. If this trend continues to hold, the method promises to substantially accelerate the pre-training of frontier models while requiring only a few lines of code to implement.

补充信息

备注

Comments welcome