SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Prompt Pre-Training with Twenty- Thousand Classes for Open-Vocabulary Visual Recognition
Shuhuai Ren†, Aston Zhang†,* Yi Zhu†, Shuai Zhang†, Shuai Zheng†, Mu Li†, Alex Smola†, Xu Sun‡
$\ddagger$ National Key Laboratory for Multimedia Information Processing, School of Computer Science, Peking University  $\dagger$ Amazon Web Services
This work proposes POMP, a prompt pre-training method for vision-language models. Being memory and computation efficient, POMP enables the learned prompt to condense semantic information for a rich set of visual concepts with over twenty-thousand classes. Once pre-trained, the prompt with a strong transferable ability can be directly plugged into a variety of visual recognition tasks including image classification, semantic segmentation, and object detection, to boost recognition performances in a zero-shot manner. Empirical evaluation shows that POMP achieves state-of-the-art performances on 21 datasets, e.g.,  $67.0\%$  average accuracy on 10 classification datasets  $(+3.1\%)$  compared to CoOp) and  $84.4\mathrm{hIoU}$  on open-vocabulary Pascal VOC segmentation  $(+6.9$  compared to ZSSeg). Our code is available at https://github.com/amazon-science/prompt-pretraining.
It has been a new norm to formulate visual recognition tasks (e.g., image classification, object detection, and semantic segmentation) as language-guided visual recognition or vision-and-language problems [42, 19, 61]. In language-guided visual recognition, categories of images are represented by natural language rather than discrete label IDs, and the semantics between images and their corresponding textual descriptions are often aligned via a constrastive loss during training [42]. Model inference also becomes an image-to-text matching problem, where text prompts like "a photo of a [CLASSNAME]" are curated as text descriptions of images. By varying the [CLASSNAME] placeholder and computing the similarity between text descriptions and images, we can identify the most suitable class name and consider it as the predicted target class. A significant benefit of this language-guided paradigm is that it supports open-vocabulary inference, that is, zero-shot recognition for arbitrary categories that may not even have been seen during training, thanks to the flexibility in modifying the class names in the textual prompt [42, 45].
As the context for class names, the text prompt plays a critical role in language-guided visual recognition models. A good prompt should holistically express the semantics of visual categories to better elicit the knowledge learned by vision-language models (VLMs) during pre-training. There are two popular types of prompts: hard prompts (e.g., a photo of a [CLASSNAME]) and soft prompts. Soft prompts are learnable token embeddings that can be fine-tuned given some input data, and have been demonstrated to be more effective and stable on downstream tasks than hard prompts [65, 35, 63]. However, traditional prompt tuning methods usually fine-tune the soft prompt on task-specific datasets with a limited number of class labels, making it difficult to generalize to novel classes and across tasks [45, 2]. For example, when transferred between the two downstream datasets of Flowers102 [40]
*Correspondence to Aston Zhang <az@astonzhang.com>
and DTD [9], the soft prompt fine-tuned on DTD only achieves  $33.4\%$  accuracy on Flowers102, significantly lower than the  $61.8\%$  accuracy of a hand-crafted prompt on Flowers102, demonstrating a severe overfitting issue [49].
In this work, we aim to learn a universal soft prompt that covers a broad set of visual concepts while being task-agnostic. Specifically, we propose PrOMpt Pre-training (POMP), a method for scaling up prompt learning on the ImageNet-21K dataset, which has over twenty-thousand classes organized by the WordNet [37] hierarchy. The set of classes in ImageNet-21K includes general and long-tail visual categories of various semantic granularities, which have been proven to provide better downstream results for large models [30, 14]. Pre-training on this large-scale dataset helps condense semantic information into the soft prompt for universal visual discrimination. Once pre-trained, this universal prompt (i) can be easily applied to downstream datasets to improve model performance in zero-shot settings; (ii) is compatible with both region-level and pixel-level visual patterns, making it useful for various vision tasks such as object detection and semantic segmentation.
Figure 1: POMP outperforms previous state-of-the-art models on a broad range of visual recognition tasks and datasets.
However, pre-training prompt with such a massive class set is challenging due to generally prohibitive computational costs. During prompt pre-training, the activation of the whole text encoder needs to be kept independently for every class and the memory consumption increases proportionally to the number of classes. In short, pre-training prompts on ImageNet-21K requires over 300 GB GPU memory with traditional methods like CoOp [65]. In POMP, we solve this issue with a simple class sampling strategy, local contrast, which reduces the GPU memory requirement dramatically to less than 16 GB. Moreover, we propose a local correction strategy to reduce the bias caused by class sampling and improve the generalization of the pre-trained prompt.
Experimental results in Figure 1 show that POMP outperforms previous state-of-the-art (SOTA) models on a broad range of visual recognition tasks and datasets. Specifically, compared to zero-shot CLIP [42], POMP improves the accuracy on ImageNet-21K by a gain of  $+2.9\%$ . It also achieves an average accuracy of  $67.0\%$  when transferred to 10 downstream image classification datasets, which is  $3.1\%$  higher than CoOp [65]. For semantic segmentation, POMP achieves  $39.1\mathrm{hIoU}$  on open-vocab COCO Stuff and  $84.4\mathrm{hIoU}$  on open-vocab Pascal VOC, outperforming ZSSeg [61] by  $+1.3$  and  $+6.9\mathrm{hIoU}$ , respectively. For object detection, POMP achieves  $57.9$  and  $22.9\mathrm{AP}_{50}$  when transferred from LVIS to COCO and Object365, surpassing Detic [67] by  $+1.9$  and  $+0.8\mathrm{AP}_{50}$ , respectively.
2.1 Language-Guided Visual Recognition
Language-guided visual recognition usually leverages VLMs as foundation models. Representative VLMs like CLIP [42] consist of an image encoder and a text encoder, which are used to encode image-text pairs into a joint feature space for learning the semantic alignment between vision and language [7, 46]. After being pre-trained on large-scale image-text pairs, CLIP-like models [28, 62, 33] are able to map images to their corresponding language descriptions, allowing visual recognition to generalize in the wild. This language-driven modeling paradigm also facilitates other vision tasks, including semantic segmentation [61, 32, 43, 13] and object detection [19, 15, 67]. These works typically designed a two-stage framework: it first leverages the pre-trained proposal network to extract features of specific visual patterns (e.g., segment mask and region) and then conducts classification in the same matching style as CLIP. The class descriptions for matching are synthesized using prompts, and in this work, we pre-train a soft prompt for VLMs on ImageNet-21K to further enhance their zero-shot generalization ability.
2.2 Prompt Tuning
In order to adapt VLMs to downstream tasks, recent research proposed a parameter-efficient tuning method named prompt tuning. CoOp [65] first proposed to replace the hand-crafted prompt with learnable vectors (also known as a soft prompt) for fine-tuning while freezing the entire pre-trained parameters. VPT [12], on the contrary, moved the learnable vectors from the text side to the image side, and proposed to concatenate the "visual soft prompt" and the patch sequence of an image as the input for fine-tuning. Prompt tuning was also used in other visual recognition tasks, such as object detection [15], semantic segmentation [43], and video recognition [39]. Over manual prompt engineering, the soft prompt optimized with few-shot data has achieved significant performance improvements, but only fitting one specific downstream dataset.
To enhance the generalization of the soft prompt to a wider range of unseen classes and datasets, CoCoOp [66] modeled the context condition on input images. A recent work MaPLe [29] appended the soft prompt to the hidden representations at each layer in both the text and image encoder. In sharp contrast to previous approaches, we propose to pre-train a universal soft prompt on large-scale datasets with massive visual categories. Such a pre-trained prompt is task-agnostic, allowing for direct transfer to various downstream datasets without fine-tuning.
We first review the process of classical prompt tuning for VLMs in § 3.1. To address the training efficiency issue of previous methods, in § 3.2 we introduce our method of prompt pre-training (POMP) that includes two key components: local contrast and local correction. Our pre-trained prompt can then be transferred to downstream datasets and tasks in a zero-shot manner as discussed in § 3.3.
Language-guided visual recognition models like CLIP formulate image classification as an image-text matching problem, where the goal is to select the correct textual class name from a predefined class set for the image query. Following [65, 29], we consider CLIP as our vision-language foundation model, for its simplicity in design and wide applicability. CLIP consists of an image encoder  $f_{I}$  and a text encoder  $f_{T}$ . Given a visual recognition dataset  $\mathcal{D}$  with a class set of  $N$  class names  $\mathcal{C} = \{c_i\}_{i=1}^N$ , CLIP manually devises a hard prompt to synthesize textual descriptions  $t_i$  for each class name  $c_i$ , e.g.,  $t_i = \text{"a photo of a } [c_i]$ . Then each class description is fed into the text encoder to generate the normalized class feature  $\mathbf{w}_i = f_T(t_i) / \|f_T(t_i)\|_2 \in \mathbb{R}^d$ , where  $d$  is the dimension of the feature. The concatenation of  $N$  class features  $[\mathbf{w}_1, \dots, \mathbf{w}_N] \in \mathbb{R}^{N \times d}$  can be considered as the class weight of a linear classifier for classifying an image. Given an input image  $x$ , the image encoder is used to extract its visual feature  $\mathbf{x} = f_I(x) / \|f_I(x)\|_2 \in \mathbb{R}^d$ . Finally, CLIP calculates the similarity between  $\mathbf{x}$  and all the class features, then predicts the class with the highest similarity as the target class.
To address the inefficient expressiveness of the manual prompt, previous research like CoOp [65] proposed to parameterize the manual prompt as a soft prompt  $\Theta$ , and fine-tune it to fit downstream datasets. The soft prompt is made up of a sequence of learnable token embeddings  $\Theta = [\theta_1, \theta_2, \dots, \theta_M] \in \mathbb{R}^{M \times e}$ , where  $M$  is a hyperparameter specifying the length of the soft prompt and  $e$  is the dimension of the token embedding. The token embeddings of each class name  $c_i$  are further appended to the soft prompt  $\Theta$  to generate the class feature  $\mathbf{w}_i^{(\Theta)}$ , and the prediction probability for the ground-truth class  $y$  is denoted as
$$
P (y \mid \mathbf {x}; \boldsymbol {\Theta}) = \frac {\exp \left(\mathbf {x} ^ {\top} \mathbf {w} _ {y} ^ {(\boldsymbol {\Theta})} / \tau\right)}{\sum_ {i = 1} ^ {N} \exp \left(\mathbf {x} ^ {\top} \mathbf {w} _ {i} ^ {(\boldsymbol {\Theta})} / \tau\right)}, \tag {1}
$$
where  $\mathbf{x}^{\top}\mathbf{w}_i$  represents the similarity score and  $\tau$  is a temperature parameter. The parameters of the soft prompt are updated by minimizing the cross-entropy loss [65]:
$$
\mathcal {L} (\boldsymbol {\Theta}) = \underset {(\mathbf {x}, y) \in \mathcal {D}} {\mathbb {E}} \left[ - \log P (y \mid \mathbf {x}; \boldsymbol {\Theta}) \right]. \tag {2}
$$
The gradient of  $\mathcal{L}(\Theta)$  is represented as:
$$
\nabla_ {\Theta} \Big (- \log P (y \mid \mathbf {x}; \Theta) \Big) = \frac {1}{\tau} \Big [ - \nabla_ {\Theta} (\mathbf {x} ^ {\top} \mathbf {w} _ {y} ^ {(\Theta)}) + \sum_ {i = 1} ^ {N} P (y _ {i} \mid \mathbf {x}; \Theta) \nabla_ {\Theta} (\mathbf {x} ^ {\top} \mathbf {w} _ {i} ^ {(\Theta)}) \Big ]. \tag {3}
$$
Figure 2: Overview of POMP. POMP pre-trains a soft prompt (learnable) on the ImgaNet-21K dataset with massive classes, and then directly transfers the learned prompt (frozen) to downstream datasets of image classification (CLS), object detection (DET), and semantic segmentation (SEG) tasks. For DET and SEG, the region and mask proposal networks require pre-training with POMP prompt on detection and segmentation source data, respectively (See Appendix B).
Figure 3: GPU memory overhead (GB) required for prompt tuning on datasets with varying numbers of classes. The memory cost of CoOp on ImageNet-21K is 316.4 GB, which is generally prohibitive. POMP reduces the cost dramatically to 15.7 GB with local contrast among the 1000 sampled classes for optimization.
It can be decomposed into positive reinforcement for the ground-truth class and negative reinforcement for every class, which are the first and the second terms inside the square brackets of (3), respectively.
Note that both the image encoder and the text encoder are frozen during the prompt tuning process, which allows adapting the soft prompt efficiently to downstream data with very few learnable parameters. Methods along this direction of soft prompt learning include CoCoOp [66] and MaPLe [29]. However, most previous works fine-tune task-specific prompts, which limits their versatility and generalization [49].
3.2 POMP: Prompt Pre-Training
Now we present our task-agnostic prompt pre-training method: POMP. Once pre-trained, the learned prompt can be directly used for downstream tasks without fine-tuning (see Figure 2). As introduced in § 1, we propose to pre-train the soft prompt on the ImageNet-21K dataset for universal visual discrimination. Although the prompt tuning methods such as COOP and CoCoOp are parameter-efficient, they still incur computationally prohibitive training costs when applied to large-scale datasets with massive number of classes. Recall that the learnable parameters  $\Theta$  are embedded in the text input, while the loss is calculated at the output layer of the text encoder. For every class description, we need to allocate nearly  $15\mathrm{MB}$  of GPU memory to preserve the state of the entire frozen encoder (Transformer-base [53] with 12 layers), and propagate the gradient back through the last layer to the first layer, to update the soft prompt  $\Theta$ . Accordingly, the computational and caching cost of prompt tuning is proportional to the number of classes  $N$ . As shown in Figure 3, for general large-scale datasets like ImageNet-21K with more than twenty-thousand classes, traditional tuning methods will allocate  $21\mathrm{K} \times 15\mathrm{MB}$  (more than  $300\mathrm{GB}$ ) of GPU memory, which is generally prohibitive.
To enable prompt tuning on massive classes and acquire the capability of global visual discrimination, we introduce a training-efficient algorithm called POMP, which reduces the GPU memory and training time of prompt tuning dramatically. POMP has two major components: local contrast and local correction. The former decreases the number of classes for contrastive learning through negative class sampling, and the latter reduces the bias caused by local contrast by adjusting the similarity scores of negative classes. We detail these two components in the following.
3.2.1 Local Contrast
Discriminating all classes during contrastive learning is the source of training inefficiency. In order to alleviate this problem, we propose to narrow the scope of contrastive learning from global to local, and only require the model to identify the ground-truth class of the input image from a subset of the full class set. The class subset is sampled at each training step, allowing the model to discriminate within an ever-changing set of categories, and gradually restoring the relationship among all categories.
Specifically, given an input image, we sample  $K$  classes ( $K$  is much smaller than the total number of classes,  $N$ ), including the ground-truth class  $y$  and  $K - 1$  negative classes. We use a straightforward yet effective proposal distribution of uniform distribution for negative class sampling, where every negative class has an equal probability, i.e.,  $p = 1 / (N - 1)$ , of being sampled. We also explore alternative types of proposal distribution, such as frequency-based and similarity-based distributions. However, our experiments reveal that POMP with the simple uniform distribution considers both common and rare classes, as well as easy and difficult classes, resulting in the best performance. Please refer to Appendix E.1 for details.
After sampling, we denote the set of the negative classes as  $\mathcal{N}$ , with  $|\mathcal{N}| = K - 1$ . By using the local contrast, we can significantly reduce the training overhead to a fraction of  $K / N$  of the original one. Upon completion of training, we can use the full class set to compute the prediction probability of each image. Overall, the motivation behind our local contrast is analogous to that in the NCE-based contrastive learning frameworks [52, 57]. In these frameworks, the contrast is performed by sampling a batch of instances due to computational limitations and computing the loss within the batch as an empirical estimation for the expected contrastive loss [25, 1, 22, 51, 68].
3.2.2 Local Correction
Given that the local contrast component necessitates a reduced number of negative classes, the negative reinforcement in the vanilla gradient in (3) is diminished to  $K / N$ . As a result, the prompt optimization direction is inevitably biased due to the absence of other negative classes. To mitigate this bias and enhance the model performance, we add a local correction term  $m$  to the logits of the sampled negative classes  $\mathbf{x}^{\top}\mathbf{w}_{i}^{(\Theta)} / \tau (i\neq y)$ , which serves as a margin [60, 45, 69] between the positive and the negative logits. Accordingly, the final prediction probability of POMP is denoted as:
$$
\tilde {P} (y \mid \mathbf {x}; \boldsymbol {\Theta}) = \frac {\exp \left(\mathbf {x} ^ {\top} \mathbf {w} _ {y} ^ {(\boldsymbol {\Theta})} / \tau\right)}{\exp \left(\mathbf {x} ^ {\top} \mathbf {w} _ {y} ^ {(\boldsymbol {\Theta})} / \tau\right) + \sum_ {i \sim \mathcal {N}} \exp \left(\mathbf {x} ^ {\top} \mathbf {w} _ {i} ^ {(\boldsymbol {\Theta})} / \tau + m\right)}. \tag {4}
$$
The local correction term  $m$  encourages the positive logit to be larger than the negative logits by a certain margin, resulting in a more stringent decision boundary:
$$
C _ {+}: \mathbf {x} ^ {\top} \mathbf {w} _ {y} ^ {(\boldsymbol {\Theta})} / \tau \geq \mathbf {x} ^ {\top} \mathbf {w} _ {i} ^ {(\boldsymbol {\Theta})} / \tau + m, \quad i \neq y.
$$
Therefore, compared to the prediction probability without local correction, (4) makes the decision boundary more robust against the unsampled negative classes and enforces the learning of more discriminative class features [60]. This will improve model regularization and robustness across datasets and domains (to be shown in § 4.4). Different from other margin-based losses that use a fixed margin [11, 54], our margin  $m$  is designed to adaptively adjust itself based on the value of  $K$ :
$$
m = - \log ((K - 1) / (N - 1)). \tag {5}
$$
It is worth noting that  $m$  in (5) is a positive scalar. When  $K = N$ , all classes are included during optimization, and  $m$  equals zero. In this case, (4) degenerates to the standard prediction probability in (1). As the value of  $K$  decreases and the number of visible classes is reduced, the margin  $m$  increases to create space for potential class features in the representation space. Our adaptive margin outperforms the fixed margins (which are specified as hyper-parameters) for various  $K$ , allowing models to maintain optimal performance under different computing budgets (to be shown in § 4.4).
3.3 Zero-Shot Transfer Learning
As shown in Figure 2, after pre-training, our POMP prompt can be used to synthesize class features for classification with an arbitrary class set, supporting zero-shot inference on downstream datasets
and tasks. In order to plug the POMP prompt into other visual tasks like semantic segmentation and object detection, we adopt a two-stage framework. In stage one, we use a pre-trained proposal network to generate a set of mask or region proposals. In stage two, we classify each proposal with the class features generated by our POMP prompt. Experiments in § 4.3 will show that the POMP prompt can handle both pixel-level and region-level visual patterns, leading to improved performance in segmentation and detection tasks.
4.1 POMP Prompt Pre-Training
We take CLIP (ViT/B-16) [42] as the backbone and conduct prompt pre-training on the ImageNet-21K dataset. The number of training samples for each class is 16 (16 shots), and the prompt length is 16. We sample 1,000 classes at each training step, i.e.,  $K = 1000$  in (4). See Appendix A for details.
4.2 End Task Setups and Implementation Details
To evaluate the generalization of the pre-trained prompt, we directly transfer it to downstream tasks and datasets. Appendix C lists the details of all the datasets. We follow previous works [58, 61, 19, 67] to designate data belonging to two class sets as source data and target data, respectively. The proposal networks are pre-trained on the source data with the source class set, while conducting zero-shot evaluation on the target data with the target class set. There are two protocols for the source-target data split. The first is the open-vocabulary protocol, where the class set of one dataset is divided into two disjoint groups for the source and target data. The second protocol is the cross-dataset protocol, in which the source and target data are from two independent datasets with potentially overlapping class sets. See Appendix B for implementation details.
4.3.1 Prompt Pre-Training on ImageNet-21K
Table 1 shows the results of POMP prompt on the ImageNet-21K test set after pre-training. The traditional prompt learning methods (e.g., CoOp and MaPLe) are trained on the ImageNet-1K dataset due to their prohibitive computational cost if trained on ImageNet-21K (more than 300 GB of GPU memory). On the contrary, our POMP prompt, pre-trained on ImageNet-21K using less than 16 GB of GPU memory, achieves the highest accuracy of  $25.3\%$  based on the CLIP (ViT-B/16) backbone, which surpasses Zeroshot-CLIP by  $3.5\%$  and Linear Probe by  $4.4\%$ . The VPT method, which uses visual prompts on the image side, does not require training overhead proportional to the number of classes, making it applicable to the ImageNet-21K dataset. VPT prepends independent learnable vectors to the hidden states of each layer in the visual backbone, surpassing linear probing and the previous
Table 1: Performance on the ImageNet-21K test set. ZeroshotCLIP and Prompt Ensemble in the top block conduct zero-shot inference. CoOp and MaPLe, indicated in gray in the middle block, are trained on the ImageNet-1K dataset due to prohibitive GPU memory consumption if trained on ImageNet-21K. The remaining methods in the bottom block are trained on ImageNet-21K.
Method	ResNet50	ViT-B/32	ViT-B/16
ZeroshotCLIP [42]	17.5	19.8	21.8
Prompt Ensemble [42]	18.8	20.9	23.5
CoOp [65]	16.6	18.1	20.8
MaPLc [29]	-	21.6	24.2
Linear Probing [42]	6.5	18.2	20.9
VPT [12]	-	21.8	24.8
POMP (Ours)	20.2	22.2	25.3
prompt tuning methods. However, its performance based on ViT-B/16 is still  $0.5\%$  worse than ours, demonstrating that our POMP prompt can better distinguish a large number of general visual categories. In addition, our method is agnostic to the backbone architectures like ResNet and ViT, and the improvement is consistent.