FedPhD: Federated Pruning with Hierarchical Learning of Diffusion Models

TL;DR

FedPhD introduces a hierarchical federated learning approach with structured pruning for diffusion models, significantly reducing communication costs and improving image generation quality in distributed settings.

Contribution

This paper presents FedPhD, a novel federated learning method that combines hierarchical aggregation and structured pruning to efficiently train diffusion models across distributed clients.

Findings

Achieves up to 88% reduction in communication costs.

Improves FID scores by at least 34% over baselines.

Reduces computation and communication resource usage by 44%.

Abstract

Federated Learning (FL), as a distributed learning paradigm, trains models over distributed clients' data. FL is particularly beneficial for distributed training of Diffusion Models (DMs), which are high-quality image generators that require diverse data. However, challenges such as high communication costs and data heterogeneity persist in training DMs similar to training Transformers and Convolutional Neural Networks. Limited research has addressed these issues in FL environments. To address this gap and challenges, we introduce a novel approach, FedPhD, designed to efficiently train DMs in FL environments. FedPhD leverages Hierarchical FL with homogeneity-aware model aggregation and selection policy to tackle data heterogeneity while reducing communication costs. The distributed structured pruning of FedPhD enhances computational efficiency and reduces model storage requirements in clients. Our experiments across multiple datasets demonstrate that FedPhD achieves high model performance regarding Fr\'echet Inception Distance (FID) scores while reducing communication costs by up to $88\%$. FedPhD outperforms baseline methods achieving at least a $34\%$ improvement in FID, while utilizing only $56\%$ of the total computation and communication resources.