This is the repo for the code and datasets used in the paper Lockdown: Backdoor Defense for Federated Learning with Isolated Subspace Training, accepted by the NeurIPS 2023. The camera ready paper is available here.
The overall procedure can be summarized into four main steps. i) Isolated subspace training. ii)Subspace searching. iii) Aggregation. iv) Model cleaning with consensus fusion. The following figure illustrates the overall process.
- PyTorch
- Numpy
- TorchVision
Dataset FashionMnist and CIFAR10/100 will be automatically downloaded with TorchVision.
The following code run lockdown in its default setting
python federated.py --method lockdown
You can also find script in directory src/script.
-
The main simulation program is in
federated.py, where we initialize the benign and poison dataset, call clients to do local training, call aggregator to do aggregation, do consensus fusion before testing, etc. -
The Lockdown's client local training logistic is in
agent_sparse.py. -
The vanilla FedAvg' client local training logistic is in
agent.py. -
The aggregation logistic is in
aggregation.py, where we implement multiple defense baselines. -
The data poisoning, data preparation and data distribution logistic is in
utils.py.
The logging files will be contained in src/logs. Benign accuracy, ASR, and Backdoor accuracy will be tested in every round.
For Lockdown, the three metrics correspond to the following logging format:
| Clean Val_Loss/Val_Acc: (Benign loss) / (Benign accuracy) |
| Clean Attack Success Ratio: (ASR loss)/ (ASR) |
| Clean Poison Loss/Clean Poison accuracy:: (Backdoor Loss)/ (Backdoor Acc)|
Model checkpoints will be saved every 25 rounds in the directory src/checkpoint.
If you have any questions, you can either open an issue or contact me (thuang374@gatech.edu), and I will reply as soon as I see the issue or email.
The codebase is modified and adapted from one of our baselines RLR.
Lockdown is completely free and released under the MIT License.