Title: Diffusion-TTA: Test-time Adaptation of Discriminative Models via Generative Feedback
Venue: NeurIPS 2023
Accepted: yes
Keywords: test-time adaptation

diffusion models

generative models

classification

segmentation

depth prediction
Abstract: 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/

Diffusion-TTA: Test-time Adaptation of Discriminative Models via Generative Feedback

Abstract

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

1 Introduction

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 ˆ y that maximizes the data likelihood of the example at hand, p ( x | ˆ 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 p θ ( y | x )+ p ϕ ( y | x ) 2 , where θ and ϕ 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 beforeadaptation 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 clasisifers 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/ .

2 Related Work

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 ϕ ( x | c ) , these methods make a prediction by finding the conditional representation c that maximizes the estimated likelihood p ϕ ( x | 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 from Imagen [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.

Pre-trained Discriminative

Model

f

θ

Input

x

Pre-trained Diffusion

Model

/uni03F5

ϕ

Noise

/u1D716

Noisy

x

Predicted

Noise

/uni03F5

Task Output

y

Condition

c

2

Actual

/u1D716

transform

Backprop

Backprop

Backprop

Diffusion Loss

↑

3.1 Diffusion Models

Adiffusion 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 { β t ∈ (0 , 1) } T t =1 , 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 = √ ¯ α t x + √ 1 -¯ α t ϵ where ϵ ∼ N ( 0 , 1 ) , is a sample from a Gaussian distribution (with the same dimensionality as x ), α t = 1 -β t , and ¯ α t = ∏ t i =1 α i . One then learns a denoising neural network ϵ ϕ ( x t ; t ) that takes as input the noisy image x t and the noise level t and tries to predict the noise component ϵ . Diffusion models can be easily extended to draw samples from a distribution p ( x | c ) conditioned on a condition c , where c can be an image category, image caption, semantic map, a depth map or other information [28, 40, 54]. Conditioning on c can be done by adding c as an additional input of the network ϵ ϕ . Modern conditional image diffusion models are trained on large collections D = { ( x i , c i ) } N i =1 of N images paired with conditionings by minimizing the loss:

3.2 Test-time Adaptation with Diffusion Models

In Diffusion-TTA, the condition c of the diffusion model depends on the output of the discriminative model. Let f θ denote the discriminative model parameterized by parameters θ , that takes as input an image x and outputs y = f θ ( 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 ∈ [0 , 1] L , y ⊤ 1 L = 1 . Given the learnt text embeddings of a text-conditional diffusion model for the L categories ℓ j ∈ R d , j ∈ { 1 ..L } , we write the diffusion condition as c = ∑ L j =1 y j · ℓ j .

For pixel labellers, y represents the set of per-pixel probability distributions over L categories, y = { y u ∈ [0 , 1] L , y u ⊤ 1 L = 1 , u ∈ x } . The diffusion condition then is c = { ∑ L j =1 y u j · ℓ j , u ∈ x } , where ℓ j is the embedding of category j .

For depth predictors y represents a depth map y ∈ R + w × h , where w,h the width and height of the image, and c = y .

As c is differentiable with respect to the discriminative model's weights θ , we can now update them via gradient descent by minimizing the following diffusion loss:

We minimize this loss by sampling random timesteps t coupled with random noise variables ϵ . 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 θ , diffusion model weights ϕ , adaptation steps N , batch size B , learning rate η .

2: for adaptation step s ∈ (1 , . . . , N ) do

3: Compute current discriminative output y = f θ ( x )

4: Compute condition c based on output y

5: Sample timesteps { t i } B i =1 and noises { ϵ i } B i =1

7: Update classifier weights θ ← θ -η ∇ θ L ( θ, ϕ )

6: Loss L ( θ, ϕ ) = 1 N ∑ B i =1 ∥ ϵ ϕ ( √ ¯ α t i x + √ 1 -¯ α t i ϵ i , c , t i ) -ϵ i ∥ 2

8: Optional: update diffusion weights ϕ ← ϕ -η ∇ ϕ L ( θ, ϕ )

9: end for

10: return updated output y = f θ ( 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.

4 Experiments

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:

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?

How well does Diffusion-TTA test-time adapt semantic segmentation and depth estimation models?

How does performance of our model vary for in- and out-of-distribution images?

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 comparisions, 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 comparision:

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 weightedaveraged classifications across multiple image augmentations predicted by a teacher model.

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.

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], CIFAR100 [26], Food101 [6], Flowers102 [32], FGVC Aircraft [30], and Oxford-IIIT Pets [34] annotated with 1000 , 100 , 101 , 102 , 100 , and 37 classes, respectively.

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.

Image

ground-truth

Before-TTA

After-TTA

STAR

nerb

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.

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 θ overfitting

Ground truth RGB

Before-TTA RGB Prediction

After-TTA top-3 classification

Before-TTA top-3 classification

After-TTA RGB Prediction

1.0

1.0

0.8

0.8

Probability

0.6

Probability

0.6

0.4

0.4

0.2

0.2

0.0

sea lion

drake

aso05

0.0

sea lion

drake

goose

1.0

1.0

0.8

0.8

0.6

0.4

0.4

0.2

0.2

0.0

0.0

burrito hippopotamus snow leopard

burrito hippopotamus snow leopard

to the generative loss and thus predicting precisely p ϕ ( 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 ϕ ( c | x ) , p θ ( c | x )) is too high, then it becomes very difficult to optimize θ to reach a location such that it predicts p ϕ ( c | x ) . To verify our conjecture, we run a few ablations in Table 6, i) We randomly initialize θ , ii) Instead of optimizing θ we directly optimize the conditioning variable 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 ( l 1 | y ) or p ( l 1 | l 2) , where l 1 and l 2 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.

5 Conclusion

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

References

I. Alhashim and P. Wonka. High quality monocular depth estimation via transfer learning. arXiv e-prints , abs/1812.11941:arXiv:1812.11941, 2018. URL https://arxiv.org/abs/1812. 11941 .

S. Azizi, S. Kornblith, C. Saharia, M. Norouzi, and D. J. Fleet. Synthetic data from diffusion models improves imagenet classification. arXiv preprint arXiv:2304.08466 , 2023.

D. Baranchuk, I. Rubachev, A. Voynov, V. Khrulkov, and A. Babenko. Label-efficient semantic segmentation with diffusion models. arXiv preprint arXiv:2112.03126 , 2021.

A. Barbu, D. Mayo, J. Alverio, W. Luo, C. Wang, D. Gutfreund, J. Tenenbaum, and B. Katz. Objectnet: A large-scale bias-controlled dataset for pushing the limits of object recognition models. Advances in neural information processing systems , 32, 2019.

A. Bartler, A. Bühler, F. Wiewel, M. Döbler, and B. Yang. Mt3: Meta test-time training for self-supervised test-time adaption. In International Conference on Artificial Intelligence and Statistics , pages 3080-3090. PMLR, 2022.

L. Bossard, M. Guillaumin, and L. Van Gool. Food-101 - mining discriminative components with random forests. In European Conference on Computer Vision , 2014.

R. Burgert, K. Ranasinghe, X. Li, and M. S. Ryoo. Peekaboo: Text to image diffusion models are zero-shot segmentors. arXiv preprint arXiv:2211.13224 , 2022.

K. Clark and P. Jaini. Text-to-image diffusion models are zero-shot classifiers. arXiv preprint arXiv:2303.15233 , 2023.

J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee, 2009.

J. Donahue, P. Krähenbühl, and T. Darrell. Adversarial feature learning. arXiv preprint arXiv:1605.09782 , 2016.

A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 , 2020.

Y. Gandelsman, Y. Sun, X. Chen, and A. Efros. Test-time training with masked autoencoders. Advances in Neural Information Processing Systems , 35:29374-29385, 2022.

R. Geirhos, P. Rubisch, C. Michaelis, M. Bethge, F. A. Wichmann, and W. Brendel. Imagenettrained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness. In International Conference on Learning Representations , 2019. URL https: //openreview.net/forum?id=Bygh9j09KX .

R. Geirhos, J.-H. Jacobsen, C. Michaelis, R. Zemel, W. Brendel, M. Bethge, and F. A. Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence , 2(11):665-673, 2020.

C. D. Gilbert and M. Sigman. Brain states: Top-down influences in sensory processing. Neuron , 54(5):677-696, 2007. ISSN 0896-6273. doi: https://doi.org/10.1016/j.neuron.2007.05.019. URL https://www.sciencedirect.com/science/article/pii/S0896627307003765 .

J.-B. Grill, F. Strub, F. Altché, C. Tallec, P. Richemond, E. Buchatskaya, C. Doersch, B. Avila Pires, Z. Guo, M. Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems , 33: 21271-21284, 2020.

K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770778, 2016.

K. He, X. Chen, S. Xie, Y. Li, P. Dollár, and R. Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 16000-16009, 2022.

D. Hendrycks and T. Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. arXiv preprint arXiv:1903.12261 , 2019.

D. Hendrycks, S. Basart, N. Mu, S. Kadavath, F. Wang, E. Dorundo, R. Desai, T. Zhu, S. Parajuli, M. Guo, et al. The many faces of robustness: A critical analysis of out-of-distribution generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 8340-8349, 2021.

D. Hendrycks, K. Zhao, S. Basart, J. Steinhardt, and D. Song. Natural adversarial examples. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 15262-15271, 2021.

G. E. Hinton. To recognize shapes, first learn to generate images. Progress in brain research , 165:535-547, 2007.

E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685 , 2021.

D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

P. W. Koh, S. Sagawa, H. Marklund, S. M. Xie, M. Zhang, A. Balsubramani, W. Hu, M. Yasunaga, R. L. Phillips, S. Beery, J. Leskovec, A. Kundaje, E. Pierson, S. Levine, C. Finn, and P. Liang. WILDS: A benchmark of in-the-wild distribution shifts. CoRR , abs/2012.07421, 2020. URL https://arxiv.org/abs/2012.07421 .

A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images. 2009.

A. C. Li, M. Prabhudesai, S. Duggal, E. Brown, and D. Pathak. Your diffusion model is secretly a zero-shot classifier. arXiv preprint arXiv:2303.16203 , 2023.

Y. Li, H. Liu, Q. Wu, F. Mu, J. Yang, J. Gao, C. Li, and Y. J. Lee. Gligen: Open-set grounded text-to-image generation. arXiv preprint arXiv:2301.07093 , 2023.

Z. Liu, H. Mao, C.-Y. Wu, C. Feichtenhofer, T. Darrell, and S. Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11976-11986, 2022.

S. Maji, J. Kannala, E. Rahtu, M. Blaschko, and A. Vedaldi. Fine-grained visual classification of aircraft. Technical report, 2013.

S. Manli, N. Weili, H. De-An, Y. Zhiding, G. Tom, A. Anima, and X. Chaowei. Test-time prompt tuning for zero-shot generalization in vision-language models. In NeurIPS , 2022.

M.-E. Nilsback and A. Zisserman. Automated flower classification over a large number of classes. In 2008 Sixth Indian Conference on Computer Vision, Graphics & Image Processing , pages 722-729. IEEE, 2008.

B. A. Olshausen. Perception as an inference problem. 2013. URL https://api. semanticscholar.org/CorpusID:1194136 .

O. M. Parkhi, A. Vedaldi, A. Zisserman, and C. V. Jawahar. Cats and dogs. In IEEE Conference on Computer Vision and Pattern Recognition , 2012.

D. Pathak, P. Krahenbuhl, J. Donahue, T. Darrell, and A. A. Efros. Context encoders: Feature learning by inpainting. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 2536-2544, 2016.

W. Peebles and S. Xie. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09748 , 2022.

M. Prabhudesai, A. Goyal, S. Paul, S. Van Steenkiste, M. S. Sajjadi, G. Aggarwal, T. Kipf, D. Pathak, and K. Fragkiadaki. Test-time adaptation with slot-centric models. In International Conference on Machine Learning , pages 28151-28166. PMLR, 2023.

A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning , pages 8748-8763. PMLR, 2021.

B. Recht, R. Roelofs, L. Schmidt, and V. Shankar. Do imagenet classifiers generalize to imagenet? In International conference on machine learning , pages 5389-5400. PMLR, 2019.

R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10684-10695, 2022.

C. Saharia, W. Chan, S. Saxena, L. Li, J. Whang, E. L. Denton, K. Ghasemipour, R. Gontijo Lopes, B. Karagol Ayan, T. Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems , 35: 36479-36494, 2022.

M. Singh, Q. Duval, K. V. Alwala, H. Fan, V. Aggarwal, A. Adcock, A. Joulin, P. Dollár, C. Feichtenhofer, R. Girshick, et al. The effectiveness of mae pre-pretraining for billion-scale pretraining. arXiv preprint arXiv:2303.13496 , 2023.

Y. Song, P. Dhariwal, M. Chen, and I. Sutskever. Consistency models. 2023.

Y. Sun, X. Wang, Z. Liu, J. Miller, A. Efros, and M. Hardt. Test-time training with selfsupervision for generalization under distribution shifts. In International conference on machine learning , pages 9229-9248. PMLR, 2020.

B. Trabucco, K. Doherty, M. Gurinas, and R. Salakhutdinov. Effective data augmentation with diffusion models. arXiv preprint arXiv:2302.07944 , 2023.

D. Wang, E. Shelhamer, S. Liu, B. Olshausen, and T. Darrell. Tent: Fully test-time adaptation by entropy minimization. arXiv preprint arXiv:2006.10726 , 2020.

Q. Wang, O. Fink, L. Van Gool, and D. Dai. Continual test-time domain adaptation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 7201-7211, 2022.

Z. Wang, A. C. Bovik, H. R. Sheikh, and E. P. Simoncelli. Image quality assessment: from error visibility to structural similarity. IEEE transactions on image processing , 13(4):600-612, 2004.

E. Xie, W. Wang, Z. Yu, A. Anandkumar, J. M. Alvarez, and P. Luo. Segformer: Simple and efficient design for semantic segmentation with transformers. Advances in Neural Information Processing Systems , 34:12077-12090, 2021.

J. Xu, S. Liu, A. Vahdat, W. Byeon, X. Wang, and S. De Mello. Open-vocabulary panoptic segmentation with text-to-image diffusion models. arXiv preprint arXiv:2303.04803 , 2023.

T. Yu, T. Xiao, A. Stone, J. Tompson, A. Brohan, S. Wang, J. Singh, C. Tan, J. Peralta, B. Ichter, et al. Scaling robot learning with semantically imagined experience. arXiv preprint arXiv:2302.11550 , 2023.

A. Yuille and D. Kersten. Vision as bayesian inference: analysis by synthesis? Trends in cognitive sciences , 10:301-8, 08 2006. doi: 10.1016/j.tics.2006.05.002.

A. Yuille and D. Kersten. Vision as Bayesian inference: analysis by synthesis? Trends in Cognitive Sciences , 10:301-308, 2006.

L. Zhang and M. Agrawala. Adding conditional control to text-to-image diffusion models. arXiv preprint arXiv:2302.05543 , 2023.

H. Zhao, Y. Liu, A. Alahi, and T. Lin. On pitfalls of test-time adaptation. arXiv preprint arXiv:2306.03536 , 2023.

A Appendix

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:

We provide comparison to Diffusion Classifier on open-Vocabulary classification in Section A.1.

We provide additional results with DiT on ObjectNet dataset in Section A.2.

We present architecture diagrams for the tasks of image classification and semantic segmentation in Section A.3

We provide a detailed analysis of the computational speed of Diffusion-TTA in Section A.4.

We detail the hyperparameters and input pre-processing in Section A.5.

We provide details of the datasets used for our experiments in Section A.6.

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.

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)

ImageNet

ObjectNet

(a) pose

(b) scene

(c) viewpoint

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.

Po

P1

2

P3

口

P4

Rabeit

Cat

puedeoTJem

Eagle

Clean Image

Classifier fe

Pretrained

Rabbit

0

Text Embeddings

Condition c

Diffusion Loss

Noised Image

Pretrained Diffusion

Model

Actual

个

Po

P2

P3

P4

Ceoch

Chair

Laop

PNRL

Toee

Clean Image

Segmenter fe

Pretrained

Classification

Per-Pixel

ono

due"

lqe

Text Embeddings

Condition c

Diffusion Loss

Noise E

Noised Image

Pretrained Diffusion

Model E

Actual

_个

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 c . For image classification, we generate a text prompt vector (condition c ) by weighted averaging over the probability distribution predicted by the pre-trained classifier. We 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 image pixel labeller. We condition c on a pre-trained segementation 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 seperately and thus the computation increases linearly

Inference time per Image (in Seconds)

1.0

100.0

10000.0

CIFAR10

Pets

Food101

ImageNetA

ImageNet-1K

Di

ff

Ours

Linear wrt Number of Categories

Constant wrt Number of Categories

(10 Categories)

(37 Categories)

(101 Categories)

(200 Categories)

(1000 Categories)

(21,000 Categories)

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 ϵ 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 ϵ 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 × 224 pixels. For Stable Diffusion, we process the images by resizing them to a resolution of 512 × 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 × 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 × 232 and then do center crop to a resolution of 224 × 224 pixels. For DiT we resize the image to 256 × 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 incuding 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 × 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

Oxford Pets

Flowers102

ImageNet

ImageNet

ImageNet-C

ImageNet-A

ImageNet-R

ImageNetV2

ImageNet-S

We consider ImageNet and its out-of-distribution variants for test-time adaptation. 1. ImageNet [9] is a large-scale generic object classification dataset, featuring 1000 object categories. We test on the validation set that consists of 50 images for each category. 2. ImageNet-C [19] applies various visual corruptions to the same validation set as ImageNet. We consider the shift due to gaussian noise (level-5) in our single-sample experiments. We consider 5 types of shift (gaussian noise, fog, pixelate, contrast, and snow) in our online adaptation experiments. 3. ImageNet-A [21] consists of real-world image examples that are misclassified by existing classifiers. The dataset convers 200 categories in ImageNet and 7450 images for testing. 4. ImageNet-R [20] is a synthetic dataset that renders 200 ImageNet classes in art , cartoons , deviantart , graffiti , embroidery , graphics , origami , paintings , patterns , plastic objects , plush objects , sculptures , sketches , tattoos , toys , and video game styles. The dataset consists of 30 K test images. 5. ImageNetV2 collects images from the same distribution as ImageNet, composed of 1000 categories, each with 10 test images. 6. ImageNet-S transforms ImageNet images into different artistic styles, while preserving the global contents of original images. For ImageNet and its C, R, V2, and S variants except A, we sample 3 images per category, as these datasets contains a lot more classes than the other datasets. For ImageNet-A we evaluate on it's whole test-set.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose combining a diffusion model with a discriminative model for test-time adaptation: They propose adapt a pretrained classifier using a diffusion model with its likelihood loss. They show results on several OOD robustness benchmarks, commonly used in TTA.\n\nSTRENGTHS:\nI like the idea of combining a generative model with a discriminative model to try to get the benefits of both. The proposed method is elegant, but unfortunately heavy on the compute side.\n\nThe paper is well-written and easy to follow. The approach is clear. \n\nIn general, I think the method is interesting and the idea looks novel to me. But I do think there are several major issues with the paper which need to be resolved. I am happy to engage in a discussion with the authors and will raise my rating if my points are addressed.\n\nWEAKNESSES:\nWeaknesses:\nMy biggest issue is with the experimental results. I find them impossible to judge (see below) and thus, do not understand whether this method works similarly or better than other methods. Given that the proposed method is very heavy on the compute side, I think it would need to make up for it in much stronger results. If e.g. this method performs similarly to TENT / COTTA, then people would still use TENT / COTTA because those methods are very cheap to use. Given that this is an empirical and not a theory paper, I also think more results / ablations should be presented.\n\n### Issues with the results section\nTable 2: 1) It is not clear which baseline models were used for TENT / COTTA / TTT here. The final numbers heavily depend on the chosen architecture and the results are only comparable if the same architecture was chosen. Please indicate the architecture for the baseline methods both in the text and in the Figure captions. 2) It is not clear whether accuracy or error are displayed in Table2. Please indicate this, and also please add \\uparrow and \\downarrow to indicate whether a higher or a lower number is better. An accuracy (if accuracy is displayed?) of 11% on ImageNet-C looks quite bad and I am wondering where this number comes from? The TENT authors report an accuracy of 44% for a ResNet50 in their paper. I think it would be necessary for the authors to add results for a ResNet50 such that the numbers are comparable.\n\nRight now, the results look strange. If we assume that the architecture is the same, then it seems TENT strongly degrades performance. E.g. on ImageNet-R, the TENT number is 24.3 for TENT, 34.6 for the baseline RN18 and 39.7 for Diffusion-TTA. Does this mean that there is an optimization issue because TENT should not degrade performance or is it a different architecture? In that case, the numbers are not comparable. Please rework Table 2 as following: Show the baseline performance without adaptation in the first row, then report numbers from the literature for this architecture, then report your results on the same architecture.\n\n\n### General comments\nLine 36: “In this paper, we take an alternative perspective. Instead of considering generative and discriminative models as competitive, we argue that they should be coupled in a way that leverages the best of both worlds: discriminative models are good at building powerful conditional density models but overfit to training distribution, and generative models generalize better but struggle to learn discriminative features.” -> minor: to the training distribution. The statement that “generative models generalize better” seems unsupported to me. As far as I see, [18] does not provide evidence for better OOD generalization for generative models. Further Carlini et al. [A] show that diffusion models actually memorize their training data and emit the memorized data at test time. Could the authors comment on how they think the findings presented in [A] square with the results in this paper? If diffusion models memorize their training data, I find it unintuitive that they would learn features that would allow them to generalize to OOD data. Put simply, if the task can be solved with memorization, there is no need to learn generalizable features.\n\n[A] Carlini et al. “Extracting Training Data from Diffusion Models”, https://arxiv.org/pdf/2301.13188.pdf\n\nIn my opinion, the paper lacks intuition and needs a better explanation why the method works. As written above, it is not evident to me that generative models should have better generalizable features compared to discriminative models. I am wondering whether the approach works because of an effect similar to model ensembling where we know that combining predictions of multiple models leads to better results than using a single model. I wonder if using a diffusion model has a similar effect, because it regularizes the weights to be compatible with both the discriminative and the generative model. Could the authors comment on this thought / better explain why their method works?\n\nDo the authors think one could train a discriminative network from scratch using their approach? If not, why do they think it (only?) helps in the test time adaptation stage? Is it maybe that one finds the “correct” loss basin with regular training and then finds a better loss value within that basin with test time adaptation? I am thinking of the paper by Wortsman et al. [B] where the authors showed that one can interpolate weights of different models because they are in the same low error basin. Could the authors maybe compare the training and test losses (regular cross-entropy) of the baseline and test-adapted models when interpolating the weights from the regular to the adapted model? I am wondering whether the loss goes down monotonically.\n\n [B] Wortsman et al. “Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time”, https://arxiv.org/abs/2203.05482\n\n\nThe text in Figure 5 is too small below the bars and largely unreadable.\n\nFigure 6: It would be helpful to repeat this analysis on some other dataset. The solo peak at K=5 somewhat looks like an artifact and I wonder whether one would observe similar behavior in other datasets. How do the authors interpret this strange “peakiness” of the curve at K=5?\n\nMore discussion on the ablation analysis would be nice. Why is it that “using a single randomly sampled timestep and noise latent (+diffusion TTA), we find a significant reduction in the classification accuracy (−2.9%)”?\n\t\nI think the authors should perform a complexity analysis since their method involves making a forward and a backward pass through a large diffusion model to make use of its loss. This is a large overhead compared to much simpler TTA methods such as TENT. While there is a complexity analysis in the Appendix, Fig.2, the overhead is not compared to e.g. TENT which should be much faster.\n\nIn a similar vein, I would argue it will be hard to actually find an application for this method since the computational requirements are so massive.  The authors write that “We conduct our experiments on a single NVIDIA-A100 40GB VRAM GPU (…). Since our GPU fits only a batch size of 20, (…)”. TTA methods are usually designed to be light-weight such that they can be applied on the fly, e.g. if an autonomous car needs to adapt to changing weather conditions. This method does not fulfil the requirements usually considered for TTA, and thus, will be of rather limited use for practitioners.\n\nLine 173: “For all experiments, we adjust all the parameters of the classifier.” I am curious about this choice. Several papers show that at least for ResNets, it is better to do TTA with affine BN layers [C, D]. Does this method not work when adapting affine BN layers or does it work better if all layers are adapted?\n\n[C] Wang et al. “Tent: Fully Test-time Adaptation by Entropy Minimization”\n\n[D] Rusak et al. “If your data distribution shifts, use self-learning” \n\nThe appendix should be referenced in the main text.\n\nQUESTIONS:\nPlease respond / fix the points I outlined above.\n\nPlease fix Table 2 to make baselines comparable to your results. Please add a RN50 baseline since it is common to report TTA results on this architecture.\n\nI would advise the authors to do a round of proof-reading of the manuscript, as there are some grammar errors and typos, e.g. line 234, line 123: adaptation\n\nLIMITATIONS:\nThe authors have not included a Broader Impact section and did not discuss the potential negative societal impact of their work. I think they could discuss how the chosen diffusion model with affect the adapted classifier. Practitioners should note that biases present in the diffusion model will likely also manifest in the adapted classifier. The classifier itself may or may not have biases of its own which in turn may or may not be corrected by the diffusion model.\n\nOne limitation of the method is that it is very compute-intense and thus will likely be of limited use to practitioners.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper discusses a method that uses generative models as test-time adapters for discriminative models. The authors propose a technique to adapt pre-trained classifiers and CLIP models to individual unlabeled images. They achieve this by modifying the text conditioning of a text-conditional pretrained image diffusion model and maximizing the image likelihood through backpropagation. The proposed approach improves classification accuracy on various datasets, including ImageNet. The authors compare their method with previous test-time adaptation techniques and demonstrate its superior performance. They highlight that this is the first work to adapt large-scale pre-trained discriminative models to individual images without requiring joint discriminative and self-supervised training objectives.\n\nSTRENGTHS:\n* The paper addresses a fundamental challenge in deep learning models, which is test-time out-of-distribution generalization. This problem is known to be difficult to solve. \n\n* The authors propose an interesting approach that combines generative and discriminative models to tackle this issue. Particularly intriguing is their idea of using a diffusion model as a guide for the classifier, by modulating the text conditioning of the model. \n\n* The experimental evaluation conducted in the paper is good, demonstrating the effectiveness of the proposed method. The authors thoroughly evaluate their approach on various datasets, including conducting ablations to analyze the discriminative and generative components. \n\nOverall, the paper provides a valuable contribution to the field with its innovative approach and robust experimental evaluation.\n\nWEAKNESSES:\n* While I agree with your premises regarding the class conditional model, I noticed that the improvements achieved are marginal on certain datasets, in particular for open-set problems. It would be beneficial to have a more in-depth analysis specifically focusing on ImageNet for open-set evaluation.\n\n* Additionally, it would be valuable to incorporate confidence intervals in the evaluation to provide a better understanding of the statistical significance of the results. Furthermore, I believe it would be advantageous to include additional metrics such as top-k accuracy and f-score (where possible), rather than solely relying on accuracy. These metrics would provide a more comprehensive evaluation of the model's performance in various scenarios. \n\n* It is worth mentioning that CLIP, even without any adaptation, seems to perform well, and the extent of improvement achieved by the proposed method is not decisively evident in most cases, especially without the inclusion of confidence intervals.\n\n* Finally, the clarity of the introduction can be improved. I think a better discussion of discriminative vs generative models would be a valuable addition to this paper.\n\nQUESTIONS:\n* How are the learned class or text embeddings obtained, specifically referring to line 156 or Fig.1? Are the embeddings the same as those employed by the classifier or CLIP models? If so, does this imply that the class/text embeddings change during test time adaptation?\n\n* What is the time required to adapt the model to a single image? Can this approach be scaled up for practical adaptation applications?\n\n* In Figure 6, have you investigated the effects of increasing the value of K for a CLIP model?\n\nLIMITATIONS:\nThe obtained results are interesting; however, it is worth noting that the observed improvements, particularly in relation to the CLIP models, appear relatively small. It is interesting that the base classifiers, in particular the CLIP models, already exhibit good performance even without any fine-tuning or test-time adaptation.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a test-time adaptation (TTA) method that adjusts pretrained image classifiers at test time by leveraging the power of pretrained, text-conditioned generative diffusion models. Given an unseen image, the proposed method Diffusion-TTA minimizes the diffusion loss with respect to the weights of the pretrained classifier and/or diffusion model via gradient descent at test time. In other words, Diffusion-TTA provides guidance to the classifier to adapt its performance particularly under the out-of-distribution setting.\n\nFrom the empirical evaluation, Diffusion-TTA shows its effectiveness in encouraging consistent improvements over the initially employed classifier. Furthermore, the method is applicable to utilize different neural net architectures for the classifiers such as ResNet, ViT, and ConvNext-Tiny.\n\n----\nAfter reading the rebuttal, I've raised my score from 6: weak accept to 7: accept.\n\nSTRENGTHS:\nSufficient novelty:\nDiffusion-TTA follows a similar approach from Diffusion Classifier (Li et al. 2023) by utilizing the same objective, i.e., minimizing the diffusion loss, but executing a different optimization approach towards the objective. Diffusion-TTA employs pretrained classifiers and updates the weights via back-propagation of the diffusion loss, while Diffusion Classifier performs a discrete optimization over categorical text prompts to the generative model without using any pretrained classifier. This approach undertaken by Diffusion-TTA, which results in a positive empirical outcome in the context of TTA, is a new perspective to me.\n\nRigorous ablation study: \nThe impact of each implementation component defining Diffusion-TTA is empirically investigated and reported so that the readers are aware about the contribution from each component.\n\nWEAKNESSES:\nEmpirical evaluation: \nI think that the added values of Diffusion-TTA over the existing SOTA methods remain a bit inconclusive. Specifically, it’s unclear to me that Diffusion-TTA is generally better than TTA-MAE given two occasions in ImageNet experiments (in-distribution and ImageNet-V2) where TTA-MAE seems to provide slightly or statistically insignificant negative outcomes. On the other hand, TTA-MAE produces larger improvements than Diffusion-TTA.\n\nI’d suggest having comparisons in which the other existing methods use the same backbone networks as Diffusion-TTA for the classifiers (ResNet18, ViT-B/32, or ConvNext-Tiny) whenever applicable.\n\nQUESTIONS:\nHow’s the performance of TTA-MAE or other existing TTA methods on the Open-Vocabulary CLIP experiment?\n\nHave the authors tried partial weight updating on the classifier, e.g., only adjusting weights of a few top layers, and observing the performance gain?\n\nIn L175, it’s mentioned that adapting the weights of both generator and classifier slightly boosts the performance of Diffusion-TTA for the ImageNet distribution shift case. By how much boost in comparison to keeping the weights of the generator fixed?\n\nLIMITATIONS:\nThe limitation in terms of the generality of the method, i.e., only shown to be effective for image classification thus far, has been addressed. Furthermore, although the inference speed is perhaps not a focus in this study, I’d recommend also addressing it (~55 seconds per example on a single NVIDIA-A100 40GB) as another limitation.\n\nI’m hoping that the code implementation is made available soon.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAuthors aim to improve test-time adaptation of CLIP and ImageNet trained models on different distribution data. They leverage text conditioned image diffusion model to adapt image classifier’s parameters by maximizing image diffusion likelihood. In particular, image classifier’s output probabilities are used as weights for class conditioning the diffusion model and diffusion likelihood gradients are backpropagated to classifier through the probabilities. Adaptation of CLIP across different datasets and adaptation of ImageNet trained models across different ImageNet variants using the diffusion loss show improved top-1 accuracy.\n\nSTRENGTHS:\n1.\tPaper is well-written and easy to understand.\n\n2.\tProposed setup is a simple plug and play method that does not require any additional training or warm-up of the models.\n\n3.\tResults demonstrate that the proposed adaptation improve the results, particularly benefit smaller models.\n\n4.\tAppreciate the analysis that shown diffusion loss and cross-entropy loss are correlated, which support the gains observed in the results.\n\nWEAKNESSES:\n1.\tProposed setup can be seen as straight forward extension of Li et al. [22] which introduced similar guidance strategy to optimize over text prompts. Here, classifier parameters are updated instead of text prompts.\n\n2.\tIncreased time complexity of test time adaptation with the diffusion model, where each sample takes 55 seconds long for adaptation.\n\n3.\tIncreased model complexity by including diffusion model with large number of parameters in the adaptation process.\n\n4.\tThe adaptation is not as effective on larger models when compared to smaller models.\n\n5.\tPrior works comparisons are included in Table 2 but the network architectures are different across these comparisons and wouldn't be a fair comparison.\n\nQUESTIONS:\n1.\tWhich diffusion model was used in Sec. 4.1?\n\n2.\tCLIP adaptation is shown across multiple datasets. Adapting CLIP with ImageNet variants would be interesting.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposed Diffusion-TTA, a test-time adaptation method that, given a test image, updates the classifier’s weights to maximize the image likelihood (i.e., the denoising score matching loss) using the pre-trained diffusion models. Specifically, the classifier uses the clean input image to predict a distribution over labels, which is then used for a weighted sum of learnt class embeddings into a prompt input of diffusion models. In experiments, Diffusion-TTA is tested for adapting both open-vocabulary CLIP models and pretrained image classifiers, where it improves the classification accuracy across different network architectures.\n\nSTRENGTHS:\n1. The idea of using pre-trained diffusion models as image priors to update classifiers at test time is novel. In particular, I like the idea of “using classifier’s output logits for a weighted sum of learnt class embeddings”, which is very intuitive to me. \n2. It shows the improvement of classification accuracy across different network architectures after using Diffusion-TTA.\n3. The ablation studies (in Table 3) are well conducted to show the importance of each component.\n\nWEAKNESSES:\n1. One of my major concerns is about the comparison with baselines in experiments. In Section 4.1 (“adapting open-vocabulary CLIP models”), only Diffusion-TTA works with CLIP, other baselines do not use CLIP. For example, Diffusion Classifier does not use any image classifier, and both Synthetic SD Data and SD Features use the ResNet-50 classifier. I think this setting causes an unfair comparison with baselines. In Section 4.2 (“adapting image classifiers”), I also don’t see what the pretrained image classifiers the baselines (COTTA, TENT and TTT-MAE) have used. It would be good to report their accuracies on the same set of network architectures (ResNet-18, ViT-B/32 and ConvNext-Tiny). In the current version, I’m not sure if the proposed method is better than the baselines under the same conditions. \n2. The claim “this is the first work that adapts pre-trained large-scale discriminative models to individual images” is wrong. From what I know, at least TPT [1] is a prior test-time adaptation method using only a single test image. I think TPT is more efficient than the proposed method in adapting CLIP models, since it only updates the input prompts instead of the CLIP weights, and it does not need to backpropagate through large diffusion models with a large batch size. \n3. In Figure 6, it looks weird to me that the accuracy curve has a spike at K=5 and does not change after K>=10. Also, a smaller K (maybe K=1?) achieves very similar accuracy with K=5, it seems that we may not need a distribution of output to do a weighted sum. Instead, we can just use the class embedding from the top-1 prediction as the prompt input.\n4. Some suggestions regarding the writing presentations: 1) In line 155 and line 156, the variable format is not consistent, such as $z_i$ and $i \\in T$. 2) Since the baselines are already discussed in the related work section, there is no need to give a full detailed description of them in experiment section again (line 234-250). They can be moved to Appendix if necessary. 3) The fonts in Figure 5 is too small.\n[1] Test-Time Prompt Tuning for Zero-Shot Generalization in Vision-Language Models, NeurIPS 2022.\n\nQUESTIONS:\n1. Can the authors provide results of Synthetic SD Data and SD Features with other image classifiers, such as CLIP-ViT models, to make the comparison more sound? Similarly, can the authors provide results of COTTA, TENT and TTT-MAE across different network architectures (ResNet-18, ViT-B/32 and ConvNext-Tiny)?\n2. I suggest the authors add a comparison with TPT in both related work and experiments (adapting CLIP models). How does Diffusion-TTA compare with TPT regarding OOD accuracy and inference efficiency? \n3. Can the authors give insights of why the accuracy curve behaves like this (a spike at K=5 and does not change after K>=10) and justify the selection of K (K=5 vs K=1)?\n\nLIMITATIONS:\nThe authors have addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-11"}