Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Stable Diffusion-based Data Augmentation for Federated Learning with Non-IID Data

Mahdi Morafah, Matthias Reisser, Bill Lin, Christos Louizos

TL;DR

This paper tackles the non IID data challenge in Federated Learning by introducing Gen-FedSD, a diffusion-based data augmentation approach. Each client generates class-specific prompts and uses a pre-trained Stable Diffusion model to synthesize data that fills local distribution gaps, resulting in locally augmented IID-like datasets prior to FL training. Empirical results on CIFAR-10 and CIFAR-100 show substantial gains in accuracy and faster convergence across multiple FL baselines, with pronounced improvements under higher heterogeneity and notable reductions in communication cost. The method offers practical privacy-preserving IID data augmentation by performing synthesis locally, and its effectiveness scales with prompt diversity and the quality of the diffusion model.

Abstract

The proliferation of edge devices has brought Federated Learning (FL) to the forefront as a promising paradigm for decentralized and collaborative model training while preserving the privacy of clients' data. However, FL struggles with a significant performance reduction and poor convergence when confronted with Non-Independent and Identically Distributed (Non-IID) data distributions among participating clients. While previous efforts, such as client drift mitigation and advanced server-side model fusion techniques, have shown some success in addressing this challenge, they often overlook the root cause of the performance reduction - the absence of identical data accurately mirroring the global data distribution among clients. In this paper, we introduce Gen-FedSD, a novel approach that harnesses the powerful capability of state-of-the-art text-to-image foundation models to bridge the significant Non-IID performance gaps in FL. In Gen-FedSD, each client constructs textual prompts for each class label and leverages an off-the-shelf state-of-the-art pre-trained Stable Diffusion model to synthesize high-quality data samples. The generated synthetic data is tailored to each client's unique local data gaps and distribution disparities, effectively making the final augmented local data IID. Through extensive experimentation, we demonstrate that Gen-FedSD achieves state-of-the-art performance and significant communication cost savings across various datasets and Non-IID settings.

Stable Diffusion-based Data Augmentation for Federated Learning with Non-IID Data

TL;DR

This paper tackles the non IID data challenge in Federated Learning by introducing Gen-FedSD, a diffusion-based data augmentation approach. Each client generates class-specific prompts and uses a pre-trained Stable Diffusion model to synthesize data that fills local distribution gaps, resulting in locally augmented IID-like datasets prior to FL training. Empirical results on CIFAR-10 and CIFAR-100 show substantial gains in accuracy and faster convergence across multiple FL baselines, with pronounced improvements under higher heterogeneity and notable reductions in communication cost. The method offers practical privacy-preserving IID data augmentation by performing synthesis locally, and its effectiveness scales with prompt diversity and the quality of the diffusion model.

Abstract

The proliferation of edge devices has brought Federated Learning (FL) to the forefront as a promising paradigm for decentralized and collaborative model training while preserving the privacy of clients' data. However, FL struggles with a significant performance reduction and poor convergence when confronted with Non-Independent and Identically Distributed (Non-IID) data distributions among participating clients. While previous efforts, such as client drift mitigation and advanced server-side model fusion techniques, have shown some success in addressing this challenge, they often overlook the root cause of the performance reduction - the absence of identical data accurately mirroring the global data distribution among clients. In this paper, we introduce Gen-FedSD, a novel approach that harnesses the powerful capability of state-of-the-art text-to-image foundation models to bridge the significant Non-IID performance gaps in FL. In Gen-FedSD, each client constructs textual prompts for each class label and leverages an off-the-shelf state-of-the-art pre-trained Stable Diffusion model to synthesize high-quality data samples. The generated synthetic data is tailored to each client's unique local data gaps and distribution disparities, effectively making the final augmented local data IID. Through extensive experimentation, we demonstrate that Gen-FedSD achieves state-of-the-art performance and significant communication cost savings across various datasets and Non-IID settings.
Paper Structure (17 sections, 3 equations, 2 figures, 7 tables, 1 algorithm)

This paper contains 17 sections, 3 equations, 2 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Works
  3. Background
  4. Diffusion Models
  5. Image Generation using Stable Diffusion
  6. Federated Learning
  7. Gen-FedSD
  8. Overview
  9. Steps 1-2: Client-level Prompt Generation
  10. Step 3: Client-level Data Generation and Augmentation
  11. Experiments
  12. Experimental Setup
  13. Performance Evaluation
  14. Communication Cost
  15. Conclusion and Future Works
  16. ...and 2 more sections

Figures (2)

  • Figure 1: Overview of Gen-FedSD.
  • Figure 2: Illustration of clients data distribution for the CIFAR-10 dataset before and after augmentation with Gen-FedSD. In this example, 20 clients are utilized, partitioning CIFAR-10 Non-IID with a Dirichlet distribution (Dir(0.5)). Prior to augmentation, the local distributions of clients are heterogeneous (a), whereas after augmentation, the unique gaps in each client's local distribution are filled, resulting in a final locally augmented dataset that is independent and identically distributed (IID) for each client (b).