Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering
- URL: http://arxiv.org/abs/2606.03899v2
- Date: Wed, 03 Jun 2026 02:06:07 GMT
- Title: Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering
- Abstract summary: We show that momentum in Muon acts as a spectral filter.<n>Applying momentum before perturbation provably stronger alignment with the signal component of gradient than either reversing this order or simply removing momentum.
- Score: 32.15111956211244
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: 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.
Related papers
- Spectral Saliency for Machine Unlearning [30.769424654889235]
Machine unlearning (MU) aims to remove the influence of specific training data while preserving model utility.<n>We propose Spectral Saliency Unlearning (SSU) to thresholds weak singular components and updates only those directions supported by a confident unlearning signal.
arXiv Detail & Related papers (2026-08-16T05:44:53Z) - Reassessing Muon for Matrix Factorization [9.691906492780353]
Muon has emerged as a strong contender in large-scale deep learning.<n>We study Muon on a simple, well-understood, and spectrally structured problem: low-rank matrix factorization.<n>We find that Muon does not consistently outperform AdamW in this setting.
arXiv Detail & Related papers (2026-07-14T20:21:38Z) - The Active Ingredient in Muon's Grokking [2.921159958223653]
We show the Muon-faster reaches the grokking threshold on modular arithmetic faster than AdamW.<n>To better understand Moun's behavior, we run multi-seed and multi-learning-rate sweeps to decompose and stress-test the effect.
arXiv Detail & Related papers (2026-07-06T22:35:57Z) - Extra-Merge: Tracing the Rank-1 Subspace of Model Merging in Language Model Pre-Training [56.323119575322146]
Extra-Merge consistently outperforms standard merging baselines in experiments across GPT-2 and LLaMA families.<n>It yields consistent zero-shot accuracy gains on Pythia-12B downstream tasks and generalizes effectively to the Muon citepjordan2024muon.
arXiv Detail & Related papers (2026-05-26T02:48:34Z) - Move on Muon : A Hamiltonian probability gradient flow perspective of Muon optimizer [6.003944904428636]
We develop a gradient flow on the space of probability measures defined on matrix-valued parameters induced by regularized Muon.<n>We prove an exact Hamiltonian dissipation identity, showing that the Hamiltonian energy decreases monotonically.<n>We extend the formulation to Hilbert-valued feature maps on product matrix spaces, yielding a blockwise Muon probability flow.
arXiv Detail & Related papers (2026-05-22T17:28:53Z) - Spectral Flattening Is All Muon Needs: How Orthogonalization Controls Learning Rate and Convergence [15.754779716816174]
We show that Muon's maximal stable step size scales with the average singular value of the gradient rather than the largest.<n>We recast Muon as a preconditioned gradient method and show, under a Kronecker-factored curvature model, that it improves the effective convergence factor.
arXiv Detail & Related papers (2026-05-13T06:54:01Z) - Sharp Capacity Scaling of Spectral Optimizers in Learning Associative Memory [77.27772368491698]
Spectrals such as Muon have recently shown strong empirical performance in large-scale language model training.<n>We study this question through the linear associative memory problem.<n>We show that the storage capacity of Muon significantly exceeds that of SGD.
arXiv Detail & Related papers (2026-03-27T16:13:18Z) - Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning [72.62839712454196]
Recent advances in spectral optimization, notably Muon, have demonstrated that constraining update steps to the Stiefel manifold can significantly accelerate training and improveimation.<n>We argue that this "egalitarian" constraint is suboptimal for Deep Neural Networks, where the curvature spectrum is known to be highly heavy-tailed and ill-conditioned.<n>We propose textbfMousse (textbfMuon textbfOptimization textbfUtilizing textbfShampoo's textbfStruct
arXiv Detail & Related papers (2026-03-10T14:03:49Z) - IDLM: Inverse-distilled Diffusion Language Models [70.5793829229702]
We extend Inverse Distillation, a technique originally developed to accelerate continuous diffusion models, to the discrete setting.<n>From a theoretical perspective, the inverse distillation objective lacks uniqueness guarantees, which may lead to suboptimal solutions.<n>We show that our method, Inverse-distilled Diffusion Language Models (IDLM), reduces the number of inference steps by 4x-64x.
arXiv Detail & Related papers (2026-02-22T06:47:04Z) - Insights on Muon from Simple Quadratics [2.8348950186890467]
Muon updates weight matrices along (approximate) polar factors of the gradients.<n>Existing attempts at explaining its performance largely focus on single-step comparisons.<n>We show that understanding Muon requires going beyond local proxies and pessimistic worst-case bounds.
arXiv Detail & Related papers (2026-02-12T13:43:58Z) - Muon in Associative Memory Learning: Training Dynamics and Scaling Laws [23.350512542598803]
We study Muon in a linear associative memory model with softmax retrieval and a hierarchical frequency spectrum over query-answer pairs.<n>We show that Muon mitigates this imbalance, leading to faster and more uniform progress.
arXiv Detail & Related papers (2026-02-05T14:49:40Z) - Preconditioning Benefits of Spectral Orthogonalization in Muon [50.62925024212989]
We study the effectiveness of a simplified variant of Muon in two case studies: matrix factorization and in-context learning of linear transformers.<n>Our analysis reveals that the Muon dynamics decouple into a collection of independent scalar sequences in the spectral domain, each exhibiting similar convergence behavior.
arXiv Detail & Related papers (2026-01-20T00:08:31Z) - Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMs [77.66717051042032]
Practitioners have consistently observed three puzzling phenomena in transformer-based large language models.
These phenomena are characterized by certain so-called "sink tokens" receiving disproportionately high attention weights.
We elucidate the mechanisms behind extreme-token phenomena.
arXiv Detail & Related papers (2024-10-17T17:54:06Z) - Losing momentum in continuous-time stochastic optimisation [42.617042045455506]
momentum-based optimisation algorithms have become particularly widespread.
In this work, we analyse a continuous-time model for gradient descent with momentum.
We also train a convolutional neural network in an image classification problem.
arXiv Detail & Related papers (2022-09-08T10:46:05Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.