PaperlayerView exact revision on arXiv

Source-linked AI summary

TransMorph: Transformer for unsupervised medical image registration

Junyu Chen, Eric C. Frey, Yufan He, William P. Segars, Ye Li, Yong Du

arXiv:2111.10480v6eess.IVcs.AIcs.CV

TL;DR

ConvNet-based registration is efficient but limited in modeling long-range spatial relationships needed to establish correspondence between moving and fixed images. TransMorph combines a Transformer with a ConvNet, adds diffeomorphic and Bayesian variants, and achieves superior registration accuracy across evaluated medical-imaging tasks.

  • Problem

    ConvNets may inadequately model long-range spatial relationships, while traditional registration methods are computationally expensive because they optimize separately for each unseen image pair.

  • Method

    TransMorph is a hybrid Transformer-ConvNet model for volumetric unsupervised deformable registration, with diffeomorphic and Bayesian variants.

  • Results

    TransMorph achieved superior registration accuracy to various traditional and learning-based methods on inter-patient brain MR and phantom-to-CT registration tasks.

  • Takeaways & Limitations

    Transformer-based modeling effectively supports medical image registration by capturing long-range spatial correspondence, while the variants extend the framework to topology-preserving deformation and uncertainty estimation.

  • Takeaways & Limitations

    Baseline hyperparameters were not extensively grid-searched because of training time and limited GPU memory, and the study emphasized architectural comparison over optimizing losses or complex training methods.

Abstract

from arXiv · show

In the last decade, convolutional neural networks (ConvNets) have been a major focus of research in medical image analysis. However, the performances of ConvNets may be limited by a lack of explicit consideration of the long-range spatial relationships in an image. Recently Vision Transformer architectures have been proposed to address the shortcomings of ConvNets and have produced state-of-the-art performances in many medical imaging applications. Transformers may be a strong candidate for image registration because their substantially larger receptive field enables a more precise comprehension of the spatial correspondence between moving and fixed images. Here, we present TransMorph, a hybrid Transformer-ConvNet model for volumetric medical image registration. This paper also presents diffeomorphic and Bayesian variants of TransMorph: the diffeomorphic variants ensure the topology-preserving deformations, and the Bayesian variant produces a well-calibrated registration uncertainty estimate. We extensively validated the proposed models using 3D medical images from three applications: inter-patient and atlas-to-patient brain MRI registration and phantom-to-CT registration. The proposed models are evaluated in comparison to a variety of existing registration methods and Transformer architectures. Qualitative and quantitative results demonstrate that the proposed Transformer-based model leads to a substantial performance improvement over the baseline methods, confirming the effectiveness of Transformers for medical image registration.

1. Introduction

Medical image registration needs efficient methods that capture long-range spatial correspondence between moving and fixed images. TransMorph addresses this gap with a hybrid Transformer-ConvNet framework and variants for diffeomorphic registration and uncertainty estimation.

  • Traditional registration methods estimate smooth mappings by solving an optimization problem separately for each unseen image pair, making them computationally expensive and slow.
  • ConvNets accelerate registration after training but have limited effective receptive fields for explicitly modeling long-range spatial relations.
  • Transformers use self-attention and larger effective receptive fields to better capture spatial correspondence between distant regions of moving and fixed images.
  • TransMorph is a hybrid Transformer-ConvNet model developed for volumetric affine and deformable image registration.
  • The proposed variants provide topology-preserving diffeomorphic registration, Bayesian registration uncertainty estimates, and publicly available code, pretrained models, and preprocessed data.
  • Experiments cover inter-patient and atlas-to-patient brain MRI registration and phantom-to-CT registration, with results demonstrating improved performance over baseline methods.

2. Related Work

Related work frames registration as an optimization problem, then surveys deep-learning, diffeomorphic, and Transformer approaches. These foundations motivate TransMorph’s use of self-attention, hybrid decoding, diffeomorphic integration, and Bayesian extensions.

  • Image registration: Deformable registration minimizes an energy combining image similarity and deformation regularization to estimate a smooth spatial mapping.The energy uses a similarity term for alignment and a regularization term for deformation-field smoothness.
  • Deep neural networks: Deep neural networks learn a global registration representation from training data and apply it to unseen image pairs, avoiding per-pair optimization.
  • Diffeomorphic registration: Diffeomorphic registration uses smooth, invertible mappings with topology preservation; TransMorph supports this through velocity-field integration and scaling-and-squaring.
  • Transformer mechanisms: Self-attention forms Query, Key, and Value representations, computes token similarities, and applies normalized attention scores to Values; multi-head attention processes several such operations in parallel.
  • Extensions: The framework can be adapted to existing registration techniques for diffeomorphism and Bayesian deep learning for registration-uncertainty estimation.

3. Methods

TransMorph uses Transformer-based networks for affine and deformable registration, combining hierarchical Swin features with convolutional decoding. The framework also supports diffeomorphic topology-preserving deformation and Bayesian uncertainty calibration.

  • Registration pipeline: The conventional pipeline first affinely aligns moving and fixed images, then warps the aligned image with a deformation field.The affine stage estimates transformation parameters, while deformable registration generates the field φ.
  • Affine transformation network: The affine Transformer takes two 3D volumes as input and generates 12 affine parameters for global alignment.The parameters include rotation angles and translations.
  • Transformer encoder: TransMorph splits moving and fixed volumes into non-overlapping 3D patches, projects them into tokens, and processes them through Swin Transformer stages.Patch merging reduces token resolution while increasing feature dimensionality across encoder stages.
  • ConvNet decoder: The decoder upsamples and convolves encoder features, using skip connections and additional convolutional layers to preserve localization and recover high-resolution information.The convolutional layers use the original and downsampled image pair to capture local information.
  • 3D Swin Transformer block: Swin self-attention operates within local 3D windows, while shifted windows connect neighboring regions across successive blocks.Hierarchical feature maps from multiple resolutions are passed to a ConvNet decoder.
  • Model variants: Diffeomorphic variants use scaling and squaring with a stationary velocity field to produce continuous, differentiable, topology-preserving deformations.The Bayesian variant inserts dropout in the Transformer encoder and calibrates appearance uncertainty against the known target image.

4. Experiments

The experiments validate TransMorph across three medical image registration applications using diverse datasets and comparisons with traditional, deep-learning, and Transformer-based methods. Additional studies examine implementation settings, model variants, and Transformer-module effects.

  • Datasets: Three registration applications and datasets, including more than 1000 image pairs, were used to validate the proposed method.The applications include inter-patient and atlas-to-patient brain MRI registration and phantom-to-CT registration.
  • Datasets: The brain MRI experiments used inter-patient, atlas-to-patient, and OASIS challenge data with anatomical label maps for evaluation.The datasets provided 30 or 35 anatomical structures for quantitative assessment where specified.
  • Comparisons: TransMorph was compared with four non-deep-learning registration methods using empirically selected settings balancing registration accuracy and running time.SyN4 and NiftyReg were among the traditional methods described.
  • Comparisons: Deep-learning comparisons used task-specific losses, including MSE with diffusion regularization for inter-patient MRI and LNCC for atlas-to-patient MRI.The XCAT-to-CT task used LNCC, bending energy, and Dice loss.
  • Comparisons: TransMorph was evaluated against customized Transformer-based registration networks, including ViT-V-Net, PVT, CoTr, and nnFormer.These networks were modified to produce three-dimensional deformation fields for warping moving images.
  • Implementation: All models were trained for 500 epochs with Adam, a learning rate of 1 × 10^-4, and batch size 1, while augmentation differed between MRI and CT data.Random-direction flipping was applied to brain MRI data, whereas no augmentation was applied to the CT dataset.
  • Ablation and variants: Experiments also evaluated TransMorph variants, a compact affine model, and the effects of Transformer modules and model complexity.The affine model contained 19.55 million parameters and had computational complexity of 0.4 GMacs.

5. Results

Across brain MRI, challenge, and phantom-to-CT registrations, TransMorph variants generally achieved accurate registration, while diffeomorphic variants produced smoother, topology-preserving deformations. Ablations further showed benefits from skip connections and limited influence of positional embeddings, while larger models improved performance at higher computational cost.

  • Inter-patient brain MRI registration: TransMorph achieved the highest inter-patient brain MRI mean Dice score of 0.745, improving by >0.2 over VoxelMorph and CycleMorph.Diffeomorphic variants slightly reduced Dice but produced almost no foldings, and Transformer-based models generally outperformed ConvNet-based models.
  • Atlas-to-patient brain MRI registration: TransMorph-bspl achieved the highest atlas-to-patient mean Dice score of 0.761 with nearly no folded voxels.TransMorph and TransMorph-Bayes reached 0.754, while TransMorph-diff reached 0.594 but remained ∼0.02 above VoxelMorph-diff.
  • Learn2Reg OASIS brain MRI registration: TransMorph-large achieved validation mean Dice of 0.862 and mean HdDist95 of 1.431, performing similarly to the challenge’s best-performing method, LapIRN.On the test set, TransMorph and TransMorph-large achieved comparable mean Dice to LapIRN, although LapIRN produced more uniform deformation fields by SDlogJ.
  • XCAT-to-CT registration: In qualitative CT registration, TransMorph and TransMorph-Bayes were more accurate, whereas diffeomorphic variants produced smoother deformations.The passage attributes artifacts in nnFormer’s displacement field to patch operations and insufficient convolutional refinement for displacements exceeding patch size.
  • Ablation studies: Skip connections improved performance across all three tasks, while positional embeddings and shuffled token positions produced similar results.Removing Transformer-layer skip connections reduced mean Dice by 0.019 in atlas-to-patient registration and by 0.016 in XCAT-to-CT registration.
  • Computational complexity: TransMorph outperformed Transformer alternatives with more parameters and improved across model sizes, but TransMorph-large nearly doubled computational cost for under 0.01 Dice improvement in brain MRI tasks.ViT-V-Net and PVT had around 2× and 1.5× more parameters than TransMorph, yet TransMorph outperformed them on all evaluated registration tasks.

6. Discussion

TransMorph’s architectural analyses link its registration behavior to Transformer receptive fields, skip connections, implicit positional learning, convergence, and uncertainty estimation. The discussion also identifies calibration and multimodal uncertainty boundaries.

  • Architecture analysis: Swin Transformer blocks provided broader spatial information, while convolutional layers preserved higher-resolution edge and boundary details through skip connections.Feature-map comparisons showed complementary roles for the two components.
  • Positional embedding: TransMorph learned token positional information implicitly, so mean Dice remained comparable without positional embedding or after token-position shuffling.This behavior follows from spatial registration losses backpropagating through the image-domain decoder output.
  • Uncertainty limitations: Predictive-variance appearance uncertainty is miscalibrated when computed from the predictive mean, while MSE-based uncertainty can be ineffective for multimodal registration.For PET-to-CT or MRI-to-CT registration, MSE may be dominated by squared appearance bias; predicted variance may be more appropriate.
  • Uncertainty quantification: The proposed uncertainty estimate achieved perfect calibration with UCE = 0 and captured more registration failures than the comparison estimate.The proposed estimate equaled expected model error and showed stronger correlation with registration failures.
  • Uncertainty quantification: Appearance and transformation uncertainty highlighted different regions: appearance uncertainty tracked intensity mismatches, whereas transformation uncertainty tracked large deformations.The two measures therefore provide substantially different uncertainty information.
  • Loss landscapes: TransMorph produced a substantially flatter loss landscape than the compared ConvNet-based models.The discussion relates this observation to reports that Transformers tend to promote flatter loss landscapes.
  • Training and performance: TransMorph exceeded 0.7 Dice within the first 20 epochs and consistently outperformed other Transformer-based models with comparable parameters and computational complexity.On average, Transformer-based models also achieved better validation scores than ConvNet-based models, except CoTr.
  • Efficiency: TransMorph tended to converge faster than ConvNet-based models, potentially reducing training time, epochs, computing resources, and carbon emissions.The paper also reports moderate computational complexity for TransMorph and TransMorph-Bayes at 687 GMACs.

7. Conclusion

The paper introduces TransMorph for unsupervised deformable image registration, with topology-preserving and Bayesian variants. Across brain MR and phantom-to-CT registration, it reports superior accuracy to traditional and learning-based methods.

  • Conclusion: TransMorph is a Transformer-based model for unsupervised deformable image registration that establishes long-range spatial correspondence between image voxels.The model is presented as a strong candidate for medical image registration tasks.
  • Conclusion: Diffeomorphic TransMorph variants provide topology-preserved deformations, while Bayesian deep learning enables deformation uncertainty estimation without degrading registration performance.These variants extend the model with deformation-topology and uncertainty capabilities.
  • Conclusion: TransMorph achieved superior registration accuracy across inter-patient brain MR and phantom-to-CT registration than various traditional and learning-based methods.The evaluation included a novel phantom-to-CT registration task.

Declaration of Competing Interest

The authors report no competing interests and no concurrent submission elsewhere.

  • The authors declare that they have no competing interests and that the manuscript is not under review at another venue.

CRediT authorship contribution statement

The contribution statement assigns conceptualization, methodology, software, validation, data curation, investigation, writing, visualization, supervision, and funding roles across the listed authors.

  • Junyu Chen contributed conceptualization, methodology, software, data curation, investigation, original drafting, and visualization.
  • Eric C. Frey contributed validation, resources, review and editing, supervision, and funding acquisition.
  • Yufan He contributed methodology, validation, investigation, and review and editing.
  • William P. Segars contributed data curation, while Ye Li contributed validation.
  • Yong Du contributed validation, resources, data curation, review and editing, supervision, and funding acquisition.

Appendix A. Affine Network Architecture

The appendix specifies a Swin-Transformer-based affine network for rigid registration and discusses how predictive variance relates to expected model error and bias.

  • Affine network architecture: The affine network outputs three rotation, three translation, three scaling, and three shearing parameters for rigid registration.
  • Affine network architecture: The network uses embedding dimension C = 12.
  • Predictive variance and error: The expected model error is characterized by mean squared error (MSE).
  • Predictive variance and error: The predictive mean is defined from deformed moving images, with I_d = I_m ◦ φ.
  • Predictive variance and error: Bias may be more noticeable on test images than training data because of supervised-algorithm overfitting.
  • Predictive variance and error: Predictive variance is systematically smaller than expected error, resulting in miscalibrated uncertainty estimates.

Appendix C. Additional Results for Inter-patient Brain MRI Registration

The appendix provides qualitative and quantitative comparisons for inter-patient brain MR registration, including deformation visualizations and Dice-score comparisons across methods and Transformer models.

  • Qualitative comparison: Qualitative comparisons show deformed moving images, deformation fields, and deformed grids for different registration methods.The x, y, and z displacement-field dimensions map to RGB channels, and color bars report magnitude ranges [p, q].
  • Quantitative comparison: Dice-score boxplots compare TransMorph with existing registration methods across different brain MR substructures.
  • Transformer comparison: Dice-score boxplots compare TransMorph, its variants, and other Transformer architectures across different brain MR substructures.

Appendix D. Additional Results for Atlas-to-patient Brain MRI Registration

The appendix compares registration methods for atlas-to-patient brain MR images using qualitative deformation visualizations and Dice-score boxplots across conventional and Transformer-based models.

  • Qualitative comparison: Qualitative comparisons show deformed moving images, deformation fields, and deformed grids for atlas-to-patient brain MR registration.The x, y, and z displacement-field dimensions map to RGB channels, and color bars report magnitude ranges [p, q].
  • Quantitative comparison: Dice-score boxplots compare TransMorph with existing registration methods across different brain MR substructures.
  • Transformer comparison: Dice-score boxplots compare TransMorph, its variants, and other Transformer architectures across different brain MR substructures.

Appendix E. Additional Results for XCAT-to-CT Registration

The appendix presents qualitative and quantitative comparisons for XCAT-to-CT registration, including deformation visualizations, additional phantom slices, and Dice scores across methods and Transformer models.

  • Qualitative comparison: Qualitative comparisons show deformed moving images, deformation fields, and deformed grids for XCAT-to-CT registration methods.The x, y, and z displacement-field dimensions map to RGB channels, and color bars report magnitude ranges [p, q].
  • Qualitative comparison: Additional coronal slices show deformed XCAT phantoms generated by various registration methods.
  • Quantitative comparison: Dice-score boxplots compare TransMorph with existing registration methods across different CT organs.
  • Transformer comparison: Dice-score boxplots compare TransMorph, its variants, and other Transformer architectures across different CT organs.

Appendix F. Additional Qualitative Results for Uncertainty Quantification

The appendix presents qualitative registration results alongside appearance and transformation uncertainty estimates from TransMorph-Bayes.

  • The fourth and fifth columns show appearance uncertainties estimated by the proposed uncertainty estimation scheme.
  • The last column shows transformation uncertainty for the deformation field.
  • Transformation uncertainty is computed as the square root of summed deformation variances across x, y, and z.
  • The displacement field’s x, y, and z components map to the RGB color channels, with color bars indicating magnitude ranges.

Appendix H. Probabilistic diffeomorphic registration

The probabilistic diffeomorphic framework models deformation uncertainty with variational inference, while alternative regularization and control-point formulations support efficient topology-preserving registration.

  • TransMorph-diff applies variational inference to estimate a posterior distribution over deformation fields.
  • The deformation prior uses a Gaussian precision matrix Λu = λL, where L is the voxel-grid neighborhood Laplacian.
  • The likelihood p(If|u, Im) is assumed Gaussian, and φu is computed as the exponential of a stationary velocity field using scaling and squaring.
  • The variational posterior is Gaussian with voxel-wise mean and variance generated by the network, and displacement fields are sampled by reparameterization.
  • Training minimizes a loss combining image similarity, optional segmentation similarity, and regularization of the mean displacement field.
  • TransMorph-bspl predicts a lattice of low-dimensional B-spline control-point displacements instead of a full-resolution dense field, then obtains the final displacement with scaling and squaring.
Loading 2111.10480v6…