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
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↑ ) flattens these magnitudes and learns modes in parallel, whereas small-batch noise ( σ↑ ) 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=∞
SignGF
L3/2
L1
GF
L1
L1
Table 1: Implicit regularization objective R for small init λ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: L2 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 η
Scaled η
B
42,927.81±17.82
42,927.81±17.82
2B
43,040.07±15.96
42,951.67±9.63
4B
43,181.65±9.73
43,047.80±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 η
Scaled η
256
40.79±0.07
40.79±0.07
512
40.48±0.06
40.71±0.09
1024
40.18±0.02
40.78±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: L2 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 L3/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: L2 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.
Courant Institute School of Mathematics, Computing, and Data Science, New York University · Gatsby Computational Neuroscience Unit, University College London
Hanoi University of Science and Technology Hanoi, Vietnam · Indiana University 107 S. Indiana Ave, Bloomington, IN 47405, USA · Monash University Clayton, VIC 3800, Australia