Federated Hybrid Training and Self-Adversarial Distillation: Towards Robust Edge Networks

TL;DR

The paper tackles the fragility of federated learning under data heterogeneity and adversarial attacks in edge networks. It proposes FedBAT, a framework that fuses FL-based hybrid adversarial training with augmentation-invariant self-adversarial distillation to simultaneously improve robustness against attacks and generalization under non-IID data while preserving clean accuracy. The approach combines a dual-branch training scheme and a feature-alignment regularizer that ties local adversarial representations to global clean representations, reducing client drift. Extensive experiments across MNIST, Fashion-MNIST, SVHN, Office-Amazon, and CIFAR-10 demonstrate that FedBAT consistently outperforms baselines in both accuracy and robustness, even in large-scale and highly heterogeneous settings, with supportive ablations and visualizations. Overall, FedBAT offers a scalable, practical solution for robust edge-network deployment of FL models.

Abstract

Federated learning (FL) is a distributed training technology that enhances data privacy in mobile edge networks by allowing data owners to collaborate without transmitting raw data to the edge server. However, data heterogeneity and adversarial attacks pose challenges to develop an unbiased and robust global model for edge deployment. To address this, we propose Federated hyBrid Adversarial training and self-adversarial disTillation (FedBAT), a new framework designed to improve both robustness and generalization of the global model. FedBAT seamlessly integrates hybrid adversarial training and self-adversarial distillation into the conventional FL framework from data augmentation and feature distillation perspectives. From a data augmentation perspective, we propose hybrid adversarial training to defend against adversarial attacks by balancing accuracy and robustness through a weighted combination of standard and adversarial training. From a feature distillation perspective, we introduce a novel augmentation-invariant adversarial distillation method that aligns local adversarial features of augmented images with their corresponding unbiased global clean features. This alignment can effectively mitigate bias from data heterogeneity while enhancing both the robustness and generalization of the global model. Extensive experimental results across multiple datasets demonstrate that FedBAT yields comparable or superior performance gains in improving robustness while maintaining accuracy compared to several baselines.

Figures (10)

• Figure 1: Illustration of the vanilla robust FL framework (FedPGD), which serves as the basic structure for robust federated systems. $\mathcal{F}_1$ and $\mathcal{F}_n$ denote the local models of clients $1$ and $n$, respectively. $\mathcal{F}_1(\boldsymbol{x}_1^{adv})$ and $\mathcal{F}_n(\boldsymbol{x}_n^{adv})$ represent the predictions using AEs, with the model updated by minimizing the distance between these predictions and the ground truth labels. FedBAT modifies both the local update process ② and the global aggregation process ④ to enhance robustness while maintaining accuracy.
• Figure 2: Illustration of the proposed FL-based hybrid-AT strategy designed to balance the CA and RA. $\mathcal{F}_i$ denotes the local model of an arbitrary client. $\mathcal{F}_i(\boldsymbol{x}_i^{adv})$ and $\mathcal{F}_i(\boldsymbol{x}_i)$ represent the model's predictions using AEs and CEs, respectively. A coefficient $\alpha$ is used to balance the trade-off between $\mathcal{L}_i^{adv}$ (CE loss) and $\mathcal{L}_i$ (CE loss) during the local training.
• Figure 3: Illustration of the proposed self-adversarial distillation strategy designed to address the non-IID challenge. For an arbitrary client, $\hat{\boldsymbol{x}}_i$ represents the augmented CEs, while $\hat{\boldsymbol{x}}_i + \delta$ denotes the augmented AEs. $f_i^{e}$ denotes the feature extractor, and $f_i(\hat{\boldsymbol{x}}_i^{adv})$ and $f(\hat{\boldsymbol{x}}_i)$ represent the model's feature embeddings obtained from AEs and CEs, respectively. $X$ denotes the averaged global clean feature embeddings. An MSE loss function is introduced to align the global clean embeddings with their semantically same local adversarial embeddings.
• Figure 4: Illustration of the key idea behind the proposed self-adversarial distillation. Different augmentations are employed for the same semantic image across various clients, and the resulting representations should be invariant. The goal is to minimize the distance between the representations of the augmented AEs from each client and the mean representation of the corresponding semantically identical CEs.
• Figure 5: Comparison of the average clean accuracy (CA) on Fashion-MNIST under different levels of data heterogeneity. The clean accuracy is reported for the adversarially trained model but evaluated on clean samples. A lower Dirichlet parameter value indicates higher heterogeneity.