# FLAT: FLow-Aligned Training of Unrolled Networks for MRI Reconstruction

arXiv: https://arxiv.org/abs/2512.03020 · HTML: https://mlsanity.com/paper/2512.03020/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Medical imaging and segmentation, one of 14 directions in Diffusion models.
- Builds on 31 of its 48 references that are in the map.
- First posted on arXiv 2 Dec 2025.

## Placed in

- Diffusion models / [Medical imaging and segmentation](https://mlsanity.com/diffusion/medical-imaging-and-segmentation/)

## Builds on (most cited first, 30 of 31)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [fastMRI: An Open Dataset and Benchmarks for Accelerated MRI](https://mlsanity.com/paper/1811.08839/) (2018)
- [MoDL: Model Based Deep Learning Architecture for Inverse Problems](https://mlsanity.com/paper/1712.02862/) (2017)
- [End-to-End Variational Networks for Accelerated MRI Reconstruction](https://mlsanity.com/paper/2004.06688/) (2020)
- [The probability flow ODE is provably fast](https://mlsanity.com/paper/2305.11798/) (2023)
- [Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems](https://mlsanity.com/paper/2303.05754/) (2023)
- [Adaptive Diffusion Priors for Accelerated MRI Reconstruction](https://mlsanity.com/paper/2207.05876/) (2022)
- [PixelFlow: Pixel-Space Generative Models with Flow](https://mlsanity.com/paper/2504.07963/) (2025)
- [Compressed Sensing MRI Reconstruction using a Generative Adversarial Network with a Cyclic Loss](https://mlsanity.com/paper/1709.00753/) (2017)
- [A Deep Cascade of Convolutional Neural Networks for MR Image Reconstruction](https://mlsanity.com/paper/1703.00555/) (2017)
- [Measurement-conditioned Denoising Diffusion Probabilistic Model for Under-sampled Medical Image Reconstruction](https://mlsanity.com/paper/2203.03623/) (2022)
- [Swin Transformer for Fast MRI](https://mlsanity.com/paper/2201.03230/) (2022)
- [ReconFormer: Accelerated MRI Reconstruction Using Recurrent Transformer](https://mlsanity.com/paper/2201.09376/) (2022)
- [High-Frequency Space Diffusion Models for Accelerated MRI](https://mlsanity.com/paper/2208.05481/) (2022)
- [STARFlow: Scaling Latent Normalizing Flows for High-resolution Image Synthesis](https://mlsanity.com/paper/2506.06276/) (2025)
- [DSFormer: A Dual-domain Self-supervised Transformer for Accelerated Multi-contrast MRI Reconstruction](https://mlsanity.com/paper/2201.10776/) (2022)
- [Flow Matching for Medical Image Synthesis: Bridging the Gap Between Speed and Quality](https://mlsanity.com/paper/2503.00266/) (2025)
- [HiFlow: Training-free High-Resolution Image Generation with Flow-Aligned Guidance](https://mlsanity.com/paper/2504.06232/) (2025)
- [Deform-Mamba Network for MRI Super-Resolution](https://mlsanity.com/paper/2407.05969/) (2024)
- [MMR-Mamba: Multi-Modal MRI Reconstruction with Mamba and Spatial-Frequency Information Fusion](https://mlsanity.com/paper/2406.18950/) (2024)
- [Reversing Flow for Image Restoration](https://mlsanity.com/paper/2506.16961/) (2025)
- [MambaRecon: MRI Reconstruction with Structured State Space Models](https://mlsanity.com/paper/2409.12401/) (2024)
- [DH-Mamba: Exploring Dual-domain Hierarchical State Space Models for MRI Reconstruction](https://mlsanity.com/paper/2501.08163/) (2025)
- [Unsupervised Accelerated MRI Reconstruction via Ground-Truth-Free Flow Matching](https://mlsanity.com/paper/2502.17174/) (2025)

## Abstract (arXiv)

Unrolled networks are widely used in Magnetic Resonance Imaging (MRI) reconstruction for their efficiency. Structured as a series of neural network stages (or cascades), an unrolled network takes a low-quality input and passes it sequentially through each stage to iteratively refine the reconstruction. However, unrolled networks typically exhibit unstable output quality across cascades, resulting in sub-optimal final reconstruction results. In this work, we address this inherent limitation of unrolled networks, drawing inspiration from recent Flow Matching paradigm. We first theoretically show that unrolled networks can be viewed as discretizations of approximate conditional probability flows. This connection shows that unrolled networks and Flow Matching are analogous in MRI reconstruction. Building upon this insight, we propose FLow-Aligned Training (FLAT), which (1) derives important cascade parameters from the Flow Matching discretization; and (2) aligns intermediate reconstructions with the ideal Flow Matching trajectory to improve cascade iteration stability and convergence. Experiments on three MRI datasets show that FLAT results in a stable trajectory across sub-networks, improving the quality of the final reconstruction.