FedSHIBU: Federated Similarity-based Head Independent Body Update

Most federated learning algorithms like FedAVG aggregate client models to obtain a global model. However, this leads to loss of information, especially when the data distribution is highly heterogeneous across clients. As a motivation for this paper, we first show that data-specific global models (where the clients are grouped based on their data distribution) produce higher accuracy over FedAVG. This suggests a potential performance improvement if clients trained on similar data have a higher importance in model aggregation. We use data representations from extractors of client models to quantify data similarity. We propose using a weighted aggregation of client models where the weight is calculated based on the similarity of client data. Similar to FedBABU, the proposed client representation similarity-based aggregation is applied only on extractors. We empirically show that the proposed method enhances global model performance in heterogeneous data distributions.