SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mihir Prabhudesai* Tsung-Wei Ke* Alexander C. Li Deepak Pathak Katerina Fragkiadaki {mprabhud, tsungwek, acl2, dpathak, katef}@cs.cmu.edu Carnegie Mellon University
(a) ImageNet-C (online)
(b) FGVC Aircraft (single-sample)
Figure 1: Diffusion-TTA improves state-of-the-art pre-trained image classifiers and CLIP models across various benchmarks. Our model adapts pre-trained image discriminative models using feedback from pre-trained image generative diffusion models. Left: Image classification performance of pre-trained image classifiers improves after online adaptation. The image classifiers are pre-trained on ImageNet and adapted on ImageNet-C in an unsupervised manner using generative feedback. As can be seen, we get a significant boost across various model architectures. Right: Accuracy of open-vocabulary CLIP classifiers improves after single-sample adaptation, where we adapt to each unlabelled sample in the test set independently. CLIP is trained on millions of image-text pairs collected from the Internet [38], here we test it on the FGVC dataset [30].
The advancements in generative modeling, particularly the advent of diffusion models, have sparked a fundamental question: how can these models be effectively used for discriminative tasks? In this work, we find that generative models can be great test-time adapters for discriminative models. Our method, Diffusion-TTA, adapts pre-trained discriminative models such as image classifiers, segmenters and depth predictors, to each unlabelled example in the test set using generative feedback from a diffusion model. We achieve this by modulating the conditioning of the diffusion model using the output of the discriminative model. We then maximize the image likelihood objective by backpropagating the gradients to discriminative model's parameters. We show Diffusion-TTA significantly enhances the accuracy of various large-scale pre-trained discriminative models, such as, ImageNet classifiers, CLIP models, image pixel labellers and image depth predictors. Diffusion-TTA outperforms existing test-time adaptation methods, including TTT-MAE and TENT, and particularly shines in online adaptation setups, where the discriminative model is continually adapted to each example in the test set. We provide access to code, results, and visualizations on our website: diffusion-tta.github.io/.
*Equal Technical Contribution
Currently, all state-of-the-art predictive models in machine learning, that care to predict an output  $y$  from an input  $x$ , are discriminative in nature, that means, they are functions trained to map directly  $x$  to (a distribution of)  $y$ . For example, existing successful image understanding models, whether classifiers, object detectors, segmentors, or image captioners, are trained to encode an image into features relevant to the downstream task through end-to-end optimization of the end objective, and thus ignore irrelevant details. Though they shine within the training distribution, they often dramatically fail outside of it [14, 25], mainly because they learn shortcut pathways to better fit  $p(y|x)$ . Generative models on the other hand are trained for a much harder task of modeling  $p(x|y)$ . At test-time, these models can be inverted using the Bayes rule to infer the hypothesis  $\hat{y}$  that maximizes the data likelihood of the example at hand,  $p(x|\hat{y})$ . Unlike discriminative models that take an input and predict an output in a feed-forward manner, generative models work backwards, by searching over hypotheses  $y$  that fit the data  $x$ . This iterative process, also known as analysis-by-synthesis [52], is slower than feed-forward inference. However, the ability to generate leads to a richer and a more nuanced understanding of the data, and hence enhances its discriminative potential [22]. Recent prior works, such as Diffusion Classifier [27], indeed find that inverting generative models generalizes better at classifying out-of-distribution images than popular discriminative models such as ResNet [17] and ViT [11]. However, it is also found that pure discriminative methods still outperform inversion of generative methods across almost all popular benchmarks, where test sets are within the training distribution. This is not too surprising since generative models aim to solve a much more difficult problem, and they were never directly trained using the end task objective.
In this paper, instead of considering generative and discriminative models as competitive alternatives, we argue that they should be coupled during inference in a way that leverages the benefits of both, namely, the iterative reasoning of generative inversion and the better fitting ability of discriminative models. A naive way of doing so would be to ensemble them, that is  $\frac{p_{\theta}(y|x) + p_{\phi}(y|x)}{2}$ , where  $\theta$  and  $\phi$  stand for discriminative and generative model parameters, respectively. However, empirically we find that this does not result in much performance boost, as the fusion between the two models happens late, at the very end of their individual predictions. Instead, we propose to leverage generative models to adapt discriminative models at test time through iterative optimization, where both models are optimized for maximizing a sample's likelihood.
We present Diffusion-based Test Time Adaptation (TTA) (Diffusion-TTA), a method that adapts discriminative models, such as image classifiers, segmenters and depth predictors, to individual unlabelled images by using their outputs to modulate the conditioning of an image diffusion model and maximize the image diffusion likelihood. This operates as an inversion of the generative model, to infer the discriminative weights that result in the discriminative hypothesis with the highest conditional image likelihood. Our model is reminiscent of an encoder-decoder architecture, where a pre-trained discriminative model encodes the image into a hypothesis, such as an object category label, segmentation map, or depth map, which is used as conditioning to a pre-trained generative model to generate back the image. We show that Diffusion-TTA effectively adapts image classifiers for both in- and out-of-distribution examples across established benchmarks, including ImageNet and its variants, as well as image segmenters and depth predictors in ADE20K and NYU Depth dataset.
Generative models have previously been used for test time adaptation of image classifiers and segmentors, by co-training the model under a joint discriminative task loss and a self-supervised image reconstruction loss, e.g., TTT-MAE [12] or Slot-TTA [37]. At test time, the discriminative model is finetuned using the image reconstruction loss. While these models show that adaptation boosts performance, this boost often stems from the subpar performance of their initial feed-forward discriminative model. This can be clearly seen in our TTT-MAE baseline [12], where the before-adaptation results are significantly lower than a pre-trained feed-forward classifier. In contrast, our approach refrains from any joint training, and instead directly adapts pre-trained discriminative models at test time using pre-trained generative diffusion models.
We test our approach on multiple tasks, datasets and model architectures. For classification, we test pre-trained ImageNet classifiers on ImageNet [9], and its out-of-distribution variants (C, R, A, v2, S). Further we test, large-scale open-vocabulary CLIP-based classifiers on CIFAR100, Food101, FGVC, Oxford Pets, and Flowers102 datasets. For adapting ImageNet classifiers we use DiT [36] as our generative model, which is a diffusion model trained on ImageNet from scratch. For adapting open-vocabulary CLIP-based classifiers, we use Stable Diffusion [40] as our generative model. We
show consistent improvements over the initially employed classifier as shown in Figure 1. We also test on the adaptation of semantic segmentation and depth estimation tasks, where segmenters of SegFormer [49] and depth predictors of DenseDepth [1] performance are greatly improved on ADE20K and NYU Depth v2 dataset. For segmentation and depth prediction, we use conditional latent diffusion models [40] that are trained from scratch on their respective datasets. We show extensive ablations of different components of our Diffusion-TTA method, and present analyses on how diffusion generative feedback enhances discriminative models. We hope our work to stimulate research on combining discriminative encoders and generative decoder models, to better handle images outside of the training distribution. Our code and trained models are publicly available in our project's website: diffusion-tta.github.io/.
Generative models for discriminative tasks Recently there have been many works exploring the application of generative models for discriminative tasks. We group them into the following three categories: (i) Inversion-Based Methods: Given a test input  $x$  and a conditional generative model  $p_{\phi}(x|\mathbf{c})$ , these methods make a prediction by finding the conditional representation  $\mathbf{c}$  that maximizes the estimated likelihood  $p_{\phi}(x|\mathbf{c})$  of the test input. Depending on the architecture of the conditional generative model, this approach can infer the text caption [8, 27], class label [27], or semantic map [7] of the input image. For instance, Li et al. [27] showed that finding the text prompt that maximizes the likelihood of an image under a text-to-image model like Stable Diffusion [40] is a strong zero-shot classifier, and finding the class index that maximizes the likelihood of an image under a class-conditional diffusion model is a strong standard classifier. (ii) Generative Models for Data Augmentation: This category of methods involves using a generative model to enhance the training data distribution [2, 45, 51]. These methods create synthetic data using the generative model and train a discriminative classifier on the augmented training dataset. For instance, Azizi et al. [2] showed that augmenting the real training data from ImageNet with synthetic data fromImagen [41], their large-scale diffusion model, improves classification accuracy. (iii) Generative Models as feature extractors: These techniques learn features with a generative model during pretraining and then fine-tune them on downstream discriminative tasks [3, 10, 18, 35, 42, 50]. Generative pre-training is typically a strong initialization for downstream tasks. However, these methods require supervised data for model fine-tuning and cannot be directly applied in a zero-shot manner.
Our work builds upon the concept of Inversion-Based Methods (i). We propose that instead of inverting the conditioning representation, one should directly adapt a pre-trained discriminative model using the likelihood loss. This strategy enables us to effectively combine the best aspects of both generative and discriminative models. Our method is complementary to other methods of harnessing generative models, such as generative pre-training or generating additional synthetic labeled images
Test-time adaptation Test-time adaptation (TTA) also referred as unsupervised domain adaptation (UDA) in this context, is a technique that improves the accuracy of a model on the target domain by updating its parameters without using any labelled examples. Methods such as pseudo labeling and entropy minimization [46] demonstrate that model accuracy can be enhanced by using the model's own confident predictions as a form of supervision. These methods require batches of examples from the test distribution for adaptation. In contrast, TTA approaches based on self-supervised learning (SSL) have demonstrated empirical data efficiency. SSL methods jointly train using both the task and SSL loss, and solely use the SSL loss during test-time adaptation [5, 12, 16, 31, 37, 44]. These methods do not require batch of examples and can adapt to each example in the test set independently. Our contribution falls within the SSL TTA framework, where we use diffusion generative loss for adaptation. We refrain from any form of joint training and directly employ a pre-trained conditional diffusion model to perform test-time adaptation of a pre-trained discriminative model.
3 Test-Time Adaptation with Diffusion Models
The architecture of Diffusion-TTA method is shown in Figure 2 and its pseudocode in shown in Algorithm 1. We discuss relevant diffusion model preliminaries in Section 3.1 and describe Diffusion-TTA in detail in Section 3.2.
Figure 2: Architecture of Diffusion-TTA. Our method consists of discriminative and generative modules. Given an image  $x$ , the discriminative model  $f_{\theta}$  predicts task output  $y$ . The task output  $y$  is transformed into condition  $\mathbf{c}$ . Finally, we use the generative diffusion model  $\epsilon_{\phi}$  to measure the likelihood of the input image, conditioned on  $\mathbf{c}$ . This consists of using the diffusion model  $\epsilon_{\phi}$  to predict the added noise  $\epsilon$  from the noisy image  $x_{t}$  and condition  $\mathbf{c}$ . We maximize the image likelihood using the diffusion loss by updating the discriminative and generative model weights via backpropagation. For a task-specific version of this figure, please refer to Figure 7 in the Appendix.
A diffusion model learns to model a probability distribution  $p(x)$  by inverting a process that gradually adds noise to the image  $x$ . The diffusion process is associated with a variance schedule  $\{\beta_t \in (0,1)\}_{t=1}^T$ , which defines how much noise is added at each time step. The noisy version of sample  $x$  at time  $t$  can then be written  $x_t = \sqrt{\bar{\alpha}_t} x + \sqrt{1 - \bar{\alpha}_t} \epsilon$  where  $\epsilon \sim \mathcal{N}(0,1)$ , is a sample from a Gaussian distribution (with the same dimensionality as  $x$ ),  $\alpha_t = 1 - \beta_t$ , and  $\bar{\alpha}_t = \prod_{i=1}^t \alpha_i$ . One then learns a denoising neural network  $\epsilon_\phi(x_t; t)$  that takes as input the noisy image  $x_t$  and the noise level  $t$  and tries to predict the noise component  $\epsilon$ . Diffusion models can be easily extended to draw samples from a distribution  $p(x|\mathbf{c})$  conditioned on a condition  $\mathbf{c}$ , where  $\mathbf{c}$  can be an image category, image caption, semantic map, a depth map or other information [28, 40, 54]. Conditioning on  $\mathbf{c}$  can be done by adding  $\mathbf{c}$  as an additional input of the network  $\epsilon_\phi$ . Modern conditional image diffusion models are trained on large collections  $\mathcal{D} = \{(x^i, \mathbf{c}^i)\}_{i=1}^N$  of  $N$  images paired with conditionings by minimizing the loss:
$$
\mathcal {L} _ {\text {d i f f}} (\phi ; \mathcal {D}) = \frac {1}{| \mathcal {D} |} \sum_ {x ^ {i}, \mathbf {c} ^ {i} \in \mathcal {D}} \left\| \epsilon_ {\phi} \left(\sqrt {\bar {\alpha} _ {t}} x ^ {i} + \sqrt {1 - \bar {\alpha} _ {t}} \epsilon , \mathbf {c} ^ {i}, t\right) - \epsilon \right\| ^ {2}. \tag {1}
$$
3.2 Test-time Adaptation with Diffusion Models
In Diffusion-TTA, the condition  $\mathbf{c}$  of the diffusion model depends on the output of the discriminative model. Let  $f_{\theta}$  denote the discriminative model parameterized by parameters  $\theta$ , that takes as input an image  $x$  and outputs  $y = f_{\theta}(x)$ . We consider image classifiers, image pixel labellers or image depth predictors as candidates for discriminative models to adapt.
For image classifiers,  $y$  represents a probability distribution over  $L$  categories,  $y \in [0,1]^L$ ,  $y^\top \mathbf{1}_L = 1$ . Given the learnt text embeddings of a text-conditional diffusion model for the  $L$  categories  $\ell_j \in \mathbb{R}^d$ ,  $j \in \{1..L\}$ , we write the diffusion condition as  $\mathbf{c} = \sum_{j=1}^{L} y_j \cdot \ell_j$ .
For pixel labellers,  $y$  represents the set of per-pixel probability distributions over  $L$  categories,  $y = \{y^{u}\in [0,1]^{L},y^{u^{\top}}\mathbf{1}_{L} = 1,u\in x\}$ . The diffusion condition then is  $\mathbf{c} = \{\sum_{j = 1}^{L}y_{j}^{u}\cdot \ell_{j},u\in x\}$ , where  $\ell_j$  is the embedding of category  $j$ .
For depth predictors  $y$  represents a depth map  $y \in \mathbb{R}^{+w \times h}$ , where  $w, h$  the width and height of the image, and  $\mathbf{c} = y$ .
As  $\mathbf{c}$  is differentiable with respect to the discriminative model's weights  $\theta$ , we can now update them via gradient descent by minimizing the following diffusion loss:
$$
L (\theta , \phi) = \mathbb {E} _ {t, \epsilon} \| \epsilon_ {\phi} \left(\sqrt {\bar {\alpha} _ {t}} x + \sqrt {1 - \bar {\alpha} _ {t}} \epsilon , \mathbf {c}, t\right) - \epsilon \| ^ {2}. \tag {2}
$$
We minimize this loss by sampling random timesteps  $t$  coupled with random noise variables  $\epsilon$ . For online adaptation, we continually update the model weights across test images that the model encounters in a streaming manner. For single-sample adaptation, the model parameters are updated to each sample in the test set independently.
Algorithm 1 Diffusion-TTA
1: Input: Test image  $x$ , discriminative model weights  $\theta$ , diffusion model weights  $\phi$ , adaptation steps  $N$ , batch size  $B$ , learning rate  $\eta$ .
2: for adaptation step  $s \in (1, \dots, N)$  do
3: Compute current discriminative output  $y = f_{\theta}(x)$
4: Compute condition  $\mathbf{c}$  based on output  $y$
5: Sample timesteps  $\{t_i\}_{i = 1}^B$  and noises  $\{\epsilon_{i}\}_{i = 1}^{B}$
6: Loss  $L(\theta, \phi) = \frac{1}{N} \sum_{i=1}^{B} \|\epsilon_{\phi}(\sqrt{\bar{\alpha}_{t_i}} x + \sqrt{1 - \bar{\alpha}_{t_i}} \epsilon_i, \mathbf{c}, t_i) - \epsilon_i\|^2$
7: Update classifier weights  $\theta \gets \theta -\eta \nabla_{\theta}L(\theta ,\phi)$
8: Optional: update diffusion weights  $\phi \gets \phi -\eta \nabla_{\phi}L(\theta ,\phi)$
9: end for
10: return updated output  $y = f_{\theta}(x)$
Implementation Details. Our experiments reveal that minimizing the variance in the diffusion loss significantly enhances our test-time adaptation performance. We achieve this by crafting a larger batch size through randomly sampling timesteps from a uniform distribution, coupled with randomly sampling noise vectors from a standard Gaussian distribution.
We conduct our experiments on a single NVIDIA-A100 40GB VRAM GPU, with a batch size of approximately 180. Since our GPU fits only a batch size of 20, we utilize a gradient accumulation of nine steps to expand our batch size to 180. Consequently, this setup results in an adaptation time of roughly 55 seconds per example. We use Stochastic Gradient Descent (SGD) with a learning rate of 0.005 or Adam optimizer [24] with a learning rate of 0.00001. We set momentum to 0.9.
For all experiments, we adjust all the parameters of the discriminative model. We freeze the diffusion model parameters for the open-vocabulary experiments in Section 4.2 with CLIP and Stable Diffusion 2.0, and otherwise, update all diffusion model parameters. For CLIP, instead of adapting the whole network we only adapt the LoRA adapter weights [23] added to the query and value projection layers of the CLIP model. For our experiments on adapting to ImageNet distribution shift in Section 4.1, we observe a significant performance boost when adapting both the discriminative and generative model. This finding suggests that our method effectively combines the distinct knowledge encoded by these two models. We present a more detailed analyses in Section 4.4.
We test Diffusion-TTA in adapting ImageNet classifiers [11, 17, 29], CLIP models [38], image pixel labellers [49], and depth predictors [1] across multiple image classification, semantic segmentation, and depth estimation datasets, for both in-distribution and out-of-distribution test images. Our experiments aim to answer the following questions:
1. How well does Diffusion-TTA test-time adapt ImageNet and CLIP classifiers in comparison to other TTA methods under online and single-sample adaptation settings?
2. How well does Diffusion-TTA test-time adapt semantic segmentation and depth estimation models?
3. How does performance of our model vary for in- and out-of-distribution images?
4. How does performance of our model vary with varying hyperparameters, such as layers to adapt, batchsize to use, or diffusion timesteps to consider?
Evaluation metrics. We report the top-1 accuracy, mean pixel accuracy, and Structure Similarity index [48] on classification, semantic segmentation, and depth estimation benchmark datasets respectively.
4.1 Test-time Adaptation of Pre-trained ImageNet Classifiers
We use Diffusion-TTA to adapt multiple ImageNet classifiers with varying backbone architectures and sizes: ResNet-18 [17], ResNet-50, ViT-B/32 [11], and ConvNext-Tiny/Large [29]. For fair comparisons, we use Diffusion Transformer (DiT-XL/2) [36] as our class-conditional generative model, which is trained on ImageNet from scratch.
Datasets We consider the following datasets for TTA of ImageNet classifiers: ImageNet [9] and its out-of-distribution counterparts: ImageNet-C [19] (level-5 gaussian noise), ImageNet-A [21], ImageNet-R [20], ImageNetV2 [39], and Stylized ImageNet [13].
Baselines We compare our model against the following state-of-the-art TTA approaches. We use their official codebases for comparison:
- TTT-MAE [12] is a per-image test-time adaptation model trained under a joint classification and masked autoencoding loss, and test-time adapted using only the (self-supervised) masked autoencoding loss. For comparison, we use the numbers reported in the paper when possible, else we test the publicly available model.
- TENT [46] is a TTA method that adapts the batchnorm normalization parameters by minimizing the entropy (maximizing the confidence) of the classifier at test time per example.
- COTTA [47] test-time adapts a (student) classifier using pseudo labels, that are the weighted-averaged classifications across multiple image augmentations predicted by a teacher model.
ImageNet	ImageNet-A	ImageNet-R	ImageNet-C	ImageNet-V2	ImageNet-S
Customized ViT-L/16 classifier [12]	82.1	14.4	33.0	17.5	72.5	11.9
+ TTT-MAE (single-sample)	82.0 (-0.1)	21.3 (+6.9)	39.2 (+6.2)	27.5 (+10.0)	72.3 (-0.2)	20.2 (+0.3)
ResNet18	69.5	1.4	34.6	2.6	57.1	7.7
+ TENT (online)	63.0 (-6.5)	0.6 (-0.8)	34.7 (+0.1)	12.1 (+9.5)	52.0 (-5.1)	9.8 (+2.1)
+ CoTTA (online)	63.0 (-6.5)	0.7 (-0.7)	34.7 (+0.1)	11.7 (+9.1)	52.1 (-5.0)	9.7 (+2)
+ Diffusion-TTA (single-sample)	77.2 (+7.7)	6.1 (+4.7)	39.7 (+5.1)	4.5 (+1.9)	63.8 (+6.7)	12.3 (+4.6)
ViT-B/32	75.7	9.0	45.2	39.5	61.0	15.8
+ TENT (online)	75.7 (0.0)	9.0 (0.0)	45.3 (+0.1)	38.9 (-0.6)	61.1 (+0.1)	10.4 (-5.4)
+ CoTTA (online)	75.8 (+0.1)	8.6 (-0.4)	45.0 (-0.2)	40.0 (+0.5)	60.9 (-0.1)	1.1 (-14.7)
+ Diffusion-TTA (single-sample)	77.6 (+1.9)	11.2 (+2.2)	46.5 (+1.3)	41.4 (+1.9)	64.4 (+3.4)	21.3 (+5.5)
ConvNext-Tiny	81.9	22.7	47.8	16.4	70.9	20.2
+ TENT (online)	79.3 (-2.6)	10.6 (-12.1)	42.7 (-5.1)	2.7 (-13.7)	69.0 (-1.9)	19.9 (-0.3)
+ CoTTA (online)	80.5 (-1.4)	13.2 (-9.5)	47.2 (-0.6)	13.7 (-2.7)	68.9 (-2.0)	19.3 (-0.9)
+ Diffusion-TTA (single-sample)	83.1 (+1.2)	25.8 (+3.1)	49.7 (+1.9)	21.0 (+4.6)	71.5 (+0.6)	22.6 (+2.4)
Table 1: Single-sample test-time adaptation of ImageNet-trained classifiers. We observe consistent and significant performance gain across all types of classifiers and distribution drifts.
We present classification results on in-distribution (ImageNet) and out-of-distribution (ImageNet-C, R, A, V2, and S) for single-sample adaptation in Table 1. Further we present online adaptation results in Table 2. We evaluate on ImageNet-C, which imposes different types of corruption on ImageNet. For both COTTA and TENT we always report online-adaptation results as they are not applicable in the single-sample setting.