0377

Variational diffusion models for blind MRI inverse problems
Julio A. Oscanoa1, Cagan Alkan2, Daniel Abraham2, Mengze Gao3, Aizada Nurdinova3, Daniel Ennis3, Kawin Setsompop3, John Pauly2, Morteza Mardani4, and Shreyas Vasanawala3
1Department of Bioengineering, Stanford University, Stanford, CA, United States, 2Department of Electrical Engineering, Stanford University, Stanford, CA, United States, 3Department of Radiology, Stanford University, Stanford, CA, United States, 4NVIDIA Inc., Santa Clara, CA, United States

Synopsis

Keywords: AI Diffusion Models, Machine Learning/Artificial Intelligence, Diffusion models

Motivation: Diffusion models have shown state-of-the-art performance in solving inverse problems. However, current solutions typically consider cases only when the forward operator is fully known, which limits their applicability to the wide variety of MRI inverse problems.

Goal(s): Develop a general method for blind MRI inverse problems with unknown forward operator parameters.

Approach: We extend the RED-diff framework, which has the key strength of not requiring training or fine–tuning for each specific task. We test our method for image reconstruction with off-resonance and motion correction.

Results: Our blind RED-diff framework can successfully approximate the unknown forward model parameters and produce accurate reconstructions.

Impact: We demonstrate the potential of current diffusion models to readily tackle a wide range of blind inverse problems in MRI without application-specific re-training or fine-tuning. Image reconstruction with motion and off-resonance correction are the first demonstration applications.

Introduction

Diffusion models have shown superior performance for solving inverse problems1,2. Pre-trained diffusion models can be used as strong data priors in plug-and-play fashion at inference time2-6. Recently, Mardani7 proposed a regularization by denoising (RED-diff) framework for solving generic inverse problems and Ozturkler8 extended it for MRI reconstruction. RED-diff uses variational inference to approximate the posterior distribution, which corresponds to minimizing a data-consistency loss and score-matching regularization via denoisers at different diffusion steps. Advantageously, these techniques do not require training or fine-tuning for each specific task. However, they currently require full knowledge of the forward model, hindering their applicability to blind problems with unknown parameters in the forward model.
Herein we extend the RED-diff framework to blind inverse problems. Using variational inference, we represent the sampling as an alternating stochastic optimization that estimates both the image and forward model parameters. We evaluate our blind RED-diff on image reconstruction with unknown field inhomogeneity map and motion parameters.

Theory

Consider the blind inverse problem: $$y=f_\gamma(x_0)+\eta,\;\eta\sim N(0,\sigma_\eta^2I),\;\text{[Eq.1]}$$
where the forward model $$$f$$$ is parameterized by the unknown parameter $$$\gamma$$$ to be estimated, $$$x_0$$$ is the ground-truth image, and $$$\eta$$$ is the measurement noise. We minimize the KL-divergence using a variational approach:
$$\min_q\;KL(q(x_0,\gamma|y)|p(x_0,\gamma|y)),\;\text{[Eq.2]}$$
where $$$q$$$ is a joint variational distribution that seeks the dominant mode of the posterior distribution $$$p$$$. When the image and forward model parameters are independent, the KL-divergence in Eq.2 can be expressed as:
$$\min_{q}KL\left(q(x_0|y)||p( x_0)\right)+KL\left(q(\gamma|y)||p(\gamma)\right)-\mathbb{E}_{q(x_0, \gamma|y)}\left[ \log p(y|x_0,\gamma)\right]+\log p(y),\;\text{[Eq.3]}$$
The first and third terms act as regularization on $$$x_0$$$ and data consistency, respectively, identically to RED-diff7. Therefore, term 1 can be represented as a score-matching regularization term implemented with a pre-trained diffusion model $$$\epsilon(x_t; t)$$$. The second term acts as regularization on $$$\gamma$$$. When $$$p(\gamma)$$$ has a specific distribution, e.g. Gaussian or Laplace, we can obtain a closed-form expression $$$R(\mu_\gamma,\sigma_\gamma)$$$. The optimization problem becomes:
$$\min_{\mu_{x},\mu_{\gamma}}\frac{1}{2\sigma_\eta^2}||y-f_{\mu_\gamma}(\mu_{x})||^2 + \mathbb{E}_{t,\epsilon}\left[ \lambda_t||\epsilon_{\theta}(x_t;t)-\epsilon||_2^2 \right]+\lambda_\gamma R(\mu_{\gamma},0),\;\text{[Eq.4]}$$
where we set $$$\sigma_\gamma=0$$$ for simplicity7. We solve Eq.4 using first-order stochastic optimization in an alternating fashion7 (Algorithm 1, Fig.2).

Methods

Datasets
MRI data was retrospectively simulated using ground-truth images from the fastMRI database9. Sensitivity maps were calculated using ESPIRiT10.For field-map correction, ground-truth field-maps were obtained from a separate brain dataset acquired with the Physical sequence11. Multi-channel k-space was simulated using a 16-shot variable-density spiral with . The temporal sampling rate was $$$4{\mu}s$$$ with total readout time of $$$15.6ms$$$ .
To demonstrate motion correction, we simulated motion-corrupted multi-channel k-space using a 3-shot EPI trajectory.

Implementation
We considered the following forward model:
$$ y=f_{\gamma}(x)+\eta=A_{\gamma}x_0+\eta=FST_{\gamma}x_0+\eta\;\text{[Eq.5],}$$
where $$$T_\gamma$$$ is the transformation operator with unknown parameter $$$\gamma$$$, $$$S$$$ is the the sensitivity map operator, and $$$F$$$ is the Fourier transform.
For field-map correction, $$$T_\gamma$$$ implements time-segmented off-resonance effects12 caused by the field inhomogeneities map $$$\psi(\gamma)$$$. We parameterize $$$\psi$$$ with a 5th-order spatial polynomial model with coefficients $$$\gamma$$$. Empirically, we observed $$$p(\gamma)$$$ approximates a Laplace $$$p_i(\gamma_i)\sim L(\tilde{\mu}_\gamma,\tilde{\sigma}_\gamma)$$$ from a dataset of 2,420 2D slices from 11 subjects. The regularizer becomes $$$\ell_1$$$-penalty13:
$$R(\mu_\gamma,0)=\frac{1}{{\tilde{\sigma}}_\gamma}\|\mu_\gamma-{\tilde{\mu}}_\gamma\|_1\;\text{[Eq.6]}$$
For motion correction, $$$T_\gamma$$$ implements time-segmented motion artifacts15. We assumed gaussian prior $$$p(\gamma)\sim{N}(\tilde{\mu}_\gamma,\tilde{\sigma}_{\gamma}I)$$$, which yielded $$$\ell_2$$$-penalty.
$$R(\mu_\gamma,0)=\frac{1}{{2\tilde{\sigma}}_\gamma}\|\mu_\gamma-{\tilde{\mu}}_\gamma\|^2_2\;\text{[Eq.7]}$$
Algorithm 1 was implemented by modifying the csgm-mri-langevin2 and SMRD14 libraries. For the image score function, we used the score function model from Jalal2. Reconstructions were run on a 24 GB NVIDIA Titan RTX. We performed three reconstructions for comparison:
  1. Linear reconstruction
  2. RED-diff7
  3. Blind RED-diff

Results

Results are shown in Fig.3 and 4. RED-diff is able to remove undersampling artifacts because of the diffusion prior, but not the off-resonance/motion artifacts due to the limitations of the model. Conversely, our proposed blind RED-diff is able to remove both undersampling and off-resonance/motion artifacts.

Discussion

We extended the RED-diff framework to blind inverse problems. Blind RED-diff requires a pre-trained diffusion model and the functional description of the forward model, which makes it applicable to multiple inverse problems without re-training or fine-tuning. We show that the blind RED-diff framework can successfully approximate the unknown forward model parameters and produce accurate reconstructions corrected for off-resonance and motion artifacts.
Two limitations of our work are the simple priors for the parameters $$$\gamma$$$ and the assumption of conditional independence between $$$x_0$$$ and $$$\gamma$$$. Future work will focus on developing much stronger diffusion-based priors for $$$\gamma$$$16, and for $$$x_0$$$ that consider conditional dependencies. Additionally, we plan to extend our framework to other MRI applications such as water-fat separation and quantitative parameter mapping.

Conclusion

Diffusion models have distinct potential for solving a wide range of MRI inverse problems. Our blind RED-diff method can produce accurate reconstructions with field-map and motion correction using a single pre-trained diffusion model and without fine-tuning or re-training.

Acknowledgements

This work was supported by NIH U01 EB029427.

References

  1. A. Jalal, S. Karmalkar, A. G. Dimakis, and E. Price, “Instance-optimal compressed sensing via posterior sampling,” International Conference on Machine Learning, 2021.
  2. A. Jalal, M. Arvinte, G. Daras, E. Price, A. G. Dimakis, and J. I. Tamir, “Robust compressed sensing mri with deep generative priors,” Advances in Neural Information Processing Systems 2021.
  3. H. Chung and J. C. Ye, “Score-based diffusion models for accelerated mri,” Medical image analysis , vol. 80, p. 102479, 2022.
  4. Y. Song, L. Shen, L. Xing, and S. Ermon, “Solving inverse problems in medical imaging with score-based generative models,” in International Conference on Learning Representations, 2021.
  5. H. Chung, J. Kim, M. T. Mccann, M. L. Klasky, and J. C. Ye, “Diffusion posterior sampling for general noisy inverse problems,” in International Conference on Learning Representations 2022.
  6. J. Song, A. Vahdat, M. Mardani, and J. Kautz, “Pseudoinverse-guided diffusion models for inverse problems,” in International Conference on Learning Representations, 2022.
  7. M. Mardani, J. Song, J. Kautz, and A. Vahdat, “A variational perspective on solving inverse problems with diffusion models,” arXiv preprint arXiv:2305.04391, 2023.
  8. B. Ozturkler, M. Mardani, A. Vahdat, J. Kautz, and J. M. Pauly, “Regularization by denoising diffusion process for mri reconstruction,” in Medical Imaging with Deep Learning, short paper track, 2023.
  9. F. Knoll, J. Zbontar, A. Sriram, M. J. Muckley, M. Bruno, A. Defazio, M. Parente, K. J. Geras, J. Katsnelson, H. Chandarana, et al., “fastmri: A publicly available raw k-space and dicom dataset of knee images for accelerated mr image reconstruction using machine learning,” Radiology: Artificial Intelligence, vol. 2, no. 1, p. e190007, 2020.
  10. M. Uecker, P. Lai, M. J. Murphy, P. Virtue, M. Elad, J. M. Pauly, S. S. Vasanawala, and M. Lustig, “Espirit—an eigenvalue approach to autocalibrating parallel mri: where sense meets grappa,” Magnetic resonance in medicine, vol. 71, no. 3, pp. 990–1001, 2014.
  11. S. S. Iyer, C. Liao, Q. Li, M. K. Manhard, A. Berman, B. Bilgic, and K. Setsompop, “Physical: A rapid calibration scan for b0, b1+, coil sensitivity and eddy current mapping,” in Proceedings of the 28th Annual Meeting of ISMRM, virtual, p. 0661, 2020B.
  12. P. Sutton, D. C. Noll, and J. A. Fessler, “Fast, iterative image reconstruction for mri in the presence of field inhomogeneities,” IEEE transactions on medical imaging, vol. 22, no. 2, pp. 178–188, 2003.
  13. G. P. Meyer, “An alternative probabilistic interpretation of the huber loss,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 5261–5269, June 2021.
  14. B. Ozturkler, C. Liu, B. Eckart, M. Mardani, J. Song, and J. Kautz, “Smrd: Sure-based robust mri reconstruction with diffusion models,” in International Conference on Medical Image Computing and Computer-Assisted Intervention, pp. 199–209, Springer, 2023.
  15. L. Cordero-Grande, et al. "Sensitivity encoding for aligned multishot magnetic resonance reconstruction." IEEE Transactions on Computational Imaging 2.3 (2016): 266-280.
  16. H. Chung, J. Kim, S. Kim, and J. C. Ye, “Parallel diffusion models of operator and image for blind inverse problems,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6059–6069, 2023

Figures

Figure 1. Blind RED-diff diagram. Our proposed algorithm extends the RED-diff framework to blind inverse problems. Blind RED-diff combines data-consistency loss with score-matching regularization from denoisers at different time-steps and forward model parameter prior.

Figure 2. Blind RED-diff Algorithm. The notation $$$sg(\cdot)$$$ indicates “stopped-gradient”, which indicates that the score is not differentiated with respect to $$$\mu_x$$$ during optimization.

Figure 3. Results off-resonance. Blind RED-diff is able to simultaneously resolve the off-resonance blurring and remove $$$R=2$$$ undersampling artifacts by concurrently estimating image and the forward model parameter. Conversely, RED-diff does not remove blurring artifacts. Blurring is mainly located in the regions of the ground truth field map where off-resonance is stronger.

Figure 4. Results motion. Blind RED-diff is able to simultaneously resolve the motion artifacts and remove $$$R=2$$$ undersampling artifacts by concurrently estimating image and the forward model parameter. Conversely, RED-diff does not remove motion artifacts.

Proc. Intl. Soc. Mag. Reson. Med. 32 (2024)
0377
DOI: https://doi.org/10.58530/2024/0377