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

1 Introduction

• Biomedical LMs often fail to capture hierarchical relations, while single-description representations can lose the coarse-to-fine clinical semantics encoded by ICD and ATC taxonomies.This matters especially for rare, fine-grained codes in long-tailed clinical vocabularies.Source: §1 · p. 1
• SMI improves ICD 1-hop prediction across hierarchy levels, with the largest gains at finer granularities, and remains stronger and more stable across concept-frequency bins.BioBERT accuracy drops from 85.1% at the coarsest level to 58.6% at the finest level.Source: Figure 1 · p. 2, §1 · p. 2
• SMI defines each concept’s Semantic ID as its expert-defined coarse-to-fine ontology path and aggregates frozen biomedical LM embeddings along that path.Diagnosis and procedure concepts use ICD; medication concepts use ATC.Source: §1 · p. 2, Figure 2 · p. 3
• SMI’s theoretical analysis finds more informative representations than hierarchy-agnostic LM embeddings, particularly for rare and fine-grained clinical concepts.The framework’s base version weights ontology levels uniformly; a variant selects global level weights on validation data.Source: §1 · p. 2
• SMI improves clinical relation prediction, mortality and readmission prediction, and drug recommendation on MIMIC-IV, with particularly strong gains in rare drug prediction.On eICU, it shows smaller cross-hospital distribution shifts than BioBERT.Source: §1 · p. 2

2 Related Work

• One line of prior work incorporates UMLS, ontology resources, or biomedical knowledge graphs into concept embeddings, typically requiring additional training.Examples include synonym sets, ontology descriptions, and knowledge graphs.Source: §2 · p. 3
• A second line uses ontology structure in supervised prediction models, which are trained for specific downstream tasks rather than general-purpose concept representation.Approaches include ontology-guided aggregation and graph-augmented patient modeling.Source: §2 · p. 3
• Unlike learned or quantized Semantic IDs, SMI derives each identifier directly from an expert-defined medical ontology path and aggregates frozen biomedical LM embeddings without training an additional quantizer.The path represents the concept’s coarse-to-fine ontology structure.Source: §2 · p. 3

3.1 Expert-Curated Medical Ontologies

• Diagnosis, procedure, and medication concepts are organized into distinct five-level hierarchies using CCS-based ICD taxonomies and the ATC system.Each code has a unique path from a broad category to a fine-grained code or chemical substance.Source: §3.1 · p. 4

3.2 SMI: Semantic IDs for Medical Concept Representation

• A concept’s Semantic ID is its ordered sequence of ancestor codes from coarse to fine, such as the five-code path assigned to ICD-10-CM E11.2.That example runs from an endocrine-disease major category to the decimal-level code E11.2.Source: §3.2 · p. 4
• SMI forms a concept representation by aggregating frozen biomedical LM embeddings for the codes along its Semantic ID, preserving coarse-to-fine semantics without additional trainable parameters.The level-specific embeddings encode textual descriptions, and the base SMI sums them uniformly.Source: §3.2 · p. 4
• SMI-W replaces uniform aggregation with global, shared level weights selected on validation data; equal weights recover the base SMI.The variant accounts for differing contributions of hierarchy levels.Source: §3.2 · p. 4

3.3 Theoretical Analysis of LM Representation Collapse and SMI Informativeness

• Under a frequency-dependent noise model, the analysis says hierarchy-agnostic embeddings of infrequent fine-grained concepts become noise-dominated and lose semantic separability.The analysis relates embedding noise to concept frequency and states that fine-grained concepts are often infrequent in long-tailed vocabularies.Source: §3.3 · p. 5
• Under the paper’s hierarchical semantic decomposition, SMI aggregates higher-frequency ancestor information to reduce estimation noise and preserve separability for infrequent fine-grained concepts.The analysis states that both uniform and weighted SMI improve pairwise SNR over hierarchy-agnostic LM embeddings under its assumptions.Source: §3.3 · pp. 5–6

4.1 SMI Embeddings Capture Clinical Concept Relations

• SMI and SMI-W led all three relation types; SMI-W reached ICD F-score 71.5%, ATC F-score 77.5%, and ICD→ATC F-score 57.6%, versus HiTs at 53.2%.For ICD and ATC, SMI-W exceeded SapBERT’s F-scores of 67.1% and 74.5%, respectively.Source: Table 1 · p. 7, §4.1 · p. 7

4.2 SMI Embeddings Capture Medical Ontology Structure

• SMI achieved stronger ontology-aligned clustering than BioBERT: NMI 0.77 vs. 0.39, ARI 0.56 vs. 0.25, Purity 0.78 vs. 0.51, and Silhouette 0.52 vs. 0.04.SMI also increased Recall from 0.41 to 0.80 and Precision from 0.35 to 0.61, while lowering Intra/Inter Ratio from 0.82 to 0.29.Source: Table 2 · p. 7, §4.2 · pp. 7–8

4.3 Evaluation on MIMIC-IV and eICU EHR Datasets

• Evaluation covered MIMIC-IV mortality, readmission, drug recommendation, and rare-drug prediction, plus eICU cross-hospital shifts; both datasets used 70%/10%/20% splits.MIMIC-IV prediction used patient-level mean pooling over tokens; eICU analysis included 10 hospitals and 24,179 patients.Source: §4.3 · p. 8, §4.4 · p. 8

4.4 SMI Improves Clinical Prediction and Rare Drug Prediction

• SMI and SMI-W improved clinical prediction on MIMIC-IV; SMI-W led drug recommendation (77.6% AUPRC, 65.3% F-score, 50.5% Jaccard) and rare-drug Recall@15 (44.2%).For rare-drug Recall@15, SMI-W significantly outperformed strongest baseline G-BERT (p = 0.0081).Source: Table 3 · p. 9, §4.4 · p. 9

4.6 Ablation Studies

• SMI consistently outperformed BioBERT across ICD, ATC, and ICD→ATC relation prediction, and SMI-W further improved on uniform SMI.The comparison suggests that ontology-path aggregation helps, while global level weighting captures unequal contributions from hierarchy levels.Source: Table 4 · p. 9, §4.6 · p. 10

4.5 SMI Mitigates Cross-Hospital Distribution Shift

• 50.77 vs. 74.06 cross-hospital distribution shift for SMI vs. BioBERT; paired testing found significantly smaller shifts with SMI (p = 1.97 × 10^-5).Most off-diagonal hospital-pair entries favored SMI.Source: Figure 4 · p. 10, §4.5 · p. 10

5 Conclusion

• SMI integrates medical ontologies with language-model embeddings to construct hierarchy-aware concept representations, with theoretical analysis and experiments supporting its broad utility.The reported evaluations cover clinical relation prediction, downstream clinical prediction, and cross-site distribution shift.Source: §5 · p. 10
• Qualitative analyses find clinically meaningful structures aligned with medical hierarchies, while experiments report utility across relation prediction, clinical prediction, and cross-site shift.Source: §5 · p. 10
• The study focuses on expert-curated ICD and ATC systems; emerging classifications and local coding standards may require ontology validation and code mapping.Future work includes ontology harmonization, missing-link discovery, broader biomedical resources, and integrating multiple ontologies.Source: §5 · p. 10

A Notation Table

• The notation table defines μ_i as a concept’s true semantic vector, ε_i as embedding-estimation noise, and n_i as its training-corpus frequency.It also identifies δ at level ℓ+1 as a child concept’s semantic offset relative to its parent.Source: A Notation Table · p. 14, Table 5 · p. 14

• The empirical embedding is computed from independent samples whose expected values equal the concept’s true semantic vector.Source: §B.1 · p. 15
• As concept frequency n_i decreases, embedding estimation becomes less stable and variance increases, motivating a frequency-dependent noise assumption.Source: §B.1 · p. 15
• The empirical estimator is unbiased, and its variance decreases at rate 1/n_i.The proof assumes isotropic noise Σ_i = σ^2I_d.Source: §B.1 · p. 15

• The analysis considers independent embedding-noise vectors for two concepts and characterizes the magnitude of their difference.Source: §B.2 · pp. 15–16
• The expected squared magnitude and RMS magnitude of the noise difference are derived as measures of its scale.Source: §B.2 · p. 16
• In high dimensions, the noise-difference norm concentrates around its RMS scale under the stated assumptions.The proof uses the normalized χ^2_d distribution of the squared norm of an isotropic Gaussian vector.Source: §B.2 · pp. 16–17

• Independent embedding noises with variances tied to concept frequencies yield a noise difference whose norm concentrates around its scale.Source: §B.3 · p. 17
• The theorem considers two low-frequency, fine-grained concepts whose semantic difference is negligible relative to embedding noise.The proof assumes SNR(i, j) ≪ 1.Source: §B.3 · p. 17
• For rare fine-grained concepts, semantic separation is asymptotically dominated by estimation noise, making hierarchy-agnostic LM embeddings geometrically inseparable.Source: §B.3 · p. 17

• The proof assumes a child concept's true semantic representation decomposes into parent semantics and a child-specific offset.Source: §B.4 · p. 18
• SMI aggregates representations along the child's hierarchical path, with its parent receiving nonzero weight.Source: §B.4 · p. 18
• The child-specific component remains bounded, leaving parent-level semantic separation as the dominant distinction from a non-parent concept.Source: §B.4 · p. 18
• SMI embeddings remain separable between a leaf concept and any non-parent concept in the long-tail regime.Source: §B.4 · p. 18

• Under the long-tail assumption, coarser concepts are observed more frequently and have lower noise variance than fine-grained concepts.Source: §B.5 · pp. 18–20
• Weighted SMI aggregates level-specific LM embeddings from a code's ancestors, with uniform SMI assigning equal weights across levels.Source: §B.5 · p. 19
• Because the semantic margin is fixed across estimators, the SNR comparison reduces to comparing their noise variances.Source: §B.5 · p. 19
• SMI-W and uniform SMI embeddings achieve higher pairwise SNR than hierarchy-agnostic LM embeddings.Source: §B.5 · p. 20

B.6 Offset Consistency Analysis for Hierarchical Semantic Decomposition

• Mean cosine similarity was 0.655 for same-modifier offsets versus 0.080 for different-modifier offsets, a mean difference of 0.575 (p = 2.0 × 10−4).Source: Table 6 · p. 21, §B.6 · p. 21, Figure 5 · p. 21

C.1 Construction of Medical Ontologies

• SMI uses unique paths through expert-defined hierarchies: CCS taxonomies for diagnoses and procedures, and the five-level ATC taxonomy for medications.Source: §C.1 · p. 22

C.2 Statistics of Medical Ontologies

• The five-level hierarchies span 13,193 ICD-9-CM and 73,066 ICD-10-CM diagnosis codes, 3,674 and 80,492 ICD-PCS procedure codes, and 5,495 ATC substances.Source: §C.2 · p. 22

• The processed MIMIC-IV cohort contains 223,452 patients and 546,028 hospital admissions.Source: §D.1 · p. 23
• Each patient is represented by chronologically ordered visits, with diagnoses and procedures treated as unordered sets within each visit.Source: §D.1 · pp. 23–24
• Medication codes are mapped from NDC through RxNorm to ATC, retaining level-5 codes and excluding unmappable entries.Source: §D.1 · p. 24
• The task setup includes next-visit mortality, 15-day readmission, and multi-label drug recommendation prediction.Source: §D.1 · p. 24

D.2 eICU Dataset Processing

• HICL-based imputation reduced unrecorded drug names in eICU from approximately one-third to around 20%, while processing standardized medication names and concepts for downstream modeling.The analysis used 24,179 patients from the 10 largest hospitals, treated records as unordered observations, and mapped diagnoses and medications to ICD-9 and ATC codes. Drug strings were cleaned to retain core medication terms, then matched to ATC substances using BioWordVec cosine similarity.Source: §D.2 · pp. 24–25

D.3 Prompt for Biomedical LMs

• The biomedical LM prompt encodes each concept’s hierarchy by concatenating descriptions of its code and ancestor codes at every level.The encoder uses a frozen biomedical LM, and the hierarchy-aware input was reported as more effective at capturing medical semantics.Source: §D.3 · p. 25

D.4 Baseline Details

• The baselines span text-based, ontology-aware, graph-based, and task-trained representations, but differ from SMI in whether they provide training-free, path-aggregated concept embeddings.For example, BoxLM is supervised and task-specific, while BioBERT does not explicitly encode medical hierarchies; GRAM learns hierarchy-aware representations in a supervised, task-specific manner.Source: §D.4 · pp. 25–27

D.5 Computational resources

• SMI requires a one-time computation of frozen LM embeddings that can be reused downstream, adding minimal hardware or computational overhead beyond embedding extraction and evaluation.Experiments were conducted on an NVIDIA RTX PRO 6000 GPU with 96GB memory.Source: §D.5 · p. 27

E.1 Long Tail Distribution of Clinical Concepts

• Across diagnoses, procedures, and medications, finer hierarchy levels contain more concepts with sharply lower patient frequencies, producing many rare codes.The rank–frequency analysis covers five hierarchy levels for ICD-9 diagnoses, ICD-9 procedures, and ATC medications; the reported sparsity motivates hierarchy-aware representations for infrequent, fine-grained concepts.Source: §E.1 · pp. 27, 29, Figure 7 · p. 28

E.2 Details about Clinical Relation Prediction

• Relation evaluation covers multi-hop subsumption within ICD and ATC hierarchies and one-hop ICD→ATC prediction, using Precision, Recall, and F-score.Multi-hop training uses 10 negative parent candidates per positive, with sibling hard negatives or random negatives; cross-ontology evaluation samples one negative ICD-10 code per positive pair.Source: §E.2 · p. 29

E.3 Details about Clustering Analysis

• SMI beats BioBERT on NMI (0.77 vs. 0.39), ARI (0.56 vs. 0.25), Purity (0.78 vs. 0.51), Recall (0.80 vs. 0.41), Precision (0.61 vs. 0.35), Silhouette (0.52 vs. 0.04), and Intra/Inter Ratio (0.29 vs. 0.82).These aggregate results indicate stronger ICD major-category alignment and better geometric separation for SMI embeddings.Source: Table 10 · p. 30, §E.3 · p. 30
• Blood diseases, Skin, and Injury & poisoning remain challenging, with low precision despite high recall, although SMI still outperforms BioBERT for these categories.The passage attributes mixed dominant clusters to K-means grouping semantically related categories together.Source: §E.3 · pp. 30–31
• SMI exceeds 80% recall for most ICD major categories and reaches 100% precision for seven, including Infectious & parasitic, Neoplasms, Mental disorders, Musculoskeletal, and Congenital anomalies.The reported 100%-precision categories also include Symptoms & ill-defined and External causes.Source: §E.3 · p. 30, Table 11 · p. 31

E.4 Details about Rare Code Prediction

• Rare-code drug recommendation uses 98 five-level ATC medications with 10–20 occurrences each and evaluates retrieval with Recall@15.The subset restricts recommendation to these long-tailed labels because threshold-based metrics are less informative for sparse targets.Source: §E.4 · pp. 31–32

E.5 Details about Distribution shift

• SMI yields lower average absolute log-likelihood difference across hospitals than BioBERT (|∆| = 50.77 vs. |∆| = 74.06), indicating reduced distribution shift.The analysis uses a VAE density model trained on a source hospital; larger absolute log-likelihood gaps indicate greater shift.Source: §E.5 · p. 32, Figure 9 · p. 33

NeurIPS Paper Checklist

• The checklist notes that SMI focuses on expert-curated ICD and ATC systems and requires further ontology validation and code mapping for emerging classifications, new drugs, and local coding standards.Source: NeurIPS Paper Checklist · p. 35
• The checklist says the paper provides implementation details but no anonymized code release, while dataset redistribution is restricted by data-use agreements.It states that dataset, ontology, split, task, metric, baseline, and evaluation-protocol information is provided for reproducibility.Source: NeurIPS Paper Checklist · p. 36
• The checklist reports bootstrap means and standard errors for clinical prediction, no variance for deterministic relation prediction, and significance tests for drug recommendation, rare drug prediction, and distribution shift.Source: NeurIPS Paper Checklist · p. 37
• Experiments used an NVIDIA RTX PRO 6000 GPU with 96GB memory; SMI reuses frozen-LM embeddings after one computation, adding minimal overhead.Source: NeurIPS Paper Checklist · p. 38
• The study uses de-identified, controlled-access EHR data without participant interaction or new data collection, and reports no deviation from the NeurIPS Code of Ethics.Dataset access follows the original data-use agreements.Source: NeurIPS Paper Checklist · pp. 38, 40
• The authors identify benefits from improved long-tailed concept representation and prediction, while noting deployment risks from ontology choice, code mapping, and the need for clinician validation.The study uses de-identified clinical datasets and does not release patient-level data.Source: NeurIPS Paper Checklist · p. 38
• The paper credits datasets, ontologies, language models, and baselines, and plans to release precomputed SMI embeddings and construction code with documentation.MIMIC-IV and eICU access follows credentialed-use requirements.Source: NeurIPS Paper Checklist · pp. 39–40
• The method uses frozen BioBERT embeddings to obtain level-specific clinical-code representations, as described in the methodology.Source: NeurIPS Paper Checklist · p. 41