SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathcal {L} _ {\text {a l i g n}} = \frac {1}{L} \sum_ {l = 1} ^ {L} \left(\| \boldsymbol {\mu} _ {l} (\mathcal {T}; \boldsymbol {p}) - \hat {\boldsymbol {\mu}} _ {l} \| _ {1} + \| \boldsymbol {\sigma} _ {l} ^ {2} (\mathcal {T}; \boldsymbol {p}) - \hat {\boldsymbol {\sigma}} _ {l} ^ {2} \| _ {1}\right). \tag {5}
$$
As shown above, we use  $L_{1}$  loss to enforce the distribution alignment of the test sample with the source distribution. The alignment loss  $\mathcal{L}_{\mathrm{align}}$  is added to the entropy loss (Eq. 1) to obtain the final objective  $\mathcal{L}_{\mathrm{final}}$ . Finally, the combined objective is optimized to update the prompts  $p$ .
$$
\mathcal {L} _ {\text {f i n a l}} = \mathcal {L} _ {\text {e n t r o p y}} + \beta \mathcal {L} _ {\text {a l i g n}}, \tag {6}
$$
where the hyperparameter  $\beta$  controls the contribution of alignment loss to the overall objective function. Our alignment loss in Eq. 6 bridges the gap between source and token distributions by updating the prompts to align them with the source distribution.
Discussion on  $\mathcal{L}_{\mathrm{final}}$ . As derived above, our test-time loss combines the entropy minimization  $(\mathcal{L}_{\mathrm{entropy}})$  and distribution alignment objectives  $(\mathcal{L}_{\mathrm{align}})$  together. The  $\mathcal{L}_{\mathrm{entropy}}$  objective enforces prediction consistency among different views of a sample, leading to robustness against various variations that could possibly occur at test time. On the other hand,  $\mathcal{L}_{\mathrm{align}}$  effectively narrows the domain shift and brings the test sample distribution closer to the pre-trained CLIP distribution space, which enforces CLIP to better understand the test sample. The combination of these loss objectives harmonically adapts CLIP to be robust to different sample variations and at the same time, enhances CLIP's understanding of the underlying test sample domain for better generalization.
Table 1: Effect of token distribution alignment strategy for domain generalization. The base model MaPLe is trained on ImageNet and evaluated on datasets with domain shifts.
Imagenet V2	Imagenet Sketch	Imagenet A	Imagenet R	OOD Avg.
MaPLe	64.07	49.15	50.90	76.98	60.28
MaPLe+TPT	64.87	48.16	58.08	78.12	62.31
PromptAlign	65.29	50.23	59.37	79.33	63.55
Datasets. For the domain generalization setting, we evaluate the four out-of-distribution (OOD) variants of ImageNet [7]; ImageNetV2 [31], ImageNet-Sketch [38], ImageNet-A [16] and ImageNet-R [15]. We also evaluate the domain generalization setting on the recent challenging Photorealistic Unreal Graphics (PUG) dataset [3], comprising different textures, backgrounds, sizes and orientations. In the cross-dataset evaluation, we follow TPT [34] and evaluate the performance of methods on 10 image classification datasets covering a wide range of visual recognition tasks. This includes one generic-objects dataset Caltech101 [10]; five fine-grained datasets OxfordPets [28], StanfordCars [20], Flowers102 [27], Food101 [4] and FGVC-Aircraft [24], which contain images of animals, flowers and transportation; and four datasets of scenes, textures, satellite imagery and human actions - SUN397 [36], DTD [6], EUROSAT [13] and UCF101 [35] respectively.
Baselines. We evaluate PromptAlign with existing few-shot prompt learning methods for adapting CLIP including CoOp [48] and CoCoOp [47], and TPT [34] method. MaPLe [18] is a multi-modal prompt learning baseline, which adapts CLIP by learning deep prompts on both the text and vision branches. TPT is a test-time prompt tuning method, that tunes the prompt at test time per input sample achieving state-of-the-art performance in prompt learning when combined with CoOp.
Implementation details. Following MaPLe [18], we train on ImageNet using 16-shot training data with 2 prompt tokens for a depth of 3 layers. We optimize the prompts on both the text and vision branches using a single test image. We obtain 63 augmented views using random resized crops and horizontal flip augmentations to construct a batch of 64 images including the original image to mimic the setting of TPT. From the 64 predictions, we obtain top  $10\%$  confident predictions based on the lowest entropy and compute the average prediction probability. We compute the token distribution alignment loss between the tokens of all 64 images. We optimize the prompts to minimize the combined loss of average prediction entropy and the token distribution alignment loss using the AdamW optimizer. We use a learning rate of  $5e^{-4}$  for the fine-grained datasets Flowers102, OxfordPets, Food101, SUN397, FGVCAircraft, and EuroSAT and a learning rate of 0.04 for the rest of the datasets, and set the loss scale factor  $\beta$  equal to 100. Further details are listed in Appendix A.
4.1 Token distribution alignment in V-L models
We first evaluate the performance of the token distribution alignment strategy over the baseline for test-time prompt tuning. Since the alignment can only be performed on the vision branch, we explicitly compare PromptAlign to vanilla MaPLe, and MaPLe with TPT using entropy minimization. Table 1 presents the results in domain generalization across ImageNet variants. TPT improves the performance of MaPLe from  $60.28\%$  to  $62.31\%$  on average, while the incorporation of token distribution alignment further improves the performance by  $1.24\%$ . This indicates that explicitly aligning the distribution of training and testing data enhances CLIP's generalization. Moreover, TPT is intended to improve test-time adaptation as a plug-in module, but its effectiveness is not uniform, as observed from the drop in performance on the ImageNet-Sketch dataset. On the other hand, PromptAlign does not show any performance degradation and improves beyond vanilla MaPLe consistently on all datasets. We further present qualitative results of PromptAlign in Appendix C.
4.2 Domain Generalization
Compared to CLIP, all test-time adaptation methods demonstrate better performance, indicating that V-L models benefit from test-time adaptation approaches (Table 2). PromptAlign achieves the highest Top-1 accuracy averaged across all ImageNet variants. In comparison to the previous state-of-the-art CoOp+TPT, PromptAlign achieves an average improvement of  $0.71\%$ . Out of the four different
Table 2: Comparison of PromptAlign in domain generalization setting. Prompt learning methods are trained on ImageNet and evaluated on datasets with domain shifts.
Imagenet V2	Imagenet Sketch	Imagenet A	Imagenet R	OOD Avg.
CLIP [30]	60.86	46.09	47.87	73.98	57.20
CLIP+TPT [34]	64.35	47.94	54.77	77.06	60.81
CoOp [48]	64.20	47.99	49.71	75.21	59.28
CoOp+TPT [34]	66.83	49.29	57.95	77.27	62.84
Co-CoOp [47]	64.07	48.75	50.63	76.18	59.91
Co-CoOp+TPT [34]	64.85	48.27	58.47	78.65	62.61
PromptAlign	65.29	50.23	59.37	79.33	63.55
Table 3: Effect of token distribution alignment strategy for domain generalization. The base model MaPLe is trained on ImageNet and evaluated on PUG-ImageNet.
Camera(Yaw/ Pitch/ Roll)	Pose(Yaw/ Pitch/ Roll)	Scale	Texture	Lighting	Worlds
MaPLe [18]	48.73/ 39.93/ 32.13	48.10/ 28.40/ 27.80	46.90	37.90	15.50	32.13
MaPLe+TPT	57.04/ 45.99/ 39.23	56.26/ 35.64/ 33.26	54.87	43.73	22.52	42.00
PromptAlign	58.14/ 46.93/ 40.45	57.43/ 36.31/ 34.32	56.18	44.97	23.06	43.24
Table 4: Comparison of PromptAlign in cross-dataset evaluation. Prompt learning methods are trained on ImageNet and evaluated on cross-datasets.
Caltech	Pets	Cars	Flowers	Food101	Aircraft	SUN397	DTD	EuroSAT	UCF101	Average
CLIP [30]	93.35	88.25	65.48	67.44	83.65	23.67	62.59	44.27	42.01	65.13	63.58
CLIP+TPT [34]	94.16	87.79	66.87	68.98	84.67	24.78	65.50	47.75	42.44	68.04	65.10
CoOp [48]	93.70	89.14	64.51	68.71	85.30	18.47	64.15	41.92	46.39	66.55	63.88
CoCoOp [47]	93.79	90.46	64.90	70.85	83.97	22.29	66.89	45.45	39.23	68.44	64.63
ProDA [46]	86.70	88.20	60.10	77.50	80.80	22.20	-	50.90	58.50	-	65.62
MaPLe	93.53	90.49	65.57	72.23	86.20	24.74	67.01	46.49	48.06	68.69	66.30
MaPLe+TPT	93.59	90.72	66.50	72.37	86.64	24.70	67.54	45.87	47.80	69.19	66.50
PromptAlign	94.01	90.76	68.50	72.39	86.65	24.80	67.54	47.24	47.86	69.47	66.92
domains, our method improves beyond the previous best, except in Imagenet-V2. We hypothesize this is due to the extensive training of CoOp on ImageNet which has a very similar distribution to ImageNet-V2. We further evaluate the effectiveness of our method in domain generalization on the challenging PUG dataset as reported in Table 3. PromptAlign consistently improves upon TPT here as well. These results demonstrate the effectiveness of PromptAlign, which adapts CLIP simultaneously for both robustness against view changes and alignment of train-test set distribution.
4.3 Cross-Dataset Transfer
In Table 4, we evaluate the performance of our method in generalizing across diverse cross-datasets in comparison with existing state-of-the-art methods using prompt learning. PromptAlign provides consistent improvements and outperforms the previous best method MaPLe combined with TPT in all cross-datasets with an average improvement of  $0.42\%$ . In comparison to prompt learning methods, it can be observed that both CoOp and CoCoOp on average are inferior to zero-shot CLIP+TPT, while PromptAlign consistently demonstrates improvements, which validates the necessity of token distribution alignment for preserving CLIP's generalizability. This further affirms that PromptAlign harmonically combines the effect of entropy minimization and token distribution alignment in V-L models. We also analyze the effect of a better proxy source dataset, using a subset of LAION400M in Appendix D.1, proving the effectiveness of distribution alignment with source dataset statistics. Furthermore, we note that as opposed to our method, existing prompt learning methods with TPT are not consistent across cross-datasets, which we discuss in detail in Appendix B.
Method	Entropy loss	Distribution alignment	Top-1 Acc.
MaPLe	X	X	50.90
MaPLe+TPT	✓	X	58.08
PromptAlign†	X	✓	50.85
PromptAlign	✓	✓	59.37
Table 5: Analysis on the alignment and entropy minimization loss. The average of Top-1 accuracy  $(\%)$  across three seeds is reported. PromptAlign denotes our method excluding the entropy minimization loss.
5 Ablation
We perform ablative analysis on various components of PromptAlign. Unless stated otherwise, we show ablations on ImageNet-A dataset, the smallest domain generalization variant for simplicity. Detailed ablation trends across other ImageNet variants and datasets are presented in Appendix D.
Token Distribution Alignment. Table 5 summarizes the effect of token distribution alignment in PromptAlign. It can be observed that removing the entropy loss results in almost the same performance as vanilla MaPLe. Whereas combining alignment and entropy losses improves the performance remarkably. This is because the distribution alignment loss is a regularizer that functions together with the self-entropy loss. Since distribution alignment does not promote any discriminative learning as opposed to entropy loss, it is not expected to improve test-time adaptation on its own.
Loss variants for distribution alignment. We show the variation in performance with different loss choices for distribution alignment objective. Figure 4 compares the three loss choices  $L_{1}$ ,  $L_{2}$ , and KL divergence. We observe that  $L_{1}$  constraint performs best for the distribution alignment objective. Further analysis is shown in Table E of supplementary material.
Loss balancing factor  $\beta$ . We ablate the loss scale factor  $\beta$  to analyze the significance of the distribution alignment loss. The performance improves over the baseline model with TPT as depicted in Figure 3, even with a smaller weight scale and gradually plateaus close to a scale of 100. We use ImageNet validation set for this ablation and choose  $\beta = 100$  across all experiments. We find a similar trend of the scaling factor on other datasets, which is discussed in Appendix D. For the MaPLe+TPT model, the value of  $\beta$  is 0 as there is no alignment objective in this case.
Higher Order Statistics. Since we align distributions computed using a small set of samples, we use relatively simpler (mean and variance) measures for our alignment strategy over higher order statistical measures. We evaluate this choice for the distribution alignment in PromptAlign using the  $5^{\text{th}}$  order statistics of Central Moment Discrepancy measure [43] in Table 6. Our observation reveals that the utilization of higher order statistics results in only a slight improvement implies that there is minimal distinction between opting for higher order statistics validating mean and variance based approach as a suitable measure for alignment in our case.
Prompt Regularization. We compare PromptAlign to a naive prompt regularization method in a continuous test-time adaptation setting, in which the model gradually adapts to the incoming distribution. We regularize the prompts in relation to the previous prompts, denoted by PromptReg in Table 7. We also lower the learning rate to  $1e^{-5}$  to prevent diverging in the continuous setting. We
Figure 3: Effect of the loss scaling factor  $\beta$  on ImageNet. Scale factor plateaus after  $\beta = 100$ .
Figure 4: Effect of performance on the choice of loss function for distribution alignment.
show that the naive regularization setting, even in a continuous setting is not robust as the proposed alignment strategy, which solely aligns based on a single test sample, individually.
Table 6: Comparison of the effect of higher order statistics for distribution alignment. PromptAlign† utilizes Central Moment Discrepancy measure for the distribution alignment loss.
Caltech	Pets	Cars	Flowers	Food101	Aircraft	SUN397	DTD	EuroSAT	UCF101	Average
PromptAlign	72.39	47.24	90.76	68.5	69.47	94.01	86.65	67.54	24.8	47.86	66.92
PromptAlign†	72.36	47.24	90.74	68.3	69.48	94.04	86.65	67.85	24.75	48.5	66.99
Table 7: Comparison of PromptAlign with naive prompt regularization.
Target
Caltech	Pets	Cars	Flowers	Food101	Aircraft	SUN397	DTD	EuroSAT	UCF101	Average
MaPLe	93.53	90.49	65.57	72.23	86.20	24.74	67.01	46.49	48.06	68.69	66.30
PromptReg	93.71	90.60	65.02	72.05	86.05	22.32	65.46	46.24	19.03	68.23	62.87
PromptAlign	94.01	90.76	68.50	72.39	86.65	24.80	67.54	47.24	47.86	69.47	66.92
Trade-off between compute resources and performance. Figure 5(a) demonstrates the improvement in performance as the number of augmented views at test time increases. PromptAlign shows better improvement with more augmented views, owing to a better approximation of the test sample distribution. While MaPLe+TPT performance plateaus around 64 augmented views, our method with alignment loss shows further scope for improvement given more augmented views. Figure 5(b) shows the variation in performance with the number of prompt update steps. We note that, with more update steps, PromptAlign can consistently better adapt to the test sample in comparison to MaPLe+TPT alone. In Figure 5(c) we analyze the computational overhead of PromptAlign. PromptAlign shows no additional compute cost in terms of latency. The average inference time per sample across three seeds with and without the distribution alignment strategy are  $0.216s$  and  $0.197s$ . Considering the memory usage and increase in latency, we use 64 views and a single-step update following TPT [34].
Figure 5: Analysis of compute resource constraints on performance. (a) The Top-1 accuracy increases with the number of augmented views. (b) The Top-1 accuracy improves consistently with the number of prompt update steps. (c) Impact on latency with the number of prompt update steps is similar for both methods.
In this paper, we introduce PromptAlign, a novel approach for enhancing test-time adaptation of Vision-Language (V-L) models for zero-shot generalization. Our proposed approach bridges the gap between the test sample and source distributions by explicitly aligning the test sample statistics with that of the source data distribution through token distribution alignment. To achieve this, we incorporate multi-modal prompting to facilitate the alignment of token distributions across the transformer layers during test time. Through extensive experiments, PromptAlign demonstrates superior performance over existing state-of-the-art CLIP zero-shot generalization methods in domain generalization and cross-dataset evaluation settings.
The computational resources were provided by the National Academic Infrastructure for Supercomputing in Sweden (NAISS), partially funded by the Swedish Research Council through grant agreement No. 2022-06725, and by the Berzelius resource, provided by Knut and Alice Wallenberg Foundation at the National Supercomputer Center.
In the following, we provide further details on our implementation under Section A. In Section B, we provide additional results of our proposed PromptAlign method on cross-dataset generalization, base-to-novel generalization and for different CLIP backbone architectures; and Section C shows additional qualitative results. Finally, Section D presents ablation trends across diverse datasets for the effect of the proxy dataset, the effect of test sample statistics and different parameter choices.
A Additional Implementation details
A.1 Hardware and Software details
We implement our method with the MaPLe [18] multi-modal prompting model with the CLIP ViT-B/16 backbone architecture. Our models were implemented on a single NVIDIA A100 40GB GPU using the PyTorch framework.
A.2 Effect of layer depth for token distribution alignment loss
The token distribution alignment loss is computed from all the corresponding tokens of 64 crops spanning the first 3 layers of the visual encoder. This is chosen in line with the baseline multimodal prompt learning model [18] where the deep prompts are present in layers 1-3. In order to analyze the impact of different layer depths on the alignment loss, we conduct an ablation study. Figure 6 displays the performance comparison of alignment losses obtained from various layer combinations using the ImageNet-A dataset. As observed in the Figure 6, enforcing the token alignment loss across the first three layers of the vision encoder of CLIP provides the best performance.
Figure 6: Effect of alignment depth on performance: We examine how the performance of PromptAlign is affected by altering the number of alignment layers that contribute to the alignment loss. Our findings indicate that the token distribution alignment loss is most impactful when applied to the lower layers (specifically, layer  $1 \rightarrow 3$ ). We hypothesize that the effectiveness of the token distribution alignment loss in these layers can be attributed to the presence of our prompts, which are located within the first three layers of the visual encoder.
B Additional results
B.1 Cross dataset evaluations
In Table 4 of the main paper, we compared the performance of PromptAlign with the current state-of-the-art prompt learning methods, CoOp and CoCoOp, for cross-dataset generalization. Additionally, in Table 8, we evaluate existing prompt learning methods that have been integrated with TPT and present further comparisons of our method with them. As shown in Table 8, PromptAlign outperforms both CoOp+TPT and CoCoOp+TPT in 9 out of 10 datasets. It is worth noting that when integrated with TPT, the performance of CoOp and CoCoOp decreases for the majority of cases. In contrast, PromptAlign consistently improves performance across different datasets.
B.2 Base to Novel Generalization
We also evaluate the performance of our method on the base-to-novel generalization benchmark setting. On average, our distribution alignment strategy improves the model performance by nearly  $1\%$  in the base classes and by  $0.79\%$  on novel classes. We observe that TPT alone drops the performance of the model in some instances such as for OxfordPets, Eurosat and UCF101. On such instances, PromptAlign regulates the drop in accuracy caused by TPT, while in most cases achieves better performance in novel classes of these datasets.
Table 8: Comparison of PromptAlign to vanilla TPT in cross-dataset evaluation. Prompt learning methods are trained on ImageNet and evaluated on cross-datasets at test-time.
Target
Caltech101	OxfordPets	StanfordCars	Flowers102	Food101	Aircraft	SUN397	DTD	EuroSAT	UCF101	Average
CoOp [48]	93.70	89.14	64.51	68.71	85.30	18.47	64.15	41.92	46.39	66.55	63.88
CoOp+TPT [34]	93.15	89.48	66.77	68.48	86.48	20.51	66.06	43.32	37.73	68.91	64.08
CoCoOp [47]	93.79	90.46	64.90	70.85	83.97	22.29	66.89	45.45	39.23	68.44	64.63
CoCoOp+TPT [34]	88.57	85.33	59.68	55.31	80.64	16.89	60.24	38.93	48.55	63.35	59.75
PromptAlign	94.01	90.76	68.50	72.39	86.65	24.80	67.54	47.24	47.86	69.47	66.92
Table 9: Comparison on Base-to-novel generalization of PromptAlign with previous methods. PromptAlign shows consistent improvement over TPT.
Dataset	ProDA [46]	MaPLe [18]	MaPLe+TPT [34]	PromptAlign (Ours)
Avg. on 11 datasets	Base	81.56	82.24	82.16	83.19 (+0.95)
Novel	72.30	75.09	74.95	75.88 (+0.79)
ImageNet	Base	75.40	76.67	77.73	78.26
Novel	70.23	70.54	72.24	72.59
Caltech101	Base	98.27	98.00	98.54	98.60
Novel	93.23	94.27	94.29	94.50
OxfordPets	Base	95.43	95.43	95.23	95.38
Novel	97.83	97.80	97.37	97.56
Stanford Cars	Base	74.70	72.90	74.00	75.02
Novel	71.20	73.97	75.20	75.71
Flowers102	Base	97.70	95.93	96.24	96.61
Novel	68.68	72.40	72.10	72.34
Food101	Base	90.30	90.70	91.13	91.63
Novel	88.57	92.07	92.03	92.68
FGVC Aircraft	Base	36.90	37.27	34.31	37.21
Novel	34.13	35.53	35.81	37.27
SUN397	Base	78.67	80.80	81.15	81.57
Novel	76.93	78.70	79.18	79.48
DTD	Base	80.67	80.30	82.20	82.60
Novel	56.48	59.23	59.91	60.55
Eurosat	Base	83.90	93.63	91.02	94.10
Novel	66.00	72.87	68.96	72.71
UCF101	Base	85.23	82.97	82.23	84.11
Novel	71.97	78.57	77.34	79.30
B.3 Different CLIP backbone architectures
We also evaluate our method on the ViT-B/32 backbone architecture variant of CLIP. Here, we train CLIP ViT-B/32 in the original MaPLe setting without any hyperparameter tuning specific to MaPLe. Our distribution alignment strategy improves upon TPT on average in both the domain generalization setting and cross-dataset settings as reported in Table 10 and Table 11.
Table 10: Comparison of PromptAlign in domain generalization setting for ViT-B/32. Prompt learning methods are trained on ImageNet and evaluated on datasets with domain shifts.
Imagenet V2	Imagenet Sketch	Imagenet A	Imagenet R	OOD Avg.
MaPLe	57.63	42.15	32.12	67.64	49.89
MaPLe+TPT	60.01	43.77	37.52	71.11	53.10
PromptAlign	60.43	44.24	38.02	71.44	53.53
Table 11: Comparison of PromptAlign in cross-dataset evaluation on ViT-B/32. Prompt learning methods are trained on ImageNet and evaluated on cross-datasets.
Caltech	Pets	Cars	Flowers	Food101	Aircraft	SUN397	DTD	EuroSAT	UCF101	Average
MaPLe [18]	92.50	88.13	59.93	65.33	81.00	17.53	65.00	41.70	40.80	63.63	61.56
MaPLe+TPT [34]	91.44	88.47	59.35	66.08	82.08	18.71	66.07	40.01	39.67	64.40	61.63
PromptAlign	92.10	88.44	63.48	66.14	82.07	18.76	66.08	42.54	39.68	65.57	62.49
C Qualitative Results
We generate attention map visualizations from the heads of the last block of the CLIP visual encoder on the test samples from the Imagenet-A dataset. In Figure 7, we compare the attention maps produced by PromptAlign and TPT. Our findings reveal that PromptAlign exhibits a stronger focus on discriminative regions while reducing attention on background regions compared to TPT. For the first test sample (on the left), PromptAlign intensifies attention on the object of interest and diminishes attention on the background regions. Similarly, in the second image (on the right), PromptAlign increases attention toward the salient object.
D Extended Ablations
D.1 Proxy source dataset
We assess the effectiveness of our alignment strategy, in Table 12 by employing LAION400M [33] to compute the source data statistics. We utilize a subset of 2 Million images* from the LAION400M dataset (denoted as PromptAlign†). We show that using LAION400M to compute the source dataset statistics enhances the distribution alignment objective, and thus the overall performance in Table 12. This emphasises the effectiveness of distribution alignment, since LAION400M has been shown as a valid training dataset for CLIP [5].