Source-linked AI summary
Test-Time Prompt Tuning for Zero-Shot Generalization in Vision-Language Models
Manli Shu, Weili Nie, De-An Huang, Zhiding Yu, Tom Goldstein, Anima Anandkumar, Chaowei Xiao
TL;DR
Prompt tuning from downstream training data can restrict generalization and requires annotations, motivating adaptation that remains genuinely zero-shot. TPT tunes prompts from a single test sample using augmented-view consistency and confidence selection, improving CLIP across distribution shifts and matching few-shot methods on cross-dataset generalization. Its trade-offs are one-step backpropagation and increased inference memory from multiple augmented views.
Problem
Downstream prompt tuning can be limited beyond its training distribution and requires annotated task-specific data, restricting zero-shot generalization.
Method
TPT tunes adaptive prompts at test time from a single sample, minimizing entropy across augmented views while filtering low-confidence views.
Results
TPT improves CLIP’s zero-shot Top-1 accuracy by 3.6% on average under natural distribution shifts and performs on par with few-shot methods for cross-dataset generalization.
Takeaways & Limitations
TPT improves CLIP’s zero-shot generalization without additional training data or annotations across natural shifts and cross-dataset settings.
Takeaways & Limitations
TPT requires one-step backpropagation at test time and increases inference memory because it generates multiple augmented views.
Abstract
from arXiv · showhide
Pre-trained vision-language models (e.g., CLIP) have shown promising zero-shot generalization in many downstream tasks with properly designed text prompts. Instead of relying on hand-engineered prompts, recent works learn prompts using the training data from downstream tasks. While effective, training on domain-specific data reduces a model's generalization capability to unseen new domains. In this work, we propose test-time prompt tuning (TPT), a method that can learn adaptive prompts on the fly with a single test sample. For image classification, TPT optimizes the prompt by minimizing the entropy with confidence selection so that the model has consistent predictions across different augmented views of each test sample. In evaluating generalization to natural distribution shifts, TPT improves the zero-shot top-1 accuracy of CLIP by 3.6% on average, surpassing previous prompt tuning approaches that require additional task-specific training data. In evaluating cross-dataset generalization with unseen categories, TPT performs on par with the state-of-the-art approaches that use additional training data. Project page: https://azshue.github.io/TPT.
1 Introduction
TPT adapts prompts at test time from a single sample, preserving zero-shot operation while improving generalization across shifts and unseen categories.
- Motivation: Prompt tuning with downstream training data can improve prompts but may limit generalization beyond the training distribution and requires annotated data.Hand-crafted prompts enable zero-shot use, but learned prompts are tied to the tasks and distributions represented in training data.
- Approach: TPT tunes an adaptive prompt on the fly using only the given test sample, without task-specific training data or annotations.The method retains the zero-shot setting because no additional training data or annotations are used.
- Tasks: TPT is explored for image classification and context-dependent visual reasoning, with task-specific test-time tuning strategies.For visual reasoning, prompts are tuned to distinguish support sets before classifying the query image.
- Approach: For image classification, TPT minimizes marginal entropy across randomly augmented views so predictions remain consistent.Confidence selection excludes high-entropy, low-confidence augmentations from the consistency optimization.
- Results: 3.6% average Top-1 accuracy improvement over a hand-crafted prompt is reported for CLIP under natural distribution shifts.TPT reaches on-par accuracy with prompt tuning methods requiring additional ImageNet training data, and achieves a maximum 6.9% improvement on ImageNet-A.
- Results: For cross-dataset generalization with possibly unseen categories, TPT performs on par with a state-of-the-art few-shot prompt tuning method without additional data or annotations.The reported experiments cover natural distribution shifts, cross-dataset generalization, and context-dependent visual reasoning.
2 Related Work
Related work studies prompting, distribution-shift generalization, and test-time optimization, motivating prompt-only adaptation from individual test samples.
- Prompting for foundation models: Prompting applies foundation-model knowledge to downstream tasks in a zero-shot manner, while prompt tuning learns prompts from downstream data in the input embedding space.CoOp applies prompt tuning to CLIP using a collection of training data.
- Generalization under data distribution shifts: Distribution shifts arise when test and training distributions differ, including environmental variation and unseen concepts.Zero-shot generalization by models such as CLIP is presented as a direction toward reliable and generic machine learning systems.
- Test-time optimization: Test-time optimization adapts models to test samples on the fly, but designing an effective objective is a central challenge.Prior approaches include test-time training with self-supervised objectives and entropy minimization of batch-wise predictions.
- Test-time optimization: TENT requires more than one test sample for a non-trivial solution, whereas augmentation-based approaches bypass that multi-sample requirement.Other methods adapt batch-normalization statistics, feature extractors, or entire models, each involving different architectural or parameter choices.
- Test-time optimization: This work addresses test-time objective and parameter-group challenges by optimizing text prompts while keeping the model intact.The approach is positioned as a prompt-based alternative to adapting model parameters or normalization layers.
3 TPT: Test-Time Prompt Tuning
TPT adapts CLIP prompts at test time using only the current sample, addressing the limited out-of-distribution generalization of prompts learned from task-specific data. It uses task-specific tuning objectives for image classification and context-dependent visual reasoning.
- 3.1 Background: CLIP supplies aligned image and text representations through parallel encoders trained with a contrastive loss.This foundation model is used as the vision-language backbone for TPT.
- 3.1 Background: In Bongard-HOI, each test sample supplies support images exemplifying a concept and a query image whose concept presence must be inferred.The task represents human-object interactions through positive and negative support examples.
- 3.2 TPT: Test-Time Prompt Tuning: TPT tunes prompts on the fly from a single unlabeled test sample, preserving zero-shot generalization without additional training data or annotations.The prompt is adapted to each task and test sample.
- 3.2 TPT: Test-Time Prompt Tuning: For image classification, TPT generates augmented views and minimizes entropy of their averaged predictions to promote prediction consistency.The class-probability vector for each view is produced using the tuned prompt.
- 3.2 TPT: Test-Time Prompt Tuning: Confidence selection filters augmented views with high prediction entropy because random augmentations can remove content needed for correct classification.The method selects views whose prediction entropy falls below a threshold τ, adapted for each test sample.
- 3.2 TPT: Test-Time Prompt Tuning: For context-dependent visual reasoning, TPT learns an optimal binary label token and prompt from the support images rather than hand-crafting positive and negative labels.This explicitly represents the visual concept in text prompts to assist reasoning with language context.
- 3.2 TPT: Test-Time Prompt Tuning: For Bongard-HOI, TPT jointly tunes the prompt and binary class tokens on support images using cross-entropy loss.The support set provides the task context, while tuning remains purely at test time without training on similar tasks.
4 Experiments
The experiments evaluate TPT across natural distribution shifts, cross-dataset classification, and Bongard-HOI, comparing it with zero-shot and task-trained prompt or reasoning baselines. TPT uses test-time augmentation and entropy-based prompt optimization, while the reported evaluations emphasize generalization without task-specific training data.
- Experimental scope: The evaluation covers robustness to natural distribution shifts, cross-dataset generalization, and context-dependent visual reasoning.The experiments include four ImageNet variants, ten cross-dataset classification benchmarks, and Bongard-HOI.
- Baselines: TPT is compared with CoOp and CoCoOp, which are trained on ImageNet using 16-shot data per category, alongside zero-shot CLIP and prompt ensembles.Baseline CLIP, prompt ensembles, and TPT do not require training data in the natural-shift setup.
- Test-time tuning: TPT adapts a four-token prompt from one test image by generating 63 augmented views, selecting the 10% most confident predictions, and minimizing marginal entropy for one optimization step.The prompt is initialized from “a photo of a” and optimized with AdamW at learning rate 0.005.
- Natural distribution shifts: TPT achieves higher accuracy than prompt ensembles and few-shot prompt tuning methods on the evaluated natural-shift benchmarks, including CoCoOp.TPT can also be applied to prompts learned by CoOp or CoCoOp to further improve in-domain and OOD performance.
- Cross-dataset generalization: In cross-dataset evaluation, TPT matches ImageNet-trained CoCoOp from ImageNet to fine-grained datasets and consistently improves over zero-shot performance across ten disjoint datasets.Few-shot prompt tuning methods have negative average relative improvement in the no-overlap setting, whereas TPT shows improvement on every dataset.
- Context-dependent visual reasoning: For Bongard-HOI, TPT tunes prompts using the support images of each test sample before classifying its query image, but its split comparison is not strictly applicable because TPT uses no Bongard-HOI training data.The reported implementation tunes prompt prefixes and class tokens from scratch for 64 steps.
5 Ablation Study
The ablations identify prompt optimization and confidence selection as key TPT components, while efficiency analysis supports a one-step default with multiple augmented views.
- Parameter-group optimization: Prompt optimization achieves the most performance gain among the four tested CLIP parameter groups.The comparison includes the entire model, text encoder, visual encoder, and text prompt.
- Parameter-group optimization: Visual-encoder optimization performs worst, consistent with prior findings that image-encoder fine-tuning can distort pre-trained features.
- Confidence selection: Confidence selection improves baseline TPT by filtering noisy augmented views with low-confidence predictions.
- Confidence selection: The top-10% confidence threshold yields the highest average accuracy, and confidence selection generalizes to other entropy-based test-time optimization methods.
- Efficiency-accuracy trade-off: 1-step TPT is the default because more than 2 optimization steps provide no significant gain while inference time increases linearly.Increasing steps from 1 to 2 raises accuracy by 0.4%; more than 2 steps add no significant performance gain.
- Efficiency-accuracy trade-off: TPT accuracy plateaus around N = 64 augmented views, while N = 8 still provides over 2% average accuracy gain over zero-shot CLIP.
6 Conclusion
The paper presents TPT as a training-data-free method for adapting prompts from a single test sample and improving CLIP’s zero-shot generalization. Its main inference costs are one-step backpropagation and increased memory from multiple augmented views.
- TPT learns adaptive prompts on the fly from a single test sample without training data or annotations.
- TPT improves CLIP’s zero-shot generalization on natural distribution shifts and cross-dataset evaluation.
- TPT requires one-step test-time backpropagation and multiple augmented views, increasing inference memory cost.
- Future work includes applying TPT to other foundation models and designing test-time objectives suited to each model and downstream task.
A.1 Broader Impact
The paper motivates TPT as a way to better leverage foundation-model knowledge for zero-shot robustness and generalization, while preserving pretrained internal representations.
- TPT uses prompt tuning to leverage pretrained foundation-model knowledge without altering the models’ internal representations.
A.2 More Experiments.
Additional experiments examine variability, prompt combinations, and augmentation-only baselines. TPT retains its comparative advantage across these analyses, with lower average variability in several evaluations.
- Error-bar analysis: TPT achieves higher OOD accuracy than few-shot prompt tuning methods and has a smaller average standard deviation on natural distribution-shift benchmarks.
- Error-bar analysis: TPT also has a smaller average standard deviation than few-shot prompt tuning methods in cross-dataset generalization.
- Model ensembles: TPT is complementary to CoOp, CoCoOp, and hand-crafted prompts because it tunes an initialized prompt solely at test time.
- Model ensembles: Combining CoOp or CoCoOp with TPT produces more substantial improvement than conventional ensembles of existing baselines.
- Data-augmentation baselines: Augmentation-only averaged-prediction and majority-vote baselines do not achieve improvement comparable to TPT.
- Data-augmentation baselines: The limited augmentation-only performance suggests that designing an effective algorithm for using augmented images is non-trivial.
A.3 Apply Confidence Selection to Other Methods
Confidence selection improves entropy-based test-time optimization beyond TPT, including MEMO on a standard ResNet-50. The benefit increases as the confidence threshold becomes stricter.
- Generalization beyond TPT: Confidence selection also benefits entropy-based test-time optimization methods beyond TPT.It is evaluated as a plug-in component for methods with different architectures and optimized parameter groups.
- MEMO evaluation: MEMO is implemented on a standard ResNet-50 using its original hyper-parameter configurations.The comparison applies confidence selection using cutoff percentile ρ.
- MEMO evaluation: Confidence selection improves MEMO, with larger gains as the confidence threshold increases.A smaller cutoff percentile ρ corresponds to a higher confidence threshold.
A.4 License information of the assets used in this work.
The supplementary material documents asset licenses, prompting-strategy comparison criteria, computational budgets, experimental settings, and qualitative analysis descriptions. It also notes TPT's test-time computation overhead and reports implementation resources and hyper-parameter procedures.
- License information: ImageNet-A, ImageNetV2, ImageNet-R, and ImageNet-Sketch are used under the MIT License.Oxford-IIIT Pets is listed under CC BY-SA 4.0, while several other datasets are research-purpose access only.
- License information: The source code uses existing implementations of CLIP, CoOp, CoCoOp, and TENT under the MIT License.The code for the paper's experiments is also provided in the supplemental materials.
- Prompting-strategy comparisons: Prompting strategies for CLIP are categorized by whether prompts are learnable, require no training data, and are input-dependent.The training and testing budgets of these strategies are compared separately.
- Computational considerations: TPT's main computational overhead is one-step optimization with backpropagation through CLIP's text encoder, plus data augmentation.The image branch does not require backpropagation, allowing augmentation to be parallelized with little memory increase.
- Experimental settings: Natural-distribution-shift experiments use one 32GB V100 GPU, while other experiments use one 16GB V100 GPU.For ImageNet-variant OOD experiments, hyperparameters are selected using the original ImageNet validation set; fine-grained datasets use their validation splits.
- Qualitative analysis: The qualitative analysis compares CLIP predictions before and after TPT and shows probability distributions across 200 ImageNet-R classes for 64 augmented views.The figures include the test image, prediction bars, and distributions before and after TPT.