DRIVE: One-bit Distributed Mean Estimation

Abstract:
We consider the problem where n clients transmit d-dimensional real-valued vectors using dp1 ` op1qq bits each, in a manner that allows the receiver to approximately reconstruct their mean. Such compression problems naturally arise in distributed and federated learning. We provide novel mathematical results and derive computationally efficient algorithms that are more accurate than previous compression techniques. We evaluate our methods on a collection of distributed and federated learning tasks, using a variety of datasets, and show a consistent improvement over the state of the art.

heading: 1 Introduction
text: In many computational settings, one wishes to transmit a d-dimensional real-valued vector. For example, in distributed and federated learning scenarios, multiple participants (a.k.a. clients) in distributed SGD send gradients to a parameter server that averages them and updates the model parameters accordingly [1]. In these applications and others (e.g., traditional machine learning methods such K-Means and power iteration [2] or other methods such as geometric monitoring [3]), sending approximations of vectors may suffice. Moreover, the vectors’ dimension d is often large (e.g., in neural networks, d can exceed a billion [4, 5, 6]), so sending compressed vectors is appealing.
Indeed, recent works have studied how to send vector approximations using representations that use a small number of bits per entry (e.g., [2, 7, 8, 9, 10, 11]). Further, recent work has shown direct training time reduction from compressing the vectors to one bit per coordinate [12]. Most relevant to our work are solutions that address the distributed mean estimation problem. For example, [2] uses the randomized Hadamard transform followed by stochastic quantization (a.k.a. randomized rounding). When each of the n clients transmits Opdq bits, their Normalized Mean Squared Error (NMSE) is bounded by O `
log d n
˘
. They also show a Op 1n q bound with Opdq bits via variable-length encoding, albeit at a higher computational cost. The sampling method of [10] yields an O `
r¨R n ˘
NMSE bound using dp1 ` op1qq bits in expectation, where r is each coordinate’s representation length and R is the normalized average variance of the sent vectors. Recently, researchers proposed to use Kashin’s representation [11, 13, 14]. Broadly speaking, it allows representing a d-dimensional vector in (a higher) dimension λ ¨ d for some λ ą 1 using small coefficients. This results in an O ´ λ2
p ? λ´1q4¨n
¯
NMSE bound, where each client transmits λ ¨ dp1` op1qq bits [14]. A recent work [15] suggested an algorithm where if all clients’ vectors have pairwise distances of at most y P R (i.e., for any client pair c1, c2, it holds that ∥∥xpc1q ´ xpc2q∥∥2 ď y), the resulting MSE is Opy2q (which is tight with respect to ˚Equal Contribution.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
y) using Op1q bits per coordinate on average. This solution provides a stronger MSE bound when vectors are sufficiently close (and thus y is small) but does not improve the worst-case guarantee.
We step back and focus on approximating d-dimensional vectors using dp1` op1qq bits (e.g., one bit per dimension and a lower order overhead). We develop novel biased and unbiased compression techniques based on (uniform as well as structured) random rotations in high-dimensional spheres. Intuitively, after a rotation, the coordinates are identically distributed, allowing us to estimate each coordinate with respect to the resulting distribution. Our algorithms do not require expensive operations, such as variable-length encoding or computing the Kashin’s representation, and are fast and easy to implement. We obtain an O `
1 n
˘
NMSE bound using dp1`op1qq bits, regardless of the coordinates’ representation length, improving over previous works. Evaluation results indicate that this translates to a consistent improvement over the state of the art in different distributed and federated learning tasks.

heading: 2 Problem Formulation and Notation
text: 1b - Vector Estimation. We start by formally defining the 1b - vector estimation problem. A sender, called Buffy, gets a real-valued vector x P Rd and sends it using a dp1 ` op1qq bits message (i.e., asymptotically one bit per coordinate). The receiver, called Angel, uses the message to derive an estimate px of the original vector x. We are interested in the quantity ‖x´ px‖22, which is the sum of squared errors (SSE), and its expected value, the Mean Squared Error (MSE). For ease of exposition, we hereafter assume that x ‰ 0 as this special case can be handled with one additional bit. Our goal is to minimize the vector-NMSE (denoted vNMSE), defined as the normalized MSE, i.e.,
Er‖x´px‖22s ‖x‖22 .
1b - Distributed Mean Estimation. The above problem naturally generalizes to the 1b - Distributed Mean Estimation problem. Here, we have a set of n PN clients and a server. Each client c P t1, . . . , nu has its own vector xpcq PRd, which it sends using a dp1`op1qq-bits message to the server. The server then produces an estimate yxavg PRd of the average xavg “ 1n řn c“1 xpcq with the goal of minimizing its NMSE, defined as the average estimate’s MSE normalized by the average norm of the clients’ original vectors, i.e., Er‖xavg´yxavg‖22s 1 n ¨ řn c“1‖xpcq‖22 .
Notation. We use the following notation and definitions throughout the paper:
Subscripts. xi denotes the i’th coordinate of the vector x, to distinguish it from client c’s vector xpcq.
Binary-sign. For a vector x P Rd, we denote its binary-sign function as signpxq, where signpxqi “ 1 if xi ě 0 and signpxqi “ ´1 if xi ă 0.
Unit vector. For any (non-zero) real-valued vector x P Rd, we denote its normalized vector by x̆ “ x‖x‖2 . That is, x̆ and x has the same direction and it holds that ‖x̆‖2 “ 1.
Rotation Matrix. A matrix R P Rdˆd is a rotation matrix if RTR “ I . The set of all rotation matrices is denoted as Opdq. It follows that @R P Opdq: detpRq P t´1, 1u and @x P Rd: ‖x‖2“ ‖Rx‖2. Random Rotation. A random rotation R is a distribution over all random rotations in Opdq. For ease of exposition, we abuse the notation and given x P Rd denote the random rotation of x by Rpxq “ Rx, where R is drawn from R. Similarly, R´1pxq “ R´1x “ RTx is the inverse rotation.
Rotation Property. A quantity that determines the guarantees of our algorithms is LdR,x “ ‖Rpx̆q‖21 d (note the use of the L1 norm). We show that rotations with high LdR,x values yield better estimates. Shared Randomness. We assume that Buffy and Angel have access to shared randomness, e.g., by agreeing on a common PRNG seed. Shared randomness is studied both in communication complexity (e.g., [16]) and in communication reduction in machine learning systems (e.g., [2, 7]). In our context, it means that Buffy and Angel can generate the same random rotations without communication.

heading: 3 The DRIVE Algorithm
text: We start by presenting DRIVE (Deterministically RoundIng randomly rotated VEctors), a novel 1b - Vector Estimation algorithm. Later, we extend DRIVE to the 1b - Distributed Mean Estimation
Algorithm 1 DRIVE
Buffy: 1: Compute Rpxq, S. `
2: Send ` S, signpRpxqq ˘ to Angel.zRpxq
Angel: 1: Compute zRpxq “ S ¨ sign ` Rpxq ˘ .
2: Estimate px “ R´1 ` zRpxq ˘ .
problem. In DRIVE, Buffy uses shared randomness to sample a rotation matrix R „ R and rotates the vector x P Rd by computing Rpxq “ Rx. Buffy then calculates S, a scalar quantity we explain below. Buffy then sends ` S, signpRpxqq ˘ to Angel. As we discuss later, sending ` S, signpRpxqq ˘
requires dp1 ` op1qq bits. In turn, Angel computes zRpxq “ S ¨ signpRpxqq P t´S,`Sud. It then uses the shared randomness to generate the same rotation matrix and employs the inverse rotation, i.e., estimates px “ R´1pzRpxqq. The pseudocode of DRIVE appears in Algorithm 1. The properties of DRIVE depend on the rotation R and the scale parameter S. We consider both uniform rotations, that provide stronger guarantees, and structured rotations that are orders of magnitude faster to compute. As for the scale S “ Spx,Rq, its exact formula determines the characteristics of DRIVE’s estimate, e.g., having minimal vNMSE or being unbiased. The latter allows us to apply DRIVE to the 1b - Distributed Mean Estimation (Section 4.2) and get an NMSE that decreases proportionally to the number of clients.
We now prove a general result on the SSE of DRIVE that applies to any random rotation R and any vector x P Rd. In the following sections, we use this result to obtain the vNMSE when considering specific rotations and scaling methods as well as analyzing their guarantees.
Theorem 1. The SSE of DRIVE is: ‖x´ px‖22 “ ‖x‖ 2 2 ´ 2 ¨ S ¨ ‖Rpxq‖1 ` d ¨ S2 .
Proof. The SSE in estimating Rpxq using zRpxq equals that of estimating x using px. Therefore, ‖x´ px‖22 “ ‖Rpx´ pxq‖ 2 2 “ ‖Rpxq ´Rppxq‖ 2 2 “ ∥∥∥Rpxq ´ zRpxq∥∥∥2 2
“ ‖Rpxq‖22 ´ 2 A Rpxq, zRpxq E ` ∥∥∥zRpxq∥∥∥2
2 “ ‖x‖22 ´ 2
A Rpxq, zRpxq E ` ∥∥∥zRpxq∥∥∥2
2 . (1)
Next, we have that, A
Rpxq, zRpxq E
“ řd i“1 Rpxqi ¨ zRpxqi “ S ¨ řd i“1 Rpxqi ¨ sign ` Rpxqi ˘ “ S ¨ ‖Rpxq‖1 , (2)∥∥∥zRpxq∥∥∥2 2 “ řd i“1 zRpxq 2 i “ d ¨ S2 . (3)
Substituting Eq. (2) and Eq. (3) in Eq. (1) yields the result.

heading: 4 DRIVE With a Uniform Random Rotation
text: We first consider the thoroughly studied uniform random rotation (e.g., [17, 18, 19, 20, 21]), which we denote by RU . The sampled matrix is denoted by RU „ RU , that is, RU pxq “ RU ¨ x. An appealing property of a uniform random rotation is that, as we show later, it admits a scaling that results in a low constant vNMSE even with unbiased estimates.

heading: 4.1 1b - Vector Estimation
text: Using Theorem 1, we obtain the following result. The result holds for any rotation, including RU . Lemma 1. For any x P Rd, DRIVE’s SSE is minimized by S “ ‖Rpxq‖1d (that is, S “ ‖Rx‖1 d is determined after R „ R is sampled). This yields a vNMSE of Er‖x´px‖ 2 2s
‖x‖22 “ 1´ E
“ LdR,x ‰ .
Proof. By Theorem 1, to minimize the SSE we require B BS ` ‖x‖22 ´ 2 ¨ S ¨ ‖Rpxq‖1 ` d ¨ S2 ˘ “ ´2 ¨ ‖Rpxq‖1 ` 2 ¨ d ¨ S “ 0 ,
leading to S “ ‖Rpxq‖1d . Then, the SSE of DRIVE becomes:
‖x´ px‖22 “ ‖x‖ 2 2 ´ 2 ¨ S ¨ ‖Rpxq‖1 ` d ¨ S2 “ ‖x‖ 2 2 ´ 2 ¨ ‖Rpxq‖21 d ` d ¨ ‖Rpxq‖21 d2
“ ‖x‖22 ´ ‖Rpxq‖21 d “ ‖x‖ 2 2 ´ ‖x‖22¨‖Rpx̆q‖ 2 1 d “ ‖x‖ 2 2 ` 1´ ‖Rpx̆q‖ 2 1 d ˘ .
Thus, the normalized SSE is ‖x´px‖ 2 2
‖x‖22 “ 1´LdR,x. Taking expectation yields the result.
Interestingly, for the uniform random rotation, LdRU ,x follows the same distribution for all x. This is because, by the definition of RU , it holds that RU px̆q is distributed uniformly over the unit sphere for any x. Therefore DRIVE’s vNMSE depends only on the dimension d. We next analyze the vNMSE attainable by the best possible S, as given in Lemma 1, when the algorithm uses RU and is not required to be unbiased. In particular, we state the following theorem whose proof appears in Appendix A.1 (all appendices appear in the Supplementary Material and the extended paper version [22]).
Theorem 2. For any x P Rd, the vNMSE of DRIVE with S “ ‖RU pxq‖1d is ` 1´ 2π ˘ ` 1´ 1d ˘ .

heading: 4.2 1b - Distributed Mean Estimation
text: An appealing property of DRIVE with a uniform random rotation, established in this section, is that with a proper scaling parameter S, the estimate is unbiased. That is, for any x P Rd, our scale guarantees that E rpxs “ x. Unbiasedness is useful when generalizing to the Distributed Mean Estimation problem. Intuitively, when n clients send their vectors, any biased algorithm would result in an NMSE that may not decrease with respect to n. For example, if they have the same input vector, the bias would remain after averaging. Instead, an unbiased encoding algorithm has the property that when all clients act (e.g., use different PRNG seeds) independently, the NMSE decreases proportionally to 1n .
Another useful property of uniform random rotation is that its distribution is unchanged when composed with other rotations. We use it in the following theorem’s proof, given in Appendix A.2.
Theorem 3. For any x P Rd, set S “ ‖x‖ 2 2
‖RU pxq‖1 . Then DRIVE satisfies Erpxs “ x.
Now, we proceed to obtain vNMSE guarantees for DRIVE’s unbiased estimate.
Lemma 2. For any x P Rd, DRIVE with S “ ‖x‖ 2 2
‖RU pxq‖1 has a vNMSE of E
„
1 LdRU,x

´1 .
Proof. By Theorem 1, the SSE of the algorithm satisfies:
‖x‖22 ´ 2 ¨ S ¨ ‖RU ¨ x‖1 ` d ¨ S 2 “ ‖x‖22 ´ 2 ¨ ‖x‖22 ‖RU ¨ x‖1 ¨ ‖RU ¨ x‖1 ` d ¨ ˜ ‖x‖22 ‖RU ¨ x‖1 ¸2
“ d ¨ ˜ ‖x‖22 ‖x‖2 ‖RU ¨ x̆‖1 ¸2 ´ ‖x‖22 “ d ¨ ‖x‖22 ‖RU ¨ x̆‖21 ´ ‖x‖22 “ ‖x‖ 2 2 ¨ ˜˜
d
‖RU ¨ x̆‖21
¸ ´1 ¸ .
Normalizing by ‖x‖22 and taking expectation over RU concludes the proof.
Our goal is to derive an upper bound on the above expression and thus upper-bound the vNMSE. Most importantly, we show that even though the estimate is unbiased and we use only a single bit per coordinate, the vNMSE does not increase with the dimension and is bounded by a small constant. In particular, in Appendix A.3, we prove the following:
Theorem 4. For any x P Rd, the vNMSE of DRIVE with S “ ‖x‖ 2 2
‖RU pxq‖1 satisfies:
piq For all d ě 2, it is at most 2.92. piiq For all d ě 135, it is at most π2 ´ 1` b p6π3´12π2q¨ln d`1 d .
This theorem yields strong bounds on the vNMSE. For example, the vNMSE is lower than 1 for d ě 4096 and lower than 0.673 for d ě 105. Finally, we obtain the following corollary, Corollary 1. For any x P Rd, the vNMSE tends to π2 ´ 1 « 0.571 as dÑ8 .
Recall that DRIVE’s above scale S is a function of both x and the sampled RU . An alternative approach is to deterministically set S to ‖x‖ 2 2
Er‖RU pxq‖1s . As we prove in Appendix A.4, the resulting
scale is ‖x‖2¨pd´1q¨Bp 1 2 , d´1 2 q
2d , where B is the Beta function. Interestingly, this scale no longer depends on x but only on its norm. In the appendix, we also prove that the resulting vNMSE is bounded by π 2 ´ 1 for any d. In practice, we find that the benefit is marginal. Finally, with a vNMSE guarantee for the unbiased estimate by DRIVE, we obtain the following key result for the 1b - Distributed Mean Estimation problem, whose proof appears in Appendix A.5. We note that this result guarantees (e.g., see [23]) that distributed SGD, where the participants’ gradients are compressed with DRIVE, converges at the same asymptotic rate as without compression.
Theorem 5. Assume n clients, each with its own vector xpcq PRd. Let each client independently sample RU,c„RU and set its scale to ‖ xpcq‖22
‖RU,c¨xpcq‖1 . Then, the server average estimate’s NMSE
satisfies: Er‖xavg´yxavg‖22s 1 n ¨ řn c“1‖xpcq‖22 “ vNMSEn , where vNMSE is given by Lemma 2 and is bounded by Theorem 4.
To the best of our knowledge, DRIVE is the first algorithm with a provable NMSE of Op 1n q for the 1b - Distributed Mean Estimation problem (i.e., with dp1 ` op1qq bits). In practice, we use only d`Op1q bits to implement DRIVE. We use the dp1` op1qq notation to ensure compatibility with the theoretical results; see Appendix B for a discussion.
5 Reducing the vNMSE with DRIVE`
To reduce the vNMSE further, we introduce the DRIVE` algorithm. In DRIVE`, we also use a scale parameter, denoted S` “ S`px,Rq to differentiate it from the scale S of DRIVE. Here, instead of reconstructing the rotated vector in a symmetric manner, i.e., zRpxq P S ¨ t´1, 1ud, we have that zRpxq P S` ¨ tc1, c2ud where c1, c2 are computed using K-Means clustering with K “ 2 over the d entries of the rotated vector Rpxq. That is, c1, c2 are chosen to minimize the SSE over any choice of two values. This does not increase the (asymptotic) time complexity over the random rotations considered in this paper as solving K-Means for the special case of one-dimensional data is deterministically solvable in Opd log dq (e.g., [24]). Notice that DRIVE` still requires dp1 ` op1qq bits as we communicate pS` ¨ c1, S` ¨ c2q and a single bit per coordinate, indicating its nearest centroid. We defer the pseudocode and analyses of DRIVE` to Appendix C. We show that with proper scaling, for both the 1b - Vector Estimation and 1b - Distributed Mean Estimation problems, DRIVE` yields guarantees that are at least as strong as those of DRIVE.

heading: 6 DRIVE with a Structured Random Rotation
text: Uniform random rotation generation usually relies on QR factorization (e.g., see [25]), which requires Opd3q time and Opd2q space. Therefore, uniform random rotation can only be used in practice to rotate low-dimensional vectors. This is impractical for neural network architectures with many millions of parameters. To that end, we continue to analyze DRIVE and DRIVE` with the (randomized) Hadamard transform, a.k.a. structured random rotation [2, 26], that admits a fast inplace, parallelizable, Opd log dq time implementation [27, 28, 29]. We start with a few definitions.
Definition 1. The Walsh-Hadamard matrix ([30]) H2k P t`1,´1u2 kˆ2k is recursively defined via:
H2k“ ˆ
H2k´1 H2k´1 H2k´1 ´H2k´1
˙
and H1“p1q. Also, p 1?dHq ¨ p 1? d HqT“I and detp 1? d Hq P r´1, 1s.
Definition 2. Let RH denote the rotation matrix HD?d P R dˆd, where H is a Walsh-Hadamard matrix and D is a diagonal matrix whose diagonal entries are i.i.d. Rademacher random variables (i.e., taking values uniformly in ˘1). Then RHpxq “ RH ¨ x “ 1?dH ¨ px1 ¨D11, . . . , xd ¨Dddq T is the randomized Hadamard transform of x and R´1H pxq “ RTH ¨ x “ DH?d ¨ x is the inverse transform.

heading: 6.1 1b - Vector Estimation
text: Recall that the vNMSE of DRIVE, when minimized using S “ ‖Rpxq‖1d , is 1 ´ E “ LdR,x ‰
(see Lemma 1). We now bound this quantity of DRIVE with a structured random rotation.
Lemma 3. For any dimension d ě 2 and vector x P Rd, the vNMSE of DRIVE with a structured random rotation and scale S “ ‖RHpxq‖1d is: 1´ E “ LdRH ,x ‰ ď 12 .
Proof. Observe that for all i, E r|RHpxqi|s “ E “ ˇ ˇ řd j“1 xj? d HijDjj ˇ ˇ ‰ . Since tHijDjj | j P rdsu are i.i.d. Rademacher random variables we can use the Khintchine inequality [31, 32] which implies that 1? 2d ¨ ‖x‖2 ď E r|RHpxqi|s ď 1?d ¨ ‖x‖2 (see [33, 34] for simplified proofs). We conclude that:
E “ LdRH ,x ‰ “ 1d ¨ E ” ‖RHpx̆q‖21 ı ě 1d ¨ E r‖RHpx̆q‖1s 2 ě 1d ¨ p řd i“1 1? 2d q2 “ 12 .
This bound is sharp since for d ě 2 we have that LdRH ,x “ 12 for x “ p 1?2 , 1? 2 , 0, . . . , 0qT .
Observe that unlike for the uniform random rotation, E “ LdRH ,x ‰
depends on x. We also note that this bound of 12 applies to DRIVE ` (with scale S` “ 1) as we show in Appendix C.2.

heading: 6.2 1b - Distributed Mean Estimation
text: For an arbitrary x P Rd and R, and in particular for RH , the estimates of DRIVE cannot be made unbiased. For example, for x “ p 23 , 1 3 q T we have that signpRHpxqq “ pD11, D11qT and thus {RHpxq “ S ¨ pD11, D11qT . This implies that px “ R´1H p{RHpxqq “ 1?2 ¨D ¨H ¨ S ¨ pD11, D11q T “
? 2¨S¨D¨pD11, 0qT “ ? 2¨S¨pD211, 0qT “ p ? 2¨S, 0qT . Therefore, Erpxs ‰ x regardless of the scale.
Nevertheless, we next provide evidence for why when the input vector is high dimensional and admits finite moments, a structured random rotation performs similarly to a uniform random rotation, yielding all the appealing aforementioned properties. Indeed, it is a common observation that the distribution of machine learning workloads and, in particular, neural network gradients are governed by such distributions (e.g., lognormal [35] or normal [36, 37]).
We seek to show that at high dimensions, the distribution of RHpxq is sufficiently similar to that of RU pxq. By definition, the distribution of RU pxq is that of a uniformly at random distributed point on a sphere. Previous studies of this distribution for high dimensions (e.g., [38, 39, 40, 41]) have shown that individual coordinates of RU pxq converge to the same normal distribution and that these coordinates are “weakly” dependent in the sense that the joint distribution of every Op1q-sized subset of coordinates is similar to that of independent normal variables for large d.
We hereafter assume that x “ px1, . . . , xdq, where the xis are i.i.d. and that Erx2j s “ σ2 and Er|xj |3s “ ρ ă 8 for all j. We show that RHpxqi converges to the same normal distribution for all i. Let Fi,dpxq be the cumulative distribution function (CDF) of 1σ ¨RHpxqi and Φ be the CDF of the standard normal distribution. The following lemma, proven in Appendix D.1, shows the convergence.
Lemma 4. For all i, RHpxqi converges to a normal variable: supxPR |Fi,dpxq ´ Φpxq| ď 0.409¨ρσ3?d .
With this result, we continue to lay out evidence for the “weak dependency” among the coordinates. We do so by calculating the moments of their joint distribution in increasing subset sizes showing that these moments converge to those of independent normal variables. Previous work has shown that a structured random rotation on vectors with specific distributions results in “weakly dependent” normal variables. This line of research [42, 43, 44] utilized the Hadamard transform for a different purpose. Their goal was to develop a computationally cheap method to generate independent normally distributed variables from simpler (e.g., uniform) distributions. We apply their analysis to our setting.
We partially rely on the following observation that the Hadamard matrix satisfies.
Observation 1. ([42]) The Hadamard product (coordinate-wise product), Hxiy ˝Hx`y, of two rows Hxiy, Hx`y in the Hadamard matrix yields another row at the matrix Hxiy ˝Hx`y “ Hx1`pi´1q‘p`´1qy. Here, pi´ 1q ‘ p`´ 1q is the bitwise xor of the plog dq-sized binary representation of pi´ 1q and p`´ 1q. It follows that
řd j“1HijH`j “ řd j“1pHxiy ˝Hx`yqj “ řd j“1H1`pi´1q‘p`´1q,j .
We now analyze the moments of the rotated variables, starting with the following observation. It follows from the sign-symmetry of D and matches the joint distribution of i.i.d. normal variables. Observation 2. All odd moments containing RHpxq entries are 0. That is,
@q P N,@i1, . . . , i2q`1 P t1, . . . , du : E “ RHpxqi1 ¨ . . . ¨RHpxqi2q`1 ‰ “ 0.
Therefore, we need to examine only even moments. We start with showing that the second moments also match with the distribution of independent normal variables. Lemma 5. For all i ‰ ` it holds that E rpRH ¨ xqi ¨ pRH ¨ xq`s “ 0, whereas E “ pRH ¨ xq2i ‰ “ σ2.
Proof. Since tDjj | j P t1, . . . , duu are sign-symmetric and i.i.d., E “ pRH ¨ xqi ¨ pRH ¨ xq` ‰ “ E “
1 d p řd j“1 xjHijDjjq ¨ p řd j“1 xjH`jDjjq ‰ “ E “ x2j ‰ ¨ 1d ¨ řd j“1HijH`j . Notice that řd j“1H1j “ d
and řd j“1Hij “ 0 for all i ą 1. Thus, by Observation 1 we get 0 if i ‰ ` and σ2 otherwise.
We have established that the coordinates are pairwise uncorrelated. Similar but more involved analysis shows that the same trend continues under the assumption of the existence of x’s higher moments. In Appendix D.2 we analyze the 4th moments showing that they indeed approach the 4th moments of independent normal variables with a rate of 1d ; the reader is referred to [42] for further intuition and higher moments analysis. We therefore expect that using DRIVE and DRIVE` with Hadamard transform will yield similar results to that of a uniform random rotation at high dimensions and when the input vectors respect the finite moments assumption.
In addition to the theoretical evidence, in Figure 1, we show experimental results comparing the measured NMSE for the 1b - Distributed Mean Estimation problem with n “ 10 clients (all given the same vector so biases do not cancel out) for DRIVE and DRIVE` using both uniform and structured random rotations over three different distributions. The results indicate that all variants yield similar NMSEs in reasonable dimensions, in line with the theoretical guarantee of Corollary 1 and Theorem 5.

heading: 7 Evaluation
text: We evaluate DRIVE and DRIVE`, comparing them to standard and recent state-of-the-art techniques. We consider classic distributed learning tasks as well as federated learning tasks (e.g., where the data distribution is not i.i.d. and clients may change over time). All the distributed tasks are implemented over PyTorch [45] and all the federated tasks are implemented over TensorFlow Federated [46]. We focus our comparison on vector quantization algorithms and recent sketching techniques and exclude sparsification methods (e.g., [47, 48, 49, 50]) and methods that involve client-side memory since these can often work in conjunction with our algorithms.
Datasets. We use MNIST [51, 52], EMNIST [53], CIFAR-10 and CIFAR-100 [54] for image classification tasks; a next-character-prediction task using the Shakespeare dataset [55]; and a next-wordprediction task using the Stack Overflow dataset [56]. Additional details appear in Appendix E.1.
Algorithms. Since our focus is on the distributed mean estimation problem and its federated and distributed learning applications, we run DRIVE and DRIVE` with the unbiased scale quantities.1
We compare against several alternative algorithms: (1) FedAvg [1] that uses the full vectors (i.e., each coordinate is represented using a 32-bit float); (2) Hadamard transform followed by 1-bit stochastic
quantization (SQ) [2, 10]; (3) Kashin’s representation followed by 1-bit stochastic quantization [11]; (4) TernGrad [8], which clips coordinates larger than 2.5 times the standard deviation, then performs 1-bit stochastic quantization on the absolute values and separately sends their signs and the maximum coordinate for scale (we note that TernGrad is a low-bit variant of a well-known algorithm called QSGD [9], and we use TernGrad since we found it to perform better in our experiments);2 and (5-6) Sketched-SGD [57] and FetchSGD [58], which are both count-sketch [59] based algorithms designed for distributed and federated learning, respectively.
We note that Hadamard with 1-bit stochastic quantization is our most fair comparison, as it uses the same number of bits as DRIVE` (and slightly more than DRIVE) and has similar computational costs. This contrasts with Kashin’s representation, where both the number of bits and the computational costs are higher. For example, a standard TensorFlow Federated implementation (e.g., see “CLASS KASHINHADAMARDENCODINGSTAGE” hyperparameters at [60]) uses a minimum of 1.17 bits per coordinate, and three iterations of the algorithm resulting in five Hadamard transforms for each vector. Also, note that TernGrad uses an extra bit per coordinate for sending the sign. Moreover, the clipping performed by TernGrad is a heuristic procedure, which is orthogonal to our work.
For each task, we use a subset of datasets and the most relevant competition. Detailed configuration information and additional results appear in Appendix E. We first evaluate the vNMSE-Speed tradeoffs and then proceed to federated and distributed learning experiments.
vNMSE-Speed Tradeoff. Appearing in Table 1, the results show that our algorithms offer the lowest NMSE and that the gap increases with the dimension. As expected, DRIVE and DRIVE` with uniform rotation are more accurate for small dimensions but are significantly slower. Similarly, DRIVE is as accurate as DRIVE`, and both are significantly more accurate than Kashin (by a factor of 4.4ˆ-5.5ˆ) and Hadamard (9.3ˆ-51ˆ) with stochastic quantization. Additionally, DRIVE is 5.7ˆ-12ˆ faster than Kashin and about as fast as Hadamard. In Appendix E.3 we discuss the results, give the complete experiment specification, and provide measurements on a commodity machine.
We note that the above techniques, including DRIVE, are more computationally expensive than lineartime solutions like TernGrad. Nevertheless, DRIVE’s computational overhead becomes insignificant for modern learning tasks. For example, our measurements suggest that it can take 470 ms for computing the gradient on a ResNet18 architecture (for CIFAR100, batch size = 128, using NVIDIA GeForceGTX 1060 (6GB) GPU) while the encoding of DRIVE (Hadamard) takes 2.8 ms. That is, the overall computation time is only increased by 0.6% while the error reduces significantly. Taking the transmission and model update times into consideration would reduce the importance of the compression time further.
Federated Learning. We evaluate over four tasks: (1) EMNIST over customized CNN architecture with two convolutional layers with«1.2M parameters [61]; (2) CIFAR-100 over ResNet-18 [54] ; (3) a next-character-prediction task using the Shakespeare dataset [1]; (4) a next-word-prediction task using the Stack Overflow dataset [62]. Both (3) and (4) use LSTM recurrent models [63] with«820K and «4M parameters, respectively. We use code, client partitioning, models, hyperparameters, and validation metrics from the federated learning benchmark of [62].
1For DRIVE the scale is S “ ‖x‖ 2 2 ‖Rpxq‖1 (see Theorem 3). For DRIVE` the scale is S` “ ‖x‖ 2 2 ‖c‖22 , where c P tc1, c2ud is the vector indicating the nearest centroid to each coordinate in Rpxq (see Section 5). 2When restricted to two quantization levels, TernGrad is identical to QSGD’s max normalization variant with clipping (slightly better due to the ability to represent 0).
The results are depicted in Figure 2. We observe that in all tasks, DRIVE and DRIVE` have accuracy that is competitive with that of the baseline, FedAvg. In CIFAR-100, TernGrad and DRIVE provide the best accuracy. For the other tasks, DRIVE and DRIVE` have the best accuracy, while the best alternative is either Kashin + 1-bit SQ or TernGrad, depending on the task. Hadamard + 1-bit SQ, which is the most similar to our algorithms (in terms of both bandwidth and compute), provides lower accuracy in all tasks. Additional details and hyperparameter configurations are presented in Appendix E.4.
Distributed CNN Training. We evaluate distributed CNN training with 10 clients in two configurations: (1) CIFAR-10 dataset with ResNet-9; (2) CIFAR-100 with ResNet-18 [54, 64]. In both tasks, DRIVE and DRIVE` have similar accuracy to FedAvg, closely followed by Kashin + 1-bit SQ. The other algorithms are less accurate, with Hadamard + 1-bit SQ being better than Sketched-SGD and TernGrad for both tasks. Additional details and hyperparameter configurations are presented in Appendix E.5. Figure 3 depicts the results.
Evaluation Summary. Overall, it is evident that DRIVE and DRIVE` consistently offer markedly favorable results in comparison to the alternatives in our setting. Kashin’s representation appears to offer the best competition, albeit at somewhat higher computational complexity and bandwidth requirements. The lesser performance of the sketch-based techniques is attributed to the high noise of the sketch under such a low (dp1`op1qq bits) communication requirement. This is because the number of counters they can use is too low, making too many coordinates map into each counter. In Appendix E.6, we also compare DRIVE and DRIVE` to state of the art techniques over K-Means and Power Iteration tasks for 10, 100, and 1000 clients, yielding similar trends.

heading: 8 Discussion
text: Proven Error Bounds. We summarize the proven error bounds in Table 2. Since DRIVE (Hadamard) is generally not unbiased (as discussed in Section 6.2), we cannot establish a formal guarantee for the 1b - DME problem when using Hadamard. It is a challenging research question whether there exists other structured rotations with low computational complexity and stronger guarantees.
Input Distribution Assumption. The distributed mean estimation analysis of our Hadamardbased variants is based on an assumption (Section 6.2) about the vector distributions. While machine learning workloads, and DNN gradients in particular (e.g., [35, 36, 37]), were observed to follow such distributions, this assumption may not hold for other applications.
For such cases, we note that DRIVE is compatible with the error feedback (EF) mechanism [65, 66] that ensured convergence and recovery of the convergence rate of non-compressed SGD. Specifically, as evident by Lemma 3, any scale ‖Rpxq‖1d ď S ď 2 ¨ ‖Rpxq‖1 d is sufficient to respect the compressor (i.e., bounded variance) assumption. For completeness, in Appendix E.2, we perform EF experiments comparing DRIVE and DRIVE` to other compression techniques that use EF.
Varying Communication Budget. Unlike some previous works, our algorithms’ guarantees with more than one bit per coordinate are not established. It is thus an interesting future work to extend DRIVE to other communication budgets and understand what are the resulting guarantees. We refer the reader to [67] for initial steps towards that direction.
Entropy Encoding. Entropy encoding methods (such as Huffman coding) can further compress vectors of values when the values are not uniformly distributed. We have compared DRIVE against stochastic quantization methods using entropy encoding for the challenging setting for DRIVE where all vectors are the same (see Table 1 for further description). The results appear in Appendix E.7, where DRIVE still outperforms these methods. We also note that, when computation allows and when using DRIVE with multiple bits per entry, DRIVE can also be enhanced by entropy encoding techniques. We describe some initial results for this setting in [67].
Structured Data. When the data is highly sparse, skewed, or otherwise structured, one can leverage that for compression. We note that some techniques that exploit sparsity or structure can be use in conjunction with our techniques. For example, one may transmit only non-zero entries or Top-K entries while compressing these using DRIVE to reduce communication overhead even further.
Compatibility With Distributed All-Reduce Techniques. Quantization techniques, including DRIVE, may introduce overheads in the context of All-Reduce (depending on the network architecture and communication patterns). In particular, if every node in a cluster uses a different rotation, DRIVE will not allow for efficient in-path aggregation without decoding the vectors. Further, the computational overhead of the receiver increases by a log d factor as each vector has to be decoded separately before an average can be computed. It is an interesting future direction for DRIVE to understand how to minimize such potential overheads. For example, one can consider bucketizing co-located workers and apply DRIVE’s quantization only for cross-rack traffic.

heading: 9 Conclusions
text: In this paper, we studied the vector and distributed mean estimation problems. These problems are applicable to distributed and federated learning, where clients communicate real-valued vectors (e.g., gradients) to a server for averaging. To the best of our knowledge, our algorithms are the first with a provable error of Op 1n q for the 1b - Distributed Mean Estimation problem (i.e., with dp1` op1qq bits). As shown in [14], any algorithm that uses Opdq shared random bits (e.g., our Hadamard-based variant) has a vNMSE of Ωp1q, i.e., DRIVE and DRIVE` are asymptotically optimal; additional discussion is given in Appendix F. Our experiments, carried over various tasks and datasets, indicate that our algorithms improve over the state of the art. All the results presented in this paper are fully reproducible by our source code, available at [29].
Acknowledgments and Disclosure of Funding
MM was supported in part by NSF grants CCF-2101140, CCF-2107078, CCF-1563710, and DMS2023528. MM and RBB were supported in part by a gift to the Center for Research on Computation and Society at Harvard University. AP was supported in part by the Cyber Security Research Center at Ben-Gurion University of the Negev. We thank Moshe Gabel, Mahmood Sharif, Yuval Filmus and, the anonymous reviewers for helpful comments and suggestions.

Review summaries:
The paper studies distributed mean estimation in the $d$-dimensional space where each client can use only $d(1+o(1))$ bits for communication. A crucial assumption is that the server and the $n$ clients can have shared randomness while each client can independently communicate with the server. Under such an assumption, the authors proposed two schemes and demonstrated the efficiency of their new approaches. 

This paper studies the problem of distributed mean estimation under a single bit (per coordinate) communication constraint. It proposes a novel compression scheme based on random rotation and sign quantization. The suggested method is applied to distributed and federated machine learning problems and shows improved results over previously studied quantization methods.

This paper studies the distributed mean estimation. In this problem, there are n clients,  each of which has a d-dimensional real-valued; the goal is to estimate their mean on the sever side using minimum communication cost. The paper focuses on 1-bit quantization methods, i.e., each client sends one bit per coordinate and d(1+o(1)) bits in total. The algorithms are standard (including the Hadamard transformation based method); the main difference is new scaling factors S. Theoretical analysis on the expectation and variance are provided. Empirical results also show some advantages of the new estimate. 

The paper considers the problem of communication-efficient distributed mean estimation. This paper proposes a novel algorithm DRIVE which uses as little as almost 1 bit for each coordinate of the vector and estimate the distributed mean with bounded error. The idea is to use a uniform random rotation to minimize the vNMSE. Besides, DRIVE can use structured random rotation to accelerate the compression for high dimensional vectors. The empirical evaluation of DRIVE shows that federated learning with gradients compressed by DRIVE can achieve faster convergence than other compression schemes.