The Newton-Muon Optimizer

TL;DR

This paper introduces Newton-Muon, a new optimizer inspired by a surrogate quadratic model that improves training efficiency for large language models, reducing iteration steps and training time.

Contribution

It derives a closed-form update rule for a new optimizer, Newton-Muon, providing a novel interpretation of Muon as an implicit Newton-type method.

Findings

Newton-Muon reaches target validation loss 6% faster in iterations.

Reduces wall-clock training time by approximately 4%.

Offers a new theoretical perspective on Muon optimizer design.

Abstract

The Muon optimizer has received considerable attention for its strong performance in training large language models, yet the design principle behind its matrix-gradient orthogonalization remains largely elusive. In this paper, we introduce a surrogate model that not only sheds new light on the design of Muon, but more importantly leads to a new optimizer. In the same spirit as the derivation of Newton's method, the surrogate approximates the loss as a quadratic function of the perturbation to a weight matrix $W$ using only three matrices: the gradient $G$, an output-space curvature matrix $H$, and the data matrix $Z$ that stacks the layer inputs. By minimizing this surrogate in one step and adopting a certain isotropic assumption on the weights, we obtain the closed-form update rule (up to momentum and weight decay) $W \leftarrow W - \eta \cdot \mathrm{msgn}(G(ZZ^\top)^{-1})$, where $\eta$ is the learning rate and $\mathrm{msgn}(X)=UV^\top$ if $X=USV^\top$ is a compact singular value decomposition. This new optimization method, which we refer to as Newton-Muon, shows that standard Muon can be interpreted as an implicit Newton-type method that neglects the right preconditioning induced by the input second moment. Empirically, on a reproduction of the earliest publicly released Modded-NanoGPT speedrun configuration using Muon for GPT-2 pretraining, Newton-Muon reaches the target validation loss in 6\% fewer iteration steps and reduces wall-clock training time by about 4\%.