Small batch noise brings simpler learning to powerful new optimizer
From Parallel to Sequential: Small-Batch Noise Restores Simplicity Bias of Muon
Machine Learning
Summary
Deep learning often works well because it learns one simple feature at a time, a process called simplicity bias. A new optimizer called Muon, which uses spectral descent, normally learns features all at once when given full data batches, which might seem confusing. The authors show that when training uses small batches that add a bit of randomness, Muon goes back to learning features one after another, restoring simplicity bias. They explain this change mathematically and confirm it with experiments on different types of models.
What this means in practice
- •For machine learning engineers: Improve training strategies for deep learning models by using small batch sizes to recover simple feature learning with the Muon optimizer.
- •For ai model deployment teams: Use noise-induced sequential learning insights to enhance performance and reliability in vision transformers and large language models.
Authors
Tom Jacobs, Rebekka Burkholz
Abstract
Simplicity bias, which can arise from sequential feature learning during training, has been proposed as a general explanation for the success of deep learning. An important driver of this phenomenon is saddle-to-saddle behavior, in which the gradient flow trajectory remains close to a sequence of saddle manifolds while learning one feature at a time. However, Muon, a state-of-the-art optimizer based on spectral descent, does not exhibit such dynamics in the full-batch setting and instead learns features in parallel. How is it then that Muon is so successful in deep learning settings? We show that stochastic gradient noise fundamentally changes this behavior: small-batch noise restores sequential feature learning, providing a mechanism for simplicity bias that is absent from full-batch Muon. We explain this transition through stochastic smoothing of the spectral descent operator underlying Muon. When gradients are small relative to the noise scale, the expected spectral update becomes approximately proportional to the actual gradient, which restores gradient-flow-like dynamics. We formalize this mechanism for deep linear models by deriving the effective continuous dynamics of noisy spectral descent. Experiments on low-rank matrix reconstruction, vision transformers, and large language models (LLMs) support these findings.