Error Feedback for Muon and Friends
Authors: Kaja Gruntkowska
Alexander Gaponov
Zhirayr Tovmasyan
Peter Richtárik
Venue: ICLR 2026 Poster (2026)
Keywords: optimization
communication efficiency
compression
error feedback
Primary area: optimization

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

Published as a conference paper at ICLR 2026
ERROR FEEDBACK FOR MUON AND FRIENDS
Kaja Gruntkowska, Alexander Gaponov, Zhirayr Tovmasyan & Peter Richtárik King Abdullah University of Science and Technology (KAUST) Thuwal, Kingdom of Saudi Arabia
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 tradi- tional algorithms cannot. Yet, no principled distributed framework exists for these methods, and communication bottlenecks remain unaddressed. The very few dis- tributed variants are heuristic, with no convergence guarantees in sight. We intro- duce EF21-Muon, the first communication-effiotte optimizer with rigorous convergence guarantees. EF21-Muon supports stochas- tie 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 off and specific norms are cho- sen, providing the first efficient distributed implementation of this powerful fam- ily. Our theory covers non-Eucliann Euclidean smooth and the more general (L", -'>- smooth setting, matching best-known Euclidean rates and enabling faster conver- gence under suitable norm choices. We further extend the analysis to layer-wise (generalized) smoothness regimes, capturing the anisotropic structure of deep net- works. Experiments on NanoGPT benchmarking EF21-Muon against uncom- pressed Muon/Scion//luun demonstrate up to 7x communication savings with no accuracy degradation.
Euclidenn LMO-based
INTRODUCTION
Over the past decade, Adam and its variants (Kingma Ba, 2015; Loshchilov Hutter, 2019) have established themselves as the comerrtone of optimization in deep learning. Yet emerging evidence suggests that this dominance may be giving way to a new class of optimizers better suited to the geometry and scale of modern deep networks. Leading this shift are Muon (Jordan al., 2024) and methods inspired by it-Scion (Pethick et al., 2025b) and Gluon (Riabinin et al., 2025b) replace Adam's global moment estimation with layer- wise, geometry-aware updates via linear min- imization oracles (LMOs) over non-Euclidean norm balls. Though relatively new, these optimizers already gaining traction-supported by growing body of theoretical insights, community adop- tion. and empirical success particularly in training large language models (LLMs) (Liu et al., 2025; Pethick et al., 2025b; Shah et al.. 2025; Thérien et al., 2025; Moonshot AI, 2025).
Despite this momentum, the development of these algorithms remains less mature than that of more established methods. Significant gaps persist-bor both theory and practice that must be addressed to fully realize their potential and make them truly competitive for the demands of ultra-scale learning.
Scaling Up. Modern machine learning (ML) thrives on scale. Today's 's state-of-the-art models rely on massive datasets and complex architectures, often requiring weeks or even months of training (Touvron et al., 2023; Comanici et al., 2025). This scale imposes new demands optimization methods, which must not only be effective at navigating complex nonconvex landscapes but also efficient in distributed. essource-constrained environments. Since training single machine is no longer feasible (Dean et al., 2012; You et al., 2017), distributed computing has become the default. Mathematically, this task is commonly modeled the (generally non-convex) optimization problem
min [(x) = (00,, Xes
Fx(X) = -EEyD, (X,(X;
(1)

---

Published as a conference paper at ICLR 2026
Table 1: Summary convergence guarantees. Algorithm: Deterministic EF21-Muon with deterministic gradients (Algorithm 2). Stochastic = EF21-Muon with stochastic gradients (Algorithms 1 and 3); Smooth: / (layer-wise) smooth setting (Assumptions 3 and 6). x= (layer- x wise) generalized smooth setting (Assumptions 4 and 8); Rate = rate of convergence to achieve min.=o.....K F. V e; Fuel. recovers the state-of-the-art guarantees in the Fuclidean case; Non-comp. guarantees.
recovers the state-of-the-art uncompressed
Algorithm
Result
Layer-wise
Smooth
Rate
Eucl.
Non-comp.
Theorem 3
Theorem 14
Deterministic
Theorem 4
K1/2
X
Theorem 17
Theorem 5
/
Stochastic
Theorem 19
O
1
Theorem 6
X
K1/4
X
Theorem 24
where X E represents the model parameters, n 2 the number of worker//cients/machines
n2. and (X) is the loss of the model (X) on the data (D₂) stored on worker j E [n] := nz]. We consider the general heterogeneous setting, where the local objectives fy may differ arbitrarily across machines, reflecting real-world scenarios such multi-datacenter pipelines federated learning (McMahan et al., 2017; Konečný et al., 2016). Here, S is a d-dimensional vector space equipped with inner product (+-) S S + IR and the standard Euclidean norm Il- Furthermore, we endow S with an arbitrary norm ||-1| : S - Rzz- The corresponding dual norm II- S + Rz0 defined via || == s (X, Z). The general framework introduced in this work gives rise variety of interesting algorithms arising from different norm choices. In matrix spaces, a particularly important class is the family of operator norms, defined for any A E mR by ||All,-,, == sup z ||
where Il-Il. and ||-11/ are some norms R and E" respectively.
Communication: the Cost of Scale. In client-server architectures, coordination is centralized, with workers performing local computations and periodically synchronizing with the coordinator (Seide et al., 2014; Alistarh et al., 2017; Khirirat et al., 2018; Stich et al., 2018; Mishchenko et al., 2019: Karimireddy et al., 2019; Mishchenko al.. 2024). While this distributed design unlocks learning at unprecedented scales, it introduces a critical bottleneck: communication. The mas- sive size of modern models places a heavy burden on the channels used to synchronize updates across machines, as each step requires transmitting large d-dimensional vectors (e.g., parameters or gradients) over links that can be far slower than local computation (Kairouz et al., 2021). With- out communication-effin mately limiting the efficiency and scalability of distributed optimization.
strategies, this imbalance makes communication dominant cost, ulti-
Bridging the Gap. The case for communication-efiicieentt distributed training is clear, as is the promise of Muon for deep learning. The natural question is: can we merge the two? Perhaps surprisingly, this intersection remains largely unexplored. Nonetheless, three recent efforts worth noting. Liu et al. (2025) propose a distributed variant of Muon based on ZeRO-1 (Rajbhan- dari et al., 2020). Thérien et al. (2025) show that Muon can be used instead of AdamW as the inner optimizer in DiLoCo. The introduced MuLoCo framework is shown to consistently converge faster than the original DiLoCo (Douillard et al., 2023) when pre-training a 220M parameter transformer language model. In parallel, Ahn al. (2025) introduce Dion, a Muon-inspired algorithm compatible with 3D parallelism that employs low-rank approximations for efficient orthonormalized updates.
Distributed Muon:
While promising empirically, these approaches lack any formal theoretical guarantees. Our goal is to bridge this gap by developing a distributed optimizer leveraging non-Euce un-Euulidean geometry that both works in practice and comes with strong convergence guarantees. Our central question
Can we efficiently distribute Muon without compromising its theoretical and practical benefits?
In this work, we provide affirmative answer through the following contributions:
2

---

Published as a conference paper at ICLR 2026
1. A framework for compressed non-Euclidean distributed optimization. We propose EF21- Muon, an LMO-based distributed optimizer based on bidirectionally compressed updates with error feedback (Seide al., 2014; Richtárik et al., 2021). It is ceaammiication-fffiiiee (never sending uncompressed messages) and practical, supporting stochasticity and momentum. Parameterized by the norm in the LMO step, EF21 recovers broad class of compressed methods, and for spectral norms yields the first commmnication-efficee distributed variants of Muon and Scion.
2. Practical deep learning variant. The main body of this paper presents a simplified version of EF21-Muon that treats all parameters jointly (Algorithm 1), consistent with standard theoretical ex- position. Our main algorithms, however, are designed for and analyzed in a layer-wise manner (see Algorithms 2 and 3 for the deterministic and stochastic gradient variants. respectively), explicitly modeling the hierarchical structure of neural networks. This allows us to better align with practice (methods like Muon are applied per layer) and to introduce anisotropic modeling assumptions.
Strong convergence guarantees. EF21 -u2n-Muon comes with strong theoretical guarantees (see Table 1) under two smoothness regimes: non-Euclidean smoothness (Theorems 3 and 5) and non- Euclidean (L", -mmoothhssssss (Theorems 4 and 6). both cases, our bounds match the state- of-the-art rates for EF21 in the Euclidean setting, while allowing for potentially faster convergence under well-chosen norms. These results are subsumed by a more general analysis of the full layer- wise methods. In Theorems 14 and 19, prove convergence under layer-wise non- non-Euclidean smoothness (Assumption 6), and extend this to layer-wise non-Fuclidean (T,, I. >-smoohhness (As- sumption 8) in Theorems 17 and 24. This refined treatment allows to better capture the geometry of deep networks, leading to tighter guarantees.
4. Non-Euclidean compressors. EF21-Muon supports standard contractive compressors as well as a new class of non-Euclidean compressors (Section D), which may be of independent interest.
5. Strong empirical performance. Experiments training NanoGPT model on the F neWeb dataset systematically compare EF21-Muon with multiple compressors against the uncompressed baseline (Muon/Scion/Gluon) and show that compression reduces worker-to-server communication by up to 7x with no loss in accuracy (Sections 5 and G).
Outline. Section 2 introduces the necessary preliminaries and reviews Muon (Jordan et al., 2024), placing within the broader class of LMO-based optimizers. This naturally raises the central ques- tion of our work: how can such methods be distributed efficiently? We highlight the main challenges and motivate compression and error feedback as practical solutions (with deeper motivation and an extensive literature review deferred to Section A). Section 3 presents our proposed method, EF21- Muon. In Section 4, we present convergence results in both deterministic and stochastic settings, under two smoothness regimes: standard (non-Fuulldean) and (T,", )-moootthsssss Finally, Sec- tion 5 provides empirical validation, demonstrating the practical benefits of our approach.
2 BACKGROUND
We frame problem (1) in an abstract vector space S. In several of our results, the specific structure of S does not matter. One may simply flatten the model parameters into a d x 1 vector and view S as RRd However, in the context of deep learning, it often useful to explicitly model the layer-wise structure (see Section B). Then, X E S represents the collection of matrices X, E S == Rm, Xn, of trainable parameters across all layers i E [p] of the network with a total number d == Σp= m.n. of parameters. Accordingly,
the d-dimensional product space S = _ S. = S1 @ 8 Sp- where each S; is associated with the trace inner product (X,, Yi)(0) == ir(X, Y.) for X,, Y2 E S., and a norm [1-11()) (not necessarily induced by this inner product). We write X =
Xp).
What is Muon? Muon, introduced by Jordan et al. (2024), is an optimizer for the hidden layers of neural networks.' For clarity of exposition, let us assume that the parameters X represent a single layer of the network (a full layer- wise description is provided in Section B. 1). In this setting, Muon updates x X (VK)T. T, where t > 0 and the matrices U V are derived from the
'The first and last layers are typically optimized using other optimizers, such as AdamW (Loshchilov & Hutter, 2019)-sce Section B. for details.
3

---

Published as a conference paper at ICLR 2026
Algorithm EF21-Muon (simplified)
1: Parameters: radii 1* > O; momentum parameter B E (0, 1]; initial iterate X E S (stored on the server); initial iterate shift W = X (stored on the server and the workers); initial gradient estimators G3 (stored on the workers); G = 1 Lj= j=1 Gj (stored on the server); initial momentum My (stored on the workers); worker compressors Ch: server compressors C*
2: for k = 0, 1..... ........ K do
3:
x LMO LLMO B(X*, ( (G*)
Take LMO-typ step
4:
sh= (X
W
Compress shifted model on the server Update model shilt
5:
Wh+1 = Wh + 5*
6:
Broadcast S to all workers
7:
for j 1 ..... in parallel do
8:
W = W + 5
Update model shift
MA++ =(1 -11- B)M5 + (W 15 )
9:
Compute momentum Compress shifted gradient
10:
R = C* (M G
11:
G G. R8tt1
12:
Broadcast R to the server
13: end for
14: Gk+1 1 -/ Gftt Gk + E R
Compute gradient estimator
15:
end for
SVD of the momentum matrix G = (V*)) This update rule is, in fact, special case of a more general one, based on the norm-constrained linear minimization oracle (LMO)
X x L*LMOB(0,1) (G*).
(2)
where t( == {z S: 2 x|| S t} and LMO (G) == arg min zeB(.)) (G. Z). Muon corresponds to the case where ||-|1 = 11-112--2 is the spectral (operator) norm, in which case LMO (G*) = -U Consequently, its recent analyses (Pethick al., 2025b; Kovalev, 2025: Riabinin al., 2025b) have shifted focus the general form (2). Among them, Pethick al. (2025b) introduce Scion, which extends the LMO update across layers, and Riabinin et al. (2025b) develop Gluon-a general LMO-based framework that subsumes Muon and Scion special cases while providing stronger convergence guarantees. We adopt this unifying viewpoint by treating all three algorithms instances of Gluon, which we the umbrella term for the entire class.
The challenges of distributing the LMO. Distributing (2) is far from trivial, as the limited liter- ature suggests. Even in the relatively well-structured special case of spectral norms, Muon relies on the Newton-Schut iteration (Kovarik, 1970: Björck & Bowie, 1971), a procedure requiring dense matrix operations that are incompatible with standard parameter-shardin schemes used in LLM training (Ahn et al., 2025). To illustrate the difficulty, consider deterministic version of (2), where G* is replaced by the exact gradient vf((*). Applied problem (1), the iteration becomes
X++
X LMO
1
Σ j=1
T
(f)(
The most basic approach to distributing this update consists the following four main steps:
1. Each worker computes its local gradient (xx*) at iteration k.
2. w2s: The workers send their gradients V (X++)) to the central server.
The server averages these gradients and computes the LMO update.
4.
s2w: The server sends X (or LMO (-)) back to the workers.
This scheme involves two potentially costly phases: (1) workers-to-server (= w2s) and (2) server- to to-workers (= s2w) communication. As each transmitted object resides in S, every iteration in- volves exchanging dense, d-dimensional data, imposing substantial communication overhead that quickly overwhelm available resources. This is where compression techniques come into play.
Compression. Compression is one of the two main strategies for improving communication effi- ciency in distributed optimization (the other being local training (Povey al., 2014; Moritz al.,
4

---

Published as a conference paper at ICLR 2026
2015: McMahan et al., 2017)), extensively studied in the Euclidean regime (Alistarh et al., 2017; Horváth et al., 2022; Richtárik et al., 2021). It is typically achieved by applying an operator C map- ping the original dense message X to more compact representation C(X). We work with general biased (or contractive) compressors.
Definition (Contractive compressor), A (possibly randomized) mapping C S - is a contrac- five compression operator with parameter a E (0, if
E [|c(()) - xII| a1
VX
Remark 2. The classical definition of contractive compressor is based on the Euclidean norm, i... -1-lll 21-1l2220 in (3). A canonical example in this setting is the Top compressor, which retains the K largest-magriude entries of the input vector. In (3), we generalize this arbitrary norms for greater flexibility. Section D provides examples of such compressors (to our knowledge, not studied in this context before). Depending on the compression objective, we apply (3) with respect to II-I1,
I1-Il, or 11- denoting the respective families of compressors as B(x)) I8 (a). and (a))).
Error Feedback. To address the communication bottleneck, a natural approach is to apply biased compressors to transmitted gradients. However, this "enhancement" can result in exponential di- vergence, even in the simple case of minimizing the average of three strongly convex quadratics (Beznosikov et al., 2020, Example 1). A remedy, Error Feedback (EF), was introduced by Seide et al. (2014) and for years remained a heuristic with limited theory. This changed with Richtárik al. (2021), who proposed EF21, the first method to achieve the desirable (1//K) rate for ex- pected gradient norms under standard assumptions. Since then, EF21 has inspired many extensions, including EF21-P (Gruntkowska al., 2023), a primal variant targeting s2w communication.
For deeper dive into compression and EF, we refer the reader to Appendices A.1 and A.2.
3 NON-EUCLIDEAN DISTRIBUTED TRAINING
Marrying geometry-aware updates of Giuon with the communication efficiency enabled by compres- sion promises potentially high-yield strategy. Yet, from a theoretical standpoint, their compatibility is far from obvious-nothing a priori ensures that these two paradigms can be meaningfully unified.
Most importantly, it is unclear what kind of descent lemma to use. The analysis of EF21 relies on a recursion involving squared Euclidean norms ||-1|2- while LMO-based methods naturally yield descent bounds in terms of first powers of norms ||-1|-a structure common to all existing analyses (Kovalev, 2025; Pethick et al., 2025b; Riabinin et al., 2025b). We initially adopted the latter ap- proach, but the resulting guarantees failed to recover those of EF21 in the Euclidean case. The pivot point came from reformulating the update via sharp operators (Nesterov, 2012; Kelner et al., 2014). For any G E S, the sharp operator is defined as G# == arg x (G, - 1 ||]]] which is connected to the LMO via the identity ||G'], LMO (G) = -G Hence, (2) is equivalent to
x x (LM00 (G*) X
||.. (G*)*
1
i.., a normalized steepest descent step with stepsize 2* := /|||=||, We alternate between the sharp operator and LMO formulations, depending the assumptions play. Theorems 3 and 5 use the former; Theorems 4 and 6, the latter. We explore this and other reformulations in Section C.
The algorithm. Working with compression in -uclideen geometry presents several chal- lenges. In addition to the lack of a standard descent lemma, further complications arise from in- teractions between gradient stochasticity and compression and unknown variance behavior under biased compression. Yet, we develop the first communication-effie tension, its special cases Muon and Scion), called EF21 that combines biased compression, gradient stochasticity, and momentum, all while enjoying strong theoretical guarantees. A simpli- fied version of the algorithm, applied globally to X, is shown in Algorithm 1. A more general, deep learning-orien oriented layer-wise variant operating in the product space S 8 - mm,x", is given in Algorithm 3. For clarity, we focus on the simplified variant throughout the main text; all theoretical results presented here are special cases of the general layer-wise guarantees provided in Section E.
variant of Gluon (and by ex-
5

---

Published as a conference paper at ICLR 2026
While the pseudocode is largely self-explanatory (for more detailed description, Section B.2), we highlight the most important components:
Role of Compression. Compression is key for reducing communication overhead in distributed training. Algorithm adheres to this principle by transmitting the compressed messages S and R
only, never the full dense updates. When compression is disabled (i.e., C.. C* are identity mappings) and in the single-node setting (n = 1), EF21 MMunn reduces exactly to Gluon, which turn recovers Muon and Scion (all of which were originally designed for non-distributed settings).
Role of Error Feedback. Even in the Euclidean setup, biased compression can break distributed GD unless some form of error feedback is used (see Section 2). To remedy this, we adopt a modern strategy inspired by EF21 (Richtárik et al., 2021) for the w2s direction. Its role is to stabilize training and prevent divergence. To reduce s2w communication overhead, we further incorporate the primal compression mechanism of EF21-P (Gruntkowska et al., 2023).
Role of Gradient Stochasticity. In large-scale ML, computing full gradients (z) is typically computationally infeasible. In practice, they are replaced with stochastic estimates, which drastically reduces per-step computational cost and makes the method scalable practical workloads.
Role of Momentum. Stochastic gradients inevitably introduce noise into the optimization pro- cess. Without further stabilization, this leads convergence to neighborhood of the solution only. Momentum mitigates this issue, reducing the variance in the updates and accelerating convergence.
4 CONVERGENCE RESULTS
To support our convergence analysis, we adopt standard lower-boundedness assumptions on the global objective f. and in certain cases, also the local functions fz-
Assumption 1. There exist f* E R such that fXX 2 f for all E S. Assumption 2. For all j E [rl,, there exist fj R such that fx(X) 2 f3 for all X E S.
We study two smoothness regimes. The first, standard I-smoothness generalized to arbitrary norms (used in Theorems and 5), the default in virtually all convergence results for Muon and Scion (Kovalev, 2025; Pethick et al., 2025b; Li & Hong, 2025).
Assumption 3. The function is L-smooth, ie.. (xV V)(Y)), S L Y for all X, E S. Moreover, the functions f, are Lm-smooth for all j E [n]]. We define" Í == Σ L3
To our knowledge, the only exception departing from this standard setting is the recent work on Gluon (Riabinin al., 2025b). The authors argue that layer-wise optimizers designed specifi- cally for deep learning, where the classical smoothness assumption is known to fail (Zhang et al., 2020). Instead, they build upon the (L., -mmosmootssssss model introduced by Zhang al. (2020)² (Assumption 4), a strictly weaker alternative motivated by empirical observations from NLP training dynamics. Riabinin et al. (2025b) introduce a layer-wise variant (Assumption 8), arguing that het- erogeneity across network layers requires smoothness constants to vary accordingly. Consistent with this line of work, we provide convergence guarantees under the layer-wise (L", L L')-mmoothness as- sumption (Theorems 17 and 24). For clarity, the main text treats the case of generic vector space S, without delving into the product space formulation (see Section B), in which case the assumption reduces to a non-Fuclidean variant of asymmetric (T.", _T") smmotthess from Chen et al. (2023).
Assumption 4. The function f - is (IO, -mooth,, iee,, there exist L L such that [Vf(()) (())),, (L + L' (V)(())).) ||X -YI VX,Y e S.
Moreover, the functions f3. j E [n), are (LD, L.) -mooshh We define L == maxje(m] L} and L =
Assumption 4 is strictly more general than Assumption 3, as it allows the smoothness constant to grow with the norm of the gradient, a key property observed in deep learning (Zhang al., 2020).
In theoretical results, i. could potentially be improved to the arithmetic mean-see Richtárik et al. (2024). The original (I.", T.') Jmmooohhssssss assumption of Zhang et al. (2020) was defined for twice--iifeeeentiable
functions via Hessian norms. This notion and our Assumption 4 are closely related-see Chen et al. (2023).

---

Published as a conference paper at ICLR 2026
Deterministic setting. As warm-up, we first present the convergence guarantees of Algorithm 2- a deterministic counterpart of Algorithm using deterministic gradients without momentum (though stochasticity may still arise from compression). The first theorem addresses the smooth setting.
Theorem 3. Let Assumptions I and 3 hold. Let {X* I k=0
K-1 K
2 1, be the iterates of Algorithm 2 j E [nl, and run with C+ E +B(ap),
(with; p 1) initialized with X = W., G3 = (X))).
сj E BB, (aD) and 0 < <h = E S 2L + 4/DD /12 12
10/07,7 00/09.1)
TThnn
K Σ
[1.]] S -11(x)-1-)
E
K
Ky
Theorem 3 is a special case of the general layer-wise result in Theorem 14. To our knowledge, no prior work analyzes comparable compressed methods under general non- non-Euclidean geometry. In the Euclidean case, our guarantees recover known results (up to constants): without primal compression (or = 1), they match the rate of Richtárik al. (2021, Theorem 1); with primal compression, they align with the rate of EF21-BC from Fatkhullin et al. (2021, Theorem 21) (though EF21 -Muon and EF21-BC differ algorithmically, and the former does not reduce to the latter in the Euclidean case).
In the generalized smooth setup, we establish convergence without primal compression. However, as we argue in Section D.I, the s2w communication can still be made efficient through appropriate norm selection. Indeed, we find that LMOs under certain norms naturally induce compression- behavior.
Theorem 4. Let Assumptions 1, and 4 hold and let {x=1
1k= K 1, be the iterates of Algo-
rithm 2 (with = 1) initialized with G = (x) j E [n]. and run with C* = II (the identity compressor), Cj5y E (aD)) and th VK+I for some 1) > 0. Then,
10c E 405
min E [[[[][X*)]] [1.] S k=O....K K
exp( (n max)
121
nVK+
VK+1
where 80 : -(x) C L' 2 2/I-a and D = 1 1-VI-aD
MT-ap 2 1-vI-ap
Theorem 4, a corollary of the layer-wise result in Theorem 17, achieves the same desirable O(1 rate for expected gradient norms as Theorem 3, but with radii t that are independent of problem- specific constants. If smoothness constants are known in advance, they can be incorporated into the choice of 1) improve the dependence on these constants in the final rate. In the Fuclidean case, our guarantee matches that of || |EF21|| || under (L", smoothness established by Khirirat et al. (2024).
Stochastic setting. We now turn to the convergence guarantees of our practical variant of EF21- Muon (Algorithms 1 and 3), which incorporates noisy gradients and momentum. We assume access to standard stochastic gradient oracles VJz(; &j ~ Dj with bounded variance.
Assumption 5. The stochastic gradient estimators (((()))) : S + S are unbiased and have bounded variance. That is, BE,,D, IV (X;E)) = fx(X) for all X E S and there exists 2 0 such that Be -VVL((:)))V((()I
o² for all X E
Note that the variance bound in Assumption 5 is expressed in terms of the Euclidean norm rather than II-II to facilitate the bias-variance decomposition. Nevertheless, since S is finie-dimennsionll, the magnitudes measured in I1-112 can be related to quantities measured in ||-I| via norm equivalence. That is, there exist p.P > 0 such that e || < || < أ |||x|| for all X E S.
As in the deterministic setting, we begin by analyzing the smooth case.
Theorem 5. Let Assumptions 1, 3 and 5 hold. Let [X*]K-1 K > be the iterates of Al- Jk-o
gorithm 1 initialized with X W G9 MD (x, (9). j [n], and run with
C² E B(ap). Cj Bz(av). any B E 1]. and 0 < 26 7 S (2V5+2L) -1 where
1² s= o²
g2 1,2
24(3+2) 1 (6( i + (2945(22+5) i? Then azab
of
aB
K-1
Σ k=o
E |[][++]|]]
15th =
24
+
128 of
op + 24
1 + (1-ap)
+ 128 o
o 32
K
Vn3
7

---

Published as a conference paper at ICLR 2026
where 8 : f(X²) f
Theorem is a special case of Theorem 19. 4/1 La Ø-
Choosing Y
(2V5+ + 2L)
and B
min
1.
8 4/2 x
Lan p
4/a
it guarantees that
K- Σ [|] k=0
OR D p'apank
5 p g
1/2
80 Soo p'vank
4/3
88 1000/102101
3/4
K
+
+
(see Corollary 2). the absence of stochasticity and momentum (o² 0, = =), Algorithm reduces to Algorithm 2 (with/ p 1), and the guarantee in Theorem 5 recovers that of Theorem 3, up to constants (Remark 22). In the Euclidean case without primal compression (7² p2 ap = 1). Theorem 5 matches the rate of EF21-SDGM established by Fatkhullin et al. (2023, Theorem 3). again up to constants (Remark 21). Finally, one may employ compressors C* E (B2(ap) instead of Ch E B (ap)), though this introduces an additional dependence on p² in the constant (Remark 23).
As in Theorem 4, in the (L", LL' -moothh setup, we set C² = I.
Theorem 6. Let Assumptions 1, 2, 4 and hold. Let {X*} KZ 2 1, be the iterates of Algo- k=0
rithm I initialized with MD (x; -g)) G. (VS (.)). j E [n], and run with C I (the identity compressor), C E (ap) 3 = 1/(K+)212 and S th = t 5/(K+19/14 where
(K+1)2/2 (1-1)" 6(L
n < min
Then
24v1 -a aapp (
min
E
3(((x)-1'
71.
16 appo
8po
k=O.....K K
+K+1
(K+1)374
+1(K+
8vi-ap 2 (K+1) (I-VI-aD)(K+1))7
(1)
Analogously Theorem Theorem 6 (a corollary of Theorem 24) establishes an O(/K'1)) convergence rate, matching state-of-the-art guarantees for SGD-type methods in the non-convex setting (Cutkosky Mehta, 2020; Sun et al., 2023). Among the terms with the worst scal- ing in K, """(""""""(((("""""""""" is standard and reflects the impact of the initial suboptimality. 8p0/ / captures gradient stochasticity, scaling linearly with the standard deviation a, but decaying with the square root of the number of clients n. The term 1 (E;=((,)) (f* ff) quan- To
tifies client heterogeneity and vanishes when local optima fj coincide with the global minimum f*, and otherwise scales with the local smoothness constants L.J. All compression-dres -driive error terms vanish when compression is disabled (a = 1). Finally, in the Euclidean case (p² = p² = 1), the rate recovers that of |EF21-SDMM] from Khirirat al. (2024, Theorem 2), up constants.
5 EXPERIMENTS
We present key experimental results below, with additional details and extended experiments avail- able in Section G.
Experimental setup. All experiments are conducted on 4 NVIDIA Tesla V100-SXM2-32GB GPUs 4 NVIDIA A A100-SXM4-80GB in a Distributed Data Parallel (DDP) setup. The dataset is evenly partitioned across workers, with one worker node acting as the master, aggregating com- pressed updates from the others. Training and evaluation are implemented in PyTorch," extending open-source codebases (Pethick et al., 2025a; Riabinin et al., 2025a; Karpathy, 2023).
We train a NanoGPT model (Karpathy, 2023) with 124M parameters on the FineWebl0K dataset (Penedo et al., 2024), using input sequences of length 1024 and a batch size of 256. Optimization performed with EF21-Muon using spectral norm LMOS for hidden layers and f.xe norm LMOs for embedding and output layers (which coincide due to weight sharing), following the approach of Pethick et al. (2025b). For spectral norm LMOs, inexact updates are computed with New- ton-Schulz iterations (Kovarik, 1970; Björck & Bowie, 1971), as in Jordan et al. (2024).
*Code for experiments is available here. 'PyTorch Documentation: https: http:/yytttttccccccc
e/ htm
8

---

Published as a conference paper at ICLR 2026
Figure 1: Left: Test loss vs. # of tokens processed. Right: Test loss vs. # of bytes sent from each worker to the server normalized by model size to reach test loss 3.31. Kank/Top.X'> = RRank/Top compres- with sparsification level X%; ID = no com- pression.
3.0
Ranuml
5.
17
Tepl
3.6
Tolorns
Cummuncatioon Coul
3.00
Following common practice in communication compression literature, we assume that broadcasting free and focus on w2s communication. Thus, the server-side compressor is fixed to Z, while worker compressors vary among TopK, RankK (Safaryan et al., 2021), Natural compressor (Horváth ct al., 2022) and combinations thereof: TopK + Natural com- pressor of selected elements, and Rank Natural com- pressor applied to all components of the low rank decom- position. These are tested under multiple compression levels and compared against an uncompressed baseline (i.e., stan- Section 3). Learning rates tuned
Fankl0%
1.6
Top15% Natural
@op115%
4.0
4.4
Rankk1 5% + Natural
428
406
Top20%
e Fank15%
2.0
Natural
dard Scion/Gluon; per optimizer and experimental setting, initialized from the values in the Gluon repository (Riabinin et al., 2025a) (see Section G.3). We adopt the same learning rate scheduler as Karpathy (2023) and fix the momentum parameter to 0.9. Model and optimizer hyperparameters are summarized in Ta- bles 3 and 5, respectively.
10%
1AK
Communication
Figure 2: Trade-off between token efficiency and communication cost for different compression setups at a target test loss of 3.31.
Results. For Rank and TopK compressors, evaluate multiple compression levels (in plots, MonkX%/TooXX denotes a Rank RankK/TopK compressor with compression level X%). We report experimental results for 5B-token training budget (> 40x model size) Figure (left), and to reach a strong loss threshold of 3.31 in Figures (right) and 2.
Table 2: Communication cost per round (in bytes), normalized relative the identity compressor.
The number tokens required to reach a target loss de- pends on the compressor. Figure provides a compari- son of the numbers of tokens used in the training run to reach a strong test loss threshold of 3.31 plotted against the communication cost (reported as the number of bits transmitted from each worker to the server normalized by the model size), plotted against the w2s communication cost. Shorter 2.5B-token runs are reported in Section G.5 to assess performance under limited training budgets.
Compressor
Relative Cost
ID
1,0000
Natural
0,5000
Rank20%
0.2687
Rank15%
0.2019
In Figure 1, we plot test loss vs. tokens processed, as well the w2s communication cost required reach the 3.31 loss threshold. For each compressor, we report its most competitive configuration (see Section G.4 for detailed ablation). As expected, compression slows convergence in terms of number of training steps, but substantially re- duces per-step communication cost (Table 2). Overall, this yields significant communication savings-as up to
R5nk Natural
0.1010
Rank 10%
0.1335
Rank Natural
0.0667
Rank5%
0.0667
Top20%
0.3625
Top15%
0.2718
1op15% + Natural
0.1969
Top10%
0.1812
7x for Rank15% Natural compressor, and roughly 4x for Top15% Natural compressor- relative the un- compressed baseline.
Top10% + Natural
0.1312
Top5%
0.0906

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
The research reported in this publication was supported by funding from King Abdullah University of Science and Technology (KAUST): i) KAUST Baseline Research Scheme, ii) CRG Grant ORFS- CRGI2- 6-2044004666 and iii) Center of Excellence for Generative AI, under award number 5940.
For computer time, this research used the resources of the Supercomputing Laboratory at King Abdullah University of Science & Technology (KAUST) in Thuwal, Saudi Arabia.
REFERENCES
Kwangjun Ahn, Byron Xu, Natalie Abreu, and John Langford. Dion: Distributed orthonormalized updates. arXiv preprint arXiv:2504.0522 2025. URL https:t/aaxiiii org/ abs 2504. 05295. (Cited page 2 and 4)
Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. QSGD:
Communication-efficcse SGD via gradient quantization and encoding. Advances in neural in- formation processing systems, 30, 2017. URL https: //arai///axvvv./ http://xxxxv.....abbba. wwtp/bagg//eeeeeeeeeeeeeeeeeeeee 02132.
(Cited on page 5, 19, and 27)
Dan Alistarh, Torsten Hoefler, Mikael Johansson, Nikola Konstantinov, Sarit Khirirat, and Cédric Renggli. The convergence of sparsified gradient methods. Advances in Neural Information Processing Systems, 31, 2018. URL https: /tappp:///axxvv... aaaxx.xxiii//ii////aa////aaaa/ra owg/o////oooooo abspaa/.....1.. 10505. (Cited on page 19)
Jeremy Bernstein, Yu-Xiang Wang. Kamyar Azizzadenesheli, and Animashree Anandkumar. signSGD: Compressed optimisation for non-convex problems. In International Conference on Machine Learning, pp. 560-569. PMLR, 2018. URL https:://// 04434. (Cited on page 21)
Aleksandr Beznosikov, Samuel Horváth, Peter Richtárik, and Mher Safaryan. On biased com- pression for distributed learning. arXiv preprint arXiv:2002. 12410, 2020. URL https: //trxiv.org/ abs 2002000022220022222002020 12410. (Cited page 19, and 27)
A. Björck and C. Bowie. An iterative algorithm for computing the best estimate of an orthogonal matrix. SIAM Journal on Numerical Analysis, 8()):388-364, 1971. URL https: / //doi. org/ 10.1137/070800 (Cited on page 4, 8, and 24)
Ziyi Chen, Yi Zhou, Yingbin Liang, and Zhaosong Lu. Generalized smooth nonconvex optimiza- tion is as efficient as smooth nonconvex optimization. In International Conference on Machine Learning, pp. 5396-5427. PMLR, 2023. URL pttettss://ax/te http:///aawww.o/aa////////so (Cited on page 6 and 21)
wwwp:/rgg///aooe..o.oo... 2303 02854.
Gheorghe Comanici, Eric Bieber, Mike Schackermann, Ice Pasupat, Noveen Sachdeva, Inderjit Dhillon, Marcel Blistein, Ori Ram, Dan Zhang, Evan Rosen, al. Gemini 2.5: Pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capa- bilities. arXiv preprint arXiv:2507.06261, 2025. URL https:///aaaetee. org /b50/2575 06261v4. (Cited on page 1)
Laurent Condat, Artavazd Maranjyan, and Peter Richtárik. LoCoDL: Communication-effiieenn dis- tributed learning with local training and compression. arXiv preprint arXiv:2403.04348, 2024. URL https: /arxiv org/ abs/2403 04348. (Cited on page 19)
Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang, Robustness to unbounded smoothness of generalized signSGD. Advances in neural information processing systems, 35:9955- 2022. URL https: //tpxvaavvv/////////////a////aaa org/abs/ 2208 11195. (Cited on page 21 and 22)
Ashok Cutkosky and Harsh Mehta. Momentum improves normalized SGD. In International confer- ence on machine learning, pp. 2260-2266 PMLR, 2020. URL https: /tt///axxxx /ariivvvivvv .org/ alss/ 2002. 53 (Cited page and 38)
10

---

Published as a conference paper at ICLR 2026
Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc' relio Ranzato, Andrew Senior, Paul Tucker, Ke Yang. Quoc Le. and Andrew Ng. Large scale distributed deep networks. In F. Percira, C.J. Burges, L. Bottou, and K.Q. Weinberger (eds.), Advances in Neural Information Processing Systems, volume 25, Curran Associates, Inc., 2012. URL https: http://oooooedeecooeooooooo ngs neuri cc wwtpaper.ppeeeeeeeeeeeeeeeee.. fil.ffiee//eeeee http//ppprrr//00 file/ http://cccc700661111222222366661
Paper paf. (Cited on page 1)
Arthur Douillard, Qixuan Feng. Andrei A Rusu, Rachita Chhaparia, Yani Donchev, Adhiguna Kuncoro, Marc Aurelio Ranzato, Arthur Szlam, and Jiajun Shen. Diloco: Distributed low- communication training of language models. arXiv preprint arXiv:2311.08105, 2023. URL 1w211.11110100 (Cited on page 2)
aitpss://axxaa org/
llyas Fatkhullin, Igor Sokolov, Eduard Gorbunov, Zhize Li, and Peter Richtárik. EF21 with bells & whistles: Practical algorithmic extensions of modern error feedback. arXiv preprint TrXiv:2110.03294 2021. URL https: http///axvvvvvv/// org/ aal/ab////11 /ttp////1l... 03299. (Cited on page 7,
19, 20, and 42)
Ilyas Fatkhullin, Alexander Tyurin, and Peter Richtárik. Momentum provably improves error feed- back! Advances in Neural Information Processing Systems, 36.76444-76495, 2023, URL https:/ /arxivv..../ org/ abss 2305.15155. (Cited on page 8 and 49)
Athanasios Glentis, Jiaxiang Li, Andi Han, and Mingyi Hong. A minimalist optimizer design for LLM pretraining, Xrviv preprint arXiv:2506. 16659, 2025. URLhttps:
abs/2506 16659, (Cited on page 23)
Eduard Gorbunov, Dmitry Kovalev, Dmitry Makarenko, and Peter Richtárik. Linearly converging error compensated SGD, Advances in Neural Information Processing Systems, 33:20889-20900, 2020. URL https: : /ttp//aaxvvvv//////////////a.av. org aaaaaaas/2000 /2000.000000 12292. (Cited on page 19)
Eduard Gorbunov, Konstantin Burlachenko, Zhize Li, and Peter Richtárik. MARINA: Faster non-
convex distributed learning with compression. In 38th International Conference on Machine Learning. 2021. URL https: https/////axxv.o/a/aa/aoao/
2202....2 07845. (Cited on page 19)
Kaja Gruntkowska, Alexander Tyurin, and Peter Richtárik. EF21-P and friends: Improved theo- retical communication complexity for distributed optimization with bidirectional compression. In International Conference on Machine Learning, pp. 11761-11800 PMLR, 2023. URL http 2209. 2//.//299..111 15218. (Cited on page 5, 6, 19, 20, and 24)
Kaja Gruntkowska, Alexander Tyurin, and Peter Richtárik. Improving the worst-case bidirec- tional communication complexity for nonconvex distributed optimization under functioon sim- ilarity. Advances in Neural Information Processing Systems, 01888877-8873,, 2024. URL atttps://axxe org/ abs/ 2402.06612. (Cited on page 19)
Samuel Horváth, Chen-Yo Yu Ho, Ludovit Horvath, Atal Narayan Sahu, Marco Canini, and Peter Richtárik. Natural compression for distributed deep learning. In Mathematical and Scientific Machine Learning, pp. 129-141. PMLR, 2022. URL https:///// http://axvvvvv//a org abs httpg//bs//.....// 1905. 10988. (Cited on page 5, 9, 19, and 27)
Junxian Huang, Feng Qian, Alexandre Gerber, Z Morley Mao, Subhabrata Sen, and Oliver Spatscheck. A close examination performance and power characteristies of 4g Ite networks. In Proceedings of the 10th international conference on Mobile systems, applications, and services, pp. 225-238, 2012. URL https://///. //dl///ll/.. acm. wwtp://o////ooooooooeo 11tp///....111//33. 2307658. (Cited on page 19)
Rustem Islamov, Xun Qian, and Peter Richtárik. and compressed communication. In International conference on machine learning, pp. 4617- 7- 4628. PMLR. 2021. URL https: /ttpsss/////vvv.///ae org/ /ttpg//////////////////////aa 22122210222222 07158. (Cited on page 27)
Distributed second order methods with fast rates
Rustem Islamov, Xun Qian, Slavomír Hanzely, Mher Safaryan, and Peter Richtárik. Distributed newton-type methods with communication compression and bernoulli aggregation. Transactions on Machine Learning Research, 2023. URL https: arxiv. org/ abs /200 03588.
(Cited on page 20)
11

---

(Cited on Published as a conference paper at ICLR 2026
Ruichen Jiang, Devyani Maladkar, and Aryan Mokhtari. Convergence analysis of adaptive gradient methods under refined smoothness and noise assumptions. arXiv preprint arXiv:2406.04592
2024. URL https: http:/axivvv.oo// org/ abs /2406 04592. (Cited page 21 and 22)
Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein Muon: An optimizer for hidden layers in neural networks, 2024. URL https: / //kellerjordan. github. wttptttttsssoooooooooooooooooost /muon/. (Cited on page 1, 3. 8, and 23)
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Ad- vances and open problems in federated learning. Foundations and Trends@ in Machine Learning, 14(1-2):1 2100, 2021. URL https: wat.iaivvvvv.///..///aaa org/ abs 1912.2222 04977. (Cited page
and 19)
Sai Praneeth Karimireddy, Quentin Rebjock, Sebastian Stich, and Martin Jaggi. Error feedback fixes signSGD and other gradient compression schemes. In International conference on machine learning, pp. 3252-3266 PMLR, 2019, URL
(Cited on page and 19)
Andrej Karpathy. nanoGPT, 2023. URL
(Cited on page and 67)
Jonathan A Kelner, Yin Tat Lee, Lorenzo Orecchia, and Aaron Sidford. An almost-linear-time algo- rithm for approximate max flow in undirected graphs, and its multicommodity generalizations. In Proceedings of the twenty-fitth annual ACM-SIAM symposium on Discrete algorithms, pp. 217- 226. SIAM, 2014. URL https: http://///aaa...occa////////222 2333. (Cited on page 5 and 25)
Sarit Khirirat, Hamid Reza Feyzmahdavian, and Mikael Johansson. Distributed learning with com- pressed gradients. arXiv preprint 1806.0880066674 2018. URL httpps//arxiv.orr abs/1880 06573. (Cited on page 2)
Sarit Khirirat, Abdurakhmon Sadiev, Artem Riabinin, Eduard Gorbunov, and Peter Richtárik. Error feedback under (t_0, (...LLL)ssmoossnesss Normalization and momentum. arXiv preprint arXiv:2410. 16871, 2024. URL 8, 20, and 40)
16871, (Cited on page 7,
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015. URL httts:://aatttttet org/ abs /142. 6980. (Cited on page
Anastasia Koloskova, Tao Lin, Sebastian U Stich, and Martin Jaggi. Decentralized deep learn- ing with arbitrary communication compression. arXiv preprins 190.v09907.09356, 2019. URL /ttps://arrtt wwt./bs////aeeoeeeeeeeeeeeeeeeee abs//907//19 110../.9377........
09356. (Cited on page 19)
Jakub Konečný, H Brendan McMahan, Felix X Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency. arXiv preprint arXiv:1610.05492, 2016. URL https: //ts:////xxvv..// httpaaiww......////////aaaaaaa/a 11600.11111..00 05492. (Cited
on page 2)
Dmitry Kovalev.
Understanding gradient orthogonalization for deep learning via non-Euclidean trust-region optimization, 2025. URL https: /attps://arxiv..o. http/axxvvvvrrr.. abt/abbs//250// page 4, 5, 23, and 24)
/ttp///xw.....////////..... 12645. (Cited on
Zdislav Kovarik. Some iterative methods for irpronnng orthonormality. SIAM Journal Numerical Analysis, 7(3):386-382) -898, 1970. URL page 4, 8, and 24)
Haochuan Li, Jian Qian, Yi Tian, Alexander Rakhlin, and Ali Jadbabaie. Convex and non-convex optimization under generalized smoothness. Advances in Neural Information Processing Systems, 0000002384402711 2023. URL https: /ttp//aaiiv///.////////////..aal org/ abs/ abbs/2300 2306 01264. (Cited on page 21)
Jaxiang Li and Mingyi Hong. A note on the convergence of Muon and further, 2025. URL https: //arxiv.org abs /25.aaaa5500 22... 02000. (Cited on page and 23)
12

---

Published as a conference paper at ICLR 2026
Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for LLM training. arXiv preprins arXiv:2502.16692, 2025. URL pttttp::/arxxi..ooo wwt.orgggabb...eoeeeeeeeee 6982. (Cited page
and 2)
Xiaorui Liu, Yao Li, Jiliang Tang, and Ming Yan. A double residual compression algorithm for effi- cient distributed learning. In International Conference Artificial Intelligence and Statistics, pp- 133-143. PMLR, 2020. URL https::////... /ttp/aiiivvv./////// org, wttp////..... 07561. (Cited on page 19)
Yuxing Liu, Rui Pan, and Tong Zhang. AdaGrad under anisotropic smoothness, 2024. URL https: / //aaxiv.cco//// abs / 240 15244. (Cited page 21 and 22)
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019. URL. httpt:ttp////c /ttps:///arxiii httpr/ivvvrg//b//////// .05101.
(Cited on page and
Maksim Makarenko, Elnur Gasanov, Rustem Islamov, Abdurakhmon Sadiev, and Peter Richtárik. Adaptive compression for communication--eficieen distributed training. arXiv:2211.0016 2022. URthttps: //axxiv arggabs/2211.....
arXiv preprint 221//22111110 001 88. (Cited on page 20)
L. O. Mangasarian. Parallel gradient distribution in unconstrained optimization. SIAM Journal on Control and Optimization, 33(6):1916-199) 1995. URL http::://c...oog///o/o/
$036301299325022 (Cited on page 19)
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication--ffiiieee learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1273-1282. PMLR, 2017.
05629. (Cited on page 2, 5, and 19)
Konstantin Mishchenko, Filip Hanzely, and Peter Richtárik. 99% of distributed optimization is waste of time: The issue and how to fix it. arXiv preprint 10X1::1901.09437, 2019, URL https: /ttpivavvv.///////////ee////.av. org/ abs/ 19199990 09437. (Cited on page 2)
Konstantin Mishchenko, Eduard Gorbunov, Martin Takáč, and Peter Richtárik. Distributed learning with compressed gradient differences. Optimization Methods and Software, pp. 1-16, 2024. URL https: / //axxvaaavvv/////////////////avv org/ abs 1901 09269. (Cited on page 2)
Moonshot AL Kimi K2: Open agentic intelligence, 2025. URL https::/noonshotal
github. i Kimi-iii---- K2/. (Cited on page 1)
Philipp Moritz, Robert Nishihara, Ion Stoica, and Michael I Jordan. Sparknet: Training deep net- works in spark. arXiv preprint 1 2015. URL https://aaxii..c org. abs/ 1511. 06011.
(Cited on page 4)
Arvind Narayanan, Xumiao Zhang, Ruiyang Zhu, Ahmad Hassan, Shuowei Jin, Xiao Zhu, Xiaoxuan Zhang, Denis Rybkin, Zhengxuan Yang, Zhuoqing Morley Mao, et al. A variegated look at 5g in the wild: performance, power, and qoe implications. In Proceedings of the 2021 ACM SIGCOMM 202/ Conference, pp. 610-625, 2021. URL https: /dl. acm. www.arrm........................ doi dood//....... /10. 1145/3452295 34722 4229223. (Cited on page 19)
Yu Nesterov.
Efficiency of coordinate descent methods on huge-scale optimization problems. SIAM Journal Optimization, 22(2):341-362 2012. URL https: /tppbs..pbb////////// siam. org/ (Cited on page 5, 21, 22, and 25)
10.
Yurii Nesterov. Introductory Lectures on Conver Optimization: Basic Course, volume 87. Springer Science & Business Media, 2003. URL https: springer....... com/book/ 10. 1w1-1007977----4-50 101-44419888---- 9. (Cited on page 19)
Feiping Nie, Heng Huang. Xiao Cai, and Chris Ding. Efficient and robust feature selection via joint -nors minimization. In J. Lafferty, C. Williams, J. Shawe-Taylor, R. Zemel, and A. Culotta (eds.), Advances in Neural Information Processing Systems, volume 23, Cur- Associates, Inc., 2010. URL https: wttptp:sppptceeeeeeeepepp /ttp:///eoceeeeeeeee//eeec neurips. cttpppp....pppaeeeeeeeee files/ pdf. (Cited
page 29)
13

---

Published as a conference paper at ICLR 2026
Julie Nutini, Issam Laradji, and Mark Schmidt. Let's make Block Coordinate Descent converge faster: Faster greedy rules, message passssing.... active-set complexity, and superlinear convergence. arXiv preprint arXiv:1712.08859, 2017. URL https: http://axiivv.ooo org www.o//gg......././o/.ooo 08859.
(Cited on page 21)
Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin A Raffel, Leandro Von Werra, Thomas Wolf, et al. The Win datasets: Decanting the web for the finest text data at scale. Advances in Neural Information Prucessing Systems, 37:30811-308499 2024. URL https: arxiv. org/ abs /406 2406.17557. (Cited page
Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, and Volkan Cevher. Scion. https://github: http:/gghhbb..com/oo//SS/e////// GitHub repository. (Cited page
2025a.
Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, and Volkan Cevher. Training deep learning models with norm-constraaned LMOs. arXiv preprint arXiv:2502.07529, 2025b. URL https: /wtp///aaavvv/aa///////e./ea.a wttp::/wwwwggbb////////// 07529. (Cited on page 4,5, 23, 24, and 69)
Thomas Pethick, Wanyun Xie, Mete Erdogan, Kimon Antonakopoulos, Tony Silveti-Falls, and Volkan Cevher. Generalized gradient norm clipping & non-Fuclidean (1 0, -msoothhsnss arXiv preprint arXiv:2506.01913, 2025c. URL http:::://aar...oog//bbssss///// (Cited on page 21 and 66)
01 91
Constantin Philippenko and Aymeric Dieuleveut. Preserved central model for faster bidirectional compression in distributed settings. Advances in Neural Information Processing Systems, 34: 2387-2399 2021. URL atttps://arxila //aaxiivv../ .rgg abs 2102. 12528. (Cited on page 19)
Vitali Pirau, Aleksandr Beznosikov, Martin Takáe, Vladislav Matyukhin, and Alexander V. Gas- nikov. Preconditioning meets biased compression for efficient distributed optimization. Com- putational Management Science, 21(1):14, 2024. URL https: /twii ort/oorg//... 2rg/1001007/ 6. (Cited page 19)
Daniel Povey, Xiaohui Zhang, and Sanjeev Khudanpur. Parallel training of DNNs with natural gradient and parameter averaging. arXiv preprint arXiv:1410.7455, 2014. URL https://
arxiv.
.0000..455555 7455. (Cited page and 19)
Xun Qian, Peter Richtárik, and Tong Zhang. Error compensated distributed SGD can be accelerated. Advances in Neural Information Processing Systems, 34:30401-30413, 2021. URL https: / //aaxxiv.rrg//aaaeeeeeeeeaa www./r///aaa/aaeeeee//.../eeeeee 00091.0011.1... (Cited on page 19)
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. ZeRO: Memory optimiza- tions toward training trillion parameter models. In SC20: International Conference for High Performance Computing. Networking, Storage and Analysis, pp. 1-16, 2020. URL https: //ieeexplore.eee leee wwwpo////oo//////e/eoeeee (ewccummnnn/35330..... (Cited page
Ahmad Rammal, Kaja Gruntkowska, Nikita Fedin, Eduard Gorbunov, and Peter Richtárik. Commu- nication compression for byzantine robust learning: New efficient algorithms and improved rates. In International Conference Artificial Intelligence and Statisticy, pp. 1207-1215, PMLR, 2024.
URL wwtphhtt////...... /ttp:///axxvv../ org/ abs/ 231 013.00000000 09804. (Cited on page 20 and 42)
Artem
Riabinin,
Egor
Shulgin,
and Peter https: /ttptuub//tttu..//// com/ artem- htt///arteerriiiit. riabinin/ 2025a. GitHub
Kaja Gruntkowska.
Richtárik.
Gluon.
inneermment-eee wet.siimainngga. inw.iiiiggsmmmmeeessssssssesesse haaprchheess-------------------- repository. (Cited on page 8, 9, and 67)
Artem Riabinin, Hgor Shulgin, Kaja Gruntkowska, and Peter Richtárik. Gluon: Making Muon & Scion great again! (Bridging theory and practice of LMO-basee optimizers for LLMs). arXiv preprint arXiv:2505 13416, 2025b. URL https: /tapsi/////vvv../// aaa/aaiiiv..///.///aaaa//aaaaava org aabs/////////e http//abb////.../// 2255. (Cited on page 1, 4, 5, 6, 21, 22, 23, 24, 65, and 66)
Peter Richtárik and Martin Takáč. Iteration complexity of randomized block-coordinate descent methods for minimizing a composite function. Mathematical Programming. 144(1):1-38, 2014. URL httpt:tpp/// /ttp///avvvvv//////////////..wvv wwtp:/ivwor..../ooooooooooo
28 B. (Cited on page 21 and 22)
14

---

Published as a conference paper at ICLR 2026
Peter Richtárik, Igor Sokolov, and Ilyas Fatkhullin. EF21: A new, simpler, theoretically better, and practically faster error feedback. In Neural Information Processing Systems, 2021., 2021. URL https::/arxiv.. org/ als /tsp///s 05203. (Cited on page 3, 5, 6. 7, 20, 24. 27, and 42)
Peter Richtárik, Igor Sokolov, Elnur Gasanov, Ilyas Fatkhullin, Zhize Li, and Eduard Gorbunov. 3PC: Three point compressors for communicationn-efficceen distributed training and better the- ory for lazy aggregation. In International Conference Machine Learning, pp. 18596-18648.
PMLR, 2022. URL https:ttt//// /ttwiaxxivv/.. org/ abs/////////// 220222222222 00998. (Cited on page 20)
Peter Richtárik, Elnur Gasanov, and Konstantin Burlachenko. Error feedback reloaded: From quadratic to arithmetic mean of smoothness constants, arXiv preprint arXiv:2402. 10774, 2024.
URL https: aawiivivvvv.////vv//aaaa///aavvv
org/ abs /2022 10774. (Cited on page 6)
Mher Safaryan, Rustem Islamov, Xun Qian, and Peter Richtárik. FedNL: Making Newton-type URthttps:
methods applicable to federated learning. arXiv preprint urXiv22/06.02966 2021.
//arxiv..ccm org abs 22106. 02969. (Cited on page 9)
Frank Seide, Hao Fu, Jasha Droppo, Gang Li, and Dong Yu. 1-bit stochastic gradient de- scent and its application to data-parallel distributed training of speech DNNs. Fif- teenth Annual Conference of the International Speech Communication Association, 2014. URL https: /tt///wwww/.. mi crosoft. com www.oonooeeeoooo eusuennuu////eeeee htt////nnsr/eeaeeeeeeeee rresearch/wee /esparch//pppcccnnee/////// wwwppp/oonentttc.coocc 201 6016//2/// TS1 40094 pdf. (Cited page 2, 3,5, 19, and 27)
Ishaan Shah, Anthony M Polloreno. Karl Stratos, Philip Monk, Adarsh Chaluvaraju, Andrew Hojel, Andrew Ma, Anil Thomas, Ashish Tanwer, Darsh Shah, et al. Practical efficiency of Muon for pretraining arXiv preprint arXiv:2505.02222, 2025. URL
abs/
25C5. 02222. (Cited on page 1)
David So, Wojeiech Manke, Hanxiao Liu, Zihang Dai, Noam Shazeer, and Quoc V Le. Searching for efficient transformers for language modeling. Advances in neural information processing systems, 34:6010-6022, 2021. URL https:/ /apps://arxiiin /tt//aaiiv////////////e////aaava wwwaaii.org/aaaiaaaaiooaoo wwt./br//////.oooooooooooo. www.////........ 08668v2. (Cited on page 66)
Sebastian U Stich and Sai Praneeth Karimireddy. The error-feedback framework: Better rates for sgd with delayed gradients and compressed communication. arXiv preprint arXiv:1909.05350 2019. URL https: arxiv. org/ abs / 909. 05350. (Cited page 19)
org
Sebastian U Stich, Jean- Baa-tttp Cordonnier, and Martin Jaggi. Sparsified SGD with memory.
Advances in neural information processing systems, 31, 2018. URL httpt:t///x/xii...tt// bbss/809..0 00909099999999.....
07599. (Cited on page 2 and 19)
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. RoFormer: En- hanced transformer with rotary position embedding. Neurocomputing, 568:127063 2024. URL
pittps://arxer org/ alss/ 2104. 09864. (Cited on page 66)
Tao Sun, Qingsong Wang, Dongsheng Li, and Bao Wang. Momentum ensures convergence of SIGNSGD under weaker assumptions. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th Inter- national Conference Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 33077-33099. PMLR, 23-29 Jul 2023. URL https:tpps... httpss///coccceegg....om mllr. press/v202/st sun2/llnn311mmm/mmm/ html. (Cited page 8)
Rafal Szlendak, Alexander Tyurin, and Peter Richtárik. Permutation compressors for provably faster distributed nonconvex optimization. In International Conference on Learning Representations, 2021. URL https: httpaaxvvvvoo/ wwwpiv.oog///o..ooooo abs 1abs/210 22100.100.0111 03300. (Cited page 19 and 27)
Hanlin Tang, Chen Yu, Xiangru Lian, Tong Zhang, and Ji Liu. DoubleSqueeze: Parallel stochas- tic gradient descent with double-pass croor-compensated compression. In International Confer- ence on Machine Learning, pp. 6155-6165. PMLR, 2019. URL http:://wxxxvvoo.eeoeeee abs/
1905. 05957. (Cited on page 19)
Benjamin Thérien, Xiaolong Huang, Irina Rish, and Eugene Belilovsky. MuLoCo: Muon is practical inner optimizer for DiLoCo. arXiv preprint arXiv:2505.2372 2025. URL https: / w/aaxiivorrg///ooo..aoooooo a2555 23725. (Cited on page and 2)
15

---

Published as a conference paper at ICLR 2026
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajwaal Bhargava, Shruti Bhosale, et al. Llama 2: Open foun- dation and fine-tuned chat models. arXiv preprint arXtv:2307.09228, 2023. URL https:
/ //arxiv.orrg org abs 2307 09288. (Cited on page 1)
Alexander Tyurin and Peter Richtárik. DASHA: Distributed nonconvex optimization with communi- cation compression, optimal oracle complexity, and no client synchronization. 11th International Conference on Learning Representations (ICLR), 2023a. URL https:://axiii///s/t org/ abs/
2202. 01 268. (Cited on page 19)
Alexander Tyurin and Peter Richtárik. 2Direction: Theoretically faster distributed training with bidirectional communication compression. Advances in Neural Information Processing Systems (NeurIPS), 2023b. URL https: /arxivv..vvv wtg/abs////.eeooooo...ooe.e /305 12379. (Cited on page 19)
Cong Xie, Shuai Zheng, Sanmi Koyejo, Indranil Gupta, Mu Li, and Haibin Lin. Cser: Communication-effiiiennn SGD with error reset. Advances in Neural Information Processing Sys- tems, 33:12593- 2020. URL https: htaps://arxiv..oo www.wawwww.aaaaaaaaaooooooooooot abs 2007 22200000122... 13221. (Cited on page 19)
Shuo Xie, Mohamad Amin Mohamadi, and Zhiyuan Li. Adam exploits -eomeetyy of loss land- scape via coordinate-wise adaptivity. arXiv preprint arXiv:2410.08/98, 2024. URL https: / // www.xii.org//aa//e//22
08198. (Cited on page 21)
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint a788881708.03888, 2017. URL https: aaaaaaxxvvv.......a /ttpa//ii....//b/////////////// 03888. (Cited on page 1)
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in neural informa- tion processing systems, 32, 2019. URL https: htgg/bbs//119g 07467. (Cited page 66)
Advances
Jingzhao Zhang, Tianxing Hc, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. In International Conference on Learning Rep- resentations, 2020. URL https: URLAttps::/arxii..oo
1905 /tt.aaa//.1....1.1.......1
(Cited page and 21)
Shuai Zheng, Ziyue Huang, and James Kwok. Communication-efficies distributed blockwise mo- mentum SGD with error-feedback. Advances in Neural Information Processing Systems, 32, 2019. URL https: //tpx/aavvv////////////////aaavl www.oo////a.........oooo. /bbg////////// 905. 10936. (Cited page 19)
16

---

Published as a conference paper at ICLR 2026
APPENDIX
CONTENTS
Introduction
2 Background
3 Non-Euclidean distributed training
Convergence results
5 Experiments
A Related work
19
A.1 Compression
19
A.2 Error feedback
19
A.3 Generalized smoothness
B Layer-wise setup
22
B.1 Muon, Scion and Gluon
23
B.2 Layer-wise EF21-Muon
24
C LMO in many guises
25
Non-Euclidean contractive compressors
27
D.1 Compression via norm selection
29
Convergence analysis
E.1 Descent Lemmas
31
E.2
Auxiliary Lemmas
................................
33
E.2.1
Smooth case
33
E.2.2
Generalized smooth case
36
E.3
Deterministic setting
42
E.3.1
Layer-wise smooth regime
42
E.3.2
Layer-wise (L", smooth regime.
45
E.4
Stochastic setting
48
E.4.1
Layer-wise smooth regime..
48
E.4.2
Layer-wise (L", L")ssooomoooth regime
55
F Useful Facts and Lemmas
65
Experiments
66
G.1 Setup details
66
17

---

Published as a conference paper at ICLR 2026
G.2 TopK compression details
66
G.3 Learning rate ablation
66
G.4 Compression level ablation
67
G.5 2.5B tokens experiment
67
G.6 MediumGPT experiment
67
G.7 Bidirectional compression
69
G.8 Limitations
70
18

---

Published as a conference paper at ICLR 2026
A RELATED WORK
A.1 COMPRESSION
The ML community has developed two dominant strategies to address the communication bottle- neck. The first is compression, implemented through techniques such as sparsification or quantiza- tion (Seide et al., 2014; Alistarh et al., 2017; Beznosikov et al., 2020; Szlendak et al., 2021; Horváth et al., 2022), which reduce communication costs by transmitting lossy representations of dense up- dates. Compression techniques have been extensively studied in the Euclidean regime. The other approach is local training (Mangasarian, 1995; Povey et al., 2014; McMahan et al., 2017), which lowers communication frequency by synchronizing with the server only periodically, after several local updates on the clients. These two approaches can be combined, yielding additional provable benefits by leveraging both mechanisms (Condat et al., 2024). In this work, we focus on compres- sion. Local training introduces distinct set of challenges and trade-offs, and is orthogonal to our approach.
There are two primary compression objectives in distributed optimization: workers-to-server (w2s) (= uplink) and server-to-workers (s2w) (= downlink) communication. A large body of prior work focuses exclusively on w2s compression, assuming that broadcasting from the server the workers is either free or negligible (Gorbunov et al., 2021; Szlendak et al., 2021; Tyurin & Richtárik, 2023a; Pirau et al., 2024). This assumption is partly due to analytical convenience, but can also be justified in settings where the server has significantly higher bandwidth, greater computational resources, or when the network topology favors fast downlink speeds (Kairouz et al., 2021). However, in many communication environments, this asymmetry does not hold. For instance, in 4G LTE and 5G networks, the upload and download speeds can be comparable, with the ratio between w2s and s2w bandwidths bounded within an order of magnitude (Huang et al., 2012; Narayanan et al., 2021). In such cases, s2w communication costs become negligibe,,,,, and optimizing for both directions is essential for practical efficiency (Zheng et al., 2019; Liu et al., 2020; Philippenko & Dieuleveut, 2021; Fatkhullin et al., 2021; Gruntkowska et al., 2023; Tyurin & Richtárik, 2023b; Gruntkowska et al., 2024).
A.2 ERROR FEEDBACK
To address the communication bottleneck, a natural approach is to apply biased compressors to the transmitted gradients. For the standard (Fuclidean) GD, which iterates
X = 3FR->"VS(X*)) x* 7*th I (EV((X))
where 7h > 0 is the stepsize, this would yield the update rule
1 n x = x 7
) j=1
Sadly, this "enhancement" can result in exponential divergence, even in simplest setting of mini- mizing the average of three strongly convex quadratic functions (Beznosikov et al., 2020, Example 1). Empirical evidence of such instability appeared much earlier, prompting Seide et al. (2014) to propose a remedy in the form of an error feedback (EF) mechanism, which we refer to as EF14.
Initial theoretical insights into EF14 were established in the simpler single-node setting (Stich al., 2018: Alistarh et al., 2018). The method was subsequently analyzed in the convex case by Karim- ireddy et al. (2019); Beznosikov et (2020); Gorbunov et al. (2020). Next, Qian et al. (2021) showed that error feedback methods can be combined with Nesterov-style acceleration (Nesterov, 2003), though at the cost of incorporating additional unbiased compression, leading to increased communication overhead per iteration. These analyses were later extended to the nonconvex regime by Stich & Karimireddy (2019). This motivated a series of extensions combining error feedback with additional algorithmic components, such as bidirectional compression (Tang et al., 2019), de- centralized training protocols (Koloskova et al., 2019), and the incorporation of momentum either on the client (Zheng et al., 2019) server side (Xie al., 2020). While these works advanced the state
19

---

Published as a conference paper at ICLR 2026
of the art, their guarantees relied on strong regularity assumptions, such as bounded gradients (BG) or bounded gradient similarity (BGS), which may be difficult to justify in practical deep learning scenarios.
The limitations of EF14 and its successors were partially overcome by Richtárik et al. (2021), who proposed a refined variant termed EF21. EF21 eliminates the need for strong assumptions such as BG and BGS, relying only on standard assumptions (smoothness of the local functions f3 and the existence of global lower bound on f). while improving the iteration complexity the desirable O(//K) in the deterministic setting. Building on this foundation, a series of extensions and gen- eralizations followed. These include adaptations to partial participation, variance-reduction. proxi- mal setting, and bidirectional compression (Fatkhullin et al., 2021), a generalization from contrac- tive to three-point compressors (Richtárik et al., 2022), support for adaptive compression schemes (Makarenko et al., 2022), and EF21-P-a modification of EF21 from gradient compression to model compression (Gruntkowska et al., 2023). Further developments used EF21 in the design of Byzan- tine robust methods (Rammal et al., 2024), applied it to Hessian communication (Islamov et al., 2023), and extended the theoretical analysis to the (L", , LL' -moothh regime (Khirirat et al., 2024).
With this historical overview in place, we now narrow our focus to two developments in the error feedback literature that are particularly relevant to this work: EF21 (Richtárik et al., 2021) and EF21-P (Gruntkowska et al., 2023).
EF21 is a method for w2s communication compression. aims to solve problem (1) via the iterative process
X = X GG*h Ghtl = G + (V) - -c),
Gh+l
where y > is the stepsize and Cj E B2(aD) are independent contractive compressors. In the EF21 algorithm, each client j keeps track of gradient estimator Gj., At each iteration, the clients compute their local gradient V VfJ(X++)) subtract the stored estimator G/. and then compresses this difference using a biased compression operator. The compressed update is sent to the server, which aggregates updates from all clients and uses them update the global model. Concurrently, each client updates its error feedback vector by using the same compressed residual. Importantly, EF21 compresses only the uplink communication (i.e., vectors sent from clients to the server), while downlink communication remains uncompressed. That is, the global model X*+1 is transmitted in full precision from the server to all clients, under the assumption that downlink communication not a bottleneck.
complementary approach is proposed in the follow-up work of Gruntkowska et al. (2023), which introduces primal variant of EF21, referred to as EF21-P. Unlike EF21, which targets uplink com- pression (from workers to server), EF21-P is explicitly designed for s2w compression. The method proceeds via the iterative scheme
X = X - -VS(W = X
[sS(W*) 1=1
W++ W + (* -W*),
where 7 >> is the stepsize and C+ E (2(())) are independent contractive compressors. Analogous to EF21, the EF21-P method employs error feedback to compensate for the distortion introduced by compression. However, rather than correcting gradient estimates, EF21-P maintains and updates an estimate of the model parameters, W The server computes the update x++1, but broadcasts only compressed difference C* (x+-1 w to the clients.
In its basic form, EF21-P assumes dense uplink communication-i.c. the clients transmit full gra- dients (W** to the server. Nonethcless, EF21-P can be naturally extended to bidirectional compression by integrating it with an uplink compression mechanism, enabling full communication efficiency (Gruntkowska et al., 2023).
20

---

Published as a conference paper at ICLR 2026
A.3 GENERALIZED SMOOTHNESS
A standard assumption in the convergence analysis of gradient-based methods is Lipschitz smooth- ness of the gradient (Assumption 3). However, many modern learning problens-especially in deep learning-violate this assumption. Empirical evidence has shown non-smoothness in a variety of architectures and tasks, including LSTM language modeling, image classification with ResNet20 (Zhang et al., 2020), and transformer models (Crawshaw et al., 2022). These observations motivated the search for alternative smoothness models that better reflect the behavior of practical objectives.
One such model is (L", >msoothhnssss, introduced by Zhang et al. (2020) for twice continuously differentiable functions in the Euclidean setting. The authors define function f Rd R to be (L", L > if
L (VS(() VX E R
This condition generalizes standard Lipschitz smoothness and has been shown empirically to capture deep learning loss landscapes more faithfully than the classical model (Zhang et al., 2020; Craw- shaw et al., 2022). Subsequent works extended the above condition beyond the twice differentiable case (Li et al., 2023; Chen et al., 2023). In particular, Chen et al. (2023) introduced asymmet- ric and symmetric variants of (L It') -sooothness, where the asymmetric form (a special case of Assumption restricted to Euclidean norms) is given by
[V - ff()I12 & (I" + I. (V YIll WX,Y E R...
This framework has since been used in the non-Euclidean setting (Pethick et al., 2025c) and adapted to the layer-wise structure of deep networks by Riabinin al. (2025b), who introduced non-Fuctidean layer-wise (TO, 1 ssmoothttssssss assumption (Assumption 8). This layer-aware view aligns naturally with LMO-based optimizers that operate on individual parameter groups.
The idea of accounting for the heterogencous structure of parameters is not unique to the work of Riabinin al. (2025b). Anisotropic smoothness conditions, where smoothness constants can vary across coordinates or parameter blocks, have been studied extensively, for example in the context of coordinate descent methods (Nesterov, 2012; Richtárik & Takáč, 2014; Nutini et al., 2017). Variants of coordinate-wise or block-wise (generalized) smoothness assumptions have also been used to an- alyze algorithms such as signSGD (Bernstein et al., 2018; Crawshaw et al., 2022), AdaGrad (Jiang et al., 2024; Liu al., 2024), and Adam (Xie et al., 2024). These works collectively reinforce the need for smoothness models that reflect the anisotropic geometry modern neural networks.
21

---

Published as a conference paper at ICLR 2026
B LAYER-WISE SETUP
So far, we have been operating in an abstract vector space S, without assuming any particular struc- ture. This is the standard approach in the vast majority of the theoretical optimization literature in machine learning, where model parameters are typically flattened into vectors in R However, modern deep networks are inherently structured objects, with clear layer-wise organization. While treating parameters as flat vectors can still yield meaningful convergence guarantees, explicitly mod- eling this layer-wise structure allows us to formulate assumptions that more accurately reflect the underlying geometry of the model Nesterov (2012); Richtárik & Takáč (2014); Crawshaw et al. (2022); Jiang et al. (2024). This, in turn, can lead improved theoretical results (Liu et al., 2024; Riabinin al., 2025b).
further motivation for adopting the layer-wise perspective is that the algorithms that inspired this work-Muon, Scion, and Gluon-are themselves layer-wise by design. Rather than operating on the entire parameter vector, they apply separate LMO updates to each layer or building block indepen- dently. This modular treatment is one of the main reasons for their strong empirical performance.
With this motivation mind, we now turn to solving the optimization problem (1) in a setting where the parameter vector X S represents collection of matrices X; S; := mm X corresponding to the trainable parameters of each layer i E {1..... .p) in a neural network. For notational con- venience, write X = [X1.... Xp] and [X) [VI((X). .... pxf(X)], where (Xf(X) is the gradient component corresponding to the ith layer. Accordingly, S is the d-dimensional product space
: - i=1 S; S ... Sp:
where d = Ei= miTt Each component space S; is equipped with the trace inner product, defied as (X,, (1) == tr(X,TV) Yi) for X,, Yi E S,, and an arbitrary norm Illl(y not necessarily induced by this inner product. We use [|-1()) denote the dual norm associated with [I-]]() (i.e., [|
sup ( (X,, ()) for any X, E S.). Furthermore, we use P.P Ps > 0 to denote the norm equivalence constants such that
2. (1 S |x S P. (I
VX, E S,
(or, equivalently, L, | S [X)l((0 S Sp: (1iil22)
Remark 7. In the case of Muon, the norms Il-1leo, are taken to be the spectral norms, i.e., I1-1 I1-1l2-12 Since for any matrix X; of rank at most r, we have
|Xill2-2 S [ S VFllXil2
in this setting, P; = and p. = /T.
Given the block structure of X across layers, the smoothness assumptions in Assumption 3 can be made more precise by assigning separate constants to each layer.
Assumption 6 (Layer-wise smoothness). The function f S R is layer-wise -mmoooht with constants L =
I.) E R i.e.,
(x)((X) V.f(Y))( S.1. Y:ll(0
for all i I,.... . and all X = [X....., E,J [V....., Ypl E. Assumption (Local layer-wise smoothness). The functions ʃ3 S IR, j [n], are layer-wise T.y-smooth with constants L. (I.......,, [.g.s) RR, i.e.,
[V) V.ff $ Lb || Yil(()
for all i = 1...., P and all = [X,...., x,1 E S, = IV....., .Yy,, E S. We define (īg)² = "E
We invoke Assumptions and in Appendices E.3.1 and E.4.1 to extend Theorems 3 and 5 to the more general setting.
22

---

Published as a conference paper at ICLR 2026
Smoothness is the standard assumption used in virtually all convergence results for Muon and Scion (Kovalev, 2025; Pethick et al., 2025b; Li & Hong, 2025) (except for the recent work on Gluon (Riabinin et al., 2025b)). However, often fails to hold in modern deep learning settings. To address this, we adopt a more flexible and expressive condition: the layer-wise (I.", T.*) ) smoothness assumption (Riabinin al., 2025b).
discussed in Section 4 and Section A.3, this assumption
Assumption 8 (Lisee-wise (T.", I.²) smoothnesss) The function f - R is layer-wise (Tº 1.-)- smooth with constants L : (L)..... Lg) E R2 and L' = (L)....,L))) E RT, i.e.
[V+ - v.f(()) - (L.- + L (V. |X,-- (Xil-))ll() for all i = I. and all X [x,,...,,,p E S. Y = [Y....,,pp] S.
Since, unlike Gluon, we operate in the distributed setting, we will also need an analogous assumption on the local functions fy.
Assumption 9 (Local layer-wise (L", L smoothnesss. The fanctions f5. j E [nl, are layer-wise (L). - smoothh with constants L3 = (Li.j...., (p.)) E R2, and L} = (Lig.... Lps) E Rt, i.e.,
[V.f((X) - ())(0))))) S (Lig + Lis (I |x,, Yill())
for all i = 1,...., and
.and all X [X Xpl E S, Y Vy....,,y))E. Y
For O E {0,13. we define L = maxie')) p)) Lo L. = maxje[m] Lij and Lş ==
Riabinin et al. (2025b) present empirical evidence showing that this more flexible, layer-wise ap- proach is essential for accurately modeling the network's underlying structure. They demonstrate that the layer-wise (L", L (mmosohhhsssss condition approximately holds along the training trajectory of Gluon in experiments on the NanoGPT language modeling task. Motivated by these findings, in Appendices E.3.2 and E.4.2, provide analysis within this generalized framework, offering full generalization of Gluon to bidirectional compression.
In the stochastic setting, will also require layer-wise analogue of Assumption 5.
Assumption 10. The stochastic gradient estimators ())))) S + S are unbiased and have bounded variance. That is, RE (X: (X, for all X and there exist o; such that
Eg DD [V.f)((:E) (xf)(X)) oof
VX S, isi=
We permit layer-dependent variance parameters of motivated by empirical evidence that variance is not uniform across layers. For example, Glentis et al. (2025) observe that, during training of LLaMA 1 330 with SGD and column-wis normalization (i.e., Gluon using the I1-112--2 norm), the final and embedding layers display significantly higher variance.
B.1 Muon, Scion AND Gluon
Muon, introduced by Jordan et al. (2024), is an optimizer for the hidden layers of neural networks (the first and last layers are trained with AdamW). Unlike traditional element wise gradient methods, updates each weight matrix as a whole. Given layer X, and the corresponding (stochastic) gradient G,, Muon selects an update that maximizes the alignment with the gradient to reduce loss, while constraining the update's 's size to avoid excessive model perturbation. This is formulated as a constrained optimization problem over the spectral norm ball:
arg min (G,, AX,) s.t. ||AX St,
(5)
where the radius t; > 0 plays a role similar to a stepsize. The optimal update AX; is obtained by orthogonalizing G₁ via its singular value decomposition G, = U.E.V?. leading
AX; -U..V.
23

---

Published as a conference paper at ICLR 2026
This yields the basic update
X Ax = X (V.)).
(6)
In practice, computing the SVD exactly at every step is expensive and not GPU-friendly. Muon instead uses Newton-Schule iterations (Kovarik, 1970; Björck Bowic, 1971) to approximate the orthogonalization. Combined with momentum, the practical update is
M (1 3,) + B,G.
X X (ewronSchulz(M,")
where 3r E (0, 1] is the momentum parameter and M is the monentum-averaaed gradient.
While Newton-Schulz iterations and momentum crucial for practical efficiency, the essence of Muon lies in solving (5)-that is, computing the linear minimization oracle (LMO) over the spectral norm ball. Recall that LMO (G) == arg min (E)(X.)) (G, Z). Then
AX; = arg min (Gi, Z,) = LMO (2 (Gt) 2.EB(-2(0.A))
where (0,t,) == {2, E S; |Z I t is the spectral norm ball of radius t around 0. Thus, the update (6) can equivalently written as
= + L-0 B(-00( (G.),
where G may be replaced with a momentum term
Crucially, nothing in this formulation ties us to the spectral norm. The same update structure can be defined over any norm ball, opening the door to an entire family of optimizers whose proper- ties depend the underlying geometry. This insight has led to several Muon-inspired methods with provable convergence guarantees (Pethick et al., 2025b; Kovalev, 2025; Riabinin et al., 2025b). Scion (Pethick et al., 2025b) removes the restriction to matrix-shaped layers by applying LMO-based updates to all layers, pairing the spectral norm for hidden layers with the I| norm elsewhere. Gluon (Riabinin et al., 2025b) expands the view ysis for LMO updates over arbitrary norm balls, supported by layer-wise (L", L (²)--moothness assumption that captures the heterogeneity of deep learning loss landscapes more accurately than standard smoothness.
further: it provides general convergence anal-
B.2 LAYER- LAYER-WISE EF21-Muon
The simplified EF21-Muon in Algorithm 1, analyzed in Section 4, omits the layer-wise treatment introduced above. The full structured variant is given in Algorithm 3. Its deterministic counterpart formalized Algorithm 2, extending the simplified version studied in Section 4.
Both Algorithms 2 and 3 operate ture. For each layer i, the parameters are updated via X LMO LMO .e))
per-layer basis. We now briefly describe their struc-
(G!) (equiva-
lently, X X 7 (G+)², where 7 ''/[[[*]] see Section C). Next, the algorithms perform the server-to-workers (s2w) compression, following technique inspired by EF21-P (Grun- tkowska et al., 2023). The resulting compressed messages S C. Cx w are sent to the workers. Each worker then updates the model shift and uses the resulting model estimate W compute the local (stochastic) gradient. This gradient then used (either directly or within a momentum term) to form the compressed message gk+1 This part of the algorithm follows the
(w2s) compression strategy of EF21 (Richtárik et al., 2021). The messages Ri++1 are sent back to the server, which updates the layer-wie wise gradient estimators via
Gk+t # Ey ck+1 G + En R8t1 This process is repeated until convergence.
24

---

Published as a conference paper at ICLR 2026
Algorithm 2 Deterministic EF21-Muon
1: Parameters: radii 1 0 stepsizes y; initial iterate X [Xp,..., XpI
(stored on the
server); initial iterate shift W = X (stored on the server and the workers); initial gradient estimators G [G....... Gp. (x), [xf( S (stored on the workers), G Σラ j=1 Gj (stored on the server); worker compressors C,P, server compressors C*
2: for 0....... K do
3: for i = 1,2..... do
X LLLLLMBBBX B(X) +5) (Gt) X 7 (c+)²
4:
Take LMO-type 110-typp step
5:
S (X W
Compress shifted
6:
Wh+1 W + S
Update model shili
7:
Broadcast 8 . .... Spi to all workers
8:
end for
9:
for j
= ...... .....,n n in parallel do for i = 1. ........ do
10:
11:
Wh+1 = W + S
Update model shift
12:
Rk+1 C., (V.5
ckty)
Compress shifted gradient
G G. + Rk+1
13:
ij
end for
Broadcast 1. = [ ..... 125 to the server
15:
16: end for
for i
1 1....,, do G
2i=1 Gh+t1 G Eド= R
Compute gradicnt estimator
end for
20: end for
C LMO IN MANY GUISES
As outlined Section 2, the update rule (2)
x++1= x + t LLLM0 B(..1) (G*)
admits several equivalent reformulations.
LMO viewpoint. over a norm ball
The original update (2) is the solution of simple linear minimization problem
x++1 -LMOB(X* (G*)
arg min (G*, x), xeB(x*,
where B(X,t) = (Z E
| - XI| t). The LMO satisfies
(G. LM.0 (G)) -t ||G]].
Sharp operator viewpoint. An equivalent perspective is obtained via the sharp operators (Nes- terov, 2012; Kelner et al., 2014). Define the function s(X) || Its Fenchel conjugate is
given by (G) = supxes (G - x(X)} = 1 || and its subdifferential 0s* coincides
with the sharp operator:
(G)
(XES: (G, X) = ||]||, |X| [G = [X]]
-I|G]|, IMO B(O.1) (G)
G
where 6 == arg max (G, X) - | | is the sharp operator. Therefore,
x = x +LLMO (G+) = x
(G+)²
G
25

---

Published as a conference paper at ICLR 2026
Algorithm 3 EF21-Muon
1: Parameters: radii 1 > 0/ stepsizes > momentum parameters 3; (0, 1]; initial iterate X [XP,... xp] E S (stored on the server); initial iterate shift W X (stored on the server and the workers); initial gradient estimators Gj [G] j........ Co p.j E S (stored the workers); G n وÿ= Cg (stored on the server); initial momentum Mj [M....... M S
(stored on the workers); worker compressors c.y: server compressors C₁
2: for = 0, 1...... K
3: for i = 1.2...., do
x =LMO (G.) X (G.)²
4:
Take LMO-type stop
5:
S (x W
Compress shifted model on the server
6:
W W + S.
Update model shift
7:
Broadcast 8th .S...... .... spi all workers
8: end for
9:
for j
1......... . n in parallel do for i = 1, ....... do
10:
w = =W? + S
11:
Update model shift
M =(1 (1-8.)M5, +
(
12:
Compute Compress shifted gradient
13:
Rk+1 (M1 G.
14:
G G R
15:
end for
16:
Broadcast R. = [ 1
pk++1 p.j
to the server
17:
end for
18:
for i 1,2, ... p do
19: G -==1 Gh+1 Gh. + + R
k++
Compute gradient estimator
20: end for
end for
ie., a normalized steepest descent step with effective stepsize 7 = 1+/
Two properties of the sharp operator used later are
(x, (x, = |
||X|_ = +|Xx
Subdifferential
viewpoint. The negative LMO direction -LMO B( (A)
arg max (A, Z) is subdifferential of the dual norm θ 1|l so (2) can also be written as
x = x + 4ALMOBO..1 -
-
for some H E 111111 (G*), where by the definition of subdifferential, for any Gh + 0,
(1* G*) |G*
|
26

---

Published as a conference paper at ICLR 2026
D NON-EUCLIDEAN CONTRACTIVE COMPRESSORS
Recall from Definition that a mapping C + S is called a contractive compression operator with parameter (0, 1] if, for all X E S,
E -l(( - x12 - (1- a1-a)
(9)
When ||-|| is the Euclidean norm, wide range of such compressors is available in the literature (Seide et al., 2014; Alistarh et al., 2017; Beznosikov et al., 2020; Richtárik et al., 2021; Szlendak et al., 2021; Horváth et al., 2022). However, when ||-I| is non-Euclidean norm, Euclidean con- tractivity does not in general imply contractivity with respect to ||-I|. Indeed, suppose that C is contractive with respect to the Euclidean norm. Then, using norm equivalence, for any X S,
-² [IC((()))
x||] - s
[[C(()- x x16]
s (1-aa a) | p| 1|x||²
Rearranging gives
1² (1-0)
E [ s x
and hence C is not contractive with respect to the norm ||-|| unless > 1 f/p. Consequently, dedicated compressors are needed when working outside the Euclidean setting.
In this section, we first present two simple examples of operators that satisfy condition (9) for any norm. These are, however, general not very practical choices. We then turn to more useful examples of non-Euclidean compressors for several matrix norms of interest.
simple deterministic example of contractive compressor the scaling or damping operator.
Definition 8 (Deterministic Damping). For any X S, the deterministic damping operator with parameter y E (0,2) 2) is defined as
C(X) = 7
For this operator,
E [I((x) -11] = 7) |x
and thus C satisfies Definition with a = (1-1-)) for any E (0,2).
Despite meeting the definition, the deterministic damping operator is of little use in communication constrained optimization: it merely scales the entire input vector by a constant, without reducing the amount of data to be transmitted. The fact that it formally satisfies the contractive compressor definition is more of a theoretical curiosity. It highlights that the definition captures a broader math- ematical property that does always align with the practical engineering goal of reducing data transmission.
The random dropout operator (whose scaled, unbiased variant appears in the literature the Bernoulli compressor (Islamov et al., 2021)) is simple yet more practically relevant example of contractive compressor that can reduce communication cost.
Definition 9 (Random Dropout). For any X E S, the random dropout operator with probability parameter p E (0, 1] defined as
c(X)= fx 0
with probability p,
C(X)
with probability 1 1-P p.
Then
E [[cc(x) x1]] = -11
and hence C E B(p).
The examples of deterministic damping and random dropout apply to any valid norm defined on the
space S. However, one can also design compressors directly for the norm interest. A natural example for both the spectral norm 1|:112-12 and the nuclear norm ||-||, is based truncated SVD.
27

---

Published as a conference paper at ICLR 2026
Definition 10 (TopK SVD compressor). Let X UEVT Rmxn be matrix of rank where
= (1. .r) contains the singular values 01 2 ... 2 or > 0. For K < 7, the Top SVD compressor is defined by
0) retains the K largest singular values, setting the rest to where Σς- = d oK. .,.....) zero.
The TopK SVD compressor can be used in conjunction with several commonly used matrix norms:
Spectral norm. The spectral norm, frequently used in LMO-based optimization methods, is defined by || 01. Under this norm, the compression residual is
|| -
aK+1-
This yields valid contractive compressor (unless o x of) and Definition is satisfied with parameter <////..
Nuclear norm. The nuclear norm, dual to the spectral norm, is given by ||x = Ei= In this case,
- c(x)ll
Σ o
1=
and Definition 1 holds with a = 1- -1-(Eiek+17
(
oi
Frobenius norm. The Euclidean norm of the matrix can be expressed as |X|ee VE of Then.
|X c(x))l=
of.
L= K+1
and so Definition is satisfied with a =
2i-1
In fact, the TopK SVD compressor is naturally well-suited for larger family of Schatten p-norm, defined in terms of the singular values of a matrix X by
1/p
| Σ
Σop
of
Important special cases include the nuclear norm (or trace norm) for p = 1 (i.c., || = || ) and the spectral norm for = oo (i.e., ||X||2-2 = | In general, is easy to show that the Topk SVD compressor satisfies Defini- tion with respect to the ||-Ils, norm with
the Frobenius norm for p = (i.e., || = || s,
Ei of
2/v
a a=1-
Σ E
Remark 11. For large-scale matrices, computing the exact SVD may be computationally pro- hibitive. In such cases, one may resort to approximate methods to obtain stochastic compressor satisfying Definition I in expectation:
|(xxx)
x|"] (1 a1 a
E
1
where 5 > 0 quantifies the approximation error and can be made urbitrarily small. Remark 12. The expressions for a above depend on the singular values of X. and hence a is gen- erally matrix- deeended rather than uniform constant. For theoretical guarantees, may take the minimum a observed over a training run. Alternatively, our framework admits a straightforward extension iteration-dppendr compression parameters.
28

---

Published as a conference paper at ICLR 2026
Beyond Schatten norms, similar ideas can be applied to other structured non-Euclidean norms. Throughout, we let Xs:, X.j, and Xis denote the ith row, jth column, and (i. (i,))th entry of the matrix X Rmx" respectively.
Definition 13 (Column-wist Top, compressor), The column-wiss Top, compressor keeps the K columns with largest Ep norm, setting the rest to zero:
Xij. j jE E
c(X);
10,
otherwise,
where Ix indexes the K columns with the largest 6p norm.
This operator naturally suited for the mixed ppg norms (p, q 2 1), defined as
9/9
1/9
1/a
m
| pl
Σ
X
where I1-11, is the standard (vector) Ep norm. The compression residual satisfies
1/a
|| - C(X)
Σ
TEIX
and hence Definition holds with
Σ Dj || ビグ=1
a a=1-
This general formulation recovers, for example, the Ê2, norm (commonly used in robust data anal- ysis (Nie et al., 2010)) and the 22 norm (Frobenius norm).
D.1 COMPRESSION VIA NORM SELECTION
useful perspective on communication reduction in distributed optimization emerges from the con- nection between compression operators and mappings such as the sharp operator and the LMO. Recall that for any norm ||-|| with dual norm ||- the sharp operator of G E TRmxn defined as
G# == arg max
(G.)) - |x
X XeRmxn
Since ||G]], LMO (G) - one can view G as the LMO over the unit ball of [|- scaled by ||G]].. For many norms, G naturally acts structured compressor. Below, we list several such examples.
Nuclear norm. For the nuclear norm (with dual norm ||:1|2-22 the operator/spectral norm), the sharp operator is
=
o T
where 01, M1, and i1 are the leading singular value and singular vectors of G, yielding Rankl compression via truncated SVD. This operator satisfies Definition with parameter 1/r, where r is the rank of G.
Element-wise & norm. For the norm |X]]] = Em Ej
(iiz (with dual || =
maxi,j (xj)), the sharp operator is
G = Topolp ||G|| B(c-))):
where (i*,5", arg max;,; GGay and E(srr*) is the matrix with a in entry (i*.j*) and zeros elsewhere. Thus, the sharp operator associated with the E norm corresponds to Top1 sparsification, which satisfies Definition with 1/mn.
29

---

Published as a conference paper at ICLR 2026
Max norm. For |x gixi Σj=1 [iyl
the dual norm is
1 = Ej=1 |X and the sharp operator yields
G# EiGIl [sign(Top1(G.))..... (.....
i.c., it keeps a single non-zero entry in each column of G, with all of these entries equal across columns.
These are only some examples of the compression capabilities of sharp operators. They open the door to compressed server-to-worker communication even in the absence of primal compression, as briefly mentioned in Section 4. Indeed, instead broadcasting the compressed messages S in Algorithms to 3, the server can compute GF, transmit this naturally compressed object, and let the workers perform the model update locally. In doing so, we preserve communication efficiency while avoiding the introduction of additional primal compressors.
30

---

Published as a conference paper at ICLR 2026
E CONVERGENCE ANALYSIS
E DESCENT LEMMAS
We provide two descent lemmas corresponding to the two smoothness regimes. The first applies to the layer-wise smooth setting.
Lemma 1 (Descent Lemma D. Let Assumption 6 hold and consider the update rule x Xp
X %/2 (G.). i = I....,P. where Xk+1 [x] '...,Xp 1 = ,
Xb1, G
[G .G...... Cy1 E and 7, > 0. Then
p
3
p
f(xttl) s f(x*) +
V
2
2
4
-=1
i=1
L
(1/0) ||cH1lo..
4
k
2
Proof. First, for any > O, we have
V+](x*) 1la0 =
[|V,[(X*) - -G + +cll
(28)
[++)][V/[(**)-- + (1 c
meaning that
- -1|c% S
1+ | - - ()) I+ (x*
1
1+ I
= V - G.lley. s
s+1 [V lo...
(10)
Then, using layer-wise smoothness of f and Lemma 14 with Li = 0, we get
f(x S s(x*) + (V x xx) + 2
x
x
x2
2
1=1
p = f(x²) +
(+f((* x, -x²)
+
p
L
| - -x x l::0
2
1=1
1=1
p
f(xt) Σ k
(vij(x*) -ot, ct.(C))')) (c))
Σ
< ) (9)
+ Σ
p
(a (c 1
2
-=1
Σ o |
(38,(4)
f(x*) -
Σ x
(v)(( - G.. (ct)*)
(2)
p
p 0
Σ
|||
2
Σ
(o:) ||o 18.0.
-
1=1 2
+
1=1
k | 1oo-
f(x*) - Σ N/ (V. c (a+)))
2
i=1
+
=||\vff(**) - 1lo
- -
k
V x
2
2 s 1
1=1
+
2 (1
|
31

---

Published as a conference paper at ICLR 2026
Therefore, applying Fenchel's inequality, we get
1(x
(29)
S 1(x
+ Σ
p
k 2r
- -
=2
+
2
|c - 2 Ic
2
i=1
+
$ 2hh
|| |+((x
- G
1| (1)+
x s
||
|| 10)= +
L
(1:))
c l(0))
2
2 s+11
2
p
1/0
f(xx)
+
Σ
+
|ea((x() -c
8
V
x 11:00
2r
2
2
s+1
im1
iml
1
L
(
|
2
2
for any r > Choosing and r 1/2 finishes the proof.
The next lemma is specific to the layer-wise smooth case.
Lemma 2 (Descent Lemma II), Let Assumption 8 hold and consider the update rule X
LMO ((x)B( (t)) (Gt), L.....P P. where X [XX+] :......t:1 xx X
[X]..... Xb1, Gh
[G),.. G E S and t > 0. Then
s(xttl) S s(x*) + 2 Vc](x*) - ct -
p
L7 L
v+f(x*) (0~ (1/)2.
2
Proof. Assumption 8 and Lemma 14 give
1(x
Li + Li v+f((*>
(8) | - - 1
s f(x*) +
+
2
P (v.((X*)))) x, - -x) +
L + I
V
(10 || x 113
)(x*)
Σ
2
1=1
f(x*) +
(V. - - x4t1 - xt) +
- x))
i=1
p
L + LA v.f(x*)) 1lox (1/)2
p
f(x*) (Vif(x*) xf+ - xt) + ((G:llew)))
p
D + L V 1100 (1992
1 + L1 [xf((*)) (+) (1))²)
< f(x*)
([Vvs((xt) - G 1lo -
+ ||G ((o)
2
where the last line follows from the Schwarz inequality and the fact that |x/+ - 11,1 t Therefore, using triangle inequality, we get
((xtt1
32

---

Published as a conference paper at ICLR 2026
p
$
f(x*) Σ
1/
Vif(xt)
10 + [V Gfll + +V(x(())))))))
p
1. + 1
V
(
2
1(x*)
Σ
2 [[, - collo) -t, V.5(x*)))))
i=1
4 + L3 (l (1:t22
2
E.2 AUXILIARY LEMMAS
Lemma 3. The iterates of Algorithm 2 and 3 run with C.²² E B((ar) satisfy
E x w [10] s (1-
x
w 2 (0) +
(
|| [
Proof. Let Be [-] denote the expectation over the randomness introduced by the compressors. Then
Ec x - W 110]
Bc: I c
- w!b) -
x ((()
(1 (p-ap) ||x,t_ - -w:$112
(1-a0 (1-ap) (1+ ) || - w (|:s + (1 (1-ap)
(1+0r)
2
| k+1 - -x,113
ap
(0)
(1 0r)
| - w 1
2 | x (20)
2
ap
remains take full expectation and use the fact that
(2) N ||G || 11 = = (c+
E.2.
SMOOTH CASE
Lemma 4. Let Assumptions 7 and 10 hold. Then, the iterates of Algorithm 3 run with Ch. E B2(ap) satisfy
E [| - GG [2]]
<
(1-2)
0D
E
[
M - 1
+ aD
6,2
||MM - Vefs(x*
2
6,8
+ agg
(T..)
(2 [|]
6
2
||x w (1 a
a
Proof. Using the definition of contractive compressors and the algorithm's momentum update rule, we get
Ec
[|M - G
Ee
[| - ct,
- -
33

---

Published as a conference paper at ICLR 2026
(I- ||MMj G 1
where Ee [ denotes the expectation over the randomness introduced by the compressors. Then, be the expectation over the stochasticity of the gradients, we have
letting E
E ||
c. [
s E
EC
|| ++1 - -.
| - [
< (1 an J
(1 aD)
BE - -B) + (
G. [b:]
(3)
(1 - aD)
[1(1 -B. + -V 112]
(1-aa))fp
|| (W 55 - [
(28)
(1 aI- (1+
02)
||aM, - G [2]
I
E
2
+(1 an)) (1
2
3E
[|ME - ve(,(W ||² + +11 - an)) +1
ap
where in the last line we used Assumption 10. Then, Assumption 7 gives
||
E
(30)3))
(1-
a.
E
||M
2
B || (. + +(1 (1-ap)
2
0D
(28)
(1-
uD)
E
[IME -%,
63
_
[ -IM.t-V.,G(XYH [
2
6,8
| - v.f(( [
aD
632
|| V.fs(WA+1]) 11/] +(1 op)
aD
65
<
(1-
E
[IME -
E || (x
2
aD -x +1 [co]
6/8
aDp? (4.,)
2E
||
6
(L)))&B
||
11,0 + (1 - 2
aDP?
Noting that || x (1)= ||))
(44
7 ||
finishes the proof.
(1)
Lemma 5. Let Assumptions 6, 7 and 10 hold. Then, the iterates of Algorithm 3 satisfy
- - [12]
E
1
1
3
E
[[+,,(x*+) - -
2
(1.
(1,))
| [
2
Bee
+
83
1 3
)
(L 'E ||| w 2 [1:]]
+ 3307
0²
and
E
Vz[(x++)) - 1,
34

---

Published as a conference paper at ICLR 2026
1
-
2
(1,))
||c [8)))
-
2
B.p
87
2
(L:)E1
[|xx7*] 11
3707
1
3.
(1)
+
n
where M.. n2 شラ=ı Mry
Proof. Using the momentum update rule and letting Eg [1] be the expectation over the stochasticity of the gradients, we get
E [1Va)) MM [[]]
Eg (X*+)) (1 -11B)) B.V. (W++1; +
(3)
[V,s((+*+) - -1- (.V.I
+ [IV.f](W**)]
(.f,(W
(28)
41 |V.f((**+)) - MMtll 112
s
(1 8,)²
[+
2
+
1+
2
(efs(x - V.fg 1
B:
RR [[ (5t1) v.f,( [13]
(30
(1- 8)) V.Js(X +ス - M 11
+ 1 B.
2
Vifi - 1|2 8.7,
where in the last line we used Assumption 5. Then, Assumption gives
E
[[
M + [2]
E
Eg
|| [xf,((
- [22]
(2)
1 Bi
Vi -Mh 2/
S
(1
- B.)
E
2
+(1
3.)
(1+ 3,
E
- vxf,(x* )|
+33 1 +
2
[|
Vess(W(++))) ||/] + +8000
E
B,
(30)(3))
(1-
A
[|V+/;(**) - M [2]
2
|||x
xx
E
+ B.
2
88 2? 3, (1 1+
2
(I 3 ||| -.o--W--110 116)
+ B.77
To prove the second part of the statement, define v.f(x: = : (x: Then
M =1 (1 B)M + B V.f g gh+1), following similar steps as above, get
Te. [[Iv M. 11
[[2]
= E
Ee
[| - -1- B.) - (w.
(E) E
(= - (1-3.MM- - -
35

---

Published as a conference paper at ICLR 2026
E
E
V
:W v.((h |1
(28)
(1 8.)²
1+
Bi
E
[| - M [13]
2
+.8
11
2
|
v.f(W++)))]]) [|[]]
Bi
- v.s(W ||]]]
EE
(30)
(1 - B,)E
V| *+ - -M111
2
V V.f((xtt) v.f(W*-1))) + 3800
+
1+
E
B
(28)
IA
(1 8.)
(1+
+
B 2
E
[xx - M9913]
+(1 - 84)
(1+)
2
E [ef((x=t)) - v.5(x**)13 ||3]
+ 1
2
E
[IVef(x*tt) V.f(W* - "116]
8309
B,
(00,31)
(1- B
E
V.,((**)
11
2 (4.)*L0
[|x
11( (1())
2
3.0º
+
B?
|+
(49)2B
||1x11] - w
+ S902
3i
n
It remains to use the fact that
x
x111(0 1t ||>>)>]||
(4) n | 10)
(8)
E.2.2 GENERALIZED SMOOTH CASE
Lemma 6. Let Assumption 9 hold. Then, the iterates of Algorithm 2 run with C+ = I (the identity compressor) and Ch E (an) satisfy
E [ef,(x x+1) G |1(0)
x Ga
s v1 an Vif,(x*)) - -ttylen + vI - an (EMS + D [V.,f((**))(())) th
Proof. The algorithm's 's update rule and Jensen's 's inequality give
E
Vif,(x+)))
G.t 1 xx+11(G G+]
Vil,(( - Gt. -Ct, (V x++1 CK
E
E
(X - Gh Cr,, (xff(( - G ,h,,) (1)+
X G*
x VI- aD
(xf,(x - -Chlleen
< V1 - aD
Vit,(** - Eislles + ap [if((x(+)) - vxfi(x
VI- AD
(x*)(**> -Ctsllor
+V1 - an (Li, + Li, (V.s((**)l(o)) -
where ||x,111 - xx1119) =
36

---

Published as a conference paper at ICLR 2026
Lemma 7. Let Assumptions 9 and 10 hold. Then, the iterates of Algorithm 3 run with C₁ I (the identity compressor) and Ch25 E ((2(( satisfy
E
|| G., 112 X M, [tu] S
v - aD M - Ghill + vT -I I Vx5,(**)
v1 -1- v1-apßi aDi
(Les + L1 [V., 1 + a
Proof. Using the definition of contractive compressors and triangle inequality, we get 1
E
|
ct3 ||a M,y' , ct..]
E
||MM Ch Ch (M G2,b1l MM ct.y
| - - Ch CM - G. G M ct,)
- = v1 11-an |[1 M. + (. cl
v1 an - l12
Hence,
EE | k+ 1.
o 11g X Mr, [
-
E
||M
G 112 MHJ | x M, Ct
v1 apE
||(1 M + 3 V.I, - Ghollal , Ot.
VI - aap
||(1 - B.)M + &V (
Gtollal x , MM [%,]
+ +Vi a
Vi (X 11 - - Ve |11| x M [t,b]
IA IV
VT - -a
MrJ - G + v1 a |M - Ve 111- + an
VI an
Mt - - + v1 ap.
|MM V+fs(x*))l
+VI- aD& [xff(x**) (xfx(((+)))
)|12 + VI - -ap30
VI aD Mt otyll + VI -aD3: |I..M.,, vefr(**))l l|1
aD (Ibs + L (1)(( ||x,* - xx "110) + ap
P.
Using the fact that
||
x = 1/ finishes the proof.
Lemma 8. Let Assumptions 8, 9 and 10 hold. Then, the iterates of Algoritium 3 run with C. I (the identity compressor) and th = t. satisfy
4LL
E M
v.[(x [|]] <1 - M1 [I
- v=
3eg.
t 1
Σ 41 Σ
(1--8) - 1 [xfs(()))loo) + ta//2
&
2 n y=1
n
1=0
and
1
E
-
5(1
-(1-8)= 11
I v.f((x5)12)
+t 2
(1- B.)
37

---

Published as a conference paper at ICLR 2026
1 1 8: 1
1 R
[
[100]]
Biors
2 n j=1
where M = n Ei= Mty-
Proof. The proof uses techniques similar to those in Cutkosky & Mehta (2020, Theorem 1). First, using the momentum update rule, we can write
M (1-3))Mf + B.V.f (.
(1- (.) - vxf((x*)) + -1- - (x(((())))
+B (X +)
and hence
Ufit = (1 - B.)UU,, (1 - B.) + +,
where we define Uf. == M Vxff(x*), U2,13 == vef,(x*) (xfx(X and U3 ==
{ (xfj(x
(X*}( Unrolling the recursion gives
k E - 1
k
1.. (1 B.) ++1 Vaa
2,0 + 3, > (1 8,)+ 3,1j
1=0
1=0
Hence, using the triangle inequality,
1
E
1+1 1,
n
< (1 - 11
1
Σ
UP
U..
E
Σ (1
= 1
Σ
Up
2,
72
n
1=0
+
Σ (1
B,)*
(11)
3,3
Let us now bound the last two terms of the inequality above. First, triangle inequality and Assump- tion 9 give
E
Σ (1
U 2,
<
(1 1,) [22.
-
8)- -
@
(1-
1 (T. + 1
|| (10))) |
x
n
1=
B,)
6+1 T0
+ to 3
Σ
| [60)
e n= 1=1
1=
1=0
4
Σ
(
3eg
38

---

Published as a conference paper at ICLR 2026
and using Jensen's inequality, the last term can be bounded
Σ
-
1 72
U++1 3,2,
E
Σ (1 3,)=
1
1++1
(10)
Σ
(1
3,)
Σ
E
|||
3,
n2 3=1
==1
1=0
k
(10)
Σ (1
8,)2(k-))
Σ
0,
01
Σ
(1 3,)2 v
S 05
n²
vn
Vn - B.)) VnB.
t=0
==1
1=0
Substituting this in (11) yields
n
1
+ 1
E
1,6
S (1 - B-))tttBB
n1
U0
1,50
B.L.
+ 4 2 T
(1 B.) I [l [0. + 3: oi
VnB,
1=0
To prove the second incquality, recall that - (1 B, 1,, + (1 - B,)U + =.U Hence, taking norms, averaging, and using the trianglé inequality,
Σ E [|o |=1
(1 - B:)
n2 E
Σ
+ + - B,))
1
Σ
E
| 1 l
n
2=1
==1
+B,- E
(12)
j=1
where the last two terms can be bounded as
E
112
E
VI - (x [|b]]
(9) 11
E
( 13 | )l(o)) ||x - -x,+1| |1 [0)o
2 n
T.0
1
Σ
[1
t
2
2 n2
and
n j=1
| = 1 Σ E
(x((( 10 v.f((*++) [11]]
(10)
S o,
It remains to substitute this in (12) to obtain
E [101131|]
S
(1-B)
1
E
| Tuiis [|]
+ t,
(1- -1-B.LL B.)L
11
Pr
j=1
1
3,
n Σ
+4
E
+ Boo..
2
n 1=1
39

---

Published as a conference paper at ICLR 2026
Lemma 9. Let Assumptions I and 8 hold. Then
S s(() f
1 + 1 (x)(())
for any X _X Xpl E.
IV IVJ Io.e H2 for some H, E
Proof. Let Y = Iy,...., Yp1 E S, where Y, = X, a 11., (V.f(X)). Then, Lemma 14 and the definition of subdifferential give
+ ((X)ll f())
P LT 1
(Vf(())
((())
f(Y) - f(X) + (Vf()), x)
|| - --il()
2
P f(X) + (.f((), Y, Xi(
P L + L. VV
||x; Yillo)
i=1
[Vv (0)
f(x) Σ
(Vif(X), H.)
L + L. V.f(())
17 + L VV)
11
L + L V.f(X)))(()
(X) 18
[V)f(X)
( f(X) +
D L. (x)))
2 2 L [X)
= f(x) - Σ
2 Li + 43 V v.(())l(o)))
f=1
and hence
II |V),(X))
S = S f(x)
L L (x))))))))
as needed.
Lemma 10. Let Assumptions and hold. Then, for any x; 0, E [p], we have
<4max(x,,)) (f(X) f*) + Σ 2, 2i
#Elpl
max (z,L,)
for all X E.
Proof. We follow an approach similar to that in Khirirat al. (2024, Lemma 2). Applying Lemma and Lemma 12 with y VV (1)g))) 21 Lọ L3 [X)]l and any positive 24, we have
(xx)) 1
(f(X)- 1t)
0
Vif(X)
x L. Dp. x7 L: [Vf(())
L 1
x L max (z L El= 2; (
40

---

Published as a conference paper at ICLR 2026
(Di (1 2 2 of El V.
if E. 2, max (xr 2 Ei= 2; [V.f(X))((
otherwise.
2 maxie (x)1,)
Therefore,
4max(x;; (z) (f(X) - -5), Σ x 0
Σ
2 Vx)(())
M max
sEby]
max (n) (x;
2 ( -M) + Σ ? L.1
xc(p]
maxie(p) (z:
Lemma 11. Let Assumptions and hold. Then, for any x; 0, i E [p]. we have
S L f*) + 4 ( -f*
ielp)
Σ x. Lo
max (sa
for all X E S.
Proof. The proof is similar to that of Lemma 10. Applying Lemma 9 and Lemma 12 with y/ [V+[5(X)) (.) 24 1. + L [V)s((X)) (1,)) and any positive 2i, we have
V. (X 1i--
2 (J)(X) 17) 2
1=1 Loj + L,,,
Σ 1 34 | (x)1)))))
Σ1 1= x D. + E x L1, (X)]((
E x; V.f,
E xt 12 + max (1; ( Ed=1 2e V.f, (X) 10),
(2 11 (x)x
if Ef 1o
2 Σ Is (x)))
" x7 17
maxich)(aa
El a1 V
otherwise.
mma (x(()t))
Therefore,
E._ x 12
(X)
max
4max(1r.1))) -303--f)) fj) iElp)
mmxxi (IIL);,))
max( (. (X)(X)- -1)) + Σ x L2
s
Ellp
max:e(M) (z;
4max(z xz (1,,) (X) -f")+ (")44max mmax(x;L)))) -f5)
sElp!
sE(p)
Σi= x LoJ
max (x:Li..)²
41

---

Published as a conference paper at ICLR 2026
E.3 DETERMINISTIC SETTING
E.3.1 LAYER-WISE SMOOTH REGIME
Theorem 14. Let Assumptions 1, 6 and 7 hold. Let (X*) IK=U run with Cj E B(ap), Cr E (. and
K 2 1, be the iterates of Algorithm 2
s オे = 74 S
1 ,....p.
2L + oD V
4 12 +
66 L
as
Then
1 K-1 P
% EE 71
E
Vef(x*)
2 (1)+
S 1 4N00
K
K p
Ei=1 7
0= 1=
where
Tp
= 1(x%) - -5
6%: 1
Vilz(x**) -Ctulies
aD
i=1
j=1
66% Σ ob
( - -1 op
2
Remark 15. Theorem 3 follows as a corollary of the more general result above by setting p = and initializing with X = W and G. = ( for all j E [n]].
Remark 16. In the Euclidean case and when p = 1, our convergence guarantees recover several existing results. When primal compression is disabled (i.e., op 11, they match the rate of Richtárik et al. (2021, Theorem 1). up to constant factors. With primal compression, the rate coincides with that of EF21-BC in Fatkhullin et al. (2021, Theorem 21). Additionally, our results match those of EF2-EF21-BC (a bidirectionally compressed method with error feedback for Byzantine-robust learning) from Rammal et al. (2024, Theorem 3.1), in the absence of Byzantine workers.
Proof of Theorem 14. Let A,, B, > 0 be some constants to be determined later, and define
p
vh = f(x+) f + Σ A V+f,(x** - c. + +Σ Bi || w 16:
Step Bounding E
|I (x
G
+ [oo]
The algorithm's update rule gives
E [[
1 1| x . [ B1[V
B
oi cht, (W.t( ot 1100) xh+t1, whit1 , c.t]
(2))
(1+
aD
V. Gt. -ct, (V.,,(WAA))) - - xh+1 wh+1 Gh.)
E
2
1+ ap
2
V.f (. (Poor 1t4y w [.
(1+ aD)
[. - 1l.. xt+2 th+1 c...
2
2
E [[V11)(x*t)) - 1:00 1 w++ [.:.
1+ an
11 oD)
| |ifj(W**t)) - 120 + (1+ 2
|V,f,(X(+|) - (ef,(W*+)))?? (1
2
ap
(1- aD)
(1+
42)
V.fj(x*) 1o
2
4
+ (1
4D)
(1 1 ap
4
[[,,( - -x,,(
42

---

Published as a conference paper at ICLR 2026
+ +(1+ 2) )
an
(30,,31) S
1 + |xff(x**) - -Galleo + aD
4 [ef((W - (x*")>
I+ 2 on |.,}(**))) V.b,(W**)] 1100
Therefore, using smoothness,
E [[V[ - (1)x x Wht ct.s
(1) [V,}((*** ch
4 (12,)) || -x:1l
aD
(1+ aD
2
(4/,) || W 2 (1()
(28) K
(1 4 ) [.fz((*)) 1 -c + aD
8 (1. )² | - -
8 [(D))) 2 || - 1 18: + (1+ ap
(1,,)² || -
+ oD
(1 4
0D) [[,}(((*) - 2 (1).
8
3 ||0%
aD
11 (L ||1 11+1 11:0
aD
Taking expectation, we obtain the recursion
E
|if((xxt))
- 1 (1)+
(1- 0D))
E
Vif - - (1)a
+ aD
(12)))92 2 [11]]
4
+
11
(
||x, - -w [()
(13)
aD
Step II: Bounding E [||x - w |1:0
-
By Lemma 3
E [ k+1
wh 1
[co]] - (1
ap)
E x - w:
+ ap
2
||
[1))]
2
(14)
2
Step III: Bounding w+l By Lemma and Jensen's 's inequality
g++
= =f(x - fr +
A
V (x - - +
Σ
B, | - 1
iml
im1
$ 1(xh)-1)) - fr Σ
3%
-
p 2 - Σ
2
4
i=1
1
L
22 | 12:0
4%
2
A
V
- G 2
+
+Σ
B
[ - 1
(8)+
1=1
1=1
1=1
S f(x
f
2
3% 1
Σ 24 Vaf( 1120
F
2 11
j=1
i=1
43

---

Published as a conference paper at ICLR 2026
P
1
n2
||
4%/
2
P
+ p + -
Σ
4 n
Σ
V
- -Ot l (0)x
- w. 13 (0)
Taking expectation and using (13) gives
E
[
B - - +
Σ
3% 1
E
[[ - -
Σ
2 E
M
2 n
4
==1
i=1
P
(1 1
1g)
n
+
Σ
A
Σ (1
aD)
E [[( 16o.]
4%/2
2
n /=1
4
1=1
i=1
+
A
n
(1.
)2 2222 E
2
+
Σ
A
Σ
11
(1
''B ||
- (13)
0D
1=1
==1
BB [|| - 1,2t11(3))
+
i=1
[[((x - = + Σ
37
+ A: (1- 0D))
))
1
Σ E
[
- -collio
=
2
4
2=1
j=1
7 E
[(00]
2
Σ
1
L
A aD
8
(19)2)
,, [|03|[o]]]
47
1=1
+
A
11 (Z))² + B. E
||
+1
11tt1112
[1cc]]
aD
Next, applying (14), we get
E
S E [ -++
+
Σ
p.
(3%4
+ Ai
(1-24)
2D
4
_ Σ E |
2
1=1
E | (1)))
p -Σ
17
A aD
8
(2:)2)
7 [|ac
4
47/
2
1=1
A
11
(IQ)² + B.)
2
1/B
[|a
[(:0]]
aD
ap
+
Σ
p
A
11
(I + -B.) (1- ap)
IE |||X - -w:1
ap
Taking A; = 6% 0D
and B;
A
11
2
-1) (19)2= = 66%
2
1) (T.)) yields
oD
ar
op
op
3% + As (1- an
= A,,
2
A _ + BB) (1- a.))
11
Bii
ap
2
and consequently,
p
1 n
E 1 ++1
s E 1
+
Σ
A
-
E |[V - Gity 1oo
n2 1=1
|
x
112
1
- Σ
(11
1
SA, (in}²)
7f [|CH|] 2
47
2 aD
1=1
44

---

Published as a conference paper at ICLR 2026
P
B
2
2
Σ
P
B,E
||
D2 2
ap
(1)+
E
V
E
V
4
(1)+
i=
1
L7
8A,
(2
4B,
G
[[:0]]
47
ap
ap -ap)
Now, note that
1
L
8A,
(1))
4B,
48
(L))
+
264
(1,)²
vi 2
4%
2
ap
(2
ap)
47
2
aZa
for 75 S 2L.+2V6
For such a choice of the stepsizes, we have
T k+
S E
78
E
(1))
4
and hence
K-1 P
K-1
ΣΣ
E
Vif
<4
(E
[
E [
S 4
(0)+
Lastly, dividing by
K
EP=1 n. we obtain p
1 K p
E %
1 4
E
K k=
(1)+
K p
Ei= m
E.3.2 LAYER-WISE (L", "T)L'' REGIME
We now consider a deterministic variant EF21-Muon (Algorithm 2) without primal compression, which iterates
x LMO LMO B(X (G+),
1 Ct, + Ct (V:lj -
1 M G c + Ct, G Gt n
-=1
This corresponds using identity compressors the server side. Theorem 17. Let Assumptions 2, and hold and let [- K 1, be the iterates of Algorithm 2 run with Cj = I (the identity compressor), Ch (a()), and
Ik20
t = VK +1
i= 1, 1,...,P.
for some > 0. Then,
E || |[()]
min k=
1=1 p DE 11
exp (4 max;e [pl.jEmm] (n/C C(L,T))
VK +1 E1
45

---

Published as a conference paper at ICLR 2026
1
1
n Σ
max C,LL.)))) (s* - fī) +
(1-
R
C.L.
+
Σ
DD
VK+I 1 E 118)
12
se(r.
n2
Tiss
im
where C == 1. 2 + T-p
3/T-abli.max
Dj = 17
2 2 1
2 +
a-p and
w= == =(x))) - f
2t,
Σ j=1
[V, - Gh
:=1 V1 alapp
aD T.
Remark 18. Theorem 4 follows corollary of the result in Theorem 17 by setting p 1 and initializing with Gj] = (X") all j E [n].
Proof. Let A; > 0 be some constants to be determined later, and define
T := ,
A Σ
1 n ==1
Vifj(x*)
By Lemma 2 and Jensen's inequality
P
R
gh+1 = f(xx+1) - fr+ + Σ
A
/rft(xb - -G G 1
1=1
j=1
I f(xt) - f + Σ
2t:
Vif
ΣEth t (xf((*)))l
1=1
p 17 + L
v+S(X*) (1)a
13 + Σ
A
1
Vefj(x - oft
2
n j=l
n
p
S s(x*) - f +
21
Vx}((** G. (0o)
Σ 4 Vi 11(0)
n
i=1
1=1
p L7 + Li
11
+ Σ ==1
A
1 n Σ
o 11o)
2
1=1
Taking expectation conditioned on [x++] xx G*1 and using Lemma 6 gives
E 4 |+| x
xx*
G*1
S 1(xx) - 1 +
24 n1
-
Σ 4 V+]((*)))
i=1
1=1
p
LJ + L
V
1100 1
P
Σ
A
E
|| (xttt)) - - (10))) x x*,
1=1
)=1
(
f(xt)- - f
F
2t
1 n
Vifj(x*) - G Σ t Vi 1lo0
p
n
T + T
V 1100 ?
2
A,11 aD
1
Σ
(
+ (ILs + 1ts [T,f(((*))l t:)
j=1
P Σ
p
= s(xt) f +
(2t + A.VI- aD)
(x+) G.alllo Σ t Vif(x
71
46

---

Published as a conference paper at ICLR 2026
4
1 a
[l +
11--ap an Σ Asli
Σ
ررا
[x,,,( (100))
2
j=1
P 0.4
p
Σ
v1 aD
AdtiL
2
Now, letting A;-
I-VI-aD" we have
21:
21;
21; + 4 10:=222 21 +
VI - ap = -Aa
1- vI aD
and consequently,
E | Xx+1 x Gh]
p
n
p
s s(xx) 1
A
-
Vil - choll
4 [
p
Vi(((**)
11 + v1- oD
Arti
Σ
L V.fr(**))lo) 11(0n
2
1=1
j=1
P Σ
p
+
+v1 aD
A
2
p
p T.
1 n
S
th
4 {/
n 1=1
P
2 I - an 11 i,aax
n
p
E.L
-
2 2/ aD P.
13
(xfs(**))lon)
+ Σ
1 V1 10--0 aD
2
1 v1 11-1ap 0D
4 V| 11
1=1
L
2 _- i.max
()) + + 2 - aD 70
L
+] Σ
+
2
1 aD
2 1 OD
i=1
c=
Taking 4
VK+I for some m/ > 0 and using Lemma 11 with x4 = 1// we get
E 1 x++1 x* c*]
n
te Vaf((*>))| (1)) K +
[,,,((**)
(0)
+
> D2t
1n
j=1 1=1
(1)
Σ
4 [.f((*))
Dut2 +
1
1
4 max iElp
max C.TLg)) (1x(+))
- 1*)
K +
1 i
1
1
n CZL -
+
4 max tElp
th C.LL)) (f* -f)) +
K
n
K
n
max (p) (nk C.L..)
4
ni Vif((**) (+) + K + 1 KElPieE(M
4
(n (
Σ (1)( - fr)
x
max
VK +
n j=1
1
1
4 max 2 C.L.,) (f* - 13) sElp
T
C.LL,
+
Σ nr D:
K ++ 1
t=1
47

---

Published as a conference paper at ICLR 2026
Now, since E== (f)()) f() f(x+) - ff I v obtain
E
[[ [pptt
x xh+1 , G*]
$
1 +
4
max (n ((LLs)) w²
V vif(X*
(5)+
K +
JelpI.EEMM
VK
1=1
1
4 max (17 C2 C2 (f* fj) + 1 sE|p
C.L.
> n
p
+
D:
K
K+1
Lis
j=
2=1
Taking expectation,
E g
P
1+
max
(n) (gLL)
E E
n
V 100m
K + 1 sElpIEE
VK +
1
Σ
n
4 4max(1 (f fj)
P
C
> T
p
+
n
sE(P
Lij
and hence, applying Lemma 15 with A E [Tp and B
E [|V,1((**)]coo))
vK+I
min O
Th
E [+f(**)
[(0]]]
K +1
(1/ +1)
exp
K+I
maxxe(pl
10
(K 1)
4 max v
13)
Σ n D.
K +1+1 1
n
EE|p
1=1
1 E di finishes the proof. Dividing by 2 VK+1
E.4 STOCHASTIC SETTING
E.4.1 LAYER-WISE SMOOTH REGIME
Theorem 19. Let Assumptions 1, 6, 7 and 10 hold. Let {X*) (x)) K 2 1, be the iterates of Algorithm 3 run with Cj E BB(ar), Cta E BB(an)) any Br E (0, 1), and
O
0 7 = <
I.....P.
2L +2V
where Gi =
12
(L + 24 2)
2
36 (6(
(1z)2
(288 aja3
(19)2)
Then
K-1
E
Vif(X** k
||201]
K
1-117 71
K=0 1= p
1
0
1 +
(1-ap)a
128 2
of Bitk
<
+
Σ
+
(15)
K p
E
1= 7
aD
a
I -E' m
1=1
p
where
p 6pF
To == - fr +
E
Vif(xo)
-M 1
3
1=1
48

---

Published as a conference paper at ICLR 2026
+ +
p
72p Pi
1
Vaf - Mbile
P
6p
||
74
78
a
n
aD
71
M,y and MR = レ M.y
Corollary 1. Let the assumptions of Theorem 19 hold and let {x*] K Σ 1, be the iterates of Tk=0
Algorithm 3 initialized with M. G v.ff (X. EJ), j E [n]. Then, the result in Theorem 19
guarantees that
K
78
E
vif(x*t
||
K
-== 71
0=
4 (f(")) )()
24
Σ
1
125,
0.p33%
$
+
KP Ei=1 71
K
VnB;
a
2
I
-E=1 Y
t=1
p
(D-a +24 +
12
0f
QD
D
2
1
Ei=
/
Remark 20. Theorem 5 follows as a corollary of the result above by setting p = 1.
Corollary 2. Let the assumptions of Theorem 19 hold and let {X*] Sk=0 K K21 1, be the iterates of
Algorithm I (Algorithm with p 1) run with C₁¹ E B(ap). C. E <Bz(a)) Choosing the stepsize
1
Pi R
"1
V25 + 21
13331
PapaD
(16)
and momentum
yoU LUn e
LOL 2
4o L9 gook
B1 min
1,
(17)
the result in Theorem 19 guarantees that
[
K
To garaDk
Po
To p091o91 1/2
To po 1 e givonk
10 2/3 L4
3/1
(gnnk
gink
+
2/3
K
Remark 21. In the Euclidean case (p = - = 1), without primal compression (ap 1). and for
== 1, the result in Theorem 19 simplifies to
K-1
yi Kr
1
3
32
Σ
E
VJ(X*
+
30
K
ODD
0
2
for f O 4
which recovers the rate of EF21-SDGM in Fatkhullin et al. (2023, Theorem
3) (up to a constant).
Remark 22. In the absence of stochasticity and momentum, i.e., when of = 0 and 3 1, and under the initialization W = = M = G9 = (x)())), Algorithm reduces to Algorithi 2. In this setting. Theorem 19 guarantees that
K
1 14 5
Σ
Σ k= ==
%
E [IV.
K
I
E
K
D 71
for
0 +/ = = 2L + 2
i= L....,p.
where G == 2. (12(1))+ (I.)
(i.))²
(f1)2).
This recovers the guarantee in
af
a
Theorem 14, up constant factor.
49

---

Published as a conference paper at ICLR 2026
Remark 23. Alternatively, one may use compressors C B2(ap) in Theorem 19. The proof is essentially the same, with the only modification being the replacement of Lemma 3 by the recursion
Tee [1| x - -
Iec || + (X - - xx 11]
(1 ap)
|| - 11
(28)
S
(1 up)
(1 2
- || - w + (1 -ap)
(1+2²)
#arr op
2
|x - x2:12
(0)
(1- ap)
| wr11/² + 2/ | -x, 100
ap
ap)
|| - - + 2p:
(7/))²
|
2
ap
The resulting convergence guarantee matches that of Theorem 19 up to a modification of the constant Sa. which now becomes
S
p3 12 e²
(1))
24p2 (3:+ 36 (8. + 4) (E9)2 + 144 (2 + 5) (i9)2)
24p
(L.)² + a
app
aZaz
where the additional norm equivalence factors highlighted in red arise due to the use of Euclidean compressors.
Proof of Theorem 19. Lemma I and Young's 's and Jensen's 's inequalities give
f(x
() 2 f(x²) +
3
| - 10 11 -
- |V,
2 1=1
1=1
L.
4%:
2
(17 - 1 -
(2)
I 1(x +3 3 Σ p:.:
Σ
M - G 11
1=1
p
( 1
1
Σ
||c/|1200
-
47/s
2
1=1
Recall that by Lemmas 3, 4 and 5, we have
E
|| x x
W.
(9)
3 (1
ar
2
E
|
W [2]]
+ ap
2 N'B
| 11
(1))
E [|
11
(4)
$
(1 AD
E
||M,
- 1
6
E
| (x+ [|]]
2
aD
(L )2
6 6 )²
||G
E |
11
aDP?
aDp?
+ (1 an) 3 o
2(L2g)2
E [|Vx,|(X - -
1
E
Vxl((** - M [2]
2
3.03
8
1 +
2
) ||xx+41 -ess-------0 w.++
(1)
+ 8.o7
2
3
E
| 1000--14131 11 (5)
1
E
V
(x - -M8111] +
2(2
BB
||o [(o]]]
B.g.
50

---

Published as a conference paper at ICLR 2026
B8 (1 1+ (L9)²E)
2 B,
2 || - 1 [3:1] + AP0?
+
where M.- = , Ej= MM To simplify the notation, let us define 8t = E [x(x)) - 1
== 7: 1 PK == xx -3= E [[ - M S
11]
i دز E ||MM - -Gty [|2] and R == RE [||T w7 3 [. Then, the above inequalities
yield
gh+1 S 8 +
P.Pk +3
Σ
2
Σ
BE (
4 1=1
i=1
1=1
1
1.D)
B
[|GT
[(0+]]
(18)
4%
2
1=1
R
1 (1
ap)
R +
2 L
|
(1))
(19)
2
$ (1 ap
5 63 pk
(..(
||o
[1
2
2
aDp?
(1))
12tt1 + (1- an 989
(20)
ang
ph+11
< (1-
ps +
2
E G [
2
Big
B²
2
+
1+
(L9)² R
(21)
0
81
ph+1
1
31
Pk
2(L
E
||c (1)
B.p²
2
+
1
+
(19)2R4
h+1
(22)
B
Now, let A,, B,, Cr, D3 > be some constants to be determined later, and define
v = 8h +
A.P.t + ΣB..
+ +Σ
c,St
Sh + D.R.
1=1
i=1
1=1
i=1
Then, applying (18), (00), (21), and (22), we have
yp++1
= gh+1_ 8.+1
A,Pl
Σ
B.P
p1+1
+
Σ
C,S,+L +
Σ
D
1=1
1=1
1=1
i=1
P
<
8k +3 3
Σ
Pe
pk
+3 +3] Σ
555
Σ
B
[ [
4
1=1
1=1
1=1
B ||C
2
4%
2
A
(11
3,
ph +
2(T.)22
3
||
2
1+
83
1+
2
(L 2 Rh+ + 87
2
(1)+
B,
n
Σ
B
(1-
3
p 2(202²
B || (1)+
+
3
1+
2
(I + of 8f
2
3.g.
B,
+ +Σ C
(1-
aD
56 +
6 ph +
6,8% (Z9)
7 ||c
2
[10x
2
aD
aDg
1=1
51

---

Published as a conference paper at ICLR 2026
P
p
Ci
R + (1 - 0
+
Σ
D
aDP?
1=1
i=1
B
3p + A (1- ()) ) P (1-2)
Bi
6 aD
8
Σ
Σ5
B:
1
+
+C,
p
2
2
Σ
P
(3pp + Oi (1- --22 55
))
1
1
Σ M
P
[V )x)))
1 [[or]]
4
Σ
A 88
1+ 2
(1.)2 B, 32
(1
1
2
(E9)² + C. (.
+ Di) ne
3:
2
3.
aDg²
2(L9)
2(09)
+ C 633(129)
B [ICh|PoN]]
A1
B.
Bi02
B,
angz
P
1
DP
B [|cctlioo]]]
p > 1=1
A
+ B, + +1(11 -an) ) 0f
4/z
n
1=1
Then, using gives urtt1
8t
+
Σ
p
3p +A 1
(1-2)
P 31) p + Σ
B,
1 - 3.)
+
+
6,3
p.t
2 )
aD
P
3pp + Ci (1-22) 3/ $ Σ
1
P
WEE
1=1
A p2
1
2
(L)) + B
3
1+
1
2 B.
(L9)² + +C, (19)
+ Dj
(1- ap)
R
3.
0²
aDP?
2
A 8
1 2
(L))² + B 8.
1+
2
(I9)² + +C, 6,3% (I9)
Dj
2
M
||cl|l
8,
B,)
aDo²
ap
A 2(L28²
20122
+ (3)(2
T ||CH [
B,
Big
aDp?
P
1
L
[IGH (0)
+ Σ
A
+ B + C,(1 - -an)) 07
B,+CC(1
47/
2
n
1=1
i=1
Taking A = 6/f
8,
B,
=
o
C.
6 and
aD
A 32
1+ 2
(L + B, 82
2
(L.)² + C. (T9
(2
Dj =
1+ 3,
3,
aDgz
ap
6 -
(3: + 2) (19)2 + 63 (28, + 5) (1.92)
(
(22 1).
=
ap
we obtain
B
6p-
3,
3p + A
1
= 3p, +
1
= A
2
3
2
3,
6,3
72 8,
(1
3
6
632
B
1
+ O
+
B
2
an
ap
2
0D
UD
3p3 + C. (1
aD
= 3p +
6
aD
= C..,
2
aD
2
and
32 A 2
1+
2
(1 + B
82 1+
2
(i0)² + C.
+ +D:) (1- 2
B,
3r
angz
52

---

Published as a conference paper at ICLR 2026
D
(1-
ap)
Dj
D,
2
2
ap
Consequently,
T
8
+
A.P.
P.P.
P
Σ
C
ch
Σ
D
1
2
nB
|
||
4
1=1
t=1
1=1
1=1
i=1
Dj
+ Di
E
2
1
17
[cc
2 -1
aP
(1)*
476
2
aP
6p, 2(4))²
32
2(.)
6Ø (3)
|G
B1 3.g
a
ap ang?
1 6p
72p
60
(1 - -aa)) 07
B
2
aD
Σ
E
V
v.f(xx>>) (0)
Σ
B
4
47:s
2
1m1
im1
Σ
12p
(Tg)²+ + 144p (I9)²+
+ 36,3
(Io)²= +
4D;
||c 1 [.)
+
abL²
abp²
ar(2 ap)
1=1
1 +
12
0 0.p%8 (1 .D)
+6
Σ
+
2
D
0D
==1
Now, note that
1
L
12pº
(193 + 144
( +
36,2pp
(i9)² + a.p(2 - a.))
4D;
4%
2
7/2
0
a
1
Li
4%
2
p
12
24 24 (2)
36 (3? + 4)
(ig)² + 144 (28, +5) (ig)²
%
0²
B2
(L
(L +
a
aBaa
2
for 91 S
For such choice of the stepsizes, we have
2/(++210-
1283
(1
al-a
mb+ < t
E
[ (1)
+
+
4
2
QD
Summing over the first K iterations gives
K-1 p R
K-1 p Σ
P 4T0 4K Σ
|.s((*>)
< 4
(v+ - ve+1) + Σ
km
En
k=
and lastly, dividing by
K p
Ei
7h, we obtain
K-1 p
Σ I 78 Em
4 p Xe
4 1 E:7
E
K
71
(1)
K
1
-' 7/6
0=
i=1
p
Substituting x WT proves the theorem statement.
53

---

Published as a conference paper at ICLR 2026
Proof of Corollary 1. Substituting the initialization, we have
E V 242112]
E
-x
π j=1
(10) a,
- -V
vn'
==1
ΣΣ
R
[IV:,((X)) - M3
1
Σ
E
[|V. - v.5((")5))) :5911p] (10)
S oz.
3=1
j-1
and hence
p
6p?
TO = -(x)) - fr
Σ
M1
Bi
1=1
32
1
E
Vits(x))
- -220112]
+
6p
7/2
1
Σ ==1
E
M - 1
aD
i=1
j=1
P
p Σ
7
5 ft
no
+
70.
Vn3,
a
i=1
Substituting this the rate, we get
P
7/2
E
[|
[. [1oN]]
K
E 71
k=0 i=1
4 (x)) f*)
24
Σ
1
125;
o.p3%
+ K i=1
+
K_ Σi= ?
/n,
a
E 7
(1 -a))3
12
+24
+
+
-1 n
aD
a
1 E
p
Proof of Corollary 2. Substituting the choice of y from (16) in (15), we have
K 1 K
E
V
|1?]
400
+
24
-
1
(p-an)81
1255
oipi3
o
Kn
n2
QD
a
Va pi
3
3i
o
+
Papank
K
2
aD
Then, choosing 31 as in (17) guarantees that R&oo into the upper bound gives
<
Substituting this
ap
a
K 1 E
[!]]
K
0=
( pink
1/2
2/3
+ (
40 2/33 2 2/ 01
)
3/0
0
Papapk
+
(pppk
+
p/Vapk
pja K
2/3
as needed.
54

---

Published as a conference paper at ICLR 2026
E.4.2 LAYER-WISE (I", - SMOOTH REGIME
As in Section E.3.2, in the generalized smooth setting we consider EF21-Muon without primal com- pression.
Theorem 24. Let Assumptions 1, 2, 8, 9 and 10 hold. Let {x* K 2 1. be the iterates of Tk=0
Algorithm 3 run with C₁ = I (the identity compressor), Ch2 E B2(aD), B3 = B = (K+1)172 and
t =
7/8
.....P
(K + 1)3/4*
(K+1)/2 (1 G(L.)=
-Tappa (K (111///2
Aming (k+1)2/2
where of S min
24p.(L. mmma )2 .1. Then
24 p. mmxx)
min
Th
E
|V
k= K
-1=1 111
1=1
3500
6
T:Pi
E Vif((x) - 11,11
(K + 1)2/4 P 2 111
(K + 1)
1=1 m
max 2. )²
8
-aD
+
(f- f7)
(K + 1)11/4
(1- /1 aD)(K + 1)3/4
-i=1 171
L.
40,L0
4p.y 1 - anLi
T/1
+
(K + 1)3// p,(K + 1)1/4
2 v1 -a) + 1)8/4
D1=1 1
1=1
+ Σ
P
::P::
4vT-ap aD
2 vn(K - 1)114
+
Σ 1
(1- vI- (K) + 1/1/22
i=1 P
where M. = 1 Eラ=ı MD and
Vo = 5(x0)
P f
2tpp:
1
ΣE
B
|| M [,
V1
aD n2
1=1
1=1
2t.p. - aD 1
-
E
V1 ap
j=1
Corollary 3. Let the assumptions of Theorem 24 hold and let (X* K-1 K 2 1, be the iterates of
1k=
Algorithm 3 initialized with MD, ;x9)), Ch CT,(( (V. (Xx: j E [n], and run
with C. = I (the identity compressor), Ch E B2(a)), 3r = 3 = (K+1)1/2 and
0 c t t, (K 1/3/4
I.....P
(K+1)//2 (I-vT-ap)p 1o, (K+1) Sminp , (K+1) ((1,)²
where ni < min
24 (L )2 1
Then, the result in Theo-
24 mmx
rem 19 guarantees that
T
[[
min k=
K
P 1=11 1
1=1
[(xx)- -
4 - apmipoo (K + 11311111-11111
(K + 1)1/41
Ei= T1 p
P
Pinjiss
(K + 111/22
Σ 21=1 1h
n max 1/ (T.)
8
8vI- -
+
(f-
(K (K1/)
(K + 1) 3) - /1 -ap)
I E m
55

---

Published as a conference paper at ICLR 2026
2
Tl
(K + 1)3/3 + 4p:L1
1
4p.v 1 aDL!
+
P,(K + 1)1/4
2 + () - vl-ap)
2I=1 1
i=1
+
Σ
P
::Pio;
4VT 0D
2
Ei=1
(K+ - v1 ap)
Vn(K + 1)1/4
i=1
Remark 25. Theorem 6 follows from Corollary 3 by setting = 1:
min
E [[|]](((*]]] [[_]
k= K
3 ( -ft)
12v - aDio
6p
+
n(K + +11141
(1 VI aD)(K (K +1)
+ n0
8
8VT -aD
(f)
+
L (K +
1122/4
(1- (D)
(K + 1/3/111//
j=1
111
4
A aD
(K + 1)3//
L
(K+ +11/4
(1- (D) (K + 1/3/4
470 - aD
2p
(1 vl-vvllayy (l1-aa)(4
(K) + 1)11/2 Vn(K 111//1
3 (f(X))-f)
16 -appp
SLO
8po
<
(K + +111/2 + (1- vI (
v(K + +1j1/1
(K + 1)3/1
VT(k + 1)1/1
7P 8
8vI-ap aD
1F) 13) + Lo n
1 Σ
+ (1-V1 ab)(K + 1)3/4
P (K+1)
Proof of Theorem 24. By Lemma 2 and Jensen's inequality
p t [
f(xhtl) s f(xx) + Σ 2t, Vif((*)- - cl
1=1
1 1 |x 1100g
2
i=1
$ 1(xx) + Σ
2/4 [,5(**) - M Mello)- + 21z 21: M - G -
1=1
Σ 4 ||
(8)+
+
Σ
L7
22
+
I
|| t
2
2
1=1
1 n
$ f(x*) + Σ
2p,ts Vef(x*) -
-Ma, +
+2pit.
Σ E [| - challla
12 5=1
1=1
p
L7
I!
4 ||
11:0 + Σ
| 11 12)
2
To simplify the notation, let 8² = E [f(*)) - 1 P. = E V.f - Mf 11ll Pk = 1/4 (1V,, MrI11 and 5t == 1/4 プーー
n
E
||| | - G [1|]
Then, Lemmas 7, 8,
and the descent inequality above yield
5141 2 VI-
5k + VI aD. tivI-
ap
| n
Pi
1=1
I I
v1 -ap3,on
(23)
2
56

---

Published as a conference paper at ICLR 2026
R-1
pe _ (1- B.)* P. +
4 Σ
Lig
Σ
(1-31:10 (1-
[1
Pi j=1
1=0
t, t
+
(24)
p
n
p1+1
(8)
(1 B.)
(111-1 1 a
n Σ
T B (1+5
P2
(1(1 - -3.)
+
+ 3,00..
(25)
&
P
84+1
S 8hh
Σ
4E
[ [0))
(26)
P
+ 1.1
C E
P: + 2t
Let A,, B, > 0 be some constants to be determined later, and define
V ; 8th +
S.S. k + Σ P.
1=1
1=1
Then, using (23), (25) and (26)
T
g++1
P A
= =8b+1
5k+
+
Σ
pk+1
1=1
P
E.L.
8h_
4E
(1)a + Σ
24spi + 21.p.
2
2
1=1
i=1
4vI aD33
Σ
A
aps 5k + VI aD3,Pk +
Σ
P₁
/=1
4411-an anBiLo
Σ
A
+ v1 aDBios
P.
Σ
B;
(1- P. +
4(1 4(1-Bi)
1
Σ
[1,,((
4(1 4(1-8.)1 + B,00
4x
n
2.
i=1
j=1
8k
4E
(. [1(0
+
Σ
(2. Avv1 - aD) S1
L
+ B, - B,)) P.
Pt.p..P +
Σ
E [
2
1=1
i=1
t=1
P
t
(AvI--ap3 -
B B.))
1 Σ
ا E
M10-
P:
j=1
P
W LDDIL
t1 - -.
A
B
2
P₂
1=1
P
+
Σ
Atv1 &D33. +
BB
1=1
57

---

Published as a conference paper at ICLR 2026
P1,P: Taking A; I-VI-ap
and B, AivI an I-vI-op
2-pyT-ap we obtain
2tpp
20,0+ Avi-ap- -t + 1- V1 - aD
VI QD = A
AvI-aD&, + B1 - = = - + A.v -ap(1 - -3,) = B,
Consequently,
1++11
p
P
P B.P.
8c
t
Vi
1100 [000]
+
A
5k
+
p
i=1
P
21 pi
P 1 - aD
+)
E
V.f( [100 +
Σ " j=1
2
VI QD)
1=1
2f ps - -apL 10
4t - aD3,o,
+
+
2
L. (1 VI - aD)
vI- aD
1=1
1=1
1=1
P
84
4E
|V: [00)
+ Σ A + Σ B.P.
i=1
i=1
+
P Σ
24ppi
(1-
(1
4
Σ
k-1
(1 - 3.) "B
Vis [100]]
4.0
+
os
3i
L n
2,3:
1=1
j=1
1=0
+
P
P.L.
E
ViI((
(1)m
+
2f PiV1 11 aD
1 Σ
L E [1,,)((+)))(
2
vI ap)
"
1=1
P
P
21. - L
4t aDBio
+
+
2
2 (1 - vI 11-aay aD)
1 vI Vl-aD aD
i=1
P Σ
1 p
T
ME
[+f( V,,(x>>)00> 111001]
Σ 21,p;(1 - P/
Σ L [|
2
1 n
+2
L Σ
(1-
(1 - B,)
Σ
D
E
||:rf(
n
j=l
+
2 I &D
p
ti
22
1
Σ
E
||
ix 11 11o0
+ Σ
0D
2
1=1
|=1
1=1
+ 28 L
2t Pi
v
LDLo To v1 aD)
4t I -
3 20.p:or
(27)
2,B,
2
V1-ap aD
Let us bound the terms involving the norms of the gradients. Using Lemma 10, we get
Σ E [l)
(10)
4max((((L)))
E [[(x)) - + (E
maxie(m) (L!)))
sElpl
p Crg.
4 EEpp
Similarly, Lemma 11 gives
P 2
1- Σ
1 n 8+)+
(1
-
Σ
D (IV.f(('))
[co
2 1=0
n J=1
58

---

Published as a conference paper at ICLR 2026
Ppp (1 (1- Ba)
- [ [
L
Lipt (1 - B.) 2 (w1,)) IE [ +1
4max TE|D
Pi
tsppi (1 Ba)* (5* -13)
4 max iE|p
Pi
Σ
H
P. (1 - 9:) k Lis) 2 LiS
[ (1 (.))-((L1))) )2
max;e(p]
(2
1. (1 (11))) 6²
max :E|pl.jE/u
&
t (1 - B,) (1 (j) +
t. (1 - B,)
4 max hE}p]
Pi
P.
i=1
and
Σ p.p:
1 E V xx) (8)
E
|| V (x**) [1(0)]]
L
(. )2
(
)2 (8-- f;)
4 max (E(p)
E (x*)) - 141 +
4 max (E(p)
L1
Li
j=1
2
P
Lin))
Li
P.
n1
maXie(nl
[10 (2)
P.
17: (I)
1
(.pp(
p
< 4 max xE(pl.Eim]
gh
4 max iE rElpl
(f* fī)
La
P2
2.
Substituting these bounds in (27), we obtain
1+++1
V
LE
V+]((* 1100 + Σ 21:p:(1 (1-(1) P8 2 m
1 Σ
(E
L2 (1- 3,) (1-8) (L)))) 8¹
+8
max
(ED))EE())
P.
t.p: (1 - B8)* (Izs)) 2 17) + Σ
Epp. (1 82) I
+2
4 max xElpl
L,
2
J=1 1=0
i=1
8v1 T aD
P (T2,)) )2
+
max 1 - v1 - an iE(P)jE[m]
8t
&
aD 1 n
C.p(L
P
12
4 4max tElp)
(f+ +
+
Σ
1 V1 V1-ap n
P.
P2
jml
P p.L.
P
2t p / I - a 0
4t.p I
0,3,0
B4
+ Σ
+
+
+
2tsp:01
2
2 (1
vT - ap)
1 1- v1 - ODD
1=1
1=1
i=1
59

---

Published as a conference paper at ICLR 2026
Since 8 I 1 it follows that
wk+1
8vI-ap
Pp fe (1:3) we
s
1 + 2 max (f (L.)))
+ 1 v1 aD (E(o]J((M]
max
EE|
Li
HE
[|V+((**) (100)
21 -B,)* P.
(AILL)))
M
+8 max VE(p).E[M]
max((1 - -1:6)
Pe
sElp!
)
max ((1- (1)-3,) 3)) (5* -1}) EE
+8
max kE(p)
&
3=1
1=U
42
C. P. R-1
Σ (1 (1- 8.)a²
8V1 - ap
max xE(p]
(5* (1)
L 1=0
v1 aD
P1
1=1
2/1 - RD
22 EppL
+ Σ
ELD
VI 1app-
0D
2
i=1
1=1
21 P.
p
28 piv I anLi 1
Al:p: apBio
B:
21api0s
L.
p
vI
VI aa
1=1
1=1
+ 2 max ( (L!)) + 8/I - ap
(P
max
EE|
vI ap vE(p)EE(n)
P
:00
E Vrf((*>)l((>
2t.p:(1 11p6(111,8)) -3,)* PD
(
A-1
+8 max [E)
((1 8m
Pi
1=0
VI a'
n
(5*-55) 15)
+8
max ic(p)
Buin
1 1 VI- aD
n
P,
4p,L
46, I aDLo
[ - 44I aD
+ +2 3,
+
L +
+
t pio;
2
(1 v1 ap))
1 1-
im1
i=1
C=
:=
= (1+CI) v Σ te v.f(x*) > (1)) + Σ 2t,p, p
1-1 +C2
Σ ((1 Pmin ( Pm - C's + Σ C.C.s + Σ t,,
1=0
Now, define a weighting sequence " 1+C++ C2
where w = 1. Then, multiplying the above
Pmin
inequality by w and summing over the first K + 1 iterations, we obtain
K
R ++1
0=
60

---

Published as a conference paper at ICLR 2026
K
K
R-1
K
W
Σ
C2C2
C2
Σ
((1 - 10-imi) ( -min)a
Σ
E Σ t
||
||
[ x
11) [00)
k=
k=
1=0
k=
i=1
K
w*
2t,p;(1 - -1:) P. +
C² +
Σ w
Css
+
Σ w s,s
k=
k=
A==
k=0
K
K
h-1
K
= (1 + 11)
Σ
W k To + C2
Σ
t
((1
9mina
Σ
w
1,
[|V
k=0
k=0
1=0
k=
K
+ Σ
w
2t, pt,p: - -1-1,1-)ppp P. + w Cs - W Σ Ca,s + W
C.s.s
k=0
1m
where WK-
Ex w Since, by definition, wº S w S w-w 1, we have
Σ w gytti
K
K k-1
K
<
(1 + C1)
Σ K=0
lokup vo + C2
ΣV Σ
(wo (1 - Bmin) )min))-bup
Σ
w
4I
[
k= 1=0
k==
1=1
K P
(1 - B. + +W* C3 + W
C.Css +
W. K
t,,s,
k=
1=1
K
20
K
K
P
S (1 C1
k Th + C2
> (1
Bmia)
)
Σ
w Σ +E
k=U
1=0
k=0
k=
Lipi p +
P
p
+2
W C3 + W.K
G.C,, + WKK
t,Os,i
Bi
C2
K
K
=
1 + Ci +
wo Ap
Σ
LE
[V [1(00]
Bmin
k=
k=0
1=1
Lipi P.. + K Cs + W
p C. + W Σ
+2
4.so,i
i=1 Bi
1=1
K
K
Lipt P B
=
w Ap
LE
+
0=
k=
1=1
CW*Cs+ W
C W
Σ C.C...
Rearranging the terms and dividing by WK gives
min
L
[[
-.......K
K
p
w
WK t
||+ [60N]]
k=
K
(w b
- ( +
Life
p + C2 +
Σ
C.
+
Σ
4Ca,
WK
WK 1=1
Bi
k=0
1=1
1=1
00
2
P
tpp
p0 + Cs + Σ Cass + Σ C,ss.
W K
WK
3
i=1
i=1
61

---

Published as a conference paper at ICLR 2026
Now, note that
w
K
w
2 (K + 1)w*
(K + 1)
+
(1+ + C1 +
C2 jk+1
exp (K+1)(0 + (2)
k=0
Amin
Buain
(K+1) "[T.))=
(
Smim, (K (I+1) 24p, maxe 1 to
Taking 44-
The (K+1)a77 where oz - min
21V1 app )mmaa))
ensure that
2(K + 1) (
1
tElp!
3
8V ODD O
(P((L)))
1
(K + 1)
1 -1- un iElp|JEiM
Pa
3
8
(PI(LE)) )
1
(K + 1)
max
Bmin EED]JE
P₂
3
have (K+ + 1)(C) +
C2 3mi
S and WK
K+1 exp(1)
K+1 Therefore,
p Σ
min
V
(100
k= K
i=1
31/00
p
tp
p0 + C3 + Σ
M.O. +
405,i
+1
K +1
Bi
i=1
i=1
350
6
P
3K(K + 1)3/4 po
TiP:
K11
K +1
8
VI- aD
1 7
(T! )2
(f-
max iElp
(K + 1)3/2
8min
1 1 V1 - 0DD
71
L.
+ 4p:
4piv 1
an
L
+
(K 1)3/2
2.3,
(1
vT an))
i=1
4/1 - 0
+ +2 Bi
K + 1)3/4 P:0;
1 v1 ap
Lastly, dividing by Ep 1/
(K+Hj/77 El 111 gives
P
7
||V
min
(1)4
K
i=1
21=1
3
P
11
Pi
pp
(K + 1) 1 >p m
K +
P
1=1
17
3
CE )
8
VI - aD
max
L.
(5)
(K + 1)3/4
3min
1 VI aD
1 D m
m
Li+ 40.L.
4p, V - LD
(K + 1)3/41
P Σ m
2
2, vI -a))
niPios
4 4411 aD&
B
+2
1=1 ΣV 11
1 /1 ap
3
6 + (K + 1)1/2
Σ
T&P
P.
(K 1 1114411
Σm_m
1=1
p
62

---

Published as a conference paper at ICLR 2026
". )²
8
(1 VI 8VI BVI-00 - aD
n maXiclpl
+
(3)
(K + 111144
aD) + 1)3/4
1/2 + 1111
j=1
Lg
4
4ps v aDLo
+ 2 + 1)2/4
Σ 1/1
(K+1)3/
2
VI (Iaa
aD)( (K + 1)3/4
P
TiPios
(1- -1 4vT-00 aDD
2 Vn(K + 111/11 an) (K + 1)1/2
|||
i=1
where in the last equality we set B, = (K+1)u/2
Proof of Corollary 3. Substituting the initialization, we have
1 n
P.D = [(I
11 =B
V.1((99)
Y1
E
- 5
S
σi
p00
E
- |
= Σ E | - v.5((x :3911] Soi
50
1
E
M.M. [bal|]] = - B [I.. -2,, (v. x [591)11]]
1
V1
aD
E
n
j=1
1 n
v.f((')),]] (10
v1 an n
E
V x 4
S v1 vn-apo,
and hence
wa = 1(x0) - f + Σ
2t.p.
1 Σ E
[IME, -
1
VI
aD
i=1
j=1
p 2tpp. - aD 1
[[ M
E
1=1 VI ap π j=1
2t.p:
p
2t I an
$ f(x0)- - fr
VI V--aDo; +
Σ
VI -aD
1 VI
f(x") 4/1 aptipoo
2
aD
Substituting this in the rate, we get
7h E
1=1 M
1=1
[((x9) - ft
4 aDMPoo
(K 1)1/41
Ei=1"h m
(K + 1)3/4(1 - v1 vl)
6
P
P:Tko
+
(K + 1)1/2
Vn p Σi=m m
63

---

Published as a conference paper at ICLR 2026
max [pl 1/ L (L )² P.
8
8v1 - ODD
1
+
(K + 1)1/4 +
(f-
(K + 1)///111111 vI -ap)
- '11
1/1
1.
4pr
4piv apl. 0
(K+1)33/ +
(KK + 1)1/4
(K + 1 vI DD)
21=1
PiPio;
4 aD
2 Vn(K - 1)1/4
E
(K+ 1) v1 ap)
64

---

Published as a conference paper at ICLR 2026
F USEFUL FACTS AND LEMMAS
For all X,Y e ES, E (where S* is the dual space of S), >> and a E (0, 1]. we have:
|X + Y1 = (1+
||
+ (1++-²) |
(28)
|
t
1
(X,2) <
(29)
21
2
(1 0)
1
< 2
(30)
2
1+ 2 2
a < a
(31)
(G. L.. (G)) - |G]|.
(32)
(x, ||x* 2
(33)
|X||. = ||
(34)
Lemma 12 (Riabinin et al. (2025b). Lemma 3). Suppose that 21,..,, Tps Mi..... max [my] > and ..,..., Zp > 0, Then
................................ Up E R,
1= Di
kiyk)
24
2
Lemma 13 (Variance decomposition). For any random vector X E S and any non-random c E S we have
E || - E [|X - E[X]] + +|
Lemma 14 (Riabinin et al. (2025b), Lemma 1). Let Assumption 8 hold. Then. for any X, Y E S,
p L. + Li IIV.f(X)ll
- ( (V/(x), x) S
Σ
| Yillo
i=1
Lemma 15. Let {A* 1120- {B: 1200 E [p] be non-negative sequences such that
A S (1 + a1) A
B
1=1
where a1,02 2 0. Then
F Σ
expla (K + +1)
min O.... K
B
A4 + a2
(K 1)
Proof. Let define weighting sequence wh
wy Itai
where w 1. Then
w Aht1 I w (1 + a1) A
# Σ
B
+ wh az = -1 w Σ B + w waaz 0z
and hence
K
p
min k= K
B
w
Bk
wh K=
A=0
t=1
K
K
kw A - - Aht1) +
1
Σ
aaaz
Σ = wk
Ekz with
k=0
k=
(w 11A0 w Wk 1 + a2.
K k= wi
Using the fact that w = 1 and Σk= th
= -kko 1
(1+01)=#1 2 (1+a KTI we get
K+1
min EB (1 (1 K+1
exp(01 (K + 1))
(A0 R A + + <
A
O= 1=1
(K +1+
(K +
which finishes the proof.
65

---

Published as a conference paper at ICLR 2026
G EXPERIMENTS
This section provides additional experimental results and setup details complementing Section 5.
G.1 SETUP DETAILS
Tables to 5 summarize the model and optimizer hyperparameters. The scale parameters (Hidden/-
Head Scale) in Table 5 specify the LMO trust-region radius as
radius scale learning rate,
following Pethick et al. (2025c); Riabinin et al. (2025b).
Table 3: NanoCPT-124M model configuration.
Hyperparameter
Value
Total Parameters
124M
Vocabulary Size
50,304
Number of Transformer Layers
12
Attention Heads
Hidden Size
768
FEN Hidden Size
3,072
Positional Embedding
RoPE (Su al., 2024)
Activation Function
Squared ReLU (So et al., 2021)
Normalization
RMSNorm (Zhang & Sennrich, 2019)
Bias Parameters
None
Table 4: Medi www.U--CPPPP model configuration.
Hyperparameter
Value
Total Parameters
335M
Vocabulary Size
50,304
Number of Transformer Layers
24
Attention Heads
16
Hidden Size
1024
FFN Hidden Size
4096
Positional Embedding
RoPE (Su al., 2024)
Activation Function
Squared ReLU (So et al., 2021)
Normalization
RMSNorm (Zhang & Sennrich, 2019)
Bias Parameters
None
G.2 TOPK COMPRESSION DETAILS
TopK compressor requires transmitting both the selected values and their corresponding indices to reconstruct the original tensors. At high compression levels, this introduces significant commu- nication overhead, especially in compositional schemes such as TopK combined with the Natural compressor, where the cost of transmitting indices can even exceed that of the quantized values. To illustrate this effect, embedding layer and the classification head, each of size 50, 304 x 768. Representing an index for any element in these matrices requires loga (50,304 768) 26 bits. We use this calculation when visualizing communication costs.
analyze the largest parameter matrices in the NanoGPT model: the token
G.3 LEARNING RATE ABLATION
To ensure a fair and robust comparison, we perform a learning rate hyperparameter sweep for each compression configuration, detailed in Figure 3. For every method, the search space is initialized
66

---

Published as a conference paper at ICLR 2026
Table 5: Optimizer configuration.
Hyperparameter
Value
Sequence Length
1024
Batch Size
256
Optimizer
EF21-Muon
Weight Decay
Hidden Layer Norm
Spectral norm
Hidden Layer Scale
50
Newton- Schulz Iterations
Embedding and Head Layers Norm Embedding and Head Layers Scale
& norm
3000
Initial Learning Rate
For non-compressed: 3.6 x 10-4 Constant followed by linear decreasing
Learning Rate Schedule
Learning Rate Constant Phase Length Momentum
40% of tokens
0.9
at the optimal learning rate of the uncompressed baseline (taken from the Gluon repository (Riabinin et al., 2025a)) and spans downward by up to an order of magnitude. We consistently observe that more aggressive compression schemes require smaller learning rate for stable convergence.
This tuning protocol applied uniformly across all experiments for models trained with 2.5B (Sec- tion G.5) and 5B token budgets.
G.4 COMPRESSION LEVEL ABLATION
This section presents ablation study on the compression ratio, governed by the parameter K. Fig- ures and illustrate the convergence curves for various compression configurations, each trained with its optimal learning rate (see Section G.3). Figure 6 summarizes the final loss function of K.
Our results show that for TopK and RandK compressors, aggressive compression ratio of K 5% quite severely impairs convergence (see Figure 6), while configurations with K 2 10% achieve satisfactory loss reduction. When these compressors are composed with the Natural compressor, convergence degradation is more pronounced for K = 10% than for the less aggressive K 15% setup.
We also examine a more challenging loss threshold of 3.28 (Figure 7). The communication cost improvement this threshold even more pronounced than for 3.31 (Figure 1), but this comes at cost: only subset of compressors can reach the threshold within the 5B token budget.
G.5 2.5B TOKENS EXPERIMENT
In Section 5, we report runs with 5B token budget (> 40x model size). Testing convergence over large number of tokens important, as the limitations compressors relative to the baseline become more pronounced after many steps. At the same time, evaluating compressed runs with smaller token budget useful for cases with limited resources. We provide a learning rate ablation in Figure 3, a summarized comparison in Figure 6, and convergence trajectories for the 2.5B-token setup in Figures and 9.
G.6 MEDIUMGPT EXPERIMENT
To assess whether the patterns observed on NanoGPT scale to larger models, we conduct experi- ments Medi umGPT (335M parameters) (Karpathy, 2023) with 2.5B token budget. The model configuration is provided Table 4. We compare the uncompressed baseline to -uon with the Natural compressor and evaluate convergence in terms of both tokens and bytes communicated.
67

---

Published as a conference paper at ICLR 2026
3.40
3.58
3.35
336
-+- 2.50
3.42
2.50
3.56
5.34
3.54
3.32
3.38
3.52
330
3.32
3:36
3.50
328
3.30
3.34
525
3.32
324
3.46
0.34
1oc
0.4
1.44
1.8x
Initial Learing Rate
Learing Rate
Inital Learing Rate
Iibill Learing Rate
Top20%
Top1 5%6
ToplON
Top5%
58
3.50
5B
3.95
58
2.58
258
2.58
258
3.90
Riad Text Loe
3.35
3.45
3.05
3:34
3.44
3.80
3.42
3:32
3.75
3.30
525
3.70
525
3.36
3.65
0.37
10x
017
Irtial Learing Rate
Intial Learing Rate
Inttal Learing Rate
Iritial Learing Rate
Rank] Natural
Rank10% Natural
Top255% Natural
Top10% Natural
3
3,475
3.54
3.42
3.52
3.40
3.425
3.8
3.48
334
3.36
3.40
3,400
3.46
332
3.38
3.350
3.44
X30
3,325
328
3.40
0.28
0,251
6.1
1.0x
0
1.00
Learing
Learing
Learing
Learing
Figure 3: Learning rate ablation. The grid spans from the optimal learning rate of the non- compressed baseline, 3.6 x 10-4 (denoted as 1.0x), down to 0.1x. Red curves correspond to experiments processing 5B tokens (Section 5), while blue curves correspond to 2.5B tokens (Sec- tion G.5).
3,9
3.9
ID
Top1 0%
Top10%
Top10%
Top10% Natural
Top1
Top15%
Top15% Natural
Top15% + Natural
Top20%
Top20%
Loss Threshold 3.33
Test Loss
3.3
3.2
3.2
0.08
1.08
2.08
3.08
4,08
5.0B
0K
2K
8K 10K
12K
14K
16K
Tokens
Communication Cost
Figure 4: Left: Test loss ys. # of tokens processed. Right: Test loss vs. # of bytes sent to the server from each worker normalized by model size reach test loss 3.33. TopX% TopK compressor with sparsification level X%; ID no compression.
68

---

Published as a conference paper at ICLR 2026
3.9
ID
Rank5%
Rark5%
Rank10%
Rank10%
Rank10% Natural
Rank10% Natural
Rank15%
Rank15%
Rank15% Natural
Rank15% Natural
Loss Threshold 3.33
Test Loss
Test
3.4
3.3
3.3
3.2
3.2
0.08
1.08
2.08
3,08
4,08
5.0B
0K
2K
4K
8K 10K 12K
14K
16K
Tokens
Communication Cost
Figure 5; Left: Test loss vs. # of tokens processed. Right: Test loss vs. # of bytes sent to the server from each worker normalized by model size reach test loss 3.33. RankX% Rank compressor with sparsification level X%; ID = no compression.
3.6
3.7
5B
3.5
3.5
Loss
3.4
3.5
3.4
3.4
3.3
3.3
3.2
3,2
3.2
Rank5%
Rank10%
Rank15%
Bank100%
Top5%
Top10%
Top15%
To020%
Tep100%
Figure 6: Final test loss compression parameter K. Results are shown after processing 5B tokens (red) and 2.5B tokens (blue) for RankK (left) and TopK (right) compressors. corresponds to the non-compressed baseline. In the TopK plot, the 2.5B setup outperforms 5B due to differences in scheduler behavior, as the runs execute a different number of steps.
100%
We adopt the learning rate obtained from the sweep described in Section G.3 and use the same optimization and training setup as in the NanoGPT experiments. The LMO step scaling mechanism (Pethick et al., 2025b) is applied ensure adaptivity across weight matrices of varying sizes.
The resulting convergence curves are shown in Figure 10. We observe qualitatively similar behavior the NanoGPT setting: Natural compression achieves close-to -o-beline loss while substantially reducing communication cost.
G.7 BIDIRECTIONAL COMPRESSION
To complement the unidirectional compression experiments, we evaluate EF21-Muon in a fully bidi- rectional setup in which both server-to-worker and worker-to-serven communication compressed. We apply the Natural compressor in both directions. The training task matches the setup in Sec- tion G.5 (NanoGPT with 2.5B token budget).
69

---

Published as a conference paper at ICLR 2026
3.9
ID
Natural
Natural
Rank15%
Rank15%.
Rank15% Natural
Rank15% Natural
Top20%
Top20%
Loss Threshold 3.28
Test
Test
3.4
3.3
3.3
3.2
3.2
0.08
1.08
2,08
3,08
4.08
5.08
0K
2K
4K
6K
8K
10K 12K
14K
16K
Tokens
Communication Cost
Figure 7: Left: Test loss vs. # of tokens processed. Right: Test loss vs. # of bytes sent to the from each worker normalized by model size to reach test loss 3.28. RanpX%%%o R/nkK/TopK compressor with sparsification level X%; ID = no compression.
3,9
ID
Top10%
Top10%
Top10% Natural
Tep10% Natural
Top15%
Top15%
Top15%
Top15%
T0p20%
Top20%
Lass Threshold 3.38
Test Loss
3.5
3.4
3.4
3.3
3.3
0.08
0.58
1.08
1.58
2.08
2.58
0K
1K
2K
3K
4K
5K
6K
7K
8K
Tokens
Communication
Figure 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.38. TopX% TopK compressor with sparsification level X%; ID no compression. "+ Natural" corresponds to applying Natural compression after TopK compressor.
We follow the same hyperparameter selection protocol described in Section G.3. The corresponding learning rate sweep is shown in Figure 11. After tuning, we find that EF21-Muon remains effec- tive in this more challenging bidirectional configuration, improving communication efficiency by approximately 2x relative to the uncompressed baseline while achieving comparable convergence, as shown in Figure 12.
G.8 LIMITATIONS
Reporting results for all compressors the same token budget (for instance, 5B) and then measuring the prefix needed to reach a given loss threshold may not be fully consistent, as results can be affected by the scheduler. To mitigate this, number of tokens are processed beyond the constant learning rate phase. Additionally, tuning the initial learning rate can help stabilize the results.
relatively strong loss threshold that ensures significant
70

---

Published as a conference paper at ICLR 2026
3.9
Rark5%
Rank10%
Rank10% Natural
Rank15%
Rank15% Natural
Loss Threshold 3.38
Test Loss
ID
Rank5%
3.4
Rank10%
Rank10% + Natural
3.3
3.3
Rank15%
Rank15% Natural
3.2
3.2
0.08
0.58
1.08
1.58
2.08
2.58
0K
1K
2K
4K 5K 6K
8K
Tokens
Communication Cost
Figure 9: Left: Test loss vs. # of tokens processed. Right: Test loss vs. # of bytes sent to the server from each worker normalized by model size reach test loss 3.38. RankX% Rank compressor with sparsification level X%; ID no compression. "+ Natural" corresponds to applying Natural compression after RankK compressor.
3.9
3.9
MediumGPT
MediumGPT Natural MediumGPT Loss Threshold 3.15
3.8
Natural MediumGFT
3.8
3.7
3.7
3.6
3.6
3.5
Test
3.3
Tokens
Communication Cost
Figure 10: Left: Test loss vs. # of tokens processed. Right: Test loss vs. # of bytes sent to the server from each worker normalized by model size reach test loss 3.15. no compression.
Note LLM Usage. Large Language Models were used to assist in polishing the writing of the manuscript. LLM assistance did not contribute to the scientific content of the paper.
71

---

Published as a conference paper at ICLR 2026
3.40
Baseline
--- Bidirectional Natural Compression
3.39
Final Test Loss
3.34
3.33
3.32
0.25
10k
Initial Learing Rate
Figure 11: Learning rate ablation for the bidirectional setup. The grid spans from the optimal learning rate of the non-compressed baseline, 3.6 x 10-4 (denoted as 1.0x), down to 0.1
3,9
3.9
ID
Natural Bidirectional
Natural Bi didrectional Threshold
Test Loss
3.2
3.2
0.08
0.58
1.08
1.58
2.08
2.58
0K
2K
4K
6K
10K
Tokens
Communication Cost
Figure 12: Left: Test loss vs. # of tokens processed. Right: Test loss vs. # of bytes sent to the server from each worker normalized by model size reach test loss 3.31. IDD compression. Both s2w and w2s directions are compressed using the Natural compressor.
72


Reviews:
[{"review_id": "cD5gMNWCeE", "reviewer": "Reviewer_y1YB", "summary": "This paper introduces a distributed optimizer that brings error‑feedback (EF) to layer‑wise linear minimization oracle (LMO) methods such as Muon/Scion/Gluon, extending EF beyond the Euclidean setting and enabling bidirectional compression with support for stochastic gradients and momentum.", "questions": "Results are on NanoGPT‑124M, but how would the conclusions change for billion‑parameter LLMs and larger clusters?\n\nThe main results and plots focus on w2s, but how is the total communication cost including s2w, and what are the gains from s2w compression?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper establishes non‑Euclidean EF convergence for LMO‑based methods, and shows how rates recover Euclidean EF21 results when specialized, which is a nontrivial generalization.\n\nOn NanoGPT-124M trained on FineWeb, they report up to 7x reduction in worker‑to‑server communication without loss of accuracy.", "weaknesses": "The novelty is rather incremental since EF21-P introduces bidirectional EF with momentum/stochasticity, Gluon introduces non‑Euclidean layer‑wise LMO analyses, and Dion introduces distributed Muon‑style optimizers.\n\nExperiments show that compression reduces uplink bytes but degrades token efficiency and final loss at a fixed 5B budget, so the claim of \"no accuracy degradation\" is unsupported under equal‑budget comparisons.\n\nThe convergence results assume exact LMO steps, yet the implementation uses inexact Newton–Schulz updates. Recent literature emphasizes that this gap matters, but the paper neither analyzes nor bounds the approximation error.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ornjludw5w", "reviewer": "Reviewer_NH9Y", "summary": "This work studies a distributed non-Euclidean LMO-based optimization method that generalizes and recovers Muon and Scion in the non-compressed regime. Specifically, the authors introduce EF21-Muon, a unified and communication-efficient algorithm that incorporates stochastic gradients, momentum, and bidirectional compression with error feedback, while encompassing several existing compressed methods as special cases. Furthermore, the authors present comprehensive convergence guarantees for multiple settings, including both deterministic and stochastic cases, as well as for non-Euclidean smooth and generalized non-Euclidean smooth functions, under both layer-wise and joint parameter treatments. Experimental results on nanoGPT under various worker compressors show significant communication savings for the proposed algorithm.", "questions": "- Can the authors elaborate a bit on why in the generalized smooth setup the convergence guarantees are established without primal compression? Can these results potentially be extended to that setting as well?\n- Could the authors clarify how the combinations of compressors are applied in the experimental setup?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This work provides rigorous and comprehensive convergence guarantees for a wide range of settings.\n- It supports bidirectional compression for the smooth regime.\n- The proposed algorithm further supports non-Euclidean contractive compressors, thus enhancing generality.\n- The analysis of the algorithms is conducted in non-Euclidean norms.\n- The proposed algorithms achieve the state-of-the-art convergence rates.\n- The paper is clearly structured and easy to follow.\n- The discussion of the contribution of each term in the convergence guarantees is insightful.", "weaknesses": "- The results in the generalized smoothness regime do not include primal compression.\n- The experimental results could additionally include some other baseline algorithms for comparison to better contextualize performance of EF21-Muon.\n- It may also have been beneficial to include wall-clock time experiments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NDDMNgJV5j", "reviewer": "Reviewer_PPbv", "summary": "This paper presents EF21-Muon, a communication-efficient distributed optimizer based on recent LMO optimizers with rigorous convergence guarantees. This method achieves up to 7x communications savings compared to uncompressed baselines.\n\nFurthermore, a few notable insights include:\n1. Bidirectional compression with error feedback\n2. Support for non-Euclidean geometries through arbitrary norm choices\n3. Layer-wise treatment of neural network parameters\n4. Theoretical guarantees under many settings", "questions": "1. How does this method scale when applied to bigger models? How does wall-clock time scale as you increase the model size? Is it a constant factor increase or is it quadratic, log, etc?\n2. How does EF21-Muon compare to other compressed distributed training variants of AdamW or SGD?\n3. The method proposed has a lot of hyperparameters, how sensitive is the training to choices beyond the limited ablations shown?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Significant theoretical contribution, rigorous convergence analysis for compressed, distributed LMO-based methods in non-Euclidean settings.\n2. Addresses a real bottleneck in the distributed training of large models, where Muon is gaining popularity.\n3. Comprehensive theory, a very long appendix with many proofs.", "weaknesses": "1. Only one model size was trained, and it is very small by modern standards (120M).\n2. Only one evaluation (Loss) was used. It's good practice to include a few downstream evaluations just in case for bold claims.\n3. Compression overhead is not clearly addressed, topk requires transmitting indices and not thoroughly analyzed for the distributed setting with varying model architectures. This is especially important in optimization as a general-purpose optimizer should work on different model shapes.\n4. Very heavy and dense notation, difficult to read.\n5. The algorithm itself is not especially novel, error feedback comes from EF21.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tqEvCp7KLo", "reviewer": "Reviewer_VseT", "summary": "The paper proposes EF21-Muon, which aims to be a communication-efficient distributed learning framework for the Muon family of optimizers. A key insight is to reformulate the LMO update using the sharp operator which recasts the update as a normalized steepest descent step. As Muon is steadily gaining popularity, this area of research is particularly timely and well-placed. Additionally, it is important to provide a principles framework for Muon, including in the setting of distributed training which suffers from communication inefficiencies.", "questions": "Please see weaknesses. Not much questions to add.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The topic is particularly timely, and addresses the constraints of distributed training which is not typically encountered in centralized settings. \n\n2. The experiment setup is quite relevant. I believe that NanoGPT on FineWeb is a good choice for optimizer evaluation in the context of model LMs.", "weaknesses": "1. Not much of a weakness, but the authors argue that due to increasing size, all training is distributed. This is true, but most frameworks tend to be parallelism (pipeline, model, etc) rather than the FL-type. \n\n2. Experiments are conducted on a single model/scale (NanoGPT 124M). Demonstrating the effectiveness of EF21-Muon on other domains (e.g., ViT) or at a larger scale would make the claims of general applicability significantly more robust.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)