Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Error Feedback for Muon and Friends

Kaja Gruntkowska, Alexander Gaponov, Zhirayr Tovmasyan, Peter Richtárik

TL;DR

This work addresses the challenge of scalable, geometry-aware distributed optimization for non-Euclidean, LMO-based neural network training. It introduces EF21-Muon, the first communication-efficient, non-Euclidean LMO-based optimizer with convergence guarantees, integrating error feedback and bidirectional compression within a layer-wise framework. Theoretical results cover deterministic and stochastic settings under standard and $(L^0,L^1)$-smoothness assumptions, matching Euclidean EF21 rates in key regimes and enabling faster convergence under suitable norms. Empirically, EF21-Muon achieves up to $7\times$ communication savings on NanoGPT-FineWeb benchmarks with no loss in accuracy, illustrating practical impact for ultra-scale, structure-aware training.

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\times$ communication savings with no accuracy degradation.

Error Feedback for Muon and Friends

TL;DR

This work addresses the challenge of scalable, geometry-aware distributed optimization for non-Euclidean, LMO-based neural network training. It introduces EF21-Muon, the first communication-efficient, non-Euclidean LMO-based optimizer with convergence guarantees, integrating error feedback and bidirectional compression within a layer-wise framework. Theoretical results cover deterministic and stochastic settings under standard and -smoothness assumptions, matching Euclidean EF21 rates in key regimes and enabling faster convergence under suitable norms. Empirically, EF21-Muon achieves up to communication savings on NanoGPT-FineWeb benchmarks with no loss in accuracy, illustrating practical impact for ultra-scale, structure-aware training.

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 -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 communication savings with no accuracy degradation.

Paper Structure

This paper contains 53 sections, 26 theorems, 202 equations, 8 figures, 4 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Scaling Up
  3. Communication: the Cost of Scale
  4. Distributed Muon: Bridging the Gap
  5. Outline
  6. Background
  7. What is Muon?
  8. The challenges of distributing the LMO.
  9. Compression.
  10. Error Feedback.
  11. Non-Euclidean Distributed Training
  12. The algorithm.
  13. Convergence Results
  14. Deterministic Setting
  15. Stochastic Setting
  16. ...and 38 more sections

Key Result

Theorem 3

Let Assumptions as:lower_bound and as:smoothness hold. Let $\{X^k\}_{k=0}^{K-1}$, $K \geq 1$, be the iterates of alg:ef_gluon_det (with $p=1$) initialized with $X^0 = W^0$, $G_j^0=\nabla f_j(X^0)$, $j\in[n]$, and run with $\mathcal{C}^k \in \mathbb{B}(\alpha_P)$, $\mathcal{C}_j^k \in \mathbb{B}_\sta

Figures (8)

  • Figure 1: Left: Test loss vs. $\#$ of tokens processed. Right: Test loss vs. $\#$ of bytes sent to the server from each worker normalized by model size to reach test loss $3.31$. Rank$X\%$/Top$X\%$ = Rank$K$/Top$K$ compressor with sparsification level $X\%$; ID = no compression.
  • Figure 2: Trade-off between token efficiency and communication cost for different compression setups at a target test loss of $3.31$.
  • Figure 2: Communication cost per round (in bytes), normalized relative to the identity compressor.
  • Figure 3: Learning rate ablation. The grid spans from the optimal learning rate of the non-compressed baseline, $3.6 \times 10^{-4}$ (denoted as $1.0\times$), down to $0.1\times$. Red curves correspond to experiments processing 5B tokens (\ref{['sec:experiments']}), while blue curves correspond to 2.5B tokens (\ref{['app:2.5B_exps']}).
  • Figure 4: Test loss vs. $\#$ of tokens processed. Top$X\%$/Rank$X\%$ = Top$K$/Rank$K$ compressor with sparsification level $X\%$; ID = no compression. "+ Natural" corresponds to applying Natural compression after Top$K$/Rank$K$ compressor. Experiments use a tokens budget of 5B.
  • ...and 3 more figures

Theorems & Definitions (62)

  • Definition 1: Contractive compressor
  • Remark 2
  • Theorem 3
  • Theorem 4
  • Theorem 5
  • Theorem 6
  • Remark 7
  • Definition 8: Deterministic Damping
  • Definition 9: Random Dropout
  • Definition 10: Top$K$ SVD compressor
  • ...and 52 more