Technical Summary: AMUSE (Anytime Muon with Stable Gradient Evaluation)

Problem Statement

Modern deep learning optimization typically relies on the AdamW optimizer combined with prescribed learning rate schedules (e.g., cosine decay). Recent advancements have challenged these components individually: Schedule-Free (SF) optimization removes explicit schedules by evaluating gradients at interpolated points between current and averaged iterates, while Muon improves update geometry for matrix-valued parameters by orthogonalizing momentum.

Despite Muon's empirical success in large language model (LLM) pretraining, its underlying mechanism is not fully understood. The paper identifies a specific instability in Muon: while its orthogonalization accelerates progress along "flat" directions (the "river" of the loss landscape), it inadvertently amplifies noise in "steep" high-curvature directions (the "valley walls"), leading to oscillatory trajectories. Conversely, while Schedule-Free methods stabilize trajectories by keeping them near the river, they may lack the rapid bulk progress offered by Muon's orthogonalization. The core problem is designing an optimizer that combines Muon's rapid convergence with the stability of Schedule-Free averaging without requiring manual learning rate schedules.

Methodology: AMUSE

The authors propose AMUSE (Anytime Muon with Stable gradient Evaluation), which integrates Muon's orthogonalized matrix updates with the Schedule-Free framework using a time-varying interpolation coefficient.

Theoretical Foundation: River-Valley Landscape

The analysis is grounded in the river-valley loss landscape perspective:

• The River: A high-dimensional, low-curvature bulk subspace where useful training progress occurs.
• The Valley Walls: A low-dimensional, high-curvature dominant subspace that induces oscillations.
• Muon's Mechanism: Orthogonalizing the momentum matrix increases the relative bulk component of updates, accelerating progress along the river. However, it also prevents dominant components from decaying, causing small noise in dominant directions to amplify into oscillations across valley walls.
• Schedule-Free Mechanism: Evaluating gradients at an averaged trajectory (closer to the river) filters out high-curvature fluctuations before the update step.

Algorithm Design

AMUSE employs a time-varying interpolation coefficient $\beta_t$ to determine the gradient evaluation point $Y_t$ between the fast base sequence $Z_t$ (Muon updates) and the stable averaged sequence $X_t$:
$$Y_t = (1 - \beta_t)Z_t + \beta_t X_t$$

The schedule for $\beta_t$ is designed to balance early adaptation and late-stage stability:

1. Warmup Phase ($t \le T_0$): $\beta_t$ is fixed at a small value $\beta_1$. This allows the optimizer to evaluate gradients near the fast Muon sequence ($Z_t$) for rapid initial adaptation.
2. Post-Warmup ($t > T_0$): $\beta_t$ gradually increases toward 1 according to:
   $$\beta_t = 1 - \left( \frac{T_0 - 1}{t - 1} \right)^\rho (1 - \beta_1)$$
   Here, $\rho \in [0, 1]$ controls the rate of shift. As $\beta_t$ increases, the gradient evaluation point shifts toward the averaged sequence $X_t$. This suppresses the high-curvature (dominant) components of the gradient before they undergo Muon's orthogonalization, thereby stabilizing the trajectory.

Crucially, this schedule is independent of the total training horizon, enabling anytime training (training can be stopped at any point without performance degradation). AMUSE requires no learning rate decay schedule.

Key Contributions

1. Geometric Analysis of Muon: The paper provides a geometric explanation for Muon's efficiency and instability. It demonstrates that orthogonalization promotes bulk-oriented updates (accelerating river progress) but amplifies dominant-direction noise (causing valley-wall oscillations).
2. Stabilization via Interpolation: The authors show that evaluating gradients closer to the averaged trajectory (via Schedule-Free) reduces dominant components prior to orthogonalization, mitigating oscillations.
3. AMUSE Algorithm: A novel optimizer that dynamically shifts the gradient evaluation point from the fast Muon trajectory to the stable averaged trajectory. This combines the rapid convergence of Muon with the stability of Schedule-Free methods.
4. Anytime Training Support: By using a time-varying schedule independent of the total steps and avoiding learning rate decay, AMUSE supports flexible training horizons.

Experimental Results

AMUSE was evaluated across diverse benchmarks, including image classification (CIFAR-10/100, SVHN, ImageNet-1k), image segmentation (ISIC 2018), and LLM pretraining (124M, 720M, and 1B parameter Llama-style models on FineWeb).

• Performance-Iteration Pareto Frontier: AMUSE consistently outperforms (Schedule-Free) AdamW and Muon across all tasks.
• Efficiency Gains:
  • On 720M LLM pretraining, AMUSE reaches Muon's final performance with 1.51$\times$ fewer steps.
  • On ImageNet with ResNet-50, it achieves the same performance with 1.12$\times$ fewer steps.
  • On ImageNet with ViT fine-tuning, it requires 3.08$\times$ fewer steps.
• Stability: Unlike vanilla Muon, which benefits significantly from post-hoc exponential weight averaging (EWA) or learning rate decay, AMUSE maintains a stable trajectory throughout training, showing minimal improvement when these post-hoc techniques are applied.
• Scalability: AMUSE scales effectively to large batch sizes (up to 2048) and large model scales (1B parameters), where other Schedule-Free variants sometimes struggle.

Significance and Claims

The paper claims that AMUSE represents a principled approach to optimizer design derived from a systematic understanding of training dynamics and loss landscape geometry. By explicitly addressing the trade-off between rapid bulk progress and valley-wall stability, AMUSE achieves superior performance without the need for manual learning rate scheduling.

The authors modestly note that while AMUSE outperforms existing optimizers, it incurs a slight memory overhead (one extra state copy for the base sequence $Z_t$) compared to vanilla Muon, though it remains comparable to AdamW and Schedule-Free AdamW. They suggest that future work could focus on achieving similar stability and performance without this additional memory cost. The paper emphasizes that the "river-valley" perspective offers a robust framework for designing optimizers that are both efficient and stable.