🤖 AI

Hierarchical Modeling of ICD Codes in EHR Foundation Models

This paper demonstrates that explicitly incorporating the hierarchical structure of ICD-10-CM codes into EHR foundation models, through either token augmentation or graph-based edge injection, significantly improves both in-domain and cross-dataset clinical prediction performance compared to treating codes as flat tokens.

Original authors: Megha Thukral, Dong Gyun Kang, Rudra Pratap Singh, Shruthi Kashinath Hiremath, Katrin Hänsel, Thomas Plötz

Published 2026-06-16

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Megha Thukral, Dong Gyun Kang, Rudra Pratap Singh, Shruthi Kashinath Hiremath, Katrin Hänsel, Thomas Plötz

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a computer how to understand a patient's medical history. The computer reads a long list of codes (like "S72.001A") that doctors use to describe illnesses.

The Problem: The "Flat" View
Currently, most computer models treat these codes like a list of random, unrelated words. It's as if the computer sees "Apple," "Car," and "Zebra" as three completely separate items with no connection. In reality, medical codes are organized like a massive, detailed family tree. A code for a specific type of broken leg is part of a "leg injury" group, which is part of a "bone injury" group, which is part of the "injury" family.

The paper argues that by ignoring this family tree, computers are missing a huge clue about how diseases are related.

The Solution: Two New Ways to Teach the Computer
The researchers tried two different methods to force the computer to respect this family tree structure:

The "Label Sticker" Method (HICD-BERT):
Imagine you have a box of toys. Instead of just writing "Toy" on the box, you add stickers that say "Toy," "Action Figure," "Superhero," and "Batman."
In this method, the computer doesn't just see the specific code; it also sees "stickers" representing the broader categories the code belongs to. It learns that a specific code is part of a bigger group, just like seeing the "Batman" sticker helps you understand the toy better.
The "Social Network" Method (HICD-Graph):
Imagine drawing a map where every disease is a city. Usually, you only draw roads between cities that often appear together in the same patient's record.
In this method, the researchers added extra roads to the map. They drew lines connecting a specific disease city to its "parent" cities (the broader groups) based on the official medical family tree. This creates a hybrid map that knows both which diseases happen together and how they are related by definition.

What They Found
The researchers tested these methods on two huge databases of real patient records (one from Boston hospitals and one from ICUs across the US).

Better Predictions: In almost every test, the models that used the "family tree" structure were better at predicting future health events (like if a patient would be readmitted to the hospital or if they needed emergency care) than the models that ignored the structure.
The "Goldilocks" Zone: They found that the "best" level of detail depends on the job.
- For the "Label Sticker" method, the most specific details (the finest level of the tree) were usually the most helpful.
- For the "Social Network" method, using all levels of the tree (from the broadest family down to the specific code) worked best.
The "Traveling" Test: They trained the computer on one set of hospital data and then asked it to predict outcomes for a completely different set of hospitals it had never seen before.
- The "Social Network" method traveled very well; it kept its accuracy because it learned the universal rules of how diseases are related.
- The "Label Sticker" method struggled to travel; it seemed to have memorized the specific habits of the first hospital rather than the general rules of medicine.

The Bottom Line
The paper concludes that medical codes aren't just random labels; they are a structured language. By explicitly teaching computers to understand the "family tree" of these codes, we can build smarter, more reliable AI that understands the context of a patient's illness, not just the isolated facts. The researchers showed that you don't need to completely rebuild the computer's brain to do this; you just need to add a few "guideposts" (stickers or extra roads) to help it navigate the medical landscape.

Technical Summary: Hierarchical Modeling of ICD Codes in EHR Foundation Models

Problem Statement

Electronic Health Record (EHR) foundation models, such as BEHRT and Med-BERT, typically treat International Classification of Diseases (ICD) diagnosis codes as flat, independent tokens. This approach overlooks the clinically meaningful hierarchical structure inherent in ICD coding systems, which organizes diseases into families, subcategories, and increasingly specific diagnoses. Consequently, existing representation learning methods fail to explicitly exploit this ontology, treating related diagnoses as isolated symbols. This is particularly problematic in clinical data where diagnoses are often sparse, long-tailed, and institution-specific. The paper hypothesizes that explicitly modeling this hierarchy as an inductive bias could improve representation learning by allowing rare codes to benefit from data associated with clinically similar codes, capturing broader care pathways, and improving robustness under domain shifts.

Methodology

The authors propose and evaluate two complementary mechanisms for injecting ICD-10-CM hierarchy into contemporary EHR representation learning, collectively termed Hierarchical ICD models (HICD). They define three nested hierarchy levels ( $G_0, G_1, G_2$ ) ranging from coarsest to finest granularity.

1. HICD-BERT (Token-Level Injection)

This approach extends the BEHRT Transformer architecture. Instead of treating ICD codes as atomic identifiers, the model augments the standard token embeddings with hierarchy-specific embeddings derived via string-prefix truncation.

Mechanism: For a diagnosis code, the model adds additive embeddings corresponding to different levels of the code string (e.g., 1-character prefix, 2-character prefix, 3-character category).
Data Source: This method is data-driven, requiring no external ontology access at training time; the hierarchy is derived directly from the code strings.
Training: Pretrained using Masked Language Modeling (MLM) on diagnosis sequences, followed by fine-tuning for downstream tasks.

2. HICD-Graph (Graph-Level Encoding)

This approach injects hierarchy at the code-graph level before initializing a Transformer encoder.

Mechanism: A hybrid diagnosis co-occurrence graph is constructed. It combines:
- Empirical edges: Weighted by Pointwise Mutual Information (PMI) based on diagnosis co-occurrence within patient visits.
- Ontological edges: Uniform-weight edges connecting leaf diagnosis nodes to their ancestor nodes in the ICD hierarchy ( $G_0, G_1, G_2$ ).
Training: A Graph Convolutional Network (GCN) is trained with a link-prediction objective to learn hierarchy-aware code embeddings. These embeddings then initialize the embedding table of a patient-level Transformer encoder.

Experimental Setup

Datasets: MIMIC-IV (used for pretraining and in-domain evaluation) and eICU (used for cross-dataset transfer evaluation).
Tasks:
1. 30-Day Readmission (MIMIC-IV): Binary classification.
2. Emergency-Admission Prediction (MIMIC-IV): Binary classification.
3. ICU Readmission (eICU): Binary classification for transfer evaluation (frozen encoder, fine-tuned classifier).
Ablation: The study systematically evaluates all eight combinations of the three hierarchy levels ( $G_0, G_1, G_2$ ) for both architectures against a no-hierarchy baseline.

Key Results

In-Domain Performance (MIMIC-IV)

Overall Improvement: Hierarchy-augmented configurations outperformed the no-hierarchy baseline in 26 of 28 pairwise comparisons across both models and tasks.
HICD-BERT: Performance gains were driven primarily by the finest granularity level ( $G_2$ ). The mid-level ( $G_1$ ) alone did not significantly improve performance but contributed when combined with other levels. The optimal configuration varied by task (e.g., all levels for readmission, single fine level for emergency prediction).
HICD-Graph: Showed a monotonic improvement trend where adding more hierarchy levels progressively increased performance. The graph-based architecture effectively integrated information from multiple granularity levels simultaneously.

Cross-Dataset Transfer (MIMIC-IV to eICU)

HICD-Graph: Demonstrated robust transferability. Six of seven hierarchy variants significantly improved over the baseline, with the all-levels configuration achieving the strongest AUROC (0.587). The graph-based relational structure appeared more invariant to dataset shifts.
HICD-BERT: Hierarchy benefits largely vanished under distribution shift. Only the coarsest single-level configuration ( $G_0$ ) showed a significant improvement over the baseline. The authors suggest that sequential token patterns learned by HICD-BERT are more tightly coupled to the specific visit structures of MIMIC-IV, whereas the ontology-anchored graph structure is more generalizable.

Embedding Analysis

Semantic Clustering: Hierarchy-aware models produced more coherent code clusters. Specifically, the $G_1+G_2$ configuration in HICD-Graph resulted in the tightest local code clusters in embedding space.
Correlation with Performance: The configurations that produced the tightest clinically meaningful clusters (e.g., codes sharing the same 2-character prefix) also achieved the strongest downstream predictive performance, suggesting that hierarchy improves performance by shaping the embedding geometry to respect clinical relationships.

Significance and Claims

The paper claims that the hierarchical structure present in ICD coding systems is an underutilized signal in EHR foundation models. The key contributions and claims include:

Hierarchy as a Robust Inductive Bias: Explicitly encoding ICD hierarchy significantly improves downstream prediction performance across different architectures (Transformer-based and Graph-based) and tasks.
Granularity Dependence: The utility of hierarchy levels is not uniform; the optimal level depends on both the specific modeling approach (token-level vs. graph-level) and the downstream task.
Transferability of Structure: Graph-based hierarchy encoding transfers more robustly across datasets than token-level injection, suggesting that ontology-derived relational structures capture more fundamental clinical relationships than sequential token patterns.
Low-Cost Integration: The benefits are achieved with minimal architectural modifications (additive embeddings or additional graph edges) without changing core pretraining objectives or pipelines.
Generalizability: The insights regarding the value of structured clinical ontologies extend beyond ICD-10-CM and can be applied to other structured vocabularies (e.g., ATC codes for medications) and representation learning methods.

The authors remain modest, noting that while the absolute AUROC improvements are modest, they are statistically robust and consistent. They acknowledge limitations regarding the generalizability to other clinical outcomes, institutions, and EHR systems, and note that their hierarchy levels were manually defined rather than learned automatically.