MorphSeek: Latent Policy for Deformable Registration

• MorphSeek is a framework that reformulates DIR as latent-space policy optimization, enabling fine-grained and efficient deformation estimation.
• It leverages a stochastic Gaussian policy head and multi-trajectory, multi-step GRPO to enhance spatial coherence and data efficiency.
• Experimental evaluations on brain MRI and CT benchmarks show consistent Dice gains of 2–4% and up to 60% NJD reduction, validating its robustness.

MorphSeek is a fine-grained latent representation-level policy optimization framework for deformable image registration (DIR) that reformulates the DIR task as a spatially continuous policy learning process in latent feature space. By introducing a stochastic Gaussian policy head on the high-resolution encoder features and leveraging Group Relative Policy Optimization (GRPO) with multi-trajectory, multi-step exploration, MorphSeek enables highly data-efficient, spatially coherent optimization of dense deformation fields with minimal step-level latency and parameter overhead. The paradigm is agnostic to backbone architecture (e.g., U-Net, TransMorph, NICE-Trans) and underlying optimizer, addressing the challenge of high-dimensional spatially varying deformations in medical imaging while substantially increasing label efficiency and achieving consistent improvements on established 3D DIR benchmarks (Zhang et al., 21 Nov 2025).

1. Architectural Paradigm and Problem Statement

Deformable image registration requires predicting a dense, often millions-dimensional, displacement field $\Phi \in \mathbb{R}^{3 \times H \times W \times D}$ that spatially warps a moving image $I_m$ to align with a fixed image $I_f$. Conventional approaches often rely on direct, per-voxel predictions (encoder–decoder networks) or treat the action space in a reinforcement learning (RL) framework as extremely coarse for tractability, which limits their ability to capture subtle, spatially varying deformations.

MorphSeek addresses this by reparametrizing the policy over the latent feature space of an encoder. Following concatenation of $I_m$ and $I_f$, a standard deep encoder $\mathcal{E}$ produces multi-resolution features $\mathbf{f}_1,...,\mathbf{f}_L$, where $\mathbf{f}_L$ is of highest semantic resolution. Instead of directly decoding $\mathbf{f}_L$, a stochastic Gaussian policy head generates per-voxel means $\boldsymbol{\mu}(\mathbf{f}_L)$ and log-standard deviations $\log\boldsymbol{\sigma}(\mathbf{f}_L)$, producing a latent vector $\mathbf{z}$ via reparametrization: $$\mathbf{z} = \boldsymbol{\mu} + \tau\, \boldsymbol{\sigma} \odot \boldsymbol{\epsilon}, \qquad \boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0},\mathbf{I}) \,,$$ which feeds into the decoder to obtain $$\Phi = \mathcal{D}(\mathbf{f}_1,...,\mathbf{f}_{L-1}, \mathbf{z})$$. The temperature $\tau$ modulates exploration during training. This explicit modeling of a high-dimensional spatially resolved latent policy makes policy gradient optimization feasible and effective, enabling fine-grained, data-driven refinement.

2. Mathematical Formulation and Training Objective

MorphSeek's two-phase training regime commences with an unsupervised warm-up and progresses to weakly-supervised GRPO fine-tuning:

• Warm-up (Unsupervised): With $\tau=0$ (deterministic), $\mathbf{z}=\boldsymbol{\mu}$. The objective combines image similarity ($\mathcal{L}_{\rm sim}$, e.g., MSE or local NCC), regularization on deformation smoothness ($\mathcal{L}_{\rm reg}$), and a KL prior constraint:

$$\mathcal{L}_{\mathrm{warm}}(\theta) = \mathcal{L}_{\mathrm{sim}}(I_f, I_{m\circ\Phi}) + \lambda_{\mathrm{reg}}\, \mathcal{L}_{\mathrm{reg}}(\Phi) + \beta_{\mathrm{KL}}\, D_{\mathrm{KL}}(q_{\theta_E}(\mathbf{z}|\mathbf{f}_L) \| \mathcal{N}(\mathbf{0},\mathbf{I}))$$

• GRPO Fine-Tuning (Weak Supervision): For each pair $(I_m,I_f)$ and corresponding labels $(S_m, S_f)$, the registration is refined over $T$ incremental steps. At each step, $J$ latent samples $\{\mathbf{z}^{(j)}\}$ yield candidate displacement fields $\phi_t^{(j)}$ and compound warps $\Phi_t^{(j)} = \Phi_{t-1} \circ \phi_t^{(j)}$. The reward incorporates incremental Dice improvement and Jacobian determinant regularity (NJD):

$$R^{(j)} = w_{\rm Dice} \left[\mathrm{Dice}(S_f, S_m \circ \Phi_t^{(j)}) - \mathrm{Dice}(S_f, S_m \circ \Phi_{t-1})\right] + w_{\rm NJD}\,\mathrm{NJD}(\Phi_t^{(j)})$$

Advantage-normalized policy gradients over each group of $J$ trajectories are then computed, with variance stabilization via Latent-Dimension Variance Normalization (LDVN):

$$\log\pi_\theta(\mathbf{z}|\cdots) = -\frac{1}{2s}\sum_{i=1}^N \left[\left(\frac{z_i - \mu_i}{\tau \sigma_i}\right)^2 + \log(2\pi \tau^2 \sigma_i^2)\right]$$

The total loss is:

$$\mathcal{L}_{\rm grpo}(\theta) = \mathcal{L}_{\rm policy}(\theta_E) + \lambda_{\rm warm}\mathcal{L}_{\rm warm}(\theta) + \lambda_{\rm Dice} \mathcal{L}_{\rm Dice}(\theta)$$

3. Group Relative Policy Optimization and Stability Mechanisms

The GRPO scheme operates by sampling multiple trajectories at each step, which stabilizes policy optimization and allows for relative ranking of policy outcomes. Each group of sampled trajectories is normalized using both reward and log-likelihood statistics, preventing bias due to high-dimensional action spaces. LDVN rescales log-likelihoods by $1/\sqrt{N}$, where $N$ is the latent dimension, thus maintaining $\operatorname{Var}\sim O(1)$ regardless of $N$.

Cascade-style, multi-step refinement allows each label pair to be reused for incremental coarse-to-fine improvement, increasing label efficiency and registration accuracy. The KL constraint maintains an implicit trust region, while multi-step, multi-trajectory exploration prevents collapse or local overfitting.

4. Experimental Evaluation and Performance Analysis

MorphSeek was systematically evaluated on OASIS (brain MRI), LiTS (liver CT), and Abdomen MR–CT registration tasks across three mainstream backbones (VoxelMorph-L, TransMorph, NICE-Trans). The principal metrics were Dice overlap (higher is better) and negative Jacobian determinant (NJD, lower is better). Across these settings, MorphSeek delivered consistent Dice gains of 2–4% and NJD reductions of 30–60% over baselines:

Dataset	Backbone	Baseline Dice (%)	MorphSeek Dice (%)	Baseline NJD (%)	MorphSeek NJD (%)
OASIS	VoxelMorph-L	84.77 (±2.49)	87.16 (±1.97)	0.15	0.10
OASIS	TransMorph	85.89	88.89	0.16	0.06
OASIS	NICE-Trans	86.79	89.02	0.02	0.02
Abdomen MR–CT	TransMorph	82.37	86.49	0.84	0.35

On cross-modality registration, the accuracy improvement with label-efficient fine-tuning was particularly pronounced. MorphSeek reached 98.5% of full-labeled performance using only ~60% of the labeled pairs, outperforming the original backbones in terms of label efficiency. Overhead was modest, with the stochastic policy head adding <3% parameters and three-step inference increasing latency from 625 ms to 2022 ms on VoxelMorph-L.

Component ablation demonstrated that both the multi-trajectory, multi-step process and LDVN are essential for optimal stability and performance. Neglecting either reduced Dice and yielded unstable or irregular deformations.

5. Label Efficiency, Scalability, and Resource Profile

MorphSeek's multi-step, multi-trajectory GRPO approach allows maximally efficient exploitation of scarce labeled registration pairs, supporting coarse-to-fine spatial adjustment not possible with either single-shot or action-space-restriction RL strategies. The fine-grained latent policy enables spatially detailed deformations at moderate computational cost: <3% increase in parameters, <5 ms per inference step (for single-step), and linear scaling with steps/trajectories.

The framework generalizes across diverse backbones and datasets and supports trade-off between accuracy and runtime at deployment via step/trajectory scheduling.

6. Innovations, Limitations, and Future Directions

Key contributions include:

• Reformulation of DIR as latent-space policy optimization, reducing action dimensionality while preserving spatial granularity and tractability for policy gradients.
• Introduction of LDVN, which is the first approach to variance stabilization for high-dimensional log-likelihoods in policy loss.
• Combination of multi-trajectory, multi-step GRPO to maximally exploit weak labels and enable progressive registration refinement with improved Dice/NJD.

Identified limitations:

• Linear inference latency growth with number of steps; large trajectory/step counts may yield out-of-memory errors.
• Trust region is enforced via fixed KL divergence penalty rather than explicit Proximal Policy Optimization (PPO) or Trust Region Policy Optimization (TRPO) constraints.
• Evaluation to date is restricted to U-Net–style and Transformer backbones and medical registration; extension to tasks such as optical flow or alternative architectures remains an open direction.
• No incorporation of biomechanical priors or alternative reward metrics beyond Dice/NJD; exploration of more adaptive reward or scheduling is suggested.

MorphSeek defines a new data-efficient, scalable, and backbone-agnostic paradigm for deformable registration in high-dimensional visual alignment, achieving state-of-the-art performance with robust spatial coherence and efficient exploitation of limited supervision (Zhang et al., 21 Nov 2025).