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

AI summary

Overview

• Research area: Biomedical natural language processing and clinical concept representation, combining medical ontologies (ICD, ATC, CCS) with frozen biomedical language model embeddings.
• Technical level: Intermediate. The framework itself is conceptually simple (path lookup plus embedding addition), but the paper includes a formal theoretical analysis using signal-to-noise ratio arguments and assumptions about frequency-dependent embedding noise.
• Scope in one sentence: The paper introduces SMI, a training-free method that represents each medical concept as its coarse-to-fine ontology path and aggregates frozen language model embeddings along that path, then evaluates it on clinical relation prediction, clustering, MIMIC-IV prediction tasks, and cross-hospital distribution shift on eICU.

What This Paper Is About

Biomedical language models such as BioBERT encode each clinical code mainly from its textual description, so they do not explicitly capture the hierarchical structure that ICD and ATC coding systems impose on diagnoses, procedures, and medications. This weakness is worst for rare, fine-grained codes in long-tailed clinical vocabularies. The paper proposes SMI (Semantic Medical Identifier), a training-free framework that defines each concept's "Semantic ID" as its coarse-to-fine path through an expert-curated ontology and then aggregates frozen biomedical LM embeddings of the codes along that path into a single hierarchy-aware representation.

Key Contributions

1. A training-free Semantic ID framework. SMI derives Semantic IDs directly from expert-defined medical ontologies — five-level ICD/CCS hierarchies for diagnoses and procedures, and the inherent five-level ATC hierarchy for medications — with no additional quantizer or tokenization model trained.
2. Hierarchy-aware concept embeddings by aggregation. For each concept, SMI sums the frozen LM embeddings of its ancestor codes from coarse to fine levels. A level-weighted variant, SMI-W, replaces uniform aggregation with global level weights selected on validation data.
3. A theoretical analysis of representation collapse. The paper models hierarchy-agnostic LM embeddings as frequency-dependent noisy estimates of latent concept semantics and proves that SMI reduces estimation noise variance, yielding strictly higher pairwise signal-to-noise ratio, particularly for rare and fine-grained concepts.
4. Evaluation across relation prediction, ontology structure, and EHR tasks. SMI is tested on ICD→ICD, ATC→ATC, and ICD→ATC relation prediction, on clustering of diagnosis embeddings, on MIMIC-IV clinical prediction, and on eICU cross-hospital distribution shift.

Main Findings

• Hierarchy-agnostic LMs degrade at finer granularity. In ICD 1-hop prediction, BioBERT reaches 85.1% at the coarsest levels but drops to 58.6% at the finest level. SMI maintains stronger and more stable accuracy across concept-frequency bins, while BioBERT accuracy declines as concept frequency decreases.
• SMI improves clinical relation prediction. On Table 1, SMI-W achieves the highest F-scores for ICD (71.5%), ATC (77.5%), and cross-ontology ICD→ATC (57.6%). The strongest baseline on ICD and ATC is SapBERT (67.1% and 74.5%), and the strongest baseline on ICD→ATC is HiTs with an F-score of 53.2%. (BioBERT's ICD F-score is listed as 44.2 in Table 1 and 43.4 in the ablation table.)
• SMI embeddings align better with ontology structure. K-means clustering of ICD diagnosis embeddings gives SMI versus BioBERT: NMI 0.77 vs. 0.39, ARI 0.56 vs. 0.25, Purity 0.78 vs. 0.51, per-category 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 (lower is better).
• UMAP visualization shows clearer group separation. With ICD codes of frequency at least 5, plotted across 13 major categories, BioBERT embeddings are highly mixed, while SMI forms distinct clusters by major category that further split into subcategory-aligned subgroups.
• MIMIC-IV prediction improves, especially for drugs. Across 123,488 patients, SMI-W reaches the best mortality AUPRC and AUROC (45.1, 88.0), readmission AUROC (72.9), and drug recommendation AUPRC, F-score, and Jaccard (77.6%, 65.3%, 50.5%), outperforming the strongest baseline on drug recommendation at p = 0.015.
• Rare drug prediction gains are the largest. Using 98 ATC codes with frequency between 10 and 20 and Recall of the top 15 predictions, SMI-W achieves Recall@15 of 44.2%, significantly above the strongest baseline G-BERT at 40.7% (p = 0.0081).
• SMI is more robust across hospitals on eICU. Fitting hospital-specific VAE density models on patient embeddings, the cross-hospital shift is smaller for SMI than BioBERT (50.77 vs. 74.06), with a Wilcoxon signed-rank paired test over off-diagonal hospital pairs giving p = 1.97 × 10^-5.
• Weighting levels helps. Ablations show SMI outperforms BioBERT on ICD, ATC, and ICD→ATC relation prediction, and SMI-W further improves over uniform SMI, indicating that hierarchy levels contribute unequally.
• Theoretical results. Under the frequency-dependent noise assumption, Theorem 3.3 states that fine-grained concepts become inseparable under hierarchy-agnostic LM embeddings; Proposition 3.6 shows SMI preserves separability against non-parent concepts as the child's frequency goes to zero; Theorem 3.7 and Corollary 3.8 show weighted and uniform SMI strictly reduce estimation noise variance and raise pairwise SNR.
• Empirical support for semantic decomposition. An offset consistency analysis on real ICD-9 parent-child pairs reports mean cosine similarity of 0.655 for offsets sharing the same clinical modifier versus 0.080 for offsets from different modifier groups (p = 2.0 × 10^-4).

Methodology in Plain English

1. Build five-level ontology paths. Diagnoses follow Major category → Sub-category → CCS code → ICD-CM code (integer) → ICD-CM code (decimal). Procedures follow major category, sub-category, CCS code, prefix ICD-PCS code (first three characters), and full ICD-PCS code. Medications use ATC's Anatomical Main Group → Therapeutic Subgroup → Pharmacological Subgroup → Chemical Subgroup → Chemical Substance.
2. Define the Semantic ID as the path. Each concept's Semantic ID is the ordered sequence of its ancestor codes from coarse to fine. For example, the ICD-10-CM code E11.2 follows the path [3, 3.3, 50, E11, E11.2].
3. Encode every code at every level. A frozen biomedical LM encodes each clinical code's textual description, and mean pooling over token embeddings produces a level-specific code embedding. These are stored in per-level embedding tables.
4. Aggregate the path. The concept representation is the sum of the level-specific embeddings of the codes along its Semantic ID. No trainable parameters are introduced.
5. Optionally reweight levels. SMI-W replaces the uniform sum with global, non-negative level weights that sum to L and are selected on the validation set; when all weights equal 1 it reduces to plain SMI.
6. Evaluate broadly. The paper runs within-ontology and cross-ontology relation prediction, UMAP visualization and K-means clustering of ICD embeddings, MIMIC-IV mortality (death within 90 days after discharge), readmission (within 15 days after discharge), drug recommendation, and rare drug prediction, then uses eICU with hospital as the grouping axis and a VAE-based density model to measure the log-likelihood gap between source and target hospitals.

Why This Matters

• Research impact: SMI offers a way to inject expert ontology structure into clinical representations without training a structure-aware encoder or a quantizer, and it provides a theoretical account of why hierarchy-agnostic embeddings collapse for long-tailed codes. This reframes ontology structure as a variance-reduction mechanism rather than an extra supervision signal.
• Real-world application: drug recommendation. A model that predicts medications from patient history, with improved performance on rare drugs, has direct clinical decision-support relevance.
• Real-world application: patient risk stratification. Improved mortality and readmission prediction supports earlier identification of high-risk patients.
• Real-world application: cross-site deployment. Because SMI reduces cross-hospital distribution shift on eICU, representations may transfer more reliably when a model trained at one institution is applied at another.
• Real-world application: clinical code analytics. Ontology-consistent concept embeddings support clustering and comparison of diagnoses, which can aid cohort definition and vocabulary curation.
• Industry relevance: EHR vendors, hospital systems, insurers, and health data platforms that must handle thousands of hierarchical billing and diagnosis codes could adopt SMI as a drop-in representation layer over existing biomedical LMs, since it requires no additional training.

Future Directions

• Ontology harmonization across coding standards. The paper notes that Canadian systems use ICD-10-CA or Drug Identification Numbers, which require careful mapping before deployment; future work will pursue clinician-validated harmonization across coding systems.
• Emerging diseases, new drugs, and local coding standards. Applying SMI to newly introduced classifications and local vocabularies requires additional ontology validation and code mapping, an open problem the paper explicitly flags.
• Missing-link discovery. The authors list discovering missing ontology links as a future direction.
• Unified multi-ontology integration. Future work will study how to jointly integrate multiple ontologies within a single SMI framework, and extend it to broader biomedical knowledge resources.

Target Audience

• Clinical NLP and medical informatics researchers working on biomedical concept embeddings, ontology-aware representation, or EHR-based prediction.
• Machine learning researchers interested in retrieval-style Semantic IDs, hierarchical representation learning, and theory linking frequency to embedding noise.
• Applied data scientists and engineers in healthcare organizations who need task-agnostic, training-free concept representations for downstream models.
• Graduate students who want a readable example of combining an expert knowledge structure with frozen language model embeddings and validating it both theoretically and empirically.
• Readers should be comfortable with basic embedding concepts, ontology terminology such as ICD and ATC, and standard evaluation metrics like AUPRC, AUROC, and clustering indices.

Note: several implementation details (appendices C, D.4, E.2, E.3, and E.5) are referenced in the paper but were truncated in the provided content, so the specific hyperparameters, encoder configuration, and mapping procedures are not reported here.

Authors’ 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.