Robust Classification via a Single Diffusion Model

TL;DR

This paper addresses adversarial robustness by reframing diffusion models as robust generative classifiers. It introduces Robust Diffusion Classifier (RDC), which uses Bayes' theorem to compute $p(y|oldsymbol{x})$ from $p(oldsymbol{x}|y)$ estimated by a pre-trained diffusion model, augmented with likelihood maximization to move inputs toward high-density regions and a multi-head diffusion backbone to reduce computation. Theoretical analysis shows that an optimal diffusion model can achieve 100% robust accuracy, while practical implementations hinge on reducing the gap between log-likelihood and diffusion loss; empirically, RDC achieves state-of-the-art robustness on CIFAR-10 under strong adaptive attacks and generalizes well to unseen threats. The combination of a generative probabilistic framework, efficiency-focused architecture, and pre-optimization steps demonstrates the potential of diffusion models for robust classification in real-world adversarial settings.

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}.

Figures (4)

• Figure 1: Illustration of our proposed Robust Diffusion Classifier (RDC). Given an input image $\boldsymbol{\mathbf{x}}$, our approach first maximizes the data likelihood (Left) and then classifies the optimized image $\hat{\boldsymbol{\mathbf{x}}}$ with a diffusion model (Right). The class probability $p(y|\hat{\boldsymbol{\mathbf{x}}})$ is given by the conditional log-likelihood $\log p_\theta(\hat{\boldsymbol{\mathbf{x}}}|y)$, which is approximated by the variational lower bound involving calculating the noise prediction error (i.e., diffusion loss) averaged over different timesteps for every class.
• Figure 2: (a): Randomness of different methods. (b-c): Ablation studies of $\eta$ and $T'$.
• Figure 3: The prediction of ELBO and BPD on CIFAR-10 test set and CIFAR-10-C.
• Figure 4: The images generated by multi-head diffusion.

Theorems & Definitions (8)

• Theorem 3.1
• Theorem 3.2
• Corollary 3.3
• proof
• proof
• proof
• Lemma 1.1
• proof