PolarGrad: A Class of Matrix-Gradient Optimizers from a Unifying Preconditioning Perspective

深度学习模型和数据集的规模不断扩大，突显了高效优化方法的关键重要性。尽管像 Adam 和 AdamW 这样的预条件梯度方法已成为训练神经网络和大型语言模型的事实标准优化器，但像 Shampoo 和 Muon 这样的结构感知预条件优化器（利用梯度的矩阵结构）已经展现出更快收敛的有希望的证据。 在本文中，我们提出了一种统一的框架来分析“矩阵感知”的预条件方法，这一框架不仅揭示了 Muon 及相关优化器的有效性，还引导出一类新的结构感知预条件方法。该框架的一个关键贡献在于它精确区分了两种预处理策略：一种是将神经网络权重视为向量（应对曲率各向异性），另一种则是考虑其矩阵结构（应对梯度各向异性）。这种视角为语言模型预训练中的几个经验现象提供了新的见解，包括 Adam 的训练不稳定性、Muon 的加速收敛以及 Adam 对学习率热启动的必要性。 基于这一框架，我们引入了 PolarGrad，这是一种基于矩阵值梯度极分解的新类预条件优化方法。作为其特殊实例，PolarGrad 包含了由梯度的核范数缩放更新的 Muon。我们提供了这些方法的数值实现，并通过高效的数值极分解算法增强了收敛性能。我们在各种矩阵优化问题和语言模型预训练任务上的广泛评估表明，PolarGrad 在表现上优于 Adam 和 Muon。

该论文试图解决深度学习模型优化过程中效率和收敛速度的问题，特别是针对大规模神经网络和语言模型的训练。它探讨了如何通过利用梯度的矩阵结构来加速优化过程，并解释了一些现有优化器（如Adam）在训练中的不稳定性和其他现象。这是一个持续研究的问题，但本文提出了新的视角和方法。

论文的关键思路是提出一个统一框架来分析‘矩阵感知’预调节优化方法，区分基于权重向量化的预调节策略与基于矩阵结构的策略。基于这一框架，作者引入了一种新的优化方法PolarGrad，它利用梯度的极分解，并以Muon作为其特例。这种方法可以更好地处理梯度各向异性问题，从而实现更快的收敛。

论文设计了广泛的实验来验证PolarGrad的有效性，涵盖了多种矩阵优化问题和语言模型预训练任务。实验结果表明PolarGrad在性能上优于Adam和Muon。此外，论文还提供了高效的数值算法实现，以及对Adam训练不稳定性、Muon加速收敛等现象的新见解。代码尚未提及是否开源，但值得期待。未来可进一步研究PolarGrad在更大规模模型上的表现及其与其他正则化技术的结合。

近期相关研究包括：1) Shampoo（Gupta et al., 2018），一种基于梯度矩阵结构的优化器；2) K-FAC（Martens & Grosse, 2015），一种利用二阶信息的优化方法；3) Muon（Anil et al., 2022），一种基于核范数的优化器；4) 基于自适应学习率的方法如Adam和AdamW。这些方法均试图通过不同的方式改善优化效率或稳定性，而本文通过引入PolarGrad为这一领域提供了新的解决方案。