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

我们介绍了FedDM，这是一个新颖的训练框架，专门设计用于联邦训练扩散模型。我们的理论分析证明了在联邦设置下训练扩散模型的收敛性，并提出了保证此收敛性的具体条件。我们提出了一套训练算法，利用U-Net架构作为我们扩散模型的骨干。这些算法包括基本的联邦平均变体FedDM-vanilla，FedDM-prox用于处理客户端之间的数据异构性，以及FedDM-quant，它包含一个量化模块，以减少模型更新大小，从而增强联邦网络中的通信效率。 我们在FashionMNIST（28x28分辨率）、CIFAR-10（32x32分辨率）和CelebA（64x64分辨率）以及LSUN Church Outdoors（256x256分辨率）上对DDPM进行了评估，在图像模态上专注于评估。我们的评估结果表明，FedDM算法在各种图像分辨率下均保持高质量的生成效果。同时，本地训练目标中使用量化更新和近端项显著增强了通信效率（高达4倍）和模型收敛性，特别是在非独立同分布数据设置中，但代价是增加FID分数（高达1.75倍）。

本篇论文旨在解决联邦学习中扩散模型的训练问题。通过提出一种新的训练框架FedDM，探讨在联邦学习环境下扩散模型的收敛性，并提出一系列算法来优化训练效果。

FedDM框架下的扩散模型训练算法，包括FedDM-vanilla、FedDM-prox以及FedDM-quant，能够有效地处理联邦学习中的数据异质性问题，并通过量化模块来提高通信效率。

论文在FashionMNIST、CIFAR-10、CelebA以及LSUN Church Outdoors数据集上进行了实验，并且开源了代码。实验结果表明，FedDM算法能够在不同的图像分辨率下保持高质量的生成效果，并且能够显著提高通信效率和模型收敛速度。

在联邦学习领域，还有一些相关研究，如Federated Learning with Non-IID Data、Federated Learning: Strategies for Improving Communication Efficiency等。