Error Feedback for Muon and Friends
2510.00643v1
cs.LG, math.OC, stat.ML
2025-10-04
Авторы:
Kaja Gruntkowska, Alexander Gaponov, Zhirayr Tovmasyan, Peter Richtárik
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.