Source-linked AI summary
Unsupervised 3D End-to-End Medical Image Registration with Volume Tweening Network
Shengyu Zhao, Tingfung Lau, Ji Luo, Eric I-Chao Chang, Yan Xu
TL;DR
Dense deformation fields are difficult to annotate, motivating an unsupervised approach to clinically important 3D medical image registration. VTN uses end-to-end CNNs with cascaded subnetworks, integrated affine registration, and invertibility loss, achieving state-of-the-art performance and an 880x speed-up over traditional optimization-based methods.
Problem
Dense medical-image deformation fields are difficult to label accurately, while traditional optimization-based registration is too time-consuming for clinical applications.
Method
VTN is an unsupervised end-to-end CNN framework that uses warping, cascaded registration subnetworks, integrated affine registration, and invertibility loss for 3D medical image registration.
Results
VTN achieves state-of-the-art performance and is 880x faster than traditional optimization-based registration, or 3.3x faster without GPU acceleration.
Takeaways & Limitations
Cascading addresses large displacement, integrated affine registration avoids separate affine alignment, and invertibility loss encourages backward consistency.
Abstract
from arXiv · showhide
3D medical image registration is of great clinical importance. However, supervised learning methods require a large amount of accurately annotated corresponding control points (or morphing), which are very difficult to obtain. Unsupervised learning methods ease the burden of manual annotation by exploiting unlabeled data without supervision. In this paper, we propose a new unsupervised learning method using convolutional neural networks under an end-to-end framework, Volume Tweening Network (VTN), for 3D medical image registration. We propose three innovative technical components: (1) An end-to-end cascading scheme that resolves large displacement; (2) An efficient integration of affine registration network; and (3) An additional invertibility loss that encourages backward consistency. Experiments demonstrate that our algorithm is 880x faster (or 3.3x faster without GPU acceleration) than traditional optimization-based methods and achieves state-of-theart performance in medical image registration.
I. INTRODUCTION
The paper motivates VTN as an unsupervised, end-to-end CNN framework for voxel-level 3D medical image registration, avoiding costly dense ground-truth deformation labels. Its cascading, affine-integration, and invertibility components target large displacements, efficient alignment, and backward consistency.
- 3D medical image registration finds nonlinear spatial correspondence between images for applications including longitudinal alignment and atlas matching.
- Traditional optimization-based registration can perform well but is too slow for many clinical applications because it iteratively updates transformation parameters.
- Dense deformation fields are nearly impossible to label manually, while synthetic datasets may not reflect realistic medical-image demands, limiting supervised methods.
- VTN trains end-to-end CNNs without deformation supervision by warping the moving image and comparing it with the fixed image through similarity and regularization losses.
- Cascaded subnetworks repeatedly warp intermediate moving images, enabling later subnetworks to register progressively transformed images and improving performance for large displacements.
- The framework integrates affine registration and adds an invertibility loss that encourages backward consistency while improving registration performance.
A. Traditional Algorithms
Traditional registration optimizes transformation parameters iteratively, which makes it time-consuming, whereas learning-based methods predict registrations more efficiently. Prior supervised and unsupervised approaches leave gaps that VTN addresses through end-to-end 3D registration and integrated affine handling.
- A. Traditional Algorithms: Traditional algorithms define transformation spaces and alignment metrics, then iteratively optimize parameters, making the process impractical for clinical applications.
- B. Supervised Learning Methods: Earlier CNN registration methods used synthetic data, pre-aligned pairs, sparse parameters, or separately trained stages rather than a fully end-to-end 3D framework.
- B. Supervised Learning Methods: Supervised methods depend on abundant ground-truth alignments or carefully designed synthetic data that resemble real images.
- C. Unsupervised Learning Methods: VoxelMorph predicts dense deformation fields without supervision but assumes affine pre-alignment and performs poorly when image displacement is large.
- C. Unsupervised Learning Methods: VTN uses warping to support unsupervised training, network cascading, and invertibility loss, producing a design intended for registration with or without large displacement.
A. Problem Formulation
The problem is posed as finding a displacement field that maps a moving 3D image toward a fixed image. Warping applies that field, and composing successive flows describes the cascaded transformation.
- A. Problem Formulation: The inputs are single-channel grayscale images defined on a 3D cuboid domain, and registration seeks a displacement field from one image to the other.
- A. Problem Formulation: The flow from I1 to I2 specifies where each voxel in I1 occurs in I2, while warping evaluates the moving image at displaced coordinates.
- A. Problem Formulation: Registration can be expressed as finding a flow that maximizes similarity between the fixed image and the warped moving image.
- A. Problem Formulation: The first image is the fixed image, the second is the moving image, and the moving image is transformed during registration.
- A. Problem Formulation: Applying two warps sequentially yields a composed flow, motivating flow composition for combining transformations across registration stages.
- A. Problem Formulation: In practice, images and flow fields are extended from lattice points by trilinear interpolation, with nearest-point interpolation for out-of-bound indices.
B. Unsupervised End-to-End Registration Network
VTN performs unsupervised end-to-end 3D registration by cascading differentiable warping stages and integrating affine alignment. Its losses combine image similarity with flow smoothness and affine-transform regularization.
- Cascaded registration: VTN cascades registration subnetworks, warping the moving image after each stage so gradients reach all preceding subnetworks.Differentiable trilinear interpolation enables back-propagation through the cascaded warps.
- Affine integration: The integrated affine subnetwork predicts affine parameters before dense-flow warping, replacing a separate preprocessing tool.The affine stage provides global alignment within the registration network.
- Training objective: Training minimizes image dissimilarity between warped moving images and the fixed image, together with regularization losses on predicted flows.The framework uses correlation coefficient similarity and total variation regularization for dense flow predictions.
- Similarity loss: Correlation coefficient measures linear image relatedness and is more robust than L2 loss because non-degenerate linear intensity changes do not alter it.Its range is [-1, 1], reaching ±1 when the images are linear functions of one another.
- Flow regularization: Total variation regularization discourages discontinuities in dense flow fields.The formulation differs from the initial total-variation definition by omitting the square root.
- Affine regularization: Orthogonality loss penalizes affine transforms whose singular values deviate from 1, while determinant loss excludes reflections under the same-chirality assumption.For an orthogonal I + A, the orthogonality loss is zero; determinant regularization requires det(I + A) > 0.
A. Cascading
VTN cascades affine and dense deformable subnetworks so each stage refines the current alignment, while invertibility loss encourages round-trip consistency.
- Cascading: Each subnetwork aligns the fixed image with the current moving image, warps the moving image, and passes the result to the next stage.Predicted flow fields are composed into the final registration estimate, with gradients back-propagated through the differentiable cascade.
- Cascading: The cascade differs from FlowNet 2.0 by sequentially feeding warped images and fixed images rather than combining two separate flow-estimation lines.FlowNet 2.0 also supplies intermediate subnetworks with the initial image, current flow, and brightness error.
- Affine registration: The affine subnetwork is used first to produce a global affine alignment before subsequent deformable registration.It outputs 12 parameters representing a 3 × 3 matrix A and a 3-dimensional displacement vector b.
- Dense deformable registration: Subsequent dense deformable subnetworks refine registration through an encoder-decoder architecture with skip connections.Strided 3D convolutions reduce resolution, while transposed convolutions recover it for dense prediction.
- Invertibility: Invertibility loss penalizes nonzero composed forward-and-backward flows, encouraging registrations that return voxels to their starting positions.Perfect round-trip registration has zero composed displacement and zero invertibility loss.
V. EXPERIMENT
The experiments evaluate VTN on liver CT and brain MRI pairwise registration against traditional and learning-based methods, using accuracy, landmark distance, and runtime metrics.
- Experimental setting: VTN is evaluated separately on liver CT and brain MRI datasets using pairwise subject-to-subject registration.This setting is described as more general than atlas-based registration, where moving images are aligned to a fixed atlas.
- Baselines: The comparison includes ANTs, Elastix, and VoxelMorph alongside the proposed algorithm.ANTs and Elastix are traditional registration systems, while VoxelMorph is a learning-based comparator.
- Results: VTN achieves state-of-the-art performance while being much faster than the compared algorithms.The experiments also report that more unlabeled training data improve the unsupervised method’s performance.
- Metrics: Evaluation metrics include segmentation IoU, landmark distance, and average registration time.Runtime is reported separately with and without GPU acceleration because some methods use GPUs.
- Implementation: Training uses TensorFlow with Adam, batches of 8 image pairs, an initial learning rate of 10^-4, and five epochs.Performance evaluation uses an NVIDIA TITAN Xp GPU, while traditional methods run on CPU; neural methods are also tested without GPU acceleration.
A. Experiments on Liver Datasets
The liver experiments use cropped 128^3 CT volumes, cascaded affine and deformable subnetworks, and separate training and test datasets with landmark and segmentation annotations.
- Model configuration: Liver CT inputs are 128^3 volumes; affine subnetworks downsample to 4^3 and dense deformable subnetworks to 2^3 before upsampling.
- Cascading: The model cascades up to four registration subnetworks to progressively align liver scans with large displacement.The ADDD configuration contains one affine and three dense deformable registration subnetworks.
- Datasets: Three liver datasets are used: LITS and BFH for training, and MICCAI’07 for testing.LITS contains 130 volumes, BFH contains 92, and MICCAI’07 contains 20 volumes.
- Loss functions: The experiment includes a table reporting ratios of loss functions.
- Evaluation data: MICCAI’07 provides liver segmentations and four anatomical landmarks annotated by three expert doctors for evaluation.The averaged expert annotations serve as landmark ground truth.
3) Comparison among Methods:
On the liver dataset, VTN is compared with traditional registration and VoxelMorph using segmentation, landmark, speed, and deformation-quality measures. The reported results show substantial speed advantages and significant accuracy gains, while deformation smoothness remains a trade-off.
- 3) Comparison among Methods:: ADDD and ADDD + inv are VTN variants evaluated against traditional methods and VoxelMorph on the liver dataset.ADDD + inv adds an invertibility-loss term with relative weight 10^-3 in the central area.
- 3) Comparison among Methods:: The Wilcoxon signed-rank test indicates that VTN methods significantly surpass state-of-the-art registration algorithms in Segmentation IoU and Landmark Distance.
- 3) Comparison among Methods:: Learning-based methods provide a vast speed-up over optimization-based methods.
- 3) Comparison among Methods:: VTN's better registration performance comes at the price of worse deformation smoothness, but its fraction of folding remains below 1%.The invertibility loss considerably reduces the fraction of folding.
- 3) Comparison among Methods:: Network cascading better aligns images with large displacement, while invertibility loss has a remarkable effect on the liver boundary.
4) Performance with Different Amount of Data:
Experiments examine how unlabeled-data quantity, preprocessing, affine integration, and cascade depth affect VTN across liver and brain registration settings. More unlabeled data and deeper cascades improve performance, while affine integration improves landmark accuracy without separate affine preprocessing.
- 4) Performance with Different Amount of Data:: Training ADDD with more unlabeled data improves liver-registration performance.The comparison uses full LITS + BFS data versus only part of the data.
- 4) Performance with Different Amount of Data:: Methods using ANTs affine alignment and methods with integrated affine registration are not directly equivalent because they transform the fixed image differently.The authors warn that non-orthogonal ANTs transforms can advantage either comparison group depending on whether the transform shrinks.
- 4) Performance with Different Amount of Data:: ADD achieves the lowest Landmark Distance with competitive speed on the brain datasets.Compared with DD, integrated affine registration significantly improves Landmark Distance and produces significantly less folding, although ANTs remains better on Segmentation IoU.
- 4) Performance with Different Amount of Data:: With more unlabeled data, DD's Landmark Distance consistently improves.
- 4) Performance with Different Amount of Data:: As the model uses more cascaded subnetworks, performance in each group improves.
- 4) Performance with Different Amount of Data:: Integrating affine registration yields better Landmark Distance, whereas ANTs affine alignment yields higher Segmentation IoU.The higher IoU comes with a 74-second preprocessing phase, while integrated affine registration avoids that separate preprocessing stage.
VI. DISCUSSION
The discussion reports strong gains over traditional methods on liver CT but a remaining advantage for ANTs on brain MRI. It attributes this difference partly to displacement characteristics and notes that alternative similarity measures could be explored.
- VI. DISCUSSION: VTN significantly outperforms traditional methods on liver CT, while ANTs performs slightly better on brain MRI.
- VI. DISCUSSION: The authors suggest that brain MRI's mainly small displacements may better suit traditional iterative methods.
- VI. DISCUSSION: The paper identifies alternative similarity measurements, such as cross correlation, as possible avenues for improving accuracy.
VII. CONCLUSION
The paper presents VTN as an unsupervised, end-to-end CNN framework for 3D medical image registration. Cascading, integrated affine registration, and invertibility loss address large displacement, end-to-end alignment, and performance, respectively.
- VII. CONCLUSION: VTN is an unsupervised end-to-end convolutional-neural-network framework for 3D medical image registration.
- VII. CONCLUSION: 880x faster, or 3.3x faster without GPU acceleration, VTN achieves state-of-the-art performance compared with traditional medical image registration methods.
- VII. CONCLUSION: Cascaded deformable subnetworks tackle registration with large displacement.
- VII. CONCLUSION: Integrated affine registration makes the method truly end-to-end and is more effective than out-of-band affine alignment.
- VII. CONCLUSION: Adding invertibility loss further enhances registration performance.
- VII. CONCLUSION: The authors state that VTN methods can potentially be applied to various other medical image registration tasks.