Prototypes guided model transformations between personalization and generalization in federated learning

Federated Learning (FL) has gained popularity due to its ability to train a collaborative model while preserving privacy. However, it still faces limitations when dealing with heterogeneous data, primarily manifesting as the performance degradation of the global model and the inadaptability of the single global model to the divergence of client data distributions. Although the above issues are summarized by researchers as goals for generalization and personalization, few studies have simultaneously addressed both goals, with most prioritizing one over the other. In this paper, it is demonstrated that the FL iteration already incorporates model transformations between personalization and generalization, with a focus on ensuring the smooth functionality of these transformations under high data heterogeneity. Specifically, a novel Federated Prototype Transformation Framework (FedPT) is proposed, which is capable of generating a well-performing generalized model as well as personalized models simultaneously. FedPT constructs local prototype classifiers that explicitly guide personalized model optimization during local training, and these can be aggregated into a global prototype classifier suitable for generic tasks. The momentum update design retains the global knowledge in local training and aligns features between clients, which results in a smoother iteration. Moreover, an improved sample-level contrastive loss is presented to dig into deeper representations, achieving high-quality prototype generation even for missing or imbalanced classes. Experimental results demonstrate the exceptional performance of FedPT in both generalization and personalization tasks, outperforming latest methods.