Federated Learning for ICD Classification with Lightweight Models and Pretrained Embeddings

Riepilogo Tecnico: Apprendimento Federato per la Classificazione ICD con Modelli Leggeri e Embedding Pre-addestrati

Enunciato del Problema
La classificazione automatizzata delle note cliniche in codici della Classificazione Internazionale delle Malattie (ICD) è fondamentale per la ricerca epidemiologica e l'amministrazione sanitaria, ma rimane un collo di bottiglia a causa dello sforzo manuale richiesto. Sebbene i Modelli Linguistici di Grande Formato (LLM) abbiano mostrato promesse nell'automazione di questo compito, il loro dispiegamento pratico in ambito sanitario è ostacolato dalle normative sulla privacy (es. GDPR) e dalle preoccupazioni etiche che impediscono la centralizzazione di dati sensibili dei pazienti. Gli approcci tradizionali di Apprendimento Federato (FL) richiedono spesso il fine-tuning di modelli transformer massicci su nodi distribuiti, il che è computazionalmente costoso, gravoso in termini di comunicazione e difficile da implementare in ambienti ospedalieri con risorse limitate. Inoltre, gli studi FL esistenti in ambito sanitario si concentrano spesso su compiti più semplici (es. previsione della mortalità) piuttosto che sulla complessa codifica ICD multi-etichetta.

Metodologia
Gli autori propongono una pipeline modulare in due fasi che disaccoppia la rappresentazione semantica dalla classificazione per affrontare le sfide relative alla privacy e alla scalabilità:

1. Estrazione Offline degli Embedding: Invece di eseguire il fine-tuning degli LLM attraverso le istituzioni, lo studio utilizza sei modelli di embedding open-source all'avanguardia (selezionati dalla classifica Massive Text Embedding Benchmark) per codificare i riassunti di dimissione clinica grezzi in vettori densi a dimensionalità fissa (768–1536 dimensioni). Questo processo avviene localmente e offline su ogni nodo, garantendo che nessun testo grezzo lasci la struttura. I modelli selezionati includono Qwen3-Embed-0.6B, gte-Qwen2-1.5B-instruct, gte-multilingual-base, snowflake-arctic-embed-l/m-v2.0 e nomic-embed-text-v1.5.
2. Classificazione Leggera: Gli embedding risultanti servono come input per semplici classificatori a Perceptron Multistrato (MLP) leggeri. Sono state valutate tre architetture: un MLP standard, un DeepMLP (tre livelli impilati con normalizzazione in batch e dropout) e un DeepResMLP (che incorpora connessioni residue).
3. Configurazione di Addestramento Federato: Lo studio simula un ambiente FL orizzontale utilizzando il framework Flower su 20 clienti virtuali. I dati sono partizionati dal dataset MIMIC-IV (nello specifico, riassunti di dimissione etichettati con codici ICD-9 e ICD-10). La strategia di addestramento impiega FedAvg per l'aggregazione su 30 round di comunicazione.
4. Disegno Sperimentale: Gli autori hanno condotto esperimenti estesi confrontando l'addestramento centralizzato (baseline) con l'addestramento federato. Hanno valutato le prestazioni su sei modelli di embedding e tre architetture di classificatore sotto due standard di codifica (ICD-9 e ICD-10). Inoltre, è stato eseguito uno studio di ablazione utilizzando 10 divisioni stratificate casuali per valutare la stabilità.

Contributi Chiave

• Architettura a Protezione della Privacy: Il documento dimostra una pipeline praticabile in cui il testo clinico grezzo non lascia mai il nodo locale. Congelando i modelli di embedding e addestrando solo testine di classificazione leggere in modo federato, l'approccio aderisce a requisiti rigorosi di sovranità dei dati.
• Disaccoppiamento di Rappresentazione e Classificazione: Lo studio mette in discussione la necessità del fine-tuning end-to-end di modelli massicci per la codifica ICD. Sostiene che embedding di alta qualità e congelati possono fornire una ricchezza semantica sufficiente affinché classificatori semplici (come gli MLP) raggiungano prestazioni competitive.
• Benchmarking Completo: Il lavoro fornisce una valutazione rigorosa di sei distinti modelli di embedding e tre profondità di classificatore sia in ambienti centralizzati che federati, offrendo un'analisi dettagliata di come la qualità dell'embedding rispetto alla complessità del classificatore impatti le prestazioni multi-etichetta.
• Analisi di Scalabilità: L'architettura proposta utilizza modelli con significativamente meno parametri (sotto i 7,1 milioni per il classificatore più grande) rispetto ai transformer end-to-end all'avanguardia (es. PLM-ICD con ~139 milioni di parametri), rendendola fattibile per il dispiegamento su hardware ospedaliero standard.

Risultati

• Dominanza della Qualità degli Embedding: Gli esperimenti rivelano che la scelta del modello di embedding è il determinante primario delle prestazioni predittive, superando la complessità del classificatore. Ad esempio, Qwen3-Embed-0.6B accoppiato a un semplice MLP ha ottenuto il punteggio micro F1 più alto per ICD-10 (0,3942), mentre gte-Qwen2-1.5B-instruct con DeepResMLP ha ottenuto i migliori risultati per ICD-9.
• Complessità del Classificatore: Architetture più profonde (DeepMLP, DeepResMLP) non hanno costantemente superato il semplice MLP. In molte configurazioni, la complessità aggiuntiva ha prodotto rendimenti decrescenti o addirittura cali di prestazioni, suggerendo che il carico semantico è efficacemente gestito dagli embedding pre-addestrati.
• Federato vs Centralizzato: I risultati dell'apprendimento federato hanno corrisposto strettamente alle baseline di addestramento centralizzato. Ad esempio, la configurazione Qwen3-Embed-0.6B + MLP ha ottenuto un micro F1 di 0,3944 in FL rispetto a 0,3942 nell'addestramento centralizzato per ICD-10, indicando che la decentralizzazione dei dati non degrada significativamente le prestazioni quando si utilizzano embedding standardizzati.
• Prestazioni su Codici Frequenti: Sui 10 e 50 codici ICD più frequenti, i modelli hanno ottenuto punteggi F1 significativamente più alti (spesso >0,60 per Top-10), dimostrando una forte capacità discriminativa per diagnosi comuni.
• Stabilità: Lo studio di ablazione su 10 divisioni casuali ha mostrato una bassa varianza nelle prestazioni (deviazioni standard spesso <0,003 per i punteggi F1), confermando la robustezza dell'approccio.

Significato e Affermazioni
Gli autori affermano che questo lavoro offre una "via praticabile verso sistemi di codifica medica scalabili e attenti alla privacy". Sostengono che sistemi di codifica medica ad alte prestazioni possono essere costruiti senza affidarsi ad architetture end-to-end costose e rigide che richiedono la centralizzazione di dati sensibili. Lo studio suggerisce che il "paradosso dei dati abbondanti ma isolati" in sanità può essere affrontato sfruttando localmente embedding standardizzati e congelati e addestrando collaborativamente solo testine leggere.

Il documento mantiene un tono modesto riguardo alle sue limitazioni. Riconosce che la dipendenza da embedding statici e non fine-tunati può limitare le prestazioni su codici rari o altamente specifici del dominio. Inoltre, la configurazione FL assume condizioni idealizzate (distribuzione omogenea dei dati, aggiornamenti sincroni) che potrebbero non riflettere pienamente l'eterogeneità delle reti ospedaliere reali. Gli autori concludono che, sebbene il loro approccio ottenga risultati competitivi con una frazione dei parametri dei modelli SOTA, il lavoro futuro dovrebbe esplorare strategie federate adattive, fine-tuning congiunto e l'integrazione di conoscenze mediche specifiche del dominio per migliorare ulteriormente la robustezza.