PolarGrad: A Class of Matrix-Gradient Optimizers from a Unifying Preconditioning Perspective
作者: Tim Tsz-Kit Lau, Qi Long, Weijie Su
分类: math.OC, cs.LG, stat.ML
发布日期: 2025-05-27 (更新: 2026-01-02)
💡 一句话要点
提出PolarGrad，一种基于矩阵梯度极分解的矩阵梯度优化器，提升语言模型预训练效果。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 优化算法 预处理梯度 矩阵梯度 极分解 深度学习 语言模型 各向异性
📋 核心要点
- 现有优化器如Adam在训练大型语言模型时存在训练不稳定性和需要学习率预热等问题。
- 论文提出PolarGrad，一种基于矩阵梯度极分解的预处理优化方法，旨在解决梯度各向异性问题。
- 实验表明，PolarGrad在多种矩阵优化问题和语言模型预训练任务中，性能优于Adam和Muon。
📝 摘要（中文）
深度学习模型和训练数据规模的不断增长，凸显了高效优化方法的重要性。虽然Adam和AdamW等预处理梯度方法是训练神经网络和大型语言模型的实际优化器，但像Shampoo和Muon这样利用梯度矩阵结构的结构感知预处理优化器，已经展示了更快的收敛速度。本文提出了一个统一的框架来分析“矩阵感知”预处理方法，不仅阐明了Muon和相关优化器的有效性，而且还产生了一类新的结构感知预处理方法。该框架的一个关键贡献是精确区分了将神经网络权重视为向量（解决曲率各向异性）与考虑其矩阵结构（解决梯度各向异性）的预处理策略。这种视角为语言模型预训练中的几个经验现象提供了新的见解，包括Adam的训练不稳定性、Muon的加速收敛以及Adam学习率预热的必要性。在此框架的基础上，我们引入了PolarGrad，一类基于矩阵值梯度极分解的新的预处理优化方法。作为特例，PolarGrad包括Muon，其更新由梯度的核范数缩放。我们提供了这些方法的数值实现，利用高效的数值极分解算法来增强收敛性。我们在各种矩阵优化问题和语言模型预训练任务中的广泛评估表明，PolarGrad优于Adam和Muon。
🔬 方法详解
问题定义：现有预处理梯度方法，如Adam和AdamW，在训练大型神经网络和语言模型时面临挑战。虽然它们是常用的优化器，但结构感知的预处理优化器，如Shampoo和Muon，通过利用梯度矩阵的结构，展现出更快的收敛速度。现有方法在处理梯度各向异性方面存在不足，导致训练不稳定和需要学习率预热等问题。
核心思路：论文的核心思路是利用矩阵梯度极分解来设计新的预处理优化器。通过将梯度分解为正交矩阵和正定矩阵的乘积，可以更好地处理梯度各向异性，从而提高训练的稳定性和收敛速度。这种方法能够更有效地利用梯度矩阵的结构信息，从而实现更高效的优化。
技术框架：PolarGrad的整体框架包括以下几个主要步骤：1）计算梯度矩阵；2）对梯度矩阵进行极分解，得到正交矩阵和正定矩阵；3）利用正交矩阵和正定矩阵对梯度进行预处理；4）更新模型参数。该框架的关键在于极分解步骤，它能够将梯度分解为具有更好性质的矩阵，从而提高优化的效率。
关键创新：PolarGrad的关键创新在于利用矩阵梯度极分解来设计预处理优化器。与现有方法相比，PolarGrad能够更有效地处理梯度各向异性，从而提高训练的稳定性和收敛速度。此外，PolarGrad还提供了一个统一的框架来分析“矩阵感知”预处理方法，为理解和改进现有优化器提供了新的视角。
关键设计：PolarGrad的关键设计包括：1）使用高效的数值极分解算法来计算梯度矩阵的极分解；2）使用梯度的核范数来缩放更新，以提高训练的稳定性；3）将Muon作为PolarGrad的一个特例，从而可以利用Muon的优点。此外，论文还提供了PolarGrad的数值实现，方便研究人员使用。
🖼️ 关键图片
📊 实验亮点
实验结果表明，PolarGrad在多种矩阵优化问题和语言模型预训练任务中，性能优于Adam和Muon。例如，在语言模型预训练任务中，PolarGrad能够更快地达到相同的性能水平，并且训练过程更加稳定。具体性能提升幅度取决于具体的任务和数据集，但总体而言，PolarGrad能够显著提高训练效率。
🎯 应用场景
PolarGrad可应用于各种需要训练大型神经网络和语言模型的领域，如自然语言处理、计算机视觉和语音识别。该方法能够提高训练的稳定性和收敛速度，从而降低训练成本，并提高模型性能。尤其适用于对计算资源要求高的场景，例如预训练大型语言模型。
📄 摘要（原文）
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.