Awesome
Decoupled Kullback-Leibler (DKL) Divergence Loss
This repository contains the implementation code for our arXiv paper Decoupled Kullback-Leibler (DKL) Divergence Loss, arXiv.
Overview
In this paper, we delve deeper into the KullbackLeibler (KL) Divergence loss and observe that it is equivalent to the Doupled Kullback-Leibler (DKL) Divergence loss that consists of 1) a weighted Mean Square Error (wMSE) loss and 2) a Cross4 Entropy loss incorporating soft labels. From our analysis of the DKL loss, we have identified two areas for improvement. Firstly, we address the limitation of DKL in scenarios like knowledge distillation by breaking its asymmetry property in training optimization. This modification ensures that the wMSE component is always effective during training, providing extra constructive cues. Secondly, we introduce global information into DKL for intra-class consistency regularization. With these two enhancements, we derive the Improved KullbackLeibler (IKL) Di11 vergence loss and evaluate its effectiveness by conducting experiments on CIFAR12 10/100 and ImageNet datasets, focusing on adversarial training and knowledge dis13 tillation tasks. The proposed approach achieves new state-of-the-art performance on both tasks, demonstrating the substantial practical merits.
Results and Pretrained Models for Knowledge Distillation
ImageNet
Method | Model-Teacher | Model-Student | Training Speed | Top-1 Acc(%) | link | log |
---|---|---|---|---|---|---|
ReviewKD | ResNet-34 | ResNet18 | 0.319 s/iter | 71.61 | - | - |
DKD | ResNet-34 | ResNet18 | - | 71.70 | - | - |
IKL-KD | ResNet-34 | ResNet18 | 0.197 s/iter | 71.91 | - | - |
Method | Model-Teacher | Model-Student | Training Speed | Top-1 Acc(%) | link | log |
---|---|---|---|---|---|---|
ReviewKD | ResNet-50 | MobileNet | 0.526 s/iter | 72.56 | - | - |
DKD | ResNet-50 | MobileNet | - | 72.05 | - | - |
IKL-KD | ResNet-50 | MobileNet | 0.252 s/iter | 72.84 | - | - |
Results and Pretrained Models for Adversarial Robustness
By 2023/05/20, with IKL loss, we achieve new state-of-the-art adversarial robustness under settings that with/without augmentation strategies on auto-attack.
CIFAR-100 with autoaug
# | Method | Model | Natural Acc | Robust Acc (AutoAttack) | link | log |
---|---|---|---|---|---|---|
1 | DAJAT | WRN-34-10 | 68.74 | 31.30 | - | - |
2 | IKL-AT | WRN-34-10 | 65.93 | 32.52 | model | log |
CIFAR-100 with basic data preprocessing (random crop and random horizontal flip)
# | Method | Model | Natural Acc | Robust Acc (AutoAttack) | link | log |
---|---|---|---|---|---|---|
1 | AWP | WRN-34-10 | 60.38 | 28.86 | - | - |
2 | LBGAT | WRN-34-10 | 62.31 | 29.33 | - | - |
3 | LAS-AT | WRN-34-10 | 62.99 | 30.77 | - | - |
4 | ACAT | WRN-34-10 | 65.75 | 30.23 | - | - |
5 | IKL-AT | WRN-34-10 | 66.51 | 31.43 | model | log |
6 | IKL-AT | WRN-34-10 | 63.40 | 31.92 | model | log |
CIFAR-100 with synthesized data
# | Method | Model | Natural Acc | Robust Acc (AutoAttack) | link | log |
---|---|---|---|---|---|---|
1 | Wang et al. (better diffusion models) 1M | WRN-28-10 | 68.06 | 35.65 | - | - |
2 | Wang et al. (better diffusion models) 50M | WRN-28-10 | 72.58 | 38.83 | - | - |
3 | IKL-AT 1M | WRN-28-10 | 68.99 | 35.89 | - | - |
4 | IKL-AT 50M | WRN-28-10 | 73.85 | 39.18 | model | log |
CIFAR-10 with basic data preprocessing (random crop and random horizontal flip)
# | Method | Model | Natural Acc | Robust Acc (AutoAttack) | link | log |
---|---|---|---|---|---|---|
1 | AWP | WRN-34-10 | 85.36 | 56.17 | - | - |
2 | LBGAT | WRN-34-20 | 88.70 | 53.57 | - | - |
3 | LAS-AT | WRN-34-10 | 87.74 | 55.52 | - | - |
4 | ACAT | WRN-34-10 | 82.41 | 55.36 | - | - |
5 | IKL-AT | WRN-34-10 | 85.31 | 57.13 | model | log |
CIFAR-10 with synthesized data
# | Method | Model | Natural Acc | Robust Acc (AutoAttack) | link | log |
---|---|---|---|---|---|---|
1 | Wang et al. (better diffusion models) 1M | WRN-28-10 | 91.12 | 63.35 | - | - |
2 | Wang et al. (better diffusion models) 20M | WRN-28-10 | 92.44 | 67.31 | - | - |
3 | IKL-AT 1M | WRN-28-10 | 90.75 | 63.54 | - | - |
4 | IKL-AT 20M | WRN-28-10 | 92.16 | 67.75 | model | log |
Training
More training scripts will be provided soon to reproduce our results on knowledge distillation and adversarial training tasks.
For the adversarial training task:
cd adv_training
bash sh/train_dkl_cifar100.sh
bash sh/train_dkl_cifar100_autoaug.sh
bash sh/train_dkl_cifar10.sh
Evaluation
before running the evaluation with auto-attack, please download the pre-trained models.
cd adv_training/auto_attacks
bash sh/eval.sh
Contact
If you have any questions, feel free to contact us through email (jiequancui@gmail.com) or Github issues. Enjoy!
BibTex
If you find this code or idea useful, please consider citing our related work:
@article{cui2023decoupled,
title={Decoupled Kullback-Leibler Divergence Loss},
author={Cui, Jiequan and Tian, Zhuotao and Zhong, Zhisheng and Qi, Xiaojuan and Yu, Bei and Zhang, Hanwang},
journal={arXiv preprint arXiv:2305.13948},
year={2023}
}
@inproceedings{cui2021learnable,
title={Learnable boundary guided adversarial training},
author={Cui, Jiequan and Liu, Shu and Wang, Liwei and Jia, Jiaya},
booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision},
pages={15721--15730},
year={2021}
}