Addressing Client Drift in Federated Learning via Class-Prototype Similarity Distillation and Adaptive Mask

Federated learning (FL) enables multiple clients to learn collaboratively in a distributed way, allowing for privacy protection. However, the real-world nonindependent and identically distributed (non-IID) data will lead to client drift, which degrades the performance of FL. Interestingly, we find that the logit difference between the local and global models increases as the model is continuously updated, which is the primary factor behind performance degradation. This is mainly due to catastrophic forgetting caused by non-IID data 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 logits of local and global models. FedCSD does not simply transfer global knowledge to local clients, as an insufficiently trained 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 leverages the similarity between local logits and the global prototype to refine the global logits, thereby enhancing its class similarity information. Furthermore, FedCSD adopts an adaptive mask to filter out the terrible soft labels of the global models, thereby preventing them from misleading local optimization. Extensive experiments demonstrate the superiority of our method over the state-of-the-art FL approaches in various non-IID settings. Code is publicly available at https://github.com/IAMJackYan/FedCSD.