Your Diffusion Model is Secretly a Zero-Shot Classifier

Alexander C. Li; Mihir Prabhudesai; Shivam Duggal; Ellis Brown; Deepak Pathak

Source-linked AI summary

Your Diffusion Model is Secretly a Zero-Shot Classifier

Alexander C. Li, Mihir Prabhudesai, Shivam Duggal, Ellis Brown, Deepak Pathak

arXiv:2303.16203v3Mar 28, 2023cs.LGcs.AIcs.CVcs.NEcs.RO

TL;DR

The paper asks whether diffusion models’ conditional density estimates can support image classification beyond generation without additional training. It introduces Diffusion Classifier, which compares conditional noise-prediction errors, and finds strong zero-shot, standard-classification, compositional-reasoning, and robustness results. However, inference remains impractical for many-class settings and robustness gains do not extend to every tested distribution shift.

Problem
Diffusion models have mainly been used for content creation, leaving their direct use for discriminative tasks such as image classification relatively less studied.
Method
Diffusion Classifier uses conditional diffusion density estimates and Monte Carlo estimates of conditional ELBOs to extract classifiers without additional training.
Results
Across zero-shot, standard-classification, and multimodal compositional-reasoning evaluations, Diffusion Classifier narrows the gap with discriminative methods, significantly outperforms them on compositional reasoning, and achieves 79.1% ImageNet accuracy with weak augmentations.
Takeaways & Limitations
The results support using generative diffusion models directly for downstream classification, with particularly strong multimodal compositional reasoning and effective robustness to distribution shift.
Takeaways & Limitations
Inference remains impractical for many classes, and improved effective robustness appears on ImageNet-A but not ImageNetV2 or ObjectNet.

Abstract

from arXiv · show

The recent wave of large-scale text-to-image diffusion models has dramatically increased our text-based image generation abilities. These models can generate realistic images for a staggering variety of prompts and exhibit impressive compositional generalization abilities. Almost all use cases thus far have solely focused on sampling; however, diffusion models can also provide conditional density estimates, which are useful for tasks beyond image generation. In this paper, we show that the density estimates from large-scale text-to-image diffusion models like Stable Diffusion can be leveraged to perform zero-shot classification without any additional training. Our generative approach to classification, which we call Diffusion Classifier, attains strong results on a variety of benchmarks and outperforms alternative methods of extracting knowledge from diffusion models. Although a gap remains between generative and discriminative approaches on zero-shot recognition tasks, our diffusion-based approach has significantly stronger multimodal compositional reasoning ability than competing discriminative approaches. Finally, we use Diffusion Classifier to extract standard classifiers from class-conditional diffusion models trained on ImageNet. Our models achieve strong classification performance using only weak augmentations and exhibit qualitatively better "effective robustness" to distribution shift. Overall, our results are a step toward using generative over discriminative models for downstream tasks. Results and visualizations at https://diffusion-classifier.github.io/

1. Introduction

This paper revisits generative models for image classification, using diffusion models’ conditional density estimates to classify images without additional training. Diffusion Classifier extracts zero-shot classifiers from Stable Diffusion and standard classifiers from class-conditional diffusion models.

Diffusion models have primarily been used for content creation, while their ability to perform discriminative tasks remains less studied.
Conditional diffusion models can estimate class-conditional likelihoods with the ELBO, enabling posterior class probabilities through Bayes’ theorem.
Diffusion Classifier extracts zero-shot classifiers from text-to-image diffusion models and standard classifiers from class-conditional diffusion models without additional training.
Across eleven benchmarks, Diffusion Classifier achieves strong zero-shot accuracy, outperforms alternative diffusion-model extraction methods, and outperforms strongest contrastive methods on Winoground.
79.1% accuracy on ImageNet is achieved with DiT using only weak augmentations, alongside better robustness to distribution shift than competing discriminative classifiers.

2. Related Work

Related work contrasts generative approaches, which model data distributions, with discriminative approaches, which directly learn task decision boundaries. The paper focuses on directly using modern diffusion models as classifiers without downstream fine-tuning.

Discriminative models learn task decision boundaries directly, whereas generative models learn the data distribution and address classification through maximum likelihood estimation.
Generative representations have supported classification, segmentation, adversarial robustness, and calibration, but many methods jointly train or fine-tune them for downstream tasks.
Diffusion models have demonstrated high-fidelity generation across images, videos, 3D, and audio from modalities including text.
Zero-shot classifiers exploit large-scale image-text data to generalize to new tasks and categories, although evaluation samples may still overlap their training distribution.

3. Method: Classification via Diffusion Models

Diffusion Classifier converts a conditional diffusion model into a classifier by comparing conditional noise-prediction errors as approximations to class-conditional likelihoods. Shared Monte Carlo samples reduce variance in relative comparisons, while timestep choices affect accuracy.

3.1. Diffusion Model Preliminaries: Diffusion Classifier estimates class-conditional densities from diffusion models and uses those estimates to classify an input image.
3.1. Diffusion Model Preliminaries: The forward process adds Gaussian noise to a clean image, while the learned reverse process denoises the image conditioned on text embeddings or class indices.
3.1. Diffusion Model Preliminaries: The ELBO approximates log pθ(x | c) using weighted noise-prediction errors, with the final approximation setting wt = 1.
3.2. Classification with diffusion models: With a uniform class prior, posterior probabilities are computed from exponentiated negative expected noise-prediction errors across conditioning inputs.
3.2. Classification with diffusion models: An unbiased Monte Carlo estimate samples timestep-noise pairs and evaluates each conditioning input on those samples.
3.2. Classification with diffusion models: Diffusion Classifier is a hyperparameter-free method that extracts zero-shot or standard classifiers from pretrained conditional diffusion models without additional training.
3.3. Variance Reduction via Difference Testing: Using the same timestep-noise samples for every conditioning input makes error differences more consistent and improves posterior estimation accuracy.
3.3. Variance Reduction via Difference Testing: Single-timestep Pets accuracy is highest at an intermediate noise level, with t = 500.

4. Practical Considerations

Practical improvements target the substantial inference cost of evaluating every class across diffusion timesteps and noise samples. Even with adaptive pruning, classification remains impractical for many-class settings.

4. Practical Considerations: Diffusion Classifier requires repeated error evaluations for every class, creating significant inference time.
4.1. Effect of timestep: Evenly spaced timestep sampling performs best among tested strategies as the number of averaged samples increases.
4.2. Efficient Classification: A naive implementation requires C × N trials, where C is the number of classes and N is the number of timestep-noise samples per conditional ELBO.
4.2. Efficient Classification: Stage-wise evaluation prunes classes with the highest average error, allocating additional computation to plausible candidates.
4.2. Efficient Classification: Classifying one 1000-class ImageNet image takes about 1000 seconds with Stable Diffusion at 512×512 resolution, even using adaptive evaluation.

5. Experimental Details

The experiments evaluate zero-shot and supervised classification settings, comparing Diffusion Classifier with diffusion-based and discriminative baselines across multiple datasets and distribution shifts.

The study provides setup details, baselines, and datasets for zero-shot and supervised classification.
Zero-shot setup: Stable Diffusion 2.0 supplies the zero-shot Diffusion Classifier, with ℓ1 or ℓ2 error chosen as a per-dataset inference hyperparameter and an adaptive procedure from Algorithm 2.
Baselines: The zero-shot comparison includes CLIP, OpenCLIP, Synthetic SD Data, and SD Features, with fairness caveats for discriminative baselines trained on different datasets and architectures.Table 1 reports average accuracy or mean-per-class accuracy.
Zero-shot evaluation: Zero-shot evaluation covers eight image-classification datasets plus the Winoground compositional-reasoning benchmark.ImageNet evaluation uses 2,000 test images because of computational constraints.
Supervised setup: The supervised setup uses DiT-XL/2 class-conditional diffusion models at 2562 and 5122 resolutions, evaluating each class 250 times per image.
Supervised evaluation: Supervised models are compared with ImageNet-trained ResNets and ViTs on ImageNet, ImageNetV2, ImageNet-A, and ObjectNet.ObjectNet uses 113 classes shared with ImageNet, and Diffusion Classifier and baselines are evaluated on the same 10,000-image ImageNet subset.

6. Experimental Results

Across zero-shot, compositional, supervised, and robustness evaluations, Diffusion Classifier performs strongly against diffusion-based and discriminative baselines. The experiments also identify data-distribution and distribution-shift caveats, while showing competitive ImageNet accuracy and improved ImageNet-A effective robustness.

Experimental scope: The experiments compare Diffusion Classifier with zero-shot classifiers, diffusion-based alternatives, supervised discriminative models, and robustness baselines across multiple benchmarks.The study evaluates five questions spanning zero-shot accuracy, alternative diffusion classifiers, compositional reasoning, supervised classification, and distribution-shift robustness.
6.1. Zero-shot Classification Results: Diffusion Classifier outperforms Synthetic SD Data and generally surpasses the supervised SD Features baseline without additional training or labels.It also outperforms CLIP ResNet-50 and is competitive with OpenCLIP ViT-H, although training-dataset differences make those comparisons less direct.
6.2. Improved Compositional Reasoning Abilities: Diffusion Classifier significantly outperforms OpenCLIP ViT-H/14 and CLIP ViT-L/14 on Winoground across Object, Relation, and Both swap categories.The advantage extends to Relation swaps, where contrastive baselines perform no better than random guessing; the authors attribute the improvement to better cross-modal concept binding.
6.3. Supervised Classification Results: 77.5% and 79.1% ImageNet accuracy are achieved at resolutions 256^2 and 512^2, respectively, outperforming ResNet-101 and ViT-L/32.This comparison uses ImageNet-trained DiT-XL/2 and discriminative models trained on the same dataset.
6.3.1 Better Out-of-distribution Generalization: 15-25% effective robustness is achieved on ImageNet-A, with Diffusion Classifier lying far above the discriminative models’ linear ID-OOD fit.The result is reported without extra training data, contrasting with hundreds of discriminative models that follow the fitted relationship.
Caveats and limitations: The robustness finding does not extend to ImageNetV2 or ObjectNet, and Stable Diffusion’s curated training distribution leaves many evaluated datasets out-of-distribution.The authors report that only 0-3% of test images in several datasets would remain after applying all three Stable Diffusion filtering criteria.

7. Conclusion and Discussion

Diffusion Classifier uses conditional diffusion density estimates to extract zero-shot and standard classifiers without additional training. It narrows the classification gap, excels at multimodal compositional reasoning, and shows stronger effective robustness, while inference cost and base-model design remain important constraints.

A simple, unbiased Monte Carlo estimate of conditional ELBO extracts classifiers from conditional diffusion models without additional training.
Diffusion Classifier narrows the gap with state-of-the-art discriminative approaches on zero-shot and standard classification and significantly outperforms them on multimodal compositional reasoning.
Diffusion Classifier exhibits far better effective robustness to distribution shift than the comparison approaches.
Accelerating Inference: Inference time is a practical bottleneck, with acceleration avenues including lower resolution, class pruning, gradient-based search, architectural changes, and parallelization.
Role of Diffusion Model Design Decisions: Because the base diffusion model is unchanged, training choices such as text encoder strength and latent-versus-pixel-space diffusion affect classifier behavior.
The results are an encouraging step toward generative models for classification, compositional reasoning, and robustness.

B. Inference Costs and Hybrid Classification Approach

Diffusion Classifier inference can be expensive because evaluation scales with image resolution and, especially on ImageNet, the number of classes. Adaptive pruning and lower resolution offer substantial speedups, with pruning also improving accuracy.

A single-image classification takes 18 seconds on Pets and 1000 seconds on ImageNet, while ImageNet inference still approximately scales linearly with the number of classes.
Reducing resolution to 128 × 128 would reduce inference time by roughly 16×, although its accuracy impact is difficult to estimate without retraining Stable Diffusion.
Hybrid Classification Approach: A weak discriminative classifier can prune unlikely classes, simultaneously increasing accuracy and reducing inference time.

C. Inference Objective Function

The theoretically justified squared ℓ2 noise-prediction loss is not uniformly optimal in practice. The ℓ1 alternative performs better on roughly half of the evaluated zero-shot datasets, despite lacking theoretical or training-objective justification.

Table 9 evaluates supervised Diffusion Classifier performance across different loss functions.
The inference objective uses the squared ℓ2 error ∥ϵ − ϵθ(xt, c)∥2 as justified by the theory.
The ℓ1 loss performs better on roughly half of the Stable Diffusion zero-shot datasets than the squared ℓ2 loss.
The ℓ1 loss is neither theoretically justified nor part of the Stable Diffusion training objective, making its empirical advantage puzzling.

D. Interpretability via Image Generation

Image reconstruction provides a visual window into the class-dependent features and failures of Diffusion Classifier. Descriptive captions align best with inputs, while confused class prompts can produce similar reconstructions and errors.

Generative samples make class-dependent features and model failures easier to visualize than in discriminative classifiers.
Human-modified BLIP captions produce reconstructions most similar to the input because they provide the most descriptive prompts.
Reconstructions using only class names align less with the input because class names are not dense image descriptions.
Stable Diffusion can generate similar reconstructions for correct Birman and incorrect Ragdoll prompts, causing Diffusion Classifier to misclassify the Birman cat.
Finetuning on 175 Birman/Ragdoll cat images raises accuracy from 45% to 85% for those two classes.
Stable Diffusion accuracy improves across SD 1.x releases but decreases from SD 2.0 to 2.1 on almost every dataset.

F. Additional Implementation Details

The paper details inference and evaluation setups for Stable Diffusion, Winoground, and ImageNet-trained DiT classifiers, alongside two Stable Diffusion baselines.

Training Data: Stable Diffusion 2.0 uses filtered LAION-5B data, excluding images below 512 × 512 pixels, unsafe samples, and low-aesthetic samples.The filtering thresholds are punsafe ≥0.1 and aesthetic score ≤4.5.
Inference Details: Inference resizes each test image’s shortest edge to 512 pixels, takes a 512 × 512 center crop, and normalizes values to [−1, 1].FP16 and Flash Attention enable batch-size-32 inference without changing test accuracy relative to FP32 without Flash Attention.
Winoground: Winoground evaluates each image-caption pair with 1000 evenly spaced timesteps using Stable Diffusion 2.0, without adaptive inference.Adaptive inference is omitted because each example contains only four image-caption pairs.
ImageNet Classifiers: The ImageNet experiments repurpose DiT-XL/2 models trained on ImageNet-1k, containing about 1.28 million images across 1,000 classes.Models are evaluated at resolutions 2562 and 5122 against discriminative models trained on the same data.
Baselines: The study compares Diffusion Classifier with Stable Diffusion synthetic-data and feature-based classification baselines.The synthetic-data baseline generates training examples, while the feature baseline builds classification directly on Stable Diffusion features.

G. Techniques that did not help

Several proposed variance-reduction and guidance techniques fail to improve Diffusion Classifier accuracy. In particular, no classifier-free guidance and no error-map cropping perform best on Pets.

Overall limitation: Diffusion Classifier requires many samples to estimate the ELBO accurately, and additional tested variance-reduction methods did not work.The paper reports these negative results to discourage repeating unsuccessful methods.
Classifier-free guidance: w = 0, meaning no classifier-free guidance, performs best on Pets classification.The authors hypothesize that guidance affects uncertain examples unpredictably, where Diffusion Classifier already fails.
Error map cropping: Any amount of cropping the noise-error map reduces Pets accuracy.The method measures the ELBO error only on a center crop of the 64 × 64 × 4 latent, but cropping does not help.
Importance sampling: Every tested importance-sampling strategy underperforms sampling noise ϵ from a standard normal distribution.The tested strategies fix ϵ, truncate samples, or rescale samples to an expected norm; the experiment uses a 10% Pets subset.