Model layered optimization with contrastive learning for personalized federated learning^☆

Abstract

In federated learning (FL), the distribution of data across different clients leads to the degradation of global model performance in training. Personalized Federated Learning (pFL) can address this problem through global model personalization. Researches over the past few years have calibrated differences in weights across the entire model or optimized only individual layers of the model without considering that different layers of the whole neural network have different utilities, resulting in lagged model convergence and inadequate personalization in non-IID data. In this paper, we propose model layered optimization for feature extractor and classifier (pFedEC), a novel pFL training framework personalized for different layers of the model. Our study divides the model layers into the feature extractor and classifier. We initialize the model's classifiers during model training, while making the local model's feature extractors learn the representation of the global model's feature extractors to correct each client's local training, integrating the utilities of the different layers in the entire model. Our extensive experiments show that pFedEC achieves 92.95% accuracy on CIFAR-10, outperforming existing pFL methods by approximately 1.8%. On CIFAR-100 and Tiny-ImageNet, pFedEC improves the accuracy by at least 4.2%, reaching 73.02% and 28.39%, respectively.

Keywords

Federated learning (FL) / Personalized federated learning (pFL) / Contrastive learning / Theoretical analysis

References

[1] W. Wang, J. Liang, C. Zhang, X. Liu, L. Zhu, S. Guo, Epass: efficient and privacy-preserving asynchronous payment on blockchain, in: IEEE Transactions on Depend-able and Secure Computing, 2025, pp. 1-14.

[2] X. Ren, H. Liang, Y. Wang, C. Zhang, Z. Xiong, L. Zhu, Besa: boosting encoder steal-ing attack with perturbation recovery, IEEE Trans. Inf. Forensics Secur. 20 (2025) 10007-10018.

[3] F. Zhao, B. Yang, C. Li, C. Zhang, L. Zhu, G. Liang, Two-layer consensus based on master-slave consortium chain data sharing for Internet of vehicles, IEEE Trans. Veh. Technol. 73 (2024) 13828-13838.

[4] J. Zhang, Y. Liu, Y. Hua, H. Wang, T. Song, Z. Xue, R. Ma, J. Cao, Pfllib: a beginner-friendly and comprehensive personalized federated learning library and benchmark, 2023.

[5] E. Zhu, L. Zhang, J. Wang, C. Hu, Q. Jing, W. Shi, Z. Xu, P. Ai, Z. Dai, D. Shan,Z. Ai, Personalized surgical recommendations and quantitative therapeutic insights for pa-tients with metastatic breast cancer: insights from deep learning, Cancer Innovation 3.

[6] B. McMahan, E. Moore, D. Ramage, S. Hampson, B.A. y Arcas, Communication-efficient learning of deep networks from decentralized data, in: Artificial Intelligence and Statistics, PMLR, 2017, pp. 1273-1282.

[7] T. Li, A.K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, V. Smith,Federated optimiza-tion in heterogeneous networks, Proc. Mach. Learn. Syst. 2 (2020) 429-450.

[8] S.P. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, A.T. Suresh, Scaffold: stochas-tic controlled averaging for federated learning,in:International Conference on Ma-chine Learning, PMLR, 2020, pp. 5132-5143.

[9] Q. Li, B. He, D. Song,Model-contrastive federated learning, in:IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2021, pp. 10708-10717.

[10] M.G. Arivazhagan, V. Aggarwal, A.K. Singh,S. Choudhary, Federated learning with personalization layers, arXiv:1912.00818.

[11] L. Collins, H. Hassani, A. Mokhtari, S. Shakkottai, Exploiting shared representations for personalized federated learning, in: Proceedings of the 38th International Con-ference on Machine Learning, PMLR, 2021, pp. 2089-2099.

[12] J. Oh, S. Kim, S.-Y. Yun, FedBABU: toward enhanced representation for federated im-age classification,in: International Conference on Learning Representations, 2022.

[13] J. Liu, J. Liu, H.-Z. Xu, Y. Liao, Z. Wang, Q. Ma, Yoga: adaptive layer-wise model aggregation for decentralized federated learning, IEEE/ACM Trans. Netw. 32 (2024) 1768-1780.

[14] Z. Pan, C.-R. Li, F. Yu, S. Wang, H. Wang, X. Tang, J. Zhao, Fedlf: layer-wise fair federated learning,in: AAAI Conference on Artificial Intelligence, 2024.

[15] Q. Liu, S. Sun, Y. Liang, J. Xue, M. Liu, Personalized federated learning for spatio-temporal forecasting: a dual semantic alignment-based contrastive approach,in: AAAI Conference on Artificial Intelligence, 2024.

[16] S. Seo, J. Kim, G. Kim, B. Han,Relaxed contrastive learning for federated learning, in:2024 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2024, pp. 12279-12288.

[17] B. Kang, S. Xie, M. Rohrbach, Z. Yan, A. Gordo, J. Feng, Y. Kalantidis, Decoupling representation and classifier for long-tailed recognition, in: CoRR, 2014.

[18] H. Yu, N. Zhang, S. Deng, Z. Yuan, Y. Jia, H. Chen,The devil is the classifier: inves-tigating long tail relation classification with decoupling analysis, arXiv:2009.07022.

[19] X. Yang, W. Huang, M. Ye, Dynamic personalized federated learning with adaptive differential privacy, in: Neural Information Processing Systems, 2023.

[20] Y. He, L. Wei, F. Chen, H. Zhang, J. Yu, H. Wang, Fedai: federated recommendation system with anonymized interactions, Expert Syst. Appl. 271 (2025) 126564.

[21] J. Wang, S. Zhang, Y. Xu, Z. Zhang, X. Yang, Y. Cheng, Ilamp: improved text extrac-tion from gradients in federated learning using language model priors and sequence beam search, Expert Syst. Appl. 271 (2025) 126592.

[22] C. Shahnazeer, G. Sureshkumar,Prediction of ailments using federated transfer learning and weight penalty-rational tanh-rnn, expert systems with applications.

[23] O. Marfoq, G. Neglia, L. Kameni, R. Vidal,Personalized federated learning through local memorization, in:AAAI Conference on Artificial Intelligence, 2021.

[24] Y. Jiang, J. Konečný, K. Rush,S. Kannan, Improving federated learning personaliza-tion via model agnostic meta learning, arXiv:1909.12488.

[25] Z. Wang, P. Yu, H. Zhang, Privacy-preserving regulation capacity evaluation for hvac systems in heterogeneous buildings based on federated learning and transfer learn-ing, IEEE Trans. Smart Grid 14 (2023) 3535-3549.

[26] H. Zhao, Z. Pan, Y. Wang, Z. Ying, L. Xu, Y. an Tan, Personalized label inference attack in federated transfer learning via contrastive meta learning, in: AAAI Confer-ence on Artificial Intelligence, 2025.

[27] A. Afzali, P. Shamsinejadbabaki,Private and heterogeneous personalized hierar-chical federated learning using conditional generative adversarial networks, expert systems with applications.

[28] Y. Jia, X. Zhang, H. Hu, K.-K.R. Choo, L. Qi, X. Xu, A. Beheshti, W. Dou, Dapperfl: domain adaptive federated learning with model fusion pruning for edge devices, in: Neural Information Processing Systems, 2024.

[29] X. Xing, Q. Zhan, X. Xie, Y. Yang, Q. Wang, G. Liu,Flexible sharpness-aware person-alized federated learning, in:AAAI Conference on Artificial Intelligence, 2025.

[30] T. Yang, S. Cen, Y. Wei, Y. Chen, Y. Chi, Federated natural policy gradient and actor critic methods for multi-task reinforcement learning, in: Neural Information Processing Systems, 2024.

[31] J. Chen, J. Xue, Y. Wang, Z. Liu, L. Huang, Classifier clustering and feature align-ment for federated learning under distributed concept drift, in: Neural Information Processing Systems, 2024.

[32] L. Wang, J. Bian, L. Zhang, C. Chen, J. Xu, Taming cross-domain representation vari-ance in federated prototype learning with heterogeneous data domains, in: Neural Information Processing Systems, 2024.

[33] C.-Y. Hong, Y.-C. Hsu, T.-L. Liu,Contrastive learning for deepfake classification and localization via multi-label ranking, in:2024 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2024, pp. 17627-17637.

[34] T. Chen, S. Kornblith, M. Norouzi, G. Hinton, A simple framework for contrastive learning of visual representations, in: Proceedings of the 37th International Confer-ence on Machine Learning, PMLR, 2020, pp. 1597-1607.

[35] Y. Liu, L. Huang, F. Giunchiglia, X. Feng, R. Guan,Improved graph contrastive learn-ing for short text classification, in:AAAI Conference on Artificial Intelligence, 2024.

[36] Y. Song, F. Wang, CoReFace: sample-guided contrastive regularization for deep face recognition, Pattern Recognit. 152 (2023) 110483.

[37] Y. Yang, W. Hu, H. Lin, H. Hu, Robust cross-domain pseudo-labeling and contrastive learning for unsupervised domain adaptation nir-vis face recognition, IEEE Trans. Image Process. 32 (2023) 5231-5244.

[38] H. Basak,Pseudo-label guided contrastive learning for semi-supervised medical im-age segmentation, in:2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2023, pp. 19786-19797.

[39] X. Wang, Y. Du, S. Yang, J. Zhang, M. Wang, J. Zhang, W. Yang, J. Huang, X. Han, RetCCL: clustering-guided contrastive learning for whole-slide image retrieval, Med. Image Anal. 83 (2023) 102645.

[40] C. Liu, Y. Zhang, H. Wang, W. Chen, F. Wang, Y. Huang, Y. Shen, L. Wang, Efficient token-guided image-text retrieval with consistent multimodal contrastive training, IEEE Trans. Image Process. 32 (2023) 3622-3633.

[41] A.M.H. Tiong, J. Li, G. Lin, B. Li, C. Xiong, S.C.H. Hoi, Improving tail-class rep-resentation with centroid contrastive learning, Pattern Recognit. Lett. 168 (2023) 123-130.

[42] Y. Wang, K. Sun, C. Guo, S. Zhong, H. Liu, Y. Ma, Multiple contrastive experts for long-tailed image classification, Expert Syst. Appl. 255 (2024) 124613.

[43] S. Xuan, S. Zhang,Decoupled contrastive learning for long-tailed recognition, in:AAAI Conference on Artificial Intelligence, 2024.

[44] C. Louizos, M. Reisser, D. Korzhenkov, A mutual information perspective on feder-ated contrastive learning, arXiv:2405.02081 [abs].

[45] T. Zhang, Z. Xue, A. Mahmood, J. Du, Y. Dong, S. Ou, L. Feng, M.-H. Yang, Y. Qi,Gen-erating synthetic data for unsupervised federated learning of cross-modal retrieval, in:AAAI Conference on Artificial Intelligence, 2025.

[46] X. Kong, W. Zhang, H. Wang, M. Hou, X. Chen, X. Yan, S.k. Das,Federated graph anomaly detection via contrastive self-supervised learning, IEEE transactions on neu-ral networks and learning systems PP.

[47] C. Dai, S. Wei, S. Dai, I.S.G. Member, I.G.K.S. Member, I. M.S.H.S. Member, M. Shamim, Federated self-supervised learning based on prototypes clustering con-trastive learning for Internet of vehicles applications, IEEE Internet Things J. 12 (2025) 4692-4700.

[48] A.M. Saxe, J.L. McClelland, S. Ganguli, Exact solutions to the nonlinear dynamics of learning in deep linear neural networks, in: CoRR, 2013.

[49] J. Lezama, Q. Qiu, P. Muse, G. Sapiro, OLE: orthogonal low-rank embedding, a plug and play geometric loss for deep learning, in: 2018 IEEE/CVF Conference on Com-puter Vision and Pattern Recognition, IEEE, Salt Lake City, UT, 2018, pp. 8109-8118.

[50] X. Glorot, Y. Bengio, Understanding the difficulty of training deep feedforward neu-ral networks, in: Proceedings of the Thirteenth International Conference on Artifi-cial Intelligence and Statistics, JMLR Workshop and Conference Proceedings, 2010, pp. 249-256.

[51] K. He, X. Zhang, S. Ren, J. Sun, Delving deep into rectifiers: surpassing human-level performance on imagenet classification,in: Proceedings of the IEEE International Conference on Computer Vision, 2015, pp. 1026-1034.

[52] K. He, X. Zhang, S. Ren, J. Sun, Deep residual learning for image recognition, in: 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), IEEE, Las Vegas, NV, USA, 2016, pp. 770-778.

[53] H. Wang, M. Yurochkin, Y. Sun, D. Papailiopoulos, Y. Khazaeni, Federated learning with matched averaging, arXiv:2002.06440.

[54] M. Yurochkin, M. Agarwal, S. Ghosh, K. Greenewald, N. Hoang, Y. Khazaeni, Bayesian nonparametric federated learning of neural networks, in: Proceedings of the 36th International Conference on Machine Learning, PMLR, 2019, pp. 7252-7261.

[55] Z. Qin, L. Yao, D. Chen, Y. Li, B. Ding, M. Cheng, Revisiting personalized feder-ated learning: robustness against backdoor attacks,in: Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining., 2023.