Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Runxun Zhang, Yizhou Liu, Li Dongrui, Bo XU, Jingwei Wei

TL;DR

Deformable image registration faces challenges from high-dimensional deformation fields and scarce voxel-level supervision. MorphSeek reframes DIR as latent-space policy optimization, placing a stochastic Gaussian head on the encoder to operate in the latent feature space and enabling coarse-to-fine refinement via Group Relative Policy Optimization (GRPO). Key contributions include Latent-Dimension Variance Normalization (LDVN) to stabilize high-dimensional policy updates, an unsupervised warm-up to shape the latent space, and a multi-trajectory, multi-step GRPO fine-tuning regime that reuses scarce labels efficiently. Across three 3D DIR benchmarks (OASIS brain MRI, LiTS liver CT, and Abdomen MR-CT) and multiple backbones, MorphSeek delivers consistent Dice gains and NJD reductions with modest parameter and latency overhead, demonstrating practical RL-based DIR under realistic label budgets.

Abstract

Deformable image registration (DIR) remains a fundamental yet challenging problem in medical image analysis, largely due to the prohibitively high-dimensional deformation space of dense displacement fields and the scarcity of voxel-level supervision. Existing reinforcement learning frameworks often project this space into coarse, low-dimensional representations, limiting their ability to capture spatially variant deformations. We propose MorphSeek, a fine-grained representation-level policy optimization paradigm that reformulates DIR as a spatially continuous optimization process in the latent feature space. MorphSeek introduces a stochastic Gaussian policy head atop the encoder to model a distribution over latent features, facilitating efficient exploration and coarse-to-fine refinement. The framework integrates unsupervised warm-up with weakly supervised fine-tuning through Group Relative Policy Optimization, where multi-trajectory sampling stabilizes training and improves label efficiency. Across three 3D registration benchmarks (OASIS brain MRI, LiTS liver CT, and Abdomen MR-CT), MorphSeek achieves consistent Dice improvements over competitive baselines while maintaining high label efficiency with minimal parameter cost and low step-level latency overhead. Beyond optimizer specifics, MorphSeek advances a representation-level policy learning paradigm that achieves spatially coherent and data-efficient deformation optimization, offering a principled, backbone-agnostic, and optimizer-agnostic solution for scalable visual alignment in high-dimensional settings.

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

TL;DR

Deformable image registration faces challenges from high-dimensional deformation fields and scarce voxel-level supervision. MorphSeek reframes DIR as latent-space policy optimization, placing a stochastic Gaussian head on the encoder to operate in the latent feature space and enabling coarse-to-fine refinement via Group Relative Policy Optimization (GRPO). Key contributions include Latent-Dimension Variance Normalization (LDVN) to stabilize high-dimensional policy updates, an unsupervised warm-up to shape the latent space, and a multi-trajectory, multi-step GRPO fine-tuning regime that reuses scarce labels efficiently. Across three 3D DIR benchmarks (OASIS brain MRI, LiTS liver CT, and Abdomen MR-CT) and multiple backbones, MorphSeek delivers consistent Dice gains and NJD reductions with modest parameter and latency overhead, demonstrating practical RL-based DIR under realistic label budgets.

Abstract

Deformable image registration (DIR) remains a fundamental yet challenging problem in medical image analysis, largely due to the prohibitively high-dimensional deformation space of dense displacement fields and the scarcity of voxel-level supervision. Existing reinforcement learning frameworks often project this space into coarse, low-dimensional representations, limiting their ability to capture spatially variant deformations. We propose MorphSeek, a fine-grained representation-level policy optimization paradigm that reformulates DIR as a spatially continuous optimization process in the latent feature space. MorphSeek introduces a stochastic Gaussian policy head atop the encoder to model a distribution over latent features, facilitating efficient exploration and coarse-to-fine refinement. The framework integrates unsupervised warm-up with weakly supervised fine-tuning through Group Relative Policy Optimization, where multi-trajectory sampling stabilizes training and improves label efficiency. Across three 3D registration benchmarks (OASIS brain MRI, LiTS liver CT, and Abdomen MR-CT), MorphSeek achieves consistent Dice improvements over competitive baselines while maintaining high label efficiency with minimal parameter cost and low step-level latency overhead. Beyond optimizer specifics, MorphSeek advances a representation-level policy learning paradigm that achieves spatially coherent and data-efficient deformation optimization, offering a principled, backbone-agnostic, and optimizer-agnostic solution for scalable visual alignment in high-dimensional settings.

Paper Structure

This paper contains 17 sections, 17 equations, 5 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. DL-Based Deformable Medical Image Registration
  4. Reinforcement Learning in DIR
  5. Method
  6. Refactoring Registration Networks for Latent Policy Learning
  7. Warm-up Priors for Stable Policy Optimization
  8. Multi-Trajectory GRPO for Step-Wise Registration
  9. Experiments
  10. Datasets
  11. Baselines and Implementations
  12. Results and Analysis
  13. Overall Performance Across Tasks and Backbones
  14. Policy & Label Efficiency Ablation
  15. Computational Overhead and the Role of Warm-up
  16. ...and 2 more sections

Figures (5)

  • Figure 1: Two Major Challenges Faced by DL-based DIR
  • Figure 2: MorphSeek Registration Framework Process
  • Figure 3: The Performance of MorphSeek Across Three Different Tasks. Labels are overlaid only for the two abdominal datasets; OASIS is left unlabeled to avoid clutter from its 35 foreground classes.
  • Figure 4: Impact of Warm-up and MorphSeek on GRPO Fine-tuning Performance with Limited Labeled Data (OASIS dataset)
  • Figure 5: Validation Dice on OASIS: Effect of Warm-up Before GRPO (TransMorph Backbone), where the dotted line marks the switch from warm-up to GRPO