Multi-Modal Masked Autoencoders for Medical Vision-and-Language Pre-Training

In plain words

Medical data comes as pictures (X-ray, CT, MRI) paired with text such as captions and reports. Labelled examples of such pairs are scarce, so this paper pretrains one model on large sets of unlabelled pairs. The trick is to hide parts of both inputs and ask the model to fill them back in. Most of the image patches are hidden and their pixels are reconstructed; some words are hidden and predicted; and the model also learns whether an image and a caption belong together.

The image is hidden much more heavily than the text, because neighbouring image patches reveal most of each other while each word carries a lot of meaning. The image reconstruction is read from a middle layer of the model, not from the last one.

After pretraining, the model is fine-tuned on medical visual question answering, on classifying image-text pairs, and on retrieving captions for images. The authors report state-of-the-art results on these tasks. The official code is public.

TL;DR

• Self-supervised pretraining of a two-stream Transformer (ViT for images, RoBERTa for text, plus cross-attention fusion) on ROCO and MedICaT image-text pairs §3.1.
• Three design choices make masked autoencoding work across modalities: different masking ratios (75% image, 15% text), reconstruction from an intermediate fusion layer, and different decoders (a Transformer for pixels, an MLP-style head for words) §2.2.
• Reported gains are on three Med-VQA sets, MELINDA classification and ROCO retrieval; these are author-reported §3.3.
• Pitfalls: several key settings (patch size, batch size, decoder sizes, fine-tuning hyperparameters) are only in the code; the code has three forward passes per step and differs from the paper in several places (see Paper vs code).

Problem

Medical vision-and-language pretraining (Med-VLP) would give transferable image-text representations, but few studies had covered it. Earlier medical VLP work used CNN visual backbones, and purely Transformer-based models were not explored; one related study used diagnosis labels rather than self-supervision and covered chest X-ray only §1. Masked image modelling was reported to hurt general VLP pretraining, so the authors argue that the objective and architecture must be designed for the two modalities §1 §2.2.

Method

Backbone. Three parts: a vision encoder, a language encoder and a multi-modal fusion module §2.1.

• Vision encoder. A ViT takes an image cut into P×P patches, linearly projects each patch, prepends a learnable token and adds learnable 1-D position embeddings §2.1 eq. 2. The paper's ViT is CLIP ViT-B, with the patch size not stated §3.1. The encoder output for the image is a sequence of contextual vectors, one per patch plus the prepended token §2.1.
• Language encoder. BERT-style WordPiece tokens, with a start token and a boundary token added, and position embeddings §2.1 eq. 3. The paper uses RoBERTa-base §3.1.
• Fusion. Two streams of N_m Transformer layers. Each layer has a self-attention sub-layer within its modality, a cross-attention sub-layer over the other modality, and a feed-forward sub-layer §2.1 eq. 4 eq. 5 eq. 6. Attention is written as softmax(QKᵀ)·V, with no scaling in the equation §2.1 eq. 4. The paper uses 6 fusion layers, hidden size 768 and 12 heads §3.1.

Objective. Pretraining minimises the sum of S pretext losses. Each loss has its own decoder on top of the shared backbone §2 eq. 1. Equation 1 is unweighted.

Masking.

• Images: random sampling of patches at a 75% masking ratio §2.2.
• Text: random tokens at a 15% masking ratio §2.2.

Representation selection. The masked-image loss is computed from the visual outputs of the k-th fusion layer, Z^{vk}. The masked-language loss uses the final text output Z^l §2.2. The chosen layer is k = 3, with layer 0 (no textual information) the worst §3.4.

Decoders.

• Vision: a Transformer maps Z^{vk} to pixel space. The loss is the mean squared error between reconstructed and original pixels, on the masked patches §2.2.
• Language: an MLP head; the loss is the negative log-likelihood of the masked tokens §2.2.
• The MLM and MIM losses run in separate forward passes §2.2, footnote 2.

Image-text matching (ITM). The paper uses the image-text matching task from UNITER by default, added to the two masked objectives §3.1. It is not described in the paper beyond that.

# M3AE pre-training step: three separate forward passes  [§2] [eq. 1]
# Input:  image (3, 288, 288); caption token ids (M,)
# Output: scalar loss, the unweighted sum of three terms  [eq. 1]

def pretraining_step(image, caption_ids):
    # 1. MLM pass: mask 15% of text tokens, keep the whole image  [§2.2]
    masked_ids, mlm_labels = mask_text(caption_ids, ratio=0.15)
    text_out, _ = fuse(encode_image(image), masked_ids)
    mlm_loss = cross_entropy(mlm_head(text_out), mlm_labels)  # masked only

    # 2. MIM pass: drop 75% of patches, keep the text unmasked  [§2.2] [§3.1]
    visible = keep_random_patches(image, ratio=0.25)
    _, image_taps = fuse(visible, caption_ids, tap_layer=3)
    # image_taps[3]: image features after the 3rd fusion layer  [§3.4]
    pred = mim_decoder(image_taps[3])  # (N, P, 3 * P * P)
    target = normalise_per_patch(patchify(image))  # [code: objectives.py#L51-L54]
    mim_loss = mse(pred, target)  # averaged over removed patches only

    # 3. ITM pass: half the batch uses a random other image  [§3.1]
    is_match, paired_image = swap_random_images(image, caption_ids)
    cls_out = fuse_cls(paired_image, caption_ids)
    itm_loss = cross_entropy(itm_head(cls_out), is_match)

    return mlm_loss + mim_loss + itm_loss

Each fusion layer, as the released code runs it:

# One fusion layer: self-attention, then cross-attention, then feed-forward
# Input:  text_x (B, M, 768), image_y (B, N, 768) from the previous layer
# Output: new text_x, new image_y  [code: m3ae/modules/m3ae_module.py#L281-L285]

def fusion_layer(text_x, image_y):
    # 1. Self-attention within each stream, residual + layer norm  [§2.1]
    text_s = self_attend(text_x)
    image_s = self_attend(image_y)
    # 2. Cross-attention: queries from own stream, keys and values from
    #    the OTHER stream's layer input, not its self-attended output
    #    (paper eq. 5-6 use the self-attended outputs)  [code: bert_model.py#L479-L487]
    text_c = cross_attend(text_s, keys_values=image_y)
    image_c = cross_attend(image_s, keys_values=text_x)
    # 3. Feed-forward sub-layer  [§2.1]
    return feed_forward(text_c), feed_forward(image_c)

Hyperparameters

Values marked "paper" are stated in the paper. Values marked "code only" come from the released configuration and are not in the paper.

Fine-tuning (code only; the paper does not state these). Learning rate 5e-6 for all rows, linear schedule with 10% warm-up, and the stated head multiplier:

Retrieval margin 0.2 m3ae_module.py. Zero-shot retrieval runs the pretrained checkpoint at image size 288 finetune_m3ae.sh.

Results

Author-reported. Pretraining on the ROCO and MedICaT training sets; each downstream task uses its official split. Accuracy for VQA and classification, Recall@K for retrieval. Each experiment is run three times with different seeds and the mean is reported §3.2. The ablation rows are on the VQA-RAD test set §3.4.

Med-VQA tab. 1

The table header says "SLACK"; the text and the citation are for SLAKE tab. 1 §3.2.

MELINDA classification (test accuracy) tab. 2

ROCO image-caption retrieval (test set) tab. 3

Ablation on VQA-RAD test tab. 4

Table 1 is the same configuration as the last row of Table 4 tab. 1 tab. 4.

MIM layer choice. fig. 2 plots results for MIM taken from layers 0 to 6. The text says layer 0 (no textual information) is the worst and layer 3 is the best §3.4 fig. 2. The figure's metric and exact values are not stated in its caption or the text.

Qualitative analysis. Four Med-VQA cases on VQA-RAD. The MLM+MIM model answers all four as described; the MLM-only model fails the second and third cases §3.5 fig. 3. The paper presents this as illustrative only.

Paper vs code

Official repository zhjohnchan/M3AE at commit 14da377255280eed547e451f368cf71a43f3ca1c. No LICENSE file. Cite only; no verbatim code is quoted.

• Cross-attention inputs. Paper eq. 5-6 use the self-attended outputs H^vs and H^ls as keys and values in cross-attention §2.1 eq. 5 eq. 6. The code passes the other stream's layer input, before that stream's self-attention, as keys and values m3ae_module.py bert_model.py.
• Attention scaling. Paper eq. 4 has no scaling factor eq. 4. The code divides scores by √(head dimension) bert_model.py.
• MIM target. Paper: mean squared error against original pixels §2.2. Code: per-patch normalised pixels (mean and variance per patch, plus 1e-6) objectives.py config.py.
• Vision encoder depth. Paper: CLIP ViT-B §3.1. The code's Transformer builds layers − 1 blocks clip_model.py. build_model sets the count from the checkpoint, which is 12 for ViT-B clip_model.py clip_model.py. By reading only, that gives 11 vision blocks, not 12. The MIM decoder uses the same class with num_layers + 1, which gives 4 blocks prediction_heads.py. I did not run the model to confirm the vision block count.
• Initialisation. The paper says the architectures are CLIP ViT-B and RoBERTa-base §3.1 and does not say whether the weights are pretrained. The code downloads OpenAI CLIP ViT-B/16 weights and loads Hugging Face RoBERTa-base weights clip_model.py m3ae_module.py.
• Patch size. Not in the paper. The released script uses ViT-B/16 (clip16), so 288 px gives 18 × 18 = 324 patches pretrain_m3ae.sh.
• Language decoder. Paper: an MLP §2.2. Code: a BERT prediction head (dense, activation, layer norm) followed by a linear projection onto the vocabulary, with its own output weights prediction_heads.py.
• Masking style. Paper: random tokens at 15% §2.2. Code: whole-word masking, with 80% [MASK], 10% random token and 10% unchanged base_datamodule.py data_collator.py.
• Layer tap. Code captures the image and text features at the input of fusion layer index 3, that is after three fusion layers. Layer 0 is before any fusion layer m3ae_module.py. This matches the paper's description of layer 0 §3.4, but the paper's "k-th Transformer layer" wording does not say which side of a layer is meant.
• ITM sampling. Paper: ITM from UNITER, no details §3.1. Code: exactly half the batch is positive; each negative image is a random index from the whole dataset objectives.py base_dataset.py.
• Data filtering. Paper: ROCO and MedICaT, with 1,000 validation and 1,000 test images sampled from MedICaT §3.1. Code: MedICaT entries are filtered to radiology figures, texts are captions plus s2orc references of at least three words, and the split uses seed 42 with the last 2,000 shuffled entries as validation and test prepro_pretraining_data.py. ROCO pairs come from the radiology folders, with texts from the ROCO captions and from GORC reference files prepro_pretraining_data.py. The paper does not mention the radiology filter.
• Image preprocessing. Paper: "center-crop to resize each image into 288×288" §3.1. Code: bicubic resize of the shorter side to 288, then centre crop transform.py.
• Batch size and precision. Not in the paper; see the Hyperparameters table.
• Checkpoint selection. Paper: not stated. Code: fine-tuning reports the checkpoint with the best validation metric main.py main.py.
• Seeds. Paper: three seeds, mean reported §3.2. Code: default seed 0 with no seed loop in the run scripts config.py.
• Loss sum. Paper eq. 1 and code agree: unweighted sum, weights 1 eq. 1 m3ae_module.py.

No discrepancy was found for: 100,000 steps, AdamW, the 1e-5 and 5e-5 learning rates, the 10% warm-up with linear decay, the 75% and 15% masking ratios, the 288 px input and the 6-layer, 768-wide, 12-head fusion.

Pitfalls

• Three forward passes per step. Each of the MLM, MIM and ITM losses calls the full fusion model again objectives.py. Training cost is roughly three forward passes plus backward, not one. The MIM pass removes 75% of image patches before the vision encoder, so it sees 81 visible patches plus the prepended token at 288 px with patch 16 m3ae_module.py m3ae_module.py.
• Unstated settings. Patch size, batch size, decoder widths and depths, optimiser betas and weight decay, text max length, precision, the MLM 80/10/10 rule, ITM negative sampling and every fine-tuning hyperparameter are not in the paper. Use the released config values listed in the Hyperparameters table, and say which ones you chose.
• Config defaults are not the released run. config.py defaults are ViT-B/32, a BERT tokenizer and max text 40, and the pretraining named config keeps 224 px and BERT. The released script overrides these to ViT-B/16, RoBERTa, 288 px and max text 64 config.py pretrain_m3ae.sh.
• MIM target normalisation. Reconstruction is compared against per-patch normalised pixels. A plain-pixel MSE, as the paper text reads, will change the loss scale and possibly the learning dynamics.
• Tap layer indexing. Implementations that count the tap after the layer rather than before it will reconstruct from a different depth. Use the code's convention (features after three fusion layers) if you want the reported setting.
• Resize versus crop. Resizing the shorter side and then centre-cropping keeps aspect ratio; a direct square resize does not. The paper's wording allows either.
• Downstream image size differs. Fine-tuning uses 384 px, while the zero-shot retrieval row uses 288 px finetune_m3ae.sh finetune_m3ae.sh.
• Broken command in the released scripts. The zero-shot retrieval line in finetune_m3ae.sh contains a stray fragment pwdper_gpu_batchsize= finetune_m3ae.sh. Fix it before use; I did not run it.
• Published checkpoints. The test script lists single checkpoints scoring 77.4, 76.9 and 76.7 on VQA-RAD test_m3ae.sh. Paper values are means over three seeds §3.2; compare like with like.
• Table header typo. "SLACK" in the table header refers to SLAKE tab. 1.
• Minor rounding in text. The text says 5.5% over CPRD-BAN on VQA-RAD closed questions; the table gives 83.46 − 77.90 = 5.56 §3.3 tab. 1.
• Equation 1 typesetting. The source omits one closing parenthesis in eq. 1 eq. 1.

Lineage

• Builds on: masked autoencoders for images, which this paper adapts to two modalities §2.2 MAE; BERT for the language encoder and its masking §2.1 §2.2.
• Related objectives: the image-text matching task is from UNITER §3.1. Other medical VLP work named in the paper: LXMERT, VisualBERT, UNITER and PixelBERT applied to medical classification, and a Med-VQA pretraining study §1.
• Compared against: ViLT and METER on retrieval tab. 3 (METER is also on wiki as 2111.02387); MFB, SAN, BAN, MEVF and CPRD-BAN on Med-VQA tab. 1; NLF, SAN and RoBERTa on MELINDA tab. 2.
• Superseded by: not known from wiki content.