Source-linked AI summary

MedKLIP: Medical Knowledge Enhanced Language-Image Pre-Training in Radiology

Chaoyi Wu, Xiaoman Zhang, Ya Zhang, Yanfeng Wang, Weidi Xie

arXiv:2301.02228v3eess.IVcs.CLcs.CV

TL;DR

Medical VLP needs domain-specific knowledge to support fine-grained diagnosis, especially for diseases unseen during training. MedKLIP extracts report triplets, translates entities into medical descriptions, and aligns them with image patches, achieving strong zero-shot and fine-tuned classification and grounding results across public benchmarks.

  • Problem

    Medical VLP must support fine-grained, explainable zero-shot diagnosis despite raw reports overlooking medical prior knowledge and new diseases requiring costly annotation and retraining.

  • Method

    MedKLIP extracts {entity, position, exist} triplets, translates entities through medical knowledge bases, and uses Transformer-based local alignment with image patches.

  • Results

    MedKLIP shows strong zero-shot classification and grounding across diseases and image distributions, with fine-tuning still outperforming previous models.

  • Takeaways & Limitations

    Entity-level medical knowledge and spatial alignment provide a foundation for diagnosis involving both seen and unseen diseases.

Abstract

from arXiv · show

In this paper, we consider enhancing medical visual-language pre-training (VLP) with domain-specific knowledge, by exploiting the paired image-text reports from the radiological daily practice. In particular, we make the following contributions: First, unlike existing works that directly process the raw reports, we adopt a novel triplet extraction module to extract the medical-related information, avoiding unnecessary complexity from language grammar and enhancing the supervision signals; Second, we propose a novel triplet encoding module with entity translation by querying a knowledge base, to exploit the rich domain knowledge in medical field, and implicitly build relationships between medical entities in the language embedding space; Third, we propose to use a Transformer-based fusion model for spatially aligning the entity description with visual signals at the image patch level, enabling the ability for medical diagnosis; Fourth, we conduct thorough experiments to validate the effectiveness of our architecture, and benchmark on numerous public benchmarks, e.g., ChestX-ray14, RSNA Pneumonia, SIIM-ACR Pneumothorax, COVIDx CXR-2, COVID Rural, and EdemaSeverity. In both zero-shot and fine-tuning settings, our model has demonstrated strong performance compared with the former methods on disease classification and grounding.

1. Introduction

The paper targets medical vision-language pre-training that can support zero-shot disease diagnosis and grounding by incorporating domain knowledge into image-report learning. It extracts structured medical information, translates entities into descriptions, and aligns those descriptions with image patches, with strong results across multiple benchmarks and settings.

  • Medical diagnosis models often require costly annotation and retraining when new diseases or categories emerge.
  • Medical VLP aims to improve zero-shot disease classification and grounding, while addressing fine-grained concepts, domain knowledge, robustness, and explainability.
  • The model extracts reports into entity-position-existence triplets, translates entities into knowledge-based descriptions, and aligns them with image patches using a Transformer.
  • The model is pre-trained on MIMIC-CXR and evaluated across ChestX-ray14, RSNA Pneumonia, SIIM-ACR Pneumothorax, COVIDx CXR-2, COVID Rural, and EdemaSeverity.
  • The approach achieves state-of-the-art zero-shot classification and grounding across diseases and image distributions, and remains significantly stronger after fine-tuning.

2. Related Work

Related work spans general vision-language pre-training, medical information extraction, medical knowledge enhancement, and concurrent medical VLP methods. Existing approaches include dual- and single-stream architectures, knowledge integration, entity-relation extraction, and contrastive alignment.

  • General Vision-Language Pre-training (VLP) Models: General VLP architectures are organized into two-stream dual encoders and single-stream methods that favor visual-language fusion.Some works also incorporate commonsense knowledge into vision-language pre-training.
  • Medical Named-Entity-Recognition (NER) Models: Early radiology-report extraction methods focused only on diseases, while later methods extracted relationships among entities without a predefined closed disease set.The later approaches aim to retain more useful information with high accuracy.
  • Medical Knowledge Enhanced Models: Medical knowledge enhancement methods are classified as model-based or input-based according to how external knowledge is used.Model-based approaches imitate radiological or diagnostic practice, whereas input-based approaches treat knowledge as an extra feature-computation input.
  • Concurrent Works in Medical VLP: Concurrent medical VLP methods generally follow a two-stream flow using contrastive learning without a fusion module.ConVIRT aligns medical scans with corresponding reports, while LoVT and GLoRIA improve local alignment.

3. Method

The method converts radiology reports into knowledge-enriched triplets and aligns entity queries with image features through a Transformer fusion module. This design supports disease existence prediction and spatial grounding, including for unseen entities.

  • Problem Scenario: At inference, the model predicts disease existence and visual evidence for queried entities, replacing unseen entities with user-provided descriptions for zero-shot inference.The existence output supports classification, while cross-attention supports grounding.
  • Report Pre-processing: Triplet extraction reformulates each report into {entity, position, exist}, removing unnecessary language-grammar complexity while retaining clinical observations, locations, and existence status.Entities describe clinical observations, positions denote anatomical locations, and exist labels encode true, false, or uncertain findings.
  • Knowledge-enhanced Triplet Encoding: Knowledge-enhanced triplet encoding translates medical entities into detailed descriptions retrieved from medical knowledge bases before ClinicalBERT embedding.The descriptions decompose professional disease terms into shared attributes intended to support more reliable zero-shot diagnosis and relationships across seen and unseen diseases.
  • Fusion Module: A Transformer-based fusion module uses entity embeddings as queries over image features, producing entity-level existence and position predictions.The entity query set contains the most common entity embeddings from training reports, while visual features provide keys and values.
  • Fusion Module: The model averages Transformer cross-attention maps into spatial heatmaps that support grounding at inference without directly applying a training loss to the heatmaps.The heatmaps are up-sampled to the input image size and function as segmentation-like visual evidence maps.

3.4. Training

Training supervises triplet-derived existence and position predictions rather than aligning images with entire reports. The resulting paradigm provides fine-grained entity-level supervision through classification and contrastive position losses.

  • Training: Existence prediction uses binary cross-entropy with triplet existence labels, while position prediction uses contrastive learning over positive and sampled negative position embeddings.Uncertain labels are passed through for the existence loss, and position embeddings are sampled from the most common position set.
  • Training: Predictions for entities without corresponding report labels are ignored during loss computation.This avoids supervising unmatched entity queries.
  • Training: Triplet-based training provides supervision at the entity level instead of relying on global image-report alignment.The position and existence components of triplets act as fine-grained supervision labels.

4. Experiment

The experiments pre-train on MIMIC-CXR and evaluate medical image-language methods across classification, grounding, and segmentation benchmarks. They compare against existing medical VLP baselines using task-specific metrics and zero-shot protocols.

  • Datasets: MIMIC-CXR v2 provides over 227k paired image-report studies, totaling 377,110 images from 65,379 patients.Studies contain one or two images from different scan views.
  • Datasets: The evaluation covers ChestX-ray14, RSNA Pneumonia, SIIM-ACR Pneumothorax, COVIDx CXR-2, COVID Rural, and Edema Severity.These datasets support disease classification, segmentation, and fine-grained pulmonary-edema classification tasks.
  • Implementation: During pre-training, triplet extraction and text encoders remain fixed, while the visual encoder and fusion module train end-to-end on image-text pairs.Fine-tuning initializes ResNet50 or ResUNet encoders with the pre-trained image encoder for classification or segmentation.
  • Baselines: The study compares against ConVIRT, GLoRIA, BioViL, and CheXzero, retraining some models on MIMIC-CXR for fair comparison.BioViL uses officially released models, while zero-shot prompts follow the described comparison protocols.
  • Metrics: Classification uses AUC, F1, and ACC; grounding uses Pointing Game accuracy, while segmentation reports Dice and IOU.Precision and recall supplement segmentation evaluation, especially when strict overlap metrics may understate differences in difficult zero-shot cases.

5. Results

The model consistently outperforms prior methods in zero-shot classification and grounding, including unseen diseases, and transfers strongly across fine-tuning tasks and data regimes.

  • Zero-shot evaluation compares classification and grounding against state-of-the-art image-text pre-training methods, while fine-tuning tests transferability with 1%, 10%, and 100% of training data.
  • Zero-shot classification: 0.87 AUC on RSNA Pneumonia and 0.89 AUC on SIIM-ACR Pneumothorax improve over 0.83 and 0.71, respectively.ChestX-ray14 average AUC also rises from 0.69 to 0.77.
  • Zero-shot classification: 0.74 AUC and 0.70 ACC for unseen COVID-19 improve over 0.66 AUC and 0.59 ACC when medical entity descriptions are used.The unseen-disease setting tests COVID-19, absent from the 2015 MIMIC-CXR reports.
  • Zero-shot grounding: Zero-shot grounding improves across seen diseases, including RSNA Pneumonia pointing game from 0.83 to 0.87 and detection precision from 0.50 to 0.64.The method also surpasses prior approaches on SIIM-ACR Pneumothorax grounding metrics.
  • Zero-shot grounding: 0.58 pointing game score for unseen COVID-19 grounding improves over 0.40, with results largely consistent with unseen-disease classification.The paper attributes more interpretable representations to knowledge-enhanced language encoding.
  • Fine-tuning: Fine-tuning yields substantial AUC improvements across classification datasets and large segmentation gains across diseases and image distributions, especially with limited data.Severity grading also achieves the best results on most severity levels.

6. Conclusion

The paper introduces a knowledge-enhanced medical VLP model that extracts and translates medical triplets, then aligns entity descriptions with image regions. It reports strong zero-shot and fine-tuned performance across diagnosis and grounding tasks.

  • The model extracts medical triplets, translates entities into detailed descriptions, and uses a Transformer-based structure for local region alignment.
  • The method shows strong zero-shot classification and grounding abilities, including for unseen diseases.
  • With fine-tuning, the method significantly outperforms state-of-the-art methods across evaluated settings.

A. The Entity Description Base and Position Set

The entity-description base supplements extracted radiology entities with medical descriptions, while the position set provides predefined anatomical locations for triplet representation.

  • The entity query set Q contains 75 entities and covers 90% of entities in reports.COVID-19 is included only for inference because it does not appear in pre-training reports.
  • Entity descriptions translate individual names into medical explanations, with examples including pleural thickening, interstitial markings, scars, and hyperinflated lungs.
  • The position set P contains 51 positive anatomical positions used to represent report locations.

B. Implementation Details

The implementation uses a ResNet50 visual encoder, ClinicalBERT-based report encoding, and Transformer decoding to process image patches and medical triplets.

  • Images are resized to 224 × 224 × 3, and the first four ResNet50 layers produce 14 × 14 × 256 visual feature maps.
  • A pre-trained NER module extracts report triplets, while ClinicalBERT supplies entity and position embeddings.
  • The model uses 75 entities, 51 positions, seven negative positions per entity, and four Transformer Decoder layers with four heads.

C. Ablation Study

The ablation study evaluates the transformer-based entity-query fusion module, PosCL, and ET encoding, finding complementary contributions to classification and grounding.

  • The final method combines transformer-based entity-query fusion, position location contrastive loss, and entity translation encoding.
  • The basic entity-level model outperforms many prior methods, supporting triplet preprocessing and entity-level supervision over raw-report processing.
  • Adding PosCL improves ChestX-ray14 AUC from 0.75 to 0.76 in the reported ablation.
  • The ablation includes zero-shot classification metrics of AUC, F1, and ACC, with ChestX-ray14 values reported as macro averages across 14 diseases.
  • Zero-shot grounding ablations are reported for RSNA Pneumonia and SIIM-ACR Pneumothorax.

D. Detailed results on ChestX-ray14

On ChestX-ray14, the method is compared disease by disease and in macro-average AUC under zero-shot and fine-tuning settings, outperforming prior methods on most diseases.

  • The method exceeds former methods for most of the 14 diseases in zero-shot ChestX-ray14 classification.
  • Under 100% fine-tuning, the method achieves similarly excellent results on ChestX-ray14.
  • ChestX-ray14 comparisons report AUC for each disease and macro-average AUC across all 14 diseases.

E. Visualization Results

Zero-shot visualizations compare ground-truth regions with model heatmaps for Pneumonia, Pneumothorax, and Covid-19, illustrating disease-specific localization behavior.

  • The model captures Pneumonia regions well despite broad ground-truth areas and highlights the thin, narrow location of Pneumothorax abnormalities.
  • For Covid-19, the visualizations suggest successful grounding despite textual similarity to Pneumonia and the challenge of an unseen disease.
  • The displayed diseases are Pneumonia, Pneumothorax, and Covid-19, with ground-truth and prediction labels shown across the visualization panels.
  • Each visualization column represents one disease, pairing ground truth on the left with the model’s prediction heatmap on the right.
  • Brighter red indicates regions the model considers more likely associated with abnormalities.
Loading 2301.02228v3…