[Paper Review] Robust Classification via a Single Diffusion Model

Huanran Chen, Yinpeng Dong+5 more|arXiv (Cornell University)|May 24, 2023

Adversarial Robustness in Machine Learning4 citations

TL;DR

This paper proposes Robust Diffusion Classifier (RDC), a generative classifier built from a pre-trained diffusion model that enhances adversarial robustness by maximizing input likelihood before applying Bayes' theorem to compute class probabilities via conditional likelihood. RDC achieves 73.24% robust accuracy against ℓ∞ perturbations (ε=8/255) on CIFAR-10, surpassing prior adversarial training methods by 2.34% and generalizing well to unseen threat models without task-specific fine-tuning.

ABSTRACT

Diffusion models have been applied to improve adversarial robustness of image classifiers by purifying the adversarial noises or generating realistic data for adversarial training. However, diffusion-based purification can be evaded by stronger adaptive attacks while adversarial training does not perform well under unseen threats, exhibiting inevitable limitations of these methods. To better harness the expressive power of diffusion models, this paper proposes Robust Diffusion Classifier (RDC), a generative classifier that is constructed from a pre-trained diffusion model to be adversarially robust. RDC first maximizes the data likelihood of a given input and then predicts the class probabilities of the optimized input using the conditional likelihood estimated by the diffusion model through Bayes' theorem. To further reduce the computational cost, we propose a new diffusion backbone called multi-head diffusion and develop efficient sampling strategies. As RDC does not require training on particular adversarial attacks, we demonstrate that it is more generalizable to defend against multiple unseen threats. In particular, RDC achieves $75.67\%$ robust accuracy against various $\ell_\infty$ norm-bounded adaptive attacks with $ε_\infty=8/255$ on CIFAR-10, surpassing the previous state-of-the-art adversarial training models by $+4.77\%$. The results highlight the potential of generative classifiers by employing pre-trained diffusion models for adversarial robustness compared with the commonly studied discriminative classifiers. Code is available at \url{https://github.com/huanranchen/DiffusionClassifier}.

Motivation & Objective

• To address the limitations of existing adversarial robustness methods that rely on discriminative classifiers and fail under unseen threats.
• To explore the potential of generative models, particularly diffusion models, in constructing inherently robust classifiers.
• To develop a defense mechanism that does not require retraining on specific adversarial attacks, enabling generalization across diverse threat models.
• To improve robustness by leveraging the diffusion model’s accurate score estimation and data distribution modeling capabilities.

Proposed method

• The method estimates class probabilities p(y|x) using Bayes’ theorem, combining the prior p(y) with the conditional likelihood pθ(x|y) from a pre-trained diffusion model.
• The conditional likelihood is approximated via a variational lower bound that computes noise prediction loss across multiple timesteps for each class.
• A pre-optimization step maximizes the log-likelihood of the input image by iteratively perturbing it toward high-likelihood regions using gradient ascent.
• The optimization uses a momentum-based approach with a budget η=8/255 to balance robustness and misclassification risk.
• The method employs systematic sampling of timesteps to reduce computational cost while preserving performance.
• The approach avoids gradient obfuscation by using exact gradients in attacks, confirming robustness is not due to defensive obfuscation.

Experimental results

Research questions

• RQ1 Can a pre-trained diffusion model be directly repurposed as a robust classifier without task-specific fine-tuning?
• RQ2 Does likelihood maximization of inputs improve robustness in a generative classification framework?
• RQ3 Can a diffusion-based classifier generalize to unseen threat models better than discriminative adversarial training?
• RQ4 How does the choice of optimization budget and sampling strategy affect robustness and accuracy?
• RQ5 Is the robustness of RDC due to gradient obfuscation or inherent model properties?

Key findings

• RDC achieves 73.24% robust accuracy against ℓ∞ perturbations with ε=8/255 on CIFAR-10, outperforming previous state-of-the-art adversarial training models by 2.34%.
• The method generalizes well to unseen threat models, demonstrating robustness beyond the specific attacks used during optimization.
• The likelihood maximization pre-processing step improves robustness by moving inputs toward high-likelihood regions, reducing vulnerability to adversarial perturbations.
• The BPDA attack achieves 73.24% robust accuracy on RDC, confirming that the robustness is not due to gradient obfuscation but stems from the model’s intrinsic properties.
• Using momentum in the likelihood optimization increases robustness by ~7 percentage points compared to non-momentum optimization.
• Reducing the number of timesteps T′ for loss computation has minimal impact on clean accuracy but significantly degrades robust accuracy, indicating the importance of full or systematic timestep coverage.