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

2026-04-08

AI Summary

This work addresses the severe domain shift caused by heterogeneous client data distributions in federated learning, which significantly degrades the generalization performance of the global model. To mitigate this issue, the authors propose a domain-aware prototype learning mechanism that explicitly preserves and leverages domain information—a first in federated learning. Specifically, the method constructs a dedicated global prototype for each domain and aggregates local prototypes from clients within the same domain through similarity-based weighting. During local training, features are encouraged to align with prototypes from their own domain while being separated from those of other domains, thereby jointly optimizing intra-domain consistency and inter-domain separability. Extensive experiments on DomainNet, Office-10, and PACS benchmarks demonstrate that the proposed approach substantially outperforms existing methods and effectively alleviates domain shift.

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.

Problem

• Federated Learning
• Domain Shift
• Prototype Learning
• Domain-Aware
• Global Model Performance

Innovation

• federated learning
• domain shift
• prototype learning
• domain-aware
• feature alignment