PaperlayerView exact revision on arXiv

Source-linked AI summary

Muon Outperforms Adam in Tail-End Associative Memory Learning

Shuche Wang, Fengzhuo Zhang, Jiaxiang Li, Cunxiao Du, Chao Du, Tianyu Pang, Zhuoran Yang, Mingyi Hong, Vincent Y. F. Tan

arXiv:2509.26030v2cs.LGcs.AImath.OC

TL;DR

The paper asks why Muon outperforms Adam despite incomplete explanations based on spectral- versus vector-norm optimization. It studies this through transformer ablations, spectral and heavy-tailed learning analyses, and associative-memory theory, finding that Muon’s outer-product-aligned updates support more balanced tail-class learning. The conclusions are bounded by a fixed learning-rate schedule assumption and by the distinct role of QK weights.

  • Problem

    The paper addresses the incomplete understanding of why Muon’s spectral-norm optimization outperforms Adam’s vector-norm optimization in transformers.

  • Method

    The paper combines transformer-component ablations, singular-spectrum and heavy-tailed-class experiments, and a one-layer associative-memory theoretical model.

  • Results

    Muon’s update rule aligns with linear associative-memory outer products, yielding more isotropic spectra and more balanced tail-class learning than Adam.

  • Takeaways & Limitations

    Muon’s associative-memory alignment provides a supported explanation for its more balanced optimization of tail classes in heavy-tailed distributions.

  • Takeaways & Limitations

    The theoretical analysis assumes a fixed learning-rate schedule, and QK weights are not expected to benefit from the associative-memory isotropy mechanism.

Abstract

from arXiv · show

The Muon optimizer is consistently faster than Adam in training Large Language Models (LLMs), yet the mechanism underlying its success remains unclear. This paper demystifies this mechanism through the lens of associative memory. By ablating the transformer components optimized by Muon, we reveal that the associative memory parameters of LLMs, namely the Value and Output (VO) attention weights and Feed-Forward Networks (FFNs), are the primary contributors to Muon's superiority. Motivated by this associative memory view, we then explain Muon's superiority on real-world corpora, which are intrinsically heavy-tailed: a few classes (tail classes) appear far less frequently than others. The superiority is explained through two key properties: (i) its update rule consistently yields a more isotropic singular spectrum than Adam; and as a result, (ii) on heavy-tailed data, it optimizes tail classes more effectively than Adam. Beyond empirical evidence, we theoretically confirm these findings by analyzing a one-layer associative memory model under class-imbalanced data. We prove that Muon consistently achieves balanced learning across classes regardless of feature embeddings, whereas Adam can induce large disparities in learning errors depending on embedding properties. In summary, our empirical observations and theoretical analyses reveal Muon's core advantage: its update rule aligns with the outer-product structure of linear associative memories, enabling more balanced and effective learning of tail classes in heavy-tailed distributions than Adam.

1 Introduction

The paper investigates why Muon outperforms Adam and links its advantage to associative-memory components and heavy-tailed learning. Empirical and theoretical analyses identify balanced updates across tail classes as the central mechanism.

  • 1 Introduction: The paper addresses an incomplete explanation of why spectral-norm optimization in Muon outperforms Adam’s vector-norm optimization.Existing steepest-descent and convergence perspectives do not account for Muon’s observed superiority.
  • 1 Introduction: Ablations show that Muon’s faster validation-loss convergence primarily comes from optimizing attention VO matrices and FFN blocks.These components are identified as the model’s primary associative-memory stores.
  • 1 Introduction: Muon’s spectral normalization balances learning rates across orthogonal outer products, reducing the dominance of frequent facts in heavy-tailed data.This hypothesis connects the optimizer’s update mechanism to associative-memory learning dynamics.
  • 1 Introduction: Muon produces more isotropic singular spectra than Adam and outperforms Adam on rare tail classes while both perform well on frequent head classes.The heavy-tailed evaluation also shows more stable and uniform convergence with Muon.
  • 1 Introduction: The paper combines transformer-component ablations, empirical spectral and tail-class analyses, and a theoretical one-layer associative-memory model.The theoretical analysis confirms balanced learning across classes under class imbalance.
  • 1 Introduction: Muon’s update rule aligns with the outer-product structure of linear associative memories, enabling more balanced learning of tail classes than Adam.The paper frames this alignment as the mechanism underlying Muon’s superiority on heavy-tailed distributions.

2 Related Works

Related work positions Muon among optimizer analyses and associative-memory interpretations of transformers. Prior studies connect transformer attention and FFN parameters to storing and retrieving knowledge.

  • 2 Related Works: Prior Adam research studies convergence through online optimization and feature learning perspectives.These perspectives examine optimization guarantees and how Adam affects learned feature properties.
  • 2 Related Works: Muon applies spectral normalization to matrix gradients and has been reported to outperform Adam across diverse model sizes and architectures.It is also interpreted as steepest descent under the matrix operator norm.
  • 2 Related Works: Associative-memory research links transformer knowledge storage to FFN output weights and attention output matrices.The cited literature motivates treating transformer components as mechanisms for retaining learned associations.

3 Preliminaries

The preliminaries define Muon’s matrix-normalized update, transformer components, and the associative-memory view of attention and FFN parameters. These definitions establish the objects analyzed in the paper.

  • 3 Preliminaries: Muon replaces a raw or momentum gradient with normalized orthogonal factors, producing a scale-invariant, norm-controlled matrix update.It forms a momentum accumulator, computes its SVD, and updates parameters using the nearest semi-orthogonal matrix.
  • 3 Preliminaries: Muon’s update can be interpreted as steepest gradient descent under the matrix operator norm, unlike Adam’s vector-norm perspective.The preliminaries note that this perspective alone does not explain Muon’s better performance.
  • 3 Preliminaries: Each transformer layer processes representations through an attention module followed by a feed-forward network.The model maps token embeddings through successive layers to predict the next token.
  • 3 Preliminaries: The FFN applies an activation-based transformation with learnable output parameters, and gated variants are common in LLMs.The preliminaries distinguish the standard FFN form from its gated variant.
  • 3 Preliminaries: A linear associative memory maps a key encoding a subject-relation pair to a value representing the associated object.The paper uses this formulation to model stored factual associations.
  • 3 Preliminaries: The attention output matrix WO and value matrix WV, together with FFN parameters, are treated as associative-memory parameters.Prior work motivates this treatment by showing that these matrices encode or store associations.

4 Main Results

Muon’s advantage is concentrated in associative-memory components, especially VO and FFN, whose updates produce more isotropic weights and improve learning of rare knowledge classes.

  • 4.1 Associative Memories Are Main Beneficiaries of Muon: VO benefits more from Muon than QK, and ablations indicate that WO is more influential than WV within the VO+FFN configuration.Applying Muon only to WV or WO already outperforms applying it to QK, while replacing either VO component with Adam degrades performance.
  • 4.1 Associative Memories Are Main Beneficiaries of Muon: Applying Muon to VO+FFN nearly recovers the full-Muon validation-loss trajectory, while QK contributes little to overall performance.This pattern holds across ungated and gated FFN settings; the remaining gap may reflect learning-rate tuning.
  • 4.1 Associative Memories Are Main Beneficiaries of Muon: Muon updates orthogonal outer-product facts at the same rate, aligning its matrix update with the structure of linear associative memories.Because singular values encode knowledge frequencies, this mechanism supports more uniform learning of frequent and infrequent facts.
  • 4.2 Muon Consistently Learns More Isotropic Weights Than Adam: Muon consistently produces more isotropic weight matrices than Adam, with spectral energy distributed broadly across directions throughout training and across random initializations.The analysis uses normalized SVD entropy, effective rank, and related singular-spectrum metrics.
  • 4.3 Muon Acquires Knowledge More Evenly Compared To Adam: On heavy-tailed knowledge tasks, Muon matches Adam on frequent classes but substantially improves convergence on rare classes and narrows the head–tail gap.Applying Muon to VO+FFN produces strong rare-class gains, whereas applying it only to QK yields limited improvement.

5 Case Study of One-Layer Models

The paper studies GD, Adam, and Muon in a one-layer linear associative-memory model under imbalanced data and varied embeddings. Across experiments and theory, Muon remains balanced and nearly isotropic, while Adam and GD can exhibit embedding-dependent imbalance.

  • 5 Case Study of One-Layer Models: The analysis assumes orthonormal embeddings and a two-class frequency imbalance, with the multi-class case extending through the same SVD calculation.
  • 5.1 Experimental Results: As loss decreases, imbalance first rises and then falls for all optimizers, becoming smaller once correct probabilities reach about 0.9.
  • 5.1 Experimental Results: Muon remains substantially more balanced across items than GD under both support-decoupled and support-coupled embeddings.SignGD is stable only in the decoupled case and resembles GD in the coupled case.
  • 5.2 Theoretical Results: Muon learns balancedly for any embeddings, whereas Adam can produce large disparities in learning error depending on embedding structure.The theoretical analysis compares one-step and multi-step associative-memory learning under class imbalance.
  • 5.2 Theoretical Results: Muon's update allocates nearly equal strength across directions, so its update singular values are nearly identical.
  • 5.2 Theoretical Results: Adam's element-wise sign normalization can disrupt matrix structure, producing spectral decay in which the smallest singular value falls below 25% of the largest.

6 Conclusion

The paper connects Muon's transformer gains to associative-memory structure and balanced learning on heavy-tailed data. Its conclusion is that Muon's outer-product-aligned updates improve tail-class learning relative to Adam.

  • 6 Conclusion: Muon’s update rule aligns with the outer-product structure of linear associative memories, enabling more balanced and effective tail-class learning than Adam.The conclusion identifies this alignment as the paper’s central explanation for Muon’s advantage.

A Steepest Descent View Understanding Muon and Adam

The steepest-descent framework explains optimizers through the norm measuring update size: Adam uses the vector infinity norm, while Muon uses the matrix spectral norm. Their resulting updates differ in whether coordinates or singular directions determine the step.

  • Steepest-descent framework: The framework formulates each update by minimizing a local quadratic loss approximation under a chosen norm and its associated dual norm.Changing the norm changes the update direction and recovers different optimizer rules.
  • Muon under the spectral norm: Muon applies steepest descent under the spectral norm, using gradient singular vectors for direction and the sum of singular values for step scaling.For G = UΣV^⊤, the direction is UV^⊤ and the scaling factor is tr(Σ).
  • Adam under the infinity norm: Adam corresponds to steepest descent under the vector infinity norm, producing an update with equal parameter magnitudes determined by gradient signs.Its step-size scaling uses the gradient’s ℓ1 norm.

B.1 Experimental Details of Training on FineWeb

The experiments combine FineWeb training, synthetic heavy-tailed knowledge evaluation, attention diagnostics, and 0.7B scaling. Across the scaled model, Muon performs best when applied broadly or specifically to VO and FFN components, while QK-only application offers little advantage.

  • Heavy-tail knowledge task: The heavy-tailed task uses synthetic biographical QA data whose class frequencies follow a power-law distribution, with head classes common and tail classes rare.The evaluation reports First Token Accuracy across frequency groups from head to tail.
  • Heavy-tail knowledge task: The dataset contains 32,768 classes, ranging from 196,608 samples for the head class to 6 samples for each of 16,384 extreme-tail classes.These settings use m = 15 and nqa = 6.
  • MaxLogit diagnostics: With RMSNorm applied to Q and K, per-layer MaxLogit values remain moderate without runaway growth, so component differences are not attributed to MaxLogit explosion.The measurement is reported for the 160M NanoGPT model trained with Muon.
  • Scaling to the 0.7B NanoGPT model: Muon applied to all components achieves the lowest 0.7B validation loss, while applying it only to VO and FFN performs nearly identically and QK-only offers little advantage over Adam.The comparison covers both non-gated and gated feed-forward configurations in Figure 5.
  • Scaling to the 0.7B NanoGPT model: For the 0.7B model, Muon produces higher SVD entropy and eRank than Adam in VO, Win, Wgate, and Wout matrices.These metrics indicate more distributed, higher-dimensional representations in the reported spectral analysis.

C.3 Additional Results about Spectral Dynamics of Transformer Weight Matrices During Training

Additional FFN spectral analyses on 160M NanoGPT models show the same Muon signature across non-gated and gated matrices: more distributed singular-value spectra than Adam.

  • FFN spectral dynamics: Across non-gated Win, gated Win, and gated Wgate, Muon increases SVD entropy and eRank while decreasing Top-k energy and the Q75/25 ratio.The consistent trends mirror the main-text spectral observation across all three FFN matrices.

C.4 Detailed Experiment Results about Heavy-Tail Imbalance Knowledge Task

Detailed tail-group results compare full Muon, Adam, SGD with momentum, and two hybrid optimizers at 2k, 5k, and 10k steps. They show faster and more balanced rare-class learning with Muon, especially when applied to VO and FFN components.

  • Tail-group performance: Muon converges rapidly on rare groups, reaching strong First Token Accuracy by 2k steps and near-saturated performance by 5k.The reported groups are g = 11, 13, and 15, representing increasingly rare frequency bands.
  • Tail-group performance: Adam retains a persistent head–tail accuracy gap across the reported checkpoints, unlike the more balanced learning pattern associated with Muon.The tables report exact FTA values for selected tail groups at 2k, 5k, and 10k steps.
  • Optimizer ablations: The VO&FFN hybrid closely tracks full Muon on tail generalization, whereas the QK-only hybrid lags behind.This identifies VO and FFN updates as the dominant contribution among the tested hybrids.

C.5 Additional Experiment Results about Heavy-Tail Imbalance Knowledge Task with Gated Feed-Forward Networks

In the gated FFN variant, Muon performs best on rare classes and reaches high accuracy earlier than Adam and SGD+Momentum. Applying Muon to VO and FFN closely matches full Muon, while applying it only to QK provides limited gains.

  • Gated FFN heavy-tail results: Muon consistently outperforms Adam and SGD+Momentum on rare classes in the gated heavy-tail knowledge task.The evaluation compares group performance at 2,000, 5,000, and 10,000 training steps.
  • Gated FFN heavy-tail results: The VO&FFN hybrid closely tracks full Muon, indicating that VO and FFN are the primary levers for tail generalization.This hybrid applies Muon to VO and FFN while using Adam for QK.
  • Gated FFN heavy-tail results: The QK-only hybrid offers limited gains relative to the configurations emphasizing VO and FFN.
  • Gated FFN heavy-tail results: The gated FFN variant does not change the qualitative conclusion about where Muon helps most.The section reports the same qualitative pattern as the non-gated setting.

D Proof of Theorem 5.3

The proof analyzes one-step updates for GD, Muon, and Adam under class-imbalanced associative-memory data. It establishes Muon's result through singular-value analysis and contrasts it with Adam under different embedding constructions.

  • GD calculation: The GD calculation shows sensitivity to data imbalance, with the imbalance quantity r(α, β) controlling disparities in correct-class probabilities.The interpretation states that r(α, β) ≪ 1 represents severe imbalance and that GD can be highly sensitive to it.
  • Muon calculation: Muon is analyzed by computing the gradient SVD and applying its normalized singular-direction update.The proof writes the gradient in SVD form and uses the resulting update quantity GMuon(W0).
  • Muon calculation: The proof bounds fastest- and slowest-learning triplets using corresponding Muon update directions.These bounds reduce the convergence analysis to two representative update directions.
  • Adam calculation: For Adam, the proof considers identity and block-wise diagonal embeddings, yielding distinct learning-dynamics analyses.The second construction uses orthogonal 3 × 3 blocks when K is divisible by 3.

E Proof of Theorem 5.4

The proof extends the associative-memory analysis to multi-step Muon optimization by preserving shared parameter and score forms along the trajectory. It then derives convergence bounds from these structured dynamics and singular-value decompositions.

  • Multi-step Muon analysis: Theorem 5.4 extends the analysis from one-step updates to multi-step Muon optimization.The proof derives shared forms for W_t and GMuon along the full optimization trajectory.
  • Trajectory structure: The induction hypothesis constrains W_t through shared coefficients and off-block terms of order O(a_t/K).It requires a_t = b_t ≥ 0 and c_ij^t = O(a_t/K).
  • Trajectory structure: Symmetry reduces the score dynamics to a small set of representative within- and across-group cases.The proof identifies five shared score forms for triplets in the two class groups.
  • Convergence analysis: The multi-step proof bounds fastest- and slowest-learning triplets by two one-step update directions.The remaining convergence analysis is stated to follow the same procedure as the one-step Muon proof.
  • Singular-value analysis: The appendix characterizes the relevant singular structure by decomposing the matrix into subspaces and a remaining 2 × 2 block.The SVD assigns singular values to the subspaces and derives the remaining values from the 2 × 2 matrix.
Loading 2509.26030v2…