# Federated Expectation Maximization with heterogeneity mitigation and variance reduction

arXiv: https://arxiv.org/abs/2111.02083 · HTML: https://mlsanity.com/paper/2111.02083/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 40 references that are in the map.
- First posted on arXiv 3 Nov 2021.

## Builds on (most cited first, 21 of 21)

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Federated Learning: Strategies for Improving Communication Efficiency](https://mlsanity.com/paper/1610.05492/) (2016)
- [Stochastic First- and Zeroth-order Methods for Nonconvex Stochastic Programming](https://mlsanity.com/paper/1309.5549/) (2013)
- [Robust and Communication-Efficient Federated Learning from Non-IID Data](https://mlsanity.com/paper/1903.02891/) (2019)
- [SAGA: A Fast Incremental Gradient Method With Support for Non-Strongly Convex Composite Objectives](https://mlsanity.com/paper/1407.0202/) (2014)
- [SPIDER: Near-Optimal Non-Convex Optimization via Stochastic Path Integrated Differential Estimator](https://mlsanity.com/paper/1807.01695/) (2018)
- [The Convergence of Sparsified Gradient Methods](https://mlsanity.com/paper/1809.10505/) (2018)
- [SARAH: A Novel Method for Machine Learning Problems Using Stochastic Recursive Gradient](https://mlsanity.com/paper/1703.00102/) (2017)
- [Federated Multi-Task Learning under a Mixture of Distributions](https://mlsanity.com/paper/2108.10252/) (2021)
- [Is Local SGD Better than Minibatch SGD?](https://mlsanity.com/paper/2002.07839/) (2020)
- [DoubleSqueeze: Parallel Stochastic Gradient Descent with Double-Pass Error-Compensated Compression](https://mlsanity.com/paper/1905.05957/) (2019)
- [A Unified Theory of SGD: Variance Reduction, Sampling, Quantization and Coordinate Descent](https://mlsanity.com/paper/1905.11261/) (2019)
- [Stochastic Distributed Learning with Gradient Quantization and Variance Reduction](https://mlsanity.com/paper/1904.05115/) (2019)
- [Bidirectional compression in heterogeneous settings for distributed or federated learning with partial participation: tight convergence guarantees](https://mlsanity.com/paper/2006.14591/) (2020)
- [A Better Alternative to Error Feedback for Communication-Efficient Distributed Learning](https://mlsanity.com/paper/2006.11077/) (2020)
- [Preserved central model for faster bidirectional compression in distributed settings](https://mlsanity.com/paper/2102.12528/) (2021)
- [Homeomorphic-Invariance of EM: Non-Asymptotic Convergence in KL Divergence for Exponential Families via Mirror Descent](https://mlsanity.com/paper/2011.01170/) (2020)
- [On the Global Convergence of (Fast) Incremental Expectation Maximization Methods](https://mlsanity.com/paper/1910.12521/) (2019)
- [Geom-SPIDER-EM: Faster Variance Reduced Stochastic Expectation Maximization for Nonconvex Finite-Sum Optimization](https://mlsanity.com/paper/2011.12392/) (2020)
- [Fast Incremental Expectation Maximization for finite-sum optimization: nonasymptotic convergence](https://mlsanity.com/paper/2012.14670/) (2020)

## Built on by (most cited first, 5 of 5)

- [Fast Heterogeneous Federated Learning with Hybrid Client Selection](https://mlsanity.com/paper/2208.05135/) (2022)
- [CoBo: Collaborative Learning via Bilevel Optimization](https://mlsanity.com/paper/2409.05539/) (2024)
- [Federated Causal Inference from Multi-Site Observational Data via Propensity Score Aggregation](https://mlsanity.com/paper/2505.17961/) (2025)
- [Graph-Regularized Learning of Gaussian Mixture Models](https://mlsanity.com/paper/2509.13855/) (2025)
- [Federated Learning-Assisted Optimization of Mobile Transmission with Digital Twins](https://mlsanity.com/paper/2602.18627/) (2026)

## Abstract (arXiv)

The Expectation Maximization (EM) algorithm is the default algorithm for inference in latent variable models. As in any other field of machine learning, applications of latent variable models to very large datasets make the use of advanced parallel and distributed architectures mandatory. This paper introduces FedEM, which is the first extension of the EM algorithm to the federated learning context. FedEM is a new communication efficient method, which handles partial participation of local devices, and is robust to heterogeneous distributions of the datasets. To alleviate the communication bottleneck, FedEM compresses appropriately defined complete data sufficient statistics. We also develop and analyze an extension of FedEM to further incorporate a variance reduction scheme. In all cases, we derive finite-time complexity bounds for smooth non-convex problems. Numerical results are presented to support our theoretical findings, as well as an application to federated missing values imputation for biodiversity monitoring.