Rethinking Model Aggregation for Federated Learning with Domain Skew: A Knowledge-Preserving Perspective

Federated learning (FL) allows distributed clients to collaboratively train a global model in a privacy-preserving manner. However, one major challenge is domain skew, where clients' data originating from diverse domains may hinder the aggregated global model from learning a consistent representation space, resulting in poor generalizable ability in multiple domains. In this paper, we argue that the domain skew is reflected in the domain-specific biased features of each client, causing the local model's representations to collapse into a narrow low-dimensional subspace. We then propose Federated Feature Decoupling and Calibration ($F^2$DC), which liberates valuable class-relevant information by calibrating the domain-specific biased features, enabling more consistent representations across domains. A novel component, Domain Feature Decoupler (DFD), is first introduced in $F^2$DC to determine the robustness of each feature unit, thereby separating the local features into domain-robust features and domain-related features. A Domain Feature Corrector (DFC) is further proposed to calibrate these domain-related features by explicitly linking discriminative signals, capturing additional class-relevant clues that complement the domain-robust features. Finally, a domain-aware aggregation of the local models is performed to promote consensus among clients. Empirical results on three popular multi-domain datasets demonstrate the effectiveness of the proposed $F^2$DC and the contributions of its two modules. Code is available at https://github.com/mala-lab/F2DC.

Federated learning (FL) enables collaborative model training without centralizing private data, but performance often degrades under feature skew: clients share labels while the conditional input distributions $p_i(x\!\mid\!y)$ vary due to latent, client-specific appearance factors. We propose Joint Domain-Class Federated Learning (JDFL), a lightweight, optimizer-agnostic extension that makes this latent domain variation usable without sharing raw data. JDFL first infers domain clusters called pseudo-domains from brief local update signals. It then expands the classifier head to output $M\times C$, joint (domain-class) logits. This allows the model to represent domain-conditioned appearance while keeping a shared backbone. To train the expanded head we introduce two complementary supervision strategies based on simple intuitions: a similarity-aware soft-labeling that transfers evidence between nearby inferred domains while allowing domain-specific specialization, and a per-sample randomized target assignment that perturbs supervision across the joint outputs and serves as a low-cost training-time regularizer. JDFL integrates with existing standard FL methods (e.g., FedAvg, SCAFFOLD) with minimal changes. Empirically, both supervision modes consistently improve global test accuracy on standard domain-shifted image benchmarks; ablations and sensitivity studies show the gains stem from the proposed supervision and parametrization rather than mere capacity increase.

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.