Deep Implicit Optimization for Robust and Flexible Image Registration

We thank the reviewer for their detailed and insightful feedback. We are excited to hear the reviewer found the paper both clear and innovative, and appreciated the paper’s technical depth, insight provided by Fig.2. We believe we have addressed all concerns, and look forward to engaging in a discussion and convincing the reviewer about the importance of the work.

A major limitation is the framework's registration accuracy, as measured by the Dice score

The focus of our paper is not solely achieving the highest Dice score on in-distribution (ID) data. Instead, our primary objective is to develop a robust and flexible method that performs well both ID and under domain shifts such as variations in preprocessing and acquisition configurations. While LKU-Net performs exceptionally well on the OASIS dataset, it completely breaks down under domain shift, which is a significant disadvantage. In contrast, a baseline like SynthMorph (SM) gets a Dice score of ~0.77 on the OASIS validation set, but performs on par with DIO under domain shifts due to its modality-agnostic training methodology. Our approach’s motivation is similar to SM, which aims to maximize performance and robustness, which we believe offers substantial practical benefits in real-world scenarios where data variations are common.

The absence of specific smoothness measurements … trade-off with the deformation field's smoothness.

These numbers are reported in Table 2. There is a slight (obvious) tradeoff between Dice score and det(Jacobian) values. Table 2 also shows the versatility of our learned feature maps, which do not overfit to the optimization algorithm used during training.

The paper lacks clarity in the reproducibility ...

There is a slight difference between the results reported in the papers and the challenge leaderboard, notably due to implementation improvements post-publication. We used the recommended parameters in the official code, and found the numbers to agree (with reasonable error margins) with that of their papers.

Additionally, in Table 4, the performance of LKU-Net with Dice supervision is paradoxically worse than without

This is not observed by us alone – [1] show in Sec 4.5.2 that most architectures like Convnet-affine and VTN-affine cannot generalize outside the training domain. Moreover, [1] observe that the semi-supervised models are inferior to their unsupervised models in the LPBA dataset, indicating anatomical knowledge injected to the model with supervision may not generalize well to unseen data beyond the training data. We observe this exact phenomenon with more baselines in Tab.4.

Baselines like SM and LapIRN perform reasonably on OASIS (way worse than LKU), but demonstrate consistent robustness to domain shift. Does it mean they show paradoxical behavior? Very likely not. In general, for registration, the validation overlap on in-distribution data tells nothing about the domain shift performance.

We’re confident that LKUNet is trained properly. We will release all code and models used in the paper for fairness and reproducibility.

the paper falls short in comparing its method to other registration methods that also combine learning with optimization

To our knowledge (see L104-L118), most methods do not perform learning and optimization end-to-end. The closest work Reviewer UTrd pointed out is KeyMorph, which cannot be extended to general warp fields which do not have closed-form solutions.

1 Could the authors clarify the primary benefits of their framework … methods would be helpful.

ConvexAdam, SAMConvex, PDDNet use 6D correlation volumes (instead of 4D features) which do not scale well with large data. Moreover, they are not learned end-to-end. PnP uses DEQ to finetune a AWGN network, but the data-fidelity term comes from the intensity images, while our features also incorporate label-fidelity. These are clear differences showing the elegance and strength of our formulation.

2 The paper claims that, "For the first time, our method allows switching, why is this an advantage over other methods?

The ability to switch between arbitrary transformation representations (e.g., from free-form to diffeomorphic) at test time without retraining is a significant advantage. Traditional DLIR methods fix the warp field representation during training, limiting flexibility and requiring retraining if the representation needs to change due to different downstream needs (e.g. model was trained with diffeomorphism for normative patients but requires non-diffeomorphism at inference for a subject with pathology). DIO allows for this flexibility, enabling the model to zero-shot adapt to different warp requirements and constraints. This capability is particularly beneficial in scenarios where the desired transformation properties may evolve or vary across applications.

3 The paper mentions that anatomical landmarks can be incorporated … within the optimization loop?

Landmark information is incorporated indirectly using a Dice loss (for labelmaps) or landmark distance (for keypoints). Similar to DLIR methods, using a Dice loss at training will learn features such that instance optimization of learned features minimizes dice alignment loss. Therefore, the learned features are label-aware (as seen by the performance gap between ours and ANTs/FireANTs).

At test time, if labels/keypoints are available, instance-optimization can simply add a Dice/landmark distance loss (sparse loss) in addition to feature image alignment loss (dense loss). Hyperparameter tuning can also be done at inference-time.

Concerns about implementation and reproducibility persist.

We have provided our entire codebase in the supplementary material. If there are additional concerns, we’re happy to clarify during the discussion period.

[1] Mok, Tony CW, and Albert Chung. "Affine medical image registration with coarse-to-fine vision transformer." CVPR 2022.

We thank the reviewer for their feedback. We are delighted to learn that the reviewer finds the paper well written and appreciates multi-scale feature optimization to be an interesting approach. We believe we have addressed all the remaining concerns and hope the reviewer increases their score and advocates for the acceptance of our paper.

W1. Possibilities of getting artifacts from different voxel sizes

This is standard practice in classical image registration, where the fixed and moving images are typically not in the same space. For applications such as unimodal and multimodal registration [1,2,3], ex-vivo to in-vivo registration [4,5], images often differ in size and voxel resolution. Instead of resampling the moving image twice (once to make its size equal to the fixed image and second during warping) thereby introducing more resampling artifacts, classical methods perform a single resampling (during warping only). This is also the motivation for introducing decoupled feature extraction and optimization – to enable registration with modalities with grossly different spatial resolutions.

Many established toolkits like ANTs, NiftyReg, Greedy, Demons, and LDDMM handle varying image sizes effectively. In contrast, current DLIR methods are limited because they require images to have matching spatial dimensions due to channel-wise concatenation. This is a problem with applications like in-vivo to ex-vivo, where downsampling the ex-vivo image loses fine detail and upsampling the in-vivo image consumes memory. Our method overcomes this limitation of existing DLIR methods, allowing for more flexible registration.

Overall, treating the input images separately from the feature extractor raises several questions regarding the credibility of the transformation field $\phi$

“this approach might introduce artifacts and lose fine details when propagating the source image to match the target image” – This is indeed the case, but for DLIR methods, which typically resample the images into a uniform spatial resolution (1mm isotropic for OASIS dataset in Learn2Reg) followed by another resampling due to warping – introducing additional resampling artifacts. This loses fine details if the image was acquired at a higher resolution (say 0.7mm in-plane resolution) and introduces resampling artifacts due to anisotropy.

W2. Motivation for using multi-scale optimization

Using multi-scale optimization is also standard practice in the classical image registration community. The rationale is that optimization at the finest scale only will converge at some bad local minima, due to the ill-conditioned and ill-posed nature of deformable registration. Therefore, registration at coarser resolutions aligns large structures like ventricles, and finer scales align small structures like gray matter sulci. Motivated by classical methods, certain DLIR methods also use multiscale optimization. Multiscale optimization in DIO is motivated by learning discriminative feature maps at different scales compared to naive downsampled versions of intensity images.

For example, intensity based registration at 4x scale achieves a Dice score of 0.6, whereas our method achieves around 0.75 in the OASIS val set, and enjoys a similar improvement for 2x downsampling. We’ve added this ablation to Appendix due to space issues.

W3. Applicabilities of learned multi-scale dense features from sparse images.

We’re not entirely sure what “DIO preserved the gradient in the loss function” and “discussed that deep networks recovered affine transform with 90% overlap” mean. We are happy to clarify in the discussion period.

W4. Experimental supports.

We’ve added an ablation on training with and without multiscale optimization. We see slightly better scores for multi-scale optimization, due to better local minima achieved during optimization at coarser scales.

Since our model learns feature maps for registration, it has to be followed by an optimization step. The feature images are indeed helpful - as is evident from the difference between ANTs and our method, wherein the underlying optimizer is the same.

Two interesting sets of experiments that validate domain shift

This is exactly what we did, i.e. trained on OASIS and tested on the rest. We did not train on other datasets due to their limited size (364 training images in OASIS versus 40 in LPBA40). Training on multiple datasets is possible, but the size of the OASIS dataset would dominate over the others. Given space considerations, we leave more nuanced ablations for future work.

How do decoupling feature learning and optimization help the model in getting improved performance over existing DLIR methods?

Our goal is to not asymptotically outperform existing methods (see Limitations of the paper), but to tradeoff a little bit of in-distribution performance (as in Table1), for generalizable and robust performance under domain shift. Moreover, small improvements in dice score can be attributed to slight misalignments for smaller anatomical regions.

[1] Murphy, Keelin, et al. "Evaluation of registration methods on thoracic CT: the EMPIRE10 challenge." IEEE transactions on medical imaging 30.11 (2011): 1901-1920.

[2] A. Klein, et al. Evaluation of nonlinear deformation algorithms applied to human brain MRI registration. NeuroImage, 46(3):786–802, July 2009.

[3] Wang, Quanxin, et al. "The Allen mouse brain common coordinate framework: a 3D reference atlas." Cell 181.4 (2020): 936-953.

[4] Khandelwal, Pulkit, et al. "Automated deep learning segmentation of high-resolution 7 tesla postmortem MRI for quantitative analysis of structure-pathology correlations in neurodegenerative diseases." Imaging Neuroscience 2 (2024): 1-30.

[5] Meyer, Charles R., et al. "A methodology for registration of a histological slide and in vivo MRI volume based on optimizing mutual information." Molecular imaging 5.1 (2006): 7290-2006.

We thank the reviewer for their insightful feedback, and we’re glad to note that they found the paper overall very interesting. The review has been immensely helpful in improving the quality of our work. We believe we have addressed all concerns in the rebuttal, and we look forward to a discussion and convincing the reviewer about the importance and relevance of our work.

Sparse keypoints (kps) don’t facilitate dense reg. is wrong - KeyMorph (KM) does deformable reg. with sparse kps

This is an oversight on our part. We have revised this incorrect characterization in the paper, and mentioned that these formulations allow linear and deformable registration. However, our motivation for DIO still holds, i.e. KM can only output deformation fields that admit closed-form warps (e.g using TPS), from learned keypoints. However, TPS is a limited class of warps, cannot be guaranteed to be diffeomorphic, and a vast majority of widely used parameterizations (free-form, SVF, geodesic, LDDMM, SyN) do not admit closed form solutions making KM inapplicable. While in KM the keypoints are trained end-to-end using backprop, we use implicit differentiation through the optimization to learn dense feature maps. DIO is therefore an elegant and generalized extension of the KM idea: “keypoints” are replaced by “dense feature maps” and “closed-form warp field” parameterizations are replaced by “arbitrary warp fields” (i.e. solutions of optimization objectives). Moreover, KM runs out of memory quickly – TPS with 512 kps runs OOM on a A6000 GPU whereas DIO outputs multi-scale feature maps equivalent to 192 * 224 * 160 * (1 + 1/8 + 1/64) * 16 / 3 ~ 41M keypoints.

acquisition and preprocessing protocols [62, 53, 70, 43]…not supported by the citations

We agree! The first three citations were supposed to reference different image acquisition and preprocessing protocols themselves (neuro/lung imaging). However, the statement itself is correct as shown in Fig.3. and Tab.4; these are very strong results to support the statement. We have cited Fig.3 in the paper instead.

DL methods dont output local minima, but instance optimization (IO) can – this is very related

We agree that these approaches are related, however there is a subtle but crucial difference. In existing approaches, the deep network is “label-aware”, but outputs a warp field that may not be a local minima of any objective. IO is performed on the intensity images that are not label aware. This IO is a small change on top of the predicted warp field. In contrast, we train feature maps to be label aware by design (since registering feature maps minimize label overlap loss), as shown in Tab.1 and Fig.3. We sidestep warp field prediction altogether and perform a single IO step on label aware features instead of label-unaware intensity images.

Synthmorph is missing in Tab1

For DLIR, we only added a few recent methods that score > 0.82 Dice on OASIS val. SM scores a ~0.77 Dice on the OASIS dataset, with the provided pretrained model.

Deformable KeyMorph is missing from the whole experiments section

We were aware only of the affine variant. However, the performance of KM is unsatisfactory. We trained KM on the OASIS train split, and observed the following dice scores for OASIS val and zero-shot on IBSR:

The in-distribution performance strongly agrees with deformable registration dice scores in Fig.5,7,8,9 in the KM paper. We observe that KM is very robust in recovering from large rotational deviations due to its training scheme (good for canonicalizing data acquired on non-standard coordinates) , but is an unsatisfactory baseline for deformable registration due to ~6M voxel displacements being parameterized by only 512 keypoints - making registration of small subcortical structures especially hard.

In the domain-shift section the authors show that their method works because of IO, but others dont use IO

This experimental design is justified by the ‘feature space’ in which IO is performed. We find the warp field as the minimizer to an objective function with feature maps. The minimizer is found using IO. However, the warp field produced is a local minimizer of IO in the “learned feature space” and not in the “intensity space”. The warp field produced by DIO may not be a minimizer in the intensity space. Moreover, implementing IO for every baseline is a highly tedious task, since different baselines have different implementations of the warp field, including scaling, different coordinate systems, etc.

A fair comparison for DIO would be to perform IO of learned features, followed by IO of intensity images. For simplicity, we do not perform IO in the intensity space for any method.

some (e.g. SM, which the authors refer to here) does not even require both images be of the same size.

This is not the case, since every DLIR method requires the fixed and moving images to be concatenated channelwise, implying matching spatial dimensions. IBSR18 volumes have different voxel sizes. Moreover, some of the methods have hardcoded image sizes – this is an implementation issue. This also highlights challenges in using DLIR methods from official implementations. Methods like TransMorph use ViT end-to-end and cannot be used with different image sizes. We will make all scripts public for fairness and reproducibility.

it seems to me that it would be an important ablation to take the features of some robust method … end-to-end training is useful.

This seems like a useful ablation, but we are not sure which method to use for feature extraction.

The claim in 4.4 is also missing some reasonable comparison

We’re not sure what that would show. Sec. 4.4 is written to show that the learned features in DIO do not “overfit” to the warp field representation during training, shown by inference-time switching to an unseen optimizer.