Personalized federated learning improves client grouping using classifier heads
From Classifier Heads to Encoder Prototypes: Personalized Federated Learning with Head-Derived Clustering
Machine Learning
Summary
Federated learning is a way for many computers to learn together without sharing their data directly. The challenge is figuring out which computers should share the parts of what they learn. The authors propose using special parts called classifier heads that each computer creates to group similar learners together. Their method combines the best shared parts from these groups while keeping each computer’s unique parts intact, letting every computer learn better for its own data. They tested this on image recognition tasks and found it performs well even for the hardest cases.
What this means in practice
- •For machine learning engineers: Improve federated learning models by grouping clients based on classifier head similarity to share encoder components more effectively.
- •For mobile device developers: Deploy federated learning algorithms that preserve personalized classification performance without costly local retraining or fine-tuning.
Authors
Amirardalan Dehghanpour, Arian Noori, Ehsan Tanghatari, Naimeh Omidvar, Hamid R. Rabiee
Abstract
Personalized federated learning must decide which clients should share representations. We introduce HDC-FL (Head-Derived Clustering for Personalized Federated Learning), which reads that decision off a signal shared-encoder training already produces: the client-specific classifi- er heads. At periodic refresh rounds the server projects the cached heads into a normalized low-rank descriptor space, groups the descriptors, builds one encoder prototype per group from the cached client encoders, and returns to each client a soft mixture of those prototypes paired with its own unchanged head. Client-specific heads are never averaged or distributed to oth- er clients; the shared-encoder trajectory is left intact; and no second optimization trajectory, additional local objective, or post-training fine-tuning stage is required. Across 11 CIFAR-10 and CIFAR-100 configurations covering label-subset, Dirichlet and rotation heterogeneity with ResNet-18 and MobileNetV2 encoders, HDC-FL is competitive with the strongest personalized and clustered baselines in macro-averaged mean accuracy and attains the numerically highest macro-averaged worst-client and bottom-20% accuracy among the evaluated methods. The tail results distinguish methods that look similar when only their means are compared. A controlled donor-swap diagnostic further finds positive donor-ranking associations between head similarity and encoder compatibility in all three settings tested, most clearly under covariate shift.