MaskCLIP: Masked Self-Distillation Advances Contrastive Language-Image Pretraining

Following [jia2021scaling, radford2021learning], we introduce a Transformer-based text encoder $E_T$ to leverage language knowledge. It aims to align the global feature representations of an image and a text with respect to some forms of similarity. Precisely, consider a given image-text pair $\{I,T\}$, besides extracting the visual feature representation $E_I\left(I\right)$ using the vision backbone as shown by Equation 1, we additionally use the text encoder $E_T$ to extract linguistic features from the text $T$. We represent these linguistic features as $E_T\left(T\right)=\{f_1^T,f_2^T,\dots ,f_M^T,f_{\text{eos}}^T\}$, where eos means the end of sequence token.

$f_{\text{eos}}^T$ from the linguistic features and mean feature of the visual features are regarded as the global representations and are fed into a projection head (implemented as a fully-connected layer) respectively to obtain the metric embeddings $e^T$ and $e^I$. Image-text contrastive loss is employed to align them during pretraining. The loss can be formulated as $L_T+L_I$, with

where $B$ stands for the number of image-text pairs within a training mini-batch, $i,j$ are indexes within the batch; $\sigma$ stands for the temperature for the loss functions, which is learned together with all other parameters during training.

Knowledge distillation is a learning paradigm where a student model is trained to match the output of a given teacher model, so that the student model can be improved by the teacher. Instead of bringing in an external teacher, self-distillation methods such as [tarvainen2017mean, caron2021dino, grill2020bootstrap] proposes using a mean teacher model that is derived from the student itself. In specific, the teacher shares the same structure with the student, while the parameters of the teacher are exponential moving averages (EMA) of the parameters from the student. In the following, we would use the term “EMA model” to represent such mean teacher model constructed from the student.

MaskCLIP leverages the mean teacher self-distillation to enhance its vision representations. Let ${\overline{E}}_I$ be the EMA model of the backbone encoder $E_I$. ${\theta }_t$ and ${\overline{\theta }}_t$ are the parameters of $E_I$ and ${\overline{E}}_I$ at training step $t$. ${\overline{\theta }}_t$ is updated with

where $\alpha$ is a hyper-parameter for smoothing updates. We propose to incorporate masked image modeling into self-distillation, resulting in masked self-distillation with asymmetric input for student model and teacher model.

In specific, consider a given input image $I$, we first feed it to the EMA model ${\overline{E}}_I$ (teacher model) to obtain the distillation targets. These target features can be represented as

In the meantime, we randomly mask a large portion of the input image patches and then feed it into the original backbone $E_I$ (student model). Following [he2021masked], we only feed the visible (unmasked) patches, denoted by $I^{\prime }$, into the original backbone $E_I$ to speed up computation and save memory. Let $M$ be the indexes of all the masked tokens. These encoded features corresponding to visible tokens can then be denoted as $E_I\left(I^{\prime }\right)=\left\{f_{\text{cls}}^{\prime }\right\}\bigcup \left\{f_{k\notin M}^{\prime }\right\}$. They are then joined with a shared and learnable feature vector, denoted as $m$, that represents mask tokens, to form a complete set of features $\{f_{\text{cls}}^{\prime },f_1^{\prime },f_2^{\prime },\dots ,f_N^{\prime }\}$, with $f_{i\in M}^{\prime }=m$. We attach positional embeddings onto all these tokens, and append a one-layer Transformer $D$ as a decoder to predict features of the masked region from the visible tokens, which could be formulated as

A distillation loss $L_{\text{Dist}}$ is imposed to make the predicted features from $(D\circ E_I)\left(I^{\prime }\right)$ match the target features generated by ${\overline{E}}_I\left(I\right)$ (Equation 4):

We utilize smooth L1 loss for feature matching and $\beta$ is the smooth factor, we set $\beta =2$ by default.

Finally, we pretrain MaskCLIP with all these losses combined:

with $\lambda$ being a hyper-parameter weighting between VL contrastive loss and masked self-distillation loss. All the components of MaskCLIP are trained from scratch, including the backbone $E_I$, the decoder $D$, as well as the text encoder $E_T$.

Model architecture. Our framework consists of three modules, the visual encoder $E_I$, the text encoder $E_T$, and a visual decoder $D$. We adopt the widely used Transformer ViT-B/16 [dosovitskiy2020image] for a fair comparison. It is composed of 12 layers, 768 width and 12 head. The input image is $224\times 224$ resolution and is further split into $14\times 14$ image patches with patch size $16\times 16$. A learnable cls token is prepended to the 196 embeddings. For the text encoder, we adopt the 12-layer, 512-width, and 8-head Transformer following CLIP [radford2021learning]. The number of input text tokens is fixed to 77 with necessary truncations or paddings. For the decoder, we directly use a one-layer Vision Transformer for both simplicity and performance considerations.

Pretraining details. We train our proposed MaskCLIP from scratch and training for 25 epochs, the batch size is fixed to 4096 for all the experiments. The masks used in the mask self-distillation branch are random mask with a mask ratio of 75%. We pretrain all the models with a randomly sampled 20M subset from the recent image-text dataset LAION-400M [schuhmann2021laion]. We denote this subset as LAION-20M.

Downstream details. We evaluated MaskCLIP on several downstream datasets, including ImageNet-1K [deng2009imagenet], ADE20K [zhou2017scene], MS-COCO [lin2014microsoftcoco], Flicr30K [young2014image] et al. For ImageNet-1K, we report zero-shot, linear probing, and finetuning performance. The zero-shot is conducted following the label prompt setting in SLIP [mu2021slip]. For linear probing, we fix the backbone and train a new linear classifier for 90 epochs. For finetuning, we follow the setting in BEiT  [bao2021beit] which finetune the model for 100 epochs with a layer decayed learning rate. The setting of the rest dataset is described in the corresponding section and please see the supplemental materials for more details.

We first present our analysis by studying different ways of boosting CLIP. The baseline is CLIP [radford2021learning] trained on LAION-20M. Besides the introduced masked self-distillation, we consider three other popular methods: (1) SimCLR [chen2020simple], a representative contrastive method; and (2) MAE [he2021masked] and BEiT [bao2021beit], two state-of-the-art masked image modeling approaches. All the compared methods are trained on LAION-20M for a fair comparison. We have the following observations.

Vision self-supervision helps VL contrastive. We evaluate the models on both vision task ImageNet-1K [deng2009imagenet] classification and vision-language task image-text retrieval on Flicker30K [young2014image] and present the comparison in Table 1. We can see that all the added vision self-supervision, regardless of contrastive or generative, improves the baseline CLIP. Among them, our proposed MaskCLIP achieves the best results in terms of all the evaluation metrics, outperforming CLIP with +5.8%, +4.3%, + 1.2% on ImageNet-1K classification for zero-shot, linear probing and finetuning respectively, and +8.7%, +8.0% on Flicker30K for image-to-text retrieval and text-to-image retrieval. We also report the training GPU memory usage and time-consuming cost in Table 1. It is worth noting that the contrastive model (CLIP+SimCLR) compares two additional views of the input image, resulting in larger GPU memory usage and longer training time.

Masked image modeling is able to learn representations for local patches. We argue that the image encoder only pays attention to the text-described objects under VL contrastive due to sparse text description and to the centric objects under image contrstive due to central-crop augmentation. In contrast, masked image modeling forces the image encoder to focus on local patches using token-wise objective by mandatorily masking a large portion of patches. Here, we provide numerical comparison for evidence. We conduct an “Annotation-free zero-shot segmentation” experiment to test the zero-shot segmentation. The results on such a dense prediction task would better reveal the ability of local patch representations than global classification. Following the design in DenseCLIP [zhou2021denseclip], we use the prompted label feature as the linear classification weight to realize segmentation, without any training procedure. We evaluate the performance on three widely used datasets: MS-COCO [lin2014microsoftcoco], ADE20K [zhou2017scene] and Pascal Context [mottaghi2014role]. The results are shown in Table.2. We can see that equipped with masked image modeling, our MaskCLIP as well as CLIP+MAE achieves better results than CLIP and CLIP+SimCLR, validating our hypothesis.

Masked self-distillation learns semantic representations for local patches. Our introduced masked self-distillation predicts visual features dynamically outputted by the visual encoder and thus implicitly gets supervision from the text side via VL contrastive. While MAE predicts fixed low-level pixels, making it inefficient to learn semantic representations (as the objective may force the representation to memorize low-level details) and thus causing conflict with VL contrastive. To show this, we select images from MS-COCO [lin2014microsoftcoco] and calculate the feature similarity between image features and its corresponding caption features. We also select objects in the caption, prompt it to a new caption, such as “a photo of teddy bears”, and calculate the similarities. An example is shown in Figure 6 (More can be found in the supplementary material). Comparing MaskCLIP with CLIP+MAE in the fourth column, we can see that CLIP+MAE uses color as evidence and fails to distinguish the white teddy bear from white snow. While our MaskCLIP successfully differentiates the two objects, suggesting ours learn more semantic features. On the other hand, the superior results of MaskCLIP shown in Table 1 and Table 2 also validate this. It is worth mentioning that CLIP and CLIP+SimCLR fail to have a correct response partition for different single objects like MaskCLIP, further justifying our second observation.

Single-Stage v.s. two-Stage. Our MaskCLIP learns the VL contrastive and masked self-distillation simultaneously and jointly in a single stage. One possible variant is to first train CLIP and then use CLIP feature from the first stage to train masked image modeling as in [wei2021masked, wei2022mvp]. We report results on three datasets in Table 3. We can see that the second stage achieves better finetuning results compared with results from the stage one, showing the effectiveness of masked image modeling. Nonetheless, such two-stage training requires longer training time and loses the transfer capability in zero-shot setting. In contrast, our MaskCLIP achieves superior results under all settings with fewer epochs.

Data scaling and training time scaling. We study the scaling performance of MaskCLIP with respect to the training data size and also the training epochs. For training data, we random sample 10M, 20M and 50M from the LAION-400M dataset and train both CLIP and our MaskCLIP for 25 epochs. For training epochs, we set the number of epochs to 25, 50, 100 and train the model on LAION-20M. The results are shown in Figure 3. We can see that MaskCLIP gets consistent improvement over CLIP with either more data or more epochs. Furthermore, MaskCLIP is efficient that MaskCLIP with only 25 epochs performs better than CLIP trained for 100 epochs in all the cases. When it comes to data size, MaskCLIP with only 20M data shows comparable performance with CLIP with 50M data on zero-shot tasks, and outperforms it on the finetuning tasks.

To show the effectiveness of MaskCLIP as a general vision-language pretrain method, we conduct experiments on both vision tasks and vision-language tasks. For vision tasks, we report results on ImageNet-1K [deng2009imagenet] classification, MS-COCO [lin2014microsoftcoco] object detection, and ADE20K [zhou2017scene] semantic segmentation. For vision-language tasks, we report zero-shot results on 25 datasets and image-text retrieval results on Flickr30K [young2014image] and MS-COCO [lin2014microsoftcoco]. In the following, we compare with the supervised baseline DeiT [touvron2021training], self-supervised methods SimCLR [chen2020simple] and MAE [he2021masked], and vision-language methods CLIP [radford2021learning] and SLIP [mu2021slip]. For fair comparison, we train SimCLR and MAE on LAION-20M with the same epochs.

Classification on ImageNet-1K. As shown in Table 4, MaskCLIP benefits from the advantages of both VL pretraining and image mask self-distillation that shows strong performance on all the metrics. For zero-shot task, MaskCLIP outperforms CLIP with $+5.8\%$ with 25 epoch training, and $+3.5\%$ better than the concurrent work SLIP. When it comes to finetune, MaskCLIP reaches $84.1\%$ top-1 accuracy with 100 epoch training, outperforms CLIP with $+1.4\%$.

Semantic segmentation on ADE20K. Then we apply our MaskCLIP on semantic segmentation task. Here we use the UperNet [xiao2018unified] framework with $512\times 512$ input and end-to-end training for 160K iterations. The evaluation metric is mean Intersection of Union (mIoU) and we report single-scale evaluation results here. The results are given in Table 4. Our method achieve 49.6 mIoU, +2.2 mIoU than supervised based methods. It is also $+3.8$ mIoU than our baseline method CLIP, and +1.4 mIoU than SLIP. When we extend the pretraining to 100 epochs, our MaskCLIP reaches a strong results of $50.8$ mIoU, consistently better than CLIP with $+3.9$ mIoU. This verifies the effectiveness of our introduced incorporation.

Object detection and instance segmentation on MS-COCO. We further investigate our transfer performance on object detection and instance segmentation in Table.4. Here we use Mask-RCNN [he2017mask] framework with single-scale input and $1\times$ schedule (12 epochs). The evaluation metric is box AP for detection and mask AP for segmentation. We find that our MaskCLIP trained on LAION-20M performs slightly worse than MAE trained on ImageNet-1K and this comparison has additional variable that the training dataset is different, we argue this may be caused by the dataset domain gap. For fair comparison with all models trained on LAION-20M, our MaskCLIP performs the best.

Zero-shot on small datasets. We also report zero-shot performance on 25 datasets in Table 5, following the setting in [mu2021slip] that prompts the label of each dataset with a different context. We find that all the methods perform poorly on some datasets such as Aircraft(1% acc for random guessing, we omit the description in the following), Fer(24.7%), Flowers(1.5%), Country211(0.5%), PCAM(50%). This might be caused by the data domain gap that the subset of LAION-400M we used does not contain related images and descriptions. For the rest of the datasets, all the methods get reasonable performance and our MaskCLIP gets the best performance on most datasets.

Zero-shot on text-image retrieval. We further report the zero-shot text-image retrieval results on two benchmark datasets, Flicr30K [young2014image] and MS-COCO [lin2014microsoftcoco]. We find that the text without any prefix or suffix works well for all the models. Table 6 shows the results. We can see that MaskCLIP exhibits a strong zero-shot performance. For example, with 25 epochs training, MaskCLIP reaches 64.9% Rank@1 image-to-text accuracy on Flickr30K, outperforming CLIP with 8.7%.