Adaptive Image Registration: A Hybrid Approach Integrating Deep Learning and Optimization Functions for Enhanced Precision

TL;DR

Air addresses the challenge of achieving accurate biomedical image registration by hybridizing learning-based deformation prediction with an adaptive, in-loop optimization stage. The method initializes optimization with learning outputs and focuses computational effort on harder image pairs through a decision module that adaptively selects iteration counts. Across multiple backbones and the IXI brain MRI dataset, Air delivers up to $1.6$ percentage point improvements in test Dice scores and smoother deformation fields, without increasing inference time. This modular framework enhances registration robustness and efficiency, with potential for deployment across diverse medical imaging tasks and datasets.

Abstract

Image registration has traditionally been done using two distinct approaches: learning based methods, relying on robust deep neural networks, and optimization-based methods, applying complex mathematical transformations to warp images accordingly. Of course, both paradigms offer advantages and disadvantages, and, in this work, we seek to combine their respective strengths into a single streamlined framework, using the outputs of the learning based method as initial parameters for optimization while prioritizing computational power for the image pairs that offer the greatest loss. Our investigations showed improvements of up to 1.6% in test data, while maintaining the same inference time, and a substantial 1.0% points performance gain in deformation field smoothness.

Figures (4)

• Figure 1: The proposed architecture for our method, Air. Divided into three main sections, Air integrates the stock learning model contained in the learning step and adds two new blocks, the decision and optimization steps. After running the outputs from the neural network into the decision module, if passed through to the optimization step, the optimizer module will loop through $\text{n}_{adp}$ or $\text{n}_{std}$ epochs (5 for $\text{n}_{std}$ and 20 for $\text{n}_{adp}$ in our experiments) and return the optimized output to the neural network for backpropagation calculation of the parameters.
• Figure 2: The inner workings of the decision module. $\text{n}_{adp}$ will be assigned if the randomly generated probability $p$ is greater than the current threshold $T$, or the value for $\mathcal{L}_\mathrm{ncc}$ is within the highest in $\Pi$. Otherwise, the standard minimum optimization value $\text{n}_{std}$ will be assigned to $\text{n}$.
• Figure 3: The inner workings of the optimizer module. For $n$ epochs, the Adam stochastic gradient descent optimizer will iteratively optimize the pair $\pi$ and return $\pi_\mathrm{opt}$
• Figure 4: Comparison of model outputs on patient-to-atlas registration on IXI for both deformation grid and field $\phi$ for all learning based methods and Air paired with its highest yielding model for the Dice Score Coefficient, TransMorph.