Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

FedDM: Enhancing Communication Efficiency and Handling Data Heterogeneity in Federated Diffusion Models

Jayneel Vora, Nader Bouacida, Aditya Krishnan, Prasant Mohapatra

TL;DR

FedDM introduces a federated learning framework for diffusion models with theoretical convergence guarantees and three variants to balance accuracy and communication efficiency under data heterogeneity. It defines FedDM-vanilla, FedDM-prox, and FedDM-Quant, combining contraction-based convergence analysis, a proximal term for non-IID data, and post-training quantization to reduce update size. Empirical results across DDPMs and LDMs on datasets from 28x28 to 256x256 demonstrate that FedDM maintains quality with notable communication savings, especially with quantized updates, albeit with some degradation at higher resolutions. The work advances privacy-preserving, scalable training of diffusion models and highlights directions for privacy analysis and extension to other modalities and conditional generation.

Abstract

We introduce FedDM, a novel training framework designed for the federated training of diffusion models. Our theoretical analysis establishes the convergence of diffusion models when trained in a federated setting, presenting the specific conditions under which this convergence is guaranteed. We propose a suite of training algorithms that leverage the U-Net architecture as the backbone for our diffusion models. These include a basic Federated Averaging variant, FedDM-vanilla, FedDM-prox to handle data heterogeneity among clients, and FedDM-quant, which incorporates a quantization module to reduce the model update size, thereby enhancing communication efficiency across the federated network. We evaluate our algorithms on FashionMNIST (28x28 resolution), CIFAR-10 (32x32 resolution), and CelebA (64x64 resolution) for DDPMs, as well as LSUN Church Outdoors (256x256 resolution) for LDMs, focusing exclusively on the imaging modality. Our evaluation results demonstrate that FedDM algorithms maintain high generation quality across image resolutions. At the same time, the use of quantized updates and proximal terms in the local training objective significantly enhances communication efficiency (up to 4x) and model convergence, particularly in non-IID data settings, at the cost of increased FID scores (up to 1.75x).

FedDM: Enhancing Communication Efficiency and Handling Data Heterogeneity in Federated Diffusion Models

TL;DR

FedDM introduces a federated learning framework for diffusion models with theoretical convergence guarantees and three variants to balance accuracy and communication efficiency under data heterogeneity. It defines FedDM-vanilla, FedDM-prox, and FedDM-Quant, combining contraction-based convergence analysis, a proximal term for non-IID data, and post-training quantization to reduce update size. Empirical results across DDPMs and LDMs on datasets from 28x28 to 256x256 demonstrate that FedDM maintains quality with notable communication savings, especially with quantized updates, albeit with some degradation at higher resolutions. The work advances privacy-preserving, scalable training of diffusion models and highlights directions for privacy analysis and extension to other modalities and conditional generation.

Abstract

We introduce FedDM, a novel training framework designed for the federated training of diffusion models. Our theoretical analysis establishes the convergence of diffusion models when trained in a federated setting, presenting the specific conditions under which this convergence is guaranteed. We propose a suite of training algorithms that leverage the U-Net architecture as the backbone for our diffusion models. These include a basic Federated Averaging variant, FedDM-vanilla, FedDM-prox to handle data heterogeneity among clients, and FedDM-quant, which incorporates a quantization module to reduce the model update size, thereby enhancing communication efficiency across the federated network. We evaluate our algorithms on FashionMNIST (28x28 resolution), CIFAR-10 (32x32 resolution), and CelebA (64x64 resolution) for DDPMs, as well as LSUN Church Outdoors (256x256 resolution) for LDMs, focusing exclusively on the imaging modality. Our evaluation results demonstrate that FedDM algorithms maintain high generation quality across image resolutions. At the same time, the use of quantized updates and proximal terms in the local training objective significantly enhances communication efficiency (up to 4x) and model convergence, particularly in non-IID data settings, at the cost of increased FID scores (up to 1.75x).
Paper Structure (24 sections, 14 equations, 3 figures, 3 tables, 2 algorithms)

This paper contains 24 sections, 14 equations, 3 figures, 3 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Communication Efficiency in FL
  4. Quantized FL Model Updates
  5. Federated Diffusion Models
  6. Method
  7. Preliminaries
  8. Diffusion Models
  9. Federated Averaging
  10. Post-Training Quantization
  11. Approach 1: FedDM-vanilla
  12. Problem Formulation
  13. Federated Aggregation Algorithm
  14. Convergence Analysis
  15. Approach 2:FedDM-prox
  16. ...and 9 more sections

Figures (3)

  • Figure 1: Simplified visualization of FedDM training algorithms presented in this work.
  • Figure 2: 64x64 images sampled using the initial random noise using models trained on the CelebA dataset with the following variants: (a) centralized (b) FedDM-vanilla with 10 clients, 6 contributing clients, 20 global epochs and 15 local epochs. (c) FedDM-vanilla with weights quantized to 8 bits, 10 clients, 6 contributing clients, 20 global epochs and 15 local epochs. (d) FedDM-quant with weights quantized to 8 bits, 10 clients, 6 contributing clients, 20 global epochs and 15 local epochs.
  • Figure 3: Performance of FedAvg-vanilla and FedAvg-prox for differing degrees of data skewness for the DDPM model trained on the CelebA dataset