FedBiP: Heterogeneous One-Shot Federated Learning with Personalized Latent Diffusion Models

TL;DR

FedBiP tackles heterogeneous One-Shot Federated Learning by personalizing pretrained Latent Diffusion Models at both the instance and concept levels to generate client-aligned synthetic data without sharing raw data. Instance-level latent vectors $z_i(T)$ capture per-sample variations, while domain and category concept vectors $S$ and $V_y$ guide asynchronous, server-side diffusion to produce diverse, label-consistent samples for training a classifier. The approach yields state-of-the-art results on OSFL benchmarks with feature-space heterogeneity and demonstrates robust performance on medical and satellite datasets with label heterogeneity, alongside favorable privacy properties and scalability. By effectively mitigating distribution shifts between pretrained LDMs and client data, FedBiP offers a practical, privacy-preserving data augmentation solution for real-world heterogeneous OSFL deployments.

Abstract

One-Shot Federated Learning (OSFL), a special decentralized machine learning paradigm, has recently gained significant attention. OSFL requires only a single round of client data or model upload, which reduces communication costs and mitigates privacy threats compared to traditional FL. Despite these promising prospects, existing methods face challenges due to client data heterogeneity and limited data quantity when applied to real-world OSFL systems. Recently, Latent Diffusion Models (LDM) have shown remarkable advancements in synthesizing high-quality images through pretraining on large-scale datasets, thereby presenting a potential solution to overcome these issues. However, directly applying pretrained LDM to heterogeneous OSFL results in significant distribution shifts in synthetic data, leading to performance degradation in classification models trained on such data. This issue is particularly pronounced in rare domains, such as medical imaging, which are underrepresented in LDM's pretraining data. To address this challenge, we propose Federated Bi-Level Personalization (FedBiP), which personalizes the pretrained LDM at both instance-level and concept-level. Hereby, FedBiP synthesizes images following the client's local data distribution without compromising the privacy regulations. FedBiP is also the first approach to simultaneously address feature space heterogeneity and client data scarcity in OSFL. Our method is validated through extensive experiments on three OSFL benchmarks with feature space heterogeneity, as well as on challenging medical and satellite image datasets with label heterogeneity. The results demonstrate the effectiveness of FedBiP, which substantially outperforms other OSFL methods.

Figures (7)

• Figure 1: Feature map visualization of client images (real), synthetic images by prompted pretrained LDM (pretrained), and our method (FedBiP) on two datasets. FedBiP mitigates the strong distribution shifts between pretrained LDM and client local data.
• Figure 2: Schematic illustration of Federated Bi-Level Personalization (FedBiP). (①) Each client executes bi-level personalization and obtains latent vectors $z^k(T)$ and concept vectors $S^k, V^k$. (②) The central server integrates the vectors into the generation process of the pretrained Latent Diffusion Model $\theta$. (③) The classification model $\phi$ is optimized using synthetic images.
• Figure 3: Validation results with varying client numbers on DomainNet.
• Figure 4: Validation results with synthesizing different numbers of images at central server.
• Figure 5: FedBiP privacy analysis: (1) Visual: The reproduced images are notably dissimilar to the original images $x_i$ and $x_{i'}$. Besides, the retrieved images exhibit visual discrepancies compared to synthetic $\tilde{x}_i$. (2) Statistical: The pixel value histogram of $z(T)$ resembles a standard Gaussian distribution more closely compared to $\overline{z}(0)$, making it hard to extract private information from $z(T)$.