AMUSE: Anytime Muon with Stable Gradient Evaluation
- URL: http://arxiv.org/abs/2605.22432v1
- Date: Thu, 21 May 2026 12:55:52 GMT
- Title: AMUSE: Anytime Muon with Stable Gradient Evaluation
- Abstract summary: We study Muon through the river-valley loss landscape, where useful training progress occurs along a flat, low-curvature bulk subspace.<n>We propose Anytime MUon with Stable gradient Evaluation (AMUSE), which integrates Muon's rapid bulk progress with the stabilizing effect of Schedule-Free averaging.
- Score: 25.647698512832648
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: Modern deep learning commonly relies on AdamW with prescribed learning rate schedules, but recent works challenge both components: Schedule-Free optimization removes explicit schedules via iterate averaging, and Muon improves the update geometry by orthogonalizing momentum for matrix parameters. Despite Muon's strong empirical performance, its underlying mechanism remains partially understood. We study Muon through the river-valley loss landscape, where useful training progress occurs along a flat, low-curvature bulk subspace (the river), while high-curvature dominant directions form steep valley walls that induce oscillations. We empirically show that while Muon's orthogonalization accelerates river progress by increasing the bulk component, it also amplifies dominant-direction noise, causing oscillatory trajectories. Building on this, we propose Anytime MUon with Stable gradient Evaluation (AMUSE), which integrates Muon's rapid bulk progress with the stabilizing effect of Schedule-Free averaging. AMUSE uses a time-varying interpolation coefficient that initially evaluates gradients near the fast Muon sequence for rapid adaptation, then gradually shifts toward the stable averaged sequence to suppress valley-wall oscillations. As a result, AMUSE requires no learning rate schedules and supports anytime training. Across vision tasks and large language model pretraining, AMUSE consistently improves the performance-iteration Pareto frontier over (Schedule-Free) AdamW and Muon.
Related papers
- Reassessing Muon for Matrix Factorization [9.691906492780353]
Muon has emerged as a strong contender in large-scale deep learning.<n>We study Muon on a simple, well-understood, and spectrally structured problem: low-rank matrix factorization.<n>We find that Muon does not consistently outperform AdamW in this setting.
arXiv Detail & Related papers (2026-07-14T20:21:38Z) - Towards Understanding the Power and Limits of the Muon Optimizer: A River-Valley Perspective [12.959531584997462]
Recently, Muon has gained substantial attention as an appealing alternative to Adam-like nuisances.<n>We introduce a mixed-spiked matrix sensing model whose sensing operator decomposes into signal, spike, and bulk components.<n>We show that Muon moves faster along the information-bearing river direction during early stages of learning.
arXiv Detail & Related papers (2026-06-19T15:10:20Z) - Timestep Rescheduling in Diffusion Inversion [79.30105376216544]
We show how the selection of the diffusion timestep in the noise scheduler influences inversion fidelity.<n>We propose a simple yet effective nonuniform timestep scheduler that integrates a global rescaling with a local dynamic programming based rescheduling.<n>Our method serves as an off-the-shelf enhancement for existing inversion techniques and requires no extra parameters or computational overhead.
arXiv Detail & Related papers (2026-06-13T16:32:01Z) - MONA: Muon Optimizer with Nesterov Acceleration for Scalable Language Model Training [32.15839901598669]
We present MONA, an acceleration term that bridges Muon's orthogonalization framework with curvature-aware acceleration.<n>We provide a detailed convergence analysis for MONA, showing that the acceleration term enables escape from sharp minima.<n>We evaluate it on general capability, mathematical reasoning, and code generation benchmarks, where MONA achieves SOTA performance.
arXiv Detail & Related papers (2026-05-26T10:56:20Z) - Extra-Merge: Tracing the Rank-1 Subspace of Model Merging in Language Model Pre-Training [56.323119575322146]
Extra-Merge consistently outperforms standard merging baselines in experiments across GPT-2 and LLaMA families.<n>It yields consistent zero-shot accuracy gains on Pythia-12B downstream tasks and generalizes effectively to the Muon citepjordan2024muon.
arXiv Detail & Related papers (2026-05-26T02:48:34Z) - DynMuon: A Dynamic Spectral Shaping View of Muon [12.175075916077702]
We consider a class of Muon-like updates, where we replace the update $M$ with $Up Vtop$ for some parameter $p$.<n>We propose DynMuon, an efficient dynamic spectral shaping method that schedules $p$ from positive to mildly negative over training.
arXiv Detail & Related papers (2026-05-16T18:30:11Z) - Spectral Flattening Is All Muon Needs: How Orthogonalization Controls Learning Rate and Convergence [15.754779716816174]
We show that Muon's maximal stable step size scales with the average singular value of the gradient rather than the largest.<n>We recast Muon as a preconditioned gradient method and show, under a Kronecker-factored curvature model, that it improves the effective convergence factor.
arXiv Detail & Related papers (2026-05-13T06:54:01Z) - Muown: Row-Norm Control for Muon Optimization [41.541871767555484]
Muon is a strong competitor to AdamW for language model pre-training.<n>Muown improves perplexity over Muon, SOAP, AdamW, and Lion.
arXiv Detail & Related papers (2026-05-11T16:26:06Z) - Sharp Capacity Scaling of Spectral Optimizers in Learning Associative Memory [77.27772368491698]
Spectrals such as Muon have recently shown strong empirical performance in large-scale language model training.<n>We study this question through the linear associative memory problem.<n>We show that the storage capacity of Muon significantly exceeds that of SGD.
arXiv Detail & Related papers (2026-03-27T16:13:18Z) - Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning [72.62839712454196]
Recent advances in spectral optimization, notably Muon, have demonstrated that constraining update steps to the Stiefel manifold can significantly accelerate training and improveimation.<n>We argue that this "egalitarian" constraint is suboptimal for Deep Neural Networks, where the curvature spectrum is known to be highly heavy-tailed and ill-conditioned.<n>We propose textbfMousse (textbfMuon textbfOptimization textbfUtilizing textbfShampoo's textbfStruct
arXiv Detail & Related papers (2026-03-10T14:03:49Z) - Muon in Associative Memory Learning: Training Dynamics and Scaling Laws [23.350512542598803]
We study Muon in a linear associative memory model with softmax retrieval and a hierarchical frequency spectrum over query-answer pairs.<n>We show that Muon mitigates this imbalance, leading to faster and more uniform progress.
arXiv Detail & Related papers (2026-02-05T14:49:40Z) - Variance-Adaptive Muon: Accelerating LLM Pretraining with NSR-Modulated and Variance-Scaled Momentum [19.385264518362472]
Large Language Models (LLMs) achieve competitive performance across diverse natural language processing (NLP) tasks.<n>We propose two variants of Muon, Muon-NSR and Muon-VS, which apply variance-adaptive normalization to momentum.<n>Experiments on GPT-2 and LLaMA pretraining demonstrate that our proposed methods accelerate convergence and consistently achieve lower validation loss than both competitive, well-tuned AdamW and Muon baselines.
arXiv Detail & Related papers (2026-01-21T02:41:56Z) - Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective [66.80315289020487]
Warmup-Stable-Decay (WSD) schedule uses a constant learning rate to produce a main branch of iterates that can continue indefinitely without a pre-specified compute budget.<n>We show that pretraining loss exhibits a river valley landscape, which resembles a deep valley with a river at its bottom.<n>Inspired by the theory, we introduce WSD-S, a variant of WSD that reuses previous checkpoints' decay phases and keeps only one main branch.
arXiv Detail & Related papers (2024-10-07T16:49:39Z) - One More Step: A Versatile Plug-and-Play Module for Rectifying Diffusion
Schedule Flaws and Enhancing Low-Frequency Controls [77.42510898755037]
One More Step (OMS) is a compact network that incorporates an additional simple yet effective step during inference.
OMS elevates image fidelity and harmonizes the dichotomy between training and inference, while preserving original model parameters.
Once trained, various pre-trained diffusion models with the same latent domain can share the same OMS module.
arXiv Detail & Related papers (2023-11-27T12:02:42Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.