Towards performant and reliable undersampled MR reconstruction via diffusion model sampling

[2203.04292] Towards performant and reliable undersampled MR reconstruction via diffusion model sampling

1 1 institutetext:
Johns Hopkins University, MD, USA 2 2 institutetext: Medical Imaging, Robotics, and Analytic Computing Laboratory and Engineering (MIRACLE) Center, School of Biomedical Engineering & Suzhou Institute for Advance Research, University of Science and Technology of China, Suzhou, China 3 3 institutetext: Key Lab of Intelligent Information Processing of Chinese Academy of Sciences (CAS), Institute of Computing Technology, CAS, Beijing, China 3 3 email: {cpeng26,pguo4,vpatel36,rchella4}@jhu.edu, s.kevin.zhou@gmail.com

Towards performant and reliable undersampled MR reconstruction via diffusion model sampling

Cheng Peng 11
  
Pengfei Guo 11
  
S. Kevin Zhou 2233
  
Vishal Patel 11
  
Rama Chellappa 11

Abstract

Magnetic Resonance (MR) image reconstruction from under-sampled acquisition promises faster scanning time. To this end, current State-of-The-Art (SoTA) approaches leverage deep neural networks and supervised training to learn a recovery model. While these approaches achieve impressive performances, the learned model can be fragile on unseen degradation, e.g. when given a different acceleration factor. These methods are also generally deterministic and provide a single solution to an ill-posed problem; as such, it can be difficult for practitioners to understand the reliability of the reconstruction. We introduce DiffuseRecon, a novel diffusion model-based MR reconstruction method. DiffuseRecon guides the generation process based on the observed signals and a pre-trained diffusion model, and does not require additional training on specific acceleration factors. DiffuseRecon is stochastic in nature and generates results from a distribution of fully-sampled MR images; as such, it allows us to explicitly visualize different potential reconstruction solutions. Lastly, DiffuseRecon proposes an accelerated, coarse-to-fine Monte-Carlo sampling scheme to approximate the most likely reconstruction candidate. The proposed DiffuseRecon achieves SoTA performances reconstructing from raw acquisition signals in fastMRI and SKM-TEA. Code will be open-sourced at www.github.com/cpeng93/DiffuseRecon .

1

Introduction

Magnetic Resonance Imaging (MRI) is a widely used medical imaging technique. It offers several advantages over other imaging modalities, such as providing high contrast on soft tissues and introducing no harmful radiation during acquisition. However, MRI is also limited by its long acquisition time due to the underlying imaging physics and machine quality. This leads to various issues ranging from patient discomfort to limited accessibility of the machines.

An approach to shorten MR scanning time is by under-sampling the signal in k-space during acquisition and recovering it by performing a post-process reconstruction algorithm. Recovering unseen signal is a challenging, ill-posed problem, and there has been a long history of research in addressing undersampled MR reconstruction. In general, this problem is formulated as:

    y  recon   =     arg  ⁡    min  y   ⁡   ∥    ℳ  ​  ℱ  ​  y   −   x  obs    ∥     +   λ  ​  R  ​   (  y  )     ,  s    .  t  .    x  obs   =   ℳ  ​   x  full      ,     formulae-sequence      subscript  𝑦   recon     subscript    𝑦         ℳ  ℱ  𝑦    subscript  𝑥   obs          𝜆  𝑅  𝑦    𝑠    𝑡      subscript  𝑥   obs       ℳ   subscript  𝑥   full        y_{\textrm{recon}}=\arg\min_{y}\lVert\mathcal{M}\mathcal{F}y-x_{\textrm{obs}}\rVert+\lambda R(y),~{}~{}s.t.~{}~{}x_{\textrm{obs}}=\mathcal{M}x_{\textrm{full}},    

(1)

where x { full,obs } subscript 𝑥 full,obs x_{{\textrm{full,obs}}} denotes the fully-sampled and under-sampled k-space signal, ℳ ℳ \mathcal{M} denotes the undersampling mask, and ℱ ℱ \mathcal{F} denotes the Fourier operator. The goal is to find an MR image y 𝑦 y such that its k-space content ℳ ​ ℱ ​ y ℳ ℱ 𝑦 \mathcal{M}\mathcal{F}y agrees with x obs subscript 𝑥 obs x_{\textrm{obs}} ; this is often known as the data consistency term. Furthermore, y recon subscript 𝑦 recon y_{\textrm{recon}} should follow certain prior knowledge about MR images, as expressed by the regularization term R ​ ( ∗ ) 𝑅 R() . The design of R 𝑅 R is subject to many innovations. Lustig et al.  [ 12 ] first proposed to use Compressed Sensing motivated ℓ 1 subscript ℓ 1 \ell_{1} -minimization algorithm for MRI reconstruction, assuming that the undersampled MR images have a sparse representation in a transform domain. Ravishankar et al.  [ 15 ] applied a more adaptive sparse modelling through Dictionary Learning, where the transformation is optimized through a set of data, leading to improved sparsity encoding. As interests grow in this field, more MR data has become publicly available. Consequently, advances in Deep Learning (DL), specifically with supervised learning, have been widely applied in MR reconstruction. Generally, DL-based methods  [ 5 , 24 , 6 , 17 , 16 , 1 , 9 , 21 , 22 , 8 , 14 , 19 , 7 ] train Convolutional Neural Networks (CNNs) with paired data { y und , y full } = { ℱ − 1 ​ x obs , ℱ − 1 ​ x full } subscript 𝑦 und subscript 𝑦 full superscript ℱ 1 subscript 𝑥 obs superscript ℱ 1 subscript 𝑥 full {y_{\textrm{und}},y_{\textrm{full}}}={\mathcal{F}^{-1}x_{\textrm{obs}},\mathcal{F}^{-1}x_{\textrm{full}}} . Following the formulation of Eq. ( 1 ), data consistency can be explicitly enforced within the CNN by replacing ℳ ​ ℱ ​ y ℳ ℱ 𝑦 \mathcal{M}\mathcal{F}y with x obs subscript 𝑥 obs x_{\textrm{obs}}   [ 17 ] . The resulting CNN serves as a parameterized R ​ ( ∗ , θ ) 𝑅 𝜃 R(,\theta) , and regularizes test images based on learned θ 𝜃 \theta from the training distribution.

While supervised DL-based methods have led to impressive results, these methods generally train CNNs under specific degradation conditions; e.g., the under-sampling mask ℳ ℳ \mathcal{M} that generates y und subscript 𝑦 und y_{\textrm{und}} follows a particular acceleration factor. As a consequence, when the acceleration factor changes, the performances of various models often degrade significantly, making R ​ ( ∗ , θ ) 𝑅 𝜃 R(*,\theta) less reliable in general. Furthermore, while the trained CNNs infer the most likely estimation of y full subscript 𝑦 full y_{\textrm{full}} given y und subscript 𝑦 und y_{\textrm{und}} , they do not provide possible alternative solutions. Since Eq.  1 is an under-constrained problem, y und subscript 𝑦 und y_{\textrm{und}} can have many valid solutions. The ability to observe different solutions can help practitioners understand the potential variability in reconstructions and make more robust decisions. As such, finding a stochastic R 𝑅 R that is generally applicable across all x obs subscript 𝑥 obs x_{\textrm{obs}} is of high interest.

We leverage the recent progress in a class of generative methods called diffusion models   [ 10 , 13 , 4 ] , which use a CNN to perform progressive reverse-diffusion and maps a prior Gaussian distribution 𝒩 𝒩 \mathcal{N} to a learned image distribution, e.g. p θ ​ ( y full ) subscript 𝑝 𝜃 subscript 𝑦 full p_{\theta}(y_{\textrm{full}}) . Based on a pre-trained θ 𝜃 \theta , we propose to guide the iterative reverse-diffusion by gradually introducing x obs subscript 𝑥 obs x_{\textrm{obs}} to the intermediate results, as shown in Fig. 1. This allows us to generate reconstructions in the marginal distribution p θ ​ ( y full | x obs ) subscript 𝑝 𝜃 conditional subscript 𝑦 full subscript 𝑥 obs p_{\theta}(y_{\textrm{full}}|x_{\textrm{obs}}) , where any sample y recon ∼ p θ ​ ( y full | x obs ) similar-to subscript 𝑦 recon subscript 𝑝 𝜃 conditional subscript 𝑦 full subscript 𝑥 obs y_{\textrm{recon}}\sim p_{\theta}(y_{\textrm{full}}|x_{\textrm{obs}}) agrees with the observed signal and lies on the MR image distribution p θ ​ ( y full ) subscript 𝑝 𝜃 subscript 𝑦 full p_{\theta}(y_{\textrm{full}}) .

We make three contributions. (i) Our proposed DiffuseRecon performs MR reconstruction by gradually guiding the reverse-diffusion process with observed k-space signal and is robust to changing sampling condition using a single model. (ii) We propose a coarse-to-fine sampling algorithm, which allows us to estimate the most likely reconstruction and its variability within p θ ​ ( y full | x obs ) subscript 𝑝 𝜃 conditional subscript 𝑦 full subscript 𝑥 obs p_{\theta}(y_{\textrm{full}}|x_{\textrm{obs}}) while leading to an approximately 40 × \times speed-up compared to naive sampling. (iii) We perform extensive experiments using raw acquisition signals from fastMRI  [ 23 ] and SKM-TEA  [ 3 ] , and demonstrate superior performance over SoTA methods.

Figure 1 : DiffuseRecon gradually incorporates x obs subscript 𝑥 obs x_{\textrm{obs}} into the denoising process through a k-Space Guidance (KSG) module; as such, we can directly generate samples from p θ ​ ( y full | x obs ) subscript 𝑝 𝜃 conditional subscript 𝑦 full subscript 𝑥 obs p_{\theta}(y_{\textrm{full}}|x_{\textrm{obs}}) . Visualizations are based on 8X undersampling.

2

DiffuseRecon

Background. Diffusion model  [ 10 , 13 , 4 ] is a class of unconditional generative methods that aims to transform a Gaussian distribution to the empirical data distribution. Specifically, the forward diffusion process is a Markov Chain that gradually adds Gaussian noise to a clean image, which can be expressed as:

    q  ​   (    y  t   |   y   t  −  1     )    =   𝒩  ​   (   y  t   ;     1  −   β  t     ​   y   t  −  1     ,    β  t   ​  𝐈   )     ;    q  ​   (    y  t   |   y  0    )    =   𝒩  ​   (   y  t   ;      α  ¯   t    ​   y  0    ,    (   1  −    α  ¯   t    )   ​  𝐈   )      ,     formulae-sequence        𝑞   conditional   subscript  𝑦  𝑡    subscript  𝑦     𝑡  1         𝒩    subscript  𝑦  𝑡            1   subscript  𝛽  𝑡      subscript  𝑦     𝑡  1         subscript  𝛽  𝑡   𝐈            𝑞   conditional   subscript  𝑦  𝑡    subscript  𝑦  0        𝒩    subscript  𝑦  𝑡          subscript   ¯  𝛼   𝑡     subscript  𝑦  0          1   subscript   ¯  𝛼   𝑡    𝐈        \displaystyle q(y_{t}|y_{t-1})=\mathcal{N}(y_{t};\sqrt{1-\beta_{t}}y_{t-1},\beta_{t}\mathbf{I});q(y_{t}|y_{0})=\mathcal{N}(y_{t};\sqrt{\bar{\alpha}_{t}}y_{0},(1-\bar{\alpha}_{t})\mathbf{I}),    

(2)

where y t subscript 𝑦 𝑡 y_{t} denotes the intermediate noisy images, β t subscript 𝛽 𝑡 \beta_{t} denotes the noise variance, α t = 1 − β t subscript 𝛼 𝑡 1 subscript 𝛽 𝑡 \alpha_{t}=1-\beta_{t} , and α ¯ t = ∏ s = 1 t α s subscript ¯ 𝛼 𝑡 superscript subscript product 𝑠 1 𝑡 subscript 𝛼 𝑠 \bar{\alpha}{t}=\prod{s=1}^{t}\alpha_{s} . For simplicity, β t subscript 𝛽 𝑡 \beta_{t} follows a fixed schedule. When T 𝑇 T is sufficiently large, y T subscript 𝑦 𝑇 y_{T} is dominated by noise. A CNN model is introduced to gradually reverse the forward diffusion process, i.e. denoise y t subscript 𝑦 𝑡 y_{t} , by estimating

    p  θ   ​   (    y   t  −  1    |   y  t    )    =   𝒩  ​   (   y   t  −  1    ;    ϵ  θ   ​   (   y  t   ,  t  )    ,    σ  t  2   ​  𝐈   )     ,           subscript  𝑝  𝜃    conditional   subscript  𝑦     𝑡  1     subscript  𝑦  𝑡        𝒩    subscript  𝑦     𝑡  1        subscript  italic-ϵ  𝜃     subscript  𝑦  𝑡   𝑡        superscript   subscript  𝜎  𝑡   2   𝐈       \displaystyle p_{\theta}(y_{t-1}|y_{t})=\mathcal{N}(y_{t-1};\epsilon_{\theta}(y_{t},t),\sigma_{t}^{2}\mathbf{I}),    

(3)

where σ t 2 = β t ¯ = 1 − α ¯ t − 1 1 − α ¯ t ​ β t superscript subscript 𝜎 𝑡 2 ¯ subscript 𝛽 𝑡 1 subscript ¯ 𝛼 𝑡 1 1 subscript ¯ 𝛼 𝑡 subscript 𝛽 𝑡 \sigma_{t}^{2}=\bar{\beta_{t}}=\frac{1-\bar{\alpha}{t-1}}{1-\bar{\alpha}{t}}\beta_{t} in this work. With fixed variance, ϵ θ subscript italic-ϵ 𝜃 \epsilon_{\theta} in Eq. ( 3 ) is trained by mean-matching the diffusion noise through an ℒ 2 subscript ℒ 2 \mathcal{L}_{2} loss. We follow  [ 13 ] with some modifications to train a diffusion model that generates MR images.

After training on MR images y full subscript 𝑦 full y_{\textrm{full}} , ϵ θ subscript italic-ϵ 𝜃 \epsilon_{\theta} can be used to gradually transform noise into images that follow the distribution p θ ​ ( y full ) subscript 𝑝 𝜃 subscript 𝑦 full p_{\theta}(y_{\textrm{full}}) , as shown in Fig.  1 . We note that diffusion models generate images unconditionally by their original design. To reconstruct high fidelity MR images conditioned on x o ​ b ​ s subscript 𝑥 𝑜 𝑏 𝑠 x_{obs} , DiffuseRecon is proposed to gradually modify y t subscript 𝑦 𝑡 y_{t} such that y 0 subscript 𝑦 0 y_{0} agrees with x o ​ b ​ s subscript 𝑥 𝑜 𝑏 𝑠 x_{obs} . DiffuseRecon consists of two parts, k-Space Guidance and Coarse-to-Fine Sampling.

K-Space Guidance. We note that R ​ ( ∗ ) 𝑅 R(*) is naturally enforced by following the denoising process with a pre-trained θ 𝜃 \theta , as the generated images already follow the MR data distribution. As such, k-Space Guidance (KSG) is proposed to ensure the generated images follow the data consistency term. In a diffusion model, a denoised image is generated at every step t 𝑡 t by subtracting the estimated noise from the previous y t subscript 𝑦 𝑡 y_{t} , specifically:

    y   t  −  1   ′   =      1    α  t      ​   (    y  t   −      1  −   α  t      1  −    α  ¯   t       ​   ϵ  θ   ​   (   y  t   ,  t  )     )    +    σ  t   ​  𝐳     ,   𝐳  ∼   𝒩  ​   (  0  ,  𝐈  )      .     formulae-sequence      subscript   superscript  𝑦  ′      𝑡  1    1      subscript  𝛼  𝑡         subscript  𝑦  𝑡            1   subscript  𝛼  𝑡          1   subscript   ¯  𝛼   𝑡       subscript  italic-ϵ  𝜃     subscript  𝑦  𝑡   𝑡          subscript  𝜎  𝑡   𝐳      similar-to  𝐳     𝒩   0  𝐈       \displaystyle y^{\prime}_{t-1}=\frac{1}{\sqrt{\alpha_{t}}}(y_{t}-\frac{1-\alpha_{t}}{\sqrt{1-\bar{\alpha}_{t}}}\epsilon_{\theta}(y_{t},t))+\sigma_{t}\mathbf{z},\mathbf{z}\sim\mathcal{N}(0,\mathbf{I}).    

(4)

For unconditional generation, y t − 1 = y t − 1 ′ subscript 𝑦 𝑡 1 subscript superscript 𝑦 ′ 𝑡 1 y_{t-1}=y^{\prime}{t-1} , and the process repeats until t = 0 𝑡 0 t=0 . For DiffuseRecon, KSG gradually mixes observed k-space signal x obs subscript 𝑥 obs x{\textrm{obs}} with y t − 1 ′ subscript superscript 𝑦 ′ 𝑡 1 y^{\prime}{t-1} . To do so, KSG first adds a zero-mean noise on x obs subscript 𝑥 obs x{\textrm{obs}} to simulate its diffused condition at step t 𝑡 t . The noisy observation is then mixed with y t − 1 ′ subscript superscript 𝑦 ′ 𝑡 1 y^{\prime}_{t-1} in k-space based on the undersampling mask ℳ ℳ \mathcal{M} . This process is expressed as follows:

   y  t   =    ℱ   −  1    ​   (     (   1  −  ℳ   )   ​  ℱ  ​   y  t  ′    +   ℳ  ​   x   obs  ,  t      )     ,        subscript  𝑦  𝑡       superscript  ℱ     1    1  ℳ   ℱ   subscript   superscript  𝑦  ′   𝑡       ℳ   subscript  𝑥    obs   𝑡         \displaystyle y_{t}=\mathcal{F}^{-1}((1-\mathcal{M})\mathcal{F}y^{\prime}_{t}+\mathcal{M}x_{\textrm{obs},t}),    

(5)

   where   ​   x   obs  ,  t     =    x  obs   +   ℱ  ​   (   𝒩  ​   (  0  ,    (   1  −    α  ¯   t    )   ​  𝐈   )    )      .           where     subscript  𝑥    obs   𝑡         subscript  𝑥   obs       ℱ     𝒩   0        1   subscript   ¯  𝛼   𝑡    𝐈         \displaystyle\text{where }x_{\textrm{obs},t}=x_{\textrm{obs}}+\mathcal{F}(\mathcal{N}(0,(1-\bar{\alpha}_{t})\mathbf{I})).    

The resulting y t − 1 subscript 𝑦 𝑡 1 y_{t-1} is iteratively denoised with decreasing t 𝑡 t until y 0 subscript 𝑦 0 y_{0} is obtained, and x obs , 0 = x obs subscript 𝑥 obs 0 subscript 𝑥 obs x_{\textrm{obs},0}=x_{\textrm{obs}} . As such, y 0 subscript 𝑦 0 y_{0} shares the same observed k-space signal and achieves the data consistency term ℳ ​ ℱ ​ y 0 = x obs ℳ ℱ subscript 𝑦 0 subscript 𝑥 obs \mathcal{M}\mathcal{F}y_{0}=x_{\textrm{obs}} . Since all generated samples fall on the data distribution p θ ​ ( y full ) subscript 𝑝 𝜃 subscript 𝑦 full p_{\theta}(y_{\textrm{full}}) , KSG allows us to stochastically sample from the marginal distribution p θ ​ ( y full | x obs ) subscript 𝑝 𝜃 conditional subscript 𝑦 full subscript 𝑥 obs p_{\theta}(y_{\textrm{full}}|x_{\textrm{obs}}) through a pre-trained diffusion model. As demonstrated in Fig. 2a, we can observe the variations in y 0 subscript 𝑦 0 y_{0} with a given x obs subscript 𝑥 obs x_{\textrm{obs}} and determine the reliability of the reconstructed results. Furthermore, KSG is parameter-free and applicable to any undersampling mask ℳ ℳ \mathcal{M} without finetuning.

Mask Init. 1 Init. 2 Init. 3 Var.

(a) (a)

(b) (b)

Figure 3 : (a) Visualizations on various under-sampled k-space signal (8 × \times , 16 × \times , 32 × \times , 1D Gaussian Mask). (b) To accelerate the sampling process for MC simulation, multiple coarse samples are generated using T k 𝑇 𝑘 \frac{T}{k} steps. These samples are averaged to y 0 avg subscript superscript 𝑦 avg 0 y^{\textrm{avg}}{0} and reintroduced to the denoising network for T refine subscript 𝑇 refine T{\textrm{refine}} steps.

Coarse-to-Fine Sampling. While sampling in p θ ​ ( y full | x obs ) subscript 𝑝 𝜃 conditional subscript 𝑦 full subscript 𝑥 obs p_{\theta}(y_{\textrm{full}}|x_{\textrm{obs}}) allows us to generate different reconstruction candidates, selecting the best reconstruction out of all candidates is a challenge. Since we do not have an analytic solution for p θ ​ ( y full | x obs ) subscript 𝑝 𝜃 conditional subscript 𝑦 full subscript 𝑥 obs p_{\theta}(y_{\textrm{full}}|x_{\textrm{obs}}) , Monte-Carlo (MC) sampling can be used to estimate 𝔼 ​ ( p θ ​ ( y full | x obs ) ) 𝔼 subscript 𝑝 𝜃 conditional subscript 𝑦 full subscript 𝑥 obs \mathbb{E}(p_{\theta}(y_{\textrm{full}}|x_{\textrm{obs}})) in an unbiased way with sufficient samples. However, as sampling from a diffusion model already requires multiple steps, generating multiple samples is computationally costly. In practice, the denoising process can be accelerated by evenly spacing the { T , T − 1 , … , 1 } 𝑇 𝑇 1 … 1 {T,T-1,...,1} schedule to a shorter schedule { T , T − k , … , 1 } 𝑇 𝑇 𝑘 … 1 {T,T-k,...,1}   [ 13 ] , where k > 1 𝑘 1 k>1 , and modifying the respective weights β , α 𝛽 𝛼 \beta,\alpha based on k 𝑘 k . However, the acceleration tends to produce less denoised results when k 𝑘 k is large.

We propose a Coarse-to-Fine (C2F) sampling algorithm to greatly accelerate the MC sampling process without loss in performance. Specifically, we note that the diffusion noise that is added in Eq. ( 2 ) is zero-mean with respect to y 0 subscript 𝑦 0 y_{0} , and can be reduced by averaging multiple samples. Since multiple samples are already required to estimate 𝔼 ​ ( p θ ​ ( y full | x obs ) ) 𝔼 subscript 𝑝 𝜃 conditional subscript 𝑦 full subscript 𝑥 obs \mathbb{E}(p_{\theta}(y_{\textrm{full}}|x_{\textrm{obs}})) , we leverage the multi-sample scenario to more aggressively reduce the denoising steps. As shown in Fig. 2b, C2F sampling creates N 𝑁 N instances of y T i , T k ∼ 𝒩 ​ ( 0 , 𝐈 ) similar-to subscript superscript 𝑦 𝑖 𝑇 𝑘 𝑇 𝒩 0 𝐈 y^{i,\frac{T}{k}}{T}\sim\mathcal{N}(0,\mathbf{I}) and applies denoising individually for T k 𝑇 𝑘 \frac{T}{k} steps based on a re-spaced schedule. The noisy results are averaged to produce y 0 avg = 1 N ​ ∑ i = 0 N y 0 i , T k subscript superscript 𝑦 avg 0 1 𝑁 subscript superscript 𝑁 𝑖 0 subscript superscript 𝑦 𝑖 𝑇 𝑘 0 y^{\textrm{avg}}{0}=\frac{1}{N}\sum^{N}{i=0}y^{i,\frac{T}{k}}{0} . Finally, y 0 avg subscript superscript 𝑦 avg 0 y^{\textrm{avg}}{0} is refined by going through additional T refine subscript 𝑇 refine T{\textrm{refine}} steps with ϵ θ subscript italic-ϵ 𝜃 \epsilon_{\theta} . To control the denoising strength in ϵ θ subscript italic-ϵ 𝜃 \epsilon_{\theta} , { T refine , T refine − 1 , … , 1 } ∈ { T , T − 1 , … , 1 } , T refine ≪ T formulae-sequence subscript 𝑇 refine subscript 𝑇 refine 1 … 1 𝑇 𝑇 1 … 1 much-less-than subscript 𝑇 refine 𝑇 {T_{\textrm{refine}},T_{\textrm{refine}}-1,...,1}\in{T,T-1,...,1},T_{\textrm{refine}}\ll T . The last refinement steps help remove blurriness introduced by averaging multiple samples and lead to more realistic reconstruction results. During the refinement steps, x obs subscript 𝑥 obs x_{\textrm{obs}} directly replaces k-space signals in the reconstructions, as is consistent with ℳ ​ ℱ ​ y 0 i , T k ℳ ℱ subscript superscript 𝑦 𝑖 𝑇 𝑘 0 \mathcal{M}\mathcal{F}y^{i,\frac{T}{k}}_{0} .

Compared to a naive approach which requires T ​ N 𝑇 𝑁 TN total steps to estimate 𝔼 ​ ( p θ ​ ( y full | x obs ) ) 𝔼 subscript 𝑝 𝜃 conditional subscript 𝑦 full subscript 𝑥 obs \mathbb{E}(p_{\theta}(y_{\textrm{full}}|x_{\textrm{obs}})) , C2F sampling introduces an approximately k 𝑘 k -time speed-up. Furthermore, while y 0 i , T k subscript superscript 𝑦 𝑖 𝑇 𝑘 0 y^{i,\frac{T}{k}}{0} are noisy compared to their fully-denoised version y 0 i , T subscript superscript 𝑦 𝑖 𝑇 0 y^{i,T}{0} , their noise is approximately Gaussian and introduces only a constant between V ​ a ​ r ​ ( y 0 i , T k ) 𝑉 𝑎 𝑟 subscript superscript 𝑦 𝑖 𝑇 𝑘 0 Var(y^{i,\frac{T}{k}}{0}) and V ​ a ​ r ​ ( y 0 i , T ) 𝑉 𝑎 𝑟 subscript superscript 𝑦 𝑖 𝑇 0 Var(y^{i,T}{0}) . As such, given a reasonable N 𝑁 N , variance in p θ ​ ( y full | x obs ) subscript 𝑝 𝜃 conditional subscript 𝑦 full subscript 𝑥 obs p_{\theta}(y_{\textrm{full}}|x_{\textrm{obs}}) can still be estimated well from coarse samples.

3

Experiments

Dataset. In our experiments, we use two large knee MR datasets that contain raw, complex-value acquisition signals. FastMRI  [ 23 ] contains 1172 subjects, with approximately 35 slices per subject and is split such that 973 subjects are used for training and 199 subjects are used for evaluation. SKM-TEA  [ 3 ] contains 155 subjects, with approximately 160 slices per subject; 134 subjects are used for training, 21 subjects are used for evaluation. Single-coil data is used for both datasets. Undersampling masks ℳ ℳ \mathcal{M} are generated by using the masking function provided in the fastMRI challenge with 6 × \times and 8 × \times accelerations. To avoid slices with little information, the first five and ten slices are removed from evaluation for fastMRI and SKM-TEA respectively. To accommodate methods based on over-complete CNNs, which require upsampling and a larger memory footprint, images from SKM-TEA are center-cropped to ℝ 448 × 448 superscript ℝ 448 448 \mathbb{R}^{448\times 448} .

Implementation Details. All baselines are implemented in PyTorch and trained with ℒ 1 subscript ℒ 1 \mathcal{L}{1} loss and Adam optimizer. Following  [ 6 ] , a learning rate is set to 1.5 × 10 − 4 1.5 superscript 10 4 1.5\times 10^{-4} and reduced by 90 % percent 90 90% every five epochs; to ensure a fair comparison, U-Net  [ 16 ] is implemented with a data consistency layer at the end. For DiffuseRecon, we follow  [ 13 ] by using a cosine noise schedule and a U-Net architecture with multi-head attention as the denoising model. The model is modified to generate two consecutive slices; as such, the input and output channel sizes are 4, accounting for complex values. For C2F Sampling, { T , k , N , T refine } = { 4000 , 40 , 10 , 20 } 𝑇 𝑘 𝑁 subscript 𝑇 refine 4000 40 10 20 {T,k,N,T{\textrm{refine}}}={4000,40,10,20} , which gives a T ​ N T ​ N k + T refine ≈ 39 𝑇 𝑁 𝑇 𝑁 𝑘 subscript 𝑇 refine 39 \frac{TN}{\frac{TN}{k}+T_{\textrm{refine}}}\approx 39 times speed-up compared to naive MC sampling.

(a) (a)

D5C5 OUCR Ours GT

fastMRI, 8 × \times → → \rightarrow 4 × \times

SKM-TEA, 6 × \times → → \rightarrow 10 × \times

(b) 33.1/.791

(c) 34.3/.827

(d) 35.6/.864

(e)

(f) 28.1/.510

(g) 28.2/.519

(h) 33.4/.736

(i)

(j) (b)

Figure 4 : (a) Ablation Study on DiffuseRecon with different parameters. (b) Visualizations on recovering from unseen acceleration factors.

Ablation Study. We examine the effectiveness of the following design choices:

•
DiffuseRecon nonoise k , N subscript superscript DiffuseRecon 𝑘 𝑁 nonoise {\textrm{DiffuseRecon}}^{k,N}{\textrm{nonoise}} : ℳ ​ ℱ ​ y t ′ ℳ ℱ subscript superscript 𝑦 ′ 𝑡 \mathcal{M}\mathcal{F}y^{\prime}{t} is directly replaced with x obs subscript 𝑥 obs x_{\textrm{obs}} in KSG instead of x obs , t subscript 𝑥 obs 𝑡 x_{\textrm{obs},t} ; k = { 40 } 𝑘 40 k={40} and N = { 1 , 2 , 5 , 10 } 𝑁 1 2 5 10 N={1,2,5,10} are tested.

•
DiffuseRecon norefine k , N subscript superscript DiffuseRecon 𝑘 𝑁 norefine {\textrm{DiffuseRecon}}^{k,N}_{\textrm{norefine}} : a combination of { k , N } 𝑘 𝑁 {k,N} are tested without the refining steps; specifically, k = { 8 , 40 } 𝑘 8 40 k={8,40} and N = { 1 , 2 , 5 , 10 } 𝑁 1 2 5 10 N={1,2,5,10} .

•
DiffuseRecon refine k , N subscript superscript DiffuseRecon 𝑘 𝑁 refine {\textrm{DiffuseRecon}}^{k,N}{\textrm{refine}} : T refine = 20 subscript 𝑇 refine 20 T{\textrm{refine}}=20 steps of further denoising is applied to the aggregate results from DiffuseRecon norefine k , N subscript superscript DiffuseRecon 𝑘 𝑁 norefine {\textrm{DiffuseRecon}}^{k,N}_{\textrm{norefine}} , { k , N } = { { 8 , 40 } , 10 } 𝑘 𝑁 8 40 10 {k,N}={{8,40},10} .

The PSNR comparisons are visualized in Fig.  4 (a) and are made based on the middle fifteen slices for all subjects in the fastMRI’s evaluation set. There are several interesting takeaways. Firstly, PSNR increases significantly with larger N 𝑁 N for all instances, demonstrating the importance of multi-sample aggregation. We note that k = 8 𝑘 8 k=8 applies 500 steps and yields sufficiently denoised results; as such, the low PSNR value for DiffuseRecon norefine 8 , 1 subscript superscript DiffuseRecon 8 1 norefine {\textrm{DiffuseRecon}}^{8,1}{\textrm{norefine}} is due to the geometric variability between sampled and groundtruth image. Secondly, when no gradual noise is added to x obs subscript 𝑥 obs x{\textrm{obs}} , the reconstruction results are significantly better when N = 1 𝑁 1 N=1 and worse when N = 10 𝑁 10 N=10 compared to the proposed KSG; such a direct replacement approach is used in a concurrent work POCS  [ 2 ] . This indicates that, while the clean x obs subscript 𝑥 obs x_{\textrm{obs}} introduced at t = T 𝑡 𝑇 t=T helps accelerate the denoising process, the aggregate results do not estimate the ground-truth well, i.e. they are more likely to be biased. Finally, the refinement steps significantly remove the blurriness caused by averaging multiple samples. We note that performances are very similar for DiffuseRecon refine 8 , 10 subscript superscript DiffuseRecon 8 10 refine {\textrm{DiffuseRecon}}^{8,10}{\textrm{refine}} and DiffuseRecon refine 40 , 10 subscript superscript DiffuseRecon 40 10 refine {\textrm{DiffuseRecon}}^{40,10}{\textrm{refine}} ; as such, k = 8 𝑘 8 k=8 leads to significant speed-up without sacrificing reconstruction quality.

Table 1 : Quantitative volume-wise evaluation of DiffuseRecon against SoTA methods. Dedicated models are trained for 6 × \times and 8 × \times acceleration; model robustness is tested by applying { 6 × , 8 × } {6\times,8\times} model on { 10 × , 4 × } {10\times,4\times} downsampled inputs. All baselines achieved similar performance compared to the original papers.

Method
fastMRI SKM-TEA

6 × \times
8 × \times
8 × \times → → \rightarrow 4 × \times
6 × \times → → \rightarrow 10 × \times
6 × \times
8 × \times
8 × \times → → \rightarrow 4 × \times
6 × \times → → \rightarrow 10 × \times

UNet  [ 16 ]
29.79 28.89 30.67 21.98 31.62 30.47 32.06 24.34

0.621 0.577 0.666 0.296 0.713 0.669 0.728 0.382

KIKI-Net  [ 5 ]
29.51 28.09 30.18 22.31 31.67 30.14 32.20 24.67

0.607 0.546 0.650 0.313 0.711 0.655 0.732 0.422

D5C5  [ 17 ]
29.88 28.63 30.90 23.07 32.22 30.86 32.99 25.99

0.622 0.565 0.675 0.349 0.732 0.683 0.763 0.512

OUCR  [ 6 ]
30.44 29.56 31.33 23.41 32.52 31.27 33.11 26.17

0.644 0.600 0.689 0.371 0.742 0.696 0.766 0.516

DiffuseRecon
30.56
29.94
31.70
27.23
32.58
31.56
33.76
28.40

0.648
0.614
0.708
0.515
0.743
0.706
0.795
0.584

Quantitative and Visual Evaluation. We compare reconstruction results from DiffuseRecon with current SoTA methods, and summarize the quantitative results in Table  1 . KIKI-Net  [ 5 ] applies convolutions on both the image and k-space data for reconstruction. D5C5  [ 17 ] is a seminal DL-based MR reconstruction work and proposes to incorporate data consistency layers into a cascade of CNNs. OUCR  [ 6 ] builds on  [ 17 ] and uses a recurrent over-complete CNN  [ 20 ] architecture to better recovery fine details. We note that these methods use supervised learning and train dedicated models based on a fixed acceleration factor. In Table  1 , DiffuseRecon is compared with these specifically trained models for acceleration factor of 6 × 6\times and 8 × 8\times in k-space. Although the model for DiffuseRecon is trained for a denoising task and has not observed any down-sampled MR images, DiffuseRecon obtains top performances compared to the dedicated models from current SoTAs; the performance gap is larger as the acceleration factor becomes higher. To examine the robustness of dedicated models on lower and higher quality inputs, we apply models trained on 6 × 6\times and 8 × 8\times acceleration on 10 × 10\times and 4 × 4\times undersampled inputs, respectively. In these cases, the performances of DiffuseRecon are significantly higher, demonstrating that 1. models obtained from supervised training are less reliable on images with unseen levels of degradation, and 2. DiffuseRecon is a general and performant MR reconstruction method.

Visualization of the reconstructed images is provided in Fig.  4 (b) and Fig.  5 to illustrate the advantages of DiffuseRecon. Due to limited space, we focus on the top three methods: D5C5  [ 17 ] , OUCR  [ 6 ] , and DiffuseRecon. We observe that many significant structural details are lost in D5C5 and OUCR but can be recovered by DiffuseRecon. The errors in D5C5 and OUCR tend to have a vertical pattern, likely because the under-sampled images suffer from more vertical aliasing artifacts due to phase encoding under-sampling. As such, it is more difficult for D5C5 and OUCR to correctly recover these vertical details, and leads to blurry reconstructions under a pixel-wise loss function. DiffuseRecon, on the other hand, outputs realistic MR images that obey the distribution p θ ​ ( y full | x obs ) subscript 𝑝 𝜃 conditional subscript 𝑦 full subscript 𝑥 obs p_{\theta}(y_{\textrm{full}}|x_{\textrm{obs}}) ; it can better recover vertical details that fit the distribution of p θ ​ ( y full ) subscript 𝑝 𝜃 subscript 𝑦 full p_{\theta}(y_{\textrm{full}}) . This is particularly pronounced in the 8 × 8\times fastMRI results in Fig.  5 , where the lost vertical knee pattern in D5C5 and OUCR renders the image highly unrealistic. Such an error is avoided in DiffuseRecon, as each of its sample has a complete knee structure based on the learned prior knowledge. The uncertainty introduced by phase-encoding under-sampling is also captured by the variance map and demonstrates that the exact placement of details may be varied. For more visualization and detailed comparisons, please refer to the supplemental material.

While DiffuseRecon holds many advantages to current SoTA methods, it does require more computation due to the multi-sampling process. The one-time inference speed is comparable to UNet  [ 16 ] at 20 ​ ms 20 ms 20\textrm{ms} ; however, DiffuseRecon performs 1000 inference steps. We note that 20s per slice still yields significant acceleration compared to raw acquisitions, and there exists much potential in accelerating diffusion models  [ 18 , 11 ] for future work. DiffuseRecon can also be used in conjunction with deterministic methods, e.g. , when variance analysis is needed only selectively, to balance speed to performance and reliability.

FastMRI SKM-TEA

D5C5  [ 17 ]

OUCR  [ 6 ]

DiffuseRecon GT D5C5  [ 17 ]

OUCR  [ 6 ]

DiffuseRecon GT

6 × \times

8 × \times

(a) 29.61/.6974

(b) 31.34/.7472

(c) 35.30/.8622

(d) Variance

(e) 30.03/.6160

(f) 30.50/.6765

(g) 31.28/.7011

(h) Variance

(i) 25.61/.5167

(j) 27.84/.6148

(k) 30.42/.7087

(l) Variance

(m) 31.59/.7399

(n) 31.97/.7513

(o) 33.26/.7896

(p) Variance

Figure 5 : Visual comparison of DiffuseRecon with other SoTA reconstruction methods. Error maps are provided for reference.

4

Conclusion

We propose DiffuseRecon, an MR image reconstruction method that is of high performance, robust to different acceleration factors, and allows a user to directly observe the reconstruction variations. Inspired by diffusion models, DiffuseRecon incorporates the observed k-space signals in reverse-diffusion and can stochastically generate realistic MR reconstructions. To obtain the most likely reconstruction candidate and its variance, DiffuseRecon performs a coarse-to-fine sampling scheme that significantly accelerates the generation process. We carefully evaluate DiffuseRecon on raw acquisition data, and find that DiffuseRecon achieves better performances than current SoTA methods that require dedicated data generation and training for different sampling conditions. The reconstruction variance from DiffuseRecon can help practitioners make more informed clinical decisions. Future works include further acceleration of the inference process and examination into different MR sampling patterns.

References

[1]

Akçakaya, M., Moeller, S., Weingärtner, S., Uğurbil, K.: Scan-specific robust artificial-neural-networks for k-space interpolation (RAKI) reconstruction: Database-free deep learning for fast imaging. Magn. Reson. Med. 81 (1), 439–453 (Jan 2019)

[2]

Chung, H., chul Ye, J.: Score-based diffusion models for accelerated mri (2021)

[3]

Desai, A.D., Schmidt, A.M., Rubin, E.B., Sandino, C.M., Black, M.S., Mazzoli, V., Stevens, K.J., Boutin, R., Re, C., Gold, G.E., et al.: Skm-tea: A dataset for accelerated mri reconstruction with dense image labels for quantitative clinical evaluation. In: Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2) (2021)

[4]

Dhariwal, P., Nichol, A.: Diffusion models beat gans on image synthesis. CoRR abs/2105.05233 (2021), https://arxiv.org/abs/2105.05233

[5]

Eo, T., Jun, Y., Kim, T., Jang, J., Lee, H.J., Hwang, D.: KIKI-net: cross-domain convolutional neural networks for reconstructing undersampled magnetic resonance images. Magn. Reson. Med. 80 (5), 2188–2201 (Nov 2018)

[6]

Guo, P., Valanarasu, J.M.J., Wang, P., Zhou, J., Jiang, S., Patel, V.M.: Over-and-under complete convolutional RNN for MRI reconstruction. In: de Bruijne, M., Cattin, P.C., Cotin, S., Padoy, N., Speidel, S., Zheng, Y., Essert, C. (eds.) Medical Image Computing and Computer Assisted Intervention

• MICCAI 2021 - 24th International Conference, Strasbourg, France, September 27 - October 1, 2021, Proceedings, Part VI. Lecture Notes in Computer Science, vol. 12906, pp. 13–23. Springer (2021)

[7]

Guo, P., Wang, P., Zhou, J., Jiang, S., Patel, V.M.: Multi-institutional collaborations for improving deep learning-based magnetic resonance image reconstruction using federated learning. In: IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2021, virtual, June 19-25, 2021. pp. 2423–2432. Computer Vision Foundation / IEEE (2021)

[8]

Hammernik, K., Klatzer, T., Kobler, E., Recht, M.P., Sodickson, D.K., Pock, T., Knoll, F.: Learning a variational network for reconstruction of accelerated mri data. Magnetic Resonance in Medicine 79 (6), 3055–3071 (2018)

[9]

Han, Y., Sunwoo, L., Ye, J.C.: ¡inline-formula¿ ¡tex-math notation=”latex”¿ k 𝑘 {k} ¡/tex-math¿¡/inline-formula¿-space deep learning for accelerated mri. IEEE Transactions on Medical Imaging 39 (2), 377–386 (2020). https://doi.org/10.1109/TMI.2019.2927101

[10]

Ho, J., Jain, A., Abbeel, P.: Denoising diffusion probabilistic models. In: Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., Lin, H. (eds.) Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual (2020)

[11]

Kong, Z., Ping, W.: On fast sampling of diffusion probabilistic models. CoRR abs/2106.00132 (2021), https://arxiv.org/abs/2106.00132

[12]

Lustig, M., Donoho, D., Pauly, J.M.: Sparse mri: The application of compressed sensing for rapid mr imaging. Magnetic Resonance in Med. 58 (6), 1182–95 (2007)

[13]

Nichol, A.Q., Dhariwal, P.: Improved denoising diffusion probabilistic models. In: Meila, M., Zhang, T. (eds.) Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event. Proceedings of Machine Learning Research, vol. 139, pp. 8162–8171. PMLR (2021), http://proceedings.mlr.press/v139/nichol21a.html

[14]

Qin, C., Schlemper, J., Caballero, J., Price, A.N., Hajnal, J.V., Rueckert, D.: Convolutional recurrent neural networks for dynamic mr image reconstruction. IEEE Transactions on Medical Imaging 38 (1), 280–290 (2019). https://doi.org/10.1109/TMI.2018.2863670

[15]

Ravishankar, S., Bresler, Y.: MR image reconstruction from highly undersampled k-space data by dictionary learning. IEEE TMI 30 (5), 1028–41 (2011)

[16]

Ronneberger, O., Fischer, P., Brox, T.: U-net: Convolutional networks for biomedical image segmentation. In: Medical Image Computing and Computer-Assisted Intervention - MICCAI 2015 - 18th International Conference Munich, Germany, October 5 - 9, 2015, Proceedings, Part III. pp. 234–241 (2015)

[17]

Schlemper, J., Caballero, J., Hajnal, J.V., Price, A.N., Rueckert, D.: A deep cascade of convolutional neural networks for dynamic MR image reconstruction. IEEE Trans. Medical Imaging 37 (2), 491–503 (2018)

[18]

Song, J., Meng, C., Ermon, S.: Denoising diffusion implicit models. In: 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net (2021)

[19]

Sriram, A., Zbontar, J., Murrell, T., Zitnick, C.L., Defazio, A., Sodickson, D.K.: Grappanet: Combining parallel imaging with deep learning for multi-coil MRI reconstruction. In: 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020. pp. 14303–14310. Computer Vision Foundation / IEEE (2020)

[20]

Valanarasu, J.M.J., Sindagi, V.A., Hacihaliloglu, I., Patel, V.M.: Kiu-net: Towards accurate segmentation of biomedical images using over-complete representations. In: Martel, A.L., Abolmaesumi, P., Stoyanov, D., Mateus, D., Zuluaga, M.A., Zhou, S.K., Racoceanu, D., Joskowicz, L. (eds.) Medical Image Computing and Computer Assisted Intervention - MICCAI 2020 - 23rd International Conference, Lima, Peru, October 4-8, 2020, Proceedings, Part IV. Lecture Notes in Computer Science, vol. 12264, pp. 363–373. Springer (2020)

[21]

Wang, S., Su, Z., Ying, L., Peng, X., Zhu, S., Liang, F., Feng, D., Liang, D.: Accelerating magnetic resonance imaging via deep learning. In: 2016 IEEE 13th International Symposium on Biomedical Imaging (ISBI). pp. 514–517 (2016). https://doi.org/10.1109/ISBI.2016.7493320

[22]

Yang, G., Yu, S., Dong, H., Slabaugh, G., Dragotti, P.L., Ye, X., Liu, F., Arridge, S., Keegan, J., Guo, Y., Firmin, D.: Dagan: Deep de-aliasing generative adversarial networks for fast compressed sensing mri reconstruction. IEEE Transactions on Medical Imaging 37 (6), 1310–1321 (2018). https://doi.org/10.1109/TMI.2017.2785879

[23]

Zbontar, J., Knoll, F., Sriram, A., Muckley, M.J., Bruno, M., Defazio, A., Parente, M., Geras, K.J., Katsnelson, J., Chandarana, H., Zhang, Z., Drozdzal, M., Romero, A., Rabbat, M.G., Vincent, P., Pinkerton, J., Wang, D., Yakubova, N., Owens, E., Zitnick, C.L., Recht, M.P., Sodickson, D.K., Lui, Y.W.: fastmri: An open dataset and benchmarks for accelerated MRI. CoRR abs/1811.08839 (2018)

[24]

Zhou, B., Zhou, S.K.: Dudornet: Learning a dual-domain recurrent network for fast MRI reconstruction with deep T1 prior. In: 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020. pp. 4272–4281. Computer Vision Foundation / IEEE (2020)

◄

 Feeling lucky? 
  
 Conversion report 
 Report an issue 
 View original on arXiv  ►