Home

Awesome

Better Diffusion Models Further Improve Adversarial Training

Code for the paper Better Diffusion Models Further Improve Adversarial Training (ICML 2023).

Environment settings and libraries we used in our experiments

This project is tested under the following environment settings:

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

pip install git+https://github.com/fra31/auto-attack
pip install git+https://github.com/ildoonet/pytorch-randaugment
datasetsizelink
CIFAR-101Mnpz
CIFAR-105Mnpz
CIFAR-1010Mnpz
CIFAR-1020Mpart1 part2
CIFAR-1050Mpart1 part2 part3 part4
CIFAR-1001Mnpz
CIFAR-10050Mpart1 part2 part3 part4
SVHN1Mnpz
SVHN50Mpart1 part2 part3 part4 part5
TinyImageNet1Mnpz
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 'dataset-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 'dataset-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 'dataset-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). Refer to argtxt below for specific hyper-parameters. Clean and robust accuracies are measured on the full test set. The robust accuracy is measured using AutoAttack.

datasetnormradiusarchitecturecleanrobustlink
CIFAR-10<sub></sub>8 / 255WRN-28-1092.44%67.31%checkpoint argtxt
CIFAR-10<sub></sub>8 / 255WRN-70-1693.25%70.69%checkpoint argtxt
CIFAR-10<sub>2</sub>128 / 255WRN-28-1095.16%83.63%checkpoint argtxt
CIFAR-10<sub>2</sub>128 / 255WRN-70-1695.54%84.86%checkpoint argtxt
CIFAR-100<sub></sub>8 / 255WRN-28-1072.58%38.83%checkpoint argtxt
CIFAR-100<sub></sub>8 / 255WRN-70-1675.22%42.67%checkpoint argtxt
SVHN<sub></sub>8 / 255WRN-28-1095.56%64.01%checkpoint argtxt
TinyImageNet<sub></sub>8 / 255WRN-28-1065.19%31.30%checkpoint argtxt

For evaluation under AutoAttack:

  1. Download checkpoint to trained_models/mymodel/weights-best.pt
  2. Download argtxt to trained_models/mymodel/args.txt
  3. Run the command:
python eval-aa.py --data-dir 'dataset-data' --log-dir 'trained_models' --desc 'mymodel'

We have uploaded CIFAR-10/CIFAR-100 models to the model zoo of RobustBench. See the tour to evaluate the performance by RobustBench.

References

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

@inproceedings{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},
  booktitle={International Conference on Machine Learning (ICML)},
  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}
}