FedDAP: Domain-Aware Prototype Learning for Federated Learning under Domain Shift

TL;DR

FedDAP introduces domain-aware prototypes in federated learning to improve model performance under domain shift by constructing and utilizing domain-specific prototypes for better feature alignment.

Contribution

It proposes a novel domain-aware prototype construction and alignment method that preserves domain information and enhances generalization in federated learning.

Findings

FedDAP outperforms existing methods on DomainNet, Office-10, and PACS datasets.

Domain-specific prototypes improve local feature alignment and global model generalization.

The approach effectively mitigates domain shift in federated learning scenarios.

Abstract

Federated Learning (FL) enables decentralized model training across multiple clients without exposing private data, making it ideal for privacy-sensitive applications. However, in real-world FL scenarios, clients often hold data from distinct domains, leading to severe domain shift and degraded global model performance. To address this, prototype learning has been emerged as a promising solution, which leverages class-wise feature representations. Yet, existing methods face two key limitations: (1) Existing prototype-based FL methods typically construct a $\textit{single global prototype}$ per class by aggregating local prototypes from all clients without preserving domain information. (2) Current feature-prototype alignment is $\textit{domain-agnostic}$, forcing clients to align with global prototypes regardless of domain origin. To address these challenges, we propose Federated Domain-Aware Prototypes (FedDAP) to construct domain-specific global prototypes by aggregating local client prototypes within the same domain using a similarity-weighted fusion mechanism. These global domain-specific prototypes are then used to guide local training by aligning local features with prototypes from the same domain, while encouraging separation from prototypes of different domains. This dual alignment enhances domain-specific learning at the local level and enables the global model to generalize across diverse domains. Finally, we conduct extensive experiments on three different datasets: DomainNet, Office-10, and PACS to demonstrate the effectiveness of our proposed framework to address the domain shift challenges. The code is available at https://github.com/quanghuy6997/FedDAP.