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

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 use Diffusion-TTA to adapt multiple ImageNet classifiers with varying backbone architectures and sizes: ResNet-18, ResNet-50, ViT-B/32, and ConvNext-Tiny/Large. For fair comparisions, we use Diffusion Transformer (DiT-XL/2) as our class-conditional generative model, which is trained on ImageNet from scratch. All models are trained on ImageNet without using additional data. We consider the following datasets for TTA of ImageNet classifiers: ImageNet and its out-of-distribution counterparts: ImageNet-C (level-5 gaussian noise), ImageNet-A, ImageNet-R, ImageNetV2, and Stylized ImageNet.

Under single-sample test-time adapation setup, our method improves CLILP classifiers of different sizes consistently over all following datasets: ImageNet, CIFAR 100, Food 101, Flowers 102, FGVC Aircraft, and Oxford-IIIT Pets datasets. CLIP classifiers have not seen test images from these dataset during training.

We use Diffusion-TTA to adapt semantic segmentors of SegFormer and depth predictors of DenseDepth. We use a latent diffusion model 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. We test adaptation performance on different distribution drifts.