Title: Investigating Training Dynamics and Scaling Laws of Muon in Associative Memory

Abstract:

Although the Muon optimizer has demonstrated significant empirical improvements by updating matrix parameters through the matrix sign of the gradient, its theoretical underpinnings and scaling behaviors remain poorly understood. This study analyzes Muon within a linear associative memory framework featuring softmax retrieval and a hierarchical frequency distribution across query-answer pairs, examining scenarios both with and without label noise. We reveal that Gradient Descent (GD) suffers from highly imbalanced learning rates across frequency components, causing convergence bottlenecks due to slow progress in low-frequency regions. Conversely, Muon alleviates this disparity, enabling more rapid and consistent advancement. Specifically, our findings indicate that Muon provides an exponential acceleration over GD in noise-free environments. In noisy conditions characterized by a power-law frequency spectrum, we establish Muon’s scaling law and confirm its superior efficiency compared to GD. Additionally, we interpret Muon as an implicit matrix preconditioner driven by adaptive task alignment and the block-symmetric nature of the gradient structure. While a preconditioner utilizing a coordinate-wise sign operator could theoretically replicate Muon’s performance with oracle knowledge of unknown task representations, such access is impractical for SignGD in real-world applications. Our theoretical insights are validated through experiments on synthetic long-tail classification tasks and pre-training simulations resembling LLaMA architectures.

Source: arXiv Generated at: 2026-06-04 00:00:00 UTC