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

Tom Jacobs; Rebekka Burkholz

cs.LGOct 12, 2026

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

Authors: Tom Jacobs, Rebekka Burkholz

Organizations: CISPA Helmholtz Center for Information Security, Germany

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.

Figures & tables

Panel of Figure 1 — Figure 1: Batch size $B$ controls the simplicity bias of Muon. Gradient descent preserves mode-dependent gradient magnitude ( $|G|$ ), producing sequential feature activation from small initialization. In contrast, full-batch Muon ( $B\uparrow$ ) flattens these magnitudes and learns modes in parallel, whereas small-batch noise ( $\sigma\uparrow$ ) smooths the spectral sign operator and restores initialization-dependent sequential learning.

Figure 2: For low-rank matrix reconstruction with whitened data, gradient flow learns singular directions sequentially, while full-batch Muon (without momentum) learns them in parallel. Small-batch noise restores sequential feature learning, thereby giving a mechanism for simplicity bias in Muon.

Figure 3: Simulation of the smoothened flows for depth

L=2

and

20

. For higher depth we observe sequential feature learning.

	$L=2$	$L=\infty$
SignGF	$L_{3/2}$	$L_{1}$
GF	$L_{1}$	$L_{1}$

Table 1: Implicit regularization objective

R

for small init

\lambda_{0}>0

, depending on depth

L

Figure 4: Noise changes the solution selected by sign descent in a diagonal linear network of depth

L=20

. We compare full-batch updates, small-batch updates, and full-batch updates with Gaussian or Cauchy noise. Left: training loss. Middle:

L_{2}

reconstruction error. Right: final predictor coefficients, with first 5 non-zero. Small-batch and injected-noise runs recover the same sparse target and thus largely overlap, whereas deterministic full-batch training reaches low loss without recovery.

Figure 5: Final validation accuracy and effective rank of a DeiT-small model trained on ImageNet. Small batches lead to lower effective rank indicating the presence of simplicity bias.

Batch size	Fixed $\eta$	Scaled $\eta$
$B$	$42{,}927.81\pm 17.82$	$42{,}927.81\pm 17.82$
$2B$	$43{,}040.07\pm 15.96$	$42{,}951.67\pm 9.63$
$4B$	$43{,}181.65\pm 9.73$	$43{,}047.80\pm 10.34$

Table 2: Total effective ranks for a 162M-parameter Llama model trained on 3.2B tokens from FineWeb-Edu. Larger batches (

B=256

) yield higher effective ranks under both fixed and linearly scaled learning rates.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Simulated flows for all

3

settings at depth

L=2

: gradient, spectral, smoothened flow.

Figure 7: Gradient descent, Muon with small and large batches. The sequential learning phenomenon persists with momentum included.

Figure 8: Matrix factorization varying the learning rate for small batch Muon. We observe that the trajectories remain similar to each other accounting for discretization time.

Batch size	Fixed $\eta$	Scaled $\eta$
$256$	$40.79\pm 0.07$	$40.79\pm 0.07$
$512$	$40.48\pm 0.06$	$40.71\pm 0.09$
$1024$	$40.18\pm 0.02$	$40.78\pm 0.05$

Appendix

Table 3: Validation accuracy for a 162M-parameter Llama model trained on 3.2B tokens from FineWeb-Edu (

B=256

). Results report means and approximate 95% confidence intervals over three seeds.

Figure 9: Effective rank of

QK

for the DeiT-small model trained on ImageNet.

Figure 10: Final validation accuracy and effective rank of a small DeiT model trained on CIFAR 10. Small batches lead to lower rank indicating the presence of simplicity bias.

Figure 11: Final validation accuracy and effective rank of a small DeiT model trained on CIFAR-100. Small batches lead to lower rank indicating the presence of simplicity bias.

Figure 12: Noise changes the solution selected by sign descent in a diagonal linear network of depth

L=2

. We compare full-batch updates, small-batch updates, and full-batch updates with Gaussian or Cauchy noise. Left: training loss. Middle:

L_{2}

reconstruction error. Right: final predictor coefficients, with first 5 non-zero. Small-batch training and both forms of injected noise recover the sparse target better but not full which is in line with the

L_{3/2}-

bias, whereas deterministic full-batch training achieves low training loss without recovery.

Figure 13: Noise does not change the solution selected by gradient descent in a diagonal linear network of depth

L=2

for small initialization. We compare full-batch updates, small-batch updates, and full-batch updates with Gaussian or Cauchy noise. Left: training loss. Middle:

L_{2}

reconstruction error. Right: final predictor coefficients, with first 5 non-zero. All recover the ground truth solution except the injected heavy tail noise due to too large variance.

Muon Optimizer

Muon optimizer uses momentum and matrix orthogonalization to shape neural-network updates, especially in large language-model training.

See all 146 papers

Feb 28, 2026

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

While Adam has long been the ubiquitous default optimizer for deep neural networks, Muon has recently seen rapid adoption due to its superior training speed. Although much of the literature focuses on validating the benefits of Muon, our work investigates the potential downsides of the mechanism driving this speedup. On the theoretical front, we analyze the learning dynamics of simplified Muon on deep linear networks and linear attention. Our analysis reveals that Muon gains speed by avoiding saddle points, but does so at the expense of the simplicity bias characteristic of Gradient Descent (GD), where the complexity of the functional solution learned grows sequentially. Experiments demonstrate the consequences of losing the simplicity bias, showing that Muon struggles to uncover common underlying structure across tasks and may be prone to fitting spurious features. More broadly, this paper serves as a reminder that faster optimization is rarely a free lunch; improvements in optimization can come at the cost of changes in the inductive biases that shape generalization.

Sara Dragutinović, Yedi Zhang, Rajesh Ranganath

Courant Institute School of Mathematics, Computing, and Data Science, New York University · Gatsby Computational Neuroscience Unit, University College London

Neural Network Generalization Muon Optimizer Deep Linear Networks Implicit Bias Neural Network Training Dynamics

Ask PDF

May 13, 2026

Spectral Flattening Is All Muon Needs: How Orthogonalization Controls Learning Rate and Convergence

Muon orthogonalizes the momentum buffer before each update, replacing its singular values with ones via Newton-Schulz iterations. This simple change lets Muon tolerate far larger learning rates and converge faster than other optimizers, but why? We show that the mechanism is spectral flattening, and develop two results around it. First, we prove that Muon's maximal stable step size scales with the average singular value of the gradient rather than the largest, which bottlenecks standard gradient descent. Second, we recast Muon as a preconditioned gradient method and show, under a Kronecker-factored curvature model, that it improves the effective convergence factor, with the improvement controlled by the spectrum of the gradient covariance. Extensive experiments validate both results: Muon remains stable at learning rates that cause SGD to diverge within the first few iterations, and reaches accuracy milestones several epochs earlier even at identical step sizes. Taken together, our results offer a principled, geometric explanation for Muon's empirical success.

Tien-Phat Nguyen, Truong Nguyen, Minh-Phuc Truong +3

Hanoi University of Science and Technology Hanoi, Vietnam · Indiana University 107 S. Indiana Ave, Bloomington, IN 47405, USA · Monash University Clayton, VIC 3800, Australia

Momentum Methods Gradient Descent Muon Optimizer Spectral Methods Stochastic Gradient Descent

Ask PDF

Oct 22, 2026

When and How Does Muon Precondition? The Role of Structured Gradient Noise

Muon has become a popular optimizer for large language model training, yet its theoretical foundations remain poorly understood. In this work, we study how gradient noise geometry shapes Muon's preconditioning benefits. We first show, theoretically and empirically, that under the classical isotropic noise assumption, Muon offers no acceleration beyond a rescaling of the learning rate. This motivates us to investigate structured noise arising from mini-batch sampling in associative memory with heterogeneous task frequencies. We establish an exact balancing threshold for two tasks and extend the analysis to multiple tasks, revealing how task coverage and gradient reliability jointly shape learning. Our analysis shows that Muon can soften the dependence on task frequency from linear to square-root, helping balance the learning of frequent and rare tasks. We further identify rotation-induced negative transfer in SignSGD that Muon avoids, and explain how momentum reduces sampling noise while tracking changing gradients. Controlled numerical experiments validate these mechanisms, and a preliminary LLM pretraining study illustrates the practical promise of fast-slow momentum.

Xiaoran Chen, Binghui Li, Liwei Wang +1

School of EECS, Peking University · Center for Machine Learning Research, Peking University · Center for Data Science, Peking University +2

Muon Optimizer

Ask PDF

Abstract

Muon Optimizer

Explore similar work

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

Spectral Flattening Is All Muon Needs: How Orthogonalization Controls Learning Rate and Convergence

When and How Does Muon Precondition? The Role of Structured Gradient Noise