SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jameel Hassan $^{1*}$  Hanan Gani $^{1*}$  Noor Hussein $^{1}$  Muhammad Uzair Khattak $^{1}$
Muzammal Naseer $^{1}$  Fahad Shahbaz Khan $^{1,2}$  Salman Khan $^{1,3}$
$^{1}$ Mohamed Bin Zayed University of AI  $^{2}$ Linköping University  $^{3}$ Australian National University
{jameel.hassan, hanan.ghani, noor.husein, uzair.khattak
muzammal.naseer, fahad.khan, salman.khan} @mbzuai.ac.ae
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/.
Deep neural networks (DNNs) have outperformed humans in numerous visual recognition tasks [12, 8]. However, their impressive performance generally holds in the case when test data originate from the same distribution as the training data. In most real-world applications, the train and test data distributions can significantly differ due to factors such as natural variations or changes in sensing equipment. The sensitivity of models to such unseen distributional shifts during inference results in their performance degradation [14, 31, 29]. A plethora of previous works [36, 39, 37] have explored test-time adaptation as a mechanism to overcome the crucial problem of distribution shift in test data. However, test-time adaptation has been minimally explored for the increasingly popular foundation models, which are large DNNs trained on massive vision-language datasets [30, 17, 44, 40, 41].
Foundation models emerging at the intersection of various modalities such as vision, language, and audio, are proving to be effective in numerous downstream applications. Among these models, Vision-Language (V-L) model CLIP (Contrastive Language-Image Pretraining) [30] has been pre-trained on large-scale image-text pairs from the web, and can generalize well to zero-shot recognition tasks. The CLIP model has parallel vision and language encoding branches and the similarities between the embeddings obtained from both the branches are used to classify an input image. At inference, a handcrafted prompt such as 'a photo of a <CLS>' is used as a query for the text encoder. However, adapting CLIP efficiently to specific downstream tasks is still a challenging problem. Naively fine-tuning such models poses the risk of losing their inherent generalization abilities [21]. Instead, recent
(a) Proposed PromptAlign method
(b) Performance comparison in cross datasets evaluation
Figure 1: (a) PromptAlign matches the distribution statistics  $\mu_l(\mathcal{T};\pmb{p})$ ,  $\sigma_l^2 (\mathcal{T};\pmb{p})$ , obtained from multiple augmented views of a single test sample, with the source data distribution statistics  $\hat{\mu}_l$ ,  $\hat{\sigma}_l^2$ . This effectively brings the test sample closer to the distribution of the source data, where the domain shift is denoted by  $\triangle_1 \rightarrow \triangle_2$ .  $\mathcal{T}$  denotes the distribution of the test sample,  $\pmb{p}$  represents the prompts that are updated and  $l$  refers to the vision-backbone layers. (b) Owing to the distribution matching via prompts, PromptAlign surpasses the existing state-of-the-art prompt learning approaches on 8 out of 10 datasets in cross-dataset generalization benchmarks.
methods have shown the promise of prompt learning on training data [48, 47, 18, 19, 42] as opposed to handcrafted prompts, allowing the model to better adapt to the training data distribution.
Existing prompt learning approaches are deployed at the training phase to learn representative prompts based on the training data for the downstream task. This conventional approach does not explicitly handle the distribution shift in the test set. Leveraging the capability of prompt learning, a recent approach TPT [34] performs test-time adaptation by tuning the text prompts on the fly to adapt the model to the test sample. For image classification, the model updates the prompts by minimizing the entropy of the top confident samples which are obtained using different augmented views. However, TPT does not explicitly align the pre-trained CLIP to become aware of the test sample distribution.
For the effective test-time adaptation of V-L foundation models, it is crucial to bridge the distribution gap between the pre-training dataset and the downstream evaluation set for high zero-shot generalization. To this end, we propose PromptAlign, a test-time token distribution alignment strategy using prompt learning. The TPT setting, which tunes the prompts on the text encoder branch alone poses an architectural limitation in performing distribution alignment of tokens. Further, at test time, there is no knowledge transfer between the text and vision branches in the intermediate stages. Since the text encoder features will be static given a dataset (as the input is the same class labels), distribution alignment of tokens can only be performed on the vision branch. Given these constraints and to further extend the strength of prompts for test time adaptation, we propose to employ a multi-modal prompt learning model (MaPLe) [18] for distribution alignment. PromptAlign explicitly aligns the mean and variances of the image token embeddings of a proxy source dataset, computed offline, with the image token embeddings of the test sample. We extend TPT with the token alignment strategy, which enforces to bridge the distribution shift in the test data (Fig. 1 a). For each input test sample, we obtain randomly augmented views that are fed into the model to obtain the token embedding statistics. Without any noticeable compute overhead, we update the prompts on both the text and vision branches of CLIP to minimize jointly the feature distribution shift and entropy of predictions.
We extensively demonstrate the effectiveness of PromptAlign by evaluating the zero-shot generalization on two representative benchmarks: domain generalization and cross-dataset generalization. In the domain generalization setting, our method improves the baseline model by  $3.08\%$  on average across the four datasets and has the highest average Top-1 accuracy compared to existing state-of-the-art methods. In the cross-dataset setting, our method achieves an absolute average improvement of  $1.82\%$  over the existing state-of-the-art method which uses test-time prompt tuning, while attaining the best Top-1 accuracy in 8 out of the 10 datasets. Our contributions can be summarized as follows:
- Given only a single test sample, we introduce a distribution alignment strategy for V-L models to improve test-time adaptation. The distribution-aware pre-trained CLIP effectively narrows the distribution gap on the test domains. To the best of our knowledge, this is the first study to explore the potential of distribution alignment in V-L models at test time.
- The proposed strategy formulates a distribution alignment loss that utilizes offline computed source data statistics to encourage the test sample token distributions to be aligned with the source data token distributions. We harmonically combine the benefits of token distribution alignment and entropy minimization using a multi-modal prompt learning approach.
- Since CLIP-pre-training data is not publicly released, we study the statistics of ImageNet as a possible candidate for the source distribution, and our empirical results show that ImageNet is an effective proxy source dataset for large-scale V-L models such as CLIP.
- We validate our method PromptAlign through extensive experiments in domain generalization and cross-dataset benchmarks. PromptAlign improves the generalization of CLIP at test time beyond existing prompt-tuning methods, achieving state-of-the-art results.
Prompting for Vision-Language models. Vision Language (V-L) foundation models [30, 17, 44, 40, 41] have emerged with the convergence of image and text modalities. Having been pre-trained on massive image-text pairs from the web in a contrastive self-supervised manner, these models like CLIP [30] and ALIGN [17] have shown strong generalizability towards downstream zero-shot recognition tasks. However, adapting them efficiently to specific downstream tasks with limited data is still a challenging problem. Prompting in CLIP-like models has been explored in the form of a text query that is usually given to the text encoder to instruct it to perform a specific task. Recent methods propose to learn these prompts by treating them as continuous learnable vectors and training them in an end-to-end manner while keeping the model parameters frozen. CoOp [48] proposes to fine-tune CLIP by learning a set of prompts in the text encoder. CoCoOp [47] highlights the inferior generalization capability of CoOp and conditions the text prompt tokens on image embeddings on the fly. MaPLe [18] proposes to jointly learn deep prompts at both vision and text encoders of CLIP. The vision prompts are further conditioned on text prompts via a V-L coupling function. While all these approaches promote the effective transfer of CLIP, they require training data to learn the prompts which restricts such adaptation on novel datasets at test time. A recent method TPT [34], attempts to learn prompts solely at test time with the objective of enforcing consistency regularization between multiple views of a test sample by minimizing their averaged entropy. However, this approach struggles to explicitly address the distribution misalignment between the pre-training data of CLIP and the downstream test data. Our method builds on multi-modal prompting variant [18] and explicitly enforces to match the distribution via a distribution alignment technique using a proxy dataset that serves as a replacement for unavailable pre-training data of CLIP. To the best of our knowledge, ours is the first work to explicitly align the distributions learned by V-L foundational models with test time data distribution.
Test-Time Adaptation (TTA). TTA [36, 37, 23] aims to bridge the distribution gap between the train and test data distributions at test time. These methods can broadly be categorized into two streams. The first stream of approaches [11, 36] typically utilizes a self-supervised proxy task such as predicting image rotations and thus alters the training phase. For example, Test-Time Training (TTT) [36] jointly trains the model for rotation prediction and image classification during training.  $\mathrm{TTT}++$  [23] adopts a self-supervised contrastive learning approach as an auxiliary task during training and test-time feature alignment module to adapt the model. By incorporating additional auxiliary tasks, these techniques aim to mitigate over-fitting and improve the generalization performance of the model. The second stream of approaches solely adopts the pre-trained model without altering the training phase. These methods improve model generalization at test time by imposing self-consistency regularization constraints or aligning statistics of training and test samples in a model. Approaches with self-consistency regularization constraints typically utilize entropy minimization within samples of a batch or multiple views of a single test sample. TENT [37] proposes entropy minimization of batch-wise prediction probability distributions. The need to have a batch of samples by generating multiple views via augmentations at test time is eliminated in [45]. Other methods, such as NORM [32] and DUA [25] adapt the model's normalization layers by matching the batch normalization statistics computed during training with the test set distribution. Both [22] and ActMAD
Figure 2: Overview of our proposed PromptAlign method for zero-shot image classification. At test time, a single test sample along with its augmented views is passed through the CLIP image encoder, and the text labels are passed to the CLIP text encoder. The token distribution statistics – mean and variance – of the test sample are aligned with the offline computed source data statistics using a distribution alignment loss. The resulting alignment loss from the distribution shift is combined with the entropy loss to update the multi-modal prompts.
[26] updates all model parameters using a batch of samples in a continuous manner, with a distribution matching loss between the source statistics and the batch statistics of activations across the network. However, it is unclear how these statistics alignment techniques can be adopted for foundational V-L models in a lightweight manner for a single sample. In this paper, we propose a multi-modal test time prompting approach without altering the pre-training phase of V-L models along with a strategy to obtain training distribution statistics. This effectively paves the way for distribution alignment together with entropy minimization for improved test-time optimization.
Test-time prompt tuning. As mentioned earlier, a plethora of works has been introduced on learning prompts to adapt V-L models using the downstream task training data. However, the learning of prompts at test time remains largely unexplored in the literature. Very recently, TPT [34] proposed a test time prompting approach by learning prompts at the text side with an entropy minimization objective. While effective, TPT struggles with addressing the explicit alignment of the pre-training data distribution of CLIP with the test data distribution. Further, such an alignment is not trivial as access to the pre-training data of CLIP is not available. Our work addresses these limitations by extending TPT and introducing a distribution alignment objective for CLIP. This objective leverages token distribution statistics from proxy source data and a single test sample using multi-modal prompt learning, and explicitly aligns the train and test sample distributions to mitigate the domain shift.
We provide an overview of CLIP [30] and prompt learning and their application for zero-shot generalization to downstream tasks in Section 3.1. We then discuss test-time prompt tuning and briefly review previous approaches. Finally, we introduce PromptAlign in Section 3.2 and provide a detailed explanation of how we apply it to achieve improved zero-shot image classification.
Contrastive Language-Image Pre-training (CLIP). CLIP consists of two parallel encoders, one for mapping the text input into a feature vector and the other for mapping the visual input. We denote the CLIP image and text encoders by  $\mathcal{F}_v$  and  $\mathcal{F}_t$  respectively, and their pre-trained parameters are represented by  $\theta_{\mathrm{CLIP}} = \{\theta_v,\theta_t\}$  respectively. The input image  $X$  is divided into  $M$  patches, which are projected to produce patch tokens, and a class token CLS is prepended to it, resulting in  $\tilde{\pmb{X}}_0 = \{\mathrm{CLS},\pmb {e}_1,\pmb {e}_2,\dots ,\pmb {e}_M\}$  where  $e_i$  is the embedding of the  $i^{th}$  patch. The image encoder  $\mathcal{F}_v$  encodes the input patches via transformer blocks to produce a latent visual feature representation  $\tilde{\pmb{f}}_v = \mathcal{F}_v(\tilde{\pmb{X}}_0,\theta_v)$ . The corresponding class label  $y$  is embedded within a text template, such as 'a photo of a  $<\mathrm{CLS}>'$ , which is formulated as  $\tilde{\pmb{Y}}_0 = \{\mathrm{SOS},\pmb {t}_1,\pmb {t}_2,\dots ,\pmb {t}_L,\pmb {c}_k,\mathrm{EOS}\}$ . Here,  $t_l|_{l=1}^L$  and  $c_k$  are the word embeddings corresponding to the text template and the class label, respectively, while SOS and EOS are the start and end token embeddings. The text encoder  $\mathcal{F}_t$  encodes  $\tilde{\pmb{Y}}$  via transformer blocks to produce the latent textual feature as  $\tilde{\pmb{f}}_t = \mathcal{F}_t(\tilde{\pmb{Y}}_0,\theta_t)$ . For zero-shot inference, each text feature with class labels  $y = \{1,2,\dots ,C\}$  is paired with the image feature to compute a similarity score  $s_i = \mathrm{sim}(\tilde{\pmb{f}}_t\cdot \tilde{\pmb{f}}_v)$ , where  $\mathrm{sim}(.)$  denotes the cosine similarity. The prediction probability on  $X$  can be denoted by  $p(y_i|X) = \frac{\exp(\sin(\tilde{f}_t\cdot\tilde{f}_v)\tau)}{\sum_{i=1}^K\exp(\sin(\tilde{f}_t\cdot\tilde{f}_v)\tau)}$ , where  $\tau$  is the temperature of the softmax.
Prompt tuning on downstream tasks. CLIP contains a plethora of knowledge leveraged from training on millions of noisy image-text pairs. To effectively extract the rich features learned by the CLIP model, recent approaches [48, 47, 18, 2, 42] append extra learnable prompts while keeping image and text encoders frozen. These prompts modify the context of the model input without distorting the pre-trained CLIP features. Prompts are appended either at the image or text encoder side and learn contextual information tailored towards a specific task. In our work, we use a recently introduced multi-modal prompting baseline [18] that learns prompt tokens on both the text and image encoders.
Specifically, we append learnable  $V$  visual and  $T$  text prompts given as  $\pmb{p_v} = \{\pmb{p}_v^1,\pmb{p}_v^2,\dots,\pmb{p}_v^V\}$ , and  $\pmb{p_t} = \{\pmb{p_t^1},\pmb{p_t^2},\dots,\pmb{p_t^T}\}$  with the visual and textual input tokens respectively. The image encoder processes the input  $\tilde{\boldsymbol{X}}_0^P = \{\mathrm{CLS},\pmb {p}_v,\pmb {e}_1,\pmb {e}_2,\dots,\pmb {e}_M\}$  to generate a prompted visual feature representation denoted as  $\tilde{f}_v^p$ . Similarly, the text encoder processes the input  $\tilde{\mathbf{Y}}_0^P = \{\mathrm{SOS},\pmb {p}_t,\pmb {t}_1,\pmb {t}_2,\dots,\pmb {t}_L,\pmb {c}_k,\mathrm{EOS}\}$  producing the textual feature representation  $\tilde{f}_t^p$ . Our approach uses deep prompting as utilized in [18], along with text prompts and conditional image prompts at subsequent transformer blocks. We jointly represent the visual and textual prompts by  $\pmb{p}$ . We refer the readers to [18] for more details on the baseline architecture.
Test-time prompt tuning. Test-time prompt tuning introduced by [34] aims to benefit from the rich knowledge of CLIP to boost its generalization in a zero-shot manner. TPT can be viewed as a means to provide the model with a context that is customized for each individual test sample in order to more accurately recall the knowledge contained within CLIP. During inference, several randomly augmented views are generated from the given test sample  $X_{\mathrm{test}}$ . The predictions having entropy below a certain threshold are kept, while other views are discarded using a confidence selection filter. The averaged entropy of the filtered predictions is then used to update the prompts  $p$  in an unsupervised fashion using the following objective function.
$$
\mathcal {L} _ {\text {e n t r o p y}} = \arg \min  _ {\boldsymbol {p}} - \sum_ {i = 1} ^ {C} \tilde {p} _ {\boldsymbol {p}} \left(y _ {i} \mid X _ {\text {t e s t}}\right) \log \tilde {p} _ {\boldsymbol {p}} \left(y _ {i} \mid X _ {\text {t e s t}}\right), \tag {1}
$$
where  $\tilde{p}_{\mathbf{p}}(y_i|X_{\mathrm{test}})$  represents the mean of vector class probabilities produced by the model across the different augmented views preserved after the confidence selection filter.
3.2 PromptAlign
Unimodal Test-time Prompt Tuning (TPT) [34] updates the text prompts on the fly at inference by minimizing the entropy loss. This approach does not explicitly handle distribution shift that arises in the test set which is sub-optimal. One solution to this problem is to align the source and target distributions by bringing the test sample into the domain of the source data distribution. However, TPT updates the prompts only on the text encoder branch with static labels which poses an architectural limitation in performing the distribution alignment of tokens. Hence, distribution alignment of tokens
can only be performed on the vision branch. We, therefore, utilize multi-modal prompt learning VL models [18] to explicitly handle the distribution shift of each test sample from the source distribution.
Given a test sample  $X_{\mathrm{test}}$ , we take multiple augmented views and pass them through the visual encoder with deep learnable prompts as shown in Fig. 2. At each layer of the visual encoder, we compute the token alignment loss between the means and variances of the test sample with that of the means and variances of a proxy dataset which mimics the behavior of the source data distribution. Our final objective combines the entropy and alignment losses to update the prompts for a given test sample.
Proxy source dataset. In order to compute the token embedding statistics on the source dataset, we require the pre-training dataset of the CLIP model. The CLIP model was trained on over 400 million image-text pairs, which is not publicly available. However, previous works have shown that LAION400M [33] can be used as a training dataset to achieve the performance of CLIP [5], leading to subsets of LAION400M being used as a proxy dataset for the CLIP training dataset. In addition to this, CLIP has also been heavily tuned to achieve excellent zero-shot performance on ImageNet [1]. Therefore, we use ImageNet as the proxy source dataset for computing the mean and variance of token distributions. Source dataset statistics are computed offline and utilized directly at test time.
Token Distribution Alignment via multi-modal prompting. We generate  $N_{k}$  random views of the test sample using a set of augmentations  $\mathcal{H}$ . The mean and variance statistics of token embeddings of the test sample are computed at the output of each transformer layer of the CLIP model's visual encoder, across the  $N_{k}$  views. Similarly, the source data statistics are pre-computed in an offline manner. We represent test sample distribution by  $(\mathcal{T})$  and source distribution by  $(\mathcal{D})$ . Specifically, we compute the token means and variances for the alignment as follows.
$$
\boldsymbol {\mu} _ {l} (\mathcal {T}; \boldsymbol {p}) = \frac {1}{N _ {k}} \sum_ {\mathrm {x} \in \mathcal {H} (X)} \tilde {\boldsymbol {X}} _ {l, \mathrm {x}} ^ {\boldsymbol {p}}, \tag {2}
$$
$$
\boldsymbol {\sigma} _ {l} ^ {2} (\mathcal {T}; \boldsymbol {p}) = \frac {1}{N _ {k}} \sum_ {\mathrm {x} \in \mathcal {H} (X)} \left(\tilde {\boldsymbol {X}} _ {l, \mathrm {x}} ^ {\boldsymbol {p}} - \boldsymbol {\mu} _ {l} (\mathcal {T}; \boldsymbol {p})\right) ^ {2}, \tag {3}
$$
where  $\mu_l(\mathcal{T};\pmb{p})$  and  $\sigma_l^2 (\mathcal{T};\pmb{p})$  are the vector means and variances of the test sample tokens at the layer  $l$  in the visual encoder and  $\tilde{\pmb{X}}_{l,\mathrm{x}}^{p}$  represents the prompted token embeddings at layer  $l$  for the augmented view input x. Similarly for each layer  $l$  in the visual encoder, we pre-compute the source data statistics as,
$$
\hat {\boldsymbol {\mu}} _ {l} = \boldsymbol {\mu} _ {l} (\mathcal {D}, \theta_ {v}) \quad \text {a n d} \quad \hat {\boldsymbol {\sigma}} _ {l} ^ {2} = \boldsymbol {\sigma} _ {l} ^ {2} (\mathcal {D}, \theta_ {v}), \tag {4}
$$
where  $\theta_v$  denotes the parameters of the visual encoder from the pre-trained CLIP model. We compute the token distribution alignment loss between the mean and variances of the test sample and the source dataset statistics as follows,