Towards Saner Deep Image Registration

To increase the discriminative power of identical pairs feeding into the model, we propose our self-sanity check as

where the mapping function $g$ learned using any models is restricted to output zero displacements for identical pairs. Such identical pairs can be filtered out using similarity measurements. However, users are unlikely to perform these filters, especially when they do not know that trained models would produce terrible predictions for identical pairs. Hence, our self-sanity check is a natural remedy.

Next, we enforce the inverse consistency on different mappings for $g$ by the model. We will search for correspondence in the fixed image for every point in the moving image such that the transformations between the two images are inconsistent. For example, suppose that a point in the space of image $a$ is registered to the space of image $b$. If we register this point from image $b$ back to image $a$, the point will arrive at the same location in image $a$ (Def. 1). We first define the backward displacement map as in optical flow studies [22, 27, 24], back-projected from $g^{b\rightarrow a}$

making it convenient for calculation. We then introduce our cross-sanity check in the form of

$\forall(a,b)\in\{(m,f),(f,m)\}$. Here, we allow the estimation errors to increase linearly with the displacement magnitude with slope $\alpha$ and intercept $\beta$. Instead of imposing zero-tolerance between forward and back-projected backward displacements [16, 31], we relax the inverse consistency with error tolerance, defined by $\alpha$ and $\beta$, to allow occlusions which is more practical. I.e., This sanity check states that every point $p$ in the moving image $a$ should be able to map back from the fixed image $b$ to its original place in image $a$ with certain error tolerance. We then prove that this error tolerance is upper bounded.

Here, $0<\alpha<1$ and $\beta>0$. $N$ is a constant, representing the total pixel/voxel numbers. Theoretically, our proposed cross-sanity checks can be viewed as a relaxed version of the strict symmetric constraint, which is commonly used. We also derive the lower/upper bound for satisfying our forward/backward consistency check as shown in Eq. 5. The derivations’ details are shown in the Appx. A.2. To sum up, our cross-sanity check allows a series of solutions to the forward/backward consistency in an enclosed set with a radius of $\sqrt{\frac{\beta(2-\alpha)N}{1-\alpha}}$.

Next, we show that there exists a unique solution for our sanity-checked optimization, in terms of a single image pair. We start with writing the standard optimization with our proposed sanity checks in the norm form as

The detailed proof for Thm. 2 is shown in Appx. A.3. In summary, for a moving-fixed pair, since the relaxation by sanity checks and data term (e.g. SSD) are convex, together with the convex search space for $g$, our problem has a unique minimizer. Thus, the optimization problem is well-conditioned by properties of our regularity constraints such that $g$ will not change dramatically.

In the next section, we further prove that the distance between the minimizer of the sanity-checked constrained optimization and the optimal minimizer for a single moving-fixed image pair can be represented and controlled by our proposed sanity checks. We first rewrite Eq. 6 as

Then, with a slight abuse of notations, we define that

We can rewrite the minimization problem as

Here, $\operatorname*{Sim}({\mathbf{g}}):=\operatorname*{Sim}(m,f\circ(g^{f\rightarrow m}+p))+\operatorname*{Sim}(f,m\circ(g^{m\rightarrow f}+p))$, $||{\operatorname*{Reg}({\mathbf{g}})}||^{2}_{2}:=||{\operatorname*{Reg}(g^{m\rightarrow f})}||^{2}_{2}+||{\operatorname*{Reg}(g^{f\rightarrow m})}||^{2}_{2}$, $||{{\mathbf{g}}_{s}}||^{2}_{2}:=||{g^{m\rightarrow m}}||^{2}_{2}+||{g^{f\rightarrow f}}||^{2}_{2}$, $||{{\mathbf{g}}_{c}+\tilde{{\mathbf{g}}_{c}}}||^{2}_{2}:=||{g^{m\rightarrow f}+\tilde{g}^{f\rightarrow m}}||^{2}_{2}+||{g^{f\rightarrow m}+\tilde{g}^{m\rightarrow f}}||^{2}_{2}$, and $||{{\mathbf{g}}_{c}}||^{2}_{2}+||{\tilde{{\mathbf{g}}_{c}}}||^{2}_{2}:=||{g^{m\rightarrow f}}||^{2}_{2}+||{\tilde{g}^{m\rightarrow f}}||^{2}_{2}+||{g^{f\rightarrow m}}||^{2}_{2}+||{\tilde{g}^{f\rightarrow m}}||^{2}_{2}$. Then, the optimal minimizer ${\mathbf{g}}_{\ast}$ can be written in the matrix equation as

Here, $\operatorname*{Reg}(\cdot):=||{\operatorname*{Reg}({\mathbf{g}})}||^{2}_{2}$, $\text{CS}(\cdot):=||{{\mathbf{g}}_{c}+\tilde{{\mathbf{g}}_{c}}}||^{2}_{2}-\alpha(||{{\mathbf{g}}_{c}}||^{2}_{2}+||{\tilde{{\mathbf{g}}_{c}}}||^{2}_{2})-2\beta N$, and $\text{SS}(\cdot):=||{{\mathbf{g}}_{s}}||^{2}_{2}$, as defined previously. ${\dagger}$ here means that we can have $<0$ solutions for the CS check as in Eq. 9, which we show will not interfere with the theoretical foundations by masking in the next section. In this way, we obtain our unique minimizer as

where $\lambda:=[2\lambda_{r},2\lambda_{c},2\lambda_{s}]$ is the vector of parameters.

The proof is presented in Appx. A.4. We then expand the right-hand side of Eq. 12 as

Empirically, the first two terms contribute relatively less (10$\times$ smaller) than the cross-sanity error (See Experiments). Thus, we focus on the CS term here. Following Thm. 1, $||{\text{CS}({\mathbf{g}}_{\rm sanity})}||^{2}_{2}$ is upper bounded in the form of

Here, $0<\alpha<1$ and $\beta>0$ ensure the inequality’s direction, where the multiplication of two accounts for cross-sanity errors from two directions, i.e., from moving image to fixed image and also fixed image to moving image. We put the full derivation of CS error bound in the Appx. A.5. In a nutshell, we prove that by satisfying the cross-sanity check, we have an upper bound for the CS error, constraining the relationship between two displacements from different directions via parameters $\alpha$ and $\beta$.

The Lemma 4 follows Thm. 3 and Eq. 14 (upper bound of CS error). That said, the loyalty w.r.t the distance of the constrained minimizer is controlled by the combination of $\alpha$, $\beta$, and the weight parameter $\lambda_{c}$.

Interpretation of derived upper bound. By Lemma 4, we prove that if we average on the total number of voxels and also two directions, the similarity distance between the optimal minimizer and our sanity-checked minimizer per pixel/voxel is upper bounded by $\lambda_{c}(1-\alpha)\beta$. That being said, to satisfy our sanity checks, and thus maintain the loyalty to the optimal minimizer, $\lambda_{c}(1-\alpha)\beta$ should be small. Numerically speaking, for example, if $\alpha=0.1$ and $\beta=10$, we have $(1-\alpha)\beta=9$. This distance of $9$ is extremely large for this delicate image registration task, causing the constrained minimizer to be untrustworthy. Therefore, we need to adopt a relatively small loss weight $\lambda_{c}$, e.g. $\lambda_{c}=0.001$, to bound the distance between two minimizers tightly. This observation from our proven theoretical upper bound also coincides with our sanity loss weight ablation study.

We show our sanity-checked training pipeline in Fig. 4. We introduce each loss in the subsection. The self-sanity check can be formulated into a loss function form as

So that the self-sanity loss penalizes the squared differences between predicted displacement maps and the ideal ones. Next, we use $m\rightarrow f$ direction as an example ($f\rightarrow m$ direction follows the same principle) to formulate the proposed cross-sanity check loss. We calculate for every voxel and define a binary mask $\mathcal{M}^{m\rightarrow f}$ in the form of

An interpretation of this binary mask $\mathcal{M}^{m\rightarrow f}$ is that it records violations of the cross-sanity check (Eq. 4) for each individual point. In this way, $<0$ solution in Eq. 10 will not challenge the theoretical formulations since these points are masked out. Thus, we can formulate the proposed cross-sanity check in the form of a loss function of

The final $\mathcal{L}_{\rm cross}=\frac{1}{2}(\mathcal{L}^{m\rightarrow f}_{\rm cross}+\mathcal{L}^{f\rightarrow m}_{\rm cross})$. Here, $\odot$ denotes element-wise multiplication so that we only retain points violating the cross-sanity check. In this case, the loss value is only calculated for those violators, visualized as occlusion masks, shown in Fig. 5. Finally, the total loss is

Here, $\mathcal{L}_{\rm sim}$ as NCC loss, $\mathcal{L}_{\rm reg}$ as $||{\nabla(u)}||^{2}_{2}$, and $\lambda_{r}=1$, following standard deep registrations [6, 9]. If not specified otherwise, $\lambda_{s}=0.1$ and $\lambda_{c}=0.001$. While training, the model optimizes the total loss $\mathcal{L}_{\rm total}$ on different moving-fixed image pairs $(m,f)$ in the training set $\mathcal{D}$ as

In this way, we fulfill our novel regularization-based sanity-enforcer formulation for training a registration model.

Atlas-to-subject registration. We split in total 576 T1–weighted brain MRI images from the Information eXtraction from Images (IXI) [1] database into 403, 58, and 115 volumes for training, validation, and test sets. The moving image is an atlas brain MRI obtained from [23]. FreeSurfer [14] was used to pre-process the MRI volumes. The pre-processed image volumes are all cropped to the size of 160×192×224. Label maps, including 30 anatomical structures, are obtained using FreeSurfer for evaluating registration. Besides CNN-based models [6, 23, 10], we also include other transformer-based deep methods [10, 51, 48, 46] as baselines following [9]. The comparison results are presented in Tab. 4, and qualitative comparisons in Fig. 6. For cross-sanity check, we set $\alpha=0.1$ and $\beta=12$. Overall, we can observe models with sanity checks achieve better diffeomorphisms without impairing any Dice performance and even improve the performance for the VM model. We show qualitative comparisons in Fig. 7. Compared to the naive counterparts, our sanity-checked model produces more regular maps from both directions and reduces huge sanity errors, comparably using error heatmaps for both self-sanity and cross-sanity errors.

Subject-to-subject registration. OASIS dataset [26, 21] contains a total of 451 brain T1 weighted MRI images, with 394/19/38 images used for training/validation/testing purposes, respectively. The pre-processed image volumes are all cropped to the size of 160×192×224. Label maps for 35 anatomical structures are provided using FreeSurfer [14] for evaluation. For the cross-sanity check, we set $\alpha=0.1$ and $\beta=10$. The results are shown in Tab. 4. In terms of main metrics, we achieve on-par performance with all the state-of-the-art methods, further certifying our sanity checks will not contaminate the performance of models. We also show our qualitative comparison in Fig. 8, indicating that our sanity-enforced image registration results in less self-sanity error from the self-sanity error map comparison and more regular maps in the deformed grid. All these comparisons prove the efficacy of our regularization-based sanity-enforcer to prevent both sanity errors.

Pre-operative and post-recurrence scans registration. BraTSReg [3] dataset contains 160 pairs of pre-operative and follow-up brain MR scans of glioma patients taken from different time points. Each time point contains native T1, contrast-enhanced T1-weighted, T2-weighted, and FLAIR MRI. Within the dataset, 140 pairs of scans are associated with 6 to 50 manual landmarks in both scans. The other 20 pairs of scans only have landmarks in the follow-up scan. Following [31], we create a five-fold schema where each fold has 122/10/28 for training/validation/test, respectively. We set $\alpha=0.01$ and $\beta=0.03$ in the cross-sanity check since DIRAC predicts normalized displacement (range 0 to 1). We take an average of five folds to report our results in Tab. 4. All baselines and our models use the same training setup. We specifically focus on the model training part and directly train all models on data downloaded from the official website [5] without pre-affine and post-processing steps mentioned in their paper [30]. The best model is selected on the validation set, which has the lowest TREs. Despite the high performance of the baseline, we can observe that the substitution by cross-sanity check and the addition of self-sanity check enable all metrics’ improvements to be claimed as state-of-the-art. We also show qualitative comparisons in Fig. 9, where the comparison between SSE error maps demonstrates the effectiveness of our self-sanity check.

We conduct a broad sanity analysis on a wide range of models on the IXI dataset, including popular models (VM [6] and TM [9]), model with cycle consistency (CM [23]), models with inverse consistency (ICNet [50] and ICON [16]), probabilistically-formulated models (VM-diff [13] and TM-diff [9]), and models with diffeomorphic add-ons (MIDIR [34] and TMBS [9]). The results are shown in Tab. 5. We find that models with probabilistic formulation or cycle consistency have higher self-sanity, whereas the other models are insensitive to this scenario and have an inferior SDice. However, these probabilistic formulated models suffer from an relatively low Dice score, refraining from practical usage. For CM, albeit low self-sanity error, the correspondence error between image pairs is not negligible. We also find it interesting that models with diffeomorphic add-ons still suffer from large sanity errors, despite their ability to produce a diffeomorphic/regular map, caused by not modeling different mappings from both directions. Overall, our behavior studies show that most models suffer from inferior Dice/SDice performance, lower diffeomorphism, or significant sanity errors. Our proposed sanity checks improve performance from every aspect, producing a more regular map and preventing sanity errors.

Cross-sanity error under strict inverse consistency.

We testify the performance under strict inverse consistency as in Def. 1, the same as CICE in [11], shown in Tab. 6. The reductions of strict inverse consistency error compared using relaxed inverse consistency are caused by many small errors taken into account for calculating the average, where our sanity-checked methods still show better performance, compared to all models, including the model trained using explicitly strict inverse consistency, e.g. [16].

Sanity checks loss weight study. We derive an upper bound for our saner registration in Lemma 4. We show in the study that we can utilize this upper bound to guide us to set appropriate loss weight $\lambda_{c}$ for the cross-sanity check. For fast verification, we randomly sample 50 subjects from the training set of the IXI dataset and validate/test in the full validation/test set. We also include the self-sanity check loss weight $\lambda_{s}$ here for comparing purposes. The experimental results are shown in Tab. 7. Compared to $\lambda_{c}$, $\lambda_{s}$ is insensitive to weight change. As we discussed earlier, if we set higher loss weight to cross-sanity check, for example, 0.01/0.001, approximately having upper bounds as 0.1/0.01, respectively. This creates a huge difference when we train our models. This high tolerance for error results in significantly degenerated performance for Dice, where our derived upper bound gives useful guidance for the setting of loss weights. More discussions can be found in Appendix.

Model ablation study. We study our proposed sanity checks on IXI and OASIS validation datasets and report the results in Tab. 8. We show that each sanity check reduces the corresponding errors without compromising other metrics’ performance compared to their naive counterparts. We also find that the bidirectional registration and cross-sanity check can also mitigate self-sanity errors to a certain level, but cannot eliminate such errors completely, proving that our self-sanity check is necessary to regulate the models’ behavior for mapping identical image pairs. More experiments such as parameters $\alpha$ and $\beta$ numerical study, sanity preservation study, statistical significance of results, ablative qualitative results, etc, can be seen in the Appendix.