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.
Access this chapter
Tax calculation will be finalised at checkout
Purchases are for personal use only
Similar content being viewed by others
References
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
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
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
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)
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)
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)
Chrabaszcz, P., Loshchilov, I., Hutter, F.: A downsampled variant of ImageNet as an alternative to the CIFAR datasets. arXiv preprint arXiv:1707.08819 (2017)
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
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
Diao, E., Ding, J., Tarokh, V.: Heterofl: Computation and communication efficient federated learning for heterogeneous clients. In: International Conference on Learning Representation (2020)
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
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
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
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
He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: IEEE Conference on Computer Vision and Pattern Recognition (2016)
Hinton, G., Vinyals, O., Dean, J.: Distilling the knowledge in a neural network (2015). https://arxiv.org/abs/1503.02531
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
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)
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
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
Kim, J., Scott, C.D.: Robust kernel density estimation. J. Mach. Learn. Res. 13(1), 2529–2565 (2012)
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
LattePanda: Lattepanda alpha 864s specs (2023). https://www.lattepanda.com/lattepanda-alpha. Accessed 05 Oct 2023
Li, D., Wang, J.: FedMD: heterogenous federated learning via model distillation. arXiv preprint arXiv:1910.03581 (2019)
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)
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)
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)
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
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)
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
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
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)
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
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
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
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)
Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., Chen, M.: Hierarchical text-conditional image generation with CLIP latents. arXiv preprint arXiv:2204.06125 (2022)
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)
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
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)
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
Sun, L., Qian, J., Chen, X.: LDP-FL: practical private aggregation in federated learning with local differential privacy. arXiv preprint arXiv:2007.15789 (2020)
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
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
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
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
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)
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
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
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
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)
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
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
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
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)
Zhang, Y., Xiang, T., Hospedales, T.M., Lu, H.: Deep mutual learning. In: IEEE Conference on Computer Vision and Pattern Recognition (2018)
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
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
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
Corresponding author
Editor information
Editors and Affiliations
1 Electronic supplementary material
Below is the link to the electronic supplementary material.
Rights and permissions
Copyright information
© 2025 The Author(s), under exclusive license to Springer Nature Switzerland AG
About this paper
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)