Technical Summary: Muon Learns More Robust and Transferable Features than Adam

Problem Statement

While the Muon optimizer (Momentum Orthogonalized by Newton-Schulz) has recently demonstrated state-of-the-art efficiency in pretraining Large Language Models (LLMs) and vision classifiers—often training nearly twice as fast as Adam and SGD—the specific nature of its feature-learning advantages remains unclear. Existing literature has largely explained Muon's success through optimization-theoretic lenses, such as steepest gradient descent under spectral norms, associative memory dynamics, or heavy-tailed data handling. However, a systematic investigation into whether Muon learns better features than Adam and SGD, and how these features manifest in model representations, has been lacking. This paper addresses the question: Does Muon learn more robust and transferable features than Adam and SGD, and what properties characterize this advantage?

Methodology

The authors employ a dual approach combining extensive empirical evaluation across vision and language domains with theoretical analysis on a stylized classification problem.

Empirical Evaluation

The study evaluates feature quality through two primary metrics: robustness and transferability.

• Robustness: Pretrained models (ResNet-18, ViT-S, GPT-2, GPT-2 Medium) are tested on corrupted inputs. Vision models are evaluated on ImageNet-C (15 corruption types), while language models are evaluated on a newly constructed benchmark, FineWeb10B-C, which applies character-level typo attacks (insertions, deletions, substitutions, swaps).
• Transferability: Pretrained features are adapted to downstream tasks. For vision, linear classifiers are trained on frozen backbones across four datasets (EuroSAT, Food-101, Flowers-102, Stanford Cars). For language, full models are fine-tuned on instruction-following datasets (Alpaca, Dolly, WizardLM).
• Hidden-State Analysis: To understand the mechanism behind these performance gains, the authors analyze layer-wise hidden states using:
  • Logit Margins: Trained linear probes (tuned-lens framework) decode hidden states to compute the margin between the correct class/token and the strongest distractor.
  • Spectral Diversity: The effective rank and Top-k energy fraction of hidden-state matrices are computed to measure the diversity and distribution of information across orthogonal directions.
• Experimental Protocol: All comparisons follow a matched-budget protocol where architecture, data, and training duration are fixed, while hyperparameters (learning rate, scheduler) are tuned independently for each optimizer to isolate representation differences from engineering effort.

Theoretical Analysis

The authors provide theoretical support using a one-layer linear classification model with multi-component features. The data model assumes features consist of a shared block-level component (coarse class) and a class-specific component (fine-grained class). The analysis compares the training dynamics of Gradient Descent (GD), Adam, and Muon under infinitesimal step sizes and zero momentum, focusing on the state where each optimizer first reaches a specific loss threshold.

Key Contributions and Results

1. Empirical Evidence of Superior Feature Quality

The paper establishes that Muon consistently learns features that are more robust and transferable than those learned by Adam and SGD:

• Robustness: Muon-pretrained models achieve the highest mean accuracy on ImageNet-C and the lowest perplexity on FineWeb10B-C across all corruption types and architectures. Notably, Muon outperforms Adam in vision transformers and language models, and outperforms SGD in language models.
• Transferability: Muon-pretrained features transfer more effectively to downstream tasks. In vision, Muon matches or surpasses Adam and SGD on all downstream benchmarks. In language, Muon-pretrained models achieve lower perplexity after instruction tuning compared to Adam-pretrained models.

2. Hidden-State Mechanisms

The empirical advantages are directly linked to specific properties of the learned hidden representations:

• Larger Logit Margins: Layer-wise probing reveals that Muon yields significantly larger logit margins across most layers compared to Adam and SGD. This indicates a wider separation between correct predictions and distractors, directly correlating with the observed robustness.
• Higher Effective Rank: Spectral analysis shows that Muon attains a higher effective rank and lower Top-10 energy fraction across layers. This suggests Muon's representations spread information more evenly across orthogonal directions, creating a richer feature subspace that facilitates better adaptation to downstream tasks.

3. Theoretical Guarantees

In the stylized multi-component feature setting, the authors prove that Muon's spectral normalization removes singular-value imbalances in the gradient matrix.

• Margin Advantage: Muon attains a strictly larger classification margin than Adam and GD at any matched training loss (outside an exponentially small terminal-loss window). This is because Muon's spectral normalization balances the contribution of orthogonal feature directions, whereas Adam's coordinate-wise normalization and GD's raw gradient flow lead to over-reliance on specific components.
• Effective Rank Advantage: Muon provably attains a strictly higher feature effective rank than Adam and GD at matched loss, supporting the empirical finding of more diverse representations.

Significance and Claims

The paper claims to "demystify" Muon's advantages by shifting the focus from optimization efficiency to feature quality. The authors argue that Muon's success is not merely a result of faster convergence but stems from its ability to learn representations that are inherently more robust to input perturbations and more adaptable to new tasks.

The significance of this work lies in:

1. Connecting Optimization to Representation: It establishes a concrete link between the matrix-aware update rule of Muon and the geometric properties of learned features (logit margins and effective rank).
2. Theoretical Validation: It provides the first theoretical proof that Muon induces systematically different and superior feature qualities (larger margins, higher effective rank) compared to Adam and GD under matched loss conditions.
3. Practical Implications: The findings suggest that for applications requiring robustness to distribution shifts or strong transferability to downstream tasks, Muon is a superior choice over standard adaptive optimizers like Adam, even when training budgets are matched.

The authors modestly note that their current analysis focuses on LLMs and vision classifiers, leaving the performance of Muon on other model families (e.g., diffusion models) as a direction for future work.