SlowGuess/NeurIPS_2023_OCR huggingface dataset:
ImageNet Corruption:	Gaussian Noise	Fog	Pixelate	Snow	Contrast
Customized ViT-L/16 classifier [12]	17.1	38.7	47.1	35.6	6.9
+ TTT-MAE (online)	37.9 (+20.8)	51.1 (+12.4)	65.7 (+18.6)	56.5 (+20.9)	10.0 (+3.1)
ResNet50	6.3	25.2	26.5	16.7	3.6
+ TENT (online)	12.3 (+6.0)	43.2 (+18.0)	41.8 (+15.3)	28.4 (+11.7)	12 (+8.4)
+ CoTTA (online)	12.2 (+5.9)	42.4 (+17.2)	41.7 (+15.2)	28.6 (+11.9)	11.9 (+8.3)
+ Diffusion-TTA (online)	19.0 (+12.7)	43.2 (+18.0)	50.2 (+23.7)	33.6 (+16.9)	2.7 (-0.9)
ViT-B/32	39.5	35.9	55.0	30.0	31.5
+ TENT (online)	38.9 (-0.6)	35.8 (-0.1)	55.5 (+0.5)	30.7 (+0.7)	32.1 (+0.6)
+ CoTTA (online)	40.0 (+0.5)	34.6 (-1.3)	54.5 (-0.5)	29.7 (-0.3)	32.0 (+0.5)
+ Diffusion-TTA (online)	46.5 (+7.0)	56.2 (+20.3)	64.7 (+9.7)	50.4 (+20.4)	33.6 (+2.1)
ConvNext-Tiny	16.4	32.3	37.2	38.3	32.0
+ TENT (online)	2.7 (-13.7)	5.0 (-27.3)	43.9 (+6.7)	15.2 (-23.1)	40.7 (-18.7)
+ CoTTA (online)	13.7 (-2.7)	29.8 (-2.5)	37.3 (+0.1)	26.6 (-11.7)	32.6 (+0.6)
+ Diffusion-TTA (online)	47.4 (+31.0)	65.9 (+33.6)	69.0 (+31.8)	62.6 (+24.3)	46.2 (+14.2)
ConvNext-Large	33.0	34.4	49.3	44.5	39.8
+ TENT (online)	30.8 (-2.2)	53.5 (+19.1)	51.1 (+1.8)	44.6 (+0.1)	52.4 (+12.6)
+ CoTTA (online)	33.3 (+0.3)	15.1 (-18.7)	34.6 (-15.3)	7.7 (-36.8)	10.7 (-29.1)
+ Diffusion-TTA (online)	54.9 (+21.9)	67.7 (+33.3)	71.7 (+22.4)	64.8 (+20.3)	55.7 (+15.9)
Table 2: Online test-time adaptation of ImageNet-trained classifiers on ImageNet-C. Our model achieves significant performance gain using online adaptation across all types of classifiers and distribution drifts.
Our conclusions are as follows:
(i) Our method consistently improves classifiers on in-distribution (ImageNet) and OOD test images across different image classifier architectures. For all ResNet, ViT, and ConvNext-Tiny, we observe significant performance gains.
(ii) Diffusion-TTA outperforms TENT and COTTA across various classifier architectures. Our results are consistent with the analysis in [55]: methods that primarily work under online settings (TENT or CoTTA) are not robust to different types of architectures and distribution shifts (see Table 4 in [55]).
(iii) Diffusion-TTA outperforms TTT-MAE even with a smaller-size classifier. ConvNext-tiny (28.5M params) optimized by Diffusion-TTA (online) achieves better performance than a much bigger custom backbone of ViT-L + ViT-B (392M params) optimized by TTT-MAE (online).
(iv) Diffusion-TTA significantly improves ConvNext-Large, in online setting. ConvNext-Large is the largest available ConvNext model with 197M parameters and achieves state-of-the-art performance on ImageNet.
4.2 Test-time Adaptation of Open-Vocabulary CLIP Classifiers
CLIP models are trained across millions of image-caption pairs using contrastive matching of the language and image feature embeddings. We test Diffusion-TTA for adapting three different CLIP models with different backbone sizes: ViT-B/32, ViT-B/16, and ViT-L/14. Experiments are conducted under single-sample adaptation setup.
We convert CLIP into a classifier following [38], where we compute the latent embedding for each text category label in the dataset using the CLIP text encoder, and then compute its similarity with the image latent to predict the classification logits across all class labels. Note that the CLIP models have not seen test images from these datasets during training.
Datasets We consider the following datasets for TTA of CLIP classifiers: ImageNet [9], CIFAR-100 [26], Food101 [6], Flowers102 [32], FGVC Aircraft [30], and Oxford-IIIT Pets [34] annotated with 1000,100,101, 102, 100, and 37 classes, respectively.
Food101	CIFAR-100	Aircraft	Oxford Pets	Flowers102	ImageNet
CLIP-ViT-B/32	82.6	61.2	17.8	82.2	66.7	57.5
+ Diffusion-TTA (Ours)	86.2 (+3.6)	62.8 (+1.6)	21.0 (+3.2)	84.9 (+2.7)	67.7 (+1.0)	60.8 (+3.3)
CLIP-ViT-B/16	88.1	69.4	22.8	85.5	69.3	62.3
+ Diffusion-TTA (Ours)	88.8 (+0.7)	69.0 (-0.4)	24.6 (+1.8)	86.1 (+0.6)	71.5 (+2.2)	63.8 (+1.5)
CLIP-ViT-L/14	93.1	79.6	32.0	91.9	78.8	70.0
+ Diffusion-TTA (Ours)	93.1 (0.0)	80.6 (+1.0)	33.4 (+1.4)	92.3 (+0.4)	79.2 (+0.4)	71.2 (+1.2)
We show test-time adaptation results in Table 3. Our method improves CLIP classifiers of different sizes consistently over all datasets, including small-scale (CIFAR-100), large-scale (ImageNet), and fine-grained (Food101, Aircraft, Pets, and Flowers102) datasets.
Table 3: Single-sample Test-time adaptation of CLIP classifiers. Our evaluation is performed across multiple model sizes and a variety of test datasets.
Corruption:	Clean	Gaussian-Noise	Fog	Frost	Snow	Contrast	Shot
Segmentation: SegFormer	66.1	65.3	63.0	58.0	55.2	65.3	63.3
+ Diffusion-TTA	66.1 (0.0)	66.4 (+1.1)	65.1 (+2.1)	58.9 (+0.9)	56.6 (+1.4)	66.4 (+1.1)	63.7 (+0.4)
Depth: DenseDepth	92.4	79.1	81.6	72.2	72.7	77.4	81.3
+ Diffusion-TTA	92.6 (+0.3)	82.1 (+3.0)	84.4 (+2.8)	73.0 (+0.8)	74.1 (+1.4)	77.4	82.1 (+0.8)
Table 4: Single-sample test-time adaptation for semantic segmentation on ADE20K and depth estimation on NYU Depth v2. We observe consistent and significant performance gain across all types of distribution drifts.
Figure 3: Semantic segmentation before and after TTA. We colorize ground-truth / predicted segmentation based on discrete class labels. Our method improves segmentation after adaptation.
4.3 Test-time Adaptation on Semantic Segmentation and Depth Estimation
We use Diffusion-TTA to adapt semantic segmentors of SegFormer [49] and depth predictors of DenseDepth [1]. We use a latent diffusion model [40] which is pre-trained on ADE20K and NYU Depth v2 dataset for the respective task. The diffusion model concatenates the segmentation/depth map with the noisy image during conditioning.
Datasets We consider ADE20K and NYU Depth v2 test sets for evaluating semantic segmentation and depth estimation. Note that both discriminative and generative diffusion models are trained on the same dataset. We evaluate with different types of image corruption under single-sample settings.
We show quantitative semantic segmentation and depth estimation results in Table 4. Further we show some qualitative semantic segmentation results before and after TTA in Figure 3. Our Diffusion-TTA improves both tasks consistently across all types of distribution drifts.
4.4 Ablations
We present ablative analysis of various design choices. We study how Diffusion-TTA varies with hyperparameters such as diffusion timesteps, number of samples per timestep and batchsize. We also study the effect of adapting different model parameters. Additionally we visualize before and after TTA adaptation results in Figure 5.
Ablation of hyperparameters In Table 5, we ablate hyperparameters of our model on ImageNet-A dataset with ConvNext as our pre-trained classifier and DiT as our diffusion model. If we perform test-time adaptation using a single randomly sampled timestep and noise latent (+diffusion TTA), we find a significant reduction in the classification accuracy  $(-2.9\%)$ . Increasing the batch size from 1 to 180 by sampling random timesteps from an uniform distribution (+ timestep aug  $BS = 180$ ) gives a significant boost in accuracy  $(+4.8\%)$ . As shown in Figure 4, correct class labels do not always have lower loss across all timesteps, and thus using only single diffusion step is a bad approximation of the likelihood. Further sampling random noise latents per timestep (+noise aug) gives an added boost of  $(+0.2\%)$ . Finally, adapting the diffusion weights of DiT (+ adapting diffusion weights) in Section 4.1, gives further  $(+1.1\%)$  boost.
Ablation of parameters/layers to adapt In Table 6, we ablate different paramters to adapt. We conduct our ablations on on ImageNet and ImageNet-R while using ResNet18 and DiT as our discriminative and generative models. We compare against following baselines: 1. we randomly initialize the classifier instead of using the pre-trained weights 2. we do not use any pre-trained classifier, instead we initialize the logits
using zeros and optimize them per example using diffusion objective (adapt logits), 3. we average the probabilities of the adapt logits baseline with the probabilities predicted by ResNet18 (adapt logits + ensemble), 4. we adapt only the batch normalization layers in ResNet18 (adapt BN), 5. we adapt only the last fully connected layer in ResNet18 (adapt last FC layer), 6. we adapt the whole ResNet18 (adapt classifier), and 7. we adapt the whole ResNet18 and DiT, which refers to our method Diffusion-TTA. Adapting the whole classifier and generative model achieves the best performance. We randomly sample one image per category for the experiment.
Figure 4: Diffusion loss of class labels at each diffusion timestep. The ground-truth class does not always have lower loss across all timesteps, compared to incorrect classes. Thus using only single diffusion timestep is a bad approximation of the likelihood.
ImageNet-A
ConvNext-Tiny [29]	22.7
+ diffusion loss TTA	19.8 (-2.9)
+ timestep aug	24.5 (+4.8)
+ noise aug	24.7 (+0.2)
+ adapting diffusion weights	25.8 (+1.1)
Table 5: Ablative analysis of Diffusion-TTA components for ConvNext classifier and DiT diffusion model. We evaluate single-sample adaptation for classification on ImageNet-A.
ImageNet	ImageNet-R
ResNet18	68.4	37.0
random init	0.7	4.0
adapt logits	71.8	31.0
adapt logits + ensemble	72.6	36.5
adapt BN	69.0	37.0
adapt last FC layer	68.6	37.5
adapt classifier	73.0	38.0
adapt classifier + adapt diffusion (Ours)	78.2	42.6
Table 6: Ablative analysis of model / layer adaptation. We evaluate single-sample adaptation for classification on ImageNet and ImageNet-R. Our Diffusion-TTA adapts both the classifier (ResNet18) and diffusion model (DiT), achieving best adaptation improvement.
4.5 Discussion - Limitations
Diffusion-TTA optimizes pre-trained discriminative and generative models using a generative diffusion loss. A trivial solution to this optimization could be the discriminative model  $f_{\theta}$  overfitting
Figure 5: Visualizing Diffusion-TTA improvement across adaptation steps. From left to right: We show input image, predicted class probabilities before adaptation (green bars indicate the ground truth category), predicted RGB image via DDIM inversion when the diffusion model is conditioned on predicted class probabilities, the predicted class probabilities after adaptation and predicted RGB image after adaptation. As can be seen, Classification ability of the classifier improves as diffusion model's ability to reconstruct the input image improves, thus indicating that there is strong correlation between the diffusion loss and the classification accuracy.
to the generative loss and thus predicting precisely  $p_{\phi}(\mathbf{c}|x)$ . Empirically we find that this does not happen: Diffusion-TTA consistently outperforms Diffusion Classifier [27], as shown in Table 7. We conjecture that the reason why the discriminative model does not overfit to the generative loss is because of the weight initialization of the (pre-trained) discriminative model, which prevents it from converging to this trivial solution. For instance, if the predictions of the generative and discriminative model are too far from one another that is: distance  $(p_{\phi}(\mathbf{c}|x), p_{\theta}(\mathbf{c}|x))$  is too high, then it becomes very difficult to optimize  $\theta$  to reach a location such that it predicts  $p_{\phi}(\mathbf{c}|x)$ . To verify our conjecture, we run a few ablations in Table 6, i) We randomly initialize  $\theta$ , ii) Instead of optimizing  $\theta$  we directly optimize the conditioning variable  $\mathbf{c}$ . In both cases, we find that results are significantly worse than when optimizing the pre-trained discriminative model.
Diffusion-TTA modulates the task conditioning of an image diffusion model with the output of the discriminative model. While this allows us to use readily available pre-trained models it also prevents us from having a tighter integration between the generative and the discriminative parts. Exploring this integration is a direct avenue of future work. Currently the generative model learns  $p(x|y)$ , where  $x$  is the input image. Training different conditional generative model for different layers such as  $p(l1|y)$  or  $p(l1|l2)$ , where  $l1$  and  $l2$  are features at different layers, could provide diverse learning signals and it is worth studying. Finally, Diffusion-TTA is significantly slow as we have to sample multiple timesteps to get a good estimate of the diffusion loss. Using Consistency Models [43] to improve the test-time-adaptation speed is a direct avenue for future work.
We introduced Diffusion-TTA, a test time adaptation method that uses generative feedback from a pre-trained diffusion model to improve pre-trained discriminative models, such as classifiers, segmenters and depth predictors. Adaptation is carried out for each example in the test-set by backpropagating diffusion likelihood gradients to the discriminative model weights. We show that our model outperforms previous state-of-the-art TTA methods, and that it is effective across multiple discriminative and generative diffusion model variants. The hypothesis of visual perception as inversion of a generative model has been pursued from early years in the field [33, 53]. The formidable performance of today's generative models and the presence of feed-forward and top-down pathways in the visual cortex [15] urges us to revisit this paradigm. We hope that our work stimulates further research efforts in this direction.
Acknowledgements We thank Shagun Uppal, Yulong Li and Shivam Duggal for helpful discussions. This work is supported by DARPA Machine Common Sense, an NSF CAREER award, an AFOSR Young Investigator Award, and ONR MURI N00014-22-1-2773. AL is supported by the NSF GRFP DGE1745016 and DGE2140739.
We present Diffusion-TTA, a test-time adaptation approach that modulates the text conditioning of a text conditional pre-trained image diffusion model to adapt pre-trained image classifiers, large-scale CLIP models, image pixel labellers, and image depth predictors to individual unlabelled images. We show improvements on multiple datasets including ImageNet and its out-of-distribution variants (C, R, A, V2, and S), CIFAR-100, Food101, FGVC, Oxford Pets, and Flowers102 over the initially employed classifiers. We also show performance gain on ADE20K and NYU Depth v2 under various distribution drifts over the pre-trained pixel labellers and depth predictors. In the Supplementary, we include further details on our work:
1. We provide comparison to Diffusion Classifier on open-Vocabulary classification in Section A.1.
2. We provide additional results with DiT on ObjectNet dataset in Section A.2.
3. We present architecture diagrams for the tasks of image classification and semantic segmentation in Section A.3
4. We provide a detailed analysis of the computational speed of Diffusion-TTA in Section A.4.
5. We detail the hyperparameters and input pre-processing in Section A.5.
6. We provide details of the datasets used for our experiments in Section A.6.
Food101	CIFAR-100	FGVC	Oxford Pets	Flowers102	ImageNet
Diffusion Classifier [27]	77.9	42.7	24.3	85.7	56.8	58.4
CLIP-ViT-L/14	93.1	79.6	32.0	91.9	78.8	70.0
+ Diffusion-TTA (Ours)	93.1 (0.0)	80.6 (+1.0)	33.4 (+1.4)	92.3 (+0.4)	79.2 (+0.4)	71.2 (+1.2)
A.1 Comparison to Diffusion Classifier on Open-Vocabulary Classification
We perform single-sample test-tiem adaptation of CLIP classifiers and compare to Diffusion Classifier [27] on open-vocabulary classification. We test on Food101, CIFAR-100, FGVC, Oxford Pets, Flowers102, and ImageNet datasets. As shown in Table 7, Diffusion-TTA consistent improves CLIP-ViT-L/14 and outperforms Diffusion Classifier over all benchmark datasets.
Table 7: Comparison to Diffusion Classifier on open-vocabulary classification. Our model outperforms Diffusion Classifier consistently over all benchmark datasets.
ObjectNet
Customized ViT-L/16 classifier [12]	37.8
+ TTT-MAE	38.9
ResNet18	23.3
+ Diffusion-TTA	26.0 (+2.7)
ViT-B/32	26.6
+ Diffusion-TTA	31.3 (+4.7)
ConvNext-Tiny	41.0
+ Diffusion-TTA	43.7 (+2.7)
Table 8: Single-sample test-time adaptation on ObjectNet. We use ImageNet pre-trained ResNet18, ViT-B/32, and ConvNext as our classifiers. Our method improves pre-trained image classifiers on out-of-distribution images of ObjectNet.
Figure 6: Comparison of image samples from ImageNet and ObjectNet. ObjectNet images have more (a) different poses, (b) complex background scenes, and (c) camera viewpoints.
A.2 Test-time Adaptation of Pre-trained ImageNet Classifiers on ObjectNet
We extend Table 1 in our paper to ObjectNet [4] dataset. ObjectNet is a variant of ImageNet, created by keeping ImageNet objects in uncommon configurations and poses as shown in Figure 6. ObjectNet has 113 common classes with ImageNet, a single class in ObjectNet can correspond to multiple classes in ImageNet. If the predicted class belongs in it's respective set of ImageNet classes, then it's considered as 'correct' if not 'incorrect'. We report the results for ObjectNet in Table 8, following the setup of Section 4.1 in the main paper. Our Diffusion-TTA consistently improves pre-trained ImageNet classifiers with different backbone architectures.
(a) Diffusion-TTA for image classification
(b) Diffusion-TTA for semantic segmentation
Figure 7: Architecture diagrams for the tasks of (a) image classification and (b) semantic segmentation. For image classification, we generate a text prompt vector (condition c) by doing a weighted average over the probability distribution predicted by the pretrained classifier. We then condition c on a pre-trained text-conditioned diffusion model via adaptive layer normalization or cross attentions. For semantic segmentation, we generate per-pixel text prompt (condition c) using a pre-trained segmentation model. We then condition c on a pre-trained segmentation conditioned Diffusion models via channel-wise concatenation. In the Figure on  $\odot$  denotes channel-wise concatenation.
A.3 Architecture Diagrams for the Tasks of Image Classification and Semantic Segmentation
We show architectural details for the tasks of image classification and semantic segmentation in Figure 7. Both tasks share the same architectural design, but differ in the way to condition on  $\mathbf{c}$ . For image classification, we generate a text prompt vector (condition  $\mathbf{c}$ ) by weighted averaging over the probability distribution predicted by the pre-trained classifier. We condition  $\mathbf{c}$  on a pre-trained text-conditioned diffusion model via adaptive layer normalization or cross attentions. For semantic segmentation, we generate per-pixel text prompt (condition  $\mathbf{c}$ ) using a pre-trained image pixel labeller. We condition  $\mathbf{c}$  on a pre-trained segmentation conditioned Diffusion models via channel-wise concatenation.
A.4 Analysis of Computation Speed
We compare the computational latency between Diffusion-TTA and Diffusion Classifier [27] on image classification. Diffusion Classifier [27] is a discrete label search method that requires a forward pass through the diffusion model for each category in the dataset. The computational time increases linearly with the number of categories. In contrast, Diffusion-TTA is a gradient-based optimization method that updates model weights to maximize the image likelihood. The computational speed depends on the number of adaptation steps and the model size, not the number of categories. As shown in Figure 8, out method is less efficient than Diffusion Classifier with fewer categories, but more efficient with more categories. Notably, discrete label search methods, like Diffusion Classifier, do not apply to online adaptation and would be infeasible for dense labeling tasks (searching over pixel labellings is exponential wrt the number of pixels in the image).
Computation compared to Diffusion Classifier. Diffusion Classifier [27] inverts a diffusion model by performing discrete optimization over categorical text prompts, instead we obtain continuous gradients to search much more effectively over a pre-trained classifier's parameter space. This design choice makes significant trade-offs in terms of computation costs. For instance, Diffusion Classifier requires to do a forward pass for each category separately and thus the computation increases linearly
Figure 8: Comparison of computation speed between Diffusion-TTA and Diffusion Classifier on image classification. We plot the required computation time with increasing numbers of image categories. Diffusion Classifier [27] requires a forward pass through the diffusion model for each category in the dataset and the computation increases linearly with the number of categories. On the other hand, our method adapts pre-trained classifiers and thus the computation is instead dependent on the throughput of the classifier. Our method becomes more computationally efficient than Diffusion Classifier as the number of categories in the dataset increases.
with the number of categories, however for us it's independent of the number of categories. We compare the per-example inference speed in Figure 8, across various datasets.
A.5 Hyper-parameters and Input Pre-Processing
For test-time-adaptation of individual images, we randomly sample 180 different pairs of noise  $\epsilon$  and timestep  $t$  for each adaptation step, composing a mini-batch of size 180. Timestep  $t$  is sampled over an uniform distribution from the range 1 to 1000 and epsilon  $\epsilon$  is sampled from an unit gaussian. We apply 5 test-time adaptation steps for each input image. We adopt Stochastic Gradient Descent (SGD) (or Adam optimizer [24]), and set learning rate, weight decay and momentum to 0.005 (or 0.00001), 0, and 0.9, respectively. To isolate the contribution of the improvement due to the diffusion model we do not use any form of data augmentation during test-time adaptation.
We use Stable Diffusion v2.0 [40] to adapt CLIP models. For the CLIP classifier, we follow their standard preprocessing step where we resize and center-crop to a resolution of  $224 \times 224$  pixels. For Stable Diffusion, we process the images by resizing them to a resolution of  $512 \times 512$  pixels which is the standard image size in Stable Diffusion models. For the CLIP text encoder in Stable Diffusion and the classifier we use the text prompt of "a photo of a <class_name>", where <class_name> is the name of the class label as mentioned in the dataset.
For the adaptation of ImageNet classifiers, we use pre-trained Diffusion Transformers (DiT) [36] specifically their XL/2 model of resolution  $256 \times 256$ , which is trained on ImageNet1K. For ImageNet classifiers we follow the standard data pre-processing pipeline where we first resize the image to  $232 \times 232$  and then do center crop to a resolution of  $224 \times 224$  pixels. For DiT we resize the image to  $256 \times 256$ , before passing it as input to their model.
A.6 Datasets
In the following section, we provide further details on the datasets used in our experiments. We adapt CLIP classifiers to the following datasets including ImageNet: 1. CIFAR-100 [26] is a compact, generic image classification dataset featuring 100 unique object classes, each with 100 corresponding test images. The resolution of these images is relatively low, standing at a size of  $32 \times 32$  pixels. 2. Food101 [6] is an image dataset for fine-grained food recognition. The dataset annotates 101 food categories and includes 250 test images for each category. 3. FGVC Airplane [30] is an image dataset for fine-grained categorization containing 100 different aircraft model variants, each with 33 or 34 test images. 4. Oxford Pets [34] includes 37 pet categories, each with roughly 100 images for
Food101
CIFAR-100
FGVC
Figure 9: Image datasets for zero-shot classification. From left to right: we show an image example obtained from Food101, CIFAR-100, FGVC, Oxford Pets, Flowers102, and ImageNet dataset. CLIP classifiers are not trained but tested on these datasets.
Oxford Pets
Flowers102
ImageNet
testing. 5. Flower102 [32] hosts 6149 test images, annotated in 102 flower categories commonly found in the United Kingdom. For all of these datasets, we sample 5 images per category for testing.