[Paper Review] Self-Supervised MRI Reconstruction with Unrolled Diffusion Models

Yılmaz Korkmaz, Tolga Çukur+1 more|arXiv (Cornell University)|Jun 29, 2023

Advanced Neuroimaging Techniques and ApplicationsMedicine3 citations

TL;DR

This paper proposes SSDiffRecon, a self-supervised unrolled diffusion model for accelerated MRI reconstruction that combines cross-attention transformers with data-consistency layers to denoise undersampled k-space data. It achieves state-of-the-art image quality and speed by training solely on undersampled k-space data, outperforming both supervised and self-supervised baselines with only five reverse diffusion steps.

ABSTRACT

Magnetic Resonance Imaging (MRI) produces excellent soft tissue contrast, albeit it is an inherently slow imaging modality. Promising deep learning methods have recently been proposed to reconstruct accelerated MRI scans. However, existing methods still suffer from various limitations regarding image fidelity, contextual sensitivity, and reliance on fully-sampled acquisitions for model training. To comprehensively address these limitations, we propose a novel self-supervised deep reconstruction model, named Self-Supervised Diffusion Reconstruction (SSDiffRecon). SSDiffRecon expresses a conditional diffusion process as an unrolled architecture that interleaves cross-attention transformers for reverse diffusion steps with data-consistency blocks for physics-driven processing. Unlike recent diffusion methods for MRI reconstruction, a self-supervision strategy is adopted to train SSDiffRecon using only undersampled k-space data. Comprehensive experiments on public brain MR datasets demonstrates the superiority of SSDiffRecon against state-of-the-art supervised, and self-supervised baselines in terms of reconstruction speed and quality. Implementation will be available at https://github.com/yilmazkorkmaz1/SSDiffRecon.

Motivation & Objective

• Address the limitations of existing MRI reconstruction methods, including poor contextual sensitivity, suboptimal data distribution modeling, and reliance on fully-sampled data for training.
• Overcome the dependency on fully-sampled acquisitions in supervised learning by introducing a self-supervision strategy based on masked k-space reconstruction.
• Improve image fidelity and reconstruction speed by integrating physics-driven data-consistency blocks within a diffusion-based unrolled architecture.
• Enhance contextual understanding and feature representation through cross-attention transformers in the reverse diffusion process.
• Demonstrate that a diffusion model can achieve competitive performance with minimal reverse sampling steps while avoiding the need for paired training data.

Proposed method

• Proposes an unrolled architecture that interleaves cross-attention transformer blocks for denoising with data-consistency projections for physical constraint enforcement during reverse diffusion.
• Employs a conditional diffusion process where the denoising network is trained via self-supervision by predicting masked k-space samples in undersampled acquisitions.
• Uses a mapper network with cross-attention to inject time-step and input-specific conditioning into the transformer-based denoiser.
• Integrates data-consistency layers after each denoising step to enforce fidelity to the acquired undersampled k-space data.
• Trains the model using a denoising loss that minimizes the difference between predicted and ground-truth noise in the diffusion process.
• Performs inference with only five reverse diffusion steps, significantly reducing reconstruction time compared to standard DDPMs.

Experimental results

Research questions

• RQ1 Can a self-supervised diffusion model achieve high-fidelity MRI reconstruction without requiring paired fully-sampled data?
• RQ2 How does integrating data-consistency layers into a transformer-based unrolled diffusion architecture affect reconstruction quality and speed?
• RQ3 To what extent does cross-attention attention improve contextual sensitivity and detail preservation in MRI reconstruction compared to CNN-based models?
• RQ4 Can a minimal number of reverse diffusion steps (e.g., five) yield competitive performance compared to standard DDPMs with thousands of steps?
• RQ5 How does the proposed self-supervision strategy based on masked k-space prediction compare to supervised and other self-supervised baselines in terms of robustness and generalization?

Key findings

• SSDiffRecon achieves a mean PSNR of 42.3 dB and SSIM of 99.3 on the IXI dataset at 4× acceleration, outperforming the second-best self-supervised method by 3.55 dB in PSNR and 1.96% in SSIM.
• On the fastMRI dataset, SSDiffRecon achieves 40.1 dB PSNR and 96.5 SSIM at 4× acceleration, outperforming supervised baselines by 0.4 dB in PSNR and 0.25% in SSIM.
• The model achieves state-of-the-art performance with only five reverse diffusion steps, significantly reducing inference time compared to standard DDPMs.
• Visual comparisons show that SSDiffRecon preserves high-frequency details better than other self-supervised methods, which often suffer from blurring or noise.
• Ablation studies confirm that both cross-attention transformers and data-consistency layers are critical for performance, with removal of either leading to significant degradation in PSNR and SSIM.
• The self-supervised training strategy enables performance on par with supervised methods while eliminating the need for paired fully-sampled data.