CLIMB: Controllable Longitudinal Brain Image Generation using Mamba-based Latent Diffusion Model and Gaussian-aligned Autoencoder

1. Il Problema

La generazione di immagini mediche longitudinali, ovvero la sintesi di scansioni di risonanza magnetica (MRI) di follow-up basate su immagini basali per simulare l'evoluzione anatomica nel tempo, è cruciale per la diagnosi precoce, la prognosi e la pianificazione del trattamento, specialmente nelle malattie neurodegenerative come l'Alzheimer.
Tuttavia, gli approcci esistenti presentano diverse limitazioni:

• Costo Computazionale: I modelli di diffusione tradizionali operano nello spazio dei pixel e richiedono processi iterativi di denoising estremamente costosi e lenti, rendendoli poco pratici per volumi 3D ad alta dimensionalità.
• Limiti dei Latent Diffusion Models (LDM) attuali: Sebbene gli LDM riducano il costo computazionale operando in uno spazio latente, spesso si basano su Autoencoder Variazionali (VAE). I VAE introducono rumore stocastico durante il campionamento dello spazio latente, il che può portare a ricostruzioni sfocate e prive di dettagli nitidi.
• Architettura inefficiente: Le architetture LDM esistenti utilizzano prevalentemente meccanismi di Self-Attention (basati su Transformer) per catturare il contesto globale. Questi meccanismi hanno una complessità quadratica, risultando inefficienti per dati medici 3D e richiedendo grandi quantità di dati di addestramento.
• Mancanza di controllo: Molti modelli non riescono a incorporare efficacemente variabili cliniche multiple (età, genere, stato della malattia, volumi cerebrali) per guidare la generazione temporale in modo controllato.

2. Metodologia Proposta: CLIMB

Gli autori propongono CLIMB (Controllable Longitudinal brain Image generation via MamBa-based latent diffusion model), un framework avanzato composto da due fasi principali di addestramento: un Autoencoder allineato al Gaussiano e un processo di diffusione latente basato su Mamba.

A. Gaussian-Aligned Autoencoder (GATE)

Per superare i limiti dei VAE tradizionali, viene introdotto il GATE.

• Obiettivo: Creare uno spazio latente continuo che aderisca a una distribuzione priorica Gaussiana senza introdurre rumore stocastico durante la codifica.
• Meccanismo: Invece del campionamento stocastico tipico dei VAE, il GATE utilizza una strategia di allineamento deterministica basata sulla Sliced Wasserstein Distance (SWD).
  • Proietta le feature latenti codificate e la distribuzione priorica Gaussiana su direzioni casuali.
  • Allinea le funzioni di distribuzione cumulativa (CDF) delle proiezioni per minimizzare la distanza tra la distribuzione latente e quella Gaussiana.
• Loss Function: Il modello è addestrato con una combinazione di:
  • Reconstruction Loss: Per la fedeltà dei pixel.
  • Sliced Wasserstein Distance (SWD): Per l'allineamento della distribuzione.
  • Perceptual Loss: Per preservare la similarità semantica e strutturale ad alto livello.
  • Adversarial Loss: Per garantire un output fotorealistico.

B. Denoising Network basata su Mamba (State Space Models)

Al posto dei tradizionali blocchi di Self-Attention nei modelli di diffusione, CLIMB integra l'architettura Mamba (un modello a spazio stato - SSM).

• Vantaggi: Mamba offre una complessità computazionale lineare rispetto alla sequenza, permettendo di catturare dipendenze a lungo raggio (contesto globale) in modo molto più efficiente rispetto ai Transformer, specialmente su dati 3D.
• Architettura: Utilizza uno schema U-Net con percorsi di codifica e decodifica. Le informazioni condizionali (dati proiettati e feature dell'immagine) vengono iniettate tramite operazioni di Cross-Attention.
• Strategia di Addestramento in Due Fasi:
  1. Addestramento iniziale del modello di diffusione usando solo le variabili proiettate (età futura, stato malattia, volumi).
  2. Congelamento dei parametri e addestramento della rete di elaborazione delle feature dell'immagine, inserendo layer convoluzionali inizializzati a zero per facilitare l'adattamento graduale.

C. Processo di Inferenza

1. Si parte da una MRI basale e si estraggono le feature tramite l'encoder GATE.
2. Un modello preesistente (IRLSTM) prevede i dati futuri (età, stato della malattia, volumi cerebrali) basandosi su variabili cliniche.
3. Questi dati proiettati e le feature dell'immagine vengono forniti alla rete di denoising basata su Mamba per campionare la rappresentazione latente del futuro.
4. Il decoder GATE ricostruisce l'immagine MRI futura.
5. Viene utilizzato il Denoising Diffusion Implicit Model (DDIM) per accelerare il campionamento (25 step) e una media di 10 campioni per robustezza.

3. Contributi Chiave

1. Framework CLIMB: Un sistema end-to-end per la generazione di MRI longitudinali controllabili, che utilizza una MRI basale, l'età di acquisizione e dati proiettati (età futura, genere, stato di malattia, volumi cerebrali).
2. Gaussian-Aligned Autoencoder (GATE): Un nuovo autoencoder che allinea lo spazio latente a una distribuzione Gaussiana tramite Sliced Wasserstein Distance, eliminando il rumore di campionamento dei VAE e preservando la nitidezza dell'immagine.
3. Integrazione di Mamba: Sostituzione dei meccanismi di Self-Attention con l'architettura Mamba (SSM) nel modello di diffusione latente, riducendo drasticamente il costo computazionale e la memoria richiesta mantenendo alte prestazioni nella cattura del contesto globale.
4. Validazione su ADNI: Addestramento e valutazione su un vasto dataset (6.306 scansioni MRI da 1.390 partecipanti) dell'Alzheimer's Disease Neuroimaging Initiative (ADNI).

4. Risultati Sperimentali

Il modello è stato testato confrontandolo con stati dell'arte come SADM e BrLP.

• Metriche Quantitative: CLIMB ha ottenuto prestazioni superiori su tutte le metriche:
  • SSIM (Structural Similarity Index): 0.9433 (il più alto), indicando un'eccellente preservazione della struttura.
  • PSNR: 27.82.
  • MSE: 2.01e-3 (riduzione significativa dell'errore).
  • LPIPS (Perceptual Similarity): 0.0587 (il più basso), dimostrando una qualità percettiva molto superiore rispetto ai baselines (es. VAE-based).
• Efficienza:
  • Tempo di Inferenza: 2.92 secondi per immagine (più veloce dei baselines).
  • Consumo di Memoria: 5699 MB (significativamente inferiore rispetto ai modelli basati su Transformer che soffrono di problemi di memoria su dati 3D).
• Ablation Study: Gli esperimenti hanno confermato che l'uso di GATE invece di VAE e di Mamba invece di Self-Attention porta a miglioramenti tangibili sia nella qualità dell'immagine che nell'efficienza computazionale.

5. Significato e Implicazioni

Il lavoro di CLIMB rappresenta un passo avanti significativo nella generazione di immagini mediche longitudinali:

• Efficienza Clinica: La riduzione del costo computazionale rende fattibile l'uso di modelli di diffusione avanzati su volumi cerebrali 3D completi, aprendo la strada a applicazioni in tempo reale o su risorse limitate.
• Qualità Diagnostica: La capacità di generare immagini nitide e strutturalmente accurate (grazie a GATE) è fondamentale per l'analisi clinica, superando il problema delle ricostruzioni sfocate tipiche dei VAE.
• Modellazione della Malattia: La capacità di controllare la generazione basandosi su variabili cliniche (stato della malattia, volumi cerebrali) permette di simulare scenari di progressione personalizzati, utili per la pianificazione terapeutica e lo studio dell'Alzheimer.
• Limitazioni e Futuro: Sebbene le metriche quantitative siano eccellenti, gli autori sottolineano la necessità di validazioni cliniche da parte di esperti (radiologi) per garantire che le immagini generate siano clinicamente interpretabili e non solo statisticamente simili.

In sintesi, CLIMB combina l'efficienza degli State Space Models (Mamba) con la stabilità di un autoencoder allineato al Gaussiano, offrendo una soluzione robusta ed economica per la previsione dell'evoluzione anatomica cerebrale.