Federated Learning for ICD Classification with Lightweight Models and Pretrained Embeddings

Technische Samenvatting: Federated Learning voor ICD-classificatie met Lichtgewicht Modellen en Voorgereed Trained Embeddings

Probleemstelling
Geautomatiseerde classificatie van klinische notities naar International Classification of Diseases (ICD)-codes is cruciaal voor epidemiologisch onderzoek en gezondheidszorgadministratie, maar blijft een knelpunt vanwege de vereiste handmatige inspanning. Hoewel Large Language Models (LLM's) veelbelovend zijn gebleken voor het automatiseren van deze taak, wordt hun praktische implementatie in de gezondheidszorg gehinderd door privacyreguleringen (bijv. AVG) en ethische bezwaren die centralisatie van gevoelige patiëntgegevens voorkomen. Traditionele Federated Learning (FL)-benaderingen vereisen vaak het fijnafstemmen van massale transformer-modellen over gedistribueerde knooppunten, wat rekenkundig duur, communicatie-intensief en moeilijk te implementeren is in hospitalen met beperkte middelen. Bovendien richten bestaande FL-studies in de gezondheidszorg zich vaak op eenvoudigere taken (bijv. mortaliteitsvoorspelling) in plaats van complexe, multi-label ICD-codering.

Methodologie
De auteurs stellen een modulaire, twee-traps pijplijn voor die semantische representatie ontkoppelt van classificatie om uitdagingen rond privacy en schaalbaarheid aan te pakken:

1. Offline Embedding-extractie: In plaats van LLM's over instellingen heen fijn te stemmen, maakt de studie gebruik van zes state-of-the-art, open-source embedding-modellen (geselecteerd uit de Massive Text Embedding Benchmark-leaderboard) om ruwe klinische ontslagrapporten te coderen in vaste dimensionale dichte vectoren (768–1536 dimensies). Dit proces vindt lokaal en offline plaats bij elk knooppunt, zodat er geen ruwe tekst de faciliteit verlaat. De geselecteerde modellen omvatten Qwen3-Embed-0.6B, gte-Qwen2-1.5B-instruct, gte-multilingual-base, snowflake-arctic-embed-l/m-v2.0 en nomic-embed-text-v1.5.
2. Lichtgewicht Classificatie: De resulterende embeddings dienen als invoer voor eenvoudige, lichtgewicht Multi-Layer Perceptron (MLP)-classificators. Drie architecturen werden geëvalueerd: een standaard MLP, een DeepMLP (drie gestapelde lagen met batch-normalisatie en dropout) en een DeepResMLP (met residual-verbindingen).
3. Federated Training-opstelling: De studie simuleert een horizontale FL-omgeving met het Flower-framework over 20 virtuele clients. Data is gepartitioneerd vanuit de MIMIC-IV-dataset (specifiek ontslagrapporten gelabeld met ICD-9- en ICD-10-codes). De trainingsstrategie maakt gebruik van FedAvg voor aggregatie over 30 communicatierondes.
4. Experimenteel Ontwerp: De auteurs hebben uitgebreide experimenten uitgevoerd waarbij centrale training (baseline) werd vergeleken met federale training. Ze evalueerden de prestaties over zes embedding-modellen en drie classifier-architecturen onder twee coderingsstandaarden (ICD-9 en ICD-10). Daarnaast werd een ablatiestudie uitgevoerd met 10 willekeurige gestratificeerde splitsingen om stabiliteit te beoordelen.

Belangrijkste Bijdragen

• Privacy-bewuste Architectuur: Het artikel demonstreert een levensvatbare pijplijn waarbij ruwe klinische tekst nooit de lokale knooppunt verlaat. Door de embedding-modellen te bevriezen en alleen lichtgewicht classificatiekoppen op een federale manier te trainen, voldoet de aanpak aan strenge eisen voor data-soevereiniteit.
• Ontkoppeling van Representatie en Classificatie: De studie daagt de noodzaak uit van end-to-end fijnafstemming van massale modellen voor ICD-codering. Het stelt dat hoogwaardige, bevroren embeddings voldoende semantische rijkdom kunnen bieden zodat eenvoudige classificatoren (zoals MLP's) concurrerende prestaties kunnen behalen.
• Uitgebreide Benchmarking: Het werk biedt een rigoureuze evaluatie van zes verschillende embedding-modellen en drie classifier-dieptes onder zowel centrale als federale settings, met een gedetailleerde analyse van hoe de kwaliteit van embeddings versus de complexiteit van de classifier de multi-label-prestaties beïnvloedt.
• Schaalbaarheidsanalyse: De voorgestelde architectuur maakt gebruik van modellen met aanzienlijk minder parameters (minder dan 7,1 miljoen voor de grootste classifier) in vergelijking met state-of-the-art end-to-end transformers (bijv. PLM-ICD met ~139 miljoen parameters), waardoor implementatie op standaard hospitalhardware haalbaar is.

Resultaten

• Dominantie van Embedding-kwaliteit: De experimenten tonen aan dat de keuze van het embedding-model de primaire determinant is van voorspellende prestaties, en dat dit de complexiteit van de classifier overtreft. Bijvoorbeeld, Qwen3-Embed-0.6B gekoppeld aan een eenvoudige MLP behaalde de hoogste micro F1-score voor ICD-10 (0,3942), terwijl gte-Qwen2-1.5B-instruct met DeepResMLP het beste presteerde voor ICD-9.
• Complexiteit van de Classifier: Diepere architecturen (DeepMLP, DeepResMLP) presteerden niet consequent beter dan de eenvoudige MLP. In veel configuraties leverde de toegevoegde complexiteit afnemende meeropbrengsten op of zelfs prestatiedalingen, wat suggereert dat de semantische last effectief wordt opgevangen door de voorgereed getrainde embeddings.
• Federated versus Centraal: De resultaten van federated learning kwamen nauw overeen met baselines voor centrale training. Bijvoorbeeld, de configuratie Qwen3-Embed-0.6B + MLP behaalde een micro F1 van 0,3944 in FL versus 0,3942 bij centrale training voor ICD-10, wat aangeeft dat data-decentralisatie de prestaties niet significant verslechtert bij gebruik van gestandaardiseerde embeddings.
• Prestaties op Frequente Codes: Op de 10 en 50 meest frequente ICD-codes behaalden de modellen aanzienlijk hogere F1-scores (vaak >0,60 voor Top-10), wat een sterke discriminerende capaciteit voor veelvoorkomende diagnoses aantoont.
• Stabiliteit: De ablatiestudie over 10 willekeurige splitsingen toonde een lage variantie in prestaties (standaardafwijkingen vaak <0,003 voor F1-scores), wat de robuustheid van de aanpak bevestigt.

Betekenis en Claims
De auteurs claimen dat dit werk een "levensvatbare weg biedt naar schaalbare, privacy-bewuste systemen voor medische codering". Zij betogen dat hoogpresterende systemen voor medische codering kunnen worden gebouwd zonder te vertrouwen op kostbare, starre, end-to-end architecturen die centralisatie van gevoelige data vereisen. De studie suggereert dat het "paradox van overvloedige maar geïsoleerde data" in de gezondheidszorg kan worden aangepakt door lokaal bevroren, gestandaardiseerde embeddings te benutten en alleen lichtgewicht koppen collaboratief te trainen.

Het artikel behoudt een bescheiden toon met betrekking tot zijn beperkingen. Het erkent dat de afhankelijkheid van statische, niet-fijnafgestemde embeddings de prestaties kan beperken voor zeldzame of zeer domeinspecifieke codes. Bovendien gaat de FL-opstelling uit van geïdealiseerde omstandigheden (homogene data-distributie, synchrone updates) die mogelijk niet volledig de heterogeniteit van real-world hospitalnetwerken weerspiegelen. De auteurs concluderen dat hoewel hun aanpak concurrerende resultaten bereikt met een fractie van de parameters van SOTA-modellen, toekomstig werk adaptieve federale strategieën, gezamenlijke fijnafstemming en integratie van domeinspecifieke medische kennis moet verkennen om de robuustheid verder te verbeteren.