PaperlayerView exact revision on arXiv

Source-linked AI summary

Musec: MomentUm SpEctral Clipping for Stable Muon-type Training

Zhuanghua Liu, Menglian Wang, Luo Luo

arXiv:2609.11655v1cs.LG

TL;DR

Muon’s spectral flattening can destabilize training, while existing clipping remedies are architecture-specific. Musec clips only excessively large momentum singular values, and Soft Musec provides an efficient implementation; the paper reports convergence guarantees and improved stability across datasets and model sizes.

  • Problem

    Muon’s spectral flattening can cause unstable weight growth, while existing weight- and logit-clipping remedies do not address all model components or the optimizer’s root cause.

  • Method

    Musec replaces spectral flattening with thresholded singular-value clipping that preserves momentum structure, with Soft Musec using smooth saturation and coupled Newton-Schulz iterations.

  • Results

    Soft Musec substantially improves training stability over existing Muon variants across multiple datasets and model sizes, while Musec has convergence guarantees in the nonconvex nonsmooth setting.

  • Takeaways & Limitations

    Musec provides an optimizer-level, architecture-agnostic approach to stabilizing Muon-type training while retaining a principled convergence analysis.

  • Takeaways & Limitations

    The convergence analysis assumes Lipschitz sample losses, weak convexity, an unbiased Clarke-subgradient oracle, bounded variance, and a finite lower objective bound.

Abstract

from arXiv · show

Muon has emerged as a highly effective optimizer for large language model training, often achieving superior convergence and performance compared with the widely adopted Adam and AdamW optimizers. Nevertheless, Muon is prone to training instability due to its spectral flattening, manifested by loss spikes and unbounded growth of model weights. Existing approaches primarily rely on weight or attention-logit clipping, which require architecture-specific modifications and do not directly address instability across all model components. We propose MomentUm SpEctral Clipping (Musec), which replaces Muon's spectral flattening with spectral clipping: rather than setting all singular values of the momentum matrix to approximately one, Musec clips singular values that exceed a threshold while preserving the underlying spectral structure of the momentum. Our strategy provides an optimizer-level, architecture-agnostic mechanism for stabilizing Muon training. We further develop Soft Musec, an efficient implementation that uses a smooth spectral saturation function approximated by coupled Newton-Schulz iterations. Theoretically, we establish convergence guarantees for Musec in nonconvex nonsmooth stochastic optimization. To the best of our knowledge, this is the first convergence guarantee for Muon-type methods in the nonconvex nonsmooth setting. We provide empirical studies to show that Soft Musec consistently improves training stability over existing Muon variants across a wide range of learning rates and model sizes. Notably, Soft Musec remains stable in settings where existing Muon variants diverge, while matching their performance under well-tuned configurations.

1 INTRODUCTION

Muon offers fast convergence but can become unstable because spectral flattening amplifies weak momentum directions. Musec replaces flattening with clipping, while Soft Musec and convergence analysis provide practical and theoretical support.

  • Motivation: Muon instability stems from flattening the momentum spectrum, which can amplify weak directions and cause exploding weights or divergent attention logits.Existing remedies target attention logits or query-key weights, leaving other components and the spectral transformation itself unaddressed.
  • Musec: Musec clips only singular values above a threshold, preserving momentum spectral structure while operating architecture-agnostically across weight matrices.This replaces Muon’s spectral flattening with an optimizer-level mechanism rather than component-specific clipping.
  • Theory: Musec is analyzed for stochastic nonconvex nonsmooth optimization, with convergence to a (δ, ϵ)-Goldstein stationary point and accuracy dependence matching an optimal stochastic first-order method.The paper presents this as the first convergence analysis of Muon-type algorithms in this setting.
  • Soft Musec: Soft Musec replaces hard clipping with smooth saturation approximated by coupled Newton-Schulz iterations for efficient implementation.The implementation avoids requiring a full singular-value decomposition.
  • Related contributions: The paper positions Musec as complementary to concurrent spectral-clipping studies by combining Muon-specific motivation with theoretical analysis and empirical stability evaluations.The related work includes optimizer-agnostic SPECTRA and MuCon, which focuses on replacing Muon’s polar step.

2 RELATED WORK

Matrix-aware optimizers exploit weight-matrix structure through preconditioning or momentum orthogonalization, while nonsmooth optimization research provides stationarity and complexity foundations.

  • Matrix-Aware Optimizers: Shampoo and SOAP use matrix structure for Kronecker-factored or eigenspace-based preconditioning that incorporates approximate second-order information.These methods improve on coordinate-wise optimization by exploiting two-dimensional weight structure.
  • Matrix-Aware Optimizers: Muon instead orthogonalizes the momentum matrix with Newton-Schulz iterations, distinguishing it from preconditioner-based matrix-aware methods.This orthogonalization underlies Muon’s characteristic spectral transformation.
  • Nonconvex Nonsmooth Optimization: Nonconvex nonsmooth optimization developed from Clarke and Goldstein subdifferentials toward non-asymptotic complexity guarantees for Goldstein stationary points.The passage places Cutkosky et al. among subsequent advances in this line of work.

3 PRELIMINARIES AND BACKGROUND

The paper formulates stochastic matrix optimization under weak convexity and nonsmoothness, reviews Muon’s spectral flattening, and motivates Musec’s selective spectral clipping.

  • 3.1 NOTATIONS: The optimization problem minimizes an expected stochastic objective over matrix-valued parameters, with spectral and Frobenius norms used throughout the analysis.The notation also includes Frobenius inner products, matrix balls, clipping operators, and convex hulls.
  • 3.2 PROBLEM FORMULATION: The analysis assumes Lipschitz sample losses, a weakly convex objective, an unbiased Clarke-subgradient oracle, bounded variance, and a finite lower objective bound.These assumptions define the stochastic nonconvex nonsmooth setting used for convergence guarantees.
  • 3.2 PROBLEM FORMULATION: Goldstein stationarity measures nonsmooth convergence through the distance from zero to the convex hull of Clarke subgradients within a δ-ball.A (δ, ϵ)-stationary point has this distance at most ϵ.
  • 3.3 THE MUON OPTIMIZER: Muon orthogonalizes its momentum matrix, replacing nonzero singular values with one and typically approximating this operation with Newton-Schulz iterations.Its algorithm updates momentum from stochastic gradients before applying the spectral transformation.
  • 3.3 THE MUON OPTIMIZER: Musec’s algorithm introduces a clipping threshold and applies spectral clipping to momentum updates instead of Muon’s polar step.Singular values above the threshold are truncated, while smaller ones remain unchanged.
  • 3.3 THE MUON OPTIMIZER: Muon’s flattening discards spectral magnitude information, giving weak directions updates comparable to strong directions and contributing to weight growth and instability.Musec is motivated by preserving momentum spectral structure while constraining only excessively large singular values.

4 METHODOLOGY

Musec replaces Muon’s spectral flattening with threshold-based spectral clipping that preserves momentum structure while bounding large singular components. It also integrates clipping into the momentum recurrence and provides convergence guarantees for nonconvex nonsmooth optimization.

  • 4.1 THE ALGORITHM: Musec clips momentum singular values above a threshold while preserving smaller values and the momentum’s spectral structure.Unlike Muon’s polar step, this avoids flattening all nonzero singular values to one.
  • 4.1 THE ALGORITHM: Musec feeds clipped momentum back into the next EMA update, keeping the momentum state itself spectrally bounded.This differs from SPECTRA, which carries the unclipped momentum state forward after post-processing.
  • 4.2 MUSEC AS CONSTRAINED STEEPEST DESCENT: Musec’s clipped update solves a Frobenius-norm-penalized quadratic subproblem under a spectral-norm constraint, giving the method a constrained-steepest-descent interpretation.The constraint clips singular values exceeding D while preserving smaller ones.
  • 4.3 CONVERGENCE ANALYSIS: Theoretical analysis establishes convergence of Musec to a (δ, ϵ)-Goldstein stationary point for nonconvex nonsmooth optimization under stated step-size and weak-convexity conditions.The guarantee applies when δ ≥ ηTD√r and uses γ = β/η with η ≤ 1/ρ and β ≤ 1/8.
  • 4.3 CONVERGENCE ANALYSIS: O(r^3/2δ^-1ϵ^-3) is the dominant oracle complexity when ρ ≤ O(rδ^-1ϵ^-1), matching optimal stochastic first-order dependence on δ and ϵ.The remaining r^3/2 factor reflects the matrix-valued update structure.

5 A PRACTICAL IMPLEMENTATION

Soft Musec replaces exact SVD-based clipping with a smooth spectral saturation function and approximates the resulting inverse square root using coupled Newton–Schulz iterations. This makes spectral clipping computationally suitable for large language model training.

  • 5 A PRACTICAL IMPLEMENTATION: Exact spectral clipping requires a full SVD at every iteration, so Soft Musec provides an efficient SVD-free approximation.The exact operation is prohibitively expensive for the large weight matrices used in LLM training.
  • 5 A PRACTICAL IMPLEMENTATION: Soft Musec replaces hard clipping with a smooth saturation function that approximates the input below the threshold and saturates near the threshold above it.Applying this function to singular values while preserving singular vectors yields soft spectral clipping.
  • 5 A PRACTICAL IMPLEMENTATION: Coupled Newton–Schulz iterations approximate the required matrix inverse square root using only matrix–matrix multiplications.The D^2I shift improves numerical stability and convergence, and five iterations perform well in bfloat16 arithmetic.

6 EXPERIMENTS

Experiments evaluate Soft Musec against Muon, MuonClip, and SPECTRA across model sizes, datasets, learning rates, and training dynamics. Soft Musec remains stable at aggressive learning rates, recovers smoothly from loss spikes, and controls spectral-norm growth across model components.

  • 6.1 EXPERIMENTAL SETUP: The experiments use NanoGPT-Small, Medium, and Wide models from 491M to 1.63B parameters across FineWeb, OpenWebText, and C4.Learning-rate sweeps on FineWeb range from 0.01 to 0.8.
  • 6.2 STABILITY ACROSS LEARNING RATES: At the largest learning rate, Muon fails across all model sizes and MuonClip fails on two sizes, while Soft Musec and SPECTRA remain stable across all configurations.Soft Musec achieves performance comparable to SPECTRA, with slight improvements in several settings.
  • 6.3 TRAINING DYNAMICS: Soft Musec and SPECTRA recover quickly and smoothly from the step-400 loss spike, whereas Muon recovers slowly and reaches substantially higher final validation loss.MuonClip improves over Muon but shows noticeable oscillations, while Soft Musec and SPECTRA converge smoothly to comparable final validation loss.
  • 6.4 WEIGHT NORM ANALYSIS: Soft Musec and SPECTRA keep spectral norms below 10 across QK, VO, and MLP components, while Muon reaches 300–400 with large oscillations.MuonClip partially reduces QK growth but leaves VO and MLP norms elevated.

7 CONCLUSION

The paper introduces Musec and Soft Musec to stabilize Muon training while preserving momentum spectral structure, and reports improved stability across datasets and model sizes.

  • Musec replaces Muon’s spectral flattening with clipping that preserves momentum structure while constraining excessively large singular values.
  • Soft Musec uses an efficient SVD-free implementation based on coupled Newton–Schulz iterations.
  • Soft Musec substantially improves training stability over existing Muon variants across multiple datasets and model sizes.
  • The appendix includes proofs, algorithmic details, update-rule comparisons, experimental configurations, and additional learning-rate sweeps.

A PROOF OF PROPOSITION 4.1

This section proves the optimization properties underlying Musec’s clipped update and develops the regret and stochastic bounds used in its convergence analysis.

  • Aligning singular vectors with the negative gradient reduces the matrix problem to singular-value optimization under a spectral constraint.
  • Each scalar quadratic subproblem is solved by clipping its unconstrained singular-value minimizer to [0, D/λ].
  • Musec’s momentum update is analyzed as online gradient descent over a convex spectral-norm constraint set.
  • The proof establishes regret bounds for the constrained online gradient descent update under β ≤ 1/8.
  • The convergence argument combines weak convexity, martingale-difference noise control, Jensen’s inequality, and bounds on iterate deviations.

B.4 PROOF OF COROLLARY 4.4

This section evaluates the corollary’s parameterized oracle-complexity bound and summarizes the SVD-free Soft Musec implementation and its distinction from SPECTRA.

  • Corollary 4.4 derives Musec’s stochastic-gradient oracle complexity by selecting parameters that satisfy the theorem’s constraints.
  • The proof verifies the parameter choices before translating Theorem 4.3’s bound into the stated complexity.
  • Soft Musec replaces exact SVD-based clipping with a smooth saturation function approximated by coupled Newton–Schulz iterations.
  • The Soft Musec procedure clips momentum through an SVD-free subroutine and updates parameters using the clipped momentum.
  • Musec propagates clipped momentum so its state remains spectrally bounded, whereas SPECTRA carries an unconstrained momentum state forward.

E.1 ARCHITECTURE CONFIGURATIONS

The experiments use three NanoGPT decoder-only Transformer configurations and evaluate optimizers across multiple datasets, model scales, and learning rates. Hyperparameters are selected by validation loss, with Soft Musec and SPECTRA additionally tuning the clipping threshold.

  • ARCHITECTURE CONFIGURATIONS: Experiments cover NanoGPT-Small, NanoGPT-Medium, and NanoGPT-Wide decoder-only Transformers evaluated on FineWeb, OpenWebText, and C4.The models use modern GPT-2 modifications, including RMSNorm, squared ReLU, rotary embeddings, Flash Attention 3, and multi-token prediction.
  • TRAINING CONFIGURATION: All methods apply their respective optimizer to attention and MLP projection matrices while using AdamW for embeddings, the language-model head, and scalar parameters.AdamW hyperparameters remain fixed across methods.
  • HYPERPARAMETERS: Weight decay is tuned over four values for every method, model configuration, and learning rate, while Soft Musec and SPECTRA additionally tune clipping threshold D over six values.The reported result for each setting is the best validation loss from these sweeps.
  • RESULTS: The evaluation compares validation loss across learning-rate sweeps for all three model scales and datasets, with complete results reported for NanoGPT-Small and NanoGPT-Medium.Runs that diverge during training are omitted from the learning-rate sweep figures.

F.2 ADDITIONAL TRAINING DYNAMICS

Across model configurations and datasets, Muon-type optimizers outperform Adam and AdamW when well tuned, but their stability gap widens at higher learning rates and larger model scales. Soft Musec and SPECTRA remain stable where Muon and MuonClip diverge.

  • ADDITIONAL TRAINING DYNAMICS: At the best-tuned learning rate, all Muon-type optimizers reach lower validation loss than Adam and AdamW.The advantage is reported across the evaluated model configurations and datasets.
  • ADDITIONAL TRAINING DYNAMICS: As learning rates increase, the stability gap widens and is amplified at larger model scales.For NanoGPT-Small, instability appears at η = 0.5; for NanoGPT-Medium and NanoGPT-Wide, instability occurs more broadly.

F.3 EFFECTIVE RANK ANALYSIS

The effective-rank analysis tests whether spectral clipping preserves less rank in optimizer updates than Muon’s polar step. Soft Musec and SPECTRA produce substantially lower effective rank than Muon, while Soft Musec adds negligible wall-clock overhead and remains robust to implementation choices.

  • EFFECTIVE RANK: Muon has near-full effective rank across QK, VO, and MLP updates, whereas Soft Musec and SPECTRA achieve substantially lower effective rank.MuonClip does not meaningfully reduce update effective rank because it clips weight matrices rather than updates.
  • COMPUTATIONAL OVERHEAD: Soft Musec introduces negligible computational overhead compared with Muon across NanoGPT-Small, NanoGPT-Medium, and NanoGPT-Wide.The reported method ranges are within 3% for Small, similar for Medium, and within 2% for Wide.
  • SENSITIVITY TO CLIPPING THRESHOLD: At η = 0.5, Soft Musec remains smoothly convergent for D values from 0.05 to 0.25, while larger thresholds produce later loss spikes but still converge.At η = 0.1 and η = 0.2, all tested thresholds from D = 0.05 to D = 1.0 give comparable validation loss and training dynamics.
  • SENSITIVITY TO NEWTON–SCHULZ STEPS: Soft Musec is largely insensitive to the number of Newton–Schulz iterations, with three iterations providing most of the stability benefit.Additional iterations yield only marginally better validation loss.
Loading 2609.11655v1…