Beyond Communication Overhead: A Multilevel Monte Carlo Approach for Mitigating Compression Bias in Distributed Learning

We thank the reviewer for their positive evaluation of our paper and for the constructive feedback. Below, we address the questions raised:

1. Scaling with number of machines

You are right that the performance gains from MLMC grow with the level of parallelization. When using only 4 machines, variance reduction is not as effective, so the gains are naturally smaller. However, as shown in our scalability plots (e.g., Figure 1 ($M=4$) compared to Figure 2 ($M=32$)), MLMC achieves significant speedups, smaller errors, and improved communication efficiency in large-scale distributed settings.

This is a key strength of our method. It retains the empirical efficiency of biased compressors while providing the unbiasedness and theoretical guarantees needed for scalable, stable learning in high-parallelism regimes. Moreover, higher parallelization induces a better variance reduction effect, which makes our MLMC estimator even better.

We will clarify this in the revised paper to better highlight the scaling advantages of MLMC.

2. Computational overhead

That is a good point. The computational overhead of MLMC is comparable to standard methods such as Top-k and AdaGrad, and is often negligible relative to the overall training time. Similar works also focus less on the computational overhead since it is often negligible compared to the overhead introduced by the communication.

Specifically, using top-$k$, for example, incurs $O(d\log(k))$ computational complexity per iteration and per machine while our adaptive MLMC method (Alg. 3) incurs $O(d\log(d))$, which is a very small difference in practical scenarios.

In more detail, using top-$k$ requires finding the $k$ largest elements, which costs $O(d\cdot\log(k))$ in each iteration and for each machine. In contrast, our adaptive MLMC method (Alg. 3) with top-$k$ requires sorting the vector first, costing $O(d\cdot\log(d))$, and computing the probabilities and constructing the MLMC estimator, which costs $O(d)$ in total (for computing the norm of the vector, similar to AdaGrad, and for picking the $l$-th largest element, which costs $O(1)$ since the vector is sorted). However, you are right that this is worth discussing to improve clarity, and we will add it to our paper.

3. Clarification on Figure 3

Thank you for pointing this out. The results in Figure 3 correspond to Algorithm 2. We will clarify this in the figure caption and text to avoid confusion.

4. Experiments

We ran new experiments, including NLP experiments using BERT on SST-2. These are anonymously available at "https://anonymous.4open.science/r/ICML2025MLMC-5346".

Thank you again for your constructive feedback!

We thank the reviewer for the thoughtful and encouraging feedback, and for pointing out areas where the theoretical analysis could be clarified. We respond to each concern below and will revise the paper accordingly.

1. Variance Analysis and Convergence Guarantee

We appreciate your observation regarding the variance and convergence analysis.

Regarding the variance, we presented the full derivations of the variance of the MLMC estimators as part of the calculation of the optimal probability distribution (which is optimized to minimize the variance) in Appendices B,C, and D (see Eq. (33, 44, 55)).

For a more practical example, in Lemma 3.6 we show that the variance of our MLMC estimator in the exponential distribution case (see Assumption 3.5) is given by $O(1/rs)$, where $r$ is the exponential decay rate of the vector's elements, which is better than the $O(d/s)$ variance of rand-$k$ when $1/r < d$ (i.e., when the decay rate is sufficient, which is the more interesting scenario). Note that when the vector is nearly uniform, we have $1/r \approx d$ and the variances will be comparable in this case, as expected.

Regarding convergence, please note that since our MLMC gradient estimators are unbiased by construction, convergence follows by the standard SGD convergence analysis, where the only difference is the additional variance introduced by the compression, as we state in lines 186-200. However, we agree that making these derivations explicit would strengthen the presentation, and we will add this to the paper as you suggested to make it clearer.

Specifically, the formal convergence theorem for Alg. 2 and Alg. 3 will hold under Assumptions 2.1 (smoothness) and 2.2 (bounded variance) and will guarantee similar error bounds (for the convex and nonconvex cases) as the ones in Theorem 2.1 and Eq. (2), only with $\sigma_{comp}+\sigma$ in place of $\sigma$, where $\sigma_{comp}$ depends on the compressor and thus on $\alpha^l, l\in[L]$ (see e.g. Eq. (60) in Appendix D).

We formalize the convergence theorem as follows.

Theorem (convex case). Under Assumptions 2.1 (smoothness) and 2.2 (bounded variance), Alg. 2 (nonadaptive MLMC compression) guarantees the following error bound: $O(\frac{1}{T}+\frac{\sigma_{comp}+\sigma}{\sqrt{MT}})$

Although the proof follows very similarly to the SGD convergence theorem, we will formalize and add it to the paper for completeness. Similar Theorem and proof follow for the nonconvex case and we will add them as well.

2. Bounds dependence on compression constants

We believe the reviewer’s concern refers to the apparent disappearance of level-specific compression constants (i.e., $\alpha_{t,i}^l$) in the final rate. This is a good point. As we mentioned in the previous point, these constants come to light in the variance introduced by compression, i.e., in $\sigma_{comp}$ (see e.g. Eq. (60) in Appendix D) which appears in the final convergence rate. We agree with the reviewer's suggestion to add this to the main paper to make it clearer, and we will incorporate this in the final version.

3. Minor Points

We thank the reviewer for the helpful suggestions regarding notation, definitions, and related clarity issues. We will revise the manuscript accordingly to improve clarity.

• You are correct that the unbiasedness of our MLMC estimator depends on the highest level, $L$. Since we define $C^L(v):=v$, i.e., there is no compression, our MLMC method produces unbiased estimates of the true gradient.

• Regarding the convergence rate of convex and non-convex SGD, we formalized the assumptions in our paper, but we will elaborate more on this and on the optimality of the bounds to improve clarity as you suggested.

• Regarding eliminating humans from the loop in experiments, we ran additional NLP experiments using the AdamW optimizer, which employs an adaptive learning rate and alleviates the need for extensive learning rate tuning. The results are anonymously available at "https://anonymous.4open.science/r/ICML2025MLMC-5346".

Thank you again for your constructive feedback!

We thank the reviewer for the positive evaluation of our contributions, including the novelty of our MLMC compression scheme and the theoretical analysis. We address the concerns below and will incorporate these improvements into the final version.

1. Implementation

We appreciate the reviewer’s suggestion to discuss real-world implementation. In our experiments, since hardware is limited, we used the "Multiprocessing" package to run multiple processes in parallel, each representing a different machine. This method accurately simulates a real-world parallel optimization scheme that runs on multiple machines. We will add this to the paper.

2. Computational Overhead

That is a good point. However, the adaptive sampling step in Algorithm 3 does not introduce significant additional computational overhead. For example, in the case of top-$k$ or $s$-top-$k$, we only need to compute each $\Delta_{t,i}^l$ (which is some norm) once per iteration, similar to what optimizers like AdaGrad already do when computing the full gradient norm. Furthermore, these norms are computed over disjoint segments, since $\sqrt{\Delta_{t,i}^l} = ||g_{t,i}^l - g_{t,i}^{l-1}||$, which is equivalent to the absolute value of the $l$-th largest element of $v_{t,i}$ (in top-$k$) or the norm of the segment of length $s$ with the $l$-th largest norm (in $s$-top-$k$). Therefore, the total computational cost is identical to that of computing the norm of the full gradient, as is done in existing adaptive methods like AdaGrad. A similar smart computation of the probabilities can be done for other compressors. For example, in bit-wise compressors, $g_{t,i}^l - g_{t,i}^{l-1}$ corresponds to the sign-bit and the $l$-th information bit.

Moreover, the number of compression levels does not have to be linear in the dimension of the compressed entity, but could be logarithmic (which is less general, but still works). This is the "classical" case of MLMC [1] in which the quality of the "levels" (which is inversely correlated with the extent of compression in our case) increases exponentially and thus induces a logarithmic number of levels.

Also, the computational overhead (even when calculating multiple compressions) is usually negligible compared to the overhead introduced by communication, which is the main motivation behind this work and prior works. This has been discussed extensively in prior work [2,3].

2. Experiments

We acknowledge that our current experiments are on modest-sized vision tasks, due to limited hardware. Although these are the standard benchmarks used in prior works, we agree that broader empirical validation is beneficial. We ran additional experiments. The results are available at "https://anonymous.4open.science/r/ICML2025MLMC-5346". We recommend downloading the repository. It includes two folders, as we elaborate below:

BERT Top-k:

We evaluated our adaptive MLMC method using Top-k compression on the SST-2 benchmark using BERT finetuning. We used the AdamW optimizer. The folder includes 2 files showcasing the accuracy vs. #Gbit communicated, and accuracy vs. iteration (#steps), both for M=4 machines and for k={0.01n, 0.05n, 0.1n, 0.5n}. We evaluated our MLMC method against EF21-SGDM, Top-k, Rand-k, and SGD (we keep this terminology, for clarity, with a slight abuse of notation, but note that the underlying optimizer for all is AdamW).

As is evident by these plots, our MLMC method enjoys the fastest convergence for the same #Gbit communicated, and it enjoys similar convergence to that of (uncompressed) SGD, and still performs better than other methods, for the same number of steps. These results (in addition to the ones in the paper) prove a strong advantage and efficiency of our method compared to existing methods, both in communication efficiency and convergence rate, across vastly different tasks like CV and NLP.

RESNET CIFAR10 Top-k:

We evaluated our adaptive MLMC method using Top-k compression on CIFAR-10 using ResNet-18 against EF21-SGDM, Top-k, Rand-k, and (uncompressed) SGD. The folder includes 4 files showcasing the accuracy vs. #Gbit and accuracy vs. iteration (#steps), both for M=4 and M=32 machines and for k={0.001n, 0.005n, 0.01n, 0.05n}.

Our method enjoys a significant advantage over comparable methods in terms of communication efficiency, convergence speed, and final accuracy. Our method's advantage grows with the number of machines. In the accuracy vs. iteration plots, uncompressed SGD eventually surpasses all methods, as expected, although our method is comparable when compression is not too extreme while other compression methods still experience performance degradation.

Thank you again for your constructive feedback!

[1] Giles, "Multilevel Monte Carlo methods", 2013.

[2] Konecny, "Federated learning: Strategies for improving communication efficiency", 2018.

[3] Wang, "A field guide to federated optimization", 2021

We thank the reviewer for the thoughtful and detailed feedback. Below, we address each concern and clarify the relationship between our MLMC framework and IS.

1. MLMC vs. Importance Sampling

We thank the reviewer for the insightful observation regarding the similarity between our MLMC construction and importance sampling (IS) in specific settings. That is an excellent point! We agree that in certain simple cases (bit-wise, Top-k) the MLMC estimator indeed reduces to an IS-like scheme, as you correctly pointed out. However, we respectfully argue that MLMC is not merely an instance of IS, but rather a significantly more general and natural framework for constructing unbiased estimators from biased compressors.

In fact, IS can be viewed as a special case of MLMC, where sampling is performed non-uniformly over coordinates, as you have stated. MLMC provides a systematic multilevel hierarchy over increasingly accurate (less compressed) estimators, and forms unbiased estimates by applying Monte Carlo sampling over the differences between successive levels. This telescoping structure is particularly well-suited to biased compressors, where compression is naturally available at varying levels of fidelity.

Importantly, MLMC offers several advantages beyond IS:

• It can be applied immediately to any sequence of biased compressors, with no need for any manual design of coordinate-level sampling probabilities or the structure of the communicated entity, which IS requires. For e.g. top-$k$, IS uses $1/p_l \cdot g_l$ w.p. $p_l |g_l|$ to achieve the same result, and while it's straightforward in this case, it might not be as straightforward or even feasible for more complex compressors, as we elaborate below. Also, please note that this intuition regarding IS with top-$k$ was enabled by MLMC.

• It is compatible with complex structured compressors that do not admit a coordinate-wise decomposition, and where IS is not naturally defined. For example, ECUQ [1] and Round-to-Nearest (RTN) [2,3] involve structured quantization (e.g., entropy constraints, grid-based rounding) for which the MLMC framework does not naturally decompose into an IS-like scheme. In such cases, it is unclear whether or how suitable IS can be defined, whereas MLMC applies seamlessly. MLMC enables these compressors to be used in a principled way to construct unbiased estimators, with automatic adaptation over compression levels.

• From a practical perspective, MLMC is also more flexible and intuitive: one can simply define a sequence of biased compressors with increasing accuracy, and MLMC provides a plug-and-play mechanism for building an unbiased gradient estimate without needing to manually tune the probabilities or derive the communicated entity.

We will revise the manuscript to clarify these points and include a detailed discussion comparing MLMC and IS, including when they coincide and when MLMC provides a strictly richer modeling framework. We thank you again for raising this important point.

2. Experiments

We acknowledge the reviewer’s concern regarding the ResNet18+CIFAR-10 setting. We chose this standard benchmark to align with prior works (e.g., EF21). Regarding accuracy, 90% requires use of Adam, LR scheduling, and more, which we did not employ as our focus was to isolate compression effects.

We fully agree that larger-scale settings will better demonstrate the advantages of our method, and we ran additional NLP experiments. The results are available at "https://anonymous.4open.science/r/ICML2025MLMC-5346".

BERT Top-k:

(See repository). We evaluated our adaptive MLMC method using Top-k compression on the SST-2 benchmark using BERT finetuning. We used the AdamW optimizer. The folder includes 2 files showcasing the accuracy vs. #Gbit communicated, and accuracy vs. iteration (#steps), both for M=4 machines and for k={0.01n, 0.05n, 0.1n, 0.5n}. We evaluated our MLMC method against EF21-SGDM, Top-k, Rand-k, and SGD (we keep this terminology, for clarity, with a slight abuse of notation, but note that the underlying optimizer for all is AdamW).

As is evident by these plots, our MLMC method enjoys the fastest convergence for the same #Gbit, and it enjoys similar convergence to that of (uncompressed) SGD, and still performs better than other methods, for the same number of steps. These results (in addition to the ones in the paper) prove a strong advantage and efficiency of our method compared to existing methods even for less aggressive compression, both in communication efficiency and convergence rate, across different tasks like CV and NLP.

Thank you again for your constructive feedback!

[1] Dorfman et al., “DoCoFL: Downlink Compression for Cross-Device Federated Learning,” ICML 2023.

[2] Gupta et al., “Quantization Robust Federated Learning for Efficient Inference on Heterogeneous Devices,” TMLR 2017.

[3] Dettmers et al., “GPT3.int8(): 8-bit Matrix Multiplication for Transformers at Scale,” NeurIPS 2022.