On-the-Fly Guidance Training for Medical Image Registration

TL;DR

Medical image registration faces label scarcity and the trade-off between weakly-supervised and unsupervised approaches. OFG introduces a two-stage, plug-and-play training framework that generates instance-specific pseudo-ground truth on-the-fly via a differentiable optimizer, supervising the predicted deformation with $L_{ofg}$ and refining deformations through $\phi_{opt}$ without affecting inference speed. The framework employs a differentiable optimizer with energy $E_{opt}(I_f, I_m, \phi) = NCC(I_f, I_m \circ \phi) + \sum_p ||\nabla \phi(p)||^2$ and uses $L_{ofg} = \frac{1}{n} \sum (\phi_{pre} - \phi_{opt})^2$ for training, enabling incremental guidance and self-improving model-optimizer dynamics. Extensive experiments on brain MRI datasets (IXI, OASIS, LPBA40) and Abdomen CT demonstrate consistent improvements over baselines across multiple models, with enhanced deformation smoothness and substantial reductions in non-diffeomorphic regions, while incurring modest training overhead and no inference burden. The code is publicly available, facilitating broad adoption as a drop-in improvement for learning-based registration methods.

Abstract

This study introduces a novel On-the-Fly Guidance (OFG) training framework for enhancing existing learning-based image registration models, addressing the limitations of weakly-supervised and unsupervised methods. Weakly-supervised methods struggle due to the scarcity of labeled data, and unsupervised methods directly depend on image similarity metrics for accuracy. Our method proposes a supervised fashion for training registration models, without the need for any labeled data. OFG generates pseudo-ground truth during training by refining deformation predictions with a differentiable optimizer, enabling direct supervised learning. OFG optimizes deformation predictions efficiently, improving the performance of registration models without sacrificing inference speed. Our method is tested across several benchmark datasets and leading models, it significantly enhanced performance, providing a plug-and-play solution for training learning-based registration models. Code available at: https://github.com/cilix-ai/on-the-fly-guidance

Figures (10)

• Figure 1: The overall structure of the proposed framework. It has two parts, the prediction stage (a), and the optimization stage (b). The framework uses the idea of on-the-fly guidance to integrate the optimizer into the training process. The optimizer will iteratively refine the deformation field predicted by the registration model (for $n$ steps), and the derived optimized deformation field will then be used as pseudo ground truth to train the registration model.
• Figure 1: Ablation on optimization steps (TransMorph on IXI). Results show that 5 to 10 steps offer optimal balance, with no notable benefits from exceeding this range. Thus, we recommend to use 5 to 10 steps.
• Figure 2: Visualization of registration results on LPBA40 lpba. Demo randomly extracted from the comparison results between baseline TransMorph, VoxelMorph (row 2) and their respective model trained with OFG (row 1). OFG shows improved smoothness.
• Figure 2: The optimizer can quickly and effectively refine the deformation field even deformation from models with random initialization (left) or random parameters (right). With DSC increasing from 0.4260 to 0.5436 (left), and 0.4260 to 0.5101 (right).
• Figure 3: Visualization comparing training progress and validation DSC on LPBA40 across models. Self-training uses pre-trained network deformation fields as pseudo labels; optimized self-training enhances this with extra optimization steps. Our method achieves the best outcome, with self-training lagging due to convergence complexities.