FedA ⁴: Federated Learning with Anti-Bias Aggregation and TrAjectory-Based Adaptation

Abstract

Non-Independent and Identically Distributed (Non-IID) data pose a fundamental challenge in Federated Learning (FL). It usually causes a severe client drift issue (various client model update directions) and thus, degrades the global model performance. Existing methods typically address this by assigning appropriate weights to client models or optimizing model update directions. However, these methods overlook client model update trends. They focus solely on the final client models to be aggregated at the server at each communication round, ignoring model optimization trajectories, which may contain richer information to aid model convergence. To address this issue, we propose FedA ⁴, a novel FL framework with Anti-bias Aggregation and trAjectory-based Adaptation, which leverages clients’ optimization trajectories, rather than only their final model snapshots. For anti-bias aggregation, by observing a phenomenon termed model collapse, where biased clients tend to predict any input data as the dominant classes in their own datasets, we quantify the class dominance and analyze the level of client drift for each client. We evaluate a prediction entropy, namely concentration, so as to assign an optimal weight to each client at each training round. To further mitigate the negative effect of clients with high levels of client drift (biased clients), we then develop a gradient adaptation mechanism termed trajectory-based adaptation, which analyzes clients’ trajectories to correct each client’s contribution to the aggregated global model. Extensive experiments on CIFAR-10, CIFAR-100, STL-10, and Fashion-MNIST demonstrate that FedA ⁴ significantly outperforms state-of-the-art baselines, particularly in scenarios with extreme data heterogeneity (high level of Non-IID).

Keywords

federated Learning / non-IID Data / client drift / anti-bias aggregation / optimization trajectory

Author Contributions

G.Z.: Conceptualization, investigation, writing, and revision. J.H.: Conceptualization, investigation, writing, and revision. Z.Y.: Conceptualization, investigation, writing, and revision. D.W.: Conceptualization, investigation, writing, and revision. All authors have read and agreed to the published version of the manuscript.

Funding

This paper is partially supported by Hong Kong Research Grants Council grants #11203523, #11209425, C1042-23GF, and Hong Kong Innovation and Technology Commission grant MHP/034/22.

Data Availability Statement

Not applicable.

Conflicts of Interest

The authors declare no conflict of interest. Given the role as Editor-in-Chief, Dapeng Oliver Wu had no involvement in the peer review of this paper and had no access to information regarding its peer-review process. Full responsibility for the editorial process of this paper was delegated to another editor of the journal.

Use of AI and AI-Assisted Technologies

During the preparation of this work, the authors used ChatGPT to check the grammar. After using this tool, the authors reviewed and edited the content as needed and take full responsibility for the content of the published article.

References

[1] McMahan, B.; Moore, E.; Ramage, D.; et al. Communication—efficient learning of deep networks from decentralized data. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics (AISTATS) 2017, Fort Lauderdale, FL, USA, 20—22 April 2017; pp. 1273-1282.

[2] Konečný, J.; McMahan, H.B.; Ramage, D.; et al. Federated optimization: Distributed machine learning for on—device intelligence. arXiv 2016, arXiv:1610.02527.

[3] Li, T.; Sahu, A.K.; Talwalkar, A.; et al. Federated learning: Challenges, methods, and future directions. IEEE Signal Process. Mag. 2020, 37, 50-60.

[4] Zhao, Y.; Li, M.; Lai, L.; et al. Federated learning with non—iid data. arXiv 2018, arXiv:1806.00582.

[5] Li, T.; Sanjabi, M.; Beirami, A.; et al. Fair resource allocation in federated learning. arXiv 2019, arXiv:1905.10497.

[6] Ji, S.; Pan, S.; Long, G.; et al. Learning private neural language modeling with attentive aggregation. In Proceedings of the 2019 International Joint Conference on Neural Networks (IJCNN), Budapest, Hungary, 14—19 July 2019; pp. 1-8.

[7] T Dinh, C.; Tran, N.; Nguyen, J. Personalized federated learning with moreau envelopes. In Proceedings of the 34th International Conference on Neural Information Processing Systems, Vancouver, BC, Canada, 6—12 December 2020.

[8] Fu, S.; Yang, Z.; Hu, C.; et al. Personalized federated learning with contrastive momentum. IEEE Trans. Big Data 2024, 11, 2184-2194.

[9] Blanchard, P.; El Mhamdi, E.M.; Guerraoui, R.; et al. Machine learning with adversaries: Byzantine tolerant gradient descent. In Proceedings of the Advances in Neural Information Processing Systems 30, Long Beach, CA, USA, 4—9 December 2017.

[10] Wang, H.; Yurochkin, M.; Sun, Y.; et al. Federated learning with matched averaging. arXiv 2020, arXiv:2002.06440.

[11] Yurochkin, M.; Agarwal, M.; Ghosh, S.; et al. Bayesian nonparametric federated learning of neural networks. In Proceedings of the 36th International Conference on Machine Learning, Long Beach, CA, USA, 9—15 June 2019; pp. 7252-7261.

[12] Lin, T.; Kong, L.; Stich, S.U.; et al. Ensemble Distillation for Robust Model Fusion in Federated Learning. In Proceedings of the 34th Conference on Neural Information Processing Systems (NeurIPS 2020), Vancouver, CA, USA, 6—12 December 2020.

[13] Li, T.; Sahu, A.K.; Zaheer, M.; et al. Federated optimization in heterogeneous networks. In Proceedings of the Third Conference on Machine Learning and Systems MLSys 2020, Austin, TX, USA, 2—4 March 2020.

[14] Karimireddy, S.P.; Kale, S.; Mohri, M.; et al. Scaffold: Stochastic controlled averaging for federated learning. In Proceedings of the 37th International Conference on Machine Learning, Virtual, 13—18 July 2020; pp. 5132-5143.

[15] Acar, D.A.E.; Zhao, Y.; Navarro, R.M.; et al. Federated learning based on dynamic regularization. arXiv 2021, arXiv:2111.04263.

[16] Gao, L.; Fu, H.; Li, L.; et al. Feddc: Federated learning with non—iid data via local drift decoupling and correction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, New Orleans, LA, USA, 18—24 June 2022; pp. 10112-10121.

[17] Zhu, Z.; Hong, J.; Zhou, J. Data—free knowledge distillation for heterogeneous federated learning. In Proceedings of the 38th International Conference on Machine Learning, Virtual, 18—24 July 2021; pp. 12878-12889.

[18] Yang, M.; Su, S.; Li, B.; et al. One—Shot Heterogeneous Federated Learning with Local Model—Guided Diffusion Models. arXiv 2025, arXiv: cs.CV/2311.08870.

[19] Goodfellow, I.J.; Pouget—Abadie, J.; Mirza, M.; et al. Generative adversarial nets. In Proceedings of the NIPS’14: Neural Information Processing Systems, Montreal, QC, Canada, 8—13 December 2014

[20] Zhu, L.; Liu, Z.; Han, S. Deep leakage from gradients. In Proceedings of the NeurIPS 2019, Vancouver, BC, Canada, 8—14 December 2019.

[21] Zhao, B.; Mopuri, K.R.; Bilen, H. Dataset condensation with gradient matching. arXiv 2020, arXiv:2006.05929.

[22] Krizhevsky, A. Learning Multiple Layers of Features from Tiny Images; University of Toronto: Toronto, ON, Canada, 2009.

[23] Coates, A.; Ng, A.; Lee, H. An Analysis of Single—Layer Networks in Unsupervised Feature Learning. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, Fort Lauderdale, FL, USA, 11—13 April 2011; Volume 15, pp. 215-223.

[24] Xiao, H.; Rasul, K.; Vollgraf, R. Fashion—mnist: A novel image dataset for benchmarking machine learning algorithms. arXiv 2017, arXiv:1708.07747.

[25] Reddi, S.; Charles, Z.; Zaheer, M.; et al. Adaptive federated optimization. arXiv 2020, arXiv: 2003.00295.

[26] Wang, R.; Chen, Y. Adaptive Model Aggregation in Federated Learning Based on Model Accuracy. IEEE Wireless Commun. 2024, 31, 200-206.

[27] Pillutla, K.; Kakade, S.M.; Harchaoui, Z. Robust aggregation for federated learning. IEEE Trans. Signal Process. 2022, 70, 1142-1154.

[28] Li, Q.; He, B.; Song, D. Model—contrastive federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, Nashville, TN, USA, 20—25 June 2021; pp. 10713-10722.

[29] Rasouli, M.; Sun, T.; Rajagopal, R. Fedgan: Federated generative adversarial networks for distributed data. arXiv 2020, arXiv:2006.07228.