MorphSeek: Fine-grained Latent Representation-Level Policy Optimization for Deformable Image Registration

Titel

MorphSeek: Fijne-granulair Latent Representatie-Niveau Beleidsoptimalisatie voor Deformabele Beeldregistratie

1. Het Probleem

Deformabele beeldregistratie (DIR) is een fundamentele maar uitdagende taak in de medische beeldanalyse. Het doel is om voxel-voor-voxel ruimtelijke correspondenties tussen twee 3D-medische beelden tot stand te brengen. Er zijn twee hoofdproblemen die bestaande methoden beperken:

1. Hoge dimensie en supervisie-schaarste: De ruimte van dense verplaatsingsvelden (deformaties) is extreem hoogdimensionaal. Tegelijkertijd is voxel-niveau supervisie (segmentatie-labels) in medische scenario's zeldzaam. Bestaande Reinforcement Learning (RL) frameworks projecteren deze ruimte vaak naar grove, laagdimensionale representaties, wat hun vermogen om ruimtelijk variabele deformaties te vangen, beperkt.
2. Beperkingen van bestaande RL-methoden: RL is natuurlijk geschikt voor stapsgewijze optimalisatie, maar directe toepassing op volledige 3D-deformabele velden als "actie-ruimte" leidt tot onaanvaardbaar hoog geheugengebruik en bemonsteringskosten. Bestaande methoden zijn vaak beperkt tot stijve transformaties of vereisen dure criticus-netwerken.
3. Eén-schot inferentie: De meeste diepe leermodellen voorspellen het vervormingsveld in één keer ("single-shot"). Dit werkt goed voor kleine vervormingen, maar faalt vaak bij complexe, grote niet-rigide vervormingen (zoals in de borstkas of buikholte) omdat lokale grenzen en fijne details verloren gaan.

2. Methodologie: MorphSeek

MorphSeek herformuleert DIR als een optimalisatieproces op latent representatieniveau in plaats van direct op het dense vervormingsveld. Het kader bestaat uit drie fasen:

A. Refactoring voor Latent Beleidsondersteuning

In plaats van een deterministische encoder te gebruiken die direct naar de decoder gaat, introduceert MorphSeek een stochastische Gaussian policy head bovenop de top-laag van de encoder.

• De encoder levert een latent feature $f_L$.
• Twee convolutiehoofden (mean $\mu$ en log-std $\log \sigma$) parametriseren een multivariate Gaussische verdeling $N(\mu, \sigma^2)$.
• Een stochastisch latent vector $z$ wordt bemonsterd via de reparameterization-trick: $z = \mu + \tau \cdot \sigma \odot \epsilon$.
• Deze $z$ wordt gebruikt als input voor de decoder om het vervormingsveld te genereren. Dit maakt de encoder een "beleid" (policy) dat kan verkennen.

B. Unsupervised Warm-up

Voordat RL wordt toegepast, ondergaat het model een fase van ongesuperviseerde warm-up op ongelabelde data.

• Hierbij wordt de temperatuur $\tau$ op 0 gezet (deterministisch: $z = \mu$).
• Dit vormt een stabiele latent ruimte en voorkomt dat het model instort (posterior collapse) voordat de stochastische exploratie begint.
• De loss bestaat uit beeldgelijkenis, regularisatie en een KL-divergentie term.

C. GRPO-based Fine-tuning (Group Relative Policy Optimization)

Na de warm-up wordt het model fijn afgestemd met beperkte labels (weakly supervised) via GRPO.

• Multi-traject & Multi-stap: Voor elk paar beelden worden meerdere trajecten (bemonsteringen) gegenereerd en door meerdere verfijningstappen ($T$) geleid.
• Beloning (Reward): De beloning is gebaseerd op de verbetering in de Dice-score (overlap van segmentaties) en een straf voor negatieve Jacobian-determinanten (NJD), wat zorgt voor fysiek plausibele vervormingen.
• Advantage Normalization: Beloningen worden genormaliseerd binnen een groep om stabiele gradiënten te garanderen.

D. Latent-Dimension Variance Normalization (LDVN)

Dit is een cruciale innovatie om GRPO stabiel te maken in hoge dimensies (tienduizenden latent dimensies).

• In hoge dimensies groeit de variantie van de log-waarschijnlijkheid lineair met de dimensie $N$, wat training instabiel maakt.
• LDVN schaalt de log-waarschijnlijkheid met een factor $s \propto \sqrt{N}$.
• Dit behoudt de richting van de beleidsgradiënt maar normaliseert de schaal, waardoor GRPO numeriek stabiel blijft in 3D-dense voorspellingstaken zonder extra parameters.

3. Belangrijkste Bijdragen

1. Nieuw Paradigma: Een latent-space beleidsoptimalisatie-paradigma voor DIR. Door het beleid te definiëren in de encoder-ruimte in plaats van op het dense veld, wordt fijne granulairheid behouden terwijl de dimensieproblematiek van RL wordt omzeild.
2. LDVN: Een statistische normalisatieschema dat GRPO stabiliseert in hoge-dimensionale ruimtes. Het biedt theoretisch en praktisch bewijs voor de toepassing van RL op dense 3D-predictie.
3. Hoog Label-efficiënt Kader: Een multi-traject, multi-stap framework dat ongesuperviseerde warm-up combineert met GRPO. Dit stelt het model in staat om schaarse labels herhaaldelijk te hergebruiken voor grof-naar-fijn verfijning, wat de prestaties aanzienlijk verbetert zonder extra parameters.

4. Resultaten

MorphSeek werd geëvalueerd op drie 3D-benchmarks: OASIS (hersenen MRI), LiTS (lever CT) en Abdomen MR-CT (kruis-modaal).

• Prestatieverbetering: MorphSeek leverde consistente verbeteringen in de Dice-score (2-4% stijging) en een significante reductie in negatieve Jacobian-determinanten (30-60% daling) ten opzichte van sterke baselines zoals VoxelMorph, TransMorph en NICE-Trans.
• Label-efficiëntie: Met slechts 60% van de trainingslabels bereikte MorphSeek 98,5% van de prestaties van het model met volledige labels, terwijl baselines 80% labels nodig hadden voor vergelijkbare resultaten.
• Efficiëntie: Het voegt minder dan 3% parameters toe aan de bestaande modellen. De inferentielatentie is lineair afhankelijk van het aantal verfijningstappen, maar blijft acceptabel voor klinische toepassingen.
• Stabiliteit: Zonder warm-up faalde GRPO in 67% van de runs; met warm-up steeg het succespercentage naar 79%.

5. Significantie

MorphSeek doorbreekt de barrière die Reinforcement Learning tot nu toe beperkte tot lage-dimensionale, stijve transformaties in de medische beeldregistratie.

• Schaalbaarheid: Het maakt RL-toepassingen mogelijk op volledige 3D-deformabele velden door de actie-ruimte te verschuiven naar een gestructureerde latent ruimte.
• Backbone-onafhankelijk: Het is een plug-in oplossing die werkt met verschillende architecturen (U-Net, Transformers).
• Praktische toepasbaarheid: Het biedt een oplossing voor het tekort aan gelabelde data in de medische wereld, waarbij complexe vervormingen nauwkeurig kunnen worden geregistreerd met minimale supervisie.

Kortom, MorphSeek introduceert een geavanceerde, datagestuurde optimalisatiestrategie die de balans vindt tussen nauwkeurige anatomische uitlijning, fysieke plausibiliteit en computerefficiëntie in complexe medische registratietaken.