LDPM: Towards undersampled MRI reconstruction with MR-VAE and Latent Diffusion Prior

TL;DR

This paper introduces LDPM, a novel MRI reconstruction method using latent diffusion models with a sketch-guided pipeline, MRI-optimized VAE, and dual-stage sampler, achieving state-of-the-art results efficiently.

Contribution

The paper proposes a new latent diffusion prior-based framework for undersampled MRI reconstruction, addressing domain gap and fidelity control challenges with innovative modules.

Findings

Achieves approximately 3.92 dB PSNR improvement over SD-VAE.

Demonstrates state-of-the-art performance on fastMRI dataset.

Shows robustness and effectiveness of each module through ablation studies.

Abstract

Diffusion models, as powerful generative models, have found a wide range of applications and shown great potential in solving image reconstruction problems. Some works attempted to solve MRI reconstruction with diffusion models, but these methods operate directly in pixel space, leading to higher computational costs for optimization and inference. Latent diffusion models, pre-trained on natural images with rich visual priors, are expected to solve the high computational cost problem in MRI reconstruction by operating in a lower-dimensional latent space. However, direct application to MRI reconstruction faces three key challenges: (1) absence of explicit control mechanisms for medical fidelity, (2) domain gap between natural images and MR physics, and (3) undefined data consistency in latent space. To address these challenges, a novel Latent Diffusion Prior-based undersampled MRI reconstruction (LDPM) method is proposed. Our LDPM framework addresses these challenges by: (1) a sketch-guided pipeline with a two-step reconstruction strategy, which balances perceptual quality and anatomical fidelity, (2) an MRI-optimized VAE (MR-VAE), which achieves an improvement of approximately 3.92 dB in PSNR for undersampled MRI reconstruction compared to that with SD-VAE \cite{sd}, and (3) Dual-Stage Sampler, a modified version of spaced DDPM sampler, which enforces high-fidelity reconstruction in the latent space. Experiments on the fastMRI dataset\cite{fastmri} demonstrate the state-of-the-art performance of the proposed method and its robustness across various scenarios. The effectiveness of each module is also verified through ablation experiments.