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

TL;DR

This paper proposes a logits calibration strategy within federated adversarial training to enhance the robustness of federated learning models against adversarial attacks, especially under non-IID data distributions.

Contribution

It introduces a simple logits calibration method to improve federated adversarial training robustness on non-IID data, addressing class imbalance and bias issues.

Findings

Improved robustness against adversarial attacks on MNIST, Fashion-MNIST, and CIFAR-10.

Effective mitigation of class imbalance in federated learning.

Achieved competitive natural and robust accuracy compared to baselines.

Abstract

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.