Source-linked AI summary
Distributed Mean Estimation with Limited Communication
Ananda Theertha Suresh, Felix X. Yu, Sanjiv Kumar, H. Brendan McMahan
TL;DR
The paper addresses distributed empirical mean estimation when communication is limited and data has no distributional assumptions. It develops quantization, random rotation, variable-length coding, and sampling schemes, showing improved MSE-communication trade-offs and a minimax-optimal strategy up to constants. The methods are applied to distributed k-means and PCA power iteration.
Problem
Communication is a bottleneck in distributed learning and optimization, motivating empirical mean estimation with limited communication and no distributional assumptions.
Method
The paper combines stochastic quantization with structured random rotation, variable-length coding, and client sampling under simultaneous independent communication.
Results
The combined quantization, variable-length coding, and sampling scheme achieves information-theoretically optimal MSE for a given communication cost in the stated minimax regime.
Takeaways & Limitations
The algorithms provide communication-efficient tools for distributed Lloyd’s algorithm for k-means and power iteration for PCA.
Takeaways & Limitations
The lower bound applies only when communication cost c < O(nd); extending the result to larger costs remains open.
Abstract
from arXiv · showhide
Motivated by the need for distributed learning and optimization algorithms with low communication cost, we study communication efficient algorithms for distributed mean estimation. Unlike previous works, we make no probabilistic assumptions on the data. We first show that for $d$ dimensional data with $n$ clients, a naive stochastic binary rounding approach yields a mean squared error (MSE) of $Θ(d/n)$ and uses a constant number of bits per dimension per client. We then extend this naive algorithm in two ways: we show that applying a structured random rotation before quantization reduces the error to $\mathcal{O}((\log d)/n)$ and a better coding strategy further reduces the error to $\mathcal{O}(1/n)$ and uses a constant number of bits per dimension per client. We also show that the latter coding strategy is optimal up to a constant in the minimax sense i.e., it achieves the best MSE for a given communication cost. We finally demonstrate the practicality of our algorithms by applying them to distributed Lloyd's algorithm for k-means and power iteration for PCA.
1 Introduction
The paper studies distributed empirical mean estimation under limited communication, without distributional assumptions, and targets optimal minimax error as a function of communication. It develops quantization, rotation, variable-length coding, and sampling methods, then applies them to distributed learning tasks.
- Applications: The algorithms are applied to distributed Lloyd’s algorithm for k-means and power iteration for PCA.In distributed power iteration, 100 clients send compressed eigenvector updates to a server that averages them.
- Motivation: Distributed mean estimation supports distributed k-means, PCA power iteration, and synchronized model training, but communication can be prohibitive for low-bandwidth devices.The paper therefore studies the optimal minimax rate for empirical mean estimation under limited communication.
- Problem setting: The model estimates the empirical mean for arbitrary data, unlike prior work that assumes i.i.d. samples and targets a statistical-model mean.Protocols are simultaneous and independent, with clients transmitting functions of their own vectors.
- Stochastic uniform quantization: Θ(d/n) MSE is achieved by naive stochastic binary quantization using one bit per dimension per client, but this error can be prohibitive when d is much larger than n.The paper shows this basic bound is tight.
- Stochastic uniform quantization: Increasing quantization to k levels lowers MSE, but communication rises to n · (d⌈log2 k⌉ + ˜O(1)) bits.This trade-off can be expensive when seeking MSE o(d/n).
- Communication-efficient algorithms: Random rotation reduces MSE at the same communication form as k-level quantization, while variable-length coding compresses frequent quantization levels.The variable-length scheme uses the same quantization as πsk and encodes levels with methods such as arithmetic coding.
- Practical considerations: Variable-length coding has stronger guarantees, whereas stochastic rotated quantization uses fixed-length coding and may offer privacy-preserving secure aggregation or better constants.The latter can have lower quantization error in some scenarios.
- Minimax MSE: Combining k-level quantization, variable-length coding, and client sampling achieves the information-theoretically optimal MSE for a given communication cost up to the stated minimax result.The product of communication cost and MSE scales linearly with the dimension.
2 Stochastic uniform quantization
The paper develops stochastic binary and k-level quantization for distributed mean estimation, analyzes their estimation error, and characterizes their communication costs. The basic binary scheme can incur error that grows with dimension, while k-level quantization offers a flexible communication–accuracy trade-off.
- Stochastic binary quantization: Stochastic binary quantization independently quantizes each coordinate using private randomness.Each client transmits a bit vector together with values representing the coordinate range.
- Stochastic binary quantization: d + ˜O(1) bits per client suffice for stochastic binary quantization.The implementation uses d bits for the quantized coordinates plus a small overhead for transmitting range information.
- Stochastic binary quantization: Θ(d/n) MSE times the average norm makes the simple binary scheme unsuitable when dimension is very large relative to the number of clients.The paper notes that for neural-network-scale dimensions, this error can exceed the norm of the mean vector.
- Stochastic k-level quantization: k-level quantization divides each coordinate’s range into k stochastic quantization levels.The scheme generalizes binary quantization and uses interval boundaries determined by the coordinate range and a scale parameter.
- Stochastic k-level quantization: Variable length coding can provide better guarantees with a higher scale value than the natural k-level choice.The paper postpones this coding improvement to a later section.
- Stochastic k-level quantization: d⌈log(k)⌉ plus overhead bits per client bound the communication cost of stochastic k-level quantization.The protocol communicates the selected quantization level for each coordinate.
3 Stochastic rotated quantization
Structured random rotation reduces the coordinate-range term driving quantization error without assuming a data distribution. The proposed Walsh–Hadamard-based rotation is fast to apply and yields an O((log d)/n) error bound for stochastic rotated quantization.
- Motivation and guarantee: O((log d)/n) error follows by reducing Xmax with a structured random rotation, without assumptions on the data distribution.The resulting method is called stochastic rotated quantization, πsrk.
- Protocol: The rotated protocol quantizes Zi = RXi and reconstructs the original mean using the inverse rotation.Clients and server generate the same random rotation using public randomness.
- Structured rotation: Random rotation preserves vector norms while changing the coordinate magnitudes relevant to quantization error.The MSE becomes smaller when the transformed coordinate extremes are reduced.
- Motivation and design: The structured rotation is motivated by the poor coordinate-range behavior of fast but naive orthogonal matrices such as block-diagonal matrices.The paper seeks an orthogonal matrix that both reduces transformed coordinate extremes and supports fast multiplication.
- Implementation: O(d log d) time and O(1) additional space suffice for applying the rotation and its inverse with an in-place algorithm.This addresses the need for fast matrix–vector products in high dimensions.
- Structured rotation: R = HD combines a Walsh–Hadamard matrix H with a random diagonal Rademacher matrix D.The diagonal entries of D are independently ±1 with probability 0.5.
4 Variable length coding
The paper replaces fixed-length level transmission with variable-length coding, compressing quantized coordinates according to their empirical frequencies while retaining the same quantization MSE. This yields O(1/n) MSE with O(nd) communication.
- 4 Variable length coding: Variable-length coding replaces fixed ⌈log2 k⌉-bit transmission with compressed encoding of quantized levels.The scheme first encodes the counts h_r, then uses arithmetic or Huffman coding based on p_r = h_r/d.
- 4 Variable length coding: The variable-length scheme retains the stochastic k-level quantizer’s MSE because both protocols quantize vectors identically.The coding strategy changes communication, not the quantization procedure itself.
- 4 Variable length coding: Arithmetic coding uses the empirical bin distribution, while the counts satisfy Σ_r h_r = d and h_r ≥ 0.The count representation and subsequent arithmetic coding together determine the compressed message length.
- 4 Variable length coding: O(1/n) MSE is achieved with d + 1 communication units, giving O(nd) total communication.The stated result corresponds to choosing k = d + 1.
5 Communication MSE trade-off
The paper reduces communication by sampling clients: only a random fraction transmits, and the server rescales the received estimates. This creates explicit MSE–communication trade-offs while preserving the protocol structure.
- 5 Communication MSE trade-off: Client sampling combines with the existing quantization protocols to trade off MSE against communication cost.The expected number of transmitting clients is reduced from n to np, and analogous analysis applies to coordinate sampling.
- 5 Communication MSE trade-off: The sampled protocol lets each client participate independently with probability p, after which the server estimates the mean from transmitting clients.The sampled set is denoted S, and the estimator is constructed from the corresponding transmitted vectors.
- 5 Communication MSE trade-off: The communication bound follows because an expected np clients transmit under independent client sampling.The MSE analysis separates sampling error from quantization error using independence and zero-mean noise.
- 5 Communication MSE trade-off: For every c ≤ nd(2 + log2(7/4)), a protocol exists with a corresponding upper bound on minimax MSE.The corollary combines the communication and error analyses of the preceding protocols.
6 Lower bounds
The lower-bound analysis establishes a communication-dependent minimax limit and explains why variable-length coding and random rotation cannot be combined to improve the asymptotic result. The lower bound is stated only below linear-in-nd communication.
- 6 Lower bounds: A distribution family supported on the bounded domain yields a lower bound for estimating an unknown distribution mean under independent protocols.The communication cost and distribution-family mean appear explicitly in the lower-bound formulation.
- 6 Lower bounds: Theorem 5 applies this statistical-estimation lower bound to empirical mean estimation for c ≤ ndt/4 and n ≥ 4/t.The proof relates the empirical mean of observed samples to the underlying distribution mean.
- 6 Lower bounds: The lower bound is established only for communication cost c < O(nd); extending it to larger communication remains open.This is the paper’s explicit scope boundary for the lower-bound result.
- 6 Lower bounds: Variable-length coding and stochastic rotated quantization exploit different data properties, so combining them does not improve the asymptotic result.Variable-length coding exploits infrequent large-index bins, whereas rotation narrows the bin range and makes bins more equally likely.
7 Practical considerations and applications
The practical section compares quantization methods in distributed Lloyd’s algorithm and power iteration, while discussing when fixed-length rotated quantization may be preferable. Variable-length coding generally performs best, with rotation competitive at low bit rates.
- 7 Practical considerations and applications: Variable-length coding provides the lowest asymptotic quantization error at a constant number of bits, while rotated quantization can have better practical constants.Rotated quantization also uses fixed-length coding, supporting privacy-preserving secure aggregation.
- 7 Practical considerations and applications: Figure 2 compares uniform, rotated, and variable-length quantization in Lloyd’s algorithm using communication cost on the x-axis and the global objective on the y-axis.The comparison tests 16 and 32 quantization levels.
- 7 Practical considerations and applications: Figure 3 compares the same quantization types in power iteration using communication cost on the x-axis and eigenvector ℓ2 error on the y-axis.The experiments again use 16 and 32 quantization levels.
- 7 Practical considerations and applications: For unbalanced data, rotated stochastic quantization has the best performance, especially at low bit rates.The experiment uses 1000 points in 256 dimensions, with the final feature drawn from N(100, 1) and the others from N(0, 1).
- 7 Practical considerations and applications: Across both applications, variable-length coding usually achieves the lowest quantization error, while rotated quantization is competitive at low bit rates.The applications are evaluated on MNIST with d = 1024 and CIFAR with d = 512.
A Proof of Lemma 7
The proof establishes an upper bound by exploiting symmetry, decomposing the analysis into two terms, and bounding their fluctuations using independence and concentration arguments.
- Symmetry in H establishes the stated equality before the upper-bound argument begins.
- The first term is analyzed through the dependence of Zmax_i on independent diagonal entries D(1), …, D(d).
- Changing one diagonal entry alters Zmax_i by at most 2X_i(j), enabling an Efron–Stein variance bound.
- For the second term, the proof introduces a bound valid for every β > 0.
- Rademacher signs and |H(k, j)| = 1 make the distributions of Z_i(k) identical across k, after which Jensen’s inequality is applied.
- The exponential bound e^a + e^−a ≤ 2e^(a^2/2) follows from independence and controls the remaining expression.