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

Summary

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

Strengths

I 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. The paper is well-written and easy to follow. The approach is clear. In 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.

Weaknesses

Weaknesses: My 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. ### Issues with the results section Table 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. Right 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. ### General comments Line 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. [A] Carlini et al. “Extracting Training Data from Diffusion Models”, https://arxiv.org/pdf/2301.13188.pdf In 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? Do 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. [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 The text in Figure 5 is too small below the bars and largely unreadable. Figure 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? More 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%)”? I 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. In 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. Line 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? [C] Wang et al. “Tent: Fully Test-time Adaptation by Entropy Minimization” [D] Rusak et al. “If your data distribution shifts, use self-learning” The appendix should be referenced in the main text.

Questions

Please respond / fix the points I outlined above. Please 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. I 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

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The 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. One limitation of the method is that it is very compute-intense and thus will likely be of limited use to practitioners.

Summary

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

Strengths

* 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. * 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. * 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. Overall, the paper provides a valuable contribution to the field with its innovative approach and robust experimental evaluation.

Weaknesses

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

Questions

* 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? * What is the time required to adapt the model to a single image? Can this approach be scaled up for practical adaptation applications? * In Figure 6, have you investigated the effects of increasing the value of K for a CLIP model?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

3 good

Limitations

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

Summary

This 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. From 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. ---- After reading the rebuttal, I've raised my score from 6: weak accept to 7: accept.

Strengths

Sufficient novelty: Diffusion-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. Rigorous ablation study: The 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.

Weaknesses

Empirical evaluation: I 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. I’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.

Questions

How’s the performance of TTA-MAE or other existing TTA methods on the Open-Vocabulary CLIP experiment? Have the authors tried partial weight updating on the classifier, e.g., only adjusting weights of a few top layers, and observing the performance gain? In 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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The 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. I’m hoping that the code implementation is made available soon.

Summary

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

Strengths

1. Paper is well-written and easy to understand. 2. Proposed setup is a simple plug and play method that does not require any additional training or warm-up of the models. 3. Results demonstrate that the proposed adaptation improve the results, particularly benefit smaller models. 4. Appreciate the analysis that shown diffusion loss and cross-entropy loss are correlated, which support the gains observed in the results.

Weaknesses

1. Proposed 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. 2. Increased time complexity of test time adaptation with the diffusion model, where each sample takes 55 seconds long for adaptation. 3. Increased model complexity by including diffusion model with large number of parameters in the adaptation process. 4. The adaptation is not as effective on larger models when compared to smaller models. 5. Prior works comparisons are included in Table 2 but the network architectures are different across these comparisons and wouldn't be a fair comparison.

Questions

1. Which diffusion model was used in Sec. 4.1? 2. CLIP adaptation is shown across multiple datasets. Adapting CLIP with ImageNet variants would be interesting.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Yes

Summary

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

Strengths

1. 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. 2. It shows the improvement of classification accuracy across different network architectures after using Diffusion-TTA. 3. The ablation studies (in Table 3) are well conducted to show the importance of each component.

Weaknesses

1. 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. 2. 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. 3. 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. 4. 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. [1] Test-Time Prompt Tuning for Zero-Shot Generalization in Vision-Language Models, NeurIPS 2022.

Questions

1. 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)? 2. 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? 3. 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)?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have addressed the limitations.

Response to authors' rebuttal

Dear authors, thank you for addressing my concerns in your rebuttal. I think you did a great job providing new results! In particular: - I appreciate reporting standardized results for your method vs TENT / CoTTA. It is nice to see that your method works best. - Thank you for adding semantic segmentation results. - I think your response to Q.6.7 is really interesting and Figure 2 in the attached pdf is insightful. I think it helps understanding the method better and would encourage you to include the results in the main manuscript, if space permits. - It is also interesting that the model cannot be trained from scratch using the method, but that it only works for finetuning. Again, I think this provides us with insight about the method and should be included in the final version. One concern remains and that is of the computational complexity of the method. This method is 2-3 orders of magnitude slower than TENT, but then TENT does not always work and also cannot be applied in the single-sample regime. Given that TTA methods typically need to be light-weight, I am not sure how widely this method will be used in practice. I don't know how big of a concern this is, as this paper can be regarded as a "proof-of-principle" contribution. Could the authors maybe try to elaborate in which settings they think this method could be used successfully? I have raised my score because I think the authors did a great job addressing my concerns. Best, Reviewer 4gPG

Response to authors’ rebuttal

Thanks for providing very detailed responses to my concerns. Since most of my concerns have been addressed, I'm happy to raise my rating to weak accept. I hope the authors can incorporate all the changes and new results into the revised version of the paper.