Navigating Heterogeneity and Privacy in One-Shot Federated Learning with Diffusion Models

TL;DR

This paper tackles data heterogeneity and privacy in one-shot federated learning by introducing FedDiff, a diffusion-model-based approach that trains class-conditioned generators on clients and synthesizes a global dataset at the server for discriminative model training. It demonstrates that diffusion models offer high sample quality and diversity, yielding substantial gains over state-of-the-art one-shot FL methods under both standard and differential privacy settings. To counteract DP-induced quality drops, the authors propose Fourier Magnitude Filtering (FMF), which improves the usefulness of generated samples for global training. The results show FedDiff achieves notable accuracy improvements (often 5%–20%) across heterogeneous data regimes and datasets, with FMF providing further boosts in challenging DP scenarios, illustrating practical impact for privacy-aware, communication-efficient FL systems.

Abstract

Federated learning (FL) enables multiple clients to train models collectively while preserving data privacy. However, FL faces challenges in terms of communication cost and data heterogeneity. One-shot federated learning has emerged as a solution by reducing communication rounds, improving efficiency, and providing better security against eavesdropping attacks. Nevertheless, data heterogeneity remains a significant challenge, impacting performance. This work explores the effectiveness of diffusion models in one-shot FL, demonstrating their applicability in addressing data heterogeneity and improving FL performance. Additionally, we investigate the utility of our diffusion model approach, FedDiff, compared to other one-shot FL methods under differential privacy (DP). Furthermore, to improve generated sample quality under DP settings, we propose a pragmatic Fourier Magnitude Filtering (FMF) method, enhancing the effectiveness of generated data for global model training.

Figures (7)

• Figure 1: Our one-shot FL approach, FedDiff. We first train a class-conditioned diffusion model on local data $\mathbf{x}$ at the clients. After completing training, the local diffusion models $D_{0}$, $D_{1}$, ..., $D_{c}$ are gathered by the server, where they are used to generate data $\mathbf{z}_0$, $\mathbf{z}_0$, ..., $\mathbf{z}_{c}$, which are combined to form the global training data $\mathbf{G}$. The global model is then trained on this synthetic dataset $\mathbf{G}$.
• Figure 2: Random sets of generated samples from FedCVAE and our FedDiff. By leveraging the intrinsic properties of diffusion models, which are well-aligned with the requirements of one-shot FL, we achieve substantial benefits in sample quality and subsequent global model performance.
• Figure 3: Histogram of distance scores for all generated samples at $\epsilon=50$ to corresponding closest training image by Eq. \ref{['eq:mem']} on each dataset. Note that the y-axis is in log scale, as there are very few samples with lower scores.
• Figure 4: Qualitative comparison of original training samples and generated samples at $\epsilon=50$. We show the closest 30 samples via the similarity metric in Equation \ref{['eq:mem']}. In each stacked row, the original samples are on top, with the corresponding nearest generated image immediately below. Even under the loosest privacy guarantee of $\epsilon=50$, we do not see blatant memorization.
• Figure 5: Results with our Fourier Magnitude Filtering under DP. FedDiff is in green and FedDiff+FMF in orange. Our FMF approach provides a simple way to boost accuracy, especially in more challenging scenarios such as lower $\epsilon$ budgets and more complex datasets. We plot the mean across three runs with different seeds for each setting. Additional $\gamma$ ablations are provided in the Supplementary Material.