Effective End-to-End Vision Language Pretraining with Semantic Visual Loss

Unlike traditional vision language methods [1, 11, 12, 13, 14, 15] that learn vision language mapping based on specific tasks, vision language pretraining models learn robust cross-modal mapping through large scale pretraining. Vision language pretraining models can be categorized into two types: the one-stream models and two-stream models. The one-stream models [5, 16] first embed the languages and images into embeddings. The language embeddings and image embeddings are then concatenated and processed with a shared transformer network. Since the model parameters are shared, the one-stream models usually have a small model size. However, since the self-attention models like transforms have quadratic time complexity, the one-stream models will have slightly larger memory usage. The two-stream models [17, 4, 18] process language embeddings and image embeddings with separate transformer encoders and then combine the two features with cross-attention transformer. Compared with the one-stream models, they usually have a larger model size.

In this paper, we briefly follow the network structure of two-stream models [17, 4, 18]. We change the region feature input to pixel input and change the visual feature encoder to feature extraction backbones.

Effective loss functions lie in the heart of visual language pretraining. A region feature based vision language model usually involves three types of losses: the language loss, the visual loss, and the matching loss. Most of the time [16, 5, 17, 4], the language loss is similar to the language modeling loss as in BERT [19], in which the language tokens are masked randomly and the models are trained to predict the masked tokens based on two-directional contexts. The visual loss is to mimic the language loss and is done by masking the visual features and predicting the masked regions’ classes or attributes. The matching loss predicts if the language input and the vision input are matched.

When training raw image input based vision language models, the visual inputs (image pixels) don’t have specific class annotations. So existing works [8] only use the language loss and the matching loss. In this paper, we show that such training regimes lead to a slow convergence speed and low finetuning performance. With our proposed visual losses, the raw input models could achieve as fast convergence speed as region feature models in early training stages.

Transformers [20], as the most widely used language models and vision language models, although perform well in terms of accuracy, take a long training time. Training and testing transformer models efficiently is a long-standing goal in both NLP and computer vision. Existing works can be roughly divided into three categories. The first category [21, 22, 23] focuses on improving network structures. They try to optimize the quadratic time complexity of the original transformer model to linear time complexity. These works directly reduce time cost on both training and testing time. The second category [24, 25, 26] attempts to extract effective sub-networks from a trained transformer model through knowledge distillations. The third category [27, 28] focuses on improving the training strategy or proposing new losses. This type of method primarily improves efficiency during training time.

This paper falls in the third category. We are not trying to propose new network structures or extract sub-networks. We improve the training effectiveness by introducing three visual losses in vision language pretraining.

The general network architecture is shown in Figure 2. We briefly follow the network structure of two-stream models.

Network Structure. The image and text inputs are first processed by two separate encoders to extract high level features. The visual features and language features are then passed to the cross-modality encoder. An illustration of the cross-modality encoder structure and the cross-attention mechanism can be found in Figure 3. The cross-modality encoder is composed of cross-modality layers. In each cross-modality layer, the language features and image features are first merged by a cross-attention sub-layer. The cross-attention sub-layer has the same structure as a self-attention layer but uses different inputs as the query and key vectors. For the language branch, the cross-attention sub-layer uses language features from the previous layer as the query vector and image features as the key vector and value vector. After the cross-attention sub-layer, the features are further processed by a self-attention sub-layer and a fully-connected layer. We add short-cut connections between the sub-layers.

The visual input of the network is a raw pixel image. For vision transformer backbone (ViT) [10], we slice the image into patches and add an extra learnable position embedding. For CNN backbone, we first run the feature extraction then combine the extracted features with position embedding.

During pretraining, we follow the common practice to use a language modeling loss and a vision language matching loss.

Masked Language Modeling (MLM). Given an input sentence, the words are randomly masked at a ratio of 0.15. Following BERT [19], at 80% of chance, the masked token is changed to [MASK] token. Otherwise, the token is changed to zero or kept as the original token. The masked language modeling loss predicts the masked tokens given bidirectional contexts.

Image-Text Matching. Given the input sentence and image, the image-text matching loss predicts whether the text could describe the image i.g. if the image and text are matched. Fake sentences and real sentences are sampled at 50%. We use the language feature at [CLS] token to perform the matching task.

For visual loss, existing pixel input vision language pretraining models [8] can be seen as “No Visual Loss”. We describe our proposed three losses in the subsection below. Also, we draw the detailed network structure for each loss in Fig 4.

Much of the success of BERT in NLP is replying on large scale self-supervised pretraining. Here, we propose to train the vision branch with a similar self-supervised loss. We first divide the input images into $N*N$ patches. Each patch is randomly masked by a ratio $p$. After that, the self-supervised loss predicts the mean color of the masked regions. To precisely predict the mean color of masked regions, the model will need to learn to extract useful information from contexts. Note that the self-supervised loss does not require any extra image annotations, hence can be easily scaled to arbitrary size of vision language training images. Formally, we define the SSUL as:

$$SSUL=\text{{L2}}(Mean(\mathbf{m})|\mathbf{w}_{\backslash\mathbf{m}},M),$$

(1)

where $\mathbf{m}$ is the masked areas, $\mathbf{w}_{\backslash\mathbf{m}}$ represents all other areas except for the masked area and $M$ represents the groundtruth mean color value of the masked area.

One of the biggest challenges in vision language pretraining is that languages and images are features of different levels. Languages usually contain high-level semantic information. For example, one language word like “car” means one car object. In contrast, image information is low-level raw information. One image pixel does not have a specific meaning. Objects and semantic information are represented by groups of pixels.

When the vision language models are trained with region feature input, the input vision features are already high-level semantic features, such that the vision and language correspondence can be learned fast. When the vision language models are trained with raw pixel input, the vision branch will be in charge of learning visual concepts.

To this end, we propose to help learning visual branch with an auxiliary Semantic Segmentation Loss (SEGL). Specifically, given vision output $X\in R^{\hat{H}\times\hat{W}}$, and groundtruth annotation $Y\in R^{H\times W}$, where $W$ and $H$ are image weight and height and $\hat{W}$ and $\hat{H}$ are feature map width and height, the SEGL is calculated by downsampling the annotation and calculating the cross-entropy loss between groundtruth and generated feature map:

$$SEGL=\text{{Cross-entropy}}(X,downsample(Y)).$$

(2)

The SEGL is the same as the commonly used semantic segmentation loss. In normal semantic segmentation task, the networks usually up-sample the feature map and use semantic segmentation loss on the high-resolution feature map, which may include additional learnable parameters during up-sampling. In our task, we use SEGL as an auxiliary training target. Instead of up-sampling the feature map, we down-sample the groundtruth label map. The detailed network structure of Semantic Segmentation Loss is shown in Fig 4 C.

Although the Semantic Segmentation Loss (SEGL) could provide semantic guidance to the network, it requires additional human annotation data, which is not preferable when training scales up. To overcome this problem, we propose the Knowledge Distillation Based Set Prediction Loss (SPL). SPL first runs a pretrained object detector to get $N$ object labels $y=\left\{y_{i}\right\}^{N}_{i=1}$ from the raw images. Unlike region feature based methods, the object detector is only run once before the pretraining step. The feature extraction process is not required during inference.

For the network part, given vision output $\hat{X}\in R^{\hat{N}\times\hat{M}}$, the network applies transformer decoder layers [19] to map the feature to $X\in R^{N}$, where $N<<\hat{N}\times\hat{M}$. The $N$ dimensional feature represents $N$ predicted objects. Then the network applies a Softmax layer and generates $N$ object predictions possibilities $\hat{p}=\left\{\hat{p}_{i}\right\}^{N}_{i=1}$. The detailed network structure is shown in Fig 4 D.

SPL supervises the network training by matching the set of predictions $\hat{p}$ with the set of distillated labels $y$. The optimal bipartite matching problem is widely studied in multi-object tracking [29, 30] and is recently also extended to object detection [31]. The calculation of SPL is done in two steps: find optimal matching between the two sets and then calculate object label cross-entropy loss.

Optimal Matching. The optimal matching between two sets is calculated by the Hungarian algorithm. Formally, given two sets $y=\left\{y_{i}\right\}^{N}_{i=1}$ and $\hat{p}=\left\{\hat{p}_{i}\right\}^{N}_{i=1}$, The Hungarian algorithm finds bipartite matching by minimizing the cost of assigning correct labels $-\hat{p}_{\sigma(i)}(c_{i})$. Formally, the optimal matching $\hat{\sigma}$ is calculated by:

$$\hat{\sigma}=\operatorname*{arg\,min}_{\sigma}\sum_{i}^{N}[-\hat{p}_{\sigma(i)}(c_{i})],$$

(3)

where $c_{i}$ is the category of object index $i$ and $\sigma$ is a permutation of the prediction $\hat{p}$ and $\hat{\sigma}$ represents the optimal matching.

Object Label Classification Loss. After the optimal bipartite matching is found, we perturb the generated labels with calculated optimal bipartite matching $\hat{\sigma}$ and continue to calculate the SPL loss as an object label cross-entropy loss between the matched prediction and distillated labels:

$$SPL=\text{{Cross-entropy}}(\hat{p}_{\hat{\sigma}},y).$$

(4)

To summarize and compare our proposed methods, we list extra computation cost, extra data, scalability, and inference speed of our methods vs baseline in Table  I.

In this section, we describe the finetuning settings on downstream tasks.

We start by building a naive version of pixel input vision language pretraining model using vision transformer [10] backbone. Same as Pixel-BERT [8], the naive version uses only a matching loss and the language modeling loss. We compare our proposed three losses with the baseline version on three tasks. A detailed ablation study is shown in Table II and Table III. For all ablation experiments, we use finetuning input size 600. Following [32], the inference time is calculated by the running speed on Nvidia P40 GPU.

We show that all of our proposed methods surpass the baseline model with a large margin. Within the proposed methods, the Semantic Segmentation Loss achieves the best performance in general. It proves that additional human annotations do help to improve pretraining vision language models. Self-supervised loss and set prediction loss also produce good results on down-stream tasks – less than 1 percent lower than semantic segmentation loss on VQA. Among all of them, the set prediction loss has a relatively faster process speed and lower memory usage. For the cross-modality module, the set prediction loss uses 36 or 100 visual features, while other vision transformer models use 324 (18*18) visual features.

By comparing ResNet [9] models and vision transformer models [10], we conclude that the pretraining difficulty exists regardless of the backbone choice and the backbones won’t affect the performance of our proposed losses.

We compare our best result (Vision Transformer + SEGL) with region feature models and pixel input models in Table IV, Table V and Table VI. For all experiments in this section, we continue to use image size 600x600 except for VQA. For VQA we use image size 800x800 to further improve performance, same as Pixel-BERT [8].

Compared with region feature models, when trained with fewer GPU hours, our method could achieve almost the same performance on various datasets. During inference, an end-to-end model could run more than 10x faster than any region feature based model.

Compared with other end-to-end models, our model is trained for only 1/10 of the pretraining GPU hours and gets competitive or even better finetuning performance.