Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering
Provides a theoretical explanation for momentum's role in Muon, addressing a gap in understanding for practitioners using this optimizer in large-scale training.
Muon with momentum improves LLM training by acting as a spectral filter that suppresses gradient perturbations and enlarges the spectral gap, leading to more reliable orthogonal updates. Experiments confirm that applying momentum before orthogonalization outperforms alternatives.
Muon has recently demonstrated strong empirical performance in large language model training, but the theoretical role of momentum in Muon remains unclear. Existing analyses of Muon either remove momentum to study spectral updates in isolation, or retain momentum without explaining why it improves empirical performance. Our work bridges this gap by showing momentum in Muon acts as a spectral filter. Under a structured signal-plus-perturbation gradient model, we prove that momentum suppresses perturbations while preserving the dominant signal, thereby enlarging the spectral gap between them. This enlarged gap stabilizes the singular subspaces of the matrix passed to Muon's orthogonalization step, making the resulting update more reliable. We further show that applying momentum before orthogonalization achieves provably stronger alignment with the signal component of the gradient than either reversing this order or simply removing momentum. Experiments across diverse tasks, including LLM pretraining, support our theoretical analysis. More broadly, our theory offers a starting point for understanding the benefits of momentum in other matrix-based optimizers.