Fast convergence of a Federated Expectation-Maximization Algorithm

TL;DR

This paper analyzes the convergence rate of a federated EM algorithm for mixture of linear regressions, showing it converges quickly under certain conditions and that data heterogeneity can even accelerate convergence.

Contribution

It provides a complete characterization of the EM algorithm's convergence in federated settings, including regimes with varying client numbers and data points, highlighting the impact of data heterogeneity.

Findings

EM converges rapidly with sufficient SNR.

Data heterogeneity can accelerate convergence.

Theoretical results are supported by synthetic experiments.

Abstract

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.