Flow Matching-enabled Test-Time Refinement for Unsupervised Cardiac MR Registration

TL;DR

FlowReg is a fast, efficient, and unsupervised cardiac MR registration method that achieves high accuracy with minimal inference steps and no need for segmentation labels, outperforming existing techniques.

Contribution

The paper introduces FlowReg, a novel flow-matching framework with warmup-reflow training that enables rapid and accurate unsupervised cardiac MR registration without pre-trained models.

Findings

Outperforms state-of-the-art on five of six tasks

Reduces LVEF estimation error by 2.58 percentage points

Achieves strong registration in as few as two steps

Abstract

Diffusion-based unsupervised image registration has been explored for cardiac cine MR, but expensive multi-step inference limits practical use. We propose FlowReg, a flow-matching framework in displacement field space that achieves strong registration in as few as two steps and supports further refinement with more steps. FlowReg uses warmup-reflow training: a single-step network first acts as a teacher, then a student learns to refine from arbitrary intermediate states, removing the need for a pre-trained model as in existing methods. An Initial Guess strategy feeds back the model prediction as the next starting point, improving refinement from step two onward. On ACDC and MM2 across six tasks (including cross-dataset generalization), FlowReg outperforms the state of the art on five tasks (+0.6% mean Dice score on average), with the largest gain in the left ventricle (+1.09%), and reduces LVEF estimation error on all six tasks (-2.58 percentage points), using only 0.7% extra parameters and no segmentation labels. Code is available at https://github.com/mathpluscode/FlowReg.