Solving Inverse Problems in Medical Imaging with Score-Based Generative Models

1 INTRODUCTION . Computed Tomography ( CT ) and Magnetic Resonance Imaging ( MRI ) are commonly used imaging tools for medical diagnosis . Reconstructing CT and MRI images from raw measurements ( sinograms for CT and k-spaces for MRI ) are well-known inverse problems . Specifically , measurements in CT are given by X-ray projections of an object from various directions , and measurements in MRI are obtained by inspecting the Fourier spectrum of an object with magnetic fields . However , since obtaining the full sinogram for CT causes excessive ionizing radiation for patients , and measuring the full k-space of MRI is very time-consuming , it has become important to reduce the number of measurements in CT and MRI . In many cases , only partial measurements , such as sparse-view sinograms and downsampled k-spaces , are available . Due to this loss of information , the inverse problems in CT and MRI are often ill-posed , making image reconstruction especially challenging . With the rise of machine learning , many methods ( Zhu et al. , 2018 ; Mardani et al. , 2017 ; Shen et al. , 2019 ; Würfl et al. , 2018 ; Ghani & Karl , 2018 ; Wei et al. , 2020 ) have been proposed for medical image reconstruction using a small number of measurements . Most of these methods are supervised learning techniques . They learn to directly map partial measurements to medical images , by training on a large dataset comprising pairs of CT/MRI images and measurements . These measurements need to be synthesized from medical images with a fixed physical model of the measurement process . However , when the measurement process changes , such as using a different number of CT projections or different downsampling ratio of MRI k-spaces , we have to re-collect the paired dataset with the new measurement process and re-train the model . This prevents models from generalizing effectively to new measurement processes , leading to counter-intuitive instabilities such as more measurements causing worse performance ( Antun et al. , 2020 ) . In this work , we sidestep this difficulty completely by proposing unsupervised methods that do not require a paired dataset for training , and therefore are not restricted to a fixed measurement process . Our main idea is to learn the prior distribution of medical images with a generative model in order to infer the lost information due to partial measurements . Specifically , we propose to train a score-based generative model ( Song & Ermon , 2019 ; 2020 ; Song et al. , 2021 ) on medical images as the data prior , due to its strong performance in image generation ( Ho et al. , 2020 ; Dhariwal & Nichol , 2021 ) . Given a trained score-based generative model , we provide a family of sampling algorithms to create image samples that are consistent with the observed measurements and the estimated data prior , leveraging the physical measurement process . Once our model is trained , it can be used to solve any inverse problem within the same image domain , as long as the mapping from images to measurements is linear , which holds for a large number of medical imaging applications . We evaluate the performance of our method on several tasks in CT and MRI . Empirically , we observe comparable or better performance compared to supervised learning counterparts , even when evaluated with the same measurement process in their training . In addition , we are able to uniformly surpass all baselines when changing the number of measurements , e.g. , using a different number of projections in sparse-view CT or changing the k-space downsampling ratio in undersampled MRI . Moreover , we show that by plugging in a different measurement process , we can use a single model to perform both sparse-view CT reconstruction and metal artifact removal for CT imaging with metallic implants . To the best of our knowledge , this is the first time that generative models are reported successful on clinical CT data . Collectively , these empirical results indicate that our method is a competitive alternative to supervised techniques in medical image reconstruction and artifact removal , and has the potential to be a universal tool for solving many inverse problems within the same image domain . 2 BACKGROUND . 2.1 LINEAR INVERSE PROBLEMS . An inverse problem seeks to recover an unknown signal from a set of observed measurements . Specifically , suppose x P Rn is an unknown signal , and y P Rm “ Ax ` is a noisy observation given by m linear measurements , where the measurement acquisition process is represented by a linear operatorA P Rmˆn , and P Rn represents a noise vector . Solving a linear inverse problem amounts to recovering the signal x from its measurement y . Without further assumptions , the problem is ill-defined when m ă n , so we additionally assume that x is sampled from a prior distribution ppxq . In this probabilistic formulation , the measurement and signal are connected through a measurement distribution ppy | xq “ q py ´Axq , where q denotes the noise distribution of . Given ppy | xq and ppxq , we can solve the inverse problem by sampling from the posterior distribution ppx | yq . Examples of linear inverse problems in medical imaging include image reconstruction for CT and MRI . In both cases , the signal x is a medical image . The measurement y in CT is a sinogram formed by X-ray projections of the image from various angular directions ( Buzug , 2011 ) , while the measurement y in MRI consists of spatial frequencies in the Fourier space of the image ( a.k.a . the k-space in the MRI community ) ( Vlaardingerbroek & Boer , 2013 ) . 2.2 SCORE-BASED GENERATIVE MODELS . When solving inverse problems in medical imaging , we are given an observation y , the measurement distribution ppy | xq and aim to sample from the posterior distribution ppx | yq . The prior distribution ppxq is typically unknown , but we can train generative models on a dataset txp1q , xp2q , ¨ ¨ ¨ , xpNqu „ ppxq to estimate this prior distribution . Given an estimate of ppxq and the measurement distribution ppy | xq , the posterior distribution ppx | yq can be determined through Bayes ’ rule . We propose to estimate the prior distribution of medical images using the recently introduced scorebased generative models ( Song & Ermon , 2019 ; Ho et al. , 2020 ; Song et al. , 2021 ) , whose iterative sampling procedure makes it especially easy for controllable generation conditioned on an observation y . Specifically , we adopt the formulation of score-based generative models in Song et al . ( 2021 ) , where we leverage a Markovian diffusion process to progressively perturb data to noise , and then smoothly convert noise to samples of the data distribution by estimating and simulating its time reversal . We provide an illustration of this generative modeling framework in Fig . 1 . Perturbation process Suppose the dataset is sampled from an unknown data distribution ppxq . We perturb datapoints with a stochastic process over a time horizon r0 , 1s , governed by a linear stochastic differential equation ( SDE ) of the following form dxt “ fptqxt dt ` gptqdwt , t P r0 , 1s , ( 1 ) where f : r0 , 1s Ñ R , g : r0 , 1s Ñ R , twt P RnutPr0,1s denotes a standard Wiener process ( a.k.a. , Brownian motion ) , and txt P RnutPr0,1s symbolizes the trajectory of random variables in the stochastic process . We further denote the marginal probability distribution of xt as ptpxq , and the transition distribution from x0 to xt as p0tpxt | x0q . By definition , we clearly have p0pxq ” ppxq . Moreover , the functions fptq and gptq are specifically chosen such that for any initial distribution p0pxq , the distribution at the end of the perturbation process , p1pxq , is close to a pre-defined noise distribution πpxq . In addition , the transition density p0tpxt | x0q is always a conditional linear Gaussian distribution , taking the form p0tpxt | x0q “ N pxt | αptqx0 , β2ptqIq where α : r0 , 1s Ñ R and β : r0 , 1s Ñ R can be derived analytically from fptq and gptq ( Särkkä & Solin , 2019 ) . Examples of such SDEs include Variance Exploding ( VE ) , Variance Preserving ( VP ) , and subVP SDEs proposed in Song et al . ( 2021 ) . We found VE SDEs performed the best in our experiments . Reverse process By reversing the perturbation process in Eq . ( 1 ) , we can start from a noise sample x1 „ p1pxq and gradually remove the noise therein to obtain a data sample x0 „ p0pxq ” ppxq . Crucially , the time reversal of Eq . ( 1 ) is given by the following reverse-time SDE ( Song et al. , 2021 ) dxt “ “ fptqxt ´ gptq2∇xt log ptpxtq ‰ dt ` gptqdw̄t , t P r0 , 1s , ( 2 ) where tw̄tutPr0,1s denotes a standard Wiener process in the reverse-time direction , and dt represents an infinitesimal negative time step , since the above SDE must be solved backwards from t “ 1 to t “ 0 . The quantity ∇xt log ptpxtq is known as the score function of ptpxtq . By the definition of time reversal , the trajectory of the reverse stochastic process given by Eq . ( 2 ) is txtutPr0,1s , same as the one from the forward SDE in Eq . ( 1 ) . Sampling Given an initial sample from p1pxq , as well as scores at each intermediate time step , ∇x log ptpxq , we can simulate the reverse-time SDE in Eq . ( 2 ) to obtain samples from the data distribution p0pxq ” ppxq . In practice , the initial sample is approximately drawn from πpxq since πpxq « p1pxq , and the scores are estimated by training a neural network sθpx , tq ( named the score model ) on a dataset txp1q , xp2q , ¨ ¨ ¨ , xpNqu „ ppxq with denoising score matching ( Vincent , 2011 ; Song et al. , 2021 ) , i.e. , solving the following objective θ˚ “ arg min θ 1 N N ÿ i “ 1 Et „ Ur0,1sExpiqt „ p0tpxpiqt |xpiqq ” ∥∥∥sθpxpiqt , tq ´∇xpiqt log p0tpxpiqt | xpiqq∥∥∥22 ı , where Ur0 , 1s denotes a uniform distribution over r0 , 1s . The theory of denoising score matching ensures that sθ˚px , tq « ∇x log ptpxq . After training this score model , we plug it into Eq . ( 2 ) and solve the resulting reverse-time SDE dxt “ “ fptqxt ´ gptq2sθ˚pxt , tq ‰ dt ` gptqdw̄t , t P r0 , 1s , ( 3 ) for sample generation . One sampling method is to use the Euler-Maruyama discretization for solving Eq . ( 3 ) , as given in Algorithm 1 . Other sampling methods include annealed Langevin dynamics ( ALD , Song & Ermon , 2019 ) , probability flow ODE solvers ( Song et al. , 2021 ) , and Predictor-Corrector samplers ( Song et al. , 2021 ) . Algorithm 1 Unconditional sampling Require : N. 1 : x̂1 „ πpxq , ∆tÐ 1N 2 : for i “ N ´ 1 to 0 do 3 : tÐ i ` 1 N 4 : x̂t´∆t Ð x̂t ´ fptqx̂t∆t 5 : x̂t´∆t Ð x̂t´∆t ` gptq2sθ˚px̂t , tq∆t 6 : z „ N p0 , Iq 7 : x̂t´∆t Ð x̂t´∆t ` gptq ? ∆t z 8 : return x̂0 Algorithm 2 Inverse problem solving Require : N , y , λ 1 : x̂1 „ πpxq , ∆tÐ 1N 2 : for i “ N ´ 1 to 0 do 3 : tÐ i ` 1 N 4 : ŷt „ p0tpyt | yq 5 : x̂t Ð T´1rλΛP´1pΛqŷt ` p1´ λq ` ΛT ¨ x̂t ` pI ´ΛqT x̂ts 6 : x̂t´∆t Ð x̂t ´ fptqx̂t∆t 7 : x̂t´∆t Ð x̂t´∆t ` gptq2sθ˚px̂t , tq∆t 8 : z „ N p0 , Iq 9 : x̂t´∆t Ð x̂t´∆t ` gptq ? ∆t z 10 : return x̂0

Summary: This paper provides an unsupervised approach to solve the inverse problem for reconstructing medical CT and MRI scans using score-based generative models. The proposed method was evaluated on the LIDC and BraTS datasets. Compared to existing supervised and unsupervised approaches, the proposed method demonstrates comparable or better performances in terms of PSNR and SSIM.