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

TL;DR

FedDM is a new federated training framework for diffusion models that improves communication efficiency and handles data heterogeneity, with proven convergence and effective performance on multiple image datasets.

Contribution

The paper introduces FedDM, a suite of federated diffusion model training algorithms with theoretical convergence guarantees and communication-efficient features.

Findings

Maintains high image generation quality across resolutions.

Significantly improves communication efficiency (up to 4x).

Handles non-IID data effectively with minimal quality loss.

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).