Paper: Accelerating Federated Learning with Quick Distributed Mean Estimation (ICLR-2024)

Reviewer comment (Writing):
The abstract is not informative at all. The reader has no idea about the techniques this work uses to improve DME and the novelty of the techniques after reading it. It should at least mention, for example, one idea QUIC-FL builds on is BSQ which sends exactly a few large coordinates and quantizes the rest small ones.

Author response:
We will revisit the abstract to make it more informative.

Reviewer comment (Experiments):
In Introduction, the paragraph starting with “For example, in Suresh et al. 2017 …”, it is mentioned the entropy encoding is “compute-intensive”. How long does the decoding time of this approach take? Table 1 does not include entropy encoding as a baseline.

Author response:
In theory, entropy encoding requires $O(d)$ time. However, this operation is not GPU-friendly and is thus much slower in practice.
In the paper “DoCoFL: Downlink Compression for Cross-Device Federated Learning” (ICML 2023), the authors performed timing measurements of entropy encoding schemes and found it to be ~ two orders of magnitude slower compared to EDEN, which runs on a GPU (QUIC-FL also runs on a GPU and is faster than EDEN as we show in the paper).

Reviewer comment (Experiments):
It is mentioned in Introduction that the decoding procedure of “entropy encoding” from a previous work is “compute-intensive”. What is the time complexity of this decoding procedure? It seems this approach is not included in Table 1 as a baseline for comparison.

Author response:
In theory, entropy encoding requires $O(d)$ time. However, this operation is not GPU-friendly and is thus much slower in practice.
In the paper “DoCoFL: Downlink Compression for Cross-Device Federated Learning” (ICML 2023), the authors performed timing measurements of entropy encoding schemes and found it to be ~ two orders of magnitude slower compared to EDEN, which runs on a GPU (QUIC-FL also runs on a GPU and is faster than EDEN as we show in the paper).

Reviewer comment (Writing):
The Lloyd-Max quantizer appears at several places in the work, e.g., in Introduction, and in serving as the Lower Bound in Section 3.5. However, this work does not introduce this quantizer properly. Can the authors give a brief introduction of this quantizer and in which cases is it optimal?

Author response:
The Lloyd-Max quantizer is a biased quantizer that yields good quantization levels for a given distribution. In some cases, the outcome may be suboptimal (see [A], Example 3.2.1 ). However, the quantizer is known to be optimal for the log-concave distributions (see, e.g., [B]), including the Guassian, which is the focus of our work.

[A] MIT lecture notes https://ocw.mit.edu/courses/6-450-principles-of-digital-communications-i-fall-2006/926689aaa62a0315473fa9b982de1b07_book_3.pdf

[B] Huang, Bormin, and Jing Ma. "On asymptotic solutions of the Lloyd-Max scalar quantization." 2007 6th International Conference on Information, Communications & Signal Processing. IEEE, 2007.
https://ieeexplore.ieee.org/document/4449824.

Reviewer comment (Writing):
In Section 2 preliminaries, it would be clearer if “unbiased algorithms and independent estimators” can be formally and clearly defined. Minor issue: “we that NMSE …” => “we want that NMSE …”

Author response:
We will revise these accordingly.

Reviewer comment (Reproducibility):
The uniform random rotation appears at several places in the work. What exactly is this rotation? Is it a uniform Gaussian random rotation?

Author response:
Random Matrix Theory (RMT) was introduced by Wishart [C] in 1928, and includes rotating the vector by multiplying it with a “rotation matrix”, which is a random orthogonal matrix (see e.g., [D] for more details and methods for its generation. We will discuss this in the paper.

[C] Wishart J. The generalised product moment distribution in samples from a normal multivariate population. Biometrika 20A, 32–52 (1928).
[D] Martin A Tanner and Ronald A Thisted. Generation of Random Orthogonal Matrices. Journal of the Royal Statistical Society. Series C (Applied Statistics), 31(2):190–192, 1982.

Reviewer comment (Writing):
In Section 3.5 “accelerating QUIC-FL with RHT”, “adversarial vectors” are mentioned but not introduced. It is confusing how the proposed approach compares against DRIVE and EDEN in terms of the application to “adversarial vectors”.

Author response:
We meant that there are inputs for which DRIVE and EDEN fail to give an $O(1/n)$ NMSE when using RHT, as we explain in Appendix A. This is not a distinct application but rather an acknowledgment that DRIVE and EDEN are not suitable for worst-case situations.

Reviewer comment (Theory):
In Section 3.5 “accelerating QUIC-FL with RHT”, it states “the result does not have the additive NMSE term is [because] we directly analyze the error for the Hadamard-rotated coordinates”. Can the author be more formal and specific how the analysis here is different from the one in Theorem 3.1?

Author response:
In Theorem 3.1., we analyze the error that arises from approximating the distribution after a uniform random rotation and show that it is manifested in the $O(\sqrt{\log d / d})$ additive error term. The analysis in Section 3.5 directly analyzes the distribution that results from applying RHT, and thus no additional error term is needed.

Reviewer comment (Presentation):
Several places in the work state in plain text the performance of QUIC-FL with different values of its hyperparameters. For example, in Section 3.5 “accelerating QUIC-FL with RHT”, it states the NMSE of QUIC-FL is bounded by … with $b = 1,2,3,4$. It’d be easier for the readers if those numbers can be turned into figures.

Author response:
We will add the figures to the paper.

Reviewer comment (Theory):
This work claims the proposed QUIC-FL has NMSE $O(1/n)$. However, by Theorem 3.1, the actual NMSE is indeed $O(1/n \cdot \sqrt{\log d / d})$. It does not seem to be OK to omit the $O(\sqrt{\log d / d})$ term and directly write $O(1/n \cdot \sqrt{\log d / d}) = O(1/n)$.

Author response:
Note that by Theorem 3.1, the NMSE is  $$NMSE = \frac{1}{n} \cdot \mathbb E[(Z {-} \widehat{Z})^2 ]+ O(\frac{1}{n} \cdot\sqrt{\frac{\log d}{d}}).$$   

We have that $\mathbb E[(Z {-} \widehat{Z})^2 ] = \Theta(1)$ (e.g., $\mathbb E[(Z {-} \widehat{Z})^2 ]=0.00982$ with $b=4$). Therefore, the error is $NMSE = \Theta(\frac{1}{n}) + O(\frac{1}{n} \cdot\sqrt{\frac{\log d}{d}}) = \Theta(\frac{1}{n})$ as $\sqrt{\frac{\log d}{d}} = o(1)$ (recall that $d$ is the dimension of the vector and thus $\frac{\log d}{d} < 1$ for any $d$).

Reviewer comment (Writing):
In Section 3.3 “distribution-aware unbiased quantization”, this work proposes two optimization problems to find the optimal quantization values to reduce NMSE. In the first optimization problem on page 4, the notations $S(z, x)$ and $R(x)$ are a bit confusing. Are $S$ and $R$ two functions to be optimized? Is $R(x)$ essentially a vector of $2^b$ variables? Is $S(z, x)$ a continuous function?

Author response:
Indeed, as shown in the program, the goal is to optimize the error as a function of $S,R$. $R$ is indeed a set of $2^b$ variables (or $2^{b+\ell}$ after the extension to shared randomness in section 3.5) and $S$ can be any probability distribution over the messages in $X$ given an input $z$. Note that we later discretize the problem (pages 5-7), effectively turning $S$ into a discrete set of variables (indeed, $m\cdot 2^b$ for the problem on page 5, increasing to $m\cdot 2^{b+\ell}$ for the problem with shared randomness in page 7).

Reviewer comment (Theory):
In the second optimization problem (i.e., the discretized version of the first problem), is $S’(i, x)$ essentially a vector of $m \cdot 2^b$ variables to be optimized?

Author response:
Indeed, as shown in the program, the goal is to optimize the error as a function of $S,R$. $R$ is indeed a set of $2^b$ variables (or $2^{b+\ell}$ after the extension to shared randomness in section 3.5) and $S$ can be any probability distribution over the messages in $X$ given an input $z$. Note that we later discretize the problem (pages 5-7), effectively turning $S$ into a discrete set of variables (indeed, $m\cdot 2^b$ for the problem on page 5, increasing to $m\cdot 2^{b+\ell}$ for the problem with shared randomness in page 7).

Reviewer comment (Theory):
Since the number of variables to be optimized is on the order of $2^b$, how efficient is the second optimization problem?

Author response:
The complexity of the solver that we use is therefore exponential in $b$ and $\ell$; however, we do not need to run it when encoding or decoding! We ran the solver once for any $b,\ell$ combination offline, allowing us to use its solution on any dimension and input vector.

Reviewer comment (Evaluation):
One concern is that the proposed QUIC-FL requires extra bits to send a few large coordinates exactly along with their indices, while the baseline quantization schemes usually allocate a fixed number of bits per coordinate. This makes QUIC-FL use more communication cost compared to the baseline. And hence it might not be fair to directly compare QUIC-FL’s NMSE against that of the baselines. How does the author compare in Table 1? Also, how does the author address this in the experiments?

Author response:
To make a fair comparison, note that we allow the baselines to send more than $b$ bits per coordinate. As QUIC’s default setting uses $p=1/512$, and we can send each exact coordinate using $64$ bit (in a simple encoding that uses $32$ bits for the value and another $32$ bits for the index), we consider QUIC-FL to be using $b+64/512=b+1/8$ bits per coordinate, e.g., in the leftmost subfigure of Figure 3, which shows the exact bit overhead for each algorithm. 

Also, note that some of the baselines in Figure 4 do not use exactly $b$ bits for coordinate either; e.g., Kashin’s representation requires at least $1.3\cdot b$ bits as it projects the vector into dimension $\lambda\cdot d$ before applying the quantization. Similarly, we did not “charge” QSGD for the sign bit that it sends in addition to each quantized coordinate (i.e., it effectively runs with $b+1$ bits in Figure 4).

Reviewer comment (Theory):
It is mentioned at several places that the optimal NMSE of any quantization is $O(1/n)$ (this lower bound is in terms of the number of clients only, I presume). This is not rigorous in the draft. Can the authors cite the theorems that indicate the optimality?

Author response:
The optimality arises from a $\Omega(1)$ information-theoretic lower bound on vNMSE of any quantization; see, e.g., [E]. We will clarify this in the paper.

[E] Mher Safaryan, Egor Shulgin, and Peter Richtárik. Uncertainty principle for communication compression in distributed and federated learning and the search for an optimal compressor. Information and Inference: A Journal of the IMA, 2020.

Reviewer comment (Theory):
The draft claims QUIC-FL achieves a “near-optimal” NMSE. However, it seems this lower bound is only empirically obtained using the Lloyd-Max quantizer in Section 3.5. “Near-optimal” specifically means one can theoretically show the algorithm achieving optimality (e.g., close to a lower bound), up to a logarithmic factor. And so it might not be appropriate to claim the “near-optimality” of QUIC-FL.

Author response:
Regarding the near-optimality claim: QUIC-FL achieves $O(1/n)$ NMSE using $b= O(1)$, which is asymptotically optimal. We use the term “near-optimal” as the constant is not tight. We also note that Lloyd-Max is known to be optimal for log-concave distributions (including Gaussian), so we get a rigorous bound on the constant gap between the optimum and our solution.

Reviewer comment (Writing):
Page 3: "have" is probably missing between "we" and "that" at the end of the paragraph "Problems and Metric".

Author response:
Regarding the requested changes: 
We will fix the first two. Thank you!

Reviewer comment (Writing):
There is confusion regarding the statement on page 4: "Empirically, sending the indices using ...as $p . \log d << 1$ in our settings, resulting in fast processing time and small bandwidth overhead."

Author response:
Regarding the overhead of encoding the indices – we do consider $p$ to be a constant (e.g., 1/512). You are correct that if we encode the indices explicitly, this will cost $O(d\log d\cdot p)$ bits ($O(d\log d)$ if $p$ is constant). In theory (and in practice), we can simply use $d$ bits to encode the indices in a 0-1 vector for $d$-bit overhead for all indices combined in all cases. In practice though, as $p$ is considerably smaller than $1/\log d$, it is cheaper to send the indices, and this was the point of our statement. 

In summary, the theoretical results can be achieved by sending a bit vector, and thus we achieve the $O(1/n)$ bound with $b=O(1)$. In practice, for reasonable $p$, $d$, explicit encoding of the $p$-fraction of the coordinates is cheaper. 

We will clarify these points in the paper.

Reviewer comment (Theory):
Does this statement mean that $p$ is not kept constant? If that is the case, then shouldn't NMSE have an order of $O(\frac{\log d}{n})$?

Author response:
Regarding the overhead of encoding the indices – we do consider $p$ to be a constant (e.g., 1/512). You are correct that if we encode the indices explicitly, this will cost $O(d\log d\cdot p)$ bits ($O(d\log d)$ if $p$ is constant). In theory (and in practice), we can simply use $d$ bits to encode the indices in a 0-1 vector for $d$-bit overhead for all indices combined in all cases. In practice though, as $p$ is considerably smaller than $1/\log d$, it is cheaper to send the indices, and this was the point of our statement. 

In summary, the theoretical results can be achieved by sending a bit vector, and thus we achieve the $O(1/n)$ bound with $b=O(1)$. In practice, for reasonable $p$, $d$, explicit encoding of the $p$-fraction of the coordinates is cheaper. 

We will clarify these points in the paper.

Reviewer comment (Theory):
I appreciate the discussion after Theorem 3.1 regarding $\mathrm{E}\big[\big(Z-\hat Z\big)^2\big]$. However, it will be still great if the authors provide an explicit error bounds in terms of $b,p,d$. In the current form, it is a bit difficult to provably show the theoretical improvement.

Author response:
Regarding 3.1, the term $\mathbb E[(Z-\widehat Z)^2]$ is independent of $d$ and only depends on $b,p$. We present it in this way to allow one to understand how further improvements in the unbiased quantization of normal random variables (e.g., if one has a more powerful solver or is able to solve the continuous variant) translate directly to improvements in the vNMSE bound. We can provide explicit tables that give bounds for meaningful values of $b,p$. One such result is given in Thm G.3 for running QUIC-FL with RHT. Here, we also give the vNMSE results for QUIC-FL with a uniform random rotation, which are closer to the error in practice even when using RHT:

p = 1/128:

b = 1, $\mathbb E[(Z-\widehat Z)^2]$ = 1.408  
b = 2, $\mathbb E[(Z-\widehat Z)^2]$ = 0.201  
b = 3, $\mathbb E[(Z-\widehat Z)^2]$ = 0.0398  
b = 4, $\mathbb E[(Z-\widehat Z)^2]$ = 0.00859  

p=1/512:

b = 1, $\mathbb E[(Z-\widehat Z)^2]$  = 1.517  
b = 2, $\mathbb E[(Z-\widehat Z)^2]$ = 0.223  
b = 3, $\mathbb E[(Z-\widehat Z)^2]$ = 0.0444  
b = 4, $\mathbb E[(Z-\widehat Z)^2]$ = 0.00982  

In order to get a general formula for uniform random rotations, we can consider uniform quantization (we do at least as good as this). The bound that depends on $b,T_p$ would be  $O(T_p^2 \cdot 4^{-b})$.  
Recall that for $p>0$ we have $T_p = \Phi^{-1}(1-p/2)$, which gives a bound in terms of $p$.

We will add this discussion to the paper.

Reviewer comment (Theory):
Could the authors provide a more explicit bound in Theorem 3.1?

Author response:
Regarding 3.1, the term $\mathbb E[(Z-\widehat Z)^2]$ is independent of $d$ and only depends on $b,p$. We present it in this way to allow one to understand how further improvements in the unbiased quantization of normal random variables (e.g., if one has a more powerful solver or is able to solve the continuous variant) translate directly to improvements in the vNMSE bound. We can provide explicit tables that give bounds for meaningful values of $b,p$. One such result is given in Thm G.3 for running QUIC-FL with RHT. Here, we also give the vNMSE results for QUIC-FL with a uniform random rotation, which are closer to the error in practice even when using RHT:

p = 1/128:

b = 1, $\mathbb E[(Z-\widehat Z)^2]$ = 1.408  
b = 2, $\mathbb E[(Z-\widehat Z)^2]$ = 0.201  
b = 3, $\mathbb E[(Z-\widehat Z)^2]$ = 0.0398  
b = 4, $\mathbb E[(Z-\widehat Z)^2]$ = 0.00859  

p=1/512:

b = 1, $\mathbb E[(Z-\widehat Z)^2]$  = 1.517  
b = 2, $\mathbb E[(Z-\widehat Z)^2]$ = 0.223  
b = 3, $\mathbb E[(Z-\widehat Z)^2]$ = 0.0444  
b = 4, $\mathbb E[(Z-\widehat Z)^2]$ = 0.00982  

In order to get a general formula for uniform random rotations, we can consider uniform quantization (we do at least as good as this). The bound that depends on $b,T_p$ would be  $O(T_p^2 \cdot 4^{-b})$.  
Recall that for $p>0$ we have $T_p = \Phi^{-1}(1-p/2)$, which gives a bound in terms of $p$.

We will add this discussion to the paper.