RAFM: Retrieval-Augmented Flow Matching for Unpaired CBCT-to-CT Translation

Il paper presenta RAFM, un metodo di traduzione non accoppiata da CBCT a CT che integra il flusso rettificato con un meccanismo di recupero guidato da un encoder DINOv3 e una banca di memoria globale per generare immagini CT sintetiche di alta qualità, superando le limitazioni dei dati non accoppiati e ottenendo risultati superiori rispetto agli stati dell'arte nel dataset SynthRAD2023.

L'essenziale

Immagina di dover trasformare una foto sbiadita, piena di "rumore" e distorsioni (come una vecchia foto scattata con una macchina fotografica economica) in un'immagine nitida, perfetta e professionale, pronta per essere usata da un architetto per calcolare esattamente quanto materiale serve per costruire una casa.

Nel mondo medico, questo è il problema che i ricercatori hanno affrontato:

• La "foto sbiadita" è la CBCT (Tomografia Computerizzata a Cono), usata spesso durante le sedute di radioterapia. È comoda da fare, ma l'immagine è piena di artefatti e i valori di densità sono inaffidabili, rendendo impossibile calcolare con precisione la dose di radiazioni da somministrare al paziente.
• La "foto perfetta" è la CT (Tomografia Computerizzata standard), che è chiara e precisa, ma non sempre disponibile o sincronizzata con la CBCT.

Il compito è creare una "CT sintetica" partendo dalla CBCT. Il problema è che spesso non abbiamo coppie perfette (la stessa persona, nello stesso momento, con entrambe le macchine). È come se avessimo un album di foto CBCT di 100 persone e un album di foto CT di altre 100 persone, e dovessimo imparare a trasformare le prime nelle seconde senza sapere quale foto corrisponde a quale persona.

La soluzione: RAFM (Il "Viaggio Rettilineo" con una Guida)

Gli autori del paper hanno creato un metodo chiamato RAFM (Retrieval-Augmented Flow Matching). Ecco come funziona, spiegato con metafore semplici:

1. Il problema dei metodi vecchi: "Il salto nel buio"

I metodi precedenti cercavano di imparare questa trasformazione facendo un "tiro alla fune" (GAN) o facendo molti piccoli passi incerti (Diffusion). Spesso, quando provavano a collegare una foto CBCT a una foto CT a caso (perché non avevano le coppie perfette), finivano per collegare un "piede" a un "occhio". Il risultato era confuso: l'immagine finale aveva la forma giusta ma i dettagli sbagliati, o viceversa.

2. L'idea di base: "Il viaggio in linea retta"

Immagina che trasformare una CBCT in una CT sia come viaggiare da una città (la CBCT) a un'altra città (la CT).
I metodi tradizionali spesso facevano un percorso tortuoso, pieno di deviazioni.
Il metodo Flow Matching (alla base di RAFM) dice: "Facciamo un viaggio in linea retta, il più diretto possibile". Matematicamente, questo significa che l'immagine si trasforma gradualmente e in modo fluido, senza salti bruschi. È come se avessimo una strada dritta che collega i due punti.

3. Il problema pratico: "La mappa sbagliata"

C'è un ostacolo: per costruire questa strada dritta, dobbiamo sapere quale punto di partenza (una specifica CBCT) corrisponde a quale punto di arrivo (una specifica CT). Se scegliamo la destinazione a caso (perché non abbiamo le coppie perfette), la strada dritta ci porterà nel posto sbagliato. Immagina di voler andare da Milano a Roma, ma per errore imposti il GPS su "Roma" mentre la tua auto è parcheggiata davanti a un negozio di Milano che assomiglia a un altro negozio di Roma. Se guidi dritto, finirai nel posto sbagliato.

4. La soluzione RAFM: "La Libreria Intelligente"

Qui entra in gioco la parte geniale del paper: Retrieval-Augmented (Arricchito dal Recupero).

Immagina di avere una biblioteca gigante (Memory Bank) piena di tutte le foto CT disponibili.
Quando il sistema deve trasformare una nuova foto CBCT, non sceglie una foto CT a caso dalla biblioteca. Invece, usa un assistente super-intelligente (un encoder chiamato DINOv3, che è come un occhio che guarda il "significato" dell'immagine e non solo i pixel).

Ecco la magia:

1. L'assistente guarda la foto CBCT e dice: "Questa è una vista del bacino di un paziente con una certa forma".
2. Va nella biblioteca e cerca la foto CT che sembra più simile a questa descrizione, anche se non è della stessa persona e non è stata scattata nello stesso momento.
3. Crea una "coppia finta" (pseudo-pair) perfetta: "Ok, trasformiamo questa CBCT in quella CT specifica".

Invece di collegare due cose a caso (come un piede a un occhio), il sistema collega due cose che hanno lo stesso "stile" e la stessa "forma" anatomica.

Perché è un successo?

• Stabilità: Poiché le coppie sono scelte con intelligenza (basate sulla somiglianza visiva), il "viaggio" verso l'immagine finale è molto più stabile e sicuro.
• Anatomia corretta: Il paziente non cambia forma. Se la CBCT mostra un'anca, la CT sintetica mostrerà la stessa anca, non un ginocchio.
• Velocità: Una volta addestrato, il sistema è veloce. Non deve fare migliaia di piccoli passi (come i metodi precedenti), ma percorre la sua strada dritta in pochi secondi.

In sintesi

RAFM è come avere un traduttore esperto che, quando deve tradurre un testo da una lingua "rumorosa" (CBCT) a una lingua "perfetta" (CT), non cerca parole a caso nel dizionario. Invece, consulta una vasta biblioteca di testi perfetti, trova quello che ha lo stesso contesto e significato del testo originale, e usa quello come guida per creare una traduzione impeccabile, anche senza avere il testo originale e quello tradotto già abbinati.

Il risultato è un'immagine medica così chiara e precisa che i medici possono usarla per calcolare con sicurezza le cure per i pazienti, salvando tempo e migliorando la qualità delle terapie.

Sommario tecnico

1. Problema e Contesto

Nella pianificazione della radioterapia, la Tomografia Computerizzata (CT) è lo standard di riferimento perché fornisce valori affidabili dell'Unità Hounsfield (HU) necessari per il calcolo della dose. Tuttavia, durante il trattamento, viene spesso utilizzata la Tomografia Computerizzata a Fascio Conico (CBCT) per la guida dell'immagine. La CBCT soffre di gravi artefatti e valori HU inaffidabili, rendendola inadatta al calcolo diretto della dose.

La generazione di una CT sintetica (sCT) a partire dalla CBCT è quindi un compito cruciale. La sfida principale risiede nel fatto che i dati CBCT-CT accoppiati (paired) sono spesso indisponibili o inaffidabili a causa di:

• Intervalli temporali tra le scansioni.
• Variazioni anatomiche del paziente.
• Errori di registrazione (registration errors).

Di conseguenza, è necessario sviluppare metodi per la tradizione non accoppiata (unpaired) da CBCT a CT. I metodi esistenti basati su GAN (Generative Adversarial Networks) soffrono di instabilità nell'addestramento, mentre quelli basati su Diffusione o Ponti di Schrödinger richiedono pipeline complesse.

2. Metodologia: RAFM

Gli autori propongono RAFM (Retrieval-Augmented Flow Matching), un framework che integra il Rectified Flow (RF) con una strategia di recupero guidato per migliorare l'apprendimento non accoppiato in contesti medici con dati limitati.

Concetti Chiave:

• Rectified Flow (RF): Modella la traduzione come un trasporto deterministico tra due distribuzioni (CBCT e CT) tramite un'Equazione Differenziale Ordinaria (ODE). L'obiettivo è imparare un campo di velocità $v_\theta$ che sposti i campioni dalla sorgente al target lungo percorsi lineari. Teoricamente, il RF non richiede corrispondenze voxel-per-voxel, ma solo che le distribuzioni marginali siano corrette.
• Il Problema Pratico: In dataset medici piccoli e con batch size ridotti (tipici per immagini ad alta risoluzione), l'accoppiamento casuale o "locale" (solo all'interno del batch) dei punti finali (endpoint) crea coppie semanticamente disallineate. Questo introduce rumore nel target di trasporto e degrada la preservazione dell'anatomia.

La Soluzione RAFM:

RAFM risolve il problema dell'accoppiamento debole introducendo un banco di memoria globale per il recupero delle caratteristiche:

1. Memory Bank: Viene mantenuto un buffer FIFO (First-In-First-Out) contenente slice CT estratte dal dominio target durante l'addestramento.
2. Encoder Congelato (Frozen): Viene utilizzato un encoder DINOv3 congelato per estrarre embedding di caratteristiche per ogni slice.
3. Recupero Guidato: Per ogni slice CBCT nel batch corrente, il sistema cerca nel banco di memoria la slice CT più simile in termini di caratteristiche (massima similarità coseno).
4. Costruzione di Coppie Pseudo-Accoppiate: La slice CT recuperata viene utilizzata come punto finale ($x_1$) per la slice CBCT ($x_0$), formando una coppia pseudo-accoppiata semanticamente coerente per l'addestramento del Flow Matching.
5. Addestramento: Il modello (una U-Net condizionata al tempo) viene addestrato per minimizzare l'errore tra il campo di velocità previsto e il vettore di trasporto reale ($x_1 - x_0$) lungo il percorso lineare.

3. Contributi Chiave

• Introduzione del Rectified Flow in ambito medico non accoppiato: Dimostrazione pratica che il RF può essere efficace per la traduzione CBCT-CT, superando le limitazioni delle GAN.
• Strategia di Recupero (Retrieval-Augmented): Proposta di un metodo innovativo per costruire accoppiamenti empirici di alta qualità in assenza di dati accoppiati reali, superando il problema dei dataset piccoli e dei batch ridotti.
• Framework Non Adversarial: Fornisce un'alternativa stabile e non basata su giochi minimax (adversarial) per la preservazione dell'anatomia, cruciale in ambito medico.
• Validazione Rigorosa: Valutazione condotta su un protocollo "vero non accoppiato" a livello di soggetto (nessuna sovrapposizione di pazienti tra CBCT e CT nel training), un setting più severo rispetto alle semplici miscele a livello di slice.

4. Risultati Sperimentali

Il metodo è stato valutato sul dataset SynthRAD2023 (bacino pelvico) e confrontato con metodi basati su GAN (CycleGAN, GcGAN, CUT) e modelli di diffusione/ponti (SynDiff, UNSB).

Metriche di Performance:
RAFM ha ottenuto risultati superiori in tutte le metriche chiave:

• MAE (Errore Assoluto Medio): 101.2 HU (migliore rispetto a SynDiff: 104.2 e UNSB: 110.8).
• SSIM (Similarità Strutturale): 80.96% (il più alto).
• PSNR: 25.15 dB (il più alto).
• FID (Fréchet Inception Distance): 53.29 (il più basso, indicando la migliore realtà distributiva).
• SegScore (Coerenza Anatomica): 75.77% (misurato tramite segmentazione degli organi pelvici), superando significativamente i metodi concorrenti.

Ablation Study:
Lo studio ha dimostrato che:

• L'uso di un semplice U-Net senza Flow Matching fallisce completamente in setting non accoppiati.
• Il RF con accoppiamento casuale è robusto ma inferiore.
• L'aumento della dimensione del pool di recupero (Memory Bank) migliora le prestazioni fino a un punto di saturazione (K=512), confermando che un pool di candidati più ampio e semanticamente rilevante è fondamentale.

5. Significato e Impatto

Il lavoro di RAFM è significativo per diversi motivi:

1. Affidabilità Clinica: Migliora la qualità delle immagini sintetiche e la preservazione dell'anatomia, rendendo la CBCT più utilizzabile per il calcolo della dose in radioterapia adattiva.
2. Efficienza Computazionale: Rispetto ai modelli di diffusione che richiedono molti passaggi di campionamento, RAFM risolve l'ODE in soli 10 passaggi, offrendo un buon compromesso tra velocità e qualità.
3. Superamento dei Limiti dei Dati: Dimostra che è possibile ottenere prestazioni vicine a quelle dei metodi supervisionati (paired) anche in assenza totale di dati accoppiati, sfruttando strategie di recupero delle caratteristiche.
4. Stabilità: Elimina l'instabilità tipica dell'addestramento GAN, offrendo un approccio più prevedibile e robusto per la ricerca e l'applicazione medica.

In sintesi, RAFM rappresenta un avanzamento significativo nella traduzione di immagini mediche non accoppiate, combinando la teoria del trasporto ottimale (Flow Matching) con tecniche di recupero delle caratteristiche per risolvere le sfide specifiche dei dataset medici di piccole dimensioni.