发表机构
The Chinese University of Hong Kong, Shenzhen; Vietnam Academy of Science and Technology（香港中文大学（深圳）; 越南科学技术研究院）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
本文针对2024年提出的MUON优化器，提出含任意NS步骤的广义变体，证明其在部分随机优化问题中无法收敛，建立误差分析并在多个实例中验证相关结论。
AI 中文摘要
随机梯度下降（SGD）优化方法是训练深度神经网络（DNN）的标准工具。在许多相关人工智能（AI）系统中，例如流行的大语言模型（LLM），所采用的优化方法并非标准SGD方案，而是合适的SGD加速变体。这类加速SGD变体中最受欢迎的方法之一，是Jordan等人于2024年提出的Newton-Schulz正交动量（MUON）优化器。MUON优化器利用DNN训练中权重参数的特殊矩阵结构，其原始形式在每次MUON迭代中采用5步Newton-Schulz（NS）矩阵运算。本研究提出并分析了MUON优化器的广义变体，该变体包含任意数量的广义NS步骤以及任意高阶多项式。所考虑的优化器涵盖了原始NS多项式形式的MUON，以及结合近期提出的Polar Express方法的MUON作为特殊情况。针对一类简单的随机优化问题（SOP），我们证明：对于几乎所有小批量大小，当梯度步骤数趋近于无穷大时，MUON无法收敛到该SOP的解。我们还为采用广义NS步骤的MUON建立了误差分析，该分析提供了关于梯度步骤数和小批量大小的收敛速率。我们针对多个具体示例说明了MUON的通用误差分析，包括二次随机优化问题（SOP）以及用于二分类的ℓ₂正则化逻辑回归。
英文摘要
Stochastic gradient descent (SGD) optimization methods are the standard instruments for the training of deep neural networks (DNNs). In many relevant artificial intelligence (AI) systems - such as popular large language models (LLMs)-not the standard SGD scheme is used as the optimization method but instead suitable accelerated variants of SGD are employed. One of the most popular methods of such accelerated SGD variants is the momentum orthogonalized by Newton-Schulz (MUON) optimizer proposed by Jordan et al. in 2024. The MUON optimizer exploits the special matrix structure of the weight parameters in the training of the DNNs and, in its original form, employs five Newton-Schultz (NS) matrix steps in each MUON iteration. In this work we propose and study a generalized variant of the MUON optimizer involving an arbitrary number of generalized NS steps with polynomials of possibly arbitrary high degree. The considered optimizer covers MUON with the original NS polynomial as well as MUON combined with the recently proposed Polar Express method as special cases. For a simple class of stochastic optimization problems (SOPs) we show for almost every mini-batch size that MUON fails to converge to the solution of the SOP as the number of gradient steps converges to infinity. We also establish an error analysis for MUON with the generalized NS steps that provides convergence rates in terms of the number of gradient steps and in terms of the size of the mini-batch. We illustrate our general error analysis for MUON in the case of several concrete examples including quadratic stochastic optimization problems (SOPs) as well as $\ell_2$ regularized logistic regression for binary classification.
Comments82 pages
