AMUSE: Anytime Muon with Stable Gradient Evaluation

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 LLM pretraining, AMUSE consistently improves the performance-iteration Pareto frontier over (Schedule-Free) AdamW and Muon.

• The paper introduces AMUSE, which fuses Muon’s geometric update with schedule-free stabilization using a dynamic, time-varying interpolation coefficient.
• It empirically outperforms traditional optimizers by efficiently navigating flat loss landscapes while suppressing high-curvature oscillations.
• Its minimal hyperparameter overhead and anytime training capability demonstrate strong potential for scalable and robust deep learning optimization.

AMUSE: Anytime Muon with Stable Gradient Evaluation

Motivation and Background

The AMUSE optimizer is designed to address two fundamental challenges in deep learning training: unstable dynamics in high-curvature directions and inefficiency in exploiting flat, low-curvature directions in the loss landscape. Standard recipes, such as AdamW with learning rate decay schedules, have dominated training for modern architectures, particularly transformers. However, recent developments in schedule-free optimization and the Muon optimizer—where momentum is orthogonalized for matrix-valued parameters—have demonstrated notable empirical improvements over traditional methods.

Muon exploits matrix structure, promoting faster progress along flat directions (“the river”), but it amplifies noise and high-curvature oscillations (“valley walls”). Schedule-Free (SF) optimizers stabilize trajectories by gradient evaluation at interpolated points between the current and averaged iterates, filtering high-curvature fluctuations and enabling anytime training without explicit learning rate schedules. AMUSE integrates Muon's geometric update with SF’s stabilizing mechanism using a time-varying interpolation coefficient that shifts gradient evaluation from Muon’s fast sequence to the stable averaged sequence, facilitating rapid early adaptation and late-stage stability.

River-Valley Loss Landscape: Analysis and Principles

The optimization landscape of deep neural networks is highly anisotropic. The Hessian spectrum typically includes a few large outlier eigenvalues (dominant subspace) and a broad, lower spectrum (bulk subspace). This geometric structure underpins the river-valley analogy: dominant directions form steep valley walls with high curvature, whereas the bulk forms a flat river conducive to stable progress.

Optimizers must traverse the river efficiently and suppress unnecessary updates along valley walls. Empirically, Muon updates contain substantially larger bulk components than SGD or AdamW, accelerating progress in low-curvature directions. Orthogonalization shifts the update further away from dominant directions:

Figure 1: Muon produces smaller dominant-subspace components compared to SGD and AdamW; AMUSE further suppresses dominant updates.

However, Muon's orthogonalization is non-selective, amplifying noise and small dominant components, which can cause oscillatory motion across valley walls and trajectory instability. This is reflected in the Hessian eigenspectrum and confirmed by measuring the dominant/bulk ratios:

Figure 2: Comparison of dominant component ratios illustrates the increased bulk orientation of Muon and AMUSE updates.

Scaling experiments demonstrate that amplifying bulk components accelerates training, while dominant components can induce instability—supporting the notion that low-curvature directions are key to optimization efficiency.

Schedule-Free Stabilization and AMUSE Mechanism

Schedule-Free optimizers interpolate between fast base and averaged sequences, controlling the gradient evaluation point via $\beta$. Moving the evaluation point towards the averaged sequence suppresses dominant high-curvature contributions prior to orthogonalization, yielding more stable updates. However, fixed $\beta$ induces a tradeoff: low $\beta$ allows rapid early progress but harms stability; high $\beta$ stabilizes late-stage progress but impairs early adaptation.

AMUSE addresses this by employing a time-varying interpolation coefficient $\beta_t$, which starts low for rapid adaptation and increases towards 1, shifting evaluation toward the stable averaged sequence. This ensures consistent filtering of valley-wall noise and maintains rapid bulk traversal.

Figure 3: Validation perplexity and update norm showcase the Pareto-optimal performance of AMUSE compared to fixed-beta SF-Muon and AMUSE’s superior training stability and speed.

The mathematical schedule for $\beta_t$ dynamically adjusts the effective averaging window without shrinking, interpolating between fixed-beta and constant-window regimes for optimal stability and progress.

Figure 4: Evolution of $\beta_t$ controlled by $\rho$ demonstrates smooth interpolation and preservation of effective averaging window size.

Empirical Evaluation

AMUSE was evaluated across a range of vision and LLM benchmarks. In image classification tasks (CIFAR-10, CIFAR-100, SVHN, ImageNet-1k), AMUSE consistently matches or outperforms all baselines, including schedule-free and Muon optimizers, both in accuracy and convergence speed. In segmentation and MAE fine-tuning contexts, AMUSE remains robust even when transferred to models pretrained with AdamW.

Figure 5: Test accuracy across image-domain experiments highlights AMUSE’s superior anytime performance, averaged over multiple seeds.

In LLM pretraining (LLama architectures at 124M, 720M, and 1B parameters, trained on FineWeb), AMUSE achieves lower validation perplexity throughout training and maintains stability and efficiency even at large batch sizes, where SF-AdamW struggles. AMUSE improves the performance–iteration Pareto frontier decisively, requiring $1.51\times$, $1.12\times$, and $\beta$0 fewer steps for final performance on 720M LLM, ResNet-50/ImageNet, and ViT/ImageNet respectively.

Figure 6: Validation perplexity on FineWeb pretraining across Llama model scales demonstrates AMUSE’s consistent superiority across all scales.

Analysis shows that adding exponential weight averaging or learning rate decay to constant learning-rate Muon improves performance but does not match AMUSE, reinforcing that its trajectory remains inherently stable and efficient.

Experimental Design and Hyperparameter Handling

AMUSE requires only one additional hyperparameter ($\beta$1) when compared to standard SF optimizers. Muon momentum is fixed and not treated as a tunable parameter. Hyperparameter sweeps demonstrate AMUSE’s mild sensitivity and robust performance across tested ranges, with every configuration outperforming tuned Muon baselines. Implementation details ensure practical memory overhead similar to AdamW and SF-AdamW, with seamless integration into both vision and language domain settings.

Theoretical and Practical Implications

AMUSE demonstrates how optimizer design can leverage loss landscape geometry—by promoting bulk traversal and suppressing valley-wall oscillations—to transcend traditional learning-rate scheduling schemes. The schedule-free, anytime paradigm facilitates flexible training and inference, enabling better resource efficiency and model adaptation. The separation of bulk and dominant subspaces, and deliberate gradient evaluation control, inform a principled approach for future optimizer development.

Theoretical analyses (e.g., quadratic bouncing with matrix-normalized updates) reinforce the importance of controlling dominant components both algorithmically and geometrically. AMUSE also shows strong potential for scale transfer and cross-domain applicability, evidenced by superior results in both vision and language tasks and efficient scaling to larger models and batches.

Speculations for Future Developments

Future research should focus on eliminating the remaining memory overhead incurred by additional averaged states, possibly by devising more selective orthogonalization and averaging mechanisms. Exploration of more adaptive, curvature-aware schedules, integration with alternative normalization or preconditioning strategies (such as Mousse or AdamNamo), and further theoretical analysis of matrix-valued optimizer dynamics in highly anisotropic spaces will continue to push the boundaries of optimizer design.

Moreover, model architecture coupling (e.g., leveraging Muon and AMUSE for hybrid transformer-conv architectures) and distributed training regimes stand to benefit from AMUSE's scale-agnostic and schedule-free properties.

Conclusion

AMUSE combines Muon's matrix orthogonalization and fast bulk progress with schedule-free stable gradient evaluation, using a time-varying interpolation parameter. It consistently outperforms conventional Adam(W), Muon, and schedule-based optimizers across diverse deep learning tasks, delivering improved stability, efficiency, and anytime capability. The geometric and algorithmic insights underlying AMUSE highlight the importance of systematic landscape analysis for principled optimizer design and robust deep learning training (2605.22432).