Abstract
Federated Learning (FL) is a rising approach towards collaborative and privacy-preserving machine learning where large-scale medical datasets remain localized to each client. However, the issue of data heterogeneity among clients often compels local models to diverge, leading to suboptimal global models. To mitigate the impact of data heterogeneity on FL performance, we start with analyzing how FL training influence FL performance by decomposing the global loss into three terms: local loss, distribution shift loss and aggregation loss. Remarkably, our loss decomposition reveals that existing local training-based FL methods attempt to reduce the distribution shift loss, while the global aggregation-based FL methods propose better aggregation strategies to reduce the aggregation loss. Nevertheless, a comprehensive joint effort to minimize all three terms is currently limited in the literature, leading to subpar performance when dealing with data heterogeneity challenges. To fill this gap, we propose a novel FL method based on global loss decomposition, called FedLD, to jointly reduce these three loss terms. Our FedLD involves a margin control regularization in local training to reduce the distribution shift loss, and a principal gradient-based server aggregation strategy to reduce the aggregation loss. Notably, under different levels of data heterogeneity, our strategies achieve better and more robust performance on retinal and chest X-ray classification compared to other FL algorithms. Our code is available at https://github.com/Zeng-Shuang/FedLD.
S. Zeng and P. Guo—Equal contributors.
This is a preview of subscription content, log in via an institution to check access.
Access this chapter
Tax calculation will be finalised at checkout
Purchases are for personal use only
Similar content being viewed by others
References
An, X., Shen, L., Hu, H., Luo, Y.: Federated learning with manifold regularization and normalized update reaggregation. Adv. Neural Inf. Process. Syst. 36 (2024)
Charles, Z., Garrett, Z., Huo, Z., Shmulyian, S., Smith, V.: On large-cohort training for federated learning. Adv. Neural. Inf. Process. Syst. 34, 20461–20475 (2021)
Dayan, I., et al.: Federated learning for predicting clinical outcomes in patients with covid-19. Nat. Med. 27(10), 1735–1743 (2021)
Geirhos, R., et al.: Shortcut learning in deep neural networks. Nat. Mach. Intell. 2(11), 665–673 (2020)
Guo, Y., Guo, K., Cao, X., Wu, T., Chang, Y.: Out-of-distribution generalization of federated learning via implicit invariant relationships (2023)
Hamilton, W.L., Ying, R., Leskovec, J.: Representation learning on graphs: methods and applications. arXiv preprint arXiv:1709.05584 (2017)
He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 770–778 (2016)
Hsu, T.M.H., Qi, H., Brown, M.: Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.06335 (2019)
Kaissis, G.A., Makowski, M.R., Rückert, D., Braren, R.F.: Secure, privacy-preserving and federated machine learning in medical imaging. Nat. Mach. Intell. 2(6), 305–311 (2020)
Kalra, S., Wen, J., Cresswell, J.C., Volkovs, M., Tizhoosh, H.: Decentralized federated learning through proxy model sharing. Nat. Commun. 14(1), 2899 (2023)
Karimireddy, S.P., Kale, S., Mohri, M., Reddi, S., Stich, S., Suresh, A.T.: Scaffold: stochastic controlled averaging for federated learning. In: International Conference on Machine Learning, pp. 5132–5143. PMLR (2020)
Koh, P.W., et al.: Wilds: A benchmark of in-the-wild distribution shifts. In: International Conference on Machine Learning, pp. 5637–5664. PMLR (2021)
Li, T., Sahu, A.K., Zaheer, M., Sanjabi, M., Talwalkar, A., Smith, V.: Federated optimization in heterogeneous networks. Proc. Mach. Learn. Syst. 2, 429–450 (2020)
Li, X., Gu, Y., Dvornek, N., Staib, L.H., Ventola, P., Duncan, J.S.: Multi-site fmri analysis using privacy-preserving federated learning and domain adaptation: abide results. Med. Image Anal. 65, 101765 (2020)
Li, X., Jiang, M., Zhang, X., Kamp, M., Dou, Q.: Fedbn: federated learning on non-iid features via local batch normalization. arXiv preprint arXiv:2102.07623 (2021)
Liu, C., Lou, C., Wang, R., Xi, A.Y., Shen, L., Yan, J.: Deep neural network fusion via graph matching with applications to model ensemble and federated learning. In: International Conference on Machine Learning, pp. 13857–13869. PMLR (2022)
Ma, X., Zhang, J., Guo, S., Xu, W.: Layer-wised model aggregation for personalized federated learning. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10092–10101 (2022)
McMahan, B., Moore, E., Ramage, D., Hampson, S., Arcas, B.A.: Communication-efficient learning of deep networks from decentralized data. In: Artificial Intelligence and Statistics, pp. 1273–1282. PMLR (2017)
Nguyen, A.T., Torr, P., Lim, S.N.: Fedsr: a simple and effective domain generalization method for federated learning. Adv. Neural. Inf. Process. Syst. 35, 38831–38843 (2022)
Pati, S., et al.: Federated learning enables big data for rare cancer boundary detection. Nat. Commun. 13(1), 7346 (2022)
Puli, A.M., Zhang, L., Wald, Y., Ranganath, R.: Don’t blame dataset shift! shortcut learning due to gradients and cross entropy. Adv. Neural Inf. Process. Syst. 36 (2024)
Qu, L., Balachandar, N., Rubin, D.L.: An experimental study of data heterogeneity in federated learning methods for medical imaging. arXiv preprint arXiv:2107.08371 (2021)
Qu, L., et al.: Rethinking architecture design for tackling data heterogeneity in federated learning. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10061–10071 (2022)
Ruder, S.: An overview of gradient descent optimization algorithms. arXiv preprint arXiv:1609.04747 (2016)
Uddin, M.P., Xiang, Y., Yearwood, J., Gao, L.: Robust federated averaging via outlier pruning. IEEE Signal Process. Lett. 29, 409–413 (2021)
Wang, H., Yurochkin, M., Sun, Y., Papailiopoulos, D., Khazaeni, Y.: Federated learning with matched averaging. arXiv preprint arXiv:2002.06440 (2020)
Wang, Z., Grigsby, J., Qi, Y.: Pgrad: learning principal gradients for domain generalization. arXiv preprint arXiv:2305.01134 (2023)
Xu, J., Tong, X., Huang, S.L.: Personalized federated learning with feature alignment and classifier collaboration. arXiv preprint arXiv:2306.11867 (2023)
Xu, J., Wang, S., Wang, L., Yao, A.C.C.: Fedcm: federated learning with client-level momentum. arXiv preprint arXiv:2106.10874 (2021)
Yan, R., et al.: Label-efficient self-supervised federated learning for tackling data heterogeneity in medical imaging. IEEE Trans. Med. Imaging 42, 1932–1943 (2023)
Yang, C., et al.: Characterizing impacts of heterogeneity in federated learning upon large-scale smartphone data. In: Proceedings of the Web Conference 2021, pp. 935–946 (2021)
Zhang, J., et al.: Flhetbench: benchmarking device and state heterogeneity in federated learning. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12098–12108 (2024)
Zhang, M., Qu, L., Singh, P., Kalpathy-Cramer, J., Rubin, D.L.: Splitavg: a heterogeneity-aware federated deep learning method for medical imaging. IEEE J. Biomed. Health Inf. 26(9), 4635–4644 (2022)
Zhang, X., Sun, W., Chen, Y.: Tackling the non-iid issue in heterogeneous federated learning by gradient harmonization. arXiv preprint arXiv:2309.06692 (2023)
Ziller, A., Usynin, D., Braren, R., Makowski, M., Rueckert, D., Kaissis, G.: Medical imaging deep learning with differential privacy. Sci. Rep. 11(1), 13524 (2021)
Acknowledgments
This work was supported by National Natural Science Foundation of China (62306253) Early career fund (27204623), Guangdong Natural Science Fund-General Programme (2024A1515010233), and UCSC hellman fellowship.
Author information
Authors and Affiliations
Corresponding author
Editor information
Editors and Affiliations
Ethics declarations
Disclosure of Interests
The authors have no competing interests to declare that are relevant to the content of this article.
1 Electronic supplementary material
Below is the link to the electronic supplementary material.
Rights and permissions
Copyright information
© 2024 The Author(s), under exclusive license to Springer Nature Switzerland AG
About this paper
Cite this paper
Zeng, S., Guo, P., Wang, S., Wang, J., Zhou, Y., Qu, L. (2024). Tackling Data Heterogeneity in Federated Learning via Loss Decomposition. In: Linguraru, M.G., et al. Medical Image Computing and Computer Assisted Intervention – MICCAI 2024. MICCAI 2024. Lecture Notes in Computer Science, vol 15010. Springer, Cham. https://doi.org/10.1007/978-3-031-72117-5_66
Download citation
DOI: https://doi.org/10.1007/978-3-031-72117-5_66
Published:
Publisher Name: Springer, Cham
Print ISBN: 978-3-031-72116-8
Online ISBN: 978-3-031-72117-5
eBook Packages: Computer ScienceComputer Science (R0)
