PaperlayerView exact revision on arXiv

Source-linked AI summary

DynMuon: A Dynamic Spectral Shaping View of Muon

Fangzhou Wu, Rikhav Shah, Sandeep Silwal, Qiuyi Zhang

arXiv:2605.17109v3cs.LGcs.AI

TL;DR

Fixed spectral shaping may not remain optimal as training dynamics evolve, motivating a method that adapts the spectral exponent. The paper develops a noise-aware local model and DynMuon, finding that positive exponents help early while mildly negative exponents improve late-stage optimization.

  • Problem

    Prior spectral-shaping methods use fixed transformations, leaving unresolved whether one shaping rule remains desirable as training dynamics evolve.

  • Method

    The paper generalizes Muon to power-law spectral shaping and develops a noise-aware local model relating the spectral exponent to residual-signal reduction, curvature, and stochastic-noise amplification.

  • Results

    DynMuon's empirical observations support the model's predictions: positive spectral exponents help early training, whereas mildly negative exponents improve late-stage optimization by emphasizing flat directions retaining useful residual signal.

  • Takeaways & Limitations

    Dynamically scheduling spectral shape lets matrix-valued updates change their emphasis across training stages, unlike fixed transformations such as Muon.

  • Takeaways & Limitations

    The dynamics cannot be fully anticipated in advance, motivating future methods that select p online from observed optimization statistics.

Abstract

from arXiv · show

In recent years, Muon has emerged as the dominant method for training large language models, and transformers more broadly. The essential difference, when compared to standard gradient descent methods, is to replace the usual update matrix $M=UΣV^\top$ with its polar factor $UV^\top$. In this work, we consider a class of Muon-like updates, where we replace the update $M$ with $UΣ^p V^\top$ for some parameter $p$. We call this a "spectral-shaping" operation, and develop a theory of how to pick $p$ which depends on (a) local curvature of the loss function, (b) noise stemming from stochastic gradients and label noise, and (c) training stage. Our theory and experimentation reveal a previously overlooked behavior: positive $p$ helps early by emphasizing high-curvature directions and accelerating signal contraction, while mildly negative $p$ helps later by reallocating update strength toward low-curvature directions that still contain useful training signals. Building on the insight, we propose DynMuon, an efficient dynamic spectral shaping method that schedules $p$ from positive to mildly negative over training. Extensive experiments across model sizes, architectures, and training settings show that DynMuon consistently achieves lower validation loss than Muon, while requiring 10.6-26.5% fewer steps to reach the same target loss. Our code is available at https://github.com/fzwark/DynMuon.

1 Introduction

The paper reframes Muon as a power-law family of spectral-shaping updates and asks how the spectral exponent should adapt as training evolves. Its model motivates DynMuon, which shifts emphasis from strong curvature directions early to flatter directions later.

  • Muon replaces each momentum-averaged gradient matrix with its polar factor, preserving singular directions while flattening singular values.
  • Prior spectral-shaping methods mainly use fixed transformations, leaving adaptation across changing training dynamics unresolved.
  • The paper generalizes Muon into a power-law family parameterized by the spectral exponent p.
  • Positive p helps early by emphasizing high-curvature directions, whereas mildly negative p helps later by refocusing updates on useful flat directions.
  • DynMuon uses a decreasing logistic schedule for p and extends Newton–Schulz approximation to retain Muon's per-step cost without full SVD.
  • Unlike scalar optimization schedules, DynMuon changes the spectral shape of matrix-valued updates across training stages.

2 Motivating Dynamic Spectral Shaping with an Idealized Model

The paper develops a noise-aware local model that treats spectral shaping as curvature-dependent reweighting of matrix updates. The model tracks residual signal and stochastic gradient noise under local curvature approximations.

  • The analysis studies how exponent p controls optimization using a local quadratic and gradient-noise approximation.
  • For a weight matrix, the stochastic gradient serves as the update matrix, which spectral shaping transforms using exponent p.
  • The model defines residual signal as the distance from current parameters to a nearby local minimizer and analyzes its evolution under shaped updates.
  • Effective local curvature is represented by modes whose eigenvalues distinguish strong directions from flat directions.
  • Stochastic gradients are decomposed into population gradients and zero-mean noise under unbiased SGD assumptions.
  • The analysis uses a curvature-aligned surrogate motivated by gradient second moments as Fisher-type local-curvature proxies.
  • The gradient-curvature alignment is an empirical assumption supporting the model's curvature-dependent reweighting interpretation.

Appendix Figure 8. Substituting this approximation into Equation (2), absorbing α

The mode-wise analysis balances residual-signal contraction against stochastic-noise amplification across curvature directions. Experiments support a late shift toward useful flat-mode signal and an early advantage for emphasizing strong modes.

  • Mode-Wise Signal–Noise Tradeoff: Larger p contracts residual signal faster in strong, high-curvature modes, while smaller p shifts contraction toward flat modes.Decreasing p also amplifies stochastic noise, especially in flat modes, creating a signal–noise tradeoff.
  • Why Residual Signal Concentrates in Flat Modes Late in Training: Residual signal decays earlier in strong modes and remains more substantial in flat modes as training progresses.For comparable projected gradient magnitudes, smaller curvature corresponds to larger residual signal.
  • Validating Residual-Signal Concentration in Flat Modes: Πt becomes mostly negative after roughly 500 steps, supporting a late-stage shift of residual signal toward flat modes.Πt compares median log residual-signal energy in the highest- and lowest-curvature mode groups.
  • Validating Flat-Mode Signal After Accounting for Noise: Flat modes retain positive noise-adjusted signal and a positive signal–noise advantage, matching the condition for late-stage mildly negative shaping.The noise-adjusted signal measures residual signal relative to noise, while the advantage compares flat and strong modes.

3 DynMuon: Dynamic Spectral Shaping

DynMuon schedules spectral shaping from a positive early exponent to a mildly negative late exponent. It uses a logistic schedule and efficient approximations designed to retain Muon-like computational complexity.

  • DynMuon: Dynamic Spectral Shaping: DynMuon monotonically decreases p from a positive early-stage value to a mildly negative late-stage value while maintaining computational efficiency.The schedule adapts spectral emphasis as training dynamics evolve.
  • Logistic Scheduling of the Spectral Exponent: A logistic schedule smoothly transitions p over training without requiring additional forward and backward passes.τ controls the transition point, w controls transition width, and the reported settings use pmax = 1 and pmin = −0.25.
  • Training Results: DynMuon achieves the lowest validation loss across three model scales trained on 10B tokens.Figure 3 compares validation-loss trajectories across the three scales.
  • Efficient Updates: DynMuon approximates varying-power spectral shaping through Newton–Schulz computations plus a polynomial correction, avoiding full SVD.The added correction has cost O(m^2n + m^3), giving the same asymptotic complexity as Newton–Schulz-based Muon.
  • Stable Anchoring for Positive Exponents: For sufficiently large positive p, DynMuon uses stable anchored operators; for mildly negative p, it applies an efficient spectral approximation.The implementation uses the original update for p_t ≥ 0.25 and standard Newton–Schulz orthogonalization for p_t ∈ [0, 0.25).

4 Evaluation

DynMuon is evaluated across model architectures, training budgets, corpora, learning rates, scheduling choices, and spectral-shaping implementations. It consistently improves validation loss over Muon, reaches target loss sooner, and closely matches exact SVD at much lower cost.

  • Main Results: DynMuon reaches Muon's target validation loss 10.6–26.5% earlier across GPT-style model scales and token budgets.The target is Muon's validation loss at 80% of training.
  • Main Results: DynMuon has a 1.003–1.025× per-step time ratio relative to Muon, indicating negligible additional runtime cost.
  • Robustness: DynMuon achieves lower validation loss than Muon across Qwen-style and GPT-style architectures, token budgets from 2.5B to 20B, and tested learning rates from 0.003 to 0.04.The learning-rate sweep also shows a flatter curve near DynMuon's optimum.
  • Robustness: DynMuon outperforms Muon on FineWeb-Edu, with the advantage becoming more pronounced late in training.
  • Ablations: Mildly negative pmin values perform best, with pmin = −0.25 outperforming Muon while pmin = −0.5 performs worse.
  • Implementation: DynMuon's spectral-shaping approximation closely matches exact SVD, while exact SVD is roughly 3× slower.
  • Ablations: The default logistic schedule outperforms abrupt and fixed-negative schedules, while performance remains reasonably robust to τ and w.The best observed parameters are around τ = 0.04 and w = 0.04.
  • Additional Results: DynMuon consistently outperforms Muon across random seeds with low variance and also improves over NorMuon; mildly negative exponents generalize across CE–Brier interpolations.Overly negative exponents degrade performance across those probability-space losses.

5 Conclusions and Future Directions

The paper concludes that dynamically scheduling the spectral exponent improves optimization over fixed choices and Muon's p = 0 update, while leaving room for online adaptation.

  • Conclusions: DynMuon schedules p from positive to mildly negative values, and experiments show lower validation loss than Muon across model scales, architectures, and training settings.
  • Future Directions: The best exponent may depend on training-time dynamics that cannot be fully anticipated in advance.The authors suggest occasionally re-estimating p online rather than at every training step.

A.1 Validation Setup

The validation setup tests the model's predictions in both last-layer-only and full-model GPT-style training on FineWeb.

  • Validation Setup: Experiments use a GPT-style model with hidden dimension 768, 12 layers, and 6 attention heads trained on FineWeb.
  • Validation Setup: The study compares last-layer-only training for 6000 steps with full-model training for 10000 steps.Last-layer-only training freezes all parameters except one selected matrix-valued parameter in the final Transformer block.

A.2 Empirical Modes and Proxies

The empirical analysis approximates curvature, noise, and residual signal along retained singular directions of Muon updates using Hessian-vector products, mini-batch variation, and fixed probe batches.

  • Empirical Modes: The analysis retains the top k = 256 singular directions of the Muon update as empirical modes representing curvature directions.Each mode is represented as B_i,t := u_i,t v_i,t^⊤.
  • Curvature Proxy: Hessian-vector products estimate the local curvature scale along each empirical mode.
  • Noise Proxy: Noise is estimated from gradient-projection variation across n_b = 32 independent mini-batches using a sample-variance proxy.
  • Residual Signal Proxy: Residual signal energy is estimated from gradient projections on a fixed probe set of 8 mini-batches.Averaging over the fixed probe set provides a low-variance proxy for the population-gradient projection, using g_i,t ≈ κ_t h_i,t δ_i,t.
  • Validation: The median log ĥ_i,t changes only mildly over steps 4k–6k, supporting approximately fixed effective local curvature over short windows.

B Additional Analysis in Section 2

Additional analyses support the model’s curvature and noise assumptions and show that the preferred spectral exponent depends on training stage and gradient-noise level. The efficient implementation applies a scheduled exponent while preserving Muon’s matrix-update structure.

  • Median effective directional curvature changes only mildly over steps 4k–6k, supporting an approximately fixed local curvature over short windows.
  • Gradient second moments remain strongly positively correlated with effective local curvature throughout training, supporting the curvature-alignment approximation.
  • Smaller batch sizes favor mildly negative exponents closer to 0, whereas larger batches can favor more negative shaping such as p = −0.5.
  • The fitted noise–curvature power law remains stable around β_t = 1.4 with consistently high R^2, reliably describing mode-wise noise structure.
  • Positive p improves early full-model training, while a positive-to-negative schedule achieves the lowest validation loss and avoids degradation from overly negative or reversed schedules.
  • DynMuon schedules p through logistic shaping, using Newton–Schulz for nonnegative exponents and Fast–Spectral for negative exponents.

D Experimental Details

The experiments cover GPT-style and Qwen-style decoder-only Transformers across multiple model scales, datasets, token budgets, and optimizer baselines. Additional analyses examine robustness across objectives and report consistently low seed variability.

  • The study evaluates GPT-style models at three scales and a Qwen-style model, using distinct decoder-only Transformer configurations.
  • DynMuon’s mean validation loss consistently outperforms Muon across three random seeds, with very small one-standard-deviation bands.
  • DynMuon outperforms both Muon and NorMuon in validation loss on the 127M model comparison.
  • Comparisons include DynMuon, Muon, and AdamW, with AdamW also used for scalar parameters outside the main matrix groups.
  • Main training uses FineWeb with a 10B-token budget, while data-scale experiments vary budgets from 2.5B to 20B tokens.

E Additional Experiments

Additional experiments test DynMuon against training randomness and a recent Muon variant. Across both analyses, DynMuon retains a validation-loss advantage.

  • DynMuon consistently outperforms Muon across three seeds, with very small across-seed variance indicating robustness to training randomness.
  • DynMuon consistently achieves lower validation loss than both Muon and NorMuon in a controlled comparison.

F Discussion on Robustness Across Loss Objectives

The paper tests whether late-stage mild-negative shaping remains useful when the loss objective changes and situates the method within broader dynamic optimization and spectral-shaping work. Mild negative exponents remain beneficial across the tested CE–Brier family.

  • The preferred spectral exponent may vary with the loss objective because losses induce different gradient, residual-signal, and noise structures across spectral directions.
  • The experiments interpolate between cross-entropy and the Brier score using λ, systematically varying how prediction errors are weighted.
  • Validation cross-entropy is used as the common evaluation metric because objective values differ across loss choices and CE is tied to next-token predictive quality.
  • Mildly negative exponents improve validation CE across the CE–Brier interpolation, whereas aggressive choices such as p = −0.75 and p = −1 degrade performance.
  • DynMuon changes only the local spectral-shaping rule, introducing no additional cross-layer or cross-device coupling beyond Muon.
  • The work extends fixed spectral-shaping studies by making the shaping rule iteration-dependent, consistent with broader phase-dependent optimizer scheduling.
Loading 2605.17109v3…