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