  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.
