[Paper Review] Spatially-varying Regularization with Conditional Transformer for Unsupervised Image Registration

Junyu Chen, Yihao Liu+2 more|arXiv (Cornell University)|Mar 10, 2023

Medical Image Segmentation Techniques4 citations

TL;DR

This paper proposes an end-to-end deep learning framework that learns spatially-varying regularization for unsupervised medical image registration using a conditional Transformer to adaptively control deformation smoothness. By conditioning the regularization hyperparameter into the network and introducing a novel log loss to enforce smoothness, the method achieves superior registration accuracy and reduced deformation irregularities across multiple datasets, outperforming state-of-the-art baselines.

ABSTRACT

In the past, optimization-based registration models have used spatially-varying regularization to account for deformation variations in different image regions. However, deep learning-based registration models have mostly relied on spatially-invariant regularization. Here, we introduce an end-to-end framework that uses neural networks to learn a spatially-varying deformation regularizer directly from data. The hyperparameter of the proposed regularizer is conditioned into the network, enabling easy tuning of the regularization strength. The proposed method is built upon a Transformer-based model, but it can be readily adapted to any network architecture. We thoroughly evaluated the proposed approach using publicly available datasets and observed a significant performance improvement while maintaining smooth deformation. The source code of this work will be made available after publication.

Motivation & Objective

• To address the limitation of spatially-invariant regularization in deep learning-based image registration, which fails to account for region-specific deformation variations.
• To develop an end-to-end trainable framework that learns spatially-varying regularization weights directly from data, improving registration accuracy and smoothness.
• To enable flexible hyperparameter tuning by conditioning the regularization strength into the network architecture, allowing control over deformation smoothness during inference.
• To ensure the method is easily integrable with existing DNN-based registration models without major architectural changes.

Proposed method

• The method introduces a weighted diffusion regularizer that applies spatially-varying regularization by learning a voxel-wise weight volume to modulate deformation smoothness.
• A novel log loss term is proposed to encourage the network to apply stronger regularization where possible, promoting overall smoothness of the deformation field.
• The regularization hyperparameter η is conditioned into a Transformer-based network, enabling dynamic adjustment of regularization strength across different image regions during inference.
• The framework is compatible with various network architectures and can be readily adapted to existing unsupervised DNN registration models.
• The weight volume is predicted end-to-end via a neural network, with the regularization strength controlled by a learnable hyperparameter η.
• The method is evaluated using both synthetic and real medical imaging datasets, demonstrating robustness and generalization across different registration tasks.

Experimental results

Research questions

• RQ1 Can a deep learning-based method effectively learn spatially-varying regularization weights from data to improve image registration performance?
• RQ2 How does conditioning the regularization hyperparameter into the network architecture affect the adaptability and control of deformation smoothness?
• RQ3 Does the proposed log loss term effectively promote global smoothness in the deformation field while allowing local variation in regularization strength?
• RQ4 Can the method be seamlessly integrated into existing unsupervised DNN registration frameworks without architectural overhaul?
• RQ5 How does the proposed method compare to state-of-the-art approaches in terms of registration accuracy and deformation quality?

Key findings

• On the OASIS dataset, the proposed TM-SPR method achieved a Dice score of 0.870 ± 0.017, significantly outperforming TM-TVF (0.869 ± 0.014) and other baselines, with p ≪ 0.0001 in paired t-tests.
• On the IXI dataset, TM-SPR achieved a Dice score of 0.769 ± 0.123, demonstrating statistically significant improvement over all compared methods, including TM-TVF.
• The method maintained low levels of deformation irregularities, with 0.36% of folded voxels (NDV) on the IXI dataset, indicating high diffeomorphic quality.
• The optimal hyperparameter η was found to be 0.45 for OASIS and 2.2 for IXI via grid search, showing the method's sensitivity to data-specific characteristics.
• Visual results confirmed that the method adaptively assigns lower regularization to ventricular regions in OASIS and higher in IXI, aligning with anatomical expectations and loss function constraints.
• The TM-SPR${}_{\text{diff}}$ variant achieved the second-best Dice score (0.762 ± 0.123) on IXI while enforcing nearly zero folded voxels, demonstrating strong diffeomorphic properties.