Registration by Regression (RbR): a framework for interpretable and flexible atlas registration

TL;DR

Registration by Regression (RbR) reframes atlas registration as voxel-wise coordinate regression, predicting atlas coordinates for every voxel and then fitting a deformation model in closed form at test time. This enables flexible, test-time choice of affine or nonlinear transforms (e.g., Bspline, Demons, log-polyaffine) with robustness from leveraging millions of voxel constraints and optional RANSAC. Trained via voxel-wise $\ell_1$ loss against ground-truth coordinates obtained from classical registrations to a chosen atlas, RbR combines a high-resolution voxel basis with a standard U-Net architecture to produce atlas-relative coordinates and a brain mask. Empirically, RbR outperforms competing keypoint methods on public datasets in Dice scores for nonlinear registration, while maintaining interpretability and enabling easy pretraining for other tasks; limitations include the need to retrain for each atlas and potential gaps to non-interpretable methods in some nonlinear settings.

Abstract

In human neuroimaging studies, atlas registration enables mapping MRI scans to a common coordinate frame, which is necessary to aggregate data from multiple subjects. Machine learning registration methods have achieved excellent speed and accuracy but lack interpretability and flexibility at test time (since their deformation model is fixed). More recently, keypoint-based methods have been proposed to tackle these issues, but their accuracy is still subpar, particularly when fitting nonlinear transforms. Here we propose Registration by Regression (RbR), a novel atlas registration framework that: is highly robust and flexible; can be trained with cheaply obtained data; and operates on a single channel, such that it can also be used as pretraining for other tasks. RbR predicts the (x, y, z) atlas coordinates for every voxel of the input scan (i.e., every voxel is a keypoint), and then uses closed-form expressions to quickly fit transforms using a wide array of possible deformation models, including affine and nonlinear (e.g., Bspline, Demons, invertible diffeomorphic models, etc.). Robustness is provided by the large number of voxels informing the registration and can be further increased by robust estimators like RANSAC. Experiments on independent public datasets show that RbR yields more accurate registration than competing keypoint approaches, over a wide range of deformation models.

Figures (3)

• Figure 1: (a) Training data are prepared by nonlinearly registering T1w scans from HCP and ADNI to the MNI template using NiftyReg. (b) A U-Net is trained to predict the MNI coordinates for every voxel of an input scan. (c) At test time, a wide array of transformations can be fitted between the input scan and predicted MNI coordinates.
• Figure 2: Coronal slice of sample fixed image and corresponding registered MNI slice, using the different approaches.
• Figure 3: (a) Box plot of average Dice between segmentations of scans and registered MNI. For RbR, the numbers indicate control point spacing (Bsplines), Gaussian $\sigma$ (demons), or supervoxel width (log-polyaffine). (b) Significance matrix for two-tailed t-tests comparing the Dice scores; pink represents p$<$0.05. (c) Membrane energy for nonlinear models. In the box plots, the center line is the median, the ends of the box are the 25$^{th}$ and 75$^{th}$ percentiles, and the whiskers extend to the furthest observations not considered outliers -- which are marked with circles. Outliers are values more than 1.5 times the interquartile range away from the ends of the box.