Personalized Federated Learning With Contrastive Momentum

In this paper, we propose pFedMo, a personalized federated learning algorithm with contrastive momentum. In pFedMo, we design a score function to personalize worker models by distilling knowledge from the aggregator's representation model so as to address the non-i.i.d. issue. To accelerate the convergence, we leverage the momentum acceleration on both the worker side and the aggregator side. However, the typical momentum without personalization does not suit well for the worker models with personalization, influencing convergence performance. To address this, we develop a personalized/contrastive momentum method for efficient momentum acceleration. We provide mathematical proof for the convergence of pFedMo on non-i.i.d. data. Extensive experiments based on real-world datasets and IoT system are conducted, verifying that pFedMo outperforms existing mainstream benchmarks, and achieves up to 35.90% accuracy increase and 3.64x training time speedup under a wide range of settings.