Muon in Associative Memory Learning: Training Dynamics and Scaling Laws

Muon updates matrix parameters via the matrix sign of the gradient and has shown strong empirical gains, yet its dynamics and scaling behavior remain unclear in theory. We study Muon in a linear associative memory model with softmax retrieval and a hierarchical frequency spectrum over query–answer pairs, with and without label noise. In this setting, we show that Gradient Descent (GD) learns frequency components at highly imbalanced rates, leading to slow convergence bottlenecked by low-frequency components. In contrast, the Muon optimizer mitigates this imbalance, leading to faster and more uniform progress. Specifically, in the noiseless case, Muon achieves an exponential speedup over GD; in the noisy case with a power-law frequency spectrum, we derive Muon's scaling law and demonstrate its superior scaling efficiency over GD. Furthermore, we show that Muon can be interpreted as an implicit matrix preconditioner arising from adaptive task alignment and block-symmetric gradient structure. In contrast, the preconditioner with coordinate-wise sign operator could match Muon under oracle access to unknown task representations, which is infeasible for SignGD in practice. Experiments on synthetic long-tail classification and LLaMA-style pre-training corroborate the theory.

While the new Muon optimizer has shown impressive real-world results in training AI models, the theoretical reasons behind its success have remained a mystery. In this paper, we bridge that gap. We prove that standard Gradient Descent (GD) is inefficient because it learns common data patterns quickly but gets severely bottlenecked by rare ones. In contrast, Muon acts as an adaptive mechanism that balances this learning process, making it much faster and more uniform. We mathematically demonstrate that Muon can achieve exponential speedups over traditional methods in clean settings, and scales vastly better in noisy, real-world environments. We validate these theoretical breakthroughs with experiments on both synthetic datasets and large language model (LLaMA) pre-training.