Struggle with Adversarial Defense? Try Diffusion

TL;DR

This work addresses adversarial vulnerability in image classification by framing diffusion models as Bayesian classifiers and introducing TMDC, a Truth Maximization-empowered diffusion classifier. By using conditional diffusion likelihoods and a ground-truth label-guided loss, TMDC achieves strong robustness against heavy white-box and Auto Attack settings on CIFAR-10, surpassing conventional discriminative models. The approach combines a probabilistic diffusion framework with efficient LoRA-based fine-tuning and demonstrates state-of-the-art performance ($82.81\%$ $l_{\infty}$, $86.05\%$ $l_2$ at $\epsilon=0.05$) without full retraining, with further gains when optimally trained under Truth Maximization. The work suggests that generative diffusion classifiers, coupled with Bayesian inference and adversarially-aware optimization, offer a promising path for robust visual recognition in security-sensitive applications.

Abstract

Adversarial attacks induce misclassification by introducing subtle perturbations. Recently, diffusion models are applied to the image classifiers to improve adversarial robustness through adversarial training or by purifying adversarial noise. However, diffusion-based adversarial training often encounters convergence challenges and high computational expenses. Additionally, diffusion-based purification inevitably causes data shift and is deemed susceptible to stronger adaptive attacks. To tackle these issues, we propose the Truth Maximization Diffusion Classifier (TMDC), a generative Bayesian classifier that builds upon pre-trained diffusion models and the Bayesian theorem. Unlike data-driven classifiers, TMDC, guided by Bayesian principles, utilizes the conditional likelihood from diffusion models to determine the class probabilities of input images, thereby insulating against the influences of data shift and the limitations of adversarial training. Moreover, to enhance TMDC's resilience against more potent adversarial attacks, we propose an optimization strategy for diffusion classifiers. This strategy involves post-training the diffusion model on perturbed datasets with ground-truth labels as conditions, guiding the diffusion model to learn the data distribution and maximizing the likelihood under the ground-truth labels. The proposed method achieves state-of-the-art performance on the CIFAR10 dataset against heavy white-box attacks and strong adaptive attacks. Specifically, TMDC achieves robust accuracies of 82.81% against $l_{\infty}$ norm-bounded perturbations and 86.05% against $l_{2}$ norm-bounded perturbations, respectively, with $ε=0.05$.

Figures (4)

• Figure 1: Simplified Illustration of Lora. Utilizing low-dimensional matrices to approximate high-dimensional ones, where pre-trained weights are frozen, and Lora tensors are employed for training. The memory require during training approaches that of the model's inference process. This configuration reduces both training time and memory overhead, while effectively mitigating catastrophic forgetting.
• Figure 2: Overview of the Inference Process of the Diffusion Classifier. Perturbed images are fed into the Diffusion model for both forward noisy processing and backward denoising, with the guiding textual prompt also inputted into the model. The model computes the posterior probabilities corresponding to each class label using Bayes' theorem, and the maximum posterior probability corresponds to the inference result of the classifier. The objective of the inference process in classification can be transformed into selecting the class corresponding to the minimum average error between the noise inferred by the diffusion model at each sampling point and the ground truth value.
• Figure 3: Comparison between Diffusion Classifier and TMDC. The PGD attack is conducted with parameters set as follows: $\epsilon$ is set to 0.05, and the attack runs for 40 iterations, in accordance with Section 4.2. As for Auto Attack, its version is uniformly designated as "plus", with $\epsilon$ set to 0.05 and the seed initialized with five sets of distinct random numbers.
• Figure 4: Study on Checkpoint Selection. For Auto Attack, the version is uniformly set to "plus", with a value of 0.05 for parameter $\epsilon$, and the seed is fixed at 2024. Throughout the Truth Maximization training process, a learning rate scheduler employing "constant with warmup" strategy is employed, wherein the learning rate is set to 1e-6, and the warm-up steps are configured to be 100. Both sets of experiments undergo optimization for 3000 steps.