Deep Implicit Optimization enables Robust Learnable Features for Deformable Image Registration

TL;DR

This paper introduces a novel deep learning framework that integrates explicit optimization layers to improve deformable image registration, achieving robustness to domain shifts and enabling flexible transformation representations without retraining.

Contribution

It proposes a method that combines deep feature learning with implicit differentiation through an optimization solver, enhancing registration accuracy and flexibility over prior DLIR approaches.

Findings

Excellent in-domain registration performance

Robust to domain shifts like anisotropy and intensity variations

Allows switching transformation models at test time without retraining

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, existing DLIR methods forego many of the benefits and invariances of optimization methods. The lack of a task-specific inductive bias in DLIR methods leads to suboptimal performance, especially in the presence of domain shift. Our method aims to bridge this gap between statistical learning and optimization by explicitly 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, we explicitly exploit invariances of the correspondence matching problem induced by the optimization, while learning registration and label-aware features, and guaranteeing the warp functions to be a 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 arbitrary test-time regularization, which is not possible with existing DLIR methods.