An Exploration of Non-Euclidean Gradient Descent: Muon and its Many Variants

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.

We study a new algorithm for training neural networks, Muon, by first identifying several design decisions implicit in Muon's design, then systematically exploring alternative choices. By making these design choices explicit, we are able to interpret Muon as a version of a very familiar algorithm (gradient descent), which enables us to adapt tools designed for gradient descent to the more powerful Muon. One such tool, called Momo, is designed to decrease the need for manual tuning of neural network training algorithms, which can be quite costly at large scale. We show how to adapt Momo for Muon (and the many variations we try), and in the end we discover new neural network training algorithms that are as powerful as Muon while requiring much less manual tuning.