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

Yunguan Fu; Wenjia Bai; Wen Yan; Matthew J Clarkson; Rhodri Huw Davies; Yipeng Hu

cs.CVMar 1, 2026

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

Authors: Yunguan Fu, Wenjia Bai, Wen Yan, Matthew J Clarkson, Rhodri Huw Davies, Yipeng Hu

Organizations: University College London · InstaDeep · Imperial College London

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.

Figures & tables

Figure 1: Overview of FlowReg. Top-left: example image pair with labels (labels are used only for evaluation). Top-right: warmup-reflow training; the teacher predicts a reference DDF from noise, and the student learns from interpolated intermediate states. Yellow and green lines indicate the gradient propagation during warm-up and reflow training phases, respectively. Bottom-left: standard inference, which updates the DDF from pure noise ( $t{=}0$ ) toward the prediction. Bottom-right: Initial Guess (IG), which replaces the noisy state after step one with the model prediction.

Figure 2

Method	Mean	Foreground	$\%\|J_{\psi}\|\leq 0$	$\text{STD}(\|J_{\psi}\|)$	Time (s)
ANTs Avants et al. (2008)	65.98 $\pm$ 12.25	83.70 $\pm$ 6.84	0.93 $\pm$ 3.01	0.08 $\pm$ 0.03
Demons Thirion (1998)	71.23 $\pm$ 9.73	83.43 $\pm$ 6.16	0.77 $\pm$ 0.42	0.40 $\pm$ 0.07
VoxelMorph Balakrishnan et al. (2019)	74.88 $\pm$ 8.69	86.08 $\pm$ 5.99	0.02 $\pm$ 0.03	0.14 $\pm$ 0.03
Fourier-Net Jia et al. (2023)	60.68 $\pm$ 12.09	83.09 $\pm$ 6.01	0.00 $\pm$ 0.00	0.03 $\pm$ 0.01
TransMorph Chen et al. (2022)	68.11 $\pm$ 11.70	82.87 $\pm$ 7.22	0.06 $\pm$ 0.01	0.15 $\pm$ 0.05
XMorpher Shi et al. (2022)	70.09 $\pm$ 9.53	84.96 $\pm$ 4.04	0.02 $\pm$ 0.04	0.13 $\pm$ 0.03

Table 1: Dice score on ACDC ES

\to

ED. Mean: average Dice score of RV, Myo, and LV. Foreground: Dice score on merged RV+Myo+LV.

\%|J_{\psi}|\leq 0

and

\mathrm{STD}(|J_{\psi}|)

measure deformation regularity (lower is better). Time: inference seconds per sample on one GPU. Bold indicates the best value per column.

Dataset	Method	LV	Myo	RV	Mean	FG	$p$
ACDC $\to$ ED	CMLP	89.59 $\pm$ 6.24	78.66 $\pm$ 5.22	83.39 $\pm$ 7.01	83.88 $\pm$ 4.50	91.60 $\pm$ 3.21	2e-4
ACDC $\to$ ED	Ours	90.49 $\pm$ 5.02	79.08 $\pm$ 4.60	84.46 $\pm$ 6.47	84.68 $\pm$ 3.80	92.15 $\pm$ 2.92	2e-4
ACDC $\to$ ES	CMLP	80.15 $\pm$ 12.57	79.58 $\pm$ 5.51	76.57 $\pm$ 8.97	78.77 $\pm$ 6.73	88.48 $\pm$ 4.66	1e-7
ACDC $\to$ ES	Ours	82.09 $\pm$ 11.24	80.70 $\pm$ 4.69	76.74 $\pm$ 8.96	79.84 $\pm$ 6.11	88.60 $\pm$ 4.58	1e-7
MM2 $\to$ ED	CMLP	82.01 $\pm$ 6.69	74.17 $\pm$ 5.01	89.79 $\pm$ 3.67	81.99 $\pm$ 3.87	89.18 $\pm$ 3.53	1e-6
MM2 $\to$ ED	Ours	83.85 $\pm$ 5.61	74.48 $\pm$ 4.89	90.23 $\pm$ 3.05	82.86 $\pm$ 3.31	90.11 $\pm$ 2.83	1e-6

Table 2: Task-wise Dice score comparison (mean

\pm

SD).

\to

MM2 denotes cross-dataset generalization. LV: left ventricle; Myo: myocardium; RV: right ventricle; Mean: average Dice score of LV/Myo/RV; FG: foreground Dice score on merged classes. Bold indicates the higher value within each task and metric.

p

: Bonferroni-corrected paired t-test on Mean Dice score (Ours

>

CMLP).

					Inference Steps
Method	SDE	Heun	IG	Guidance	1	2	5	10	20
CMLP					79.81	79.83	79.54	79.42	79.37
Ours		Y			78.92	79.76	80.12	80.24	80.30
	Y				78.92	79.62	80.17	80.33	80.38
	Y	Y			78.92	79.76	80.28	80.38	80.37
	Y	Y	Y		78.92	80.14	80.34	80.38	80.37

Table 3: Mean Dice score (%) under iterative inference and optimisation. Top: performance versus inference steps. Bottom: instance optimisation steps applied to the predicted DDF. Bold indicates the highest value in each column.

Dataset	Method	LVEF	RVEF	MT	$p_{\text{LV}}$	$p_{\text{RV}}$
ACDC $\to$ ED	CMLP	5.06 $\pm$ 3.78	14.83 $\pm$ 7.79	1.01 $\pm$ 1.18
ACDC $\to$ ED	Ours	4.00 $\pm$ 2.83	12.01 $\pm$ 6.92	0.86 $\pm$ 0.96	2e-4	1e-8
ACDC $\to$ ES	CMLP	10.98 $\pm$ 9.02	21.41 $\pm$ 10.61	1.60 $\pm$ 1.98
ACDC $\to$ ES	Ours	8.23 $\pm$ 7.39	20.31 $\pm$ 10.39	1.29 $\pm$ 1.73	2e-8	2e-6
MM2 $\to$ ED	CMLP	13.21 $\pm$ 7.99	4.87 $\pm$ 3.47	0.62 $\pm$ 0.55
MM2 $\to$ ED	Ours	8.74 $\pm$ 6.71	3.83 $\pm$ 3.00	0.64 $\pm$ 0.58	2e-20	5e-4

Table 4: Mean absolute errors for left/right ventricular ejection fraction (LVEF/RVEF, %) and myocardial thickness (MT, mm).

\to

MM2 denotes cross-dataset generalization. Bold indicates lower MAE.

p_{\text{LV}}

p_{\text{RV}}

: Bonferroni-corrected paired

t

-tests on LVEF/RVEF MAE (Ours

<

CMLP).

Explore similar work

Set-Based Groupwise Registration for Variable-Length, Variable-Contrast Cardiac MRI
May 11, 2026Yi Zhang, Yidong Zhao, Tijmen Toxopeus +3Medical Image Registration Magnetic Resonance Imaging
Department of Imaging Physics, Delft University of Technology, The Netherlands
Recursive Uncertainty-Gated Image Registration for Learning-based Algorithms
Sep 23, 2026Clara Rodrigo González, Oscar Bates, Fu Siong Ng +1Medical Imaging Medical Image Registration
Department of Bioengineering, Imperial College London, London, UK · Department of Earth Science and Engineering, Imperial College London, London, UK · National Heart and Lung Institute, Imperial College London, London, UK
One Frame, Full Heartbeat: ECG-Free 4D Cardiac Cine MRI Synthesis via Radial-Decomposed Flow Matching
Oct 6, 2026Shiyi Wang, Ruochen Sun, Xiang Li +2Medical Image Synthesis Cardiac MRI
Department of Electrical and Computer Engineering, Johns Hopkins University, Baltimore, MD, USA · Gordon Center for Medical Imaging, Department of Radiology, Massachusetts General Hospital and Harvard Medical School, Boston, MA, USA · Center for Advanced Medical Computing and Analysis, Department of Radiology, Massachusetts General Hospital and Harvard Medical School, Boston, MA, USA

May 11, 2026J/K move · Enter open · S save

Set-Based Groupwise Registration for Variable-Length, Variable-Contrast Cardiac MRI

Yi Zhang, Yidong Zhao, Tijmen Toxopeus, Maša Božić-Iven+2

Department of Imaging Physics, Delft University of Technology, The Netherlands

Medical Image Registration Magnetic Resonance Imaging Deformable Image Registration Cardiac MRI

Open paper AskPDF arXiv

Quantitative cardiac magnetic resonance imaging (MRI) enables non-invasive myocardial tissue characterization but relies on robust motion correction within these variable-length, variable-contrast image sequences. Groupwise registration, which simultaneously aligns all images, has shown greater robustness than pairwise registration for motion correction. However, current deep-learning-based groupwise registration methods cannot generalize across MRI sequences: the architecture typically encodes input data as a fixed-length channel stack, which rigidly couples network design to protocol-specific sequence length, input ordering, and contrast dynamics. At inference time, any change in imaging protocols will render the network unusable. In this work, we introduce \AnyTwoReg, a new set-based groupwise registration framework that takes a quantitative MRI sequence as an unordered set. This set formulation fundamentally decouples network design from sequence length and input ordering. By utilizing a shared encoder and correlation-guided feature aggregation, \AnyTwoReg constructs a permutation-invariant canonical reference for registration, and learns a permutation-equivariant mapping from images to deformation fields. Additionally, we extract contrast-insensitive image features from an existing foundation model to handle extreme contrast variations. Trained exclusively on a single public

T_1

mapping dataset (STONE, sequence length

L=11

), \AnyTwoReg generalizes to two unseen quantitative MRI datasets (MOLLI, ASL) with variable lengths (

L \in [11, 60]

) and different contrast dynamics. It achieves strong cross-protocol generalization in a zero-shot manner, and consistently improves downstream quantitative mapping quality. Notably, while designed for quantitative MRI sequences, our framework is directly applicable to Cine MRI sequences for inter-cardiac-phase registration.

Abstract

Explore similar work

Set-Based Groupwise Registration for Variable-Length, Variable-Contrast Cardiac MRI

Recursive Uncertainty-Gated Image Registration for Learning-based Algorithms

One Frame, Full Heartbeat: ECG-Free 4D Cardiac Cine MRI Synthesis via Radial-Decomposed Flow Matching