VLM O : Unified Vision-Language Pre-Training with Mixture-of-Modality-Experts
Abstract
We present a unified V ision-
L anguage pretrained Mo del ( VLM O ) that jointly learns a dual encoder and a fusion encoder with a modular Transformer network.
Specifically, we introduce Multiway Transformer, where each block contains a pool of modality-specific experts and a shared self-attention layer.
Because of the modeling flexibility of Multiway Transformer, pretrained VLM O can be fine-tuned as a fusion encoder for vision-language classification tasks, or used as a dual encoder for efficient image-text retrieval.
Moreover, we propose a stagewise pre-training strategy, which effectively leverages large-scale image-only and text-only data besides image-text pairs.
Experimental results show that VLM O achieves state-of-the-art results on various vision-language tasks, including VQA,
NLVR2 and image-text retrieval.
The code and pretrained models are available at http://aka.ms/vlmo .
Introduction
Vision-Language (VL) pre-training [31, 42, 36, 27, 21, 24] learns generic cross-modal representations from large-scale image-text pairs.
Previous models usually employ image-text matching, image-text contrastive learning, masked region classification/feature regression, word-region/patch alignment and masked language modeling to aggregate and align visual and linguistic information.
Then the pretrained models can be directly fine-tuned on downstream vision-language tasks, such as VL retrieval and classification (visual question answering, visual reasoning, etc.)
Two mainstream architectures are widely used in previous work.
CLIP [36] and ALIGN [19] adopt a dual-encoder architecture to encode images and text separately.
Modality interaction is handled by the cosine similarity of the image and text feature vectors.
The dual-encoder architecture is effective for retrieval tasks, especially for masses of images and text.
Feature vectors of images and text can be pre-computed and stored.
However, the shallow interaction between images and text is not enough to handle complex VL classification tasks.
ViLT [21] finds that CLIP gives a relatively low accuracy on visual reasoning task.
Another line of work [31, 42, 44, 4, 21, 24] relies on a fusion encoder with cross-modal attention to model image-text pairs.
Multi-layer Transformer [46] networks are usually employed to fuse image and text representations.
The fusion-encoder architecture achieves superior performance on VL classification tasks.
But it requires to jointly encode all possible image-text pairs to compute similarity scores for retrieval tasks.
The quadratic time complexity leads to a much slower inference speed than the dual-encoder models whose time complexity is linear.
In order to take advantage of the two types of architectures, we propose a unified V ision-
L anguage pretrained Mo del ( VLM O ) that can be used as either a dual encoder to separately encode images and text for retrieval tasks, or used as a fusion encoder to model the deep interaction of image-text pairs for classification tasks.
This is achieved by introducing Multiway Transformer that can encode various modalities (images, text, and image-text pairs) within a Transformer block.
Multiway Transformer employs a pool of modality experts to replace the feed-forward network in standard Transformer.
It captures modality-specific information by switching to different modality experts, and uses the shared self-attention across modalities to align visual and linguistic information.
Specifically, Multiway
Transformer consists of three modality experts, namely vision expert for image encoding, language expert for text encoding, and vision-language expert for image-text fusion.
Thanks to the modeling flexibility, we can reuse Multiway Transformer with the shared parameters for different purposes, i.e., text-only encoder, image-only encoder, and image-text fusion encoder.
VLM O is jointly learned with three pre-training tasks, namely image-text contrastive learning, imagetext matching, and masked language modeling.
In addition, we propose a stagewise pre-training strategy to effectively leverage large-scale image-only and text-only corpus besides image-text pairs in VLM O pre-training.
We first pretrain vision experts and self-attention modules of Multiway
Transformer on image-only data using masked image modeling proposed in BE I T [3].
We then pretrain language experts on text-only data using masked language modeling [11].
Finally, the model is used to initialize vision-language pre-training.
By getting rid of the limited size of image-text pairs and their simple and short captions, stagewise pre-training on large amounts of image-only and text-only data helps VLM O to learn more generalizable representations.
Experimental results demonstrate that VLM O achieves state-of-the-art results on vision-language retrieval and classification tasks.
Our model, used as a dual encoder, outperforms fusion-encoderbased models [4, 15, 21, 24] while enjoying a much faster inference speed on retrieval tasks.
Moreover, our model also achieves state-of-the-art results on visual question answering (VQA) and natural language for visual reasoning (NLVR2), where VLM O is used as a fusion encoder.
Our main contributions are summarized as follows:
• We propose a unified vision-language pretrained model VLM O that can be used as a fusion encoder for classification tasks, or fine-tuned as a dual encoder for retrieval tasks.
• We introduce a general-purpose multimodal Transformer for vision-language tasks, namely Multiway Transformer, to encode different modalities.
It captures modality-specific information by modality experts, and aligns contents of different modalities by the self-attention module shared across modalities.
• We show that stagewise pre-training using large amounts of image-only and text-only data greatly improves our vision-language pretrained model.
Related Work
Pre-training with Transformer [46] backbone networks has substantially advanced the state of the art across natural language processing [35, 11, 29, 23, 12, 37, 2, 8, 9, 5–7, 32], computer vision [13, 45, 3] and vision-language [44, 42, 4, 51, 36, 19, 21, 24, 47] tasks.
The approaches of vision-language pre-training can be divided into two categories.
The first category utilizes a dual encoder to encode images and text separately, and uses cosine similarity or a linear projection layer to model the interaction between images and text [36, 19].
Image-text contrastive learning is usually employed to optimize the model.
Dual-encoder models are effective for visionlanguage retrieval tasks.
However, the simple interaction is not enough to handle tasks that require complex reasoning, such as visual reasoning and visual question answering (VL classification tasks).
The second category models the interaction of images and text using a deep fusion encoder with cross-modal attention [44, 31, 42, 25, 53, 4, 27, 26, 15, 51, 17, 18, 21, 24, 48].
Image-text matching, masked language modeling, word-region/patch alignment, masked region classification and feature regression are widely used to train fusion-encoder-based models.
These models achieve better performance for vision-language classification tasks, while the joint encoding of all image-text pairs leads to a slow inference speed for retrieval tasks.
A large portion of fusion-encoder-based models rely on an off-the-shelf object detector like Faster R-CNN [38] to obtain image region features.
Generating region features slows down the inference speed and renders the approach less scalable.
Recently, Pixel-BERT [17] removes object detector and encodes images into grid features by convolutional neural networks.
ALBEF [24] employs image Transformer [13, 45] to obtain the representations of images, and uses text Transformer [11] to learn the contextualized representations
[Figure]
Shared Parameters
[Figure]
of text.
These representations are then fused by cross-modal attention.
ViLT [21] encodes images into patch embeddings, and then feed the concatenation of image patch embeddings and word embeddings into a Transformer network to learn contextualized representations and model the interaction of images and text.
Different from previous work, our unified pre-training using shared Multiway Transformer enables the model perform separate encoding for retrieval tasks, and jointly encode image-text pairs to capture deeper interaction for classification tasks.
Our model achieves competitive performance, while enjoying a faster inference speed for both retrieval and classification tasks.
Methods
Given image-text pairs, VLM O obtains image-only, text-only and image-text pair representations by the Multiway Transformer network.
As shown in Figure 1, the unified pre-training optimizes shared Multiway Transformer with image-text contrastive learning on image-only and text-only representations, image-text matching and masked language modeling on image-text pair representations.
Thanks to the modeling flexibility, the model can be used as a dual encoder for retrieval tasks to encode images and text separately during fine-tuning.
It can also be fine-tuned as a fusion encoder to model deeper modality interaction of images and text for classification tasks.
[Figure]
3.1 Input Representations
Given an image-text pair, we encode the pair into image, text and image-text vector representations.
These representations are then fed into the Multiway Transformer to learn contextualized representations and align image and text feature vectors.
Image Representations Following vision Transformers [13, 45, 3], the 2D image v ∈ R H × W × C is split and reshaped into N = HW/P 2 patches v p ∈ R N ×
( P 2 C ) , where C is the number of channels, ( H, W ) is the resolution of the input image, and ( P, P ) is the patch resolution.
The image patches are then flattened into vectors and are linearly projected to obtain patch embeddings.
We also prepend a learnable special token
[I_CLS] to the sequence.
Finally, image input representations are obtained via summing patch embeddings, learnable 1D position embeddings V pos ∈ R ( N +1)
× D and image type embedding V type ∈ R D : H v 0 =
[ v [I_CLS] , V v pi , . .
. , V v pN ]
+ V pos + V type , where
H v 0 ∈ R ( N +1) × D , linear projection V ∈ R
( P 2 C ) × D .
Text Representations Following BERT [11], we tokenize the text to subword units by WordPiece [49].
A start-of-sequence token ( [T_CLS] ) and a special boundary token ( [T_SEP] ) are added to the text sequence.
Text input representations H w 0 ∈ R ( M +2)
× D are computed via summing the corresponding word embedding, text position embedding and text type embedding
H i Image-Text Representations We concatenate image and text input vectors to form the image-text input representations H vl 0 =
[Equation]
H
[Equation]
[Equation]
H v 0 ]
3.2 Multiway Transformer
Inspired by mixture-of-experts networks [41, 14], we propose a general-purpose multimodal Transformer for vision-language tasks, namely Multiway Transformer, to encode different modalities.
Multiway Transformer introduces mixture of modality experts as a substitute of the feed forward network of standard Transformer.
Each modality expert is also the feed forward network which consists of two linear transformations and an activation.
Given previous layer’s output vectors
H l − 1 , l ∈ [1 , L ] , each Multiway Transformer block captures modality-specific information by switching to different modality expert, and employs multi-head self-attention (MSA) shared across modalities to align visual and linguistic contents.
LN is short for layer normalization.
[Equation]
Multiway - FFN selects an expert among multiple modality experts to process the input according to the modality of the input vectors H ′ l and the index of the Transformer layer.
Specifically, there
[Figure]
are three modality experts: vision expert (V-FFN), language expert (L-FFN) and vision-language expert (VL-FFN).
If the input is image-only or text-only vectors, we use vision expert for encoding images and language expert for encoding text.
If the input consists of vectors of multiple modalities, such as the vectors of image-text pair, we employ vision expert and language expert to encode the respective modality vectors at the bottom Transformer layers.
Vision-language expert is then used at the top layers to capture more modality interaction.
Compared with conventional mixture-of-experts networks [41, 14], Multiway Transformer conducts hard routing according the input modality.
Given the three types of input vectors, we obtain image-only, text-only and image-text contextualized representations.
3.3 Pre-Training Tasks
VLM O is jointly pretrained by image-text contrastive learning on the image and text representations, masked language modeling and image-text matching on the image-text pair representations with shared parameters.
Image-Text Contrast Given a batch of N image-text pairs, image-text contrastive learning aims to predict the matched pairs from N × N possible image-text pairs.
There are N 2 − N negative image-text pairs within a training batch.
The final output vectors of [I_CLS] token and [T_CLS] token are used as the aggregated representation of the image and text, respectively.
Followed by a linear projection and normalization, we obtain image vectors { ˆ h vi } Ni =1 and text vectors { ˆ h wi } Ni =1 in a training batch to compute image-to-text and text-to-image similarities:
[Equation]
Where s i 2 t i,j represents image-to-text similarity of image of i -th pair and text of j -th pair, s t 2 i i,j is the text-to-image similarity.
ˆ h wi ∈ R D and ˆ h vj ∈ R D indicate the normalized vectors of i -th text and j -th image, σ is a learned temperature parameter.
p i
2 t i and p t 2 i
i are the softmax -normalized similarities.
Cross-entropy losses over image-to-text and text-to-image similarities are used to train the model.
Masked Language Modeling Following BERT [11], we randomly choose tokens in the text sequence, and replace them with the [MASK] token.
The model is trained to predict these masked tokens from all the other unmasked tokens and vision clues.
We use 15 % masking probability as in BERT.
The final output vectors of masked tokens are fed into a classifier over the whole text vocabulary with cross-entropy loss.
Image-Text Matching Image-text matching aims to predict whether the image and text is matched.
We use the final hidden vector of the [T_CLS] token to represent the image-text pair, and feed the vector into a classifier with cross-entropy loss for binary classification.
Inspired by ALBEF [24], we sample hard negative image-text pairs based on the contrastive image-to-text and text-to-image similarities.
3.4 Stagewise Pre-Training
We introduce a stagewise pre-training strategy, which leverages large-scale image-only and text-only corpus to improve the vision-language model.
As present in Figure 2, we first perform vision pretraining on image-only data, and then perform language pre-training on text-only data to learn general image and text representations.
The model is used to initialize the vision-language pre-training to learn the alignment of visual and linguistic information.
For vision pre-training, we train the attention module and vision expert of Multiway Transformer as in BE I T [3] on image-only data.
We directly utilize the pretrained parameters of BE I T to initialize the attention module and vision expert.
For language pre-training, we freeze parameters of the attention module and vision expert to avoid catastrophic forgetting of vision knowledge learned in the first stage, and utilize masked language modeling [11] to optimize the language expert on text-only data.
Compared with image-text pairs, image-only and text-only data are easier to collect.
In addition, text data of image-text pairs is usually short and simple.
Pre-training on image-only and text-only corpus improves the generalization on complex pairs.
3.5 Fine-Tuning VLM O on Downstream Tasks
As present in Figure 3, our model can be fine-tuned to adapt to various vision-language retrieval and classification tasks.
Vision-Language Classification For classification tasks such as visual question answering and visual reasoning, VLM O is used as a fusion encoder to model modality interaction of images and text.
We use the final encoding vector of the token [T_CLS] as the representation of the image-text pair, and feed it to a task-specific classifier layer to predict the label.
Vision-Language Retrieval For retrieval tasks, VLM O can be used as a dual encoder to encode images and text separately.
During fine-tuning, our model is optimized for the image-text contrastive loss.
During inference, we compute representations of all images and text, and then use dot product to obtain image-to-text and text-to-image similarity scores of all possible image-text pairs.
Separate encoding enables a much faster inference speed than fusion-encoder-based models.
Experiments
We pretrain our model using large-scale image-text pairs and evaluate the model on visual-linguistic classification and retrieval tasks.
4.1 Pre-Training Setup
Following previous work [4, 21], our pre-training data consists of four image captioning datasets:
Conceptual Captions (CC) [40], SBU Captions [33], COCO [28] and Visual Genome (VG) [22] datasets.
There are about 4 M images and 10 M image-text pairs in the pre-training data.
Our models adopt the same network configuration as ViT [13] and BE I T [3].
VLM O -Base consists of 12 -layer Transformer blocks with 768 hidden size and 12 attention heads.
VLM O -Large is a 24 -layer Transformer network with 1024 hidden size and 16 attention heads.
VLM O -Base uses vision-language expert on the top two Transformer layers, and VLM O -Large introduces visionlanguage expert on the top three layers.
VLM O -Base consists of 175M parameters and VLM O -Large contains 562M parameters.
For images, the input resolution is 224 × 224 and the patch size is 16 ×during pre-training.
We apply RandAugment [10] to the input images.
The tokenizer of the uncased version of BERT is employed to tokenize the text.
The maximum text sequence length is set to 40 .
We also employ whole word masking for the masked language modeling pre-training task.
We pretrain the models for 200 k steps with 1024 batch size.
We utilize AdamW [30] optimizer with β 1 = 0 .
9 , β 2 = 0 .
The peak learning is 2e-4 for the base-size model, 5e-5 for the large-size model.
Weight decay is set to 0 .
We use linear warmup over the first 2 .
5 k steps and linear decay.
[Table]
VLM O -Large++ is the model trained on one billion noisy image-text pairs with a larger batch size.
We report vqa-score on VQA test-dev and test-standard split, and report accuracy for
NLVR2 development and public test set (test-P).
4.2 Training on Larger-scale Datasets
We scale up vision-language representation learning by training VLM O -Large on one billion noisy web image-text pairs with a larger batch size.
We first pretrain the model for 200k steps with 16k batch size, and then continue train the model for 100k steps with 32k batch size.
The other hyperparameters are the same as the training on 4M data.
Please refer to the supplementary material for more details of hyper-parameters used for pre-training and fine-tuning.
4.3 Evaluation on Vision-Language Classification Tasks
We first conduct fine-tuning experiments on two widely used classification datasets: visual question answering [16] and natural language for visual reasoning [43].
The model is fine-tuned as a fusion encoder to model deeper interaction.
Visual Question Answering (VQA) For VQA, a natural image and a question are given, the task is to generate/choose the correct answer.
We train and evaluate the model on VQA 2.0 dataset [16].
Following common practices, we convert VQA 2.0 to a classification task, and choose the answer from a shared set consists of 3 , 129 answers.
We use the final encoding vector of the [T_CLS] token as the representation of the image-question pair and feed it to a classifier layer to predict the answer.
Natural Language for Visual Reasoning (NLVR2)
The NLVR2 [43] dataset requires the model to predict whether a text description is true about a pair of images.
Following OSCAR [27] and
VinVL [51], we convert the triplet input to two image-text pairs, each containing the text description and one image.
We concatenate the final output vectors of the [T_CLS] token of the two input pairs.
The concatenated vector is then fed into a classification layer to predict the label.
We present the results of VL classification tasks in Table 1.
VLM O achieves state-of-the-art performance and substantially outperforms previous methods.
Our large-size model even outperforms
SimVLM-Huge [48] and Florence-Huge [50] by a large margin, which consists of more parameters and are also trained on larger-scale image-text pairs.
Our model uses a simple linear projection to embed images as in ViLT [21].
This leads to a significant speedup compared with previous models using image region features, which are extracted by an off-the-shelf object detector [31, 42, 4, 15, 26, 51].
[Table]
† : ALIGN, Florence and our model encode images and text separately, and then employ a shallow interaction (dot product) to obtain the similarity scores.
‡ : ALBEF first encodes images and text separately to obtain the top- k candidates, and then feed these representations into a fusion encoder to rerank the candidates.
The others require to encode all image-text combinations by a fusion encoder.
VLM O -Large++ represents the model trained on one billion noisy image-text pairs with a larger batch size.
[Table]
4.4 Evaluation on Vision-Language Retrieval Tasks
The retrieval tasks contain image-to-text retrieval and text-to-image retrieval.
We evaluate the model on the widely used COCO [28] and Flickr30K [34] datasets, and use the Karpathy split [20] for both datasets.
The model is used as a dual encoder for retrieval tasks.
We encode images and text separately and compute their similarity scores by the dot product of image and text vectors.
As present in Table 2, VLM O achieves competitive performance with previous fusion-encoder-based models while having a much faster speed.
Fusion-encoder-based models need to jointly encode all possible image-text pairs to compute their similarity scores, which requires quadratic time complexity.
Moreover, our large-size model even outperforms the huge-size model of Florence [50], which also trained on massive image-text pairs using a larger batch size.
VLM O pre
-training can effectively leverage larger-scale noisy pairs and benefit from large batch training.
4.5 Evaluation on Vision Tasks
As shown in Table 4, we use VLM O as an image-only encoder and evaluate it on image classification (ImageNet [39]) and semantic segmentation (ADE20K [52]) tasks.
The model also achieves competitive performance, even slightly better than the BE I T model used for the initialization of VLM O .
The image resolution is 224 × 224 for ImageNet, and 512 × 512 for ADE20K.
We perform intermediate fine-tuning [3] on ImageNet-21k for all three models.
[Table] [Table]
4.6 Ablation Studies
Stagewise Pre-Training We first conduct ablation experiments of stagewise pre-training.
ViLT [21] shows that using the ViT [13] model pretrained on image-only data as the initialization achieves better performance than the BERT model pretrained on text-only data.
Therefore we start experiments with image-only pre-training.
We compare using image-only pre-training, and image-only pre-training plus text-only pre-training as the initialization.
For image-only pre-training, we directly use the parameters of BE I T-Base to initialize the self-attention module and all modality experts.
For imageonly pre-training plus text-only pre-training, we use pretrained parameters of BE I T-Base to initialize the vision expert and self-attention module of Multiway Transformer, and then pretrain its language expert on text corpora.
As shown in Table 3, image-only pre-training plus text-only pre-training improves our vision-language model.
We also have tried to perform vision-language pre-training with random initialization but obtain a relatively low accuracy on downstream tasks.
Stagewise pre-training effectively leverages large-scale image-only and text-only corpus, and improves our vision-language pre-training.
Moreover, given the limited size of image-text pairs we used during pretraining, stage-wise pre-training on image-only and text-only data alleviates the need for image-text pair data.
We have tried to perform multitask training on image-only and text-only data to combine the first two stages and observe similar performance.
Stage-wise pre-training effectively leverages the pretrained weights to reduce the computation cost.
Multiway Transformer
We also conduct ablation experiments of Multiway Transformer.
We employ ViT-Base to initialize the models for the ablation experiments.
As present in Table 5, using
Multiway Transformer achieves better performance than standard Transformer for both retrieval and classification tasks.
In addition, we also analyse the contribution of vision-language expert (VL-FFN) used in Multiway Transformer.
We remove the vision-language expert used in the top Transformer layers.
Experimental results demonstrate that the introduction of vision-language expert improves the model.
Using vision-language expert captures more modality interaction.
Pre-Training Tasks We perform ablation studies to analyse the contribution of different pre-training tasks, and the results are presented in Table 5.
Compared with the model trained only using image-text contrastive loss, our unified training performs much better across classification and retrieval tasks.
Introducing image-text matching with hard negative mining also greatly improves the model.
This demonstrates the effectiveness of our unified-training framework with Multiway Transformer.
In addition, experimental results show that masked language modeling positively contribute to our model.
Please refer to the supplementary material for more ablation studies.
[Table]
Global Hard Negative Mining Different from ALBEF [
, which samples hard negatives from training examples of the single GPU (named as local hard negative mining).
We perform hard negative mining from more candidates by gathering training examples of all GPUs (named as global hard negative mining).
As shown in Table 6, our global hard negative mining brings significant improvements.
Conclusion
In this work, we propose a unified vision-language pretrained model VLM O , which jointly learns a dual encoder and a fusion encoder with a shared Multiway Transformer backbone.
Multiway
Transformer introduces a pool of modality experts to encode modality-specific information, and aligns different modalities using the shared self-attention module.
The unified pre-training with
Multiway Transformer enables the model to be used as a dual encoder for efficient vision-language retrieval, or as a fusion encoder to model cross-modal interactions for classification tasks.
We also show that stagewise pre-training that leverages large-scale image-only and text-only corpus greatly improves vision-language pre-training.
Experimental results demonstrate that VLM O outperforms previous state-of-the-art models on various vision-language classification and retrieval benchmarks.
In the future, we would like to work on improving VLM O from the following perspectives:
• We will scale up the model size used in VLM O pre-training.
• We are also interested in fine-tuning VLM O for vision-language generation tasks, such as image captioning, following the method proposed in UniLM [12].
• We are going to explore to what extent vision-language pre-training can help each other modality, especially as the shared Multiway Transformer backbone naturally blends in text and image representations.
• We can extend the proposed model to integrate more modalities (e
g., speech, video, and structured knowledge), supporting general-purpose multimodal pre-training.