🤖 machine learning

Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering

This paper theoretically demonstrates that momentum in the Muon optimizer acts as a spectral filter to suppress gradient noise and enlarge the spectral gap, thereby stabilizing the orthogonalization step and improving signal alignment, a finding supported by experiments in large language model training.

Original authors: Xianliang Li, Zihan Zhang, Weiyang Liu, Han Bao

Published 2026-06-03

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xianliang Li, Zihan Zhang, Weiyang Liu, Han Bao

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: Tuning a Radio in a Storm

Imagine you are trying to tune a radio to hear a specific song (the signal) while driving through a heavy thunderstorm (the noise). The radio is your AI model, and the song is the correct direction it needs to learn to get smarter.

In modern AI training, we use a tool called Muon to help the radio tune in faster and better. Muon has two main steps:

Momentum: Smoothing out the static.
Orthogonalization: Aligning the antenna perfectly.

For a long time, researchers knew Muon worked great, but they didn't know why the "smoothing" step (momentum) was so important, or why the order of operations mattered. This paper solves that mystery.

The Core Discovery: You must clean the signal first (Denoise), and then align the antenna (Orthogonalize). If you try to align the antenna while it's still covered in static, you will point it in the wrong direction.

The Three Ways to Tune the Radio

The authors tested three different ways to handle the radio signal (the gradient) to see which one worked best:

The "Clean First" Method (Pre-polar / Muon's way):
- Action: You take the noisy signal, run it through a filter to smooth out the static (Momentum), and then you align the antenna (Orthogonalization).
- Result: Best performance. The antenna points exactly at the song.
The "Align First" Method (Post-polar):
- Action: You try to align the antenna immediately on the noisy signal, and then you try to smooth it out.
- Result: Poor performance. Because you aligned the antenna to the static noise first, smoothing it later doesn't fix the fact that the antenna is pointing in the wrong direction.
The "No Filter" Method (Polar-only):
- Action: You just align the antenna on the raw, noisy signal without any smoothing.
- Result: Worst performance. The antenna is jittery and confused by the storm.

The Analogy: Imagine trying to draw a straight line on a piece of paper while someone is shaking the table (noise).

Muon (Clean First): You wait for the table to stop shaking (smoothing), then you draw your straight line.
Post-polar (Align First): You try to draw a straight line while the table is shaking, and then you try to "smooth" the drawing afterward. The line is already crooked; smoothing the paper doesn't fix the crooked line.

Why Does "Cleaning First" Work? (The Spectral Filter)

The paper explains that momentum acts like a spectral filter. Think of the signal as a deep, steady bass note (the truth) and the noise as high-pitched, chaotic squeaks.

The Problem: If you look at the raw data, the squeaks are so loud they drown out the bass. If you try to find the "direction" of the music immediately, you might think the squeaks are the song.
The Solution: Momentum acts like a low-pass filter. It averages out the rapid, chaotic squeaks (noise) over time, but it keeps the steady, deep bass note (the signal) intact.
The Result: Once the noise is filtered out, the "direction" of the music becomes very clear. The gap between the loud bass (signal) and the quiet squeaks (noise) becomes huge. This makes it very easy for the next step (Orthogonalization) to lock onto the correct direction.

The "Spectral Gap"

The authors prove mathematically that when you use momentum before aligning, it creates a Spectral Gap.

Imagine the signal and noise are two different colors of light mixed together.
Without momentum, the colors are muddy and hard to separate.
With momentum, the noise fades away, leaving the signal color bright and distinct.
This "gap" makes the final alignment step much more reliable. It's like trying to find a needle in a haystack: if you first remove all the hay (noise), finding the needle (signal) is easy.

What the Experiments Showed

The authors didn't just do math; they tested this on real AI models (like NanoGPT and LLaMA).

They found that the "Clean First" method (Muon) consistently beat the other methods.
They visualized the data and saw that as they increased the "smoothing" (momentum), the noise disappeared, and the signal became clearer, exactly as their theory predicted.
They confirmed that if you try to align the antenna before smoothing, the AI gets confused and learns slower.

The Takeaway

The paper concludes with a simple design rule for future AI optimizers: Denoise First, Orthogonalize Later.

If you are building a system that needs to find a clear direction in a noisy world, do not try to make perfect decisions while the world is chaotic. First, smooth out the chaos to find the truth, and then make your decision. In the world of AI training, this means letting the momentum buffer clean up the noise before you try to orthogonalize (align) the update.

Technical Summary: "Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering"

Problem Statement
The Muon optimizer has recently demonstrated strong empirical performance in training large language models (LLMs) by approximately orthogonalizing a momentum matrix using a Newton–Schulz iteration. However, the theoretical role of momentum within Muon remains unclear. Existing analyses either strip momentum to study spectral updates in isolation or retain it without explaining its empirical benefits. Furthermore, theoretical literature on Muon often suggests that heavy momentum may degrade convergence rates or offer no benefit, contradicting practical observations where momentum is essential. Additionally, the specific ordering of operations—applying momentum before orthogonalization (Pre-polar) versus orthogonalizing before momentum (Post-polar)—lacks a rigorous explanation, despite Pre-polar significantly outperforming Post-polar in practice.

Methodology
The authors propose a structured signal-plus-perturbation gradient model to analyze Muon. They assume the gradient $G_t$ at iteration $t$ decomposes into a coherent, low-rank signal component $G^{\text{sig}}_t$ and a bounded-variance, mean-zero orthogonal (BVMZOS) perturbation $\Xi_t$ .

The core theoretical contribution is analyzing the momentum buffer $M_t = \beta M_{t-1} + (1-\beta)G_t$ as a spectral filter. The authors investigate three pipeline variants for the weight update $\Delta W_t$ :

Pre-polar (Standard Muon): $\Delta W_t \propto O(M_t)$ , where momentum is applied first, followed by the polar factor $O(\cdot)$ .
Post-polar (Orthogonal-SGDM): $\Delta W_t \propto \tilde{M}_t$ , where the polar factor is applied to the raw gradient first, then momentum is accumulated.
Polar-only: $\Delta W_t \propto O(G_t)$ , a baseline with no momentum.

The analysis relies on the effective window size $T = 1/(1-\beta)$ and utilizes tools from random matrix theory, including Wedin's $\sin \Theta$ theorem for singular subspace perturbation, to quantify how momentum affects the singular values and subspaces of the gradient matrix.

Key Contributions

Momentum as a Spectral Filter (Theorem 1 & Corollary 1):
The authors prove that under their gradient model, momentum acts as a low-pass filter that suppresses the perturbation (noise) while preserving the coherent signal. Specifically, as the effective window size $T$ increases, the tail singular values (dominated by noise) decay at a rate of $O(T^{-1/4})$ , while the leading singular values (signal) remain bounded away from zero. This creates an enlarged spectral gap between the signal and noise components. Consequently, the top- $r$ singular subspaces of the momentum buffer $M_t$ align more reliably with the true signal subspaces than the raw gradient's subspaces.
Non-Commutativity and Pipeline Ordering (Theorem 2):
The paper establishes that momentum and orthogonalization do not commute.
- Pre-polar Dominance: When momentum is applied first, it denoises the gradient stream, allowing the subsequent polar factor to operate on a matrix with a stabilized singular subspace structure. The authors prove that Pre-polar achieves a higher expected signal alignment than both Post-polar and Polar-only baselines.
- Failure of Alternatives: For Post-polar, orthogonalizing the noisy gradient first "erases" the amplitude gap between signal and noise directions. The polar factor replaces all singular values with ones, making it impossible to distinguish signal from noise in subsequent momentum steps. Theoretically, Post-polar and Polar-only pipelines fail to recover the signal direction effectively, especially in low Signal-to-Noise Ratio (SNR) regimes.
Quantitative Separation in Low SNR (Theorem 3):
In a rank-1 spiked Gaussian model with low SNR, the authors show that the signal alignment of the Polar-only pipeline vanishes as the matrix dimension $m \to \infty$ . In contrast, the Pre-polar pipeline continues to recover the target signal provided the momentum window $T$ is sufficiently large.

Results
The theoretical findings are validated through extensive experiments:

Spectral Gap Visualization: Using stationary probes on NanoGPT, the authors demonstrate that increasing the momentum coefficient $\beta$ widens the gap between the leading and tail singular values of the momentum buffer, consistent with Theorem 1.
Subspace Reliability: Empirical measurements of subspace alignment errors show that the deviation between the momentum buffer's singular subspaces and the true signal subspaces decreases as $T$ increases, validating Corollary 1.
Pipeline Comparison: Across synthetic models, CIFAR-10, and NanoGPT/LLaMA 350M training, the Pre-polar pipeline consistently outperforms Post-polar and Polar-only in terms of signal alignment and validation loss. Notably, Post-polar performance remains flat regardless of $\beta$ , confirming the theoretical prediction that orthogonalizing before averaging cannot recover signal directions lost to noise.
End-to-End Training: In LLM pretraining tasks (NanoGPT and LLaMA 350M), the Pre-polar pipeline (standard Muon) achieves lower validation loss compared to the other variants, supporting the claim that momentum is an essential component for signal recovery, not merely an acceleration trick.

Significance
The paper provides the first theoretical justification for why Muon requires momentum and why the order of operations (denoising before orthogonalizing) is critical. It reframes momentum in matrix-based optimizers not just as an acceleration mechanism, but as a spectral filter that stabilizes singular subspaces against noise. The guiding principle "denoise first, then orthogonalize" offers a design rule for future matrix-aware optimizers that employ nonlinear spectral steps, suggesting that separating signal from noise must precede any spectral normalization or orthogonalization. The authors note that while their theory explains the benefits of momentum in Muon, it also highlights limitations in current theoretical bounds that fail to predict these benefits, pointing to a gap between standard optimization theory and the practical behavior of spectral optimizers.