ProtoEHR: Hierarchical Prototype Learning in EHR

• ProtoEHR is a hierarchical, prototype-based framework that models EHR data as codes, visits, and patient sequences to boost interpretability and predictive accuracy.
• It employs an LLM-constructed medical knowledge graph and CompGCN to generate robust code-level embeddings and extract meaningful relations.
• The framework’s multi-level prototype infusion and hierarchical fusion significantly improve prediction tasks such as mortality, readmission, and length-of-stay.

Searching arXiv for ProtoEHR and closely related prototype-based EHR papers. Searching for "ProtoEHR hierarchical prototype learning EHR". Searching arXiv for prototype-based EHR modeling and ProtoEHR. ProtoEHR is a hierarchical, prototype-based framework for prediction from electronic health records (EHRs) that treats EHR data as a three-level structure—medical codes within visits, and visits within patients—and combines a large-language-model-constructed medical knowledge graph with prototype learning at each hierarchy level (Cai et al., 23 Aug 2025). It was proposed for EHR-based healthcare prediction tasks including mortality prediction, readmission prediction, length-of-stay prediction, drug recommendation, and phenotype prediction, with the stated aim of improving both predictive performance and interpretability by modeling cross-level hierarchy and within-level similarity jointly (Cai et al., 23 Aug 2025).

1. Conceptual scope and problem formulation

ProtoEHR starts from the observation that existing EHR methods often underuse the full structure of EHRs. Some methods emphasize temporal patient trajectories, some emphasize code relations, and some emphasize patient similarity, but this isolated treatment of only one aspect of EHR data is presented as limiting both predictive performance and interpretability (Cai et al., 23 Aug 2025). ProtoEHR addresses this by jointly modeling code-level, visit-level, and patient-level structure, while also introducing prototypes at all three levels rather than only at the patient level.

The formal dataset definition is:

$\mathcal D = (\mathcal C, \mathcal V),$

where $\mathcal C = \{c_i\}_{i=1}^{N_{\text{code}}}$ is the set of all unique medical codes and $\mathcal V = \{\mathcal V_i\}_{i=1}^{N}$ is the set of patient visit sequences. For patient $i$,

$$\mathcal V_i = (\mathcal V_{i,1}, \cdots, \mathcal V_{i,|\mathcal V_i|}),$$

and each visit is a set of codes

$$\mathcal V_{i,j} = \{c_{i,j,k}\in\mathcal C\}_{k=1}^{|\mathcal V_{i,j}|}.$$

The downstream prediction problem is written as

$\hat y_i = f(\mathcal D, i, \mathcal V_i).$

Five prediction tasks are evaluated. Mortality prediction is binary and predicts whether the patient dies in the final visit and whether it occurs within 30 days from the previous visit. Readmission prediction is binary and predicts whether the final visit occurs within 30 days of the previous one. Length-of-stay prediction is a 10-class task with labels $y_i \in \{0,1,\dots,9\}$. Drug recommendation is a multi-label prediction problem with output $\hat y_i \in \mathbb R^{202}$, and phenotype prediction is a multi-label problem with output $\hat y_i \in \mathbb R^{26}$ (Cai et al., 23 Aug 2025).

The paper evaluates ProtoEHR on MIMIC-III and MIMIC-IV. After preprocessing, MIMIC-III contains 5,453 patients, 14,330 visits, and 657 codes, with 2.63 average visits per patient and 39.7 average codes per visit. MIMIC-IV contains 51,473 patients, 167,042 visits, and 708 codes, with 3.25 average visits per patient and 21.6 average codes per visit. Both datasets are split into train:validation:test = 6:2:2, and for MIMIC-III a sliding window is used to augment sample size for all tasks except mortality prediction (Cai et al., 23 Aug 2025).

2. Knowledge graph construction and code-level semantics

A distinctive component of ProtoEHR is its use of a dataset-specific medical knowledge graph as the code-level knowledge source. The graph is defined as

$\mathcal C = \{c_i\}_{i=1}^{N_{\text{code}}}$0

where $\mathcal C = \{c_i\}_{i=1}^{N_{\text{code}}}$1 is the entity set of medical codes, $\mathcal C = \{c_i\}_{i=1}^{N_{\text{code}}}$2 is the relation set, and

$\mathcal C = \{c_i\}_{i=1}^{N_{\text{code}}}$3

The stated motivation is that manually identifying meaningful relations among diagnosis, procedure, and prescription codes is difficult, and external biomedical knowledge graphs such as UMLS-KG may not align directly with the dataset vocabulary (Cai et al., 23 Aug 2025).

The knowledge-graph construction pipeline has three stages: retrieval, cleaning, and refinement. In retrieval, Llama3-70B is used to produce candidate relations for code pairs through

$\mathcal C = \{c_i\}_{i=1}^{N_{\text{code}}}$4

In cleaning, GPT-4 labels a subset of triplets as true or false through

$\mathcal C = \{c_i\}_{i=1}^{N_{\text{code}}}$5

and those labels are used to train a classifier that filters the remaining candidates. The appendix reports 398,555 candidate triplets, 30,000 GPT-4-labeled triplets, 14,378 GPT-4 positives, and 122,422 classifier positives from the remaining candidates at threshold $\mathcal C = \{c_i\}_{i=1}^{N_{\text{code}}}$6. The final graph is built from all GPT-verified positives plus the top 71,890 classifier-scored positives, yielding 86,268 triplets and 713 entities including one padding entity (Cai et al., 23 Aug 2025).

Refinement reduces lexical redundancy in relation phrases. The procedure embeds relations with BERT, clusters relation phrases using agglomerative clustering with Ward linkage, and then uses an open-source LLM to detect semantically contradictory clusters and refine them. This reduces the number of unique relation types from 2,330 to 269 (Cai et al., 23 Aug 2025).

The paper also reports a stability analysis of LLM-based relation extraction. Using 100 sampled codes, all code pairs, five Llama-70B extractions per pair, BERT relation embeddings, and cosine similarity across runs, the mean cosine similarity is reported as 0.93 with standard deviation 0.05. The paper interprets this as suggesting stable relation extraction (Cai et al., 23 Aug 2025).

At the representation-learning stage, the graph is processed by CompGCN. The code-level update is:

$\mathcal C = \{c_i\}_{i=1}^{N_{\text{code}}}$7

where $\mathcal C = \{c_i\}_{i=1}^{N_{\text{code}}}$8 is circular correlation. Directed edges are then reversed for the CompGCN layer. In ProtoEHR, this graph encoder provides global semantic context before prototype infusion (Cai et al., 23 Aug 2025).

3. Hierarchical architecture and prototype learning

ProtoEHR is organized as medical KG construction, hierarchical representation learning, and hierarchical fusion. The overall processing order is medical KG $\mathcal C = \{c_i\}_{i=1}^{N_{\text{code}}}$9 code encoder $\mathcal V = \{\mathcal V_i\}_{i=1}^{N}$0 visit encoder $\mathcal V = \{\mathcal V_i\}_{i=1}^{N}$1 patient encoder $\mathcal V = \{\mathcal V_i\}_{i=1}^{N}$2 hierarchical fusion $\mathcal V = \{\mathcal V_i\}_{i=1}^{N}$3 prediction head (Cai et al., 23 Aug 2025).

At the code level, CompGCN generates code representations from the medical KG, after which a code-level prototype encoder $\mathcal V = \{\mathcal V_i\}_{i=1}^{N}$4 enriches them. The output for the $\mathcal V = \{\mathcal V_i\}_{i=1}^{N}$5-th code in the $\mathcal V = \{\mathcal V_i\}_{i=1}^{N}$6-th visit of patient $\mathcal V = \{\mathcal V_i\}_{i=1}^{N}$7 is $\mathcal V = \{\mathcal V_i\}_{i=1}^{N}$8. At the visit level, the local encoder aggregates codes within a visit by average pooling:

$\mathcal V = \{\mathcal V_i\}_{i=1}^{N}$9

The resulting $i$0 is then passed into the visit-level prototype encoder $i$1 to obtain $i$2 (Cai et al., 23 Aug 2025).

At the patient level, ProtoEHR uses a Transformer-based encoder over the sequence of visit representations $i$3. Because recent visits are described as more influential for future prediction, the representation of the last visit after the Transformer encoder is taken as the prototype-free patient representation $i$4, which is then passed into the patient-level prototype encoder $i$5 to produce $i$6 (Cai et al., 23 Aug 2025).

The prototype mechanism is shared across levels. If object representations are

$i$7

and learnable prototypes are

$i$8

prototype learning updates prototypes via cross-attention:

$i$9

Prototype infusion then enriches each object representation:

$$\mathcal V_i = (\mathcal V_{i,1}, \cdots, \mathcal V_{i,|\mathcal V_i|}),$$0

with

$$\mathcal V_i = (\mathcal V_{i,1}, \cdots, \mathcal V_{i,|\mathcal V_i|}),$$1

Thus, the infusion stage uses cosine-similarity-based soft assignment, while the prototype-learning stage uses scaled dot-product attention (Cai et al., 23 Aug 2025).

ProtoEHR interprets prototypes as representative embeddings of shared structure at each hierarchy level: code prototypes capture intrinsic similarity among medical codes, visit prototypes capture shared visit patterns, and patient prototypes capture cohort structure and shared clinical trajectories. Prototypes are randomly initialized, globally learned at each level, and optimized through the downstream task loss. The paper does not introduce separate compactness, diversity, sparsity, or separation losses for prototype learning (Cai et al., 23 Aug 2025).

4. Hierarchical fusion, prediction, and optimization

After code-, visit-, and patient-level prototype learning, ProtoEHR fuses patient representation with prototype information from all three levels through a dedicated hierarchical fusion module. For each level $$\mathcal V_i = (\mathcal V_{i,1}, \cdots, \mathcal V_{i,|\mathcal V_i|}),$$2, cross-attention between the patient representation $$\mathcal V_i = (\mathcal V_{i,1}, \cdots, \mathcal V_{i,|\mathcal V_i|}),$$3 and the level-specific prototypes $$\mathcal V_i = (\mathcal V_{i,1}, \cdots, \mathcal V_{i,|\mathcal V_i|}),$$4 yields

$$\mathcal V_i = (\mathcal V_{i,1}, \cdots, \mathcal V_{i,|\mathcal V_i|}),$$5

The final patient representation is then

$$\mathcal V_i = (\mathcal V_{i,1}, \cdots, \mathcal V_{i,|\mathcal V_i|}),$$6

where

$$\mathcal V_i = (\mathcal V_{i,1}, \cdots, \mathcal V_{i,|\mathcal V_i|}),$$7

Prediction is made with a linear head:

$$\mathcal V_i = (\mathcal V_{i,1}, \cdots, \mathcal V_{i,|\mathcal V_i|}),$$8

The $$\mathcal V_i = (\mathcal V_{i,1}, \cdots, \mathcal V_{i,|\mathcal V_i|}),$$9 coefficients are later used as interpretability signals, because they quantify the relative contribution of code-, visit-, and patient-level prototype information (Cai et al., 23 Aug 2025).

For optimization, the paper states that ProtoEHR is trained using a task-specific prediction loss. Binary cross-entropy is used for binary and multi-label tasks, and cross-entropy is used for multi-class tasks. No auxiliary losses are specified for the knowledge graph, prototypes, or fusion module (Cai et al., 23 Aug 2025).

Implementation and selection details are reported through grid search. The search ranges are: CompGCN layers $$\mathcal V_{i,j} = \{c_{i,j,k}\in\mathcal C\}_{k=1}^{|\mathcal V_{i,j}|}.$$0, Transformer depth $$\mathcal V_{i,j} = \{c_{i,j,k}\in\mathcal C\}_{k=1}^{|\mathcal V_{i,j}|}.$$1, code-level prototypes $$\mathcal V_{i,j} = \{c_{i,j,k}\in\mathcal C\}_{k=1}^{|\mathcal V_{i,j}|}.$$2, visit-level prototypes $$\mathcal V_{i,j} = \{c_{i,j,k}\in\mathcal C\}_{k=1}^{|\mathcal V_{i,j}|}.$$3, patient-level prototypes $$\mathcal V_{i,j} = \{c_{i,j,k}\in\mathcal C\}_{k=1}^{|\mathcal V_{i,j}|}.$$4, dropout $$\mathcal V_{i,j} = \{c_{i,j,k}\in\mathcal C\}_{k=1}^{|\mathcal V_{i,j}|}.$$5, and learning rate $$\mathcal V_{i,j} = \{c_{i,j,k}\in\mathcal C\}_{k=1}^{|\mathcal V_{i,j}|}.$$6. The optimizer is Adam, the scheduler is ExponentialLR, and early stopping uses patience 20, with validation AUPRC for all tasks except length-of-stay, for which validation AUROC is used (Cai et al., 23 Aug 2025).

5. Empirical performance and interpretability

ProtoEHR is compared against Deepr, AdaCare, GRASP, StageNet, GraphCare, and KerPrint. For baselines that use a knowledge graph, the same medical KG is used for consistency. The evaluation metrics are AUPRC and AUROC for mortality and readmission, AUROC and F1 for length-of-stay, and AUPRC, AUROC, and F1 for drug recommendation and phenotype prediction. Reported statistics are means and standard deviations from 100 bootstrap runs (Cai et al., 23 Aug 2025).

The paper states that ProtoEHR achieves best or second-best performance on 24/24 metrics across all five tasks and two datasets. A compact summary of the reported ProtoEHR results is as follows.

Task	MIMIC-III	MIMIC-IV
Mortality	AUPRC 11.1 ± 3.7; AUROC 71.4 ± 3.1	AUPRC 17.0 ± 3.3; AUROC 89.2 ± 1.3
Readmission	AUPRC 33.0 ± 2.5; AUROC 61.7 ± 2.2	AUPRC 54.5 ± 0.9; AUROC 70.9 ± 0.5
Length of stay	AUROC 75.3 ± 0.9; F1 23.6 ± 1.2	AUROC 83.2 ± 0.2; F1 31.9 ± 0.5
Drug recommendation	AUPRC 70.6 ± 0.6; AUROC 91.7 ± 0.3; F1 54.4 ± 0.8	AUPRC 72.4 ± 0.2; AUROC 95.8 ± 0.1; F1 60.0 ± 0.2
Phenotype prediction	AUPRC 68.2 ± 0.7; AUROC 86.1 ± 0.4	AUPRC 78.5 ± 0.3; AUROC 93.2 ± 0.1; F1 64.2 ± 0.4

For mortality prediction, the paper highlights a 20.7% AUPRC improvement on MIMIC-III and a 17.2% AUPRC improvement on MIMIC-IV relative to the best competing baselines. For length-of-stay prediction on MIMIC-IV, it highlights a 7.4% F1 improvement. Readmission is the main exception to outright dominance on every metric: on MIMIC-III, ProtoEHR has second-best AUPRC to StageNet but best AUROC (Cai et al., 23 Aug 2025).

The ablation study on MIMIC-IV removes the medical KG, code prototypes, visit prototypes, patient prototypes, and hierarchical fusion. For mortality, the full model yields AUPRC 17.0 and AUROC 89.2. Removing the medical KG reduces performance to AUPRC 8.9 and AUROC 86.7; removing code prototypes yields AUPRC 16.1 and AUROC 88.8; removing visit prototypes yields AUPRC 13.5 and AUROC 86.8; removing patient prototypes yields AUPRC 13.6 and AUROC 86.3; and removing hierarchical fusion yields AUPRC 14.0 and AUROC 88.3. The paper emphasizes that removing the medical KG causes a 47.6% decrease in mortality AUPRC, while removing code, visit, and patient prototypes causes decreases of 5.3%, 20.6%, and 20.0%, respectively (Cai et al., 23 Aug 2025).

A separate KG edge-type ablation shows that removing procedure-procedure edges has the strongest effect among same-type edge removals, especially for mortality with a 30% decrease in AUPRC, while removing diagnosis-procedure edges has the greatest impact on three of five tasks among cross-type removals. The paper interprets this as evidence that procedure-procedure links encode severity and co-occurring interventions, while diagnosis-procedure links reveal motivations behind procedures and help infer diagnoses and phenotypes (Cai et al., 23 Aug 2025).

Interpretability is analyzed in two complementary ways. First, the hierarchical fusion weights $$\mathcal V_{i,j} = \{c_{i,j,k}\in\mathcal C\}_{k=1}^{|\mathcal V_{i,j}|}.$$7 reveal level contributions. The paper reports that patient-level prototypes are most important for mortality prediction, while visit-level information contributes more to phenotype prediction than to drug recommendation. To support that interpretation, it reports average Jaccard similarity between previous-visit labels and task labels of 0.287 for drug recommendation and 0.439 for phenotype prediction (Cai et al., 23 Aug 2025).

Second, the paper analyzes prototype importance directly. For length-of-stay prediction on MIMIC-IV, patient prototype ID 0 increases in importance from 0.0034 to 0.62 as stay duration increases; its associated diagnosis and procedure codes include Complications of Surgical Procedures or Medical Care, Other Vascular Catheterization; Not Heart, Extracorporeal Circulation Auxiliary to Open Heart Procedures, and Coronary Artery Bypass Graft. Patient prototype ID 13 decreases in importance from 0.18 to 0.0097 as stay duration increases and is associated with Nonspecific Chest Pain, Alcohol-related Disorders, Routine Chest X-ray, Electrocardiogram, and Other Diagnostic Procedures. The paper interprets the first as reflecting severe or invasive clinical situations and the second as reflecting milder or more routine cases (Cai et al., 23 Aug 2025).

Prototype effectiveness is further quantified by clustering prototype-importance vectors with K-Means and computing silhouette scores. On MIMIC-IV mortality, the silhouette scores are 0.49 for code, 0.73 for visit, and 0.81 for patient prototypes. On MIMIC-IV length-of-stay, the scores are 0.56, 0.56, and 0.54, respectively. The paper interprets this as showing especially strong organization by patient prototypes for mortality, whereas length-of-stay depends more evenly on all three levels (Cai et al., 23 Aug 2025).

6. Relation to adjacent prototype-based EHR research and limitations

ProtoEHR belongs to a broader family of prototype-based EHR methods, but its specific design center differs from earlier patient-level or retrieval-guided approaches. The Progressive Prototypical Network selects typical patients as prototypes and uses progressive prototype memory for sparse-record risk prediction, but its main prototype mechanism is centered on patient-level representations (Yu et al., 2023). EHR-RAGp introduces learned prototypes as latent anchors for retrieval refinement over long EHR histories, yet its prototypes mediate retrieval relevance rather than functioning as a fully hierarchical prototype system (Shurrab et al., 12 May 2026). By contrast, ProtoEHR explicitly learns prototypes at the code, visit, and patient levels and fuses all three in the final predictor (Cai et al., 23 Aug 2025).

This distinction is also relevant to interpretability. In PPN, interpretability is strongly case-based through typical patients and prototype cohorts (Yu et al., 2023). In EHR-RAGp, interpretability mainly appears through prototype-weighted retrieval behavior over historical chunks (Shurrab et al., 12 May 2026). ProtoEHR instead provides level-specific explanations: it exposes whether prediction is being driven mainly by code-level semantics, visit structure, or patient-level cohort information, and it links influential prototypes to associated diagnosis and procedure patterns (Cai et al., 23 Aug 2025).

A separate neighboring line of work addresses heterogeneity through code-agnostic representation learning. DescEmb replaces hospital-specific code embeddings with description-based event embeddings and improves zero-shot transfer and pooled multi-hospital learning without explicit code harmonization (Hur et al., 2021). ProtoEHR does not pursue that route directly. A plausible implication is that ProtoEHR and description-based embedding are complementary rather than mutually exclusive: ProtoEHR focuses on hierarchical prototype learning and LLM-derived code relations, whereas DescEmb addresses cross-site vocabulary shift through text-grounded input representations.

The stated limitations of ProtoEHR concern both data scope and knowledge construction. The paper notes that future work is needed to improve medical KG construction and learning, especially to better capture correlations between medical codes and task labels. It also states that the current model does not incorporate additional clinical information such as laboratory test results, and that more validation on a more diverse set of datasets is needed (Cai et al., 23 Aug 2025). The setup also leaves the method dependent on the quality and coverage of LLM-extracted relations, and cross-institutional generalization is not tested beyond MIMIC-III and MIMIC-IV. The paper does not discuss scalability or memory complexity in detail. These constraints do not negate the framework’s contribution, but they delimit its present empirical scope.

Within prototype-oriented EHR modeling, ProtoEHR’s main significance lies in formalizing hierarchical prototype learning across the natural code-visit-patient nesting of longitudinal records, grounding code-level representations in an LLM-constructed knowledge graph, and exposing interpretable task-specific contributions from each hierarchy level (Cai et al., 23 Aug 2025). This suggests a shift from prototype use as only patient similarity or retrieval alignment toward a broader view in which prototypes are structural objects distributed throughout the EHR hierarchy.