Diffusion Modeling with Domain-conditioned Prior Guidance for Accelerated MRI and qMRI Reconstruction

DDPMs are generative models which are highly effective in learning complex data distributions. The forward diffusion process adds noise to the input data, gradually increasing the noise level until the data is transformed into pure Gaussian noise. This process systematically perturbs the structure of the data distribution. The reverse diffusion process, also known as the denoising process, is then applied to recover the original structure of the data from the perturbed data distribution.

Forward Process DDPM[24] presents the forward diffusion mechanism as a Markov Chain, wherein Gaussian noise is incrementally introduced across several steps to yield a collection of perturbed samples. Consider the uncorrupted data distribution that is characterized by density $q(x_{0})$, which undergoes incremental transformations through the addition of Gaussian noise at various stages, resulting in a spectrum of modified samples. If we draw a data sample $x_{0}$ from $q(x_{0})$, the forward diffusion process modifies this data point through integration of Gaussian perturbations at each time step $t,$ which can be mathematically represented as

Here, $T$ is the total number of diffusion steps, while the noise scaling sequence $\beta_{1},...,\beta_{T}\in[0,1)$ defines the magnitude of variance introduced at each step. For isotropic Gaussian noise $\epsilon\sim\mathcal{N}(\textbf{0},\textbf{I})$, at each time step,

With a large number of forward iterations, $x_{t}$ converges to an isotropic Gaussian sample, so we have $q(x_{T})\approx\mathcal{N}(x_{T};\textbf{0},\textbf{I})$. Let $\alpha_{t}:=1-\beta_{t}$ and $\bar{\alpha}_{t}:=\prod^{t}_{s=0}\alpha_{s}$, then given the initial sampled data $x_{0}$ and sampling step $t$, we can obtain

Reverse Process The reverse process aims to retrieve $x_{0}$ by stripping away the noise from $x_{T}$. Specifically, starting with a noise vector $x_{T}\sim\mathcal{N}(x_{T};\textbf{0},\textbf{I})$, we iteratively sample from the learnable transition kernel $x_{t-1}\sim p_{\theta}(x_{t-1}|x_{t})$ until $t=1$ through a learnable Markov chain in the reverse time direction.

The mean $\mu_{\theta}(x_{t},t)$ is parametrized by deep neural networks and $\sigma_{t}^{2}=\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t}}\beta_{t}$. The objective during training is to minimize the discrepancy between the true data distribution and the estimated distribution, which can be achieved by minimizing the variational bound on log likelihood:

Ho et al.[24] proposed using deep neural networks to estimate added noise $\epsilon$ and perform parameterization on $\mu_{\theta}$:

Accordingly, the loss function simplifies to

Accelerated MRI considers the following measurement model

where $x\in\mathbb{C}^{n}$ is the MR image of interest consisting of n pixels, $f\in\mathbb{C}^{m}$ is its corresponding undersampled k-space measurement, and $\varepsilon\in\mathbb{C}^{n}$ is measurement noise. The parameterized forward measurement matrix $A\in\mathbb{C}^{m\times n}$ is defined as

where $\mathcal{S}:=[\mathcal{S}_{1},...,\mathcal{S}_{j}]$ are the coil sensitivity maps for $j$ different coils, $\mathcal{F}\in\mathbb{C}^{n\times n}$ is the 2D discrete Fourier transform, and $\mathcal{P}_{\Omega}\in\mathbb{N}^{u\times n}(u\ll n)$ is the binary undersampling mask corresponding to $u$ sampled data points using undersampling pattern $\Omega$. The classical approach to solve for $x$ using undersampled k-space $f$ is to solve the following optimization problem:

The training of Static DiMo starts with the input of fully sampled k-space data, which we denote as $\hat{f}$. We can write $x=\mathcal{F}^{-1}\hat{f}$, where $\mathcal{F}^{-1}$ is the inverse discrete Fourier transform. The task of Static DiMo is to minimize $\hat{f}$, thus we can rewrite equation (10) as:

The minimizer $\hat{f}$ can be solved by applying the gradient decent algorithm, which is straightforward to compute:

Typical diffusion models aim to estimate $q(x|\mathcal{P}_{\Omega},f,\mathcal{S})$. Since Static DiMo is defined in the frequency domain and $x=\mathcal{F}^{-1}\hat{f}$, the problem can be interpreted as generating samples of $q(\hat{f}|\mathcal{P}_{\Omega},f,\mathcal{S})$ which is conditioned on undersampling mask $\mathcal{P}_{\Omega}$, scanned measurement $f$ and coil sensitivity maps $\mathcal{S}$.

According to DDPM[24], the diffusion process can be represented with a noise schedule $\beta_{1},...,\beta_{T}$ as:

Let $\alpha_{t}=1-\beta_{t},\bar{\alpha}_{t}=\sum^{t}_{s=1}\alpha_{s}$. We can track the forward process conditioned on the initial:

We choose $\alpha_{t}$ so that $\bar{\alpha}_{T}\approx 0$, then we obtain

which approaches an isotropic Gaussian distribution.
Sample $\hat{f_{0}}\sim q(\hat{f_{0}})$, the diffusion model has the following form:

The sampling of $p_{\theta}(\hat{f_{0}}|\mathcal{P}_{\Omega},f,\mathcal{S})$ is the reverse diffusion process.

The training of $p_{\theta}(\hat{f_{0}}|\mathcal{P}_{\Omega},f,\mathcal{S})$ is to optimize the variational bound of log likelihood in (5), and this can break down into optimizing random terms of $L$ with stochastic gradient descent. Now we analyse $L_{t-1}$ using (14b) and (17b), we can write

where $C$ is constant that independent from $\theta$. From (3) we get $\hat{f}_{0}=\frac{1}{\sqrt{\alpha_{t}}}(\hat{f}_{t}(\hat{f}_{0},\epsilon)-\sqrt{1-\bar{\alpha}_{t}}\epsilon)$, then apply the forward posterior formula (14d), we have

We can derive that $\mu_{\theta}$ should be parametrized to predict $\frac{1}{\sqrt{\alpha_{t}}}(\hat{f}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}}\epsilon)$, so we can put

where $\hat{f}_{t}=\sqrt{\bar{\alpha}_{t}}\hat{f}_{0}+\sqrt{1-\bar{\alpha}_{t}}\epsilon,\epsilon\sim\mathcal{N}(\textbf{0},\textbf{I})$. Then the simplified loss function is

In equation (21), we can interpret it as $\hat{f}_{t}(\sqrt{\bar{\alpha}_{t}}\hat{f}_{0}+\sqrt{1-\bar{\alpha}_{t}}\epsilon,\mathcal{P}_{\Omega},f,\mathcal{S})$ being a function that is dependent on the components of $\epsilon_{\theta}$. Then we can write our simplified loss function as

Now we can derive the function $\hat{f}_{t}$ so that it can provide a prior guidance for the diffusion during training and sampling. The diffusion process starts with $\hat{f}_{0}\sim q(\hat{f})$, we have $\hat{f}_{t}=\sqrt{\bar{\alpha}_{t}}\hat{f}_{0}+\sqrt{1-\bar{\alpha}_{t}}\epsilon$ by (3). An intuitive way of designing the function $\hat{f}_{t}$ is to incorporate a data consistency (DC) layer which perturbs the k-space data $\hat{f}_{t}$ with a linear combination of partially scanned data $f$. Then we can further fine-tune the DC layer using gradient descent (GD) as described in (12). The proposed algorithm synergizes DDPM, MRI data consistency enforcement and gradient descent optimization, which is listed in Alg. 1:

In step 2, $\lambda_{t}$ is a scheduling factor to guide the data consistency dynamics, which follows the exponential function to schedule the perturbation dynamics:

Prior knowledge of scanned signal decreases exponentially as $t$ grows. In doing this, the schedule parameter can balance network training and stabilize parameter learning. $\mathbbm{1}$ denotes the matrix with values of one and $\eta_{k}$ is the learnable step size for gradient descent.

Sampling $\hat{f}_{t-1}\sim p_{\theta}(\hat{f}_{t-1}|\hat{f}_{t},\mathcal{P}_{\Omega},f,\mathcal{S})$ leverages the Langevin dynamics with $\epsilon_{\theta}$ as learned gradient of data density which provide the local information of the data distribution. Here we compute $\hat{f}_{t-1}=\frac{1}{\sqrt{\alpha_{t}}}(\hat{f}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}}\epsilon_{\theta}(\hat{f}_{t},t,\mathcal{P}_{\Omega},f,\mathcal{S}))+\sigma_{t}z$ where $z\sim\mathcal{N}(\textbf{0},\textbf{I})$. Alg. 2 shows the sampling process which is the reverse of the diffusion process.

To extend from static MRI reconstruction to qMRI reconstruction, we first define the MR parameter maps as $\Delta=\{\delta_{i}\}_{i=1}^{N}$ where $\delta_{i}$ symbolizes each MR parameter and $N$ is the total number of MR parameters to be estimated. Given the MR signal model $\mathbf{M}$, which is a function of $\Delta$, the MR parameters $\Delta$ can be estimated by minimizing the following:

Quantitative DiMo is defined in the parameter domain, therefore we focus on estimating $q(\hat{\Delta}|\mathcal{P}_{\Omega},f,\mathcal{S},\mathbf{M})$ with an additional condition, the MR signal model $\mathbf{M}$. Since the diffusion processe and reverse process are highly similar to Static DiMo, we omit the derivation of Quantitative DiMo.

In contrast to the algorithm design of Static DiMo, a major difference is the function $\hat{\Delta}_{t}(\bar{\alpha}_{t}\hat{f}_{0}+\sqrt{1-\bar{\alpha}_{t}}\epsilon,\mathcal{P}_{\Omega},f,\mathcal{S},\mathbf{M})$. Following the same framework as Static DiMo, Quantitative DiMo starts with equation (3), which is then input to the quantitative DC (QDC) layer, followed by fine-tuning with gradient descent. The Alg. 3 displays the training process of Quantitative DiMo and Alg. 4 shows the sampling process:

QDC layer is presented in steps 2-4 in the above algorithm. The addition of the QDC step is to perturb quantitative parameters and guide the parameter space denoising and generation process. Steps 5-7 further optimize the perturbed $\hat{\Delta}$ through iterating $K$ steps of gradient descent. Alg. 4 displays the sampling process which is the reverse of the diffusion process.

To demonstrate the feasibility of Quantitative DiMo for reconstructing accelerated qMRI, we used $T_{1}$ mapping using the vFA model [13] as an example. For MR images are acquired at multiple flip angles $\phi_{i}$, where $i=1,...,M$, the MR signal model $\mathbf{M}$ reads as:

where $T_{1}\in\mathbb{R}^{n}$ denotes the spin-lattice relaxation time map and $I_{0}\in\mathbb{C}^{n}$ is proton density map. In this model, the MR parameters estimated are encapsulated in the set $\Delta=\{T_{1},I_{0}\}$. Other imaging parameters such as flip angles and repetition time TR are known. Fig. 1 shows the framework of Static DiMo and Quantitative DiMo and Fig. 2 illustrates the U-Net architecture for $\epsilon_{\theta}(x_{t},t)$.

We acquired knee data from 25 subjects, using 20 for training and 5 for testing. These fully-sampled data were obtained from a 3T GE Premier scanner equipped with an 18-element knee coil array. The data was acquired using two-dimensional fast spin-echo (FSE) sequences along the coronal plane, with both proton density-weighting (PD-weighting) and $T_{2}$-weighting. All images were resized resulting in in-plane matrix size of 320 $\times$ 320 with 3mm slice thickness. Number of slices varied from 30 to 38 in each dataset, amounting to 875 total slices for both training and testing. The experiments were conducted retrospectively using three different acceleration factors: AFs = 4$\times$, 5$\times$ and 6$\times$ using 1D variable density Cartesian undersampling masks[30] where the 20 central k-space lines were fully sampled. We evaluated the performance of the proposed method against several state-of-the-art methods which include Variational Network (VN)[6], ISTA-NET[31], and pMRI-Net[32] over two contrasts: PD- and $T_{2}$-weighting. Results were compared in terms of peak signal to noise ratio (PSNR), structure similarity (SSIM) and normalized mean squared error (NMSE). The equations for assessing PSNR, SSIM, and NMSE between reconstruction $\mathrm{v}$ and the ground truth $\mathrm{v}^{*}$ are:

where$\mu_{\mathrm{v}},\mu_{\mathrm{v}^{*}}$ represent the local mean of pixel intensities with standard deviations $\sigma_{\mathrm{v}},\sigma_{\mathrm{v}^{*}}$ for images $\mathrm{v},\mathrm{v}^{*}$, respectively. $\sigma_{\mathrm{v}\mathrm{v}^{*}}$ denotes the covariance between $\mathrm{v}$ and $\mathrm{v}^{*}$, and $C_{1}=(K_{1}L)^{2},C_{2}=(K_{2}L)^{2}$ are two constant variables introduced to stabilize the division. We chose the total number of sampling steps as $T=1000$ using a linear schedule from $\beta_{1}=10^{-5}$ to $\beta_{T}=10^{-2}$ with training batch size of 4. The total number of epochs used was 7000. The initial parameter for Alg.1 was chosen to be $\eta_{0}=10^{-4}$.

Quantitative DiMo experiments were carried out using in-vivo brain data of healthy volunteers, obtained from a Siemens 3T Prisma scanner equipped with a 20-channel head coil. The five subject fully sampled vFA brain data was acquired along the sagittal plane using a spoiled gradient echo sequence with imaging parameters TE/TR = $12/40$ ms, FA = $5^{\circ},10^{\circ},20^{\circ},40^{\circ}$, in-plane matrix size = $176\times 176$ with 3mm slice thickness and a total number of 48 slices per subject. Leave-One-Out cross-validation was used for all five subjects to train and test our method.

Two undersampling schemes were retrospectively used: (1) 1D variable density Cartesian undersampling[30] with acceleration factor AF= $4\times$, where the 16 central k-space lines were fully sampled and (2) 2D Poisson disk undersampling at AF= $4\times$ with the central 51x51 k-space portion fully sampled. The undersampling patterns were varied for each flip angle, like in previous studies[14, 16]. We compared Quantitative DiMo with two advanced non-deep learning qMRI reconstruction techniques Locally Low Rank (LLR)[33], Model-TGV[34], and a self-supervised deep learning method RELAX[16]. For LLR and Model-TGV, settings recommended in their original research papers were used along with available code. RELAX and Quantitative DiMo were trained through cross-validation. For the diffusion hyperparameters, we maintained the total number of sampling steps at $T=1000$ using linear scheduling from $\beta_{1}=10^{-6}$ to $\beta_{T}=0.05$. The training batch size used was 8 with 5500 total epochs. The initial parameter for gradient descent in Alg. 3 was set to $\tau_{0}=10^{-5}.$

All in-vivo studies were carried out under a protocol approved by our institution’s institutional review board. Another experiment condition was as follows: The coil sensitivity maps were estimated from ESPIRiT[35]. Separate $B_{1}^{+}$ maps were acquired[36] to correct for $B_{1}^{+}$ bias in estimating $T_{1}$. All the programming in this study was implemented using Python language and PyTorch package, and experiments were conducted on one NVIDIA A100 80GB GPU and an Intel Xeon 6338 CPU at Centos Linux system.

An example of reconstruction results obtained from different methods at AF = $6\times$ are shown for PD- (Fig. 3) and T2-weighted (Fig. 4) images along the coronal plane are displayed. Zero-filled results show significant blurring and artifacts due to undersampling. Although VN, ISTA-Net, and pMRI-Net are able to remove these artifacts, compared with the fully sampled reference images, blurring characteristic persists with VN showing the most, while ISTA-Net and pMRI-Net are comparable. Static DiMo’s reconstruction not only removes the undersampling artifacts, but also significantly suppresses blurring showing clarity and sharpness, thus outperforming both ISTA-Net and pMRI-Net. Furthermore, Static DiMo shows superior denoising capability, rooted in its denoising diffusion modeling. This not only proficiently removes noise and artifacts, but also retains the integrity of high-frequency image detail. This is further illustrated in the zoomed-in images of Fig. 3(b) and Fig. 4(b). Static DiMo clearly distinguishes tissue boundaries across various tissue types, including the cartilage, meniscus, bone and muscle. In addition, the fidelity of tissue texture and sharpness are preserved. The pixel-wise error maps shown in Fig. 3(c) and Fig. 4(c) also demonstrate that Static DiMo produces the least reconstruction errors with respect to the fully sampled reference. For both contrasts, Static DiMo consistently outperforms the other methods. Overall, in agreement with the qualitative observation in Fig. 3 and Fig. 4, reconstruction results are summarized in Table 1 for all testing subjects using quantitative metrics, showing the superior performance of Static DiMo over other methods in reconstruction fidelity, structure and texture preservation and noise suppression.

$T_{1}$ and $I_{0}$ maps estimated from 4$\times$ 1D variable density Cartesian undersampling using various methods are shown in Fig. 5 The zero-filled $T_{1}$ map derived from pixel-wise fitting the undersampled images (Fig. 5(a)) exhibits noise and ripple artifacts due to aliasing. While LLR manages to mitigate these artifacts to an extent, the noisy signature persists. In its attempt to neutralize the noise, Model-TGV yields a more refined tissue appearance with enhanced $T_{1}$ contrast. However, the resulting maps are overly smooth, appearing blurry. Conversely, RELAX effectively suppresses both noise and artifacts, delivering sharper maps, albeit with a slightly blocky texture. This blockiness is speculated to arise from challenges in achieving convergence in the end-to-end network using limited data. Meanwhile, Quantitative DiMo generates clear and sharp $T_{1}$ maps. Its proficiency in removing noise translates to its superior map quality, both in terms of appearance and contrast. This is further witnessed in the zoom-ins (Fig. 5 (b)) which show that Quantitative DiMo clearly distinguishes the boundary between white matter and grey matter, appearing nearly identical to the fully sampled reference map. This is quantitatively confirmed in the error maps (Fig. 5(c)) where the zero-filled incurs the largest error, followed by LLR, Model-TGV, and RELAX.

Mean $T_{1}$ values obtained from representative white matter and grey matter regions are presented in Table 2 for all subjects. Overall, LLR shows least agreement with respect to the fully sampled reference method, followed by Model-GTV. On the other hand, both RELAX and Quantitative DiMo show good agreement. However, as shown by the Wilcoxon signed rank test results, Quantitative DiMo shows better correlation with the reference than RELAX.

The $I_{0}$ maps in Fig 5(d-f) exhibit a similar signature in reconstruction quality as the $T_{1}$ maps, where LLR and Model-TGV show aliasing artifacts with the most blur, further confirmed by the zoom-ins, reflecting the undersampling pattern of many high frequency data points being undersampled with low frequency data points fully sampled. The deep learning method RELAX shows its ability to denoise through artifacts removal and detail preservation. However, it still lacks sharpness resulting in smoothened edges. Quantitative DiMo produces the sharpest maps preserving detail, which again is confirmed by its error maps, showing least error.

The $T_{1}$ and $I_{0}$ maps estimated from 4$\times$ 2D Poisson disk undersampling are shown in Fig. 6. Examining Fig. 6(a) and Fig. 6(d), it is apparent that the main differences compared to the 4x 1D undersampling case are the different artifact patterns, which in turn stems from the different undersampling patterns. As seen in the zoom-in views of the $T_{1}$ (Fig. 6(b)) and $I_{0}$ (Fig. 6(e)) maps, compared to Fig. 5, LLR provides sharper maps details but still retain artifacts and noise, whereas Model-TGV preserves detail but exhibits blurred edges due to its general averaging of noise. RELAX captures sharp structures but is not able to attain details in some fine regions. Quantitative DiMo produces artifact free $T_{1}$ and $I_{0}$ maps with superior performance compared to other methods. This is again confirmed by the error maps (Fig. 6(c) and Fig. 6(f)), with Quantitative DiMo showing the least error. This is likely achieved through integrating the unrolling gradient descent algorithm and diffusion denoising network, prioritizing noise suppression without compromising the fidelity and clarity of the underlying tissue structure.

An ablation study investigating model uncertainty was conducted by sampling the well-trained Static DiMo on PD data 100 times. From these 100 samples, we derived mean images, error maps, and variance maps. We then compared three distinct acceleration factors (4$\times$, 5$\times$, 6$\times$) using 1D Cartesian masks, shown in Fig.7. When compared with fully sampled images, the error and variance maps predominantly capture the edge and boundary details. Notably, higher acceleration factors increases the uncertainty with amplified errors, especially around the tissue interfaces. Incrementally increasing our sampling data counts, starting from 10, increasing to 50, and finally reaching 100, we witnessed a progression in the variance maps. Namely, they began displaying smoother transitions, with the high uncertainty edge regions gradually diminishing. This suggests that augmenting the number of sampling instances can be instrumental in attenuating noise within the mean images.