[Paper Review] AIRNet: Self-Supervised Affine Registration for 3D Medical Images using Neural Networks

Evelyn Chee, Zhenhua Wu|arXiv (Cornell University)|Oct 5, 2018

Medical Image Segmentation Techniques19 references53 citations

TL;DR

AIRNet directly regresses 3D affine transformation parameters between two images using a self-supervised CNN, achieving faster and often more accurate registration than conventional methods, and generalizing across modality and patient differences.

ABSTRACT

In this work, we propose a self-supervised learning method for affine image registration on 3D medical images. Unlike optimisation-based methods, our affine image registration network (AIRNet) is designed to directly estimate the transformation parameters between two input images without using any metric, which represents the quality of the registration, as the optimising function. But since it is costly to manually identify the transformation parameters between any two images, we leverage the abundance of cheap unlabelled data to generate a synthetic dataset for the training of the model. Additionally, the structure of AIRNet enables us to learn the discriminative features of the images which are useful for registration purpose. Our proposed method was evaluated on magnetic resonance images of the axial view of human brain and compared with the performance of a conventional image registration method. Experiments demonstrate that our approach achieves better overall performance on registration of images from different patients and modalities with 100x speed-up in execution time.

Motivation & Objective

• Motivate fast, end-to-end affine registration without predefined similarity metrics.
• Leverage abundant unlabeled data by generating synthetic affine transformations for training.
• Learn discriminative image features tailored for registration that generalize across modalities.
• Enable one-shot registration between two images by directly estimating the transformation parameters.
• Evaluate against conventional registration to demonstrate speed and accuracy improvements.

Proposed method

• Two-branch encoder architecture with weight sharing processes I_ref and I_t, followed by concatenation and fully connected regression to output a 12-dimensional affine parameter vector.
• Encoder blocks combine DenseNet-inspired densely connected layers with a mixture of 2D (early) and 3D (later) convolutions to handle anisotropic 3D MR data.
• Training uses self-supervised synthesis: apply random affine T to a brain image to create I_t and use the known t (flattened top three rows of T) as label.
• Pixel values are normalized to [0,1] and images are resized to 20x320x320 for network input.
• Loss is the mean squared error between predicted transformation parameters and ground-truth t, optimized with Adam.
• Baseline comparison uses SimpleITK affine registration with Mattes mutual information as the similarity metric.

Experimental results

Research questions

• RQ1 Can a CNN learn to regress 12D affine transformation parameters directly from two 3D MR images?
• RQ2 Does self-supervised synthetic data enable accurate and fast registration across different patients and modalities?
• RQ3 How does AIRNet compare to conventional affine registration in terms of accuracy (Jac, d_H) and runtime across modalities and patients?
• RQ4 Do learned features generalize to unseen modalities and to cross-patient registration tasks?

Key findings

• AIRNet variants achieve higher mean Jaccard indices than the conventional method for the temporal lobe region under across-patients and across-patients-with-modality settings.
• AIRNet- S1,2,c consistently yields strong performance with Jac 0.613 and d_H 1.431 in E_p and 0.615 and 1.459 in E_p,m, while SimpleITK reaches 0.561/1.876 and 0.524/1.978 respectively.
• All AIRNet variants run ~0.75–0.79 seconds per registration, offering roughly 100x speed-up over the conventional method (SimpleITK ~166–217 seconds).
• AIRNet exhibits better registration performance across ten brain regions and demonstrates robust generalization to unseen modalities and across-patient data.
• Encoder visualizations (PCA/t-SNE) indicate the learned representations emphasize geometric transformation information rather than modality.
• The method remains effective even when trained on same-patient, same-modality pairs but tested on cross-patient/modality registrations.