提出新型矩阵梯度优化器PolarGrad,提升大模型训练效率。
PolarGrad: A Class of Matrix-Gradient Optimizers from a Unifying Preconditioning Perspective
- 从矩阵结构视角统一分析梯度预处理方法,区分曲率与梯度各向异性。
- 新方法PolarGrad在语言模型训练中收敛更快,超越Adam和Muon。
- 适合追求高效训练的大模型研究者,尤其关注优化器设计者。
深度学习模型规模不断增大,高效优化方法愈发关键。尽管Adam和AdamW已成为神经网络与大语言模型训练的默认优化器,但利用梯度矩阵结构的结构感知预处理方法(如Shampoo和Muon)已展现出更快收敛的潜力。本文提出一个统一框架,用于分析‘矩阵感知’预处理方法,不仅解释了Muon等方法的有效性,还导出一类新的结构感知优化方法。该框架的关键贡献在于精确区分将权重视为向量(应对曲率各向异性)与考虑其矩阵结构(应对梯度各向异性)的预处理策略。这一视角揭示了大语言模型预训练中的多个经验现象:Adam的训练不稳定性、Muon的加速收敛以及Adam对学习率预热的依赖。基于此框架,我们引入PolarGrad,一类基于矩阵梯度极分解的预处理优化方法。作为特例,PolarGrad使用梯度核范数缩放更新。我们实现了高效的数值算法,利用快速极分解提升收敛速度。在多种矩阵优化问题和语言模型预训练任务上的广泛评估表明,PolarGrad优于Adam和Muon。
原文摘要 · Abstract (English)
The ever-growing scale of deep learning models and training data underscores the critical importance of efficient optimization methods. While preconditioned gradient methods such as Adam and AdamW are the de facto optimizers for training neural networks and large language models, structure-aware preconditioned optimizers like Shampoo and Muon, which utilize the matrix structure of gradients, have demonstrated promising evidence of faster convergence. In this paper, we introduce a unifying framework for analyzing "matrix-aware" preconditioned methods, which not only sheds light on the effectiveness of Muon and related optimizers but also leads to a class of new structure-aware preconditioned methods. A key contribution of this framework is its precise distinction between preconditioning strategies that treat neural network weights as vectors (addressing curvature anisotropy) versus those that consider their matrix structure (addressing gradient anisotropy). This perspective provides new insights into several empirical phenomena in language model pre-training, including Adam's training instabilities, Muon's accelerated convergence, and the necessity of learning rate warmup for Adam. Building upon this framework, we introduce PolarGrad, a new class of preconditioned optimization methods based on the polar decomposition of matrix-valued gradients. As a special instance, PolarGrad includes Muon with updates scaled by the nuclear norm of the gradients. We provide numerical implementations of these methods, leveraging efficient numerical polar decomposition algorithms for enhanced convergence. Our extensive evaluations across diverse matrix optimization problems and language model pre-training tasks demonstrate that PolarGrad outperforms both Adam and Muon.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。