[Paper Review] Personalized Federated Learning under Mixture of Distributions

Yue Wu, Shuaicheng Zhang+6 more|arXiv (Cornell University)|May 1, 2023

Privacy-Preserving Technologies in Data6 citations

TL;DR

This paper proposes FedGMM, a personalized federated learning framework that models joint data distribution heterogeneity using Gaussian Mixture Models (GMMs) via a federated Expectation-Maximization algorithm. It achieves superior performance in personalized FL and out-of-distribution (OOD) detection by enabling uncertainty quantification and efficient adaptation to new clients without gradient similarity assumptions.

ABSTRACT

The recent trend towards Personalized Federated Learning (PFL) has garnered significant attention as it allows for the training of models that are tailored to each client while maintaining data privacy. However, current PFL techniques primarily focus on modeling the conditional distribution heterogeneity (i.e. concept shift), which can result in suboptimal performance when the distribution of input data across clients diverges (i.e. covariate shift). Additionally, these techniques often lack the ability to adapt to unseen data, further limiting their effectiveness in real-world scenarios. To address these limitations, we propose a novel approach, FedGMM, which utilizes Gaussian mixture models (GMM) to effectively fit the input data distributions across diverse clients. The model parameters are estimated by maximum likelihood estimation utilizing a federated Expectation-Maximization algorithm, which is solved in closed form and does not assume gradient similarity. Furthermore, FedGMM possesses an additional advantage of adapting to new clients with minimal overhead, and it also enables uncertainty quantification. Empirical evaluations on synthetic and benchmark datasets demonstrate the superior performance of our method in both PFL classification and novel sample detection.

Motivation & Objective

• To address the limitation of existing personalized federated learning (PFL) methods that focus only on concept shift (conditional distribution heterogeneity) while neglecting covariate shift (marginal distribution heterogeneity).
• To develop a PFL framework that can effectively model joint distribution heterogeneity across clients, where both input and label distributions vary.
• To enable uncertainty quantification and out-of-distribution (OOD) detection for unseen clients and novel samples in real-world federated settings.
• To design a communication-efficient, scalable method that supports fast adaptation to new clients with minimal retraining overhead.

Proposed method

• FedGMM models the joint data distribution across clients as a mixture of Gaussian components, capturing both input (covariate) and label (concept) shift.
• It employs a federated Expectation-Maximization (EM) algorithm to estimate GMM parameters in a distributed manner, maximizing the log-likelihood of observed data across clients.
• The EM algorithm is solved in closed form without requiring gradient similarity assumptions, enabling stable and efficient convergence.
• The method supports uncertainty quantification by estimating the probability of a sample under the learned GMM, enabling OOD detection.
• Personalized models are derived from the GMM components, allowing each client to select or combine components based on its local data distribution.
• The framework allows for rapid adaptation to new clients by reusing the global GMM structure and fine-tuning component weights with minimal data.

Experimental results

Research questions

• RQ1 Can a federated learning framework effectively model joint distribution heterogeneity, including both covariate and concept shift, across clients?
• RQ2 How can uncertainty quantification be integrated into personalized federated learning to detect out-of-distribution samples and new clients?
• RQ3 Can a GMM-based approach achieve better personalization performance than existing PFL methods that only model conditional distribution shifts?
• RQ4 To what extent can the proposed method adapt to new clients with minimal retraining and communication overhead?
• RQ5 Does the federated EM algorithm used in FedGMM ensure convergence and stability in non-iid, heterogeneous FL settings?

Key findings

• FedGMM significantly outperforms state-of-the-art PFL baselines on three benchmark federated learning datasets and a synthetic dataset in terms of personalized accuracy.
• The method achieves superior performance in out-of-distribution (OOD) detection, with higher AUROC and AUPRC scores than existing uncertainty-aware PFL methods.
• FedGMM enables effective uncertainty quantification at the sample level, allowing reliable detection of novel or anomalous inputs.
• The federated EM algorithm converges reliably and efficiently, even under strong data heterogeneity, without requiring gradient similarity assumptions.
• The framework supports fast adaptation to new clients by reusing the global GMM structure, reducing training cost and communication overhead.
• Empirical results demonstrate that modeling joint distribution heterogeneity via GMMs leads to more robust and personalized models compared to methods assuming stable marginal distributions.