Unrolled Networks are Conditional Probability Flows in MRI Reconstruction

2025-12-02

AI Summary

MRI accelerated reconstruction faces dual challenges: unstable evolution in unrolled networks and high computational cost of diffusion models. This work establishes, for the first time, the theoretical equivalence between unrolled networks and discrete implementations of conditional probability flow ordinary differential equations (ODEs). Building on this insight, we propose Flow-Aligned Training (FLAT), a novel framework that enforces trajectory alignment—constraining intermediate-layer state evolution to follow the underlying ODE solution path—to achieve stable and efficient reconstruction. FLAT integrates ODE discretization, diffusion-based modeling, and unrolled architecture design. Evaluated on three MRI datasets, it achieves high-fidelity reconstructions with convergence three times faster than standard diffusion models and significantly improved stability and accuracy over conventional unrolled networks. Our core contributions are: (i) establishing the formal theoretical connection between unrolled networks and probability flow ODEs, and (ii) introducing the first trajectory-aligned training paradigm for MRI reconstruction.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Unsupervised & Self-Supervised Learning
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

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 imes$ fewer iterations than diffusion-based generative models and significantly greater stability than unrolled networks.

Problem

• Improves MRI reconstruction stability in unrolled networks
• Connects unrolled networks to conditional probability flow ODEs
• Reduces computational cost compared to diffusion models

Innovation

• Connects unrolled networks to probability flow ODEs
• Proposes Flow-Aligned Training for parameter derivation
• Aligns intermediate reconstructions with ODE trajectory