Sintesi Tecnica: CONFLUX

Definizione del Problema

I modelli generativi controllabili per immagini mediche 3D mirano a sintetizzare volumi con specifici attributi clinici per compiti quali l'aumento di coorti, la condivisione di dati nel rispetto della privacy e la progettazione di studi controfattuali. Tuttavia, ottenere una sintesi 3D nativa ad alta fedeltà che sia simultaneamente fedele al condizionamento richiesto rimane una sfida significativa. I modelli esistenti spesso faticano a garantire che i singoli campioni generati esibiscano i reperti clinici richiesti, poiché gli standard obiettivi di addestramento (come il flow matching) ottimizzano la realtà distribuzione aggregata piuttosto che la realizzazione dell'attributo per singolo campione. Questo divario è particolarmente acuto nella TC del torace, che richiede un'alta risoluzione e coerenza volumetrica pur essendo descritta da metadati radiologici strutturati.

Metodologia

CONFLUX è un framework di diffusione latente a tre stadi progettato per la sintesi di TC toraciche, integrando un autoencoder variazionale (VAE), un transformer a flusso rettificato (rectified-flow transformer) e una fase di post-addestramento tramite apprendimento per rinforzo (RL).

1. Autoencoder Latente (Stadio 1)

Per mitigare il costo computazionale della generazione diretta nello spazio dei voxel, un VAE convoluzionale 3D comprime i volumi TC preelaborati ($216 \times 176 \times 200$) in uno spazio latente a risoluzione inferiore ($16 \times 27 \times 22 \times 25$).

• Architettura: L'encoder e il decoder utilizzano blocchi residui con normalizzazione di gruppo e self-attention 3D alla risoluzione più bassa.
• Addestramento: Il modello minimizza una perdita di ricostruzione $\ell_1$, un regolarizzatore Kullback–Leibler e una perdita percettiva LPIPS triplanare.
• Normalizzazione: Post-addestramento, lo spazio latente è normalizzato a una varianza prossima all'unità per facilitare il flow matching.

2. Transformer a Flusso Rettificato Condizionale (Stadio 2)

Un transformer a singolo stream genera embedding latenti da rumore gaussiano all'interno dello spazio latente appreso.

• Flow Matching: Il modello è addestrato sotto un obiettivo di rectified-flow, imparando a predire il campo di velocità costante che connette il rumore ai latenti dei dati lungo percorsi rettilinei.
• Condizionamento: Metadati radiologici strutturati (18 anomalie binarie, sesso, decennio di età e kernel di ricostruzione) vengono concatenati in un vettore a 4 di dimensioni. Questo vettore modula il transformer tramite normalizzazione adattiva degli strati (adaLN-zero).
• Campionamento: La generazione prevede l'integrazione dell'ODE di flusso di probabilità da $t=1$ a $t=0$ utilizzando una griglia Euler con spostamento temporale, seguito dalla decodifica con il decoder VAE congelato.

3. Post-addestramento con Apprendimento per Rinforzo (Stadio 3)

Per affrontare il disconnessione tra il matching della distribuzione aggregata e la fedeltà dell'attributo per singolo campione, CONFLUX impiega la Group-Relative Policy Optimization (GRPO).

• Motivazione: Il flow matching standard non premia esplicitamente se un reperto specifico richiesto appare in un singolo volume generato.
• Meccanismo: Il campionatore ODE deterministico è sostituito da un campionatore stocastico (SDE) per fornire log-probabilità trattabili per gli aggiornamenti della policy gradient.
• Segnale di Ricompensa: Un classificatore indipendente, congelato e addestrato su rappresentazioni latenti reali, valuta i volumi generati in base all'accordo tra i reperti richiesti e quelli realizzati. La ricompensa è definita come la cross-entropy pesata negativa tra le predizioni del classificatore e il vettore di condizionamento.
• Ottimizzazione: Il modello viene perfezionato per massimizzare i vantaggi relativi al gruppo (sottraendo la media della ricompensa di un gruppo di rollout per ogni prompt e dividendo per la deviazione standard). Gli aggiornamenti utilizzano un surrogato PPO con clipping e una penalità di divergenza KL rispetto alla policy di riferimento pre-RL per prevenire il drift.

Contributi Chiave

1. Modello a Flusso Rettificato Latente 3D Controllabile: Un'architettura 3D nativa per la TC toracica che condiziona direttamente sui metadati radiologici strutturati tramite adaLN, raggiungendo una qualità di sintesi competitiva con i forti baseline di diffusione 3D.
2. Post-addestramento RL per la Fedeltà: La prima applicazione di GRPO a un modello di flow medico 3D. Questa fase ottimizza esplicitamente l'affidabilità dei reperti richiesti nei volumi generati, verificata da un modello giudice indipendente.
3. Rilascio Open: Il rilascio del modello addestrato e di un dataset di TC toraciche sintetiche di circa 200.000 volumi. Questo dataset è condizionato sulle proporzioni reali dei metadati, coprendo una vasta gamma di reperti clinici, destinato all'aumento di coorti e alla progettazione controllata di studi.

Risultati

Qualità della Sintesi

Valutato su un set di validazione tenuto separato ($N=501$) utilizzando la distanza di Fréchet Inception (FID) triplanare, CONFLUX supera significativamente i baseline esistenti:

• FID Triplanare: 32.3 (CONFLUX) vs. 74.6 (MAISI) e 145.4 (GenerateCT).
• Copertura e Densità: I campioni di CONFLUX sono 3.5$\times$ più densi e coprono 6$\times$ più del manifold reale rispetto a MAISI.
• Diversità: Il modello mantiene un'alta diversità (MS-SSIM 0.507), comparabile ai dati reali (0.41), evitando il mode collapse.

Fedeltà del Condizionamento

Il post-addestramento con GRPO migliora l'accordo tra gli attributi richiesti e quelli realizzati, misurato da un classificatore indipendente nello spazio delle immagini:

• Guadagno di Performance: L'Average Precision (AP) è aumentata da 0.330 a 0.344, e l'AUROC da 0.684 a 0.699.
• Recupero del Gap: Sebbene i guadagni assoluti sembrino modesti, essi recuperano il 47% del gap di performance tra il modello base e il limite superiore dei dati reali (AP 0.360).
• Specificità: I miglioramenti sono stati specifici per il gruppo di reperti ottimizzati; i gruppi non pesati (sesso, età, kernel) sono rimasti invariati entro il rumore.

Performance del Classificatore di Ricompensa

Il classificatore utilizzato per il segnale di ricompensa RL ha raggiunto un macro AUROC di 0.793 sul split di validazione reale, superando le varianti CT-CLIP e CT-Net nel benchmark CT-RATE, validando la sua utilità come segnale di ricompensa nonostante operi su latenti compressi.

Significato e Rivendicazioni

L'articolo sostiene che CONFLUX rappresenti un passo avanti significativo nella sintesi di immagini mediche 3D controllabili, disaccoppiando l'alta fedeltà della generazione dal controllo degli attributi. Integrando il post-addestramento RL, gli autori dimostrano che è possibile ottimizzare direttamente la "fedeltà" degli attributi clinici nei singoli campioni, una proprietà non garantita dagli standard obiettivi di flow-matching. Il rilascio di un dataset sintetico su larga scala e condizionalmente diversificato mira a facilitare la ricerca in scenari in cui i dati reali sono scarsi o soggetti a vincoli di privacy. Il lavoro si posiziona come il primo ad adattare GRPO per i modelli di flow medico 3D, offrendo un percorso verso strumenti generativi più affidabili e controllabili per la progettazione di studi clinici.