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

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.

• The paper introduces FedGMI, a framework that leverages VAEs for modeling client data as convex mixtures of inherent distributions and estimating mixture coefficients.
• It decouples the federated learning objective into learning generative VAEs, specialized classifiers, and accurate mixture inference, ensuring both personalization and shared structure.
• The framework demonstrates robust convergence with bounded partition error and superior performance on benchmarks like MNIST, CIFAR-10, and EMNIST.

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

Introduction

This work addresses the federated learning (FL) challenge of modeling client heterogeneity, particularly for the case where local data distributions are convex mixtures of latent, shared distributions. Standard approaches to FL, namely Clustered Federated Learning (CFL) and Personalized Federated Learning (PFL), have well-known trade-offs: CFL leverages shared structure but ignores client-specific properties, and PFL offers personalization but underutilizes shared information. FedGMI proposes a structured, generative approach to this mixture scenario by estimating the inherent distributions with Variational Autoencoders (VAEs) and inferring the precise mixture proportions for each client. The resulting framework enables more expressive client representation, structured personalization, and theoretically robust learning.

Figure 1: High-level architecture of FedGMI. Local datasets are modeled as mixtures of inherent distributions; data division is performed, VAEs/classifiers are updated, and server synchronization involves aggregation of both VAEs and classifiers.

Core Methodology

FedGMI is formulated for an FL setting with $N$ clients, each sampling their data locally from a distribution $p_i$, which is itself a convex mixture of $M$ unknown, shared "inherent" distributions $\{q_j\}_{j=0}^{M-1}$:

$$p_i = \sum_{j=0}^{M-1} \alpha_i^j q_j; \quad \alpha_i^j \geq 0,\, \sum_j \alpha_i^j = 1.$$

The framework decouples the problem into three interlinked tasks:

1. Learning the Inherent Distributions: Each $q_j$ is modeled by a VAE, acting as both a generative model and density estimator.
2. Distribution-Expert Classifiers: For each inherent distribution, a specialized classifier $G_{q_j}$ is learned.
3. Mixture Proportion Inference: For each client, the mixture coefficients $\{\alpha_i^j\}$ are estimated via samplewise assignment using an affinity score closely tied to the VAE log-likelihood.

To assign a sample $(\mathbf{x}, y)$ to an inherent distribution, a normalized affinity score $l(\mathbf{x},y; q_j)$, proportional to the posterior $p_i$0, is computed using VAE-based likelihoods and current mixture estimates. Efficient data partitioning is then achieved via a MAP principle, and classifier/vae training is performed on these partitions.

A key architectural feature is stable VAE initialization—instead of random parameters, VAEs are initialized to maximize pairwise KL-divergence, ensuring sufficient separation and stable data partitioning. This initialization step is theoretically shown to be necessary for convergence and bounded partition error.

Theoretical Results

A rigorous convergence analysis is provided for the FedGMI framework. Under mild smoothness, convexity, and stable-initialization assumptions, the probability of assignment error for any sample—assigning it to the wrong inherent distribution—is tightly bounded as a function of key system parameters (e.g., number of distributions $p_i$1, data variance, separation parameter $p_i$2 between optima of $p_i$3). The analysis extends via a two-phase argument: first, showing convergence of the data division (partitioning) mechanism, then bounding the global FL objective heterogeneity after partition stabilization.

A critical insight is that as $p_i$4 increases, the optimal locations for the inherent distributions in parameter space approach each other, decreasing $p_i$5 and thus increasing partitioning error. This limitation motivates keeping $p_i$6 small in practical setups and motivates future work on adaptive $p_i$7 selection.

Empirical Results

Experiments extensively validate FedGMI on synthetic and real-world data, using standard FL benchmarks (MNIST, CIFAR-10, EMNIST) under carefully controlled mixture distributions.

Mixture Proportion Estimation:

For $p_i$8 inherent distributions with linear variation of mixture coefficients across clients, FedGMI robustly reconstructs the ground-truth mixture proportions, whereas baselines (IFCA, FedSoft) either binarize client assignment or produce noisy stepwise estimates.

Figure 3: Estimated and ground-truth proportions of $p_i$9 for $M$0. FedGMI closely tracks true values, especially for simpler data (MNIST) and in both high and low communication regimes.

Specialist Classifier Accuracy:

FedGMI yields classifiers that are highly specialized to their assigned inherent distributions, maintaining very high accuracy on target distributions and large performance gaps on non-target data. This effect is prominent on both MNIST and more complex datasets like CIFAR-10, outperforming alternatives especially in challenging (low-communication, high-heterogeneity) settings.

Figure 5: Classification accuracy trajectories (CIFAR-10, M=2). Each line shows performance of a cluster-specialist classifier; a pronounced and sustained accuracy gap appears for cluster-specialized models in FedGMI.

Generative Visualization:

VAE samples generated from each estimated inherent distribution are qualitatively consistent with the true underlying data generative structure (e.g., original vs. rotated images).

Figure 2: Samples synthesized by the two VAEs learned for MNIST ($M$1): One generates canonical digits, one generates $M$2-rotated digits, mirroring the dataset construction.

Mixture Estimation under Natural Heterogeneity:

On EMNIST, where inherent distributions correspond to uppercase/lowercase letter morphology rather than synthetic transformations, FedGMI still accurately infers mixture proportions per client and yields specialist classifiers with high and differentiated accuracy.

Figure 4: Estimated proportion of uppercase letters per client for EMNIST. FedGMI tracks the ground-truth trend more reliably than alternatives.

Implications and Future Directions

FedGMI introduces a highly structured, generative foundation for probabilistic mixture modeling in federated learning, supporting:

• Accurate client-level heterogeneity modeling and interpretability via mixture decomposition.
• Improved personalized inference and cluster-specific modeling without loss of collaborative efficiency.
• Robustness under communication constraints, benefiting practical deployment to edge devices.

The explicit use of generative models opens further research paths, such as leveraging more powerful generative models (e.g., diffusion models) for richer density estimation, integrating dynamic/adaptive selection of $M$3, and exploring regularization schemes for higher stability and lower sample complexity.

Conclusion

FedGMI presents a generative, mixture-aware FL framework that moves beyond hard clustering or naive personalization by capturing the shared and client-specific structure jointly. Theoretical guarantees, along with strong empirical performance under both synthetic and real heterogeneity, position FedGMI as an advanced platform for mixture inference and structured personalization in federated systems. Future work should extend to adaptive mixture complexity and to evaluating generative model advances in the federated context.

Reference: "FedGMI: Generative Model-Driven Federated Learning for Probabilistic Mixture Inference" (2605.08760)