Source-linked AI summary

Align Your Prompts: Test-Time Prompting with Distribution Alignment for Zero-Shot Generalization

Jameel Hassan, Hanan Gani, Noor Hussein, Muhammad Uzair Khattak, Muzammal Naseer, Fahad Shahbaz Khan, Salman Khan

arXiv:2311.01459v2cs.CV

TL;DR

Test-time prompt tuning can adapt vision-language models using entropy minimization, but it does not explicitly align test and source distributions. PromptAlign uses a single test sample to tune multimodal prompts while aligning token statistics to proxy-source statistics. It reports gains on domain-generalization and cross-dataset benchmarks.

  • Problem

    Test-time prompt tuning with entropy minimization does not explicitly align CLIP's pre-training distribution with the test distribution, despite distribution shift being a source of performance degradation.

    Source: §1 · pp. 1–2
  • Method

    PromptAlign uses a single test sample to tune multimodal prompts by aligning its image-token means and variances with offline proxy-source statistics while minimizing prediction entropy.

    Source: §1 · pp. 2–3
  • Results

    PromptAlign improves the baseline by 3.08% on average across four domain-generalization datasets and achieves the best Top-1 accuracy on 8 of 10 cross-dataset datasets.

    Source: §1 · p. 2
  • Takeaways & Limitations

    The reported results show that proxy-source token-distribution alignment can be used with test-time multimodal prompting to improve zero-shot generalization across the evaluated benchmarks.

    Source: §1 · p. 2, §6 · p. 10

    PromptAlign relies on proxy-source dataset statistics and uses a single runtime update step; adaptation to object detection and segmentation remains a future direction.

    Source: §1 · p. 3, F Limitations and Future Directions · p. 18

Abstract

from arXiv · show

The promising zero-shot generalization of vision-language models such as CLIP has led to their adoption using prompt learning for numerous downstream tasks. Previous works have shown test-time prompt tuning using entropy minimization to adapt text prompts for unseen domains. While effective, this overlooks the key cause for performance degradation to unseen domains -- distribution shift. In this work, we explicitly handle this problem by aligning the out-of-distribution (OOD) test sample statistics to those of the source data using prompt tuning. We use a single test sample to adapt multi-modal prompts at test time by minimizing the feature distribution shift to bridge the gap in the test domain. Evaluating against the domain generalization benchmark, our method improves zero-shot top- 1 accuracy beyond existing prompt-learning techniques, with a 3.08% improvement over the baseline MaPLe. In cross-dataset generalization with unseen categories across 10 datasets, our method improves consistently across all datasets compared to the existing state-of-the-art. Our source code and models are available at https://jameelhassan.github.io/promptalign.

1 Introduction

PromptAlign addresses test-time distribution shift in vision-language models by aligning test-sample token statistics with proxy source statistics while tuning multimodal prompts. It reports improvements across domain-generalization and cross-dataset benchmarks.

1 Introduction

  • Prior test-time prompt tuning minimizes entropy but does not explicitly align pre-trained CLIP with the test-sample distribution.The introduction identifies distribution shift between training and test data as a source of performance degradation.Source: §1 · pp. 1–2
  • PromptAlign aligns test-sample image-token means and variances with offline proxy-source statistics, then jointly minimizes distribution-alignment and prediction-entropy losses by updating prompts in both CLIP branches.The method uses multiple augmented views of a single test sample and performs alignment across vision-backbone layers.Source: §1 · pp. 2–3
  • PromptAlign improves the baseline by 3.08% on average across four domain-generalization datasets and achieves the best Top-1 accuracy on 8 of 10 cross-dataset benchmarks.In cross-dataset generalization, it also reports an absolute average improvement of 1.82% over the existing state-of-the-art method using test-time prompt tuning.Source: §1 · p. 2
  • Because CLIP pre-training data are unavailable, the authors use ImageNet statistics as a proxy source distribution and report it as effective for CLIP.The paper frames this as a possible candidate for the source distribution.Source: §1 · p. 3

2 Related Work

Prior work learns prompts with downstream training data or adapts prompts at test time using entropy minimization. PromptAlign extends test-time prompting with multimodal token-distribution alignment using proxy-source statistics.

2 Related Work

  • CoOp learns text prompts, CoCoOp conditions text prompts on image embeddings, and MaPLe jointly learns deep prompts in CLIP's vision and text encoders; these approaches require downstream training data.TPT instead tunes text prompts at test time by minimizing averaged entropy across augmented views, but does not explicitly address distribution misalignment.Source: §2 · p. 3
  • Test-time adaptation methods either use auxiliary self-supervised tasks or adapt a pre-trained model through consistency constraints or source–test statistic alignment.Examples include entropy minimization and matching training and test batch-normalization statistics.Source: §2 · p. 3
  • Prior statistic-alignment methods use batches and may update all model parameters, leaving lightweight single-sample alignment for vision-language models unresolved.PromptAlign proposes multimodal test-time prompt tuning with training-distribution statistics, without altering pre-training.Source: §2 · p. 4
  • PromptAlign extends TPT with a distribution-alignment objective that uses proxy-source and single-test-sample token statistics to address domain shift.At test time, augmented views and text labels pass through the image and text encoders; alignment and entropy losses update multimodal prompts.Source: §2 · p. 4, Figure 2 · p. 4

3 Methodology

The methodology reviews CLIP and test-time prompt tuning, then introduces PromptAlign, which adapts multimodal prompts by aligning test-sample token statistics with source statistics.

3.1 Preliminaries

  • CLIP classifies images by comparing features from parallel image and text encoders, while prompt tuning adds learnable context to the encoders.The method uses deep, multimodal prompts on both branches.Source: §3.1 · p. 5
  • TPT adapts prompts for each test sample by minimizing entropy across confidence-filtered predictions from augmented views.The filtered views' mean class probabilities supply the entropy objective.Source: §3.1 · p. 5

3.2 PromptAlign

  • Because text-only TPT cannot align tokens across the vision branch, PromptAlign uses multimodal prompts to address test-time distribution shift.It aligns each test sample with the source distribution.Source: §3.2 · pp. 5–6
  • PromptAlign matches layer-wise means and variances of test tokens across augmented views to source statistics, then combines alignment and entropy losses to update prompts.Source statistics are precomputed offline using ImageNet as a proxy for CLIP's unavailable pretraining data.Source: §3.2 · p. 6
  • Table 1 evaluates MaPLe trained on ImageNet on datasets with domain shifts.Source: Table 1 · p. 7

4 Experiments

PromptAlign combines multimodal test-time prompting with token distribution alignment to evaluate zero-shot generalization across domain- and cross-dataset shifts. It improves on prompt-learning baselines across the reported benchmarks, with some variation by dataset and comparison.

4 Experiments

  • The evaluation covers four ImageNet OOD variants and PUG for domain generalization, plus ten cross-dataset classification benchmarks.The setup trains prompt-learning methods on ImageNet; PromptAlign updates prompts on both text and vision branches using a single test image and 64 views.Source: §4 · p. 7

4.1 Token distribution alignment in V-L models

  • Adding token distribution alignment to MaPLe+TPT improves average performance by 1.24% over MaPLe+TPT, while PromptAlign improves beyond vanilla MaPLe on every ImageNet variant.TPT alone raises average performance from 60.28% to 62.31%, but drops on ImageNet-Sketch.Source: §4.1 · p. 7

4.2 Domain Generalization

  • PromptAlign achieves the highest average Top-1 accuracy across ImageNet variants, improving 0.71% over CoOp+TPT; it also improves over the previous best on three of four variants and on PUG-ImageNet versus TPT.The reported exception among the four ImageNet variants is ImageNet-V2.Source: §4.2 · pp. 7–8, Table 2 · p. 8, Table 3 · p. 8

4.3 Cross-Dataset Transfer

  • Across cross-datasets, PromptAlign improves over MaPLe+TPT on every dataset, with an average improvement of 0.42%.The evaluation covers ten classification datasets spanning varied recognition tasks.Source: §4.3 · p. 8, Table 4 · p. 8

5 Ablation

The ablations examine the alignment loss, its scaling and statistics, prompt regularization, and test-time compute choices. They support combining alignment with entropy minimization and show a performance–resource trade-off in views and update steps.

Ablation scope

  • Most ablations use ImageNet-A for simplicity; trends across other ImageNet variants and datasets are presented in Appendix D.Source: §5 · p. 9

Token Distribution Alignment

  • Combining token distribution alignment with entropy minimization improves performance markedly, while alignment without entropy performs almost like vanilla MaPLe.The alignment loss acts as a regularizer alongside entropy loss, rather than promoting discriminative learning on its own.Source: §5 · p. 9

Loss variants for distribution alignment

  • L1 performs best among the tested L1, L2, and KL divergence choices for the distribution alignment objective.Further analysis is reported in supplementary Table E.Source: §5 · p. 9, Figure 4 · p. 9

Loss balancing factor β

  • ImageNet performance gains flatten near β = 100; the experiments therefore use β = 100, while MaPLe+TPT uses β = 0.Performance exceeds the TPT baseline even at smaller alignment-loss weights.Source: §5 · p. 9, Figure 3 · p. 9

Higher Order Statistics

  • Using fifth-order Central Moment Discrepancy statistics yields only a slight improvement, supporting mean- and variance-based alignment for this small-sample setting.The comparison concerns distribution alignment computed from a small set of samples.Source: §5 · p. 9, Table 6 · p. 10

Prompt Regularization

  • PromptAlign is more robust than naive prompt regularization in continuous test-time adaptation, despite alignment using each test sample individually.The regularization comparison uses a reduced learning rate of 1e−5 to prevent divergence.Source: §5 · pp. 9–10, Table 7 · p. 10

Trade-off between compute resources and performance

  • More augmented views and prompt-update steps improve accuracy, while latency is similar for both methods; average inference time is 0.216s with alignment and 0.197s without it.The reported configuration uses 64 views and one update step, considering memory use and latency.Source: §5 · p. 10, Figure 5 · p. 10

6 Conclusion

The conclusion presents PromptAlign as a test-time adaptation approach for zero-shot generalization that aligns test-sample statistics with source-data statistics using multimodal prompting. The appendix outlines additional results and ablation analyses.

6 Conclusion

  • PromptAlign aligns test-sample statistics with source-distribution statistics through multimodal prompting across transformer layers, and shows superior performance to existing CLIP zero-shot generalization methods in domain-generalization and cross-dataset evaluations.Source: §6 · p. 10

Appendix

  • The appendix reports additional cross-dataset, base-to-novel, and CLIP-backbone results, qualitative results, and ablations across datasets and parameter choices.Source: Appendix · p. 14

A Additional Implementation details

The implementation uses MaPLe with CLIP ViT-B/16. The layer-depth ablation finds that aligning tokens across the first three visual-encoder layers performs best on ImageNet-A.

A.1 Hardware and Software details

  • The implementation uses MaPLe with CLIP ViT-B/16 and runs on a single NVIDIA A100 40GB GPU using PyTorch.Source: §A.1 · p. 14

A.2 Effect of layer depth for token distribution alignment loss

  • The alignment loss uses corresponding tokens from 64 crops across the first three visual-encoder layers, the best-performing layer combination on ImageNet-A.This follows MaPLe's placement of deep prompts in layers 1–3.Source: §A.2 · p. 14

B.1 Cross dataset evaluations

Across cross-dataset, base-to-novel, and backbone evaluations, PromptAlign improves on competing test-time prompt-tuning approaches. Its alignment loss is most impactful in the lower visual-encoder layers.

B.1 Cross dataset evaluations

  • Alignment loss is most impactful in visual-encoder layers 1–3, where the prompts are located.The authors attribute this pattern to prompt placement in the first three layers.Source: Figure 6 · p. 14
  • PromptAlign outperforms CoOp+TPT and CoCoOp+TPT on 9 of 10 datasets in cross-dataset evaluation.The comparison is reported alongside the finding that integrating TPT decreases CoOp and CoCoOp performance in most cases.Source: §B.1 · p. 14, Table 8 · p. 15

B.2 Base to Novel Generalization

  • PromptAlign shows consistent improvement over TPT in the base-to-novel generalization benchmark.The table caption describes the comparison as base-to-novel generalization against previous methods.Source: Table 9 · p. 15

B.3 Different CLIP backbone architectures

  • On ViT-B/32, distribution alignment improves on TPT on average in both domain-generalization and cross-dataset settings.The cited tables compare methods trained on ImageNet and evaluated on shifted-domain or cross-dataset benchmarks.Source: §B.3 · p. 15, Table 10 · p. 16, Table 11 · p. 16

C Qualitative Results

On ImageNet-A test samples, PromptAlign’s visual attention favors discriminative regions and reduces attention to background regions compared with TPT.

C Qualitative Results

  • PromptAlign focuses more strongly on discriminative regions and less on background regions than TPT in the visual encoder’s last-block attention maps.The comparison covers two ImageNet-A test images; attention to the salient object increases in both examples.Source: C Qualitative Results · p. 16

D.1 Proxy source dataset

The analyses examine proxy-source choice, the quality of test-sample statistics, and parameter trends. ImageNet is generally a stronger proxy than the tested target-dataset training sets, while more test samples and more prompt updates or augmented views improve reported accuracy.

D.1 Proxy source dataset

  • A subset of LAION400M used for source statistics improves the distribution-alignment objective and overall performance.The authors describe LAION400M as a valid CLIP training dataset.Source: §D.1 · p. 16, Table 12 · p. 17
  • ImageNet as the proxy source gives better performance across the tested datasets except UCF101, where performance is the same.The comparison also shows a similar trend with EVA-CLIP and different training recipes.Source: §D.1 · p. 16, Table 13 · p. 17

D.2 Effect of test sample statistics

  • Using a bag of five same-class samples improves alignment performance by better approximating test-sample statistics, but actual test-time adaptation has only one sample.The class labels are kept unknown in the bag-of-samples experiment.Source: §D.2 · p. 17, Table 14 · p. 17

D.3 Trend across parameter choices

  • Accuracy increases with more augmented views and prompt-update steps, while the loss-scaling factor β plateaus close to 100.β is set to 100 using the ImageNet validation set and held fixed for all experiments.Source: §D.3 · pp. 17–18, Figure 8 · p. 18

E Analysis of losses for distribution alignment

The loss comparison favors L1 for test-time distribution alignment: it improves over the other tested losses, with larger gradients than L2 or KL divergence.

E Analysis of losses for distribution alignment

  • The L1 alignment loss yields better alignment than L2 or KL divergence, and all three losses improve over test-time prompting alone.The L1 gradient is nearly 80 times the L2 gradient and about 4 times the KL-divergence gradient; the authors attribute improved alignment primarily to its higher gradient values.Source: E Analysis of losses for distribution… · p. 18, Table 15 · p. 18

F Limitations and Future Directions

PromptAlign relies on proxy source-dataset statistics and uses a single runtime update step; adapting it to detection and segmentation is proposed as future work.

F Limitations and Future Directions

Loading 2311.01459v2…