Applying weight decay (WD) to matrix layers is standard practice in large-language-model pretraining. Prior work suggests that stochastic gradient noise induces a Brownian-like expansion of the weight matrices W, whose growth is counteracted by WD, leading to a WD-noise equilibrium with a certain weight norm ||W||. In this work, we view the equilibrium norm as a harmful artifact of the training procedure, and address it by introducing learnable multipliers to learn the optimal scale. First, we attach a learnable scalar multiplier to W and confirm that the WD-noise equilibrium norm is suboptimal: the learned scale adapts to data and improves performance. We then argue that individual row and column norms are similarly constrained, and free their scale by introducing learnable per-row and per-column multipliers. Our method can be viewed as a learnable, more expressive generalization of muP multipliers. It outperforms a well-tuned muP baseline, reduces the computational overhead of multiplier tuning, and surfaces practical questions such as forward-pass symmetries and the width-scaling of the learned multipliers. Finally, we validate learnable multipliers with both Adam and Muon optimizers, where it shows improvement in downstream evaluations matching the improvement of the switching from Adam to Muon.


翻译:在大型语言模型预训练中,对矩阵层应用权重衰减(WD)是标准做法。先前的研究表明,随机梯度噪声会引发权重矩阵 W 的类布朗运动式扩张,其增长受到 WD 的抑制,从而形成具有特定权重范数 ||W|| 的 WD-噪声平衡态。在本工作中,我们将该平衡范数视为训练过程的有害产物,并通过引入可学习乘子来学习最优尺度以解决此问题。首先,我们在 W 上附加一个可学习的标量乘子,并证实 WD-噪声平衡范数是次优的:学习到的尺度能够适应数据并提升性能。接着,我们认为单个行范数与列范数同样受到约束,因此通过引入可学习的逐行与逐列乘子来释放其尺度。我们的方法可视为一种可学习的、更具表达能力的 muP 乘子泛化形式。该方法优于经过良好调优的 muP 基线,减少了乘子调优的计算开销,并引发了诸如前向传播对称性以及学习到的乘子的宽度缩放等实际问题。最后,我们在 Adam 和 Muon 优化器上均验证了可学习乘子的有效性,其在下游评估中展现出的改进与从 Adam 切换到 Muon 所带来的改进相当。

0
下载
关闭预览

相关内容

【ICML2023】SEGA:结构熵引导的图对比学习锚视图
专知会员服务
23+阅读 · 2023年5月10日
【ICML2021】因果匹配领域泛化
专知
12+阅读 · 2021年8月12日
【NeurIPS2019】图变换网络:Graph Transformer Network
MNIST入门:贝叶斯方法
Python程序员
23+阅读 · 2017年7月3日
国家自然科学基金
46+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
6+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
4+阅读 · 2014年12月31日
VIP会员
相关资讯
【ICML2021】因果匹配领域泛化
专知
12+阅读 · 2021年8月12日
【NeurIPS2019】图变换网络:Graph Transformer Network
MNIST入门:贝叶斯方法
Python程序员
23+阅读 · 2017年7月3日
相关基金
国家自然科学基金
46+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
6+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
4+阅读 · 2014年12月31日
Top
微信扫码咨询专知VIP会员