Robust Federated Learning on Edge Devices with Domain Heterogeneity

TL;DR

This paper presents FedAPC, a novel federated learning framework that uses prototype augmentation and contrastive learning to improve model robustness and generalization across heterogeneous edge device domains.

Contribution

Introduces FedAPC, a prototype-based federated learning framework that enhances feature diversity and robustness under domain heterogeneity using augmented prototypes.

Findings

Outperforms state-of-the-art methods on Office-10 and Digits datasets.

Improves generalization of federated models across heterogeneous domains.

Reduces overfitting to specific domains.

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.