提出新型优化器REG,提升大模型训练稳定性和兼容性。
REG: A Regularization Optimizer for Robust Training Dynamics
- 用行列缩放替代矩阵符号函数,平滑梯度更新
- 在大模型训练中性能超越AdamW,且微调时更稳定
- 与AdamW训练范式兼容,适合迁移学习场景
优化器对大型语言模型的高效训练至关重要。尽管AdamW是主流选择,近期出现的结构感知优化器如Muon通过作用于整个权重矩阵来正则化梯度更新,平衡各方向的更新。然而,Muon依赖矩阵符号函数导致训练不稳定,且在使用AdamW预训练的模型上微调时存在不兼容问题。为此,我们提出新优化器REG,将Muon的激进矩阵符号算子替换为行列缩放(RACS)算子。理论基于矩阵平衡,RACS以更温和的方式正则化更新步长,实现更简单、更兼容的实现。在大模型训练上的大量实验证明,REG不仅在性能和稳定性上优于AdamW,还保持了与AdamW训练范式的高度一致性,尤其在微调阶段避免了Muon的性能下降。
原文摘要 · Abstract (English)
Optimizers are crucial for the efficient training of Large Language Models (LLMs). While AdamW is the de facto standard, recent structure-aware optimizers like Muon have emerged, which regularize gradient updates by operating on entire weight matrices. The Muon optimizer balances the gradient updates along all the directions. However, Muon's reliance on the matrix sign function can lead to training instability, exhibits incompatibility when fine-tuning models pre-trained with AdamW. To address these limitations, we propose \textbf{REG}, a novel optimizer that replaces Muon's aggressive matrix sign operator with the Row-and-Column-Scaling (RACS) operator. Theoretically grounded in balancing a matrix, the RACS operator regularizes the update steps in a less drastic manner, making it simpler to implement and more compatible with established training dynamics. Through extensive empirical experiments on LLM training, we demonstrate that our REG optimizer not only achieves superior performance and stability over AdamW, but also maintains consistency with the AdamW training paradigm. This consistency is particularly evident during the fine-tuning stage, where REG optimizer avoids the performance degradation observed with Muon.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。