GRAPE: Graph-Augmented Prototype Explanations for Interactive Medical Image Diagnosis

Rasul Khanbayov, Erchin Serpedin, Hasan Kurban

Accepted to NeurIPS 2026 🎉

Prototype-based medical image classifiers treat clinical findings as independent, silently amplify unsafe physician feedback, and require full retraining to add a new finding. GRAPE targets all three limitations in one architecture: a nonlinear task head implemented with graph-attention layers, a concept-mismatch safety check that flags erroneous clinician corrections before they are applied, and open-vocabulary prototype anchoring for adding a new finding from a single labeled image, with only +1ms added latency at interactive batch size.

Pointing Game localization comparison across methods

Pointing Game on TBX11K chest X-rays (Active TB concept). Each marker shows the max-activation point of a method. These images were selected because GRAPE's point (orange circle) falls inside the ground-truth bounding box (cyan); overall Pointing Game values remain low.

Three Modules, One Architecture

GRAPE architecture diagram showing Modules A, B and C

An X-ray is encoded by ResNet-50 and an MLP projector into spatial features. Module A matches them against a learned prototype atlas, and the score matrix feeds a GAT-style head built over a co-occurrence graph. Module B compares mean concept similarity inside a clinician's box against the claimed concept's, flagging a mismatch before feedback is applied. Module C adds concepts from a text anchor via a frozen BioViL-T encoder. Dashed arrows mark training-only or concept-addition-only paths.

Module A — Graph Attention Task Head

Builds a data-driven concept co-occurrence graph from training labels and refines prototype similarity scores with a two-layer GAT-style head, boosting macro-F1 by +13.8 pp over the prototype baseline on TBX11K but not on NIH ChestX-ray14. In the evaluated implementation, the graph does not affect the head's outputs at inference, so it does not contribute to this gain.

Module B — Concept-Mismatch Safety Check

To our knowledge the first such mechanism in prototype-based medical classifiers: warns when the model's dominant finding inside a doctor-drawn region conflicts with the claimed label, catching 85% of simulated erroneous annotations at no extra inference cost at a favorable operating point, though sensitivity falls substantially at clinically realistic false-positive budgets.

Module C — Open-Vocabulary Prototype Anchoring

Aligns visual prototypes to frozen BioViL-T text embeddings, letting a new clinical finding be added from a single labeled image without modifying any other component of the model. Evaluated for localization only; with one example, Pointing Game stays near or below chance.

Results

+13.8 pp

Macro F1 vs. CSR (TBX11K)

2.1×

Better Lesion Localization

85%

Safety Check Catch Rate

+1 ms

Added Latency

Dataset	Method	Macro F1 ↑	Pointing Game ↑	Δ F1	p-value
TBX11K	CSR (baseline)	0.728 ± 0.050	0.094 ± 0.033	—	—
TBX11K	GRAPE (GNN + Unc.)	0.866 ± 0.042	0.080 ± 0.032	+13.8 pp	0.0013
NIH ChestX-ray14	CSR (baseline)	0.6048 ± 0.0170	—	—	—
NIH ChestX-ray14	GRAPE (full)	0.6031 ± 0.0174	—	−0.2 pp	0.894 (n.s.)

Multi-seed comparison (5 seeds {0,1,2,3,42}) under the same interpretability-constrained training regime (frozen backbone, concept supervision, staged curriculum). On TBX11K, the task head yields a large and significant gain, to which the co-occurrence graph does not contribute; on the 14-concept NIH task, GRAPE and CSR are statistically comparable under the rigorous patient-level split protocol.

Prototype similarity maps across CSR, CBM, PIP-Net, and GRAPE on selected Active TB examples. Percentage badges report the fraction of total heatmap energy inside the ground-truth box.

Cite

@inproceedings{khanbayov2026grape,
  title={GRAPE: Graph-Augmented Prototype Explanations for Interactive Medical Image Diagnosis},
  author={Khanbayov, Rasul and Serpedin, Erchin and Kurban, Hasan},
  booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
  year={2026}
}