Rethinking Client Drift in Federated Learning: A Logit Perspective

TL;DR

This paper identifies that client drift in federated learning is caused by increasing logit differences due to data heterogeneity and catastrophic forgetting, and proposes FedCSD, a class prototype similarity distillation method, to align local and global models effectively.

Contribution

The paper introduces FedCSD, a novel federated learning algorithm that uses class prototype similarity distillation to mitigate client drift caused by data heterogeneity.

Findings

FedCSD outperforms state-of-the-art methods in heterogeneous settings.

Global logits quality is improved with adaptive masking.

Alignment of local and global models reduces client drift.

Abstract

Federated Learning (FL) enables multiple clients to collaboratively learn in a distributed way, allowing for privacy protection. However, the real-world non-IID data will lead to client drift which degrades the performance of FL. Interestingly, we find that the difference in logits between the local and global models increases as the model is continuously updated, thus seriously deteriorating FL performance. This is mainly due to catastrophic forgetting caused by data heterogeneity between clients. To alleviate this problem, we propose a new algorithm, named FedCSD, a Class prototype Similarity Distillation in a federated framework to align the local and global models. FedCSD does not simply transfer global knowledge to local clients, as an undertrained global model cannot provide reliable knowledge, i.e., class similarity information, and its wrong soft labels will mislead the optimization of local models. Concretely, FedCSD introduces a class prototype similarity distillation to align the local logits with the refined global logits that are weighted by the similarity between local logits and the global prototype. To enhance the quality of global logits, FedCSD adopts an adaptive mask to filter out the terrible soft labels of the global models, thereby preventing them to mislead local optimization. Extensive experiments demonstrate the superiority of our method over the state-of-the-art federated learning approaches in various heterogeneous settings. The source code will be released.