Search papers, labs, and topics across Lattice.
This paper explores a generalized variant of the MUON optimizer, which utilizes an arbitrary number of generalized Newton-Schultz (NS) steps, extending the original method proposed by Jordan et al. in 2024. The authors demonstrate that MUON fails to converge to solutions in a broad class of stochastic optimization problems as the number of gradient steps increases, highlighting a critical limitation in its application. Additionally, they provide a detailed error analysis that establishes convergence rates based on mini-batch size and gradient steps, illustrated through various optimization scenarios including logistic regression.
MUON's generalized variant reveals a surprising non-convergence behavior, challenging its effectiveness in deep learning applications.
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.
