SkyMask: Attack-agnostic Robust Federated Learning with Fine-grained Learnable Masks

TL;DR

SkyMask introduces a fine-grained, learnable masking approach for federated learning that effectively detects malicious updates at the parameter level, significantly improving robustness against sophisticated Byzantine attacks.

Contribution

It proposes SkyMask, a novel attack-agnostic defense mechanism using learnable masks trained on a small dataset to identify malicious model updates at the parameter level.

Findings

Achieves up to 14% higher accuracy than state-of-the-art defenses.

Effectively defends against high fractions (up to 80%) of malicious clients.

Demonstrates robustness across multiple models and datasets.

Abstract

Federated Learning (FL) is becoming a popular paradigm for leveraging distributed data and preserving data privacy. However, due to the distributed characteristic, FL systems are vulnerable to Byzantine attacks that compromised clients attack the global model by uploading malicious model updates. With the development of layer-level and parameter-level fine-grained attacks, the attacks' stealthiness and effectiveness have been significantly improved. The existing defense mechanisms solely analyze the model-level statistics of individual model updates uploaded by clients to mitigate Byzantine attacks, which are ineffective against fine-grained attacks due to unawareness or overreaction. To address this problem, we propose SkyMask, a new attack-agnostic robust FL system that firstly leverages fine-grained learnable masks to identify malicious model updates at the parameter level. Specifically, the FL server freezes and multiplies the model updates uploaded by clients with the parameter-level masks, and trains the masks over a small clean dataset (i.e., root dataset) to learn the subtle difference between benign and malicious model updates in a high-dimension space. Our extensive experiments involve different models on three public datasets under state-of-the-art (SOTA) attacks, where the results show that SkyMask achieves up to 14% higher testing accuracy compared with SOTA defense strategies under the same attacks and successfully defends against attacks with malicious clients of a high fraction up to 80%. Code is available at https://github.com/KoalaYan/SkyMask.