[Paper Review] ViT-V-Net: Vision Transformer for Unsupervised Volumetric Medical Image Registration

Junyu Chen, Yufan He+3 more|arXiv (Cornell University)|Apr 13, 2021

Advanced Neural Network ApplicationsComputer Science20 references139 citations

TL;DR

ViT-V-Net introduces a hybrid ConvNet-Transformer architecture for unsupervised volumetric medical image registration, achieving superior Dice performance over top methods on brain MRI data.

ABSTRACT

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.

Motivation & Objective

• Motivate deformable image registration (DIR) and address limitations of convolutional networks in modeling long-range spatial relations.
• Propose a hybrid ViT-ConvNet architecture to enable long-range feature learning for 3D image registration.
• Demonstrate that ViT-V-Net improves registration accuracy (Dice) and maintains localization via long skip connections.
• Evaluate against state-of-the-art registration methods on a brain MRI dataset and provide implementation details.

Proposed method

• Encode high-level features from fixed and moving 3D images through ConvNet blocks and pooling to reduce resolution.
• Partition high-level features into N patches and apply a Vision Transformer to learn long-range relations.
• Embed patches with a linear projection and add learnable position embeddings for spatial information.
• Pass Transformer outputs through a V-Net–style decoder with long skip connections to preserve localization details.
• Predict a dense displacement field u, warp the moving image with a spatial transformer, and optimize a loss combining MSE similarity and a diffusion regularizer.

Experimental results

Research questions

• RQ1 Can a hybrid ConvNet-Transformer architecture improve unsupervised 3D image registration compared to fully ConvNet-based registries?
• RQ2 Do Vision Transformer–based encodings enhance long-range spatial relationships critical for accurate volumetric alignment?
• RQ3 Is the ViT-V-Net architecture able to achieve higher Dice scores than leading DIR methods on brain MRI data?

Key findings

• ViT-V-Net achieves higher Dice scores than several top registration methods in the tested setup.
• Reported Dice on the main comparison table: ViT-V-Net 0.726 ± 0.130 vs. others (Affine 0.569 ± 0.171, NiftyReg 0.713 ± 0.134, SyN 0.688 ± 0.140, VoxelMorph-1 0.711 ± 0.135, VoxelMorph-2 0.707 ± 0.137).
• ViT-V-Net trained with long skip connections retains localization information and shows lower training loss and higher validation Dice.
• Statistical tests (paired t-test) show ViT-V-Net significantly better than several competitors (p-values listed in the paper).
• The method runs on GPUs with reported times, highlighting feasibility of the approach for practical use.