← Back to Home

MedGNN

Five graph models predicting the emergency diagnosis of a patient visit under one identical evaluation contract. CEI-GNN, the model I designed, explains itself as a sum of evidence and scores best.

Mar 2026 – now · MSc thesis

PyTorch
PyTorch Geometric
GraphXAI

Overview

Graph neural networks are used for clinical prediction but are almost always explained after the fact, so an attribution map need not describe how the model decided. This thesis predicts the 10 most frequent emergency department diagnoses from multi-visit patient graphs built from MIMIC-IV-ED, runs five graph models under one identical protocol, and adds CEI-GNN, a model whose prediction is a sum of readable evidence contributions.

Example 03 · Long historyDiagnosis label 09

Index visit
Patient
Chief complaint
Vital sign
Lab result
Earlier lab result
Prior diagnosis
Structural link
Clinical relation

Synthetic graphs, not patient records. Read them left to right in time: an earlier lab value links to the same lab now with its change, earlier diagnoses link to the visit, and findings recorded together are joined. Every visit carries 6 triage vitals; real graphs hold more labs (median 9) than drawn here, and each lab also links to a shared lab type node.

What I built

CEI-GNN v3: each prediction is the sum of contributions from single findings, pairs of findings and tests that were not measured, with lab and vital values encoded piecewise linearly. It has 94,380 parameters, about 4 times fewer than ProtGNN (399,884).
One shared protocol for all five models: the same 10,000 training visits, a separate 5,000-visit development split for checkpoint selection, 40 epochs without early stopping and three seeds; validation scored once, test fold never opened.
A runtime guard that checks node, edge and subject fingerprints match across methods and stops the run on a mismatch, plus per-run manifests with input and split checksums, command, seed and package versions; result folders are never overwritten.
A GraphXAI explanation study of CEI-GNN against ProtGNN: 15,000 explanation records on 500 visits.

Five models, one contract

Model · where the explanation comes from	How it reasons
CEI-GNN v3Built in	Additive evidence. Adds up contributions from findings, finding pairs and unmeasured tests; every term can be read directly.
ProtGNNBuilt in	Case-based. Compares each graph to learned class archetypes and classifies by similarity.
GSATBuilt in	Stochastic attention. Learns to sample an informative subgraph and predicts from it.
GraphCareBuilt in	Visit attention. Weights visits and the concepts inside them, then predicts from the weighted sum.
GCHM-PNAPost-hoc only	Relation-aware multi-aggregator message passing with gated concept hubs, added as a stronger graph encoder.

CEI-GNN, ProtGNN, GSAT and GraphCare explain themselves; GCHM-PNA needs an explainer applied afterwards.

Results

CEI-GNN v3 ranks first on validation macro-F1 at 0.639, ahead of ProtGNN at 0.628: +0.0115 with a 95% patient bootstrap interval of +0.0049 to +0.0176, and it wins on all three seeds.
CEI-GNN also leads on accuracy, balanced accuracy and top-3 accuracy, and on 8 of the 10 diagnoses; head injury is slightly lower.
Gradient explanations of CEI-GNN beat a random mask clearly (fidelity-minus lift 0.34 and 0.38); ProtGNN's barely do (0.007 and 0.008). Every CEI-GNN minus ProtGNN interval excludes zero.
Under equal conditions CEI-GNN trains slightly faster: 2.7 against 3.0 seconds per epoch.

Benchmark results

Method	Macro-F1	Accuracy	Balanced acc.	Top-3
CEI-GNN v3	Best in column: 0.639±0.006	Best in column: 0.630±0.005	Best in column: 0.642±0.011	Best in column: 0.853±0.004
ProtGNN	0.628±0.003	0.612±0.004	0.627±0.005	0.844±0.001
GCHM-PNA	0.618±0.001	0.606±0.004	0.621±0.005	0.839±0.004
GraphCare	0.606±0.008	0.595±0.006	0.608±0.005	0.824±0.006
GSAT	0.595±0.002	0.584±0.004	0.599±0.005	0.822±0.004

Validation fold (4,254 visits, 10 diagnoses), mean ± standard deviation over three seeds; checkpoints chosen on the development split; best per column underlined.

Explanation faithfulness

Explainer	Fidelity-plus lift		Fidelity-minus lift
Explainer	CEI-GNN	ProtGNN	CEI-GNN	ProtGNN
Grad	0.401	0.201	0.340	0.007
IntegratedGrad	0.431	0.226	0.384	0.008
GNNExpl.	0.109	0.072	0.086	0.002
Native	0.446	n/a	0.379	n/a

Gain over each model's own random mask on 500 visits (498 patients), three seeds, 15,000 records; higher is better. Fidelity-plus: how much more the prediction drops than with a random mask when the important inputs are removed. Fidelity-minus: how much less it drops when only the important inputs are kept.

Limits

One public dataset, one fixed training sample and validation fold only; no external validation and no claim of clinical readiness.
The CEI-GNN lead is conditional on three training seeds; the test fold stays closed until the final evaluation.
Explanation scores measure faithfulness to each model, not clinical correctness, and come from the study checkpoints rather than the validation runs above. Research snapshot, not clinical advice.