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

Jayneel Vora,Nader Bouacida,Aditya Krishnan,Prasant Mohapatra

2024-07-20

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

Machine Learning,Computer Vision and Pattern Recognition,Distributed, Parallel, and Cluster Computing

What problem does this paper attempt to address?

The paper attempts to address the issues of communication inefficiency and data heterogeneity when training diffusion models in a federated learning environment.

Specifically:

1. **Communication Efficiency**: In federated learning, frequent model updates lead to significant communication overhead, especially when dealing with high-resolution image data. The paper proposes a quantized update mechanism (FedDM-quant) to improve communication efficiency by reducing the size of model updates.

2. **Data Heterogeneity**: Data in federated learning is typically distributed across multiple clients, and each client's data distribution may differ, leading to data heterogeneity issues. The paper introduces a proximal term to the local objective function (FedDM-prox) to mitigate the impact of data heterogeneity on model training.

3. **Model Convergence**: The paper provides theoretical analysis proving the convergence of diffusion models in a federated learning environment and proposes specific conditions to ensure this convergence.

4. **Generation Quality**: The paper evaluates the proposed algorithms' generation quality across different datasets and image resolutions, demonstrating that these algorithms can effectively improve communication efficiency and model convergence while maintaining high-quality generation.

In summary, the paper aims to address the communication efficiency and data heterogeneity issues faced by federated diffusion models in practical applications by proposing new federated learning frameworks and algorithms, thereby achieving efficient and stable model training.