Better Diffusion Models Further Improve Adversarial Training

Code for the paper Better Diffusion Models Further Improve Adversarial Training.

Environment settings and libraries we used in our experiments

This project is tested under the following environment settings:

OS: Ubuntu 20.04.3
GPU: NVIDIA A100
Cuda: 11.1, Cudnn: v8.2
Python: 3.9.5
PyTorch: 1.8.0
Torchvision: 0.9.0

Acknowledgement

The adversarial training codes are modifed based on the PyTorch implementation of Rebuffi et al., 2021. The generation codes are modifed based on the official implementation of EDM. For data generation, please refer to edm/README.md for more details.

Requirements

Install or download AutoAttack:

pip install git+https://github.com/fra31/auto-attack

Install or download RandAugment:

pip install git+https://github.com/ildoonet/pytorch-randaugment

Download EDM generated data to ./edm_data/cifar10 and ./edm_data/cifar100. Since 20M and 50M data files are too large, we split them into several parts:

dataset	size	link
CIFAR-10	1M	npz
CIFAR-10	5M	npz
CIFAR-10	10M	npz
CIFAR-10	20M	part1 part2
CIFAR-10	50M	part1 part2 part3 part4
CIFAR-100	1M	npz
CIFAR-100	50M	part1 part2 part3 part4

Merge 20M and 50M generated data:

python merge-data.py

Training Commands

Run train-wa.py for reproducing the results reported in the papers. For example, train a WideResNet-28-10 model via TRADES on CIFAR-10 with the 1M additional generated data provided by EDM (Karras et al., 2022):

python train-wa.py --data-dir 'cifar-data' \
    --log-dir 'trained_models' \
    --desc 'WRN28-10Swish_cifar10s_lr0p2_TRADES5_epoch400_bs512_fraction0p7_ls0p1' \
    --data cifar10s \
    --batch-size 512 \
    --model wrn-28-10-swish \
    --num-adv-epochs 400 \
    --lr 0.2 \
    --beta 5.0 \
    --unsup-fraction 0.7 \
    --aux-data-filename 'edm_data/cifar10/1m.npz' \
    --ls 0.1

Evaluation Commands

The trained models can be evaluated by running eval-aa.py which uses AutoAttack for evaluating the robust accuracy. Run the command (taking the checkpoint above as an example):

python eval-aa.py --data-dir 'cifar-data' \
    --log-dir 'trained_models' \
    --desc 'WRN28-10Swish_cifar10s_lr0p2_TRADES5_epoch400_bs512_fraction0p7_ls0p1'

To evaluate the model on last epoch under AutoAttack, run the command:

python eval-last-aa.py --data-dir 'cifar-data' \
    --log-dir 'trained_models' \
    --desc 'WRN28-10Swish_cifar10s_lr0p2_TRADES5_epoch400_bs512_fraction0p7_ls0p1'

Pre-trained checkpoints

We provide the state-of-the-art pre-trained checkpoints of WRN-28-10 (Swish) and WRN-70-16 (Swish). Clean and robust accuracies are measured on the full test set. The robust accuracy is measured using AutoAttack.

dataset	norm	radius	architecture	clean	robust	link
CIFAR-10	ℓ_∞	8 / 255	WRN-28-10	92.44%	67.31%	checkpoint argtxt
CIFAR-10	ℓ_∞	8 / 255	WRN-70-16	93.25%	70.69%	checkpoint argtxt
CIFAR-10	ℓ₂	128 / 255	WRN-28-10	95.16%	83.63%	checkpoint argtxt
CIFAR-10	ℓ₂	128 / 255	WRN-70-16	95.54%	84.86%	checkpoint argtxt
CIFAR-100	ℓ_∞	8 / 255	WRN-28-10	72.58%	38.83%	checkpoint argtxt
CIFAR-100	ℓ_∞	8 / 255	WRN-70-16	75.22%	42.67%	checkpoint argtxt

Downloading checkpoint to trained_models/mymodel/weights-best.pt
Downloading argtxt to trained_models/mymodel/args.txt

For evaluation under AutoAttack, run the command:

python eval-aa.py --data-dir 'cifar-data' --log-dir 'trained_models' --desc 'mymodel'

References

If you find the code useful for your research, please consider citing

@article{wang2023better,
  title={Better Diffusion Models Further Improve Adversarial Training},
  author={Wang, Zekai and Pang, Tianyu and Du, Chao and Lin, Min and Liu, Weiwei and Yan, Shuicheng},
  journal={arXiv preprint arXiv:2302.04638},
  year={2023}
}

and/or our related works

@inproceedings{pang2022robustness,
  title={Robustness and Accuracy Could be Reconcilable by (Proper) Definition},
  author={Pang, Tianyu and Lin, Min and Yang, Xiao and Zhu, Jun and Yan, Shuicheng},
  booktitle={International Conference on Machine Learning (ICML)},
  year={2022}
}

@inproceedings{pang2021bag,
  title={Bag of Tricks for Adversarial Training},
  author={Pang, Tianyu and Yang, Xiao and Dong, Yinpeng and Su, Hang and Zhu, Jun},
  booktitle={International Conference on Learning Representations (ICLR)},
  year={2021}
}

xaddwell / adv_train