Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Muon in Associative Memory Learning: Training Dynamics and Scaling Laws

Binghui Li, Kaifei Wang, Han Zhong, Pinyan Lu, Liwei Wang

TL;DR

This work derives Muon's optimization scaling law and demonstrates its superior scaling efficiency over Gradient Descent, and shows that Muon can be interpreted as an implicit matrix preconditioner arising from adaptive task alignment and block-symmetric gradient structure.

Abstract

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-decay frequency spectrum, we derive Muon's optimization 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.

Muon in Associative Memory Learning: Training Dynamics and Scaling Laws

TL;DR

This work derives Muon's optimization scaling law and demonstrates its superior scaling efficiency over Gradient Descent, and shows that Muon can be interpreted as an implicit matrix preconditioner arising from adaptive task alignment and block-symmetric gradient structure.

Abstract

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-decay frequency spectrum, we derive Muon's optimization 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.
Paper Structure (81 sections, 28 theorems, 254 equations, 3 figures)

This paper contains 81 sections, 28 theorems, 254 equations, 3 figures.

Table of Contents

  1. Introduction
  2. Related Work
  3. Associative memory.
  4. The Muon optimizer.
  5. Theoretical analysis of Muon.
  6. Theory of Scaling laws.
  7. Preliminaries
  8. Notations.
  9. Associative memory learning
  10. Knowledge and embeddings.
  11. Knowledge structure.
  12. Linear softmax model.
  13. Training
  14. Label noise.
  15. Cross-entropy loss.
  16. ...and 66 more sections

Key Result

Proposition 3.3

When $\alpha = 0$, we consider a specific weight $\mathbf{W} = \gamma \widetilde{\mathbf{E}} \mathbf{E}^\top$, where $\gamma > 0$ denotes a scale coefficient. Then, it holds that $\operatorname{lim}_{\gamma\to\infty}\mathcal{L}(\mathbf{W}) = 0$.

Figures (3)

  • Figure 1: (a)(b): Numerical simulations; (c)(d): Synthetic imbalanced classification.
  • Figure 2: Data scaling behavior in language model pre-training. To extract the scaling envelope, we sweep over learning rates for each data budget and plot the minimum validation loss.
  • Figure 3: Numeriical simulations (K=1000)

Theorems & Definitions (52)

  • Proposition 3.3: Noiseless case
  • Proposition 3.4: Noisy case
  • Proposition 3.5: Gradient decomposition
  • Theorem 4.1
  • Theorem 4.2
  • Theorem 5.1
  • Theorem 5.2
  • Definition 5.3: Linear stability strogatz2001nonlinear
  • Proposition 5.4: Stability at minimizer
  • Theorem 5.5
  • ...and 42 more