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

1 Introduction

• A code’s meaning and predictive value can vary with patient history, visit context, and treatment trajectory, motivating patient-personalized representations.The same diagnosis may indicate a stable chronic condition for one patient and acute deterioration for another.Source: §1 · p. 1
• Fixed KG expansion rules ignore patient- and code-specific context needs, while unnecessary neighbors can add noise, computation, and dilute salient signals.REFINE therefore uses a sequential reinforcement learning policy to select an expansion budget for each observed code.Source: §1 · pp. 1–2
• LLM semantic refinement must account for co-occurring codes, visit position, temporal history, and selected KG context rather than encode codes in isolation.The introduction motivates a compact mechanism for exposing patient-specific structure to the LLM.Source: §1 · p. 2
• REFINE builds patient-specific graphs from a global TKG and uses sequential reinforcement learning to select code-specific KG expansion budgets.The policy accounts for interactions among observed codes and the cumulative relational context added to the graph.Source: §1 · p. 2
• A heterogeneous GNN captures relation-aware dependencies, while a frozen LLM uses graph-aware soft prompts to refine code semantics before the representations are fused.The prompts combine task and visit-code hard tokens with soft tokens for observed codes and selected KG context.Source: §1 · p. 2
• REFINE’s personalized representations are integrated into standard EHR prediction backbones and evaluated on sequential diagnosis prediction.Source: §1 · p. 2

2.1 Notation and Problem Definition

• REFINE encodes each observed code from a patient’s history for use by a downstream EHR prediction model.The history is a temporally ordered sequence of visits, each containing a set of observed codes.Source: §2.1 · p. 2

2.2 Method Summary

• A five-step pipeline builds a global text-attributed KG, extracts a patient graph, selects code-specific context, applies parallel GNN and LLM branches, and fuses their outputs.Figure 1 depicts the overview, with fused representations used for downstream EHR prediction.Source: §2.2 · pp. 2–3, Figure 1 · p. 3

2.3 Global TKG Construction

• The global TKG links medical-code nodes with typed semantic relations and textual descriptions and rationales.Candidate relations are filtered using EHR co-occurrence and next-visit transition statistics before a type-constrained LLM assigns directed relations.Source: §2.3 · p. 3

2.4 Patient-Personalized Graph Construction

• For each patient, REFINE selects a patient-dependent KG subgraph for distinct observed codes and maps selected KG edges into the local patient graph.The same medical code may receive different relational context across patients.Source: §2.4 · pp. 3–4

2.5 Personalized Concept Sub-KG Extraction

• A sequential policy chooses each observed code’s KG expansion depth and per-hop fanout, with decisions conditioned on context already selected for other codes.The policy’s reward favors prediction improvement over a no-expansion graph while penalizing unnecessary KG context.Source: §2.5 · pp. 4–5

2.6 Graph-Aware LLM Refinement for Personalized Concept Representation

• The heterogeneous GNN captures graph structure while a frozen LLM uses mixed prompts containing visit structure, code tokens, and KG-context tokens to refine code semantics.A gated fusion combines structural and LLM-refined representations for downstream prediction.Source: §2.6 · pp. 5–6

2.7 Integration with Downstream Models

• REFINE’s patient-personalized code representations are aggregated within visits and passed to an EHR backbone for next-visit diagnosis prediction.In Stage II, the selected patient graphs are fixed while the full model is trained using multi-label binary cross-entropy against next-visit diagnosis labels.Source: §2.7 · p. 6

• Evaluation uses MIMIC-III and MIMIC-IV for next-visit diagnosis prediction over label spaces of 515 and 562 codes, respectively.Source: §3 · p. 6
• Reported metrics are AUPRC, frequency-stratified AUPRC, Acc@k, and F1, with results averaged over five folds.Source: §3 · p. 6
• Implementation details, including LLM/GNN configurations and a two-stage training pipeline, are provided in Appendix A.3.Source: §3 · p. 6

4.1 RQ1: Plug-in Encoder Evaluation

• REFINE consistently improves AttPool, AdaCare, Transformer, RETAIN, and TCN when used as a plug-in concept encoder.Source: §4.1 · p. 7, Figure 3 · p. 7

4.2 RQ2: Baseline Comparison

• REFINE achieves the best overall performance and the best AUPRC in six of eight diagnosis-frequency quartile groups.Source: Table 1 · p. 7, §4.2 · p. 8

4.3 RQ3: Ablation Study

• Ablations identify personalized GNN removal as the largest drop; removing LLM refinement, KG soft tokens, or RL budgeting also degrades performance.Source: Table 3 · p. 7, §4.3 · p. 8

4.4 RQ4: Personalized KG Selection

• Sequential code-level RL budgeting outperforms fixed expansion and visit-level VoI selection, whose performance varies across fixed actions.Source: §4.4 · p. 8, Figure 4 · p. 8
• The policy selects diverse no-, one-, and two-hop actions and assigns heterogeneous budgets, including two-hop expansion for rare codes but none for a frequent medication.Source: §4.4 · pp. 8–9, Figure 5 · p. 8, Table 4 · p. 9

• Many EHR methods improve patient-level prediction architectures, whereas this work focuses on patient-personalized medical concept representations.Source: §5 · p. 9
• Ontology and KG encoders use hierarchical or richer relational knowledge, while GraphCare uses fixed neighborhood expansion rather than learned sequential patient- and code-specific budgets.Source: §5 · p. 9
• Rel-LLM uses graph-aware soft prompting, while REFINE compresses selected code-specific neighborhoods into visit-aware soft context; MedCo instead learns global concept representations.Source: §5 · p. 9

• REFINE uses sequential reinforcement learning to select code- and patient-specific KG budgets and builds compact personalized graphs for EHR code representations.Source: §6 · p. 9
• A heterogeneous GNN and a frozen LLM are combined through graph-aware soft prompts to refine patient-personalized concept representations.Source: §6 · p. 9
• Experiments on MIMIC-III and MIMIC-IV show consistent gains across diverse backbones, with analyses supporting component contributions and adaptation across codes and patient contexts.Source: §6 · p. 9

7 Limitations and Future Work

• The KG selection module uses a predefined discrete action space with fixed hop and neighbor budgets.The policy does not directly select exact nodes or adapt budgets continuously at each hop.Source: §7 · p. 10
• Neighbors within a selected budget are chosen by evidence-based edge rankings rather than random sampling.Source: §7 · p. 10
• Future work will explore RL strategies that jointly select, rank, and budget KG nodes while remaining efficient for large-scale EHR prediction.Source: §7 · p. 10

8 Ethical Considerations

• The experiments use publicly available, de-identified EHR datasets, which provide strong privacy protections.Source: §8 · p. 10
• No patient-level records or identifiable clinical histories are sent to external LLM services; the REFINE LLM runs locally and remains frozen.LLM-based KG construction uses concept-level information and aggregated EHR-derived statistics.Source: §8 · p. 10
• REFINE is intended for EHR representation-learning research, not clinical decision-making without prospective validation, fairness evaluation, and expert oversight.Source: §8 · p. 10

A.1 Global Text-Attributed KG Construction

• The work reuses a clinically validated relation-induction protocol and does not introduce a new KG induction pipeline.Source: §A.1 · p. 12
• The global graph represents diagnosis, medication, and procedure codes as typed nodes, with directed semantic-relation edges and textual descriptions and rationales.Source: §A.1 · p. 12
• Candidate code pairs come from within-visit co-occurrences and next-visit transitions; support, conditional probability, smoothed PMI, and χ2 statistics are computed using training-fold data.Source: §A.1 · p. 12
• Pairs with insufficient support, weak association, or non-significant dependence are filtered before type-constrained LLM relation induction.Source: §A.1 · p. 12
• The LLM assigns directed relations from type-specific allowed sets and produces rationales; ambiguous pairs can receive conservative abstention labels.Source: §A.1 · pp. 12–13
• Node descriptions supply concept semantics and edge rationales supply relational text; inverse computational edges support bidirectional GNN message passing without changing semantic direction.These attributes are used to initialize or generate graph-contextualized code and KG soft prompts in patient-specific temporal graphs.Source: §A.1 · p. 13

A.2 LLM-Based Node Description Construction

• A type-specific LLM prompt generates a concise clinical description for each diagnosis, medication, or procedure code.Source: §A.2 · p. 13
• The prompts tailor description content to each code type and exclude patient-level facts, URLs, dosage thresholds, hospital policies, and uncertain claims.Source: §A.2 · p. 13
• Parsed descriptions become node text attributes that provide fixed clinical semantics for initializing code representations and constructing graph-aware soft prompts.Together with edge rationales, they make the graph relation-typed and text-attributed.Source: §A.2 · p. 13

A.3 Implementation Details

• Training separates RL graph budgeting with a GNN-only proxy and cached patient graphs from full REFINE training with a frozen LLM and trainable components.The action grid varies KG depth and per-hop fanout and includes a no-expansion action.Source: §A.3 · pp. 14–15

A.4 KG-Context Pooling and Prompt Design

• KGPool summarizes each selected neighborhood as one context token, while triplet prompts increase prompt length, memory and computation costs, and perform worse.The comparison motivates KGPool as a balance between relational expressiveness and prompt efficiency.Source: §A.4 · p. 15

A.5 Value-of-Information KG Selection

• VoI scores hop–fanout actions against a no-expansion baseline using sparse label-derived targets, then selects independently by visit without labels at inference, unlike RL’s sequential code-root decisions.Source: §A.5 · p. 15

A.6 RQ5: Data Insufficiency Analysis

• REFINE improves AUPRC and Acc@20 across training-set ratios, datasets, and evaluated backbones, with gains especially meaningful in low-data regimes.The gains also persist at full data; Figure 6 compares base and REFINE-enhanced backbones.Source: Figure 6 · p. 14, §A.6 · p. 16

A.7 Sensitivity to LLM Backbone

• Replacing Llama-3.2-1B with Qwen2.5-0.5B moderately decreases AUPRC on both datasets, but performance remains competitive; larger backbones remain untested.Figure 7 reports the comparison, which is limited to relatively small models.Source: §A.7 · p. 16, Figure 7 · p. 16