Error Feedback for Muon and Friends

2025-10-01

AI Summary

To address the lack of theoretical convergence guarantees and communication efficiency in distributed non-Euclidean optimizers (e.g., Muon, Scion, Gluon), this paper proposes EF21-Muon—the first distributed non-Euclidean LMO optimizer supporting bidirectional gradient compression and error feedback. Its key contribution is the rigorous integration of error feedback into the non-Euclidean optimization framework, establishing convergence under $(L^0, L^1)$-smoothness and hierarchical generalized smoothness assumptions. EF21-Muon unifies and extends existing approaches, seamlessly accommodating both stochastic gradients and momentum. Experiments on NanoGPT demonstrate a 7× reduction in communication volume with zero accuracy loss, convergence speed matching state-of-the-art Euclidean methods, and further acceleration achievable via norm selection.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Search and Optimization: Non-convex Optimization
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization

Abstract

Recent optimizers like Muon, Scion, and Gluon have pushed the frontier of large-scale deep learning by exploiting layer-wise linear minimization oracles (LMOs) over non-Euclidean norm balls, capturing neural network structure in ways traditional algorithms cannot. Yet, no principled distributed framework exists for these methods, and communication bottlenecks remain unaddressed. The very few distributed variants are heuristic, with no convergence guarantees in sight. We introduce EF21-Muon, the first communication-efficient, non-Euclidean LMO-based optimizer with rigorous convergence guarantees. EF21-Muon supports stochastic gradients, momentum, and bidirectional compression with error feedback-marking the first extension of error feedback beyond the Euclidean setting. It recovers Muon/Scion/Gluon when compression is off and specific norms are chosen, providing the first efficient distributed implementation of this powerful family. Our theory covers non-Euclidean smooth and the more general $(L^0, L^1)$-smooth setting, matching best-known Euclidean rates and enabling faster convergence under suitable norm choices. We further extend the analysis to layer-wise (generalized) smoothness regimes, capturing the anisotropic structure of deep networks. Experiments on NanoGPT benchmarking EF21-Muon against uncompressed Muon/Scion/Gluon demonstrate up to $7 imes$ communication savings with no accuracy degradation.

Problem

• Develops distributed optimizer for non-Euclidean LMO methods
• Addresses communication bottlenecks in large-scale deep learning
• Provides convergence guarantees for bidirectional compression with error feedback

Innovation

• EF21-Muon enables communication-efficient non-Euclidean optimization
• It supports bidirectional compression with error feedback
• The method provides rigorous convergence guarantees for distributed training