To Use or not to Use Muon: How Simplicity Bias in Optimizers Matters
This work highlights a critical consideration for optimizer development in deep learning, revealing that biases introduced by new optimizers like Muon can fundamentally alter model behavior, potentially leading to negative consequences in practical applications.
The paper investigates potential downsides of the Muon optimizer, showing that while it provides superior training speed, it removes a simplicity bias preserved by methods like SGD, which can cause Muon-optimized models to struggle with uncovering common structure and be more prone to fitting spurious features.
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.