Class-Wise Federated Averaging for Efficient Personalization

TL;DR

This paper introduces cwFedAvg, a personalized federated learning framework that performs class-wise model aggregation to improve personalization, using a novel regularizer to enhance class-specific encoding, achieving better results with minimal communication overhead.

Contribution

The paper proposes cwFedAvg, a class-wise federated averaging method with a new regularizer, enabling effective personalization without extra communication or local training.

Findings

cwFedAvg outperforms existing personalized FL methods.

WDR improves class-specific information encoding.

Achieves personalization with communication costs similar to FedAvg.

Abstract

Federated learning (FL) enables collaborative model training across distributed clients without centralizing data. However, existing approaches such as Federated Averaging (FedAvg) often perform poorly with heterogeneous data distributions, failing to achieve personalization owing to their inability to capture class-specific information effectively. We propose Class-wise Federated Averaging (cwFedAvg), a novel personalized FL (PFL) framework that performs Federated Averaging for each class, to overcome the personalization limitations of FedAvg. cwFedAvg creates class-specific global models via weighted aggregation of local models using class distributions, and subsequently combines them to generate personalized local models. We further propose Weight Distribution Regularizer (WDR), which encourages deep networks to encode class-specific information efficiently by aligning empirical and approximated class distributions derived from output layer weights, to facilitate effective class-wise aggregation. Our experiments demonstrate the superior performance of cwFedAvg with WDR over existing PFL methods through efficient personalization while maintaining the communication cost of FedAvg and avoiding additional local training and pairwise computations.