Abstract
The label distribution skew has been shown to be a significant obstacle that limits the model performance in federated learning (FL). This challenge could be more serious when the participating clients are in unstable network circumstances and drop out frequently. Previous works have demonstrated that the classifier head is particularly sensitive to the label skew. Therefore, maintaining a balanced classifier head is of significant importance for building a good and unbiased global model. To this end, we propose a simple yet effective framework by introducing a calibrated softmax function with smoothed prior for computing the cross-entropy loss, and a prototype-based feature augmentation scheme to re-balance the local training, which provide a new perspective on tackling the label distribution skew in FL and are lightweight for edge devices and can facilitate the global model aggregation. With extensive experiments on two benchmark classification tasks of Fashion-MNIST and CIFAR-10, our numerical results demonstrate that our proposed method can consistently outperform the baselines, 2 8% of accuracy over FedAvg in the presence of severe label skew and client dropout.
Similar content being viewed by others
Explore related subjects
Discover the latest articles, news and stories from top researchers in related subjects.Data Availability
The used datasets can be publicly accessed from the Internet with details provided in reference.
Materials Availability
Not applicable
Code Availability
Code is available at https://github.com/JianXu95/ReBaFL
References
Hinton G, Deng L, Yu D, Dahl GE, Mohamed A, Jaitly N, Senior A, Vanhoucke V, Nguyen P, Sainath TN, Kingsbury B (2012) Deep neural networks for acoustic modeling in speech recognition. IEEE Signal Process Magaz 29(6):82–97
Hayyolalam V, Aloqaily M, Özkasap Ö, Guizani M (2022) Edge-assisted solutions for iot-based connected healthcare systems: A literature review. IEEE Internet Things J. 9(12):9419–9443
Ullah Z, Al-Turjman F, Mostarda L, Gagliardi R (2020) Applications of artificial intelligence and machine learning in smart cities. Comput Commun 154:313–323
Sharma S (2019) Data privacy and gdpr handbook
McMahan B, Moore E, Ramage D, Hampson S, Arcas BA (2017) Communication-efficient learning of deep networks from decentralized data. In: International Conference on Artificial Intelligence and Statistics, AISTATS
Yang Q, Liu Y, Chen T, Tong Y (2019) Federated machine learning: Concept and applications. ACM Trans Intell Syst Technol 10(2)
Li T, Sahu AK, Talwalkar A, Smith V (2020) Federated learning: Challenges, methods, and future directions. IEEE Signal Process Mag 37(3):50–60
Li T, Sahu AK, Zaheer M, Sanjabi M, Talwalkar A, Smith V (2020) Federated optimization in heterogeneous networks. In: Proceedings of Machine Learning and Systems 2020, MLSys
Wahab OA, Mourad A, Otrok H, Taleb T (2021) Federated machine learning: Survey, multi-level classification, desirable criteria and future directions in communication and networking systems. IEEE Commun Surv Tutorials 23(2):1342–1397
Kairouz P, McMahan HB, Avent B, Bellet A, Bennis M (2021) al: Advances and open problems in federated learning. Found Trends Mach Learn 14(1–2):1–210
Wang S, Tuor T, Salonidis T, Leung KK, Makaya C, He T, Chan K (2019) Adaptive federated learning in resource constrained edge computing systems. IEEE J Sel Areas Commun 37(6):1205–1221
Karimireddy SP, Kale S, Mohri M, Reddi SJ, Stich SU, Suresh AT (2020) SCAFFOLD: stochastic controlled averaging for federated learning. In: Proceedings of the 37th International Conference on Machine Learning, ICML
Acar DAE, Zhao Y, Navarro RM, Mattina M, Whatmough PN, Saligrama V (2021) Federated learning based on dynamic regularization. In: 9th International Conference on Learning Representations, ICLR
Li Q, He B, Song D (2021) Model-contrastive federated learning. In: IEEE Conference on Computer Vision and Pattern Recognition, CVPR
Li Q, Diao Y, Chen Q, He B (2022) Federated learning on non-iid data silos: An experimental study. In: 2022 IEEE 38th International Conference on Data Engineering (ICDE), pp 965–978
Ren J, Yu C, Sheng S, Ma X, Zhao H, Yi S, Li H (2020) Balanced meta-softmax for long-tailed visual recognition. In: Annual Conference on Neural Information Processing Systems, NeurIPS
Luo M, Chen F, Hu D, Zhang Y, Liang J, Feng J (2021) No fear of heterogeneity: Classifier calibration for federated learning with non-iid data. Adv Neural Inf Process Syst 34:5972–5984
Zhang X, Hong M, Dhople S, Yin W, Liu Y (2021) Fedpd: A federated learning framework with adaptivity to non-iid data. IEEE Trans Signal Process 69:6055–6070
Ozfatura E, Ozfatura K, Gündüz D (2021) Fedadc: Accelerated federated learning with drift control. In: IEEE International Symposium on Information Theory, ISIT
Hsu TH, Qi H, Brown M (2020) Federated visual classification with real-world data distribution. In: Computer Vision - ECCV 2020 - 16th European Conference
Wang L, Xu S, Wang X, Zhu Q (2021) Addressing class imbalance in federated learning. In: Thirty-Fifth AAAI Conference on Artificial Intelligence, AAAI
Yoon T, Shin S, Hwang SJ, Yang E (2021) Fedmix: Approximation of mixup under mean augmented federated learning. In: 9th International Conference on Learning Representations, ICLR
Wang H, Kaplan Z, Niu D, Li B (2020) Optimizing federated learning on non-iid data with reinforcement learning. In: IEEE Conference on Computer Communications, INFOCOM
Tang M, Ning X, Wang Y, Sun J, Wang Y, Li H, Chen Y (2022) Fedcor: Correlation-based active client selection strategy for heterogeneous federated learning. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp 10102–10111
Fraboni Y, Vidal R, Kameni L, Lorenzi M (2021) Clustered sampling: Low-variance and improved representativity for clients selection in federated learning. In: Proceedings of the 38th International Conference on Machine Learning, ICML
Chen H, Chao W (2021) Fedbe: Making bayesian model ensemble applicable to federated learning. In: International Conference on Learning Representations, ICLR
Lin T, Kong L, Stich SU, Jaggi M (2020) Ensemble distillation for robust model fusion in federated learning. In: Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems, NeurIPS
Zhu Z, Hong J, Zhou J (2021) Data-free knowledge distillation for heterogeneous federated learning. In: Proceedings of the 38th International Conference on Machine Learning, ICML
Zhang L, Shen L, Ding L, Tao D, Duan L (2022) Fine-tuning global model via data-free knowledge distillation for non-iid federated learning. In: IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR, pp 10164–10173
Kulkarni V, Kulkarni M, Pant A (2020) Survey of personalization techniques for federated learning. In: Fourth World Conference on Smart Trends in Systems, Security and Sustainability (WorldS4)
Dinh CT, Tran NH, Nguyen TD (2020) Personalized federated learning with moreau envelopes. In: Conference on Neural Information Processing Systems, NeurIPS
Fallah A, Mokhtari A, Ozdaglar AE (2020) Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. In: Advances in Neural Information Processing Systems, NeurIPS
Collins L, Hassani H, Mokhtari A, Shakkottai S (2021) Exploiting shared representations for personalized federated learning. In: Proceedings of the 38th International Conference on Machine Learning, ICML
Sattler F, Müller K, Samek W (2021) Clustered federated learning: Model-agnostic distributed multitask optimization under privacy constraints. IEEE Trans Neural Networks Learn Syst 32(8):3710–3722
Ghosh A, Chung J, Yin D, Ramchandran K (2022) An efficient framework for clustered federated learning. IEEE Trans Inf Theory 68(12):8076–8091
Chen Y, Ning Y, Slawski M, Rangwala H (2020) Asynchronous online federated learning for edge devices with non-iid data. In: IEEE International Conference on Big Data (BigData)
Wu W, He L, Lin W, Mao R, Maple C, Jarvis SA (2021) SAFA: A semi-asynchronous protocol for fast federated learning with low overhead. IEEE Trans Comput 70(5):655–668
Gu X, Huang K, Zhang J, Huang L (2021) Fast federated learning in the presence of arbitrary device unavailability. In: Annual Conference on Neural Information Processing Systems, NeurIPS
Zhu L, Lin H, Lu Y, Lin Y, Han S (2021) Delayed gradient averaging: Tolerate the communication latency for federated learning. In: Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems, NeurIPS
Reisizadeh A, Tziotis I, Hassani H, Mokhtari A, Pedarsani R (2022) Straggler-resilient federated learning: Leveraging the interplay between statistical accuracy and system heterogeneity. IEEE J Sel Areas Inf Theory 3(2):197–205
Luo B, Xiao W, Wang S, Huang J, Tassiulas L (2022) Tackling system and statistical heterogeneity for federated learning with adaptive client sampling. In: IEEE INFOCOM 2022 - IEEE Conference on Computer Communications
Ma Z, Xu Y, Xu H, Meng Z, Huang L, Xue Y (2021) Adaptive batch size for federated learning in resource-constrained edge computing. IEEE Trans Mobile Comput 22(1):37–53
Li L, Duan M, Liu D, Zhang Y, Ren A, Chen X, Tan Y, Wang C (2021) Fedsae: A novel self-adaptive federated learning framework in heterogeneous systems. In: International Joint Conference on Neural Networks, IJCNN
Jhunjhunwala D, Sharma P, Nagarkatti A, Joshi G (2022) Fedvarp: Tackling the variance due to partial client participation in federated learning. In: Proceedings of the Thirty-Eighth Conference on Uncertainty in Artificial Intelligence, UAI
Wang H, Xu J (2024) Friends to help: Saving federated learning from client dropout. In: ICASSP 2024-2024 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp 8896–8900 . IEEE
Tian J, Liu Y, Glaser N, Hsu Y, Kira Z (2020) Posterior re-calibration for imbalanced datasets. In: Annual Conference on Neural Information Processing System, NeurIPS
Li X-C, Zhan D-C (2021) FedRS: Federated learning with restricted softmax for label distribution non-iid data. In: Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, pp 995–1005
Yin X, Yu X, Sohn K, Liu X, Chandraker M (2019) Feature transfer learning for face recognition with under-represented data. In: IEEE Conference on Computer Vision and Pattern Recognition, CVPR
Xiao H, Rasul K, Vollgraf R (2017) Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. CoRR arXiv:1708.07747
Krizhevsky A, Hinton G (2009) Learning multiple layers of features from tiny images. University of Toronto
Oh J, Kim S, Yun S (2022) Fedbabu: Toward enhanced representation for federated image classification. In: The Tenth International Conference on Learning Representations, ICLR
Chen H, Chao W (2022) On bridging generic and personalized federated learning for image classification. In: The Tenth International Conference on Learning Representations, ICLR
Zhang J, Li Z, Li B, Xu J, Wu S, Ding S, Wu C (2022) Federated learning with label distribution skew via logits calibration. In: International Conference on Machine Learning, ICML
Smith SL, Elsen E, De S (2020) On the generalization benefit of noise in stochastic gradient descent. In: Proceedings of the 37th International Conference on Machine Learning, ICML
Dosovitskiy A, Brox T (2016) Inverting visual representations with convolutional networks. In: IEEE Conference on Computer Vision and Pattern Recognition, CVPR
Kingma DP, Welling M (2014) Auto-encoding variational bayes. In: 2nd International Conference on Learning Representations, ICLR
Ren H, Deng J, Xie X (2022) GRNN: generative regression neural network - A data leakage attack for federated learning. ACM Trans Intell Syst Technol 13(4):65–16524
Liu R, Cao Y, Chen H, Guo R, Yoshikawa M (2021) FLAME: differentially private federated learning in the shuffle model. In: Thirty-Fifth AAAI Conference on Artificial Intelligence, AAAI
Girgis AM, Data D, Diggavi SN, Kairouz P, Suresh AT (2021) Shuffled model of differential privacy in federated learning. In: The 24th International Conference on Artificial Intelligence and Statistics, AISTATS
Acknowledgements
The research of Shao-Lun Huang is supported in part by National Key R&D Program of China under Grant 2021YFA0715202, Shenzhen Key Laboratory of Ubiquitous Data Enabling under Grant ZDSYS20220527171406015 and the Shenzhen Science and Technology Program under Grant KQTD20170810150821146.
Funding
The research of Shao-Lun Huang is supported in part by National Key R&D Program of China under Grant 2021YFA0715202?Shenzhen Key Laboratory of Ubiquitous Data Enabling under Grant ZDSYS20220527171406015 and the Shenzhen Science and Technology Program under Grant KQTD20170810150821146.
Author information
Authors and Affiliations
Contributions
All authors contributed to the study conception and design. Experiment execution and paper writing were performed by Jian Xu and Meilin Yang. The first draft of the manuscript was written by Jian Xu and all authors commented on previous versions of the manuscript. All authors read and approved the final manuscript. This paper has a pre-printed version at https://export.arxiv.org/abs/2303.06314.
Corresponding author
Ethics declarations
Conflict of Interest/Competing Interests
the authors have no competing interests to declare that are relevant to the content of this article.
Ethics Approval and Consent to Participate
Not applicable
Consent for Publication
Not applicable
Additional information
Publisher's Note
Springer Nature remains neutral with regard to jurisdictional claims in published maps and institutional affiliations.
Rights and permissions
Springer Nature or its licensor (e.g. a society or other partner) holds exclusive rights to this article under a publishing agreement with the author(s) or other rightsholder(s); author self-archiving of the accepted manuscript version of this article is solely governed by the terms of such publishing agreement and applicable law.
About this article
Cite this article
Xu, J., Yang, M., Ding, W. et al. Stabilizing and improving federated learning with highly non-iid data and client dropout. Appl Intell 55, 216 (2025). https://doi.org/10.1007/s10489-024-05956-3
Accepted:
Published:
DOI: https://doi.org/10.1007/s10489-024-05956-3