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

arXiv:2610.16868 · cs.LG · Submitted 2026-10-12 · Read on arXiv

Listen

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.

Tom Jacobs, Rebekka Burkholz

CISPA Helmholtz Center for Information Security, Germany

cs.LG

Submitted: 2026-10-12

Updated: 2026-10-12

Comments: 27 pages, 13 figures

Code: https://github.com/nelaturuharsha/TurboPrune

License: http://creativecommons.org/licenses/by/4.0/

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

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.

Summary

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 absent from full-batch Muon

How it works

The paper investigates how small-batch noise fundamentally alters the behavior of the state-of-the-art optimizer Muon, which typically learns features in parallel in full-batch settings. The core mechanism is that stochastic smoothing of the spectral descent operator underlying Muon restores sequential feature learning by making the update approximately proportional to the actual gradient when gradients are small relative to the noise scale

The transition from parallel learning in full-batch Muon to sequential feature learning in small-batch spectral descent is explained through stochastic smoothing of the spectral descent operator. 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

Key Findings and Mechanisms

The authors derive effective continuous dynamics for deep linear models by showing how stochastic smoothing changes Muon's behavior. They demonstrate that small-batch spectral descent can recover a clear ordering in which stronger features are learned before weaker ones, favoring low-rank intermediate predictors

The expected dynamics under noisy gradient for the i-th singular value satisfy dθi/dτ = pθi EN(0,σ2Lτ) [Sign(λi − θi + ξτ)]. For large noise strength σ, the leading-order dynamics recover the gradient-proportional flow dθi/dt = pθi (λi − θi)

Implicit Regularization and Depth Dependence

The stochastic smoothing mechanism induces a non-Euclidean implicit regularization that depends on depth. For the diagonal restriction of the dynamics, the induced regularization is of L1+1/L-type at depth L, whereas gradient flow retains an L1 bias

The resulting mirror potential remains convex at arbitrary depth, contrasting with gradient flow whose corresponding potential loses convexity beyond depth two in the balanced limit. This suggests that stochastic smoothing restores sequential feature learning through mode-dependent growth and induces a limiting implicit regularization approaching L1 with increasing depth

Experimental Validation

Experiments on low-rank matrix reconstruction, vision transformers, and large language models (LLMs) support these findings. In the matrix factorization example, full-batch Muon learns singular modes in parallel, whereas decreasing the batch size produces sequential dynamics

Furthermore, experiments in vision and language models show that smaller batches yield lower effective rank, consistent with a stronger simplicity bias. This links batch size directly to learned complexity in these architectures

Robustness to Noise Types

The smoothing mechanism extends to heavy-tailed noise, even with infinite variance as in Example 4.7. The expected update remains gradient-proportional flow even with infinite-variance Cauchy noise because the spectral sign update is bounded under such noise

This contrasts with Levy-driven dynamics that can arise for SGD under heavy-tailed gradient noise, highlighting Muon's resilience in these regimes

Conclusion

The paper concludes that the simplicity bias of an optimizer cannot generally be inferred from its deterministic optimization geometry alone. Stochastic smoothing affects both transient dynamics and the selected solution, inducing a depth-dependent L1+1/L-type implicit regularization that approaches L1 with increasing depth

The results point to an interaction between optimizer geometry, parameterization, and stochasticity where the same deterministic update rule can exhibit qualitatively different simplicity biases depending on the training regime. The theory suggests that batch size can determine whether Muon exhibits the sequential feature learning that favors low-rank intermediate predictors

The paper's findings establish a link between batch size and learned complexity in vision and language models, showing that larger batches tend to reduce validation accuracy. The overall conclusion is that stochastic smoothing restores sequential feature learning through mode-dependent growth and induces a limiting implicit regularization approaching L1 with increasing depth

The paper's findings establish a link between batch size and learned complexity in vision and language models, showing that larger batches tend to reduce validation accuracy.

Improvements for AI systems

textbfSmall-Batch Noise Restoration of Sequential Learning Bias: Implement Stochastic Smoothing in Muon Optimizer for Feature Ordering Stability. By smoothing the nonlinear sign operation underlying Muon, the system can restore initialization-dependent sequential learning, ensuring that dominant features are learned first, which is crucial for low-rank intermediate predictors.
textbfDepth-Amplified Simplicity Bias: Utilize Large Network Depth to Strengthen Feature Separation. As shown in Corollary 4.3, increasing depth amplifies the separation between feature-learning times in the smoothened spectral flow dynamics, leading to more pronounced sequential feature learning and stronger sparsity bias of the smoothened dynamics.
textbf Learned Complexity Control via Batch Size: Tune Batch Size for Desired Model Simplicity. Experiments show that smaller batches yield lower effective rank in both vision and language models, allowing for better generalization by favoring simpler learned representations, a phenomenon directly linked to the batch size parameter.
textbf Implicit Regularization for Model Selection: Leverage Depth-Dependent L1+1/L-type Regularization. The smoothed spectral dynamics induce an implicit regularization of L1+1/L-type at depth L, which can be exploited to select sparse or low-rank solutions during training, as the potential remains convex for every finite depth.
textbf Robustness to Heavy-Tailed Gradients: Enhance Training Stability with Noise Resilience. The smoothing mechanism extends to heavy-tailed noise, even with infinite variance, allowing the system to maintain sequential learning dynamics under noise distributions that substantially alter SGD dynamics.

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.