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

2024-11-05

arXiv.org

AI Summary

This work addresses three key challenges in low-dimensional latent-space undersampled MRI reconstruction: (1) difficulty in controlling medical fidelity, (2) domain mismatch between MR physics and natural image priors, and (3) lack of latent-space data consistency. We propose a novel Latent Diffusion Prior Model (LDPM) framework. Our contributions are: (1) the first MRI-customized variational autoencoder (MR-VAE), explicitly encoding MR physical constraints; (2) a two-stage latent-space sampler incorporating a differentiable latent-space data consistency module; and (3) a sketch-guided two-stage reconstruction paradigm jointly optimizing anatomical fidelity and perceptual quality. Evaluated on fastMRI, our method achieves a 3.92 dB PSNR improvement over prior art, setting a new state-of-the-art. It significantly outperforms pixel-level diffusion-based methods while maintaining high computational efficiency and robustness across diverse undersampling rates.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Reasoning under Uncertainty: Graphical Models

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Large language models for search
• Economics, Online Markets and Human Computation: LLM based quality controls for crowd work

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 datasetcite{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.

Problem

• Reduces computational cost in MRI reconstruction using latent diffusion models.
• Addresses domain gap between natural images and MR physics.
• Ensures medical fidelity and data consistency in latent space.

Innovation

• Sketch-guided pipeline for balanced reconstruction
• MRI-optimized VAE improves PSNR significantly
• Dual-Stage Sampler ensures high-fidelity in latent space