Official implementation and reproducibility resources for:
KGP-Med: Knowledge-Graph-Grounded Prototype Learning for Explainable Multimodal Chest X-Ray Diagnosis
Authors: Moinak Bose and Longwei Wang
Affiliation: Department of Computer Science, University of South Dakota, Vermillion, SD, USA
KGP-Med is an explainable multimodal chest X-ray diagnosis framework that integrates:
- chest radiographs,
- leakage-controlled pre-interpretation clinical context,
- anatomical localization,
- structured biomedical knowledge,
- class-specific relational prototypes, and
- prototype-level counterfactual interventions.
The central objective is to provide diagnostically useful multimodal predictions while exposing a structured reasoning pathway whose prototype contributions can be inspected and intervened upon.
Rather than representing prototypes as unconstrained latent vectors, KGP-Med associates each prototype with an explicit:
anatomy → finding/concept → diagnostic target
path.
The model combines a conventional visual-clinical prediction pathway with a sparse knowledge-graph-grounded prototype reasoning pathway through a learned class-specific residual gate.
KGP-Med consists of three main stages.
The model processes:
- Chest X-ray: encoded using an ImageNet-pretrained DenseNet-121.
- Clinical context: leakage-controlled pre-interpretation text encoded using BioClinicalBERT.
- Knowledge sources: Chest ImaGenome scene-graph information and UMLS-linked concepts.
Only pre-interpretation clinical information is provided to the diagnostic text pathway. Findings, Impression text, and localized post-interpretation phrases are not used as diagnostic inputs.
Learned anatomy queries extract region-specific visual representations under Chest ImaGenome bounding-box supervision.
A relational biomedical graph represents anatomy, finding, and diagnostic concepts using typed relations. These graph representations are used to instantiate class-specific relational prototypes.
The final configuration uses:
- 29 anatomy queries
- 6 diagnostic targets
- 8 prototypes per target
- 48 relational prototypes in total
- Top-3 prototype evidence selection per target
- 3 relational graph layers
- 8-head prototype cross-attention
- 384-dimensional shared embedding space
The final prediction combines:
- a direct visual-clinical prediction pathway, and
- a sparse structured prototype reasoning pathway.
For class (c),
[ z_c = z_c^{VC} + \rho_c z_c^{reason}, ]
where (z_c^{VC}) is the direct visual-clinical logit, (z_c^{reason}) is the structured reasoning contribution, and (\rho_c) is a learned class-specific residual gate.
This design allows conventional multimodal information to reach the classifier directly while retaining an interpretable structured component.
For each prediction, the explanation can expose:
- selected anatomical region,
- finding/concept,
- anatomy-finding-target prototype path,
- prototype identifier,
- prototype contribution, and
- prediction change following prototype removal.
The reported experiments consider six chest X-ray findings:
- Atelectasis
- Cardiomegaly
- Consolidation
- Edema
- Pleural Effusion
- Pneumothorax
Unknown or unavailable labels are masked from the diagnostic loss rather than automatically treated as negatives.
The experiments use radiographs and structured annotations derived from:
Used for chest radiographs, image identifiers, patient identifiers, and associated study information.
Used for:
- structured yes/no target attributes,
- anatomical bounding boxes,
- anatomy-centered scene graphs,
- localized phrases,
reason_for_exam,- construction of relational prototype paths, and
- grounding supervision.
Used for semantic normalization and biomedical concept grounding where mappings are available.
After image/scene-graph matching and removal of samples without usable target labels, the reported cohort contains:
- 161,938 frontal chest radiographs
- 44,274 unique patients
A deterministic patient-disjoint 80/10/10 split is used:
| Split | Images | Patients |
|---|---|---|
| Training | 129,154 | 35,216 |
| Validation | 16,288 | 4,511 |
| Test | 16,496 | 4,547 |
Split seed: 42
Important: This repository does not redistribute MIMIC-CXR or other restricted source data. Users are responsible for obtaining the datasets through their official access mechanisms and complying with their respective licenses and data-use requirements.
The primary Kaggle implementation is available at:
Kaggle Notebook:
https://www.kaggle.com/code/moinakboss/jbhi-kgp
The notebook contains the implementation used for development, training, evaluation, and analysis of KGP-Med.
Training artifacts are available for the following experimental seeds.
Seed 123:
https://www.kaggle.com/datasets/moinakboss/kgp-med-v2-seed123-resume
Seed 2026:
https://www.kaggle.com/datasets/moinakboss/kgp-med-v2-seed2026-resume
These resources contain saved experiment artifacts associated with the corresponding training runs.
An additional development/execution notebook is available at:
https://colab.research.google.com/drive/1PfCteo-AKZb52YWGGgaAutP4vRsOFHTx
Access to the Colab notebook depends on its Google Drive sharing permissions.
KGP-Med/
│
├── README.md
├── KGP_Med.ipynb
│
├── requirements.txt
│
├── figures/
│ ├── architecture/
│ ├── localization/
│ └── qualitative_examples/
│
└── docs/
└── experiment_notes.md
The main notebook should contain or link to the complete preprocessing, model construction, training, evaluation, ablation, localization, and prototype-intervention workflow used for the reported experiments.
The final KGP-Med experiments use the following principal configuration:
| Component | Configuration |
|---|---|
| Image encoder | ImageNet-pretrained DenseNet-121 |
| Clinical encoder | BioClinicalBERT |
| Image resolution | 384 × 384 |
| Shared embedding dimension | 384 |
| Anatomy queries | 29 |
| Diagnostic targets | 6 |
| Prototypes per target | 8 |
| Total prototypes | 48 |
| Prototype selection | Top-3 per target |
| Relational graph layers | 3 |
| Prototype cross-attention | 8 heads |
| Clinical-context dropout | 0.30 |
| Batch size | 6 |
| Gradient accumulation | 2 |
| Effective batch size | 12 |
| Maximum training epochs | 5 |
| Non-language learning rate | (2\times10^{-4}) |
| Language learning rate | (2\times10^{-5}) |
| Weight decay | (10^{-4}) |
| Gradient clipping | 5.0 |
| Model seeds | 42, 123, 2026 |
The implementation uses validation-based checkpoint selection and validation-set temperature scaling.
Residual-gated experiments were run using two NVIDIA Tesla T4 GPUs in model-parallel configuration.
KGP-Med is evaluated along several complementary dimensions.
- AUROC
- AUPRC
- F1 score
- Brier score
- Expected Calibration Error (ECE)
- Pointing-game accuracy
- Attention mass within the annotated anatomy bounding box
Prototype interventions measure the effect on the final diagnostic logit after removing:
- the highest-contributing supporting prototype, and
- an irrelevant control prototype.
The model is additionally evaluated after removing clinical context at inference time.
Three-seed performance on the patient-disjoint test cohort:
| Model | Macro AUROC | Macro AUPRC | Macro F1 | Brier | ECE |
|---|---|---|---|---|---|
| Image-only DenseNet-121 | 0.8891 ± 0.0010 | 0.8440 ± 0.0012 | 0.7410 ± 0.0048 | 0.1219 ± 0.0017 | 0.1358 ± 0.0045 |
| DenseNet-121 + BioClinicalBERT | 0.8992 ± 0.0008 | 0.8667 ± 0.0009 | 0.7704 ± 0.0044 | 0.1106 ± 0.0015 | 0.1145 ± 0.0041 |
| KGP-Med | 0.8978 ± 0.0010 | 0.8688 ± 0.0007 | 0.7772 ± 0.0053 | 0.1098 ± 0.0019 | 0.1187 ± 0.0052 |
KGP-Med therefore maintains diagnostic discrimination comparable to the conventional multimodal baseline while introducing an explicit structured explanation pathway.
The final KGP-Med model obtains:
- Pointing-game performance: 0.9671 ± 0.0008
- Attention mass inside anatomy box: 0.8589 ± 0.0011
Removing anatomy supervision reduces these values to:
- Pointing game: 0.268 ± 0.016
- In-box attention mass: 0.204 ± 0.013
This ablation demonstrates that anatomical localization is specifically learned through the anatomy-supervision objective rather than emerging automatically from diagnostic classification.
For the complete KGP-Med model:
-
Supporting prototype removal:
(\Delta z = 0.81 \pm 0.03) -
Irrelevant prototype removal:
(|\Delta z| = 0.0065 \pm 0.0015)
The mean effect of removing the selected supporting prototype is therefore approximately 125× larger than removal of the irrelevant control prototype.
These interventions test whether the structured evidence displayed by the prototype pathway is functionally connected to the model output.
The intervention analysis should not be interpreted as claiming that the prototype pathway exhaustively explains the direct visual-clinical branch.
Removing knowledge-graph grounding produces only small changes in conventional diagnostic and localization metrics.
The experiments therefore support the primary role of the knowledge graph as semantic organization and grounding of prototype identities, rather than as an independent source of classification improvement.
The graph allows otherwise latent prototypes to remain associated with explicit anatomy-finding-target relations that can be inspected during inference.
The implementation is notebook-based.
For the most direct reproduction path:
- Open the public Kaggle notebook.
- Attach the required dataset resources according to the notebook configuration.
- Ensure access to the required MIMIC-CXR-derived and Chest ImaGenome resources.
- Execute the notebook cells in order.
- Select the desired experiment seed and configuration.
- Train or resume the corresponding model.
- Run the evaluation sections to reproduce diagnostic, localization, calibration, ablation, and intervention outputs.
Because dataset paths may differ between Kaggle, Colab, and local environments, update the corresponding input paths before running the notebook.
The manuscript reports the main comparisons using seeds:
42
123
2026
Available resume artifacts:
Seed 123:
https://www.kaggle.com/datasets/moinakboss/kgp-med-v2-seed123-resume
Seed 2026:
https://www.kaggle.com/datasets/moinakboss/kgp-med-v2-seed2026-resume
The same patient partition, target construction, image resolution, training budget, checkpoint-selection procedure, and calibration protocol should be retained when comparing experimental variants.
Preventing diagnostic label leakage is an important part of the experimental design.
The diagnostic text branch uses only leakage-controlled pre-interpretation information such as:
reason_for_exam
The following information is not supplied to the diagnostic input pathway:
Findings
Impression
augmented report text
localized post-interpretation Chest ImaGenome phrases
Localized report-derived phrases may be used as training-time semantic grounding supervision but are not diagnostic model inputs.
For an individual positive prediction, KGP-Med can associate the prediction with:
Anatomical region
Finding / concept
Prototype identifier
Anatomy → finding → target path
Prototype contribution
Prediction change after prototype removal
The qualitative examples in the accompanying paper show how these components can be presented together with anatomy-query localization maps.
KGP-Med: Knowledge-Graph-Grounded Prototype Learning for Explainable Multimodal Chest X-Ray Diagnosis
Moinak Bose and Longwei Wang
Department of Computer Science
University of South Dakota
Vermillion, SD, USA
Journal submission: IEEE Journal of Biomedical and Health Informatics (JBHI)
If you use this work, please cite the corresponding paper.
BibTeX information will be updated after publication.
@article{bose2026kgpmed,
title={KGP-Med: Knowledge-Graph-Grounded Prototype Learning for Explainable Multimodal Chest X-Ray Diagnosis},
author={Bose, Moinak and Wang, Longwei},
journal={IEEE Journal of Biomedical and Health Informatics},
year={2026},
note={Submitted}
}The citation above describes the submitted manuscript and should be updated with the final publication metadata if the article is accepted and published.
MIMIC-CXR, MIMIC-CXR-JPG, Chest ImaGenome, and UMLS remain subject to the licenses and access conditions established by their respective data providers.
This repository provides implementation and experiment resources but does not redistribute restricted clinical data.
The repository and linked notebooks are intended to support inspection and reproduction of:
- multimodal cohort construction,
- leakage-controlled text processing,
- anatomy-query learning,
- relational knowledge-graph construction,
- KG-grounded prototype generation,
- prototype cross-attention,
- sparse prototype reasoning,
- residual-gated multimodal prediction,
- multi-label diagnostic evaluation,
- calibration,
- anatomical localization,
- component ablations,
- missing-context evaluation, and
- prototype intervention analysis.
Where possible, experiment seeds and saved training artifacts are provided to improve reproducibility.
For questions regarding the implementation or experiments, please open an issue in this repository.
This work uses resources derived from MIMIC-CXR/MIMIC-CXR-JPG, Chest ImaGenome, UMLS, DenseNet-121, and BioClinicalBERT. Please cite the corresponding original resources when using these datasets or models.
KGP-Med is a research system developed for investigation of multimodal medical-image learning and explainable artificial intelligence.
It is not a clinical diagnostic device and should not be used for patient-care decisions without appropriate clinical validation, regulatory review, and expert oversight.