🤖 machine learning

Muon learns balanced solutions in matrix factorization without slow saddle-to-saddle dynamics

This paper demonstrates that the Muon optimizer accelerates matrix factorization by avoiding slow saddle-to-saddle dynamics, enabling stable training with high learning rates, and achieving balanced solutions through distinct conserved quantities and rapid weight alignment.

Original authors: Mark Rhee, Jamie Simon, Dhruva Karkada

Published 2026-06-30

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mark Rhee, Jamie Simon, Dhruva Karkada

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: A New Way to Teach a Machine

Imagine you are trying to teach a robot to recreate a complex picture (the "Target") by combining two simpler, transparent sheets of plastic (the "Weights"). The robot's job is to slide these sheets around until the image they create matches the Target perfectly. This is called Matrix Factorization.

Usually, we teach robots using a method called Gradient Descent. Think of Gradient Descent like a hiker trying to find the bottom of a valley in thick fog. The hiker takes small steps downhill. If the valley has a flat, bumpy plateau (a "saddle"), the hiker gets stuck, shuffling slowly from one bump to the next. This is slow, especially if the valley is long and narrow.

This paper introduces a new optimizer called Muon. Instead of a hiker shuffling in the fog, Muon is like a magnetic levitation train. It doesn't care about the bumps or the narrowness of the track. It glides smoothly and quickly to the destination.

Three Superpowers of Muon

The authors found that Muon behaves very differently from the standard method in three key ways:

1. No More "Stuck in the Middle" (The Saddle Problem)

The Old Way (Gradient Descent): Imagine the robot starts with tiny, almost invisible sheets. To learn the big picture, it has to grow the sheets one layer at a time. It learns the biggest, most obvious parts of the picture first, then gets stuck for a long time before it can learn the smaller details. It's like trying to fill a bathtub with a teaspoon; you fill the bottom, then the next inch, then the next. It takes forever.
The Muon Way: Muon doesn't wait. It learns all parts of the picture at the same speed. It's like pouring water from a hose; the whole tub fills up uniformly. The smaller details are learned just as fast as the big ones, skipping the long, boring waiting periods.

2. The "Speed Limit" Doesn't Exist

The Old Way: In the old method, if you tell the robot to move too fast (a high "learning rate"), it gets dizzy and crashes. The speed limit is set by how "sharp" or tricky the problem is. If the problem is hard, you must drive very slowly.
The Muon Way: Muon is like a car with adaptive suspension. No matter how bumpy the road is, Muon keeps its wheels on the ground. You can drive it at high speeds without it crashing. This means you can start training very fast and only slow down at the very end to fine-tune the result.

3. The "Perfect Balance" Dance

The Old Way: As the robot learns, the two plastic sheets (the weights) try to stay perfectly balanced in size. If one gets too big, the other has to shrink to compensate. This creates a specific, curved path (like a hyperbola) that the robot must follow.
The Muon Way: Muon has a different rule for balance. It keeps the square root of the size balanced. This changes the path the robot takes from a curved slide to a straight, parallel highway. It moves in a straight line toward the solution, which is much more efficient.

The "Alignment" Secret

Before the robot can start moving fast, the two plastic sheets need to be lined up correctly with the picture. The paper shows that Muon aligns these sheets spontaneously and very quickly, almost instantly, even if you start with random, messy sheets.

The authors calculated exactly how fast this happens. They found that if the sheets are slightly misaligned, Muon fixes the alignment so fast that it's almost like magic. In fact, they discovered a "cheat code" (a special learning rate schedule) where you can align the sheets perfectly in just two steps.

The "Cheat Code" (The Spiked Learning Rate)

The paper ends with a practical trick. Usually, you have to slowly increase the robot's speed. But Muon is so stable that you can do this:

Step 1: Give it a tiny nudge to line up the sheets perfectly.
Step 2: Give it a huge, massive boost to slam the sheets into the correct position and start filling in the details.
Step 3: Slowly ease off the gas to finish the job.

This allows the robot to reach a near-perfect solution in record time.

Summary

In short, this paper proves that Muon is a smarter way to train machines on certain types of problems. It avoids the slow, stumbling steps of traditional methods, handles high speeds without crashing, and aligns its internal parts almost instantly. It turns a slow, tedious hike into a high-speed, smooth ride.

Technical Summary: Muon Learns Balanced Solutions in Matrix Factorization Without Slow Saddle-to-Saddle Dynamics

Problem Statement
The paper investigates the learning dynamics of the Muon optimizer in the context of matrix factorization, defined as minimizing $L(P, Q) = \frac{1}{2}\|M^* - PQ^\top\|_F^2$ . Matrix factorization serves as a minimal model for studying nonlinear parameter dynamics and representation learning, exhibiting behaviors relevant to deep learning such as low-rank bias and representation alignment. The authors specifically compare Muon against standard Gradient Descent (GD) to understand how Muon's unique update rule alters the trajectory of optimization, particularly regarding the speed of convergence, stability, and the nature of the solutions found from small random initialization.

Methodology
The study focuses on the Muon update rule, which orthogonalizes the gradient before applying the learning rate. For a weight matrix $W$ , the update is $\Delta W = -\eta \cdot \widehat{\nabla_W L}$ , where $\widehat{G} = UV^\top$ is the polar decomposition (orthogonalization) of the gradient $G=USV^\top$ . This ensures all updates are (semi-)orthogonal matrices scaled by $\eta$ .

The authors analyze the dynamics through three lenses:

Scalar Reduction: Assuming weights are aligned, the problem decouples into independent scalar problems where the residual $r = m^* - pq$ dictates the update.
Alignment Dynamics: The authors derive theoretical rates for how singular vectors of $P$ and $Q$ align with each other (internal alignment) and with the target matrix $M^*$ (target alignment) under different parameter regimes (under-parameterized, exactly-parameterized, and over-parameterized).
Conserved Quantities: The paper identifies specific matrix quantities conserved during the flow of Muon versus GD.

Key Contributions and Results

1. Avoidance of Slow Saddle-to-Saddle Dynamics
Standard Gradient Descent from small initialization exhibits "saddle-to-saddle" dynamics. Singular values grow multiplicatively, meaning smaller modes grow slower than larger ones. The optimizer lingers near low-rank saddles, learning modes sequentially (stepwise learning).
In contrast, Muon exhibits additive dynamics. Once alignment occurs, all singular values of the weight matrices grow at a uniform rate ( $\eta$ per step), regardless of their current magnitude. Consequently, smaller modes converge before larger ones, and the optimizer learns all top modes of $M^*$ simultaneously without intermediate saddle points.

2. Stability Beyond Critical Learning Rates
Gradient Descent is constrained by a critical learning rate threshold determined by the local loss sharpness (specifically $2/\lambda_{\max}(\nabla^2 L)$ ). Exceeding this causes divergence, and the convergence rate is limited by the problem's condition number.
Muon remains stable even when the learning rate exceeds this critical threshold. Because the update magnitude is fixed by the orthogonalization step, the iterates cannot diverge. Instead, with large learning rates, the optimization proceeds rapidly to the solution manifold and oscillates there. This decouples the learning rate from the condition number, enabling rapid convergence via exponential (geometric) learning rate annealing.

3. Distinct Conserved Quantities and Balanced Solutions
Both optimizers find "balanced solutions" (where $P^\top P = Q^\top Q$ ) from vanishing initialization, but they conserve different quantities:

Gradient Flow: Conserves the matrix quantity $\Delta_2 = P^\top P - Q^\top Q$ . The flow lines in parameter space are hyperbolic.
Muon Flow: Conserves the matrix quantity $\Delta_1 = \sqrt{P^\top P} - \sqrt{Q^\top Q}$ . The flow lines are affine (parallel rays).
The authors derive that the Muon flow is generated by a Hamiltonian $H(P, Q) = \text{Tr}(\Delta_1) = \|P\|_* - \|Q\|_*$ (the difference in nuclear norms). This conservation law ensures that if weights start balanced, they remain balanced, leading to the balanced solution.

4. Alignment Dynamics and Rates
The paper derives alignment rates for the "silent alignment" phase (where singular vectors align while the loss is flat):

Under-parameterized ( $d < n$ ): Internal alignment is trivial; target alignment decays polynomially as $(t/\alpha_{\text{eff}})^{-2(1-s^{-1})}$ , where $s$ is the relative spectral gap.
Over-parameterized ( $d > n$ ): Target alignment is trivial; internal misalignment decays rapidly as $(t/\alpha_{\text{eff}})^{-4}$ .
General: Alignment occurs earlier with smaller initialization scales. The authors show that these theoretical rates predict empirical alignment in general settings.

5. Spiked Learning Rate Schedule
Leveraging the understanding of alignment dynamics, the authors construct a specific learning rate schedule that achieves near-perfect alignment in exactly two optimization steps (for $d \ge n$ ):

Step 1: Set $\eta^{(0)} = \alpha$ (initialization scale) to align $P$ with $Q$ .
Step 2: Set a large spike $\eta^{(1)} \gg \eta^{(0)}$ to simultaneously align the model with the target $M^*$ and drive the weights to the solution manifold.
Annealing: Geometrically anneal the learning rate to achieve linear convergence.
This schedule demonstrates that theoretical insights into the flow can be translated into practical optimization strategies, though the authors note this is a proof-of-concept for specific regimes.

Significance and Claims
The paper claims that Muon fundamentally alters the optimization landscape of matrix factorization by replacing multiplicative, saddle-laden dynamics with uniform, additive growth. This allows Muon to:

Learn all modes of the target matrix at the same rate, avoiding the slow, sequential learning of GD.
Operate stably with learning rates far exceeding the critical thresholds that constrain GD, freeing the learning rate from the problem's condition number.
Reach balanced solutions via a distinct conservation law ( $\sqrt{P^\top P} - \sqrt{Q^\top Q}$ ) that maintains weight balance without requiring the hyperbolic geometry of GD.

The authors emphasize that while their analysis is restricted to unmodified spectral gradient descent (ignoring momentum, weight decay, and Newton-Schulz iterations) and simple matrix factorization, the mechanistic understanding of alignment and uniform growth provides a foundation for understanding Muon's success in larger, more complex models. The work serves as a theoretical explanation for why Muon accelerates training in ill-conditioned or heterogeneous spectral settings.