Source-linked AI summary
Networks for Joint Affine and Non-parametric Image Registration
Zhengyang Shen, Xu Han, Zhenlin Xu, Marc Niethammer
TL;DR
Existing learned image-registration approaches can require pre-alignment, face memory limits on full-resolution 3D maps, or omit iterative refinement. AVSM combines affine registration with a momentum-based vSVF model in an end-to-end framework. On longitudinal and cross-subject knee MRI, it achieves comparable and sometimes better performance than popular registration tools, with reduced computation time and strong deformation regularity and symmetry.
Problem
Existing deep-learning registration approaches may require pre-alignment, be constrained by full-resolution 3D memory demands, or lack iterative refinement.
Source: §1 · p. 2Method
AVSM combines a multi-step affine network with a momentum-based vSVF deformation network, composing transformation maps in one end-to-end framework.
Source: §1 · p. 2Results
On longitudinal and cross-subject knee MRI, AVSM achieves comparable and sometimes better performance than popular registration tools, with dramatically reduced computation time and excellent deformation regularity and symmetry.
Source: §4 · p. 8Takeaways & Limitations
The reported knee-MRI results support AVSM as a single-forward-pass approach combining affine pre-registration and non-parametric deformation.
Source: §4 · p. 8Evaluation covers longitudinal and cross-subject knee MR images; evaluation on other registration tasks is left for future work.
Source: §4 · p. 8
Abstract
from arXiv · showhide
We introduce an end-to-end deep-learning framework for 3D medical image registration. In contrast to existing approaches, our framework combines two registration methods: an affine registration and a vector momentum-parameterized stationary velocity field (vSVF) model. Specifically, it consists of three stages. In the first stage, a multi-step affine network predicts affine transform parameters. In the second stage, we use a Unet-like network to generate a momentum, from which a velocity field can be computed via smoothing. Finally, in the third stage, we employ a self-iterable map-based vSVF component to provide a non-parametric refinement based on the current estimate of the transformation map. Once the model is trained, a registration is completed in one forward pass. To evaluate the performance, we conducted longitudinal and cross-subject experiments on 3D magnetic resonance images (MRI) of the knee of the Osteoarthritis Initiative (OAI) dataset. Results show that our framework achieves comparable performance to state-of-the-art medical image registration approaches, but it is much faster, with a better control of transformation regularity including the ability to produce approximately symmetric transformations, and combining affine and non-parametric registration.
1. Introduction
Registration establishes spatial correspondences for anatomical analysis, but existing learned approaches have limitations in pre-alignment, 3D memory use, and iterative refinement. AVSM combines affine and vSVF registration in one end-to-end framework, using map compositions and multi-step refinement.
1. Introduction
- Registration establishes spatial correspondences for anatomical analysis, while conventional methods optimize image mismatch subject to transformation regularity.Such optimization can be computationally expensive in high-dimensional parameter spaces.Source: §1 · p. 1
- Existing deep-learning methods may require pre-alignment, face memory limits for full-resolution 3D maps, or lack iterative refinement.Pre-alignment can make an approach non-end-to-end or become a computational bottleneck.Source: §1 · p. 2
- AVSM combines affine and vSVF registration in one end-to-end framework, with momentum-based smoothness control, inverse-consistency losses, and multi-step refinement.The framework also avoids unnecessary image interpolations by composing transformation maps.Source: §1 · p. 2
- A multi-step affine network initializes vSVF registration, whose momentum is smoothed into a velocity field and refined through map compositions.The source image is interpolated once; down-sampled transformation maps reduce computation and memory use.Source: §1 · p. 2, Figure 1 · p. 2
- AVSM is compared with publicly available optimization-based methods on longitudinal and cross-subject 3D OAI image pairs.The framework’s contributions include controlling smoothness through momentum prediction and encouraging similar correspondences when image order is reversed.Source: §1 · p. 2
2. Methods
The methods combine a multi-step affine network with a momentum-parameterized stationary velocity field (vSVF) component for local deformations. The affine network refines global transformations, while vSVF uses smoothed momentum to produce a transformation map.
2. Methods
- The framework pairs multi-step affine registration with a vSVF component that accounts for local deformations.The vSVF component can also be applied in multiple steps.Source: §2 · p. 3
2.1. Multi-step Affine Network
- Because non-parametric methods are typically not affine-invariant, the affine network directly predicts registration parameters to handle global displacement and rotation.Composing several affine steps improves accuracy and stability over training a single affine network.Source: §2.1 · p. 3
- At each recurrent affine step, the network refines parameters from the target and currently warped source, updating parameters rather than resampling intermediate images.This avoids numerical instabilities and dissipation from successive trilinear interpolations.Source: §2.1 · p. 3
- The affine network combines image similarity, transformation regularization, and a symmetry loss that encourages inverse consistency.The similarity measure can use multi-kernel LNCC, and the affine regularization weight is initially large before gradually decaying to zero.Source: §2.1 · pp. 3–4
2.2. Vector Momentum-parameterized SVF
- In vSVF, the network predicts momentum that is smoothed to obtain velocity, allowing spatial smoothness to be controlled separately from parameter prediction.The formulation is intended to guarantee sufficient velocity smoothness and permit diffeomorphic transformations even for large displacements.Source: §2.2 · p. 4
- The vSVF unit takes predicted momentum and a down-sampled initial map, produces a low-resolution transformation map, and upsamples it to full resolution.The initial map may be the affine map or a map from a previous vSVF step; low-resolution computation reduces memory consumption.Source: §2.2 · p. 4
- The vSVF loss combines image similarity, velocity-field regularization, and a symmetry term encouraging inverse consistency.Its similarity loss uses the same multi-kernel LNCC formulation as the affine component.Source: §2.2 · p. 5
3. Experiments and Results
Experiments on OAI knee MRI evaluate AVSM in longitudinal and cross-subject registration against optimization-based and public methods. AVSM performs best in cross-subject registration, is slightly below optimization methods longitudinally, and completes registration much faster than iterative optimization.
Evaluation setup
- The OAI evaluation uses 176 labeled longitudinal pairs and 300 cross-subject test pairs, with average Dice score as the metric.The dataset includes 176 manually labeled MR images from 88 patients and 22,950 unlabeled images from 2,444 patients.Source: §3 · p. 5
Training details
- The affine network is trained with three steps for longitudinal registration and five for cross-subject registration, then tested with seven steps in both settings.The momentum network is trained for two vSVF steps, with the affine component fixed during that training.Source: §3 · p. 5
Baselines
- AVSM is compared with optimization-based affine and vSVF methods and public tools including SyN, Demons, NiftyReg, and VoxelMorph.VoxelMorph is initialized with the proposed affine network because it assumes pre-aligned images.Source: §3 · pp. 6–7
Results
- AVSM shows high performance, small variance, and fewer registration outliers across longitudinal and cross-subject tasks.Registration uses one forward pass for both affine and vSVF stages, making AVSM much faster than iterative numerical optimization.Source: Figure 5 · p. 7, §3 · p. 7
- AVSM performs best in cross-subject registration and slightly exceeds optimization-based methods for affine and non-parametric registration, while NiftyReg with LNCC is similar.For longitudinal registration, AVSM performs well but is slightly below optimization-based methods, including vSVF-opt.Source: §3 · p. 7, Table 1 · p. 7
Ablation study
- Ablations show gains from vSVF registration, multi-step training, and mk-LNCC, while adding vSVF symmetry loss slightly worsens registration performance.The single-step affine network is difficult to train without regularization; multi-step training and inverse consistency improve its performance.Source: §3 · p. 7, Table 2 · p. 8
Symmetry
- AVSM obtains low symmetry-measure values and a balance between accuracy and symmetry; Demons and SyN also encourage symmetry.The symmetry measure is evaluated within the image volume, excluding a 10-voxel boundary margin.Source: §3 · pp. 7–8
Iterative refinement
- Testing a two-step-trained vSVF model for more than two steps can increase Dice, but average fold counts also rise, mainly near boundaries and anatomical inconsistencies.The additional iterations act as registration refinements.Source: §3 · p. 8
4. Conclusions and Future Work
AVSM combines affine and momentum-based non-parametric registration in one framework. On knee MRI registrations, it achieves comparable and sometimes better performance, with reduced computation time and strong deformation regularity and symmetry.
4. Conclusions and Future Work
- AVSM achieves comparable and sometimes better performance than popular registration tools, with dramatically reduced computation time and excellent deformation regularity and symmetry.The framework outputs a transformation map combining affine pre-registration and vSVF deformation in a single forward pass.Source: §4 · p. 8
- AVSM shows good symmetry results in both longitudinal and cross-subject registration evaluations.The evaluation compares symmetry across methods, with lower values indicating better symmetry.Source: Figure 7 · p. 8
- With two-step vSVF training, performance increases as test steps increase, while the number of folds also rises.Source: Figure 8 · p. 8
A. Supplementary material
The supplementary material adds details on affine regularization, step-count performance, momentum-network structure, and registration examples.
A. Supplementary material
- The supplement covers epoch-dependent affine regularization, affine-network performance across step counts, momentum-generation network structure, and additional registration examples.Source: A. Supplementary material · p. 9
A.1. Affine regularization factor
An epoch-dependent regularization factor discourages large affine transformations early in training, then decays until its influence becomes negligible.
A.1. Affine regularization factor
- The affine regularization factor discourages large transformations at the start of training to help the network converge.Source: §A.1 · p. 9
- The regularization factor decays to zero over the epochs, making its training influence negligible.Source: §A.1 · p. 9, Figure 9 · p. 9
- For both registration settings, K_ar is 4 and C_ar is 10; K_ar controls decay rate and C_ar is a constant.Source: §A.1 · p. 9
A.2. Dice over steps in Multi-step Affine Network
The multi-step affine network improves Dice performance as test steps increase, then saturates; it is trained with three longitudinal steps and five cross-subject steps.
A.2. Dice over steps in Multi-step Affine Network
- The affine network is trained with three steps for longitudinal registration and five for cross-subject registration.Source: §A.2 · p. 9, Figure 10 · p. 9
- Affine-network Dice performance improves with more test steps and eventually saturates at a high level.Source: Figure 10 · p. 9, §A.2 · p. 9
A.3. Structure of Momentum Generation Network
The Momentum Generation Network takes a pair of images and outputs a low-resolution initial momentum. Its four-level U-net uses residual links and omits the last decoder level.
A.3. Structure of Momentum Generation Network
- The network takes a pair of images as input and outputs a low-resolution initial momentum.Source: §A.3 · p. 9
- Its four-level U-net has residual links and removes the last decoder level to produce the low-resolution momentum.Source: §A.3 · p. 9
- No activation follows the last two convolutional layers because the momentum can be positive or negative.Source: §A.3 · p. 9
A.4. Visualization
The visualizations examine folds in six-step vSVF, registration changes over vSVF steps, and additional AVSM examples.
Folds
- Folds in six-step vSVF mostly occur at anatomical inconsistencies or image boundaries, where large predicted momenta and discretization artifacts can affect advection integration.Source: §A.4 · p. 9
Multi-step in vSVF registration
- Registration results improve over vSVF steps, while transformation maps remain largely well regularized despite possible folds in very large-deformation cases.Source: §A.4 · p. 9
More AVSM examples
- Additional AVSM examples show smooth transformation maps for both large and small deformations.Source: §A.4 · pp. 9–10