MedFormer-UR: Uncertainty-Routed Transformer for Medical Image Classification

1. Il Problema

L'integrazione clinica dei modelli di deep learning, in particolare i Vision Transformer (ViT) come MedFormer, è ostacolata dalla mancanza di affidabilità e trasparenza. Sebbene questi modelli raggiungano alte prestazioni in termini di accuratezza, presentano limiti critici per l'uso medico:

• Sovraconfidenza: Tendono a produrre previsioni eccessivamente sicure anche su input ambigui o fuori distribuzione (out-of-distribution).
• Scarsa calibrazione: Le probabilità in uscita non riflettono accuratamente la probabilità reale di correttezza (alto Expected Calibration Error - ECE).
• Natura "Black Box": Mancano di meccanismi interpretabili che spieghino il ragionamento alla base di una decisione.
• Dati clinici complessi: Le immagini mediche sono spesso rumorose, sbilanciate e presentano ambiguità intrinseche (es. tessuti sovrapposti, basso contrasto), che i modelli deterministici standard non gestiscono bene.

2. Metodologia

Il paper propone MedFormer-UR, un'estensione dell'architettura MedFormer che integra tre componenti principali per gestire l'incertezza e migliorare l'interpretabilità, senza modificare il meccanismo di attenzione Dual Sparse Selection Attention (DSSA) di base.

A. Stima dell'Incertezza Evidenziale per Token (Per-Token Evidential Uncertainty)

Invece di utilizzare l'output softmax standard, il modello stima l'incertezza a livello di ogni token spaziale (patch dell'immagine) utilizzando una distribuzione Dirichlet.

• Un "evidential head" leggero prevede valori di evidenza non negativi per ogni classe.
• Questi valori parametrizzano una distribuzione Dirichlet, permettendo di decomporre l'incertezza in:
  • Aleatoria (Aleatoric): Rumore intrinseco nei dati (es. basso contrasto).
  • Epistemica (Epistemic): Mancanza di conoscenza del modello (es. dati di addestramento insufficienti).
• L'incertezza è calcolata come l'inverso della forza della distribuzione Dirichlet: più bassa è l'evidenza accumulata, più alta è l'incertezza.

B. Instradamento Guidato dall'Incertezza (Uncertainty-Guided Routing)

L'incertezza non è solo un output diagnostico, ma agisce come un segnale di controllo attivo durante il training e l'inferenza.

• Meccanismo: Ogni blocco Transformer include un ramo di raffinamento parallelo. Un "routing network" predice una maschera spaziale per decidere quali token devono essere raffinati.
• Filtraggio: L'incertezza a livello di token modula questa maschera. I token ad alta incertezza vengono soppressi o meno raffinati, mentre quelli a bassa incertezza ricevono aggiornamenti mirati.
• Obiettivo: Evitare che le caratteristiche ambigue o inaffidabili influenzino negativamente l'apprendimento, forzando il modello a concentrarsi sulle regioni informative.

C. Apprendimento Basato su Prototipi (Prototype-Based Learning)

Per migliorare l'interpretabilità, il classificatore finale sostituisce (o affianca) il layer lineare standard con un approccio basato sui prototipi.

• Il modello impara prototipi specifici per classe (pattern di riferimento rappresentativi).
• Le previsioni sono basate sulla similarità (cosine similarity) tra le feature spaziali dell'immagine e questi prototipi.
• Questo permette di spiegare le decisioni identificando quali pattern visivi (prototipi) hanno guidato la classificazione, rendendo il processo più trasparente.

Funzione di Obiettivo (Loss Function)

L'addestramento ottimizza una funzione di perdita composita:
$$L = L_{CE} + \lambda_{route}L_{routing} + L_{prototype}$$
Dove $L_{CE}$ è la cross-entropy per la classificazione, $L_{routing}$ supervisiona le maschere di instradamento (se disponibili maschere anatomiche) e $L_{prototype}$ include termini di cluster e diversità per mantenere i prototipi distinti e rappresentativi.

3. Contributi Chiave

1. Estensione MedFormer con Teste Evidenziali: Introduzione di un meccanismo per decomporre esplicitamente l'incertezza aleatoria ed epistemica a livello di token.
2. Instradamento Guidato dall'Incertezza: Un meccanismo innovativo che usa l'incertezza stimata per modulare dinamicamente il flusso delle feature, riducendo l'impatto delle caratteristiche ambigue durante l'addestramento.
3. Integrazione di Prototipi: Utilizzo di prototipi apprendibili con una loss di contrasto per strutturare lo spazio delle embedding, permettendo un ragionamento basato sulla similarità visiva.
4. Valutazione Completa: Test su quattro modalità diverse (Mammografia, Ultrasuoni, MRI, Istopatologia) che dimostrano miglioramenti nella calibrazione e nella previsione selettiva.

4. Risultati Sperimentali

Il modello è stato valutato su quattro dataset pubblici: CBIS-DDSM (Mammografia), BUSI (Ultrasuoni), IDC (Istopatologia) e Brain MRI.

• Prestazioni di Classificazione:

  • Su dataset con pattern visivi chiari (Ultrasuoni BUSI, Istopatologia), il modello proposto ha mostrato miglioramenti significativi nell'accuratezza (es. da 0.800 a 0.893 su BUSI) e nell'AUROC.
  • Su mammografia (CBIS-DDSM), dove l'ambiguità è intrinseca, i guadagni in accuratezza sono stati più modesti ma statisticamente rilevanti (da 0.629 a 0.673), dati i limiti fisici del dataset (basso contrasto, risoluzione).

• Calibrazione e Affidabilità (Punto di Forza):

  • Il miglioramento più significativo riguarda la calibrazione. Il modello ha ridotto l'Expected Calibration Error (ECE) fino al 35% rispetto alla baseline (da 0.1836 a 0.0688 nel caso ottimale su mammografia).
  • Questo indica che le probabilità previste dal modello corrispondono molto meglio alla realtà, riducendo la sovraconfidenza.

• Previsione Selettiva (Selective Prediction):

  • Utilizzando le stime di incertezza per rifiutare le previsioni meno affidabili, il modello ha migliorato l'accuratezza su sottoinsiemi di dati selezionati (es. al 50% o 70% di copertura), dimostrando una capacità superiore di identificare i casi difficili.

• Analisi del Parametro $\beta$ (Forza di Routing):

  • È stato osservato che la massima discriminazione (AUROC) e la massima affidabilità (calibrazione) non coincidono sempre. Un'intensità di routing moderata (es. $\beta \approx 0.7 - 0.8$) offre il miglior compromesso, mentre un routing troppo aggressivo può degradare l'affidabilità probabilistica.

5. Significato e Impatto Clinico

Il lavoro di MedFormer-UR è significativo perché sposta il focus dalla sola "massimizzazione dell'accuratezza" alla costruzione di sistemi AI affidabili e interpretabili per la medicina.

• Sicurezza Clinica: La capacità di quantificare e localizzare l'incertezza permette ai clinici di sapere quando fidarsi della diagnosi e quando richiedere un secondo parere umano.
• Trasparenza: L'uso dei prototipi e delle mappe di incertezza fornisce spiegazioni visive (es. "il modello è incerto perché i bordi della lesione sono ambigui" o "la previsione si basa sul prototipo X"), riducendo l'effetto "scatola nera".
• Adattabilità: Il framework dimostra che l'integrazione dell'incertezza nel processo di apprendimento (non solo post-hoc) è cruciale per gestire la variabilità e il rumore dei dati medici reali.

In conclusione, MedFormer-UR rappresenta un passo avanti verso l'IA clinica di fiducia, fornendo non solo una diagnosi, ma anche una stima affidabile della propria certezza e una spiegazione visiva del ragionamento.