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 FedA4, 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 FedA4 significantly outperforms state-of-the-art baselines, particularly in scenarios with extreme data heterogeneity (high level of Non-IID).
- Open Access
- Article
FedA4: Federated Learning with Anti-Bias Aggregation and TrAjectory-Based Adaptation
- Guanyi Zhao 1,
- Juntao Hu 1,2,
- Zhengjie Yang 1,3,*,
- Dapeng Wu 1
Author Information
Received: 02 Jan 2026 | Revised: 29 Jan 2026 | Accepted: 02 Feb 2026 | Published: 09 Feb 2026
Abstract
Keywords
federated Learning | non-IID Data | client drift | anti-bias aggregation | optimization trajectory
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ˇcn`y, 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 36 th 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.
This work is licensed under a Creative Commons Attribution 4.0 International License.
Suite 4002 Level 4, 447 Collins Street, Melbourne, Victoria 3000, Australia
General Inquiries: info@sciltp.com