GDA: Generalized Diffusion for Robust Test-time Adaptation

TL;DR

This work tackles robustness to unseen distribution shifts by proposing Generalized Diffusion Adaptation (GDA), a diffusion-based test-time approach that does not modify model weights. GDA guides the reverse diffusion process with structural losses—marginal entropy, style transfer via CLIP, and content preservation via patch-wise contrastive signals—to pull OOD samples back toward the source domain. Using an efficient DDIM-style sampling strategy, GDA improves accuracy across ImageNet-C, Rendition, Sketch, and Stylized-ImageNet for multiple backbones, with notable gains (up to about 5 percentage points on ImageNet-C and 2.5–7.4 points on stylized-style benchmarks) and reduced adaptation cost. These results demonstrate that diffusion-based test-time adaptation with targeted guidance can generalize across diverse OOD types, offering a practical route to robust deployment without retraining models. The findings suggest promising extensions to other vision tasks and broader guidance mechanisms for diffusion in OOD settings.

Abstract

Machine learning models struggle with generalization when encountering out-of-distribution (OOD) samples with unexpected distribution shifts. For vision tasks, recent studies have shown that test-time adaptation employing diffusion models can achieve state-of-the-art accuracy improvements on OOD samples by generating new samples that align with the model's domain without the need to modify the model's weights. Unfortunately, those studies have primarily focused on pixel-level corruptions, thereby lacking the generalization to adapt to a broader range of OOD types. We introduce Generalized Diffusion Adaptation (GDA), a novel diffusion-based test-time adaptation method robust against diverse OOD types. Specifically, GDA iteratively guides the diffusion by applying a marginal entropy loss derived from the model, in conjunction with style and content preservation losses during the reverse sampling process. In other words, GDA considers the model's output behavior with the semantic information of the samples as a whole, which can reduce ambiguity in downstream tasks during the generation process. Evaluation across various popular model architectures and OOD benchmarks shows that GDA consistently outperforms prior work on diffusion-driven adaptation. Notably, it achieves the highest classification accuracy improvements, ranging from 4.4\% to 5.02\% on ImageNet-C and 2.5\% to 7.4\% on Rendition, Sketch, and Stylized benchmarks. This performance highlights GDA's generalization to a broader range of OOD benchmarks.

Figures (9)

• Figure 1: Sample OOD data and adaptations via existing diffusion method and our GDA method. The leftmost column shows OOD samples under different style changes, including sketch, painting, and sculpture. The middle column shows samples adapted by traditional diffusion. The rightmost column shows samples adapted with our GDA method. The visualization shows that GDA can generate samples with multiple visual effects, such as re-colorization for the sketch sample, texture enhancement for the painting sample, and object highlighting for the sculpture sample. All three GDA-adapted samples are correctly classified by ResNet50, whereas all others are misclassified.
• Figure 2: The flow of GDA. We guide the diffusion model with our novel structural guidance that includes marginal entropy, style loss, and content preservation loss. Given the corrupted samples $x_0$, when going through the reverse process at step $t$, our structural guidance will first (1) Generate the sample $x^g_{t-1}$ for the next reverse time step $t-1$. (2) Update the $x^g_{t-1}$ with the gradient calculated from the losses. Our loss is computed by the reference image $x_0$ and its corresponding denoised image $\hat{x}^g_{0,t}$ conditioned on $x^g_{t}$ at reverse time step $t$.
• Figure 3: Comparison of the performance for our method with baselines under 15 types of corruption in ImageNet-C for three model architectures, including ResNet50, ConvNext-T, and Swin-T. GDA shows better improvement on all corruption types for ImageNet-C.
• Figure 4: Entropy loss measurement for different corruptions on ImageNet-C. From left to right, the x-axis shows different adaptation methods. The y-axis shows the entropy loss values. The lower value means the model has higher confidence on the sample. In each subfigure, from left to right, we show the loss distribution for original sample (green), corrupted samples (orange), samples adapted by Diffpure nie2022DiffPure (blue), samples adapted by DDA gao2022back (pink), and samples adapted by our method (light green).
• Figure 5: Sensitivity analysis on the reverse sampling steps. We compare our method with DDA under different sampling steps from 1 to 50. We evaluate on the ResNet50 model and show the standard accuracy with green color line.