Source-linked AI summary
Clinical Graph-JEPA: Predictive Patient-State Knowledge Graphs for Cognitive Decision Support
Kushagra Yadav, Nalin Prabhath, Amit Lamba, James E. Schrager, Goeun Han, Yining Mao
TL;DR
Clinical records contain structured events and narrative reasoning, but LLM-drafted graphs mix certain facts with uncertain inferred relations. The paper builds evidence-scored patient-state graphs, refines them with Graph-JEPA, and finds that localized discharge-note features improve overall leave-one-out MRR from 0.364 to 0.477 under matched evaluation.
Problem
Clinical records encode patient events and reasoning in different sources, making reliable structured knowledge graphs difficult when inferred relations lack table-level provenance.
Method
The framework combines evidence-scored graph construction with a masked-latent Graph-JEPA refiner and compares note-free graphs with graphs receiving localized discharge-note embeddings.
Results
0.364 to 0.477 overall leave-one-out MRR is achieved by entity-grounded note features over the note-free configuration under a matched protocol.
Takeaways & Limitations
Clinical Graph-JEPA treats LLM-drafted graphs as evidence-scored representations to refine, with note context most useful when injected into the entities it grounds.
Takeaways & Limitations
The study uses MIMIC-IV from a single hospital system and evaluates recovery within evidence-scored drafts rather than independently adjudicated graphs or prospective clinical outcomes.
Abstract
from arXiv · showhide
Clinical records contain rich evidence about patient state, but converting that evidence into reliable, structured knowledge graphs remains difficult because extraction errors, ontology mismatch, missing relations, and temporal ambiguity can propagate into downstream systems. We propose a clinical knowledge graph construction and refinement framework that combines multi-agent relation proposal, ontology-aware normalization, deterministic evidence scoring, and JEPA-based latent refinement. Rather than treating a clinical knowledge graph as a static extraction artifact, we treat it as a predictive patient-state representation. For each admission, the system constructs an evidence-scored graph from structured MIMIC-IV records and inferred clinical cross-links, then learns to recover held-out clinical relations from the observed graph context. We evaluate the refiner with leakage-free leave-one-out edge recovery (MRR and Hits@k) and held-out batch-mask evaluation (AUC and MRR). To isolate the contribution of discharge-note context, we compare a note-embedding-free configuration with a note-augmented configuration that injects real discharge-note representations only into note-grounded entities. Under the same cohort and evaluation protocol, entity-grounded note injection improves overall leave-one-out MRR by 31% relative improvement.
1 Introduction
Clinical Graph-JEPA treats LLM-generated clinical knowledge graphs as evidence-scored drafts and refines them to predict which relations belong. Localized discharge-note representations improve inferred-edge recovery compared with diffuse or absent note context.
- Structured records capture diagnoses, medications, and procedures but not the clinician reasoning connecting them.Free-text notes encode relations such as why a β-blocker was started or how chest pain indicated acute coronary syndrome.
- LLM-inferred relations lack table-level provenance and therefore vary in quality within the drafted graph.The resulting graph mixes certain structural facts with uncertain inferred relations.
- The proposed refiner predicts which edges belong from an evidence-scored graph rather than relying on one-shot extraction.Clinical Graph-JEPA uses masked-latent prediction with a lightweight readout for held-out edge recovery.
- Localized note representations substantially improve inferred-edge recovery, whereas a diffuse admission-level note vector reaches the no-note level.
- 0.364 to 0.477 overall leave-one-out MRR is the reported improvement from Option A to entity-grounded Option B.Both options use the same consolidated graph, while Option B adds note information locally to grounded entities.
2 Related Work
Prior clinical KG systems decompose extraction into specialized stages but commonly treat the extracted graph as final. This work connects evidence-aware graph refinement with self-supervised latent prediction and frozen edge recovery.
- Clinical KG systems commonly separate entity identification, normalization, relation extraction, and schema or evidence validation.Recent LLM-based and multi-agent approaches specialize these stages to improve coverage and distinguish recognition from relation reasoning.
- These systems generally treat extracted graphs as final representations despite mixing observed facts with uncertain text-inferred relations.
- JEPA and BYOL-style methods learn representations by predicting masked latents with an EMA target and stop-gradient, without negative examples.The paper adapts this latent-prediction objective to graphs as pre-training for a clinical-state world model.
- The edge-recovery readout uses a bilinear score, frozen world-model encoder, InfoNCE training, and type-matched negative sampling.This setup makes recovery quality reflect the pretrained representation rather than a jointly tuned decoder.
- Clinical-ModernBERT supplies note representations, but the paper’s contribution concerns where those representations enter the graph.
3 Method: Clinical Graphs as Predictive World Models
The method represents each admission as an auditable, evidence-scored patient-state graph and trains Graph-JEPA to refine it through masked latent prediction. Two configurations isolate the contribution of locally injected discharge-note context.
- Clinical Graph-JEPA models refinement as world modeling over structured patient-state memory, not autonomous clinical decision-making.The typed graph organizes diagnoses, medications, procedures, microbiology, and admission context for downstream retrieval and decision support.
- Figure 1 combines staged graph extraction, schema-aware masked-patch representation learning, and inference-time scoring with schema validation.
- Each admission graph fuses deterministic MIMIC-IV table edges with narrative-mediated inferred clinical cross-links.The deterministic backbone includes patient links to diagnoses, medications, procedures, microbiology, and services.
- Every inferred edge receives a deterministic support vector based on model confidence, biomedical resources, ontology proximity, and note provenance.
- 3.2 Clinical Graph-JEPA world model: The refiner first predicts EMA-target latents for masked graph states, then freezes its encoder and trains edge recovery with eight type-matched negatives.Node features include type, entity, numeric, and—under Option B—localized note information.
- The experiment evaluates localized discharge-note representations against a note-free configuration to isolate their contribution.
- Option A excludes discharge-note embeddings, while Option B injects a single Clinical-ModernBERT representation only into note-grounded entities.Ungrounded nodes receive zero note vectors; the same graph, schema, losses, and evaluation protocol are otherwise retained.
4 Results
On the same cohort and evaluation protocol, note-augmented refinement consistently outperformed note-free refinement, with localized note placement driving the improvement. The deterministic backbone supplied most recoverable context, while inferred cross-links added remaining headroom.
- Evaluation setup: 4,000 admission graphs were evaluated on an identical 80/10/10 split and recipe, with only the note feature differing between options.The evaluation included 8,283 leave-one-out edges.
- Relation-level recovery: Option B improved every inferred relation, while Option A remained above chance on every relation and served as a fallback for admissions without notes.The smallest gain was MANAGED_FOR at +0.097, where Option A already recovered comparatively well at 0.326.
- Global metrics: Option B raised overall LOO Hits@1 from 0.182 to 0.303, Hits@10 from 0.821 to 0.917, and non-obvious batch-mask AUC from 0.712 to 0.80.The undifferentiated batch-mask AUC was saturated by trivial patient-hub edges and was not the meaningful measure.
- Note placement: A global note representation performed near the no-note level at 0.352 versus 0.364, whereas entity-grounded injection reached 0.477 overall leave-one-out MRR.The matched note-placement ablation identifies localization, rather than note presence alone, as the decisive design choice.
- Context analysis: The deterministic backbone alone supplied most recoverable context, exemplified by INDICATES MRR of 0.491 at the floor versus 0.637 at the ceiling.Remaining headroom came from other inferred cross-links.
- Limitations: The findings are limited to MIMIC-IV and recovery within evidence-scored draft graphs, without independent adjudication, prospective utility, causal validity, or outcome evaluation.The authors position the model as a research and decision-support graph-refinement component rather than an autonomous clinical decision-maker.
5 Discussion
Entity-grounded note context preserves which graph entities the narrative describes, whereas a global summary broadcasts diffuse and redundant information. Reporting note-free and note-augmented options also reflects that some admissions lack notes and enables a controlled estimate of note contribution.
- Why entity-grounding works: Entity grounding preserves note-to-entity correspondence, distinguishing note-discussed entities from purely structural ones.The note and graph overlap because graph entities were extracted from the note.
- Why entity-grounding works: A global note summary broadcasts context across the graph, making it diffuse and largely redundant.Localized note features instead act as a prior on the specific entities they describe.
- Why two options: Option A avoids overstating deployable accuracy because a third of admissions have no note, while Option B measures localized note context under otherwise identical conditions.Together, the options support deployment with or without note access.
6 Conclusion
Clinical Graph-JEPA refines evidence-scored, LLM-drafted clinical graphs by recovering held-out schema-valid relations from patient-graph context. Localized discharge-note features improved overall leave-one-out MRR under matched conditions, while broader validation remains future work.
- Conclusion: Clinical Graph-JEPA treats an LLM-drafted clinical knowledge graph as an evidence-scored representation to refine rather than accept as final.The refiner recovers held-out, schema-valid clinical relations from surrounding patient-graph context.
- Conclusion: 0.364→0.477 overall leave-one-out MRR was achieved by entity-grounded note features under a matched training and evaluation protocol.Option A used no direct discharge-note embeddings; Option B injected localized representations into note-grounded entities.
- Future work: Future work will test independently clinician-annotated relations, external institutions and populations, patient-disjoint and temporally separated evaluation, calibrated revision policies, and downstream utility.These directions extend validation beyond the current evaluation setting.
A.1 MIMIC Graph Schema
MIMIC-derived admission graphs use typed nodes and directed typed edges, combining deterministic structured-record relations with inferred clinical cross-links. Laboratory-test nodes are possible in the schema but were unavailable as measured result nodes in this extract.
- Graph representation: Each admission graph is represented as Gi = (Vi, Ei) with typed nodes and directed typed edges.The supplied schema passage begins by defining admission-level graph structure.
- Schema scope: Laboratory-test nodes can be included by schema, but laboratory values were unavailable in the MIMIC-IV extract and were not used as measured result nodes.This is a scope boundary of the extract rather than a general schema restriction.
- Relation schema: The relation set combines deterministic backbone relations with inferred clinical cross-links such as MANAGED_FOR, CONFIRMS, COMPLICATED_BY, and INDICATES.The deterministic relations connect patients to observed structured-record entities, while cross-links are not stored directly as relational columns.
- Inferred cross-links: Presenting symptoms are represented as diagnosis-layer nodes with a presenting attribute, so INDICATES preserves symptom role in node metadata.This uses the same node type for symptoms and diagnoses while retaining their distinct role.
A.2 MIMIC Record Normalization
Each graph begins with a flattened MIMIC-IV admission record containing clinical, administrative, and outcome information. Diagnoses are separated into clinically distinct categories to prevent historical or external codes from being treated as active admission diagnoses.
- Flattened MIMIC-IV admission records supply identifiers, demographics, admission and discharge metadata, diagnoses, medications, procedures, microbiology, services, and outcomes.
- Diagnoses are partitioned into presenting symptoms, active diagnoses, medical history, and other or external codes.
- Medication nodes retain dose, unit, route, and count metadata when available.
- Microbiology nodes retain specimen, organism, antibiotic, and susceptibility interpretations, while services and transfer-derived locations preserve admission context.
A.3 Graph Construction
Graph construction combines a deterministic node and table-derived relation backbone with inferred clinical cross-links proposed from a source-faithful synthetic narrative. Validation restricts additions to supported, well-typed relations between existing nodes.
- Normalized admission records first produce a deterministic node set for graph construction.
- The deterministic backbone adds direct table-derived relations involving diagnoses, medications, procedures, microbiology, services, and susceptibility-supported antibiotic–organism links.
- A source-faithful synthetic narrative rendered from each structured record enables relation-specific agents to propose clinical cross-links among existing graph nodes.
- Validation removes candidate edges with unmatched endpoints, invalid type signatures, self-loops, duplicates, or unsupported evidence.
A.4 Provenance and Edge Support
The framework preserves provenance and support signals in clinical graph edges, localizes available note context to grounded entities, and evaluates latent edge recovery under controlled protocols. Results compare note representations, external validation, and downstream graph-question answering.
- Provenance and support: Each edge retains source and target nodes, relation type, confidence, evidence or source field, and provenance metadata.Backbone edges are grounded in MIMIC tables, while inferred cross-links retain narrative evidence and deterministic support signals.
- Provenance and support: Continuous support scores combine semantic similarity, treatment evidence, ontology proximity, and provenance checks rather than hard-coded rules.These scores remain edge attributes for downstream models.
- Note grounding: Note embeddings are localized to note-grounded entities, with zero note input when no admission-level note embedding is available.Grounding can use provenance support, surface-name matching, or uniform assignment to non-patient nodes.
- Graph-JEPA encoding: The encoder uses typed message passing over a schema-sanitized graph, partitions graphs into balanced BFS patches, and predicts hidden patch states from visible context.Patch features include relative size, degree, and random-walk return features; the EMA target receives the full graph.
- Evaluation: Leave-one-out recovery removes one trusted edge, ranks same-type candidate targets with the source and relation fixed, and filters other known true targets.Evaluation keeps the encoder and readout fixed and reports MRR and Hits@k.
- External validation: Entity-grounded note injection achieved the highest ACI-Bench MRR and Hits@1, Hits@3, and Hits@10, improving MRR over global injection by approximately +0.009.The external evaluation used 207 encounters without retraining and supports localizing note information to supported entities.
- Note representation ablation: Uniform local-span pooling reached LOO MRR 0.478 versus 0.468 for entity-grounded global mean, but its +0.010 paired improvement had a confidence interval including zero.Entity-conditioned attention performed worse, with LOO MRR 0.458 and a paired difference of −0.020 versus uniform local-span pooling.
- Graph-QA ablation: Option B improved macro entity F1 by 3.6 percentage points and question-level Hits@3 by 5.0 percentage points, while strict exact-match accuracy remained 65.0%.The benchmark contained 20 questions.