DRIVE: One-bit Distributed Mean Estimation

1 Introduction . 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 . 2 Problem Formulation and Notation . 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 . 3 The DRIVE Algorithm . 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 . 4 DRIVE With a Uniform Random Rotation . 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 . 4.1 1b - Vector Estimation . 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 ˘ .

Summary: 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.