🤖 machine learning

Muown: Row-Norm Control for Muon Optimization

The paper introduces Muown, a drop-in optimizer that explicitly controls row-magnitude vectors to prevent spectral norm drift in Muon, thereby improving training stability, reducing sensitivity to weight decay, and achieving superior perplexity across various model scales.

Original authors: Kai Lion, Florian Hübler, Bingcong Li, Antonio Orvieto, Niao He

Published 2026-05-12

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Kai Lion, Florian Hübler, Bingcong Li, Antonio Orvieto, Niao He

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

Imagine you are training a giant digital brain (a Large Language Model) to write stories, solve math problems, or chat with you. To do this, the brain has to learn by adjusting millions of tiny knobs (weights) inside its circuits.

For a long time, the best way to turn these knobs was a method called AdamW. Recently, a new, faster method called Muon arrived. It was like switching from a bicycle to a sports car; it learned faster and better. But there was a catch: the sports car had a tendency to speed out of control.

The Problem: The "Drifting" Engine

The paper identifies a specific problem with Muon. As the model trains, the "size" of its internal connections (specifically, the strength of the rows in its weight matrices) starts to grow uncontrollably.

Think of it like a balloon being inflated.

Muon is great at figuring out which way to blow the balloon (the direction).
However, it keeps pumping in more air than necessary, causing the balloon to get bigger and bigger until it might pop (numerical errors) or become unstable.
The standard fix was to tie a string around the balloon and pull it tight every time it got too big (Weight Decay). But the authors found this was like using a sledgehammer to crack a nut: it slowed down the learning process because it squeezed the whole balloon, even the parts that were fine.

The Diagnosis: Two Parts to the Problem

The authors, led by Kai Lion, decided to look inside the balloon to see what was actually growing. They realized the "size" of the connection is made of two distinct parts:

The Magnitude (The Volume): How big the row of numbers is overall.
The Coherence (The Shape): How the numbers inside that row are arranged relative to each other.

They discovered that only the Magnitude was drifting out of control. The Shape (Coherence) was actually behaving perfectly fine and staying stable. The balloon wasn't changing its shape; it was just getting bigger.

The Solution: Muown (The Smart Pump)

The authors created a new optimizer called Muown.

Instead of treating the whole connection as one big lump, Muown splits the job into two separate tasks:

The Direction Team (Muon): This team keeps doing what Muon does best: figuring out the perfect direction to update the weights. They leave this part unchanged.
The Volume Team (New): This is the new part. They take the "Magnitude" (the volume) and treat it as a separate, explicit variable. They give this volume a special set of rules (a specific mathematical geometry called $\ell_\infty$ ) to keep it from growing too fast.

The Analogy:
Imagine you are driving a car where the steering wheel and the gas pedal are stuck together. If you turn the wheel, the gas pedal also gets pushed, making the car speed up uncontrollably.

Muon was like that stuck car.
Muown is like uncoupling the steering wheel from the gas pedal. You can still steer perfectly (using Muon's logic), but now you have a separate, smart cruise control for the gas pedal (the magnitude) that keeps the speed exactly where it needs to be without slowing you down.

The Results

The paper tested this on models ranging from small (124 million parameters) to very large (2.7 billion parameters). Here is what they found:

Better Performance: Muown consistently produced better results (lower "perplexity," which is a measure of how confused the model is) than Muon, AdamW, and other competitors.
More Forgiving: With Muon, you had to be very careful with your settings (learning rate). If you picked the wrong one, the model would fail. Muown is much more robust; it works well across a wide range of settings, like a car with a wider "sweet spot" for driving.
No Extra Cost: Usually, adding a new control system slows a car down. But because the authors integrated this new "volume control" directly into the engine's existing workflow, it added almost no extra time to the training process.
Stability: The "balloon" stopped drifting. The spectral norm (the size of the connections) stayed stable, preventing the numerical errors that plagued the original Muon.

Summary

The paper argues that the instability in the popular Muon optimizer wasn't a flaw in its direction-finding ability, but rather a lack of control over the "volume" of its updates. By separating the "volume" from the "direction" and controlling them with different, specialized tools, Muown keeps the training fast and stable without needing heavy-handed fixes like weight decay. It's a drop-in replacement that makes the training process smoother, faster, and more reliable.

Technical Summary: Muown - Row-Norm Control for Muon Optimization

Problem Statement

Muon has emerged as a competitive optimizer for large language model (LLM) pre-training, often outperforming AdamW by leveraging the matrix structure of Transformer linear layers. However, recent observations indicate that Muon's behavior at scale is sensitive to weight decay and prone to numerical instabilities, such as bfloat16 overflows and exploding attention logits.

A critical issue identified is the consistent upward drift of the spectral norm ( $\|W_t\|_{S_\infty}$ ) of weight matrices during training when Muon is used without decoupled weight decay. This drift violates the spectral condition required for stable feature learning ( $\|W\|_{S_\infty} = \Theta(\sqrt{m/n})$ ). While decoupled weight decay can mitigate this by uniformly shrinking the weight matrix, it is a non-targeted intervention that scales the entire spectrum regardless of which specific directions drive the growth, potentially slowing convergence in early training phases. The paper seeks to identify the specific empirical driver of this spectral norm drift and propose a more targeted control mechanism.

Methodology: Muown

The authors introduce Muown, a drop-in replacement for Muon that addresses spectral norm drift through a reparameterization strategy rather than aggregate constraints.

1. Spectral Norm Decomposition

The core insight is a decomposition of the spectral norm into two factors (Proposition 1):
$\|W_t\|_{S_\infty}^2 = \|g_t\|_\infty^2 \cdot \lambda_{\max}(P_t C_t P_t)$
Where:

$g_t = \|W_t\|_{row}$ is the vector of row magnitudes (Euclidean norms of rows).
$D_t$ is the row-normalized matrix ( $W_t = \text{Diag}(g_t)D_t$ ).
$C_t = D_t D_t^\top$ is the Gram matrix of the unit-norm rows.
$P_t$ is a diagonal matrix of normalized row magnitudes.
$\lambda_{\max}(P_t C_t P_t)$ represents the row-coherence factor.

Empirical analysis on 500M models reveals that the row-magnitude factor ( $\|g_t\|_\infty$ ) is the primary driver of spectral norm drift, moving in lock-step with $\|W_t\|_{S_\infty}$ . In contrast, the row-coherence factor remains well-behaved and bounded throughout training.

2. Reparameterization and Optimization

Based on this diagnosis, Muown reparameterizes the weight matrix $W$ as:
$W(g, R) = \text{Diag}\left(\frac{g}{\|R\|_{row}}\right) R$
Here, $g \in \mathbb{R}^m$ represents the explicit row-magnitude vector, and $R \in \mathbb{R}^{m \times n}$ represents the direction component.

Direction Component ( $R$ ): Muon is applied unchanged to $R$ , treating it under the spectral norm geometry (steepest spectral descent).
Magnitude Component ( $g$ ): The row magnitudes $g$ are promoted to explicit optimizer variables. They are updated under the $\ell_\infty$ geometry, which is the natural dual norm for the operator norm of diagonal matrices.
Update Rule: The update for $g$ utilizes an Adam optimizer (acting as a smoothed SignSGD under $\ell_\infty$ geometry) to handle stochastic noise, while the update for $R$ follows the standard Muon procedure with Newton-Schulz orthogonalization.

This approach is implemented implicitly within the optimizer state, leaving the forward pass of the model architecture unchanged.

Key Contributions

Diagnosis of Spectral Drift: The paper decomposes the spectral norm to identify the maximal row magnitude ( $\|g_t\|_\infty$ ) as the empirical driver of spectral norm growth under Muon, distinguishing it from the stable row-coherence factor.
Muown Algorithm: A new optimizer that treats row magnitudes as explicit trainable variables updated under $\ell_\infty$ geometry, while preserving Muon's directional updates. This avoids the indiscriminate shrinkage of weight decay and the orthogonality priors of spectral reparameterizations.
Convergence Guarantees: The authors prove that Muown attains optimal non-convex convergence rates in both deterministic and stochastic regimes. The analysis is conducted in a mixed dual norm ( $\ell_1$ for magnitudes, Schatten-1 for directions) aligned with the update geometries. The theoretical bound suggests a reduction in the effective stochastic noise coefficient compared to Muon.
Empirical Validation: Extensive experiments on FineWeb-Edu across model sizes from 124M to 2.7B parameters demonstrate consistent improvements.

Experimental Results

Perplexity Improvements: Muown consistently achieves lower perplexity than Muon, SOAP, AdamW, and Lion across all tested model scales (124M to 2.7B) and learning rates.
Robustness to Hyperparameters:
- Learning Rate: Muown widens the plateau of near-optimal learning rates across different model widths, reducing the sensitivity to per-scale learning rate tuning.
- Weight Decay: Muown alleviates the need for weight decay tuning; it performs well without weight decay ( $\lambda=0$ ), whereas Muon requires careful tuning (optimal $\lambda=0.1$ ) to prevent divergence or suboptimal performance.
Stability: Muown prevents the systematic spectral norm drift observed in Muon. In high learning rate regimes (e.g., $4 \times 10^{-3}$ ), Muon often diverges, while Muown remains stable and near-optimal.
Gradient Properties: Analysis shows Muown maintains a higher normalized effective rank of minibatch gradients and exhibits a lower stochastic noise coefficient throughout training compared to Muon.
Computational Overhead: The method introduces negligible step-time overhead. When the optimizer is sharded across ranks (distributed setting), the step times for Muown and Muon become indistinguishable. The memory overhead is minimal (linear in the number of rows), adding only a few megabytes per layer.

Significance and Claims

The paper claims that the spectral norm drift observed in Muon is not an irreducible property of the optimizer but a consequence of uncontrolled row-magnitude growth. By isolating this driver and promoting it to an explicit optimizer variable, Muown achieves:

Improved Convergence: Better perplexity and stability without the need for aggressive weight decay.
Theoretical Soundness: Optimal convergence rates in a geometry-aware framework that explicitly accounts for the mixed nature of the updates.
Practical Efficiency: A drop-in replacement that matches Muon's runtime in distributed settings while offering superior robustness to hyperparameter choices.

The authors position Muown as a solution that bridges the gap between the theoretical stability requirements of feature learning and the practical needs of large-scale LLM training, offering a more targeted alternative to decoupled weight decay.