Systems and methods for communication efficient distributed mean estimation
The present disclosure provides systems and methods for communication efficient distributed mean estimation. In particular, aspects of the present disclosure can be implemented by a system in which a number of vectors reside on a number of different clients, and a centralized server device seeks to estimate the mean of such vectors. According to one aspect of the present disclosure, a client computing device can rotate a vector by a random rotation matrix and then subsequently perform probabilistic quantization on the rotated vector. According to another aspect of the present disclosure, subsequent to quantization but prior to transmission, the client computing can encode the quantized vector according to a variable length coding scheme (e.g., by computing variable length codes).
1. A computing system that performs stochastic quantization to facilitate efficient transmission of machine-learned model updates from client devices to a centralized server computing device, the computing system comprising:
one or more client computing devices, wherein each client computing device comprises one or more processors and one or more non-transitory computer-readable media that store instructions that, when executed by the one or more processors cause the client computing device to perform operations, the operations comprising:
determining an update to a machine-learned model based at least in part on a local dataset stored at the client computing device, wherein the update comprises a plurality of coordinate values;
performing stochastic quantization of the update to obtain a quantized update, wherein performing stochastic quantization of the update comprises:
establishing a plurality of quantization levels;
identifying a minimum value among the plurality of coordinate values of the update;
setting a first quantization level of the plurality of quantization levels equal to the minimum value among the plurality of coordinate values; and
respectively quantizing each coordinate value of the plurality of coordinate value into one of the plurality of quantization levels according to a respective probability determined based on the coordinate value; and
transmitting the quantized update to the centralized server computing device.
2. The computing system of claim 1 , wherein:
the plurality of quantization levels comprises two quantization levels; and
transmitting the quantized update to the centralized server computing device comprises transmitting:
two real values that correspond to the two quantization levels; and
a bit vector of length equal to the number of coordinate values in the update, wherein each bit in the bit vector corresponds to a respective one of the plurality of coordinate values.
3. The computing system of claim 2 , wherein performing stochastic quantization further comprises:
identifying a maximum value among the plurality of coordinate values of the update; and
setting a second quantization level of the plurality of quantization levels equal to the maximum value among the plurality of coordinate values.
4. The computing system of claim 1 , wherein the plurality of quantization levels comprises three or more quantization levels.
5. The computing system of claim 4 , wherein performing stochastic quantization further comprises:
setting each of the quantization levels according to the following expression:
B
i
(
r
)
=
def
X
i
min
+
rs
i
k
-
1
,
where r represents a level counter, k represents a total number of the quantization levels, X i min represents the minimum value among the plurality of coordinate values, and s i is greater than or equal to a difference between a maximum value among the plurality of coordinate values and the minimum value among the plurality of coordinate values.
6. The computing system of claim 5 , wherein s i equals the square root of two times a magnitude of the update.
7. The computing system of claim 6 , wherein the operations further comprise:
prior to transmitting the quantized update to the centralized server computing device, encoding the quantized update by performing variable length coding on the quantized update.
8. The computing system of claim 1 , wherein the operations further comprise:
prior to transmitting the quantized update to the centralized server computing device, encoding the quantized update by performing variable length coding on the quantized update.
9. The computing system of claim 1 , wherein respectively quantizing each coordinate value of the plurality of coordinate value into one of the plurality of quantization levels comprises, for each coordinate value:
determining a first candidate quantization value and a second candidate quantization value;
computing the respective probability for the coordinate value as equal to a numerator divided by a denominator, wherein the numerator comprises the coordinate value minus the second candidate quantization value, and wherein the denominator comprises the first candidate quantization value minus the second candidate quantization value;
setting the coordinate value equal to the first candidate quantization value according to the respective probability and else setting the coordinate value equal to the second candidate quantization value.
10. The computing system of claim 1 , wherein the update to the machine-learned model comprises a gradient associated with training of the machine-learned model.
11. The computing system of claim 1 , further comprising:
the centralized server computing device, the centralized server computing device comprising one or more processors and one or more non-transitory computer-readable media that store instructions that, when executed by the one or more processors cause the centralized server computing device to perform second operations, the second operations comprising:
receiving the one or more quantized updates respectively transmitted by the one or more client computing devices; and
determining a mean update of all quantized updates obtained for all of the one or more client computing devices;
updating a global machine-learned model based at least in part on the mean update.
12. A computer-implemented method, comprising:
determining, by a computing system comprising one or more computing devices, an update to a machine-learned model based at least in part on a local dataset stored at a client computing device, wherein the update comprises a plurality of coordinate values;
performing, by the computing system, stochastic quantization of the update to obtain a quantized update, wherein performing stochastic quantization of the update comprises:
establishing, by the computing system, a plurality of quantization levels;
identifying, by the computing system, a minimum value among the plurality of coordinate values of the update;
setting, by the computing system, a first quantization level of the plurality of quantization levels equal to the minimum value among the plurality of coordinate values; and
respectively quantizing, by the computing system, each coordinate value of the plurality of coordinate value into one of the plurality of quantization levels according to a respective probability determined based on the coordinate value; and
transmitting, by the computing system, the quantized update to a centralized server computing device.
13. The computer-implemented method of claim 12 , wherein:
the plurality of quantization levels comprises two quantization levels; and
transmitting the quantized update to the centralized server computing device comprises transmitting:
two real values that correspond to the two quantization levels; and
a bit vector of length equal to the number of coordinate values in the update, wherein each bit in the bit vector corresponds to a respective one of the plurality of coordinate values.
14. The computer-implemented method of claim 13 , wherein performing stochastic quantization further comprises:
identifying a maximum value among the plurality of coordinate values of the update; and
setting a second quantization level of the plurality of quantization levels equal to the maximum value among the plurality of coordinate values.
15. The computer-implemented method of claim 12 , wherein the plurality of quantization levels comprises three or more quantization levels.
16. The computer-implemented method of claim 15 , wherein performing stochastic quantization further comprises:
setting each of the quantization levels according to the following expression:
B
i
(
r
)
=
def
X
i
min
+
rs
i
k
-
1
,
where r represents a level counter, k represents a total number of the quantization levels, X i min represents the minimum value among the plurality of coordinate values, and s i is greater than or equal to a difference between a maximum value among the plurality of coordinate values and the minimum value among the plurality of coordinate values.
17. The computer-implemented method of claim 16 , wherein s i equals the square root of two times a magnitude of the update.
18. The computer-implemented method of claim 17 , wherein the operations further comprise:
prior to transmitting the quantized update to the centralized server computing device, encoding the quantized update by performing variable length coding on the quantized update.
19. The computer-implemented method of claim 12 , wherein the operations further comprise:
prior to transmitting the quantized update to the centralized server computing device, encoding the quantized update by performing variable length coding on the quantized update.
20. The computer-implemented method of claim 12 , wherein respectively quantizing each coordinate value of the plurality of coordinate value into one of the plurality of quantization levels comprises, for each coordinate value:
determining a first candidate quantization value and a second candidate quantization value;
computing the respective probability for the coordinate value as equal to a numerator divided by a denominator, wherein the numerator comprises the coordinate value minus the second candidate quantization value, and wherein the denominator comprises the first candidate quantization value minus the second candidate quantization value;
setting the coordinate value equal to the first candidate quantization value according to the respective probability and else setting the coordinate value equal to the second candidate quantization value.