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

Probleemstelling

De generatie van longitudinale medische beelden (het simuleren van anatomische veranderingen in de hersenen over tijd) is cruciaal voor het modelleren van ziekteprogressie, zoals bij de ziekte van Alzheimer, en voor het ondersteunen van klinische besluitvorming. Bestaande methoden, met name Latent Diffusion Models (LDM), hebben echter enkele beperkingen:

1. Rekenkosten: Traditionele LDM's gebruiken vaak Transformer-architecturen met self-attention mechanismen. Deze hebben een kwadratische complexiteit wat leidt tot hoge rekenkosten en een groot geheugengebruik, vooral bij 3D-medische beelden.
2. Kwaliteit van Latente Ruimte: De meeste LDM's maken gebruik van Variational Autoencoders (VAE) om de data naar een latente ruimte te comprimeren. Het stochastische bemonsteren (sampling) uit een Gaussische verdeling tijdens het trainen van VAE's introduceert ruis, wat vaak resulteert in wazige of minder scherpe reconstructies.
3. Beperkte Controle: Veel bestaande modellen modelleren statische synthese of modality-vertaling, maar missen geavanceerde mechanismen om de tijdsontwikkeling te controleren op basis van meerdere klinische variabelen (zoals leeftijd, geslacht, ziektestatus en hersenvolumes).

Methodologie: Het CLIMB-raamwerk

De auteurs stellen CLIMB (Controllable Longitudinal brain Image generation via MamBa-based latent diffusion model) voor. Dit is een tweestapsraamwerk bestaande uit een aangepaste autoencoder en een latente diffusiemodel.

1. Gaussian-aligned Autoencoder (GATE)

In plaats van een standaard VAE, introduceren de auteurs de GATE om een latente ruimte te creëren die strikt voldoet aan een Gaussische verdeling zonder stochastische ruis.

• Deterministische Uitlijning: GATE gebruikt Sliced Wasserstein Distance (via sliced CDF-alignment) om de verdeling van de gecodeerde latente features direct uit te lijnen met een standaard Gaussische prior. Dit elimineert de noodzaak voor stochastisch bemonsteren (reparameterization trick) en behoudt zo de scherpte en fideliteit van de beelden.
• Verliesfuncties: Het model wordt getraind met een combinatie van reconstructieverlies, sliced Wasserstein-verlies, perceptueel verlies (voor semantische consistentie) en een adversariaal verlies (voor fotorealistische output).

2. Mamba-based Latent Diffusion Model

Het diffusiemodel werkt in de door GATE gegenereerde latente ruimte en gebruikt Mamba (een State Space Model - SSM) in plaats van Transformers.

• Efficiëntie: Mamba vervangt de self-attention mechanismen. Waar self-attention kwadratische complexiteit heeft ($O(N^2)$), biedt Mamba lineaire complexiteit ($O(N)$) door recurrente hidden states te gebruiken. Dit maakt het veel efficiënter voor het verwerken van hoge-resolutie 3D-beelden.
• Conditionering: Het model neemt als input een ruisige latente representatie, een baselijn MRI-scan, de verworven leeftijd, en voorspelde projectie-data (zoals toekomstige leeftijd, ziektestatus en hersenvolumes). Deze conditionele variabelen worden ingebracht via Cross-Attention.
• Tweestaps Trainingsstrategie: Eerst wordt het model getraind op projectie-variabelen alleen, waarna de parameters worden bevroren en het beeldverwerkingsnetwerk wordt toegevoegd en finetuned. Dit stabiliseert het trainingsproces.

3. Inference

Tijdens inferentie worden de huidige MRI en klinische data gebruikt om via het IRLSTM-model (een eerder ontwikkeld model van de auteurs) toekomstige ziektestatus en volumes te voorspellen. Deze data stuurt het Mamba-denoising netwerk om een latente representatie van een toekomstige MRI te genereren, die vervolgens door de GATE-decoder wordt omgezet in een realistische MRI-scan.

Belangrijkste Bijdragen

1. CLIMB Architectuur: Een nieuw raamwerk voor gecontroleerde longitudinale hersenbeeldgeneratie dat gebruikmaakt van een baselijn MRI en projectie-data om anatomische veranderingen te simuleren.
2. Gaussian-aligned Autoencoder (GATE): Een innovatieve autoencoder die sliced Wasserstein distance gebruikt om een deterministische, Gaussische latente ruimte te creëren, wat de beeldkwaliteit en scherpte ten opzichte van traditionele VAE's verbetert.
3. Mamba-integratie: Het toepassen van State Space Models (Mamba) als backbone voor het diffusiemodel, wat de rekenkosten en geheugengebruik aanzienlijk verlaagt ten opzichte van Transformer-gebaseerde modellen, zonder in te leveren op prestaties.
4. Meervoudige Conditionering: Het integreren van diverse factoren (leeftijd, geslacht, ziektestatus, genetische informatie en hersenvolumes) voor nauwkeurigere en persoonlijkere ziekteprogressiemodellering.

Resultaten

Het model is getraind en geëvalueerd op het ADNI-dataset (Alzheimer's Disease Neuroimaging Initiative), bestaande uit 6.306 MRI-scans van 1.390 deelnemers. Vergelijkend met state-of-the-art baselines zoals SADM en BrLP:

• Kwantitatieve Prestaties: CLIMB behaalde de beste scores op alle metrieken:
  • SSIM (Structural Similarity Index): 0.9433 (hoogste, wat aangeeft uitstekende structurele overeenkomst).
  • PSNR: 27.82.
  • MSE: 2.01e-3 (laagste, wat wijst op nauwkeurige pixelreconstructie).
  • LPIPS (Perceptual Similarity): 0.0587 (laagste, wat betekent dat de gegenereerde beelden perceptueel het dichtst bij de echte scans staan).
• Efficiëntie: CLIMB is sneller (2.92 seconden per afbeelding) en verbruikt minder geheugen (5699 MB) dan de concurrenten, mede dankzij de Mamba-architectuur.
• Kwalitatieve Analyse: Visualisaties tonen aan dat CLIMB minder artefacten introduceert en de structurele veranderingen in de hersenen over tijd realistischer simuleert dan SADM en BrLP, vooral in centrale hersengebieden.

Significantie

Deze studie biedt een doorbraak in de generatieve AI voor medische beeldvorming door twee fundamentele problemen op te lossen: de rekenintensiteit van 3D-diffusiemodellen en de vervagingseffecten van VAE's.

• Klinische Toepassing: CLIMB kan helpen bij het voorspellen van de progressie van neurodegeneratieve ziekten, wat essentieel is voor vroege diagnose en behandelplanning.
• Efficiëntie: Door Mamba te gebruiken, wordt het mogelijk om complexe longitudinale modellen te trainen op datasets die anders te groot zouden zijn voor Transformer-architecturen, wat de toepasbaarheid in klinische omgevingen vergroot.
• Toekomstige Richting: Hoewel de kwantitatieve resultaten veelbelovend zijn, benadrukken de auteurs dat validatie door medische experts (radiologen) noodzakelijk is om de klinische relevantie van de gegenereerde beelden volledig te bevestigen.