[Paper Review] Attention for Image Registration (AiR): an unsupervised Transformer approach

Zihao Wang, Hervé Delingette|arXiv (Cornell University)|May 5, 2021

Advanced Neural Network Applications17 references4 citations

TL;DR

This paper proposes AiR, an unsupervised Transformer-based framework for deformable image registration that treats image registration as a sequence-to-sequence translation task. By leveraging multi-scale self-attention mechanisms without convolutional backbones, AiR achieves state-of-the-art performance on benchmark datasets, outperforming both single-scale Transformers and CNN-based methods in deformation accuracy and structural similarity, particularly on large deformations.

ABSTRACT

Image registration is a crucial task in signal processing, but it often encounters issues with stability and efficiency. Non-learning registration approaches rely on optimizing similarity metrics between fixed and moving images, which can be expensive in terms of time and space complexity. This problem can be exacerbated when the images are large or there are significant deformations between them. Recently, deep learning, specifically convolutional neural network (CNN)-based methods, have been explored as an effective solution to the weaknesses of non-learning approaches. To further advance learning approaches in image registration, we introduce an attention mechanism in the deformable image registration problem. Our proposed approach is based on a Transformer framework called AiR, which can be efficiently trained on GPGPU devices. We treat the image registration problem as a language translation task and use the Transformer to learn the deformation field. The method learns an unsupervised generated deformation map and is tested on two benchmark datasets. In summary, our approach shows promising effectiveness in addressing stability and efficiency issues in image registration tasks. The source code of AiR is available on Github.

Motivation & Objective

• To address the instability and high computational cost of traditional non-learning image registration methods based on similarity metric optimization.
• To overcome the limitations of CNN-based deep learning methods in capturing long-range dependencies and global structural deformations.
• To explore the potential of the Transformer architecture—previously dominant in NLP—for end-to-end unsupervised image registration without ground-truth deformation fields.
• To develop a multi-scale attention mechanism that improves cross-level feature representation in the Transformer for better modeling of complex deformations.
• To demonstrate that self-attention mechanisms can effectively learn deformation fields in a weakly supervised manner, avoiding the need for annotated training data.

Proposed method

• The framework treats image registration as a sequence-to-sequence translation task, using patches of fixed and moving images as input tokens.
• A Transformer encoder processes fixed image patches using multi-head self-attention to learn memory representations of the source image.
• A Transformer decoder uses both the encoded memory features and self-attention over moving image patches to predict deformation fields.
• The deformation field is used to warp the moving image via differentiable spatial sampling, enabling backpropagation through the resampling operation.
• A multi-scale parallel attention (MAPT) module is introduced to enhance feature representation across different receptive field sizes, improving performance on large deformations.
• The model is trained in an unsupervised manner using a reconstruction loss based on pixel-wise L2 distance between the warped moving image and the fixed image.

Experimental results

Research questions

• RQ1 Can a pure Transformer architecture effectively learn deformation fields for image registration without convolutional feature extractors?
• RQ2 How does the performance of a single-scale Transformer compare to multi-scale attention mechanisms in modeling complex deformations?
• RQ3 Can an unsupervised Transformer-based method achieve competitive registration accuracy without ground-truth deformation labels?
• RQ4 Does the self-attention mechanism in Transformers outperform CNNs in capturing global structural relationships in image registration tasks?
• RQ5 What is the impact of multi-scale feature learning on the quality of generated deformation maps and image reconstruction?

Key findings

• The multi-scale parallel attention Transformer (MAPT-AiR) achieved the best performance across all metrics on the MNIST dataset, with a mean PSNR of 17.311 ± 2.522 and a Smooth DICE score of 0.827 ± 0.034.
• Single-scale AiR outperformed the CNN-based DIRNets in PSNR (16.966 ± 2.897 vs. 12.941 ± 2.218) but underperformed compared to the multi-scale variant.
• MAPT-AiR successfully captured large deformations such as in digits '9', '7', and '8', while Single-AiR failed to align such complex shapes.
• The model achieved a mean MSE of 0.027 ± 0.015 on MNIST, indicating high reconstruction fidelity with minimal pixel-level error.
• Despite strong structural alignment, all Transformer-based methods exhibited geometric detail defects and noise in warped images, suggesting limitations in pixel-level precision.
• The results confirm that while Transformers are effective for global deformation modeling, they still lag behind CNNs in fine-grained texture and edge preservation.