Deep Implicit Optimization for Robust and Flexible Image Registration

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-03 14:43:34

Abstract: Deep Learning in Image Registration (DLIR) methods have been tremendously successful in image registration due to their speed and ability to incorporate weak label supervision at training time. However, DLIR methods forego many of the benefits of classical optimization-based methods. The functional nature of deep networks do not guarantee that the predicted transformation is a local minima of the registration objective, the representation of the transformation (displacement/velocity field/affine) is fixed, and the networks are not robust to domain shift. Our method aims to bridge this gap between classical and learning methods by incorporating optimization as a layer in a deep network. A deep network is trained to predict multi-scale dense feature images that are registered using a black box iterative optimization solver. This optimal warp is then used to minimize image and label alignment errors. By implicitly differentiating end-to-end through an iterative optimization solver, our learned features are registration and label-aware, and the warp functions are guaranteed to be local minima of the registration objective in the feature space. Our framework shows excellent performance on in-domain datasets, and is agnostic to domain shift such as anisotropy and varying intensity profiles. For the first time, our method allows switching between arbitrary transformation representations (free-form to diffeomorphic) at test time with zero retraining. End-to-end feature learning also facilitates interpretability of features, and out-of-the-box promptability using additional label-fidelity terms at inference.

Main paper: 
 introduction: 
 Deformable Image Registration (DIR) refers to the local, non-linear alignment of images by estimating a dense displacement field. Many workflows in medical image analysis require images to be in a common coordinate system for comparison, analysis, and visualization, including comparing intersubject data in neuroimaging [53, 104, 97, 38, 89, 94] , biomechanics and dynamics of anatomical structures including myocardial motions, airflow and pulmonary function in lung imaging, organ motion tracking in radiation therapy [78, 77, 11, 70, 29, 105, 50, 18, 71, 84] , and life sciences research [112, 104, 99, 80, 98, 72, 17] .Classical DIR methods are based on solving a variational optimization problem, where a similarity metric is optimized to find the best transformation that aligns the images. However, these methods are typically slow, and cannot leverage learning to incorporate a training set containing weak supervision such as anatomical landmarks or expert annotations. The quality of the registration is therefore limited by the fidelity of the intensity image. Deep Learning for Image Registration (DLIR) is an interesting paradigm to overcome these challenges. DLIR methods take a pair of images as input to a neural network and output a warp field that aligns the images, and their associated anatomical landmarks. The neural network parameters are trained to minimize the alignment loss over image pairs and landmarks in a training set. A benefit of this method is the ability to incorporate weak Submitted to 38th Conference on Neural Information Processing Systems (NeurIPS 2024). Do not distribute. supervision like anatomical landmarks or expert annotations during training, which performs better landmark alignment without access to landmarks at inference time.Motivation. However, DLIR methods face several limitations. First, the prediction paradigm of deep learning implies the feature learning and amortized optimization steps are fused; transformations predicted at test-time may not even be a local minima of the alignment loss between the fixed and moving image. The end-to-end prediction also implies that the representation of the transformation is fixed (as a design choice of the network), and the model cannot switch between different representations like free-form, stationary velocity, geodesic, LDDMM, B-Splines, or affine at test time without additional finetuning, in sharp contrast to the flexibility of classical methods. Typical registration workflows require a practitioner to try different parameterizations of the transformation (free-form, stationary velocity, geodesic, LDDMM, B-Splines, affine) to determine the representation most suitable for their application and additional retraining becomes expensive. Moreover, design decisions like sparse keypoint learning for affine registration [103, 16, 69, 40] do not facilitate dense deformable registration. Furthermore, DLIR methods do not allow interactive registration using additional landmarks or label maps at test time, which is crucial for clinical applications. Hyperparameter tuning for regularization is also expensive for DLIR methods. Although recent methods propose conditional registration [44, 67] to amortize over the hyperparameter search during training, the family of regularization is fixed in such cases, and space of hyperparameters becomes exponential in the number of hyperparameter families considered. Lastly, current DLIR methods are not robust to minor domain shifts like varying anisotropy and voxel resolutions, different image acquisition and preprocessing protocols [62, 53, 70, 43] . Robustness to domain shift is imperative to biomedical and clinical imaging where volumes are acquired with different scanners, protocols, and resolutions, where the applicability of DLIR methods is limited to the training domain.Contributions. We introduce DIO, a generic differentiable implicit optimization layer to a learnable feature network for image registration. By decoupling feature learning and optimization, our framework incorporates weak supervision like anatomical landmarks into the learned features during training, which improves the fidelity of the feature images for registration. Feature learning also leads to dense feature images, which smoothens the optimization landscape compared to intensity-based registration due to homogenity present in most medical imaging modalities. Since optimization frameworks are agnostic to spatial resolutions and feature distortions, DIO is extremely robust to domain shifts like varying anisotropy, difference in sizes of fixed and moving images, and different image acquisition and preprocessing protocols, even when compared to models trained on contrast-agnostic synthetic data [43] . Moreover, our framework allows zero-cost plug-andplay of arbitrary transformation representations (free-form, geodesics, B-Spline, affine, etc.) and regularization at test time without additional training and loss of accuracy. This also paves the way for practitioners to perform quick and interactive registration, and use additional arbitrary 'prompts' such as new landmarks or label maps out-of-the-box at test time, as part of the optimization layer. 
related work: 
 Deep Learning for Image Registration DIR refers to the alignment of a fixed image I f with a moving image Im using a transformation φ ∈ T where T is a family of transformations. Classical methods formulate a variational optimization problem to find the optimal φ that aligns the images [15, 4, 7, 5, 6, 2, 15, 25, 24, 23, 27, 39, 63, 102, 101, 100, 46, 60, 61, 76, 33, 32, 12] . In contrast, earliest DLIR methods used supervised learning [19, 55, 82, 88] to predict the transformation φ.Voxelmorph [13] was the first unsupervised method utilizing a UNet [83] for unsupervised registration on brain MRI data. Recent works considered different architectural designs [21, 56, 48, 66] , cascade-based architectures and loss functions [116, 115, 49, 26, 68, 114, 79, 20] , and symmetric or inverse consistency-based formulations [65, 51, 52, 92, 116] . [67, 44] inject the hyperparameter as input and perform amortized optimization over different values of the hyperparameter. Domain randomization and finetuning [43, 96, 73, 30] are also proposed to improve robustness of registration to domain shift, that is a core necessity in medical imaging since different institutions follow varying acquisition and preprocessing pipelines. Foundational models are also proposed to improve registration accuracy [57, 93] . Another line of work propose to use the implicit priors of deep learning [95] within an optimization framework [110, 106, 49, 45] . We refer the reader to [36, 41, 28] for other detailed reviews.Iterative methods for DLIR Owing to the success of iterative optimization methods, few DLIR methods propose emulating the iterative optimization within a network. [115, 116] use a cascade of networks to iteratively predict a warp field, and use the warped moving image as the input to the next layer in the cascade. TransMorph-TVF [20] uses a recurrent network to predict a time-dependent velocity field. [114] use a shared weights encoder to output feature images at multiple scales, and a deformation field estimator utilizing a correlation layer. RAFT [91] similarly builds a 4D correlation volume from two 2D feature maps, and updates the optical flow field using a recurrent unit that performs lookup on the correlation volume. However, such recursive formulations have a large memory footprint due to explicit backpropagation through the entire cascade [8] , and are not adaptive or optimal with respect to the inputs. In contrast, DIO uses optimization as a layer -guaranteeing convergence to a local minima, and implicit differentiation avoids storing the entire computation graph making the framework both memory and time efficient.Feature Learning for Image Registration [103, 16, 69, 40] learn keypoints from images which is then used to compute the optimal affine transform using a closed form solution. However, these methods are restricted to transformations that can be represented by differentiable closed-form analytical solutions, making backpropagation trivial. These sparse keypoints cannot be reused for dense deformable registration either. On the other hand, dense deformable registration (diffeomorphic or otherwise) is almost universally solved using iterative optimization methods. This motivates the need to perform implicit differentiation through an iterative optimization solver to perform feature learning for registration. Other approaches learn image features to perform registration [108, 59, 107, 81] , but do not perform feature learning and registration end-to-end, i.e., the features obtained are not task-aware and may not be optimal for registration, especially for anatomical landmarks. Learned features are either fed into a functional form to compute the transformation end-to-end, or are learned using unsupervised learning in a stagewise manner. In contrast, by implicitly differentiating through a black-box iterative solver, and minimizing the image and label alignment losses end-to-end, DIO learns features that are registration-aware, label-aware, and dense. The optimization routine also guarantees that the transformation is a local minima of the alignment of high-fidelity feature images.Deep Equilibrium models Deep Equilibrium (DEQ) models [9, 34] have emerged as an interesting alterative to recurrent architectures. DEQ layers solve a fixed-point equation of a layer to find its equilibrium state without unrolling the entire computation graph. This leads to high expressiveness without the need for memory-intensive backpropagation through time [10, 8, 31, 75, 37, 111] . PIRATE [45] uses DEQ to finetune the PnP denoiser network for registration, but unlike our work, the data-fidelity term comes from the intensity images. However, these methods use DEQ to emulate an infinite-layer network, which typically consists of learnable parameters within the recurrent layer.Conceptually, our work does not aim to simply emulate such an infinite cascade, but rather use DEQ to decouple feature learning and optimization in an end-to-end registration framework.This inherits all the robustness and agnosticity of optimization-based methods, while retaining the fidelity of learned features. DEQ allows us to avoid the layer-stacking paradigm for cascades, and use optimization as a black box layer without storing the entire computation graph, leading to constant memory footprint and faster convergence. This allows learnable features to be registration-aware since gradients are backpropagated to the feature images through the optimization itself.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: An implicit bias of the representation and optimization algorithm is discussed in the Discussion and Appendix.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.• The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.• The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.• The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
methods: 
 The registration problem is formulated as a variational optimization problem:EQUATIONwhere I f and Im are fixed and moving images respectively, L is a loss function that measures the dissimilarity between the fixed image and the transformed moving image, and R is a suitable regularizer that enforces desirable properties of the transformation φ. We call this the image matching objective. If the images I f and Im are supplemented with anatomical label maps L f and Lm, we call this the label matching objective. Classical methods perform image matching on the intensity images, but the label matching performance is bottlenecked by the fidelity of image gradients with respect to the label matching objective, and dynamics of the optimization algorithm. Deep learning methods mitigate this by injecting label matching objectives (for example, Dice score) into the objective Eq. ( 1) and using a deep network with parameters θ to predict φ for every image pair as input. In essence, learning-based problems solve the following objective: The warped image/label are compared with the fixed image/label using a similarity metric.EQUATIONwhere φ θ (I f , Im) is abbreviated to φ θ . This leads to learned transformations φ θ that perform both good image and label matching. However, the feature learning and optimization are coupled, and the learned features are optimized only for a specific training domain. This limitation primarily marks the difference between DIO and existing DLIR methods.Fig. 1 shows the overview of our method. Our goal is to learn feature images such that registration in this feature space corresponds to both image and label matching performance, by disentangling feature learning and optimization. We do this by using a feature network to extract dense features from the intensity image, that are used to solve Eq. (1) using a black-box optimization solver, and obtain an optimal transform φ * . Once φ * is obtained, this is plugged into Eq. ( 2) to obtain gradients with respect to φ * . Since φ * is a function of the feature images, we implicitly differentiate through the optimization to backpropagate gradients to the feature images and to the deep network.We discuss the details of our method in the following sections.The first component of our framework is a feature network that extracts dense features from the intensity images. This network is parameterized by θ, and takes an image I ∈ R H×W ×D×Cin as input and outputs a feature map F ∈ R H×W ×D×C , where C is the number of feature channels, i.e. F = g θ (I). Unlike existing DLIR methods where moving and fixed images are concatenated and passed to the network, our feature network processes the images independently. This allows the fixed and moving images to be of different voxel sizes. The feature network can also output multi-featurefeature maps F = g θ (I) = [F 0 , F 1 , . . . , F N ],whereF k ∈ R H/2 k ×W/2 k ×D/2 k ×C k ,which can be used by multi-scale optimization solvers. The feature network is agnostic to architecture choice, and we ablate on different architectures in the experiments.Given the feature maps F f and Fm extracted from the fixed and moving images, an optimization solver optimizes Eq. ( 1) to obtain the transformation φ * . This can be written by modifying Eq. (1) to use the feature maps F ; i.e. φ * = arg minφ C(F f , Fm • φ). A local minima of this equation satisfies:EQUATIONThis φ * is used to compute the loss Eq. ( 2) to minimize image and label matching objective. To propagate derivatives from φ * to the feature images F f , Fm, we invoke the Implicit Function Theorem [54] :Theorem 1 For a function ϱ : R n × R m1+m2 → R n that is continuously differentiable, if ϱ(φ * , F f , Fm) = 0 and ∂ϱ ∂φ | φ * ̸ = 0, then there exist open sets U, V f , Vm containing φ * , F f , Fm,and a function φ * (F f , Fm) defined on these open sets such that ϱ(φ * (F f , Fm), F f , Fm) = 0.Given the Implicit Function Theorem, we write ϱ(φ * (F f , Fm), F f , Fm) = 0 and differentiate with respect to F f to obtain:EQUATIONThe gradients of φ come from Eq. (2) (i.e. ∂T ∂φ ), and the gradients of F f w.r.t. Eq. ( 2) are obtained as∂T ∂F f = -∂T ∂φ ∂ϱ ∂φ -1 ∂ϱ∂F f . The gradients of Fm are obtained similarly.This design ensures that optimal registration in the feature space corresponds to optimal registration both in the image and label spaces. Furthermore, the optimization layer ensures that the φ * is a local minima of this high-fidelity feature matching objective, i.e., the features obtained by the network.Jacobian-Free Backprop In practice, the Jacobian ∂ϱ ∂φ is expensive to compute, given the high dimensionality of φ and ϱ. Following [31] , we substitute the Jacobian to identity, and compute∂T ∂F f ≈ -∂T ∂φ ∂ϱ∂F f . This leads to much less memory and stable training dynamics compared to other estimates of Jacobian like phantom gradients, damped unrolling, or Neumann series [35, 34] . Top row shows the intensity image with the corresponding multiscale features predicted by the deep network, where the L th level denotes a feature of size H/2 k ×W/2 k ×C k . Bottom row shows the loss landscape as a function of the relative translation between the squares in the fixed and moving image. Note the flat maxima which occurs when there is no overlap between the fixed and moving image, making optimization impossible if there is no overlap of the squares. On the contrary, the loss landscape for learned features is smooth, even at the finest scale, leading to much faster convergence even when there is no overlap between the intensity images. Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: Code is provided in the supplemental material. Data is publicly available and instructions are provided in the supplemental material.Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details.• While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).• The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nips.cc/public/guides/CodeSubmissionPolicy) for more details.• The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.• The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.• At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).• Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
dio learns dense features from sparse images: 
 A key strength of DIO is the ability to learn interpretable dense features from sparse intensity images for accurate and robust image matching. This is especially relevant for medical image registration, which typically contain a lot of homogenity in the intensity images, making registration difficult.We design a toy task to isolate and demonstrate this behavior. The fixed and moving images are generated by placing a square of size 32×32 pixels on an image of 128×128 pixels. The squares in the fixed and moving images overlap with a 50% chance. The task is to find an affine transformation to align the two images. However, classical optimization methods will fail this task 50% of the time, because when the squares do not overlap, there is no gradient of the loss function, illustrated by the flat loss landscape in Fig. 2 . However, deep networks discover features that significantly flatten this loss landscape in the feature matching space. To show this, we train a network to output multi-scale feature maps that is used to optimize Eq. ( 1) to recover an affine transform. We choose a 2D UNet architecture, and the multi-scale feature maps are recovered from different layers of the decoder path of the UNet. Since the features are trained to maximize label matching, the loss landscape is much flatter, and the network is able to recover the affine transform with > 99% overlap (Appendix A.4).End-to-end learning enables learning of features that are most conducive to registration, unlike existing work [108, 59, 107, 81] that may not contain discriminative registration-aware features about anatomical labels due to lack of task-awareness.Setup: We evaluated our method on inter-subject registration on the OASIS dataset [62] . The OASIS dataset contains 414 T1-weighted MRI scans of the brain with label maps containing 35 subcortical structures extracted from automatic segmentation with FreeSurfer and SAMSEG. We use the preprocessed version from the Learn2Reg challenge [42] where all the volumes are skull-stripped, intensity-corrected and center-cropped to 160×192×224. We use the same training and validation sets as provided in the Learn2Reg challenge to enable fair comparison with other methods.Table 1 : Performance on OASIS validation set.DIO is highly competitive with state-of-the-art DLIR methods in the in-distribution setting. Our feature learning incorporates label-aware features, which is evident from the superior performance compared to four SOTA optimization-based classical methods.Validation Method Dice HD95 ANTs [5] 0.786 ± 0.033 2.209 ± 0.534 NiftyReg [64] 0.775 ± 0.029 2.382 ± 0.723 LogDemons [100] 0.804 ± 0.022 2.068 ± 0.448 FireANTs [46] 0.791 ± 0.028 2.793 ± 0.602 Progressive C2F [58] 0.827 ± 0.013 1.722 ± 0.318 Little learning [87] 0.846 ± 0.016 1.500 ± 0.304 CLapIRN [67] 0.861 ± 0.015 1.514 ± 0.337 Voxelmorph-huge [14] 0.847 ± 0.014 1.546 ± 0.306 TransMorph [22] 0.858 ± 0.014 1.494 ± 0.288 TransMorph-Large [22] [48] . To extract multi-scale features from the networks, we attach single convolutional layers to the feature of the desired scales from the decoder path. For each of these architectures, we also consider "Encoder-Only" versions by discarding the decoder path, and creating independent encoders for each scale Fig. 9 , denoted as UNet-E and LKU-E. We choose Encoder-Only versions to ablate the performance using shared features from the decoder path versus independent feature extraction at each scale.Results: We compare our method with existing methods on the Learn2Reg OASIS challenge (Table 1 ). We compare with state-ofthe-art classical methods [5, 46, 64, 100] , and deep networks [58, 87, 67, 14, 22, 48] . DIO is highly competitive with existing methods, especially with TransMorph which uses up to two orders of magnitude more trainable parameters than DIO to achieve a similar performance. We note that the Large Kernel UNet architecture performs better than the standard UNet architecture, which is consistent with the findings in [48] , even for dense feature extraction. This is due to the larger receptive field of LKUNet, which is able to capture more context in the image. Moreover, the Encoder-Only versions of the network perform slightly worse than the full networks, showing that sharing features across scales is beneficial for the task.A key requirement of registration algorithms is to generalize over a spectrum of acquisition and preprocessing protocols, since medical images are rarely acquired with the same configuration.Existing DLIR methods are extremely sensitive to domain shift, and catastrophically fail on other brain datasets. On the contrary, DIO inherits the domain agnosticism of the optimization solver, and is robust under feature distortions introduced by domain shift.We evaluate the robustness of the trained models on three brain datasets: LPBA40, IBSR18, and CUMC12 datasets [85, 1, 53 We compare our method with a variety of DLIR baselines, trained with and without label supervision (the former denoted as 'w/ Dice sup.' in Fig. 3 ). Our method performs substantially better than all the baselines with a significantly narrower interquartile range on the IBSR18 and CUMC12datasets. The differences are significant -on IBSR18 and CUMC12, our median performance is higher than the third quartile of almost all baselines. The sturdy performance against domain shift provides a strong motivation for using optimization-in-the-loop for learnable registration. sufficient details to perform label-aware registration. Since our method learns dense features to minimize label matching losses, we can observe which features are necessary to enable label-aware registration. Fig. 4 highlights differences between scale-space images and features learned by our network. At all scales, the features introduces heterogeneity using a watershed effect and enhanced contrast to improve label matching performance.DLIR methods have been very popular due to their fast inference time by performing amortized optimization [14] . Classical methods generally focus on robustness and reproducubility, and do have GPU implementations for fast inference. However, modern optimization toolkits [60, 46] utilize massively parallel GPU computing to register images in seconds, and scale very well to ultrahigh resolution imaging. A concern with optimization-in-the-loop methods is the inference time. Table Table 3 shows the inference time for our method for all four architectures. These inference times are fast for a lot of applications, and the plug-and-play nature of our framework makes DIO amenable to rapid experimentation and hyperparameter tuning. Conclusion DLIR methods provide several benefits such as amortized optimization, integration of weak supervision, and the ability to learn from large (labeled) datasets. However, coupling of the feature learning and optimization steps in DLIR methods limits the flexibility and robustness of the deep networks. 
conclusion and limitations: 
 In this paper, we we introduce a novel paradigm that incorporates optimization-as-a-layer for learningbased frameworks. This paradigm retains all the flexibility and robustness of classical multi-scale methods while leverging large scale weak supervision such as anatomical landmarks into high-fidelity, registrationaware feature learning. Our paradigm allows "promptable" registration out-of-the-box as part of the plug-and-play optimization, where additional supervision such as labelmaps or landmarks can be added to the optimization loss at test time. Our fast implementation allows for implementation of optimization-as-a-layer in deep learning, which was previously thought to be infeasible, due to existing optimization frameworks being prohibitively slow. Densification of features from our method also leads to better optimization landscapes, and our method is robust to unseen anisotropy and domain shift. To our knowledge, our method is the first to switch between transformation representations (free-form to diffeomorphic) at test time without any retraining. This comes with fast inference runtimes, and interpretability of the features used for optimization. Potential future work can explore multimodal registration, online hyperparameter tuning and few-shot learning.Limitations The first limitation is unlike existing DLIR methods that concatenate the fixed and moving images to feed into the network, DIO processes the images independently. The features extracted from an image are therefore trained to marginalize the label matching performance over all possible moving images, and cannot adapt to the moving image. This leads to slightly asymptotically lower in-domain performance than methods like [48] . The second limitation is the implicit bias of the optimization algorithm. Implicit bias in SGD restricts the space of solutions for optimization problems that are overparameterized, such as deep networks [113, 90, 47, 74, 109] . In deformable registration, the implicit bias of SGD restricts the direction of the gradient of the particle at φ(x), which is always parallel to ∇Fm(φ(x)), independent of the fixed image and dissimilarity function.This limits the degrees of freedom of the optimization by N -fold for N -D images. This is unlike DLIR methods where the warp is not constrained to move along ∇Fm(φ(x)). This behavior is explored in more detail in Appendix A.1. Future work aims to mitigate this implicit bias for better performance. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: Implementation details are provided in Appendix and supplemental material. Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.• The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: All results are reported either with an error bar of one standard deviation, or boxplots with interquartile ranges and outliers are reported.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: Compute resources are provided in the Appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [No] Justification: Medical image registration has no immediate negative societal impact necessitating a dedicated discussion.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.• The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: [NA] Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.• Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.• We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: Appropriate citations are provided for existing code and data.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.• If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [Yes]Justification: Code is reasonably commented for a new reader to understand the implementation.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.• The paper should discuss whether and how consent was obtained from people whose asset is used.• At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA]Justification: [NA] Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.• According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. 
 
 Appendix: Model based systems, such as deep networks are not immune to inductive biases due to architecture, loss functions, and optimization algorithms used to train them. Functional forms of the deep network induce constraints on the solution space, but optimization algorithms are not excluded from such biases either. The implicit bias for Gradient Descent is a well-studied phenomena for overparameterized linear and shallow networks. Gradient Descent for linear systems leads to an optimum that is in the span of the input data starting from the initialization [113, 90, 47, 74, 109] .This bias is also dependent on the chosen representation, since that defines the functional relationship of the gradients with the parameters and inputs. This limits the reachable set of solutions by the optimization algorithm when multiple local minima exist.In the case of image registration, the optimization limits the space of solutions (warps) that can be obtained by the SGD algorithm. To show this, we consider the transformation φ as a set of particles in a Langrangian frame that are displaced by the optimization algorithm to align the moving image to the fixed image. Consider a regular grid of particles, whose locations specify the warp field. Let the location of i-th particle at iteration t be φ (t) (x i ). For a fixed feature image F f , moving image Fm and current iterate φ (t) , the gradient of the registration loss with respect to particle i at iteration t is given byEQUATIONwhereC ′ i (F f , Fm • φ (t) ) = ∂C(F f , Fm • φ (t) ) ∂M (φ (t) (x i ))is the (scalar) derivative of scalar loss C with respect to the intensity of i-th particle computed at the current iterate, and ∇Fm(φ (t) (x i )) is the spatial gradient of the moving image at the location of the particle. Note that the direction of the gradient of particle i is independent of the fixed image, loss function, and location of other particles -it only depends on the spatial gradient of the moving image at the location of the particle. This restricts the movement of a particle located at any given location along a 1D line whose direction is the spatial gradient of the moving image at that location.Since F f and Fm are computed independently of each other (and therefore no information of F f andFm is contained in each other), the space of solutions of φ is restricted by this implicit bias. This is restrictive because the similarity function and fixed image do not influence the direction of the gradient, and the optimization algorithm is biased towards solutions that are in the direction of the gradient of the moving image. The plot shows the loss curves for a multi-scale optimization of two feature images. Each plot also shows the absolute cosine similarity of per-pixel gradients obtained by C and Csurrogate at each iteration. Note that over the course of optimization, the cosine similarity is always 1 -demonstrating the implicit bias of the optimization for registration.We show this bias empirically -we perform multi-scale optimization algorithm using feature maps obtained from the network. We keep track of two gradients, one obtained by the loss function, and another obtained by the gradient of a surrogate loss Csurrogate(Fm, φ (t) ) = i Fm(φ (t) (x i )). Note that Csurrogate does not depend on the fixed image or the loss function. The gradient of Csurrogate with respect to the i-th particle is given by ∇Fm(φ (t) (x i )). At each iteration, we compute the magnitude of cosine similarly between the gradients of C and Csurrogate. Fig. 5 shows that the loss converges, and the per-pixel gradients can be predicted by Csurrogate alone, as depicted by the magnitude and standard deviation of cosine similarity between C and Csurrogate. This limits the movement of each particle along a 1D line in an N -D space, and limits the degrees of freedom of the optimization by N -fold for N -D images. Future work will aim at alleviating this implicit bias to allow for more flexible solutions.DIO is a learnable framework that leverages implicit differentiation of an arbitrary black-box optimization solver to learn features such that registration in this feature space corresponds to good registration of the images and additional label maps. This additional indirection leads to learnable features that are registration-aware, interpretable, and the framework inherits the optimization solver's versatility to variability in the data like difference in contrast, anisotropy, and difference in sizes of the fixed and moving images. We contrast our approach with a typical classical optimization-based registration algorithm in Fig. 6 . A classical multi-scale optimization routine indiscriminately downsamples the intensity images, and does not retain discriminative information that is useful for registration. Since our method is trained to maximize label alignment from all scales, multi-scale features obtained from our method are more discriminative and registration-aware. We also compare DIO with a typical DLIR method in Fig. 7 . Note that the fixed end-to-end architecture and functional form of a deep network subsumes the representation choice into the architecture as well, limiting its ability to switch to arbitrary transformation representations at inference time without additional retraining. Our framework therefore combines the benefits of both classical (robustness to out-of-distribution datasets, and zero-shot transfer to other optimization routines) and learning-based methods (high-fidelity, label-aware, and registration-aware).For all experiments, we use downsampling scales of 1, 2, 4 for the multi-scale optimization. All our methods are implemented in PyTorch, and use the Adam optimizer for learning the parameters of the feature network. Note that in Eq. ( 3), ϱ is the partial derivative of the loss function C with respect to the transformation φ, which contains a ∇(Fm • φ) term, which is the backward transform of the grid_sample operator in PyTorch. Since this operation is not implemented using PyTorch primitives, a backward pass for the gradient operation does not exist in PyTorch. We use the gridsample_grad2 library [86] to compute the gradients of the backward pass of the grid_sample operator, used in Eq. (3) . All experiments are performed on a single NVIDIA A6000 GPU.Fig. 8 shows the loss curves for the toy dataset described in Section 4.1. An image-based optimization algorithm would correspond to the green curve being a flat line at 1 due to the flat landscape of the intensity-based loss function.Table 4 shows the quantitative results of our method for out-of-distribution performance on the IBSR18, CUMC12, and LPBA40 datasets. In 9 out of 10 cases, DIO demonstrates the best accuracy with fairly lower standard deviations, highlighting the robustness of the model. DIO therefore serves as a strong candidate for out-of-distribution performance, and can be used in a variety of settings where the training and test distributions differ.We consider four brain MRI datasets in this paper: OASIS dataset for in-distribution performance, and LPBA40, IBSR18, and CUMC12 datasets for out-of-distribution performance [85, 1, 53, 62] .More details about the datasets are provided below.• OASIS. Initialize i = 0 7: Initialize I l f , I l m = downsample(I f , s l ), downsample(Im, s l ) 8:while i < T l do 9: L i = C(I l f , I l m • φ i ) 10: Compute ∇φL 11: Update φ (i+1) = Optimize(φ i , ∇φL i ) ▷ OptimizationInput: Fixed features F f = [F 1 f , F 2 f . . . F n f ], Moving features F f = [F 1 f , F 2 f . . . F n f ] 2: Scales [s 1 ,Initialize i = 0 8: Initialize I l f , I l m = F l f , F l m 9:while i < T l do 10:L i = C(I l f , I l m • φ i ) 11:Compute ∇φLUpdate φ (i+1) = Optimize(φ i , ∇φL i ) ▷ Optimization algorithm Algorithm 2 shows our algorithm (red text highlights differences compared to Algorithm 1) that uses a separate scale-space feature at each level. Unlike classical methods, the scale-space feature can capture different discriminative features at each level to maximize label alignment and the multi-scale nature helps avoid local minima.• LPBA40. 40 brain images and their labels are used to construct the LONI Probabilistic Brain Atlas (LPBA40) dataset at the Laboratory of Neuroimaging (LONI) at UCLA [85] . All volumes are preprocessed according to LONI protocols to produce skull-stripped volumes. These volumes are aligned to the MNI305 atlas -this is relevant since existing DLIR methods may be biased towards Our feature network recovers dense multi-scale features (see Fig. 2 ) which allows all subsets to be registered with >0.99 Dice score. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper.• The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.• The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: The paper advocates a novel differentiable implicit optimization layer for image registration networks. While all reviewers agree that this is a quite interesting idea, major issues remain even after the rebuttal and discussion phase. I agree with the overall consensus that adjusting the current manuscript to reflect the required changes (e.g., inclusion of omissions) would require a major revision, which cannot be guaranteed within the NeurIPS review cycle. For that reason, I am recommending rejection at this point, but I also would like to encourage the authors to continue this work and revise accordingly for a resubmission.

Reviews: review: paper_summary: paper_summary: The authors introduce the idea of implicit optimization, and coupled feature extraction for images, to achieve robust image registration. I liked the overall idea and was eager to gain insight into how implicit optimization
strength_weakness: strength_weakness: Unfortunately, I found the rest of the paper (beyond the core idea) lacking and having several weaknesses. 

Importantly, the authors mischaracterize important relevant literature and conceptual ideas, that I think are incorrectly described in the motivation, and not compared to properly in the experiments. This makes it challenging to assess and gain insights from the contribution.

Incorrect characterization: the authors say a few statements that to me seem directly false (and are important to the paper). For example, on line 50 they say "Moreover, design decisions like sparse keypoint learning for affine registration [103, …] do not facilitate dense deformable registration" -- they repeat this in several parts of the paper (e.g. line 107). This is wrong -- their first citation, for example, 103, uses keypoints for *deformable* registration (along with affine). A few lines lower they say "Current DLIR methods are not robust to minor domain shift like varying anisotropy and voxel resolutions, different image acquisition and preprocessing protocols [62, 53, 70, 43]." -- which is incorrect and not supported by the citations. First, citations 62, 53, 70 are from before 2012 and do not discuss DLIR methods at all (but just general registration), whereas citation 43 *is explicitly tackling and achieves robustness to domain shift*. It may well be that the authors' method does better (I am not sure, see below), but the claim is incorrect. Many other papers tackle distribution shift in DL registration -- see Mok et al, 2023. Another crucial omission is related to the authors' claim that DL methods may not output local minima results -- which is true, but plenty of works propose to take the output of neural networks and perform a bit of instance-specific optimization of the resulting field to get to that local minima -- essentially 'fine-tuning' the field for a couple of seconds (e.g. VoxeMorph TMI 2019, but plenty of other methods as well after that for example from Matthias Heinrich's group). This is crucial to the current paper, since the proposed method essentially does the same thing at inference -- runs a forward neural network (albeit just as a feature extractor), and then performs an optimization for the image pair -- and while it's done differently (and more elegant in some sense) than the existing literature, these approaches are very related. Overall, I was excited about the method but overall found the motivation/related-works either misleading or lacking rigor -- perhaps the authors are simply not aware of the abilities of existing literature mentioned above, but this does  limit the novelty and insight substantially

In the experiments, it seems to me that some obvious results are missing:
. I am not sure why methods used in Figure 3 (e.g. [43]) are missing from Table 1. It seems like a crucial comparison. Deformable KeyMorph [103] is missing from the whole experiments section, and is close to the existing method in that it separates the feature extraction (there via keypoints, but using a parallel net) from the optimization. Training keymorph on oasis and testing it seems like an important comparison if we are to extract some insights into how to decisions in the proposed method (the feature extractor and the implicit optimization) improve our insights in the field.
. Overall the results in Table 1 do not seem impressive -- comparable at best with existing methods. This is totally okay in my book, if the authors are able to communicate other interesting insights. Unfortunately, I do not believe this is the case.
. In the domain-shift section the authors show that their method tens to outperform the DL methods. However, their method gets the benefit of doing instance-specific optimization (the proposed layer) after feature extraction, at the cost of some GPU work for each pair. This is what instance-specific optimization does at the end of DL methods (as discussed below), which was employed in several papers, but this is not included in the comparison! This is a peculiar omission to me -- it should be included for completeness, but importantly it is also crucial to understand whether the proposed method behaves differently -- perhaps there is some advantage, in several situations, to the proposed method, or perhaps it offers more guarantees, etc -- we simply don't know. Minor: it would also be interesting to understand what are the limits of this domain shift of this model -- does it generalize to more substantial variations in modality, or 7mm slice spacing found in clinical sequences?
. I also find the claim that DL methods do not work without crops peculiar - most DL methods are convolutional and hence not size specific, and some (e.g. SynthMorph, which the authors refer to here) does not even require both images be of the same size.
. Since one of the contributions of the paper is the parallel feature extraction (to be used with the optimizer), it seems to me that it would be an important ablation to take the features of some robust method (does not have to be registration, even a domain-shift-robust segmentation network will do, or a 'robust foundation model', etc) and see if that can be combined with an optimizer. This would help provide insights if the formulation of the proposed model and the end-to-end training is useful. 
. The claim in 4.4 is also missing some reaosnable comparison -- would it not be possible to take the displacement field of any other DL method, initialize your favorite parametrization (freeform, diffeomorphic, etc) with that, and run the (any) optimizer? It seems like this is easily doable and a reasonable comparison?
questions: questions: Please see comments above. It would be important to understand why the omissions in the motivation and experiments happened -- did the authors not know some of the keypoint based methods are deformable? Were they not aware of instance-specific optimization after the output of a DL method (which is discussed in several papers)? etc.
limitations: limitations: Yes the paper has a limitations section. It would be nice if the authors could comment on how this method can be used on CPUs -- by far the standard hardware available to non-ML users (neuroscientists, clinicians, etc).
paper_summary: paper_summary: This paper introduced DIO, a differentiable implicit optimization layer to a registration network that aimed to bridge the gap of classical-learning-based image registration, considering the incorporation of weak supervision like anatomical landmarks into the learned features. The authors decoupled feature learning and optimization and trained a deep network to predict multi-scale dense features registered through a black box iterative optimization solver.
strength_weakness: strength_weakness: **W1. Possibilities of getting artifacts from different voxel sizes.** How do different voxel sizes ensure that the velocity or transformation field is differentiable or invertible? I believe this approach might introduce artifacts and lose fine details when propagating the source image to match the target image. And how processing the image features independently will preserve the diffeomorphism property in generating the transformation field? Can the image matching term efficiently capture the intensity differences likewise treating the input images as pairwise? Overall, treating the input images separately from the feature extractor raises several questions regarding the credibility of the transformation field $\phi^*$. 

**W2. Motivation for using multi-scale optimization.** I found the motivation of using multi-scale is somewhat underdeveloped. What is the rationale behind using such kind of optimization considering different source/target image features? 


**W3. Applicabilities of learned multi-scale dense features from sparse images.** In Sec. 4 the authors tried to show that DIO learned interpretable dense features and compared to the classical methods DIO preserved the gradient in the loss function. On the other hand, the authors also discussed that deep networks recovered affine transform with $~90$% overlap. I wonder what is the advantage of capturing multi-scale dense features compared to existing DLIR methods such as VoxelMorph, TransMorph, etc. 

**W4. Experimental supports.** Though the authors are getting comparable performance in image registration (Tab. 1), they are achieving improved results in testing out of domain/distribution datasets. However, the authors might want to show their model's performance without adapting their proposed multi-scale optimization. Basically, is the optimization scheme or the multi-scale features helping the complete registration model in achieving those bits of improvements? and the important question is why? Two interesting sets of experiments that validate the domain shift hypothesis can be the following - 
*(i) train on some of the other datasets (excluding OASIS) and test on the rest,* and 
*(ii) train on multiple datasets, including OASIS, and test on the rest.* 


Overall, I appreciate the authors for working in the domain of image registration which is very relevant as well as important in the medical imaging domain. However, the current version of the manuscript lacks some important experimental justification and further experiments. With that being said, the current version of the manuscript is under the threshold of acceptance. However, I am open to reconsidering the initial rating if the above concerns are adequately justified.
questions: questions: How do decoupling feature learning and optimization blocks help the model in getting improved performance over existing DLIR methods?

I tried to summarize all the findings, concerns, and questions in the Weaknesses section.
limitations: limitations: Limitations have been discussed in the paper.
paper_summary: paper_summary: This paper presents a novel image registration framework that aims to bridge the gap between classical and learning-based approaches. It incorporates fidelity optimization directly into the neural network as a layer. The framework employs end-to-end implicit differentiation through an iterative optimization solver, ensuring that the features learned are both registration and label-aware. Additionally, the warp functions derived are guaranteed to represent local minima of the registration objective within the feature space. The authors report that this framework performs exceptionally well on in-domain datasets and remains robust against domain shifts, such as anisotropy and variations in intensity profiles. Furthermore, the framework is designed to allow seamless switching between different transformation representations at test time without the need for retraining.
strength_weakness: strength_weakness: 1. A major limitation is the framework's registration accuracy, as measured by the Dice score, which does not demonstrate clear advantages over neural-network-only methods. For instance, in Table 1 on the OASIS dataset, the proposed LKU-Net variant achieves a Dice score similar to that of TransMorph-Large. This calls into question the practical benefit of integrating classical optimization into the network.

2. The absence of specific smoothness measurements in comparison with other methods, such as the percentage of negative Jacobian determinants (||J||<0) or the standard deviation of the logarithm of the Jacobian determinant, is a significant oversight. Without these metrics, it's unclear whether the observed increase in Dice score represents a genuine improvement or merely a trade-off with the deformation field's smoothness.

3. The paper lacks clarity in the reproducibility of results compared to other methods. For example, the learn2reg OASIS leaderboard indicates that LKU-Net can achieve a Dice score of 88.5 on the OASIS dataset without explicitly optimizing for smoothness. Similarly, TransMorph scores 88.5. Additionally, in Table 4, the performance of LKU-Net with Dice supervision is paradoxically worse than without, which is counterintuitive and raises concerns about the implementation fidelity and methodological consistency.

4. While the introduction of an optimization layer with a fidelity function is a key contribution, the paper falls short in comparing its method to other registration methods that also combine learning with optimization. This lack of comparative analysis leaves unanswered questions regarding the true effectiveness and novelty of the proposed method compared to existing approaches.
questions: questions: 1 Could the authors clarify the primary benefits of their framework compared to existing learning frameworks that incorporate optimization, such as ConvexAdam [1], SAMConvex [2], PDD-Net [3], and PRIVATE [4]? Specific examples or quantitative comparisons highlighting how the framework outperforms these methods would be helpful.

2 The paper claims that, "For the first time, our method allows switching between arbitrary transformation representations (from free-form to diffeomorphic) at test time with zero retraining." Could the authors elaborate on why this capability is a significant advantage over other methods? For example, registration networks that separate feature extraction from deformation field estimation using multi-level image pyramids might achieve similar flexibility. What makes the framework distinct or superior in this context?

3 The paper mentions that anatomical landmarks can be incorporated into the framework. Are these landmarks based on image segmentation or keypoints, or do they refer to a more general concept of landmarks? If possible, could the authors demonstrate or provide insights on how these keypoints or landmarks are integrated within the optimization loop? 

[1] Siebert, Hanna, Lasse Hansen, and Mattias P. Heinrich. "Fast 3D registration with accurate optimisation and little learning for Learn2Reg 2021." International Conference on Medical Image Computing and Computer-Assisted Intervention. Cham: Springer International Publishing, 2021.

[2] Li, Zi, et al. "Samconvex: Fast discrete optimization for ct registration using self-supervised anatomical embedding and correlation pyramid." International Conference on Medical Image Computing and Computer-Assisted Intervention. Cham: Springer Nature Switzerland, 2023.

[3] Heinrich, Mattias P. "Closing the gap between deep and conventional image registration using probabilistic dense displacement networks." Medical Image Computing and Computer Assisted Intervention–MICCAI 2019: 22nd International Conference, Shenzhen, China, October 13–17, 2019, Proceedings, Part VI 22. Springer International Publishing, 2019.

[4] J. Hu, W. Gan, Z. Sun, H. An, and U. S. Kamilov. A Plug-and-Play Image Registration Network, Mar. 2024. arXiv:2310.04297
limitations: limitations: The paper commendably integrates a fidelity function with an optimization layer into a neural network, but there are several limitations:
1. The optimization layer has not demonstrated significant performance improvements compared to traditional neural network methods on OASIS leaderboard, and the implementation issues in the other datasets/methods raises concerns about the completeness of the result. Clarifying scenarios where this integration is advantageous could enhance the paper's impact.
2. The absence of quantitative smoothness metrics for registration performance is a critical gap. Introducing these metrics would strengthen the comparative analysis with existing methods.
3. The paper lacks a comprehensive comparison with major competitors combines optimization with learning. Expanding this analysis would provide clearer insights into the framework's unique contributions.
4. Concerns about implementation and reproducibility persist.

Comments: comment: The authors' rebuttal has addressed some concerns, showing the paper's potential for excellence. 
However, in its current form, it is not yet suitable for publication at NeurIPS. 
The reviewer encourages the authors to conduct more comprehensive experiments, particularly those involving fidelity loss in the post-optimization phases, rather than simply asserting differences from existing methods.
I read the rebuttal carefully and thank the authors.

However, I believe the various omissions require a new submission, since they would require a new revision. They are simply too central to the work, and cannot be explained in a single rebuttal paragraph -- in a normal journal submission, this would be a major revision.  For example, the KM omission requires a bit more than a simple run through and a 2x2 table -- this assumption that there is no related method is repeated throughout the paper and used to justify decisions. The KM method would need to be more thoroughly tested. Similarly, while IO is indeed not exactly the current method, the entire paper ignores IO as if it doesn't exist -- when in fact the difference is subtle at best. This not only requires reframing, but rigorous experimentation. This line, for example, i think substantially weakens the paper and I would argue is an important omission, still: "For simplicity, we do not perform IO in the intensity space for any method."

For a future submission, some more minor comments: 
- the citations example was more minor, but please note that there are several statements back to back like this -- where the claim is false and the citations used (that make it look like you are using related works to back up a claim) do not support the claim. I don't understand how the authors stand by the truthfulness of the (in my opinion over-reaching) statement -- there are definitely DLIR methods that are robust to domain shift, or at least some types of domain shift. In my opinion, Figure 3 does *not* support the statement -- yes the current method might improve on some methods, but only marginally -- making the claim way too strong. 
- but please note that while it's true that most DLIR methods concatenate the inputs in the architecture, some available implementations take in different-sized images, resample to the same (highest) voxel size to be isotropic, pad, and only then concatenate.  I believe SM does this, but am not an expert. Either way, it's trivial to do. This enables the method to work with differently-sized inputs. I am not using this point in my decision, though -- just a note for your submission.

Unfortunately, overall I believe the core omissions, which still exist (e.g. see the IO comment), are central to the work and would require a new submission to bre properly reviewed. I believe my original score stands. I do wish the authors good luck with this, I think there is something good here to be communicated, just not in its current form.