Text-Attributed Knowledge Graph Enrichment with Large Language Models for Medical Concept Representation

Mohsen Nayebi Kerdabadi; Arya Hadizadeh Moghaddam; Chen Chen; Dongjie Wang; Zijun Yao

cs.LGApr 14, 2026

Text-Attributed Knowledge Graph Enrichment with Large Language Models for Medical Concept Representation

Authors: Mohsen Nayebi Kerdabadi, Arya Hadizadeh Moghaddam, Chen Chen, Dongjie Wang, Zijun Yao

Organizations: University of Kansas, USA

Abstract

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)
	Model	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.

Representation Learning

Representation learning develops data transformations that expose useful structure for prediction, transfer, generation, and control.

See all 773 papers

REFINE: LLM Refinement over Budgeted Text-Attributed Graphs for Personalized Medical Concept Representation
Sep 3, 2026Mohsen Nayebi Kerdabadi, Arya Hadizadeh Moghaddam, Dongjie Wang +1Text-Attributed Graph Learning LLMs for Graph Learning
University of Kansas, USA
SMI: Ontology-Grounded Semantic IDs for Hierarchy-Aware Medical Concept Representation
Oct 13, 2026Ziyang Song, Lia Shen, Yixuan Li +6Representation Learning Hierarchical Representation Learning
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
HCOE: Hyperbolic Clinical Ontology Embeddings from Biomedical Language Models
Sep 25, 2026Yixuan Li, Weihao Li, Ziyang SongMedication Recommendation Hierarchical Representation Learning
McGill University · Northwestern University · Ohio University

Sep 3, 2026J/K move · Enter open · S save

REFINE: LLM Refinement over Budgeted Text-Attributed Graphs for Personalized Medical Concept Representation

Mohsen Nayebi Kerdabadi, Arya Hadizadeh Moghaddam, Dongjie Wang, Zijun Yao

University of Kansas, USA

Text-Attributed Graph Learning LLMs for Graph Learning Longitudinal EHR Modeling Biomedical KGs

Open paper AskPDF arXiv

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.

Abstract

Representation Learning

Explore similar work

REFINE: LLM Refinement over Budgeted Text-Attributed Graphs for Personalized Medical Concept Representation

SMI: Ontology-Grounded Semantic IDs for Hierarchy-Aware Medical Concept Representation

HCOE: Hyperbolic Clinical Ontology Embeddings from Biomedical Language Models