On the 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. We fully characterize the convergence rate of the EM algorithm under all regimes of $m/n$ where m is the number of clients and n is the number of data points per client. We show that with a signal-to-noise-ratio (SNR) of order Ω(√K), the well-initialized EM algorithm converges within the minimax distance of the ground truth under each of the regimes. Interestingly, we identify that when m grows exponentially in n, the EM algorithm only requires a constant number of iterations to converge. We perform experiments on synthetic datasets to illustrate our results. Surprisingly, the results show that rather than being a bottleneck, data heterogeneity can accelerate the convergence of federated learning algorithms.