In electronic health record (EHR) mining, learning high-quality representations of medical concepts (e.g., standardized diagnosis, medication, and procedure codes) is fundamental for downstream clinical prediction. However, ro bust concept representation learning is hindered by two key challenges: (i) clinically important cross-type dependencies (e.g., diagnosis medication and medication-procedure relations) are often missing or incomplete in existing ontology resources, limiting the ability to model complex EHR patterns; and (ii) rich clinical semantics are often missing from structured resources, and even when available as text, are difficult to integrate with KG structure for representation learning. To address these challenges, we present MedCo, an LLM empowered graph learning framework for medical concept representation. MedCo first builds a global knowledge graph (KG) over medical codes by combining statistically reliable associations mined from EHRs with type-constrained LLM prompting to infer semantic relations. It then utilizes LLMs to enrich the KG into a text-attributed graph by generating node descriptions and edge rationales, providing semantic signals for both concepts and their relationships. Finally, MedCo jointly trains a LoRA-tuned LLaMA text encoder with a heterogeneous GNN, fusing text semantics and graph structure into unified concept embeddings. Extensive experiments on MIMIC-III and MIMIC-IV show that MedCo consistently improves prediction performance and serves as an effective plug-in concept encoder for standard EHR pipelines.
Figures & tables
Figure 1: Overview of MedCo . (1) Extract intra-visit co-occurrence and next-visit transition evidence from EHRs, and retain well-supported code pairs. (2) Use type-constrained, evidence-conditioned LLM prompting to assign directed relation types (with confidence and rationales), thereby constructing a heterogeneous KG. (3) Enrich nodes with LLM-generated concept descriptions and edges with relation metadata. (4) Jointly train a LoRA-tuned LLaMA encoder and a relation-aware GNN to learn medical concept embeddings for downstream prediction.
Figure 2: Prompt for Medical Relationship Induction.
Figure 3: Performance gains from integrating MedCo into four diagnosis prediction backbones (plug-in analysis) on MIMIC-III and MIMIC-IV.
Model
General Performance
Label Category Performance (AUPRC)
AUPRC
F1
Acc@15
Acc@20
Acc@30
0–25%
25–50%
50–75%
75–100%
MIMIC-III
Base
41.00
33.16
47.20
50.40
58.80
40.60
47.30
72.80
78.40
GRAM
41.70
34.60
48.60
51.90
59.80
42.20
50.10
74.10
79.30
MMORE
42.60
35.30
49.20
52.60
60.40
42.80
51.20
74.80
80.00
KAME
42.18
35.07
48.97
52.06
60.06
41.84
49.76
74.06
79.14
G-BERT
42.47
35.38
49.26
52.33
60.27
42.06
50.04
74.24
79.28
Table 1: Performance comparison on MIMIC-III and MIMIC-IV. We report overall AUPRC, F1, and Acc@k, and stratified AUPRC by diagnosis label frequency quartiles (0–25%, 25–50%, 50–75%, 75–100%). The first row corresponds to the base Transformer, and subsequent rows denote plug-in variants (e.g., GRAM = Base + GRAM).
Metric
MIMIC-III
MIMIC-IV
# Patients
7,515
18,829
# Visits (samples)
12,430
25,028
# Labels/sample
12.31
10.56
# Unique conditions (ICD)
515
562
# Conditions/sample
26.85
59.50
# Drugs/sample
70.10
118.16
Table 2: Data statistics for MIMIC-III and MIMIC-IV.
Figure 4: Performance evaluation across different training set sizes using the MIMIC-III and MIMIC-IV datasets. An asterisk () next to each encoder indicates the integration of MedCo to that model, e.g., Transformer = Transformer + MedCo .
Model
PRAUC
F1
Acc@20
MIMIC-IV
Base (Transformer)
45.10
39.10
58.10
Base + KG
48.55
42.60
60.90
Base + KG + Edge feat.
48.72
42.85
61.10
Base + KG + Edge feat. + LLM (freeze)
49.38
43.71
61.47
Base + KG + Edge feat. + LLM (LoRA)
50.00
45.25
62.54
MIMIC-III
Base (Transformer)
41.00
33.16
50.40
Table 3: Component-wise ablation of MedCo on MIMIC-IV and MIMIC-III.
Figure 5: Ablation study showing the impact of removing each edge type on PRAUC for MIMIC-III. The dashed vertical line indicates the full-model baseline.
LLM backbone
AUPRC
Acc@20
Llama-3.2-1B
50.00
62.54
Gemma-2-2B
50.31
62.89
Qwen2.5-1.5B
50.16
62.67
Qwen2.5-0.5B
49.31
61.44
SmolLM2-1.7B
50.11
62.61
Table 4: Sensitivity of MedCo to the choice of LLM backbone on MIMIC-IV.
Variant
AUPRC
Acc@20
Train(s)
Infer(s)
Peak train mem (MiB)
Peak infer mem (MiB)
Base (no KG, no LLM/GNN)
45.10
58.10
22.0
1.30
169
106
GNN only (KG, no LLM)
48.72
61.10
42.7
3.19
1,273
518
LLM frozen ( K=0 ) + GNN
49.38
61.47
40.0
3.19
1,273
518
LLM LoRA + GNN ( K=5 )
49.70
62.35
183.4
3.19
14,477
518
LLM LoRA + GNN ( K=10 )
50.00
62.54
275.4
3.19
24,260
518
Table 5: Cost-accuracy trade-off on MIMIC-IV with batch size 128. Experiments were run on a machine with an Intel Xeon Silver 4214R CPU (24 cores), 256 GiB RAM, and one RTX A6000 GPU (48 GB). For compactness, Train(s) denotes training time per epoch, and Infer(s) denotes inference time per epoch for the test set.
Appendix figures & tables13 assets
Supplementary material from the paper’s appendix.
Appendix
Resource
Structure / Goal
Scope
Cross-Type
Key Limitations
ICD / CCS
Coding taxonomy
dx, px
No
Primarily within-dx or within-px hierarchical categories with only parent-child relations; lacks explicit cross-domain relations.
ATC
Drug classification taxonomy
rx
No
Drug taxonomy; primarily within-rx hierarchical categories with only parent-child relations; lacks explicit cross-domain relations.
SNOMED CT
Definitional concept model
dx, px, rx
Limited
Rich definitional semantics (e.g., finding site, causative agent, method, active ingredient); limited coverage of within-domain relations (dx-due-to-dx); no coverage of cross-domain relations at predictive granularity.
UMLS
Metathesaurus integration / mapping
dx, rx, px
Limited
Strong normalization/mapping; cross-type relations are incomplete/uneven (e.g., limited coverage for px-dx indications and px-rx semantics).
Appendix
Table 6: Common medical resources (ICD World Health Organization (2025) , CCS Healthcare Cost and Utilization Project (HCUP) (2025) , ATC World Health Organization Collaborating Centre for Drug Statistics Methodology (2025) , SNOMED CT Stearns et al. (2001) , UMLS Bodenreider (2004) ) and their limitations.
Figure 6: Distribution of final evidence-supported code pairs by category (MIMIC-III). Percentage values are shown on top of each bar.
Figure 7: Distribution of final evidence-supported code pairs by category (MIMIC-IV). Percentage values are shown on top of each bar.
Predicted relationship
Mean
Std
causes_adverse_event
4.70
0.27
co_occurs_with
4.80
0.45
co_prescribed_with
4.80
0.27
complicates
4.90
0.22
interacts_with
4.90
0.22
leads_to
4.90
0.22
Appendix
Table 7: Clinical expert audit of sampled LLM-inferred edges. We sampled 50 edges by selecting 5 edges from each of the 10 most frequent predicted relation types. Each edge was independently rated by two frontline clinicians on a 1-5 correctness scale, and the reported values are computed from the per-edge mean of the two ratings.
Figure 8: Frequency of LLM-predicted relation labels over evidence-filtered candidate code pairs in MIMIC-III. The y-axis is log-scaled to highlight the long-tailed distribution of relation types.
Figure 9: Frequency of LLM-predicted relation labels over evidence-filtered candidate code pairs in MIMIC-IV. The y-axis is log-scaled to highlight the long-tailed distribution of relation types.
Relation
Description
causes
codeA is an etiologic cause of codeB (direct causal or pathophysiologic link).
risk_factor_for
codeA increases the likelihood of codeB (epidemiologic association; not necessarily causal).
leads_to
codeA typically precedes codeB as a downstream condition or stage (temporal progression without requiring strict causality).
complicates
codeA occurs as a complication during the course of codeB (arises secondary to codeB or its treatment).
co_occurs_with
codeA and codeB are frequently observed together without implied direction or clear causality (contextual association).
no_significant_relation
Clinical prior and available data indicate no clinically meaningful association between codeA and codeB.
Appendix
Table 8: Diagnosis-Diagnosis (DX-DX) relation types. Here, codeA and codeB are diagnosis codes.
Relation
Description
treats
codeA treats or manages codeB (therapeutic use).
prevents
codeA reduces risk or recurrence of codeB (prophylaxis).
causes_adverse_event
codeA may induce codeB as an adverse reaction.
contraindicated_for
codeA should be avoided when codeB is present.
co_occurs_with
codeA and codeB often co-occur without a clear causal link.
no_significant_relation
Clinical prior and available data indicate no clinically meaningful association between codeA and codeB.
Appendix
Table 9: Medication-Diagnosis (RX-DX) relation types. codeA is a medication and codeB is a diagnosis.
Relation
Description
diagnostic_of
codeA is performed to diagnose, confirm, or rule in/out codeB (diagnostic evaluation).
treats
codeA is a procedure used to treat, correct, or palliate codeB (therapeutic intervention).
monitors
codeA is performed to monitor, follow, or assess disease activity or treatment response for codeB.
contraindicated_for
codeA should be avoided when codeB is present due to safety or unfavorable risk-benefit considerations.
co_occurs_with
codeA and codeB frequently appear together without a clear diagnostic, therapeutic, or causal link (contextual association).
no_significant_relation
Clinical prior and available data indicate no clinically meaningful association between codeA and codeB.
Appendix
Table 10: Procedure-Diagnosis (PX-DX) relation types. codeA is a procedure and codeB is a diagnosis.
Relation
Description
co_prescribed_with
codeA and codeB are intentionally prescribed together in practice.
contraindicated_with
Concomitant use of codeA and codeB is contraindicated due to serious safety risk.
interacts_with
codeA and codeB have a clinically meaningful drug-drug interaction (pharmacokinetic and/or pharmacodynamic) that may impact safety or efficacy.
substitute_for
codeA is commonly used as a therapeutic alternative to codeB for similar indications (typically not co-prescribed).
combination_therapy_with
codeA and codeB are used together as an established combination therapy.
co_occurs_with
codeA and codeB frequently appear together in data without a known therapeutic relationship or interaction.
Appendix
Table 11: Medication-Medication (RX-RX) relation types. codeA and codeB are medications.
Relation
Description
sequential_care
codeA is typically followed by codeB as the next procedural step (ordered workflow; not necessarily causal).
prerequisite_for
codeA is commonly required or preparatory for performing codeB (e.g., access, imaging guidance, setup).
alternative_to
codeA and codeB are procedural alternatives for a similar clinical purpose (mutually substitutable options).
performed_same_session
codeA and codeB are commonly performed during the same procedural session or time block within an episode (intentional bundling).
co_occurs_with
codeA and codeB occur within the same episode without implied order or intentional bundling (contextual association).
no_significant_relation
Clinical prior and available data indicate no clinically meaningful association between codeA and codeB.
Appendix
Table 12: Procedure-Procedure (PX-PX) relation types. codeA and codeB are procedures.
Relation
Description
requires_medication
codeA routinely requires administration of codeB as part of the procedure (e.g., sedation, anesthesia, anticoagulation).
premedication_for
codeB is typically given before codeA to enable or optimize the procedure (e.g., anxiolysis, antibiotic prophylaxis).
post_procedure_medication_for
codeB is commonly given after codeA for recovery, prophylaxis, or symptom control (e.g., analgesia, anticoagulation).
prophylaxis_with
codeB is used to prevent complications associated with codeA (e.g., peri-procedural antibiotics, DVT prophylaxis).
adjunct_medication_for
codeB is an adjunct given with codeA to improve efficacy, safety, or tolerability (non-essential but commonly used).
contraindicates_medication
When codeA is planned or present, codeB should be avoided due to safety or risk-benefit concerns.
Appendix
Table 13: Procedure-Medication (PX-RX) relation types. codeA is a procedure and codeB is a medication.
Figure 10: Prompt template used to generate LLM-based node-level clinical descriptions.
Learning rich medical concept representations is essential for EHR prediction. Text-attributed knowledge graphs (TKGs) provide a natural foundation by organizing heterogeneous medical relations together with textual semantics. However, most existing encoders process concepts uniformly across patients, despite the fact that a code's meaning and predictive value depend on patient-specific clinical context and trajectory. Learning patient-personalized concept representations from TKGs introduces two key challenges: (1) deciding how much KG context to incorporate for each observed code, and (2) aligning semantic information with the patient-specific relational structure. We propose REFINE, a KG-aware budgeted LLM graph refinement framework for patient-personalized medical concept encoding. Starting from a global TKG, REFINE constructs patient-specific temporal graphs. A sequential reinforcement learning policy selects a personalized KG expansion budget for each observed code. The resulting patient graph is processed by a heterogeneous GNN to capture relation-aware structural dependencies, while a frozen LLM uses graph-aware soft prompts to semantically refine concept representations. Experiments on MIMIC-III and MIMIC-IV show that REFINE consistently improves diverse EHR backbones, outperforms strong baselines, and demonstrates robust gains across component ablation, KG selection, and data insufficiency.
Mohsen Nayebi Kerdabadi, Arya Hadizadeh Moghaddam, Dongjie Wang +1
Biomedical language models (LMs) provide strong representations for clinical tasks, but they often fail to capture the hierarchical structure among medical concepts. This limitation is especially pronounced for rare and fine-grained clinical codes in long-tailed vocabularies, where hierarchy-agnostic LM representations become less informative as concept frequency decreases. We introduce Semantic Medical Identifier (SMI), a training-free framework for constructing ontology-grounded Semantic IDs and hierarchy-aware medical concept representations. For each diagnosis, procedure, or medication concept, SMI defines its Semantic ID as the coarse-to-fine path through an expert-defined medical ontology. SMI then aggregates frozen biomedical LM embeddings across the corresponding ontology levels to form hierarchy-aware medical concept representations that preserve coarse-to-fine clinical semantics. The base SMI uses uniform aggregation across ontology levels, while a level-weighted variant selects global level weights on validation data. Our theoretical analysis shows that SMI yields more informative medical concept representations than hierarchy-agnostic LM embeddings, particularly for rare and fine-grained clinical concepts. Experiments show that SMI captures meaningful hierarchical structure and improves clinical relation prediction. On MIMIC-IV, SMI improves mortality prediction, readmission prediction, drug recommendation, and rare drug prediction. On eICU, SMI yields more robust patient representations under cross-hospital distribution shift.
Ziyang Song, Lia Shen, Yixuan Li +6
School of Electrical Engineering and Computer Science, Ohio University · Department of Quantitative and Computational Biology, University of Southern California · School of Computer Science, McGill University +2
Biomedical language models (LMs) encode textual semantics but do not explicitly preserve medical code hierarchies. We present Hyperbolic Clinical Ontology Embeddings (HCOE) for hierarchy-aware clinical concept representation. HCOE maps frozen BioBERT embeddings into a Poincare ball, combining parent-side and child-side ontology-guided contrastive learning with coarse-to-fine ontology-path aggregation. It uses International Classification of Diseases (ICD) codes organized by Clinical Classifications Software (CCS) and Anatomical Therapeutic Chemical (ATC) medication hierarchies. Evaluations show that HCOE performs best on ICD/ATC clinical relation prediction and CCS-to-PheCode hierarchy transfer. On the MIMIC-IV dataset, HCOE also achieves the best performance on mortality prediction, readmission prediction, medication recommendation, and rare drug prediction.
Yixuan Li, Weihao Li, Ziyang Song
McGill University · Northwestern University · Ohio University