GDA: Generalized Diffusion for Robust Test-time Adaptation

The trade-off between preserving content while translating domains or style has been studied by DDA [5, 48]. When the noise variance $\sigma$ is more extensive, it is challenging to preserve the content information. Therefore, the structural guidance allows the diffusion model to generate samples conditioned on the predefined objectives. In particular, the structural guidance iteratively refines the latent for the input images during the reverse process so that the content information in the sample can be preserved while translating the style or shifting the domain.

Due to the sampling process of DDPM being a Markov chain, it requires all past denoising steps to obtain the next denoised image. The long stochastic operations can lead to huge distortion of the content information. To guide the diffusion more efficiently with structural guidance, we speed up the sampling process with DDIM [38] by skipping several reverse steps. The reverse process can be redefined as:

where $x^{g}_{0,t}$ is the predicted denoised image for $x_{0}$ conditioned on $x^{g}_{t}$ at the time step $t$ and is defined as:

Our structural guidance has two steps: (1) At time step $t$, generate the sample $x^{g}_{t-1}$ for the next step $t-1$. (2) Update $x^{g}_{t-1}$ with the gradient calculated from our structural guidance. To avoid the conflicting with the original reverse sampling step at each time step in the diffusion, our structural guidance is computed by the reference image $x_{0}$ and its corresponding denoised image $x^{g}_{0,t}$ at reverse time step $t$. The updated process is defined as:

where $\ell_{guided}$ is our objective function for structural guidance, and the inputs of the objective are $x^{g}_{0,t}$ and $x_{0}$.

In Algorithm 1, we present GDA. Our proposed structural guidance incorporates the marginal entropy loss into the objective function to ensure the output behavior of the model has consistent predictions on generated samples and their augmented version. Inspired by [48], we combine text-driven style transfer using CLIP and content preservation using zero-shot contrastive loss. Our objective function is:

where $\ell_{marginal}$ denotes the marginal entropy loss. $\ell_{style}$ and $\ell_{content}$ denote the style and content preservation loss. We further discuss the details for each loss component.

We notice the stochastic nature of the diffusion model in the reverse process, where the noise $\epsilon$ can lead to the distortion of content information in the input image and cannot correctly generate samples close to the source domain that diffusion has been trained on. Given a model $f_{\theta}$ which is trained on the source domain, we add the marginal entropy loss for guiding the diffusion reverse process. In particular, the loss will force the whole diffusion process to generate samples that can decrease the model’s uncertainty for $f_{\theta}$. At timestep $t$, given a generated sample $x^{g}_{t}$ and a set of augmentation functions $\mathcal{A}=\{A_{1},A_{2},...,A_{n}\}$, we augment the sample $x^{g}_{t}$ by choosing subset of augmentation functions from $\mathcal{A}$. We denote the image sequence of augmented data as $A_{1}(x^{g}_{t}),A_{2}(x^{g}_{t}),...,A_{k}(x^{g}_{t})$, where $k\leq n$. The marginal output distribution for the given generated sample $x^{g}_{t}$ is defined as:

where $p_{\theta}$ is the output prediction of each augmented sample and $\bar{p}_{\theta}$ is the average on all augmented samples. Our intuition lies in that $f_{\theta}$ is trained on the source domain $\mathcal{X}_{S}=[x_{1},x_{2},...,x_{N}]$ and should learn the invariance between the augmented samples $\hat{x_{1}},\hat{x_{2}},...,\hat{x_{N}}$ and $\mathcal{X}_{S}$. When generating a sample $x^{g}_{t}$ at time step $t\in[T,...,1]$ from diffusion, if the sample is close to the source domain, the output prediction of its augmented versions will be consistent, and the marginal entropy loss will become small. Thus, we can utilize this loss to ensure the diffusion process generates samples close to the source domain. Here, the entropy of marginal output distribution is defined as:

To better control the sample quality from the diffusion model, the uncertainty estimation on original and adapted samples is then applied to the sampling strategy. The uncertainty score function is $H(x)=-\sum_{y\in\mathcal{Y}}p_{\theta}(y|x))\log p_{\theta}(x)$, where the input can be the original sample $x$ or adapted samples $x^{g}_{0}$.

To transfer samples from one style to another without content distortion, prior work proposed guided-loss for the diffusion model [48]. Inspired by them, We use the CLIP model to calculate the style loss. By injecting a text prompt related to the source domain (e.g., photo-realistic, real), the CLIP model calculates the similarity between the features extracted from the input image and the text prompt. Our style loss is defined as:

where $Enc_{img}$ and $Enc_{txt}$ are the image and text encoder in the CLIP model.

To avoid content distortion, we use patch-wise contrastive loss to ensure the generated sample’s content information is consistent with the original sample. In [31], they show contrastive unpaired image-to-image translation loss can preserve the content information by maximizing the mutual information between the input and output patches. To compute the content preservation loss, we extract the spatial features from the UNet component of the diffusion model. The content preservation loss is:

where $\hat{z}$ and $z$ are the corresponding patch-wise features of $x^{g}_{0,t}$ and $x_{0}$ extracted from UNet $h(\cdot)$. $\tau$ is the temperature scaling value. $y_{i,j}$ is a 0-1 vector for indicating the positive pairs and negative pairs. If $y_{i,j}$ is 1, the $i$-th feature $\hat{z}_{i}$ and $j$-th feature $z_{j}$ are at the same location from the $x^{g}_{0,t}$ and $x_{0}$ samples. Otherwise, they are from different locations.

We do the quantitative measurement of our method by showing the entropy loss distribution for different corruptions. Our entropy is defined as $H(x)=-\sum_{y\in\mathcal{Y}}p_{\theta}(y|x))\log p_{\theta}(x)$, where it measures the ambiguity of the data with respect to the given target classifier. The lower entropy loss means the model has the higher confidence in the samples. As Fig. 4 shows, the different colors represent different adaptation methods. The dark green color represents the original sample and the orange color represents the loss distribution of corrupted samples. We show that the entropy loss distribution has a massive shift between corrupted and original samples, which means the model has lower confidence in most of the corrupted samples than the original samples. We then show the entropy loss of samples after adapting with three diffusion-driven adaptation methods, including Diffpure (blue), DDA (red), GDA (light green). As every subfigure in Fig. 4 shows, for every corruption type, the loss distribution of samples generated from GDA moves toward the entropy loss distribution of original samples, which means that our method indeed shifts the OOD samples back to the source domain. However, DDA and Diffpure do not have excessive shifting on the entropy loss distribution.

In Figure 5, we show the effect of different reverse steps on the performance of the diffusion model. In our experimental results in Section 4, we fix the reverse step number as 10 for every baseline. Here, we compare different reverse sampling steps for DDA and ours from small to large (1 to 50). As Fig. 5 shows, GDA has a more significant improvement under a small number of reverse steps (e.g., 10) and is more effective compared to the DDA baseline. When increasing the reverse sampling steps to 50, GDA slightly improves but still outperforms the DDA baseline on every OOD benchmark.

To show how our structural guidance can guide the diffusion model, we visualize the samples generated from GDA and their corresponding gradient classification activation maps (GradCAM). In Fig. 6, the corrupted images after adaptation are visually de-corrupted, and the saliency map from GradCAM demonstrates how our objective function can guide the model during the adaptation. In Appendix Fig. 7 and 8, we show the samples from Rendition and Stylized with wrong predictions before adaptation and their corresponding adapted models with correct predictions.

In Table 4, we show the adaptation cost under different adaptation methods, including DDA, Diffpure, without marginal guidance, and GDA. For GDA, the run time depends on the number of augmented samples. Thus, we select the number with the best accuracy (16) for comparison. Compared to DDA and Diffpure, our method outperforms them by $\sim$7% on ImageNet-Rendition and reduces 3.85x run time.

We apply the CLIP model with model architecture ViT-Base/16 for calculating the style loss. By leveraging the rich semantic information of CLIP, we are able to shift the OOD sample to the source domain. It has been used in [C2] for style transfer. The input images are presented to the model as a sequence of fixed-size patches, where the patch size is 16*16). We get the corresponding image embedding for all image patches from the output of the visual encoder of CLIP model. We then calculate the similarity between the image embeddings and the text token embedding extracted from language encoder of CLIP model. The text prompts we use for style loss are the words related to photo-realistic or real photo. We assume partially knowing the source domain information is allowable in domain generalization.

We provide a more detailed of the contrastive loss for content preservation. The input of content loss is a batch of features extracted from generated sample itself $x^{g}_{t}$ and the corresponding source sample $x_{0}$. For example, $v$ is the $i^{th}$ patch in sample $x^{g}_{t}$, the $i^{th}$ patch $p$ in sample $x_{0}$ is its positive pair $p+$, and all the other patches except the $i^{th}$ patch in sample $x_{0}$ will be the negative pair $p-$. The purpose of the contrastive loss is to force the feature distance between a patch $p$ and its corresponding positive patch $p+$ to become closer to each other under the latent space. Meanwhile, the loss forces $p$ and $p-$ apart from each other.

We adopt AugMix [12], a data augmentation tool from Pytorch, which randomly select several augmentation functions (e.g., posterize, rotate, equalize) to augment the data. The augmention set $\mathcal{A}={A_{1},A_{2},...,A_{k}}$ excludes operations that overlap with corruption types in ImageNet-C. For generating one augmented sample $x_{aug}$, we set the mixing weight $w_{1},w_{2},...w_{3}$ for every augmentation in $\mathcal{A}$. The mixing weight, which is a $k$-dimensional vector of convex coefficients, is randomly sampled from a Dirichlet distribution. The augmented sample $x_{aug}$ equals to $w_{n}*A_{n}(...(W_{2}*A_{2}(w1*A_{1}(x_{orig}))$.

We conduct the sensitivity analysis on hyperparameters for every loss. We follow the range of hyperparameters used in [48]. In Table 5, we whoe the results of ImageNet-R under different combination of loss terms.

We show the impact of different loss term by removing content preservation loss or style loss in Table 6. The result of using only style loss is better than content loss on ImageNet-Rendition and Sketch.

In GDA, the weights for each loss function are hyperparameters that need to be chosen by users. We combine the three loss terms as a joint optimization, with their Lagrange multipliers as hyperparameters. The hyperparameter values for each benchmark are shown in Table 7.