Skip to main content

FedTSA: A Cluster-Based Two-Stage Aggregation Method for Model-Heterogeneous Federated Learning

  • Conference paper
  • First Online:
Computer Vision – ECCV 2024 (ECCV 2024)

Part of the book series: Lecture Notes in Computer Science ((LNCS,volume 15141))

Included in the following conference series:

Abstract

Despite extensive research into data heterogeneity in federated learning (FL), system heterogeneity remains a significant yet often overlooked challenge. Traditional FL approaches typically assume homogeneous hardware resources across FL clients, implying that clients can train a global model within a comparable time frame. However, in practical FL systems, clients often have heterogeneous resources, which impacts their training capacity. This discrepancy underscores the importance of exploring model-heterogeneous FL, a paradigm allowing clients to train different models based on their resource capabilities. To address this challenge, we introduce FedTSA, a cluster-based two-stage aggregation method tailored for system heterogeneity in FL. FedTSA begins by clustering clients based on their capabilities, then performs a two-stage aggregation: conventional weight averaging for homogeneous models in Stage 1, and deep mutual learning with a diffusion model for aggregating heterogeneous models in Stage 2. Extensive experiments demonstrate that FedTSA not only outperforms the baselines but also explores various factors influencing model performance, validating FedTSA as a promising approach for model-heterogeneous FL.

This is a preview of subscription content, log in via an institution to check access.

Access this chapter

Subscribe and save

Springer+ Basic
$34.99 /Month
  • Get 10 units per month
  • Download Article/Chapter or eBook
  • 1 Unit = 1 Article or 1 Chapter
  • Cancel anytime
Subscribe now

Buy Now

Chapter
USD 29.95
Price excludes VAT (USA)
  • Available as PDF
  • Read on any device
  • Instant download
  • Own it forever

Tax calculation will be finalised at checkout

Purchases are for personal use only

Institutional subscriptions

Similar content being viewed by others

References

  1. Alam, S., Liu, L., Yan, M., Zhang, M.: FedRolex: model-heterogeneous federated learning with rolling sub-model extraction. In: Advances in Neural Information Processing Systems (2022). https://openreview.net/forum?id=OtxyysUdBE

  2. Blalock, D., Gonzalez Ortiz, J.J., Frankle, J., Guttag, J.: What is the state of neural network pruning? In: Dhillon, I., Papailiopoulos, D., Sze, V. (eds.) Proceedings of Machine Learning and Systems, vol. 2, pp. 129–146 (2020). https://proceedings.mlsys.org/paper_files/paper/2020/file/6c44dc73014d66ba49b28d483a8f8b0d-Paper.pdf

  3. Bonawitz, K., et al.: Towards federated learning at scale: system design. In: Talwalkar, A., Smith, V., Zaharia, M. (eds.) Proceedings of Machine Learning and Systems, vol. 1, pp. 374–388 (2019). https://proceedings.mlsys.org/paper_files/paper/2019/file/7b770da633baf74895be22a8807f1a8f-Paper.pdf

  4. Caldarola, D., Mancini, M., Galasso, F., Ciccone, M., Rodola, E., Caputo, B.: Cluster-driven graph federated learning over multiple domains. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, pp. 2749–2758 (June 2021)

    Google Scholar 

  5. Chang, H., Shejwalkar, V., Shokri, R., Houmansadr, A.: Cronus: robust and heterogeneous collaborative learning with black-box knowledge transfer. arXiv preprint arXiv:1912.11279 (2019)

  6. Chen, Y., Ning, Y., Slawski, M., Rangwala, H.: Asynchronous online federated learning for edge devices with non-IID data. In: 2020 IEEE International Conference on Big Data (Big Data), pp. 15–24. IEEE (2020)

    Google Scholar 

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

  8. Croitoru, F.A., Hondru, V., Ionescu, R.T., Shah, M.: Diffusion models in vision: a survey. IEEE Trans. Pattern Anal. Mach. Intell. 45(9), 10850–10869 (2023). https://doi.org/10.1109/TPAMI.2023.3261988

    Article  Google Scholar 

  9. Dhariwal, P., Nichol, A.: Diffusion models beat GANs on image synthesis. In: Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., Vaughan, J.W. (eds.) Advances in Neural Information Processing Systems, vol. 34, pp. 8780–8794. Curran Associates, Inc. (2021). https://proceedings.neurips.cc/paper_files/paper/2021/file/49ad23d1ec9fa4bd8d77d02681df5cfa-Paper.pdf

  10. Diao, E., Ding, J., Tarokh, V.: Heterofl: Computation and communication efficient federated learning for heterogeneous clients. In: International Conference on Learning Representation (2020)

    Google Scholar 

  11. Fan, B., Jiang, S., Su, X., Hui, P.: Model-heterogeneous federated learning for internet of things: enabling technologies and future directions (2023). https://arxiv.org/abs/2312.12091

  12. Fan, B., Su, X., Tarkoma, S., Hui, P.: Behave differently when clustering: a semi-asynchronous federated learning approach for IoT. ACM Trans. Sen. Netw. 20(3) (2024). https://doi.org/10.1145/3639825

  13. Gong, X., et al.: Preserving privacy in federated learning with ensemble cross-domain knowledge distillation. In: AAAI, vol. 36, no. 11, pp. 11891–11899 (2022). https://doi.org/10.1609/aaai.v36i11.21446, https://ojs.aaai.org/index.php/AAAI/article/view/21446

  14. Goodfellow, I., et al.: Generative adversarial nets. In: Ghahramani, Z., Welling, M., Cortes, C., Lawrence, N., Weinberger, K. (eds.) Advances in Neural Information Processing Systems, vol. 27. Curran Associates, Inc. (2014). https://proceedings.neurips.cc/paper_files/paper/2014/file/5ca3e9b122f61f8f06494c97b1afccf3-Paper.pdf

  15. He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: IEEE Conference on Computer Vision and Pattern Recognition (2016)

    Google Scholar 

  16. Hinton, G., Vinyals, O., Dean, J.: Distilling the knowledge in a neural network (2015). https://arxiv.org/abs/1503.02531

  17. Ho, J., Jain, A., Abbeel, P.: Denoising diffusion probabilistic models. In: Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., Lin, H. (eds.) Advances in Neural Information Processing Systems, vol. 33, pp. 6840–6851. Curran Associates, Inc. (2020). https://proceedings.neurips.cc/paper_files/paper/2020/file/4c5bcfec8584af0d967f1ab10179ca4b-Paper.pdf

  18. Ilhan, F., Su, G., Liu, L.: ScaleFL: resource-adaptive federated learning with heterogeneous clients. In: IEEE Conference on Computer Vision and Pattern Recognition, pp. 24532–24541 (2023)

    Google Scholar 

  19. Imteaj, A., Thakker, U., Wang, S., Li, J., Amini, M.H.: A survey on federated learning for resource-constrained IoT devices. IEEE Internet Things J. 9(1), 1–24 (2022). https://doi.org/10.1109/JIOT.2021.3095077

    Article  Google Scholar 

  20. Jeong, E., Oh, S., Kim, H., Park, J., Bennis, M., Kim, S.L.: Communication-efficient on-device machine learning: Federated distillation and augmentation under non-IID private data (2023). https://arxiv.org/abs/1811.11479

  21. Kim, J., Scott, C.D.: Robust kernel density estimation. J. Mach. Learn. Res. 13(1), 2529–2565 (2012)

    MathSciNet  Google Scholar 

  22. Krizhevsky, A., Hinton, G.: Learning multiple layers of features from tiny images. Technical report, University of Toronto (2009). https://www.cs.toronto.edu/~kriz/learning-features-2009-TR.pdf

  23. LattePanda: Lattepanda alpha 864s specs (2023). https://www.lattepanda.com/lattepanda-alpha. Accessed 05 Oct 2023

  24. Li, D., Wang, J.: FedMD: heterogenous federated learning via model distillation. arXiv preprint arXiv:1910.03581 (2019)

  25. Li, J., Sharma, N.K., Ports, D.R., Gribble, S.D.: Tales of the tail: hardware, OS, and application-level sources of tail latency. In: Proceedings of the ACM Symposium on Cloud Computing, pp. 1–14 (2014)

    Google Scholar 

  26. Li, Q., Diao, Y., Chen, Q., He, B.: Federated learning on non-IID data silos: an experimental study. In: 2022 IEEE 38th International Conference on Data Engineering (ICDE), pp. 965–978. IEEE (2022)

    Google Scholar 

  27. Li, Q., et al.: A survey on federated learning systems: vision, hype and reality for data privacy and protection. IEEE Trans. Knowl. Data Eng. (2021)

    Google Scholar 

  28. Li, T., Sahu, A.K., Zaheer, M., Sanjabi, M., Talwalkar, A., Smith, V.: Federated optimization in heterogeneous networks. In: Dhillon, I., Papailiopoulos, D., Sze, V. (eds.) Proceedings of Machine Learning and Systems, vol. 2, pp. 429–450 (2020). https://proceedings.mlsys.org/paper_files/paper/2020/file/1f5fe83998a09396ebe6477d9475ba0c-Paper.pdf

  29. Liao, D., Gao, X., Zhao, Y., Xu, C.Z.: Adaptive channel sparsity for federated learning under system heterogeneity. In: IEEE Conference on Computer Vision and Pattern Recognition, pp. 20432–20441 (2023)

    Google Scholar 

  30. Lim, W.Y.B., et al.: Federated learning in mobile edge networks: a comprehensive survey. IEEE Commun. Surv. Tutor. 22(3), 2031–2063 (2020). https://doi.org/10.1109/COMST.2020.2986024

    Article  Google Scholar 

  31. Lin, T., Kong, L., Stich, S.U., Jaggi, M.: Ensemble distillation for robust model fusion in federated learning. In: Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., Lin, H. (eds.) Advances in Neural Information Processing Systems, vol. 33, pp. 2351–2363. Curran Associates, Inc. (2020). https://proceedings.neurips.cc/paper_files/paper/2020/file/18df51b97ccd68128e994804f3eccc87-Paper.pdf

  32. McMahan, B., Moore, E., Ramage, D., Hampson, S., y Arcas, B.A.: Communication-efficient learning of deep networks from decentralized data. In: Artificial Intelligence and Statistics, pp. 1273–1282. PMLR (2017)

    Google Scholar 

  33. 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, MobiSys 2021, pp. 54–66. Association for Computing Machinery, New York (2021). https://doi.org/10.1145/3458864.3467681

  34. Park, J., Han, D.J., Choi, M., Moon, J.: SageFlow: robust federated learning against both stragglers and adversaries. In: Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., Vaughan, J.W. (eds.) Advances in Neural Information Processing Systems, vol. 34, pp. 840–851. Curran Associates, Inc. (2021). https://proceedings.neurips.cc/paper_files/paper/2021/file/076a8133735eb5d7552dc195b125a454-Paper.pdf

  35. Phuong, M., Lampert, C.: Towards understanding knowledge distillation. In: Chaudhuri, K., Salakhutdinov, R. (eds.) International Conference on Machine Learning. Proceedings of Machine Learning Research, vol. 97, pp. 5142–5151. PMLR (2019). https://proceedings.mlr.press/v97/phuong19a.html

  36. Qiao, T., Zhang, J., Xu, D., Tao, D.: MirrorGAN: learning text-to-image generation by redescription. In: IEEE Conference on Computer Vision and Pattern Recognition (2019)

    Google Scholar 

  37. Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., Chen, M.: Hierarchical text-conditional image generation with CLIP latents. arXiv preprint arXiv:2204.06125 (2022)

  38. Rombach, R., Blattmann, A., Lorenz, D., Esser, P., Ommer, B.: High-resolution image synthesis with latent diffusion models. In: IEEE Conference on Computer Vision and Pattern Recognition, pp. 10684–10695 (2022)

    Google Scholar 

  39. Ronneberger, O., Fischer, P., Brox, T.: U-net: convolutional networks for biomedical image segmentation. In: Navab, N., Hornegger, J., Wells, W.M., Frangi, A.F. (eds.) MICCAI 2015. LNCS, vol. 9351, pp. 234–241. Springer, Cham (2015). https://doi.org/10.1007/978-3-319-24574-4_28

    Chapter  Google Scholar 

  40. Ruiz, N., Li, Y., Jampani, V., Pritch, Y., Rubinstein, M., Aberman, K.: DreamBooth: fine tuning text-to-image diffusion models for subject-driven generation. In: IEEE Conference on Computer Vision and Pattern Recognition, pp. 22500–22510 (2023)

    Google Scholar 

  41. Snell, J., Swersky, K., Zemel, R.: Prototypical networks for few-shot learning. In: Guyon, I., et al. (eds.) Advances in Neural Information Processing Systems, vol. 30. Curran Associates, Inc. (2017). https://proceedings.neurips.cc/paper_files/paper/2017/file/cb8da6767461f2812ae4290eac7cbc42-Paper.pdf

  42. Sun, L., Qian, J., Chen, X.: LDP-FL: practical private aggregation in federated learning with local differential privacy. arXiv preprint arXiv:2007.15789 (2020)

  43. Tan, C.M.J., Motani, M.: DropNet: reducing neural network complexity via iterative pruning. In: III, H.D., Singh, A. (eds.) International Conference on Machine Learning. Proceedings of Machine Learning Research, vol. 119, pp. 9356–9366. PMLR (2020). https://proceedings.mlr.press/v119/tan20a.html

  44. Tan, Y., et al.: FedProto: federated prototype learning across heterogeneous clients. In: AAAI, vol. 36, no. 8, pp. 8432–8440 (2022). https://doi.org/10.1609/aaai.v36i8.20819, https://ojs.aaai.org/index.php/AAAI/article/view/20819

  45. Tun, Y.L., Thwal, C.M., Yoon, J.S., Kang, S.M., Zhang, C., Hong, C.S.: Federated learning with diffusion models for privacy-sensitive vision tasks. In: 2023 International Conference on Advanced Technologies for Communications (ATC), pp. 305–310 (2023). https://doi.org/10.1109/ATC58710.2023.10318858

  46. Wang, J., Liu, Q., Liang, H., Joshi, G., Poor, H.V.: Tackling the objective inconsistency problem in heterogeneous federated optimization. In: Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., Lin, H. (eds.) Advances in Neural Information Processing Systems, vol. 33, pp. 7611–7623. Curran Associates, Inc. (2020). https://proceedings.neurips.cc/paper_files/paper/2020/file/564127c03caab942e503ee6f810f54fd-Paper.pdf

  47. Wang, Z., Li, C., Wang, X.: Convolutional neural network pruning with structural redundancy reduction. In: IEEE Conference on Computer Vision and Pattern Recognition, pp. 14913–14922 (2021)

    Google Scholar 

  48. Wu, C., Li, Z., Wang, F., Wu, C.: Learning cautiously in federated learning with noisy and heterogeneous clients. In: 2023 IEEE International Conference on Multimedia and Expo (ICME), pp. 660–665 (2023https://doi.org/10.1109/ICME55011.2023.00119

  49. Wu, C., Zhu, Y., Zhang, R., Chen, Y., Wang, F., Cui, S.: FedAB: truthful federated learning with auction-based combinatorial multi-armed bandit. IEEE Internet Things J. 10(17), 15159–15170 (2023). https://doi.org/10.1109/JIOT.2023.3264677

    Article  Google Scholar 

  50. Wu, W., He, L., Lin, W., Mao, R., Maple, C., Jarvis, S.: SAFA: a semi-asynchronous protocol for fast federated learning with low overhead. IEEE Trans. Comput. 70(5), 655–668 (2021). https://doi.org/10.1109/TC.2020.2994391

    Article  MathSciNet  Google Scholar 

  51. Yang, M., Su, S., Li, B., Xue, X.: Exploring one-shot semi-supervised federated learning with a pre-trained diffusion model. arXiv preprint arXiv:2305.04063 (2023)

  52. Ye, M., Fang, X., Du, B., Yuen, P.C., Tao, D.: Heterogeneous federated learning: state-of-the-art and research challenges. ACM Comput. Surv. 56(3) (2023). https://doi.org/10.1145/3625558

  53. Zhang, R., Chen, Y., Wu, C., Wang, F.: Multi-level personalized federated learning on heterogeneous and long-tailed data. IEEE Trans. Mob. Comput. 1–14 (2024). https://doi.org/10.1109/TMC.2024.3409159

  54. Zhang, R., Chen, Y., Wu, C., Wang, F., Liu, J.: Optimizing efficient personalized federated learning with hypernetworks at edge. IEEE Netw. 37(4), 120–126 (2023). https://doi.org/10.1109/MNET.008.2200654

    Article  Google Scholar 

  55. Zhang, T., Gao, L., Lee, S., Zhang, M., Avestimehr, S.: TimelyFL: heterogeneity-aware asynchronous federated learning with adaptive partial training. In: IEEE Conference on Computer Vision and Pattern Recognition, pp. 5064–5073 (2023)

    Google Scholar 

  56. Zhang, Y., Xiang, T., Hospedales, T.M., Lu, H.: Deep mutual learning. In: IEEE Conference on Computer Vision and Pattern Recognition (2018)

    Google Scholar 

  57. Zhao, Z., Yang, F., Liang, G.: Federated learning based on diffusion model to cope with non-IID data. In: Liu, Q., et al. (eds.) PRCV 2023. LNCS, vol. 14433, pp. 220–231. Springer, Singapore (2023). https://doi.org/10.1007/978-981-99-8546-3_18

    Chapter  Google Scholar 

  58. Zhu, Z., Hong, J., Zhou, J.: Data-free knowledge distillation for heterogeneous federated learning. In: Meila, M., Zhang, T. (eds.) International Conference on Machine Learning. Proceedings of Machine Learning Research, vol. 139, pp. 12878–12889. PMLR (2021). https://proceedings.mlr.press/v139/zhu21b.html

Download references

Acknowledgements

This work is partially funded by Nordic University Cooperation on Edge Intelligence (168043) and the Guangzhou Municipal Nansha District Science and Technology Bureau under Contract No. 2022ZD012.

Author information

Authors and Affiliations

Authors

Corresponding author

Correspondence to Boyu Fan .

Editor information

Editors and Affiliations

1 Electronic supplementary material

Below is the link to the electronic supplementary material.

Supplementary material 1 (pdf 545 KB)

Rights and permissions

Reprints and permissions

Copyright information

© 2025 The Author(s), under exclusive license to Springer Nature Switzerland AG

About this paper

Check for updates. Verify currency and authenticity via CrossMark

Cite this paper

Fan, B., Wu, C., Su, X., Hui, P. (2025). FedTSA: A Cluster-Based Two-Stage Aggregation Method for Model-Heterogeneous Federated Learning. In: Leonardis, A., Ricci, E., Roth, S., Russakovsky, O., Sattler, T., Varol, G. (eds) Computer Vision – ECCV 2024. ECCV 2024. Lecture Notes in Computer Science, vol 15141. Springer, Cham. https://doi.org/10.1007/978-3-031-73010-8_22

Download citation

  • DOI: https://doi.org/10.1007/978-3-031-73010-8_22

  • Published:

  • Publisher Name: Springer, Cham

  • Print ISBN: 978-3-031-73009-2

  • Online ISBN: 978-3-031-73010-8

  • eBook Packages: Computer ScienceComputer Science (R0)

Publish with us

Policies and ethics