Training Deep Learning Models with Norm-Constrained LMOs

2025/02/11 by Thomas Pethick, Pethick, Thomas, Wanyun Xie +9 · 3 voices · 39 citations
Computer Science · #Advanced Neural Network Applications #Neural Networks and Applications #Stochastic Gradient Optimization Techniques #cs.LG #math.OC

paper · pdf · doi:10.48550/arxiv.2502.07529

openalex publication_date 2025/02/11 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

In this work, we study optimization methods that leverage the linear minimization oracle (LMO) over a norm-ball. We propose a new stochastic family of algorithms that uses the LMO to adapt to the geometry of the problem and, perhaps surprisingly, show that they can be applied to unconstrained problems. The resulting update rule unifies several existing optimization methods under a single framework. Furthermore, we propose an explicit choice of norm for deep architectures, which, as a side benefit, leads to the transferability of hyperparameters across model sizes. Experimentally, we demonstrate significant speedups on nanoGPT training using our algorithm, Scion, without any reliance on Adam. The proposed method is memory-efficient, requiring only one set of model weights and one set of gradients, which can be stored in half-precision. The code is available at https://github.com/LIONS-EPFL/scion .

Cited by

• Smoothing DiLoCo with Primal Averaging for Faster Training of LLMs
• Ky Fan Norms and Beyond: Dual Norms and Combinations for Matrix Optimization
• Correction of Decoupled Weight Decay
• FOAM: Blocked State Folding for Memory-Efficient LLM Training
• Greedy Alignment Principle for Optimizer Selection
• Hyperparameter Transfer Enables Consistent Gains of Matrix-Preconditioned Optimizers Across Scales
• Turbo-Muon: Accelerating Orthogonality-Based Optimization with Pre-Conditioning
• Dion2: A Simple Method to Shrink Matrix in Muon
• A Tale of Two Geometries: Adaptive Optimizers and Non-Euclidean Descent
• Non-Euclidean SGD for Structured Optimization: Unified Analysis and Improved Rates
• Muon is Not That Special: Random or Inverted Spectra Work Just as Well
• A Proof of Learning Rate Transfer under μP
• Isotropic Curvature Model for Understanding Deep Learning Optimization: Is Gradient Orthogonalization Optimal?
• FedMuon: Accelerating Federated Learning with Matrix Orthogonalization
• How Muon's Spectral Design Benefits Generalization: A Study on Imbalanced Data
• Beyond the Ideal: Analyzing the Inexact Muon Update
• Weight Decay may matter more than muP for Learning Rate Transfer in Practice
• MARS-M: When Variance Reduction Meets Matrices
• MuonBP: Faster Muon via Block-Periodic Orthogonalization
• Robust Layerwise Scaling Rules by Proper Weight Decay Tuning
• Noise-Adaptive Layerwise Learning Rates: Accelerating Geometry-Aware Optimization for Deep Neural Network Training
• Don't Be Greedy, Just Relax! Pruning LLMs via Frank-Wolfe
• Cautious Weight Decay
• Adaptive Conditional Gradient Descent
• Nonlinearly Preconditioned Gradient Methods: Momentum and Stochastic Analysis
• Preconditioned Norms: A Unified Framework for Steepest Descent, Quasi-Newton and Adaptive Methods
• An Exploration of Non-Euclidean Gradient Descent: Muon and its Many Variants
• Optimal Scaling Needs Optimal Norm
• On Provable Benefits of Muon in Federated Learning
• Non-Euclidean Broximal Point Method: A Blueprint for Geometry-Aware Optimization
• DeMuon: A Decentralized Muon for Matrix Optimization over Graphs
• FedMuon: Federated Learning with Bias-corrected LMO-based Optimization
• Conda: Column-Normalized Adam for Training Large Language Models Faster
• Beyond Outliers: A Study of Optimizers Under Quantization
• FORGE: Fused On-Register Gradient Elimination for Memory-Efficient LLM Training
• Constrained Stochastic Spectral Preconditioning Converges for Nonconvex Objectives
• On the Convergence of Muon and Beyond
• Low-rank Orthogonalization for Large-scale Matrix Optimization with Applications to Foundation Model Training
• AdaGrad Meets Muon: Adaptive Stepsizes for Orthogonal Updates

Discussions

• Also, a shoutout to this amazing paper by @tonysf.bsky.social and collaborators, which is well worth reading: arxiv.org/abs/2502.07529 [bsky, 5 points, 0 comments]
• arxiv.org/abs/2502.07529 🚀 Key results: - Based on conditional gradient method - Beats Muon+Adam on NanoGPT (tested up to 3B params) - Zero-shot learning rate transfer across model size - Uses WAY le [bsky, 4 points, 1 comments]
• arxiv.org/abs/2502.07529 [bsky, 0 points, 0 comments]

Related

• Gluon: Making Muon & Scion Great Again! (Bridging Theory and Practice of LMO-based Optimizers for LLMs)
• Optimization Over Trained Neural Networks: Taking a Relaxing Walk
• Learning Gradient Descent: Better Generalization and Longer Horizons
• Open Problems in Applied Deep Learning
• LDAdam: Adaptive Optimization from Low-Dimensional Gradient Statistics