Exploring Structured Semantic Priors Underlying Diffusion Score for Test-time Adaptation

Q1: Existing art, e.g. [1], show there may be some benefit to using a few time steps. It would have been nice to see some exploration of the impact of using more than 1 step, insofar as it would likely make the approximation in Equation 9 more accurate. Was any exploration of ensembling time steps attempted?

A1: Thanks for the valuable advice. The inspiring work [1] unveils the discriminativeness of U-Net features in diffusion models and achieves superior unsupervised unified representation learning performance across diverse tasks, where ensembling of timesteps is found beneficial. Our DUSA shares a similar spirit of exploring the discriminativeness in generative diffusion models, thus the ensembling of timesteps is definitely worth trying out.

To explore the impact of ensembling timesteps as suggested, we experiment by performing DUSA on more than one timestep simultaneously and averaging their losses, and the results are as follows:

It can be seen that, the ensembling of different timesteps does bring a further improvement in performance, thanks to the knowledge from multiple timesteps. Furthermore, our DUSA has its special advantage of draining knowledge from every single timestep, and can still get good results with reduced computational overhead. We appreciate the insightful advice and will add the discussions in the revision.

[1] Mukhopadhyay et al., Do text-free diffusion models learn discriminative visual representations?

Q2: Considering the related work is positioned at the end, the introduction should do a little more to introduce and motivate the task.

A2: Thanks for the good suggestion. In test-time adaptation, a pre-trained model is updated on the fly to make accurate predictions on the incoming target samples without label access, which is challenging as the target data distribution might drift from that in pre-training. Being a competitive family of generative models, diffusion models exhibit great capability in modeling data distributions and even show discriminative potential in learned features [1,2,3,4], making them a strong candidate to provide guidance for discriminative tasks. In this work, we aim to extract discriminative priors from diffusion models to facilitate the challenging task of test-time adaptation.

We thank the reviewer's efforts in improving the readability of our manuscript and will further polish up the expression and integrate the above discussions into the introduction in the revision.

[1] Mukhopadhyay et al., Do text-free diffusion models learn discriminative visual representations?

[2] Zhang et al., Diffusionengine: Diffusion model is scalable data engine for object detection.

[3] Zhao et al., Unleashing text-to-image diffusion models for visual perception.

[4] Xu et al., Open-vocabulary panoptic segmentation with text-to-image diffusion models.

Q3: What is the latency of the method? Limitations mentions that it would not work for real time (e.g. autonomous driving) applications, but how far off is it?

A3: Thanks for the question. As shown by the results in the table below, our DUSA predicts at a rate of 0.11s/image for classification with Diffusion Transformer (DiT), which is -99.6% in time compared to Diffusion-TTA. For segmentation, the time is 4.5s/image for segmentation with ControlNet.

The latency is largely related to the selection of the diffusion model as well as the task complexity. With this said, we can freeze the diffusion model (DUSA-Frozen) and only update the task model in DUSA, which will still provide decent performance on a handful of tasks with further reduced latency and nearly real time, and the results are as follows:

We will continue researching how to reduce the latency for discriminative tasks further while utilizing the diffusion model's strong power of generative modeling.

Q1: How the joint update for the two networks is performed is unclear.

A1: Thanks. Before adaptation, a single fixed timestep is selected and other timesteps are dismissed. Hereafter, the diffusion model is only updated on this single timestep, i.e., we added a same level of noise to all images and denoise them at a same timestep. Note that the task model predicts on clean images, thus not affected by timestep. The diffusion noise estimations and the task model predictions are then integrated into Eq. 10 for a joint update.

Q2: The timestep selection should impact the diffusion model's update schedule and the output logits.

A2: Thanks. In our DUSA, the diffusion model is only updated on a single fixed timestep. We agree that timestep selection should impact the amount of noise in $x_t$ and the timestep condition $t$ for noise estimation, but we argue that the timestep in DUSA is fixed throughout adaptation and how we select it makes no difference to DUSA's pipeline. Besides, the task model predicts logits from clean images and is not affected by timestep selection.

Q3: The computational overhead needed for the approach and the gain over the prior work are not discussed.

A3: Thanks. Firstly, with $K$ classes and $T$ diffusion timesteps, the prior work Diffusion-TTA has a computational complexity of $\mathcal{O}(T)$, indeed 180 timesteps for each image. By contrast, our DUSA has an $\mathcal{O}(K)$, further reduced to $\mathcal{O}(b)(b\ll K)$ by practical designs (lines 136-139, 193). We list the computation time with gain as follows:

It can be seen that our DUSA shows significant gain over Diffusion-TTA with a tremendous reduction (-99.6%) in computational overhead.

Q4: Class embedding as conditional is generally required for the diffusion model. This should limit the discriminative settings that can be considered for TTA with diffusion models.

A4: Thanks. We agree that diffusion models generally require class embeddings. However, the specific form of class embeddings can vary. For classes with semantic meanings, we can construct a prompt containing class names and use a text model as embedder. For classes without explicit meanings, we can discretize them to IDs and map them to embedding vectors. In this way, we believe any discriminative task can be taken into consideration.

Q5: How are the conditionals computed for the segmentation task? The increase in computational overhead should also be discussed with the increased complexity of the task.

A5: Thanks. As we stick to text-to-image diffusion models, the conditionals for segmentation are computed at class level. Specifically, we construct a prompt "a photo of a {class}" and use the diffusion model's text encoder to obtain class embedding. Please note that a single conditional noise estimation in diffusion could also be viewed as a noise estimation for every single pixel. Therefore, the computation of conditionals for the segmentation task is as simple as in classification: given a class, predict the noise. The only difference is the way of ensembling noise estimations with predictions as in Eq. 12, which introduces little overhead.

The increase in computational overhead is first raised by a large network to adapt in segmentation. Besides, the overhead is partially due to using a larger diffusion model (ControlNet, 1.7B params) than that in classification (DiT, 675M params) and the sliding-window strategy for non-square images. Lastly, the overhead also comes from the unpolished use of our CSM module for this task. We find that our method is versatile and competitive across tasks, and it's majorly the task complexity that adds to overhead. We will add these discussions in the revision.

Q6: Single timestep for adapting diffusion models on vision tasks has been discussed in [a] and [b]. These works can be discussed in the related work.

A6: Good suggestion. [a] gets strong results in finding semantic correspondences by obtaining transferable prompt embeddings at a single timestep. [b] achieves significant gains in the one annotation-based image segmentation task at any granularity by learning reusable text embeddings at a single timestep. While also using a single timestep, our DUSA focuses on utilizing the diffusion model to adapt a discriminative model. We find the works relevant and will discuss them in the revision.

[a] Unsupervised semantic correspondence using stable diffusion. NeurIPS 2023.

[b] Slime: Segment like me. ICLR 2024.

Q7: How is the update performed for the diffusion model as in [26] updating the diffusion model is optional. Does it incur any additional overhead?

A7: Thanks. As described in A1, we update the diffusion model on a single fixed timestep. Like [26], updating the diffusion model is optional in DUSA, and we show the performance of freezing the diffusion model in Table R6 in uploaded PDF. It can be seen that when not updating diffusion model, our method still outperforms [26]. Besides, with updating, our DUSA takes the lead by a large margin, with a vastly reduced overhead of -99.6% adaptation time compared to [26], as shown in A3.

Q8: Do different tasks, for example, segmentation and classification, have a single timestep where they perform well, or depending on the granularity of the task the choice of the timestep is different?

A8: Thanks. We vary the selected timestep for segmentation and report the results in Figure R2 in uploaded PDF. As the timestep being too large or too small is not preferred (lines 169-172), we focus on a narrower range here. It can be seen that there is no global best for all scenarios, but our $t=100$ proves effective across tasks.

Q1: Though Eq. 10 is well-supported from theoretical perspective, the alternative Eq. 11 needs further elucidation. Insights or intuitions on why the modified objective (Eq. 11) works as well as the original objective (Eq. 10)?

A1: Thanks for the valuable question. Please note that we can view Eq. 10 from two perspectives: (a) the task model learns from the diffusion model, (b) weighted optimization is performed so that $\epsilon_\phi(x_t,t,c_y)$ with a larger weight $p_\theta(y|x_0)$ is more responsible for estimating true noise. Therefore, Eq. 10 encourages the conditional diffusion model to adapt to test-time data, under instruction of task model predictions.

Similarly, two perspectives of Eq. 11 can be shown. The conditional part adapts the task model as in Eq. 10. For the unconditional part, we start from Eq. 9: $\epsilon\approx\sum_yp(y|x_t)\epsilon_\phi(x_t,t,c_y).\quad(\text{Eq. 9})$ Without loss of generality, we ignore the estimation error below. In a CFG-based diffusion model, the true noise in Eq. 9 can be unconditionally estimated: $\epsilon_\phi(x_t,t,\varnothing)=\sum_yp(y|x_t)\epsilon_\phi(x_t,t,c_y).\quad(1)$ Please note that all noise estimations are from a same diffusion model, and the weights $p(y|x_t)$ are the implicit priors in the diffusion model. Therefore, we can view (1) as an internal constraint of the diffusion model. Intuitively, the unconditional part implicitly enforces conditional adaptation of the diffusion model to test-time data, now under instruction of the implicit priors from the diffusion model itself, explaining why Eq. 11 works as well as Eq. 10. We will add the discussions in the revision.

Q2: The unconditional version of CFG-based models is much weaker than the conditional version, and they usually do not agree well with each other. Do the unconditional training really help the conditional noise estimation?

A2: Thanks. We agree that the unconditional version is weaker and there can be disagreement. However, this version is still decent for noise estimation, as it also undergoes extensive diffusion training. Besides, in A1 we show that the unconditional training helps adapt conditional noise estimations to test-time data in CFG-based models. Therefore, we prefer to view the two versions as cooperative, and their agreement is not forced.

We compare unconditional training with freezing the diffusion model in Table R1 in uploaded PDF. We can see that the unconditional results are comparable to conditional training. This strengthens our belief that the unconditional training does help the conditional noise estimations, with reduced overhead.

Q3: The conditional part in Eq. 11 does not involve tuning the diffusion weights, why does the diffusion model (unconditional version) still need to be constrained by a regular CFG objective?

A3: Thanks. The regular CFG objective explicitly enhances both conditional and unconditional noise estimations for better generation. In contrast, Eq. 11 focuses on utilizing the diffusion model to facilitate task model adaptation (conditional part) and leveraging the implicit priors to adapt the conditional noise estimations (unconditional part), as shown in A1.

We constrain the diffusion model with unconditional training to handle a potential distribution shift between diffusion model's training set and test-time data. Results are shown in Table R1 in uploaded PDF. If the diffusion model is generalizable to unseen data, unconditional training can be removed with limited sacrifice in performance (DUSA-Frozen for Gaussian). Otherwise, should the test-time data be out-of-distribution for the diffusion model, the absence of unconditional training might cause degradation (DUSA-Frozen for Defocus). Therefore, the constraints are preferred for a more consistent performance gain.

Q4: Whether a large text-to-image model pre-trained on web images is better than a model pre-trained on the same distribution for segmentation.

A4: Thanks. Actually, the diffusion model we adopt for semantic segmentation is a ControlNet (based on Stable Diffusion v1.5) fine-tuned on ADE20K, same distribution as the task model. Although not equal to training from scratch, we assume the fine-tuning makes the diffusion model aware of the task model's training data distribution.

We compare the results of adapting with ControlNet against Stable Diffusion models in Table R3 in uploaded PDF. It can be seen that the diffusion model's training on the same distribution as the task model generally merits more performance gain, as it narrows the gap between the task model and the diffusion model for better cooperation.

Q5: How does the quality of diffusion models affect TTA performance?

A5: Thanks. We compare the quality of popular diffusion models in Table R4 in uploaded PDF. We apply our DUSA to them and show classification in Table R2 and segmentation in Table R3 in uploaded PDF. It can be seen that while a diffusion model of better quality fosters TTA performance (SD series), one training on the same distribution as the task model is generally more favorable (blue lines).

Q6: The method also relies on the fine-tuning of diffusion models, rather than only evaluating off-the-shelf models. This makes the TTA process more costly, since diffusion models are usually much larger than discriminative task models.

A6: Thanks. We agree that the fine-tuning of diffusion models makes TTA more costly, but we highlight the substantial performance gain obtained by updating the diffusion model (DUSA) against freezing it (DUSA-Frozen), as shown in Table R1 in uploaded PDF.

Besides, our DUSA is more than 100x faster in comparison with Diffusion-TTA, and approaches the strong traditional TTA method EATA in speed. The results are in Table R5 in uploaded PDF, where our DUSA shows leading performance with little extra cost.

Q1: The algorithm in Appendix should be integrated into the main paper.

A1: Thanks. Following your advice, we will integrate the algorithm into the main body of the paper in the revision.

Q2: The term "fresh" is used ambiguously in lines 66-68.

A2: Good suggestion. We agree that "a fresh proposition" might be ambiguous, and we clarify by restating it as "a novel proposition" and will update it in the revision.

Q3: The verification of this reduction and a clear comparison with standard Diffusion-TTA remain unclear. Whether averaging out the loss function $\mathcal{L}_{DUSA}(\theta, \phi; t)$ over T=180 under the same conditions as Diffusion-TTA would affect performance.

A3: Thanks. Diffusion-TTA is built on the Monte Carlo estimation of likelihood, and demands randomly sampling up to 180 timesteps to boost performance. By contrast, our DUSA utilizes semantic priors from a single timestep, and needs only a single fixed timestep to achieve superior performance.

For a clear comparison with the standard Diffusion-TTA, we put our DUSA under the same conditions of Diffusion-TTA: (a) we change the timestep selection from fixed to random sampling, (b) we increase the timesteps number from 1 to 180. The comparison results are as follows:

It can be seen that our DUSA significantly outperforms Diffusion-TTA even under its same conditions.

Q4: Comparision of Diffusion-TTA and DUSA with only a single timestep, as shown in Figure 3.

A4: Thanks. To compare under only a single fixed timestep, modifications are made to Diffusion-TTA: (a) we reduce the utilized number of timesteps from 180 to 1, (b) we change the timestep selection from random sampling to a fixed timestep. The comparison results for $t=100$ are as follows:

Our DUSA surpasses Diffusion-TTA by a large margin, justifying the superiority of DUSA on only a single timestep. More results on varied timesteps can be found in Figure R1 in uploaded PDF.

Q5: The paper mentions a biased approximation but lacks concrete proof. Whether the proposed DUSA method can demonstrate that it does not rely on biased approximations.

A5: Thanks. Indeed, our original claim "... for a biased approximation of likelihood" may cause confusion, and a better expression is "... to estimate a biased approximation of likelihood". Actually, the bias comes from Diffusion-TTA's theoretical approximation of the likelihood instead of the Monte Carlo method, and we provide a proof below.

Diffusion-TTA is built on the evidence lower bound (ELBO) of log-likelihood in diffusion: $\log p_\phi(x_0|c)\geq \mathbb{E}\_q\Big[\log\frac{p\_\phi(x\_{0:T}|c)}{q(x\_{1:T}|x_0)}\Big],\quad(1)$ where $q$ is the forward process and $p_\phi$ is the backward process. We can further derive the ELBO in (1) with denoising network $\epsilon_\phi$: $\mathbb{E}_q\Big[\log\frac{p\_\phi(x\_{0:T}|c)}{q(x\_{1:T}|x_0)}\Big]=-\mathbb{E}\_\epsilon\Big[\sum\_{t=2}^Tw_t\\|\epsilon-\epsilon\_\phi(x_t,t,c)\\|_2^2-\log p\_\phi(x_0|x_1,c)\Big]+C\approx-T\mathbb{E}\_{\epsilon,t}[\\|\epsilon-\epsilon\_\phi(x_t,t,c)\\|_2^2]+C,\quad(2)$ where the theoretical approximation is made by simplifying weights $w_t$ to 1 and ignoring the $\log p\_\phi(x_0|x_1,c)$ term. Note that Diffusion-TTA actually estimates the simplified objective in (2), thus the result is theoretically biased from real likelihood because the approximation in (2) just prevents the equality sign in (1) from holding.

In contrast, our DUSA objective in Eq. 10 is built on the theory in Eq. 4 with no theoretical approximation. The task model thus directly utilizes semantic priors from the diffusion model. We will add the discussions in the revision.

Q6: Advantages over Diffusion-TTA in dense prediction tasks. Whether Diffusion-TTA is also applicable and effective for dense prediction tasks.

A6: Thanks. Diffusion-TTA is applicable to dense prediction tasks, as depicted in [26]. However, a pixel-level conditioning capability is assumed on the diffusion model for Diffusion-TTA to be effective, which requires re-training of diffusion models and reduces versatility. We re-train the diffusion model on ADE20K according to [26] to fit segmentation task, and comparison results are as follows:

It can be seen that our DUSA consistently outperforms Diffusion-TTA on the dense prediction task.

Q7: The effects of using a diffusion model trained on ImageNet instead of one trained on a large dataset like Stable Diffusion. Given that Stable Diffusion benefits from a larger dataset, not comparing it to other models might be seen as a limitation.

A7: We appreciate the reviewer's insights in exploring our DUSA on diffusion models varying in training set. We would clarify that, for image classification, an ImageNet pre-trained Diffusion Transformer (DiT) is already used throughout our paper. We only use Stable Diffusion-based ControlNet for the segmentation task. As suggested, we further compare to other diffusion models and provide results for classification in Table R2 in uploaded PDF and segmentation in Table R3 in uploaded PDF.

We can find that a diffusion model trained on the same distribution as the task model is generally more beneficial for adaptation. Intuitively, the diffusion model's being aware of the pre-training data of the task model eliminates the disagreement between them, fostering their cooperation under a narrowed distribution gap.