Multi-Granularity Cross-modal Alignment for Generalized Medical Visual Representation Learning

In plain words

Chest X-ray datasets often come with a written report for each image. Labelling images by hand is slow, but the reports already describe what is in the picture. This paper trains an image model by matching each X-ray to its report, so it learns from the reports without extra labels.

Plain matching is not enough, because a disease usually covers a small part of the picture and a report mentions only a few key findings. The method therefore matches at three levels: the whole image with the whole report, small image patches with individual report words, and groups of cases with the same disease across both kinds of input. The groups are found by clustering, and the two sides are asked to agree on them.

The trained image model is then reused on other chest X-ray tasks such as classification, detection and segmentation, where only a small amount of labelled data is available. The authors report that the gain is largest when fine-tuning uses 1% of the labels.

TL;DR

MGCA combines three contrastive or cluster-consistency losses to pre-train a ViT-B/16 (or ResNet-50) image encoder and a BioClinicalBERT text encoder on MIMIC-CXR image and report pairs. The image encoder is then frozen for linear classification, a YOLOv3 detector or a U-Net segmenter. Author-reported results beat GLoRIA and ConVIRT in the settings they report, with the largest margins at 1% of training data. Use this card if you need a medical chest X-ray pre-training recipe with image-report pairs, but expect to reconcile several hyperparameter conflicts between paper and code (see Paper vs code).

Problem

Medical image models need many labelled examples, and labels are expensive to produce §1. Radiology reports come with the images, but earlier image-text methods use only part of the correspondence: instance-level contrast (ConVIRT) or local region-word contrast that runs in one direction (GLoRIA) §1, §2. Pathologies usually occupy a small portion of the image, and only a few report tags describe the key condition §3.2, token-wise alignment. The authors argue that disease-level correspondence between patients is also left unused, because standard instance-level contrast pushes apart pairs that share a diagnosis §3.2 Disease-level Cross-modal Alignment. The paper assumes that the frozen pre-trained encoder should transfer to small labelled downstream sets §4.2.

Method

MGCA has one image encoder f_v and one text encoder f_t. Each produces a global embedding and a set of token embeddings §3.1. Three losses are summed with weights \lambda_1,\lambda_2,\lambda_3 §3.3, eq. 10.

Shared front end §3.1, §3.2.

1. The image encoder gives a sequence of visual token embeddings \mathbf{R}_i and a global vector \mathbf{v}_i. The text encoder gives token embeddings \mathbf{Z}_i and a global vector \mathbf{t}_i §3.1.
2. Two non-linear projection heads g_v and g_t map the globals to d-dimensional normalised embeddings \tilde{\mathbf{v}}_i, \tilde{\mathbf{t}}_i §3.2, eq. 1. The paper sets d=128 §4.1.

Instance-wise alignment (ITA) §3.2, eq. 1-3. 3. Compute the cosine similarity of each image-report pair eq. 1. 4. For each pair in a minibatch of size B, compute an image-to-text InfoNCE loss and a text-to-image InfoNCE loss, both at temperature \tau_1 eq. 2. 5. \mathcal{L}_{\mathrm{ITA}} is the average of the two losses over the N pairs, written with a factor 1/(2N) eq. 3.

Token-wise alignment (CTA) §3.2, eq. 4-6. 6. For each visual token j, compute a soft-attention summary \mathbf{o}_i^j of the text tokens. The paper writes this with learnable Q, K, V matrices and scaling by \sqrt{d} eq. 4. 7. The local image-to-text loss \mathcal{L}_{\mathrm{LIA}} pulls each visual token toward its summary and pushes it away from the other summaries. It is an InfoNCE in both directions, at temperature \tau_2, with each token weighted by w_i^j eq. 5. 8. The weight w_i^j is the last-layer attention from visual token j to the [CLS] token, averaged over heads §3.2, after eq. 5. 9. The text-to-image loss \mathcal{L}_{\mathrm{LTA}} is built the same way, in the other direction §3.2. 10. \mathcal{L}_{\mathrm{CTA}} is the mean of \mathcal{L}_{\mathrm{LIA}} and \mathcal{L}_{\mathrm{LTA}} eq. 6.

# Token-wise loss (CTA), one direction shown  [§3.2] [eq. 4] [eq. 5]
# Input:  visual_tokens (B, S, d), text_tokens (B, L, d), weights (B, S)
# Output: scalar loss

def token_alignment_loss(visual_tokens, text_tokens, weights):
    # 1. Soft-match each visual token to text tokens  [eq. 4]
    summaries = attend(visual_tokens, text_tokens)   # (B, S, d)
    # 2. Contrast each visual token against the summaries of its image
    #    with a per-token InfoNCE weighted by importance  [eq. 5]
    loss_forward = weighted_infonce(visual_tokens, summaries, weights)
    loss_backward = weighted_infonce(summaries, visual_tokens, weights)
    return (loss_forward + loss_backward) / 2

Disease-level alignment (CPA) §3.2, eq. 7-9. 11. Keep K trainable prototypes \mathbf{c}_k of dimension d §3.2. The paper uses K=500 §4.1. 12. For each image and each report embedding, compute a softmax over cosine similarities to the prototypes, at temperature \tau_3 eq. 7. 13. Compute a soft cluster code for each image and each report with iterative Sinkhorn-Knopp §3.2. The paper does not give the number of iterations or the balancing parameter. 14. Cross-predict: the image prediction is trained with the report's code as target, and the report prediction with the image's code as target eq. 8. 15. \mathcal{L}_{\mathrm{CPA}} averages the two cross-entropies over the N pairs eq. 9.

Overall objective §3.3, eq. 10. 16. \mathcal{L} = \lambda_1 \mathcal{L}_{\mathrm{ITA}} + \lambda_2 \mathcal{L}_{\mathrm{CTA}} + \lambda_3 \mathcal{L}_{\mathrm{CPA}}, with \lambda_1=\lambda_2=\lambda_3=1 in the main setting §4.1.

The three modules are trained jointly, not in stages. The authors name this as a limitation and a direction for future work §5 Limitations and Future Work.

Hyperparameters

Values from the paper are cited to the paper. Code values are cited to the repository at code_commit. Where they conflict, both are shown in Paper vs code.

Results

Author-reported. Setup: MIMIC-CXR pre-training, ViT-B/16 unless marked ResNet-50, linear or frozen-backbone protocols §4.2, §4.3. Baselines are GLoRIA (pre-trained on CheXpert), and ConVIRT and GLoRIA-MIMIC (pre-trained on MIMIC-CXR). Values are the paper's own numbers.

Classification, AUROC % for CheXpert and RSNA, accuracy % for COVIDx tab. 1:

The paper states that the ViT-B/16 row is best in all nine classification settings §4.3. Table 1 marks best and second-best values with colours, not bold.

Error bars over three runs, ViT-B/16 linear protocol tab. 6: CheXpert 88.7 ± 0.18 (1%), 89.13 ± 0.16 (10%), 89.5 ± 0.21 (100%); RSNA 89.03 ± 0.11, 89.92 ± 0.14, 90.77 ± 0.04; COVIDx 73.9 ± 0.64, 84.75 ± 0.22, 92.85 ± 0.50. The 1% CheXpert mean is 88.7 here but 88.8 in tab. 1.

Object detection, mAP % on RSNA and Object CXR, ResNet-50 YOLOv3 frozen backbone tab. 2. Values marked "-" in the paper are below 1%:

Semantic segmentation, Dice % on SIIM and RSNA, ResNet-50 U-Net tab. 3:

Ablation, linear classification on CheXpert and RSNA (AUROC) and U-Net on SIIM (Dice), ViT-B/16 for the first two and ResNet-50 for SIIM tab. 4 §4.4:

Natural-domain VLP comparison: BLIP (pre-trained on 14M natural image-text pairs), fine-tuned in the same linear setting, scores 69.1 / 74.9 / 77.7 on CheXpert and 53.7 / 82.0 / 84.1 on RSNA at 1% / 10% / 100% tab. 5. The authors use this to argue that medical pre-training is needed §4.4.

Figures: fig. 1 shows multi-level correspondences between an image and its report. fig. 2 is the architecture with the three modules. fig. 3 visualises learned token correspondences from the attention. fig. 4 is a t-SNE of image embeddings coloured by disease and by cluster. These are described only; no figure is reproduced here.

Paper vs code

Repository: https://github.com/HKU-MedAI/MGCA at commit 3a2b4f25d47efc62e7d01d9204632a23d453afae. The paper links https://github.com/fuying-wang/MGCA, which redirects to this repository; the README cites the paper. The owner is the HKU-MedAI organisation. The repository does not show that the organisation is the authors' lab, so ownership is only partly confirmed; the paper's own link is the basis for using it. The repository has no LICENSE file and GitHub reports no license, so code_license is recorded as not stated. Code is cited by path and line only; nothing is quoted. The paper says code "will be released later" [checklist, §Checklist], so the repository is the release referred to.

Discrepancies, each with both citations:

• Instance temperature swapped. Paper: \tau_1 = 0.1 §4.1. Code applies softmax_temperature = 0.07 to the instance loss mgca_module.py mgca_module.py.
• Token temperature swapped. Paper: \tau_2 = 0.07 §4.1. Code applies local_temperature = 0.1 to all token losses mgca_module.py mgca_module.py.
• Instance loss scale. Paper: average over two losses with factor 1/(2N) §3.2, eq. 3. Code: sum of two cross-entropies, each a batch mean, with no factor 1/2 mgca_module.py. This doubles the instance term relative to the paper.
• Token attention has no learned projections by default. Paper: Q, K, V are learnable matrices §3.2, eq. 4. Code: use_local_atten defaults to False, so the attention is a plain softmax of the normalised embeddings with no Q, K, V mgca_module.py mgca_module.py. The learned multi-head attention branch exists but is off by default mgca_module.py.
• Token weights are clipped. Paper: w_i^j is the last-layer attention to [CLS], averaged over heads §3.2. Code does this, but also clips each weight to its 10th and 90th percentiles, and does the same for word-side weights taken from the text encoder mgca_module.py mgca_module.py. Not stated in the paper.
• Padding is masked for patch-to-word attention mgca_module.py. Not stated in the paper.
• Loss weights differ between constructor and argparse. Paper: \lambda = 1, 1, 1 §4.1. Code argparse defaults are 1, 1, 1 mgca_module.py, but constructor defaults are 1, 0.7, 0.5 mgca_module.py. Which applies depends on the entry point.
• Batch size. Paper and README: 144 §4.1 README.md. Argparse default 72 mgca_module.py; constructor default 64 mgca_module.py.
• Prototype normalisation. Code re-normalises the prototype weights to unit length on every step mgca_module.py. Not stated in the paper.
• Prototype assignments use exponentiated, \varepsilon-scaled prototype scores with 3 Sinkhorn iterations on the batch mgca_module.py mgca_module.py. The paper gives neither \varepsilon nor an iteration count §3.2.
• Pre-training length and warm-up match: 50 epochs and warm-up over the first 40% of steps, which is 20 epochs mgca_module.py mgca_module.py.
• Fine-tuning scripts. run_cls_finetune.sh is empty in the repository. run_det_finetune.sh exists but the classification run command is in the README README.md. The README uses --data_pct 0.01 for CheXpert and a detector learning rate of 5e-4 README.md README.md. The finetuner default is data_pct=0.01 and batch_size=48 mgca_finetuner.py.
• Embedding dimension matches: 128 mgca_module.py.

Pitfalls

• Temperatures. Paper and code disagree on which temperature goes with which loss (see above). The paper does not say which setting produced its tables, so the reported numbers cannot be tied to either choice.
• Instance loss scale. The paper's 1/(2N) and the code's sum differ by a factor of two on that term. With \lambda_1 = 1, the effective weight of the instance loss differs by two.
• Downstream details are missing from the main text. Learning rate, batch size, optimiser and number of epochs for linear classification and for the downstream detection and segmentation heads are not in the paper body; the appendix that the paper cites is not in the source. The README gives 5e-4 for the detector and segmenter, but nothing for the linear head.
• Dataset split arithmetic. The SIIM split is written as 70% / 30% / 30% §4.2. Those sum to 130%, so the split is ambiguous. The per-dataset preprocessing scripts in mgca/preprocess were not checked for the actual split.
• Sinkhorn is unspecified. Without \varepsilon, iteration count, and whether the codes are computed over the full batch or across GPUs, CPA is not reproducible from the paper. The code uses a distributed variant when DDP is active mgca_module.py.
• Prototype freezing. The code freezes prototypes for one epoch mgca_module.py mgca_module.py; the paper does not mention this.
• Batch size strongly affects the contrastive losses. InfoNCE uses in-batch negatives, so 144 (paper) and 64 or 72 (code defaults) give different negatives per step.
• Error bars are narrow. Reported standard deviations are over three re-runs of the downstream linear protocol tab. 6 §4.4. The pre-training itself is not repeated, and the checklist answers "No" to whether error bars were reported.
• Data licences. The paper's checklist states that the datasets are under a non-commercial licence [checklist]. The cards here do not redistribute those datasets, but check each dataset's terms before using pre-trained weights or data commercially.
• Compute. Two RTX 3090 GPUs for 50 epochs at batch 144 §4.1. The README says about one day README.md.

Lineage

• Builds on: ConVIRT (image-report contrastive pre-training on MIMIC-CXR) and GLoRIA (region-word contrast on CheXpert) §1, §2; the InfoNCE and cross-attention formulation §3.2; Sinkhorn-Knopp clustering for prototype codes §3.2. These papers are not compiled cards in this wiki, so nothing from them is used here.
• Compared against: ConVIRT, GLoRIA, GLoRIA-MIMIC (the authors' re-implementation on MIMIC-CXR), BLIP for natural-image VLP, and ImageNet and random initialisation tab. 1-5.
• Superseded by: not known from content in this wiki.