# Towards Robust Federated Learning via Logits Calibration on Non-IID Data

arXiv: https://arxiv.org/abs/2403.02803 · HTML: https://mlsanity.com/paper/2403.02803/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 32 references that are in the map.
- First posted on arXiv 5 Mar 2024.

## Builds on (most cited first, 23 of 23)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [Towards Deep Learning Models Resistant to Adversarial Attacks](https://mlsanity.com/paper/1706.06083/) (2017)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Towards Evaluating the Robustness of Neural Networks](https://mlsanity.com/paper/1608.04644/) (2016)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Adversarial examples in the physical world](https://mlsanity.com/paper/1607.02533/) (2016)
- [Model-Contrastive Federated Learning](https://mlsanity.com/paper/2103.16257/) (2021)
- [Adversarial Training for Free!](https://mlsanity.com/paper/1904.12843/) (2019)
- [Long-tail learning via logit adjustment](https://mlsanity.com/paper/2007.07314/) (2020)
- [Federated Learning on Non-IID Data: A Survey](https://mlsanity.com/paper/2106.06843/) (2021)
- [Privacy and Robustness in Federated Learning: Attacks and Defenses](https://mlsanity.com/paper/2012.06337/) (2020)
- [Federated Learning with Label Distribution Skew via Logits Calibration](https://mlsanity.com/paper/2209.00189/) (2022)
- [FedProc: Prototypical Contrastive Federated Learning on Non-IID data](https://mlsanity.com/paper/2109.12273/) (2021)
- [FAT: Federated Adversarial Training](https://mlsanity.com/paper/2012.01791/) (2020)
- [CalFAT: Calibrated Federated Adversarial Training with Label Skewness](https://mlsanity.com/paper/2205.14926/) (2022)
- [MP-FedCL: Multiprototype Federated Contrastive Learning for Edge Intelligence](https://mlsanity.com/paper/2304.01950/) (2023)
- [Delving into the Adversarial Robustness of Federated Learning](https://mlsanity.com/paper/2302.09479/) (2023)
- [Adversarial training in communication constrained federated learning](https://mlsanity.com/paper/2103.01319/) (2021)
- [Generative AI-driven Semantic Communication Framework for NextG Wireless Network](https://mlsanity.com/paper/2310.09021/) (2023)
- [Boosting Federated Learning Convergence with Prototype Regularization](https://mlsanity.com/paper/2307.10575/) (2023)

## Built on by (most cited first, 5 of 5)

- [Non-IID data in Federated Learning: A Survey with Taxonomy, Metrics, Methods, Frameworks and Future Directions](https://mlsanity.com/paper/2411.12377/) (2024)
- [DeepSeek-Inspired Exploration of RL-based LLMs and Synergy with Wireless Networks: A Survey](https://mlsanity.com/paper/2503.09956/) (2025)
- [Logit Calibration and Feature Contrast for Robust Federated Learning on Non-IID Data](https://mlsanity.com/paper/2404.06776/) (2024)
- [Federated Hybrid Training and Self-Adversarial Distillation: Towards Robust Edge Networks](https://mlsanity.com/paper/2412.19354/) (2024)
- [Towards Communication-Efficient Adversarial Federated Learning for Robust Edge Intelligence](https://mlsanity.com/paper/2501.15257/) (2025)

## Abstract (arXiv)

Federated learning (FL) is a privacy-preserving distributed management framework based on collaborative model training of distributed devices in edge networks. However, recent studies have shown that FL is vulnerable to adversarial examples (AEs), leading to a significant drop in its performance. Meanwhile, the non-independent and identically distributed (non-IID) challenge of data distribution between edge devices can further degrade the performance of models. Consequently, both AEs and non-IID pose challenges to deploying robust learning models at the edge. In this work, we adopt the adversarial training (AT) framework to improve the robustness of FL models against adversarial example (AE) attacks, which can be termed as federated adversarial training (FAT). Moreover, we address the non-IID challenge by implementing a simple yet effective logits calibration strategy under the FAT framework, which can enhance the robustness of models when subjected to adversarial attacks. Specifically, we employ a direct strategy to adjust the logits output by assigning higher weights to classes with small samples during training. This approach effectively tackles the class imbalance in the training data, with the goal of mitigating biases between local and global models. Experimental results on three dataset benchmarks, MNIST, Fashion-MNIST, and CIFAR-10 show that our strategy achieves competitive results in natural and robust accuracy compared to several baselines.