Distance-Aware Muon: Adaptive Step Scaling for Normalized Optimization

TL;DR

This paper introduces three adaptive step scaling algorithms for Muon optimizers, improving their robustness and performance across various tasks and geometries, with theoretical guarantees and empirical validation.

Contribution

The paper develops novel adaptive scaling rules for Muon optimizers, including Distance-Adaptive Muon, Scale-Calibrated Muon, and Distance-Free Muon, with theoretical analysis and practical benefits.

Findings

Distance-Adaptive Muon guarantees stationarity under bounded trajectories.

Scale-Calibrated Muon achieves O(1/T) objective-gap bounds in star-convex settings.

Experiments show adaptive rules reduce tuning sensitivity and match or outperform fixed-scale baselines.

Abstract

Muon and related normalized optimizers decouple the choice of update direction from the choice of step scale, but their practical performance remains sensitive to the scale of the normalized step. We study adaptive scaling rules for Muon in general norm geometries and develop three complementary algorithms. For smooth non-convex objectives, we introduce Distance-Adaptive Muon, whose trust-region radius is set from the radius explored by the trajectory, and prove a stationarity guarantee under a bounded-trajectory assumption. We then turn to star-convex objectives, a tractable model of the favorable global geometry often used to reason about the empirical loss landscapes of deep neural networks, where objective-gap guarantees are possible. In this setting, we first introduce Scale-Calibrated Muon, which keeps Muon's exponential moving average but sets the step length from a local descent certificate computed from the current gradient and momentum. For this method, we prove a last-iterate O(1/T) objective-gap bound under a bounded initial sublevel-set assumption, where the corresponding radius parameter appears only in the analysis and not in the algorithm. Finally, we develop Distance-Free Muon, a recentered trust-region method that uses a scalar distance certificate and a majorized one-dimensional search to select the trust-region radius without requiring the unknown distance from the initialization to a global minimizer. Experiments on Transformer language modeling (GPT-124M/WikiText-103) and image classification (ViT-Tiny/CIFAR-100) show that the proposed adaptive scaling rules reduce sensitivity to manual scale tuning and match or improve tuned fixed-scale Muon baselines under the tested budgets.