To Use or not to Use Muon: How Simplicity Bias in Optimizers Matters

TL;DR

This paper analyzes the biases introduced by the Muon optimizer in deep learning, highlighting potential downsides such as reduced simplicity bias and increased risk of fitting spurious features, despite its training speed advantages.

Contribution

It provides a theoretical analysis of Muon's biases, illustrating how they affect learning trajectories and model solutions, and emphasizes the importance of considering biases in optimizer development.

Findings

Muon removes simplicity bias present in SGD

Muon may struggle to learn common structures across tasks

Muon is more prone to fitting spurious features

Abstract

For a long period of time, Adam has served as the ubiquitous default choice for training deep neural networks. Recently, many new optimizers have been introduced, out of which Muon has perhaps gained the highest popularity due to its superior training speed. While many papers set out to validate the benefits of Muon, our paper investigates the potential downsides stemming from the mechanism driving this speedup. We explore the biases induced when optimizing with Muon, providing theoretical analysis and its consequences to the learning trajectories and solutions learned. While the theory does provide justification for the benefits Muon brings, it also guides our intuition when coming up with a couple of examples where Muon-optimized models have disadvantages. The core problem we emphasize is that Muon optimization removes a simplicity bias that is naturally preserved by older, more thoroughly studied methods like Stochastic Gradient Descent (SGD). We take first steps toward understanding consequences this may have: Muon might struggle to uncover common underlying structure across tasks, and be more prone to fitting spurious features. More broadly, this paper should serve as a reminder: when developing new optimizers, it is essential to consider the biases they introduce, as these biases can fundamentally change a model's behavior -- for better or for worse.