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

Problem Statement

While deep learning models have achieved radiologist-level accuracy in medical imaging, their clinical adoption is hindered by a lack of interpretability and safety risks associated with interactive feedback. Existing prototype-based models, which offer interpretability by matching images to learned visual examples, suffer from three specific structural limitations:

1. Independent Concept Modeling: Current task heads treat clinical findings as independent, ignoring known anatomical co-occurrences (e.g., active tuberculosis and cavity formation), thereby discarding clinically meaningful structure.
2. Unsafe Spatial Interaction: When physicians provide corrective feedback (e.g., drawing a bounding box), the system blindly applies this feedback. If the model is confused and associates a different concept with the indicated region, the feedback silently amplifies errors, creating patient safety risks.
3. Closed Vocabulary: Adding a new clinical finding (e.g., a rare variant) requires full retraining of the prototype atlas, which is prohibitively expensive in clinical settings.

Methodology: The GRAPE Architecture

GRAPE proposes a unified four-stage training framework that integrates concept supervision, contrastive prototype learning, Vision-Language Model (VLM) alignment, and graph-structured classification. The architecture consists of three core modules:

Module A: Graph Attention Task Head

To address the independence of concepts, GRAPE constructs a data-driven co-occurrence graph from training labels.

• Graph Construction: A $K \times K$ conditional probability matrix is derived from the training set, representing $P(k | k')$. This matrix is thresholded to create a directed adjacency graph with self-loops.
• Graph Attention Network (GAT): Instead of a flat linear layer, the model employs two layers of Graph Attention Networks. The initial node features are the prototype similarity scores ($s_{k,:}$). The GAT refines these scores via learned attention mechanisms before passing them to a Layer-Norm–Dropout–Linear readout for classification. This allows the model to leverage anatomical correlations between findings.

Module B: Concept-Mismatch Safety Check

This module introduces a safety mechanism to prevent the amplification of erroneous physician feedback.

• Mechanism: When a physician draws a bounding box ($B$) to reinforce a specific concept ($k^*$), the system computes the mean similarity ($\bar{s}_k(B)$) for all concepts within that region using the spatial mean of prototype similarity maps.
• Safety Logic: The system identifies the dominant concept ($k^\dagger$) inside the box. If $k^\dagger \neq k^*$ and the difference in mean similarity exceeds a threshold ($\eta = 0.05$), a safety warning is issued. The feedback is withheld until the physician confirms the correction.
• Variance Signal: The system utilizes prototype variance maps ($U_k$) to detect ambiguous evidence, guiding prototype masking during training and serving as a deterministic signal for safety checks at inference (avoiding the sampling overhead of Bayesian uncertainty methods).

Module C: Open-Vocabulary Prototype Anchoring

To eliminate the closed-vocabulary bottleneck, GRAPE aligns visual prototypes to frozen text embeddings.

• Alignment: A learnable projection ($W_{txt}$) maps text embeddings from a frozen BioViL-T encoder into the visual prototype space. An alignment loss ($L_{align}$) pulls visual prototypes toward their corresponding text anchors.
• Zero-Shot Addition: At test time, new concepts can be added without retraining the backbone or other components. A clinician initializes prototypes for a new concept using the text anchor perturbed by Gaussian noise. These prototypes are immediately usable for inference or can be adapted with a few labeled examples via concept-local fine-tuning.

Key Contributions

1. Graph-Structured Reasoning: The introduction of a Graph Attention Task Head that models concept co-occurrence, improving macro-F1 by +13.8 percentage points over prototype baselines on the TBX11K dataset.
2. Safety Mechanism: The first "Concept-Mismatch Safety Check" in prototype-based medical classifiers. It detects when a physician's intended label conflicts with the model's dominant visual evidence in a region, catching 85% of erroneous annotations (vs. 51% for MC-Dropout) with no extra inference cost.
3. Open-Vocabulary Extension: A mechanism to add new clinical findings from a single labeled image (or text description) without modifying the existing model components, addressing the retraining bottleneck.
4. Spatial Faithfulness: Prototype maps achieve 2.6$\times$ better lesion localization (Pointing Game) than end-to-end baselines on TBX11K, demonstrating that interpretability does not require sacrificing spatial accuracy.
5. Efficiency: All three capabilities add only +1 ms latency at interactive batch sizes, satisfying real-time interaction requirements.

Experimental Results

The authors evaluated GRAPE on TBX11K (3-class TB detection) and NIH ChestX-ray14 (14-class finding detection).

• Classification Performance: On TBX11K, GRAPE (with GNN and Uncertainty modules) achieved a Macro F1 of 0.866 ± 0.042, significantly outperforming the CSR baseline (0.728 ± 0.050) with a p-value of 0.0013. The gain is attributed primarily to the GNN head.
• Localization (Pointing Game): GRAPE achieved a Pointing Game score of 0.140 on TBX11K, compared to 0.055 for Concept Bottleneck Models (CBM) and 0.053 for PIP-Net. This indicates superior alignment between the model's attention and ground-truth lesion locations.
• Safety Check Efficacy: In simulated error scenarios, the safety check caught 85.1% of misdrawn boxes with a False Positive rate of 16.4%.
• Zero-Shot Capability: On NIH ChestX-ray14, a single labeled image for the "Effusion" concept allowed the model to recover 100% of the full-supervision Pointing Game accuracy. On TBX11K, 1-shot learning recovered 41% of full supervision accuracy.
• Faithfulness: Insertion/Deletion AUC evaluations confirmed that the prototype similarity maps capture genuinely predictive spatial regions, outperforming standard baselines in Insertion AUC.

Significance and Claims

The paper positions GRAPE as a step toward safer, more adaptable, and clinically viable interactive AI.

• Clinical Trust: By preventing the silent amplification of errors through the Concept-Mismatch Safety Check, the system addresses a critical failure mode in human-in-the-loop medical AI.
• Adaptability: The Open-Vocabulary capability lowers the barrier to entry for resource-limited settings, allowing new concepts to be introduced without the prohibitive cost of full retraining.
• Interpretability vs. Accuracy: The work demonstrates that graph-structured reasoning can improve classification accuracy (macro-F1) while maintaining or improving spatial interpretability, challenging the trade-off often assumed between these two goals.

The authors note limitations, including the absolute Pointing Game scores remaining low due to resolution mismatches (7x7 maps vs. 224x224 boxes) and the need for further validation on demographic subgroups. They emphasize that the system is a research prototype and not yet deployed in clinical settings.