TrasMuon: Trust-Region Adaptive Scaling for Orthogonalized Momentum Optimizers
This addresses optimization instability for users of orthogonalized momentum methods, but it is incremental as it builds on existing Muon-style optimizers.
The paper tackled the sensitivity to step-size hyperparameters and vulnerability to high-energy bursts in Muon-style optimizers by introducing TrasMuon, which stabilizes magnitudes through global RMS calibration and energy-based trust-region clipping, resulting in faster convergence and superior stability without warmup stages in vision and language models.
Muon-style optimizers leverage Newton-Schulz (NS) iterations to orthogonalize updates, yielding update geometries that often outperform Adam-series methods. However, this orthogonalization discards magnitude information, rendering training sensitive to step-size hyperparameters and vulnerable to high-energy bursts. To mitigate this, we introduce TrasMuon (\textbf{T}rust \textbf{R}egion \textbf{A}daptive \textbf{S}caling \textbf{Muon}). TrasMuon preserves the near-isometric geometry of Muon while stabilizing magnitudes through (i) global RMS calibration and (ii) energy-based trust-region clipping. We demonstrate that while reintroducing adaptive scaling improves optimization efficiency, it typically exacerbates instability due to high-energy outliers. TrasMuon addresses this by defining a trust region based on relative energy ratios, confining updates to a stable zone. Empirical experiments on vision and language models demonstrate that TrasMuon converges faster than baselines. Furthermore, experiments without warmup stages confirm TrasMuon's superior stability and robustness.