Solving Inverse Problems in Medical Imaging with Score-Based Generative Models
Abstract
Solving inverse problems with a small number of measurements has important applications in medical imaging, including image reconstruction for undersampled
MRI and sparse-view CT.
With the progress of machine learning, traditional image reconstruction methods have been outperformed by models that learn to directly map measurements to medial images.
However, these models require both ground truth medical images and their measurements for training, which complicates data collection and harms their generalization performance to unknown measurement processes.
To address these issues, we propose a fully unsupervised technique for inverse problem solving, leveraging the recently introduced score-based generative models.
Speciﬁcally, we train a score-based generative model to capture the prior distribution of medical images, which is subsequently combined with a given physical measurement process to sample images consistent with measurements at the test time.
Our method makes no assumption on the measurement process during training, and can be ﬂexibly adapted to any linear measurement processes.
Empirically, we observe comparable or better performance to supervised learning techniques, with better generalization to unknown measurement processes on several MRI/CT datasets.
Introduction
Magnetic Resonance Imaging (MRI) and Computed Tomography (CT) are commonly used imaging tools for medical diagnosis.
Reconstructing MRI/CT images from raw measurements (k-space for MRI and sinograms for CT) is a well-known linear inverse problem.
In order to speed up the hour-long process of MRI acquisition and reduce the dose of harmful ionizing radiation needed for
CT scans, we aim to reduce the number of measurements required for solving these inverse problems by estimating and leveraging the prior information of MRI/CT images.
Speciﬁcally, we consider image reconstruction tasks for undersampled MRI and sparse-view CT.
Many machine learning algorithms [1, 2, 3, 4, 5, 6] have been proposed to perform medical image reconstruction given a small number of measurements.
However, most of these methods are supervised learning techniques.
They learn to map measurements directly to medical images, by training on a large dataset comprising pairs of medical images and MRI/CT measurements.
This complicates the data collection procedure and moreover, limits their generalization capability to different measurement processes that may appear at the test time, leading to counterintuitive instabilities such as more measurements causing worse performance [7].
We sidestep this difﬁculty by proposing unsupervised methods that make no assumption on the measurement process during training.
Speciﬁcally, we train a score-based generative model [8, 9, 10] on medical images as the data prior, and modify the sampling algorithm as a general way to solve inverse problems with linear measurement processes.
As an emerging family of techniques that have achieved top performance in image generation [8, 10, 11, 12], score-based generative models allow us to match or outperform the performance of supervised learning techniques, while being more resilient to the change of measurement processes, such as using a different number of measurements.
Background
The latest development of score-based generative models [10] leverages a Markovian diffusion process to progressively perturb data to noise, and then learns to smoothly convert noise to samples by solving its time reversal.
Importantly, reversing the diffusion process requires estimating scores , i.e ., gradients of log probability density functions of the data distribution, from a training dataset.
Perturbation process Suppose the dataset is sampled from an unknown data distribution p p x q .
For any datapoint x ˚ P R n „ p p x q , we perturb it with a stochastic process over a time horizon r 0 , 1 s , governed by a stochastic differential equation (SDE) of the following form
[Equation]
where f p¨ , t q : R n Ñ R n and g p t q P R are called the drift and diffusion coefﬁcients, t w t P R n u t Pr 0 , 1 s denotes a standard Wiener process (a.k.a., Brownian motion), and t x t u t Pr 0 , 1 s symbolizes the trajectory of random variables in the stochastic process.
We further denote the marginal probability distribution of x t as p t p x q , and the transition distribution from x 0 to x t as p 0 t p x t | x 0 q .
By deﬁnition, we clearly have x 0 “ x ˚ and p
p x q ” p p x q .
The drift and diffusion coefﬁcients are typically hand-crafted such that the distribution at the end of the perturbation process, p 1 p x q , is close to a pre-deﬁned prior distribution π p x q .
Examples of such SDEs include Variance Exploding (VE), Variance Preserving (VP), and subVP SDEs [10].
Sampling process By reversing the perturbation process in Eq.
(1), we can start from a sample x 1 „ p 1 p x q and gradually remove the noise therein to obtain a data sample x 0 „ p
p x q ” p p x q .
Crucially, the time reversal of Eq.
(1) is given by the following reverse-time SDE [13, 14]
[Equation]
where t ¯ w t u t Pr 0 , 1 s is now a standard Wiener process in the reverse-time direction, and d t represents an inﬁnitesimal negative time step, since the above SDE must be solved backwards from t “ 1 to t “ 0 .
When x 1 „ p 1 p x q , the trajectory of the reverse stochastic process given by Eq.
(2) is again t x t u t Pr 0 , 1 s , same as the one obtained by the forward SDE in Eq.
Once given an initial sample from p 1 p x q , as well as scores at each intermediate time step, ∇ x log p t p x q , we can use any numerical
SDE solver to solve the reverse-time SDE in Eq.
The initial sample can be approximately drawn from π p x q since p 1 p x q « π p x q , while the scores can be estimated by training a neural network s θ p x , t q (named the time-dependent score model) with denoising score matching [15, 10], such that s θ ˚ p x , t q « ∇ x log p t p x q .
Here θ ˚ denotes the optimal model parameters.
In Algorithm 1, we provide one such sampling method based on the classic Euler-Maruyama solver.
Method
By incorporating data consistency constraints into the sampling process, we can use score-based generative models to solve linear inverse problems.
We consider a compressed sensing [16, 17, 18] setting, while replacing sparsity with a data prior learned by our score models.
Suppose A P R m ˆ n is a linear operator with full rank, i.e ., rank p
A q “ min p n, m q .
Let x P R n be a signal, and y P R m be a corresponding observation obtained by m measurements.
In an inverse problem, we are given an observation y and aim to ﬁnd x such that A x “ y .
In order for the inverse problem to have at least one solution, we clearly require m ď n .
Since there might be multiple solutions, we further assume a prior distribution over x , denoted as p p x q .
With this assumption, solving the inverse problem becomes equivalent to ﬁnding the posterior distribution p p x | y q .
To facilitate our theoretical discussion, we introduce an alternative formulation of A below:
Lemma 1.
If A has full rank, we can always write it as A “ C Λ T , where T P R n ˆ n is an invertible matrix, Λ P R n ˆ n is a diagonal matrix taking values in t 0 , 1 u that satisﬁes tr p Λ q “ m , and
Require: N 1: ˆ x 1 „ π p x q , ∆ t Ð 1 N 2: for i “ N ´ 1 to 0 do 3: t Ð i ` 1 N 4: ˆ x t ´ ∆ t Ð ˆ x t ` g p t q 2 s θ ˚ p
ˆ x t , t q ∆ t 5: z „ N p 0 , I q 6: ˆ x t ´ ∆ t Ð ˆ x t ´ ∆ t ` g p t q?
∆ t z 7: return ˆ x
Require: N , y , λ 1: ˆ x 1 „ π p x q , ∆ t Ð 1 N 2: for i “ N ´ 1 to 0 do 3: t Ð i ` 1 N 4: ˆ y t „
p 0 t p y t | y q 5: ˆ x t Ð λ T ´ 1 Λ C ´ 1 Λ ˆ y t ` p 1 ´ λ q T ´ 1 Λ T ˆ x t `
T ´ 1 p I ´ Λ q T ˆ x t 6: ˆ x t ´ ∆ t Ð ˆ x t ` g p t q 2 s θ ˚ p ˆ x t , t q ∆ t 7: z „ N p 0 , I q 8: ˆ x t ´ ∆ t Ð ˆ x t ´ ∆ t ` g p t q? ∆ t z 9: return ˆ x C Λ : R n Ñ R m is a linear operator that takes any vector a P R n and reduces its dimensionality to m by removing each i -th element of a if Λ ii “ 0 .
For example, T corresponds to the Fourier transform for undersampled MRI reconstruction, and the
Radon transform for sparse-view CT reconstruction.
Given a dataset t x 1 , x 2 , ¨ ¨ ¨ , x N u i.i.d. „
p p x q , the approach outlined in Algorithm 1 allows us to train a time-dependent score model s θ ˚ p x , t q to generate samples from p p x q .
To sample from p p x | y q , however, we will need to incorporate the additional information of y into the sampling process.
In inverse problems, y and x are connected through y “ A x .
By introducing y t “ A x t , we generalize this connection to stochastic processes t y t u t Pr 0 ,T s and t x t u t Pr 0 ,T s .
When the drift coefﬁcient f p¨ , t q in Eq.
(1) is linear, the transition density for t x t u t Pr 0 ,T s is always a Gaussian distribution [10], taking the form p 0 t p x t | x 0 q “ N p x t | α p t q x 0 , β p t
q
I q where α p t q and β p t q can be computed from f p¨ , t q and g p t q .
The transition density of t y t
u t Pr 0 ,T s is therefore given by p 0 t p y t | y [-
-]0 q “ N p y t | α p t q y 0 , β p t q AA T
I q .
When y 0 is ﬁxed to y ˚ , the marginal distributions of y t equals p 0 t p y t | y
“ y ˚ q and can be computed in closed form.
Suppose ˆ x t is an intermediate sample obtained by running one step of the sampler in Algorithm 1, and let ˆ y t „ p 0 t p y t | y q .
Our goal is to steer ˆ x t towards ˆ y t to promote data consistency.
To this end, we propose to ﬁnd a steering function h : R n Ñ R n , such that ˆ x 1 t “ h p ˆ x t q satisﬁes both ˆ x 1 t « ˆ x t and A ˆ x 1 t « ˆ y t .
Speciﬁcally, this steering function needs to minimize the distance between ˆ x 1 t and ˆ x t , as well as the distance between ˆ x 1 t and the hyperplane
Ax “ ˆ y t .
We propose to balance these two objectives by solving the following optimization problem
[Equation] [Equation]
where
ă λ ă 1 controls the strength of the data-consistency loss, and we choose the norm (cid:107) a (cid:107) 2 T : “ (cid:107) T a (cid:107) 22 to simplify our theoretical analysis.
The following result gives a closed-form solution to the optimization problem in Eq.
Theorem 1.
The optimal solution ˆ x 1 t is given by
[Equation]
where C ´ 1 Λ : R m Ñ R n denotes any right inverse of C Λ .
By transforming ˆ x t to ˆ x 1 t after each sampling step, we can incorporate y to generate approximate samples from p p x | y q .
Our method requires minimal modiﬁcation to the original unconditional sampling method of score-based generative models.
For example, we can convert the sampler in
Algorithm 1 to an inverse problem solver in Algorithm 2 by adding just two lines of pseudo-code.
A concurrent work in [19] proposes a different method to solve inverse problems with score-based generative models, leveraging posterior sampling and Langevin dynamics.
In comparison, our method is applicable to a larger family of sampling methods in score-based generative models, such as numerical SDE solvers [10], and achieves better empirical performance in our experiments.
[Table]
Experiments
We consider two important examples of inverse problem solving in medical imaging: image reconstruction for undersampled MRI and sparse-view CT.
We compare the performance of our method against traditional reconstruction algorithms as well as other deep learning based methods.
Datasets For undersampled MRI, we use the Brain Tumor Segmentation (BraTS) 2021 dataset [20, 21].
We slice 2D images of resolution 240 ˆ 240 from original 3D MRI volumes.
We simulate measurements by computing the k-space with Fast Fourier Transform, and follow [22, 23] to undersample the k-space with an equispaced Cartesian mask.
For sparse-view CT, we use the Lung
Image Database Consortium image collection (LIDC) [24, 25] dataset.
We obtain 2D images of resolution 320 ˆ 320 by slicing the 3D CT volumes.
We simulate measurements (sinograms) with a parallel-beam geometry using projection angles equally distributed across 180 degrees.
Evaluation Our baselines include the traditional CT reconstruction method FISTA-TV [26], supervised learning methods Cascade DenseNet [27], DuDoRNet [28], cGAN [5], Neumann [29],
SIN-4c-PRN [6], and a concurrent unsupervised learning method based on Langevin dynamics sampling from score-based generative models [30], which we denote as “Langevin”.
We evaluate the performance of different methods using the peak signal-to-noise ratio (PSNR) and the structural similarity metric (SSIM [31]).
For deep learning based methods, we train and tune the corresponding models with 30 measurements for MRI and 23 projections for CT.
We use the same score models in both Langevin and our method.
Results We provide experimental results in Table 1 and Fig.
For undersampled MRI, our method achieves comparable performance to the best supervised learning technique DuDoRNet.
However,
DuDoRNet cannot generalize well to different number of measurements since, as a supervised learning technique, it is speciﬁcally trained using 30 measurements.
In contrast, our method is unsupervised and can generalize to an arbitrary number of measurements.
As shown in Fig.
1, our method uniformly outperforms competing methods when using measurements other than 30.
For sparse-view CT, we are able to outperform all baselines, even including supervised learning methods tested with the same number of measurements (23 projections) for training.
Conclusion
This paper describes a new method to solve linear inverse problems with score-based generative models.
Our method is fully unsupervised, requires no paired or labeled data for training, and can ﬂexibly adapt to different measurement processes at the test time.
Preliminary results demonstrate that our method can match or outperform existing supervised learning counterparts for undersampled
MRI and sparse-view CT reconstruction, while being more robust to the change of the number of measurements.
Author Contributions
Yang Song designed the project, wrote the paper, and ran all experiments for score-based generative models.
Liyue Shen ran baseline experiments, and helped write the paper. Lei Xing and Stefano Ermon supervised the project, provided valuable feedback, and helped edit the paper. Acknowledgments
YS is supported by the Apple PhD Fellowship in AI/ML.
LS is supported by the Stanford Bio-X
Graduate Student Fellowship.
This research was supported by NSF (#1651565, #1522054, #1733686),
ONR (N000141912145), AFOSR (FA95501910024), ARO (W911NF-21-1-0125), Sloan Fellowship, and Google TPU Research Cloud.
This research was also supported by NIH/NCI (1R01 CA and 1R01 CA227713)
A Related Work
The idea of solving inverse problems with generative models has been proposed in the Compressed
Sensing with Generative Models (CSGM) framework [18].
However, the original CSGM method is centered around latent variable models such as Generative Adversarial Networks (GANs [32]) and
Variational Auto-Encoders (VAEs [33, 34]).
It is unclear how to apply CSGM to other families of generative models like Energy-Based Models (EBMs [35]) and score-based generative models.
Solving linear inverse problems with unsupervised learning and denoising score matching [15] has also been explored in [36].
In contrast with our approach, their method requires projecting an noisy sample onto the kernel space of a linear operator, which can be expensive to compute for high dimensional data.
They did not consider applications in medical imaging either.
Concurrently, ref.
[30] extends the NCSNv2 model [9], a score-based generative model sampled by annealed Langevin dynamics [8], to solve inverse problems in downsampled MRI reconstruction.
Our method is applicable to more general sampling methods [10] for score-based generative models, such as numerical SDE solvers, Predictor-Corrector samplers, and probability ﬂow ODE samplers.
We also obtained uniformly better empirical performance in our experiments.
B Proofs
Lemma 1.
If A has full rank, we can always write it as A “ C Λ T , where T P R n ˆ n is an invertible matrix, Λ P R n ˆ n is a diagonal matrix taking values in t 0 , 1 u that satisﬁes tr p Λ q “ m , and
C Λ : R n Ñ R m is a linear operator that takes any vector a P R n and reduces its dimensionality to m by removing each i -th element of a if Λ ii “ 0 .
Proof.
Let A “ p a T 1 , a T 2 , ¨ ¨ ¨ , a T m q P R m ˆ n .
Since A has full rank, the row vectors t a 1 , a 2 , ¨ ¨ ¨ , a m u are linearly independent.
We can therefore extend them to a total of n linearly independent vectors, i.e ., t a 1 , a 2 , ¨ ¨ ¨ , a m , b 1 , ¨ ¨ ¨ , b n ´ m u .
Due to the linear independence, we know T “ p a T 1 , a T 2 , ¨ ¨ ¨ , a T m , b T 1 , ¨ ¨ ¨ , b T n ´ m q P R n ˆ n has full rank and is invertible. Next, we deﬁne Λ “ diag p 1 , 1 , ¨ ¨ ¨ , 1 loooomoooon m , 0 , 0 , [Equation]
¨ [Equation] 0 loooomoooon n ´ m q , where diag converts a vector to a diagonal matrix.
Clearly tr p Λ q “ m and A “ C Λ T , which completes the proof.
Lemma 2.
Let C ´ 1 Λ : R m Ñ R n be any right inverse of C Λ : R n Ñ R m .
For any u P R n and ˆ y t P R m , we have
[Equation]
Proof.
By the deﬁnition of C Λ , we have C Λ “ C Λ Λ , and
[Equation]
To prove the “if” direction, we note that
[Equation]
To prove the “only if” direction, we have
[Equation]
where (i) is due to the property in Eq.
This completes the proof for both directions.
Theorem 1.
The optimal solution ˆ x 1 t is given by
[Equation]
where C ´ 1 Λ : R m Ñ R n denotes any right inverse of C Λ .
Proof.
The optimization objective function in Eq.
(3) can be written as
[Equation]
Since Au “ ˆ y t [Equation] we have C Λ T u “ ˆ y t and equivalently Λ T u “ Λ C ´ 1 Λ ˆ y t due to Lemma 2.
This constraint does not restrict the value of p I ´ Λ q T u .
Therefore, when Au “ ˆ y t , we have
[Equation]
[Table]
This simpliﬁes the optimization problem in Eq.
(3) to min
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation] [Table] [Equation] [Table]
[Equation] [Table] [Equation] [Table]
[Equation]
[Table]
which is minimizing a quadratic function of z .
The optimal solution z ˚ is thus in closed form:
[Equation]
According to the deﬁnition, ˆ x 1 t “ z ˚ , whereby the proof is completed.
C Additional details
Training and sampling We use the NCSN++ model architecture in [10], and perturb the data with the Variance Exploding (VE) SDE.
Instead of generating samples according to the numerical SDE solver in Algorithm 1, we use the Predictor-Corrector (PC) sampler as described in [10] since it generally has better performance for VE SDEs.
In PC samplers, the predictor refers to a numerical solver for the reverse-time SDE while the corrector can be any Markov chain Monte Carlo (MCMC) method that only depends on the scores.
One such MCMC method considered in this work is
Langevin dynamics, whereby we transform any initial sample x p 0 q to an approximate sample from p t p x q via the following procedure:
[Equation]
Here N P N ą 0 , (cid:15) ą 0 , and z p i q „ N p 0 , I q .
The theory of Langevin dynamics guarantees that in the limit of N Ñ 8 and (cid:15) Ñ 0 , x p N q is a sample from p t p x q under some regularity conditions.
Note that
Langevin dynamics only requires the knowledge of ∇ x log p t p x q , which can be approximated using the time-dependent score model s θ ˚ p x , t q .
In PC samplers, each predictor step immediately follows multiple consecutive corrector steps, all using the same s θ ˚ p x , t q evaluated at the same t .
This jointly ensures that our intermediate sample at t is approximately distributed according to p t p x q .
As shown in [10], PC sampling often outperforms numerical solvers for the reverse-time SDE, especially when the forward SDE in Eq.
(1) is a VE SDE.
In order to use PC samplers for inverse problem solving, our modiﬁcation is similar to the change made in Algorithm 2 for Algorithm 1.
Speciﬁcally, we run line 4 & 5 in Algorithm 2 before every corrector or predictor step.
The step size (cid:15) in the Langevin corrector, as well as λ in Eq.
(3) are all tuned by running Bayesian optimization for 100 steps.
Datasets The Lung Image Database Consortium image collection (LIDC) [24, 25] consists of diagnostic and lung cancer screening thoracic computed tomography (CT) scans for lung cancer detection and diagnosis, which contains 1018 cases.
We convert the Hounsﬁeld units loaded from dicom ﬁles to the attenuation coefﬁcients and set the background as zero.
Then, 2D CT images are sliced from 3D CT cases.
The sinogram are simulated from 2D CT images based on parallel-beam geometry with different number of projection angles that are equally distributed across 180 degrees.
We conduct experiments of 2D MRI image reconstruction on a multi-modality MR image dataset.
The Brain Tumor Segmentation (BraTS) 2021 dataset [20, 21] collected for a image segmentation challenge contains 2000 cases (8000 MRI scans), where each case has four different MR contrasts: native (T1), post-contrast T1-weighted (T1Gd), T2-weighted (T2), and T2 Fluid Attenuated Inversion
Recovery (T2-FLAIR).
For each 3D MR volume, we extract 2D slices from 3D volumes and simulate k-space data by Fourier Transform.
To reconstruct MR images, k-space data is under-sampled by a equispaced Cartesian mask, where the center k-space is fully sampled while the left k-space is under-sampled by equispaced columns.
Baselines To evaluate the proposed method, we compare with baseline models including the traditional CT reconstruction method FISTA-TV [26], supervised learning methods for CT reconstruction:
cGAN
[5], Neumann [29], SIN-4c-PRN [6], and for MRI reconstruction: Cascade DenseNet [27], and DuDoRNet [28].
CT Reconstruction: FISTA-TV [26] is a fast iterative shrinkage-thresholding algorithm (FISTA) for solving linear inverse problems in image processing.
It uses a total variation term as the regularization in the optimization procedure.
Each optimization iteration involves a matrix-vector multiplication followed by a shrinkage/soft-threshold step.
Conventional iterative CT reconstruction algorithms like
FISTA-TV are typically slow due to their iterative nature.
In ref.
[5], authors propose to cast sparseview CT reconstruction as a sinogram inpainting problem.
Speciﬁcally, they propose to use conditional
GANs to complete the projection data (sinogram) prior to reconstructing CT images, thereby avoiding the costly iterative tomographic inversion.
However, the imperfect sinogram inpainting may cause other image artifacts.
To reduce such artifacts, SIN-4c-PRN [6] proposed a 2-step sparse-view CT reconstruction model, which contains a sinogram inpainting network (SIN) to generate super-resolved sinograms and then a post-processing reﬁning network (PRN) to further remove image artifacts.
Both networks are connected through a ﬁltered back-projection operation (FBP) operation.
Meanwhile, in another parallel direction, researchers proposed to learn the regularizer used in optimization from training data, outperforming traditional regularizers.
Speciﬁcally, ref.
[29] presented an end-to-end, data-driven method for solving inverse problems inspired by the Neumann series, called a Neumann network, which directly solves the linear inverse problem with a data-driven nonlinear regularizer.
Note that except for iterative reconstruction algorithm FISTA-TV, all the other deep-learning-based algorithms are supervised learning methods that require paired measurements (sinogram) and ground truth medical images for training networks.
MRI Reconstruction: To reconstruct de-aliased MR images from under-sampled k-space data, ref.
[27] proposed a cascaded dilated dense network (CDDN) for MRI reconstruction, based on stacked dense blocks with residual connections while using the zero-ﬁlled MR image as inputs.
Speciﬁcally, they use a two-step data consistency layer for k-space correction, and replace corresponding phasecoding lines of generated image with the original sampled k-space data after each block.
Based on this model, ref.
[28] proposed a dual domain recurrent network (DuDoRNet) to to simultaneously recover k-space and images for MRI reconstruction, in order to address aliasing artifacts in both frequency and image domains.
The original model in [28] also embedded a deep T1 prior to make use of fully-sampled short protocol (T1) as complementary information.
For a fair comparison with other supervised learning approaches, in our experiments, we do not include this additional information but train the DuDoRNet model without T1 prior.
In this setting, we also observe that cascaded densenet generalizes better to more measurements than DuDoRNet.