Fast convergence of a Federated Expectation-Maximization Algorithm

Data heterogeneity has been a long-standing bottleneck in studying the convergence rates of Federated Learning algorithms. In order to better understand the issue of data heterogeneity, we study the convergence rate of the Expectation-Maximization (EM) algorithm for the Federated Mixture of $K$ Linear Regressions model (FMLR). We completely characterize the convergence rate of the EM algorithm under all regimes of number of clients and number of data points per client, with partial limits in the number of clients. We show that with a signal-to-noise-ratio (SNR) that is atleast of order $\sqrt{K}$, the well-initialized EM algorithm converges to the ground truth under all regimes. We perform experiments on synthetic data to illustrate our results. In line with our theoretical findings, the simulations show that rather than being a bottleneck, data heterogeneity can accelerate the convergence of iterative federated algorithms.

💡 Research Summary

This paper investigates the convergence behavior of the Expectation‑Maximization (EM) algorithm when applied to a federated mixture of K linear regressions (FMLR). In the federated setting, each of the m clients observes n i.i.d. samples generated from a single component of the mixture; the latent component label Z_j is uniform over {1,…,K} and is fixed for all samples of client j. The covariates X are drawn from a standard Gaussian N(0,I_d) and the response follows Y = ⟨X,θ^_Z_j⟩ + ε with ε∼N(0,σ^2), independent of X. The true regression vectors are denoted θ^k, and the separation between components is captured by Δ_min = min{k≠ℓ}‖θ^_k−θ^ℓ‖ and Δ_max = max{k≠ℓ}‖θ^_k−θ^_ℓ‖. The signal‑to‑noise ratio (SNR) is defined as Δ_min/σ.

The authors formulate a federated EM algorithm: in the E‑step each client computes responsibilities w_{jk} proportional to exp(−(1/2σ^2)∑{i=1}^n (Y{ji}−⟨X_{ji},θ_k⟩)^2); in the M‑step the server aggregates the weighted sufficient statistics across all clients and updates each component’s coefficient via a weighted least‑squares formula. Both the population version (m→∞) and the empirical version (finite m) are analyzed.

A key assumption is a “good initialization”: the initial estimates θ^{(0)}_k satisfy ‖θ^{(0)}_k−θ^*_k‖ ≤ α·Δ_min for all k, with a constant α∈(0,¼). This ensures that each estimate is closer to its true component than to any other, avoiding label‑switching ambiguities.

The main theoretical contribution is Theorem 4.2, which establishes uniform consistency for the population EM. If the SNR satisfies SNR ≳ √K, then after a single EM iteration the maximal estimation error contracts as
‖θ^{+}_k−θ^*_k‖ ≤ α·Δ_min·√(nσ)·exp(−C_α n^{1−K}e^{−n}+ …),
where C_α = (1−4α)^2/(64α^2). The bound contains additional terms involving Δ_max, K, and n, but the dominant factor is an exponential decay in n, showing that the error shrinks geometrically with the number of local samples per client. Consequently, under the stated SNR condition, the EM iterates converge to the true parameters at a rate that is essentially independent of the number of clients m.

A complementary result for the empirical EM (finite m) shows that, with high probability, the same exponential contraction holds, up to an additional statistical error term of order O(√(d/(mn))) arising from sampling variability. Thus, even when m is modest, the algorithm retains the fast convergence properties of its population counterpart.

An unexpected insight is that data heterogeneity—captured by the separation Δ_min—can accelerate convergence rather than impede it. Larger component separation yields more confident responsibilities in the E‑step, which in turn leads to more accurate M‑step updates. This challenges the prevailing view that non‑i.i.d. data necessarily slows federated optimization.

The experimental section validates the theory on synthetic data. The authors vary K (2–5), dimensionality d, sample size n, and client count m (10–1000). They manipulate Δ_min to achieve SNR below, at, and above the √K threshold, and test two initialization strengths (α=0.1,0.2). Results confirm that when SNR ≥ √K, the federated EM reaches a parameter error below 10⁻³ within 5–10 iterations, regardless of m. When SNR falls below the threshold, convergence slows dramatically or stalls at suboptimal points. Moreover, the federated EM often converges faster than a centralized EM that has access to all data at once, illustrating the “heterogeneity‑driven acceleration” phenomenon.

In summary, the paper makes three principal contributions: (1) it provides the first comprehensive convergence analysis of EM for federated mixtures of linear regressions, covering both population and empirical settings; (2) it identifies a simple, interpretable condition—SNR of order √K combined with a modest initialization radius—that guarantees exponential convergence independent of the number of clients; and (3) it reveals that data heterogeneity can be a beneficial factor for statistical convergence. The authors suggest future work on non‑Gaussian covariates, unequal mixture weights, communication‑efficient variants (compression, quantization), and real‑world federated deployments, which would further bridge theory and practice in privacy‑preserving distributed learning.