DACOA: Diffusion-Aligned Coherent Augmentation and Consistency Constraint Strategies for Federated Domain Generalization

Guangshuo Wang, Yuesheng Zhu, Guibo Luo

Abstract: Federated learning is a privacy-preserving, decentralized machine learning approach, which faces the challenge of non-identically distributed (non-iid) data between train-test data and across client domains. Previous methods generally exchange domain information between clients to perform data augmentation. While bringing privacy risks and communication costs, these methods also destroy the coherence of images. To tackle these issues, we propose Diffusion-Aligned Coherent Augmentation (DACOA), a diffusion-based and text-guided style transfer method. By composing different domains and labels as prompts, clients are guided to perform cross-domain image augmentation with high quality, thereby learning robust representation against domain shift. To better utilize the augmentation results and help the model focus on semantic information, we conduct alignment on both the feature dimensions and prediction results. We introduce the Domain Aligning Contrastive Learning (DaCon) loss, which brings the feature similarity of the same label closer. Also, we introduce the Semantic-Consistency KL (SCKL) loss, aligning the prediction results of the augmented images with the original classification results. Our model outperforms state-of-the-art FedDG methods through comprehensive experiments. What’s more, we achieve 3.21%, 1.76% and 5.01% improvement on PACS, Office-Home, and Digits-DG benchmarks. Ablation study validates the efficacy of each module.