KD-VLP: Improving End-to-End Vision-and-Language Pretraining with
Object Knowledge Distillation

We adopt a CNN backbone to extract image features $\mathcal{V}$ = $\{\mathbf{v}_{i}\}_{i=1}^{L}$ for each image $I$ where $L$ is the size of feature grids and $\mathbf{v}_{i}\in\mathcal{R}^{d_{v}}$ is a feature vector with dimension $d_{v}$. In addition, each feature is further concatenated with its 2-D sine position embedding Carion et al. (2020). Following SOHO, we use a ResNet-101He et al. (2016) as the visual backbone, followed by additional 1x1 Conv and 2x2 strides Max-pooling to reduce the memory footprint.

For the language $D$, we first tokenize the sentence into a sequence of word tokens using WordPiece Wu et al. (2016), then encode them into word embeddings $\mathcal{W}$ = $\{\mathbf{w}_{j}\}_{j=1}^{T}$ where $\mathbf{w}_{j}\in\mathcal{R}^{d_{w}}$ is the feature vector. Similarly, an index position Devlin et al. (2018) embedding is supplemented to each word embedding.

After obtaining image and linguistic embeddings, we assemble them into a sequence of tokens $\{\mathcal{V},\mathbf{[sep]},\mathcal{W},\mathbf{[cls]}\}$, and adopt a multi-layer Transformer to compute their representations encoded by the final-layer states $\{\mathcal{H}_{\mathcal{V}},\mathbf{h}_{\mathbf{sep}},\mathcal{H}_{\mathcal{W}},\mathbf{h}_{\mathbf{cls}}\}$ where $\mathcal{H}_{\mathcal{V}}$ = $\{\mathbf{h}_{\mathbf{v}_{i}}\}_{i=1}^{L}$ and $\mathcal{H}_{\mathcal{W}}$ = $\{\mathbf{h}_{\mathbf{w}_{j}}\}_{j=1}^{T}$ represent the states for visual and language part respectively. Finally, those representations are sent into each pretext task head to compute the supervision signals.

The first task aims to learn more explicit object concepts in the E2E pretraining. Specifically, we sample an object each time and mask out its features in the Transformer input, and enforce the network to generate external object RoI features and semantic labels. To learn better cross-modal alignment, we propose a knowledge-guided masking strategy, which samples noun phrase-related object regions to mask based on the (normalized) similarity score $\alpha_{z,n}$. The selected object region is denoted with its binary mask, category and RoI features, as $(\mathbf{m}^{*},c^{*},\mathbf{f}^{*})$.

We design two learning objectives, Masked Region Classification (MRC) and Masked Region Feature Regression (MRFR) as below

To calculate the losses $L_{\rm MRC}$ and $L_{\rm MRFR}$, we first compute the object representation $\mathbf{h}_{\mathbf{m}^{*}}$ for the masked region at the final layer, which is average-pooled over $\mathcal{H}_{\mathcal{V}}$ based on its binary mask $\mathbf{m}^{*}$. For MRC, a multi-layer FC network $\Phi_{{\rm MRC}}$ is adopted to predict its object category. Thus, $L_{{\rm MRC}}$= ${\rm CE}(\Phi_{\rm{MRC}}(\mathbf{h}_{\mathbf{m}^{*}}),c^{*})$ is the standard cross-entropy loss. In addition, we take another FC network $\Phi_{{\rm MRFR}}$ to learn the object concept in feature space directly by minimizing the L2 distance, $L_{{\rm MRFR}}$ = $||\Phi_{{\rm MRFR}}(\mathbf{h}_{\mathbf{m}^{*}})-\mathbf{f}^{*})||_{2}^{2}$.

The second task, PRA, mainly focuses on learning cross-modal alignment at object-level, which aims to pull positive phrase-region pairs closer and push negative pairs away. Here we utilize the similarity $\alpha_{z,n}$ between the noun phrase and object category in the linguistic space as a guidance.

Concretely, we first compute the object representation ${\mathbf{h}_{\mathbf{m}_{n}}}$ for each proposal and the phrase representation $\mathbf{h}_{p_{z}}$, both of which are obtained from the final layer states of the Transformer. Specifically, ${\mathbf{h}_{\mathbf{m}_{n}}}$ is average-pooled over $\mathcal{H}_{\mathcal{V}}$ based on binary mask ${\mathbf{m}}_{n}$ while $\mathbf{h}_{p_{z}}$ = $\frac{1}{|p_{z}|}\sum_{j\in p_{z}}$ $\mathbf{h}_{\mathbf{w}_{j}}$ represents average states of word tokens within $p_{z}$. We define the cross-modal similarity as $\hat{\alpha}_{z,n}$ = ${\rm Cos}(\mathbf{h}_{p_{z}},{\mathbf{h}_{\mathbf{m}_{n}}})$.

The task PRA minimizes the KL-divergence between the cross-modal similarities $\hat{\alpha}_{z}$ = $\{{\rm Softmax}(\hat{\alpha}_{z,n})\}_{n=1}^{N}$ and the phrase-label similarities $\alpha_{z}$ = $\{{\rm Softmax}(\alpha_{z,n})\}_{n=1}^{N}$ as below:

Finally, denoting the mask set $\mathcal{M}=\{\mathbf{m}_{n}\}_{n=1}^{N}$, we have the overall PRA loss function as follows:

We take the same masking strategy (15% prob. to mask) as in BERT Devlin et al. (2018) to randomly mask out the input word tokens. Here, MLM aims to predict the original word index in vocabulary space for each masked token based on the whole image and its surrounding language context via the Transformer. Hence a cross-entropy loss is adopted:

In ITM, the multi-layer Transformer is trained to distinguish whether the input image-text pairs are semantically matched based on the final layer [cls] token representation $\mathbf{h}_{\mathbf{cls}}$. To construct the training samples, we randomly replace the text for each image-text pair with another text from dataset with a probability of 0.5. Thus, the output label can be defined as $y\in\{0,1\}$ where $y=1$ indicates matched pair. The training objective for the ITM task is to minimize binary cross-entropy loss:

Following the E2E pretraining strategy Huang et al. (2021, 2020); Xu et al. (2021), we take indomain datasets: MSCOCO  Lin et al. (2014) and VG Krishna et al. (2016) as pretraining datasets since it is widely used in literature. In total, two datasets comprise about 200K images and 5.6M image-text pairs, where each image is associated with multiple captions.

We follow BERT to tokenize caption into word tokens by using WordPiece, and resize the image into (800, 1333) as prior works. For model architecture, a widely-used ResNet101 for visual encoding and 12-layer Transformer for multi-modal fusion are adopted for a fair comparison. Both networks are initialized with ImageNet and BERT pretrained parameters. Besides, following the majority of two-step methods, we apply the widely-used object detector BUTD Anderson et al. (2018) to generate object proposals as well as their RoI embeddings as our supervision.

For model learning, we optimize the entire network by using SGD for CNNs with a learning rate of 1e-2 and AdamW for Transformer with a learning rate of 1e-4, as suggested in SOHO. The training iterations are up to 100K with batch-size 512 in each. The learning rate decays 10 times at 20K, 40K respectively. All experiments are conducted on 16 NVIDIA V100 GPUs with mixed-precision training to reduce memory cost about 7 days.

The baseline takes standard ITM and MLM to train the entire model. In Tab.3, it still achieves decent results over various VL tasks.

As in Tab.3, compared with baseline, OMVM presents a clearly consistent improvement over all downstream tasks. It suggests that OMVM can enhance the end-to-end multi-modal representations with explicit object concepts learning. In addition, the knowledge-guided masking strategy further helps establish cross-modal correspondence.

To further investigate the OMVM task, we randomly mask a box region with 15% probability rather than sampling a region based on the normalized similarity score $\alpha_{z,n}$, denoted as RandomMVM. The other pretraining details are the same as in OMVM. We observe a significant performance drop over all downstream tasks, especially in image-text retrieval and NLVR². It indicates that simple RandomMVM will result in inefficient multi-modal representation learning because there is a high probability that the selected region has no relationship with the associated description.

In addition, we also explore the similar masked feature regression task as in UNITER by randomly masking out the image grid features as in BERT and then requiring the Transformer to reconstruct its original features rather than the external object RoI embeddings, denoted as StandardMVM. The results show that such StandardMVM fails to facilitate multi-modal representation learning in the E2E framework.

The OMVM above mainly focuses on instance-level knowledge distillation by absorbing external object RoI features and semantic labels. Different from that, PRA aims to establish positive object-phrase correspondence while suppressing the negative ones under the guidance of similarities between noun phrases and object labels in linguistic space. As in Tab. 3, we significantly improve 0.78% R@1 of MSCOCO-TR and 1.87% in MSCOCO-IR. In addition, PRA shows slight improvements for more challenging fine-grained reasoning tasks, like VE, NLVR², and VQA. The results indicate that PRA is beneficial to multi-modal representation learning.

In Fig.2(a), our knowledge-guided masking strategy always masks out the phrase-related image regions, which can facilitate multi-modal learning. On the contrary, previous works, like SOHO, VILLA …, mask out background regions or part of the object region with a high probability, which have no relationship with the corresponding description and result in inefficient cross-modal alignment. Fig.2(b) demonstrates the word-to-image attention maps. Compared to SOHO, our method can attend more accurately to image regions for the corresponding word. Surprisingly, even the word "smiling" can locate the baby’s face correctly, which suggests that our approach not only learns better noun-region alignment but also helps establish high-order correspondence, like actions. (see Suppl. for more visualization.)

We adopt the default BUTD detector in a typical 2-step pretraining method for a largely fair comparison. To investigate the influence of object detectors, we also conduct pretraining with objects knowledge extracted from FRCNN-RN101 pretrained on COCO. In Tab.4, we observe a performance drop compared with the model pretrained with BUTD, which suggests large object knowledge space will facilitate multimodal pretraining. Besides, although with COCO detector, we still outperform SOHO by a clear margin, indicating the superiority of object knowledge in E2E pretraining framework.

In Tab.5, we show the individual contributions of our proposed tasks. MRC, MRFR, PRA pretext tasks all help facilitate multi-modal representation learning and improve the performance compared with the baseline model as a result.

We take SOHO as a strong baseline and compare it at different model sizes (ResNet18 + 3-layer Transformer, ResNet101 + 12-layer Transformer) to investigate the impact of object knowledge distillation. Fig.3 demonstrates the performance gains over some representative vision-language tasks. It shows that object concepts learning always helps multi-modal representation learning no matter what model size it is. In VE and text-retrieval, the larger model even improves significantly than the light-weighted model and shows more capacities to learn external object semantics knowledge.