An Empirical Study of Training End-to-End Vision-and-Language Transformers

In plain words

Vision-and-language models answer questions about pictures, check whether a caption fits an image, and search images by text. Most of the strong ones have so far looked at an image through a separate object detector that first cuts out regions. That detector is frozen during training and slow to run, so it limits the whole system.

This paper asks what happens if you drop the detector and use a plain image transformer, the same kind of model used for text, that splits the picture into patches. It then tests many design choices one at a time: which pre-trained image and text models to start from, how the two streams should share information, whether to use an encoder or an encoder-decoder, and which extra training tasks help. The resulting model, METER, reaches strong question-answering scores while using far fewer training images than the best rivals.

The paper also reports what did not help. Adding a task that asks the model to rebuild hidden image patches lowered scores in its tests, which contradicts results from detector-based models.

TL;DR

• Fully transformer vision-language models can match detector-based ones if the design is chosen carefully. METER is the model that results from the study §1.
• The best setup pairs a CLIP or Swin image encoder with a RoBERTa text encoder, fused by co-attention layers. It is pre-trained with masked language modeling and image-text matching only §3.3 §4.4.
• Author-reported VQAv2 test-std accuracy is 77.64 with about 4M pre-training images, and 80.54 for a 14M-image, larger-backbone version tab. 8 tab. 10.
• Pre-training the image patches to be reconstructed did not help in the authors' setup tab. 7.

Problem

Region-feature vision-language models use a frozen, pre-trained object detector to extract features. The authors note two problems: the detector is not perfect and is usually kept frozen during pre-training, which caps the model's capacity, and extracting region features is time-consuming §1 §2. Earlier fully transformer models (ViLT, ALBEF, visual parsing) reportedly lagged behind region-feature models on VQA §2. The paper asks whether a ViT-based model trained end to end can close that gap §1.

The study covers only 4 standard image-caption datasets and English text. The authors list scalability and non-English data as open limitations app. I.

Method

METER has four parts §3:

1. Text encoder §3.1. The sentence is split into subwords and passed through a pre-trained text transformer (BERT, RoBERTa, ELECTRA, ALBERT or DeBERTa were tried). The default is RoBERTa §4.1.1. A linear layer maps its output to the fusion width of 768 meter_module.py.
2. Vision encoder §3.1. The image is cut into patches and encoded by a pre-trained vision transformer. The paper compares the original ViT, DeiT, Distilled DeiT, CaiT, VOLO, BEiT, Swin and CLIP-ViT §3.1. The default is CLIP-ViT-224/32, and the best reported models use CLIP-ViT-224/16 or Swin §4.1.1 §4.5. A linear layer maps its output to 768 meter_module.py.
3. Multimodal fusion §3.1 §3.3. The default fusion is co-attention. Each stream has its own stack of M = 6 layers. Each layer contains one self-attention block, one cross-attention block that attends to the other stream, and one feed-forward block bert_model.py. Both streams are updated from the states of the previous layer, so the two updates do not see each other's new values meter_module.py. The merged-attention alternative (one shared transformer over the concatenated tokens) was tested and scored lower §4.2 tab. 6.
4. Heads §3.3. The first token of each stream goes through a pooler (dense plus tanh). The two pooled vectors are concatenated meter_module.py meter_module.py. Task heads sit on top: an MLP for VQA over 3129 answers, a two-way head for image-text matching, and a head for the NLVR2 and SNLI-VE tasks meter_module.py.

The default text and vision encoders are not decoders. The paper tests an encoder-decoder variant with a 3-layer encoder and 3-layer decoder, and finds it weaker on the discriminative tasks §4.3 tab. 6.

Pre-training objectives

The default objectives are masked language modeling (MLM) and image-text matching (ITM) §3.3.

• MLM §3.2. Mask 15% of caption tokens and predict the original tokens from the remaining text and the image.
• ITM §3.2 §3.3. Each sample is a matched or mismatched image-caption pair with equal probability. A binary head decides whether they match.

The paper also tests masked image modeling (MIM) in two forms. Both mask 15% of the image patches §3.2:

• In-batch negatives. The model must pick the original patch from the set of all patches in the batch. The probability is given by equation 1 §3.2 eq. 1.
• Discrete codes. Each patch is replaced by a token from a VQ-VAE codebook, borrowed from DALL-E. The model predicts that token §3.2.

Neither variant is in the released code at this commit (see Paper vs code).

p(\mathbf{v}_i^k \mid [\mathbf{v}^{k,mask}; \mathbf{l}^{k}]) = \frac{e^{h(\mathbf{v}_i^k)^{\mathrm{T}} c(\mathbf{v}_i^k)}}{\sum_{j, k'} e^{h(\mathbf{v}_i^k)^{\mathrm{T}} c(\mathbf{v}_j^{k'})}} eq. 1

The loss at masked positions is the negative log of this probability §3.2. The paper says the model is trained with a loss similar to noise-contrastive estimation §3.2.

# Pre-training loss: MLM + ITM, weight 1 each in the code  [§3.2] [§3.3]
# Input:  batch of matched image-caption pairs
# Output: scalar loss

def pretrain_loss(images, captions):
    # MLM pass: mask 15% of caption tokens, keep the image  [§3.2]
    masked_captions, targets = mask_tokens(captions, rate=0.15)
    text_feats, _, _ = meter_forward(masked_captions, images)
    # Predict the original token at each masked position  [§3.2]
    mlm_logits = mlm_head(text_feats)                  # (B, N, vocab)
    mlm_loss = cross_entropy(mlm_logits, targets)      # masked positions
    # ITM pass: keep half the images, swap in a random one for the rest
    # [§3.2] [code: meter/modules/objectives.py#L44-L50]
    labels = random_half_ones(batch_size)              # 1 = match
    mixed_images = swap_mismatched(images, labels)
    _, _, cls_feats = meter_forward(captions, mixed_images)
    itm_logits = itm_head(cls_feats)                   # (B, 2)
    itm_loss = cross_entropy(itm_logits, labels)
    return mlm_loss + itm_loss

# Masked patch classification with in-batch negatives  [§3.2] [eq. 1]
# Follows the paper only: this MIM variant is not in the official code.

def mim_in_batch_loss(patches, masked_outputs, masked_index):
    # Original patch representations c(v) for every patch in the batch
    # [§3.2]
    candidates = original_patch_embedding(patches)     # (B * P, d)
    # Model outputs h(v) at the masked positions  [§3.2]
    queries = masked_outputs                           # (M, d)
    # Dot-product scores against all batch patches, then softmax  [eq. 1]
    scores = queries @ candidates.T                    # (M, B * P)
    # The target is the index of the original patch in the batch  [eq. 1]
    return cross_entropy(scores, masked_index)

Pre-training data

Four image-caption datasets: COCO, Visual Genome, Conceptual Captions and SBU Captions. Together they make about 4M images and about 9M image-caption pairs §3.3 app. A tab. 11. VQAv2 training uses its training and validation data, and its test-dev split is used as the development set. Visual Genome's VQA annotations are not used app. A.

# LR per step  [§3.3] [code: meter/modules/meter_utils.py#L280-L295]
# base_lr = 1e-5 for the pre-trained encoders; cross-modal layers use 5x

def learning_rate(step, total_steps, base_lr):
    # Linear warm-up over the first 10% of steps  [§3.3]
    warmup_steps = int(0.1 * total_steps)
    if step < warmup_steps:
        return base_lr * step / warmup_steps
    # Linear decay to 0 at the last step  [§3.3] [code: meter/config.py#L64-L68]
    remaining = (total_steps - step) / (total_steps - warmup_steps)
    return base_lr * remaining

Downstream fine-tuning

Fine-tuning uses a learning-rate grid search over {1e-6, 2e-6, 5e-6, 1e-5} and image resolutions {288, 384, 576} app. A. RandAugment is applied during fine-tuning app. A. Downstream tasks are VQAv2, NLVR2, SNLI-VE, and Flickr30k and COCO image-text retrieval. Retrieval is tested both zero-shot (Flickr30k) and after fine-tuning §3.3.

Hyperparameters

Fine-tuning settings in the released named configs config.py:

The paper gives the grid of learning rates and image resolutions app. A but does not say which grid point was chosen for each task. The code's named configs use one value per task, all within that grid. The head multipliers and the 15 negatives per image for COCO and Flickr retrieval config.py are not in the paper.

Inference cost: about 53 ms per VQA instance (CLIP-ViT) and 59 ms (Swin) on one V100 GPU app. B tab. 12.

Results

Author-reported. Setup: models pre-trained on image-caption data, fine-tuned or zero-shot on the listed tasks. Ablations use VQAv2 test-dev and Flickr30k zero-shot retrieval §4.1 §4.2 §4.3 §4.4. Headline comparison is in tab. 8 and tab. 9.

Ablation: text encoder without VLP tab. 2

Vision encoder CLIP-ViT-224/32, no vision-language pre-training. Results are on test-dev or validation sets. SQuAD and MNLI are copied from their papers tab. 2 caption.

All other rows are base-size models tab. 2.

Ablation: vision encoder without VLP tab. 3

RoBERTa text encoder tab. 3.

ImageNet numbers are copied from their papers tab. 3.

Ablation: encoders with VLP tab. 4

Results on VQAv2 test-dev and Flickr30k zero-shot retrieval tab. 4.

The text says CLIP-ViT-224/16 reaches 77.19 and 77.20 on test-dev and test-std §4.1.2. The table lists only test-dev tab. 4.

Ablation: learning rates tab. 5

VQAv2 test-dev and Flickr30k zero-shot retrieval tab. 5.

Image resolution during fine-tuning matters: raising it from 224 to 576 improves VQA by about 3 points for CLIP-ViT-224/32 and about 1 point for CLIP-ViT-224/16 §4.1.2 fig. 5.

Ablation: multimodal fusion tab. 6

VQAv2 test-dev and Flickr30k zero-shot retrieval tab. 6.

Ablation: pre-training objectives tab. 7

VQAv2 test-dev and Flickr30k zero-shot retrieval tab. 7.

Comparison with prior models tab. 8 tab. 9

Table 8 covers VQAv2, NLVR2, SNLI-VE and zero-shot Flickr30k retrieval. Rows are grouped by pre-training size. A "-" means the source table has no value.

Table 9 covers fine-tuned Flickr30k and COCO retrieval. Columns are IR@1/5/10 and TR@1/5/10.

The paper's headline: METER-CLIP-ViT BASE reaches 77.64 on VQAv2 test-std, 1.04 above VinVL (76.60) and 1.6 above ALBEF (14M) (76.04), using 4M images §4.5 tab. 8.

Scaling tab. 10

Pre-training on about 14M images (COCO, CC, CC12M, SBU, VG; about 20M pairs) with a CoSwin-Huge vision backbone and RoBERTa-base §4.5. The VQAv2 results:

Analysis tab. 11 to tab. 17

Pre-training dataset statistics tab. 11:

Inference time, ms per VQA instance on one V100 app. B tab. 12. The "Time (ViLT)" column is copied from the ViLT paper and used different hardware app. B:

Image captioning, COCO, maximum-likelihood training app. C tab. 13:

For generation, the generated tokens attend to each other without the causal mask, and decoding uses beam size 5 app. C.

Multi-scale feature fusion (gated sum over encoder layers, see equation 2) app. D tab. 14:

o(x_i) = h(x_i^N) + \sum_{j=0}^{N-1} g(h(x_i^j))\, h(x_i^j) eq. 2, app. D

Fusion helps without pre-training and hurts after it app. D tab. 14. The Flickr columns are reported only with VLP.

Unimodal results after VLP app. F.

Linear probe on CIFAR, image encoder before and after VLP (text encoder RoBERTa-base) app. F tab. 16:

The authors report that VLP helps Swin's image features and not CLIP's app. F.

Pre-training dataset ablation, CLIP-ViT-224/32 and BERT-base-uncased, 50k steps per dataset app. G tab. 17:

The authors credit data size and domain similarity: CC and VG are the largest, and COCO matches the downstream domain most closely app. G.

Correlation between vision-language and single-modality performance app. E fig. 6: Pearson r = -0.09 (p = 0.88) against language tasks and r = 0.41 (p = 0.36) against vision tasks, which the authors read as little or no correlation app. E.

Paper vs code

Official code: github.com/zdou0830/METER, linked from the abstract. Cloned at commit f4f09345b26ee21add0a756d06598e3c04726345. License MIT, from the repository LICENSE file. The README also names ViLT (Apache-2.0) as the source of borrowed parts README.md.

Checked against the paper: AdamW optimizer, 100k steps, batch 4096, 1e-5 base and 5x cross-modal LR, 10% warm-up with linear decay to 0, hidden 768, 12 heads, 6 cross-modal layers, 15% MLM, 50/50 ITM §3.3 [config values as cited above]. No discrepancy found in these.

Discrepancies:

• Pre-training image size. The paper says images are resized to 224 or 384 depending on the ViT §3.3. The pre-training named config uses 224 config.py. The README's pre-training example uses 288 for CLIP-ViT/16 with RoBERTa README.md, and the released pre-trained checkpoint is also listed at 288 README.md. The comparison section also gives 288 for the best CLIP model §4.5.
• ITM head input. The paper describes a classifier on a [CLS]-style token of the joint sequence §3.2. The code classifies the concatenation of the pooled text and pooled image vectors meter_module.py with an input width of 2 x 768 meter_module.py.
• Head learning rate. The paper specifies 1e-5 and 5e-5 for encoders and the cross-modal module §3.3. The code also multiplies MLM and ITM heads by 5 config.py meter_utils.py. The paper does not mention this.
• RandAugment. The paper says RandAugment is applied during fine-tuning app. A. The default transform in the code has no RandAugment transform.py. The augmented version is selected by the clip_randaug named config config.py, and the README's fine-tuning commands use it README.md.
• Whole-word masking. The paper does not mention it. The pre-training named config enables it config.py. A code comment says whole-word masking does not work with RoBERTa config.py, yet the README pairs this named config with RoBERTa README.md.
• Masked image modeling. Neither MIM variant (in-batch negatives, discrete codes) is in this repository at this commit. A search of the objectives module finds no MIM code. The results in tab. 7 cannot be reproduced from this release.

The authors read the MIM result as a sign that conclusions from region-based models may not carry over to patch-based models §4.4.

• Merged attention and encoder-decoder. The released model runs only the co-attention path meter_module.py. The config has no option for merged attention, a decoder, or the embedding-only text encoder config.py.
• 14M scaling run. No CoSwin-Huge or 14M-image config exists in this repository. The released configs cover Swin base and large, and CLIP ViT-B/16 and /32 config.py.
• VQA loss. The paper does not state the VQA loss. The code uses binary cross-entropy with logits, scaled by the number of answers (3129) objectives.py.
• Retrieval loss. The paper does not state it. The code uses a softmax over one positive and 15 negative captions per image, with the positive at index 0 objectives.py config.py.
• Position embeddings at new resolutions. The compiled paper does not say how positional embeddings change with image size. The code interpolates them when a checkpoint is loaded at a new resolution meter_module.py meter_module.py.

Pitfalls

• Table and figure numbers are assigned by source order. The compiled PDF was not available for checking, so the numbers used here (for example tab. 8 for the main comparison) follow the LaTeX order of the source. Check them against the PDF before citing externally.
• Appendix letters are also assigned by source order. Implementation Details is app. A and Limitations is app. I. These are inferred from the \appendix sections.
• Test-dev and test-std differ. Ablations use test-dev; the headline 77.64 uses test-std. Table 8 gives the METER-CLIP test-dev value as 77.68. Compare numbers only within one table and one split.
• Equation numbering. Equation 1 is the in-batch MIM probability §3.2. Equation 2 is in appendix D. Its number is inferred from source order.
• Pre-training data totals. Table 11 sums to about 4.2M images and about 9.9M captions. The text says about 4M images and about 9M pairs app. A.
• Two forward passes per step. The code runs one forward pass for MLM (masked text, matched image) and a separate one for ITM (unmasked text, mismatched image) objectives.py objectives.py. This doubles the encoder cost per step relative to a single pass. The paper does not say whether it does this.
• ITM pairing. The code uses an exact half-and-half split within each batch, drawn with a random permutation objectives.py. The paper says each sample is matched or mismatched with equal probability §3.2, a per-sample statement.
• Layer-wise learning rates. Implement the parameter groups explicitly. The paper gives 1e-5 and 5e-5. The code also applies 5x to heads and uses separate weight-decay groups for biases and LayerNorm meter_utils.py.
• Simultaneous stream updates. Both streams in each cross-modal layer read the previous layer's states. A sequential implementation (text first, then image on the new text states) would differ meter_module.py.
• Resolution changes. Fine-tuning at 288, 384 or 576 requires interpolating the CLIP positional embeddings, which the paper does not describe in its main text. The code does this when loading a checkpoint meter_module.py.
• Text inconsistency. The text gives CLIP-ViT-224/16's VQA as 77.19 (test-dev) and 77.20 (test-std) §4.1.2. Table 4 lists only test-dev tab. 4. Table 8's 77.64 test-std is the headline model, a different comparison, so do not mix the two.
• Scaling run. The 14M-image result tab. 10 has no released config (see Paper vs code). The paper does not state its reproduction cost.
• Compute. 8 A100 GPUs for 3 to 8 days per model app. A. The paper gives no per-run cost for the ablations.
• Retrieval negatives. The code uses 15 negatives per image config.py. The paper does not state the number.

Lineage

• Builds on the region-feature vision-language line the paper reviews (for example LXMERT, ViLBERT, UNITER, VinVL) and on fully transformer models ViLT, ALBEF and visual parsing §2 tab. 1. None of these has a card in this wiki yet.
• Uses pre-trained encoders CLIP-ViT, Swin, RoBERTa, BERT and DeBERTa, all cited in §3.1.
• Compared against: VinVL and ALBEF on VQAv2 and retrieval tab. 8 tab. 9; SimVLM on VQAv2 and captioning tab. 10 tab. 13.
• The MIM variants follow BEiT's use of discrete codes §3.2, but the paper reports that they hurt in its setting tab. 7.