Federated Representation Learning for Heterogeneous Data

Abstract

Federated learning provides a promising paradigm for collaborative model training among mutually untrusted parties without sharing local data. However, data distributions in real-world federated scenarios are usually heterogeneous, which can significantly degrade global model performance. Existing approaches mainly address the non-independent and identically distributed (Non-IID) problem by optimizing aggregation strategies or sharing auxiliary data. Nevertheless, these methods often exhibit limited effectiveness under severe heterogeneity, insufficient adaptability to highly skewed data distributions, and additional privacy concerns. To address these challenges, this paper proposes AE-FRL, a federated representation learning framework for heterogeneous data. Specifically, AE-FRL employs autoencoders to extract latent representations from local training data and utilizes a representation sharing mechanism to mitigate client drift caused by Non-IID data, thereby improving global model accuracy. To overcome the limited expressiveness of latent representations, a jointly trained supervised autoencoder is introduced, which incorporates downstream classification supervision during sample reconstruction. This design enhances both the discriminative capability and representation quality of the learned latent features. Furthermore, a representation mixup mechanism is proposed to reduce privacy leakage risks during representation sharing and improve robustness against potential inference attacks. Experimantal results on four real-world datasets demonstrate that under various Non-IID settings, AE-FRL consistently outperforms baseline methods including FedAvg, FedProx, SCAFFOLD, FedNova, and FedMix, achieving higher model accuracy. In highly heterogeneous scenarios, AE-FRL achieves over 53.94% accuracy improvement without sacrificing communication efficiency or privacy preservation.