DRCFL: Representation Driven Head Clustering for Federated Learning on Edge Devices

Clustered Federated Learning (CFL) effectively mitigates the negative impact of data heterogeneity among clients on model accuracy by grouping clients with similar data distributions for training. However, existing methods require communication of model parameters between clients and the server, leading to high communication costs. Additionally, current clustering strategies restrict the flexibility of cluster partitioning. To address these issues, we propose a Dynamic Representation Clustering for Federated Learning (DRCFL) algorithm. It decouples the model into a Representation Generation Module (RGM) and a Shared header (S-header). The former calculates client data representations, while the server clusters these representations and trains a shared header for each cluster. Clients only need to transmit data representations to the server, remarkably reducing communication costs. Additionally, the server re-clusters based on the most recent data representations in each training round, enhancing the flexibility of cluster partitions. In terms of communication cost, Extensive experimental results indicate that DRCFL achieves a reduction of 22.70%−83.33% compared to state-of-the-art (SOTA) personalized FL methods.