From Parallel to Sequential: Small-Batch Noise Restores Simplicity Bias of Muon

summary

Video file (mp4)

The gist

The gist The behavior of Muon changes qualitatively once gradients become stochastic, restoring sequential feature learning through small-batch noise, which provides a mechanism for simplicity bias

In short

The study investigates how adding small-batch noise fundamentally changes Muon's optimization behavior, which usually learns features in parallel. Stochastic smoothing of the spectral descent operator restores sequential feature learning by making updates proportional to actual gradients when noise is present. This mechanism introduces a depth-dependent implicit regularization that favors low-rank intermediate predictors.

Key concepts

Muon Optimizer
A state-of-the-art optimizer known for learning features in parallel during full-batch training. The paper examines how its behavior shifts when subjected to small-batch noise, revealing a transition from parallel to sequential feature learning dynamics.
Stochastic Smoothing
The process where small noise is added to the spectral descent operator underlying Muon. When gradients are small compared to this noise, the update approximates gradient flow dynamics, allowing the optimizer to recover sequential feature learning instead of parallel updates.
Simplicity Bias
A tendency in optimizers to favor simpler solutions, such as those with lower rank or fewer features. The paper shows that stochastic smoothing induces a depth-dependent regularization that favors low-rank intermediate predictors, linking batch size directly to this learned complexity.

Terminology used across episodes

This episode discusses

The paper

From Parallel to Sequential: Small-Batch Noise Restores Simplicity Bias of Muon · Read on arXiv

Tom Jacobs, Rebekka Burkholz

CISPA Helmholtz Center for Information Security, Germany

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.

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "From Parallel to Sequential".

Jane: The gist The behavior of Muon changes qualitatively once gradients become stochastic, restoring sequential feature learning through small-batch noise,

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So, looking back at the paper "From Parallel to Sequential: Small-Batch Noise Restores Simplicity Bias of Muon," it seems the key takeaway is that the way we look at optimization geometry alone isn't enough to understand why some deep learning models are so successful. It’s not just about the deterministic update rule itself.

Jane: The authors argue that stochastic smoothing, specifically from small batch noise, fundamentally alters this behavior by restoring sequential feature learning through mode-dependent growth and inducing a limiting implicit regularization that approaches L1 with increasing depth. This is a mechanism for simplicity bias that was absent in full-batch Muon but is present when you smooth the spectral sign operator.

Lu: It really points to an interaction between optimizer geometry, parameterization, and stochasticity where the same deterministic update rule can exhibit different simplicity biases depending on the training regime, depending on whether you have small batch noise or not.

Meng: From an engineering standpoint, this means that we might be able to tune our training regimes by controlling the batch size to explicitly favor lower-rank intermediate predictors when we are building those large vision and language models.

Tom: Exactly. The results establish a direct link between batch size and learned complexity in vision and language models, showing that larger batches tend to reduce validation accuracy because they suppress this sequential learning effect.

Jane: So what the paper ultimately suggests is that if we want those simpler, lower-rank representations that might be more robust, controlling the training noise—specifically using smaller batches—is a way to steer the optimization dynamics towards those features.

Conclusion: Tom: So we’ve seen how adding small batch noise changes things for Muon, and now we’re looking at what the authors actually concluded about this whole paper, "From Parallel to Sequential: Small-Batch Noise Restores Simplicity Bias of Muon."

Jane: It really boils down to this idea that the way the optimizer behaves isn't just set by its math on paper, but also by how you actually train it. The authors are saying that if you introduce small batch noise, it restores a kind of sequential feature learning we didn't see before.

Lu: That sequencing comes from this stochastic smoothing of the spectral descent operator. When the gradients are small compared to the noise level, the expected update starts acting like a true gradient flow again.

Meng: So what does that actually mean for us when we’re trying to deploy these models at scale? Does it just mean we have to use smaller batch sizes to get better results?

Tom: That’s exactly where they land. They show this mechanism induces an implicit regularization, something that depends on how deep the model is. It pushes the learned features toward a simpler structure, specifically one that favors low-rank intermediate predictors.

Jane: And it's interesting because this effect isn't just a weird quirk of small batches; it’s tied to the depth of the neural network. They found this regularization term changes depending on whether you look at the dynamics at shallow versus deep layers.

Lu: The authors showed that under heavy noise, even with infinite variance Cauchy noise, we still get this gradient-proportional flow update. That's a strong result because it means this mechanism is quite robust to different kinds of noise.

Tom: Robustness is important for us because we can’t always control the exact type of noise in production. But there’s a trade-off here, too. They mentioned that larger batches tend to reduce validation accuracy, which suggests that if you want those simpler representations, you might actually need to keep your batch sizes small.

Jane: So they aren't just giving us a neat trick for training; they’re pointing out that the very way we set up our optimization—the noise we use—fundamentally shifts what the AI learns about the data structure.

Lu: Yeah, it suggests that optimizer geometry and stochasticity interact in a way that dictates complexity, which is something we need to keep investigating for representation learning.

More episodes

← Home