Federated Averaging as Expectation Maximization

Christos Louizos, Matthias Reisser, Joseph Soriaga, Max Welling

Keywords: federated, learning, sparsity, expectation, maximization, efficient, FedAvg, FedSparse, EM

Abstract: Federated averaging (FedAvg), despite its simplicity, has been the main approach in training neural networks in the federated learning setting. In this work, we show that the algorithmic choices of the FedAvg algorithm correspond to optimizing a single objective function that involves the global and all of the shard specific models using a hard version of the well known Expectation-Maximization (EM) algorithm. As a result, we gain a better understanding of the behavior and design choices of federated averaging while being able to provide interesting connections to recent literature. Based on this view, we further propose FedSparse, a version of federated averaging that employs prior distributions to promote model sparsity. In this way, we obtain a procedure that leads to reductions in both server-client and client-server communication costs as well as more efficient models.

One-sentence Summary: We provide a view of federated averaging through Expectation-Maximization and propose FedSparse, an extension that learns sparse models in the federated setting.

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics

Reviewed Version (pdf): https://openreview.net/references/pdf?id=L8_C7Z8nt0