Learning a Sampling-Free Variational DNN Plugin from Tiny Training Sets to Refine OOD Segmentation With Uncertainty Estimation

Sintesi Tecnica: Apprendimento di un Plugin DNN Variazionale Sampling-Free da Piccoli Set di Addestramento per Affinare la Segmentazione OOD con Stima dell'Incertezza

Problematica
Le reti neurali profonde (DNN) soffrono frequentemente di un significativo degrado delle prestazioni quando applicate a immagini mediche fuori distribuzione (OOD - out-of-distribution). Questo problema deriva dalle inevitabili variazioni nei protocolli di imaging, nell'hardware degli scanner e negli schemi di ricostruzione tra i diversi siti clinici. Sebbene il riaddestramento dei modelli su dati del dominio target sia una soluzione potenziale, esso è spesso impraticabile a causa degli elevati costi associati all'acquisizione e all'annotazione di nuovi dataset medici. Inoltre, i metodi esistenti di raffinamento della segmentazione richiedono spesso l'accesso a dati OOD durante l'addestramento o si affidano a un esteso pre-addestramento, rendendoli inadatti per scenari in cui i dati OOD non sono disponibili durante la fase iniziale di sviluppo del modello. La sfida consiste nel correggere le segmentazioni degradate su immagini OOD senza accesso a dati di addestramento OOD, fornendo simultaneamente stime dell'incertezza affidabili e mantenendo l'efficienza computazionale, in particolare quando sono disponibili solo piccoli set di addestramento in-distribuzione (ID - in-distribution).

Metodologia: VarDeepPCA
Gli autori propongono VarDeepPCA, un framework DNN variazionale leggero progettato come plugin per ripristinare o affinare le mappe di segmentazione degradate generate da DNN esistenti su dati OOD. L'innovazione principale risiede nello sfruttamento dei priori geometrici intrinseci dell'oggetto anatomico di interesse, che rimangono ampiamente invarianti rispetto ai cambiamenti di dominio nell'acquisizione delle immagini.

1. Apprendimento dei Priori Geometrici: Invece di apprendere da immagini grezze, VarDeepPCA apprende i modi principali di variazione direttamente da un piccolo set di mappe di segmentazione ID (100–200 coppie). L'ipotesi è che le geometrie anatomiche valide giacciano su un manifold a bassa dimensionalità.
2. Rappresentazione Latente Basata su Simplex: Il framework impiega un'architettura encoder-decoder. L'encoder mappa una mappa di segmentazione in input in un vettore di caratteristiche $K$-dimensionale, che viene poi trasformato tramite una funzione softmax in un vettore di probabilità situato su un simplex $(K-1)$-dimensionale. Questo vettore rappresenta una miscela di $K$ modi principali di variazione.
3. Apprendimento Variazionale Sampling-Free: Un contributo teorico chiave è la reinterpretazione della mappatura softmax all'interno di un quadro bayesiano. Gli autori dimostrano che l'output softmax esegue implicitamente una marginalizzazione bayesiana esatta di una distribuzione categorica latente parametrizzata da un prior di Dirichlet. Ciò consente un apprendimento e un'inferenza sampling-free (senza campionamento). A differenza dei tradizionali Autoencoder Variazionali (VAE) che richiedono costosi campionamenti Monte Carlo, VarDeepPCA deriva la media e la varianza della distribuzione di output analiticamente in forma chiusa.
   • Il decoder produce una mappa di segmentazione media ($M$) e una mappa di varianza per pixel ($V$).
   • La varianza è calcolata utilizzando il Jacobiano del decoder e la varianza della distribuzione categorica latente.
   • La mappa di incertezza è derivata come radice quadrata della mappa di varianza.
4. Algoritmo di Raffinamento a Due Stadi: Per correggere una segmentazione OOD $\tilde{Z}$ prodotta da una DNN di base:
   • Filtraggio: La mappa degradata viene passata attraverso l'encoder-decoder per rimuovere le componenti di variabilità non principali, producendo una mappa filtrata $Z$.
   • Proiezione: La mappa filtrata viene esplicitamente proiettata sul manifold appreso delle geometrie valide. Ciò viene ottenuto ottimizzando il vettore delle caratteristiche latenti tramite l'ascesa del gradiente per massimizzare la verosimiglianza della mappa filtrata sotto la distribuzione gaussiana appresa, seguito dal decoding per ottenere la segmentazione ripristinata $M^*$ e la relativa incertezza $U^*$.

Contributi Chiave

• Framework Innovativo: Introduzione di VarDeepPCA, un plugin che affina le segmentazioni OOD utilizzando solo piccoli dataset ID, eliminando la necessità di dati del dominio target o di un esteso pre-addestramento.
• Inferenza Variazionale Sampling-Free: Una reinterpretazione teorica della funzione softmax che consente la modellazione esatta della distribuzione e un apprendimento e un'inferenza computazionalmente efficienti e sampling-free, fornendo stime dell'incertezza in forma chiusa.
• Architettura Leggera: Il modello è progettato per apprendere da tiny training set (100–200 immagini) e possiede un basso numero di parametri (da 1.02M a 2.72M), rendendolo adatto ad ambienti clinici con dati limitati.
• Validazione Completa: Ampia validazione empirica attraverso quattro distinte applicazioni cliniche (miocardio, bordo neuroretinico, prostata, testa fetale) utilizzando 14 dataset pubblici e confronti con 15 metodi esistenti, inclusi transformer allo stato dell'arte, modelli di diffusione e modelli foundation.

Risultati
L'articolo riporta estesi risultati quantitativi e qualitativi che dimostrano come VarDeepPCA migliori costantemente i metodi esistenti:

• Accuratezza della Segmentazione: In tutte e quattro le applicazioni e rispetto ai 15 baseline, l'aggiunta di VarDeepPCA alle DNN esistenti ha ridotto significativamente gli errori di confine (misurati tramite HD95 e ASD) sui dati OOD. Ad esempio, nella segmentazione del miocardio, la media HD95 sui dataset OOD è scesa da una media di 13.2–13.6 pixel a 7.2–8.0 pixel, portando le prestazioni in un intervallo clinicamente accettabile.
• Plausibilità Anatomica: Il metodo ha ripristinato efficacemente geometrie anatomicamente implausibili (ad esempio, blob disconnessi) comuni nei fallimenti OOD, proiettandole sul manifold delle forme valide apprese.
• Calibrazione dell'Incertezza: VarDeepPCA ha migliorato significativamente la calibrazione delle stime dell'incertezza (misurate tramite NCC, US e TACE) per i baseline probabilistici come ProbUNet e PHISeg. Le mappe di incertezza generate erano più localizzate e allineate con i confini anatomici rispetto ai baseline.
• Efficienza: Il plugin aggiunge un overhead minimo ai tempi di addestramento e inferenza e non richiede dati di addestramento aggiuntivi oltre a quelli utilizzati dai metodi di base.
• Robustezza: L'analisi di sensibilità ha mostrato che il metodo è robusto alle variazioni nella dimensione latente ($K$) e nella dimensione del set di addestramento ($S$), performando bene anche con soli 100 campioni di addestramento.

Significatività e Rivendicazioni
Gli autori affermano che VarDeepPCA affronta una barriera critica per l'impiego clinico delle DNN: l'incapacità di generalizzare a nuovi siti di imaging senza un costoso riaddestramento. Modellando esplicitamente la distribuzione delle geometrie anatomiche valide da piccoli dataset ID, il framework fornisce una soluzione fondata e computazionalmente efficiente per affinare le segmentazioni OOD.

Il lavoro sottolinea che, a differenza degli approcci esistenti che possono richiedere dati OOD per l'addestramento o soffrire di alti costi computazionali (come i modelli di diffusione o i grandi transformer), VarDeepPCA offre un'alternativa leggera e sampling-free che migliora sia l'accuratezza della segmentazione che la stima dell'incertezza. Gli autori riconoscono con modestia alcuni vincoli, notando che il metodo si basa sulla coerenza geometrica e potrebbe avere difficoltà con patologie altamente amorfe (come i tumori) o casi in cui l'errore iniziale è così grave che la geometria corretta si trova al di fuori del manifold appreso. Tuttavia, anche in tali casi, il framework proietta con successo l'output verso una geometria anatomica valida. Questo lavoro posiziona VarDeepPCA come un plugin versatile capace di potenziare una vasta gamma di segmentatori esistenti, con il potenziale di estendersi anche a metodi non basati su DNN.