An Exploration of Non-Euclidean Gradient Descent: Muon and its Many Variants
Michael Crawshaw, Chirag Modi, Mingrui Liu, Robert M. Gower
TL;DR
This work unifies non-Euclidean gradient methods for neural networks by formulating a general steepest-descent framework operating on the full parameter space with per-layer norms and a product-norm aggregation. It introduces MuonMax, a robust variant derived from a new product norm, and shows how model truncation (Momo) can be integrated with arbitrary norms to further boost stability and hyperparameter-tuning robustness. By recasting Muon, Scion, and related methods within this framework, the authors derive closed-form update rules and demonstrate that MuonMax-Momo consistently matches or outperforms prior baselines while exhibiting far greater resilience to learning-rate choices, across 1B-token FineWeb and 6B-token SlimPajama tasks. The practical impact lies in enabling robust hyperparameter tuning and scalable training of large language models with reduced manual tuning, thanks to the combination of MuonMax and Momo within a principled norm-driven optimization scheme.
Abstract
To define a steepest descent method over a neural network, we need to choose a norm for each layer, a way to aggregate these norms across layers, and whether to use normalization. We systematically explore different alternatives for aggregating norms across layers, both formalizing existing combinations of Adam and the recently proposed Muon as a type of non-Euclidean gradient descent, and deriving new variants of the Muon optimizer. Through a comprehensive experimental evaluation of the optimizers within our framework, we find that Muon is sensitive to the choice of learning rate, whereas a new variant we call MuonMax is significantly more robust. We then show how to combine any non-Euclidean gradient method with model based momentum (known as Momo). The new Momo variants of Muon are significantly more robust to hyperparameter tuning, and often achieve a better validation score. Thus for new tasks, where the optimal hyperparameters are not known, we advocate for using Momo in combination with MuonMax to save on costly hyperparameter tuning.