Learning Open-vocabulary Semantic Segmentation Models
From Natural Language Supervision

Image-caption Alignment. For each image-text pair, the objective is to align their visual and textual embeddings. The visual embedding $z^{\text{I}}$ is the average of group tokens, and the textual embedding $z^{\text{T}}$ is obtained by taking the [EOT] token feature of the caption embedding $\mathcal{T}^{\text{cap}}$, both projected to a 256-d joint feature space followed by normalisation. The image-caption contrastive loss $\mathcal{L}_{\text{contrast}}$ is formulated as:

Here, we omit the temperature parameter for simplicity.

Masked Entity Completion. The goal of masked entity completion is to infer all the masked entities in the sentence given the group tokens. In specific, we adopt a Transformer decoder layer, where a projection of the masked caption embedding is treated as query, and two linear transformations of group tokens are treated as key and values, respectively.

where $W_{q}^{\text{dec}}$, $W_{k}^{\text{dec}}$ and $W_{v}^{\text{dec}}$ are linear transformations, and $\tilde{\mathcal{T}}^{\text{m-cap}}$ denotes the updated vector sequence, with the masked entities being completed by querying group tokens. For training, we extract the [EOT] token features $z^{\text{M}}$ and $z^{\text{E}}$ from $\tilde{\mathcal{T}}^{\text{m-cap}}$ and $\mathcal{T}^{\text{entity}}$, and construct a contrastive loss:

Intuitively, the entity completion task enables better alignment between the groups and entities.

Cross-image Mask Consistency. To encourage visual invariance, we enforce consistent mask predictions between images that contain shared entities. Specifically, for each entity of interest, we can easily source multiple image-caption pairs from the dataset by text searching. Given one sampled image-caption-entity triplet, $(I_{1},T_{1},E_{1})$, we can easily search for another sample $(I_{2},T_{2},E_{2})$, with both triplet sharing one entity, i.e., $e\in E_{1}\cap E_{2}$. Given encoded groups $\mathcal{G}_{1}$, $\mathcal{G}_{2}$, two sets of subgroups $\mathcal{G}_{1}^{\text{sub}},\mathcal{G}_{2}^{\text{sub}}\in\mathbb{R}^{K^{\prime}\times D}$ that represent the common entity $e$ (e.g., cat in Fig. 2) are obtained by choosing groups with higher similarity to the entity embedding, where $K^{\prime}=rK$ and $r$ is the selection ratio. The masks of co-attentive entity in image $I_{1}$ can be grounded by both entity-specific subgroups $\mathcal{G}_{1}^{\text{sub}},\mathcal{G}_{2}^{\text{sub}}$ as:

where $\sigma$ is the sigmoid activation; both $\mathcal{M}_{1}=\{m_{k}^{1}\}_{k=1}^{K^{\prime}}$ and $\hat{\mathcal{M}}_{1}=\{\hat{m}_{k}^{1}\}_{k=1}^{K^{\prime}}$ consist of $K^{\prime}$ unordered masks.

Training Objective. We adopt a combination of three different loss functions:

where $\lambda$ is the weight for balancing the mask consistency.

Pre-training Dataset. Following [52, 32], we use Conceptual Captions 12M [7] for training, which is originally constructed with over 12M image-text pairs collected from the Internet. However, due to some links have been expired, we have downloaded about 10M image-text pairs. The constructed entity set in Sec. 3.1.2 includes a total number of 100 frequently appeared entities while abstract nouns (e.g., art, view) are discarded. After filtering CC12M, we obtain 4.3 million image-text pairs for pre-training, which is termed as CC4M. Examples of entities include people, car, cup, chair, T-shirt, house, bed, cat, ball, pizza, etc. Please refer to the supplementary material for the full entity set.

Downstream Evaluation Datasets. We evaluate our model on three benchmarks, namely, PASCAL VOC 2012 [18], PASCAL Context [38] and COCO Object [31] with 20, 59 and 80 foreground classes, respectively. An extra background class is considered in all three datasets. We ignore their training sets and directly evaluate our method on the validation sets without any finetuning, including 1449, 5105, and 5000 images, respectively. In general, we report the mean Intersection-over-Union (mIoU) on all the classes.

Implementation Details. In our model, the self-attention layers in the visual encoder are initialised with DINO [6] pre-trained on ImageNet. The text encoder is initialised with BERT [16] model pre-trained on BookCorpus and English Wikipedia. Our decoder with one randomly initialised Transformer Decoder layer performs reasonably well. The input image is randomly cropped to 224$\times$224 at training time, and the batch size is set to 2048 with an initial learning rate $3.2\times 10^{-4}$. We train our model for 40 epochs using the Adam optimizer with weight decay set to 0.5. The coefficient for updating the momentum model is 0.99. As the generated masks are unreliable in early epochs, we set the mask consistency coefficient $\lambda$=0 for the first 30 epochs and $\lambda$=0.1 for the remaining epochs. The group selection ratio $r$ is 0.5. As for the threshold in mask consistency loss, we use $\delta$=0.65. At inference time, the image is resized with a shorter length of 448. We follow [52] to set a threshold for the background class, which is $0.9$, $0.5$ and $0.9$ on PASCAL VOC, PASCAL Context and COCO Object, respectively.

Comparison with Finetuning Transfer. We compare our method with DeiT [45], MoCo [10] and DINO [6], which are pre-trained on ImageNet [15] or CC12M+YFCC15M datasets with class labels [45] or self-supervision [10, 6], and finetuned on the training set from downstream benchmarks, with a randomly initialised convolution head appended on the backbone network. As shown in Table 1, our model achieves competitive performance on PASCAL Context, and outperforms the self-supervised methods by over 10% on PASCAL VOC, with zero-shot transfer.

Comparison with Zero-shot Transfer. Here, we compare with existing works under the zero-shot transfer scenario, including GroupViT [52], ViL-Seg [32] and CLIPpy [42], with the pre-training data ranging from CC12M [7] to 134M in-house dataset HQITP-134M [42]. For fair comparison, we re-train GroupViT [52] with their official codebase on our CC4M and CC12M datasets, with the same pre-trained weights as ours being adopted on CC4M. However, we observe no further performance gain of either applying pre-trained weights to GroupViT or using ViT-B as the backbone on CC12M. The results on CC12M match the reported ones (40.2 vs 41.1). Under the same pre-trained ViT-B backbone and CC4M dataset, our method surpasses the original GroupViT by 35.7%. Additionally, by only using 4M pre-training data, our model yields the best segmentation performance on PASCAL VOC, even outperforming CLIPpy, which is a concurrent work to ours, and pre-trained on 134M data, indicating the effectiveness and training efficiency of our proposed model.

Comparison with Zero-shot Segmentation Methods. In this line of research [50, 22, 17], the idea is to train the model with full mask labels (obtained either from manual groundtruth or pseudo-labelling [58]) on the seen classes and transfer the model to unseen classes, and the task is thus dubbed as zero-shot semantic segmentation (ZSS). To be specific, 5 classes (potted plant, sheep, sofa, train and tv-monitor) in PASCAL VOC and 4 classes (cow, motorbike, sofa and cat) in PASCAL Context are considered unseen while the remaining classes belong to seen. For comparison, we also report the zero-shot transfer performance of our model on these unseen classes, however, note that, we do not use any manual mask annotations for training. As observed in Table 2, our model surpasses majority of the models trained under the ZSS scenario, except for [17, 58] that adopted CLIP model pre-trained on 400M data. In terms of transfer efficiency, our model excels at zero-shot transfer ability without the need of training on seen classes.

Ablation Study on Proxy Tasks. Here, we aim to understand the effects of our proposed proxy tasks, i.e., masked entity completion and cross-image mask consistency. As shown in Table 3, the baseline model uses the image-text contrastive loss $\mathcal{L}_{\text{contrast}}$ only, while adding the entity completion task, $\mathcal{L}_{\text{entity}}$, we observe a significant improvement by 8.4% and 4.8% on PASCAL VOC and PASCAL Context, respectively. The performance gain is due to the ability of better aligning the pixel groups with the visual entities. Additionally, the mask consistency also brings improvements, and combing both leads to the best performance. Qualitative results can be seen in Fig. 3. We refer the readers for more visualisations in the supplementary material.

On the Choice of Masking Objectives. We compare our masked entity completion with a series of variants as shown in Table 4. Our proposed objective of masking all entities is listed in the first row. (1) All entities vs one entity: the masked language modeling (MLM) in prior works [16, 35] normally choose 15% of the token positions in the sentence for prediction, which results in one entity in most of our cases. We observe that masking all entities is 2.9% mIoU better than single entity masking, as it forces the network to infer all possible object categories in the image, that is potentially beneficial for the group to category alignment. (2) Entities vs nouns: masking and predicting noun phrases in the sentence [20] is one feasible option to learn fine-grained vision-text matching. However, noun masking leads to 3.5% lower mIoU than our entity masking strategy. This is because not all of the nouns in the sentence are visually corresponding to the objects in the image (e.g., illustration, night, etc), thus the group tokens are not expected to align with these nouns. Our method avoids this issue by only masking the visual entities. (3) Masked entity completion vs multi-label contrastive loss: comparing with the multi-label contrastive loss used in GroupViT [52], our proposed strategy shows superior performance. (4) Masked entity completion vs masked language modeling: MLM [16] originally predicts the masked token over the entire vocabulary via a cross-entropy loss. Here, we restrict the vocabulary to our constructed entity set for fair comparison. Our masking strategy surpasses both MLM variants by a clear margin, which we conjecture is because: (1) MLM classifies each masked token individually, and the model can easily refer to the context words without relating to groups. (2) MLM focuses on word-level representations, which is not in accordance with the sentence-level representation of the class embeddings we used during inference. Fig. 4 shows the effect of the masked entity completion in improving visual grouping (left) and group-text alignment (right).

On Cross-image Mask Consistency. Here, we analyze two key factors in cross-image consistency: (1) group numbers and selection ratios. As shown in Table 6, 8 groups perform comparably well for all selection ratios, and we pick 0.5 as the default. Smaller ratios miss the entities encoded in remaining groups while larger ratios introduce entity-irrelevant information (e.g., background). However, while increasing the group number to 16, it brings over-segmentation for large objects, deteriorating the performance. (2) the objective for cross-image consistency. We study another variant of cross-image consistency by directly aligning two sets of group tokens encoded with shared entity. We adopt the contrastive loss (NCE in Table 6) to pull two sets of group tokens closer, while other group tokens in each mini-batch are pushed farther. Table 6 reveals the superiority of our proposed mask consistency over group consistency, which we believe is because mask consistency involves the image content to realize visual invariance.

Performance on Unseen Entities. One might question whether the performance gain is mainly attributed to the selected (seen) entities in the entity set. We measure the mean IoU for 65 objects within entity set (e.g., person, bus) and 16 objects out of entity set (e.g., frisbee, stop sign) on COCO. As shown in Table 7, the mIoU of unseen entities is comparable to that of seen entities, indicating our model retains strong open-vocabulary segmentation ability without being affected by the choice of entities in pre-training. The text encoder pre-trained on large text corpus remains its ability to encode the semantic concept of objects out of entity set.