NPB-REC: A Non-parametric Bayesian Deep-learning Approach for Undersampled MRI Reconstruction with Uncertainty Estimation

TL;DR

This paper tackles the challenge of reconstructing MRI images from undersampled k-space while providing calibrated uncertainty. It introduces NPB-REC, a non-parametric fully Bayesian framework that uses Stochastic Gradient Langevin Dynamics to sample the posterior over network parameters, enabling posterior averaging and uncertainty maps. Evaluated on fastMRI multi-coil data with an E2E-VarNet backbone, NPB-REC achieves higher PSNR/SSIM than the baseline and shows uncertainty measures that better correlate with reconstruction error and generalize to anatomical and sampling-pattern shifts. The approach offers a principled path toward safer clinical deployment of DL-based MRI reconstruction, with public code and models available for reuse.

Abstract

The ability to reconstruct high-quality images from undersampled MRI data is vital in improving MRI temporal resolution and reducing acquisition times. Deep learning methods have been proposed for this task, but the lack of verified methods to quantify the uncertainty in the reconstructed images hampered clinical applicability. We introduce "NPB-REC", a non-parametric fully Bayesian framework, for MRI reconstruction from undersampled data with uncertainty estimation. We use Stochastic Gradient Langevin Dynamics during training to characterize the posterior distribution of the network parameters. This enables us to both improve the quality of the reconstructed images and quantify the uncertainty in the reconstructed images. We demonstrate the efficacy of our approach on a multi-coil MRI dataset from the fastMRI challenge and compare it to the baseline End-to-End Variational Network (E2E-VarNet). Our approach outperforms the baseline in terms of reconstruction accuracy by means of PSNR and SSIM ($34.55$, $0.908$ vs. $33.08$, $0.897$, $p<0.01$, acceleration rate $R=8$) and provides uncertainty measures that correlate better with the reconstruction error (Pearson correlation, $R=0.94$ vs. $R=0.91$). Additionally, our approach exhibits better generalization capabilities against anatomical distribution shifts (PSNR and SSIM of $32.38$, $0.849$ vs. $31.63$, $0.836$, $p<0.01$, training on brain data, inference on knee data, acceleration rate $R=8$). NPB-REC has the potential to facilitate the safe utilization of deep learning-based methods for MRI reconstruction from undersampled data. Code and trained models are available at \url{https://github.com/samahkh/NPB-REC}.

Figures (6)

• Figure 1: Schematic illustration of our proposed NPB-REC system. At the inference phase, we use a set of models with the same backbone network but different parameters, which resulted from the SGLD-based training. Then, these models predict a set of reconstructed images, $\left\{ \hat{x}_{t}\right\} _{t_{b}}^{N}$ by propagating the under-sampled k-space input data,$\{\tilde{k_i}\}_{i=1}^{N_c}$, through each one of the backbone models. Lastly, the averaged reconstructed image and the pixel-wise std., $\bar{x}$ and $\Sigma$, are calculated. The average is used as the most probable reconstruction prediction and the $\Sigma$ is utilized for uncertainty assessment.
• Figure 2: Normalized MSE (NMSE) and SSIM vs. $N-t_b$, the number of models used in the averaging, obtained on a subset from the test set (32 images sampled randomly). Lower NMSE curves correspond to better prediction quality. However, a higher SSIM is better. See red and blue arrows.
• Figure 3: Reconstructions results. Rows 1 and 3: Examples of comparing the ground truth (GT) fully sampled image to the reconstructed images obtained by the three models (1-3), NPB-REC, baseline, E2E-VarNet trained with Dropout, and the NPB-REC std. map at accelerations $R=4$, $R=8$, respectively. Rows 2 and 4: The corresponding annotated ROIS of Nonspecific white matter lesions.
• Figure 4: Uncertainty assessment. Scatter plots of the mean value of Std. estimate versus the MSE metric, calculated between the reconstructed and the ground truth, in log scale, for our NPB-REC method \ref{['fig4:a']} and Monte Carlo Dropout \ref{['fig4:b']}. \ref{['fig4:c']} Our measure of uncertainty versus the acceleration rate.
• Figure 5: Uncertainty Assessment. Box plots of the mean value of Std. estimate obtained by NPB-REC, in log scale, for different anatomical shifts \ref{['fig5:a']} and \ref{['fig5:b']} and two types of masks \ref{['fig5:c']} and \ref{['fig5:d']} for both NPB-REC and Dropout methods, respectively.