SegWithU: Uncertainty as Perturbation Energy for Single-Forward-Pass Risk-Aware Medical Image Segmentation

1. Il Problema

La segmentazione delle immagini mediche è fondamentale per la quantificazione anatomica, la pianificazione terapeutica e il supporto alle decisioni cliniche. Tuttavia, l'implementazione clinica richiede non solo una segmentazione accurata, ma anche una stima dell'incertezza affidabile per identificare quando e dove le previsioni potrebbero essere inaffidabili.

Le sfide principali identificate nel paper sono:

• Costo computazionale: I metodi di incertezza più robusti (come Deep Ensembles, Monte Carlo Dropout e Test-time Augmentation) richiedono inferenze multiple, rendendoli costosi e difficili da integrare in flussi di lavoro clinici in tempo reale.
• Limitazioni dei metodi a singola passata: Le alternative efficienti a singola passata (single-forward-pass) spesso offrono una classificazione degli errori (ranking) debole o si basano su assunzioni restrittive nello spazio delle caratteristiche (feature space) che richiedono il riaddestramento del modello di base (backbone).
• Integrazione post-hoc: In molti contesti medici, il modello di segmentazione è già validato e integrato nei flussi di lavoro istituzionali. Sostituire o riaddestrare l'intero modello è spesso impraticabile. C'è quindi un bisogno critico di un modulo di "auto-audit" che possa essere aggiunto a posteriori (post-hoc) senza modificare il backbone pre-addestrato.

2. Metodologia: SegWithU

SegWithU è un framework post-hoc che arricchisce un backbone di segmentazione pre-addestrato e "congelato" (frozen) con una testina di incertezza leggera. L'intuizione centrale è che le delineazioni anatomiche inaffidabili sono quelle che risultano instabili sotto piccole perturbazioni latenti della rappresentazione delle caratteristiche.

L'architettura si compone di tre elementi principali:

1. Acquisizione delle caratteristiche (Feature Tapping):
   Il sistema estrae le mappe di caratteristiche intermedie dal backbone pre-addestrato (prima del blocco di output) senza modificare la struttura originale del modello. Supporta sia l'uso di una singola mappa di caratteristiche che l'aggregazione multi-scala.

2. Sonde Posteriori di Rango 1 (Rank-1 Posterior Probes):
   Questo è il cuore del metodo. Invece di stimare densità complesse nello spazio delle caratteristiche, SegWithU proietta l'incertezza in uno spazio di perturbazione compatto utilizzando un insieme di sonde apprese.

   • Le caratteristiche vengono mappate in risposte di sonda ($v$) tramite una convoluzione $1\times1$.
   • Queste sonde generano perturbazioni sui logit di classe ($\Delta z$).
   • L'incertezza è modellata come energia di perturbazione: la varianza delle probabilità di classe risultanti da un insieme di pattern di perturbazione fissi genera una mappa di incertezza epistemica ($U_{epi}$).

3. Testine Duali (Dual Heads):
   Il framework genera due mappe di incertezza distinte per scopi diversi:

   • Mappa Orientata alla Calibrazione ($U_{cal}$): Utilizzata per "temperare" (addolcire) le probabilità predittive. Combina segnali epistemici, aleatorici (rumore nei dati) e di ambiguità per migliorare la coerenza numerica delle probabilità (es. per il punteggio Brier).
   • Mappa Orientata al Ranking ($U_{rnk}$): Ottimizzata specificamente per rilevare errori e supportare la predizione selettiva. Combina l'entropia, l'energia di residuo e l'incertezza epistemica per ordinare i voxel in base alla probabilità di errore.

Funzione di Perdita (Loss Function):
Il modello viene addestrato congelando il backbone e ottimizzando solo la testina di incertezza. La perdita totale è una somma ponderata di diversi termini:

• NLL (Negative Log-Likelihood): Per la calibrazione delle probabilità.
• Error-Correlation Loss: Per correlare l'incertezza con gli errori reali.
• Pairwise Ranking Loss: Per garantire che i voxel errati abbiano un punteggio di incertezza superiore a quelli corretti.
• Tail Loss: Per penalizzare l'assegnazione di bassa incertezza ai voxel errati (errori ad alto rischio).
• Trust Loss: Per regolare l'energia di perturbazione e mantenere i cambiamenti nei logit controllati.

3. Contributi Chiave

1. Framework Post-Hoc Leggero: Introduzione di SegWithU, che aggiunge una testina di incertezza a un segmentatore pre-addestrato senza riaddestrare il backbone, preservando le prestazioni originali.
2. Modellazione dell'Incertezza come Energia di Perturbazione: Una nuova prospettiva che utilizza sonde posteriori di rango 1 per misurare l'instabilità delle delineazioni anatomiche nello spazio delle caratteristiche, producendo mappe separate per calibrazione e ranking.
3. Prestazioni Superiori in Singola Passata: Dimostrazione che questo approccio raggiunge risultati di incertezza competitivi o superiori rispetto ai metodi multi-passata (come gli ensemble), mantenendo un costo computazionale di inferenza pari a quello di una singola passata.

4. Risultati Sperimentali

Il metodo è stato valutato su tre dataset medici pubblici: ACDC (cuore), BraTS2024 (tumori cerebrali) e LiTS (tumori al fegato).

• Metriche Principali:

  • Dice Score: SegWithU mantiene prestazioni di segmentazione competitive (es. 0.9035 su ACDC, 0.6275 su BraTS24, 0.7821 su LiTS), paragonabili ai metodi multi-passata.
  • AUROC (Area Under ROC): Misura la capacità di distinguere voxel errati da quelli corretti. SegWithU ottiene i migliori risultati tra i metodi a singola passata e spesso supera anche gli ensemble (es. 0.9946 su BraTS24).
  • AURC (Area Under Risk-Coverage Curve): Misura l'utilità per la predizione selettiva (scartare i casi incerti). SegWithU ottiene sistematicamente i valori più bassi (migliori) su tutti i dataset (es. 0.2660 su BraTS24), superando anche Deep Ensembles e Test-time Augmentation.

• Analisi Qualitativa:
  Le mappe di incertezza generate da SegWithU sono più focalizzate sulle regioni anatomiche ambigue e sui confini, riducendo le risposte diffuse o "distaccate" (off-target) osservate in altri metodi come Deep Ensembles o DUQ.

• Studi di Ablazione:
  Gli esperimenti confermano che la separazione tra calibrazione e ranking è cruciale. L'uso esclusivo di una sola delle due testine degrada le prestazioni. Inoltre, la perdita Pairwise Ranking e la meccanica delle sonde di perturbazione sono i driver principali del successo del modello.

5. Significato e Implicazioni

Il lavoro di SegWithU rappresenta un passo significativo verso l'adozione clinica dell'IA medica:

• Praticità: Offre una soluzione di controllo qualità che può essere applicata a modelli esistenti senza costi di riaddestramento massicci o infrastrutture di inferenza complesse.
• Affidabilità Clinica: Sposta il focus dall'incertezza come semplice quantità ausiliaria a un segnale di controllo qualità "azioneabile" (actionable), permettendo di identificare casi che richiedono revisione umana.
• Efficienza: Dimostra che un'accurata stima dell'incertezza non richiede necessariamente inferenze multiple, rendendo possibile l'uso di sistemi "risk-aware" in ambienti con risorse computazionali limitate.

In sintesi, SegWithU dimostra che la modellazione dell'incertezza basata sulla perturbazione è un percorso efficace e pratico per rendere la segmentazione medica più sicura e affidabile, trasformando un segmentatore pre-addestrato in un sistema di auto-audit senza alterarne il nucleo predittivo.