Introduction

Prior models are crucial for solving Magnetic Resonance Imaging (MRI) inverse problems, though existing models are not perfected. Prior models are utilized to yield desirable outputs due to the ill-posed nature of MRI inverse problems. However, traditional priors, being hand-crafted, often lead to overly smoothed reconstructions because of their simplistic nature⁶, while learning-based priors offer flexibility but are reliant on large datasets. The Deep Image Prior¹¹(DIP) method emerges as a potential solution, yet is hampered by its slow sampling speeds. Recently, Diffusion Model (DM) based methods^10,2 have showcased their generalizability in conventional retrospective experiments. This study aims to explore DM-based methods' generalizability into three other clinical scenarios: noisy measurements, varying anatomies, and complex datasets. Through experiments, this study highlights the potential for DM-based methods to reduce reliance on large datasets, thereby moving a step closer to practical clinical applications.

Methods

In this section, we introduce the MRI inverse problem and Diffusion models, along with two related approaches: projection-based and score-based.

MRI Inverse Problem: The MRI inverse problem's forward model is denoted by \(\mathbf{y} = \mathbf{M}\mathbf{F}\mathbf{x} + \boldsymbol{\epsilon}\), with \(\mathbf{M}\) as the binary under-sampling mask, \(\mathbf{F}\) the forward Fourier transform, and \(\boldsymbol{\epsilon}\) Gaussian noise during signal acquisition.

Diffusion Models: Diffusion models consist of a fixed forward process and a learned reverse process. The forward process transforms an unknown data distribution \(p_{\text{data}}\) towards a noise distribution like Gaussian, by gradually perturbing signals with Gaussian noise across many steps, with the noise standard deviation given by a schedule \(\beta(t)\), where \(t = 1, \ldots, T\). For reverse process, a U-Net is trained to estimate the noise at \({\bf x}_t\) and time \(t\), denoted as \(\boldsymbol{\epsilon}_{\boldsymbol{\theta}}({\bf x}_{t}, t)\). Diffusion models enable random sampling from \(p_{\text{data}}\) by initially sampling from the noise distribution to get \({\bf x}_T\), then iteratively removing estimated noise \(\boldsymbol{\epsilon}_{\boldsymbol{\theta}}({\bf x}_{t}, t)\) to obtain a clean image \({\bf x}_0\), a sample from \(p_{\text{data}}\).

Projection-based Methods: The projection \({\bf x}_{t}\) onto the set of images consistent with \({\bf y}_{t}\) ensures data consistency, which denoted as \({\bf P}_{{\bf y}_{t}}({\bf x}_{t})={\bf F}^{-1}[{\bf M}{\bf y}_{t}+({\bf I}-{\bf M}){\bf F}{\bf x}_{t}]\), where \({\bf y}_{t}=\alpha(t){\bf y}+\beta(t){\bf A}{\bf z}_t\), \({\bf z}_t\sim \mathcal{N}({\bf 0},{\bf I})\), and \(\alpha(t)\) related to \(\beta(t)\). In DDPM⁵, \(\alpha(t)\) and \(\beta(t)\) are \(\sqrt{\bar{\alpha}_t}\) and \(\sqrt{1-\bar{\alpha}_t}\), respectively. Both Song's method¹⁰ and Prediction-Projection-Noising (PPN) are projection-based; Song's method performs projection at \({\bf x}_{t}\), while PPN at \({\bf x}_{0}\). PPN's denoising is defined as \({\bf x}_{t-1}=\sqrt{\bar{\alpha}_{t-1}}\cdot{\bf P}_{{\bf y}}\left({\bf x}_{t}-\sqrt{1-\bar{\alpha}_{t}}\cdot\boldsymbol{\epsilon}_{\boldsymbol{\theta}}({\bf x}_{t},t)/\sqrt{\bar{\alpha}_{t}}\right)+\sqrt{1-\bar{\alpha}_{t-1}}\cdot{\bf z}_t\), illustrated in Figure 1.

Score-based Method: Diffusion Posterior Sampling (DPS)¹ utilizes a Langevin dynamics-style sampling defined by \({\bf x}_{t+1}={\bf x}_{t}+\eta_{t}\nabla_{{\bf x}_{t}}\log p({\bf x}_{t}|{\bf y})+\sqrt{2\eta_{t}}{\bf z}_t\), with \(\eta_{t}\) as the step size and \({\bf z}\sim \mathcal{N}({\bf 0},{\bf I})\). The posterior score \(\nabla_{{\bf x}_{t}}\log p({\bf y}|{\bf x}_{t})\) can be approximated by \(-\frac{1}{\sigma^{2}}\nabla_{{\bf x}_{t}}\|{\bf y}-{\bf A}\hat{{\bf x}}_{0}\|_{2}^{2}\), with \(\hat{{\bf x}}_{0}=\left({\bf x}_{t}-\sqrt{1-\bar{\alpha}_{t}}\boldsymbol{\epsilon}_{\boldsymbol{\theta}}({\bf x}_{t},t)\right)/\sqrt{\bar{\alpha}_{t}}\) using Tweedie’s approach^4,7.

In this study, a pretrained Diffusion Model (DM) is utilized as a prior for solving medical inverse problems. We test three reconstruction methods: PPN, Song's method¹⁰, and DPS¹, all sharing a similar structure. Starting with a Gaussian-drawn sample \({\bf x}_T\) at noise level \(T\), we iteratively apply a denoising and a data-consistent step, either projection or score-based, to obtain a noise-free, \({\bf y}\)-consistent image \({\bf x}_0\).

Experiments

We utilize ADM^8,3 as the DM implementation, training it unsupervisedly on the fastMRI¹² knee dataset, then test it on two datasets: fastMRI knee and brain, without fine-tuning. Our baselines include zero-filled, total variation (TV), and U-Net methods⁹, with a 0.1 regularization parameter for TV, and U-Net trained supervisedly on the fastMRI knee dataset. Additionally, we incorporate three diffusion model-based methods: PPN, DPS¹ and Song's method¹⁰, with a step size of 1.0 for DPS and \( \lambda \) set to 1.0 for Song's method. For results, refer to Figures 2 to 5.

Discussion

Overall, DM-based methods outperformed the baselines in most scenarios except for complex measurements, notably excelling in retrospective subsampling experiments and those involving varying anatomies. Reconstructions with Poisson masks yielded better results during accelerated sampling. While the performance gap diminished in noisy measurements, DM-based methods generated fewer artifacts compared to the baselines. However, DM-based methods did not perform well in the context of complex measurements.

Conclusion

This study investigates the utility of Diffusion Models in addressing distribution shift challenges in medical imaging. Our results demonstrate that DM-based methods adapt well to retrospective subsampling experiments, noisy measurements, and diverse anatomies. However, they underperform on complex test sets. Future work should explore complex in-vivo datasets to advance DM-based methods towards clinical applications.