Muon Learns More Robust and Transferable Features than Adam

2026-06-08

AI Summary

This work systematically investigates the robustness and transferability advantages of the Muon optimizer over Adam and SGD in representation learning. By evaluating pretrained models on corrupted images and text—combined with layer-wise probing, logit margin analysis, linear classification on downstream tasks, and full-model fine-tuning—the study introduces effective rank metrics and theoretical analysis of multi-component feature structure to rigorously establish Muon’s superiority from the perspectives of robustness and transferability. Experiments across diverse architectures demonstrate that Muon consistently learns features that are more discriminative, exhibit higher diversity, and generalize better, thereby significantly enhancing both model robustness and transfer performance.

Abstract

Muon has recently emerged as a state-of-the-art optimizer for pretraining Large Language Models (LLMs) and vision classifiers. Despite its efficiency advantage over Adam and SGD, the feature-learning advantage of Muon remains unclear. This paper investigates Muon's feature-learning advantage through the lens of robustness and transferability. First, by evaluating pretrained models on corrupted images and texts, we show that features learned by Muon are consistently more robust than those learned by Adam and SGD across different architectures, including transformers and Convolutional Neural Networks (CNNs). Using trained layer-wise probes, we further show that this robustness advantage is reflected in larger logit margins across layers. Second, by training linear classifiers or fine-tuning full models from pretrained parameters on downstream tasks, we demonstrate that Muon-learned features transfer more effectively than those learned by Adam and SGD. This transferability advantage is further supported by the diversity of hidden states across layers, as measured by effective rank. Finally, in a representative classification problem with multi-component features, we prove that Muon attains larger margins and higher effective rank than Adam and SGD, providing theoretical support for our empirical findings.

Problem

• robustness
• transferability
• feature learning
• optimizers
• large language models

Innovation

• Muon optimizer
• feature robustness
• transferability
• logit margin
• effective rank