Mitigating Domain Shift in Federated Learning via Intra- and Inter-Domain Prototypes

TL;DR

This work tackles domain shift in federated learning by introducing I$^2$PFL, which jointly leverages intra-domain prototypes gathered locally via MixUp-based augmentation and inter-domain generalized prototypes computed with a prototype reweighting scheme. The method employs Generalized Prototypes Contrastive Learning (GPCL) and Augmented Prototype Alignment (APA) losses to transfer cross-domain knowledge while enriching local feature diversity, with EMA updates ensuring stable prototype evolution. Empirical results on Digits, Office-10, and PACS show improved generalization across domains and strong performance in unseen-domain settings, with ablations confirming the contributions of intra-domain augmentation, inter-domain reweighting, and prototype-based contrastive learning. The approach offers a practical pathway to robust FL under real-world multi-domain heterogeneity, while preserving privacy through prototype averaging operations.

Abstract

Federated Learning (FL) has emerged as a decentralized machine learning technique, allowing clients to train a global model collaboratively without sharing private data. However, most FL studies ignore the crucial challenge of heterogeneous domains where each client has a distinct feature distribution, which is popular in real-world scenarios. Prototype learning, which leverages the mean feature vectors within the same classes, has become a prominent solution for federated learning under domain shift. However, existing federated prototype learning methods focus soley on inter-domain prototypes and neglect intra-domain perspectives. In this work, we introduce a novel federated prototype learning method, namely I$^2$PFL, which incorporates $\textbf{I}$ntra-domain and $\textbf{I}$nter-domain $\textbf{P}$rototypes, to mitigate domain shift from both perspectives and learn a generalized global model across multiple domains in federated learning. To construct intra-domain prototypes, we propose feature alignment with MixUp-based augmented prototypes to capture the diversity within local domains and enhance the generalization of local features. Additionally, we introduce a reweighting mechanism for inter-domain prototypes to generate generalized prototypes that reduce domain shift while providing inter-domain knowledge across multiple clients. Extensive experiments on the Digits, Office-10, and PACS datasets illustrate the superior performance of our method compared to other baselines.

Figures (8)

• Figure 1: Illustration of domain shift challenge in Federated Learning (FL). We address this from two perspectives: inter-domain, referring to distribution shifts across domains, and intra-domain, capturing variations within the same domain. In the figure (right), we measure the inter-domain prototype variance $\overline{\mathrm{Var}}$, which is defined as the average distance between local prototypes and the generalized prototype for each class, averaged across all clients and classes: $\mathrm{Var}_k =\frac{1}{M}\sum^M_{m=1}\bigl\lVert p^k_m - g^k\bigr\rVert_2^2$, $\overline{\mathrm{Var}}=\frac{1}{K}\sum_{k=1}^{K}\mathrm{Var}_k$, where $p^k_m$ is the local prototype of class $k$ from client $m$, and $g^k$ is the generalized prototype for class $k$. It can be seen that our method reduces inter-domain prototype variance, highlighting the generalization across domains.
• Figure 2: Illustration of I$^2$PFL. Clients first upload their local prototypes based on Eq. \ref{['eq:localproto']} to the server. We introduce the prototype reweighting scheme to generate the Generalized Prototypes $G^{t+1}$ based on Eq. \ref{['eq:gen_proto']} and update them with $G^{t}$ from the previous round using the Exponential Moving Average from Eq. \ref{['eq:ema_gen_proto']}. We provide inter-domain knowledge from the Generalized Prototypes with $\mathcal{L}_{GPCL}$ from Eq. \ref{['loss_gpc']} and enhance the local feature diversity with $\mathcal{L}_{APA}$ based on Eq. \ref{['loss_APA']} using the Augmented Prototype from Eq. \ref{['eq:augment_proto']}.
• Figure 3: Illustration of Prototype Reweighting scheme. We present the prototype reweighting scheme of the prototypes from different domains in the same semantic class.
• Figure 4: t-SNE Visualization of features in the Digits dataset.
• Figure 5: Visualization of training curves of average test accuracy on three datasets under the domain shift setting.