SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kevin Clark*
Google DeepMind Toronto
kevclark@google.com
Priyank Jaini*
pjaini@google.com
The excellent generative capabilities of text-to-image diffusion models suggest they learn informative representations of image-text data. However, what knowledge their representations capture is not fully understood, and they have not been thoroughly explored on downstream tasks. We investigate diffusion models by proposing a method for evaluating them as zero-shot classifiers. The key idea is using a diffusion model's ability to denoise a noised image given a text description of a label as a proxy for that label's likelihood. We apply our method to Stable Diffusion and Imagen, using it to probe fine-grained aspects of the models' knowledge and comparing them with CLIP's zero-shot abilities. They perform competitively with CLIP on a wide range of zero-shot image classification datasets. Additionally, they achieve state-of-the-art results on shape/texture bias tests and can successfully perform attribute binding while CLIP cannot. Although generative pre-training is prevalent in NLP, visual foundation models often use other methods such as contrastive learning. Based on our findings, we argue that generative pre-training should be explored as a compelling alternative for vision-language tasks.
Large models pre-trained on internet-scale data can adapt effectively to a variety of downstream tasks. Increasingly, they are being used as zero-shot learners with no task-specific training, such as with CLIP (Radford et al., 2021) for images and GPT-3 (Brown et al., 2020) for text. In natural language processing, many successful pre-trained models are generative (i.e., language models). However, generative pre-training is less commonly used for visual tasks. Until recently, the usual practice for vision problems was to pre-train models on labeled datasets such asImagenet (Deng et al., 2009), or JFT (Sun et al., 2017). Later research in visual and vision-language problems has led to image-text models pre-trained primarily using either contrastive losses (Radford et al., 2021; Jia et al., 2021; Yuan et al., 2021) or autoencoding tasks (Vincent et al., 2010; He et al., 2022).
On the other hand, generative text-to-image models based on denoising diffusion probabilistic models (Ho et al., 2020) such as Imagen (Saharia et al., 2022a), Dalle-2 (Ramesh et al., 2022), and Stable Diffusion (Rombach et al., 2022) can generate realistic high-resolution images and generalize to diverse text prompts. Their strong performance suggests that they learn effective representations of image-text data. However, their ability to transfer to downstream discriminative tasks and how they compare to other pre-trained models has not been explored thoroughly.
In this paper, we investigate these questions by transferring Imagen and Stable Diffusion (SD) to discriminative tasks. While previous studies have used representations from diffusion models for downstream tasks (Brempong et al., 2022; Burgert et al., 2022; Zhao et al., 2023), we instead propose
Figure 1: Zero-Shot Classification using Diffusion Models. We first compute denoising scores for each label prompt across multiple time-steps to generate a scores matrix. We then classify an image by aggregating the scores for each class using a weighting function over the time-steps. The image is assigned the class with the minimum aggregate score. In Section 3.1, we discuss how efficiency can be improved only computing a subset of the full scores matrix.
a way of using text-to-image diffusion models directly as zero-shot image classifiers. Our method essentially runs the models as generative classifiers (Ng & Jordan, 2001), using a re-weighted version of the variational lower bound to score images since diffusion models do not produce exact likelihoods. More specifically, the method repeatedly noises and denoises the input image while conditioning the model on a different text prompt for each possible class. The class whose text prompt results in the best denoising ability is predicted. This procedure is expensive because it requires denoising many times per class (with different noise levels). To make it usable in practice, we present improvements that increase the method's sample efficiency by up to  $1000\mathrm{x}$ , such as pruning obviously-incorrect classes early. While still requiring too much compute to be an easily-deployable classifier, our method allows us to quantitatively study fine-grained aspects of a diffusion model's learned knowledge through evaluation on classification tasks (as opposed to qualitatively examining model generations).
We compare Imagen and SD against  $\mathrm{CLIP}^2$  (Radford et al., 2021), a widely used model for zero-shot image-text tasks trained with contrastive learning. A high-level goal of the experiments is to see the strengths and weaknesses of generative and contrastive pre-training for computer vision. First, we demonstrate that diffusion models have strong zero-shot classification accuracies (competitive with CLIP) on several diverse vision datasets. Next, we show both Imagen and SD performs remarkably well on the Cue-Conflict dataset (Geirhos et al., 2019), where images have been stylized with textures conflicting with their labels. For example, Imagen achieves  $>50\%$  error reduction over CLIP and even outperforms the much larger ViT-22B (Dehghani et al., 2023) model. This finding is particularly interesting because, unlike supervised classifiers, humans are known to be much more reliant on shape than texture when identifying images. Lastly, we study attribute binding using the synthetic data from Lewis et al. (2022), and find that, unlike CLIP, diffusion models can successfully bind together attributes in some settings.
The main contributions of this paper are:
- We show text-to-image diffusion models can be used as effective zero-shot classifiers. While using too much compute to be very practical on downstream tasks, the method provides a way of quantitatively studying what the models learn.
- We develop techniques that hugely lower the compute cost of these zero-shot classifiers, making them usable (although still slow) on datasets with many classes.
- We demonstrate the strong generalization capabilities of Imagen and Stable Diffusion, resulting in good zero-shot performance on vision datasets (comparable to CLIP).
- We show that diffusion models are robust to misleading textural cues, achieving state-of-the-art results on Cue-Conflict.
- We use our framework to study attribute binding in diffusion models and find that they can perform some binding tasks while CLIP cannot.
2We use ViT-L/14, the largest public CLIP model
Together, our results suggest that text-to-image diffusion models learn powerful representations that can effectively be transferred to tasks beyond image generation.
Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020; Song & Ermon, 2020) are latent variable generative models defined by a forward and reverse Markov chain. Given an unknown data distribution,  $q(\pmb{x}_0)$ , over observations,  $\pmb{x}_0 \in \mathbb{R}^d$ , the forward process corrupts the data into a sequence of noisy latent variables,  $\pmb{x}_{1:T} := \{\pmb{x}_1, \pmb{x}_2, \dots, \pmb{x}_T\}$ , by gradually adding Gaussian noise with a fixed schedule defined as:
$$
q \left(\boldsymbol {x} _ {1: T} \mid \boldsymbol {x} _ {0}\right) := \prod_ {t = 1} ^ {T} q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {t - 1}\right) \tag {1}
$$
where  $q(\pmb{x}_t|\pmb{x}_{t-1}) \coloneqq \mathrm{Normal}(\pmb{x}_t; \sqrt{1 - \beta_t}\pmb{x}_{t-1}, \beta_t\pmb{I})$ . The reverse Markov process gradually denoises the latent variables to the data distribution with learned Gaussian transitions starting from  $\mathrm{Normal}(\pmb{x}_T; 0, \pmb{I})$  i.e.
$$
p _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {0: T}\right) := p \left(\boldsymbol {x} _ {T}\right) \cdot \prod_ {t = 0} ^ {T - 1} p _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} _ {t}\right)
$$
$p_{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t) \coloneqq \text{Normal}\big(\boldsymbol{x}_{t-1}; \boldsymbol{\mu}_{\theta}(\boldsymbol{x}_t, t), \boldsymbol{\Sigma}_{\theta}(\boldsymbol{x}_t, t)\big)$ . The aim of training is for the forward process distribution  $\{\boldsymbol{x}_t\}_{t=0}^T$  to match that of the reverse process  $\{\tilde{\boldsymbol{x}}_t\}_{t=0}^T$  i.e., the generative model  $p_{\theta}(\boldsymbol{x}_0)$  closely matches the data distribution  $q(\boldsymbol{x}_0)$ . Specifically, these models can be trained by optimizing the variational lower bound of the marginal likelihood (Ho et al., 2020; Kingma et al., 2021):
$$
- \log p _ {\theta} (\boldsymbol {x} _ {0}) \leq - \mathrm {V L B} (\boldsymbol {x} _ {0}) := \mathcal {L} _ {\text {P r i o r}} + \mathcal {L} _ {\text {R e c o n}} + \mathcal {L} _ {\text {D i f f u s i o n}}
$$
$\mathcal{L}_{\text {Prior }}$  and  $\mathcal{L}_{\text {Recon }}$  are the prior and reconstruction loss that can be estimated using standard techniques in the literature (Kingma & Welling, 2014). The (re-weighted) diffusion loss can be written as:
$$
\mathcal {L} _ {\mathrm {D i f f u s i o n}} = \mathbb {E} _ {\boldsymbol {x} _ {0}, \boldsymbol {\varepsilon}, t} \left[ \boldsymbol {w} _ {t} \| \boldsymbol {x} _ {0} - \tilde {\boldsymbol {x}} _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {t}, t) \| _ {2} ^ {2} \right]
$$
with  $\pmb{x}_0 \sim q(\pmb{x}_0)$ ,  $\varepsilon \sim \mathrm{Normal}(0, I)$ , and  $t \sim \mathcal{U}([0, T])$ . Here,  $\pmb{w}_t$  is a weight assigned to the timestep, and  $\tilde{\pmb{x}}_{\theta}(\pmb{x}_t, t)$  is the model's prediction of the observation  $\pmb{x}_0$  from the noised observation  $\pmb{x}_t$ . Diffusion models can be conditioned on additional inputs like class labels, text prompts, segmentation masks or low-resolution images, in which case  $\tilde{\pmb{x}}_{\theta}$  also takes a conditioning signal  $\pmb{y}$  as input.
3 Zero-Shot Classification using Diffusion Models
In this section, we show how to convert the generation process of a text-to-image diffusion model into a zero-shot classifier to facilitate quantitative evaluation on downstream tasks. Figure 1 shows an overview of our method.
Diffusion Generative Classifier: We begin with a dataset,  $\{(x^1,y^1),\ldots ,(x^n,y^n)\} \subseteq \mathbb{R}^{d_1\times d_2}\times$ $[y_K]$  of  $n$  images where each image belongs to one of  $K$  classes  $[\mathbf{y}_K]\coloneqq \{\mathbf{y}_1,\mathbf{y}_2,\dots ,\mathbf{y}_K\}$ . Given an image  $\pmb{x}$ , our goal is to predict the most probable class assignment
$$
\tilde {y} = \underset {y _ {k}} {\arg \max } p (y = y _ {k} | \boldsymbol {x}) = \underset {y _ {k}} {\arg \max } p (\boldsymbol {x} | y = y _ {k}) \cdot p (y = y _ {k}) = \underset {y _ {k}} {\arg \max } \log p (\boldsymbol {x} | y = y _ {k}).
$$
where we assume a uniform prior  $p(y_{i} = y_{k}) = \frac{1}{k}$  that can be dropped from the arg max. A generative classifier (Ng & Jordan, 2001) uses a conditional generative model with parameters  $\theta$  to estimate the likelihood as  $p_{\theta}(\pmb{x}|y = y_k)$ .
Using a text-to-image diffusion model as a generative classifier requires two modifications. First, the models are conditioned on text prompts rather than class labels. Thus we convert each label,  $y_{k}$ , to text using a mapping  $\phi$  with a dataset-specific template (e.g.  $y_{k} \to A$  photo of a  $y_{k}$ ). Second,
3 For simplicity, we use  $\pmb{x}$  in place of  $\pmb{x}_0$  to refer to an image.
4We can't use a learned prior in the zero-shot setting.
diffusion models do not produce exact log-likelihoods (i.e. we cannot compute  $\log p_{\theta}(\pmb{x}|y = \mathbf{y}_k)$  directly). Our key idea for a solution is to use the VLB (more specifically  $\mathcal{L}_{\mathrm{Diffusion}}$  as Imagen and SD are not trained with the other losses) as a proxy. Thus we have:
$$
\begin{array}{l} \tilde{y} = \operatorname *{arg  max}_{y_{k}}\log p_{\theta}(\boldsymbol {x}|y = y_{k})\approx \operatorname *{arg  min}_{y_{k}}\mathcal{L}_{\text{Diffusion}}(\boldsymbol {x},y_{k}) \\ = \underset {y _ {k} \in [ y _ {K} ]} {\arg \min } \mathbb {E} _ {\epsilon , t} \left[ \boldsymbol {w} _ {t} \| \boldsymbol {x} - \tilde {\boldsymbol {x}} _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {t}, \phi \left(y _ {k}\right), t\right) \| _ {2} ^ {2} \right] \tag {2} \\ \end{array}
$$
Note that for SD,  $\pmb{x}$  and  $\tilde{\pmb{x}}_{\theta}$  are latent representations, with  $\pmb{x}$  obtained by encoding the image using a VAE. With Imagen on the other hand,  $\pmb{x}$  consists of the raw image pixels.
Estimating the Expectation: We approximate the expectation in Equation (2) using Monte-Carlo estimation. At each step, we sample a  $t \sim \mathcal{U}([0,1])$  and then a  $\pmb{x}_t$  according to the forward diffusion process (Equation (1)):  $\pmb{x}_t \sim q(\pmb{x}_t | \pmb{x}_0)$ . Next, we denoise this noisy image using the model (i.e. we use it to predict  $\pmb{x}$  from  $\pmb{x}_t$ ), obtaining  $\hat{\pmb{x}} = \tilde{\pmb{x}}_{\pmb{\theta}}(\pmb{x}_t, \phi(\mathbf{y}_k), t)$ . We call the squared error of the prediction,  $\| \pmb{x} - \hat{\pmb{x}} \|_2^2$ , a score for  $(\pmb{x}, \mathbf{y}_k)$ . We score each class  $N$  times, obtaining a  $K \times N$  scores matrix for the image. Finally, we weight the scores according to the corresponding  $\pmb{w}_t$  and take the mean, resulting in an estimate of  $\mathcal{L}_{\text{Diffusion}}$  for each class.
Choice of Weighting Function: Imagen and SD are trained with the "simple" loss, where  $w_{t} = \mathrm{SNR}(t)$ , the signal-to-noise ratio (Kingma et al., 2021) for timestep  $t$ . However, we found other weighting functions can improve results. First, we experimented with learning  $w_{t}$  by binning the times into 20 buckets and training a 20-features logistic regression model to learn weights for the buckets that maximize classification accuracy. However, using that weighting is not truly zero-shot since it requires label information to learn. We thus, also handcrafted a weighting function that can be used across datasets. We designed  $w_{t}$  by finding a simple function that looked close to our learned weighting function on CIFAR-100 (we did not look at other datasets to preserve zero-shot protocol). Interestingly, we found that the simple function  $w_{t} \coloneqq \exp(-7t)$  works well for both Imagen and SD across tasks and used it for our experiments. As it is monotonic,  $\mathcal{L}_{\text{Diffusion}}$  with this weighting can still be viewed as a likelihood-based objective that maximizes the variational lower bound under simple data augmentations (Kingma & Gao, 2023). We provide details on learning  $w_{t}$  and an empirical comparison of different weighting functions in Appendix B.
3.1 Improving Efficiency
Computing  $\tilde{y}$  with naive Monte-Carlo estimation can be expensive because  $\mathcal{L}_{\mathrm{Diffusion}}$  has fairly high variance. Here, we propose techniques that reduce the compute cost of estimating the arg min over classes. The key idea is to leverage the fact that we only need to compute the arg min and do not require good estimates of the actual expectations.
Shared Noise: Differences between individual Monte-Carlo samples from  $\mathcal{L}_{\mathrm{Diffusion}}$  can of course be due to different  $t$  or forward diffusion samples from  $q(\pmb{x}_t|\pmb{x}_{t-1})$ , whereas we are only interested in the effect of the text conditioning  $\phi(y_k)$ . We find far fewer samples are necessary when we use the same  $t$  and  $\pmb{x}_t$  across different classes, as shown in Figure 1. After sampling a  $t \sim \mathcal{U}([0,1])$  and  $\pmb{x}_t \sim q(\pmb{x}_t|\pmb{x}_0)$ , we score all classes against this noised image instead of a single one. As a result, the differences between these estimates are only due to the different text conditioning signals.
Candidate Class Pruning: Rather than using the same amount of compute to estimate the expectation for each class, we can further improve efficiency by discarding implausible classes early and dynamically allocating more compute to plausible ones. In particular, we maintain a set of candidate classes for the image being classified. After collecting a new set of scores for each candidate class, we discard classes that are unlikely to become the lowest-scoring (i.e. predicted) class with more samples. Since we are collecting paired samples (with the same  $t$  and  $\hat{x}_{i,t}$ ), we use a paired student's t-test to identify classes that can be pruned. This pruning can be viewed as a successive elimination algorithm for best-arm identification in a multi-armed bandit setting (Paulson, 1964; Even-Dar et al., 2002). Of course, scores do not exactly follow the standard assumptions of a student's t-test, so we
5 Later we discuss how we can avoid computing the full matrix for efficiency.
use a small p-value  $(2\mathrm{e}^{-3}$  in our experiments) and ensure each class is scored a minimum number of times (20 in our experiments) to minimize the chance of pruning the correct class. The full procedure is shown in Algorithm 1.
Comparison: Figure 2 compares the number of samples needed to accurately classify CIFAR-100 images for different efficiency strategies. Using shared noise and pruning greatly improves efficiency, requiring up to  $1000\mathrm{x}$  less compute than naive scoring. Nevertheless, classifying with a diffusion model still typically takes 10s of scores per class on average, making the diffusion classifier expensive to use for datasets with many classes.
4 Empirical Analysis and Results
In this section, we detail our analysis for the diffusion zero-shot classifiers on a variety of tasks. These include classification on various vision datasets to study generalization capabilities on diverse domains, evaluating model robustness to conflicting cues between texture and shape, and studying attribute binding ability through targeted evaluation on synthetic data.
We mainly compare Imagen and SD with CLIP (Radford et al., 2021). We chose CLIP because it is a well-studied and widely-used model, as our primary aim is to study diffusion models rather than push state-of-the-art zero-shot accuracies. Our experiments reveal strengths and weaknesses of image-text representations learned via generative training vs. CLIP's contrastive training.
Algorithm 1 Diffusion model classification with pruning.
given: Example to classify  $\mathbf{x}$ , diffusion model w/ params  $\theta$ , weighting function  $\mathbf{w}$ , hyperparameters min Scores, max scores, cutoff_pval. //Map from classes to diffusion model scores. scores = {y_i : [] for y_i ∈ [y_K]}  
n = 0  
while |scores| > 1 and n < max Scores:  
    n = n + 1  
    //Noise the image  
    t ~ U([0, 1])  
    x_t ~ q(x_t|x)  
//Score against the remaining classes.  
for y_i ∈ scores:  
    add w_t \| x - \tilde{x}_\theta(x_t, \phi(y_i), t) \|_2^2  
    to scores[y_i]  
//Prune away implausible classes.  
 $\tilde{y} = \arg \min_{y_i} \text{scores}[y_i]. \text{mean()}$   
if n ≥ min Scores:  
    for y_i ∈ scores:  
        if paired_ttest_pval(scores[tilde_y], scores[y_i]) < cutoff_pval:  
            remove y_i from scores.  
return  $\tilde{y}$
Model details: Imagen is a cascaded diffusion model (Ho et al., 2022) consisting of a  $64 \times 64$  low-resolution model and two super-resolution models. We only use the  $64 \times 64$  model for our experiments because we found the high-resolution models performed poorly as classifiers. Combining the  $64 \times 64$  model's scores with scores from the higher-resolutions models did not improve results either (see Appendix E for details). The issue is that high-resolution models condition strongly on their low-resolution inputs and are therefore less sensitive to the text prompt. Unlike with Figure 3, high-resolution denoising with different text prompts produces images imperceptibly different to the human eye because they all agree with the same low resolution image.
We use version 1.4 of Stable diffusion for our experiments. It uses a pre-trained text encoder
from CLIP to encode the text and a pre-trained variational autoencoder to map images to a latent space.
CLIP consists of vision and text transformers trained with contrastive learning. We use the largest CLIP model (ViT-L/14@224px). We provide more details on all the models in Appendix A.
Figure 2: Comparison of efficiency improvements for Imagen on CIFAR-100. Shared noise improves sample efficiency by roughly  $100\mathrm{x}$  and pruning by an additional  $8 - 10\mathrm{x}$ .
Figure 3: Example predictions fromImagen when denoising the same image with different text prompts. Each set of images shows the original, noised, and denoised images for the two classes. The top two rows use ImageNet images and the bottom row uses Cue-Conflict.
Experiment details: For each experiment, we obtain scores using the heuristic timestep weighting function and the efficient scoring method in Algorithm 1. Due to the method's still-substantial compute cost, we use reduced-size datasets (4096 examples) for our experiments. We preprocess each dataset by performing a central crop and then resizing the images to  $64 \times 64$  resolution for Imagen,  $512 \times 512$  for SD, and  $224 \times 224$  for CLIP. We use min Scores  $= 20$ , max scores  $= 2000$ , and cutoff_pval  $= 2 \times e^{-3}$ . These are simply reasonable choices that keep the method efficient to run; changing the values does effect the model behavior in expectation but trades off compute for reduction in variance. We use a single prompt for each image. While we could have used an ensemble of prompts (i.e., made the expectation in Equation (2) also over different prompt templates), we chose not to be the sake of simplicity, as our goal is to better understand models rather than achieve state-of-the-art zero-shot performance. Therefore, our reported results are slightly lower than in the CLIP paper, which uses prompt ensembling. We found in our experiments that diffusion models were quite robust to the choice of prompt. For example, we tried four different prompts from the CLIP templates for CIFAR-100 and found accuracies to all be within  $1.5\%$  of each other.
Comparing models: Imagen, SD and CLIP have different model sizes, input resolutions, and are trained on different datasets for different amounts of time, so the comparison is not direct. While ideally we would train models of the same size on the same data, this would be very expensive and challenging in practice; we instead use these strong existing pre-trained models. Our comparisons are geared towards highlighting the strengths and weaknesses of text-image diffusion models.
Setup: We first evaluate the performance at zero-shot classification. For this purpose, we consider 13 datasets from Radford et al. (2021) as reported in Table 1. We use the prompt templates and class labels used by Radford et al. (2021), which renames some classes that confuse models (e.g. "crane  $\rightarrow$  "crane bird" in Imagenet) (OpenAI, 2021b). We use the first prompt from the list, except for Imagenet, where we use "A bad photo of a label" since this is a good prompt for Imagen, SD and CLIP (OpenAI, 2021a).
Since we use the low-resolution Imagen model, we obtain results using CLIP under two settings for a more thorough comparison. In the first setting, we resize all the datasets to  $64 \times 64$  which serves as the base low-resolution dataset. Imagen uses this dataset directly. For CLIP, we subsequently upsample the images, resizing them to  $224 \times 224$  resolution. In the second setting, we directly resize all datasets to  $224 \times 224$  resolution to obtain the best results possible using CLIP where it can take advantage of its higher input resolution.