Technische Samenvatting: Hiërarchische Modellering van ICD-codes in EHR Foundation Models

Probleemstelling

Foundation models voor Elektronische Gezondheidsdossiers (EHR), zoals BEHRT en Med-BERT, behandelen International Classification of Diseases (ICD) diagnosecodes doorgaans als platte, onafhankelijke tokens. Deze aanpak negeert de klinisch betekenisvolle hiërarchische structuur die inherent is aan ICD-coderingssystemen, die ziekten organiseert in families, subcategorieën en steeds specifiekere diagnoses. Bijgevolg falen bestaande methoden voor representatieleer erin om deze ontologie expliciet te benutten, waarbij gerelateerde diagnoses als geïsoleerde symbolen worden behandeld. Dit is bijzonder problematisch in klinische data waar diagnoses vaak schaars, 'long-tailed' en instellingsspecifiek zijn. Het artikel veronderstelt dat het expliciet modelleren van deze hiërarchie als een inductieve bias de representatieleer kan verbeteren door zeldzame codes voordeel te laten trekken uit data die geassocieerd is met klinisch vergelijkbare codes, bredere zorgtrajecten te vangen en de robuustheid onder domeinverschuivingen te verbeteren.

Methodologie

De auteurs stellen twee complementaire mechanismen voor en evalueren deze voor het injecteren van de ICD-10-CM hiërarchie in hedendaagse EHR-representatieleer, gezamenlijk aangeduid als Hierarchical ICD-modellen (HICD). Ze definiëren drie geneste hiërarchie-niveaus ($G_0, G_1, G_2$) variërend van de grofste tot de fijnste granulariteit.

1. HICD-BERT (Token-niveau Injectie)

Deze aanpak breidt de BEHRT Transformer-architectuur uit. In plaats van ICD-codes als atomaire identificatoren te behandelen, verrijkt het model de standaard token-embeddings met hiërarchie-specifieke embeddings afgeleid via string-prefix truncatie.

• Mechanisme: Voor een diagnosecode voegt het model additieve embeddings toe die overeenkomen met verschillende niveaus van de codestring (bijv. een 1-karakter prefix, 2-karakter prefix, 3-karakter categorie).
• Databron: Deze methode is datagedreven en vereist geen externe toegang tot de ontologie tijdens de training; de hiërarchie wordt direct afgeleid van de codestrings.
• Training: Pre-trained middels Masked Language Modeling (MLM) op diagnose-sequenties, gevolgd door fine-tuning voor downstream taken.

2. HICD-Graph (Graph-niveau Codering)

Deze aanpak injecteert de hiërarchie op code-graph niveau voordat een Transformer encoder wordt geïnitialiseerd.

• Mechanisme: Een hybride diagnose co-occurrence graph wordt geconstrueerd. Deze combineert:
  • Empirische randen: Gewogen door Pointwise Mutual Information (PMI) gebaseerd op diagnose co-occurrence binnen patiëntbezoeken.
  • Ontologische randen: Uniform-gewicht randen die blad-diagnose-nodes verbinden met hun voorouder-nodes in de ICD-hiërarchie ($G_0, G_1, G_2$).
• Training: Een Graph Convolutional Network (GCN) wordt getraind met een link-prediction objectief om hiërarchie-bewuste code-embeddings te leren. Deze embeddings initialiseren vervolgens de embedding tabel van een patiënt-niveau Transformer encoder.

Experimentele Opzet

• Datasets: MIMIC-IV (gebruikt voor pretraining en in-domain evaluatie) en eICU (gebruikt voor cross-dataset transfer evaluatie).
• Taken:
  1. 30-daagse Heropname (MIMIC-IV): Binaire classificatie.
  2. Emergency-Admission Voorspelling (MIMIC-IV): Binaire classificatie.
  3. ICU Heropname (eICU): Binaire classificatie voor transfer evaluatie (frozen encoder, fine-tuned classifier).
• Ablatie: De studie evalueert systematisch alle acht combinaties van de drie hiërarchie-niveaus ($G_0, G_1, G_2$) voor beide architecturen tegenover een no-hierarchy baseline.

Belangrijkste Resultaten

In-Domain Prestaties (MIMIC-IV)

• Algehele Verbetering: Hiërarchie-verrijkte configuraties presteerden beter dan de no-hierarchy baseline in 26 van de 2be 28 pairwise vergelijkingen over beide modellen en taken heen.
• HICD-BERT: Prestatiewinsten werden primair gedreven door het fijnste granulariteitsniveau ($G_2$). Het middenniveau ($G_1$) alleen verbeterde de prestaties niet significant, maar droeg wel bij wanneer het gecombineerd werd met andere niveaus. De optimale configuratie varieerde per taak (bijv. alle niveaus voor heropname, enkel fijn niveau voor emergency prediction).
• HICD-Graph: Toonde een monotone verbeteringstrend waarbij het toevoegen van meer hiërarchie-niveaus progressief de prestaties verhoogde. De graph-gebaseerde architectuur integreerde effectief informatie van meerdere granulariteitsniveaus simultaan.

Cross-Dataset Transfer (MIMIC-IV naar eICU)

• HICD-Graph: Demonstreerde robuuste transfereerbaarheid. Zes van de zeven hiërarchie-varianten vertoonden een significante verbetering ten opzien van de baseline, waarbij de all-levels configuratie de sterkste AUROC (0.587) behaalde. De graph-gebaseerde relationele structuur leek meer invariant voor datasetverschuivingen.
• HICD-BERT: De voordelen van de hiërarchie verdwenen grotendeels onder distributieverschuiving. Alleen de meest grove single-level configuratie ($G_0$) vertoonde een significante verbetering ten opzien van de baseline. De auteurs suggereren dat de sequentiële tokenpatronen geleerd door HICD-BERT nauwer gekoppeld zijn aan de specifieke bezoekstructuren van MIMIC-IV, terwijl de ontologie-verankerde graph-structuur algemener is.

Embedding Analyse

• Semantische Clustering: Hiërarchie-bewuste modellen produceerden coherenteerdere code-clusters. Specifiek resulteerde de $G_1+G_2$ configuratie in HICD-Graph in de meest compacte lokale code-clusters in de embedding ruimte.
• Correlatie met Prestaties: De configuraties die de meest compacte klinisch betekenisvolle clusters produceerden (bijv. codes die dezelfde 2-karakter prefix delen) behaalden ook de sterkste downstream voorspellende prestaties, wat suggereert dat hiërarchie de prestaties verbetert door de embedding geometrie zo te vormen dat deze klinische relaties respecteert.

Betekenis en Claims

Het artikel claimt dat de hiërarchische structuur aanwezig in ICD-coding systemen een onderbenut signaal is in EHR foundation models. De belangrijkste bijdragen en claims zijn:

1. Hiërarchie als een Robuuste Inductieve Bias: Het expliciet coderen van de ICD-hiërarchie verbetert de downstream voorspellende prestaties aanzienlijk over verschillende architecturen (Transformer-gebaseerd en Graph-gebaseerd) en taken.
2. Granulariteitsafhankelijkheid: De bruikbaarheid van hiërarchie-niveaus is niet uniform; het optimale niveau hangt af van zowel de specifieke modelleringsaanpak (token-niveau versus graph-niveau) als de downstream taak.
3. Transfereerbaarheid van Structuur: Graph-gebaseerde hiërarchie-codering is robuuster transfereerbaar over datasets dan token-niveau injectie, wat suggereert dat ontologie-afgeleide relationele structuren fundamentele klinische relaties vastleggen die verder gaan dan sequentiële tokenpatronen.
4. Lage Kosten Integratie: De voordelen worden behaald met minimale architecturale wijzigingen (additieve embeddings of extra graph-randen) zonder de kern van de pretraining objectieven of pipelines te veranderen.
5. Generaliseerbaarheid: De inzichten met betrekking tot de waarde van gestructureerde klinische ontologieën strekken zich uit voorbij ICD-10-CM en kunnen worden toegepast op andere gestructureerde vocabulaire (bijv. ATC-codes voor medicatie) en representatieleer methoden.

De auteurs blijven bescheiden en merken op dat hoewel de absolute AUROC-verbeteringen bescheiden zijn, ze statistisch robuust en consistent zijn. Ze erkennen beperkingen met betrekking tot de generaliseerbaarheid naar andere klinische uitkomsten, instellingen en EHR-systemen, en merken op dat hun hiërarchie-niveaus handmatig zijn gedefinieerd in plaats van automatisch geleerd.