Introduction

MRI serves as a crucial test in diagnostic imaging. Nonetheless, the prolonged scan time of MRI compels a compromise between the speed of acquisition and the quality of the reconstructed images. Deep learning models, such as U-nets^1–3 and Vision Transformers^4,5, have been proposed to improve MR image reconstruction. Recently, Denoising Diffusion Probabilistic Models (DDPM)⁶ have demonstrated remarkable performance for high-quality image reconstruction. However, the original DDPM method is computationally intensive, especially with high-resolution inputs, and fraught by slow inference times. In this work, we propose a novel patch-based conditioned Denoising Diffusion Probabilistic model that aims at optimizing computational efficiency and accelerating the inference process.

Methods

Patch-based Conditioned Denoising Diffusion Probabilistic (PC-DDPM) enables faster reverse processing than standard DDPM. This patch-based model also uses a guided denoising process for size-independent reconstruction during inference.

Denoising Diffusion Probabilistic Model (DDPM):
Given a sample $$$\mathbf{x}_\mathbf{0}$$$ from the real data distribution $$$q(\mathbf{x})$$$, the forward process simulates the diffusion process by gradually adding noise to the data until pure Gaussian noise is reached: $$q\left(\mathbf{x}_{\mathbf{1}:\mathbf{T}}\middle|\mathbf{x}_\mathbf{0}\right)=\prod_{t=1}^{T}{q\left(\mathbf{x}_\mathbf{t}\middle|\mathbf{x}_{\mathbf{t}-\mathbf{1}}\right)},$$ where $$$T$$$ is the total number of steps.
The inference process (e.g., reverse process) approximates the reverse distribution $$$q\left(\mathbf{x}_{t-1}\middle|\mathbf{x}_t\right)$$$ with the parameterized model $$$p_\theta$$$ through each time step $$$t$$$: $$p_\theta\left(\mathbf{x}_{0:T}\right)=p\left(\mathbf{x}_T\right)\prod_{t=1}^{T}\, p_\theta\left(\mathbf{x}_{t-1}\mid\mathbf{x}_t\right).$$

Conditioned DDPM:
To accurately represent complex distributions, conditioned diffusion models⁷ incorporate auxiliary conditions into the generating process. Different to DDPM, the reverse process of a Conditioned DDPM becomes: $$p_\theta\left(\mathbf{x}_0,\mathbf{x}_1,\ldots,\mathbf{x}_T\mid\mathbf{x}_c\right):=p\left(\mathbf{x}_T\right)\prod_{t=1}^{T}\, p_\theta\left(\mathbf{x}_{t-1}\mid\mathbf{x}_t,\mathbf{x}_c\ \right),$$ where $$$\mathbf{x}_c$$$ is the auxiliary condition, and the noise distribution estimate becomes: $${\epsilon_\theta\left(\mathbf{x}_t,t\right)\ \rightarrow\epsilon}_\theta\left(\mathbf{x}_t,\mathbf{x}_c,t\right).$$
However, the reverse diffusion process of generating a high-resolution sample from DDPM through the Markov chain can be considerably slow, given that $$$T$$$ can range from one to many thousands.

Patched-based Conditioned DDPM (PC-DDPM):
Inspired by the patched-based Denoising Diffusion model⁸, we propose a class of non-Markovian diffusion processes built from the seminal work of Song et al⁹.
For high-resolution MRI image reconstruction, our patch-based technique allows size-independent image reconstruction and faster reverse processing process than typical DDPM. A guided denoising approach that applies smoothed noise estimates across overlapping patches during inference can improve image restoration quality and efficiency.
Given a patch is sampled from the fully sampled data distribution $$$q(\mathbf{x}_\mathbf{0}^i)$$$ and with corresponding undersampled MRI as condition $$$\mathbf{x}_\mathbf{c}^i$$$. The transition probability for the patch-based conditional diffusion can be formulated as: $$p_\theta\left(\mathbf{x}_{t-1}^i\mid\mathbf{x}_t^i,\ \mathbf{x}_c^i\right)=\mathcal{N}\left(\mathbf{x}_{t-1}^i;\mathbf{\mu}_\theta\left(\mathbf{x}_t^i,\ \mathbf{x}_c^i,t\right),\mathbf{\Sigma}_\theta\left(\mathbf{x}_t^i,\ \mathbf{x}_c^i,t\right)\right).$$
The objective of the model is to minimize the difference between the expected noise during the diffusion process and the one generated by the model parameterized by $$$\theta$$$: $$\mathbf{E}_{x_0,t,\epsilon_t\sim N(0,I)}\left[\left|\epsilon_t-\epsilon_\theta\left(\mathbf{x}_t^i,\ \mathbf{x}_c^i,t\right)\right|^2\right].$$
At each time step, the model progressively refines the image patch: $$\mathbf{x}_{t-1}^i=\sqrt{{\bar{\alpha}}_{t-1}}\left(\frac{\mathbf{x}_t^i-\sqrt{1-{\bar{\alpha}}_t}\epsilon_\theta\left(\mathbf{x}_t^i,\ \mathbf{x}_c^i,t\right)}{\sqrt{{\bar{\alpha}}_t}}\right) +\sqrt{1-\alpha_{t-1}}\epsilon_\theta\left(\mathbf{x}_t^i,\ \mathbf{x}_c^i,t\right).$$
Implicit sampling allows accelerating the sampling process by only requiring a subset of the timesteps $$$\tau_i$$$ from complete sequence $$${1,...,T}$$$ where $$$i<T$$$.

Experiments:
We performed experiments using fastMRI¹⁰ brain MRI dataset from 30 patients (1,000 training, 200 testing), undersampled at 4x and 8x speeds with center acquisition rates of 0.08 and 0.04, respectively. Multi-coil images were combined using sum of squares. Then, for each subject, the combined images (320x320 pixels) were divided into 16 patches of 64x64 pixels each. Our model was trained with the principal diffusion loss, followed the original implementation of DDPM¹¹ that consists of ResNet blocks, group normalization and self-attention.
To evaluate the quality of our approach, we measured reconstructed images using the PSNR and SSIM metrics. The proposed method was compared against several baseline reconstruction methods, including zero-filled, standard U-Net, Vision Transformer, and conditioned DDPM.

Result

Table 1 shows that the proposed patch-based conditioned DDPM (PC-DDPM) outperforms other comparison methods on PSNR and SSIM. For brain MRI, PC-DDPM demonstrates superior reconstruction performance and outperforms other methods, achieving top SSIM (0.817 and 0.796) and PSNR (31.04 and 30.85) scores at both x4 and x8 acceleration rates, respectively.

Figure 1 indicates the difference maps between the reconstructed image and the fully-sampled image. Our proposed PC-DDPM shows fewer errors than other reconstruction models. The PC-DDPM model demonstrated five times better efficiency, reconstructing MRI images in just 2 seconds, compared to the 10 seconds in the conditioned DDPM.

Discussion

We proposed a PC-DDPM model which outperforms other reconstruction methods, such as U-Net, Vision Transformer, and conditioned DDPM. PC-DDPM produces high-resolution reconstruction results at different acceleration rates (x 4 and x8). This is due to the inherent properties of diffusion models that seek to sample from the true distribution and refine the image over time. Together with patch-based input and non-Markovian diffusion processes, our PC-DDPM approach demonstrates the potential of reconstructing high-fidelity MRI images in a short time.

Acknowledgements

No acknowledgement found.