SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shuli Jiang
Robotics Institute
shulij@andrew.cmu.edu
Pranay Sharma
pranaysh@andrew.cmu.edu
Gauri Joshi
gaurij@andrew.cmu.edu
We study the problem of communication-efficient distributed vector mean estimation, a commonly used subroutine in distributed optimization and Federated Learning (FL). Rand- $k$  sparsification is a commonly used technique to reduce communication cost, where each client sends  $k < d$  of its coordinates to the server. However, Rand- $k$  is agnostic to any correlations, that might exist between clients in practical scenarios. The recently proposed Rand- $k$ -Spatial estimator leverages the cross-client correlation information at the server to improve Rand- $k$ 's performance. Yet, the performance of Rand- $k$ -Spatial is suboptimal. We propose the Rand-Proj-Spatial estimator with a more flexible encoding-decoding procedure, which generalizes the encoding of Rand- $k$  by projecting the client vectors to a random  $k$ -dimensional subspace. We utilize Subsampled Randomized Hadamard Transform (SRHT) as the projection matrix and show that Rand-Proj-Spatial with SRHT outperforms Rand- $k$ -Spatial, using the correlation information more efficiently. Furthermore, we propose an approach to incorporate varying degrees of correlation and suggest a practical variant of Rand-Proj-Spatial when the correlation information is not available to the server. Experiments on real-world distributed optimization tasks showcase the superior performance of Rand-Proj-Spatial compared to Rand- $k$ -Spatial and other more sophisticated sparsification techniques.
In modern machine learning applications, data is naturally distributed across a large number of edge devices or clients. The underlying learning task in such settings is modeled by distributed optimization or the recent paradigm of Federated Learning (FL) [1, 2, 3, 4]. A crucial subtask in distributed learning is for the server to compute the mean of the vectors sent by the clients. In FL, for example, clients run training steps on their local data and once-in-a-while send their local models (or local gradients) to the server, which averages them to compute the new global model. However, with the ever-increasing size of machine learning models [5, 6], and the limited battery life of the edge clients, communication cost is often the major constraint for the clients. This motivates the problem of (empirical) distributed mean estimation (DME) under communication constraints, as illustrated in Figure 1. Each of the  $n$  clients holds a vector  $\mathbf{x}_i \in \mathbb{R}^d$ , on which there are no distributional assumptions. Given a communication budget, each client sends a compressed version  $\widehat{\mathbf{x}}_i$  of its vector to the server, which utilizes these to compute an estimate of the mean vector  $\frac{1}{n} \sum_{i=1}^{n} \mathbf{x}_i$ .
Quantization and sparsification are two major techniques for reducing the communication costs of DME. Quantization [7, 8, 9, 10] involves compressing each coordinate of the client vector to a given precision and aims to reduce the number of bits to represent each coordinate, achieving a constant reduction in the communication cost. However, the communication cost still remains  $\Theta(d)$ .
Sparsification, on the other hand, aims to reduce the number of coordinates each clinet sends and compresses each client vector to only  $k \ll d$  of its coordinates (e.g. Rand-  $k$  [11]). As a result, sparsification reduces communication costs more aggressively compared to quantization, achieving better communication efficiency at a cost of only  $O(k)$ . While in practice, one can use a combination of quantization and sparsification techniques for communication cost reduction, in this work, we focus on the more aggressive sparsification techniques. We call  $k$ , the dimension of the vector each client sends to the server, the per-client communication budget.
Most existing works on sparsification ignore the potential correlation (or similarity) among the client vectors, which often exists in practice. For example, the data of a specific client in federated learning can be similar to that of multiple clients. Hence, it is reasonable to expect their models (or gradients) to be similar as well. To the best of our knowledge, [12] is the first work to account for spatial correlation across individual client vectors. They propose the Rand- $k$ -Spatial family of unbiased estimators, which generalizes Rand- $k$  and achieves a better estimation error in the presence of cross-client correlation. However, their ap
Figure 1: The problem of distributed mean estimation under limited communication. Each client  $i \in [n]$  encodes its vector  $\mathbf{x}_i$  as  $\widehat{\mathbf{x}}_i$  and sends this compressed version to the server. The server decodes them to compute an estimate of the true mean  $\frac{1}{n} \sum_{i=1}^{n} \mathbf{x}_i$ .
proach is focused only on the server-side decoding procedure, while the clients do simple Rand-  $k \times k$  encoding.
In this work, we consider a more general encoding scheme that directly compresses a vector from  $\mathbb{R}^d$  to  $\mathbb{R}^k$  using a (random) linear map. The encoded vector consists of  $k$  linear combinations of the original coordinates. Intuitively, this has a higher chance of capturing the large-magnitude coordinates ("heavy hitters") of the vector than randomly sampling  $k$  out of the  $d$  coordinates (Rand- $k$ ), which is crucial for the estimator to recover the true mean vector. For example, consider a vector where only a few coordinates are heavy hitters. For small  $k$ , Rand- $k$  has a decent chance of missing all the heavy hitters. But with a linear-maps-based general encoding procedure, the large coordinates are more likely to be encoded in the linear measurements, resulting in a more accurate estimator of the mean vector. Guided by this intuition, we ask:
Can we design an improved joint encoding-decoding scheme that utilizes the correlation information and achieves an improved estimation error?
One naive solution is to apply the same random rotation matrix  $G \in \mathbb{R}^{d \times d}$  to each client vector, before applying Rand-  $k$  or Rand-  $k$ -Spatial encoding. Indeed, such preprocessing is applied to improve the estimator using quantization techniques on heterogeneous vectors [13, 10]. However, as we see in Appendix A.1, for sparsification, we can show that this leads to no improvement. But what happens if every client uses a different random matrix, or applies a random  $k \times d$ -dimensional linear map? How to design the corresponding decoding procedure to leverage cross-client correlation? As there is no way for one to directly apply the decoding procedure of Rand-  $k$ -Spatial in such cases. To answer these questions, we propose the Rand-Proj-Spatial family estimator. We propose a flexible encoding procedure in which each client applies its own random linear map to encode the vector. Further, our novel decoding procedure can better leverage cross-client correlation. The resulting mean estimator generalizes and improves over the Rand-  $k$ -Spatial family estimator.
Next, we discuss some reasonable restrictions we expect our mean estimator to obey. 1) Unbiased. An unbiased mean estimator is theoretically more convenient compared to a biased one [14]. 2) Non-adaptive. We focus on an encoding procedure that does not depend on the actual client data, as opposed to the adaptive ones, e.g. Rand-  $k$  with vector-based sampling probability [11, 15]. Designing a data-adaptive encoding procedure is computationally expensive as this might require using an iterative procedure to find out the sampling probabilities [11]. In practice, however, clients often have limited computational power compared to the server. Further, as discussed earlier, mean estimation is often a subroutine in more complicated tasks. For applications with streaming data [16],
the additional computational overhead of adaptive schemes is challenging to maintain. Note that both Rand- $k$  and Rand- $k$ -Spatial family estimator [12] are unbiased and non-adaptive.
In this paper, we focus on the severely communication-constrained case  $nk \leq d$ , when the server receives very limited information about any single client vector. If  $nk \gg d$ , we see in Appendix A.2 that the cross-client information has no additional advantage in terms of improving the mean estimate under both Rand- $k$ -Spatial or Rand-Proj-Spatial, with different choices of random linear maps. Furthermore, when  $nk \gg d$ , the performance of both the estimators converges to that of Rand- $k$ . Intuitively, this means when the server receives sufficient information regarding the client vectors, it does not need to leverage cross-client correlation to improve the mean estimator.
1. We propose the Rand-Proj-Spatial family estimator with a more flexible encoding-decoding procedure, which can better leverage the cross-client correlation information to achieve a more general and improved mean estimator compared to existing ones.
2. We show the benefit of using Subsampled Randomized Hadamard Transform (SRHT) as the random linear maps in Rand-Proj-Spatial in terms of better mean estimation error (MSE). We theoretically analyze the case when the correlation information is known at the server (see Theorems 4.3, 4.4 and Section 4.3). Further, we propose a practical configuration called Rand-Proj-Spatial(Avg) when the correlation is unknown.
3. We conduct experiments on common distributed optimization tasks, and demonstrate the superior performance of Rand-Proj-Spatial compared to existing sparsification techniques.
Quantization and Sparsification. Commonly used techniques to achieve communication efficiency are quantization, sparsification, or more generic compression schemes, which generalize the former two [17]. Quantization involves either representing each coordinate of the vector by a small number of bits [8, 9, 10, 18, 19, 20], or more involved vector quantization techniques [21, 22]. Sparsification [15, 23, 24, 25, 26], on the other hand, involves communicating a small number  $k < d$  of coordinates, to the server. Common protocols include Rand-  $k$  [11], sending  $k$  uniformly randomly selected coordinates; Top-  $k$  [27], sending the  $k$  largest magnitude coordinates; and a combination of the two [28]. Some recent works, with a focus on distributed learning, further refine these communication-saving mechanisms [29] by incorporating temporal correlation or error feedback [14, 25].
Distributed Mean Estimation (DME). DME has wide applications in distributed optimization and FL. Most of the existing literature on DME either considers statistical mean estimation [30, 31], assuming that the data across clients is generated i.i.d. according to the same distribution, or empirical mean estimation [10, 32, 33, 12, 11, 34, 35], without making any distributional assumptions on the data. A recent line of work on empirical DME considers applying additional information available to the server, to further improve the mean estimate. This side information includes cross-client correlation [12, 13], or the memory of the past updates sent by the clients [36].
Subsampled Randomized Hadamard Transformation (SRHT). SRHT was introduced for random dimensionality reduction using sketching [37, 38, 39]. Common applications of SRHT include faster computation of matrix problems, such as low-rank approximation [40, 41], and machine learning tasks, such as ridge regression [42], and least square problems [43, 44, 45]. SRHT has also been applied to improve communication efficiency in distributed optimization [46] and FL [47, 48].
Notation. We use bold lowercase (uppercase) letters, e.g.  $\mathbf{x}(G)$  to denote vectors (matrices).  $\mathbf{e}_j\in \mathbb{R}^d$  , for  $j\in [d]$  , denotes the  $j$  -th canonical basis vector.  $\| \cdot \| _2$  denotes the Euclidean norm. For a vector  $\mathbf{x}$ $\mathbf{x}(j)$  denotes its  $j$  -th coordinate. Given integer  $m$  , we denote by  $[m]$  the set  $\{1,2,\dots,m\}$
Problem Setup. Consider  $n$  geographically separated clients coordinated by a central server. Each client  $i \in [n]$  holds a vector  $\mathbf{x}_i \in \mathbb{R}^d$ , while the server wants to estimate the mean vector  $\bar{\mathbf{x}} \triangleq \frac{1}{n} \sum_{i=1}^{n} \mathbf{x}_i$ . Given a per-client communication budget of  $k \in [d]$ , each client  $i$  computes  $\hat{\mathbf{x}}_i$  and sends it to the central server.  $\hat{\mathbf{x}}_i$  is an approximation of  $\mathbf{x}_i$  that belongs to a random  $k$ -dimensional
subspace. Each client also sends a random seed to the server, which conveys the subspace information, and can usually be communicated using a negligible amount of bits. Having received the encoded vectors  $\{\widehat{\mathbf{x}}_i\}_{i=1}^n$ , the server then computes  $\widehat{\mathbf{x}} \in \mathbb{R}^d$ , an estimator of  $\bar{\mathbf{x}}$ . We consider the severely communication-constrained setting where  $nk \leq d$ , when only a limited amount of information about the client vectors is seen by the server.
Error Metric. We measure the quality of the decoded vector  $\widehat{\mathbf{x}}$  using the Mean Squared Error (MSE)  $\mathbb{E}\left[\| \widehat{\mathbf{x}} -\bar{\mathbf{x}}\| _2^2\right]$ , where the expectation is with respect to all the randomness in the encoding-decoding scheme. Our goal is to design an encoding-decoding algorithm to achieve an unbiased estimate  $\widehat{\mathbf{x}}$  (i.e.  $\mathbb{E}[\widehat{\mathbf{x}}] = \bar{\mathbf{x}}$ ) that minimizes the MSE, given the per-client communication budget  $k$ . To consider an example, in rand- $k$  sparsification, each client sends randomly selected  $k$  out of its  $d$  coordinates to the server. The server then computes the mean estimate as  $\widehat{\mathbf{x}}^{(\mathrm{Rand - }k)} = \frac{1}{n}\frac{d}{k}\sum_{i = 1}^{n}\widehat{\mathbf{x}}_i$ . By [12, Lemma 1], the MSE of Rand- $k$  sparsification is given by
$$
\mathbb {E} \left[ \| \widehat {\mathbf {x}} ^ {(\text {R a n d -} k)} - \bar {\mathbf {x}} \| _ {2} ^ {2} \right] = \frac {1}{n ^ {2}} \left(\frac {d}{k} - 1\right) \sum_ {i = 1} ^ {n} \| \mathbf {x} _ {i} \| _ {2} ^ {2} \tag {1}
$$
The Rand- $k$ -Spatial Family Estimator. For large values of  $\frac{d}{k}$ , the Rand- $k$  MSE in Eq. 1 can be prohibitive. [12] proposed the Rand- $k$ -Spatial family estimator, which achieves an improved MSE, by leveraging the knowledge of the correlation between client vectors at the server. The encoded vectors  $\{\hat{\mathbf{x}}_i\}$  are the same as in Rand- $k$ . However, the  $j$ -th coordinate of the decoded vector is given as
$$
\widehat {\mathbf {x}} ^ {(\text {R a n d - k - S p a t i a l})} (j) = \frac {1}{n} \frac {\bar {\beta}}{T (M _ {j})} \sum_ {i = 1} ^ {n} \widehat {\mathbf {x}} _ {i} (j) \tag {2}
$$
Here,  $T: \mathbb{R} \to \mathbb{R}$  is a pre-defined transformation function of  $M_j$ , the number of clients which sent their  $j$ -th coordinate, and  $\bar{\beta}$  is a normalization constant to ensure  $\hat{\mathbf{x}}$  is an unbiased estimator of  $\mathbf{x}$ . The resulting MSE is given by
$$
\mathbb {E} \left[ \| \widehat {\mathbf {x}} ^ {(\text {R a n d - k - S p a t i a l})} - \bar {\mathbf {x}} \| _ {2} ^ {2} \right] = \frac {1}{n ^ {2}} \left(\frac {d}{k} - 1\right) \sum_ {i = 1} ^ {n} \| \mathbf {x} _ {i} \| _ {2} ^ {2} + \left(c _ {1} \sum_ {i = 1} ^ {n} \| \mathbf {x} _ {i} \| _ {2} ^ {2} - c _ {2} \sum_ {i = 1} ^ {n} \sum_ {l \neq i} \langle \mathbf {x} _ {i}, \mathbf {x} _ {l} \rangle\right) \tag {3}
$$
where  $c_{1}, c_{2}$  are constants dependent on  $n, d, k$  and  $T$ , but independent of client vectors  $\{\mathbf{x}_i\}_{i=1}^n$ . When the client vectors are orthogonal, i.e.,  $\langle \mathbf{x}_i, \mathbf{x}_l \rangle = 0$ , for all  $i \neq l$ , [12] show that with appropriately chosen  $T$ , the MSE in Eq. 3 reduces to Eq. 1. However, if there exists a positive correlation between the vectors, the MSE in Eq. 3 is strictly smaller than that for Rand- $k$  Eq. 1.
4 The Rand-Proj-Spatial Family Estimator
While the Rand- $k$ -Spatial family estimator proposed in [12] focuses only on improving the decoding at the server, we consider a more general encoding-decoding scheme. Rather than simply communicating  $k$  out of the  $d$  coordinates of its vector  $\mathbf{x}_i$  to the server, client  $i$  applies a (random) linear map  $\mathbf{G}_i \in \mathbb{R}^{k \times d}$  to  $\mathbf{x}_i$  and sends  $\widehat{\mathbf{x}}_i = \mathbf{G}_i \mathbf{x}_i \in \mathbb{R}^k$  to the server. The decoding process on the server first projects the encoded vectors  $\{\mathbf{G}_i \mathbf{x}_i\}_{i=1}^n$  back to the  $d$ -dimensional space and then forms an estimate  $\widehat{\mathbf{x}}$ . We motivate our new decoding procedure with the following regression problem:
$$
\widehat {\mathbf {x}} ^ {(\text {R a n d - P r o j})} = \arg \min  _ {\mathbf {x}} \sum_ {i = 1} ^ {n} \| G _ {i} \mathbf {x} - G _ {i} \mathbf {x} _ {i} \| _ {2} ^ {2} \tag {4}
$$
To understand the motivation behind Eq. 4, first consider the special case where  $\mathbf{G}_i = \mathbf{I}_d$  for all  $i \in [n]$ , that is, the clients communicate their vectors without compressing. The server can then exactly compute the mean  $\bar{\mathbf{x}} = \frac{1}{n} \sum_{i=1}^{n} \mathbf{x}_i$ . Equivalently,  $\bar{\mathbf{x}}$  is the solution of  $\arg \min_{\mathbf{x}} \sum_{i=1}^{n} \| \mathbf{x} - \mathbf{x}_i \|_2^2$ . In the more general setting, we require that the mean estimate  $\hat{\mathbf{x}}$  when encoded using the map  $\mathbf{G}_i$  should be "close" to the encoded vector  $\mathbf{G}_i \mathbf{x}_i$  originally sent by client  $i$ , for all clients  $i \in [n]$ .
We note the above intuition can also be translated into different regression problems to motivate the design of the new decoding procedure. We discuss in Appendix B.2 intuitive alternatives which,
Unfortunately, either do not enable the usage of cross-client correlation information, or do not use such information effectively. We choose the formulation in Eq. 4 due to its analytical tractability and its direct relevance to our target error metric MSE. We note that it is possible to consider the problem in Eq. 4 in the other norms, such as the sum of  $\ell_2$  norms (without the squares) or the  $\ell_{\infty}$  norm. We leave this as a future direction to explore.
The solution to Eq. 4 is given by  $\widehat{\mathbf{x}}^{(\mathrm{Rand - Proj})} = (\sum_{i = 1}^{n}\pmb{G}_{i}^{T}\pmb{G}_{i})^{\dagger}\sum_{i = 1}^{n}\pmb{G}_{i}^{T}\pmb{G}_{i}\mathbf{x}_{i}$ , where  $\dagger$  denotes the Moore-Penrose pseudo inverse [49]. However, while  $\widehat{\mathbf{x}}^{(\mathrm{Rand - Proj})}$  minimizes the error of the regression problem, our goal is to design an unbiased estimator that also improves the MSE. Therefore, we make the following two modifications to  $\widehat{\mathbf{x}}^{(\mathrm{Rand - Proj})}$ : First, to ensure that the mean estimate is unbiased, we scale the solution by a normalization factor  $\bar{\beta}^1$ . Second, to incorporate varying degrees of correlation among the clients, we propose to apply a scalar transformation function  $T:\mathbb{R}\to \mathbb{R}$  to each of the eigenvalues of  $\sum_{i = 1}^{n}\pmb{G}_{i}^{T}\pmb{G}_{i}$ . The resulting Rand-Proj-Spatial family estimator is given by
$$
\widehat {\mathbf {x}} ^ {(\text {R a n d - P r o j - S p a t i a l})} = \bar {\beta} \left(T \left(\sum_ {i = 1} ^ {n} \boldsymbol {G} _ {i} ^ {T} \boldsymbol {G} _ {i}\right)\right) ^ {\dagger} \sum_ {i = 1} ^ {n} \boldsymbol {G} _ {i} ^ {T} \boldsymbol {G} _ {i} \mathbf {x} _ {i} \tag {5}
$$
Though applying the transformation function  $T$  in Rand-Proj-Spatial requires computing the eigen-decomposition of  $\sum_{i=1}^{n} G_i^T G_i$ . However, this happens only at the server, which has more computational power than the clients. Next, we observe that for appropriate choice of  $\{G_i\}_{i=1}^n$ , the Rand-Proj-Spatial family estimator reduces to the Rand- $k$ -Spatial family estimator [12].
Lemma 4.1 (Recovering Rand- $k$ -Spatial). Suppose client  $i$  generates a subsampling matrix  $\mathbf{E}_i = [\mathbf{e}_{i_1}, \ldots, \mathbf{e}_{i_k}]^\top$ , where  $\{\mathbf{e}_j\}_{j=1}^d$  are the canonical basis vectors, and  $\{i_1, \ldots, i_k\}$  are sampled from  $\{1, \ldots, d\}$  without replacement. The encoded vectors are given as  $\widehat{\mathbf{x}}_i = \mathbf{E}_i\mathbf{x}_i$ . Given a function  $T$ ,  $\widehat{\mathbf{x}}$  computed as in Eq. 5 recovers the Rand- $k$ -Spatial estimator.
The proof details are in Appendix C.5. We discuss the choice of  $T$  and how it compares to Rand- $k$ -Spatial in detail in Section 4.3.
Remark 4.2. In the simple case when  $G_{i}$ 's are subsampling matrices (as in Rand-k-Spatial [12]), the  $j$ -th diagonal entry of  $\sum_{i=1}^{n} G_{i}^{T} G_{i}$ ,  $M_{j}$  conveys the number of clients which sent the  $j$ -th coordinate. Rand-k-Spatial incorporates correlation among client vectors by applying a function  $T$  to  $M_{j}$ . Intuitively, it means scaling different coordinates differently. This is in contrast to Rand-k, which scales all the coordinates by  $d / k$ . In our more general case, we apply a function  $T$  to the eigenvalues of  $\sum_{i=1}^{n} G_{i}^{T} G_{i}$  to similarly incorporate correlation in Rand-Proj-Spatial.
To showcase the utility of the Rand-Proj-Spatial family estimator, we propose to set the random linear maps  $\pmb{G}_i$  to be scaled Subsampled Randomized Hadamard Transform (SRHT, e.g. [38]). Assuming  $d$  to be a power of 2, the linear map  $\pmb{G}_i$  is given as
$$
\boldsymbol {G} _ {i} = \frac {1}{\sqrt {d}} \boldsymbol {E} _ {i} \boldsymbol {H} \boldsymbol {D} _ {i} \in \mathbb {R} ^ {k \times d} \tag {6}
$$
where  $E_{i} \in \mathbb{R}^{k \times d}$  is the subsampling matrix,  $H \in \mathbb{R}^{d \times d}$  is the (deterministic) Hadamard matrix and  $D_{i} \in \mathbb{R}^{d \times d}$  is a diagonal matrix with independent Rademacher random variables as its diagonal entries. We choose SRHT due to its superior performance compared to other random matrices. Other possible choices of random matrices for Rand-Proj-Spatial estimator include sketching matrices commonly used for dimensionality reduction, such as Gaussian [50, 51], row-normalized Gaussian, and Count Sketch [52], as well as error-correction coding matrices, such as Low-Density Parity Check (LDPC) [53] and Fountain Codes [54]. However, in the absence of correlation between client vectors, all these matrices suffer a higher MSE.
In the following, we first compare the MSE of Rand-Proj-Spatial with SRHT against Rand-  $k$  and Rand-  $k$ -Spatial in two extreme cases: when all the client vectors are identical, and when all the client vectors are orthogonal to each other. In both cases, we highlight the transformation function  $T$  used in Rand-Proj-Spatial (Eq. 5) to incorporate the knowledge of cross-client correlation. We define
$$
\mathcal {R} := \frac {\sum_ {i = 1} ^ {n} \sum_ {l \neq i} \left\langle \mathbf {x} _ {i} , \mathbf {x} _ {l} \right\rangle}{\sum_ {i = 1} ^ {n} \| \mathbf {x} _ {i} \| _ {2} ^ {2}} \tag {7}
$$
to measure the correlation between the client vectors. Note that  $\mathcal{R} \in [-1, n - 1]$ .  $\mathcal{R} = 0$  implies all client vectors are orthogonal, while  $\mathcal{R} = n - 1$  implies identical client vectors.
<sup>1</sup>We show that it suffices for  $\bar{\beta}$  to be a scalar in Appendix B.1.
4.1 Case I: Identical Client Vectors  $(\mathcal{R} = n - 1)$
When all the client vectors are identical  $(\mathbf{x}_i \equiv \mathbf{x})$ , [12] showed that setting the transformation  $T$  to identity, i.e.,  $T(m) = m$ , for all  $m$ , leads to the minimum MSE in the Rand- $k$ -Spatial family of estimators. The resulting estimator is called Rand- $k$ -Spatial (Max). Under the same setting, using the same transformation  $T$  in Rand-Proj-Spatial with SRHT, the decoded vector in Eq. 5 simplifies to