[Paper Review] Learning Fourier-Constrained Diffusion Bridges for MRI Reconstruction

Muhammad Usama Mirza, Onat Dalmaz+6 more|arXiv (Cornell University)|Aug 2, 2023

Advanced Neuroimaging Techniques and Applications11 citations

TL;DR

The paper introduces Fourier-Constrained Diffusion Bridges (FDB) as a diffusion-bridge prior that maps between moderately undersampled and fully sampled MRI data using stochastic Fourier frequency removal, improving reconstruction over prior methods.

ABSTRACT

Deep generative models have gained recent traction in accelerated MRI reconstruction. Diffusion priors are particularly promising given their representational fidelity. Instead of the target transformation from undersampled to fully-sampled data required for MRI reconstruction, common diffusion priors are trained to learn a task-agnostic transformation from an asymptotic start-point of Gaussian noise onto the finite end-point of fully-sampled data. During inference, data-consistency projections are injected in between reverse diffusion steps to reach a compromise solution within the span of both the trained diffusion prior and the imaging operator for an accelerated MRI acquisition. Unfortunately, performance losses can occur due to the discrepancy between target and learned transformations given the asymptotic normality assumption in diffusion priors. To address this discrepancy, here we introduce a novel Fourier-constrained diffusion bridge (FDB) for MRI reconstruction that transforms between a finite start-point of moderately undersampled data and an end-point of fully-sampled data. We derive the theoretical formulation of FDB as a generalized diffusion process based on a stochastic degradation operator that performs random spatial-frequency removal. We propose an enhanced sampling algorithm with a learned correction term for soft dealiasing across reverse diffusion steps. Demonstrations on brain MRI indicate that FDB outperforms state-of-the-art methods including non-diffusion and diffusion priors.

Motivation & Objective

• Motivate improved MRI reconstruction with task-agnostic priors that generalize across sampling patterns and datasets.
• Address mismatches between diffusion priors (Gaussian start-points) and MRI degradation by proposing a finite, stochastic start-point diffusion bridge.
• Develop a forward degradation process in k-space that progressively removes frequencies in peripheral-to-central order.
• Derive training objectives and an enhanced sampling algorithm with a learned correction term for soft dealiasing.

Proposed method

• Define a generalized diffusion process where forward steps remove random k-space frequencies in a peripheral-to-central order.
• Use a finite start-point X_Tf corresponding to R'-fold undersampling, with X0 representing fully-sampled Fourier data.
• Train a recovery operator G_theta via L_FDB-ub to predict x0 from x_t, reflecting a diffusion-like denoising in image domain.
• Interleave reverse diffusion steps with data-consistency projections, imputing frequency components with a cumulative mask C_t.
• Introduce a sampling correction term with weight w_t to achieve soft dealiasing, where w_t is learned from energy differences in Fourier space.
• Provide an algorithm (with equations 15-18) that governs forward degradation, reverse updates, and data fidelity steps.

Experimental results

Research questions

• RQ1 Can Fourier-constrained diffusion bridges map a finite-start, moderately undersampled data distribution to fully-sampled MRI data more effectively than traditional diffusion priors?
• RQ2 Does introducing stochastic, frequency-based forward degradation and a learned soft-dealiasing correction improve MRI reconstruction across varying sampling patterns and datasets?
• RQ3 How does the FDB prior perform relative to task-specific, diffusion-based, and diffusion-bridge methods under within-domain and cross-domain conditions?
• RQ4 What is the impact of the learned correction term on soft dealiasing and reconstruction quality across different acceleration rates?

Key findings

• FDB outperforms state-of-the-art methods including non-diffusion priors, diffusion priors, and diffusion bridges in within-domain MRI reconstruction tasks.
• In IXI (R=4,8) and fastMRI (R=4,8), FDB achieves highest PSNR and SSIM across contrasts compared to LORAKS, D5C5, rGAN, DDPM, CDiffMR, I2SB, and DB blur.
• Across datasets and sampling patterns, FDB demonstrates strong cross-domain generalization, retaining superior performance under shifts in sampling density (2D vs 1D) and dataset (multi-coil vs single-coil).
• On average, FDB improves over LORAKS by ~5.6 dB PSNR and ~30.1% SSIM in cross-domain tests, and over other baselines by multiple dB in PSNR and several percentage points in SSIM.
• The learned weighting schedule w_t for the correction term follows an exponential trend with t, supporting effective soft dealiasing during reverse diffusion.