Gradient-based optimizers choose each parameter update by (approximately) minimizing a linear model of how the loss changes, subject to a constraint on the update’s size measured in some norm. Stochastic gradient descent and Adam implicitly use a Euclidean or entrywise norm on the flattened parameter vector, which does not distinguish a weight matrix’s rows and columns from one another. Muon instead treats a layer’s weight update as a matrix and constrains its size by the spectral norm (its largest singular value) rather than the Euclidean norm. The resulting steepest-descent direction for a momentum matrix with singular value decomposition is (up to scale) the orthogonalized matrix — a matrix with the same singular directions as but all singular values flattened to one — which is approximated efficiently in practice using a few steps of Newton–Schulz iteration rather than computing a full SVD.
Muon has seen rapid practical adoption in pretraining large language models, with reported wall-clock training speedups over Adam on the matrix-shaped parameters (attention and MLP weights) it is designed for. Extending the same spectral-steepest-descent idea to parameters that must additionally stay on the Stiefel manifold — matrices constrained to have orthonormal columns, which arise in low-rank adaptation and other orthogonality-constrained layers — requires solving a version of the same per-step subproblem restricted to the manifold’s tangent space, which is a strictly harder optimization problem than the unconstrained case.
Related papers
- Jordan, “Muon: An optimizer for hidden layers in neural networks” (2024) — the original description of Muon and its Newton–Schulz-based orthogonalization update.