Technical Summary: DynMuon: A Dynamic Spectral Shaping View of Muon

Problem Statement

Muon has recently emerged as a dominant optimizer for training Large Language Models (LLMs) and transformers. Its core mechanism involves replacing the standard gradient update matrix $M = U\Sigma V^\top$ with its polar factor $UV^\top$, effectively "flattening" the singular values while preserving singular directions. While this improves convergence and stability, the paper identifies a critical limitation: Muon employs a static spectral transformation (specifically, the case where the spectral exponent $p=0$ in the family $U\Sigma^p V^\top$).

The central unresolved question addressed by this work is whether a single, fixed spectral shaping rule remains optimal throughout the entire training process. Prior work has explored fixed variations of spectral shaping (e.g., fixed power-law shapings or suppressing dominant subspaces) but lacks a theoretical model of how the relative influence of different spectral directions evolves as training progresses. The authors hypothesize that the optimal spectral exponent $p$ should adapt dynamically based on the changing landscape of the loss function, the nature of gradient noise, and the training stage.

Methodology

1. Theoretical Framework: Noise-Aware Local Modeling

The authors develop a noise-aware local model to analyze the training dynamics of spectral shaping operations defined as $D(p) := U\Sigma^p V^\top$.

• Signal-Noise Trade-off: The model decomposes the update into local curvature directions (modes). It reveals a fundamental trade-off controlled by the exponent $p$:
  • High $p$ (Positive): Accelerates the contraction of residual signals in high-curvature (strong) modes but suppresses updates in flat modes.
  • Low $p$ (Negative): Increases contraction in low-curvature (flat) modes but simultaneously amplifies stochastic gradient noise in those same directions.
• Stage-Dependent Dynamics:
  • Early Training: Residual signal is concentrated in high-curvature directions. A positive $p$ is beneficial as it prioritizes these directions, accelerating signal contraction with limited noise amplification.
  • Late Training: As training progresses, residual signal in strong modes decays, shifting concentration to flat modes. A mildly negative $p$ becomes advantageous, reallocating update strength to these flat directions which still contain useful signal. However, $p$ must not be too negative, or noise amplification will degrade optimization.

2. DynMuon Algorithm

Based on the theoretical insight that $p$ should transition from positive to mildly negative, the authors propose DynMuon, a dynamic spectral shaping algorithm.

• Scheduling Strategy: DynMuon employs a simple logistic schedule to monotonically decrease the spectral exponent $p$ from a positive early-stage value ($p_{max} = 1$) to a mildly negative late-stage value ($p_{min} = -0.25$).
• Efficient Implementation: Calculating exact SVD for arbitrary $p$ is computationally prohibitive. DynMuon extends the Newton-Schulz (NS) approximation used in Muon:
  • For $p \ge 0.25$, it uses the original Muon update (or standard NS orthogonalization for $p \in [0, 0.25)$).
  • For $p < 0$, it approximates the fractional spectral shaping $A^{p/2}$ (where $A = X_n X_n^\top$) using a second-order Taylor expansion around the identity. This adds only a polynomial correction term to the standard NS computation, maintaining the same asymptotic complexity as Muon.
• Stability Anchoring: To ensure stability when $p$ is positive and large, the algorithm anchors the positive regime to stable operators, avoiding unreliable Taylor approximations in that region.

Key Contributions

1. Dynamic Spectral Perspective: Reframing Muon-style optimization from a fixed spectral operation to an adaptive problem of choosing a spectral exponent $p$ that evolves with training dynamics.
2. Noise-Aware Local Model: Developing a theoretical model that explains the stage-dependent preference for $p$ through a trade-off between residual-signal reduction and stochastic-noise amplification across local curvature directions.
3. Discovery of Stage-Dependent Regimes: Uncovering that positive exponents benefit early training (emphasizing strong modes), while mildly negative exponents (previously overlooked) improve late-stage optimization by emphasizing flat modes that retain useful residual signal.
4. DynMuon Algorithm: Proposing an efficient, dynamic spectral shaping algorithm that schedules $p$ from positive to mildly negative, achieving consistent improvements over fixed Muon with negligible computational overhead.

Experimental Results

The authors evaluated DynMuon across various model sizes (127M to 1.1B parameters), architectures (GPT-style and Qwen-style), and training settings (FineWeb and FineWeb-Edu datasets).

• Performance: DynMuon consistently achieves lower validation loss than the standard Muon ($p=0$) and AdamW baselines.
• Efficiency: To reach the same target validation loss (defined as the loss Muon reaches at 80% of training), DynMuon requires 10.6% to 26.5% fewer training steps.
• Overhead: The per-step time cost of DynMuon is only 1.003–1.025× that of Muon, confirming the efficiency of the Taylor-approximated spectral shaping.
• Robustness:
  • Architecture: Benefits transfer to Qwen-style architectures.
  • Training Budget: Improvements hold across token budgets from 2.5B to 20B.
  • Learning Rate: DynMuon shows lower sensitivity to learning rate choices compared to Muon.
  • Loss Objectives: The benefit of mild negative shaping persists across a family of CE-Brier loss interpolations, suggesting the mechanism is robust to different probability-space losses.
• Ablations: Experiments confirm that a smooth logistic transition outperforms abrupt switching or fixed negative schedules. The optimal $p_{min}$ is found to be $-0.25$; more aggressive negative values (e.g., $-0.5$) degrade performance due to noise amplification.

Significance and Claims

The paper claims that DynMuon demonstrates that the "best" spectral bias is not static but evolves during training. By dynamically shifting emphasis from high-curvature directions to flat directions, DynMuon captures optimization dynamics that fixed transformations like Muon miss.

The authors position this work as a step toward more adaptive optimizers. They modestly note that while their logistic schedule outperforms fixed choices, it may not be the final word; the optimal $p$ might depend on real-time optimization statistics. They suggest future work could focus on online estimation of $p$, potentially using more expensive signals applied infrequently, given that the useful range of $p$ appears relatively small. The method is presented as a general-purpose optimization improvement for LLM training that enhances efficiency and performance without significant computational cost.