MMBERT: Multimodal BERT Pretraining for Improved Medical VQA

In plain words

Answering a question about a medical scan, such as which organ is shown or whether something is abnormal, needs a model that understands both the picture and the words. Labelled medical question-answer data is small and expensive to make, so the model cannot learn much from it alone.

The authors first teach a small text model to fill in blanked-out medical words in the captions of radiology images. While it does this, it also gets a few numerical summaries of the picture as extra input. The idea is that this pretraining gives the model a better sense of medical images and their words. Then the same model is trained on the question-answer data. It also produces attention maps that show which image regions it used, which helps a reader judge its answers.

TL;DR

• Pretrain a 4-layer, BERT-style transformer with masked language modelling on ROCO radiology image-caption pairs, using image features as extra tokens §2.2.
• Fine-tune on VQA-Med 2019 (one model per question category, or one for all) and on VQA-RAD (one model for all questions) §3.2.
• Author-reported accuracy beats previous results on both datasets, including an ensemble on VQA-Med 2019 tab. 1 tab. 2.
• The official code is MIT-licensed. The VQA-RAD script does not reach training as shipped (see Pitfalls), and several loss, pooling and evaluation details differ from the paper. I did not run the code.

Problem

Models trained on general-domain VQA do not transfer well to radiology images, and annotating medical images needs experts, so labelled VQA sets are small §1. The authors argue that a model pretrained on image-caption pairs learns richer medical image-text features than one trained only on VQA data, and that earlier medical VQA methods do not use such large multimodal medical datasets for pretraining §1.

Method

Overview. The model is a transformer encoder over one token sequence that holds a question or caption together with five image vectors §2.2. Pretraining uses masked language modelling (MLM) on ROCO captions, where the masked words are medical keywords §2.2. Fine-tuning reuses the pretrained weights and replaces the output layer with a classifier over answers §2.3.

Self-attention. Each head computes scaled dot-product attention. The query, key and value matrices are Q, K, V, and the scale is \sqrt{d_k}, the query and key dimension §2.1. This display equation is unnumbered in the paper:

\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{QK^{T}}{\sqrt{d_{k}}}\right)V §2.1

Several heads run in parallel and their outputs are concatenated §2.1.

Image features. Each image is resized to 224 x 224 §3.2 and passed through a ResNet-152 §2.2 roco_utils.py. Five feature maps are taken from different depths: the output of layers 4, 3, 2 and 1, and the output of the stem before max-pooling roco_utils.py. Each map goes through a 1 x 1 convolution to the hidden size, a ReLU, and global average pooling, which gives one vector per map roco_utils.py. The five vectors are the image tokens §2.2 fig. 2.

Text input. Captions are split into BERT WordPiece tokens §2.2. The sequence is [CLS], five image positions, [SEP], the caption tokens and a final [SEP] fig. 2 roco_utils.py. Segment id 0 marks [CLS], the image positions and the first [SEP]; segment id 1 marks the caption fig. 2 roco_utils.py.

Masking. Only medical keywords from a provided list are candidates for masking; common words are left alone §2.2. Each WordPiece of a keyword is masked with a fixed probability and replaced by [MASK]. The original WordPiece is the label roco_utils.py. The probability is a required argument with no default in the code roco_train.py, and the paper does not state its value.

Encoder. The five image vectors overwrite the embedding output at positions 1 to 5. The code initialises the text embedding layers from BERT-base; the paper does not say so roco_utils.py roco_utils.py. The encoder has four transformer layers with 12 heads each, and the paper describes it as "only 4 BERT layers and a total of 12 attention heads" §2.2 roco_train.py.

Output head. The last hidden states go through a dense layer, a tanh, and a vocabulary head. The vocabulary head is a linear layer, a LayerNorm and a second linear layer to the vocabulary size roco_utils.py.

Pretraining objective. Masked positions are predicted from the image and the visible caption words §2.2 fig. 2. The loss is a negative log-likelihood over log-softmax outputs roco_utils.py roco_train.py.

Fine-tuning. The pretrained weights are loaded, and the final linear layer of the vocabulary head is replaced by one with as many outputs as there are answers §2.3 train.py. The paper says the [CLS] token is not used. It averages the last layer's token representations, then passes the result through dense layers §2.3. The VQA-Med code does this utils.py. The VQA-RAD code averages over all four layers instead utils_vqarad.py.

Three VQA-Med settings §3.2:

• MMBERT General: one model for all question categories, pretrained §3.2.
• MMBERT Exclusive: one fine-tuned model per question category, pretrained §3.2.
• MMBERT NP: one model per category, no pretraining §3.2.

At inference, the per-category settings first predict the question category with a BERT-base classifier, and then use that category's model §3.2. The repository does not contain this category classifier README.md.

# Build one masked-caption input  [§2.2]
# [code: pretrain/roco_utils.py#L132-L154]
# Input:  caption (str), keywords (set of str), mask_prob (float)
# Output: tokens, segment_ids, mask, labels; each of length max_len

def build_input(caption, keywords, mask_prob, max_len):
    # Mask each WordPiece of a keyword with probability mask_prob  [§2.2]
    # [code: pretrain/roco_utils.py#L105-L130]
    pieces, labels = mask_keywords(caption, keywords, mask_prob)
    # Keep at most max_len - 8 WordPieces  [code: pretrain/roco_utils.py#L139]
    pieces = pieces[: max_len - 8]
    labels = labels[: max_len - 8]
    # Layout: [CLS], 5 image slots, [SEP], caption, [SEP]  [fig. 2]
    tokens = [CLS] + [0] * 5 + [SEP] + to_ids(pieces) + [SEP]
    # Labels are 0 except at masked caption positions
    # [code: pretrain/roco_utils.py#L143]
    labels = [0] * 7 + labels + [0]
    # Segment 0 for [CLS], image slots and first [SEP]; 1 for caption  [fig. 2]
    segment_ids = [0] * 7 + [1] * (len(pieces) + 1)
    # 1 for real tokens, 0 for padding  [code: pretrain/roco_utils.py#L146-L151]
    mask = [1] * len(tokens) + [0] * (max_len - len(tokens))
    tokens = pad_with_zero(tokens, max_len)
    segment_ids = pad_with_zero(segment_ids, max_len)
    labels = pad_with_zero(labels, max_len)
    return tokens, segment_ids, mask, labels

# Five image vectors from one ResNet-152  [§2.2]
# [code: pretrain/roco_utils.py#L489-L524]
# Input:  image (B, 3, 224, 224)
# Output: five vectors, each (B, d)

def image_slots(image):
    # Stage outputs, deepest first: layer 4, 3, 2, 1, then the stem  [§2.2]
    stage_maps = resnet152_stages(image)
    slots = []
    for fmap in stage_maps:
        # 1x1 conv to the hidden size, ReLU, then global average pool
        # [code: pretrain/roco_utils.py#L508-L524]
        slots.append(global_avg_pool(relu(conv1x1(fmap))))
    return slots

# Encoder and MLM head  [§2.2] [code: pretrain/roco_utils.py#L616-L657]
# Input:  image (B, 3, 224, 224); tokens, segment_ids, mask (B, max_len)
# Output: logits (B, max_len, vocab_size)

def mlm_logits(image, tokens, segment_ids, mask):
    slots = image_slots(image)
    # Embeddings from the BERT-base embedding layers
    # [code: pretrain/roco_utils.py#L629]
    hidden = bert_embeddings(tokens, segment_ids)
    # Overwrite positions 1..5 with the five image vectors
    # [code: pretrain/roco_utils.py#L630-L639]
    for k in range(5):
        hidden[:, k + 1] = slots[k]
    # Four pre-norm blocks, 12 heads each  [§2.2]
    # [code: pretrain/roco_utils.py#L640-L641]
    for block in blocks:
        hidden = block(hidden, mask)
    # Dense + tanh, then vocabulary head
    # [code: pretrain/roco_utils.py#L653-L656]
    return vocab_head(tanh(fc1(hidden)))

# One pre-norm block as coded  [code: pretrain/roco_utils.py#L591-L614]
# Input:  hidden (B, max_len, d), mask (B, max_len)

def block(hidden, mask):
    # Self-attention on the normalised input, then a residual add  [§2.1]
    update = proj(attention(norm(hidden), mask))
    hidden = hidden + dropout(update)
    # Feed-forward (width 4d, GELU) on the same norm module
    # [code: pretrain/roco_utils.py#L599-L601]
    update = feed_forward(norm(hidden))
    return hidden + dropout(update)

# Pretraining loss as coded  [code: pretrain/roco_train.py#L67]
#   [code: pretrain/roco_utils.py#L181-L187]
# Input:  logits (B, max_len, V); labels (B, max_len) integer ids

def mlm_loss(logits, labels):
    log_probs = log_softmax(logits, dim=-1)
    # Averaged over every position, not only masked ones; label 0
    # ([PAD]) is an ordinary class here
    return nll_loss(log_probs, labels)

# Fine-tuning head  [§2.3] [code: vqamed2019/utils.py#L598-L602]
# Input:  hidden (num_layers, B, max_len, d)

def answer_logits(hidden):
    # Mean over all positions, padding included, of the last layer
    pooled = tanh(fc1(mean_over_positions(hidden[-1])))
    # answer_head: Linear, LayerNorm, Linear(d -> num_answers)
    return answer_head(pooled)

Hyperparameters

Values are from the paper where it states them, and from the code otherwise. "not stated" means neither source gives a value. Code defaults may be overridden on the command line; the repository has no logs that show the values used for the reported numbers.

Pretraining trains on ROCO images and captions §3.1 roco_utils.py. The paper says "all the images" and uses the keywords for masking §3.1.

Results

Author-reported. The paper evaluates on the VQA-Med 2019 test set and on VQA-RAD §3.1 §3.3. Accuracy is exact-match on the answer. BLEU is described in the paper as unigram-based in one place and 4-gram in another; the code uses unigram BLEU §3.3 utils.py.

VQA-Med 2019 tab. 1. Columns: % is accuracy and B is BLEU, per category, with overall (All). Dedicated models are per category. "-" means no value in the paper.

Source: tab. 1. Bold is as in the paper.

VQA-RAD tab. 2. Accuracy only.

Source: tab. 2.

Qualitative. The paper shows attention maps for example questions, with ground truth and the model's answer fig. 1 fig. 3. It notes that the Modality example attends to the wrong cue, and that an Abnormality answer is more detailed than the reference answer §3.4.

Paper vs code

The code is the authors' own repository (VirajBagal, an author of the paper) at 9471148, MIT-licensed LICENSE. Discrepancies, each with both citations:

• Pretraining loss covers every position. The paper says the model predicts the masked tokens §2.2. The code takes the NLL over all positions, and unmasked positions carry label 0, which is [PAD] in BERT's vocabulary roco_utils.py roco_utils.py roco_train.py. Accuracy is computed only on masked positions roco_utils.py.
• Pretraining uses only ROCO train and validation. The paper says "all the images" §3.1. The code loads the train and validation CSVs; the test split is not loaded roco_utils.py.
• Heads: 12 per layer, not 12 in total. The paper says "a total of 12 attention heads" §2.2. The code builds 12 heads in each of the 4 layers roco_train.py roco_utils.py.
• VQA-RAD pooling. The paper pools the last layer §2.3. VQA-Med does utils.py; the VQA-RAD code averages all four layers and all positions utils_vqarad.py.
• VQA-Med val and test images are not resized to 224. The paper says all images are resized to 224 x 224 §3.2. VQA-Med validation and test skip any resize train.py eval.py. Only the training crop outputs 224. VQA-RAD resizes at test time train_vqarad.py.
• Scheduler metric for VQA-RAD. The paper says the learning rate drops when validation loss stalls §3.2. The VQA-RAD script steps the scheduler on training loss train_vqarad.py, and the script has no validation split train_vqarad.py.
• Early stopping. The paper does not mention early stopping in the text. The VQA-Med script stops after more than 20 epochs without a gain in validation accuracy train.py. The LaTeX source contains this sentence only as a comment §3.2.
• Checkpoint selection. The VQA-Med script keeps the checkpoint with the best validation accuracy train.py. The VQA-RAD script keeps the checkpoint with the best test accuracy train_vqarad.py. The paper states no protocol §3.2.
• BLEU. The paper gives two definitions: 4-gram §3.3, then unigram §3.3. The code uses unigram BLEU utils.py utils_vqarad.py.
• Image normalisation. The paper does not state it. The code uses mean 0.5 and std 0.5, not the ImageNet statistics that the ResNet was trained with roco_train.py train.py.
• Hidden size defaults. Pretraining defaults to 768 roco_train.py, while the VQA-Med training and eval scripts default to 312 train.py eval.py. The fine-tuning checkpoint must match the pretraining size.
• Text numbers do not match the table. The paper says pretraining raises Accuracy and BLEU by 7.2 and 9 points §3.3. The table's Exclusive vs NP rows differ by 7.0 (Accuracy) and 6.3 (BLEU) tab. 1. The text also calls a 67.2 vs 64.0 gain "5%" §1; that is a relative gain of about 5%, or 3.2 points.
• README copy error. The README lists CGMVQA's overall score as 62.4/64.4 [README.md], which is the VGG16 + BERT value. The paper's value is 60.0/61.9 tab. 1. This card uses the paper's values.
• Attention maps. The code returns per-layer attention scores utils.py utils.py, but the repository has no script that turns them into the paper's region maps README.md.

Searched: the abstract links github.com/VirajBagal/MMBERT, which is the repository used here. The repository was cloned and checked: pretraining, VQA-Med training, VQA-RAD and eval code were read; vqamed2019/utils.py was read in its model, data and loss sections and grepped elsewhere; LICENSE and README were read.

Pitfalls

• The VQA-RAD script does not run as shipped. vqarad/train_vqarad.py uses val_df and val_tfm at line 131, but neither is defined in the script, and the load_data call returns only two frames train_vqarad.py train_vqarad.py. Fix this before running it, and also decide on a validation split, since the paper does not describe one for VQA-RAD.
• Two dependencies are missing from the repo. Pretraining needs vocab/med_vocab.pkl under the data directory roco_utils.py and the ROCO and VQA-Med data, none of which are in the repository. The default paths point to the authors' machines roco_train.py.
• Layer norm is reused. In each pre-norm block the code applies norm1 twice, to the attention input and to the feed-forward input. norm2 is created but used only in the post-norm branch roco_utils.py roco_utils.py. Reproduce this exactly if you want the paper's numbers.
• Image slots are not normalised. The five image vectors replace the embedding output after the embedding LayerNorm, so they reach the first block without normalisation roco_utils.py. Each vector is non-negative because of the ReLU before pooling roco_utils.py.
• Padding is averaged in. Pooling for fine-tuning averages over all max_len positions, including padding. The attention mask only affects attention, not the pool utils.py utils_vqarad.py.
• Channel order. VQA-Med images are read with OpenCV, which gives BGR utils.py. The code does not swap the channels before they reach the RGB-trained ResNet. Pretraining uses PIL with RGB conversion roco_utils.py. This is an inconsistency to be aware of; the paper does not discuss it.
• Captions and questions are truncated. Pretraining keeps at most 67 WordPieces with the default length of 75 roco_utils.py. Fine-tuning keeps at most 20 WordPieces with the default length of 28 utils.py. The paper does not state either limit.
• Masking is always [MASK]. Masked positions are always replaced by [MASK], with no random-token or unchanged option roco_utils.py. The paper does not say so.
• Per-category runs drop Yes/No answers. With --category, the VQA-Med script removes answers yes and no from every split train.py. The repository does not show how the Yes/No category is named in the data.
• The answer vocabulary includes the test split. The answer-to-index map is built from train, validation and test answers together train.py. The class head therefore has test-only answers as outputs. This is not label leakage, but it affects the head size.
• Checkpoint choice. The VQA-RAD test-set selection train_vqarad.py means its reported numbers can be optimistic if the best checkpoint is taken from the test curve. The VQA-Med selection on validation accuracy avoids this train.py.
• Hidden-size mismatch. A pretrained 768-dimensional checkpoint will not load into a 312-dimensional model, and the VQA-Med training script's default is 312 train.py. Pass --hidden_size 768 when you fine-tune a 768-dimensional pretraining run.
• Costs. Pretraining runs a ResNet-152 and a transformer on every batch, and the ResNet is trainable end-to-end; the freezing lines are commented out roco_utils.py. The paper does not give run times.

Lineage

• Builds on: BERT [bert] for the WordPiece tokenizer and transformer encoder; ViLBERT and UNITER for pretraining with image-text pairs §1 [chen2020uniter] [vilbert]. None of these are carded in wikillm, and their arXiv identifiers are not in the paper's bibliography.
• Compared against: VGG16 + BERT, the VQA-Med 2019 challenge winner [winner]; CGMVQA and its ensemble [cgmvqa]; MEVF + SAN and MEVF + BAN [nguyen2019overcoming]; conditional reasoning [zhan].
• Superseded by: not known from content already in this wiki.