Muon Optimizes Under Spectral Norm Constraints

TL;DR

This paper provides a theoretical foundation for the Muon optimizer by linking it to the Lion-$\\mathcal{K}$ family and showing it enforces spectral norm constraints, enhancing understanding of its regularization effects.

Contribution

It establishes a theoretical connection between Muon and Lion-$\mathcal{K}$ optimizers, revealing Muon's implicit spectral norm regularization and proposing generalizations.

Findings

Muon corresponds to Lion-$\mathcal{K}$ with nuclear norm

Muon's implicit regularization enforces spectral norm constraints

Generalizations via different convex maps expand optimizer design

Abstract

The pursuit of faster optimization algorithms remains an active and important research direction in deep learning. Recently, the Muon optimizer [JJB+24] has demonstrated promising empirical performance, but its theoretical foundation remains less understood. In this paper, we bridge this gap and provide a theoretical analysis of Muon by placing it within the Lion-$\mathcal{K}$ family of optimizers [CLLL24]. Specifically, we show that Muon corresponds to Lion-$\mathcal{K}$ when equipped with the nuclear norm, and we leverage the theoretical results of Lion-$\mathcal{K}$ to establish that Muon (with decoupled weight decay) implicitly solves an optimization problem that enforces a constraint on the spectral norm of weight matrices. This perspective not only demystifies the implicit regularization effects of Muon but also leads to natural generalizations through varying the choice of convex map $\mathcal{K}$, allowing for the exploration of a broader class of implicitly regularized and constrained optimization algorithms.