On the Convergence of a Federated Expectation-Maximization Algorithm
<|reference_start|>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 $\Omega(\sqrt{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.<|reference_end|>
@article{tao2024on,
    title={On the Convergence of a Federated Expectation-Maximization Algorithm},
    author={Zhixu Tao, Rajita Chandak, Sanjeev Kulkarni},
    journal={arXiv preprint arXiv:2408.05819},
    year={2024},
    archivePrefix={arXiv},
    eprint={2408.05819},
    primaryClass={stat.ML cs.LG}
}