通过动态调整权重衰减,让模型训练更快更稳,无需改架构。
Scale Weight Decay and Train Better
- 按学习率比例动态调整权重衰减,避免参数持续缩小。
- 在93000万参数模型上,验证损失下降速度提升30%。
- 只需少量代码修改,适合大规模模型预训练场景。
Scaling laws的发现推动了神经网络在更大规模数据上的训练。通常采用恒定的解耦权重衰减,导致网络权重随训练过程持续缩小。受Robbins–Monro条件启发,我们提出将权重衰减按峰值学习率比例η/η_max进行缩放。理论上证明,该方法对随机梯度下降和非欧几何谱优化器Muon均保持渐近平稳性,避免了恒定解耦权重衰减引入的额外偏差。这在不改变最终优化目标的前提下,保留了权重衰减的稳定性优势。通过稳态分析,我们解释了为何标准权重衰减使权重范数持续下降,而缩放后则趋于稳定。应用于混合专家模型训练时,使用缩放权重衰减的Muon(Muon-SW)在相同超参数下始终优于原版,在7200万至9.3亿参数模型、每活跃参数约600个标记的规模下,达到相同验证损失快30%。若趋势持续,该方法有望显著加速前沿模型的预训练,且仅需数行代码即可实现。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。