Muon in Associative Memory Learning: Training Dynamics and Scaling Laws
Not on the map. Why?
Abstract (arXiv)
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.
Key facts. Cited by 7 papers in the map (data through 7 Oct 2026). First posted on arXiv 5 Feb 2026.
Builds on
0 of 0 references.
No references here yet.
Built on by
Citing papers per year: 2026: 7.
7 papers.
- 2026Optimal Learning Rate Schedules under Functional Scaling Laws: Power Decay and Warmup-Stable-Decaycited by 4
- 2026Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Lawscited by 4
- 2026Why Muon Outperforms Adam: A Curvature Perspectivecited by 4
- 2026Muon Learns More Robust and Transferable Features than Adamcited by 1
- 2026Muonp: Muon with Fractional Spectral Powerscited by 0
- 2026Muon learns balanced solutions in matrix factorization without slow saddle-to-saddle dynamicscited by 0
- 2026Convergence of Practical Muoncited by 0
- 2026Convergence of Practical Muon27 Sep 2026
- 2026Muon learns balanced solutions in matrix factorization without slow saddle-to-saddle dynamics29 Jun 2026
- 2026Muonp: Muon with Fractional Spectral Powers11 Jun 2026
- 2026Muon Learns More Robust and Transferable Features than Adam8 Jun 2026
- 2026Why Muon Outperforms Adam: A Curvature Perspective3 Jun 2026
- 2026Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Laws15 Feb 2026
- 2026Optimal Learning Rate Schedules under Functional Scaling Laws: Power Decay and Warmup-Stable-Decay6 Feb 2026