[Paper Review] Robust Federated Learning with Attack-Adaptive Aggregation

Ching Pui Wan, Qifeng Chen|arXiv (Cornell University)|Feb 10, 2021

Adversarial Robustness in Machine Learning51 references4 citations

TL;DR

This paper proposes an attack-adaptive aggregation mechanism for robust federated learning that uses a neural network with an attention mechanism to detect and reweight corrupted model updates from malicious clients. By learning vulnerable regions of model updates under various attacks in a self-supervised manner, the method achieves state-of-the-art defense performance against model poisoning and backdoor attacks across image and text benchmarks.

ABSTRACT

Federated learning is vulnerable to various attacks, such as model poisoning and backdoor attacks, even if some existing defense strategies are used. To address this challenge, we propose an attack-adaptive aggregation strategy to defend against various attacks for robust federated learning. The proposed approach is based on training a neural network with an attention mechanism that learns the vulnerability of federated learning models from a set of plausible attacks. To the best of our knowledge, our aggregation strategy is the first one that can be adapted to defend against various attacks in a data-driven fashion. Our approach has achieved competitive performance in defending model poisoning and backdoor attacks in federated learning tasks on image and text datasets.

Motivation & Objective

• To address the vulnerability of federated learning to model poisoning and backdoor attacks, especially under non-IID data distributions.
• To overcome limitations of existing aggregation rules—such as median, geometric median, and similarity-based methods—that can be bypassed by adaptive attackers.
• To develop a data-driven, adaptive defense mechanism that learns to detect attack patterns without prior knowledge of attack types.
• To improve robustness in federated learning by dynamically reweighting client updates based on learned vulnerability scores from an attention module.

Proposed method

• Train an attention-based neural network to learn a low-dimensional representation of model update vectors, identifying regions sensitive to various attacks.
• Use self-supervised learning on collected update vectors from simulated attacks to train the attention module without requiring labeled attack data.
• Compute alignment scores between latent representations and update vectors to assign adaptive weights, emphasizing benign updates and downweighting corrupted ones.
• Apply a robust aggregation rule that combines reweighted updates using a scale factor c and a truncation function to suppress outliers.
• Leverage the attention mechanism to approximate similarity measures between updates, enabling detection of subtle attacks like edge-case backdoors.
• Integrate the learned aggregation strategy into the federated averaging framework, allowing dynamic adaptation to unknown or evolving attack patterns.

Figure 1: The overview of our approach for estimating the robust mean ${\bm{\mu}}_{robust}$ of the update vectors.

Experimental results

Research questions

• RQ1 Can a data-driven, adaptive aggregation strategy outperform fixed aggregation rules in defending against diverse and adaptive attacks in federated learning?
• RQ2 How effective is an attention-based neural network in detecting subtle or edge-case model poisoning and backdoor attacks without prior attack knowledge?
• RQ3 To what extent can the proposed method generalize across different datasets, client numbers, and attack parameters?
• RQ4 What is the theoretical error bound of the proposed robust aggregation under adversarial conditions, and how does it scale with the fraction of attackers?
• RQ5 How does the attention mechanism’s approximation accuracy affect the convergence and robustness of the global model?

Key findings

• The proposed attack-adaptive aggregation strategy achieves competitive performance in defending against both model poisoning and backdoor attacks across MNIST, CIFAR-10, Tiny-ImageNet, and IMDb datasets.
• The method outperforms classical robust estimators like median and geometric median, as well as similarity-based defenses such as FoolsGold and Krum, especially under non-IID data and high attack rates.
• The attention-based model successfully detects edge-case backdoor attacks where only a small region of the model update is altered, which previous defenses often fail to detect.
• Theoretical analysis shows that the error bound of the robust estimate is controlled by $2e^{-2c}$ for less than 50% attackers, achieving a relative error bound of about $10^{-9}$ when $c=10$, indicating high robustness.
• The method generalizes well across different datasets and client configurations, maintaining strong defense performance even when attack parameters or data distributions vary.
• Using a conservative scale factor $c$ ensures robustness against previously unseen attacks by preventing false detection and rejection of benign clients.

Figure 2: The performance of aggregation strategies in the MNIST (top), CIFAR(middle), ImageNet(bottom) classification task. The ‘no attack’ scenario is combined with the omniscient attack in the left plots.