Flow Matching with Optimized Subclass Priors for Medical Image Augmentation

Riepilogo Tecnico: Flow Matching con Priors di Sottoclasse Ottimizzati per l'Augmentation di Immagini Mediche

Enunciato del Problema
I dataset di imaging medico sono caratterizzati da gravi squilibri a coda lunga, dove le condizioni comuni dominano mentre le patologie rare e clinicamente critiche (ad esempio, piccoli pneumotoraci, bronchiectasie) sono sottorappresentate. Questo squilibrio causa un overfitting dei classificatori discriminativi sulle classi principali e un fallimento sulle condizioni rare. Sebbene l'augmentation dei dati generativa offra una potenziale soluzione, i modelli di flow matching (FM) condizionati standard e i modelli di diffusione affrontano due difficoltà combinate quando addestrati su etichette di malattia grezze:

1. Aggregazione Multi-modale: Una singola etichetta diagnostica spesso aggrega sottotipi diversi, protocolli di acquisizione e comorbidità. Ciò crea condizioni multi-modali che spingono i generatori verso sottomodi dominanti, troncando efficacemente variazioni rare ma clinicamente valide.
2. Geometria di Trasporto Inefficiente: Il tracciamento di tutte le traiettorie di generazione da una fonte gaussiana standard condivisa forza sottopopolazioni specifiche attraverso percorsi di trasporto sproporzionatamente lunghi. Ciò aumenta la difficoltà di regressione e degrada la fedeltà della generazione delle classi di coda.

Metodologia
Gli autori propongono una strategia offline che introduce priors informativi a due livelli per affinare il processo generativo senza modificare il ciclo di addestramento principale del modello di flow matching.

1. Induzione di Sottoclassi tramite Modellazione a Mixture di Gaussiane (GMM):

   • Il metodo partiziona ogni etichetta grezza $c$ in sottomodi coerenti (sottoclassi $k$) nello spazio latente del modello generativo.
   • Utilizzando un VAE congelato (FLUX1.dev), vengono calcolati latenti deterministici. Per ogni classe $c$, viene adattato un modello a mixture di gaussiane diagonale ai latenti residui (dopo aver sottratto il centro della classe $\mu_c$).
   • Il numero di componenti $K_c$ viene selezionato utilizzando il Criterio di Informazione Bayesiano Esteso (EBIC). I campioni ricevono assegnazioni rigide alle sottoclassi, creando una partizione con varianza condizionale inferiore.
   • Teoricamente, gli autori dimostrano che la condizione su queste etichette raffinate non può aumentare il rischio di flow matching ottimale di Bayes; il miglioramento è pari alla variabilità attesa tra le sottoclassi dello spostamento medio condizionale.

2. Ottimizzazione della Fonte Dipendente dalla Sottoclasse:

   • Invece di una fonte gaussiana condivisa, il metodo apprende una distribuzione di fonte specifica $\pi_0(\cdot | c, k)$ per ogni sottoclasse. Queste fonti sono parametrizzate come gaussiane apprendibili $N(\mu_{c,k}, \text{diag}(\sigma^2_{c,k}))$ che ridefiniscono e ridimensionano la distribuzione di partenza.
   • Ciò accorcia le traiettorie di trasporto tipiche e riduce la dispersione all'interno della sottoclasse.
   • Controllo Geometrico: Per prevenire soluzioni degeneri (come l'inflazione della norma o percorsi eccessivamente concentrati), gli autori ottimizzano un obiettivo proxy offline che impone:
     • Compattezza Direzionale: Incoraggiando le direzioni di spostamento normalizzate ad allinearsi moderatamente con prototipi apprendibili.
     • Limitazione della Lunghezza del Percorso: Imponendo una penalità fissa sugli outlier utilizzando un limite specifico per la sottoclasse ($r_{cap}$), definito come il 99° percentile delle distanze dal centro della classe. Ciò impedisce al modello di aggirare la perdita inflazionando le norme di spostamento.
     • Regolarizzazione della Scala: Regolarizzando i log-scali delle distribuzioni di fonte verso un riferimento unitario.

Principali Contributi

1. Formulazione Condizionata alla Sottoclasse: Un approccio innovativo per la generazione di immagini mediche a coda lunga che partiziona le condizioni multi-modali tramite assegnazioni rigide indotte da GMM, riducendo efficacemente l'ambiguità condizionale.
2. Fonti Apprese Dipendenti dalla Sottoclasse: L'introduzione di priors informativi che rimodellano il trasporto condizionale ridefinendo e ridimensionando la distribuzione di fonte per ogni sottoclasse, riducendo così la dispersione all'interno della sottoclasse.
3. Obiettivo Geometricamente Regolarizzato: Un obiettivo proxy pratico che bilancia una compattezza angolare moderata con un controllo esplicito degli outlier di lunghezza del percorso, prevenendo scorciatoie di inflazione della norma e collasso direzionale.
4. Validazione Empirica: Dimostrazione di un miglioramento della copertura della coda e delle prestazioni di classificazione a coda lunga a valle su benchmark di radiografie toraciche (MIMIC-LT, NIH-LT) e TC (CT-RATE).

Risultati Sperimentali
Il metodo è stato valutato su tre benchmark di imaging medico a coda lunga:

• Classificazione a Valle: Il metodo proposto ha costantemente ottenuto le più alte accuratezza bilanciata (bAcc) e punteggi macro-F1 su tutti e tre i dataset rispetto ai baseline senza augmentation e ad altre strategie di augmentation (Vanilla FM, CBFM, CPD). In particolare, è stata l'unica strategia a migliorare in modo affidabile le prestazioni su entrambe le modalità a raggi X e TC.
• Qualità Generativa: Su MIMIC-LT, il metodo ha ottenuto la migliore distanza Fréchet Inception (FID) e il miglior punteggio Inception Ratio (IRS), indicando guadagni simultanei in fedeltà e diversità. Su CT-RATE, sebbene la FID fosse leggermente superiore a quella di un baseline, il metodo ha mantenuto un IRS elevato e una classificazione a valle superiore, suggerendo una migliore copertura distribuzionale piuttosto che una semplice memorizzazione.
• Studi di Ablazione: I risultati hanno confermato che entrambe le fasi della pipeline (sottoclassificazione e ottimizzazione della fonte) sono necessarie. La sola sottoclassificazione ha migliorato le prestazioni, ma la combinazione ha prodotto i migliori risultati. L'ottimizzazione senza sottoclassificazione ha talvolta degradato le prestazioni mediando su modalità conflittuali.
• Validità delle Sottoclassi: L'analisi ha mostrato che le sottoclassi indotte catturano strutture rilevanti per la patologia piuttosto che solo metadati di acquisizione. Le sonde per i confondenti hanno indicato che i metadati erano rilevabili ma non dominanti, e la purezza kNN delle sottoclassi apprese ha superato significativamente le partizioni casuali.

Significato e Affermazioni
Il documento afferma che la progettazione della fonte e il raffinamento condizionale sono leve efficaci per un'augmentation generativa robusta in contesti clinici a coda lunga. Partizionando etichette eterogenee e apprendendo fonti dipendenti dalla sottoclasse, il metodo riduce l'ambiguità condizionale e stabilizza le traiettorie di trasporto a un costo computazionale trascurabile (l'intero processo è offline e richiede meno di un'ora di GPU per dataset). Gli autori lo presentano come una strategia principiale che migliora la fedeltà e la diversità della generazione delle classi di coda, portando a classificatori a valle più affidabili per condizioni mediche rare. Il lavoro evidenzia che la generazione condizionale ingenua spesso fallisce sulle sottopopolazioni rare a causa di problemi di geometria di trasporto, che possono essere mitigati attraverso priors strutturati.