Probabilistic Feature Imputation and Uncertainty-Aware Multimodal Federated Aggregation

1. Il Problema: Eterogeneità delle Modalità nell'Apprendimento Federato

L'apprendimento federato (Federated Learning - FL) è fondamentale in ambito sanitario per addestrare modelli collaborativi preservando la privacy dei dati (es. immagini radiologiche e referti testuali) senza trasferirli centralmente. Tuttavia, sorge una sfida critica nota come eterogeneità delle modalità:

• Disparità delle risorse: Molti siti clinici (es. piccole cliniche o ospedali rurali) dispongono solo di un sottoinsieme di modalità (es. solo immagini a raggi X), mentre altri (centri accademici) possiedono dati multimodali completi (immagini + referti testuali strutturati).
• Limitazioni degli approcci esistenti: Le soluzioni attuali utilizzano reti di imputazione delle caratteristiche deterministiche per sintetizzare le modalità mancanti. Il difetto principale è che queste reti producono stime puntuali (valori fissi) senza fornire misure di affidabilità. In contesti medici critici, un modello deterministico può "allucinare" un referto testuale per un'immagine ambigua con alta confidenza, portando a diagnosi errate senza alcun avviso di incertezza.

2. Metodologia Proposta: P-FIN e Fed-UQ-Avg

Gli autori propongono un framework innovativo composto da due pilastri principali: la Probabilistic Feature Imputation Network (P-FIN) e la strategia di aggregazione Fed-UQ-Avg.

A. Probabilistic Feature Imputation Network (P-FIN)

Invece di prevedere un vettore fisso per la modalità mancante, P-FIN modella la distribuzione condizionale $p(z_T | z_I)$ come una Gaussiana eteroschedastica.

• Architettura: Utilizza un encoder per le immagini (ResNet-50) e uno per il testo (BERT). La rete di imputazione prende le caratteristiche dell'immagine e le mappa su una distribuzione definita da una media ($\mu$) e una varianza ($\sigma^2$).
• Addestramento Calibrato: Per evitare che il modello minimizzi la perdita semplicemente prevedendo varianze infinite (un problema noto come "variance explosion"), viene utilizzata una funzione di perdita specializzata chiamata $\beta$-NLL (Negative Log-Likelihood). Questa funzione applica un operatore stop-gradient per forzare il modello a ridurre l'errore di previsione invece di inflazionare l'incertezza.
• Fusione Locale Consapevole dell'Incertezza:
  • Gating: Prima della fusione delle caratteristiche, viene calcolato un "gate" basato sulla varianza prevista: $g = \text{sigmoid}(-\log \sigma^2)$. Se l'incertezza è alta, il gate si avvicina a zero, attenuando le caratteristiche imputate inaffidabili.
  • Attenzione Incrociata: Le caratteristiche dell'immagine e quelle del testo (filtrate dal gate) vengono fuse tramite meccanismi di attenzione bidirezionale, permettendo al modello di pesare dinamicamente il contributo di ciascuna modalità.

B. Aggregazione Federata Consapevole dell'Incertezza (Fed-UQ-Avg)

La strategia di aggregazione standard (FedAvg) pesa i contributi dei client solo in base alla dimensione del dataset locale. Questo è rischioso se un client con molti dati ha un'imputazione di bassa qualità.

• Nuova Strategia: Fed-UQ-Avg introduce un peso combinato che tiene conto sia della dimensione dei dati ($w_{data}$) sia della fiducia nell'imputazione ($w_{conf}$).
• Calcolo della Fiducia: La fiducia è calcolata esponenzialmente in base all'incertezza media del client ($\bar{\sigma}^2_k$): $conf_k = \exp(-\bar{\sigma}^2_k / T)$.
• Risultato: I client con bassa incertezza (alta qualità di imputazione) ricevono un peso maggiore nell'aggiornamento del modello globale, mentre quelli con alta incertezza contribuiscono meno, proteggendo il modello globale da gradienti rumorosi.

3. Contributi Chiave

1. Imputazione Probabilistica: Sostituzione delle stime deterministiche con distribuzioni probabilistiche calibrate, fornendo una misura esplicita dell'affidabilità delle caratteristiche sintetizzate.
2. Meccanismo di Gating Locale: Un metodo per sopprimere dinamicamente le caratteristiche imputate inaffidabili prima della classificazione, prevenendo la propagazione di errori.
3. Aggregazione Fed-UQ-Avg: Una nuova strategia di aggregazione federata che integra la qualità dell'imputazione nel processo di fusione globale, superando i limiti del semplice dimensionamento basato sui dati.
4. Validazione Clinica: Dimostrazione che l'approccio "imparare a non fidarsi" (learning to distrust) è superiore ai metodi deterministici in scenari con dati multimodali incompleti.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su tre dataset pubblici di radiografie toraciche: CheXpert, NIH Open-I e PadChest, simulando ambienti federati con diversi rapporti di clienti unimodali/multimodali (8:2, 6:4, 4:6).

• Prestazioni: Il modello proposto (P-FIN + Fed-UQ-Avg) ha mostrato miglioramenti costanti rispetto a tutti i baseline deterministici.
  • Nella configurazione più difficile (80% di clienti con dati unimodali, rapporto 8:2), si è registrato un guadagno assoluto di +5.36% nell'AUC rispetto al metodo deterministico più avanzato (FIN + FedAvg).
  • L'aggiunta della componente di aggregazione Fed-UQ-Avg ha portato ulteriori guadagni significativi (circa +2.4% - +2.8%) rispetto all'uso di P-FIN con FedAvg standard.
• Calibrazione: L'analisi ha confermato che le stime di incertezza sono ben calibrate (Expected Calibration Error - ECE = 0.0422) e correlate positivamente con l'errore di ricostruzione: alta incertezza predice correttamente bassa qualità dell'imputazione.
• Robustezza: Il modello ha dimostrato una maggiore resilienza nella gestione di dati eterogenei e scarsità di modalità, evitando che i client con dati incompleti e di bassa qualità degradassero il modello globale.

5. Significato e Implicazioni

Questo lavoro è significativo per diversi motivi:

• Sicurezza Clinica: Allinea lo sviluppo di AI medica alle normative (es. FDA) che richiedono la quantificazione dell'incertezza. Permette ai radiologi di sapere quando fidarsi della sintesi automatica e quando intervenire manualmente.
• Affidabilità Federata: Risolve il problema della "spazzatura in entrata, spazzatura in uscita" (garbage in, garbage out) nell'apprendimento federato multimodale, garantendo che i client con dati incompleti non danneggino il modello collettivo.
• Futuro dell'AI Medica: Fornisce una base solida per sistemi di apprendimento federato più resilienti e affidabili, essenziali per l'implementazione su larga scala di strumenti di supporto alla decisione clinica in contesti reali con risorse disomogenee.

In sintesi, il paper introduce un paradigma in cui l'incertezza non è un difetto da nascondere, ma una risorsa gestibile che migliora sia la fusione locale delle caratteristiche sia la collaborazione globale tra istituzioni mediche.