FedGMI: Generative Model-Driven Federated Learning for Probabilistic Mixture Inference

TL;DR

FedGMI introduces a generative model-driven federated learning framework that models client data as mixtures of shared distributions, improving personalization and robustness in heterogeneous data environments.

Contribution

It leverages VAEs to infer and utilize inherent data distributions, enabling structured personalization in federated learning.

Findings

FedGMI accurately characterizes and discriminates inherent data distributions.

It estimates mixture proportions effectively across clients.

The framework maintains robust performance under communication constraints.

Abstract

Federated Learning (FL) facilitates collaborative model training across decentralized clients while preserving data privacy by avoiding raw data exchange. Despite its potential, FL performance is often compromised by data heterogeneity across clients. To address this, Clustered Federated Learning (CFL) groups clients with similar data distributions to improve model performance, but constrained by intra-cluster heterogeneity. Conversely, Personalized Federated Learning (PFL) tailors models to individual clients, but usually neglects the underlying structural similarities among clients. In this work, we investigate a probabilistic mixture (PM) scenario, where each client's local data distribution is modeled as a convex combination of several shared inherent distributions. To effectively model this structure, we propose FedGMI, a framework that utilizes Variational Autoencoders (VAEs) as generative density estimators to represent these inherent distributions and infer the mixture components of clients' local data distributions. This approach enables structured personalization without sacrificing the benefits of collaborative learning. Extensive experiments demonstrate that FedGMI effectively characterizes and discriminate the inherent distributions, as well as accurately estimates mixture proportions. Furthermore, FedGMI maintains robust performance even under communication cost constraints.