Muon in Associative Memory Learning: Training Dynamics and Scaling Laws

AI summary

Overview

Research area: Machine learning theory — optimization dynamics and neural scaling laws for the Muon optimizer, studied through the lens of associative memory learning.

Technical level: Advanced. The paper is written in the style of an ICML theory paper, with formal assumptions, propositions and asymptotic bounds (Θ, 𝒪, Ω, Õ, Ω̃, ω_K(1), o_K(1)) and matrix-algebraic arguments. Readers need comfort with singular value decomposition, spectral norms, softmax cross-entropy and asymptotic analysis.

Scope in one sentence: The paper gives a theoretical account of how Muon (momentum-free, matrix-sign updates) learns query–answer associations faster and more evenly than gradient descent, with explicit convergence rates, a scaling law under a power-law frequency spectrum, and a preconditioning interpretation.

Note: the supplied paper content is truncated partway through Section 6.1, so the SignGD comparison section, the experiments section, and any conclusion are only inferable from the abstract and contribution list, not from detailed text.

What This Paper Is About

Muon updates matrix-valued parameters using the matrix sign of the gradient rather than the raw gradient, and it has shown strong empirical gains in large-scale training, but its theoretical underpinnings have been thin — most prior work proves static convergence bounds rather than characterizing the training trajectory. This paper asks what Muon actually does during learning: how it distributes progress across knowledge items of different frequencies, how fast it converges relative to gradient descent (GD), and how its loss scales with training time. To answer this, the authors build a tractable linear associative memory model with softmax retrieval and a hierarchical (and later power-law) frequency spectrum over query–answer pairs, with and without label noise.

Key Contributions

1. Training dynamics and acceleration. In the noiseless associative memory model the authors show GD learns different frequency components at highly imbalanced rates and is bottlenecked by low-frequency classes, while Muon equalizes progress across frequency groups and achieves an exponential speedup. Under label noise they characterize Muon's three-phase training dynamics (later, per sub-task, a two-phase descent/oscillation pattern) and show an Ω(C)-fold speedup, where C is the knowledge-group size.

2. Scaling laws. When group frequencies follow a power law (p̃_i ∝ i^{−β}, β > 1), they derive Muon's optimization scaling law, showing loss decays as Õ(1/T²), whereas GD admits a lower bound of Ω̃(1/T^{1−1/β}), yielding a substantially steeper scaling exponent for Muon. The authors state this is the first theoretical analysis to characterize the scaling laws of Muon.

3. Mechanism and connection to SignGD. They provide a preconditioning perspective in which Muon acts as an implicit matrix preconditioner: the matrix-sign update implicitly aligns with the underlying task representations and exploits a block-symmetric gradient structure. By contrast, a coordinate-wise SignGD could match Muon only under oracle access to the unknown task representations, which is infeasible in practice.

4. Experiments. Experiments on synthetic long-tail (imbalanced) classification and LLaMA-style pre-training are reported to corroborate the theory, showing improved long-tail learning and stronger scaling efficiency consistent with the predictions. (Detailed experimental tables and setups are not present in the truncated content provided.)

Main Findings

• GD is frequency-bottlenecked, Muon is not. In the noiseless case, the sub-task cross-entropy loss for the j-th knowledge item under GD satisfies ℒ_j^GD(t) ≂ 1/(p_j t), so low-frequency items converge much more slowly; the total loss scales as ℒ^GD(t) ≂ K/t. Under Muon, ℒ_j^Muon(t) ≂ K·e^{−(1+o_K(1))t} for every sub-task, and the total loss follows the same form — implying uniform treatment of frequency components, since ℒ_{j₁}^Muon(t)/ℒ_{j₂}^Muon(t) → 1 as K → ∞.

• Exponential acceleration in the noiseless regime. To reach loss precision ℒ(t) ≤ ε, GD requires 𝒪(1/ε) time whereas Muon requires 𝒪(log(1/ε)).

• Noisy sub-task dynamics have two phases. Under Muon, the j-th sub-task loss is bounded by Ke^{−η(1+o_K(1))t} + ηt when t ≤ T_j*, and by η² + ℒ_j* when t > T_j*, where T_j* = Θ(log K / η) is the critical time and ℒ_j* = ℒ* is the irreducible loss. The descent phase contains an exponential term plus a noise-accumulation term ηt; the oscillation phase fluctuates at scale ℒ_j* + 𝒪(η²). This creates a learning-rate trade-off: larger η speeds descent but enlarges the oscillation.

• Noisy total dynamics have three phases. For the full loss: all sub-tasks descending (t ≲ log K/η); a mixed phase (t ≂ log K/η); then all sub-tasks oscillating (t ≳ log K/η). The excess risk at the end of phase 1 is 𝒪(M² log K / K).

• A concrete GD lower bound and the C-fold speedup. Under the linear-stability condition ηp₁ ≲ 1, GD satisfies ℒ_j^GD(t) − ℒ_j* ≳ e^{−ηp_j t}(log K)², giving ℒ^GD(t) ≥ ℒ* + Ω(e^{−ηp₁t}(log K)²). To reach excess risk 𝒪(M² log K/K), GD needs Ω(C log(K/M²)/η) steps while Muon needs 𝒪(log K/η) — a speedup by a factor of C, treating M as a relatively fixed constant.

• Scaling-law separation. Under the power-law spectrum p̃_i ∝ i^{−β} with β > 1 and the proportional regime cM^β ≤ T ≤ M^β (0 < c < 1, (log K)^{1/β} ≤ M ≪ K^{1/2}), GD obeys ℒ^GD(T) − ℒ* ≳ log K / T^{1−1/β}, whereas Muon with η = Θ(log K/T) obeys ℒ^Muon(T) − ℒ* ≲ (log K / T)². Muon therefore achieves Õ(T^{−2}) scaling versus GD's Ω̃(T^{−(1−1/β)}) lower bound. The power-law decay of GD arises from task accumulation under a heavy-tailed frequency distribution; a single sub-task still decays exponentially.

• Muon as an implicit preconditioner. In task-representation space (Ŵ = ẼᵀWE, G_t = Ẽᵀ∇ℒ(W_t)E), the Muon update is Ŵ_{t+1} = Ŵ_t − η·msgn(G_t), and msgn(G_t) = msgn(P − P̂_t) where (P){i,j} = p_j p{i|j} and (P̂_t){i,j} = p_j p̂{i|j}(W_t). Proposition 6.1 shows that throughout the global descent phase, ‖msgn(P − P̂_t) − I_K‖_max = o_K(1) (and in the noiseless case for any finite t ≥ 0). Consequently msgn(G_t) ≈ I_K and Ŵ_t ≈ t·I_K: Muon preconditions updates to identify the desired task representations and moves near-isotropically across them.

• Why the preconditioner works: block symmetry. Starting from W₀ = 0, Muon updates preserve a block symmetry induced by the frequency groups. The residual decomposes as P − P̂_t = R_t⁺ − R_t⁻, where R_t⁺ is block diagonal with M blocks each proportional to I_C, and R_t⁻ is block-wise constant with M² blocks each proportional to J_C. On the (C−1)-dimensional within-group contrast subspaces 𝒮_i, the block-wise constant term vanishes and R_t⁺ acts as a scalar multiple of the identity, making the union 𝒮 = 𝒮₁ ⊕ ⋯ ⊕ 𝒮_M an M(C−1)-dimensional singular subspace.

Methodology in Plain English

The authors do not analyze a real Transformer. Instead they construct a deliberately simplified "associative memory" problem that still contains the ingredients they care about: a set of K atomic knowledge items, each a query (subject–relation pair) mapped to an answer. Query and answer embeddings are assumed to be orthonormal (Assumption 3.1) with equal norm, and the embedding dimension is set to d = K.

Knowledge items are grouped into M groups of size C = K/M, and every item in a group shares the same frequency (Assumption 3.2). Frequencies decay across groups, mimicking the head/tail structure of real corpora — common facts like ("Paris", "capital_of") → "France" versus rare facts like ("Thulium", "boiling_point") → 1950°C. The model is a linear softmax map: a memory matrix W ∈ ℝ^{K×K} produces a softmax probability over answers, and training minimizes the standard cross-entropy loss.

They analyze two optimizers: plain gradient descent (W_{t+1} = W_t − η∇ℒ(W_t)) and a momentum-free Muon (W_{t+1} = W_t − η·msgn(∇ℒ(W_t))), where msgn uses the SVD to replace singular values by their signs. This momentum-free variant coincides with Spectral GD (SpecGD), so the results transfer to that method as well. Both start from zero initialization.

Label noise is modeled with a parameter α ∈ [0,1): the observed label equals the true index with probability 1 − α, and is uniform over all K items with probability α, giving an explicit conditional distribution p_{i|j}. The noiseless case (α = 0) is treated first as a warmup with η = 1 fixed; the noisy case (0 < α < 1) is then analyzed, including a linear-stability condition on the learning rate. Two asymptotic regimes are used: a fixed-K, fixed-M, t → ∞ regime for the dynamics results, and a joint K, M, T → ∞ proportional regime (T ≂ M^β) for the scaling laws. The final section re-expresses everything in the task-representation space to explain Muon's behavior as preconditioning. The theory is complemented by synthetic long-tail classification and LLaMA-style pre-training experiments; the paper also notes an empirical comparison of normalized GD versus GD and Muon (Appendix G.3), because global normalization couples all columns and makes a sharp analysis beyond the paper's scope.

Why This Matters

• It closes a theory gap for a widely used optimizer. Muon is already deployed in large-scale training, yet prior theory focused on static convergence bounds. This work characterizes the training trajectory — which frequency components get learned when, and how the loss decays with training time — and, per the authors, is the first to characterize Muon's scaling laws.

• It explains the mechanism, not just the rate. The preconditioning view shows Muon's advantage is not merely "normalization." The authors explicitly note that while normalized GD is faster than GD, it remains more imbalanced and slower than Muon, suggesting Muon's acceleration is not solely due to normalization.

• It draws a sharp boundary against SignGD. A coordinate-wise sign method could match Muon only with oracle access to unknown task representations, which is infeasible in practice — an important distinction between matrix-sign and coordinate-sign preconditioning.

• It yields actionable hyperparameter guidance. The scaling-law result prescribes η = Θ(log K/T) to balance descent speed against oscillation magnitude, and the noisy analysis gives an explicit trade-off between learning rate and final oscillation scale.

Real-world applications (implications of the model, not claims verified on deployed systems):

• Large-scale LLM pre-training, where compute and data efficiency translate directly into cost.
• Long-tail knowledge retention — keeping rare or specialized facts from being crowded out by high-frequency ones.
• Optimizer and learning-rate-schedule design for matrix-valued parameters in deep networks, including dense Transformers and Mixture-of-Experts (MoEs) as cited in the paper.
• Retrieval-style or memory-based systems where query–answer association quality on tail items determines user-visible accuracy.

Industry relevance: Muon is described as consistently surpassing Adam across scales and architectures, and the paper's central claims — a C-fold step-count speedup to a fixed excess-risk target and Õ(T^{−2}) versus Ω̃(T^{−(1−1/β)}) scaling — speak directly to compute budgets in industrial training runs. The theory also informs when Muon's advantage should be expected to be large (strongly heavy-tailed frequency spectra, many items per group) rather than uniform.

Future Directions

• Beyond momentum-free Muon. The analysis deliberately removes Muon's momentum term. Extending the dynamics and scaling results to the momentum variant actually used in practice is an open question.
• Beyond exact orthogonal, equinorm embeddings. The paper states its techniques and results hold without the equinormality condition and that the strict orthogonality can be relaxed to near-orthogonality, but the fully general case and its quantitative consequences are not settled here.
• Beyond the linear softmax associative memory proxy. The model is offered as "analytically tractable" and a proxy for Transformer knowledge retrieval; carrying the three-phase dynamics and the T^{−2} scaling law over to multi-layer, nonlinear architectures remains open.
• Understanding normalized GD. The authors state that global normalization yields M coupled nonlinear systems, and that the noisy case introduces oscillations, making a sharp theoretical analysis substantially more involved and beyond the scope of this work — so a theory of when normalization alone suffices is still missing.

Target Audience

Optimization theorists and machine learning researchers working on training dynamics, preconditioned or spectral optimizers, and neural scaling laws; practitioners training large models who want principled guidance on optimizer choice and learning-rate scheduling for matrix parameters; and graduate students with a background in convex/asymptotic analysis and matrix algebra who want a worked example of deriving a scaling law from first principles. Familiarity with Muon, SignGD/SignSGD, softmax cross-entropy and asymptotic notation will make the paper substantially easier to read.

Authors’ 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-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.