Fast and Robust Diffusion Posterior Sampling for MR Image Reconstruction Using the Preconditioned Unadjusted Langevin Algorithm

TL;DR

The paper tackles slow, tunable sampling in diffusion-based Bayesian MRI reconstruction by introducing a preconditioned unadjusted Langevin algorithm (pULA) that uses an exact likelihood term across all diffusion scales. A theoretically grounded preconditioner M_t = (A^H A + σ_t^{-2} I)^{-1} enables larger updates in data-poor directions and smaller updates where the data strongly constrain the image, allowing a fixed step size and faster convergence. Empirical results on Cartesian and non-Cartesian brain MRI demonstrate faster sampling and improved reconstruction quality with less parameter tuning than annealed-likelihood methods, across varying undersampling patterns and coil counts. This approach yields rapid, robust posterior sampling with reliable uncertainty estimates, facilitating practical diffusion-model-based MRI reconstruction without extensive tuning.

Abstract

Purpose: The Unadjusted Langevin Algorithm (ULA) in combination with diffusion models can generate high quality MRI reconstructions with uncertainty estimation from highly undersampled k-space data. However, sampling methods such as diffusion posterior sampling or likelihood annealing suffer from long reconstruction times and the need for parameter tuning. The purpose of this work is to develop a robust sampling algorithm with fast convergence. Theory and Methods: In the reverse diffusion process used for sampling the posterior, the exact likelihood is multiplied with the diffused prior at all noise scales. To overcome the issue of slow convergence, preconditioning is used. The method is trained on fastMRI data and tested on retrospectively undersampled brain data of a healthy volunteer. Results: For posterior sampling in Cartesian and non-Cartesian accelerated MRI the new approach outperforms annealed sampling in terms of reconstruction speed and sample quality. Conclusion: The proposed exact likelihood with preconditioning enables rapid and reliable posterior sampling across various MRI reconstruction tasks without the need for parameter tuning.

Figures (5)

• Figure 1: A: Analytical diffusion process for a 2D toy model with a 1D linear measurement $A=(1\quad-1)^T$. The prior distribution is a mixture of 2D Gaussians forming a circle, which models the data manifold. Likelihood modifications corresponding to the diffused posterior and annealing are compared to the exact likelihood. All methods have the same posterior distribution at the minimum noise scale. B: Samples of the same 2D toy model with annealed (left) and exact (right) likelihood sampled with ULA for 10 and 500 iterations and pULA for 10 iterations. Samples are shown for noise levels $\sigma = [1, 0.2, 0.05, 0.01]$
• Figure 2: Reconstructions of a T2-weighted brain image for different undersampling patterns using $\ell_1$-Wavelet regularization and diffusion posterior sampling with annealed and exact likelihood. Error maps and PSNR/SSIM values are computed relative to the fully-sampled $\ell_1$-Wavelet reconstruction. Per noise level, $K=8$ ULA or $K=4$ pULA iterations were performed.
• Figure 3: Reverse diffusion process with exact (A) and annealed (B) likelihood for a T2-weighted brain image sampled from 8x random undersampled data.
• Figure 4: Reconstructions of a T1-weighted brain image for selected undersampling patterns and different numbers of virtual coils after coil compression. Uncertainty maps show the standard deviation over drawn samples.
• Figure 5: Brain image reconstructed from radially acquired FLASH data using different undersampling factors. A: Reconstruction with $\ell_1$-Wavelet regularization, a single sample from the posterior with exact likelihood and the average of ten samples. B: Pixel-wise standard deviation map in image space and k-space.