权重衰减比muP更关键,能稳定大模型学习率迁移。
Weight Decay may matter more than muP for Learning Rate Transfer in Practice
- 用权重衰减替代muP实现学习率跨规模迁移
- 训练中后期权重衰减主导内部表示更新稳定性
- 适合关注大模型高效训练的实践者
将最优学习率从小型神经网络迁移至大型网络,可显著降低大规模训练中的超参数调优成本。为此,最大更新参数化(muP)提出了一种学习率缩放规则,旨在保持不同模型宽度下内部表示的更新动态一致。然而,muP的缩放规则依赖于强假设,尤其是层输入与权重、梯度更新之间的几何对齐性。在大规模实证研究中,我们发现这些假设仅在训练初期短暂成立,而实际中学习率迁移最需要的场景(如大语言模型训练)中,权重衰减而非muP才能真正稳定跨宽度的内部表示更新动态,从而支持学习率迁移。这表明muP的缩放主要起到隐式学习率预热作用,可被改进后的预热策略替代。这些发现从根本上挑战了现有学习率迁移认知,并可解释为何muP需配合独立权重衰减变体才能实现良好迁移。
原文摘要 · Abstract (English)
Transferring the optimal learning rate from small to large neural networks can enable efficient training at scales where hyperparameter tuning is otherwise prohibitively expensive. To this end, the Maximal Update Parameterization (muP) proposes a learning rate scaling designed to keep the update dynamics of internal representations stable across different model widths. However, the scaling rules of muP rely on strong assumptions, particularly about the geometric alignment of a layer's inputs with both its weights and gradient updates. In this large-scale empirical investigation, we show that these assumptions hold only briefly at the start of training in the practical setups where learning rate transfer is most valuable, such as LLM training. For the remainder of training it is weight decay rather than muP that correctly stabilizes the update dynamics of internal representations across widths, facilitating learning rate transfer. This suggests muP's scaling primarily acts as a form of implicit learning rate warmup, allowing us to largely replace it with modified warmup schedules. Together these findings fundamentally challenge prevailing beliefs about learning rate transfer and can explain empirical observations such as why muP requires the independent weight decay variant for good transfer.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。