Robust Federated Learning on Edge Devices with Domain Heterogeneity

2025-05-15

AI Summary

To address poor generalization and slow convergence of global models in federated learning caused by domain heterogeneity, this paper proposes FedAPC, a prototype-augmented contrastive learning framework. Its core innovation is a novel prototype-augmentation mechanism: it constructs transferable global prototypes from mean features of augmented local samples, enabling alignment between local representations and global semantics and mitigating overfitting induced by domain shift. FedAPC integrates prototype learning, contrastive learning, and cross-domain representation alignment, while remaining compatible with FedAvg-style aggregation. Evaluated on Office-10 and Digits benchmarks, FedAPC significantly outperforms state-of-the-art methods, achieving an average classification accuracy gain of 3.2%, accelerating convergence by 28%, and enhancing model robustness and cross-domain generalization capability.

Technology Category

• Machine Learning: Transfer, Domain Adaptation, Multi-Task Learning
• Computer Vision: Learning & Optimization for CV
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Vertical and domain-specific search
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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.

Problem

• Addressing domain heterogeneity in Federated Learning on edge devices
• Enhancing global model generalization via prototype augmentation
• Improving feature diversity and robustness in FL frameworks

Innovation

• Federated Learning with prototype augmentation
• Enhances feature diversity and robustness
• Aligns local features with global prototypes