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

2026-03-01

AI Summary

This work addresses the high computational cost and multi-step inference of diffusion models in unsupervised cardiac MR image registration by proposing FlowReg, a flow-matching-based displacement field registration framework. FlowReg employs a warmup-reflow two-stage training strategy combined with an initial guess feedback mechanism, enabling high-quality registration in just two inference steps. It further supports test-time refinement from any intermediate state without requiring pretraining or segmentation labels. Evaluated on six tasks across the ACDC and MM2 datasets, FlowReg outperforms existing methods on five, achieving an average Dice score improvement of 0.6% (up to 1.09% for the left ventricle) and reducing left ventricular ejection fraction (LVEF) estimation error by 2.58 percentage points, with only a 0.7% increase in model parameters.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Calibration & Uncertainty Quantification
• Intelligent Robots: Multimodal Perception & Sensor Fusion

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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. Anonymized code is available at https://github.com/mathpluscode/FlowReg.

Problem

• unsupervised registration
• cardiac MR
• diffusion models
• test-time refinement
• computational efficiency

Innovation

• flow matching
• unsupervised registration
• test-time refinement
• warmup-reflow training
• cardiac MR