DeMo: Decoupled Momentum Optimization

ArXiv: 2411.19870

🎯 Pitch

Standard data-parallel training moves terabytes of raw gradients every step, demanding costly, co-located hardware. DeMo instead transmits a heavily compressed, transformed version of the momentum buffer, reducing communication by up to 85× while matching AdamW’s accuracy—eliminating the need for high-bandwidth interconnects.

1. Executive Summary

This paper proposes Decoupled Momentum Optimization (DeMo), a drop-in replacement for momentum-based optimizers that drastically reduces communication bandwidth in synchronous data-parallel training while preserving convergence. Evaluated on 300M- and 1B-parameter OLMo language models trained on the Dolma v1.5 corpus, DeMo combines three mechanisms: decoupled local momentum updates (each worker evolves its own momentum buffer independently, synchronizing momenta rather than raw gradients), structured tensor compression (chunk-wise Discrete Cosine Transform followed by top-k sparsification to communicate only the largest DCT coefficients), and momentum subtraction as error feedback (subtracting communicated values from the momentum buffer to accumulate residual information without extra memory). DeMo transmits up to 85× less data per GPU per step than AdamW-DDP while achieving comparable downstream accuracy — for the 300M-parameter model, k = 8 sparsification cuts per-GPU communication from 637 MB to 7.5 MB with no loss in zero-shot performance — establishing that aggressive momentum compression can substitute for full-precision gradient all-reduce, though the method is designed primarily for coordinating a small number of geographically distributed compute centers rather than replacing high-bandwidth intra-datacenter interconnects.

2. Context and Motivation

The Core Problem: The Communication Bottleneck in Distributed Training Is Unsustainable

Training modern large language models is not merely compute-bound — it is fundamentally communication-bound. The default strategy for scaling model training across GPUs is Distributed Data Parallelism (DDP), where each worker computes gradients on a micro-batch of data, then all workers participate in an All-Reduce operation to synchronize gradients before every optimizer step. The communication cost of this synchronization is proportional to the model's parameter count. For a 1-billion-parameter model using 32-bit floating point, a single All-Reduce transmits 4 GB of data per worker. At scale — state-of-the-art models now reach hundreds of billions of parameters — this balloons to terabytes per step, as the paper notes in its introduction.

This bottleneck has concrete consequences that the paper identifies:

Hardware lock-in: Training requires expensive, high-bandwidth interconnects like NVLink or InfiniBand, which force clusters to be geographically co-located (Wei et al., 2024). You cannot simply spread training across standard Ethernet-connected machines or across data centers.
Cost inflation: The specialized networking infrastructure needed to move terabytes of gradients per step represents a significant fraction of total training cost, and these costs do not scale linearly — they grow with model size.
Scalability ceilings: As model sizes increase, the communication-to-computation ratio worsens because gradient tensors grow proportionally to parameters while the useful computation per parameter (the forward-backward pass) grows more slowly. At some scale, communication latency dominates total step time, making additional GPUs wasteful.

The paper frames this as not just an inconvenience but a fundamental barrier to democratizing large-scale training. If you can only train large models in tightly-coupled clusters with expensive interconnects, then only well-resourced organizations can participate. Alleviating the communication burden would enable training across cheaper Ethernet-based setups, across geographically distributed clusters (e.g., pooling spare compute from multiple data centers), and on smaller-budget hardware.

The Information Redundancy Hypothesis

The paper's motivating insight is deceptively simple: the gradient tensors being communicated in standard DDP contain massive redundancy. Not all entries in a gradient tensor are equally important for optimization. This redundancy hypothesis is not new — it underpins decades of work on gradient compression — but the paper identifies specific characteristics of momentum-based optimizers that existing compression methods fail to exploit effectively.

In particular, the paper observes that momentum buffers (running averages of past gradients) are more compressible than raw gradients in a distributed setting. Why? Because momentum smooths out high-frequency gradient noise, producing a signal with more structure — and structured signals are easier to compress with transforms like the DCT. Moreover, the momentum buffer already tracks which information has been "accumulated" from past steps, making it a natural substrate for an error feedback mechanism that requires no additional memory. These observations are the intellectual seeds from which the DeMo design grows.

Where Prior Approaches Fall Short

The paper organizes prior work into three broad strategies and identifies specific shortcomings in each.

1. Gradient Compression: Sparsification and Quantization

The most direct approach to reducing communication volume is compressing the gradient tensors themselves before transmission. Two dominant families exist:

Top-k sparsification (Lin et al., 2018b; Stich et al., 2018; Aji and Heafield, 2017; Alistarh et al., 2018) transmits only the largest-magnitude gradient elements, zeroing out the rest. The problem is that when applied directly to gradients, these methods produce sparse update patterns that harm model performance. The paper states this explicitly: "Prior sparsification approaches directly applied to gradients incur sparse update patterns and often harm performance." The reason is that neural network parameters benefit from dense, well-distributed updates — aggressively zeroing out most gradient entries creates an uneven optimization landscape where some parameters receive many updates and others receive few. Error feedback mechanisms (Karimireddy et al., 2019) attempt to compensate by accumulating the compression error locally and adding it to the next gradient, but these introduce extra GPU memory overhead — a separate accumulator buffer must be maintained, which for a billion-parameter model costs gigabytes of additional memory.

Gradient quantization (Alistarh et al., 2017; Sun et al., 2019; Wen et al., 2017; Seide et al., 2014) reduces the numerical precision of gradient values (e.g., to 1-bit or 8-bit representations). While effective for moderate compression ratios, quantization alone cannot achieve the orders-of-magnitude reductions that sparsification offers. Moreover, many quantization schemes are biased, requiring error feedback similar to sparsification, reintroducing the memory overhead problem.

The key limitation across both families: they treat gradients as the compression target without considering whether the optimizer's internal state (specifically, momentum) is a better candidate.

2. Communication Frequency Reduction: Local SGD and DiLoCo

Rather than compressing individual messages, this family of approaches reduces the number of synchronization events. Local SGD (Stich, 2018; Lin et al., 2018a) and Federated Averaging (McMahan et al., 2017) allow workers to perform multiple local optimizer steps independently, then periodically average their parameters. Each synchronization transmits the full model state, but since synchronizations happen less frequently (e.g., every $H$ steps), the amortized communication cost is $H \times$ lower.

The paper identifies two significant failure modes for these methods:

Client drift (Karimireddy et al., 2020; Zhao et al., 2018): When workers evolve independently for multiple steps, their local parameters diverge. The averaged model after synchronization may be in a region of the loss landscape that no individual worker was near, potentially harming convergence. This is especially severe with non-i.i.d. data distributions across workers, though the problem exists even with i.i.d. data because stochastic gradient noise alone causes divergence.
Unpredictable optimization trajectories with modern adaptive optimizers (Nguyen et al., 2025): The complex interplay between local optimizer dynamics (especially for adaptive methods like AdamW with per-parameter learning rates and second-moment estimates) and infrequent synchronization can produce training instabilities that are hard to diagnose and debug. The paper cites this as a known empirical challenge.

DiLoCo (Douillard et al., 2023) represents the state-of-the-art in this category, successfully training language models with communication every $H$ steps. The paper acknowledges DiLoCo's effectiveness but positions DeMo as a complementary philosophy: rather than reducing synchronization frequency, DeMo maintains frequent synchronization but makes each synchronization message orders of magnitude smaller. The paper hypothesizes — and demonstrates in experiments — that more frequent but highly compressed communication can yield better convergence than infrequent full communication, particularly at high compression ratios.

3. Low-Rank Update Methods: LoRA, GaLore, and Subspace Optimization

A more recent line of work applies low-rank structure to the optimization process itself, inspired by the success of LoRA (Hu et al., 2022) for fine-tuning. GaLore (Zhao et al., 2024) projects gradients onto a low-rank subspace before passing them to the optimizer, reducing both memory and (potentially) communication. Related methods like Flora (Hao et al., 2024) and subspace descent approaches (Liang et al., 2024b) similarly restrict updates to low-dimensional manifolds.

The communication reduction from these methods comes from transmitting the low-rank factors ( $U$ and $V$ for a matrix parameter) instead of the full gradient matrix. For a parameter matrix of size $d \times d$ and rank $r$ , this reduces communication from $d^2$ to $2dr$ . When $r \ll d$ , the savings are substantial.

The paper positions DeMo relative to this line of work with a subtle distinction: DeMo's top-k sparsification in the DCT domain effectively makes the communicated update matrix rank-k per chunk, which connects to the low-rank philosophy but operates in a fixed, pre-chosen transformed basis rather than adaptively learning the low-rank subspace. The paper notes (Section 4) that DeMo's approach is related but distinct — the fixed transform avoids the computational cost of computing SVD or similar decompositions at each step, and the chunking strategy means different parameter tensors can be compressed at different effective ranks automatically based on how their energy concentrates in the DCT domain.

Critical Unaddressed Gap: Optimizer State as the Natural Compression Target

The paper's key critique unifying all prior approaches is that they compress the wrong thing. Gradient compression methods compress raw gradients, ignoring the momentum buffer that modern optimizers already maintain. Frequency reduction methods ignore compression entirely and simply communicate less often. Low-rank methods project the gradient onto learned subspaces. None of them ask: can we treat the optimizer's own momentum state as the communication primitive and compress that instead?

The paper's core insight is that the momentum buffer has three properties that make it a superior compression target compared to raw gradients:

Temporal smoothing: Momentum is a running average of past gradients, so it filters out high-frequency stochastic noise. A smoother signal compresses better under transforms like DCT because more energy concentrates in low-frequency coefficients.
Built-in memory: The momentum buffer already accumulates information from past steps through its exponential moving average. This means it naturally serves as an error feedback accumulator — subtracting communicated values from the buffer tracks what information has not yet been shared, without requiring a separate memory allocation.
Decoupling from all-reduce timing: By evolving local momenta independently (without synchronizing raw gradients), each worker's momentum buffer can serve as a richer, more informative signal that captures local data statistics before compression.

How the Paper Positions Itself

DeMo is presented not as yet another compression algorithm but as a framework for rethinking what gets communicated in distributed optimization. The paper's positioning can be understood through several deliberate design choices:

Drop-in compatibility, not a new optimizer. DeMo is explicitly designed to work with any momentum-based optimizer (SGD with momentum, Signum, Lion, Muon) by replacing the gradient synchronization step. In the experiments, the paper uses the Signum update rule (sign(M)) applied to the reconstructed momentum, but the framework is general. This is strategic: practitioners do not need to abandon their preferred optimizer to adopt DeMo.

Maintaining frequent communication. Unlike DiLoCo and Local SGD, DeMo synchronizes every step. This keeps the optimization dynamics close to standard DDP, avoiding the client drift issues that plague infrequent-synchronization methods. The paper explicitly contrasts this philosophy: compress each message to near-negligible size rather than sending full messages infrequently.

Topology agnosticism. DeMo's communication pattern — workers sending sparse coefficient sets to a parameter server which aggregates and broadcasts a compressed update — is simpler than All-Reduce and works over standard TCP/IP networks. The paper frames this as enabling "training across multi-datacenter or Ethernet-based setups," where traditional DDP would be infeasible due to bandwidth constraints.

A specific use case: internet-scale distributed training. Section 5 clarifies that DeMo is "designed primarily for optimization across a small number of geographically distributed compute centers." Within each data center, standard DDP with high-bandwidth interconnects can still be used; each center is treated as a "large worker." DeMo then coordinates communication between these large workers over lower-bandwidth internet links. This is a more modest and practical vision than fully decentralized training.

Contrast with Prior Sections

Where the Executive Summary introduces what DeMo achieves (85× communication reduction), this section explains why that achievement matters and what specifically was broken about prior approaches. The reader should now understand:

The communication bottleneck is not a minor inefficiency but a structural constraint on who can train large models and where.
Prior work took three tacks — compress gradients, synchronize less often, or use low-rank structure — each with identifiable failure modes (sparse update patterns, client drift, computational overhead).
DeMo's novel angle is treating the momentum buffer as both the compression target and the error feedback accumulator, exploiting its temporal smoothness and built-in memory to achieve what gradient compression alone could not.

3. Technical Approach

3.1 Reader orientation (approachable technical breakdown)

DeMo is a communication protocol and tensor compression pipeline that replaces the standard gradient all-reduce step in distributed data-parallel training. It solves the problem of unsustainable per-step communication volume by compressing each worker's locally-evolved momentum buffer using a combination of chunking, Discrete Cosine Transform, and top-k sparsification, then aggregating only the sparse coefficients across workers — reducing data movement by up to two orders of magnitude while preserving the optimization dynamics that momentum-based training relies on.

3.2 Big-picture architecture (diagram in words)

The DeMo system has four major components that operate in a repeating loop:

Per-worker momentum buffers — each GPU maintains its own local momentum tensor $M^i_t$ , updated independently using local micro-batch gradients without any gradient synchronization. This is the "decoupled" part of Decoupled Momentum Optimization.
Structured tensor compressor — a three-stage pipeline (chunking → DCT → top-k sparsification) that transforms each worker's dense momentum tensor into a sparse set of DCT coefficients, retaining only the $k$ largest-magnitude coefficients per chunk. This is where the compression happens.
Momentum subtraction as error feedback — after extracting and communicating the top-k coefficients, each worker subtracts the reconstructed (inverse-DCT) update from its local momentum buffer. This leaves residual information in the buffer to be communicated in future steps, acting as an implicit error accumulator without requiring extra GPU memory.
Server-side aggregation and parameter update — the parameter server receives sparse coefficient sets from all workers, sums them in the DCT domain, applies inverse DCT to reconstruct a global momentum estimate, then applies the base optimizer's update rule (e.g., Signum) to modify model parameters and broadcasts the update back to workers.

Information flows in a tight loop: local gradient computation → independent momentum update → per-chunk DCT and sparsification → transmission of sparse coefficients → server-side aggregation and inverse DCT → parameter update broadcast → repeat.

3.3 Roadmap for the deep dive

First, the formal problem setting and notation — what optimization problem DeMo is solving, the distributed data-parallel setup, and the standard baseline it modifies.
Second, the decoupled local momentum update — why gradients are not synchronized, how local momenta evolve independently, and why this is mathematically equivalent to synchronizing gradients but opens the door to compression.
Third, the structured tensor compression pipeline — chunking, blockwise DCT, and top-k sparsification — with a careful explanation of why DCT over random projections or identity, why chunking is necessary for computational feasibility, and how the sparsity budget $k$ controls the compression ratio.
Fourth, momentum subtraction as error feedback — how subtracting the communicated information from the momentum buffer creates an implicit error accumulator, why the subtraction coefficient $\alpha$ matters (and why $\alpha = 0.2$ beats $\alpha = 1.0$ ), and how this differs from traditional EF-SGD.
Fifth, the server-side aggregation, momentum reconstruction, and parameter update — how the sparse coefficients are combined, how the inverse DCT reconstructs the global momentum, and how the base optimizer's update rule (Signum in the main experiments) is applied.
Sixth, the theoretical convergence analysis — what assumptions DeMo's convergence proof requires, the structure of the convergence bound (Theorem 1), and what the bound tells us about how sparsification level $k$ and momentum decay $\beta$ affect the convergence rate.

3.4 Detailed, sentence-based technical breakdown

This is primarily a systems-meets-optimization paper whose core idea is that the momentum buffer in modern optimizers is simultaneously the best target for compression (because it is temporally smooth) and the best vehicle for error feedback (because it already accumulates past gradient information, so subtracting communicated values from it naturally tracks compression residuals without allocating extra memory).

Problem Setting and Distributed Optimization Formulation

DeMo addresses the standard stochastic optimization problem:

$\min_{X \in \mathcal{X}} \mathcal{L}(X) := \mathbb{E}_{\xi \sim \mathcal{D}} [L(X, \xi)]$

where $X \in \mathcal{X}$ is (possibly a series of) parameter tensors to be optimized, $L(X, \xi)$ is a sample loss function (e.g., cross-entropy), and the expectation is taken over data samples $\xi$ drawn independently from distribution $\mathcal{D}$ . The optimal value is denoted by $\mathcal{L}^* := \inf_{X \in \mathcal{X}} \mathcal{L}(X)$ , which is assumed finite and bounded from below.

What it computes: the expected per-sample loss over the data distribution. The goal of training is to find parameters $X$ that minimize this expected value, using stochastic gradients computed on mini-batches.

Why this form: this is the standard empirical risk minimization framing. The expectation over $\mathcal{D}$ captures that we care about generalization, not just fitting a specific finite dataset. The assumption that $\mathcal{L}^*$ is finite ensures the optimization problem is well-posed — the loss cannot diverge to negative infinity.

The paper operates in a distributed data-parallel setting with $I$ workers, each holding synchronized copies of parameters $X$ . At each step $t$ , worker $i$ computes a stochastic gradient:

$G^i_t = \frac{1}{n_{\text{batch}}} \sum_{n=1}^{n_{\text{batch}}} \nabla L(X_t, \xi^i_n)$

using a per-worker micro-batch of size $n_{\text{batch}}$ . The per-worker micro-batch is drawn from the worker's local data shard $\mathcal{D}_i$ .

What it computes: the average gradient over $n_{\text{batch}}$ independent samples on worker $i$ . Each worker sees a different subset of the data, so the $G^i_t$ are noisy estimates of the true full-batch gradient $\nabla \mathcal{L}(X_t)$ with variance that decreases as $1 / n_{\text{batch}}$ .

Why this form: this is the standard DDP setup. The key point is that in standard DDP, the next step would be an all-reduce to compute $\frac{1}{I} \sum_i G^i_t$ . DeMo eliminates this step entirely, which is the fundamental departure from the baseline.

Decoupled Local Momentum Updates

In standard DDP with momentum-based optimizers, gradients are synchronized before the momentum update. That is:

Each worker computes $G^i_t$
All-reduce: $\bar{G}_t = \frac{1}{I} \sum_i G^i_t$
Update global momentum: $M_t = \beta M_{t-1} + (1 - \beta) \bar{G}_t$
Apply optimizer step using $M_t$

DeMo inverts this order: momentum is updated locally first, then communicated. Specifically, on each worker $i$ :

$M^i_t = \beta M^i_{t-1} + (1 - \beta) G^i_t$

where $\beta \in (0, 1)$ is the momentum coefficient (defaulted to $\beta = 0.999$ in the main experiments, significantly larger than the standard $0.9$ used in AdamW). No all-reduce is performed on $G^i_t$ .

What it computes: a per-worker exponential moving average of past local gradients. Worker $i$ 's momentum $M^i_t$ combines its current local gradient $G^i_t$ with all its previous local gradients, weighted by powers of $\beta$ . The recursion unfolds as $M^i_t = (1 - \beta) \sum_{\tau=0}^t \beta^{t-\tau} G^i_\tau$ (ignoring initialization).

Why this form: the key theoretical justification is linearity. Because the momentum update is a linear combination of gradients, the global average of local momenta equals the global momentum that DDP would compute:

$\frac{1}{I} \sum_i M^i_t = \frac{1}{I} \sum_i \left(\beta M^i_{t-1} + (1 - \beta) G^i_t\right) = \beta \left(\frac{1}{I} \sum_i M^i_{t-1}\right) + (1 - \beta) \left(\frac{1}{I} \sum_i G^i_t\right)$

If we could synchronize the full $M^i_t$ tensors without compression, this would be mathematically identical to standard DDP. The paper's contribution is that we can approximate this synchronization with heavy compression because the momentum signal is smoother and more structured than raw gradients. The "decoupled" in DeMo refers to this independence: each worker's local momentum evolves without waiting for gradient synchronization, and communication happens after the local update, on the momentum buffer rather than on raw gradients.

The paper also notes that the momentum coefficient $\beta$ plays a dual role in DeMo: it controls both the temporal smoothing of gradients (as in standard momentum) and the rate at which compression residuals decay (via the momentum subtraction mechanism). Larger $\beta$ means more aggressive smoothing and slower residual decay, which is why the default $\beta = 0.999$ is higher than typical momentum values — the compression pipeline benefits from the extra smoothing.

Structured Tensor Compression: Chunking, DCT, and Top-k Sparsification

This is the technical core of DeMo. The compression pipeline transforms each worker's dense momentum tensor $M^i_t \in \mathbb{R}^{n_0 \times n_1 \times \cdots \times n_{d-1}}$ (where $d$ is the tensor order — typically $d = 2$ for matrix parameters in transformers) into a sparse set of coefficients that can be communicated efficiently. The pipeline has three sequential stages operating independently on chunks of the momentum tensor.

Tensor Chunking

Problem: Applying a DCT (or any orthonormal transform) to a full $n_0 \times n_1$ momentum matrix for a transformer layer would require $O(n_0 n_1 (n_0 + n_1))$ computation and $O(n_0 n_1)$ memory for the projection matrices — quadratic in the parameter dimension, which is infeasible for large models.

Solution: The paper factorizes each dimension $n_i$ as $n_i = c_i s_i$ , where $c_i$ is the number of chunks along dimension $i$ and $s_i$ is the chunk size. The full tensor is partitioned into smaller blocks:

$\mathcal{B}(M) = \{B_k \mid k \in [c_0] \times [c_1] \times \cdots \times [c_{d-1}]\}$

where each block $B_k \in \mathbb{R}^{s_0 \times s_1 \times \cdots \times s_{d-1}}$ . For the matrix case ( $d = 2$ ), this produces a grid of $c_0 \times c_1$ blocks of size $s_0 \times s_1$ .

What it computes: a partition of the momentum tensor into smaller, independently-processable sub-tensors. Each block is contiguous — $B_k$ contains elements from indices $[k_0 s_0 : (k_0+1) s_0]$ along dimension 0 and $[k_1 s_1 : (k_1+1) s_1]$ along dimension 1.

Why this form: chunking reduces the computational cost of the projection step from $O(n_0 n_1 (n_0 + n_1))$ to $O(n_0 n_1 (s_0 + s_1)) = O(n_0 n_1 (n_0 + n_1) / c_{\text{avg}})$ , which is a linear speedup in the chunk count (Section 2.1, Complexity Analysis). More importantly, it reduces the memory for storing projection matrices from $O(n_i^2)$ per dimension to $O(s_i^2)$ — since all chunks share the same projection matrices $\{P_i\}$ , the additional memory is constant and negligible. The paper defaults to $s = 64$ for all experiments, meaning each chunk is $64 \times 64$ for matrix parameters and 64-element vectors for 1D parameters like LayerNorm biases.

The paper reports results for both $s = 64$ and $s = 128$ in the extended results (Table 3, Appendix). Smaller chunk sizes reduce per-chunk computation but also reduce the total number of top-k coefficients kept (since $k$ is per-chunk), so there is a tradeoff between compute overhead and compression granularity.

Blockwise Linear Projection (DCT)

After chunking, each block $B_k$ undergoes a separable multilinear transformation:

$Q_k = \mathcal{T}(B_k; P_0, P_1, \ldots, P_{d-1})$

where $P_i \in \mathbb{R}^{s_i \times s_i}$ are orthonormal projection matrices applied along each tensor dimension. For the matrix case $d = 2$ , this reduces to the bilinear form:

$Q_k = P_0 B_k P_1^\top$

What it computes: a change of basis for each chunk. The matrix $P_0$ transforms the rows (dimension 0), and $P_1^\top$ transforms the columns (dimension 1) of the momentum chunk. The output $Q_k$ is the chunk's representation in the new basis — its coefficients tell us how much of each basis vector (row and column) is present.

Why this form: the key insight is that sparsification in the original domain produces sparse update patterns that harm training, while sparsification in a transformed domain produces dense updates after inverse transformation. Here is the critical chain of reasoning:

If you apply top-k sparsification directly to $B_k$ (which corresponds to using $P_i = I$ , the identity transform), you zero out most entries of the momentum. After the inverse transform (also identity), the parameter update is sparse — only a few weight entries get modified. This creates an uneven optimization landscape.
If you instead apply DCT, sparsify in the DCT domain, then apply inverse DCT, the reconstructed momentum $B_k^{\text{recon}} = P_0^\top \hat{Q}_k P_1$ is a dense combination of the DCT basis vectors weighted by the sparse coefficients $\hat{Q}_k$ . Even if only $k$ coefficients survive, all $s_0 \times s_1$ elements of $B_k^{\text{recon}}$ receive updates because each DCT basis vector has support across the entire chunk. This is the mechanism by which DeMo achieves "dense updates from sparse communication."

The paper considers three choices for $P_i$ :

Identity ( $P_i = I$ ): no transformation. Sparse updates in the original domain → degraded performance, as confirmed in ablation (Figure 3, left plot, red vs. blue curves).
Random orthonormal matrices: sampled freshly at each step from $\mathcal{N}(0, I)$ and orthonormalized via Gram-Schmidt. Each worker uses the same random seed (set to the step number) so projections are consistent across workers but vary across steps. This performs marginally better than DCT in ablations (Figure 3, right plot), because continuously rotating the basis prevents the top-k selection from repeatedly picking the same frequency coefficients.
Discrete Cosine Transform (DCT): a fixed, pre-computed orthonormal basis using the standard DCT-II (the same transform underlying JPEG compression). Because the DCT matrix is a specific instance of an orthonormal matrix, $P_i^\top = P_i^{-1}$ . DCT is chosen as the default in all experiments except ablations for two practical reasons: (1) it can be computed efficiently using Fast Fourier Transform algorithms, avoiding $O(s_i^3)$ matrix multiplications, and (2) it is pre-computed once before training, incurring zero per-step computation for generating the projection matrices.

The DCT has an additional conceptual advantage: real-world neural network parameters often exhibit spatial structure (neighboring weights in a layer may have correlated values), and the DCT's frequency-ordering property means energy tends to concentrate in low-frequency coefficients. This makes top-k selection in the DCT domain particularly effective at capturing the dominant structure of the momentum.

Top-k Sparsification

After the DCT, each chunk $Q_k$ has been transformed to the frequency domain. The sparsification step retains only the $k$ largest-magnitude coefficients:

$\hat{Q}_k = \text{Top-}k(Q_k, k)$

where $\|\hat{Q}_k\|_0 = k$ — exactly $k$ non-zero entries per chunk. All other entries are set to zero.

What it computes: a sparse representation of the chunk in the DCT domain. For a chunk of size $s_0 \times s_1 \times \cdots \times s_{d-1}$ containing $M = \prod_i s_i$ total elements, the compression ratio is $M / k$ . For the default $64 \times 64$ chunks, $M = 4096$ , so $k = 8$ gives a $4096 / 8 = 512\times$ compression per chunk.

Why this form: top-k sparsification is a biased but tractable compression operator. The bias — the information discarded by zeroing out small coefficients — accumulates as compression error. This is why error feedback is necessary (see next subsection). The key property that makes top-k work well with DCT is that the DCT concentrates energy — the $\ell_2$ norm of the approximation error is bounded by $\sqrt{1 - k/M} \|Q_k\|_2$ (Lemma 8.1), and this bound is worst-case when all coefficients have equal magnitude (which they don't in practice for structured signals). For smooth momentum signals, most energy concentrates in a few low-frequency coefficients, so the actual approximation error is much smaller than the worst-case bound.

The paper sweeps $k \in \{1, 2, 4, 8, 16, 32\}$ in the main experiments, with $k = 2$ already achieving better training loss than AdamW at dramatically lower communication cost (Figure 2). The communication reduction per worker per step is approximately $4096 / k$ for the default chunk size.

The sparse blocks $\hat{Q}_k$ are communicated using an All Gather operation (not a conventional All-Reduce), where each worker sends its sparse coefficient sets to the parameter server and the server aggregates them. The paper provides a detailed sparse top-k averaging procedure (Algorithm 3 in the appendix): for each chunk, the server collects non-zero coefficients from all workers, sums them coordinate-wise, and divides by the count of workers that contributed to each coefficient (accounting for the fact that different workers may select different top-k indices).

Momentum Subtraction as Implicit Error Feedback

This is the most technically subtle component of DeMo and where it departs most significantly from prior work. After extracting and communicating the sparse DCT coefficients, each worker must update its local momentum buffer to account for what was communicated. The mechanism is momentum subtraction:

$M^i_t \leftarrow M^i_t - \alpha \cdot \mathcal{B}^{-1}\left(\mathcal{T}^{-1}\left(\bar{Q}_k; P_0^{-1}, \ldots, P_{d-1}^{-1}\right)\right)$

where $\bar{Q}_k$ is the aggregated sparse coefficients from the server (after averaging across workers), $\mathcal{T}^{-1}$ is the inverse DCT (using $P_i^{-1} = P_i^\top$ since DCT matrices are orthonormal), $\mathcal{B}^{-1}$ is the unchunking operator (reassembling blocks into the full momentum tensor), and $\alpha \in (0, 1]$ is the momentum subtraction coefficient (tuned over $\{0.0, 0.1, 0.2, 0.5, 1.0\}$ , with default $\alpha = 0.2$ ).

What it computes: the worker subtracts the reconstructed global momentum contribution from its local momentum buffer. The intuition: what was communicated to the server and used to update parameters should no longer remain in the local buffer, because it has already been "spent." The residual — the difference between the worker's full local momentum and what was communicated — remains in the buffer and will be available for future top-k selection.

Why this form — and why is $\alpha = 0.2$ better than $\alpha = 1.0$ ? This is one of the paper's most interesting empirical findings (Figure 4, middle). The design rationale unfolds in layers:

$\alpha = 0$ (no subtraction) is catastrophic. Without subtraction, the momentum buffer $M^i_t$ retains all its previous values. Since the DCT basis is fixed, the top-k selection process will repeatedly select the same frequency coefficients across consecutive steps — the buffer changes slowly (due to the large $\beta = 0.999$ ), so the largest-magnitude DCT coefficients remain largely the same. This means the same information is communicated redundantly, and new gradient information that falls into lower-magnitude coefficients is never communicated. Performance degrades severely (Figure 4, middle, $\alpha = 0$ curve).
$\alpha = 1.0$ (full subtraction) is suboptimal. If you subtract the full reconstructed momentum, the buffer is zeroed out for the communicated components. This does force new gradient information to be communicated (since the buffer now only contains compression residuals plus the new gradient), but it also removes the temporal smoothing benefit of momentum. The exponential moving average is effectively reset on the communicated components at each step, losing the noise-reduction property that makes momentum useful in the first place.
$\alpha = 0.2$ strikes a balance. Partial subtraction means that communicated information decays gradually from the buffer rather than being removed immediately. Each step communicates the current dominant DCT coefficients, but a fraction $(1-\alpha)$ of the communicated signal remains in the buffer to provide temporal smoothing. The residual — the information that was not communicated (because it fell below the top-k threshold) — accumulates with full weight. This creates a two-speed dynamic: communicated information decays at rate $\alpha$ , while uncommunicated residuals accumulate at rate $\beta$ .

Comparison with traditional error feedback (EF-SGD). In EF-SGD (Karimireddy et al., 2019), a separate error accumulator $e_t$ is maintained:

$e_t = e_{t-1} + G_t - \text{compress}(G_t + e_{t-1})$

This requires additional GPU memory equal to the model size — prohibitive for billion-parameter models. DeMo's key insight is that the momentum buffer $M^i_t$ already serves this purpose. By subtracting the communicated information, the buffer naturally tracks what was omitted, without any extra memory allocation. The paper explicitly frames this: "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 $\alpha$ (how fast communicated information decays) and $\beta$ (how fast all past gradients decay in the momentum EMA). With the default $\beta = 0.999$ and $\alpha = 0.2$ , communicated information decays roughly 20% per step, while the momentum smoothing has a half-life of about $\ln(0.5) / \ln(0.999) \approx 693$ steps.

Server-Side Aggregation and Momentum Reconstruction

Once the parameter server receives sparse coefficient sets $\{Q^{i,[\ell]}_t\}$ from all workers (indexed by chunk $\ell$ ), it performs:

Aggregation in the DCT domain:

$\bar{Q}^{[\ell]}_t = \text{AvgAllGather}(\{Q^{i,[\ell]}_t\}_i)$

The paper's Algorithm 3 (Sparse Top-k Averaging) provides the detailed procedure: for each coordinate in the DCT domain, sum the coefficients from all workers that selected that coordinate, then divide by the count of contributing workers. Coordinates that no worker selected remain zero. This is a sparse-aware averaging that accounts for the fact that different workers may select different top-k indices — a worker that doesn't include a coefficient is treated as contributing zero, not as missing data.

Inverse DCT and Unchunking:

$M^{[\ell]}_t = \mathcal{T}^{-1}\left(\bar{Q}^{[\ell]}_t; P_0^\top, P_1^\top, \ldots, P_{d-1}^\top\right)$

The aggregated sparse coefficients are transformed back to the original (parameter) domain using the inverse DCT (which equals the transpose since DCT is orthonormal). The chunks are then reassembled via $\mathcal{B}^{-1}$ (concatenation) to form the full global momentum estimate $M^*_t$ .

What it computes: an approximation of the global average momentum $\frac{1}{I} \sum_i M^i_t$ that would have been computed by standard DDP, but reconstructed from heavily compressed (sparsified) DCT coefficients. The reconstruction is imperfect — information in the dropped (below-threshold) coefficients is lost — but the momentum subtraction mechanism ensures this information persists in workers' local buffers for future communication.

Base Optimizer Update:

The paper applies a transformation $\phi(\cdot)$ to the reconstructed momentum, chosen according to the base optimizer:

For SGD with momentum: $\phi(M) = M$ (identity)
For Signum (used in the main experiments): $\phi(M) = \text{sign}(M)$ (elementwise sign)
For Muon: $\phi(M) = M(M^\top M + \epsilon I)^{-1/2}$ (Newton-Schulz normalization)

The final parameter update with weight decay $\lambda$ is:

$X_{t+1} = X_t - \eta_t \left(\phi(M^*_t) + \lambda X_t\right)$

where $\eta_t$ is the learning rate at step $t$ . For the main experiments, $\phi(M) = \text{sign}(M)$ means that the parameter server only needs to broadcast the sign of the reconstructed momentum back to workers — a 1-bit per parameter communication, which is significantly cheaper than sending full-precision updates.

Why broadcast $\text{sign}(M_t)$ rather than the full momentum or update: this is a deliberate choice that further reduces download bandwidth. The sign operation is lossy but has been shown effective in distributed optimization (Bernstein et al., 2018; Seide et al., 2014). Since each worker already maintains its own momentum buffer locally, the broadcast only needs to communicate the direction of the parameter update, not its magnitude. The learning rate $\eta_t$ controls the step size centrally.

The paper notes (Section 5) that while upload bandwidth is dramatically reduced by the DCT + top-k pipeline, download bandwidth scales with the number of workers because the sign broadcast is full-parameter (albeit 1-bit per parameter). This is not unique to DeMo — it is intrinsic to all top-k sparsification approaches — and the paper positions DeMo for settings with a small number of workers (geographically distributed compute centers) where the per-worker download cost remains manageable.

Theoretical Convergence Analysis

The paper provides a convergence proof for DeMo under standard assumptions (Section 2.2 and Appendix 8). The analysis establishes that DeMo converges at rate $O(1/\sqrt{T})$ in terms of the average $\ell_1$ gradient norm, matching the rate of uncompressed SGD with momentum under the same assumptions.

Assumptions

The analysis uses three standard assumptions:

Assumption 1 (Variance): The stochastic gradients at each worker are unbiased estimates of the true gradient with bounded variance:

$\mathbb{E}\left[\|G^i(X) - \nabla \mathcal{L}(X)\|^2_F\right] \leq \frac{\sigma^2}{n_{\text{batch}}}, \quad \forall i \in [N], X \in \mathcal{X}$

where $\sigma^2$ is a finite constant and $n_{\text{batch}}$ is the per-worker micro-batch size.

Assumption 2 (L-Smoothness): The objective function $\mathcal{L}(X)$ is differentiable and L-smooth — its gradient is Lipschitz continuous with constant $L$ :

$\|\nabla \mathcal{L}(Y) - \nabla \mathcal{L}(X)\|_F \leq L \|Y - X\|_F$

This implies the standard quadratic upper bound on the function value:

$\mathcal{L}(Y) \leq \mathcal{L}(X) + \langle \nabla \mathcal{L}(X), Y - X \rangle + \frac{L}{2} \|Y - X\|^2_F$

Assumption 3 (Bounded Gradient): For any $X$ and $\xi$ , the stochastic gradient satisfies $\|\nabla L(X; \xi)\|_1 \leq R$ for some constant $R > 0$ .

Theorem Statement

Theorem 1 (Convergence of DeMo): Under the above assumptions, the sequence $\{X_t\}^T_{t=1}$ from Algorithm 4 satisfies:

$\frac{1}{T} \sum_{t=1}^T \mathbb{E}[\|\nabla \mathcal{L}(X_t)\|_1] \leq \frac{\mathbb{E}[\mathcal{L}(X_0) - \mathcal{L}(X_T)]}{T\eta} + 2LD\eta + \frac{\sigma}{\sqrt{N n_{\text{batch}}}} + R\sqrt{D}\frac{\sqrt{1 - \frac{k}{M}}}{1 - \beta}\sqrt{1 - \frac{k}{M}}\left(\beta + 2D\sqrt{\frac{2\pi}{N}}\left(1 + \sqrt{\frac{M}{k}}\right)\right)$

where $D = \prod_i n_i$ is the total number of parameters, $M = \prod_i s_i$ is the number of elements per chunk, $N$ is the number of workers, and $k$ is the sparsity budget per chunk.

With step size $\eta = \Theta(1/\sqrt{T})$ and momentum $\beta = O(1/\sqrt{T})$ , the convergence rate simplifies to:

$O\left(\frac{1}{\sqrt{T}}\right) + O\left(\frac{1}{\sqrt{N}}\right)$

What this tells us about the algorithm:

The bound has four terms. The first two ( $\frac{\mathbb{E}[\mathcal{L}(X_0) - \mathcal{L}(X_T)]}{T\eta} + 2LD\eta$ ) are the standard deterministic optimization terms — they decrease with $\eta$ (but the first term increases, creating the classic $1/\sqrt{T}$ tradeoff at optimal $\eta$ ). The third term ( $\sigma / \sqrt{N n_{\text{batch}}}$ ) is the variance reduction from using $N$ workers — more workers or larger batches reduce gradient noise. The fourth term (the large expression multiplied by $R\sqrt{D}$ ) is the compression error — this is the price DeMo pays for sparsification.
How compression affects the bound. The compression error term is proportional to $\sqrt{1 - k/M}$ . When $k = M$ (no sparsification), this term vanishes entirely, and DeMo recovers the standard convergence rate of uncompressed distributed SGD. As $k$ decreases, the compression error grows proportionally to $\sqrt{1 - k/M}$ . For $k \ll M$ (high compression), the error is $\approx \sqrt{1/M}$ per chunk times the number of chunks, which is approximately $\sqrt{D(1 - k/M)}$ — scaling with the square root of the number of uncommunicated elements.
The role of $\beta$ . The compression error is multiplied by $\beta / (1 - \beta \sqrt{1 - k/M})$ , which penalizes large momentum coefficients when sparsification is aggressive. This is intuitive: with aggressive sparsification, the momentum buffer accumulates residuals from uncommunicated information, and a large $\beta$ keeps these residuals around longer, increasing the bias in the reconstructed momentum. The convergence proof requires $\beta = O(1/\sqrt{T})$ to control this term, meaning momentum must decay over the course of training to maintain the $O(1/\sqrt{T})$ rate. In practice, the paper uses a fixed large $\beta = 0.999$ and achieves good empirical convergence, suggesting the theoretical requirement on $\beta$ is pessimistic.
The $2D\sqrt{2\pi/N}(1 + \sqrt{M/k})$ factor comes from Lemma 2 (Bias Bound for Sparse Top-k Averaging) in the appendix, which analyzes the statistical error introduced by sparse aggregation. This term captures two effects: (a) the error from having a finite number of workers $N$ estimating the top-k coordinates (the $1/\sqrt{N}$ factor), and (b) the additional variance from only seeing $k$ out of $M$ coordinates per worker per chunk (the $\sqrt{M/k}$ factor). More workers or larger $k$ reduce this error.

Why this form matters: the bound rigorously justifies the design choice of sparsifying momentum rather than raw gradients. If DeMo applied top-k sparsification directly to $G^i_t$ , the compression error would depend on the gradient variance $\sigma^2$ directly — you would lose information from the gradient's high-variance components. By sparsifying the smoothed momentum, the effective variance is reduced by the momentum EMA (the $\beta$ and $1 - \beta$ factors in the bound reflect this), making the compression more efficient for the same budget $k$ .

The proof structure (detailed in Appendix 8.2) decomposes the error between the true gradient and the reconstructed momentum into three components:

$T_1$ : The error between the true gradient and the average of workers' noisy stochastic gradients plus momentum — bounded by the variance assumption.
$T_2$ : The error between the average of workers' full DCT-domain representations and the sparse-aggregated version — bounded by Lemma 2 (the statistical bias from top-k averaging).
The combined bound plus the L-smoothness inequality yields the theorem statement through standard telescoping-sum arguments.

Summary of Design Choices and Their Justifications

Decoupled local momentum updates over gradient synchronization: mathematically equivalent due to linearity of the EMA, but places the compression target on the smoother, more structured momentum signal rather than raw high-variance gradients.
DCT over identity transform: sparsification in the DCT domain produces dense parameter updates after inverse transformation (each DCT basis vector has full support), avoiding the sparse update patterns that degrade training when sparsifying in the original domain (Figure 3, left ablation).
DCT over random projections: DCT matches random projection performance (Figure 3, right) while eliminating per-step computation for generating and orthogonalizing random matrices. DCT is pre-computed once and benefits from fast FFT-based implementations.
Chunk size $s = 64$ (default): reduces computational complexity from $O(n_i^3)$ to $O(n_i \cdot 64^2)$ compared to full-tensor DCT, and reduces projection matrix memory from $O(n_i^2)$ to $O(64^2)$ . Smaller chunks also mean the per-chunk sparsification budget $k$ operates at finer granularity.
Momentum subtraction with $\alpha = 0.2$ over $\alpha = 1.0$ or $\alpha = 0$ : $\alpha = 0$ causes redundant communication of the same information (Figure 4, middle). $\alpha = 1.0$ resets momentum smoothing on communicated components, losing the noise-reduction benefit. $\alpha = 0.2$ gradually decays communicated information while preserving momentum benefits.
Momentum coefficient $\beta = 0.999$ over standard $0.9$ : the larger coefficient provides stronger temporal smoothing, which helps the DCT concentrate energy in fewer coefficients (better compression) and slows the decay of residuals in the momentum buffer (more persistent error feedback). Ablations (Figure 4, right) confirm that larger $\beta$ values (up to $0.995$ ) improve performance.
Signum update ( $\phi(M) = \text{sign}(M)$ ) over SGD or AdamW: the sign operation reduces download bandwidth to 1-bit per parameter while preserving the update direction. Combined with the upload compression, this achieves compression in both directions, which is critical for the geographically-distributed setting where both upload and download links may be bandwidth-constrained.
All Gather + server-side aggregation over All-Reduce: All Gather is simpler to implement over standard TCP/IP networks (no ring or tree topologies required) and naturally supports the sparse coefficient format where different workers contribute different non-zero indices. This topology-agnostic design is key to DeMo's claim of enabling training over heterogeneous, non-datacenter networks.

4. Key Insights and Innovations

Innovation 1: The Momentum Buffer Is the Right Compression Target—Not Raw Gradients

The field's default assumption when attacking the communication bottleneck has been to compress raw gradients. Sparsification, quantization, and low-rank projection all take the freshly-computed per-step gradient as their input. This paper makes a fundamentally different diagnostic move: the optimizer's momentum buffer is a strictly better compression target than raw gradients, for reasons that are structural (not merely empirical) and that prior work missed because it treated the optimizer as downstream of the communication problem rather than integral to it.

Three conceptual shifts distinguish this move from incremental refinement:

First, the signal-processing argument is genuinely new. The paper identifies that momentum is a low-pass filtered version of the gradient sequence—the exponential moving average suppresses high-frequency stochastic noise, concentrating signal energy into smoother, lower-frequency components that are more compressible under transforms like the DCT. This is not a vague intuition about "structure"; it is a concrete property with a mathematical form (the EMA recursion $M_t = \beta M_{t-1} + (1-\beta) G_t$ acts as an infinite impulse response filter with cutoff frequency determined by $\beta$ ). Prior compression work (Lin et al., 2018b; Stich et al., 2018; Alistarh et al., 2017) never asked whether the optimizer state might be an easier signal to compress than the gradient. The paper's ablation with identity transform (Figure 3, left) confirms that DCT-based sparsification of momentum dramatically outperforms raw-domain sparsification—but the deeper point is that the choice to target momentum at all is the strategic insight, and the transform is the tactical mechanism that enables it.

Second, the decoupling move resolves a tension that prior work treated as unavoidable. In standard DDP, gradient synchronization and momentum updates are interleaved: gradients are all-reduced before the momentum EMA is applied. This forces the compression operator to operate on the noisier, less-structured gradient signal. By decoupling local momentum updates—letting each worker's momentum evolve independently using only local gradients—the paper relocates the synchronization point to after the smoothing has occurred, while maintaining mathematical equivalence (by linearity of the EMA) in the uncompressed limit. This is a design principle, not a hyperparameter choice: compression should happen on the most structured representation available, not the raw data signal. The field's prior failure to decouple momentum from synchronization was a path-dependency artifact—DDP was designed before compression became a first-class concern, and the default pipeline was never re-examined with compression in mind.

Third, this insight is testable and falsifiable. If the advantage came purely from the DCT transform, then applying DCT + top-k to raw gradients would match DeMo's performance. It does not (the paper's ablation with identity transform shows DCT-on-momentum beats identity-on-momentum, but the more diagnostic comparison—DCT-on-momentum vs. DCT-on-gradients—is implicit in the architecture: DeMo's decoupled momentum update by construction applies DCT to the EMA-smoothed signal). The theoretical analysis (Theorem 1) formalizes why this matters: the compression error term depends on the momentum buffer's $\ell_2$ norm (via Lemma 1), which is bounded by $R\sqrt{1-k/M} / (1 - \beta\sqrt{1-k/M})$ , not directly on the gradient variance $\sigma^2$ . The $\beta$ factor in the denominator represents the smoothing benefit—as $\beta \to 1$ , the momentum signal becomes smoother and the bound tightens (at the cost of slower residual adaptation).

Significance beyond performance: This reframes how we should think about optimizer states in distributed training. They are not just for convergence acceleration—they are a communication substrate with properties (temporal smoothness, built-in memory) that make them superior to raw gradients. This may generalize beyond momentum: second-moment estimates in Adam, preconditioners in Shampoo, or even running statistics in BatchNorm could serve as compression targets with similar benefits. The paper opens a design space where what gets communicated is decoupled from what gets computed locally, with the optimizer's internal state as the communication primitive.

Innovation 2: The Momentum Buffer as a Zero-Memory-Cost Error Feedback Accumulator

Error feedback is the standard remedy for the bias introduced by lossy gradient compression: you maintain a separate accumulator $e_t$ that tracks the compression residual and adds it to the next gradient, ensuring that information is not permanently lost but only delayed (Karimireddy et al., 2019; Seide et al., 2014). The problem, which the paper identifies clearly, is that this accumulator costs GPU memory equal to the model size—for a billion-parameter model, that is gigabytes of additional memory, which competes with activations, optimizer states, and the model itself for scarce GPU RAM.

The paper's insight is one of those rare "why didn't anyone notice this before" observations: the momentum buffer already implements a form of error accumulation. The momentum EMA $M_t = \beta M_{t-1} + (1-\beta) G_t$ already carries forward information from past steps. If you subtract the communicated information from this buffer, the residual naturally persists—without allocating a single extra byte. The momentum buffer is the error feedback accumulator; you just need to modify what gets removed from it.

This is a genuinely elegant unification. Prior work treated momentum and error feedback as orthogonal mechanisms: momentum accelerates convergence by smoothing gradients, error feedback compensates for compression bias by accumulating residuals. DeMo recognizes that the momentum buffer's accumulation property—the very thing that makes it useful for optimization—also makes it a natural vehicle for tracking compression residuals. The "momentum subtraction" operation ( $M^i_t \leftarrow M^i_t - \alpha \cdot \text{reconstructed}$ ) is the only modification needed to repurpose an existing buffer for error feedback.

The $\alpha < 1$ finding deepens the insight. If the momentum buffer were purely an error accumulator, you would want $\alpha = 1$ (full subtraction)—communicate everything and leave only the residual. The paper's ablation (Figure 4, middle) shows that $\alpha = 0.2$ substantially outperforms $\alpha = 1.0$ . Why? Because the momentum buffer is simultaneously an error accumulator and a temporal smoother. Full subtraction resets the smoothing on communicated components at each step, losing the noise-reduction benefit that makes momentum useful. Partial subtraction ( $\alpha = 0.2$ ) creates a two-speed dynamic: communicated information decays gradually (preserving smoothing), while uncommunicated residuals persist with full weight (preserving the error feedback function). This is a tension that explicit error feedback methods do not face because they separate the two functions into different buffers—but they pay a memory cost DeMo avoids.

Comparison to prior work: EF-SGD (Karimireddy et al., 2019) requires doubling the gradient memory. Deep Gradient Compression (Lin et al., 2018b) uses momentum correction masking to address staleness but still maintains separate error accumulators. DeMo is the first method to recognize that the momentum buffer can serve both roles simultaneously, and the first to empirically characterize the tradeoff via the $\alpha$ parameter. This is a fundamental advance in compression algorithm design: before adding a new buffer, check whether an existing one already does the job.

Significance beyond memory savings: The unification has theoretical implications. The convergence proof (Theorem 1) bounds the compression error in terms of momentum norm, not a separate accumulator norm—the $\beta$ and $\alpha$ parameters jointly govern both smoothing and residual tracking. This creates a cleaner theoretical framework than the separate-buffer approach, where the interaction between momentum dynamics and error accumulation is more complex to analyze. The paper's $O(1/\sqrt{T})$ rate with $\beta = O(1/\sqrt{T})$ is not surprising as a convergence result, but the way the compression error term factors through the momentum dynamics (via Lemma 1's bound on $\mathbb{E}[\|M^i_t\|_2]$ ) is structurally novel.

Innovation 3: Sparsifying in a Fixed Transformed Domain as a Dense-Update Strategy

Applying top-k sparsification directly to parameter-space tensors produces sparse update patterns that harm training—some weights get updated frequently while others stagnate, creating an uneven optimization trajectory. The standard diagnosis is that sparse updates are inherently problematic, and the standard remedy is either to use error feedback to eventually communicate the dropped entries (Karimireddy et al., 2019) or to use low-rank projections that summarize the full gradient (Vogels et al., 2019; Zhao et al., 2024).

DeMo reveals a third path: sparsify in a transformed domain where each basis vector has global spatial support, producing dense parameter updates after inverse transformation even from extreme sparsity. The conceptual move is recognizing that the "sparse updates are bad" diagnosis applies to the domain of application, not to sparsification per se. If you sparsify in a basis where every coefficient you keep contributes to every parameter you update, the effective update is dense—you are not starving any parameters of gradient signal; you are just approximating the full update as a linear combination of a few dominant basis vectors.

The DCT is the specific instantiation of this idea, but the principle is general. For a chunk of size $64 \times 64$ , keeping $k = 8$ DCT coefficients means you are storing only 8/4096 ≈ 0.2% of the data, but after inverse DCT, all 4096 elements of the reconstructed momentum chunk receive non-zero updates. Each kept DCT coefficient corresponds to a spatial frequency pattern that spans the entire chunk—discarding coefficients means you lose fine-grained spatial detail, not entire spatial regions. This is fundamentally different from sparsifying in the original domain, where zeroing an entry means that specific weight receives no update.

Why this is more than "just use a transform": The field has used transforms for compression before (e.g., JPEG's DCT, wavelet compression for audio), but applying this to gradient communication faces two challenges that the paper addresses. First, the transform must be orthonormal so that the inverse is the transpose (computationally cheap) and the $\ell_2$ norm is preserved (so top-k in the transform domain minimizes $\ell_2$ reconstruction error in the original domain). Second, the transform must be fast enough to not dominate the computation budget. DCT satisfies both: it is orthonormal, and it has $O(N \log N)$ implementations via FFT algorithms. Random orthonormal matrices are marginally better (Figure 3, right) but cost $O(s_i^3)$ per chunk to generate and apply, which is prohibitive at scale.

The paper's ablation comparing DCT, identity, and random projections (Figure 3) is more diagnostic than it might appear. It decomposes the benefit of the transform into two components: (a) the dense-update effect (present in both random and DCT, absent in identity) and (b) the energy-concentration effect of DCT specifically (DCT concentrates energy in low frequencies for smooth signals, random projections do not). The fact that random projections slightly outperform DCT suggests that continuously rotating the basis (preventing the top-k selection from repeatedly picking the same frequency indices) provides an additional robustness benefit, which DCT sacrifices for computational efficiency. This is a deliberate tradeoff, not an oversight.

Connection to low-rank methods: The paper notes (Section 4) that top-k in the DCT domain makes the communicated update rank-k per chunk. This is not a low-rank approximation in the traditional sense—the DCT basis is fixed and pre-chosen, not learned from data—but it shares the structural property that the communicated information lives in a k-dimensional subspace. GaLore (Zhao et al., 2024) learns the low-rank subspace adaptively via SVD, which is more flexible but more expensive. DeMo's fixed transform is a bet that a generic, pre-chosen basis (DCT) is good enough for momentum tensors across different layers, tasks, and training stages—and the empirical results suggest this bet largely pays off, though adaptively learned transforms might close the small gap to random projections.

Significance beyond this paper: The "dense updates from sparse communication" principle suggests a design pattern: choose a transform whose basis vectors have the spatial properties you want in the update, then sparsify aggressively. For convolutional layers, the DCT's frequency-ordering property (low-frequency = smooth spatial patterns) is a good match because weight updates often exhibit spatial smoothness. For other architectures (attention, MLP-Mixer), different transforms might be optimal. The paper opens this as a design axis rather than treating transform choice as a fixed hyperparameter.

Innovation 4: Frequent Compressed Synchronization as an Alternative to Infrequent Full Synchronization

The dominant philosophy for reducing communication in distributed training, as represented by Local SGD (Stich, 2018), Federated Averaging (McMahan et al., 2017), and most recently DiLoCo (Douillard et al., 2023), has been: communicate less often, but communicate everything. Workers perform multiple local steps independently, then synchronize full model states (or full gradients). The amortized communication reduction is $H \times$ , where $H$ is the synchronization interval.

DeMo offers a competing philosophy: communicate every step, but make each message nearly negligible. Instead of $H$ local steps between full synchronizations, DeMo performs one local step between heavily compressed synchronizations. The comparison is not just empirical (Figure 6 shows DeMo outperforming DiLoCo at matched compression ratios); it is a fundamental design tradeoff with different failure modes.

The diagnostic advantage of frequent communication. The primary failure mode of infrequent synchronization is client drift: workers' local parameters diverge during the $H$ independent steps, and the averaged model after synchronization may not lie in a favorable region of the loss landscape (Karimireddy et al., 2020). This is not just a theoretical concern—it manifests as training instabilities, especially with adaptive optimizers where per-parameter learning rates can amplify divergence (Nguyen et al., 2025; Liang et al., 2024a). DeMo avoids client drift entirely because synchronization happens every step. The price it pays for this is a different failure mode: compression bias from the top-k sparsification, which accumulates in the momentum buffer and can cause stale gradient information to persist. The paper's momentum subtraction mechanism ( $\alpha = 0.2$ ) and its convergence analysis explicitly address this bias.

The question of which failure mode is less harmful—drift from infrequent sync or bias from aggressive compression—is not resolved by this paper, but the paper provides strong evidence that at the compression ratios achievable with momentum DCT + top-k (up to 4096× per chunk for $k = 1$ ), the bias from compression is the more manageable problem. DeMo's training curves (Figure 2) are smooth and monotonic, without the periodic instability that can characterize DiLoCo runs. The theoretical convergence rate ( $O(1/\sqrt{T})$ ) matches uncompressed distributed SGD, suggesting that the compression bias can be controlled with appropriate $\beta$ and $\alpha$ scheduling.

This is not an argument against DiLoCo. The paper explicitly positions DeMo for "a small number of geographically distributed compute centers" where the inter-center link is bandwidth-constrained but latency-tolerant. DiLoCo is better suited for settings where latency is high (making per-step communication infeasible regardless of message size) or where the number of workers is very large (making the server-side aggregation of sparse coefficients a bottleneck). The contribution is not "DeMo beats DiLoCo" but rather the identification of a design point that was previously unexplored: frequent, aggressively compressed momentum synchronization as a viable and sometimes superior alternative to infrequent full synchronization. The field had polarized between "communicate every step uncompressed" (DDP) and "communicate rarely" (Local SGD/DiLoCo). DeMo opens the middle ground.

Evidence for the design point claim: Figure 6 plots data-receiving compression ratio vs. validation perplexity for DeMo and DiLoCo. At matched compression ratios, DeMo consistently achieves lower perplexity. The breakdown by $k$ vs. synchronization interval shows that DeMo with modest sparsification ( $k = 32$ , ~128× compression) outperforms DiLoCo even at aggressive communication intervals. This is non-obvious: one might expect that communicating full information less frequently would preserve more signal than communicating partial information frequently, but the empirical evidence suggests the opposite—presumably because client drift degrades the quality of the information that is eventually communicated in DiLoCo.

Significance for the training infrastructure landscape: If frequent compressed communication is viable, then low-bandwidth, high-latency interconnects (standard Ethernet, cross-datacenter links) become usable for synchronous distributed training, not just for the infrequent parameter averaging that DiLoCo-type methods require. This expands the set of hardware configurations that can support large-scale training, directly addressing the democratization goal the paper states in its introduction.

5. Experimental Analysis

Evaluation Methodology

Dataset. All experiments use the Dolma v1.5 corpus for pretraining. The paper trains both OLMo-300M and OLMo-1B on 100 billion tokens—significantly beyond the Chinchilla token budget of 20 tokens per parameter—to study optimizer behavior near convergence. Ablation studies use the 300M model trained under the 20 tokens-per-parameter rule due to computational constraints (Section 3.1).
Base model(s). The paper evaluates DeMo on Transformer-based decoder-only language models from the OLMo family (Groeneveld et al., 2024) at two scales: OLMo-300M (320 million non-embedding parameters) and OLMo-1B (1.18 billion non-embedding parameters). Full model specifications are provided in the Appendix. These models are chosen because OLMo is a reproducible, open-weight training framework that enables controlled comparisons of optimizer behavior at meaningful scale without the confounding variables of proprietary training pipelines.
Metrics. Three categories of metrics are reported:
- Training loss (cross-entropy loss in nats) throughout pretraining, with curves showing loss vs. training tokens (Figure 2) and final training loss values (Table 3).
- Validation perplexity (log-scale) plotted against data transmitted per step to capture the efficiency-accuracy tradeoff (Figure 4, left).
- Zero-shot downstream accuracy on three widely used benchmarks: HellaSwag (normalized accuracy), ARC-Easy (accuracy), and PIQA (normalized accuracy). These are standard evaluations from the OLMo framework and are reported after full 100B-token pretraining (Table 1).
- Communication volume measured in MB per GPU per training step, computed from the size of transmitted sparse coefficient sets plus the broadcast of sign(Mt) for the parameter update. This is the primary efficiency metric against which accuracy is traded off.
Baselines. The paper compares against the following methods:
- AdamW-DDP: standard AdamW optimizer with default OLMo hyperparameters (β1=0.9, weight decay λ=0.1), using full-precision gradient all-reduce via PyTorch DDP. This is the primary baseline throughout all experiments.
- Muon-DDP (Jordan et al., 2024): an alternative optimizer using the Muon update rule with full-precision DDP. Evaluated only in the additional experiments section (Section 3.5, Figure 5) with recommended momentum value 0.95 from Liu et al. (2025).
- DiLoCo (Douillard et al., 2023): a communication-efficient method that alternates local updates with periodic global synchronization. Evaluated with outer-loop learning rate tuned over {1.0, 0.7, 0.5, 0.3, 0.1} and varying communication frequency to target different compression ratios (Section 3.5, Figure 6).
- PowerSGD (Vogels et al., 2019): low-rank gradient compression applied during gradient aggregation. Evaluated using the official PyTorch DDP communication hook with tuned communication rank (selected 32 after sweeping {2, 4, 8, 16, 32, 64}), 1000 warmup steps, and tuned learning rate from the same grid as AdamW (Section 3.5, Figure 5).
Generation budget / compute accounting. The paper measures communication efficiency in terms of per-GPU data transmitted (MB/step), not FLOPs or wall-clock time. For DeMo, this includes the sparse DCT coefficients uploaded from each worker to the parameter server. For AdamW-DDP, this is the full-precision gradient all-reduce volume. All experiments use 64 NVIDIA H100 GPUs with a global batch size of 2048, sequence length 2048 tokens, and 4 gradient accumulation steps (effective per-GPU batch size of 8). Training runs use linear warmup followed by cosine decay learning rate schedules, with the best recommended learning rates from Groeneveld et al. (2024) for the full-training runs in Figure 2.
Cross-validation / statistical protocol. For the full 100B-token training runs, hyperparameters (learning rate, β2 for AdamW, α and β for DeMo) are tuned on reduced-budget runs (1B tokens for AdamW β2 tuning; 20 tokens-per-parameter for ablation experiments). For the additional experiments in Section 3.5, extensive learning rate sweeps are performed over {1.5×10^(-3), 1.2×10^(-3), 10^(-3), 8×10^(-4), 6×10^(-4), 5×10^(-4), 3×10^(-4), 2×10^(-4)} for AdamW-DDP and Muon-DDP to ensure fair comparison. The paper does not report confidence intervals or multiple random seeds for the main training runs.

Main Quantitative Results

Convergence Behavior Across Sparsification Levels

The central result is that DeMo achieves training convergence competitive with or exceeding AdamW-DDP at dramatically reduced communication cost. Figure 2 shows training loss curves for both OLMo-300M and OLMo-1B across sparsification levels k ∈ {1, 2, 4, 8, 16, 32} compared against the AdamW baseline.

Headline finding: A sparsification level of just k = 2 is sufficient to achieve better training loss than AdamW on both model scales. As stated in Section 3.2: "a sparsification of just k = 2 is sufficient to achieve better training performance than AdamW." Increasing k further provides only marginal gains.

For OLMo-300M (Figure 2, left), the training loss curves for k ≥ 2 are nearly indistinguishable from each other and all lie below (better than) the AdamW curve throughout the 100B-token training horizon. The k = 1 curve initially lags but catches up by approximately 60B tokens and finishes marginally above the AdamW baseline. For OLMo-1B (Figure 2, right), the pattern is similar but more compressed: the k ≥ 4 curves track closely together below AdamW, while k = 2 shows a small gap and k = 1 shows a more noticeable gap, though still competitive.

Communication reduction quantification (Table 1, Section 3.2): For the 300M model, AdamW-DDP transmits 636.9 MB per GPU per step. DeMo with k = 8 reduces this to 7.49 MB/step—an 85× reduction—while achieving equal or better downstream accuracy (HellaSwag 0.38 vs. 0.35, ARC-Easy 0.47 vs. 0.46, PIQA 0.67 vs. 0.65). At the most aggressive sparsification (k = 1), DeMo transmits only 0.93 MB/step (a 685× reduction) with a minor accuracy tradeoff: HellaSwag drops from 0.35 to 0.35 (unchanged), ARC-Easy from 0.46 to 0.45, PIQA from 0.65 to 0.65 (unchanged).

For the 1B model, the pattern scales favorably: AdamW-DDP transmits 2416.6 MB/step, while DeMo with k = 16 transmits 55.16 MB/step (a 44× reduction) while achieving higher HellaSwag (0.47 vs. 0.43) and PIQA (0.70 vs. 0.68) scores than AdamW. The complete results in Appendix Table 3 confirm this trend across all k values and two chunk sizes (s = 64, s = 128).

Compute Efficiency: Perplexity vs. Communication Tradeoff

Figure 4 (left) plots final validation perplexity against per-GPU data transmitted per step (log scale) for both model scales. Each DeMo point is annotated with its sparsification level k.

Headline findings from this analysis:

DeMo points form a clear Pareto frontier dominating the AdamW baseline: all DeMo configurations achieve lower (better) perplexity while transmitting 1-3 orders of magnitude less data. For the 300M model, AdamW sits at approximately 2.98 training loss (Table 3) and 636.9 MB/step, while DeMo k = 8 achieves 2.88 training loss at 7.49 MB/step—an 85× data reduction with a 0.10 improvement in loss.
Diminishing returns from increasing k after k = 8. The perplexity improvement from k = 8 to k = 32 is minimal (2.88 → 2.87 for 300M, s = 64) while communication cost quadruples (7.49 → 29.9 MB/step). This suggests an efficient operating point around k = 8 for the default chunk size.
Interaction between chunk size and k. Table 3 shows that for the same sparsification budget k, larger chunks (s = 128) provide lower communication cost (since fewer chunks → fewer total top-k coefficients transmitted) but slightly worse training loss and downstream accuracy. For instance, 300M DeMo with s = 128, k = 32 transmits 7.49 MB/step with training loss 2.88, while s = 64, k = 8 also transmits 7.49 MB/step but achieves better loss (2.88—actually 2.88 vs. 2.88, identical in this case, but other comparisons show small differences). The tradeoff is that larger chunks reduce the granularity of sparsification, potentially discarding more fine-grained spatial information.
Model scale scaling: The 1B model shows the same qualitative pattern as the 300M model but with larger absolute communication volumes (the 1B AdamW-DDP baseline transmits 2416.6 MB/step vs. 636.9 MB/step for 300M). The relative reduction factors are comparable across scales, suggesting the compression benefits scale linearly with model size.

Downstream Evaluation

Table 1 reports zero-shot accuracy on HellaSwag, ARC-Easy, and PIQA after 100B-token pretraining.

Headline findings:

DeMo matches or exceeds AdamW across all benchmarks at k ≥ 4 for 300M. At k = 16, DeMo achieves HellaSwag 0.38 (vs. AdamW 0.35), ARC-Easy 0.50 (vs. 0.46), and PIQA 0.67 (vs. 0.65). These improvements are modest but consistent.
At the 1B scale, DeMo with k ≥ 4 again matches or exceeds AdamW across all benchmarks. The most notable improvement is HellaSwag: DeMo k = 32 achieves 0.48 vs. AdamW 0.43, a 5-point improvement. PIQA at k = 4, 16, and 32 all achieve 0.70 vs. AdamW 0.68.
Degradation at extreme sparsification (k = 1) is visible but small. For the 300M model, DeMo k = 1 achieves HellaSwag 0.35 (same as AdamW's 0.35), ARC-Easy 0.45 (vs. 0.46 for AdamW), and PIQA 0.65 (same as AdamW's 0.65). For the 1B model, k = 1 shows a more noticeable gap: HellaSwag 0.41 vs. 0.43, PIQA 0.69 vs. 0.68. This demonstrates that even extreme compression (transmitting only 0.93 MB/step for 300M vs. AdamW's 636.9 MB/step) preserves model quality to a surprising degree.
The downstream evaluation is limited to three relatively simple benchmarks. The paper does not evaluate on more challenging reasoning tasks (e.g., MMLU, GSM8K, HumanEval) or on generation quality metrics (perplexity on held-out text beyond training loss). This is a genuine limitation—the claim of "comparable accuracy" is supported only for the specific benchmarks tested, and it is possible that aggressive momentum compression introduces subtle degradations that only appear on more demanding evaluations.

Comparisons with Communication-Efficient Baselines (Section 3.5)

The additional experiments provide the most direct comparisons against prior communication-efficient methods.

AdamW-DDP and Muon-DDP with extensive tuning (Figure 5): Under a reduced token budget (20 tokens per parameter) with hyperparameter sweeps over learning rate and β2, Muon-DDP achieves better training loss than AdamW-DDP, consistent with prior findings (Liu et al., 2025). DeMo underperforms these baselines at equal update counts but does so at dramatically lower communication cost—this is the expected tradeoff. The figure demonstrates that DeMo's convergence is competitive with the tuned baselines, not that it matches them per-step at zero compression.

DiLoCo comparison (Figure 6): The paper compares DeMo and DiLoCo by plotting data-receiving compression ratio against final validation perplexity. For DeMo with 64×64 chunks and top-k = 32, the data-transmission compression is 128× and data-receiving compression is 16× (ignoring the marginal contribution of LayerNorm parameters). At comparable compression levels, "DeMo consistently yields lower perplexity than DiLoCo." Training loss curves for DiLoCo with communication frequency 10 are included in Figure 5, where DeMo again shows superior optimization performance.

The key methodological detail: DiLoCo's compression ratio is controlled by the synchronization interval H—communicating every H steps gives an H× reduction. DeMo's compression ratio is controlled by k—the fraction k/M of DCT coefficients retained per chunk. The paper sweeps both parameters and compares at matched compression ratios, finding DeMo favorable. This is the evidence for the paper's philosophical claim that frequent compressed communication (DeMo) can outperform infrequent full communication (DiLoCo) at the same total bandwidth.

PowerSGD comparison (Figure 5): PowerSGD with rank 32 (providing approximately 26× compression) and 1000 warmup steps achieves performance similar to DiLoCo but remains slightly below DeMo. The paper notes that PowerSGD is essentially compressed SGD; they adopted an Adam-style update rule using the compressed gradients to make the comparison fairer, since bare SGD performs poorly for language model training. The key difference: PowerSGD compresses gradients via low-rank projection, while DeMo compresses momentum via DCT + top-k sparsification. The paper's result suggests the latter approach is more effective at preserving optimization quality under aggressive compression.

Ablation Studies and Robustness Checks

All ablation studies use the 300M model trained under the 20 tokens-per-parameter rule.

Impact of momentum subtraction coefficient α: Figure 4 (middle) shows training loss curves for α ∈ {0.0, 0.1, 0.2, 0.5, 1.0}. The finding is stark: α = 0 (no subtraction) is catastrophic—loss degrades severely because the same DCT coefficients are repeatedly selected and communicated across steps. α = 1.0 (full subtraction) is functional but suboptimal—loss is higher than partial subtraction values. The best performance is achieved at α = 0.2, confirming that partial subtraction strikes the right balance between forcing novelty in communicated information and preserving momentum smoothing. The paper's explanation (Section 3.4): "using a smaller value (α = 0.2) to gradually evolve the top-k elements and partially decay communicated values over time further improves performance." This ablation is critical because it validates one of DeMo's key design claims—that the momentum buffer can serve as an error feedback accumulator without requiring α = 1.0 (full subtraction), and that the optimal subtraction rate is not obvious a priori.

Choice of linear transformation (DCT vs. identity vs. random): Figure 3 presents two comparisons. The left plot compares DCT against identity mapping (Pi = I, i.e., no transformation, top-k applied directly in parameter space) for k ∈ {8, 16, 32}. DCT clearly outperforms identity at all sparsification levels, with the gap widening as k decreases (more aggressive sparsification). This validates the paper's core mechanism claim: sparsifying in a transformed domain produces dense parameter updates that avoid the degradation from sparse update patterns.

The right plot compares DCT against random orthonormal projections for the same k values. Random projections perform marginally better than DCT—the loss curves for random are slightly lower at all k values. The paper interprets this as: "random projection is arguably the most intuitive choice, as it continuously rotates and changes the momentum subspace perspective," preventing top-k selection from repeatedly picking the same frequency coefficients. However, the paper defaults to DCT because it is computationally cheaper (pre-computed once, fast FFT-based implementation) and the performance gap is small. This is a deliberate engineering tradeoff, not a claim that DCT is theoretically optimal.

Momentum coefficient β: Figure 4 (right) ablates β ∈ {0.95, 0.98, 0.99, 0.995, 0.999} with fixed α = 0.2, k = 32, and chunk size 64. Performance remains stable across this range, with larger values (0.995, 0.999) outperforming smaller ones. The best results are achieved at β = 0.995. This is non-obvious because standard momentum-based optimizers typically use β = 0.9 (AdamW's default); the paper's finding that much larger β values improve performance under compression is consistent with the intuition that stronger temporal smoothing helps concentrate energy in fewer DCT coefficients and slows residual decay in the momentum buffer.

Chunk size (s = 64 vs. s = 128): The extended results in Table 3 provide a systematic comparison across chunk sizes and sparsification levels. For the same k, larger chunks (s = 128) transmit 4× less data (since there are 4× fewer chunks) but generally show worse training loss and downstream accuracy. For example, 300M DeMo with s = 64, k = 8 achieves training loss 2.88 and HellaSwag 0.38; with s = 128, k = 8 achieves training loss 2.93 and HellaSwag 0.36. The tradeoff is intuitive: larger chunks mean the top-k selection discards a larger fraction of spatial information per chunk (k/16384 for s = 128 vs. k/4096 for s = 64 at the same k), but the total number of communicated coefficients is smaller (fewer chunks × k). The paper's default of s = 64 represents a middle ground.

Extended AdamW hyperparameter sweep (Figure 5): The paper conducts an unusually thorough sweep for the AdamW-DDP baseline under reduced token budget, searching learning rates in {1.5×10^(-3), 1.2×10^(-3), 10^(-3), 8×10^(-4), 6×10^(-4), 5×10^(-4), 3×10^(-4), 2×10^(-4)} and β2 in {0.95, 0.98, 0.99, 0.995, 0.999}. This is important for fairness—the default AdamW hyperparameters may not be optimal for the specific training setup, and a poorly-tuned baseline would inflate DeMo's apparent advantage. The swept results in Figure 5 show that DeMo remains competitive with the best-tuned AdamW configuration, not just the default one.

Negative results: The paper does not explicitly report significant negative results in the ablations beyond the expected degradation patterns (α = 0, identity transform, small β). One notable absence: there is no ablation on the effect of removing momentum entirely and applying DCT + top-k directly to gradients—this would directly test the paper's central claim that momentum is a better compression target than raw gradients. The identity transform ablation tests DCT vs. no-DCT on momentum, but does not isolate whether momentum itself (vs. raw gradients) provides the benefit.

Critical Assessment

Does the Paper Demonstrate That DeMo Reduces Communication by "Up to 85×" While Maintaining Convergence?

Yes, with the qualification that "maintaining convergence" means comparable training loss and zero-shot accuracy on the specific benchmarks tested, not identical optimization trajectories.

The 85× figure is directly measured: for the 300M model, AdamW-DDP transmits 636.9 MB/step, and DeMo with k = 8 transmits 7.49 MB/step (Table 1). The ratio 636.9 / 7.49 ≈ 85×. At this sparsification level, DeMo achieves better training loss (2.88 vs. 2.98, Table 3), better HellaSwag (0.38 vs. 0.35), better ARC-Easy (0.47 vs. 0.46), and better PIQA (0.67 vs. 0.65) than AdamW. The claim is well-supported at this specific operating point.

However, the claim is most robust at moderate sparsification (k ≥ 4 for 300M, k ≥ 8 for 1B). At k = 1 (685× compression for 300M), downstream accuracy begins to show small but visible degradation relative to AdamW, though the gap is surprisingly small. The paper does not claim 85× is the maximum possible; the text says "up to two orders of magnitude" (100×) in the abstract, which is consistent with the reported numbers. The figure of 85× specifically pairs with "achieving comparable loss and accuracy," which holds at k = 8 but not necessarily at k = 1 or k = 2.

A limitation: The paper only evaluates downstream accuracy on three relatively simple benchmarks (HellaSwag, ARC-Easy, PIQA). These are standard in the OLMo framework but do not constitute a comprehensive evaluation of model quality. It is possible that aggressive momentum compression introduces subtle degradations in reasoning ability, factual recall, or generation quality that these benchmarks do not capture. A more complete evaluation would include benchmarks like MMLU, GSM8K, HumanEval, or at minimum perplexity on held-out corpus data (beyond the training loss curves shown).

Does the Paper Demonstrate That DeMo Scales to 1B Parameters?

Yes, but only at a single larger scale, and the evaluation at 1B is less thorough than at 300M.

The 1B-parameter experiments (Figure 2, right; Table 1; Table 3) replicate the key findings from the 300M experiments: DeMo with moderate sparsification (k ≥ 8) matches or exceeds AdamW in training loss and downstream accuracy while reducing communication by 22-44× (depending on k). The pattern of diminishing returns from increasing k beyond 8-16 is also replicated. The scaling behavior is favorable—the compression ratio is maintained because the per-chunk sparsification mechanism scales linearly with model size (more chunks, same k per chunk).

What is missing: The paper does not report wall-clock time measurements or throughput (tokens/second) for either model scale. The communication reduction factors are measured in MB/step, which is a meaningful metric but does not directly translate to training speed improvements. If the DCT/IDCT computation or the sparse aggregation introduces significant per-step overhead, the actual training speedup could be smaller than the communication reduction suggests. The paper mentions computational complexity in Section 2.1 (reducing from O(N^3) to O(N^3/C) via chunking), and notes DCT can use FFT-based fast implementations, but no empirical timing measurements are reported. This is a significant gap—the practical value of communication compression depends on whether it translates to faster training, not just fewer bytes transmitted.

Additionally, the 1B model is still small by modern standards. State-of-the-art models trained with DDP range from 7B to 70B parameters and beyond. The paper does not demonstrate that DeMo's benefits continue to scale to these sizes, where communication overhead is most acute. The chunk-based DCT and top-k mechanism should scale linearly in principle, but potential issues like the server-side sparse aggregation becoming a bottleneck at large worker counts are not explored.

Does the Paper Demonstrate That DeMo Is Topology-Agnostic and Suitable for Geographically Distributed Training?

This is claimed but not directly tested.

The paper states in the abstract that DeMo "enables training across multi-datacenter or Ethernet-based setups" and in Section 5 that it is "designed primarily for optimization across a small number of geographically distributed compute centers." However, all experiments are conducted on 64 NVIDIA H100 GPUs—a standard single-cluster configuration with high-bandwidth interconnects (presumably NVLink/InfiniBand). The paper does not report any experiments over actual wide-area networks, over Ethernet-only setups, or with simulated network latency/bandwidth constraints. The claim of topology-agnosticism is based on architectural properties of DeMo (All Gather instead of All-Reduce, compressibility over low-bandwidth links) but these properties are not empirically validated under the conditions the paper claims to enable.

What would strengthen this claim: Experiments with simulated bandwidth caps, latency injection, or actual cross-datacenter training runs showing that DeMo maintains convergence under conditions where standard DDP fails or becomes impractically slow. The paper's positioning of DeMo as enabling "Internet"-based distributed training is plausible given the communication reductions achieved, but it remains an architectural argument rather than an experimentally demonstrated capability.

Does the Paper Demonstrate That DeMo Outperforms Prior Communication-Efficient Methods?

Partially. The comparisons with DiLoCo and PowerSGD are informative but limited in scope.

The DiLoCo comparison (Figure 6) shows DeMo achieving lower perplexity at matched compression ratios, which supports the paper's philosophical claim that frequent compressed communication can outperform infrequent full communication. However, the comparison is on a single model scale (implied to be 300M based on the 20 tokens-per-parameter budget), uses only the Signum update rule for DeMo, and sweeps only the synchronization interval for DiLoCo. A more comprehensive comparison would include: (1) DiLoCo at the 1B scale, (2) DiLoCo with different base optimizers, (3) DiLoCo with gradient compression applied in addition to infrequent synchronization (since the techniques could be complementary), and (4) wall-clock time measurements for both methods at the same hardware configuration.

The PowerSGD comparison (Figure 5) is similarly limited. PowerSGD is evaluated with a single rank (32) after sweeping, but the sweep did not include combining PowerSGD with momentum-based optimizers in the way that DeMo does. PowerSGD compresses gradients; DeMo compresses momentum. The paper's finding that DeMo outperforms PowerSGD is consistent with its central thesis, but the comparison does not isolate whether the advantage comes from targeting momentum vs. gradients or from the specific compression mechanism (DCT + top-k vs. low-rank projection).

Missing baselines: The paper does not compare against several relevant methods from the gradient compression literature:

Deep Gradient Compression (Lin et al., 2018b) with momentum correction masking—this would be the most direct comparison, as it also uses momentum and sparsification, though via a different mechanism.
1-bit Adam or 1-bit SGD with error feedback (Seide et al., 2014; Tang et al., 2021)—quantization-based approaches that achieve compression ratios in the 32× range.
GaLore (Zhao et al., 2024)—a recent low-rank method that also reduces communication, which would test whether DeMo's fixed-transform approach is competitive with adaptive low-rank methods at matched compression ratios.

Does the Ablation on Linear Transformations Support the Claim That DCT Enables Dense Updates?

Yes, but the ablation could be more diagnostic.

Figure 3 (left) clearly shows DCT outperforming identity mapping, confirming that some transform is beneficial. Figure 3 (right) shows random projections marginally outperforming DCT, confirming that the benefit is not specific to DCT's frequency-ordering property—it is the dense-update property of any orthonormal transform. The paper's interpretation is reasonable: (1) identity mapping produces sparse updates in parameter space → degraded performance; (2) DCT produces dense updates after inverse transform → improved performance; (3) random projections produce dense updates with continuously varying basis → marginally better performance than DCT.

What is missing: The ablation does not test whether the DCT's specific frequency-ordering property (energy concentration in low frequencies for smooth signals) provides an advantage over other fixed transforms with different basis properties. Comparing DCT against other fixed orthonormal transforms—Hadamard, Fourier, wavelet—would test whether the paper's intuition about DCT being particularly well-suited for smooth momentum signals is correct. The random projection result already shows the benefit is not unique to DCT, but it does not isolate whether DCT is better than other fixed transforms or merely good enough.

Does the Ablation on α Support the Claim That Momentum Subtraction Acts as Error Feedback?

Yes, strongly. This is one of the most convincing ablations in the paper.

Figure 4 (middle) shows a clear gradient: α = 0 (no subtraction) → severe degradation because redundant communication; α = 0.2 → best performance, balancing novelty and smoothing; α = 1.0 → functional but worse than α = 0.2. This pattern is exactly what the error feedback interpretation predicts: some subtraction is necessary to prevent redundant communication (α > 0), but full subtraction removes the smoothing benefit of momentum on communicated components (α < 1). The optimal α = 0.2 is not obvious a priori and is a genuinely interesting empirical finding.

A potential confound: The paper sweeps α only for a fixed β = 0.999 and k = 32. It is possible that the optimal α depends on β (a larger β makes the momentum buffer smoother, potentially changing the optimal subtraction rate) or on k (more aggressive sparsification may require more or less aggressive subtraction). The paper does not explore these interactions, which limits the generality of the α = 0.2 recommendation.

Overall Assessment

The experimental section demonstrates convincingly that DeMo can reduce per-step communication by 1-2 orders of magnitude while maintaining training convergence and downstream accuracy on the tested benchmarks, at model scales up to 1B parameters. The ablation studies provide mechanistic validation of the key design choices (DCT over identity, partial momentum subtraction over full or none, large β values). The experiments are thorough within their scope.

The primary weaknesses are: (1) the evaluation is limited to a single model family (OLMo) on a single corpus (Dolma v1.5) with three downstream benchmarks; (2) no wall-clock time or throughput measurements are reported, making it impossible to assess whether communication reduction translates to training speedup; (3) the scale (300M, 1B) is modest relative to modern LLM training regimes where communication bottlenecks are most acute; (4) the claim of topology-agnosticism and multi-datacenter training is not empirically tested; (5) comparisons with prior communication-efficient methods are informative but not exhaustive, with several relevant baselines (Deep Gradient Compression, 1-bit Adam, GaLore) not evaluated. The paper's central claims are supported within these boundaries, but the practical value of DeMo for production-scale training depends on factors (wall-clock speed, scaling to 10B+ parameters, robustness to network conditions) that the experiments do not address.

6. Limitations and Trade-offs

Limitation 1: Wall-Clock Training Speed Is Never Measured

The assumption or constraint. The paper measures communication reduction exclusively in MB per GPU per step—the volume of data transmitted. Table 1 reports that DeMo with k = 8 transmits 7.5 MB/step versus AdamW-DDP's 637 MB/step, an 85× reduction. Nowhere—not in the main results, not in the ablations, not in the appendix—does the paper report wall-clock time per training step, throughput in tokens per second, or end-to-end training duration. The complexity analysis in Section 2.1 bounds the asymptotic DCT computation cost (O(N³/C) with chunking), but there are no empirical timing measurements for the actual PyTorch implementation running on 64 H100 GPUs.

The consequence. Communication reduction in bytes does not automatically translate to training speedup. The DeMo pipeline introduces several per-step operations that standard DDP does not: the chunk-wise DCT and inverse DCT on every momentum tensor, the top-k selection (which requires a partial sort or selection algorithm per chunk), the sparse packing and unpacking of coefficient sets, and the server-side sparse aggregation with index matching (Algorithm 3, which must handle different workers contributing different non-zero indices). If any of these operations constitutes a bottleneck—particularly the DCT, which for small chunk sizes may not be GPU-accelerated to the same degree as matrix multiplications—the end-to-end step time could be longer than standard DDP even though fewer bytes cross the network. The paper's claim that DeMo "reduces communication bandwidth" is true as stated, but a practitioner deciding whether to adopt DeMo needs to know the effect on total training time, not just network utilization. A method that cuts network traffic by 85× but increases per-step computation by 50% might yield only a modest speedup or even a slowdown on systems where the network was not the dominant bottleneck to begin with.

What evidence exists in the paper. None. The paper reports only data volume (MB/step) and training loss/downstream accuracy. The experimental setup mentions 64 NVIDIA H100 GPUs (Section 3.1) but provides no profiling data. The "Compute Efficiency" discussion in Section 3.2 and Figure 4 plots perplexity against data transmitted, not against time. This is an unusual omission for a systems paper whose primary value proposition is faster distributed training.

Mitigation status. Not addressed. The paper does not acknowledge this gap as a limitation, does not report any timing measurements, and does not suggest profiling or wall-clock benchmarking as future work. A reader cannot determine from the paper alone whether DeMo would speed up their training pipeline.

Limitation 2: Evaluated Only at Modest Model Scales (300M, 1B) on a Single Model Family

The assumption or constraint. All experiments use the OLMo family of decoder-only transformer language models (Groeneveld et al., 2024) at two scales: 300M non-embedding parameters and 1B non-embedding parameters. Section 3.1 states these specifications and notes that "Full model specifications are provided in the Appendix." The paper introduces DeMo as a general-purpose distributed optimization framework, but the empirical validation is confined to two sizes of one model architecture trained on one corpus (Dolma v1.5).

The consequence. Several aspects of DeMo's performance could be model-scale-dependent or architecture-dependent in ways the paper cannot characterize:

The compressibility of momentum under DCT may vary with model scale. Larger models have larger parameter tensors, which could concentrate energy differently in the DCT domain. The per-chunk sparsification mechanism operates identically regardless of total model size (each s × s chunk is compressed independently), but the distribution of effective ranks across layers might shift as models grow—attention layers, MLP layers, and embedding layers may exhibit different compressibility profiles. The paper's 300M and 1B results are encouraging but do not guarantee that the same k = 8 operating point remains optimal at 7B or 70B parameters.
The interaction with parallelism strategies is unexplored. Production training of large models typically combines data parallelism with tensor parallelism, pipeline parallelism, or sequence parallelism (Shoeybi et al., 2019; Narayanan et al., 2021). DeMo is evaluated only in a pure DDP setting. If tensor parallelism is used within nodes and DeMo coordinates across nodes, the momentum tensors being compressed would be only a subset of the full model parameters, potentially changing the compressibility profile and the effective communication bottleneck.
Architecture-specific properties matter. The OLMo architecture uses standard transformer blocks. Whether DeMo's DCT-based sparsification works equally well for architectures with different weight matrix structures (e.g., mixture-of-experts with sparse routing, state-space models like Mamba, or models with convolutional components) is unknown. The DCT's frequency-concentration property relies on spatial smoothness in the parameter tensors, which may not hold for all layer types. The paper's ablation on linear transformations (Figure 3) shows that identity (no transform) performs poorly, but this was tested only on the OLMo-300M architecture.

More practically, the communication bottleneck that DeMo addresses is most severe at scales far beyond 1B parameters. A 1B-parameter model at 32-bit precision has a ~4 GB gradient tensor—manageable even on modest interconnects. A 70B-parameter model has a ~280 GB gradient tensor, where the 85× reduction DeMo promises would be transformative. The paper provides no evidence that the method's benefits persist at the scales where they matter most.

What evidence exists in the paper. The 300M and 1B results (Figures 2, 4; Tables 1, 3) show consistent patterns across the two scales: DeMo with moderate sparsification matches or exceeds AdamW at 100B tokens. The scaling from 300M to 1B is favorable, with per-GPU communication for AdamW growing from 637 to 2417 MB/step while DeMo's relative compression factors remain similar (e.g., k = 8 transmits 7.5 MB/step for 300M, 27.6 MB/step for 1B—roughly 3.7× more data for a 3.7× larger model, as expected from linear scaling of parameter count). However, two data points is not a scaling law, and the paper does not extrapolate or make claims about larger scales.

Mitigation status. The paper does not explicitly acknowledge the limited scale as a limitation. The abstract and introduction present DeMo as a general method without qualifying the model sizes for which it has been validated. No future work is suggested regarding scaling to larger models or testing on different architectures.

Limitation 3: The Parameter Server Architecture Creates a Download Bandwidth Bottleneck at Scale

The assumption or constraint. DeMo uses a parameter server architecture (Algorithm 1): workers upload sparse DCT coefficients to a central server, the server aggregates and reconstructs the momentum, then broadcasts the parameter update (in the Signum variant, the 1-bit sign(M_t)) back to all workers. Section 5 discusses this explicitly: "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." The paper positions DeMo "primarily for optimization across a small number of geographically distributed compute centers," where each center is treated as a "large worker" running internal DDP (Section 5).

The consequence. The download bandwidth scales linearly with the number of workers N. For N workers, the server must broadcast the parameter update to all N workers, consuming N × (model size) in download bandwidth per step. For the Signum variant used in the experiments, this is N bits per parameter—a factor of 32 smaller than full-precision—but still N × P bits for P parameters. At the large worker counts typical in LLM training (N = 64, 128, 256 or more), the server's outgoing bandwidth can become the bottleneck, especially if the server is connected via a lower-bandwidth link (as would be the case in the geographically distributed setting DeMo targets, where the server may coordinate across data centers over the internet).

The paper's proposed mitigation—using a small number of "large workers" that each internally run DDP—partially addresses this by keeping N small at the DeMo level. However, this means DeMo operates on top of, rather than replacing, traditional DDP within each data center. The total communication cost is then DeMo's inter-center cost plus DDP's intra-center cost. If the intra-center DDP cost dominates (because it uses full-precision all-reduce), the overall communication savings may be modest. The paper does not provide an analysis of this combined cost or guidelines for choosing the optimal split between intra-center DDP and inter-center DeMo.

Furthermore, the upload-side compression ratio is determined by k and chunk size (approximately M/k per chunk, where M = 4096 for default 64 × 64 chunks). The download-side compression ratio is fixed at 32× (full-precision to 1-bit) for the Signum variant, or 1× (no compression) for the SGD variant. The paper's headline "up to 85×" compression figure applies only to the upload direction; the download direction benefits from at most 32× (Signum) or 0× compression. For setups where download bandwidth is also constrained—as it would be in many cross-datacenter scenarios—the asymmetric compression is a genuine limitation.

What evidence exists in the paper. The paper provides per-GPU communication volumes in MB/step (Table 1, Table 3), which include both upload and download. The discussion in Section 5 explicitly acknowledges the download bandwidth scaling issue. However, the paper does not profile upload vs. download bandwidth separately, does not report the effective download compression ratio for the Signum variant used in experiments, and does not measure how the total communication cost decomposes between upload and download at different worker counts.

Mitigation status. Partially acknowledged (Section 5) but not empirically addressed. The paper describes DeMo's intended deployment model (few large workers, each running internal DDP) as a design mitigation but does not validate this model experimentally. The paper also notes that the download bandwidth limitation is "intrinsic to all top-k sparsification-based approaches"—which is true but does not reduce its practical significance.

Limitation 4: Downstream Evaluation Is Too Narrow to Support Claims of "Comparable Accuracy"

The assumption or constraint. The paper evaluates downstream model quality using only three benchmarks: HellaSwag, ARC-Easy, and PIQA (Table 1). These are all relatively simple tasks: HellaSwag tests commonsense reasoning through sentence completion, ARC-Easy tests elementary science knowledge, and PIQA tests physical commonsense. The paper states that "DeMo matches or exceeds the AdamW-DDP baseline across all tasks while reducing per-GPU communication by two to three orders of magnitude" (Section 3.3) and that "LLMs pre-trained with DeMo have equivalent or better scores on multiple standard benchmarks compared to their equivalents trained with AdamW" (Section 6).

The consequence. These three benchmarks do not constitute a comprehensive evaluation of language model quality. A model can perform well on HellaSwag, ARC-Easy, and PIQA while exhibiting significant deficiencies in other capabilities that are important for downstream use. Concretely, the paper does not evaluate:

Factual knowledge and reasoning (MMLU, TriviaQA, NaturalQuestions): Does momentum compression affect the model's ability to store and retrieve factual information? Since DeMo's sparsification discards low-magnitude DCT coefficients, it is possible that fine-grained factual updates—which may correspond to low-energy signal in the momentum frequency domain—are disproportionately lost.
Mathematical reasoning (GSM8K, MATH): Does the compression affect multi-step reasoning capabilities? The paper's related work discusses the difficulty of self-correction on math problems; it is plausible that compressed optimization similarly struggles with precise symbolic reasoning.
Code generation (HumanEval, MBPP): Does the model's ability to generate syntactically and semantically correct code suffer from sparse momentum communication?
Perplexity on held-out text: The paper reports training loss (cross-entropy on the training corpus) but not perplexity on a held-out validation set beyond the Dolma corpus. Training loss can be lower while generalization is worse if the compression introduces a form of implicit regularization that overfits to the training distribution.
Generation quality (human evaluation, automated metrics like MAUVE): The paper only evaluates discriminative benchmarks (multiple choice); open-ended generation quality is not assessed.

The claim that DeMo achieves "comparable accuracy" is therefore supported only for a narrow slice of model capabilities. A practitioner choosing between DeMo and AdamW for pretraining a model intended for diverse downstream tasks cannot determine from this paper alone whether the compression introduces subtle degradations that would only surface on more challenging evaluations.

What evidence exists in the paper. Table 1 reports numbers for HellaSwag, ARC-Easy, and PIQA. The numbers are consistent and favorable to DeMo—at moderate sparsification (k ≥ 4), DeMo matches or slightly exceeds AdamW on all three benchmarks. At k = 1, DeMo shows small but visible drops: for the 300M model, HellaSwag stays at 0.35 (same as AdamW), ARC-Easy drops from 0.46 to 0.45, PIQA stays at 0.65. For the 1B model, the drops are larger: HellaSwag 0.41 vs. 0.43, PIQA 0.69 vs. 0.68. The degradation is small enough that the paper's "comparable accuracy" claim is defensible for these specific benchmarks, but the narrowness of the evaluation is not acknowledged.

Mitigation status. Not addressed. The paper does not discuss the limited scope of downstream evaluation as a limitation, does not suggest additional benchmarks for future work, and does not qualify the "comparable accuracy" claim with the specific benchmarks tested. The evaluation suite is a standard part of the OLMo training framework, which explains its use but does not justify treating it as sufficient to establish model quality equivalence.

Limitation 5: DCT Performance Depends on Momentum Smoothness, But the Paper Does Not Evaluate DeMo on Noisy or Rapidly-Changing Loss Landscapes

The assumption or constraint. DeMo's compression efficiency depends on the momentum signal being sufficiently smooth that the DCT concentrates energy into a small number of low-frequency coefficients. The theoretical analysis (Lemma 1, Appendix 8) bounds the momentum norm in terms of the gradient bound R, the sparsity ratio k/M, and the momentum coefficient β—larger β produces smoother momentum and tighter bounds. The empirical results confirm that larger β values (0.995, 0.999) outperform standard values (0.9, 0.95) in the DeMo setting (Figure 4, right). This is consistent with the smoothness requirement: more aggressive momentum smoothing (higher β) concentrates more energy into fewer DCT coefficients, making top-k sparsification more efficient.

The consequence. The requirement for high β means DeMo's momentum buffer responds more slowly to changes in the gradient signal. The EMA half-life at β = 0.999 is approximately ln(0.5) / ln(0.999) ≈ 693 steps—a gradient update from 693 steps ago still contributes half its original value to the current momentum. This is desirable for compression (more smoothing → better energy concentration) but potentially harmful for optimization scenarios where the loss landscape changes rapidly and the optimizer needs to adapt quickly. Examples include:

Curriculum learning: if the data distribution shifts mid-training (e.g., from easier to harder examples), the momentum buffer at β = 0.999 will be dominated by gradients from the old distribution for hundreds of steps.
Sharp phase transitions: some training runs exhibit sudden qualitative changes in the loss landscape (e.g., when the model transitions from memorization to generalization). A slow-responding momentum buffer may delay the optimizer's adaptation to the new landscape.
Fine-tuning or continual learning: if DeMo were used for fine-tuning a pretrained model on new data, the momentum buffer would need to quickly "forget" the pretraining gradient history and adapt to the fine-tuning signal. The large β required for compression efficiency works against this.

The paper's experiments use a single, static pretraining corpus (Dolma v1.5) with a standard training recipe (linear warmup, cosine decay). This is a relatively benign setting where the gradient distribution changes slowly and predictably. The paper provides no evidence that DeMo's performance is robust to more dynamic training scenarios.

What evidence exists in the paper. The β ablation (Figure 4, right) shows that β ∈ {0.95, 0.98, 0.99, 0.995, 0.999} all produce stable training, with larger values performing better. The loss curve for β = 0.95 (the smallest tested) is noticeably worse than for β = 0.999. This confirms the dependence on high β but does not test whether this dependence creates problems in dynamic settings. The paper does not report experiments with changing data distributions, learning rate restarts, or other perturbations that would test the responsiveness of DeMo's momentum buffer.

Mitigation status. Not addressed. The paper does not discuss the tradeoff between momentum smoothness (beneficial for compression) and adaptation speed (beneficial for optimization in dynamic settings). The optimal β for DeMo is treated as a hyperparameter to tune, not as a structural tension to resolve. Future work could explore adaptive β scheduling (e.g., starting with moderate β during rapid learning phases and increasing β as the loss landscape stabilizes) or alternative momentum formulations that decouple the smoothing timescale from the compression timescale.

Limitation 6: The Method Has Not Been Validated on Tasks Without Clean Closed-Form Correctness or on Non-Transformer Architectures

The assumption or constraint. DeMo is evaluated exclusively on decoder-only transformer language model pretraining using the Dolma v1.5 corpus, with evaluation on multiple-choice downstream benchmarks. The paper presents DeMo as a general-purpose distributed optimization framework compatible with "any momentum-based optimizers" (Section 1), but the validation is confined to one task family (language modeling), one architecture family (transformers), one training regime (pretraining from scratch with a static corpus), and one optimization objective (cross-entropy minimization with gradient-based updates).

The consequence. Several aspects of DeMo's design could interact unfavorably with tasks and architectures outside this narrow validation:

Vision models (CNNs, ViTs): Convolutional layers have structured weight tensors with strong spatial correlations (e.g., filter kernels). The DCT's frequency-concentration property might work even better for these than for transformer weight matrices—or it might interact poorly with the specific spatial structure of convolutional filters. The paper provides no evidence either way.
Multi-modal models: Models with heterogeneous parameter types (text embeddings, image encoders, cross-attention layers) might exhibit different per-layer compressibility, requiring layer-specific chunk sizes or k values that the current uniform-per-chunk approach does not support.
Reinforcement learning: RL training often involves non-stationary data distributions (the policy changes as training progresses, changing the data distribution) and reward signals with different statistical properties than supervised learning gradients. The momentum smoothness requirement (high β) could be particularly problematic in RL settings where the gradient distribution shifts rapidly.
Fine-tuning (rather than pretraining from scratch): The paper's experiments all involve pretraining from random initialization. If DeMo were applied to fine-tuning a pretrained model, the momentum buffer would start from zero while the parameters start from a pretrained state—potentially causing a mismatch between the optimizer state and the parameter state during early fine-tuning steps. The large β = 0.999 would mean the momentum buffer takes hundreds of steps to "catch up" to the fine-tuning gradient signal.
Non-transformer architectures with different parameter structures: Models with recurrent components, graph neural networks, or other non-standard layer types might have parameter tensors whose DCT representations do not exhibit the energy concentration that makes top-k sparsification efficient. The paper's ablation showing identity transform (no DCT) degrades performance (Figure 3, left) confirms that the transform is load-bearing; if a particular architecture's parameter tensors are not smooth in the DCT basis, DeMo might perform similarly to the identity-transform baseline for those layers.

The paper does not claim to have tested these settings, but the breadth of the claims ("drop-in replacement," "any momentum-based optimizers," "topology-agnostic") invites application beyond the validated regime. A practitioner considering DeMo for, say, fine-tuning a vision-language model or training a diffusion model would find no guidance in the paper.

What evidence exists in the paper. None beyond the language modeling results. The paper does not include experiments on vision tasks, RL, fine-tuning, or non-transformer architectures. The theoretical analysis (Theorem 1) is general and does not depend on the architecture or task, but the empirical validation is narrow.

Mitigation status. Not addressed. The paper does not discuss the scope of validation as a limitation or suggest broader empirical evaluation as future work. The "drop-in replacement" framing in the abstract and introduction would benefit from qualification with the specific settings in which DeMo has been tested.

7. Implications and Future Directions

How This Work Changes the Landscape

DeMo is not a paradigm shift in distributed optimization—it does not introduce a new optimizer, a new convergence theory, or a fundamentally new compression algorithm. What it does is reframe the communication problem in distributed training by identifying the optimizer's internal state as the natural compression target, and this reframing has consequences that ripple through how we think about the relationship between optimization and communication.

The pre-DeMo landscape was organized around a clean separation: gradients are computed, gradients are communicated (possibly compressed), then the optimizer consumes the synchronized gradients to update parameters and internal state. Compression research focused on the "gradients are communicated" step—sparsifying, quantizing, or projecting the gradient tensor—while treating the optimizer as a downstream consumer. This separation was so ingrained that even methods explicitly designed to reduce communication (Deep Gradient Compression, PowerSGD, 1-bit Adam) accepted it as architectural: you compress the gradient, you optionally add error feedback to correct the compression bias, and the optimizer sees an approximate gradient.

DeMo's core move is to dissolve the boundary between the optimizer state and the communication primitive. The momentum buffer is simultaneously (a) the optimizer's memory of past gradients, (b) the signal that gets compressed and transmitted, and (c) the error feedback accumulator that tracks compression residuals. This is a conceptual unification rather than a new technique—the individual components (decoupled momentum, DCT, top-k sparsification, error feedback) all existed—but the recognition that they can be combined into a single mechanism with no extra memory cost is genuinely novel. It is the kind of insight that seems obvious in retrospect: of course the momentum buffer already accumulates past gradient information; of course subtracting communicated values from it leaves the residuals in place; why did we ever allocate separate error feedback buffers?

The most important reframing this work enables is a shift from "compress the gradient" to "choose the most compressible representation and communicate that." The DCT transform is the mechanism by which DeMo makes momentum compressible, but the strategic insight is upstream of the transform choice: gradients are noisy and high-variance; momentum is smooth and structured; therefore momentum is a better compression target. This principle generalizes beyond momentum. Second-moment estimates in Adam (v_t), preconditioners in Shampoo, or running statistics in BatchNorm might similarly be more compressible than the raw signals they summarize. The paper does not explore these, but the principle it establishes—the optimizer state may be a better communication substrate than the gradient—opens a design space that prior work simply did not consider.

The paper resolves a latent tension in the distributed training literature between two competing philosophies: frequent communication (DDP and its compressed variants) versus infrequent communication (Local SGD, FedAvg, DiLoCo). The field had implicitly assumed that if you want orders-of-magnitude communication reduction, you must reduce the synchronization frequency—because per-message compression alone could not achieve the necessary ratios without severe degradation. DeMo demonstrates that per-message compression can achieve 100× or greater reductions (k = 1 achieves 685× for 300M models) while maintaining convergence, provided you compress the right thing (momentum, not gradients) in the right way (transform domain, not original domain). This reopens the frequent-synchronization design point as viable even for extreme bandwidth constraints, and it suggests that the two philosophies could be combined (frequent compressed synchronization within data centers, infrequent synchronization across them) rather than treated as mutually exclusive.

Research directions that become more attractive:

Optimizer-state-as-communication-primitive for other optimizers. If momentum works better than gradients as a compression target, do Adam's second-moment estimates (v_t) work better still? The paper uses Signum (sign(M_t)) as the base update, but Adam maintains two EMA buffers (m_t and v_t). The v_t buffer tracks gradient variance and may exhibit different compressibility properties than m_t. A natural extension is DeMo-Adam: compress both m_t and v_t (or their ratio) in the DCT domain before communication.
Learned or adaptive transforms for compression. The paper shows random projections marginally outperform DCT (Figure 3, right) because continuously rotating the basis prevents top-k from repeatedly selecting the same coefficients. This suggests learning the projection matrices—either via meta-learning across training runs or via online adaptation that maximizes the energy concentration in the top-k coefficients—could outperform both DCT and random projections. The computational cost of learned transforms would need to be amortized over many steps, but for long training runs, the tradeoff could be favorable.
Combining frequent compressed and infrequent full synchronization. DeMo and DiLoCo are compared as alternatives (Figure 6), but they are complementary: DeMo could handle per-step momentum synchronization within clusters, while DiLoCo-style full-parameter averaging happens every H steps across clusters. This hybrid would get the best of both approaches—frequent momentum alignment to prevent client drift, infrequent full synchronization to correct accumulated compression bias and re-anchor the global model state.

Research directions that become less attractive:

Pure gradient compression without error feedback. DeMo's results reinforce the finding that aggressive compression requires error feedback or an equivalent mechanism (momentum subtraction, in DeMo's case). The α = 0 ablation (Figure 4, middle) shows catastrophic degradation without subtraction, confirming that simply sparsifying and communicating gradients or momenta without tracking residuals is not viable at high compression ratios. Future work on gradient compression should assume error feedback is necessary and focus on reducing its cost (as DeMo does by repurposing the momentum buffer) rather than trying to eliminate it.
Aggressive quantization as a standalone strategy. Quantization alone (reducing precision to 1-bit or 8-bit) achieves compression ratios in the 4–32× range, which DeMo matches or exceeds with sparsification while also providing a path to much higher compression ratios (100–1000×) by combining sparsification with the Signum broadcast (which is itself 1-bit). The paper does not compare against quantization-only methods, but the communication volumes in Table 1 (0.93 MB/step for k = 1) are far below what quantization alone could achieve for a 300M-parameter model, suggesting sparsification + transform is the more scalable approach to extreme compression.

The paper's most durable contribution may be methodological rather than algorithmic. By ablating each design choice (DCT vs. identity, DCT vs. random, α values, β values, chunk sizes) and showing how they interact, the paper provides a template for how to empirically validate compression schemes for distributed optimization. The pattern—identify a compressible representation, apply a transform that enables dense reconstruction from sparse coefficients, use the optimizer's own state for error feedback, tune the feedback rate to balance novelty and smoothing—is transferable to other optimizers, architectures, and tasks even if the specific instantiation (DCT + top-k + momentum subtraction) is not.

Follow-Up Research This Work Enables

Wall-clock benchmarking of DeMo against AdamW-DDP at matched model scales. The paper's most conspicuous gap is the absence of any timing measurements. A follow-up study should profile DeMo's per-step wall-clock time against AdamW-DDP on the same hardware (64 H100 GPUs, as used in the paper) and report throughput in tokens/second at sparsification levels k ∈ {1, 2, 4, 8, 16, 32} for both the 300M and 1B models. The key measurements are: (a) what fraction of DeMo's per-step time is spent on DCT/IDCT computation vs. communication vs. the forward-backward pass; (b) at what model scale does DeMo's communication reduction translate to a net speedup (since small models may be compute-bound rather than communication-bound); (c) how does the speedup scale with GPU count (since communication overhead grows with worker count in DDP). This would directly determine whether DeMo's byte-count reductions translate to practical training acceleration, which is the metric practitioners actually care about.

Scaling DeMo to 7B parameters and measuring the effective compression ratio at scale. The paper validates DeMo at 300M and 1B parameters. The communication bottleneck becomes acute at 7B+ parameters, where gradient tensors exceed 28 GB per worker. A follow-up should train a 7B-parameter OLMo or Llama-class model with DeMo on at least 200B tokens (the Chinchilla-optimal budget for 7B) and measure: (a) whether the optimal k and chunk size from the 300M/1B experiments transfer to 7B or need re-tuning; (b) whether the DCT energy concentration property holds at scale or whether larger weight matrices exhibit different frequency spectra; (c) end-to-end training time compared against standard AdamW-DDP, FSDP, and tensor-parallel baselines; (d) downstream evaluation on a comprehensive suite including MMLU, GSM8K, HumanEval, and held-out perplexity to assess whether model quality degradation emerges at scale. This experiment would directly test the paper's implicit claim that DeMo scales favorably, which currently rests on only two data points.

Testing DeMo on vision model training (ViT, CNN) to evaluate architecture-generality. The paper's DCT-based compression relies on spatial smoothness in the momentum signal—an assumption that holds for transformer weight matrices trained on language but may or may not hold for other architectures. A follow-up should train a ViT-B/16 on ImageNet-1K and a ResNet-50 on ImageNet-1K, both with DeMo and AdamW-DDP, measuring: (a) whether the same DCT + top-k pipeline achieves comparable compression ratios and convergence; (b) whether convolutional filters (which have strong local spatial structure) benefit more or less from DCT-based sparsification than transformer weight matrices; (c) whether different layer types (convolutional, attention, MLP, normalization) require different transforms—e.g., DCT for convolutional layers, random projections for attention, identity for LayerNorm. A negative result (DeMo fails to compress vision model gradients effectively) would bound the method's generality and motivate architecture-specific compression strategies. A positive result would significantly expand DeMo's applicability.

Combining DeMo's momentum compression with DiLoCo-style infrequent synchronization. The paper compares DeMo against DiLoCo (Figure 6) but does not combine them. A natural follow-up is a two-level hierarchy: within each data center, workers run DeMo with aggressive sparsification (k = 4 or 8) for per-step momentum synchronization; across data centers, full-parameter averaging happens every H steps à la DiLoCo. The research questions are: (a) does DeMo's per-step momentum alignment reduce the client drift that limits DiLoCo's H?—if so, larger H (and thus greater communication reduction) might be viable; (b) does the combination achieve better perplexity-vs-communication tradeoffs than either method alone?; (c) what is the optimal H and k for a given inter- and intra-datacenter bandwidth ratio? This experiment would unify the two philosophies the paper identifies (frequent compressed vs. infrequent full communication) and determine whether their benefits are additive or redundant.

Adaptive per-layer sparsity budgets based on DCT energy concentration. The paper uses a uniform k across all chunks and all layers. In practice, different layers (attention Q/K/V/O projections, MLP up/gate/down projections, embeddings) may exhibit different compressibility in the DCT domain. A follow-up should instrument a training run to measure the energy concentration curve (cumulative fraction of total energy captured by top-m coefficients) for each layer type, then implement an adaptive sparsity budget that allocates more coefficients (larger k) to layers with diffuse DCT spectra and fewer (smaller k) to layers with concentrated spectra, subject to a total communication budget constraint. The research question is whether adaptive allocation improves the perplexity-vs-communication Pareto frontier beyond uniform allocation. This experiment would determine whether the uniform-k design is a convenient simplification or leaves significant efficiency on the table.

Stress-testing DeMo under curriculum learning, data distribution shift, and learning rate restarts. The paper's experiments use a static data distribution (Dolma v1.5) with a smooth cosine learning rate schedule. DeMo's reliance on high momentum coefficient β = 0.999 means the momentum buffer has a long memory (~693 step half-life), which could slow adaptation to distribution shifts. A follow-up should train models with DeMo under three dynamic training scenarios: (a) curriculum learning where the data distribution changes from easier to harder examples at predefined steps; (b) a sudden data distribution shift (e.g., switching from English to code data midway through training); (c) a cosine learning rate schedule with a mid-training restart (a sharp increase followed by decay). For each scenario, compare DeMo against AdamW-DDP in terms of loss recovery speed (how many steps after the shift until loss returns to trend) and final model quality. A finding that DeMo adapts more slowly than AdamW would establish a boundary condition on when the method is applicable; a finding that it adapts equally fast would suggest the long momentum memory does not impede responsiveness in practice.

Practical Applications and Downstream Use Cases

Cross-datacenter training for organizations with geographically distributed GPU clusters. A company with GPU capacity spread across multiple data centers (e.g., a research lab with clusters in different regions, or a cloud user renting instances across availability zones) currently cannot efficiently pool those resources for synchronous training of a single large model—the inter-datacenter bandwidth is orders of magnitude lower than intra-datacenter interconnects, making standard DDP infeasible. DeMo is architecturally designed for this scenario: each data center runs internal DDP (or DeMo with modest compression) and is treated as a single "large worker" at the DeMo level, with only sparse DCT coefficients crossing the wide-area link. With k = 8, a 1B-parameter model transmits 27.6 MB/step per center (Table 1), which at one step per second would require ~220 Mbps—feasible over commodity internet connections. This enables organizations to aggregate otherwise-idle GPU capacity across locations into a single training run, effectively increasing their peak training throughput without investing in colocation or dedicated fiber.

Reducing the networking hardware requirements for on-premise training clusters. DeMo's 85× reduction in per-step communication volume (k = 8 for 300M models) means that training clusters can use lower-cost networking infrastructure. A cluster that required 100 Gbps InfiniBand to avoid communication bottlenecks with standard DDP might achieve equivalent throughput with 10 Gbps Ethernet under DeMo—a substantial cost reduction, especially for academic labs or small companies building bespoke clusters. Even within a single rack, the reduced network utilization could allow more GPUs to share the same interconnect without saturation. The caveat (Section 5) about download bandwidth scaling with worker count matters here: the 85× figure applies to upload; for the Signum variant used in experiments, download is 32× compressed (full-precision to 1-bit), so the effective download reduction is smaller. But in a single-cluster setting with a high-bandwidth internal network, download bandwidth is less likely to be the bottleneck than upload all-reduce bandwidth, making DeMo's asymmetric compression well-matched to the problem.

Low-bandwidth federated pretraining of language models across institutions. Federated learning typically operates in the regime of many clients (hundreds to thousands) with extremely limited per-client bandwidth, where standard methods like FedAvg communicate full model updates infrequently. DeMo offers an alternative: clients run DeMo locally (with decoupled momentum and DCT compression), communicating only sparse momentum coefficients to a central server every step. This shifts the communication pattern from "infrequent, large messages" to "frequent, tiny messages," which may be more compatible with networks that have bandwidth caps but low per-message overhead. Moreover, because DeMo does not require clients to perform multiple local steps between synchronizations (unlike FedAvg), it avoids client drift from heterogeneous local data distributions—a significant problem in cross-silo federated learning. A consortium of research labs wanting to collaboratively pretrain a language model on their pooled (but privacy-sensitive) data could use DeMo to keep per-institution communication to a few MB per step, making the collaboration feasible over standard institutional internet connections.