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

Ziyang Song; Lia Shen; Yixuan Li; Qincheng Lu; Weihao Li; Haoran Xu; He Zhu; Xinyu Wang

cs.CLOct 13, 2026

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

Authors: Ziyang Song, Lia Shen, Yixuan Li, Qincheng Lu, Weihao Li, Haoran Xu, He Zhu, Xinyu Wang,

Organizations: 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 · School of Engineering, Northwestern University ·

Abstract

Biomedical language models (LMs) provide strong representations for clinical tasks, but they often fail to capture the hierarchical structure among medical concepts. This limitation is especially pronounced for rare and fine-grained clinical codes in long-tailed vocabularies, where hierarchy-agnostic LM representations become less informative as concept frequency decreases. We introduce Semantic Medical Identifier (SMI), a training-free framework for constructing ontology-grounded Semantic IDs and hierarchy-aware medical concept representations. For each diagnosis, procedure, or medication concept, SMI defines its Semantic ID as the coarse-to-fine path through an expert-defined medical ontology. SMI then aggregates frozen biomedical LM embeddings across the corresponding ontology levels to form hierarchy-aware medical concept representations that preserve coarse-to-fine clinical semantics. The base SMI uses uniform aggregation across ontology levels, while a level-weighted variant selects global level weights on validation data. Our theoretical analysis shows that SMI yields more informative medical concept representations than hierarchy-agnostic LM embeddings, particularly for rare and fine-grained clinical concepts. Experiments show that SMI captures meaningful hierarchical structure and improves clinical relation prediction. On MIMIC-IV, SMI improves mortality prediction, readmission prediction, drug recommendation, and rare drug prediction. On eICU, SMI yields more robust patient representations under cross-hospital distribution shift.

Figures & tables

Figure 1: SMI better distinguishes hierarchical relations than BioBERT on ICD 1-hop prediction. (a) SMI consistently improves hierarchical relation prediction across levels, with the largest gains at finer ICD granularities. (b) SMI maintains stronger and more stable performance than BioBERT across concept-frequency bins, suggesting that medical ontology structure helps mitigate representation degradation for long-tailed clinical codes.

Figure 2: Overview of SMI and hierarchy-aware medical concept embedding. a . An expert-curated medical ontology organizes clinical concepts from coarse to fine levels. b . A frozen biomedical LM encodes clinical codes at each ontology level into level-specific embeddings. c . SMI represents each medical concept by its coarse-to-fine ontology path as a Semantic ID and aggregates the corresponding embeddings to form a hierarchy-aware medical concept representation.

Method	ICD			ATC			ICD $\rightarrow$ ATC
Method	Precision	Recall	F-score	Precision	Recall	F-score	Precision	Recall	F-score
SMI	70.6	67.9	69.2	83.8	67.6	74.9	53.2	61.4	57.0
SMI-W	77.2	66.6	71.5	86.1	70.5	77.5	53.6	62.2	57.6
BioBERT	30.4	75.8	44.2	50.7	45.7	48.1	49.3	57.0	50.9
SapBERT	69.9	64.5	67.1	84.0	66.9	74.5	50.2	55.7	52.9
HiTs	69.3	61.7	65.3	83.0	66.2	71.5	50.5	56.3	53.2

Table 1: Clinical relation prediction results for ICD

\rightarrow

ICD, ATC

\rightarrow

ATC, and ICD

\rightarrow

ATC, comparing SMI and SMI-W against baselines. Best F-scores are highlighted in bold. We reported the optimized result for each method as the decision threshold is selected via deterministic grid search.

Method	NMI $\uparrow$	ARI $\uparrow$	Purity $\uparrow$	Recall $\uparrow$	Precision $\uparrow$	Sil. $\uparrow$	Intra/Inter $\downarrow$
SMI	0.77	0.56	0.78	0.80	0.61	0.52	0.29
BioBERT	0.39	0.25	0.51	0.41	0.35	0.04	0.82

Table 2: Aggregate clustering quality of ICD diagnosis embeddings by major disease category. NMI, ARI, Purity, Recall, Precision, and Silhouette are higher better, while Intra/Inter Ratio is lower better. SMI achieves stronger cluster-label alignment and better geometric separation than BioBERT, indicating more ontology-consistent diagnosis embeddings.

Figure 3: Visualization of ICD diagnosis embeddings across 13 major categories using UMAP. (a) BioBERT embeddings. (b) SMI embeddings. ICD codes embedded by SMI form more distinct category-level clusters that are better aligned with the medical ontology.

Method	Mortality		Readmission		Drug			Rare Drug
	AUPRC	AUROC	AUPRC	AUROC	AUPRC	F-score	Jaccard	Recall@15
SMI	44.5 (0.5)	87.4 (1.0)	48.9 (0.2)	72.5 (0.5)	75.6 (0.1)	62.4 (0.1)	48.6 (0.1)	41.4 (0.2)
SMI-W	45.1 (0.3)	88.0 (0.7)	48.5 (0.3)	72.9 (0.3)	77.6 (0.1)	65.3 (0.2)	50.5 (0.1)	44.2 (0.3)
GRAM	38.5 (0.7)	79.1 (0.9)	43.2 (0.8)	68.2 (0.4)	72.0 (0.3)	61.3 (0.4)	48.6 (0.4)	34.5 (0.2)
cui2vec	39.0 (0.4)	80.4 (0.8)	43.7 (0.6)	68.9 (0.3)	75.2 (0.1)	64.3 (0.3)	49.9 (0.4)	35.6 (0.1)
RotatE	39.2 (0.3)	79.6 (0.6)	43.6 (0.4)	69.2 (0.6)	74.1 (0.2)	62.7 (0.4)	48.7 (0.4)	36.4 (0.1)

Table 3: Clinical prediction results on MIMIC-IV for mortality, readmission, and (rare) drug recommendation, comparing SMI with two groups of baselines. Metrics are reported as mean (standard error) from bootstrap. The best results in each column are highlighted in bold.

Method	ICD			ATC			ICD $\rightarrow$ ATC
Method	Precision	Recall	F-score	Precision	Recall	F-score	Precision	Recall	F-score
BioBERT	30.4	75.8	43.4	50.7	45.7	48.1	49.3	57.0	52.9
SMI	70.6	67.9	69.2	83.8	67.6	74.9	53.2	61.4	57.0
SMI-W	77.2	66.6	71.5	86.1	70.5	77.5	53.6	62.2	57.6

Table 4: Ablations on clinical relation prediction comparing the LM baseline, SMI, and SMI-W.

Figure 4: Cross-hospital distribution shift comparison between BioBERT and SMI. Each cell shows $|\Delta|_{\text{BioBERT}}-|\Delta|_{\text{SMI}}$ ; red colors indicate that SMI is more cross-site robust. P-value indicates that SMI achieves significantly lower distribution shift.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

SMI and Embedding
Notation	Description
$L$	Total number of hierarchy levels
$l$	Index of a specific hierarchy level
$C^{(\ell)}$	Set of clinical codes at hierarchy level $\ell$
$c_{j}^{(\ell)}$	The $j$ -th clinical code at hierarchy level $\ell$
$\mathcal{P}(i)$	Semantic ID of concept $i$

Appendix

Table 5: Summary of Notations and Variables used in SMI

Pair type	Mean cosine	Median cosine
Same modifier	0.655	0.730
Different modifier	0.080	0.068

Appendix

Table 6: Strict ICD-9 offset consistency analysis. Same-modifier offsets are substantially more aligned than different-modifier clinical offsets, with a mean cosine difference of

0.575

and a one-sided

p

-value of

2.0\times 10^{-4}

, supporting the hierarchical semantic decomposition assumption.

Figure 5: Distribution of cosine similarities between ICD-9 parent-child offset vectors in the strict modifier analysis. Same-modifier pairs compare offsets sharing the same extracted clinical modifier, while different-modifier pairs compare offsets from different modifier groups.The same-modifier distribution is shifted strongly to the right, indicating stronger offset consistency than random clinical offsets and supporting the hierarchical semantic decomposition assumption.

Level	Description	Diagnosis	Procedure	Medication
Level 1	Name	Major category	Major category	Anatomical main group
Level 1	# Concepts	18	16	14
Level 2	Name	Sub-category	Sub-category	Therapeutic subgroup
Level 2	# Concepts	34	200	90
Level 3	Name	CCS code	CCS code	Pharmacological subgroup
Level 3	# Concepts	265	224	248

Appendix

Table 7: Hierarchical taxonomy of diagnosis, procedure, and medication concepts used in this study. Each level contains two attributes: Name , which specifies the semantic grouping at that level, and # Concepts , the number of unique codes at that level. Diagnosis categories are derived from ICD-CM codes mapped through the CCS diagnostic hierarchy; procedure categories use ICD-PCS codes organized by the CCS procedure hierarchy; and medications follow the five-level ATC ontology. For ICD-CM and ICD-PCS levels, the # Concepts is reported as the total number of ICD-9 and ICD-10 codes shown in parentheses, respectively.

Dataset	#patients	#visits	#visits/patient	#conditions/patient	#procedures/patient	#drugs/patient
MIMIC-IV	223,452	546,028	2.44	28.48	3.85	48.32

Appendix

Table 8: Basic statistics of the pre-processed MIMIC-IV EHR dataset. “#” denotes counts and “/patient” denotes per-patient averages.

Figure 6: The drug rates before and after imputation of the eICU data. Since some drugs may lack recorded names, drug rates indicate the proportion of administered drugs with recorded names over all unique drugs.

Level	Diagnosis	Procedure	Medication
Level 1	Diseases of circulatory system	Miscellaneous diagnostic and therapeutic procedures	Alimentary tract and metabolism
Level 2	Hypertension	Magnetic resonance imaging	Drugs used in diabetes
Level 3	Hypertension with complications and secondary hypertension	Magnetic resonance imaging	Insulins and analogues
Level 4	Secondary hypertension	Imaging, Axial Skel, Exc Skull Facial, MRI Imaging, Axial Skeleton, Except Skull and Facial Bones, Magnetic Resonance Imaging (MRI)	Insulins and analogues for injection, long-acting
Level 5	Hypertension secondary to other renal disorders	Magnetic Resonance Imaging (MRI) of Lumbar Spine	Insulin glargine

Appendix

Table 9: Level-specific textual descriptions used by the biomedical LM. The example includes an ICD-10 diagnosis ( I15.1 Hypertension secondary to other renal disorders ), an ICD-10 procedure ( 79658 Magnetic Resonance Imaging (MRI) of Lumbar Spine ), and an ATC medication ( A10AE04 Insulin glargine ). For each hierarchy level, we encode only the textual description of the code at that level, rather than concatenating descriptions along the full ontology path.

Figure 7: Rank–frequency distributions (log–log scale) of clinical concepts in MIMIC-IV across five hierarchy levels: (a) ICD-9 diagnosis codes, (b) ICD-9 procedure codes, and (c) ATC medication codes. For each hierarchy level, concepts are ranked by patient frequency within that level.

Method	NMI $\uparrow$	ARI $\uparrow$	Purity $\uparrow$	Recall $\uparrow$	Precision $\uparrow$	Sil. $\uparrow$	Intra/Inter $\downarrow$
SMI	0.77	0.56	0.78	0.80	0.61	0.52	0.29
BioBERT	0.39	0.25	0.51	0.41	0.35	0.04	0.82

Appendix

Table 10: Aggregate clustering quality of ICD diagnosis embeddings by major disease category. NMI, ARI, Purity, Recall, Precision, and Silhouette are higher better, while Intra/Inter Ratio is lower better. SMI achieves stronger cluster-label alignment and better geometric separation than BioBERT, indicating more ontology-consistent diagnosis embeddings.

Major	Major Category	Recall (%)		Precision (%)
		SMI	BioBERT	SMI	BioBERT
1	Infectious & parasitic	73.3	33.1	100.0	45.4
2	Neoplasms	62.8	53.7	100.0	91.5
3	Endocrine & metabolic	85.3	19.8	21.8	8.6
4	Blood diseases	91.7	41.7	10.6	4.8
5	Mental disorders	89.9	58.1	100.0	58.7

Appendix

Table 11: Per-major clustering recall and precision by ICD major disease category. Recall measures whether codes from the same major category are concentrated in a dominant cluster, while precision measures the purity of that dominant cluster. SMI achieves higher recall and precision than BioBERT for most ICD major categories, indicating more coherent and less mixed ontology-aligned clusters.

Figure 8: Cluster assignment heatmaps for ICD diagnosis embeddings. Rows denote ICD major disease categories and columns denote K-means clusters ordered by their dominant major category. Each cell shows the percentage of codes from a major category assigned to a cluster. Compared with BioBERT, SMI exhibits a stronger diagonal pattern, indicating better alignment between embedding structure and ICD major categories.

Figure 9: Cross-hospital distribution shift analysis based on average absolute log-likelihood differences ( $|\Delta|$ ) between source and target hospitals. Each heatmap illustrates the distribution shift of patient embedding using SMI (left) and BioBERT (right), both evaluated with a VAE-based density model. SMI embeddings exhibit smaller distribution shifts (average $|\Delta|=50.77$ ) compared to BioBERT ( $|\Delta|=74.06$ ), indicating that SMI is more robust to distribution shift.

Representation Learning

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

See all 773 papers

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
Text-Attributed Knowledge Graph Enrichment with Large Language Models for Medical Concept Representation
Apr 14, 2026Mohsen Nayebi Kerdabadi, Arya Hadizadeh Moghaddam, Chen Chen +2Representation Learning Biomedical KGs
University of Kansas, USA
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

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

HCOE: Hyperbolic Clinical Ontology Embeddings from Biomedical Language Models

Yixuan Li, Weihao Li, Ziyang Song

McGill University · Northwestern University · Ohio University

Medication Recommendation Hierarchical Representation Learning Hyperbolic Representation Learning

Open paper AskPDF arXiv

Biomedical language models (LMs) encode textual semantics but do not explicitly preserve medical code hierarchies. We present Hyperbolic Clinical Ontology Embeddings (HCOE) for hierarchy-aware clinical concept representation. HCOE maps frozen BioBERT embeddings into a Poincare ball, combining parent-side and child-side ontology-guided contrastive learning with coarse-to-fine ontology-path aggregation. It uses International Classification of Diseases (ICD) codes organized by Clinical Classifications Software (CCS) and Anatomical Therapeutic Chemical (ATC) medication hierarchies. Evaluations show that HCOE performs best on ICD/ATC clinical relation prediction and CCS-to-PheCode hierarchy transfer. On the MIMIC-IV dataset, HCOE also achieves the best performance on mortality prediction, readmission prediction, medication recommendation, and rare drug prediction.

Abstract

Representation Learning

Explore similar work

HCOE: Hyperbolic Clinical Ontology Embeddings from Biomedical Language Models

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

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