Unsupervised learning of spatially varying regularization for diffeomorphic image registration

TL;DR

This work introduces an unsupervised, end-to-end framework for learning spatially varying regularization in diffeomorphic image registration. By modeling a voxelwise regularization field with a hyperprior and predicting it via a lightweight decoder, the method achieves subject- and region-specific deformation control without relying on anatomical labels. The approach unifies a diffusion-like prior over regularization with either a normal or beta prior on per-voxel weights and optimizes hyperparameters through Bayesian optimization, achieving improved Dice scores and lower deformation irregularities across whole-body, brain, cardiac, and lung datasets while preserving diffeomorphism. The resulting spatial weight maps also provide interpretability into registration behavior, enabling adaptive, discontinuity-aware registration suitable for complex anatomical motions in medical imaging.

Abstract

Spatially varying regularization accommodates the deformation variations that may be necessary for different anatomical regions during deformable image registration. Historically, optimization-based registration models have harnessed spatially varying regularization to address anatomical subtleties. However, most modern deep learning-based models tend to gravitate towards spatially invariant regularization, wherein a homogenous regularization strength is applied across the entire image, potentially disregarding localized variations. In this paper, we propose a hierarchical probabilistic model that integrates a prior distribution on the deformation regularization strength, enabling the end-to-end learning of a spatially varying deformation regularizer directly from the data. The proposed method is straightforward to implement and easily integrates with various registration network architectures. Additionally, automatic tuning of hyperparameters is achieved through Bayesian optimization, allowing efficient identification of optimal hyperparameters for any given registration task. Comprehensive evaluations on publicly available datasets demonstrate that the proposed method significantly improves registration performance and enhances the interpretability of deep learning-based registration, all while maintaining smooth deformations.

Figures (25)

• Figure 1: The probabilistic dependencies of the model parameters are graphically illustrated in the diagram. Random variables are depicted as circles, while rounded squares represent model parameters. Shaded quantities indicate observed elements at test time, and the plate indicates replication given different samples.
• Figure 2: Weighted graph Laplacian on a 2D image grid centered at $\mathbf{p}_i$. Red edges represent forward adjacency. Blue and green edges indicate backward adjacency. Each color corresponds to a distinct edge weight.
• Figure 3: The plots of $p^{\mathcal{B}e}(\pmb{\lambda})$ and $\log p^{\mathcal{B}e}(\pmb{\lambda})$ under the assumption of beta distribution, with varying shape parameter $\alpha$.
• Figure 4: The overall framework of the proposed method for unsupervised learning of spatially varying regularization for diffeomorphic image registration. The method introduces lightweight ConvNet blocks that take the feature maps from the bottleneck layers of the registration network to generate a spatial weight volume $\pmb{\lambda})$, assigning regularization at the voxel level. The spatial weight volume is used to regularize the deformation during training. Additionally, the spatial weight volume is itself regularized by an auxiliary term (i.e., $\mathcal{L}^{\mathcal{N} \mathrm{or} \mathcal{B}e}(\pmb{\lambda})$), promoting smoother deformations when appropriate. While the framework is demonstrated using TransMorph as the registration backbone, it is agnostic to the choice of registration networks and can be easily adapted to other architectures.
• Figure 5: The schematic of the ConvNet block that produces the spatial weight volume, $\pmb{\lambda}$.