Source-linked AI summary
Fast Predictive Image Registration
Xiao Yang, Roland Kwitt, Marc Niethammer
TL;DR
Non-parametric image registration can be computationally demanding, and predicting parameters is difficult when image appearance and high-dimensional parameters are nonlinearly related. This paper predicts LDDMM momentum from image patches, reporting faster computation while maintaining high prediction accuracy and diffeomorphic transformations.
Problem
Non-parametric registration can be computationally demanding, and high-dimensional registration parameters are difficult to predict from image appearance because their relationship is nonlinear.
Source: §1 · p. 1Method
Results
The method reports dramatic speedups over direct optimization while maintaining high prediction accuracy and diffeomorphic transformations.
Source: §1 · p. 2Takeaways & Limitations
The atlas-based formulation keeps predictions in a consistent tangent space, and deformation uncertainty highlights ambiguous areas.
Source: §1 · p. 2, §4 · p. 8The authors identify landmark- and segmentation-overlap accuracy comparisons with optimization-based techniques as future evaluation needs.
Source: §4 · p. 8
Abstract
from arXiv · showhide
We present a method to predict image deformations based on patch-wise image appearance. Specifically, we design a patch-based deep encoder-decoder network which learns the pixel/voxel-wise mapping between image appearance and registration parameters. Our approach can predict general deformation parameterizations, however, we focus on the large deformation diffeomorphic metric mapping (LDDMM) registration model. By predicting the LDDMM momentum-parameterization we retain the desirable theoretical properties of LDDMM, while reducing computation time by orders of magnitude: combined with patch pruning, we achieve a 1500x/66x speed up compared to GPU-based optimization for 2D/3D image registration. Our approach has better prediction accuracy than predicting deformation or velocity fields and results in diffeomorphic transformations. Additionally, we create a Bayesian probabilistic version of our network, which allows evaluation of deformation field uncertainty through Monte Carlo sampling using dropout at test time. We show that deformation uncertainty highlights areas of ambiguous deformations. We test our method on the OASIS brain image dataset in 2D and 3D.
1 Introduction
Image registration provides spatial correspondences but can be computationally demanding for non-parametric models. Predicting registration parameters is challenging because they are high-dimensional and relate nonlinearly to image appearance; this paper explores patch-based deep regression for that task.
1 Introduction
- Non-parametric registration models, including LDDMM, can make optimization computationally demanding, while image registration provides spatial correspondences.Atlas-to-image registration is commonly used for atlas-based segmentations and population analyses.Source: §1 · p. 1
- Existing prediction approaches are constrained by affine or low-rank parameterizations, key-point selection, or an overly restrictive linear appearance–parameter relationship.The cited approaches include a multi-scale linear regressor, key-point matching with interpolation, and semi-coupled dictionary learning.Source: §1 · pp. 1–2
- The paper predicts LDDMM momentum from image patches to retain diffeomorphic transformations while pursuing faster computation, high prediction accuracy, atlas-coordinate consistency, and uncertainty estimates.The authors describe the atlas viewpoint as keeping predictions within a consistent tangent space and propose a Bayesian model for parameter and deformation-map uncertainty.Source: §1 · p. 2
2 Initial Momentum LDDMM Parameterization
In LDDMM shooting, the initial momentum is the registration parameter used to determine the deformation from source and target images. The paper predicts this parameter patch by patch, retaining LDDMM properties while avoiding difficulties that alternative parameterizations can face in homogeneous regions.
2 Initial Momentum LDDMM Parameterization
- The initial momentum m0 is the LDDMM registration parameter from which the deformation map is computed by integrating the geodesic equations.The source image S is mapped toward the target image T at time t = 1.Source: §2 · p. 2
- Predicting momentum leaves smoothness to be imposed afterward through K, which can preserve LDDMM’s diffeomorphic transformations.The momentum may be compactly supported at image edges, while K smooths it to obtain velocity.Source: §2 · p. 3
- Patch-wise prediction of velocity or displacement can be difficult in homogeneous regions, whereas momentum is zero in those regions.The paper identifies this as a reason momentum avoids those prediction difficulties.Source: §2 · p. 3
3 Network Structure
The network predicts LDDMM initial momentum from paired image patches using an encoder-decoder architecture, with dropout enabling probabilistic predictions. Patch pruning and large strides reduce whole-image prediction work.
3 Network Structure
- Paired fixed and moving 15 × 15 patches feed a VGG-style encoder, which predicts initial momentum patches for the spatial directions.The 2D input has two image layers; the 3D network adds volumetric layers and a decoder for the third direction.Source: §3 · p. 3
- Separate decoders share encoder features and predict momentum independently by spatial dimension; overlapping patch outputs are averaged for whole-image prediction.The decoder uses max-unpooling with the pooling indices, and training evaluates output with an L1 difference.Source: §3 · pp. 3–4
- Dropout makes the network an approximate Bayesian model, and test-time Monte Carlo samples yield a mean momentum prediction and deformation variance.The sampled initial momenta are integrated separately to calculate deformation variance.Source: §3 · p. 4
- 99.5%/99.97% fewer patches are predicted for 128×128 2D/128 × 128 × 128 3D images using background pruning and large strides, with negligible accuracy loss.Background patches are pruned because their initial momentum should be zero; the stride is 14 for 15 × 15 patches.Source: §3 · pp. 3–5
4 Experiments
On OASIS 2D and 3D brain images, momentum-based prediction retains diffeomorphic transformations and is substantially faster than GPU-based LDDMM optimization. The experiments also compare parameterizations, examine uncertainty, and discuss where the approach may generalize or be limited.
4 Experiments
- The OASIS evaluation used 100 training targets and 50 test targets, with atlas-based LDDMM momenta as network training targets.The probabilistic network used 50 samples per test case to estimate predictions and deformation-field variation.Source: §4 · pp. 5–6
- In 2D, momentum networks outperform SCDL and improve accuracy over velocity and displacement prediction at both tested stride settings.Both deterministic and probabilistic momentum networks show comparatively small sensitivity to stride and patch pruning; their transformations remain diffeomorphic in all scenarios.Source: §4 · pp. 6–7
- Uncertainty maps highlight areas of drastic or ambiguous deformation, including the anterior ventricle edge and posterior cortex in the 2D example.The displayed uncertainty is derived from deformation variance across directions; red denotes high and blue low uncertainty.Source: Figure 2 · p. 6, §4 · p. 7
- 1500x/66x speedups over GPU-based PyCA optimization were achieved for 2D/3D prediction, with subvoxel accuracy for about 99.8% of 3D voxels.With stride 14 and patch pruning, prediction took 0.19s for a 2D image and 7.68s for a 3D image.Source: §4 · p. 7
- The model may extend to other pixel- or voxel-wise registration methods, while B-spline decoders could be replaced by fully connected layers.The discussion notes that large strides and patch pruning may be unsuitable when parameter locality is not guaranteed.Source: §4 · pp. 7–8
- Landmark and segmentation-overlap accuracy against optimization-based techniques remains for future study.The authors also identify multi-atlas segmentation, multimodal registration, and multi-patch-scale networks as possible extensions.Source: §4 · p. 8