FedMHC: Overcoming Dimensionality and Communication Challenges for Personalized Federated Learning Using Model Head Clustering

In real Internet of Things (IoT) environments, IoT devices vary widely in data types and needs. IoT devices participating in Personalized Federated Learning (PFL) all have their own unique data characteristics, but there exist some similarities. Previous work enhances Personalized Local Models (PLMs)’ performance by clustering IoT devices’ PLM while ignoring dimensionality and communication volume, resulting in lower Global Model (GM) accuracy and PLMs’ performance. To this end, we propose a Personalized Federated Learning method based on Model Head Clustering (FedMHC). Specifically, FedMHC groups IoT devices with similar data characteristics and distributes different GMs to different device groups. FedMHC allows each IoT device to obtain a GM that best fits its local data characteristics and guides PLM training. FedMHC only clusters model head parameters on the server side. Thus, the edge server only needs to transmit the head parameters and a single shared feature extractor parameters during communication with IoT devices. The improvement can effectively address the issues of dimensionality and high communication volume. Experiments on CIFAR-100, Tiny-ImageNet, and AG News datasets demonstrate that FedMHC enhances the model accuracy by 1.79% and 5.9% in pathological heterogeneous scenarios, and by 1.43%, 0.92%, and 0.94% in practical heterogeneous scenarios, compared to the topperforming methods among 9 baselines.