MedDIFT: Multi-Scale Diffusion-Based Correspondence in 3D Medical Imaging

MedDIFT is een trainingsvrij framework voor 3D medische beeldregistratie dat gebruikmaakt van multi-schaal diffusiële kenmerken uit een vooraf getraind model om nauwkeurige ruimtelijke correspondenties te vinden zonder taakspecifiek trainen.

De kern

Stel je voor dat je twee foto's van dezelfde persoon hebt: één foto is gemaakt terwijl de persoon diep inademde (longen vol), en de andere terwijl hij uitademde (longen leeg). Als je wilt weten waar precies hetzelfde puntje in de longen zit op beide foto's, is dat best lastig. De vorm verandert, de contrasten zijn soms vaag, en er zijn geen duidelijke lijnen om te volgen.

In de medische wereld heet dit het vinden van "correspondentie": het koppelen van punt A op foto 1 aan punt B op foto 2. Dit is cruciaal voor artsen om ziektes te volgen of behandelingen te plannen.

Deze paper introduceert een nieuwe, slimme manier om dit te doen, genaamd MedDIFT. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het Probleem: De "Kale" Vergelijking

Tot nu toe keken computers naar medische scans (zoals CT-schermen) en probeerden ze te matchen op basis van helderheid.

• De analogie: Stel je voor dat je twee landschappen vergelijkt door alleen te kijken naar de kleur van de grond. Als er een stukje mist is of als de grond er anders uitziet door de wind, raak je de weg kwijt. De computer ziet dan alleen "grijze vlekken" en weet niet dat dit eigenlijk een "long" of een "bloedvat" is.

2. De Oplossing: De "Geest" van de Afbeelding

De auteurs gebruiken een heel nieuw type kunstmatige intelligentie, een diffusiemodel. Dit is hetzelfde type AI dat tegenwoordig gebruikt wordt om prachtige nieuwe afbeeldingen te genereren.

• De analogie: In plaats van alleen naar de "huidskleur" (de helderheid) van de afbeelding te kijken, laat MedDIFT de AI de afbeelding "dromen" of "ontleden".
• Het model is getraind om 3D-longscans te maken. Tijdens het proces van het maken (of het "dichtmaken" van een wazige afbeelding), bouwt de AI een intern begrip op van wat er te zien is. Het weet: "Ah, dit is een long, dit is een rib, en dit is een bloedvat," zelfs als het beeld wazig is.

3. Hoe MedDIFT Werkt: De "Geheime Code"

MedDIFT pakt deze interne gedachten van de AI en gebruikt ze als identiteitskaarten voor elk klein puntje (voxel) in de scan.

1. Het Oude Model: De AI is al getraind (het is een "pre-trained" model). Je hoeft er dus niets extra's voor te leren. Het is alsof je een ervaren gids meeneemt die de stad al kent.
2. De Code: Voor elk puntje in de scan haalt de AI een code op die vertelt: "Dit puntje zit in de bovenste longkwab, dicht bij een groot vat."
3. De Match: Vervolgens zoekt de computer in de tweede scan naar het puntje met de exactzelfde code.
   • Vergelijking: Het is alsof je twee mensen in een drukke menigte zoekt. In plaats van te kijken naar hun kleren (die kunnen veranderen), kijk je naar hun DNA (de diepe, onveranderlijke structuur). Als het DNA matcht, weet je dat het hetzelfde persoon is, ongeacht of ze nu een hoed op hebben of niet.

4. Waarom is dit speciaal?

• Geen extra training: De meeste nieuwe methodes moeten maandenlang getraind worden op duizenden scans. MedDIFT doet het direct, "zonder training" (training-free). Het gebruikt de wijsheid die het model al heeft.
• Meer lagen: De auteurs ontdekten dat je niet alleen naar de "hoofdgedachte" moet kijken, maar ook naar de details. Ze combineren verschillende niveaus van inzicht (zoals een kaart die zowel de hele stad toont als de kleine straatjes). Dit werkt beter dan alleen kijken naar één detail.
• De "Zoekruimte": Soms helpt het om te zeggen: "Het puntje zit waarschijnlijk niet 10 meter verderop, maar hooguit 2 centimeter." Door de zoektocht te beperken tot een klein vakje, wordt het resultaat nog nauwkeuriger.

Het Resultaat

Op een publieke dataset van longscans bleek MedDIFT net zo goed (en soms beter) te presteren dan geavanceerde methodes die wel getraind moesten worden.

• Conclusie: Het is alsof je een oude, ervaren detective (de AI) hebt die de stad kent, en die je kunt sturen om verbanden te leggen tussen twee verschillende momentopnamen, zonder dat je hem eerst een nieuwe les hoeft te geven.

Kortom: MedDIFT maakt het makkelijker voor artsen om te zien hoe ziektes zich verplaatsen of veranderen in het lichaam, door gebruik te maken van de "diepe kennis" van moderne AI, zonder dat ze zelf zware rekenwerk hoeven te doen.

Technische samenvatting

Probleemstelling

In de medische beeldvorming is het nauwkeurig vaststellen van ruimtelijke correspondentie tussen afbeeldingen (bijvoorbeeld tussen verschillende tijdstippen, patiënten of modaliteiten) cruciaal voor longitudinale analyse, het volgen van laesies en therapieplanning. Traditionele registratiemethoden, zowel klassiek als op leer gebaseerd, vertrouwen voornamelijk op lokale intensiteitsgebonden vergelijkingsmaten (zoals kruiscorrelatie of wederzijdse informatie).

Deze methoden hebben echter beperkingen:

• Ze falen vaak in gebieden met laag contrast, artefacten of grote anatomische variabiliteit.
• Ze missen het vermogen om globale semantische structuren te vangen, omdat ze puur opereren op lokale verschijning.

Hoewel recente vooruitgang in diffusiemodellen heeft aangetoond dat hun tussenliggende representaties rijke geometrische en semantische informatie bevatten, zijn bestaande frameworks (zoals DIFT voor natuurlijke 2D-afbeeldingen) niet direct toepasbaar op 3D medische data of zijn ze niet specifiek getraind op medische beelden.

Methodologie: MedDIFT

De auteurs presenteren MedDIFT, een trainingsvrij (training-free) raamwerk voor het vinden van voxel-correspondenties in 3D medische beelden. Het systeem maakt gebruik van een vooraf getraind 3D medisch diffusiemodel, specifiek MAISI (een latent diffusion model voor 3D CT-beelden), zonder dat er fijnafstelling (fine-tuning) of taakspecifiek trainen nodig is.

Het proces verloopt in drie fasen:

1. Extractie van Diffusie-features:

   • Een paar 3D-medische beelden ($A$ en $B$) wordt via de variational autoencoder (VAE) van MAISI gecodeerd naar een latente representatie $z_0$.
   • Er wordt Gaussisch ruis toegevoegd om een verstoord latent $z_t$ te verkrijgen op een specifiek timestap $t$.
   • Dit verstoord latent wordt door de bevroren (frozen) diffusie U-Net van MAISI gevoerd voor één stap ontruisen.
   • Tussenliggende activaties ($F_{l,t}$) worden geëxtraheerd uit meerdere decoder-blokken ($l$) op verschillende tijdstippen ($t$). Deze activaties bevatten semantische informatie die varieert van lokaal tot globaal, afhankelijk van het niveau en de hoeveelheid ruis.

2. Constructie van Multi-Schaal Descriptoren:

   • De geëxtraheerde feature maps hebben verschillende ruimtelijke resoluties (bijv. 1/16, 1/8, 1/4 van de originele grootte).
   • Om een uniforme descriptor te krijgen, worden alle feature maps via trilineaire upsampling teruggebracht naar de originele beeldresolutie.
   • Vervolgens worden ze genormaliseerd ($L_2$) en geconcateneerd over de verschillende schaalniveaus. Dit resulteert in een rijke, voxel-voor-voxel descriptor die zowel fijne details als globale semantiek combineert.

3. Correspondentie Matching:

   • Voor een query-voxel $p$ in beeld $A$ wordt de corresponderende voxel $q^*$ in beeld $B$ gevonden door de cosinus-similariteit tussen de diffusie-descriptoren te maximaliseren.
   • Optioneel kan de zoekruimte worden beperkt tot een lokaal gebied rondom de verwachte positie (MedDIFT-Box) om de rekentijd te verlagen en onwaarschijnlijke matches uit te sluiten, wat nuttig is bij vooraf gerigiditeerde beelden.

Belangrijkste Bijdragen

• Eerste 3D Medische Toepassing: MedDIFT is het eerste framework dat diffusie-features van een vooraf getraind 3D medisch diffusiemodel gebruikt voor het vaststellen van voxel-correspondenties.
• Trainingsvrij: Het systeem vereist geen gewichtsoptimalisatie of taakspecifiek trainen, wat het een krachtige, directe oplossing maakt.
• Multi-Schaal Fusie: De auteurs tonen aan dat het combineren van features van meerdere decoder-niveaus essentieel is voor de prestaties in medische beelden.
• Competitieve Prestaties: Het bereikt vergelijkbare nauwkeurigheid met geavanceerde deep learning registratiemodellen, maar zonder de noodzaak van training.

Resultaten

De methode is geëvalueerd op het Learn2Reg Lung CT-dataset (intra-patiënt inspiratoire en expiratoire CT-scans met geannoteerde sleutelpunten).

• Vergelijking: MedDIFT werd vergeleken met:
  • NiftyReg: Een conventionele B-spline free-form deformation (FFD) methode.
  • UniGradICON: Een recente deep learning foundation model voor medische registratie.
• Prestatiemetingen: De fout werd gemeten als de Euclidische afstand (in mm) tussen voorspelde en ground-truth sleutelpunten.
  • NiftyReg behaalde de laagste gemiddelde fout (5.98 mm case mean).
  • UniGradICON had een hogere fout (10.03 mm).
  • MedDIFT behaalde een vergelijkbare fout met UniGradICON (10.47 mm case mean), maar met een lagere standaarddeviatie, wat wijst op grotere stabiliteit.
  • De variant met beperkte zoekruimte (MedDIFT-Box) verbeterde de prestaties further (9.97 mm), wat dicht in de buurt kwam van de conventionele NiftyReg.
• Ablatiestudies:
  • Het combineren van features van alle vier decoder-niveaus leverde de beste resultaten op.
  • Een matige hoeveelheid ruis (timestep $t=20$) bleek optimaal; te veel ruis (hoge $t$) degradeerde de prestaties.

Betekenis en Toekomstperspectief

MedDIFT bewijst dat semantische representaties afgeleid van diffusiemodellen een veelbelovend alternatief zijn voor traditionele intensiteitsgebaseerde vergelijkingsmaten in de medische beeldregistratie. Het overbrugt de kloof tussen conventionele lokale registratie en rijke, geleerde features.

Hoewel het niet consistent alle bestaande methoden overtreft, biedt het een unieke waarde als trainingsvrij alternatief dat direct inzetbaar is op nieuwe data zonder extra berekeningskosten voor training. De resultaten suggereren dat voor medische beelden de combinatie van grove semantische en fijne ruimtelijke informatie via multi-scale fusie cruciaal is.

Toekomstig werk zal zich richten op het fijnafstellen van de feature-extractoren, het verbeteren van de fusiestrategieën en het integreren van MedDIFT in bredere registratie- of multimodale correspondentie-frameworks.