Source-linked AI summary
ViT-V-Net: Vision Transformer for Unsupervised Volumetric Medical Image Registration
Junyu Chen, Yufan He, Eric C. Frey, Ye Li, Yong Du
TL;DR
ConvNets are limited in explicitly modeling long-range spatial relations, while ViT lacks detailed localization information for registration. ViT-V-Net combines convolutional encoding, Transformer processing, and V-Net-style decoding for volumetric registration, evaluated on brain MRI volumes using Dice-based comparisons.
Problem
ConvNets have limited explicit long-range spatial modeling, whereas ViT’s low-resolution emphasis reduces detailed localization information needed for image registration.
Method
ViT-V-Net encodes fixed and moving volumes with convolutional layers, processes high-level 3D patches using a Vision Transformer, and decodes a dense displacement field.
Results
The method was evaluated on brain MRI registration using 260 T1-weighted scans and Dice-score comparisons with other registration methods.
Takeaways & Limitations
The study investigates a hybrid ConvNet-Transformer architecture for volumetric medical image registration.
Abstract
from arXiv · showhide
In the last decade, convolutional neural networks (ConvNets) have dominated and achieved state-of-the-art performances in a variety of medical imaging applications. However, the performances of ConvNets are still limited by lacking the understanding of long-range spatial relations in an image. The recently proposed Vision Transformer (ViT) for image classification uses a purely self-attention-based model that learns long-range spatial relations to focus on the relevant parts of an image. Nevertheless, ViT emphasizes the low-resolution features because of the consecutive downsamplings, result in a lack of detailed localization information, making it unsuitable for image registration. Recently, several ViT-based image segmentation methods have been combined with ConvNets to improve the recovery of detailed localization information. Inspired by them, we present ViT-V-Net, which bridges ViT and ConvNet to provide volumetric medical image registration. The experimental results presented here demonstrate that the proposed architecture achieves superior performance to several top-performing registration methods.
1. Introduction
Deformable registration needs efficient spatial correspondence modeling, but ConvNets have limited explicit long-range modeling. ViT-V-Net introduces a hybrid ConvNet-Transformer approach for volumetric medical image registration.
- Deformable image registration establishes spatial correspondences between fixed and moving images through a spatially varying deformation model.Traditional methods optimize image similarity while enforcing deformation smoothness separately for each image pair.
- ConvNet registration methods remain limited in modeling explicit long-range relations because convolution operations are intrinsically local.
- Vision Transformer uses pure self-attention to model long-range spatial relations and achieved state-of-the-art image-recognition performance.
- ViT-V-Net is proposed as a hybrid ConvNet-Transformer architecture for self-supervised volumetric medical image registration.The work investigates ViT usage for volumetric registration because 2D images do not fully exploit spatial correspondences in 3D volumes.
2. Methods
ViT-V-Net encodes paired volumes with convolutional layers, models high-level patches using a Vision Transformer, and decodes them into a dense deformation field. Training combines image similarity with deformation smoothness, and evaluation uses brain MRI data.
- The network encodes fixed and moving volumes into high-level features before dividing them into vectorized 3D patches for ViT processing.This avoids applying ViT directly to full-resolution volumetric images, which would have large computational complexity.
- The Transformer output is reshaped and decoded with a V-Net-style decoder using long skip connections.
- The final network output is a dense displacement field used by a spatial transformer to warp the moving image.
- Training uses mean squared error for image similarity together with a diffusion regularizer weighted by λ to impose displacement-field smoothness.
- The brain MRI experiment uses 260 T1-weighted scans split into 182 training, 26 validation, and 52 test volumes.Volumes were cropped to 160 × 192 × 224, with 29 anatomical structures used for evaluation and Dice score comparisons against registration methods.
Appendix A. Overview of Vision Transformer
The Vision Transformer embeds volumetric patches with positional information before processing them through repeated Transformer encoder blocks. The encoded representation is then available for downstream decoding.
- Volumetric features are represented as vectorized patches and mapped into a latent D-dimensional space through a trainable linear projection.
- Learnable positional embeddings are added to patch embeddings to retain the patches’ positional information.
- The Transformer encoder consists of 12 blocks alternating Multihead Self-Attention and Multi-Layer Perceptron layers.Layer normalization is applied before each MSA and MLP layer.
- The encoder produces an encoded image representation denoted by zℓ.
Appendix B. Loss Functions
The training loss combines image similarity with deformation-field smoothness, using MSE for alignment and a diffusion regularizer for smooth displacements.
- Loss Function: The loss function combines mean squared error between deformed moving and fixed images with a diffusion regularizer weighted by λ.The regularization term enforces smoothness in the deformation field.
- Image Similarity Measurement: MSE measures image similarity between the deformed moving image and the fixed image over the image domain Ω.The passage identifies Ω as the image domain.
- Deformation Field Regularization: The diffusion regularizer is used to enforce smoothness in the deformation field.The displacement field u is the network output.
Appendix C. Hyperparameters Settings
The appendix describes shared training settings for learning-based models and a separately selected NiftyReg regularization parameter motivated by over-smoothing with the default.
- Learning-Based Models: Learning-based models used the ADAM optimizer and shared training hyperparameters, except ViT-V-Net used 0.1 dropout in its ViT linear layers.Each model took about 3 days to train on a single GPU.
- NiftyReg and SyN: NiftyReg used a 0.0002 regularization parameter because its default 0.005 produced over-smoothed, suboptimal deformations.The setting is reported alongside SyN hyperparameters.
Appendix D. Additional Results
Additional results present training curves, anatomical-structure Dice distributions, qualitative deformation visualizations, and quantitative comparisons across registration methods.
- Quantitative Comparisons: Quantitative comparisons evaluate Dice score, folded voxels, and computational time across registration methods.NiftyReg and SyN ran on CPUs, whereas VoxelMorph and ViT-V-Net ran on GPUs.
- Statistical Comparisons: Paired t-tests were used to compute p-values for Dice-score differences between ViT-V-Net and other registration methods.The comparison concerns Dice scores between the proposed method and competing methods.
- Training Curves: Training curves show lower loss values and higher validation Dice scores for ViT-V-Net during training.The figure reports these trends per epoch.
- Anatomical Results: Boxplots compare Dice scores across anatomical structures and registration methods, with left and right hemispheres averaged.Orange triangles denote the means.
- Qualitative Results: Qualitative results display moving and fixed images, deformed images, and displacement fields from different registration methods.Displacement values are clamped to [−10, 10] before mapping spatial dimensions to RGB channels.