Test-Time Adaptation with CLIP Reward for Zero-Shot Generalization in Vision-Language Models

TL;DR

The paper addresses zero-shot generalization under distribution shift in vision-language models by proposing Reinforcement Learning with CLIP Feedback (RLCF), a universal test-time adaptation framework that uses CLIP-based rewards to refine outputs from a single test sample via policy-gradient updates. It defines a CLIP-based reward $ ext{R}(m{t},m{v})$ with a baseline and employs REINFORCE to adapt the model parameters, enabling task-specific pipelines for image classification, text-image retrieval, and image captioning. Key contributions include a novel, practical reward function that can operate without ground-truth labels, task-tailored TTA pipelines, and a suite of tricks (ensembles, episodic TTA, momentum buffers) that improve robustness and generalization. The approach yields significant improvements over zero-shot baselines and prior TTA methods on multiple benchmarks, demonstrating flexible, generation-capable, test-time generalization for VLMs with practical runtime characteristics.

Abstract

One fascinating aspect of pre-trained vision-language models~(VLMs) learning under language supervision is their impressive zero-shot generalization capability. However, this ability is hindered by distribution shifts between the training and testing data. Previous test time adaptation~(TTA) methods for VLMs in zero-shot classification rely on minimizing the entropy of model outputs, tending to be stuck in incorrect model predictions. In this work, we propose TTA with feedback to rectify the model output and prevent the model from becoming blindly confident. Specifically, a CLIP model is adopted as the reward model during TTA and provides feedback for the VLM. Given a single test sample, the VLM is forced to maximize the CLIP reward between the input and sampled results from the VLM output distribution. The proposed \textit{reinforcement learning with CLIP feedback~(RLCF)} framework is highly flexible and universal. Beyond the classification task, with task-specific sampling strategies and a proper reward baseline choice, RLCF can be easily extended to not only discrimination tasks like retrieval but also generalization tasks like image captioning, improving the zero-shot generalization capacity of VLMs. According to the characteristics of these VL tasks, we build different fully TTA pipelines with RLCF to improve the zero-shot generalization ability of various VLMs. Extensive experiments along with promising empirical results demonstrate the effectiveness of RLCF. The code is available at https://github.com/mzhaoshuai/RLCF.

Figures (9)

• Figure 1: Feedback mechanism in zero-shot generalization with CLIP as the reward model.
• Figure 2: Fully TTA for zero-shot image classification with CLIP reward.Left: image encoder tuning. Right: prompt tuning. The pipelines of the two are the same except for the learnable parameters. A single test image is first augmented to produce multiple views, then only confident views with low-entropy predictions are selected. For each selected view, we sample the top-$K$ classes, calculate their rewards, and update the parameters using policy gradient.
• Figure 3: Fully TTA for zero-shot text-image retrieval with CLIP reward.
• Figure 4: Fully TTA for image captioning with CLIP reward.
• Figure 5: ECE and average accuracy on ImageNet-A/V2/R. Prompt tuning with CLIP-ViT-B/16.