Heterogeneity‐Aware Federated Learning via Heterogeneity Identification and Adaptive Optimization

ABSTRACT

Data heterogeneity remains a major challenge in building robust and scalable federated learning systems for real‐world distributed applications, while many existing methods are designed for predefined heterogeneity settings and therefore have limited adaptability when multiple types of heterogeneity coexist across participants. To enable heterogeneity‐aware federated learning, we propose an integrated framework that combines parameter‐oriented heterogeneity identification with adaptive optimization. Specifically, a Transformer‐based classifier identifies the heterogeneity type of each participant using uploaded model parameters. The attention weights learned during identification are further interpreted as heterogeneity‐related parameter sensitivities and reused to guide subsequent optimization. For local updating, we introduce an attention‐guided prior loss that imposes stronger constraints on heterogeneity‐sensitive parameter regions while preserving local adaptability in less sensitive regions. For global aggregation, a weighted Gaussian‐product scheme is employed to approximate the global posterior and mitigate aggregation errors caused by mixed heterogeneous data. The proposed identification method achieves a test accuracy of 98.22%. Extensive experiments on benchmark datasets and real‐world medical image classification tasks demonstrate the effectiveness of the proposed framework under mixed data heterogeneity. Furthermore, sparse parameter exchange based on the attention weights reduces communication cost compared with FedAVG, supporting the practical applicability of the proposed framework in heterogeneous distributed learning environments.