Unrolled Networks are Conditional Probability Flows in MRI Reconstruction

TL;DR

This work reframes MRI reconstruction as a flow-based problem by proving that unrolled networks are discrete realizations of energy-based conditional probability flow ODEs. It derives concrete implications for cascade scheduling and parameterization, and introduces Flow-Aligned Training (FLAT), which enforces ODE-consistent timesteps and velocity alignment for intermediate outputs. Across three public MRI datasets, FLAT achieves high reconstruction quality with up to 3x fewer iterations than diffusion-based methods and substantially greater stability than vanilla unrolled networks. The approach delivers a principled, faster alternative to diffusion models while retaining strong image fidelity, offering meaningful practical impact for real-time MRI applications.

Abstract

Magnetic Resonance Imaging (MRI) offers excellent soft-tissue contrast without ionizing radiation, but its long acquisition time limits clinical utility. Recent methods accelerate MRI by under-sampling $k$-space and reconstructing the resulting images using deep learning. Unrolled networks have been widely used for the reconstruction task due to their efficiency, but suffer from unstable evolving caused by freely-learnable parameters in intermediate steps. In contrast, diffusion models based on stochastic differential equations offer theoretical stability in both medical and natural image tasks but are computationally expensive. In this work, we introduce flow ODEs to MRI reconstruction by theoretically proving that unrolled networks are discrete implementations of conditional probability flow ODEs. This connection provides explicit formulations for parameters and clarifies how intermediate states should evolve. Building on this insight, we propose Flow-Aligned Training (FLAT), which derives unrolled parameters from the ODE discretization and aligns intermediate reconstructions with the ideal ODE trajectory to improve stability and convergence. Experiments on three MRI datasets show that FLAT achieves high-quality reconstructions with up to $3\times$ fewer iterations than diffusion-based generative models and significantly greater stability than unrolled networks.

Figures (7)

• Figure 1: (I) Illustrating the MRI reconstruction task: from an under-sampled, aliased input (a), the task is to recover the clean, fully-sampled image (b). (II) Comparison of reconstruction approaches: c) Vanilla unrolled networks achieve fast results with few iterations but suffer from unstable intermediate steps that degrade image quality. d) Diffusion models achieve high image quality but require many iterations, adversely impacting speed. e) Our FLAT, grounded in probability flow ODEs, achieves high reconstruction quality with a low iteration count compared with diffusion-based generative models, yielding both stability and speed.
• Figure 2: I) Vanilla unrolled networks vs. our Flow-Aligned Training (FLAT). Vanilla unrolled networks iteratively refine reconstructions step-by-step with supervision only at the final output. Our theory reformulates unrolling (orange) as a discretized flow ODE (blue); in FLAT, each step predicts a velocity field, with intermediate supervision that aligns predicted and ideal velocities. II) Trajectory comparison. Without intermediate supervision, vanilla unrolled networks exhibit unstable (oscillatory) trajectories that "under-run" or "overshoot" the target. FLAT supervises intermediate steps to follow stable, straight-line paths guided by flow ODE theory.
• Figure 3: Qualitative results on Brainweb (rows 1-2), MRBrainS13 (rows 3-4) and fastMRI knee dataset (rows 5-6). For each dataset, the first row shows reconstructions; the second row shows the squared-error map relative to the ground truth to visualize the error magnitude.
• Figure 4: PSNR curves for 12-step iterations.
• Figure 5: PSNR curves for timestep $t=1 \to 0$.

Theorems & Definitions (2)

• Theorem 1
• proof