Traceback of Data Poisoning Attacks in Neural Networks

ABOUT

This repository contains code implementation of the paper “Traceback of Data Poisoning Attacks in Neural Networks“, at USENIX Security 2022. This traceback tool for poison data is developed by researchers at SANDLab, University of Chicago.

Note

The code base currently only support CIFAR10 dataset and BadNets attack. Adapting to new dataset is possible by editing the data loading and model loading code.

DEPENDENCIES

Our code is implemented and tested on Keras with TensorFlow backend. Following packages are used by our code.

keras==2.4.3
numpy==1.19.5
tensorflow-gpu==2.4.1

Our code is tested on Python 3.6.8

Steps to train a backdoored model and run traceback

This code base includes code to train a backdoored model and traceback code to identify malicious data points from the training set.

Step 1: Train/download backdoor model

To train a backdoored model, simply run:

python3 inject_backdoor.py --config cifar2 --gpu 0

The --config argument will load a config file under the config directory. Feel free to modify the config file for a different set of parameters. --gpu argument will specify which GPU to use. We do not recommend running this code on CPU-only machines.

Step 2: Run traceback

Running traceback leverages two scripts. gen_embs.py will generate the embedding needed for clustering and run_traceback.py performs traceback using generated embeddings and unlearning.

First, run python3 gen_embs.py --config cifar2 --gpu 0 to generate embedding, which will be saved under results/. Then run python3 run_traceback.py --config cifar2 --gpu 0. It should perform clustering and pruning. The unlearning process will be printed out. At the end, it will output the final precision and recall.

Things to keep in mind when adapting to new dataset/model

During unlearning process, make sure to set the training parameters identical to ones used during the initial model training. This includes augmentation, learning rate, and optimizers. If you found the unlearning process to be unstable, please reduce the learning rate or add gradient clipping. If the issue persist, please reach out to us via [email protected].

Citation

@inproceedings{shan2022traceback,
  title={Traceback of Data Poisoning Attacks in Neural Networks},
  author={Shan, Shawn and Bhagoji, Arjun Nitin and Zheng, Haitao and Zhao, Ben Y},
  journal={Proc. of USENIX Security},
  year={2022}
}