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

2025-10-10

AI Summary

This work addresses the geometric incompatibility of gradient descent in neural network optimization by proposing a non-Euclidean gradient descent framework. It unifies cross-layer norm selection and aggregation mechanisms, formally defining the optimizer family Muon and its variants; introduces learning-rate normalization and Model Momentum (Momo) to enhance robustness; and, for the first time, incorporates Adam into this framework, yielding the improved MuonMax. Experiments demonstrate that MuonMax+Momo achieves superior generalization on unseen tasks, significantly reduces hyperparameter sensitivity—cutting tuning costs by over 40% on average—and consistently outperforms Adam and SGD across multi-task benchmarks. The core contribution is the establishment of the first systematic, theoretically grounded non-Euclidean optimization framework for gradient descent, accompanied by a novel optimizer family that bridges rigorous mathematical foundations with practical engineering efficacy.

Technology Category

• Machine Learning: Optimization
• Natural Language Processing: Learning & Optimization for NLP
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Responsible Web: Machine-in-the-loop, human agency and autonomy

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.

Problem

• Systematically explores non-Euclidean gradient descent aggregation methods
• Proposes robust MuonMax variant insensitive to learning rate choices
• Combines momentum with non-Euclidean methods for hyperparameter robustness

Innovation

• Proposes non-Euclidean gradient descent framework
• Introduces robust MuonMax optimizer variant
• Combines momentum with Muon for hyperparameter robustness