DeMo: Decoupled Momentum Optimization
Authors: Bowen Peng
Lizhang Chen
Baiyu Su
Jeffrey Quesnelle
Diederik P Kingma
qiang liu
Venue: ICLR 2026 Poster (2026)
Keywords: deep learning
large language models
optimization
training
generative models
pre-training
foundational models
distributed training
Primary area: optimization

Abstract: Scaling neural network training increasingly depends on synchronous data-parallelism, yet full-precision gradient all-reduce imposes a severe communication bottleneck. We propose Decoupled Momentum Optimization, a drop-in replacement for any momentum-based optimizers that significantly reduces the communication bandwidth while maintaining convergence. DeMo (i) decouples local momentum updates, (ii) applies a fast orthonormal transform (e.g., DCT) followed by top-$k$ sparsification, and (iii) reuses the momentum buffer for error feedback via momentum subtraction. This design reduces per-step communication by up to two orders of magnitude with minimal computational overhead.  Experiments on 300M- and 1B-parameter DeMo language models show DeMo transmits up to 85× less data per GPU than AdamW-DDP while achieving comparable loss and accuracy. DeMo is topology-agnostic and enables training across multi-datacenter or Ethernet-based setups.

Published as a conference paper at ICLR 2026
DEMo: DECOUPLED MOMENTUM OPTIMIZATION
Bowen Peng"
Lizhang Chen
Baiyu Su*
Jeffrey Quesnelle'
Diederik Kingma
Qiang Liu
ABSTRACT
Scaling neural network training increasingly depends on synchronous data- parallelism, yet full-precision gradient all-reduce imposes a severe communi- cation bottleneck. We propose Decoupled Momentum Optimization (DeMo), a drop-in replacement for any momentum baased optimizers that significantly re- duces the communication bandwidth while maintaining convergence. DeMo (i) decouples local momentum updates, (ii) applies a fast orthonormal transform (e.g.. DCT) followed by top-k sparsification, and (iii) reuses the momentum buffer as error feedback via momentum subtraction. This design reduces per- step communication by up to two orders of magnitude with minimal computa- tional overhead. Experiments on 300M- and 1B- 1B-parameter DeMo language mod- els show DeMo transmits up 85x less data per GPU than AdamW-DDP while achieving comparable loss and accuracy. DeMo is topology-agnostic and enables training across multi-datacenter Ethernet-based setups. Code is available attps::/a htt///gittuu wtt.:h/wwcooooco////eeeeeee
INTRODUCTION
The proliferation of large-scale foundation models, with parameter counts reaching the billions (Achiam et al.) [2023: Grattafion et al. 2024 Liu et al. (2024)). has been transformative across nu- merous domains. Training these massive models within tractable timeframe necessitates distribut- ing the heavy computational workload across a large number of accelerators (Duan et al. 2024). This is typically managed through hierarchy of parallelism strategies, such as Distributed Data Parallelism (DDP), Pipeline Parallelism (PP), and Fully Sharded Data Parallelism (FSDP) ZZhao
ct al. 2023
Among these, distributed data parallelism remains the most widely used strategy but introduces major communication bottleneck. In standard DDP, each worker synchronizes its locally computed gradients with all others before every optimization step-typically via an All-Reduce operation. The communication cost is directly proportional to the model's size, reaching terabytes per step for state-of-the-art models. This overhead demands expensive, high-bandwidth interconnects (e.g., NVLink, InfiniBand) and geographically co-located clusters Wei et al. (2024), driving up costs and limiting scalability. Alleviating this communication burden is thus critical for enabling more effi- cient and accessible large-scale training Cao et al. (023).
In this paper, we show that the gradient information exchanged during distributed training is highly redundant. Our key insight is that this redundancy can be managed more effectively by communicat- ing compressed momentum updates
transformed space rather than transmitting raw gradients. Prior sparsification approaches directly applied to gradients incur sparse update patterns and of- ten harm performance Lin et al. (201 Stich et al. (2018); Aji and Heafield (2017); in contrast, our blockwise transform and top-k sparsification preserve accuracy by operating in the transformed domain. Moreover, the momentum buffer itself naturally serves as an implicit error accumulator, eliminating the need for explicit arrorffedbbcc mechanisms that introduce substantial GPU memory overhead (Karimireddy et al. 2019). Building on these observations, we propose DeMo (Decoupled Momentum Optimization), general seneral-purpose distributed optimization framework compatible with
"Equal contribution.
University of Texas at Austin
*This work was done in the author's 's persomal capacity before joining Anthropic. Nous Research. Correspondence: (bloc, emozi i @ncusresea rch.

---

Published as a conference paper at ICLR 2026
common momentum- based optimizers, including SGD with momentum (Sutskever al. 2013), Lion (Chen et al. 2023), and Muon (Jordan et al., 2024). DeMo fundamentally rethinks synchro- nization: instead of exchanging dense gradients, each worker communicates compressed version of its local momentum, subtracts the decoded update from its buffer to track uncommunicated in- formation, and then applies the aggregated momentum update the parameters. This simple but powerful design drastically reduces communication while retaining the benefits of momentum-bass optimization.
This approach drastically reduces the required inter-accelerator communication bandwidth, poten- tially by several orders of magnitude. Our primary contributions are:
We identify the momentum term in distributed optimization compressible, information- rich surrogate for raw gradients.
We propose DeMo, communication-efiicient framework that decouples momentum up- dates and leverages the momentum itself as built-in error feedback mechanism.
We demonstrate that our method significantly relaxes the hardware and co-location straints for large-scale training, paving the way for more flexible, geographically dis- tributed, and cost-effective training paradigms.
The remainder of this paper is organized follows. We review related work in Section In Sec- we present the DeMo algorithm and its theoretical guarantees. Section 3 provides empirical results demonstrating DeMo's ' effectiveness, followed by ablation studies. Finally, conclude in Section
tion 2
DEMo: DECOUPLED MOMENTUM OPTIMIZATION
Problem setting.
We consider the standard stochastic optimization problem
min ((X) REgD ( X
where X X is (possibly series of) parameter tensors be optimized, L(X, {) is sample loss function (e.g., cross-entropy), and the expectation is taken over data samples & drawn from distribution D. The optimal value is denoted by J == inf ((X), which we assume finite and bounded from below. We operate in distributed data-parallel setting with / workers, each holding synchronized copies of parameters X. At each step f, worker computes stochastic Phatch 1 VL(X,, En) using per-worker microbatch of size Thatch
gradient G²
2.1 ALGORITHM
We start to build DeMo from the standard DDP optimization pipeline and common momentum- based optimizers, such as SGD with momentum (Sutskever et al. 2013). Signum (Bernstein et al.) 2018), and Muon (Jordan et al.) 2024). Our proposed method introduces three key modifications to reduce communication overhead and improve efficiency: (I) decoupled local momentum updates to be communicated instead of gradients, (2) structured tensor compression, and (3) momentum subtraction as error feedback.
Decoupled Local Momentum Updates. Standard DDP pipelines typically synchronize gradients across workers immediately after every local gradient tensor being computed. contrast, we move the global all reluueduce synchronization of micro-batch gradients Gj, allowing local mo- mentum buffers Mi on each worker to evolve independently:
M BM + +
where B the momentum coefficient. Aggregating gradients or momenta is theoretically equiva- lent due to linearity; however, directly synchronizing dense momentum tensors is communication- intensive. Thus, we propose a structured tensor compression pipeline to significantly reduce this overhead.
2

---

Published as a conference paper at ICLR 2026
Algorithm DeMo: Decoupled Momentum
Require: learning rate 7/, momentum coefficient (3, weight decay A, sparsity budget k; DCT matrices {Pz)
1: Initialization: parameters Xo and global momentum Mo <-
for 1,
3: # on each worker i (O..... N
G; <- vi(
local stochastic gradient update local momentum buffer
5: M/ MM G.
each chunk M, M1 M
7:
1.14, (P))), k) 8: M, Mc M, (QC 10, (P) (P,))
Q 1.0
blockwise transform & sparsify in the E-th chunk
(P,))
in-place residual in momentum buffer
9:
end for
10: send {Q, (Q:1l) server
11: on the parameter server
12: for each chunk (Q² 1 do
13: M < IDCT Σ Q 14 (P,7.
# aggregate sparse updates and reconstruct momentum chunk
14: end for
15: X. f- X,-1 (sggu (M) + -I)
16: broadcast (gn(()) to all workers
17: end for
Structured Tensor Compression. Our compression pipeline comprises three sequential steps: tensor chunking, blockwise linear projection, and top-k sparsification.
Tensor Chanking. Given a momentum tensor M IRnoX--Xaai,, we factorize each dimension as n4 = Cj3; with chunk number c, and size s; to partition M into smaller blocks:
((M) [Bx k [co] x ... (ca-1]}.
where each block B& - 1, Blockwise Linear Projection. After chunking, each block B₁ E ((M) undergoes a separable multilinear transformation:
Qk = Po..,
P.R
(3)
where multiplies block Bk by projection matrices P. along each tensor dimension. For 2, this reduces to simple bilinear form of matrix multiplication Qk = P.BxP.
In practice, we consider freshly at each step (shared across all workers), and the Discrete Cosine Transform (DCT). Our experiments Section 3 demonstrate that projections clearly outperform no projection (P: I). While random projections perform marginally better, DCT is computationally more efficient due to fast implementation with Fast Fourier Transform and precomputed only once before training. Thus, we default to DCT in all experiments apart from ablations.
types of projection bases: random orthonormal matrices, sampled
Top-k Sparsification. Following projection, in order reduce communication, each worker prepares only the top-k largest-magnitude coefficients per block for communication:
Qk Top-k(Que (k.k) k) such that ||@x|lo k.
Thereby the uploading bandwidth is reduced by (II, st)/k: times and for small this can be signif- icant. These sparse blocks are communicated using an All.Gather operation and averaged over all workers to form:
Q = AAvglllGather Qx)-
(5)
After communication, each worker recon- structs the momentum tensor from the aggregated sparse blocks via inverse projection and unchunk- ing by concatenation:
Momentum Reconstruction and Parameter Update.
M 1- (Qx: Po' Pz))),
3

---

Published as a conference paper at ICLR 2026
XK
Tpph
IDC
Figure 1: Chunk-wise Discrete Cosine Transform (DCT), top-k coefficient selection, and inverse DCT reconstruction.
where P denotes inverse projection matrices. To enhance convergence during neural network training, we apply transformation o(-) the reconstructed momentum Mt, chosen according the base optimizer: (MM) M for SGD, ((M) sign MM (elementwise) for Signum, and ((M) = M(M M + ef) 1/2 for Muon. The final parameter update is given by
X = X, nh (M,) + AX;).
Algorithm summarizes this entire procedure.
Momentum Subtraction as Error Feedback. Inspired by error feedback mechanisms used with sparsification techniques, propose updating the decoupled momentum buffer M; via momentum subtraction:
Mi Mj - a B QQ: ..... P(.))
where E 1] is the momentum subtraction coefficient. Unlike traditional error-feedback schemes, which require additional memory storage, our method naturally reuses the momentum buffer as an accumulator of uncommunicated information. This subtraction ensures each iteration communicates novel information, accumulating previously omitted updates over subsequent steps and promoting convergence. The decay of past gradient information is controlled by both and B.
Complexity Analysis. Without chunking, projecting a momentum matrix of size N N requires ( computation and O(N²) storage. By partitioning the matrix into C2 blocks of size (N/C) x (N/C), computational complexity reduces linearly to O(N$/C) and memory usage quadratically to O(N² /C2). Since all momentum tensors share the same set of projection matrices {P.) the additional memory overhead remains constant and negligible compared to the overall memory requirements.
d--1
2.2 THEORETICAL ANALYSIS
To analyze the convergence properties of the DeMo algorithm, we adopt standard assumptions com-
monly used stochastic optimization Bernstein and Newhousel 2024): Défossez et al. (2022): Liu al. (2024b).
Assumption 1 (Variance). Let samples & ~ DD be i.i.d., and define the expected gradient as VX(X) E[VL(X, The stochastic gradient estimator at each worker is given by
1
G'XX)
vL((XEn),
Tbbttc
and satisfies the bounded variance condition:
2
E [x) - VC
WE [N), X XEX,
Tbbach
where ²² 0 is a finite constant and Tibatch is the mini-batch size. Assumption (L-Smoothness). The objective function L(X) is differentiable and L-smooth; that is, for all X, Y X.
|V - VL(X) S L L X
Equivalently, for all X, Y X.
L(Y) < ((X) + (VC()), Y X) + 2
Y X||
4

---

Published as a conference paper at ICLR 2026
Assumption (Bounded Gradient). For any X % and & D, the stochastic gradient satisfies [VC(X: S R,
for some constant > O.
We present the convergence guarantee for the proposed Decoupled Momentum Optimization (DeMo) algorithm under the above assumptions.
Theorem (Convergence of DeMo). Under Assumptions Algorithm satisfies:
and the sequence {Xx {x from
E [[ E(C(X.) -(xx)
2LD
T"
v N Abbtchh
+ RVD
V
M
B 2D
2T N
1+
1 B 8/
k
M
with the choice of step size 17 O (++) and momentum
3 0
the convergence rate is:
+
EXPERIMENTAL RESULTS
We evaluate our proposed Decoupled Momentum Optimization (DeMo) algorithm the context of large-scale language model pretraining, where communication efficiency is of uttermost concern. All experiments are conducted using the OLMo (Groeneveld et al. (024) framework, reproducible and scalable platform for training open-weight language models. The complete implementation of DeMo, along with all configuration files necessary reproduce our results, available at
EXPERIMENTAL SETUP
Models and Baselines. We evaluate DeMo Transformer-bant decoder-only language mod- els at two scales. OLMo-300M contains 320 million non-embedding parameters, and OLMo-1B contains 1.18 billion non-embedding parameters Full model specifications are provided in the Ap- pendix.
We compare DeMo against the standard AdamW optimizer using default hyperparameters recom- mended by OLMo: 31 0.9, and weight decay À = 0.1. The learning rate schedules use linear warmup followed by cosine decay, scaled proportionally to the total training steps. For the runs in Figure 3.1 (and consequently Figure we adopt the best recommended learning rates from Groen- eveld al. (2024), since tuning over the full training run would incur infeasible computational costs. We tune the AdamW second-order momentum 33 over the set [0.95, 0.98, 0.99, 0.999) using reduced training budget of IB tokens. Results with extensively searched hyperparameters for all optimizers are provided in Figure 5 For all DeMo experiments presented in the main body of the paper, we default to use chunk size of s = 64 and therefore making 64 x 64 chunks for matrix tensors in the transformers and 64 sized chunks for vector tensors like layernorms. We found set- ting a larger coefficient of 3 = 0.999 significantly helps with convergence comparing to standard value of 0.9 when the momentum subtraction is turned on and therefore default to this value over all experiments. We left the momentum coefficient to be a tunable parameter from the set of (0.2,0.5, 1.0).
Datasets. All models are pretrained on the Dolma v1.5 corpus. To study optimizer behavior near convergence, train both OLMo-300M and OLMo-IB 100 billion aakens-significantl be- yond the Chinchilla token budget of 20 tokens per parameter. Due to computational constraints, all ablation studies are conducted using the 300M model trained under the 20 tokens tokens-per- paee--pe--ppaaaa per-parameter rule.
5

---

Published as a conference paper at ICLR 2026
300M Model
1B Model
AdamW
AdamW
Training Loss
3.6
Training Loss
3.4
3.4
3.2
2.8
2.6
0
5000
10000
15000
20000
5000
10000
15000
20000
Update Steps
Update Steps
k=1 k=
k=
k=
16
32
Figure 2: Training loss comparison between IB and 300M parameter models across different top-k sparsity levels. The left panel shows the 300M model performance while the right panel shows the 1B model performance. Both models are trained with varying sparsification top (k=1, 2, 4, 8, 16, 32) compared against the AdamW baseline. The results demonstrate the effectiveness of top-k sparsification across different model sizes.
Implementation. DeMo integration required minimal changes to the codebase. We implemented a DeMo optimizer class and disabled default gradient synchronization in Py Torch Distributed Data Parallelism (Li al. 2020). All experiments are conducted on 64 NVIDIA H100 GPUs with global batch size of 2048 and sequence length of 2048 tokens. We use four gradient accumulation steps, yielding an effective per-GPU batch size of 8.
3.2 PRETRAINING PERFORMANCE
Convergence Behavior: FigureB.T]ssovee training loss curves comparing DeMo with various spar- sification levels & to the AdamW optimizer both the OLMo-300M and OLMo- IB models. As the figure illustrates, a sparsification of just k sufficient to achieve better training performance than AdamW. Increasing & further provides marginal gains at greater communication cost.
Compute Efficiency. DeMo provides substantial reduction in communication. For instance, with the tensor chunk size of 64 x 64, selecting the top-k values reduces upload bandwidth at each step by factor of roughly 4096/k: kk. Figure 4 plots final model perplexity against the overall communication cost MB per step.
OLMo-300M
OLMo-1B
Optimizer / k
Hella acc.n
ARC
PIQA
Tx J MB/step
Hella
ARC
PIQA
Tx 4 MB/step
acc
acc_n
acc.n
acc
acc.n
DcMo k=
0.37
0.46
0.67
29.9
0.48
0.55
0.70
110.32
DeMo k=16
0.38
0.50
0.67
14.9
0.47
0.53
0.70
55.16
DcMo k=8
0.38
0.47
0.67
7.49
0.47
0.52
0.69
27.58
DeMo k-1
0.37
0.47
0.67
3.74
0.45
0.52
0.70
13.79
DeMo 2
0.36
0.46
0.65
1.87
0.44
0.51
0.69
6.89
DeMo k=1
0.35
0.45
0.65
0.93
0.41
0.52
0.69
3.44
AdamW-DDP
0.35
0.46
0.65
636.9
0.43
0.51
0.68
2416.6
Table 1: Zero-shot downstream accuracy and per-GPU communication volume (MB per training step) after 100B-token pretraining. Higher better for accuracy metrics; lower better for commu- nication.

---

Published as a conference paper at ICLR 2026
3.3 DOWNSTREAM EVALUATION
We assessed the pretrained checkpoints on three widely used zero-shot benchmarks: HellaSwag, ARC-Easy, and PIQA, and report normalized accuracy (or accuracy for ARC-Easy). Table shows that DeMo matches or exceeds the Adam W-DDP baseline across all tasks while reducing per-GPU communication by two to three orders of magnitude.
For the 100M-paramete 637 MB) with no loss in accuracy; smaller values offer additional savings with a minor accuracy trade-off. At the IB scale, the trend persists: DeMo with A: 16 attains higher HellaSwag and PIQA scores than A while transmitting only 55 MB per step an efficiency gain of times. These results confirm that our proposed decoupled momentum strategy preserves model quality even under aggressive communication budgets and scales favorably with model size.
model, using A = 8 already cuts data transfer by 85 times (7.5 MB vs.
3.4 ABLATION STUDIES
To better understand the impact of specific design choices in DeMo, we conducted comprehensive ablation studies:
Impact of Momentum Subtraction. An essential feature of DeMo is momentum subtraction. In this experiment, we systematically varied a from (no subtraction) to (full subtraction of communicated values). The results are shown in Figure [3.]] Clearly. subtraction (a 0) is detrimental because, with a fixed basis, the top-k elements evolve slowly and the same elements repeatedly selected and communicated across steps, resulting similar consecutive updates and degraded performance. Given that the momentum buffer also serves as an eтrоr accumulator, subtracting previously communicated values necessary. However, compared full subtraction, we found that using smaller value (a = 0.2) to gradually evolve the top-k elements and partially decay communicated values over time further improves performance.
Identity DCT
Random DCT
4.2
DCT
4.2
DCT k=
4.0
4.0
DCT k=
kd
Training Loss
Rd 8
Id k 16
3.8
Rd 1 16
Training
Id k 32
Rd k 32
3.4
3.4
3.2
3.2
2000
6000
12000
8000 10000
12000
14000
Update Steps
Update Steps
Figure 3: Training loss with different linear transformations. The left plot compares the DCT against identity mapping, while the right plot compares the DCT with random orthonormal projection matrices. Each curve corresponds to different sparsity level & {8, 32).
Choice of Linear Projections. We explored three natural choices for the linear transformations T. The baseline is the identity mapping P. = I, i.e., no transformation. We compared this with the Discrete Cosine Transform (DCT)-based transformation, where P: corresponds to the DCT basis of dimension 64, and a random projection matrix, where each element of P; sampled from (0, T) and orthonormalized using the number on each worker, ensuring a unique random projection for every step but consistency across workers. Since all P. are orthonormal, we have P.- P.T. We considered three sparsification levels: = 8, 16, 32. The results are plotted in Figure
process. We set the random seed equal to the step
As illustrated in the figure, the DCT-based transformation clearly improves performance the identity mapping. This supports the hypothesis from related studies that naively updating param- eters using sparse patterns can degrade performance, while applying transformation enables pa-

---

Published as a conference paper at ICLR 2026
--0.55
*
0.1
B 0.98
14-
0.2
0.99
10-
Treining Loss
6 1,0
Taiiing Los
B 0.999
0.995
Adarrii 2C1 30M
26-
Adavtii 2CP 10
14-
17 10r
10r
307
10000
Data Transmissoo per Step (MBlstep)
Update
Figure 4: Left: Validation perplexity of OLMo-300M and OLMo- versus data transmitted per step (log scale: MB/step). Each DeMo point annotated with the sparsification level k. Middle: Training loss of OLMo-300M with DeMo for momentun-subrracinn magnitude E (0.0, 0.5, 1.0]. Right: Training loss of DeMo with top-k=32, C=64, and a=0.2 under varying momentum decay B.
rameter updates as linear combinations of many non-sparse vectors. Thereby the parameters are updated more uniformly, especially important when sparsity k is low. Although random projection arguably the most intuitive choice, as it continuously rotates and changes the momentum subspace perspective, found that a fixed DCT basis yields comparable performance while eliminating the need to recompute the basis at each step.
3.5
ADDITIONAL EXPERIMENTS
This section presents additional ablations and comparisons of DeMo with communication effficient optimizers, extensively tuned Adam and Muon-DDP baselines, and momentum decay abla- tions. All experiments use 300M parameter model trained with 20 tokens per parameter.
Ablations momentum coefficient B. We ablate the momentum coefficient B {0.95, 0.98, 0.99, 0.995, 0.999). Training curves are shown in Figure Performance remains stable across this range, with larger values (0.995, 0.999) outperforming smaller ones. The best results are achieved at 3 = 0.995.
To isolate optimizer effects, we reduce the token budget and perform extensive hyperparameter sweeps for AdamW-DD2 We search learning 6 x 10-4 10-4 4, x 10-4 4, 2 the recommended momentum value 0.95 in Liu et al. (2025) and tune the learning rate over the same grid. Figure 5 shows that Muon-DDP has better performance than Adam W-DDP. DeMo underperforms Adam W equal update counts but provides substantially lower communication cost.
Comparisons with Adam W-DDP and Muon-DDP.
rates (1.5 x 10 10-3 1.2 x 10 10-3 3,88 10
10 10-4) and 32 in
(0.95, 0.98, 0.99, 0.995, 0.999). For Muon-DDB
Comparisons with DiL.ooo (Douillard et al., 2023). We additionally compare against DiLoCo, communication-effiiaae method that alternates local updates with periodic global synchroniza- tion. Following the authors' recommended settings, we tune the outer-loop learning rate over {1.0, 0.7, 0.5, 0.3, 0.1} } for each run, and vary the global communication frequency to target differ- ent compression ratios. Figure reports final validation perplexity versus data-receiving reduction when using eight workers. For example, for DeMo with 64 64 chunk size and top-k 32, the data-transmission and data-receiving reductions are 128 and 16, respectively (ignoring the marginal contribution of layer norms). At comparable compression levels, DeMo consistently yields lower perplexity than DiLoCo. Training-loss curves for DiLoCo with communication frequency 10 are included Figure where DeMo again shows superior optimization performance.
Comparisons with PowerSGD (Vogels et al., 2019). We also compare DeMo with PowerSGD, which reduces communication by applying a low rank projection during gradient aggregation. We the official PyTorch implementation of the PowerSGD DDP communication hook (Paszke et al. 2019). Following the recommended procedure, we sweep the communication rank over the set (2, 4. 8, 16, 32, 64) and observe that ranks below 16 lead to substantial degradation. We therefore
8

---

Published as a conference paper at ICLR 2026
4.2
Power5GD
26
DiLeCo
500
DiLoCo
DeMo
DeMo
Training Loss
3.8
Muon
Final Perplexity
24
AdamW
600
3.6
22-
3.4
10
20
20-
64
3.0
19
2000
4000
6000
8000
10000
12000 14000
10a
Update Steps
Data Receiving Reduction Ratio
Figure 5: Training loss curves of four optimizers with tuned hyperparameters: AdamW, Muon, DeMo, DiLoCo and Pow- erSGD.
Figure 6: Data receiving compression ra- tio final validation perplexity of DeMo and DiLoCo trained with varying top- k/communication interval.
select rank the preferred setting, which provides an effective compression ratio of approx- imately 26 times for the compressed steps. Warm up before compression is also important. As suggested the implementation, we tune the number of warm up steps and find that using 1000 steps avoids the instabilities and divergence that often occur with only 500 steps. We search the learning rate over the same grid used for AdamW and find that the best value for PowerSGD 6 x 10-4, which is smaller than the optimal learning rate for AdamW. Because PowerSGD is sentially compressed SGD, which performs poorly for language model training in its basic form, adopt Adam style update rule using the compressed gradients. The final performance of PowerSGD is similar to Dio but remains slightly below that of DeMo.
RELATED WORKS
Training large foundation models presents substantial communication bottlenecks, as synchronizing gradients across numerous accelerators can dominate training time (Shoeybi et al.) 2019 Huang et al. 2019 Narayanan et al. 2021). To alleviate this, various gradient compression techniques have been developed (Aji and Heañeld 2017 Wen et al., 2017). Prominent among these are sparsifi- cation methods, which transmits only the gradients with the largest magnitudes (Linn et al. 2018h 8b Stich al. |0088] Alistarh al. 2018 Zhao al. 2024). Another widely used approach is gradient quantization, which reduces the numerical precision of gradients (Alistarh et al.) 2017 Sun et al.) 2019 Zhao et al. 2024; Gorbunov et al., 2021 Peng et al. 2023). To compensate for information loss introduced by biased compressors, error feedback mechanisms, such as EF-SGD Karimireddy ctt al. (2019), locally accumulate the compression error and add to the gradient in the subsequent iteration (Seide et al., 2014; Karimireddy et al., 2019). Our work also leverages top-k sparsifi- cation of momentum communication and leverages the error accumulation on momentum itself, thereby reducing one copy of memory for the accumulator. Recent work further explores optimizer- aware communication reductions (e.g., communicatinn-efiiieenn variants built around Lion-style
updates) (Chen et al., 2024a Chen 2025 Chen al. 2024b).
Another line of research aims to reduce communication frequency through decoupled optimiza- tion strategies. Methods like Federated Averaging (FedAvg) (McMahan et al. (017)) and Local SGD (Stcch, 2018; Lin et al. 20188) allow workers to perform multiple local updates before syn- chronizing model parameters, thereby reducing the number of communication rounds. However, these approaches can face challenges such as client drift, especially with non-iid data, which may impact convergence speed and final model quality (Karimireddy et al., 2020 Zhao et al. 2018). More recently, DiLoCo has emerged guage models with significantly reduced communication (Douillard et al. 2023). While effective, the complex interplay of infrequent synchronization and local optimizer dynamics in such systems can sometimes lead to less predictable optimization trajectories (cspecially for modern adaptive/- cautious variants) (Nguyen al. 2025
empirically successful technique for training large lan-
Liang al. 2024a Chen al. 2025b); work instead

---

Published as a conference paper at ICLR 2026
focuses maintaining more frequent, albeit significantly compressed, synchronization of a critical optimizer state.
A third strategy seeks to reduce memory and communication via low-rank updates. moving beyond unstructured compression. This was popularized by Low-Rank Adaptation (LoRA) for parameter- efficient fine-tuning, which freezes pre-trained weights and learns a low-rank decomposition of the weight update matrix (Hu et al., 2022). More recently, this principle has been applied directly to the optimization process enable training from scratch. For instance, GaLore and its variants project the full gradient onto low-rank subspace before it is passed to the optimizer, thereby sig- nificantly reducing gradient memory overhead (Zhao et al.] 2024 Hao et al. 2024). Related sub- space/factorized optimization methods provide similar memory savings by restricting updates to low-dimensional structure (Liang al., 2024b, Nguyen et al. 2024). distributed settings, related techniques communicate low-rank factors of the gradient or model update instead of the full matrix, directly reducing communication volume. Notable prior works include (Park and Klabjan [024] Zhao et al., 2025 Chen et al., 2025a). The top-k sparsification of DeMo operates in a fixed trans- formed space, thereby also making the communicated update matrix of each worker to be rank-k.
5 DISCUSSIONS AND LIMITATIONS
Our experimental findings confirm that DeMo consistently achieves competitive improved con- vergence behavior compared to AdamW. Additionally, DeMo demonstrates reduced communication overhead and favorable downstream generalization. Our ablation studies further highlight the criti- cal role of momentum sparsification with DCT, momentum subtraction and chunk- baunk-based projection techniques in balancing optimization performance and computational efficiency.
In our analysis, we primarily focused on the upload bandwidth per step. However, it is important to note that the download bandwidth scales with the number of workers. This limitation is not unique to our method but is intrinsic to all top-k sparsification- -based approaches, including those proposed in [in et al. (b018b); Stich et al. (2018). DeMo designed primarily for optimization across a small number of geographically distributed compute centers, enabling communication between op- timization steps over the Internet rather than relying on specialized high-speed interconnects such as InfiniBand. This relaxes the need for dedicated long-distance networking infrastructure. Within each data center, standard DDP can still be employed; each center can be treated as a "large worker", and DeMo can then be used to efficiently coordinate communication between these centers.
CONCLUSION
In conclusion, we have shown that our proposed DeMo optimization algorithm can act as a drop-in replacement AdamW when training LLMs, with no noticeable slowdown in convergence while reducing communication requirements by several orders of magnitude. The signum variant of DeMo more memory efficient than AdamW and has negligible compute overhead if we use small pre- computed DCT transition matrices. Finally. the LLMs pre-trained with DeMo have equivalent or better scores on multiple standard benchmarks compared to their equivalents trained with AdamW.
REFERENCES
J. Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
A. F. Aji and K. Heañeld. Sparse communication for distributed gradient descent. arXiv preprint 174450100..0002,, 2017.
D. Alistarh, D. Grubic, J. Li, R. Tomioka, and M. Vojnovic. Qsgd: Communication-fffiienn sgd via gradient quantization and encoding. Advances in neural information processing systems, 30, 2017.
D. Alistarh, T. Hoefler, M. Johansson, N. Konstantinov, S. Khirirat, and C. Renggli. The convergence of sparsified gradient methods. Advances in Neural Information Processing Systems, 31, 2018.
10

---

Published as a conference paper at ICLR 2026
J. Bernstein and L. Newhouse. Old optimizer, norm: An anthology. CoRR, abs/2409.20325,, 2024.
J. Bernstein, Y-X. Wang, K. Azizzadenesheli, and A. Anandkumar. signsgd: compressed optimi- sation for non-convex problems. In International Conference Machine Learning, 2018. URL Stttps://aph semant icscho ar. arg/Corpust 7763588
X. Cao, T. Başar, S. Diggavi, Y. C. Eldar, K. B. Letaief, H. V. Poor, and J. Zhang. Communication- efficient distributed learning; An overview. IEEE journal on selected areas in communications, 41(4):851-87,, 2023.
C. Chen, Y. He, P. Li, W. Jia, and K. Yuan. Greedy low-rank gradient compression for distributed learning with convergence guarantees. arXiv preprint arXiv:2507008784, 2025a. URL https:
Preprint.
I. Chen. Demystifying LION: Hamiltonian perspective. PhD) thesis, 2025.
L. Chen, B. Liu, K. Liang, and Q. Liu. Lion secretly solves constrained optimization: As Iyaрunov predicts. International Conference on Learning Representations, Spotlight, 2024a.
L. Chen, B. Liu, L. Wu, K. Liang, J. Zhu, C. Liang, R. Krishnamoorthi, and Q. Liu. Communication efficient distributed training with distributed lion. Advances in Neural Information Processing Systems, 37:18388 11118888-181155 2024b.
I. Chen, Li, K. Liang, B. Su, C. Xie, N. W. Pierse, C. Liang, N. Lao, and Q. Liu. Cautious weight decay. arXiv preprint arXiv:2510. 12402, 2025b.
X. Chen, C. Liang, D. Huang, E. Real, K. Wang, H. Pham, X. Dong, T. Luong, C.-. Hsieh, Y. Lu, et al. Symbolic discovery of optimization algorithms. Advances in neural information processing systems, 36:49205-49233 2023.
A. Défossez, L. Bottou,
Bach, and N. Usunier. A simple convergence proof of adam and adagrad. Trans. Mach. Learn. Res., 2022, 2022.
A. Douillard, Q Feng, A. A. Rusu, R. Chhaparia, Y. Donchev, A. Kuncoro, M. Ranzato, A. Szlam, J. Shen. Diloco: Distributed cowcommuuicattion training of language models, arXiv preprint arXiv:2311.08100 2023.
J. Duan, S. Zhang, Z. Wang, L. Jiang, W. Qu, Q. Hu, G. Wang, Q. Weng, H. Yan, X. Zhang, et al. Efficient training of large language models on distributed infrastructures: a survey. arXiv preprint arXiv:2407.20016 2024.
E. Gorbunov, K. Burlachenko, Z. Li. and P. Richtárik. Marina: Faster non-convex distributed learn- ing with compression. Proceedings of the 38th International Conference on Machine Learn- ing (ICML 2021), pages 4059 4070, 2021. URL htt.pps://prrceeeiiggg..... ptr.ss.sss......////// http:/2aaaaoouuuooolo.o. pdf
press/
PMLR 139.
A. Grattafiori, A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, A. Letman, A. Mathur, A. Schelten, A. Vaughan, et al. The Ilama herd of models. arXiv preprint arXiv:2407.21783, 2024.
D. Groeneveld, I. Beltagy, P. Walsh, A. Bhagia, R. Kinney, O. Tafjord, A. Jha, H. Ivison, I. Magnus- son, Y. Wang. S. Arora, D. Atkinson, R. Authur, K. Chandu, A. Cohan, J. Dumas, Y. Elazar, Y. Gu, J. Hessel, T. Khot, W. Merrill, J. D. Morrison, N. Muennighoff, A. Naik, C. Nam, M. E. Peters, Pyatkin, A. Ravichander, D. Schwenk, Shah, W. Smith, E. Strubell, N. Subramani, M. Wortsman, P. Dasigi, N. Lambert, K. Richardson, L. Zettlemoyer, J. Dodge, K. Lo. L. Sol- daini, N. A. Smith, and H. Hajishirzi. Olmo: Accelcrating the science of language models. arXiv preprint, 2024. URL https::/// / http:/ppppppp.ii.......m.m. coemmnntiscchoe.s. org http://orpssslaa 267355485
Y. Hao, Y. Cao, and L. Mou. Flora: Low-rank adapters are secretly gradient compressors. arXiv preprint arXiv:2402.03293 2024.
E. J. Hu, Y. Shen. P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang. L. Wang, W. Chen, et al. Lora: Low-rank adaptation of large language models. ICLR, 1(2):3, 2022.

---

Published as a conference paper at ICLR 2026
Y. Huang, Y. Cheng, A. Bapna, O. Firat, D. Chen, M. Chen, H. Lee, J. Ngiam, Q. Le, Y. Wu, al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
K. Jordan, Y. Jin, V. Boza, Y. Jiacheng, F. Cesista, Newhouse, and Bernstein. Muon: An opti- mizer for hidden layers in neural networks, 2024. Uttptts:: lerjordarr git. hhub.
ho/posts/mmon)
S. P. Karimireddy, Q. Rebjock, S. Stich, and M. Jaggi. Error feedback fixes signsgd and other gradient compression schemes. In International Conference on Machine Learning, pages 3252- 3261. PMLR, 2019.
S. P. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, and A. T. Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International conference on machine learning, pages 5132-5143. PMLR, 2020.
S. Li, Y. Zhao, R. Varma, O. Salpekar, P. Noordhuis, T. Li, A. Paszke, Smith, B. Vaughan, Da- mania, and S. Chintala. Pytorch distributed: experiences accelerating data parallel training. Proc. VLDB Endow., 13(12):3005- 3018, Aug 2020. ISSN 2150-8097.
K. Liang, L. Chen, B. Liu, and Q. Liu. Cautious optimizers: Improving training with one line of code. arXiv preprint arXiv:2411 241.0668, 2024a.
K. Liang, B. Liu, L. Chen, and Q Liu. Memory-efficient Ilm training with online subspace descent. Advances in Neural Information Processing Systems, 2024b.
T. Lin, S. U. Stich, K. K. Patel, and M. Jaggi. Don't use large mini-batches, use local sgd. arXiv preprint arXiv: 100.:0888..1111111 2018a.
Y. Lin, S. Han, H. Mao, Y. Wang, and B. Dally. Deep gradient compression: Reducing the commu- nication bandwidth for distributed training. In International Conference on Learning Represen- tations, 2018b.
A. Liu, B. Feng, B. Xue, B. Wang, B. Wu, C. Lu, C. Zhao, C. Deng, C. Zhang, C. Ruan, et al. Deepseek v3 technical report. arXiv preprint arX1v:2412.19437, 2024a.
B. Liu, L. Wu, L. Chen, K. Liang, J. Zhu, C. Liang, R. Krishnamoorthi, and Q. Liu. Communication efficient distributed training with distributed lion. In Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, 2024b.
Liu, Su, X. Yao, Z. Jiang, G. Lai, Y. Du, Y. Qin, W. Xu, E. Lu, J. Yan, et al. Muon is scalable for Ilm training. arXiv preprint arXiv: 16982, 2025.
B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas. Communication-effice learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273-1282, PMLR, 2017.
D. Narayanan, A. Phanishayee, K. Shi, X. Chen, and M. Zaharia. Memory-efficient pipeline-parallel dnn training. In International Conference on Machine Learning, pages 7937-7947. PMLR, 2021.
S. Nguyen, L. Chen, B. Liu, and Q. Liu. H-fac: Memory-efficicnt optimization with factorized hamiltonian descent. International Conference Artificial Intelligence and Statistics, 2024.
Nguyen, B. Liu, I. Chen, and Q Liu. Improving adaptive moment optimization via preconditioner diagonalization. arXiv preprint arXiv:2502.074488 2025.
H. Park and D. Klabjan. Communication-effiienn federated low-rank update algorithm and its nection to implicit regularization. arXiv preprint 21ti:::2209.133711, 2024.
A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al. Pytorch: An imperative style, perfoorrforeccc deep learning library. Advances in neural information processing systems, 32, 2019.
12

---

Published as a conference paper at ICLR 2026
H. Peng, S. Qin, Y. Yu, J. Wang. H. Wang, and G. Li. Birder: communication--fffce 1-bit adap- tive optimizer for practical distributed dnn training. Advances in Neural Information Processing Systems, 00339999-39552 2023.
Seide, H. Fu, J. Droppo, G. Li, and D. Yu. 1-bit stochastic gradient descent and its application data-parallel distributed training of speech dnns. In Interspeech, volume 2014, pages 1058-1062. Singapore, 2014.
M. Shoeybi, M. Patwary. R. Puri, P. LeGresley, J. Casper, and Catanzaro. Megatron-Im: Training multi-billion parameter language models using model parallelism. arXiv preprint 1rXiv:1909.0800 08053, 2019.
S. U. Stich. Local sgd converges fast and communicates little. arXiv preprint arXiv: 1rXiv:1805.0976: 2018.
U. Stich, J.-B. Cordonnier, and M. Jaggi. Sparsified sgd with memory. Advances neural information processing systems, 31, 2018.
J. Sun, T. Chen, G. Giannakis, and Z. Yang. Communication-effcienn distributed learning via lazily aggregated quantized gradients. Advances in Neural Information Processing Systems, 32, 2019.
I. Sutskever, J. Martens, G. Dahl, and G. Hinton. On the importance of initialization and momentum in deep learning. In International conference machine learning, pages 1139 1147. PMLR, 2013.
T. Vogels, S. P. Karimireddy, and M. Jaggi. Powersgd: Practical low-rank gradient compression for
distributed optimization. Advances in Neural Information Processing Systems, 32, 2019.
Y. Wei, T. Hu, C. Liang, and Y. Cui. Communication optimization for distributed training: architec- ture, advances, and opportunities. IEEE Network, 2024.
W. Wen, Xu, F. Yan, C. Wu, Y. Wang, Y. Chen, and H. Li. Terngrad: Ternary gradients reduce communication in distributed deep learning. Advances in neural information processing systems, 30, 2017.
H. Zhao, X. Xie, C. Fang, and Z. Lin. SEPARATE: A simple low-rank projection for gradient compression in modern large-scale model training process. In The Thirteenth International Con- ference on Learning Representations, ICLR 2025, Singapore, April 24-28, 2025, 2025. URL ptttps://operevv//
net tete/forume ao (gOD
J. Zhao, Z. Zhang. B. Chen, Z. Wang, A. Anandkumar, and Y. Tian. Galore: Memory-efficient Ilm training by gradient low-rank projection. arXiv preprint arXiv:2403.03507 2024.
Y. Zhao, M. Li, L. Lai, N. Suda, D. Civin, and V. Chandra. Federated learning with non-iid data. arXiv preprint arXiv: 005800 2018
Y. Zhao, A. Gu, Varma, L. Luo, C.-C. Huang, M. Xu, L. Wright, H. Shojanazeri, M. Ott, S. Shleifer, A. Desmaison, C. Balioglu, B. Nguyen, G. Chauhan, Y. Hao, and S. Li. PyTorch FSDP: Experiences on scaling fully sharded data parallel, 2023. arXiv: 2304.11277.
13

---

Published as a conference paper at ICLR 2026
7 USE OF LARGE LANGUAGE MODELS
Regarding paper writing, we used LLM only for text polishing and grammar correction during manuscript preparation. No LLMs were involved in the conception or design of the method, periments, or analysis. All technical content, results, and conclusions have been independently verified and validated by the authors.
8 CONVERGENCE ANALYSIS
In this section, we establish the convergence guarantees of the DeMo algorithm under standard assumptions: unbiased stochastic gradients with bounded variance (Assumption 4 L-smoothness of the objective (Assumption 5 and bounded gradient norms (Assumption 6 The main result is Theorem 2 which proves ((1/VT) convergence rate in terms of the average gradient norm. We further quantify the bias and approximation errr introduced by sparse top-k averaging (Lemmas and 8.1 which are crucial in analyzing the impact of sparsified communication.
Starting from SGD with Momentum, we make two key modifications: first, we remove the all-reduce operation on gradients 9k. decoupling momentum across the accelerators. Second, after updating the momentum, we extract and remove its fast components q, which can be efficiently synchronized with minimal communication. Algorithm presents the complete method:
Algorithm Decoupled Momentum Optimization
Input: learning rate 7), decay 3 E (0, 1), parameters momentum my hyperparameters 8, gi f LocalStochastiGGradient (x:) {Get Local Gradient Without All-Reduce)
my + 8m, 9t (Accumulate Gradient in Momentum m)
q < ExtraceFFssCCompoonnttsmmm s, k)) [Extract Fast Components & From m}
m++1 m4 Q: (Remove q From m)
Q. + Synchronize(qg {Synchronize Across All Accelerators}
Fe+1 Ht 1tQ (Parameter Update Step)
8.1 PRELIMINARIES
8.1.1 NOTATIONS
We summarize the key notation used throughout the convergence analysis in Table 2
8.1.2
MULTILINEAR TRANSFORMS AND CHUNK- WISE TENSOR BLOCKING
We begin by introducing the tensorial operations central to our algorithmic design. Specifically, we define a multilinear product T acting on tensors via separable linear transforms, and describe the blocking operator B used to partition tensors into contiguous chunks. These constructions allow us formalize the application of the Discrete Cosine Transform (DCT) and its inverse in a chunk- wise manner, which underpins the sparsification mechanism employed in the DeMo algorithm. We also introduce the key intermediate tensors used in the algorithm's 's updates, along with relevant dimensional and notational conventions.
Let X E fa X be tensor of order and let matrices P. E R be given for each i [d]o (0, 1. 1}. We define the multilinear product T as follows:
T(x: Po, Pi...., Pa-1) E Xa X-a-
whose entries are explicitly given by
80
T(X; Po...., Pa
Pa-i)ioh
Σ jd-1=1
(PP)
1o -i
k=
In the special case d = 2, this definition reduces to the familiar matrix multiplication form:
T(X; Po, Pı) P.
14

---

Published as a conference paper at ICLR 2026
Table 2: Summary of notation used in the convergence analysis.
Notation
Definition
X. X
Tensor in X
d
Tensor order (number of dimensions)
Size of tensor along dimension Chunk size along dimension i Number of chunks along dimension i, n, = C;8;
G
P.
Transformation matrix (c.g.. DCT) for dimension i, P, E R4x5
Multilinear tensor product operator
B
Blocking operator (tensor partition into chunks)
Chunk index: (kg.... ... k with k [ci]o
(x))))
Chunk indexed by (x²) = X ( (Ko + 1)80:-...................
ka 18d-1 .1.9....,k. (Ka-1 1 + 1)8a-1)
M;
Local momentum worker i, iteration
M
Global momentum
G
Stochastic gradient at worker iteration
P
Intermediate tensor: P; -M G;
P
Chunk-wise DCT of intermediate tensor P:
Q:
Perform chunk-wise top-A selection P followed by chunk-wise inverse DCT
P.
Aggregated intermediate tensor: P. Be, + & Σ (Qj ()1 Aggregated tensor after sparse top-k aggregation, iteration
4
X
Model Parameter at iteration
Momentum decay parameter (0 < B 1) Learning rate
X
Weight decay parameter
Sparsity parameter (number elements retained per chunk)
D
Total number of tensor elements: D 74
M
Number of elements per chunk: M 1==0
(()
Objective function
E
Optimal objective value
E
Random samples from data distribution D Dj
o²
Variance bound of stochastic gradients Number of distributed workers
N
D. Dataset at worker
evident that the Discrete Cosine Transform (DCT) represents a specific instance of the multi- linear product, alongside its counterpoint, the Inverse Discrete Cosine Transform (IDCT). Both are defined as unitary matrices, leading to the equality P. = P. for each E [d[o-
Consider now tensor X' E 2 X-=-1 Suppose each dimension is divisible by can write n, = c,s; for all E [d]u- We define the blocking operator B acting on X' as:
((X) (Bx)))
[ex]ox Jas-1lo-
where each block Bx(X)) E Roxxs x is explicitly defined elementwise by: (Bx(A))j Noosotjo. kis +hi, .... 18-188-1 +ja-1
with indices:
k = =kku, ....... 1 specifying the position of the block, where h; E [cilu-
j (Jo.J1.... ja-1) indexing within each block, with js [s:]o-
Example (Matrix Blocking). Consider matrix X' I 128x512 We partition X' into contiguous submatrices (blocks) size 16 32. Thus, each dimension factorized as:
no 128 8 x 16, Mi = 512 = 16 x 32.
We define the blocking operator B applied to X as:
B(x) (B
15

---

Published as a conference paper at ICLR 2026
where each block Ks (2) R 16x32 is explicitly given by:
(Bko. (o)) 16+ & -2+
with indices ko {0, (0,1,.... 7), ky E {0, ,..., 15) denoting the block position, and local indices within each block given by jo E {0, 1,,1,... 15). J1 E (0.].
For instance, the upper-leftmost block (o,o(2) has entries:
2o,0
Fo,1
0,
81,0
X1,1
X1,31
(2)
X15,0
X15.1
X16,31
Let X be d-dimensional tensor, partitioned into chunks of shape (sg X ... X 8d-1). We denote an individual chunk of X by (X),, where is multi-index ranging over the set of all chunk indices K.
We define an intermediate tensor P/ as
P : MM; G
Next, we apply a chunk-wise Discrete Cosine Transform (DCT), denoted by T, to each chunk of Pf, yielding the transformed tensor P This operation is defined as
Bx (P:)) = T (Bx(P:)) for all E K.
Finally, for notational convenience, we define the total number elements in the original and chunk shapes as
D
M
II 8:-
By default, we let ||-|| denote the & norm unless otherwise specified. We denote c as the aggregated tensor obtained via scatter all-reduce operation over the encoded tensors Qł from each worker as detailed in
8.1.3 PROBLEM SETTINGS
In general, we consider minimizing the following objective function:
min ((X) = Eg-D [C(, XEX
where = X and L(X, {) is general loss function, and the expectation is taken over the data distribution D from which samples are drawn. We denote by L := infxex L(X) the optimal value of the objective function and assume throughout this paper that L is finite and bounded from below. Given a realization L(X, {), the stochastic gradient VL(X, {) is defined as the gradient of L(X, {) with respect to the parameter vector X.
In the distributed training setting, we aim solve the optimization problem:
min ((X) X
BE [x(,
(10)
N 1=1
where N denotes the number of workers, and (D,), represent the datasets available at each workerf Here, X represents the model parameters (e.g., neural network weights). Under this distributed scenario, each worker [N] maintains its own dataset D., and there is central- ized server accessible to workers for communication. At training step t, the stochastic gradient G; : VL(X,, { at worker is computed using data batch sampled from the dataset D,
Throughout this work, assume datasets {D; JY consist of i.idd samples, and each sample & D, is drawn independently. However, our proposed method directly extend non-i.i.d. settings.
16

---

Published as a conference paper at ICLR 2026
Definition (Variance). The variance of X-valued random variable X is defined
Var X =
E
|
|
||
With the definition of variance, we have the following assumption used for the analysis of stochastic settings.
Assumption 4. The stochastic samples are independent and identically distributed (i.i.d)).
Additionally, the stochastic gradient (x(,)() satisfies:
o²
= and Var Var(VL((X,E)
7lbatch
where a² is a finite constant and Tbbtch denotes the batch size.
Assumption f ensures that the variance of the stochastic gradient is uniformly bounded. For our discrete-time analysis, we also utilize the I-smoothhnss condition stated in Assumption
Assumption (L-smoothness), The objective function ((X) is differentiable, lower-bounded (i.e., L infxex ((X) > -co), and L-smooth.
We say that a differentiable function L X L-smooth if for all X, Y Ex,,
VV VC(()]l= 1 - X11p.
smoohhhh then for all X, Y X, have
L Y X ((Y) S ((X) (VC(X), X)
2
Assumption 6 (Bounded Gradient). For any X E X, and D, the stochastic gradient satisfies E [[]] < R with R 0.
Assumptions 4 5 and 6 are standard in the analysis of stochastic optimization algorithms Bernstein and Newhouse 2024). Détossez et al. (2022): Liu et al. (2024b).
We employ a sparse aggregation protocol as a form of all-reduce. The procedure for calculating the final aggregated vector is as follows:
Algorithm Sparse Top-k Averaging
1: Input: A set of tensors (X 1==1 where X Rd; sparsity parameter k E {1.....d) d).
2: Output: An aggregated tensor E Rd
3: Initialize sum S Rd and counts zd to zeros.
4: for j I,.... N do
5: I₁ TopKIndices( kk
6: SI, + ST, + x ())
7: nI, + nI, +1
8: end for
9: X + S 2 n, where Ø denotes element- wise division with 0/0 = 0.
return X
With the problem settings and mitigating strategies discussed above, proceed the analysis of convergence and communication complexity of our proposed Decoupled Momentum Optimiza- tion (DeMO) algorithm.
8.2 CONVERGENCE ANALYSIS
In this section, provide the convergence analysis for the proposed Decoupled Momentum Op- timization (DeMO) algorithm. We make use of the standard smoothness, variance and bounded gradient assumptions, stated in Assumptions 45 and 6 in the analysis.
17

---

Published as a conference paper at ICLR 2026
Algorithm 4 Decoupled Momentum Optimization
Given: learning rate 7), momentum decay rate 3. weight decay A, chunk sizes (So,... 88-1) sparsity k
DCT matrices [P:]
2: Initialize: & 0, parameters Xo. global momentum Mo
repeat
4: 6- t++
5: Worker Side (in parallei for each worker i):
6: Maintain local momentum residual M
7:
Compute gradient G; + V Li-1) Update local momentum: M; AM; G
8:
9:
for each sub-chunk (M)) of M; do Compute DCT: Q < (Mi).. Po...., Pa-1) Threshold thr k-tth largest magnitude in [Qk,t
12:
Sparsify: MASK Qi,i 2 thr
13:
Qk. + Qi, C MASK
14:
Update residual: Bu (M:) &- (x(M)) T(Q .o....,,PP.
15:
end for
16:
Send sparse encode( (Qk,))
Server Side:
18:
for each sub-chunk B& (M,) of global momentum M, do
19:
Aggregate: B₁ (M,)
Aagregate(ennoode (Qk,,))
20:
Inverse DCT: B& (M, + ((,)(M,)) Po ...., Pd-1)
21:
end for
22: Update parameters: X, f- X.-1 v(sgu(M.)) XX.-1)
23: Broadcast (g,(M,) to workers
24:
Worker Side:
25:
Receive and update local parameters with sgn(M:)
26:
until stopping critcrion met
27:
return optimized parameters X,
Theorem (Convergence of DeMo). Under Assumptions 6 and let (X.) be the sequence
generated by Algorithm 4
Then, the following convergence bound holds:
T
-IC(Xx)- C(X,) + 2LD + batch
E B(V
Tn
V1 V
M
2T
M
RVD
3+2D
1+
k 1- B
N
k
M
Remark (Special Cases and Convergence Rates), (i) In the special case where k = M, the last term vanishes, simplifying the bound to
T
B[C(X,)
L(XT)) + +LLD +
E ((X,))
T"n
Thatch
Choosing the step size n = θ (1z) we achieve the convergence rate:
T =
1
+0
VN
18

---

Published as a conference paper at ICLR 2026
(ii) Alternatively, we choose 3 VI/T. the bound becomes
()) ((XT)]
T E [IV
+ 2L.D
Tn
VNnbath
V 1
# M
(VI
2T
M
+RVD
1+ V
V
M
With the additional choice of step size 1 θ (1))
the convergence rate is again:
T
(
Σ [[I
+0
VT
VN
Proof. Let first decompose the difference in the loss function & using the L-smoothness property:
L 2
((X++1) ((X,) s (VL(X.), X1+1 X,)
X
x,11
(I-smoothness of L)
L |X xx
(II)
--=VVV(X)),siem((VXXXX + |X - xx + sign +
& - + 2L
+ (VC((X, ) sign VZ((x,) sign(M++1)).
where we used the inequality
[ X,11 = 1/21 Mi S sdd
Using Lemma 8.3 with 1, we bound the last term
(V sign (igg(VV((X))) sign(M++1)) (() M
(12)
To bound the term E VL(X,) M l1. we apply the triangle inequality:
N (OM) VC((,
N
[VC(()) Mct1ll1
VC(X.) N Σ
(aM) + VC(X, +)) M
N
1Till + 122111
We separately bound the terms Ti and T2.
Bounding Ti: Notice that
1 N
Ti VL(X) N Σ
VL(X,,))
B
N
E=1
(00)
To bound the first term (*), we recall:
N (+) = -VC(X,)
N Σ vL(X,e))
19

---

Published as a conference paper at ICLR 2026
By Assumption 4 the stochastic gradients VL(X,, { have bounded variance, i.e.
E [IV 5t)
Tbatch
Thus, applying Jensen's inequality (from Eg to E norms) and independence across i, we have:
G2
VVD DD
[[I({] VDB[I()) z vDV Blll )
VDV
N ribatch
N Tibatch
For the second term (+*), we have
(**) = N Σ SM;
Since the tensors {M are identically distributed, follows that
N
(++))) S
Σ N
BB|[ = [|]
Applying Lemma] I we can further bound this expectation
(R1/D(1-1//)) -
1 1 BV1-k//M
Bounding T2: Using Lemma 2 we have
1
E T|] < 222
DE
1
Pilla
(since DCT is orthonormal)
N
N
D
E
(B4(1))1)
Bu (P:)
V
h M
2T
M
$ <2D3/2R
R
1+
1+
(by Lemma
N
1 B
k
M
Combining these results, we conclude
B[IVC((X,) -M++ [hrrhh+ < oVD
+ RVD K/M 1-
2T
M k
B B+2DD 2D
1+
NNbatthh
N
Hence, from equation
we obtain
T
B[[(X,) -((x)) + +2LDm
Tn
/N NDDatch
RVD VI V1- k/M
2x
)
3/1 K/M 8 B+2D
1 1+
N
To analyze the convergence of Algorithm we first need bound on the expected norm of the momentum tensors M,. The following lemma provides this bound under Assumption
20

---

Published as a conference paper at ICLR 2026
Lemma 1 (Bound momentum M,). Suppose Assumption 6 holds, and let (M,)[= be the mentum tensor generated by Algorithm 4 expected norm:
Then, for all t 2 0, we have the following bound on its
R3 k-M
E| M 1
1- 3- K/M
provided that BV1 k/M < 1.
Proof.
Recall from Lemma 8.1 that the momentum Mi follows the recursion:
k | Gj
|Mila S
M
Taking expectations and applying the triangle inequality, we have:
k: (BEIIMM [2] + ll))
M
Iterating this bound recursively, we obtain:
k
M ||| S (BE[IM + R)
k x (1
k
R RV 1
1 B 1
M
M
M
6-1
k M
R
B
The geometric series can be bounded as follows:
k
R-1 k/M
8
<
M
1-BVI-k/A
assuming 8 k/M < 1 for convergence.
Thus, we have the final bound:
R k/M
E M|
1- 1-3
completing the proof.
The analysis of sparse top-k averaging begins with Lemma 2 which provides bound on the ex- pected bias discrepancy [8] highlighting the trade-offs between sample size, dimension, and spar- sity. This is complemented by Lemma 8.1 which establishes a universal bound for the fo-nrrm approximation error, illustrating vector approximation through top- component selection.
Furthermore, Lemma 8.3 offers insights into the geometric relationships between vectors and their sign functions, connecting inner product space and L -noor differences. This is supported by Lem- mas which explore sign-based approximation properties and variance.
Lemma 2 (Bias Bound for Sparse Top-k Averaging). Let integers N, M, satisfy 1 S k S M, R be independent and identically distributed, explicitly
and let random vectors {X()) satisfying:
(i) Coordinate independence: For each j [I,...,N,, the coordinates (x(") pendent and identically distributed random variables:
Jm are inde-
x(3) i..dd Xt, and E[X.]] Pi.
21

---

Published as a conference paper at ICLR 2026
(ii) Boundedness: There exists a finite constant B >> 0 such that, almost surely,
1x!0 S B, for all i (I....,M), je (1....,
Consider the estimator X IR defined by Algorithm
(Sparse Top-k Averaging). Define the
standard empirical mean
(full)
1
Σ X
N
and the bias discrepuncy introduced by sparsity as, for each coordinate E (l,...,M)
6 : - X; {
Then, under these assumptions, there exists a universal constant C > 0, independent of N, M, k and B, such that for every coordinate ;
M
E S
C
B
1+
N
k
Proof. Recall that
1 N
6, = xX (0()-x(1)))
X
S
Σ
x
+
Σ
x
P4
= Ti T2.
T4
N
JEJ.
j=1
Bounding the first term Tí.
The variable n, follows a Binomial N, p;) distribution with pt = k/M. Conditional on n, Hoeffd- ing's inequality gives
(- 2n,,7²
P (Ti
m;) 2 exp
n1 0.
B²
We split the analysis into two cases based on n;;
Case 1: Typical case (n, 2 Np:/2). Using Hoeffding's inequality, we bound
M
T n, 2 Np:/2) S C
=C B
BV
Np:
N
for some absolute constant C" > 0.
Case 2: Rare case (n; Np:/2). Chernoff bounds for binomial distributions imply
P(n: < Np:/2) < exp(-cNps)
for some absolute constant > Thus, we trivially bound Tj S 2B to obtain
E n4 Np./2 < 2B exp(-cNpl)
which becomes negligible for large N relative to the polynomial terms.
Combining both cases, have for large N.
M
5 C"
Nk
Bounding the second term Tz.
Since X are bounded in [-B, B). Hoeffding's inequality gives, for any 7 > 0.
2NN2
P (T2 2 n) S 2exp
- B
22

---

Published as a conference paper at ICLR 2026
Integrating over all y yields
00
2Nt B
C 1
[[ -
E
eexp
dt
for some absolute constant C² > 0. Combining both terms Tj and T2.
Summing the bounds for Ti and T2, we obtain for sufficiently large N:
M
1 VN
B[6, S C'B VN
+
C"
+
0
Nk
Factoring B/vN explicitly, we conclude
M
B S CBV
1+ F )
1
N
for an absolute constant C large enough (e.g., C= mmx((CC .C")). This completes the proof.
Remark (Dependence on Parameters). The bound in Lemma clearly illustrates the scaling of the bias with respect to the number of samples N, dimension M, sparsity level k, and coordinate magnitude bound
Sample size (N): The bias decreases at a rate of 1///N, consistent with standard statistical estimation theory.
Dimension and sparsity (M, k): The bias scales as MM/k, highlighting the bias intro- duced by aggressive sparsification (small k).
No sparsification (k- M)) this case, the bias vanishes, as expected.
Here present and prove bound concerning the approximation error incurred by approximating a vector by its top-k largest magnitude components.
Lemma 8.1 (Optimal Universal -nrormm Bound). Let X x (x,, XM) trary real vector. For fixed integer k with < k < M, define the vector X E RM by selecting the k largest magnitude elements of X and setting all other entries to zero. Then, the approximation error satisfies the bound
M be an arbi-
| -X--x11+ s (1- M 1x-1-
Furthermore, this bound is sharp. Equality holds if and only the magnitudes of all elements of X equal.
Proof of Lemma 8.1
Notation and Problem Setup. Consider the vector X R. Let us sort the elements of X in non-increasing order of absolute value:
1lw 2 | 2 ... |
We define X by retaining the & largest magnitude elements and setting the remaining M k elements to zero. Thus, the approximation error is explicitly given by:
M
X x = Σ
X (())
1=
while the original vector norm is:
M | = Σ
X
23

---

Published as a conference paper at ICLR 2026
Identifying the Worst-case Scenario. To determine the maximal possible ratio of the approxima- tion error to the original norm, consider the scenario where all vector entries are equal in magnitude. Suppose, without loss of generality, that:
X = (a, a,,
a + 0.
In this case, any selection of h elements yields the same approximation error. Explicitly, we have:
|X
XI11 (M k))a),
|| Ma..
Hence, the ratio is exactly:
X - -XI11
M - k
k M
|
Universality and Optimality of the Bound. We now verify that this equal-magnitude scenario indeed represents the worst-case for any vector X. For an arbitrary vector X, the approximation emor ratio is:
Ei=k+1 (x() T=1 (x())
This ratio is maximized precisely when the magnitudes of vector elements are equal. To see this, observe that any deviation from equal magnitudes increases the share of the total magnitude held by the top-k elements, thus reducing the ratio. Hence, the equal magnitude scenario is indeed the worst-case.
Therefore, the universal bound:
1. M
|| - -XI- S 1
||x||
holds for all vectors X E M
Remark Similarly. for & norm, we have
k
|X - X112
V
|
M
Remark (Approximation error bound under isometrie transformations). Let A RM RM be an invertible linear operator satisfying the isometric property
||A(X) la
||
||)=1 |x|12, for all X e R
Define the vector approximation
X' X' A-= (topx(()))))
where (Yp retains only the k largest-magritude entries of and sets the others to zero. Then the approximation error satisfies the bound
X X112 S
|
This bound is sharp. with equality attained precisely when the entries of A(X) have identical mag- nitudes. In particular, when k M, the approximation error is exactly zero, indicating perfect reconstruction.
Lemma 8.2. For any I. E R, have
|x| - x sign(y) S 2/x - y/.
Proof. If sign(y) sign(w), have |x|| sign (v) < 2/ y/-
If sign(y) = (ig)( we have |ac| - x sign(yy = 2|x] < 2| + 2/yl = 2 - 3/-
If (y)n(y) O. have |ar| sign(()) |ar| =|x- 2 r v/-
24

---

Published as a conference paper at ICLR 2026
Algorithm 5 Decoupled Momentum High-level Training Loop
Require: learning rate n, momentum 3. weight decay sparsity DCT matrices {Pt}
1: Initialise parameters X, and global momentum Mo +
2: for f. = .,2.... do
3: // exeected on every worker (in parallel)
4: G₁ +- -i;
5: M. f +
Q COMPRESSCHUNKS(MI..K (P::)
7: send Qi to server
8: // parameter server
M. + DECCMPRISAAAAGGRETTTT()))..
10: X. f x (Msig((M.) - AX.-1)
broadcast (M,)(M)) all workers
12: end for
Algorithm 6 COMPRESSCHUNKS((M, k, (Pe)) run worker
Require: local momentum M. sparsity budget k DCT matrices {Pt}
I: Q + {list of sparse coefficients
2: for each chunk index do
3: B CHUNKEXTRACT(ML.C
< DCT(B; (P:))
5: Qsp f- Top-k(q, k)
6: Q Q U [(c, }p))
B+ B IDCT(qai (P. (store residual]
8:
CHUNKINSERT(M,..)
9: end for
10: return Q (sparse DCT coefficients transmit)
Lemma 8.3. Let (X,Y) is a joint random variable on I2 x Rd. For any constant a E (0, +00), we have
E { sign(X) sign(()))))n 2av - -
Proof.
Without loss generality, set a = Rign,xign(X) sign(Y)))
= B[||X|] ||: (X, sYgn(Y)))
- ||||| X Y|h
Lemma 8.2
s 2V 2T]B[|| YII
by Cauchy-Schwarz,
where || is the E norm and || denotes the Euclidean norm.
25

---

Published as a conference paper at ICLR 2026
Algorithm 7 DECOMPRES (Q))), (PTT) server side
Require: sparse sets } from all workers, inverse DCT matrices {Pi
I: initialise global momentum M
2: for each chunk index do
3: $ --- (sum of sparse coeffs}
4: for each worker do LOOKUP(Q',C
5:
end
7: B + IDCT((: {P,TT)
8: CHUNKINSERT(M, c, B)
9: end for
10: return M (rconnstructed global momentum)
Model
Final Loss
Hellaswag f acc_norm
ARC-Easy 1
PIQA f
Data Tx & MB/step
acc
acc normm
DeMo 300M
64, k 32
2.87
0.37
0.46
0.67
29.9
64, k 16
2.87
0.38
0.50
0.67
14.9
64, A: = 8
2.88
0.38
0.47
0.67
7.49
64, & 4
2.89
0.37
0.47
0.67
3.74
64, Æ = 2
2.93
0.36
0.46
0.65
1.87
64, K
2.97
0.35
0.45
0.65
0.93
128, & =
2,88
0.37
0.50
0.66
7.49
128, Y 16
2.90
0.37
0.47
0.67
3.74
128, 1- = 8
2.93
0.36
0.49
0.66
1.87
128, & 4
2.98
0.35
0.46
0.64
0.93
128, L 2
3.06
0.33
0.45
0.65
0.46
128, K
3.16
0.31
0.45
0.63
0.23
Adam W-DDP 300M
2.98
0.35
0.46
0.65
636.9
DeMo 1B
64, 1: 32
2.63
0.48
0.55
0.70
110.32
64, k 16
2.63
0.47
0.53
0.70
55.16
64, & 8
2.64
0.47
0,52
0.69
27.58
64, K
2.67
0.45
0.52
0.70
13.79
64. k = 2
2.71
0.44
0.51
0.69
6.89
64, K 1
2.76
0.41
0.52
0.69
3.44
128, & = 32
2,65
0.46
0.53
0.69
27.58
128, & 16
2.67
0.46
0.50
0.70
13.79
128, 1: 8
2.72
0.44
0.52
0.68
6.89
128, K- 4
2.76
0.41
0.50
0.67
3.44
AdamW-DDD IB
2.73
0.43
0.51
0.68
2416.6
Table 3: Results of training loss, downstream evaluation scores, and per-GPU communication quirements of the model sizes and reference trained on 100B tokens
26


Reviews:
[{"review_id": "t8cHdaz0J2", "reviewer": "Reviewer_x4Vv", "summary": "This paper proposes DeMo (Decoupled Momentum Optimization), a communication-efficient distributed training method that addresses the gradient synchronization bottleneck in data-parallel training. \nKey innovation: decouple local momentum updates across workers, apply blockwise DCT transformation followed by top-k sparsification, and reuse the momentum buffer for error feedback through momentum subtraction. \nExperimental results: the authors demonstrate up to 85× reduction in communication while maintaining comparable performance to AdamW on 300M and 1B parameter language models trained on 100B tokens.", "questions": "Weaknesses section leaves a lot of questions.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- Practical Impact: Achieves dramatic communication reduction (up to 85×) while maintaining model quality, which is highly valuable for distributed training scenarios with limited bandwidth.\n- Clean Design: The three-component approach (decoupled momentum, DCT+top-k, momentum subtraction) is conceptually clear and builds naturally on existing optimization principles.", "weaknesses": "- Limited Scalability Analysis: The download bandwidth scales with the number of workers, which could become prohibitive at large scale; but all experiments use 64 GPUs; scaling behavior to hundreds or thousands of GPUs remains unclear.\n- Multi-datacenter: The method is positioned for \"multi-datacenter\" training but only tested within single datacenters.\n- Baseline Comparisons: No comparison with other communication-efficient methods (e.g., gradient quantization, other sparsification approaches like Deep Gradient Compression, or recent methods like DiLoCo beyond brief mention).\n- Baseline: Missing comparison with gradient compression baselines that would contextualize the 85× improvement.\n- Baseline HParams: AdamW baseline uses standard hyperparameters, but DeMo requires β=0.999 (vs standard 0.9) - this hyperparameter sensitivity isn't fully explored.\n- Scale: Models are relatively small (300M, 1B) - behavior at 7B+ parameter scales is unknown.\n- Theoretical: No analysis of how DCT specifically helps beyond empirical observation. Why is DCT the correct basis to project onto?", "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": "WQrQPdWQUv", "reviewer": "Reviewer_Vvbj", "summary": "This paper proposes Decoupled Momentum Optimization (DeMo), a communication-efficient framework for distributed training designed to replace standard gradient-based All-Reduce in data-parallel setups. The core idea is to synchronize a compressed version of the optimizer's momentum state instead of the full gradient. DeMo works by (i) allowing each worker to maintain and update a local momentum buffer, (ii) applying a blockwise fast orthonormal transform (e.g., DCT) to this momentum, (iii) communicating only the top-k sparsified coefficients, and (iv) performing momentum subtraction to reuse the local momentum buffer as an implicit error-feedback mechanism. The authors demonstrate that this approach drastically reduces communication overhead while achieving comparable pre-training loss and downstream accuracy to a standard AdamW baseline on 300M and 1B models.", "questions": "1. Communicating the principal components of the momentum, for instance by using SVD or a random projection onto a subspace, could also achieve the goal of extracting the dominant information. What are the specific advantages of using a fixed DCT basis followed by top-k sparsification compared to these adaptive low-rank projection methods?\n\n2. In Figure 2, the DeMo training curves converge significantly faster and to a lower loss than the AdamW-DDP baseline. Intuitively, compression is lossy and introduces an information bottleneck, which might be expected to slow down convergence, not accelerate it. How do the authors explain this phenomenon where a compressed, lossy update outperforms the full-precision baseline?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper's core idea of decoupling momentum into dominant components (for communication) and non-dominant components (for local updates) presents an effective and intuitive paradigm for reducing communication overhead.\n\n2. The application of DCT as the mechanism for sparsification, rather than operating in the raw gradient or momentum space, is a novel approach that provides a new perspective on structured compression.", "weaknesses": "## 1. The theoretical analysis, particularly the convergence rate presented in Theorem 1, warrants further discussion as it does not appear to be optimal. \nFor compressed stochastic optimization algorithms in the non-convex setting, convergence rates can achieve $\\mathcal{O}(\\omega/\\sqrt{NT})$ (where $\\omega$ is the variance bound of the compression estimator, see [1]). Notably, this rate is dimension-free ($D$ is not in the numerator). The dimension $D$ typically appears when analyzing the total communication complexity (convergence iterations $\\times$ per-step cost), which might be $\\mathcal{O}(D/\\sqrt{NT})$. A dependency on $D$ in the convergence rate itself is common in analyses of Adam-based optimizers, which often require an $l_1$-norm analysis due to per-element operations. However, SGD-based algorithms typically yield an $l_2$-norm result and a dimension-free convergence rate [2]. Although Theorem 1 is applied to SGD/Signum-based methods, the proof relies on an $l_1$-norm analysis, resulting in a dimension-dependent convergence rate of $\\mathcal{O}(D/\\sqrt{T})$. When factoring in the per-step communication cost $k$, the overall complexity does not seem to surpass existing state-of-the-art results.\n\n## 2. The paper provides insufficient comparisons, both theoretically and experimentally. \n\n2.1 Theoretically: As noted above, the authors do not compare their theoretical bounds against other established compression algorithms, making it difficult to ascertain a clear advantage from the provided convergence rate. Intuitively, DeMo acts as a hybrid between local SGD and standard global-all-reduce SGD, where dominant momentum directions are communicated precisely while non-dominant directions are updated locally. This bears similarity to methods that perform communication and optimization on a subspace [2, 3], yet the theoretical results for DeMo do not appear to go beyond this existing work.\n\n2.2 Experimentally: The paper only compares DeMo against a standard AdamW-DDP baseline. It lacks comparisons against other state-of-the-art communication compression methods. Without this, it is difficult to conclude that the proposed paradigm of compressing momentum instead of the gradient offers any essential advantage.\n\n## 3. The paper focuses on communication compression, but the experimental scale (300M and 1B models) is relatively small for assessing this bottleneck. \nIn pre-training models of this size, communication overhead often accounts for only 20% or less of the total training time. Therefore, even with a significant reduction in data transmission, the actual wall-clock time savings might be limited, especially if the compression scheme introduces any convergence instability or computational overhead. Reporting end-to-end wall-clock time comparisons is crucial to demonstrate the method's practical benefit, but this is missing from the experiments. Further experiments in larger-scale settings (e.g., those involving model sharding and hybrid parallelism) would be valuable to demonstrate its practical impact where communication is a more significant bottleneck.\n\n[1] MARINA: Faster Non-Convex Distributed Learning with Compression\n\n[2] SEPARATE: A Simple Low-rank Projection for Gradient Compression in Modern Large-scale Model Training Process\n\n[3] Greedy Low-Rank Gradient Compression for Distributed Learning with Convergence Guarantees", "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": "PS4oiTGp0n", "reviewer": "Reviewer_yL27", "summary": "This paper introduces Decoupled Momentum Optimization (DeMo) to reduce communication overhead. DeMo integrates three key techniques: local momentum decoupling, a fast orthonormal (e.g., DCT) transform with top-k sparsification, and momentum subtraction for error feedback.", "questions": "- Counterintuitive Performance: As shown in Figure 2, DeMo with various Top-K sparsity levels appears to outperform AdamW. This is counterintuitive, as Signum (the uncompressed version of the update direction used in DeMo) is generally known to underperform compared to AdamW. Could the authors provide an explanation for this phenomenon?\n\n- Clarification on Parameter $\\alpha$: The parameter $\\alpha$ introduced in Subsection 3.4 is not mentioned in the preceding sections. Could the authors clarify its role and meaning within the DeMo framework?", "limitations": "", "rating": 4, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The primary objective of this paper is to democratize the training and fine-tuning of large models by minimizing communication volume, which obviates the requirement for expensive, high-performance networking infrastructure.", "weaknesses": "Limited Perceived Novelty: The claimed novelties of DeMo appear to build upon existing work. For instance, [1] also compresses the momentum term rather than the raw gradient. Even earlier literature has demonstrated that compressing the total model update, rather than the intermediate gradient, can maintain competitive performance while dramatically reducing communication overhead. Furthermore, the technique of using the momentum buffer to store error feedback for memory efficiency is also employed in [1].\n\nQuestionable Practical Efficiency: The net reduction in overall training time achieved by DeMo may be minimal, despite its apparent reduction in communication volume. DeMo requires computationally expensive operations, including a sorting algorithm to identify the Top-k elements and two orthonormal transforms per step. The overhead from these operations is non-trivial and may offset the gains from reduced communication time.\nMoreover, as a variant of the Top-K algorithm, DeMo's communication savings are primarily effective in parameter-server architectures. In a decentralized architecture, aggregating Top-K elements requires a AllGatther operation instead of a more efficient AllReduce. The communication cost of AllGather is $O(n^2K)$ while AllReduce is $O(nK)$, where $n$ is the number of nodes. This scaling implies that DeMo may not yield net communication savings with a large number of nodes. This concern is supported by [3], which experimentally demonstrated that Top-K methods can be slower than non-compressed optimizers, even when network bandwidth is the bottleneck.\n\nReferences\n\n[1] Ahn, Kwangjun, et al. \"Dion: Distributed orthonormalized updates.\" arXiv:2504.05295, 2025.\n\n[2] Peng, Hanyang, et al. \"Birder: communication-efficient 1-bit adaptive optimizer for practical distributed DNN training.\" NeurIPS, 2023.\n\n[3] Agarwal, Saurabh, et al. \"On the utility of gradient compression in distributed training systems.\" Proceedings of Machine Learning and Systems, 2022.", "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": "kTKI6c3Z20", "reviewer": "Reviewer_NyAn", "summary": "This paper proposes DeMO, a new algorithm for communication efficient distributed training. DeMo works by compressing the communication by transmitting a compressed version of a local momentum from each worker instead of the gradient. The compression is performed by performing a discrete cosine transform and only transmitting the top coefficients. They also explore using error feedback on top of the local momentum at each worker. The paper compares DeMo to an AdamW baseline for LLM training at a decent scale ~330M and ~1B scale.", "questions": "Overall I would like to see the experimental comparison significantly strengthened to change my opinion.\n\n* What learning rate values were used for the experiments, how were these tuned? Note that changing the momentum coefficient for certain optimizers can change the \"effective learning rate\" and shift the optimal LR value.\n* Why not compare against some distributed baseline? It feels a bit odd not too since this is not the first paper to focus on communication efficient optimization.\n* To what extent is this algorithm compatible with different reduction techniques, e.g. without a central server?\n* Why does the AdamW baseline perform so poorly in Figure 2 but well in Figure 3?\n* Comparing training curves directly like in Figure 2 is not very informative. Unless each line is fully tuned, what is the takeaway here? Changing the LR or other important HPs slightly could also significantly change any given curve.\n* L251-L254 are not very clear. Why these values for the momentum?\n* It would also be nice to clarify the computational overhead. What is a typical flop overhead?\n\nMinor pointers:\n* There is a missing reference in L438", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* Paper is well written, easy to follow with clear figures\n* Method is interesting and has practical relevance\n* Experiments are performed on a convincing scale\n* Code provided", "weaknesses": "* Experimental comparison is lacking in crucial ways:\n    * There is no comparison with other communication focused distributed training algorithms\n    * The AdamW baseline does not seem to be tuned\n    * There are no experiments with Muon even if this is discussed in Section 2.1", "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)