FedAIMS: adaptive intermediate supervision for personalized federated learning

Abstract

Personalized Federated Learning (PFL) enables the training of customized deep models on decentralized, heterogeneous data while preserving privacy. However, existing PFL methods primarily optimize the final layer, overlooking intermediate layers, which degrades backbone training, especially in non-IID settings. In this work, we propose FedAIMS (Federated Adaptive Intermediate Supervision), a novel PFL framework that incorporates intermediate supervision to enhance model training. FedAIMS adopts prototype-based feature alignment to provide effective intermediate supervision and adaptive supervision sampling to reduce computational overhead for resource-limited clients. Experiments on diverse datasets show that FedAIMS outperforms state-of-the-art PFL baselines by up to 36.76% in accuracy.

Graphical abstract

Keywords

personalized federated learning / intermediate supervision / data heterogeneity

References

[1] Tan A Z, Yu H, Cui L, Yang Q . Towards personalized federated learning. IEEE Transactions on Neural Networks and Learning Systems, 2023, 34( 12): 9587–9603

[2] McMahan B, Moore E, Ramage D, Hampson S, Arcas B A Y. Communication-efficient learning of deep networks from decentralized data. In: Proceedings of the 20th International Conference on Artificial Intelligence and Statistics. 2017, 1273–1282

[3] Leroy D, Coucke A, Lavril T, Gisselbrecht T, Dureau J. Federated learning for keyword spotting. In: Proceedings of 2019 IEEE International Conference on Acoustics, Speech and Signal Processing. 2019, 6341–6345

[4] Ouyang X, Xie Z, Zhou J, Huang J, Xing G. ClusterFL: a similarity-aware federated learning system for human activity recognition. In: Proceedings of the 19th Annual International Conference on Mobile Systems, Applications, and Services. 2021, 54–66

[5] Ouyang X, Shuai X, Li Y, Pan L, Zhang X, Fu H, Cheng S, Wang X, Cao S, Xin J, Mok H, Yan Z, Yu D S F, Kwok T, Xing G. ADMarker: a multi-modal federated learning system for monitoring digital biomarkers of Alzheimer’s disease. In: Proceedings of the 30th Annual International Conference on Mobile Computing and Networking. 2024, 404–419

[6] Sun N, Wang W, Tong Y, Liu K . Blockchain based federated learning for intrusion detection for internet of things. Frontiers of Computer Science, 2024, 18( 5): 185328

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

[8] Liang P P, Liu T, Liu Z, Allen N B, Auerbach R P, Brent D, Salakhutdinov R, Morency L P. Think locally, act globally: Federated learning with local and global representations. 2020, arXiv preprint arXiv: 2001.01523

[9] Collins L, Hassani H, Mokhtari A, Shakkottai S. Exploiting shared representations for personalized federated learning. In: Proceedings of the 38th International Conference on Machine Learning. 2021, 2089–2099

[10] Chen H Y, Chao W L. On bridging generic and personalized federated learning for image classification. In: Proceedings of the 10th International Conference on Learning Representations. 2022

[11] Oh J, Kim S, Yun S Y. FedBABU: toward enhanced representation for federated image classification. In: Proceedings of the 10th International Conference on Learning Representations. 2022

[12] McLaughlin C, Su L. Personalized federated learning via feature distribution adaptation. In: Proceedings of the 38th Conference on Neural Information Processing Systems. 2024

[13] Luo M, Chen F, Hu D, Zhang Y, Liang J, Feng J. No fear of heterogeneity: classifier calibration for federated learning with non-IID data. In: Proceedings of the 35th Conference on Neural Information Processing Systems. 2021, 5972–5984

[14] Li Z, Shang X, He R, Lin T, Wu C. No fear of classifier biases: neural collapse inspired federated learning with synthetic and fixed classifier. In: Proceedings of 2023 IEEE/CVF International Conference on Computer Vision. 2023, 5296–5306

[15] Lee C Y, Xie S, Gallagher P, Zhang Z, Tu Z. Deeply-supervised nets. In: Proceedings of the 18th International Conference on Artificial Intelligence and Statistics. 2015, 562–570

[16] Szegedy C, Liu W, Jia Y, Sermanet P, Reed S, Anguelov D, Erhan D, Vanhoucke V, Rabinovich A. Going deeper with convolutions. In: Proceedings of 2015 IEEE Conference on Computer Vision and Pattern Recognition. 2015, 1–9

[17] Li T, Hu S, Beirami A, Smith V. Ditto: fair and robust federated learning through personalization. In: Proceedings of the 38th International Conference on Machine Learning. 2021, 6357–6368

[18] Zhang J, Hua Y, Wang H, Song T, Xue Z, Ma R, Guan H. FedALA: adaptive local aggregation for personalized federated learning. In: Proceedings of the 37th AAAI Conference on Artificial Intelligence. 2023, 11237–11244

[19] Zhang C J, Tong Y, Chen L . Where to: crowd-aided path selection. Proceedings of the VLDB Endowment, 2014, 7( 14): 2005–2016

[20] Cao C C, Tong Y, Chen L, Jagadish H V. WiseMarket: a new paradigm for managing wisdom of online social users. In: Proceedings of the 19th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining. 2013, 455–463

[21] Li T, Sahu A K, Zaheer M, Sanjabi M, Talwalkar A, Smith V. Federated optimization in heterogeneous networks. In: Proceedings of the 3rd Conference on Machine Learning and Systems. 2020, 429–450

[22] Li Q, He B, Song D. Model-contrastive federated learning. In: Proceedings of 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021, 10708–10717

[23] Wang Y, Tong Y, Shi D, Xu K. An efficient approach for cross-silo federated learning to rank. In: Proceedings of the 37th IEEE International Conference on Data Engineering. 2021, 1128–1139

[24] Wang Y, Tong Y, Zhou Z, Ren Z, Xu Y, Wu G, Lv W. FED-LTD: towards cross-platform ride hailing via federated learning to dispatch. In: Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2022, 4079–4089

[25] Tong Y, Zeng Y, Zhou Z, Liu B, Shi Y, Li S, Xu K, Lv W . Federated computing: query, learning, and beyond. IEEE Data Engineering Bulletin, 2023, 46( 1): 9–26

[26] Wei S, Tong Y, Zhou Z, Xu Y, Gao J, Wei T, He T, Lv W . Federated reasoning LLMs: a survey. Frontiers of Computer Science, 2025, 19( 12): 1912613

[27] Dinh C T, Tran N H, Nguyen T D. Personalized federated learning with Moreau envelopes. In: Proceedings of the 34th Conference on Neural Information Processing Systems. 2020, 21394–21405

[28] Guo W, Zhuang F, Zhang X, Tong Y, Dong J . A comprehensive survey of federated transfer learning: challenges, methods and applications. Frontiers of Computer Science, 2024, 18( 6): 186356

[29] Zhang J, Guo S, Ma X, Wang H, Xu W, Wu F. Parameterized knowledge transfer for personalized federated learning. In: Proceedings of the 35th Conference on Neural Information Processing Systems. 2021, 10092–10104

[30] Sattler F, Müller K R, Samek W . Clustered federated learning: Model-agnostic distributed multitask optimization under privacy constraints. IEEE Transactions on Neural Networks and Learning Systems, 2021, 32( 8): 3710–3722

[31] Liu B, Ma Y, Zhou Z, Shi Y, Li S, Tong Y. CASA: clustered federated learning with asynchronous clients. In: Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2024, 1851–1862

[32] Huang Y, Chu L, Zhou Z, Wang L, Liu J, Pei J, Zhang Y. Personalized cross-silo federated learning on non-IID data. In: Proceedings of the 35th AAAI Conference on Artificial Intelligence. 2021, 7865–7873

[33] Ye R, Ni Z, Wu F, Chen S, Wang Y. Personalized federated learning with inferred collaboration graphs. In: Proceedings of the 40th International Conference on Machine Learning. 2023, 1661

[34] Zhang W, Zhou Z, Wang Y, Tong Y. DM-PFL: Hitchhiking generic federated learning for efficient shift-robust personalization. In: Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2023, 3396–3408

[35] Chen D, Yao L, Gao D, Ding B, Li Y. Efficient personalized federated learning via sparse model-adaptation. In: Proceedings of the 40th International Conference on Machine Learning. 2023, 5234–5256

[36] Sun Y, Wang X, Tang X. Deeply learned face representations are sparse, selective, and robust. In: Proceedings of 2015 IEEE Conference on Computer Vision and Pattern Recognition. 2015, 2892–2900

[37] Wu X, Finnegan D, O’Neill E, Yang Y L. HandMap: robust hand pose estimation via intermediate dense guidance map supervision. In: Proceedings of the 15th European Conference on Computer Vision. 2018, 246–262

[38] Wang C, Wu Y, Chen S, Liu S, Li J, Qian Y, Yang Z. Improving self-supervised learning for speech recognition with intermediate layer supervision. In: Proceedings of 2022 IEEE International Conference on Acoustics, Speech and Signal Processing. 2022, 7092–7096

[39] Gou J, Yu B, Maybank S J, Tao D . Knowledge distillation: a survey. International Journal of Computer Vision, 2021, 129( 6): 1789–1819

[40] Romero A, Ballas N, Kahou S E, Chassang A, Gatta C, Bengio Y. FitNets: hints for thin deep nets. In: Proceedings of the 3rd International Conference on Learning Representations. 2015

[41] Zagoruyko S, Komodakis N. Paying more attention to attention: improving the performance of convolutional neural networks via attention transfer. In: Proceedings of the 5th International Conference on Learning Representations. 2017

[42] Passalis N, Tefas A. Learning deep representations with probabilistic knowledge transfer. In: Proceedings of the 15th European Conference on Computer Vision. 2018, 283–299

[43] Kim J, Park S, Kwak N. Paraphrasing complex network: network compression via factor transfer. In: Proceedings of the 32nd International Conference on Neural Information Processing Systems. 2018, 2765–2774

[44] Tun Y L, Thwal C M, Park Y M, Park S B, Hong C S. Federated learning with intermediate representation regularization. In: Proceedings of 2023 IEEE International Conference on Big Data and Smart Computing. 2023, 56–63

[45] Li R, Wang X, Huang G, Yang W, Zhang K, Gu X, Tran S N, Garg S, Alty J, Bai Q. A comprehensive review on deep supervision: theories and applications. 2022, arXiv preprint arXiv: 2207.02376

[46] Tan Y, Long G, Liu L, Zhou T, Lu Q, Jiang J, Zhang C. FedProto: federated prototype learning across heterogeneous clients. In: Proceedings of the 36th AAAI Conference on Artificial Intelligence. 2022, 8432–8440

[47] Xu J, Tong X, Huang S L. Personalized federated learning with feature alignment and classifier collaboration. In: Proceedings of the 11th International Conference on Learning Representations. 2023

[48] Yang J, Duan Y, Qiao T, Zhou H, Wang J, Zhao W . Prototyping federated learning on edge computing systems. Frontiers of Computer Science, 2020, 14( 6): 146318

[49] Qu L, Li S, Zhou Z, Liu B, Xu Y, Tong Y. DarkDistill: difficulty-aligned federated early-exit network training on heterogeneous devices. In: Proceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining V.2. 2025, 2374–2385

[50] Li X, Huang K, Yang W, Wang S, Zhang Z. On the convergence of FedAvg on non-IID data. In: Proceedings of the 8th International Conference on Learning Representations. 2020

[51] Zeiler M D, Fergus R. Visualizing and understanding convolutional networks. In: Proceedings of the 13th European Conference on Computer Vision. 2014, 818–833

[52] Yi L, Yu H, Ren C, Wang G, Liu X, Li X. Federated model heterogeneous matryoshka representation learning. In: Proceedings of the 38th Conference on Neural Information Processing Systems. 2024

[53] Cohen G, Afshar S, Tapson J, Van Schaik A. EMNIST: Extending MNIST to handwritten letters. In: Proceedings of 2017 International Joint Conference on Neural Networks. 2017, 2921–2926

[54] Krizhevsky A. Learning multiple layers of features from tiny images. 2009

[55] Chrabaszcz P, Loshchilov I, Hutter F. A downsampled variant of ImageNet as an alternative to the CIFAR datasets. 2017, arXiv preprint arXiv: 1707.08819

[56] Hsu T M H, Qi H, Brown M. Measuring the effects of non-identical data distribution for federated visual classification. 2019, arXiv preprint arXiv: 1909.06335

[57] He K, Zhang X, Ren S, Sun J. Deep residual learning for image recognition. In: Proceedings of 2016 IEEE Conference on Computer Vision and Pattern Recognition. 2016, 770–778

RIGHTS & PERMISSIONS

The Author(s) 2025. This article is published with open access at link.springer.com and journal.hep.com.cn