Multi-Ontology Integration with Dual-Axis Propagation for Medical Concept Representation

TL;DR

This paper introduces LINKO, a novel framework that integrates multiple medical ontologies using dual-axis knowledge propagation and LLM-augmented initialization to improve medical concept representations in EHRs.

Contribution

The paper presents a unified approach for cross-ontology medical concept learning by combining LLM-based initialization with dual-axis knowledge propagation, addressing limitations of prior single-ontology methods.

Findings

LINKO outperforms state-of-the-art baselines in experiments.

Enhanced robustness in limited data and rare disease scenarios.

Effective integration of multiple ontologies improves concept representations.

Abstract

Medical ontology graphs map external knowledge to medical codes in electronic health records via structured relationships. By leveraging domain-approved connections (e.g., parent-child), predictive models can generate richer medical concept representations by incorporating contextual information from related concepts. However, existing literature primarily focuses on incorporating domain knowledge from a single ontology system, or from multiple ontology systems (e.g., diseases, drugs, and procedures) in isolation, without integrating them into a unified learning structure. Consequently, concept representation learning often remains limited to intra-ontology relationships, overlooking cross-ontology connections. In this paper, we propose LINKO, a large language model (LLM)-augmented integrative ontology learning framework that leverages multiple ontology graphs simultaneously by enabling dual-axis knowledge propagation both within and across heterogeneous ontology systems to enhance medical concept representation learning. Specifically, LINKO first employs LLMs to provide a graph-retrieval-augmented initialization for ontology concept embedding, through an engineered prompt that includes concept descriptions, and is further augmented with ontology context. Second, our method jointly learns the medical concepts in diverse ontology graphs by performing knowledge propagation in two axes: (1) intra-ontology vertical propagation across hierarchical ontology levels and (2) inter-ontology horizontal propagation within every level in parallel. Last, through extensive experiments on two public datasets, we validate the superior performance of LINKO over state-of-the-art baselines. As a plug-in encoder compatible with existing EHR predictive models, LINKO further demonstrates enhanced robustness in scenarios involving limited data availability and rare disease prediction.

Tables

Table 1. Table 2. Performance comparison and ablation study on MIMIC-III and MIMIC-IV. General Performance includes AUPRC, F1 score, and Acc@k. Diagnosis labels are grouped into four frequency bands (0–25%, 25–50%, 50–75%, 75–100%), and corresponding AUPRC scores are reported under Label Category Performance. The first row shows the base Transformer model, while subsequent rows represent variants with added medical concept encoders (e.g., GRAM = Base + GRAM). Reported values are means with 95% confidence intervals.

	Model	General Performance					Label Category Performance (AUPRC)
		AUPRC	F1	Acc@15	Acc@20	Acc@30	0-25%	25-50%	50-75%	75-100%
MIMIC-IV	Base	${28.54}_{\pm 0.39}$	${22.90}_{\pm 0.10}$	${37.88}_{\pm 0.53}$	${39.53}_{\pm 0.60}$	${44.10}_{\pm 0.52}$	${28.33}_{\pm 0.44}$	${58.73}_{\pm 0.86}$	${46.54}_{\pm 0.30}$	${67.55}_{\pm 0.18}$
	GRAM	${29.96}_{\pm 0.45}$	${24.17}_{\pm 0.18}$	${39.42}_{\pm 0.25}$	${41.23}_{\pm 0.48}$	${45.90}_{\pm 0.61}$	${29.73}_{\pm 0.53}$	${60.62}_{\pm 0.27}$	${47.67}_{\pm 0.27}$	${68.45}_{\pm 0.21}$
	MMORE	${30.06}_{\pm 0.25}$	${25.11}_{\pm 1.60}$	${39.56}_{\pm 0.18}$	${41.46}_{\pm 0.32}$	${46.19}_{\pm 0.34}$	${29.84}_{\pm 0.22}$	${60.60}_{\pm 0.37}$	${48.84}_{\pm 0.34}$	${68.07}_{\pm 1.3}$
	KAME	${29.13}_{\pm 0.32}$	${23.39}_{\pm 0.32}$	${39.18}_{\pm 0.32}$	${40.28}_{\pm 0.32}$	${45.67}_{\pm 0.25}$	${28.84}_{\pm 0.32}$	${59.62}_{\pm 0.32}$	${47.21}_{\pm 0.42}$	${68.08}_{\pm 0.32}$
	G-BERT	${30.13}_{\pm 0.36}$	${25.19}_{\pm 0.33}$	${39.65}_{\pm 0.73}$	${41.66}_{\pm 0.81}$	${46.29}_{\pm 0.95}$	${29.91}_{\pm 0.59}$	${60.83}_{\pm 0.69}$	${49.88}_{\pm 0.74}$	${69.03}_{\pm 1.0}$
	HAP	${30.02}_{\pm 0.23}$	${24.22}_{\pm 1.30}$	${39.55}_{\pm 0.23}$	${41.37}_{\pm 0.37}$	${46.22}_{\pm 0.33}$	${29.64}_{\pm 0.34}$	${60.76}_{\pm 0.35}$	${49.53}_{\pm 0.33}$	${68.94}_{\pm 1.4}$
	ADORE	${30.11}_{\pm 0.34}$	${25.21}_{\pm 0.87}$	${39.87}_{\pm 0.24}$	${41.51}_{\pm 0.42}$	${46.25}_{\pm 0.34}$	${30.15}_{\pm 0.30}$	${60.89}_{\pm 0.45}$	${49.44}_{\pm 1.02}$	${68.18}_{\pm 0.93}$
	KAMPNet	${30.21}_{\pm 0.23}$	${25.32}_{\pm 1.10}$	${40.18}_{\pm 0.42}$	${41.88}_{\pm 0.54}$	${46.41}_{\pm 0.56}$	${30.21}_{\pm 0.34}$	${61.05}_{\pm 0.89}$	${48.93}_{\pm 0.53}$	${70.30}_{\pm 1.0}$
	${\text{LINKO}}_{\text{w/ GAT}}$	${32.13}_{\pm 0.12}$	${28.56}_{\pm 0.09}$	${42.60}_{\pm 0.14}$	${43.56}_{\pm 0.14}$	${48.12}_{\pm 0.18}$	${31.83}_{\pm 0.13}$	${63.03}_{\pm 0.26}$	${50.96}_{\pm 0.40}$	${71.35}_{\pm 0.29}$
	${\text{LINKO}}_{\text{w/ HAT}}$	${32.38}_{\pm 0.32}$	${30.02}_{\pm 0.13}$	${42.05}_{\pm 0.65}$	${43.70}_{\pm 0.66}$	${48.63}_{\pm 0.73}$	${32.31}_{\pm 0.41}$	${63.25}_{\pm 0.73}$	${51.17}_{\pm 0.23}$	${71.58}_{\pm 0.12}$
	w/o HGIP	${30.65}_{\pm 0.28}$	${26.19}_{\pm 0.22}$	${40.16}_{\pm 0.20}$	${42.03}_{\pm 0.21}$	${46.81}_{\pm 0.40}$	${30.36}_{\pm 0.45}$	${61.34}_{\pm 0.13}$	${49.15}_{\pm 0.19}$	${68.63}_{\pm 0.55}$
	w/o LLM	${30.86}_{\pm 0.59}$	${26.69}_{\pm 0.52}$	${40.40}_{\pm 0.83}$	${42.06}_{\pm 0.75}$	${46.54}_{\pm 0.88}$	${30.66}_{\pm 0.68}$	${61.42}_{\pm 0.10}$	${48.35}_{\pm 0.18}$	${67.10}_{\pm 0.32}$
	w/o leaf-HMP	${30.70}_{\pm 0.33}$	${26.89}_{\pm 0.88}$	${40.27}_{\pm 0.43}$	${42.09}_{\pm 0.38}$	${46.76}_{\pm 0.60}$	${30.41}_{\pm 0.45}$	${61.30}_{\pm 0.99}$	${48.77}_{\pm 0.79}$	${69.46}_{\pm 0.99}$
	w/o parent-HMP	${30.32}_{\pm 0.69}$	${25.49}_{\pm 0.25}$	${39.83}_{\pm 0.59}$	${41.72}_{\pm 0.92}$	${46.47}_{\pm 0.90}$	${30.07}_{\pm 0.78}$	${61.40}_{\pm 0.72}$	${46.46}_{\pm 0.23}$	${68.06}_{\pm 0.29}$
	w/o HMP	${30.10}_{\pm 0.39}$	${25.26}_{\pm 0.88}$	${39.67}_{\pm 0.41}$	${41.50}_{\pm 0.34}$	${46.25}_{\pm 0.53}$	${29.85}_{\pm 0.51}$	${60.94}_{\pm 0.75}$	${47.39}_{\pm 0.28}$	${67.33}_{\pm 0.62}$
MIMIC-III	Base	${28.23}_{\pm 0.24}$	${22.36}_{\pm 0.33}$	${36.15}_{\pm 0.22}$	${38.03}_{\pm 0.34}$	${43.19}_{\pm 0.43}$	${28.07}_{\pm 0.28}$	${54.28}_{\pm 0.10}$	${49.89}_{\pm 0.29}$	${73.34}_{\pm 0.26}$
	GRAM	${28.99}_{\pm 0.34}$	${24.10}_{\pm 0.48}$	${37.17}_{\pm 0.36}$	${39.13}_{\pm 0.52}$	${44.32}_{\pm 0.51}$	${28.79}_{\pm 0.41}$	${54.81}_{\pm 1.3}$	${50.58}_{\pm 0.34}$	${74.44}_{\pm 0.77}$
	MMORE	${29.11}_{\pm 0.38}$	${23.67}_{\pm 0.70}$	${37.15}_{\pm 0.46}$	${39.14}_{\pm 0.53}$	${44.36}_{\pm 0.50}$	${28.87}_{\pm 0.43}$	${54.92}_{\pm 0.87}$	${51.29}_{\pm 0.25}$	${74.42}_{\pm 0.30}$
	KAME	${28.52}_{\pm 0.32}$	${23.18}_{\pm 0.09}$	${36.76}_{\pm 0.47}$	${38.36}_{\pm 0.45}$	${43.75}_{\pm 0.36}$	${28.19}_{\pm 0.38}$	${55.13}_{\pm 0.16}$	${50.09}_{\pm 0.32}$	${73.79}_{\pm 0.23}$
	G-BERT	${29.22}_{\pm 0.81}$	${24.92}_{\pm 0.83}$	${37.36}_{\pm 0.12}$	${39.55}_{\pm 0.12}$	${45.15}_{\pm 0.44}$	${29.16}_{\pm 0.10}$	${54.47}_{\pm 0.99}$	${52.45}_{\pm 0.62}$	${75.40}_{\pm 0.18}$
	HAP	${29.28}_{\pm 0.41}$	${23.25}_{\pm 0.88}$	${37.47}_{\pm 0.56}$	${39.64}_{\pm 0.55}$	${45.05}_{\pm 0.57}$	${29.10}_{\pm 0.45}$	${55.11}_{\pm 0.99}$	${52.19}_{\pm 0.27}$	${75.50}_{\pm 0.24}$
	ADORE	${29.31}_{\pm 0.29}$	${23.61}_{\pm 0.81}$	${37.53}_{\pm 0.31}$	${39.68}_{\pm 0.32}$	${45.01}_{\pm 0.58}$	${29.18}_{\pm 0.73}$	${55.14}_{\pm 0.42}$	${52.14}_{\pm 0.45}$	${75.41}_{\pm 0.27}$
	KAMPNet	${29.38}_{\pm 0.65}$	${25.01}_{\pm 0.72}$	${37.83}_{\pm 0.39}$	${39.91}_{\pm 0.61}$	${45.26}_{\pm 0.53}$	${29.31}_{\pm 0.45}$	${55.20}_{\pm 0.87}$	${52.56}_{\pm 0.34}$	${75.91}_{\pm 0.29}$
	${\text{LINKO}}_{\text{w/ HAT}}$	${30.78}_{\pm 0.99}$	${27.67}_{\pm 0.23}$	${39.11}_{\pm 0.92}$	${41.18}_{\pm 0.73}$	${46.65}_{\pm 0.83}$	${30.72}_{\pm 0.85}$	${54.74}_{\pm 0.19}$	${57.16}_{\pm 0.26}$	${79.23}_{\pm 0.10}$
	${\text{LINKO}}_{\text{w/ GAT}}$	${31.79}_{\pm 0.11}$	${28.66}_{\pm 0.18}$	${39.95}_{\pm 0.11}$	${41.84}_{\pm 0.14}$	${46.96}_{\pm 0.96}$	${31.62}_{\pm 1.09}$	${55.68}_{\pm 0.15}$	${56.90}_{\pm 0.59}$	${80.76}_{\pm 0.47}$
	w/o HGIP	${30.25}_{\pm 0.37}$	${25.75}_{\pm 0.14}$	${38.54}_{\pm 0.51}$	${40.47}_{\pm 0.55}$	${46.00}_{\pm 0.66}$	${30.04}_{\pm 0.44}$	${55.61}_{\pm 0.12}$	${53.36}_{\pm 0.31}$	${77.00}_{\pm 0.22}$
	w/o LLM	${30.38}_{\pm 0.10}$	${25.88}_{\pm 0.90}$	${38.23}_{\pm 0.57}$	${40.38}_{\pm 0.66}$	${45.82}_{\pm 0.90}$	${30.33}_{\pm 0.10}$	${55.05}_{\pm 0.99}$	${55.03}_{\pm 0.64}$	${79.16}_{\pm 0.33}$
	w/o leaf-HMP	${30.09}_{\pm 0.43}$	${25.44}_{\pm 0.15}$	${38.25}_{\pm 0.54}$	${40.35}_{\pm 0.54}$	${45.92}_{\pm 0.59}$	${29.88}_{\pm 0.51}$	${55.57}_{\pm 0.57}$	${53.62}_{\pm 0.30}$	${76.68}_{\pm 0.17}$
	w/o parent-HMP	${29.76}_{\pm 0.43}$	${26.02}_{\pm 0.10}$	${37.97}_{\pm 0.39}$	${39.93}_{\pm 0.56}$	${45.43}_{\pm 0.52}$	${29.53}_{\pm 0.48}$	${55.24}_{\pm 0.88}$	${52.48}_{\pm 0.26}$	${75.78}_{\pm 0.17}$
	w/o HMP	${29.41}_{\pm 0.40}$	${25.28}_{\pm 0.88}$	${37.55}_{\pm 0.41}$	${39.61}_{\pm 0.52}$	${45.25}_{\pm 0.60}$	${29.26}_{\pm 0.50}$	${54.52}_{\pm 0.63}$	${52.58}_{\pm 0.20}$	${75.46}_{\pm 0.21}$

Table 2. Table 3. Performance analysis of different prompting strategies on MIMIC-III and MIMIC-IV. The reported values include means and 95% confidence intervals.

Prompt Information	MIMIC-IV					MIMIC-III
	AUPRC	F1	Acc@15	Acc@20	Acc@30	AUPRC	F1	Acc@15	Acc@20	Acc@30
no LLM	${30.86}_{\pm 0.59}$	${26.69}_{\pm 0.52}$	${40.40}_{\pm 0.83}$	${42.06}_{\pm 0.75}$	${46.54}_{\pm 0.88}$	${30.38}_{\pm 0.10}$	${25.88}_{\pm 0.90}$	${38.23}_{\pm 0.57}$	${40.38}_{\pm 0.66}$	${45.82}_{\pm 0.90}$
type-code-noise	${31.28}_{\pm 0.10}$	${28.46}_{\pm 0.30}$	${40.99}_{\pm 0.75}$	${43.25}_{\pm 0.10}$	${47.39}_{\pm 0.11}$	${30.57}_{\pm 0.74}$	${27.88}_{\pm 0.21}$	${38.99}_{\pm 0.46}$	${40.96}_{\pm 0.57}$	${46.57}_{\pm 0.63}$
type-code	${31.96}_{\pm 0.40}$	${27.77}_{\pm 0.67}$	${41.57}_{\pm 0.35}$	${43.24}_{\pm 0.42}$	${47.92}_{\pm 0.50}$	${31.15}_{\pm 0.98}$	${27.89}_{\pm 0.10}$	${39.28}_{\pm 0.69}$	${41.26}_{\pm 0.44}$	${46.79}_{\pm 0.53}$
type-code-concept	${32.00}_{\pm 0.25}$	${29.07}_{\pm 0.56}$	${41.73}_{\pm 0.74}$	${43.52}_{\pm 0.83}$	${48.17}_{\pm 0.78}$	${31.18}_{\pm 1.00}$	${28.35}_{\pm 0.65}$	${39.39}_{\pm 0.91}$	${41.55}_{\pm 0.83}$	${46.90}_{\pm 0.94}$
type-code-concept-children	${32.05}_{\pm 1.00}$	${29.03}_{\pm 1.02}$	${41.57}_{\pm 0.11}$	${43.41}_{\pm 0.13}$	${47.94}_{\pm 0.14}$	${31.09}_{\pm 1.00}$	${27.67}_{\pm 1.02}$	${39.23}_{\pm 0.13}$	${41.18}_{\pm 0.10}$	${46.66}_{\pm 0.99}$
type-code-concept-parent	${32.35}_{\pm 0.50}$	${28.27}_{\pm 0.48}$	${42.01}_{\pm 0.73}$	${43.59}_{\pm 0.38}$	${48.49}_{\pm 0.41}$	${31.21}_{\pm 0.95}$	${28.48}_{\pm 0.29}$	${39.46}_{\pm 0.79}$	${41.64}_{\pm 0.10}$	${46.91}_{\pm 0.97}$
type-code-concept-parent-task	${32.38}_{\pm 0.32}$	${30.02}_{\pm 0.13}$	${42.05}_{\pm 0.65}$	${43.70}_{\pm 0.66}$	${48.63}_{\pm 0.73}$	${31.79}_{\pm 0.11}$	${28.66}_{\pm 0.18}$	${39.95}_{\pm 0.11}$	${41.84}_{\pm 0.14}$	${46.96}_{\pm 0.96}$
w/ LLM-emb-freezed	${22.13}_{\pm 0.77}$	${13.35}_{\pm 0.89}$	${30.79}_{\pm 0.90}$	${32.91}_{\pm 0.92}$	${37.70}_{\pm 0.97}$	${25.75}_{\pm 0.37}$	${19.96}_{\pm 0.36}$	${33.82}_{\pm 0.39}$	${36.28}_{\pm 0.38}$	${42.00}_{\pm 0.38}$

Table 3. Table 4. Performance analysis of training on different combinations of medical concept types with/without Multi-level integration of ontologies on MIMIC-III and MIMIC-IV. The reported values include means and 95% confidence intervals.

	Concept Type	w/ Multi-level Integration					w/o Multi-level Integration
		AUPRC	F1	Acc@15	Acc@20	Acc@30	AUPRC	F1	Acc@15	Acc@20	Acc@30
MIMIC-IV	rx,px	${22.74}_{\pm 0.04}$	${16.92}_{\pm 0.17}$	${31.56}_{\pm 0.20}$	${33.57}_{\pm 0.27}$	${38.75}_{\pm 0.26}$	${21.35}_{\pm 0.27}$	${15.15}_{\pm 0.26}$	${30.03}_{\pm 0.53}$	${32.21}_{\pm 0.65}$	${37.23}_{\pm 0.39}$
	dx,px	${30.67}_{\pm 0.44}$	${25.87}_{\pm 0.13}$	${40.25}_{\pm 0.62}$	${42.21}_{\pm 0.59}$	${46.99}_{\pm 0.42}$	${29.28}_{\pm 0.31}$	${22.49}_{\pm 0.15}$	${38.62}_{\pm 0.31}$	${40.68}_{\pm 0.37}$	${45.56}_{\pm 0.24}$
	dx,rx	${30.86}_{\pm 0.10}$	${25.34}_{\pm 0.23}$	${40.44}_{\pm 0.10}$	${42.53}_{\pm 0.12}$	${47.31}_{\pm 0.10}$	${30.17}_{\pm 0.10}$	${25.08}_{\pm 0.14}$	${39.77}_{\pm 0.57}$	${41.71}_{\pm 0.70}$	${46.59}_{\pm 0.64}$
	dx,rx,px	${32.38}_{\pm 0.32}$	${30.02}_{\pm 0.13}$	${42.05}_{\pm 0.65}$	${43.70}_{\pm 0.66}$	${48.63}_{\pm 0.73}$	${30.10}_{\pm 0.39}$	${25.26}_{\pm 0.23}$	${39.67}_{\pm 0.41}$	${41.50}_{\pm 0.34}$	${46.25}_{\pm 0.53}$
MIMIC-III	rx,px	${24.24}_{\pm 0.97}$	${18.29}_{\pm 0.97}$	${31.92}_{\pm 0.86}$	${34.43}_{\pm 0.14}$	${39.87}_{\pm 0.13}$	${22.98}_{\pm 0.22}$	${16.55}_{\pm 0.14}$	${30.81}_{\pm 0.31}$	${32.95}_{\pm 0.49}$	${38.61}_{\pm 0.33}$
	dx,px	${30.19}_{\pm 0.11}$	${26.40}_{\pm 0.26}$	${38.03}_{\pm 0.81}$	${40.18}_{\pm 0.79}$	${46.03}_{\pm 0.11}$	${29.01}_{\pm 0.67}$	${24.48}_{\pm 0.10}$	${37.04}_{\pm 0.74}$	${39.16}_{\pm 0.61}$	${44.89}_{\pm 0.68}$
	dx,rx	${30.64}_{\pm 0.40}$	${27.06}_{\pm 0.15}$	${39.31}_{\pm 0.58}$	${41.27}_{\pm 0.64}$	${46.62}_{\pm 0.57}$	${29.56}_{\pm 0.38}$	${24.44}_{\pm 0.65}$	${37.88}_{\pm 0.61}$	${39.98}_{\pm 0.65}$	${45.46}_{\pm 0.60}$
	dx,rx,px	${31.79}_{\pm 0.11}$	${28.66}_{\pm 0.18}$	${39.95}_{\pm 0.11}$	${41.84}_{\pm 0.14}$	${46.96}_{\pm 0.96}$	${29.41}_{\pm 0.40}$	${25.28}_{\pm 0.88}$	${37.55}_{\pm 0.41}$	${39.61}_{\pm 0.52}$	${45.25}_{\pm 0.60}$

Equations

$$x_i^{(l)}=\mathcal{LLM}(\text{Prompt}(c_i^{(l)})),l=1,\dots ,L,i=1,\dots ,N^{(l)}$$

$$Q_{pq}^{(l)}=\sum _{c_i^{(L)}\in \mathcal{C}(p)}\sum _{c_j^{(L)}\in \mathcal{C}(q)}{Q}_{ij},l=1,\dots ,L$$

$${\mathbf{X}}_{(k+1)}^{(l)}={\text{GAT}}_{{\mathcal{G}}_h^{(l)}}\left({\mathbf{X}}_{(k)}^{(l)},{\mathcal{A}}_h^{(l)}\right),l=1,\dots ,L$$

$${\mathbf{X}}_{(k+1)}^{(L)}={\text{HAT}}_{{\mathcal{G}}_h^{(L)}}\left({\mathbf{X}}_{(k)}^{(L)},{\mathcal{H}}_h^{(L)}\right)$$

$$[{\mathbf{X}}_{(k+1)}^{(L-s)},{\mathbf{X}}_{(k+1)}^{(L-s+1)}]={\text{GAT}}_{{\mathcal{G}}_v^{(L-s)}}\left([{\mathbf{X}}_{(k)}^{(L-s)},{\mathbf{X}}_{(k)}^{(L-s+1)}],{\mathcal{A}}_v^{(L-s)}\right)$$

$$\text{GRAM}:z_i=\sum _{l=1}^L{\alpha }_{il}{\mathcal{P}}^l(x_i^{(L)}),{\alpha }_{il}\ge 0,\text{for}l=1,\dots ,L$$

$${\alpha }_{il}=\frac{\exp (f(x_i^{(L)},{\mathcal{P}}^l(x_i^{(L)}))}{{\sum }_{k=1}^L\exp (f(x_i^{(L)},{\mathcal{P}}^k(x_i^{(L)}))}$$

$$\mathbf{Z}=[{\mathbf{z}}_1,{\mathbf{z}}_2,\dots ,{\mathbf{z}}_{N^{(L)}}]$$

$${\mathbf{v}}_1,{\mathbf{v}}_2,\dots ,{\mathbf{v}}_t$$

$${\mathbf{h}}_p$$

$${\hat{\mathbf{y}}}_{t+1}$$

Full text

Representation Learning; Health Informatics; Knowledge Graphs; Ontology; Large Language Models; Graph-Augmented Prompting

††copyright: none††ccs: Computing methodologies Knowledge representation and reasoning††ccs: Applied computing Health informatics

1. Introduction

With the ubiquity of electronic health records (EHRs) in modern healthcare systems, developing machine learning models to analyze comprehensive medical histories has shown great potential in enhancing a wide range of predictive tasks (Choi et al., 2016a; Poulain and Beheshti, 2024; Jiang et al., 2023; Moghaddam et al., 2024; Nayebi Kerdabadi et al., 2023; Hadizadeh Moghaddam et al., 2025). Yet, the inherent complexity of medical concepts, i.e. EHR codes, characterized by diversity, sparsity, and temporal variability, poses a challenge for learning expressive and robust representations, particularly for infrequent codes, such as rare diseases.

A promising direction to this challenge is to incorporate domain knowledge into the concept representation learning, particularly the structured medical knowledge provided by ontology graphs. By capturing rich domain context of medical codes and their interrelationships, ontologies serve as valuable knowledge bases that support deeper interpretation of codes in EHRs (World Health Organization, 2025). Typically, these ontologies offer hierarchical classifications that represent medical concepts from general to specific definitions. Leveraging these hierarchies enables learning algorithms to better capture associations among medical codes, yielding more robust representations and more accurate predictions of rare concepts.

Accordingly, recent research has placed growing emphasis on augmenting EHR representations through the incorporation of supplementary ontology graphs (Choi et al., 2017; Shang et al., 2019; Ye et al., 2021; Zhang et al., 2020; An et al., 2023; Cheong et al., 2023). However, a critical limitation of current literature is the lack of a unified learning framework that can effectively accommodate and integrate multiple ontology systems. Specifically, existing methods typically treat each ontology as an isolated structure, supporting only vertical message passing (e.g., parent-child concept aggregation). As a result, concept representation learning is confined to intra-ontology relationships among homogeneous concepts (e.g., parent disease to child disease). To address this, there is a need to fuse multiple ontologies into a unified heterogeneous multi-knowledge graph that supports the incorporation of cross-ontology relationships (e.g., disease-drug, disease-procedure) into healthcare representation learning. These relationships should span all levels of the hierarchy, enabling horizontal message passing both within ontologies (e.g., parent disease to parent disease at the same level) and across ontologies (e.g., parent disease to parent drug, or child disease to child drug). Such a unified framework can enable comprehensive utilization of rich and diverse medical knowledge bases for more effective medical concept representation learning.

Moreover, large language models (LLMs), pre-trained on extensive corpora such as biomedical literature, textbooks, and websites, provide a rich source of domain-specific knowledge that, when carefully prompted, can significantly enhance various predictive tasks. A substantial body of research highlights their potential as knowledge bases (Petroni et al., 2019; Jiang et al., 2023; Xu et al., 2024; AlKhamissi et al., 2022). However, augmenting healthcare tasks with LLMs presents a significant challenge of factual inconsistency, that LLMs may generate inconsistent or misleading content in free-text responses (Huang et al., 2025; Zhang et al., 2023). In a high-stakes domain like healthcare, relying on such outputs for decision-making poses significant risks. This underscores the need for better strategies to leverage LLMs as expert knowledge sources in healthcare predictive modeling while ensuring reliability and minimizing risk.

To address the aforementioned gaps, we propose an integrative graph learning architecture, named LINKO: LLM-augmented INtegrative Knowledge Propagation over Ontology Graphs. LINKO enables knowledge propagation across multiple ontologies along both vertical and horizontal axes. Specifically, it facilitates (1) intra-ontology concept associations through a vertical knowledge propagation within each ontology graph, and (2) inter-ontology concept interactions via horizontal knowledge propagation across diverse ontology graphs. By integrating the dual-axis message passing in a unified structure, LINKO is advantageous in capturing both homogeneous and heterogeneous concept interactions at all levels of EHR concept granularity, from the coarsest (highest) to the finest (lowest) levels, therefore enabling a synergistic learning approach that integrates EHR mining (observational patterns) with ontology structure learning (external domain knowledge).

Furthermore, LINKO employs LLMs through a graph-retrieval-augmented initialization, where each concept and its local ontology context (e.g., connected concepts and their descriptions) form a prompt to retrieve dense embeddings from the LLM. These embeddings, which capture prior knowledge encoded in the LLM, initialize node representations in ontology graphs and are subsequently refined through dual-axis propagation. This use of LLM-derived embeddings serves solely as an augmentation step rather than a direct decision-making component (e.g., predictions, clustering) in EHR modeling. By encoding semantic concept information and then being refined within the graph, these embeddings enhance predictive performance while reducing the risk of hallucinations typically associated with free-text LLM outputs.

Lastly, our concept encoder can function as a plug-in module to existing healthcare predictive models, enhancing concept learning and improving performance. To validate our work, we conducted extensive experiments on two widely used EHR datasets, MIMIC-III and MIMIC-IV, performing sequential diagnosis prediction, plug-in enhancement evaluation, baseline comparisons, data insufficiency tests, and interpretative case studies. The results demonstrate that LINKO significantly improves the encoding phase of healthcare predictive models, leading to enhanced predictive performance.

2. Methodology

2.1. Notation and Problem Definition

EHR Dataset. Denoted by $\mathcal{D}=\{\mathcal{X}_{j}\mid j\in\mathcal{J}\}$, an EHR dataset consists of the medical histories from a collection of patients $\mathcal{J}$. For each patient, their history $\mathcal{X}_{j}$ consists of a sequence of $T_{j}$ clinical visits so that $\mathcal{X}_{j}=\{V_{j,t}\}_{t=1}^{T_{j}}$. Each visit $V_{j,t}$ is a set of $N_{t}$ medical codes thereby $V_{j,t}=\{c_{i}\}_{i=1}^{N_{t}}$, where each code $c_{i}$ represents a diagnosis (dx), a prescription (rx), or a procedure (px). Each medical code $c_{i}$ can also be associated with a descriptive name $\mathcal{S}(c_{i})$, which is typically a short text provided by the ontology. For brevity, we omit the subscripts $j$ and $t$ denoting each patient and visit.

Medical Ontology. A medical ontology is a hierarchical tree-like structure that organizes clinically related concepts from general categories at the upper levels, to specific concepts at the lower levels. Each group contains clinically related codes organized under a shared parent code, representing a more general concept. For different types of medical concepts (e.g., diseases, drugs, procedures, etc.), there exists a unique ontology (e.g., ICD (World Health Organization, 2025), ATC (World Health Organization Collaborating Centre for Drug Statistics Methodology, 2025)). A medical concept in an ontology is denoted as $c^{(l)}_{i}\in\bm{c}^{(l)}$, where $l\in[1:L]$ indexes the ontology level (from the highest to lowest), $i\in[1:N^{(l)}]$ indexes the concept at the $l$-th level with $N^{(l)}$ total number of code, and $\bm{c}^{(l)}$ represents the set of all concepts at level $l$. Last, we define the function $\mathcal{P}^{k}(c_{i}^{(l)})$, which maps a concept $c_{i}^{(l)}$ at level $l$ to its ancestor or descendant concepts at level $k$. If $k>l$, it returns the set of descendants; if $k<l$, it returns the ancestor; and if $k=l$, it returns the concept itself. Medical concepts (codes) in EHRs are typically located at the leaf level of the ontology ($\bm{c}=\bm{c}^{(L)}$).

Healthcare Predictive Tasks. Given a patient’s visit sequence $\mathcal{X}_{j}=\{V_{1},V_{2},\ldots,V_{T_{j}}\}$, the objective is to predict clinical outcomes. These may involve binary tasks (e.g., mortality, readmission) or multi-label classification (e.g., prescription recommendation, diagnosis prediction). In this paper, we focus on the latter–predicting diagnosis codes for the next visit $V_{T_{j}+1}$, a comprehensive task aimed at identifying potential diseases or conditions for future encounters.

2.2. Method Overview

We present an overview of our proposed method, LINKO, a hierarchically integrative multi-knowledge graph encoder designed for medical concept representation learning. This encoder can be added to different EHR predictive models to enhance medical concept representation and improve performance. The framework depicted in Figure 1 is summarized in three key steps:

Step 1: Formulating the Meta-KG denoted by $\mathcal{G}$, a multifaceted knowledge graph (KG) that joins multiple ontology graphs through every level of hierarchy. The initialization of node embeddings in Meta-KG is generated through a curated LLM prompting strategy and dense embedding retrieval. This Meta-KG consists of $L$ horizontal inter-ontology graphs, one at each level, and two vertical intra-ontology graphs per ontology (bottom-up and top-down).

Step 2: Carrying out multi-level Horizontal Message Passing (HMP) over inter-ontology concepts enabled by co-occurrence-based edges. We construct horizontal graphs in Meta-KG using two approaches: (1) regular graph structure, where edges are defined by concept co-occurrence probabilities; and (2) hypergraph structure, where each edge connects all concepts within a visit.

Step 3: Performing Vertical Message Passing (VMP) over intra-ontology concept edges derived from parent-child relationships in the ontology hierarchy. Within each ontology, the process begins with bottom-up propagation, where embeddings flow from child concepts to their parents, followed by top-down propagation, where updated embeddings are passed back to the children, yielding the final representations for leaf nodes.

2.3. Step 1: Graph-Augmented LLM Initialization

We construct Meta-KG nodes by collecting medical concepts across all hierarchical levels of each ontology. For each ontology, we first identify leaf nodes at the lowest level and then progressively map them to higher levels, thus defining nodes across the hierarchy. We use embedding $x^{(l)}_{i}\in\mathbb{R}^{d}$ to represent each medical concept $c^{(l)}_{i}$ in the Meta-KG. We extract the concept name associated with code $c^{(l)}_{i}$, denoted by $\mathcal{S}(c^{(l)}_{i})$. We develop a Graph-Augmented prompting strategy specifically tailored for EHRs to retrieve embeddings from LLMs for each concept. These embeddings are then used to initialize $x^{(l)}_{i}$, which is further refined through end-to-end graph learning. This approach leverages the invaluable knowledge of LLMs while minimizing the risk of hallucination, as we rely solely on their embeddings rather than their final natural language responses. We empirically found the following prompting strategy (Pr) most effective:

• •

For $l=1$, prompt: *¡¡ For the task of [task name], provide a semantic representation for [type] code [$c^{(l)}_{i}$] which represents [$\mathcal{S}(c^{(l)}_{i})$], a general medical concept. ¿¿ *

• •

For $l>1$, prompt: *¡¡ For the task of [task name], provide a semantic representation for [type] code [$c^{(l)}_{i}$] which represents [$\mathcal{S}(c^{(l)}_{i})$]. It falls under the broader [type] categories of [$\mathcal{P}^{l-1}(c^{(l)}_{i})$]

([$\mathcal{S}(\mathcal{P}^{l-1}(c^{(l)}_{i}))$]), …, [$\mathcal{P}^{1}(c^{(l)}_{i})$] ([$\mathcal{S}(\mathcal{P}^{1}(c^{(l)}_{i}))$]). ¿¿*

where “task name” can be any predictive task, e.g. diagnosis prediction, “type” refers to name of the ontology concept system (e.g., ICD-9 diagnosis/procedure, and ATC drug). For each code, we provide the code and its descriptive concept, followed by its broader categories or EHR ancestors and their descriptions, identified at higher levels of the ontology graph using the mapping function $\mathcal{P}$. An example of an LLM prompt for a diagnosis code is as follows: ICD-9 Diagnosis 250.7: Prompt: ¡¡ For the task of diagnosis prediction, provide a semantic representation for ICD-9 diagnosis code 250.7, which represents Diabetes with peripheral circulatory disorders. It is a specific medical concept under the broader ICD-9 Diagnosis categories of 250 (Diabetes mellitus), 249-259 (Diseases of Other Endocrine Glands), and 240-279 (Endocrine, Nutritional, and Metabolic Diseases, and Immunity Disorders).¿¿

We employ the OpenAI off-the-shelf model, GPT text-embedding-3-small (OpenAI, 2023), denoted as $\mathcal{LLM}$ to generate a semantic embedding, containing clinical knowledge and context background from LLMs. We initialize the vector representation $x^{(l)}_{i}\in\mathbb{R}^{d}$ as follows:

[TABLE]

2.4. Step 2: Horizontal Message Passing (HMP)

We formulate horizontal information propagation that links medical concepts within and across ontologies at all levels. Leveraging observational co-occurrence patterns in EHR visits, we establish edges between frequently co-occurring concepts and apply GNN-style message passing. Horizontal graph edges at higher hierarchy levels are established by linking ancestor concepts mapped from observed codes at the lowest level. This operation achieves three key goals: 1) hierarchically fusing diverse ontologies (diagnosis, drug, and procedure) to capture heterogeneous code interactions, 2) utilizing information at all levels of granularity for EHR representation. While leaf-level EHR codes offer detailed but sparse insights, mapping to higher-level concepts reduces the graph complexity, allowing us to leverage both fine-grained and coarse-grained information simultaneously for representation learning, and 3) enabling a unified approach that combines EHR mining (clinical statistics) with ontology-driven learning (defined expert knowledge).

To construct graph edges based on EHR co-occurrence information, we first create a leaf (child) level co-occurrence count matrix $Q^{(L)}\in\mathbb{R}^{N^{(L)}\times N^{(L)}}$, where $Q_{ij}$ denotes the number of occurrences of leaf code $c_{j}^{(L)}$ given the presence of leaf code $c_{i}^{(L)}$ within a visit. We then derive the co-occurrence matrices at higher levels by aggregating the co-occurrence counts of their children as follows:

[TABLE]

where $\mathcal{C}(p)=\mathcal{P}^{L}(c_{p}^{(l)})$ and $\mathcal{C}(q)=\mathcal{P}^{L}(c_{q}^{(l)})$ denote the sets of leaf-level children of parent-level nodes $c_{p}^{(l)}$ and $c_{q}^{(l)}$, respectively. Next, we derive the co-occurrence conditional probability matrix $P^{(l)}$ from the count matrix $Q^{(l)}$ by normalizing each entry with the total occurrences of the corresponding node $p$, expressed as: $P^{(l)}_{pq}=Q^{(l)}_{pq}/\sum_{j}Q^{(l)}_{pj},\quad\text{for }l=1:L$. The co-occurrence probability matrix is then used to define edges in the graph. Specifically, an edge between nodes $p$ and $q$ is included if the co-occurrence probability exceeds a threshold $\tau^{(l)}$. This binarization generates the adjacency matrix $\mathcal{A}_{h}^{(l)}$ from $P^{(l)}$ as: $\mathcal{A}_{h}^{(l)}=\mathbb{I}(P^{(l)}\geq\tau)$, where indicator function $\mathbb{I}(\cdot)$ returns 1 if true and 0 otherwise. Consequently, we construct the horizontal graphs $\mathcal{G}_{h}^{(l)}=(\mathcal{V}^{(l)},\mathcal{E}^{(l)})$, where $\mathcal{V}^{(l)}=\mathbf{c}^{(l)}$ and $\mathcal{E}^{(l)}=\mathcal{A}_{h}^{(l)}$, with $N^{(l)}$ nodes and $M^{(l)}$ edges at the $l$-th level of the ontology. Note that since $P^{(l)}_{pq}\neq P^{(l)}_{qp}$ in general, $\mathcal{A}_{h}^{(l)}$ is not necessarily a symmetric matrix. Therefore, $\mathcal{G}_{h}^{(l)}$ represents a directed graph. We employ a regular graph neural operator, such as GAT (Veličković et al., 2017), leveraging the multihead attention mechanism, to encode medical codes in each level:

[TABLE]

where $\mathbf{X}^{(l)}_{(k+1)}\in\mathbb{R}^{N^{(l)}\times d}$ and $\mathbf{X}^{(l)}_{(k)}\in\mathbb{R}^{N^{(l)}\times d}$ denote the node features at the $(k+1)$-th and $k$-th layers of $\mathcal{G}_{h}^{(l)}$, respectively. Alternatively, for the leaf level of the ontology ($l=L$), we can utilize a hypergraph structure due to its robust capability to capture the high-order complex relationships between visits and medical codes (Xu et al., 2023, 2024; Cai et al., 2022). In this approach, visits are treated as hyperedges $\mathcal{E}^{(L)}=V$, and leaf-level medical codes are treated as nodes $\mathcal{V}^{(L)}=\bm{c}^{(L)}$. This allows us to construct $\mathcal{G}_{h}^{(L)}=(\mathcal{V}^{(L)},\mathcal{E}^{(L)})$ at the leaf level of the ontology with $N^{(L)}$ nodes and $M^{(L)}$ hyperedges. We employ the Hypergraph Attention Network (HAT) (Bai et al., 2021) to encode this leaf-level hypergraph:

[TABLE]

where $\mathcal{H}^{(L)}_{h}\in\mathbb{R}^{N^{(L)}\times M^{(L)}}$ is the hypergraph incidence matrix mapping nodes to edges.

We avoid using hypergraphs for the ancestral levels, which involve fewer nodes. Employing hypergraphs at these levels would necessitate defining hyperedges with visits, resulting in an excessive number of hyperedges relative to the fewer nodes. This leads to a dense graph where nodes are redundantly connected, causing embeddings to become overly similar, despite representing distinct entities. Instead, we employ a regular graph structure, allowing for one-to-one edges based on co-occurrence, with edge inclusion controlled by a predefined co-occurrence probability threshold.

2.5. Step 3: Vertical Message Passing (VMP)

Furthermore, we leverage vertical information propagation across the hierarchical levels within the nodes of each individual ontology separately in the Meta-KG. The message passing over ancestor to descendant levels and vice versa enables the information sharing across concepts at different levels of granularity. Inspired by ideas of attention propagation in (Zhang et al., 2020) and the “two-round propagation” approach in (Pearl, 2022), we design the VMP module, a two-round hierarchy-aware graph-based encoding technique that integrates information across all ontology levels.

The first round, the bottom-up propagation, called Hierarchical Graph Information Propagation (HGIP), adaptively updates each concept node in the ontology as a convex combination of itself and its child concepts using multi-head attention mechanism. This process begins by constructing a series of sequential directed vertical subgraphs consisting of a pair of adjacent ontology levels, starting with $\mathcal{G}_{v}^{(L-1)}$ (connecting level $L$ to $L-1$) and continuing up to $\mathcal{G}_{v}^{(1)}$ (connecting level 2 to level 1). Edges in each subgraph are defined by parent-child relationships, with directed edges from nodes in level $l$ to their parent nodes in level $l-1$, forming the adjacency matrix $\mathcal{A}_{v}^{(l)}$ for the vertical subgraph $\mathcal{G}_{v}^{(l)}$. We then apply a graph attention operator, such as GAT, to each subgraph, encoding the hierarchical structure in a bottom-up manner. Starting with $\mathcal{G}_{v}^{(L-1)}$, parent node embeddings in level $L-1$ are updated by aggregating information from their children in level $L$ using multi-head attention. These updated results are then incorporated into $\mathcal{G}_{v}^{(L-2)}$, where the child embeddings are the updated parent nodes from $\mathcal{G}_{v}^{(L-1)}$. This sequential process continues to the root level subgraph $\mathcal{G}_{v}^{(1)}$, propagating information throughout the ontology. Defining step $s=0,\dots,L-1$, the sequential bottom-up learning process of HGIP is expressed as:

[TABLE]

where $\mathbf{X}^{(l-s)}_{(k)}\in\mathbb{R}^{N^{(l-s)}\times d}$ and $\mathbf{X}^{(l-s+1)}_{(k)}\in\mathbb{R}^{N^{(l-s+1)}\times d}$ represent the embeddings of parent and child nodes at the $(l-s)$-th and $(l-s+1)$-th level of the ontology, which reside in the subgraph $\mathcal{G}_{v}^{(l-s)}$. The operator $[,]$ denotes the concatenation of these embeddings, forming the node set of $\mathcal{G}_{v}^{(L-s)}$.

Unlike the previous approaches defining a single graph for vertical message passing, we represent an ontology as a series of sequential subgraphs, where each corresponds to a pair of adjacent levels. This approach, while ensuring that each parent node in the hierarchy contains the curated distilled information of all its descendants in lower levels, respects the order of node embedding updates across the main graph, incorporating the hierarchical order. This enhances the bottom-up HAP (Zhang et al., 2020) in two ways: 1) it employs a sequential GNN structure for efficient, parallel node updates, and 2) it integrates a multi-head attention mechanism to compute attention weights, enabling expressive multi-view representations and addressing the potential inconsistencies between EHR co-occurrences and ontologies (Song et al., 2019).

Second round, we apply GRAM (Choi et al., 2017) to compute the final representation of the leaf-level nodes by adaptively aggregating information from their ancestors using attention mechanism. The final representation $z_{i}\in\mathbb{R}^{d_{c}}$ of each leaf-level code $c_{i}^{(L)}$, where $i=1:N^{(L)}$, is computed as a convex combination of child embedding $x_{i}^{(L)}$ and all its ancestors’ embeddings:

[TABLE]

where $\alpha_{il}\in\mathbb{R}^{+}$ denotes the attention weight for the code embedding $\mathcal{P}^{l}(x_{i}^{(L)})$ in computing $z_{i}$. The attention weight $\alpha_{il}$ is computed using the softmax function as:

[TABLE]

where $f(a,b)$ is an MLP producing a scalar energy for the raw attention between $a$ and $b$, which the softmax normalizes into values between 0 and 1.

2.6. Integrating LINKO into EHR Models

We introduce LINKO as a medical concept encoder designed to be connected to different EHR predictive models and trained end-to-end, enhancing concept representation learning and improving predictive performance. The final medical concept (EHR code) embeddings produced by LINKO, $\mathbf{z}_{1},\mathbf{z}_{2},\dots,\mathbf{z}_{N^{(L)}}$ form the embedding matrix $\mathbf{Z}\in\mathbb{R}^{d\times N^{(L)}}$, where $\mathbf{z}_{i}$ is the $i$-th column of $\mathbf{Z}$. This embedding matrix is then used in a downstream task. For instance, for sequential diagnosis prediction, which maps a sequence of visits to the predicted diagnoses of the next visit, we have $f:\{V_{1},V_{2},\ldots,V_{t}\}\rightarrow\mathbf{\hat{y}}_{t+1}$, where $\mathbf{\hat{y}}_{t+1}\in\mathbb{R}^{N^{(L)}_{dx}}$ is a multi-hot vector, with $N^{(L)}_{dx}$ denoting the total number of diagnosis codes:

[TABLE]

For each visit $V_{t}$, we obtain a representation $\mathbf{v}_{t}\in\mathbb{R}^{d}$ by multiplying the final embedding matrix $\mathbf{Z}$ with a multi-hot vector $\mathbf{u}_{t}=\{0,1\}^{N^{(L)}}$, which represents clinical events existence in the visit. The sequence of visit representations $\{\mathbf{v}_{1},\mathbf{v}_{2},\ldots,\mathbf{v}_{t}\}$ serves as input to a main health model, $\text{Model}(\cdot)$ (e.g., Retain (Choi et al., 2016b)), producing the patient representation $\mathbf{h}_{p}\in\mathbb{R}^{d}$. The final prediction is computed by applying a Sigmoid function to the linear transformation of $\mathbf{h}_{p}$, with $\mathbf{W}\in\mathbb{R}^{N^{(L)}\times d}$ and $\mathbf{b}\in\mathbb{R}^{N^{(L)}}$ as the weight and bias, respectively. The output $\mathbf{\hat{y}}_{t+1}$ is the predicted vector in $\mathbb{R}^{N^{(L)}}$. The loss at each timestamp is the cross-entropy between the ground truth $y_{t+1}$ and prediction $\hat{y}_{t+1}$.

3. Experimental Setting

**Datasets: **We utilize two publicly available datasets: MIMIC-III (Johnson et al., 2016) and MIMIC-IV (Johnson et al., 2023). MIMIC-III (2001–2012) uses ICD-9 codes, while MIMIC-IV (2008–2019) includes both ICD-9 and ICD-10 and provides more comprehensive longitudinal data. Prescription codes in both datasets follow the National Drug Code (NDC) system, which we map to the Anatomical Therapeutic Chemical (ATC) Classification. Table 1 presents cohort statistics. The task is a multi-label sequential prediction involving identifying the ICD-9 diagnosis codes for the next visit from a highly imbalanced and vast label space, including 4,283 unique codes in MIMIC-III and 8,818 in MIMIC-IV.

**Implementations: **We report the mean and confidence intervals of the results after 5-fold experimentation. LINKO uses 4 attention heads for horizontal graphs, 2 for vertical graphs, dropout rates of 0.1 and 0.2, respectively, and a shared embedding dimension of $d=256$ for all nodes in the Meta-KG. We use a 3-level hierarchy ($L=3$) for the ICD-9 diagnosis, ICD-9 procedure, and ATC drug ontologies in our experiments.

Evaluation Metrics: (1) AUPRC: Measures the area under the precision-recall curve, capturing the trade-off between precision and recall across different thresholds. We also report AUPRC scores stratified by label frequency to assess performance across code rarity levels. (2) Acc@k: The number of correct diagnosis codes among the top $k$ predictions, divided by $\min(k,\|y_{t}\|)$, where $\|y_{t}\|$ is the number of ground-truth labels in the $(t+1)$-th visit. (3) F1-score: The harmonic mean of precision and recall, providing a balanced measure of model performance.

4. Evaluation Results

To evaluate our proposed model, we investigate the following research questions: RQ1: Does our method enhance the performance of EHR models as a complementary medical concept encoder? RQ2: How does LINKO perform compared to existing medical code encoders? RQ3: What is the impact of each component of LINKO on performance? RQ4: How do the prompting strategy and its included information affect performance? RQ5: How does our method mitigate data insufficiency limitations and rare disease predictions? RQ6: How does LINKO learn representations for individual codes?

4.1. RQ1: Plug-in Enhancement Evaluation

We propose that incorporating our medical concept encoder, LINKO, into existing EHR models boosts downstream performance through concept representation enhancement. To verify this, we integrated LINKO into four diverse EHR models: (1) AdaCare (Ma et al., 2020), an explainable model modeling biomarker variations to represent health status across time scales, (2) Transformer (Vaswani, 2017), which leverages the power of the self-attention mechanism; (3) RETAIN (Choi et al., 2016b), an RNN-based model for EHRs that utilizes a two-level reverse time attention mechanism; and (4) TCN (Bai et al., 2018), which uses causal convolutions to capture temporal dependencies in sequential data. We conducted experiments with each model both with and without LINKO. Figure 2 illustrates that LINKO consistently improves predictive accuracy when integrated with any of the models, validating its effectiveness in concept representation learning.

4.2. RQ2: Baseline Comparison

We compare our method with several existing ontology-based medical concept encoders: GRAM (Choi et al., 2017), MMORE (Song et al., 2019), KAME (Ma et al., 2018), HAP (Zhang et al., 2020), G-BERT (Shang et al., 2019), ADORE (Cheong et al., 2023), and KAMPNet (An et al., 2023). These encoders are designed for ontology-based augmentation of medical code representation and are added to predictive models as an extension to boost performance. We used the Transformer (Vaswani, 2017) as the base diagnosis prediction model. We compared the following setups: (1) the base model with no medical concept encoder; (2) the base model with each of the four existing encoders; and (3) the base model with LINKO. The general performance section in Table 2 (right side) shows that LINKO outperforms the existing encoders in enhancing the predictive performance of the base model, demonstrating its effectiveness as a complementary concept encoder. We also test two different graph encoding techniques for $\mathcal{G}_{h}^{(L)}$ in LINKO. Both techniques outperformed baselines: HAT achieved the best performance on MIMIC-IV (8,818 unique dx codes) by capturing higher-order dependencies in its large, sparse code space. In contrast, GAT excelled on MIMIC-III (4,283 unique dx codes), where denser pairwise co-occurrence enabled more effective neighborhood attention.

To assess the model’s ability to represent and predict rare concepts, diagnosis codes are divided into four frequency bands based on their occurrence: 0–25%, 25–50%, 50–75%, and 75–100%, with the 0–25% band indicating the rarest codes. This stratification highlights LINKO’s robustness under data scarcity, as it effectively propagates information to improve rare code representation. As shown in Table 2 (Label Category Performance), LINKO consistently outperforms baselines, especially for rarer code groups.

4.3. RQ3: Ablation Study

As shown in Table 2, we conduct an ablation study evaluating LINKO by removing key components: (1) w/o leaf-HMP: no horizontal message passing at the last level, (2) w/o parent-HMP: no horizontal message passing at parent levels, (3) w/o HMP: no horizontal message passing at any level, (4) w/o HGIP: no hierarchical graph information propagation in vertical message passing, (5) w/o LLM: random initialization for concept embeddings (removing LLM initialization). All ablated versions showed performance degradation. The drop in w/o leaf-HMP highlights the importance of ontology fusion at the last (leaf/child) level, while w/o parent-HMP exhibits an even greater decline, emphasizing the significance of ontology fusion at coarser-grained (parent) levels, which has been largely overlooked in prior studies. The largest drop occurs in w/o HMP, indicating the critical role of multi-ontology connections across all levels. The decline in w/o HGIP underscores the impact of vertical propagation through HGIP in structured hierarchical subgraphs. Finally, the performance drop in w/o LLM highlights the effectiveness of LLM-based embedding initialization, which injects clinical knowledge into the model, provides a strong training foundation, and guides it toward better performance.

To further assess the impact of ontology integration (fusion) on concept representation learning and overall performance, we train our model using different combinations of concept types (diagnosis (dx), drugs (rx), and procedures (px)), each time drop one concept type to compare the performance with when training on all. Also, we repeat the same experiments with LINKO w/o HMP, an ablated version of our model which HMP is removed, meaning ontological multi-level integration is disabled. Based on Table 4, we draw the following conclusions: (1) Removing ontology integration (i.e., HMP) in any combination cases (comparing results on the left and right sides of the table) leads to a noticeable drop in performance. This highlights the importance of HMP in enabling multi-ontology integration and capturing both heterogeneous and homogeneous code interactions across all hierarchical levels. (2) Dropping any concept type (ontology) results in a performance decline when ontology integration is enabled (right-side of the table), with dx having the highest impact and px the least for diagnosis prediction. However, when ontology integration is removed (left-side of the table), this is not always the case. For example, training on (dx, rx) yields better results than (dx, rx, px) in the absence of ontology integration. This occurs because px information is significantly sparser than dx and rx, and when learned independently, it not only fails to contribute to diagnosis prediction but even degrades performance. Without ontology integration, each ontology is forced to be learned independently, making it insufficient to extract useful information for diagnosis prediction. However, with HMP, px concepts are connected to dx and rx across all hierarchical levels, facilitating pattern extraction in concept representation learning.

4.4. RQ4: Prompt Analysis

As shown in Table 3, we also experiment with different prompting strategies to determine the most effective information to include. The best results are achieved when the prompt includes the code type and name (e.g., ICD-9 Diagnosis 250.7), its corresponding descriptive concept (e.g., Diabetes with peripheral circulatory disorders), its parent codes and concepts, and the task name (e.g., diagnosis prediction). Notably, incorporating child information for parent codes slightly decreases performance, likely due to excessive prompt length, as each parent has many children. In contrast, adding parent information maintains a manageable prompt length since each code has only one parent per level. The second row of the table shows that adding irrelevant noise to the prompt degrades performance. In the last row of the table, we freeze the LLM embedding initialization and trained the model with fixed embeddings. The poor performance demonstrates that raw LLM embeddings without further refinement are ineffective. Lastly, we noticed that the LLM initialization technique accelerates training convergence, requiring about $\sim 60$ epochs compared to $\sim 250$ with random initialization.

4.5. RQ5: Data Insufficiency Analysis

To assess the robustness of LINKO under data insufficiency, we vary the size of the training dataset to simulate limited data scenarios. As shown in Figure 3, our model consistently outperforms its components in data-scarce settings, demonstrating its effectiveness. Furthermore, in the Plug-in Enhancement setting, Figure 4 shows that even with reduced training data, our model significantly enhances the performance of EHR models.

4.6. RQ6: Case Study Analysis

Figure 5 illustrates how LINKO learns the representation of the ICD-9 code 428.0, which denotes “Congestive heart failure, unspecified” (CHF), through a dual-axis message passing paradigm for rich medical concept representations. Vertically, LINKO retrieves all ancestors of this code across levels, and horizontally, it gathers co-occurring codes (red for procedures, green for diagnosis, and blue for drugs) for each parent and the target code across all ontologies. The figure shows the extracted sub-KG for ICD-9 code 428.0 within the Meta-KG. Representation learning begins by first initializing each node embedding using LLM prompting and dense embedding retrieval. For example, for ICD-9 428.0 code, the prompt is ¡¡ For the task of diagnosis prediction, provide a semantic representation for ICD-9 code 428.0 representing “Congestive heart failure, unspecified”. It falls under the broader ICD-9 categories of ‘420-429.99’ (other forms of heart disease) and ‘390-459.99’ (disease of the circulatory system)¿¿. LINKO then performs horizontal propagation, applying graph attention to aggregate information from neighboring nodes across all levels. Next, the HGIP propagates information upward, updating each parent node using its children’s embeddings via the sequence of graph attention operators. Finally, the node embedding is refined through a convex combination of its own representation and those of its ancestors.

Larger circles indicate higher-level codes, representing more general concepts compared to lower-level codes. Within each hierarchical level, solid blue lines connect the codes, representing co-occurrence-driven edges for horizontal message passing. Dashed lines illustrate parent-child relationships for bottom-up vertical message passing (HGIP), with color coding (red for procedures, green for diagnoses, and blue for drugs). Parent-child edges for top-down GRAM propagation are depicted using solid black lines. The weights for all graph edges during horizontal and vertical propagation are learned through attention techniques, with edge thickness in the figure indicating the relative attention assigned to each edge.

5. Related Work

EHR Predictive Models. The widespread adoption of EHRs has enabled the development of machine learning models for healthcare prediction, starting with sequential models (Choi et al., 2016a; Ashfaq et al., 2019), followed by attention-based models (Choi et al., 2016b; Hu et al., 2025), transformer-based approaches (Li et al., 2020; Choi et al., 2020; Nayebi Kerdabadi et al., 2023), and recently, graph neural networks (Su et al., 2020; Lu et al., 2021a; Xu et al., 2022; Yang et al., 2023; Poulain and Beheshti, 2024).

Ontology-Driven Concept Representation. These works aim to enhance medical concept representation learning by augmenting structured EHR data (sequence of discrete medical codes) with hierarchical medical ontology, without using any other data modalities (e.g., unstructured clinical text, numerical measurements). We compare our method with these baselines in our experiments. For instance, GRAM (Choi et al., 2017) leverages the ontology hierarchy to represent a medical concept as a convex combination of itself and its ancestors. Building on GRAM, MMORE (Song et al., 2019) enhances GRAM by enabling multiple representations for each parent concept, addressing discrepancies between EHR data and medical ontologies. KAME (Ma et al., 2018) further proposes incorporating ontology knowledge throughout the entire prediction process on top of code representation learning. While GRAM-based methods improve performance, they limit expressiveness by treating ancestors as unordered. HAP (Zhang et al., 2020) addresses this by propagating attention hierarchically, enabling top-down and bottom-up information sharing across the ontology. Further, G-BERT (Shang et al., 2019) combines GNNs on ontology and BERT for medical code representation. Later, ADORE (Cheong et al., 2023) utilizes the relational ontology SNOMED to integrate multi-source medical codes, which was ignored in previous studies. However, it suffers from loss of information in code conversion from ICD (hierarchical) to SNOMED (relational), especially for medication codes. Also, KAMPNet (An et al., 2023) employs graph contrastive learning for effective EHR representation, but first fails to fully account for the hierarchical order of node calculation, and second, it only captures the interaction of multi-source codes (diagnosis, medication) in the last level of ontologies.

Multi-Source-Augmented EHR Representation (Lu et al., 2021b, 2024; Safaeipour and Hashemi, 2024; Jiang et al., 2024). GCL (Lu et al., 2021b) is a collaborative graph learning model that jointly learns patient and disease representations, incorporating unstructured text with attention regulation. RAM-EHR (Xu et al., 2024) improves EHR predictions by retrieving external textual medical knowledge from multiple online sources, augmenting the local model co-trained with consistency regularization. EMERGE (Zhu et al., 2024), a RAG framework, enhances multimodal health representations by encoding clinical notes with an LLM and time-series data with a GRU. KARE (Jiang et al., 2024) integrates structured EHR codes and unstructured biomedical text via multi-source knowledge graphs and community-level retrieval. It enhances clinical prediction by augmenting patient context and fine-tuning LLMs with reasoning chains. ClinicalRAG (Lu et al., 2024) uses a multi-agent RAG pipeline to integrate structured (e.g., knowledge graphs) and unstructured (e.g., online text) medical knowledge into LLMs, improving its diagnostic accuracy. In experiments, we do not compare with such models that incorporate external descriptive texts, as they operate in a multimodal setting and focus on general patient representation. In contrast, our model targets concept-level representation using only structured code.

6. Conclusion

We introduced LINKO, a framework for LLM-augmented multi-ontology integration via dual-axis knowledge propagation, designed to enhance medical concept representation in EHR models. LINKO extracts cross-ontology relationships through message passing in two dimensions: vertical and horizontal, and initializes concept embeddings with graph-retrieval-augmented LLM prompting. By enhancing medical concept representations, LINKO achieves superior performance on downstream tasks compared to other medical concept encoder baselines. Additionally, we showcase the robustness of LINKO in data-limited scenarios and validate its plug-in compatibility to enhance existing healthcare models.

7. GenAI Usage Disclosure

We used generative AI tools (e.g., ChatGPT) solely for the purpose of grammatical and stylistic editing of the manuscript. No content generation, data analysis, or experimental design was performed using these tools.