[Tech Insight] MUON+: Simple Normalization is the Secret Sauce for Faster Pre-training

TL;DR

The community recently pivoted towards Muon, an optimizer that orthogonalizes gradients to prevent rank collapse. However, a new challenger, MUON+, proves that adding a single normalization step after orthogonalization can yield massive gains. Across GPT and LLaMA models up to 1B parameters, MUON+ consistently lowers perplexity and makes training significantly more robust to learning rate swings.

The Problem: Orthogonalization is Not Enough

Standard optimizers like AdamW rely on vector-wise scaling. Muon changed the game by treating layers as matrices and enforcing orthogonality, effectively keeping the "rank" of the gradient high.

However, the authors of MUON+ noticed a gap: while the gradients were orthogonal, their structural magnitudes were still suboptimal. While other recent variants tried to fix this with "second-moment scaling" (NorMuon) or "manifold optimization" (Mano), these additions were computationally heavy. The researchers at UCSB asked a critical question: Is the complexity necessary, or is it just the normalization that matters?

Methodology: The Power of One Extra Step

The core of MUON+ is deceptively simple. After computing the momentum $M_t$ and orthogonalizing it into $O_t$ using Newton-Schulz iterations, MUON+ applies a row-wise or column-wise $L_2$ normalization:

${\mathbf{O}}_t={\mathrm{Norm}}_{(d)}(\mathrm{Ortho}({\mathbf{M}}_t))$

Model Architecture and Update Flow

The authors tested four directions:

1. Column-wise: Normalizing the scale of each output feature.
2. Row-wise: Normalizing the contribution of each input feature.
3. Col-Row / Row-Col: Bi-directional normalization.

Interestingly, they found that Row-wise and Col-Row normalization consistently provided the best results, suggesting that balancing the input-side influence is vital for stable transformer training.

Experimental Results: SOTA Performance

MUON+ was tested in "compute-optimal" settings (where token count scales with parameters) and "overtraining" settings (up to 200 tokens per parameter).

Performance Gains across Architectures

Stability and Robustness

One of the most striking findings is the learning rate sensitivity. In the figure below, MUON+ (colored lines) stays at a lower perplexity across a much wider range of learning rates compared to the original Muon (None).

Deep Insight: Decoupling Complexity

The paper includes a brilliant ablation study against NorMuon. By setting ${\beta }_2=0$ (removing the second-moment/Adam-like scaling), they proved that the performance didn't drop. This suggests that the "Adaptive" part of "Adaptive Optimizers" might be less important than simple "Structural Normalization" when gradients are already orthogonal.

Conclusion & Future Outlook

MUON+ demonstrates that in the world of LLM optimization, Less is More. By adding a simple normalization step that costs almost zero extra FLOPs, we can achieve better convergence than much more complex manifold-based methods.

Limitations: The study primarily focuses on models up to 1.3B parameters. Scaling this to 70B+ models will be the next frontier to see if these structural benefits hold at the "Extreme Scale."

Paper cited: Zhang et al. (2025). MUON+: Towards Better Muon via One Additional Normalization Step. UCSB.