Robust Federated Learning on Edge Devices with Domain Heterogeneity

TL;DR

The paper tackles domain heterogeneity in federated learning by introducing FedAPC, a framework that uses augmented prototypes and a prototype-contrastive loss to align local features with global class prototypes and diversify representations. Local augmented prototypes are generated via multi-view augmentations and aggregated into global prototypes $\mathcal{G}^m$, which are broadcast to clients to guide alignment through $\mathcal{L}_{APC}$. Empirical results on Digits and Office-10 show FedAPC delivering higher average accuracy and faster convergence than baselines like FedAvg, MOON, and FedProto, with particularly large gains in challenging domains such as SYN and DSLR. The work demonstrates improved robustness to domain shifts, highlighting FedAPC’s practical potential for edge deployments in heterogeneous environments such as healthcare, IoT, and personalized services.

Abstract

Federated Learning (FL) allows collaborative training while ensuring data privacy across distributed edge devices, making it a popular solution for privacy-sensitive applications. However, FL faces significant challenges due to statistical heterogeneity, particularly domain heterogeneity, which impedes the global mode's convergence. In this study, we introduce a new framework to address this challenge by improving the generalization ability of the FL global model under domain heterogeneity, using prototype augmentation. Specifically, we introduce FedAPC (Federated Augmented Prototype Contrastive Learning), a prototype-based FL framework designed to enhance feature diversity and model robustness. FedAPC leverages prototypes derived from the mean features of augmented data to capture richer representations. By aligning local features with global prototypes, we enable the model to learn meaningful semantic features while reducing overfitting to any specific domain. Experimental results on the Office-10 and Digits datasets illustrate that our framework outperforms SOTA baselines, demonstrating superior performance.

Figures (3)

• Figure 1: Problem illustration of federated learning under domain heterogeneity.
• Figure 2: Illustration of FedAPC.
• Figure 3: Visualization of average accuracy across with $100$ communication rounds on different datasets.