Robust Compressed Sensing MRI with Deep Generative Priors
Abstract
The CSGM framework (Bora-Jalal-Price-Dimakis’17) has shown that deep generative priors can be powerful tools for solving inverse problems.
However, to date this framework has been empirically successful only on certain datasets (for example, human faces and MNIST digits), and it is known to perform poorly on out-of-distribution samples.
In this paper, we present the ﬁrst successful application of the CSGM framework on clinical MRI data.
We train a generative prior on brain scans from the fastMRI dataset, and show that posterior sampling via Langevin dynamics achieves high quality reconstructions.
Furthermore, our experiments and theory show that posterior sampling is robust to changes in the ground-truth distribution and measurement process.
Introduction
Compressed sensing [12, 8] has enabled reductions to the number of measurements needed for successful reconstruction in a variety of imaging inverse problems.
In particular, it has led to shorter scan times for magnetic resonance imaging (MRI) [25, 39], and most MRI vendors have released products leveraging this framework to accelerate clinical workﬂows.
More recently, deep learning techniques have been used as powerful data-driven reconstruction methods for inverse problems [20, 29].
There are two broad families of deep learning inversion techniques [29]: end-to-end supervised and distribution-learning approaches.
End-to-end supervised techniques use a training set of measured images and deploy convolutional neural networks (CNNs) and other architectures to learn the inverse mapping from measurements to image see e.g. [16, 2, 26].
End-to-end methods are trained for speciﬁc imaging anatomy and measurement models and show excellent performance in these tasks.
However, reconstruction quality can suﬀer when applied out of distribution, such as under various types of natural measurement and anatomy perturbations [3, 11].
In this paper we study deep learning inversion techniques based on distribution learning.
The most common family of such techniques, known also as Compressed Sensing with Generative Modeling (CSGM) [6] uses pre-trained generative models as priors.
These methods have only recently been applied to MRI and have not yet been shown to be competitive with supervised end-to-end methods.
The very recent work [22] trains a StyleGAN for magnitude-only DICOM images but requires the presence of side-information and studies simulated measurements for reconstruction.
The deviation from the true MRI measurement model and the use of magnitude images are known to be problematic when evaluating performance [31].
Untrained and unamortized generators [17] have also been recently explored [11], showing promising results in some cases.
Further, [10] studies the harder problem of learning a generative model for a class of images using only partial observations, as ﬁrst proposed in AmbientGAN [7]. [Figure] [Figure] [Figure]
and on knee & abdomen scans. Our reconstructions are competitive with state-of-the-art methods, and introduce fewer artifacts out of distribution. All measurements are multicoil k-space from the NYU fastMRI dataset and the supervised baselines are trained from scratch on MVUE targets for a fair comparison. 1. Contributions
• We successfully train a score-based deep generative model (NCSNv2 [34]) for complex-valued, T2-weighted brain MR images without any assumptions on the measurement scheme.
When applied to multi-coil MRI reconstruction under the CSGM framework, we achieve competitive performance compared to end-to-end deep learning methods when the test-time data are sampled within distribution.
This is shown in the top row of Figure 1 and quantitatively in Figure 2.
• We give evidence that posterior sampling should give high-quality reconstructions.
First, we show that for any measurements (including the Fourier measurements in MRI) that posterior sampling with the correct prior is within constant factors of the optimal recovery method; second, even if the prior is wrong but gives α mass to the true distribution, we show that posterior sampling for Gaussian measurements is nearly optimal with just an additive O (log(1 /α )) loss.
These results are in Appendix A. • We empirically show that our approach is robust to test-time distribution shifts including diﬀerent sampling patterns and imaging anatomy.
The former is unsurprising given that our model was trained without knowledge of the measurement scheme.
As a consequence, our approach provides a degree of ﬂexibility in choosing scan parameters – a common situation in routine clinical imaging.
In contrast, in Figureand 6, we empirically show that end-to-end methods do not always enjoy the same robustness guarantees, in some cases leading to severe degradation in reconstruction quality when applied out-of-distribution.
• As shown in Figure 3, our method can be used to obtain multiple samples from the posterior by running
R
P S
N R ( d
B )
Langevin (Ours)
E2E-Varnet
MoDL
L1-Wavelet
ConvDecoder 12 8 6 4
R 0.
SS I
M
Langevin dynamics with diﬀerent random initializations.
This allows us to get multiple reconstructions which can be used to obtain conﬁdence intervals for each reconstructed voxel and visualize our reconstruction uncertainty on a voxel-by-voxel resolution.
Uncertainty quantiﬁcation can be incorporated into end-to-end methods, e.g., using variational auto-encoders [13], but this requires changes to the architecture.
Our method does not require any modiﬁcation and multiple reconstruction samplers can be run in parallel.
Our main results are succinctly summarized in Figure 1: we achieve equivalent reconstruction performance using a reduced training set when evaluated in-distribution and is robust when evaluated out-of-distribution.
System Model and Algorithm 2. Multi-coil Magnetic Resonance Imaging
MRI is a medical imaging modality that makes measurements using an array of radio-frequency coils placed around the body.
Each coil is spatially sensitive to a local region, and measurements are acquired directly in the spatial frequency, or k-space , domain.
To decrease scan time, reduce operating costs, and improve patient comfort, a reduced number of k-space measurements are acquired in clinical use and reconstructed by incorporating explicit or implicit knowledge of the spatial sensitivity maps [32, 30, 14].
Formally, the vector of measurements
y i ∈
C L acquired by the i th coil can be characterized by the forward model [30]:
[Equation]
where x ∗ ∈ C N is the image containing N pixels, S i is an operator representing the point-wise multiplication of the i th coil sensitivity map, F is the spatial Fourier transform operator, P represents the k-space sampling operator, and we assume w i ∼ N c (cid:0) 0 , σ 2 I (cid:1) for simplicity.
Importantly, note that the same under-sampling operator is applied to all N c coils.
The acceleration factor R denotes the degree of under-sampling in the k -space domain, i.e., R = N/L .
Note that we use the true acceleration factor R , and this does not match the values in fastMRI [23] 1 .
Given multi-coil measurements y , sensitivity maps represented by S and the sampling operator P , the goal of MR image reconstruction is to estimate the underlying image variable x ∗ . [Figure]
2. Posterior Sampling
The algorithm we consider is posterior sampling [19, 33].
That is, given an observation of the form y =
Ax ∗ + w , where y ∈ C M , A ∈ C M × N , w ∼ N c
(0 , σ 2 I ), and x ∗ ∼ µ , the posterior sampling recovery algorithm outputs b x according to the posterior distribution µ ( ·| y ).
In order to sample from the posterior, we use Langevin Dynamics [5].
Assuming we have access to ∇ x log µ ( x | y ), we can sample from µ ( x | y ) by running noisy gradient ascent:
[Equation]
Prior work [5] has shown that as t → ∞ and η t → 0, Langevin dynamics will correctly sample from µ ( x | y ).
In practice, vanilla Langevin Dynamics are slow to converge.
Hence, the work in [33] proposes annealed
Langevin Dynamics, where the marginal distribution of x at iteration t is modelled as µ t = µ ∗ N (0 , β 2 t ) and the generative model is trained to estimate the score function f ( x t ; β t ) :
= ∇ x t log(( µ ∗ N (0 , β 2 t )( x t )).
Since the distribution of y | x ∗ is Gaussian in Eqn (1), our ﬁnal algorithm with annealed Langevin dynamics is: for x 0 ∼ N c (0 , I ) and for all t = 0 , · ·
[Equation]
Note that the parameters T, { β t } T − 1 t =0 were ﬁxed during training of the generative model, and hence the only hyperparameters during inference are { η t } T − 1 t =0 and { γ t } T − 1 t =0 .
Appendix G describes hyper-parameter values used in our experiments.
Conclusions
This paper reports the ﬁrst wide-scale, successful application of the CSGM framework for multi-coil MR image reconstruction under realistic sampling conditions, and provides theoretical evidence for the robustness of posterior sampling.
Our score-based model was trained on a subset of brain MRI scans without explicit information about the sampling scheme (unlike end-to-end baselines) or the reconstruction task and shows a considerable degree of generalization to out-of-distribution samples such as abdomen and knee MRI.
These scans were acquired using diﬀerent MRI vendors with diﬀerent pulse sequence parameters and at diﬀerent institutions.
We postulate that adding a small set of diverse training samples to the generative model could further improve robustness, and that these samples may not necessarily be restricted to MR images.
Though promising, our initial results are still limited to fast spin-echo imaging only and all data were retrospectively under-sampled.
Further study is required to demonstrate prospective performance in a larger body of heterogeneous MRI data.
Our method also currently requires a high compute cost at inference time, as well as the need for a pre-trained generative model.
Clinical use requires fast reconstruction in addition to fast scanning.
Future work should investigate whether score-based models can be trained without a fully-sampled training set as well as investigate approaches to reducing computation time.  Acknowledgements
Ajil Jalal, Giannis Daras and Alex Dimakis have been supported by NSF Grants CCF 1763702, 1934932,
AF 1901281, 2008710, 2019844, the NSF IFML 2019844 award as well as research gifts by Western Digital,
Interdigital, WNCG and MLL, computing resources from TACC and the Archie Straiton Fellowship.
Eric Price has been supported by NSF Award CCF-1751040 (CAREER), NSF Award CCF-2008868, and NSF IFML 2019844. Marius Arvinte and Jon Tamir have been supported by NSF IFML 2019844 award, ONR grant N0001419-1-2590, NIH Grant U24EB029240, and an AWS Machine Learning Research Award. A Appendix: Theoretical Results
Background and Notation.
We ﬁrst introduce background and notation required for our theoretical results.
k · k refers to the ‘ 2 norm.
In this section alone, for simplicity of exposition, we will assume that all matrices and vectors are real valued.
For two probability distributions µ, ν on some normed space Ω, and for any q ≥ 1, the Wasserstein- q [40, 4] and Wasserstein- ∞ [9] distances are deﬁned as:
[Equation]
where Π( µ, ν ) denotes the set of joint distributions whose marginals are µ, ν .
The above deﬁnition says that if W ∞ ( µ, ν ) ≤ ε , and ( u, v ) ∼ γ , then k u − v k ≤ ε almost surely.
The ( ε, δ ) − approximate covering number [19], is deﬁned as the smallest number of ε -radius balls required to cover 1 − δ mass under a distribution.
Deﬁnition A.1 (( ε, δ )-approximate covering number) .
Let µ be a distribution on R N .
For some parameters ε > 0 , δ ∈ [0 , 1] , the ( ε, δ )-approximate covering number of µ is deﬁned as
[Equation]
where B ( x, ε ) is the ‘ 2 ball of radius ε centered at x .
Distributional robustness under Gaussian measurements.
First, we consider mismatch between the ground-truth distribution, denoted by µ , and the generator distribution, denoted by ν .
Prior work [19] has shown that if (i) W q ( µ, ν ) ≤ ε for some q ≥ 1 and (ii) we are given M ≥ O (log Cov ε,δ ( µ ))
Gaussian measurements, then posterior sampling with respect to ν will recover x ∗ ∼ µ up to an error of ε/δ 1 /q with probability 1 − δ .
Closeness in Wasserstein distance is a reasonable assumption in certain examples, such as when µ is the distribution of celebrity faces and ν is the distribution of a generator trained on FlickrFaces [21].
However, this assumption is unsatisfactory when we consider distributions of abdominal and brain MR scans, for example, since images of these anatomies look entirely diﬀerent.
We deﬁne the following weaker notion of divergence between distributions.
Informally, this new deﬁnition tells us that ν and µ are “close” if they can each be split into components which are close in W ∞ distance, such that the close components contain a suﬃciently large fraction under ν and µ .
Formally, this is deﬁned as:
Deﬁnition A.2 (( δ, α )- W ∞ divergence) .
For two probability distributions ν and µ , and parameters δ, α ∈ [0 , 1] , the ( δ, α ) - W ∞ divergence is deﬁned as
[Equation]
Lemma A.3 highlights that this is a strict generalization of Wasserstein distances, in the sense that closeness in Wasserstein distance implies closeness in this new divergence.
Lemma A.3 ( W q implies ( δ, α )- W ∞ ) .
If two distributions µ and ν satisfy W q ( µ, ν ) ≤ ε for some q ≥ 1 , then they satisfy ( δ, δ ) - W ∞ ( µ, ν ) ≤ ε/δ 1 /q .
Futhermore, there exist distributions that satisfy ( δ, δ ) - W ∞ ( µ, ν ) ≤ ε , but W q ( µ, ν ) = ∞ for all q ≥ 1 .
Since the ( δ, α )- W ∞ divergence is a generalization of Wasserstein distances, it is not clear that the main
Theorem in [19] holds for distributions that are close in this new divergence.
The following result shows a rather surprising fact: if ( δ, α )- W ∞ ( µ, ν ) ≤ ε then posterior sampling with M = O (cid:16) log (cid:16) 1 1 − α (cid:17) + log Cov ε,δ ( µ ) (cid:17) measurements will still succeed with probability ≥ 1 − O ( δ ).
Theorem A.4.
Let δ, α ∈ [0 , 1] , and ε > 0 be parameters.
Let µ, ν be arbitrary distributions over R N satisfying ( δ, α ) - W ∞ ( µ, ν ) ≤ ε .
Let x ∗ ∼ µ and suppose y =
Ax ∗ + w , where A ∈ R M × N and w ∈ R M are i.i.d.
Gaussian normalized such that A ij ∼ N (0 , 1 /M ) and w i ∼ N (0 , σ 2 /M ) , with σ (cid:38) ε .
Given y and the ﬁxed matrix A , let b x be the output of posterior sampling with respect to ν .
Then for M ≥ O (cid:16) log (cid:16) 1 1 − α (cid:17) + min(log Cov σ,δ ( µ ) , log Cov σ,δ ( ν )) (cid:17) , there exists a universal constant c >such that with probability at least 1 − e − Ω( M ) over A, w ,
[Equation]
For our running example of ν being a generator trained on brain scans, and µ the distribution of abdominal scans, we can set ν 0 to be the distribution of our generator restricted to abdominal scans, and we can let µbe the distribution restricted to “inliers” in µ .
This shows that even if our generator places an exponentially small probability mass(i.e.
, 1 − α (cid:28) 1) on the set of abdominal scans, we can still recover abdominal scans with a polynomial additive increase in the number of measurements (i.e., log(1 / (1 − α ))).
Near-optimality under arbitrary measurement processes.
The previous result required Gaussian matrices to handle the distribution shift.
Our next result shows that for an arbitrary measurement process, and assuming that there is no distribution shift between the generator and the ground truth distribution, posterior sampling is almost the best algorithm for this ﬁxed measurement process.
Theorem A.5.
Let x ∗ ∼ µ and let y
=
A ( x ∗ ) be measurements generated from x ∗ for some arbitrary forward operator A : R N → R M .
Then if there exists an algorithm that uses y as inputs and outputs x 0 such that
[Equation]
, then posterior sampling b x ∼ µ ( ·| y ) will satisfy k x ∗ − b x k ≤ 2 ε with probability ≥ 1 − 2 δ.
Remark on combining these results.
Our theoretical results above show that posterior sampling is (1) highly robust to distribution shift under Gaussian measurements, and (2) accurate with arbitrary measurements without distribution shift.
A natural hope would be to combine these two results and show that it is robust to distribution shift under Fourier measurements.
Unfortunately, this is not true for general distributions: for example, if µ and ν are both random distributions over Fourier-sparse signals, then Fourier measurements will usually give zero information about the signal, so cannot convince the sampler to sample near µ rather than ν .
A. Proof of Lemma A.
Proof.
Let Γ be a coupling between µ, ν such that E ( u,v ) ∼ Γ [ k u − v k q ] ≤ ε q .
Then an application of Markov’s inequality gives
[Equation]
Now, we can split the distribution Γ into two unnormalized components Γ 0 , Γ 00 deﬁned as
[Equation]
Using Γ 0 , Γ 00 , we can deﬁne measures µ 0 , µ 00 , ν 0 , ν 00 , via
[Equation]
where B is any measurable set and Ω is the state-space.
Since Γ is a valid coupling between µ, ν , and Γ 0 , Γ 00 are disjoint distributions, for any measurable B ⊆ Ω, we have:
[Equation]
µ 00 (Ω) . Using Eqn (4), we can conclude that µ 0 (Ω) ≥ 1 − δ, µ 00 (Ω) ≤ δ .
Setting µ 0 ← µ 0 /µ {+
+}
(Ω) and µ 00 ← µ 00 /µ {+
+}
(Ω), we can now rewrite µ as µ = (1 − δ ) µ 0 + δµ 00 .
A similar argument for ν gives ν = (1 − δ ) ν 0 + δν 00 .
By construction, µ 0 , ν 0 can be W ∞ coupled via Γ 0 to within a distance of ε/δ 1 /q .
This shows that ( δ, δ )- W ∞ ( µ, ν ) ≤ ε/δ 1 /q .
Now we need to show that two distributions can be close in ( δ, δ )- W ∞ , but W q = ∞ for all q .
Consider two scalar distributions µ, ν deﬁned as
[Equation]
Clearly, these distributions satisfy ( δ, δ )- W ∞ ( µ, ν ) ≤ ε , but W q ( µ, ν ) ≈ r for all q .
As r → ∞ , we get
W q ( µ, ν ) → ∞ for all q ≥ 1.
A. Proof of Theorem A.
In order to prove the Theorem, we make use of the following three Lemmas from [19].
Lemma A.6.
For c ∈ [0 , 1] , let H :
=
(1 − c ) H 0 + cH 1 be a mixture of two absolutely continuous distributions
H 0 , H 1 admitting densities h 0 , h 1 .
Let y be a sample from the distribution H , such that y | z ∗ ∼ H
z ∗ where z ∗ ∼ Bernoulli
( c ) .
[-y = ch 1 ( y ) (1 − c ) h 
-]
Deﬁne b c y = ch 1 ( y ) (1 − c ) h 0 ( y )+ ch 1 ( y ) , and let b z | y ∼ Bernoulli ( b c y ) be the posterior sampling of z ∗ given y .
Then we have
[Equation]
Lemma A.7.
Let y be generated from x ∗ by a Gaussian measurement process with noise rate σ .
For a ﬁxed ˜ x ∈ R n , and parameters η > 0 , c ≥ 4 e 2 , let P out be a distribution supported on the set
S ˜ x,out := { x ∈ R n : k x − ˜ x
k ≥ c ( η + σ ) } .
Let P ˜
x
be a distribution which is supported within an η − radius ball centered at ˜ x .
For a ﬁxed A , let H ˜ x denote the distribution of y when x ∗ ∼ P ˜ x .
Let H out denote the corresponding distribution of y when x ∗ ∼ P out .
Then we have:
[Equation]
Lemma A.8.
Let R, P, denote arbitrary distributions over R n such that W ∞ ( R, P ) ≤ ε .
Let x ∗ ∼ R and z ∗ ∼ P and let y and u be generated from x ∗ and z ∗ via a Gaussian measurement process with m measurements and noise rate σ .
Let b x ∼ P ( ·| y, A ) and b z ∼ P ( ·| u, A ) .
For any d > 0 , we have
[Equation]
Theorem A.4.
Let δ, α ∈ [0 , 1] , and ε > 0 be parameters.
Let µ, ν be arbitrary distributions over R N satisfying ( δ, α ) - W ∞ ( µ, ν ) ≤ ε .
Let x ∗ ∼ µ and suppose y =
Ax ∗ + w , where A ∈ R M × N and w ∈ R M are i.i.d.
Gaussian normalized such that A ij ∼ N (0 , 1 /M ) and w i ∼ N (0 , σ 2 /M ) , with σ (cid:38) ε .
Given y and the ﬁxed matrix A , let b x be the output of posterior sampling with respect to ν .
Then for M ≥ O (cid:16) log (cid:16) 1 1 − α (cid:17) + min(log Cov σ,δ ( µ ) , log Cov σ,δ ( ν )) (cid:17) , there exists a universal constant c >such that with probability at least 1 − e − Ω( M ) over A, w ,
[Equation]
Proof.
We know from ( δ, α )- W ∞ ( µ, ν ) ≤ ε that there exist µ 0 , ν 0 , µ 00 , ν 00 and a ﬁnite distribution Q supported on a set S such that
W ∞ ( µ 0 , ν 0 ) ≤ ε , 2.
min {W ∞ ( ν 0 , Q ) , W ∞ ( µ 0 , Q ) } ≤ σ , 3.
µ = (1 − δ ) µ 0 + δµ 00 and ν = (1 − α )
ν {+
+}
+ αν 00 .
Suppose W ∞ ( ν 0 , Q ) ≤ σ .
If not, then W ∞ ( µ 0 , Q ) ≤ σ , and by (1), we see that W ∞ ( ν 0 , Q ) ≤ σ + ε , and we will use this in the proof instead.
By decomposing µ = (1 − δ ) µ 0 + δµ 00 , we have
[Equation]
We now bound the second term on the right hand side of the above equation.
For this term, consider the joint distribution over x ∗ , A, w, b x .
By Lemma A.8, we can replace x ∗ ∼ µ 0 with z ∗ ∼ ν 0 , replace y =
Ax ∗ + w with u =
Az ∗ + w, and replace b x ∼ ν ( ·| A, y ) with b z ∼ ν ( ·| A, u ) to get the following bound
Pr
[Equation]
We now bound the second term in the right hand side of the above inequality.
Let Γ denote an optimal
W ∞ − coupling between ν 0 and Q .
For each ˜ z ∈ S , the conditional coupling can be deﬁned as
[Equation]
By the W ∞ condition, each Γ( ·| ˜ z ) is supported on a ball of radius σ around ˜ z .
Let E = { z ∗ , b z ∈ R n : k z ∗ − b z k ≥ (2 c + 1) σ } denote the event that z ∗ , b z are far apart.
By the coupling, we can express ν 0 as
[Equation]
This gives
[Equation]
For each ˜ z ∗ ∈ S , we now bound Q (˜
z ∗
[Equation]
For each ˜ z ∗ ∈ S , we can write ν as ν =
(1 − α ) Q ˜
z ∗ ν ˜ z ∗ , 0 + c ˜ z ∗ , 1 ν ˜ z ∗ , 1 + c ˜ z ∗ , 2 ν ˜ z ∗ , 2 , where the components of the mixture are deﬁned in the following way.
The ﬁrst component ν ˜ z ∗ , 0 is Γ( ·| ˜ z ∗ ), the second component is supported within a 2 cσ radius of ˜ z ∗ , and the third component is supported outside a 2 cσ radius of ˜ z ∗ .
Formally, let B ˜ z ∗ denote the ball of radius cσ centered at ˜ z ∗ , and let B c
˜
z ∗ be its complement.
The constants are deﬁned via the following Lebesque integrals, and the mixture components for any Borel measurable B are deﬁned as
[Equation]
Z
B ˜ z ∗ dν − (1 − α ) Q ˜ z ∗
Z
B ˜ z ∗ d Γ( ·| ˜ z ∗ ) , c ˜ z ∗ , 2 := Z B c ˜ z ∗ dν − (1 − α ) Q ˜
z ∗ Z
B c ˜ z ∗ d Γ( ·| ˜ z ∗ ) , ν ˜ z ∗ , 0 ( B ) : = Γ( B ∩ B ˜ z ∗ | ˜ z ∗ ) =
Γ( B | ˜ z ∗ ) since supp(Γ( ·| ˜ z ∗ ))
⊂ B ˜ z ∗ , ν ˜ z ∗ , 1 ( B ) :
= ( 1 c ˜ z ∗ , 1 ν ( B ∩ B ˜ z ∗ ) − 1 − α c ˜
z ∗ , 1 Q ˜ z ∗ Γ( B ∩ B ˜ z ∗ | ˜ z ∗ ) if c ˜ z ∗ , 1 > 0 , do not care otherwise.
, ν ˜ z ∗ , 2 ( B ) : = ( 1 c ˜ z ∗ , 2 ν ( B ∩ B c ˜ z ∗ ) −
[Equation]
α c ˜ z ∗ , 2 Q ˜ z ∗ Γ( B ∩ B c ˜ z ∗ | ˜ z ∗ ) if c ˜ z ∗ , 2 > 0 , do not care otherwise. .
Notice that if z ∗ is sampled from Γ( ·| ˜ z ∗ ), then by the W ∞ condition, we have k z ∗ − ˜ z ∗ k ≤ σ .
Furthermore, if b z is (2 c + 1) σ far from z ∗ , an application of the triangle inequality implies that it must be distributed according to ν ˜ z ∗ , 2 .
That is,
[Equation]
where H ˜ z ∗ , 0 , H ˜
z ∗ , 2 are the push-forwards of ν ˜ z ∗ , 0 , ν ˜ z ∗ , 2 for A ﬁxed and the last inequality follows from Lemma A.6.
Notice that if we sum over all ˜ z ∗ ∈ S , then the LHS of the above inequality is an expectation over z ∗ ∼ ν 0 .
This gives:
[Equation]
Notice that ν ˜ z ∗ , 0 is supported within an σ − ball around ˜ z ∗ , and ν ˜ z ∗ , 2 is supported outside a 2 cσ − ball of ˜ z ∗ .
By Lemma A.7 we have
[Equation]
This implies
[Equation]
where the last inequality is satisﬁed if m ≥ 4 log (cid:16) 1 1 − α (cid:17) + 4 log ( | S | ) .
Substituting in Eqn (6), if c > 4 exp (cid:16) 2 + 8 ε ( ε +2 σ ) σ(cid:17) , we have
[Equation]
This implies that there exists a set S A,w over A, w satisfying Pr A,w [ S A,w ] ≥ 1 − e − Ω( m ) , such that for all
A, w ∈ S A,w , we have
[Equation]
Substituting in Eqn (5), we have
[Equation]
Rescaling c gives us our result.
At the beginning of the proof, we had assumed that W ∞ ( ν 0 , Q ) ≤ σ .
If instead W ∞ ( µ 0 , Q ) ≤ σ , then we need to replace σ in the above bound by σ + ε .
Rescaling c in the above bound gives us the Theorem statement.
A. Proof of Theorem A.
Theorem A.5.
Let x ∗ ∼ µ and let y
=
A ( x ∗ ) be measurements generated from x ∗ for some arbitrary forward operator A : R N → R M .
Then if there exists an algorithm that uses y as inputs and outputs x 0 such that
[Equation]
probability [Equation] , then posterior sampling b x ∼ µ ( ·| y ) will satisfy [Equation]
Proof.
Without loss of generality, let the optimal algorithm output x 0 that is a deterministic function of y .
Yao’s lemma tells us that this is as good as a randomized function of y .
Then by the statement of the Lemma, we have
[Equation]
Now, the probability that k x ∗ − b x k ≤ 2 ε can be expressed as
[Equation]
where the second line follows from a triangle inequality, the third line follows since x ∗ , b x are independent conditioned on y , the fourth line follows since b x | y is distributed according to x ∗ | y , and the ﬁfth line follows from Jensen’s inequality.
B Appendix: Experimental Results
We perform retrospective under-sampling in all experiments, i.e., given fully-sampled k-space measurements from the NYU fastMRI [23, 42] and Stanford MRI [1] datasets, we apply sampling masks and evaluate the performance of all considered algorithms on the reconstructed data.
Depending on scan parameters (e.g., 3D scans for the Stanford knee data in Appendix F), we appropriately slice and sample the data in the proper dimension so as to not commit any inverse crime [15, 31].
We ﬁrst highlight that an advantage of the proposed approach is the invariance to the sampling scheme during training.
In contrast, this is a design choice that must be made for supervised end-to-end methods, which here were trained on equispaced, vertical sampling masks, following the fastMRI 2020 challenge guidelines [42, 28].
As our results show, this aﬀords us a signiﬁcant degree of robustness across a wide distribution of sampling masks during inference.
We train a score-based model, NCSNv2 [34], on a small subset of scans from the NYU fastMRI brain dataset.
Speciﬁcally, we train using T2-weighted images at a ﬁeld strength of 3 Tesla for a total of 14, 2D training slices.
We calculate the MVUE from the fully sampled data and use the ESPIRiT algorithm [37, 18] applied to the fully-sampled central portion of k-space to estimate the sensitivity maps.
The backbone network for our model is a ReﬁneNet [24].
Since the generator’s output is expected to be complex-valued, we treat the real and imaginary parts as separate image channels.
Details about the architectures are given in
Appendix G.
We use an ‘ 1 -Wavelet regularized reconstruction algorithm [25] as a parallel imaging and compressed sensing baseline.
We use the publicly available implementation from the BART toolbox [38, 36] and optimize the regularization hyper-parameter using the same subset of samples from the brain dataset that was used to train our method.
We ﬁnd that λ = 0 .
01 performs the best on the training data and use this value for all experiments.
Weconsider three diﬀerent deep learning baselines: MoDL [2], E2E-VarNet [35], and the ConvDecoder architecture [11].
The ﬁrst two methods, MoDL and E2E-VarNet, use supervised learning and belong to the general class of unrolled optimization methods, while the ConvDecoder is a self-supervised method that optimizes an individual model per scan.
We train the MoDL baseline using a residual network backbone on the same training dataset as our method, at an acceleration factor R = { 3 , 6 } and equispaced under-sampling, with a supervised SSIM loss on the MVUE image.
[Figure] [Figure]
[Figure]
For the E2E-VarNet baseline, we train a U-Net with a depth of four stages, with 18 hidden channels in the ﬁrst stage, for a total of 29 million learnable parameters.
This model also include a smaller deep neural network that is used to estimate the sensitivity maps.
This is also a U-Net, with four stages, but only eight hidden channels after the ﬁrst stage, for an additional 0 .
7 million parameters.
The model is trained for a number of 12 unrolls, and separate image networks are used at each unroll.
We train this model forepochs using an Adam optimizer with default PyTorch parameters and a learning rate of 2e − 4, as well as gradient clipping to a maximum magnitude of 1, on the fully-sampled MVUE reconstructions from the brain
T2 contrast used to train all methods.
We use a batch size of 1 for training and use a supervised SSIM loss between the absolute values of ground truth and reconstructed MVUE images at accelerations R = { 3 , 6 } , using a vertical, equispaced sampling pattern, same as all other baselines.
For the ConvDecoder baseline, we use the architecture for brain data in [11] and optimize the number of ﬁtting iterations on a subset of samples from the training data.
We ﬁnd that 10000 iterations are suﬃcient to reach a stable average performance at R = 3.
We evaluate reconstruction performance using the MVUE of the fully sampled data as a reference image and measure the structural similarity index (SSIM)
[41] between the absolute values of the reconstruction and ground-truth MVUE images.
[Figure]
B. In-Distribution Performance
In this experiment, we test all models using the same forward model that matches the training conditions for the baselines: vertical, equispaced sampling patterns.
Examples of various sampling patterns are shown in
Appendix C.
Figure 1 (top row) and Figure 2 show qualitative and quantitative results, respectively, for the case when there is no distributional mismatch between the training and inference sampling patterns.
From Figure 2, we notice that our method surpasses all baselines in the PSNR and SSIM metric at higher acceleration factors and is competitive at lower acceleration factors.
Note that at accelerations of R ≥ 6, MoDL and E2E-Varnet attain good SSIM and PSNR scores, but Figures 1, 10, 11 show that these methods suﬀer from heavy aliasing.
This highlights the importance of qualitative evaluations in medical image reconstruction and the limitations of existing image quality metrics [27].
We ﬁnd that ‘ 1 -Wavelet suﬀers both qualitatively and quantitatively at high acceleration factors, while the
ConvDecoder is also a competitive architecture, but incurs a large computational cost.
When benchmarked on an NVIDIA RTX 2080Ti GPU, our method takes 16 minutes and 0 .
95 GB of memory to reconstruct a highresolution brain scan, whereas the ConvDecoder takes longer than 80 minutes and 6 . 6 GB of memory.
While our method is limited by the inference time and is not in the range of end-to-end models (where reconstruction takes at most on the order of seconds and 3 .
5 GB of memory), multiple scans can be reconstructed in parallel due to the reduced memory footprint.
B. Out-of-Distribution Performance
Test-time anatomy shifts.
We now consider the much more diﬃcult problem of reconstructing diﬀerent anatomies than the ones seen during training.
This was previously investigated in [11], which concluded that all methods suﬀer a drastic shift due to the various changes in scan parameters between body parts.
In contrast to prior work, our main ﬁnding is that the proposed score-based model retains a signiﬁcant degree of robustness under these shifts, and outputs excellent qualitative reconstructions.
In some cases, end-to-end methods retain robustness as well.
Figure 4c shows SSIM scores obtained on reconstructed abdominal scans obtained from [1] at an acceleration factor of R = 4.
This represents both an anatomy and sampling pattern shift, and it can be seen that our method, MoDL and the ‘ 1 -Wavelet algorithm retain their competitive advantage, while the ConvDecoder and E2E-VarNet suﬀer severe performance losses.
Figure 6 further shows a qualitative comparison of a reconstructed abdominal scan, with highlighted artifacts.
Appendix E shows another abdomen scan.
Finally, Figure 4d shows SSIM scores obtained on fastMRI knee reconstructions, while Figure 1 (bottom row) shows the accompanying qualitative plots.
This anatomy is challenging especially because of the poor signal-to-noise ratio conditions, which can be seen even in the ground-truth image.
It can be noticed that this is the most severe shift for all methods, but our approach still shows the best performance at R = 2 , 4 and a signiﬁcantly lower variance.
Appendix D shows more examples of knee reconstructions with and without fat suppression.
B. Uncertainty Estimation
Our method can also provide uncertainty estimates for each reconstructed pixel by running multiple reconstruction samplers.
For a given observation y , we can obtain independent samples b x 1 , · · · , b x K ∼ µ ( ·| y ), for K suﬃciently large.
Now, using the conditional mean estimate ¯ x = P Ki =1 b x i /K , we can compute the pixel-wise standard deviation qP Ki =1 | b x i − ¯ x | 2 /K , and this gives an estimate of the error in each pixel.
As shown in
Fig 7, the pixel-wise standard deviation is a good estimate of the ground truth error | x ∗ − ¯ x | .
Additionally, notice that the reconstructions are able to recover ﬁne details such as the annotated meniscus tear 2 in Figand predict low uncertainty for these features.
We again emphasize that our model was not trained on knee scans.
Figure 15 in Appendix D shows another example of an annotated meniscus tear.
C Appendix: fastMRI Brain C. Examples of Sampling Masks
Figure 8 shows example of some of the masks used throughout the experiments in the paper and their corresponding reconstructions.
Note that the type of mask used is coupled with the scan parameters (e.g., two-dimensional slices from a three-dimensional scan will use a 2D grid of points).
We also highlight that, in all cases, a central region of the k-space is kept fully sampled and is used to estimate the coil sensitivity maps for all methods.
The bottom row of Figure 8 shows naive reconstructions of a single coil image using the zero-ﬁlled k-space.
This shows that diﬀerent types of masks lead to diﬀerent types of aliasing patterns in the image domain, motivating the need for robust image reconstruction algorithms.
C. More Exemplar Reconstructions
Figures 9 throughout 13 show detailed qualitative reconstructions on diﬀerent brain scans from the fastMRI dataset.
We highlight Figures 12 and 13, which represent a contrast shift from the in-distribution data (T and FLAIR vs. T2, respectively).
Our method still produces excellent qualitative reconstructions.
[Figure] [Figure]
[Figure] [Figure] [Figure]
Ground
[Figure]
L1-Wavelet
ConvDecoder
MoDL
E2E-Varnet
(Ours)
D Appendix: fastMRI Knee
Figure 14 shows further examples of knee reconstructions for the fastMRI dataset.
We highlight that the rows alternate between a fat suppressed (FS) and non-fat suppressed (NFS) contrast scan.
The FS scans are diﬃcult to reconstruct for the baselines, since they also exhibit a much lower signal-to-noise ratio than the brain training data, whereas our approach is robust to this type of shift as well.
����������������� ������������������� ������������������ ���������������������
E Appendix: Abdomen
Figure 16 shows an additional example of a reconstructed abdominal scan.
This is obtained from the same volume as the ﬁgure in the main text, and has a resolution of 164 × 320 voxels, but a much larger ﬁeld of view, leading to a resolution shift for all models.
[Figure]
F Appendix: Stanford Knee
Figures 17 and 18 show quantitative and qualitative reconstruction under an anatomy shift induced by testing axial knee scans.
In this case, we ﬁrst obtain a complete three-dimensional fast spin echo (3D-FSE) knee scan from the publicly available repository at mridata.org .
To obtain two-dimensional slices, we apply an IFFT operator on the readout axis and select 24 equally spaced slices for evaluation.
Each slice has a resolution of 320 × 256 pixels.
G Appendix: Implementation G. Score-Based Generative Model
Training the model We use the implementation from https://github.com/ermongroup/ncsnv2 .
As
MRI data are complex valued, we changed the generator such that the output and input have two channels, one each for the real and imaginary components.
We did not change the architecture otherwise.
We used the FlickrFaces (FFHQ) conﬁgs ﬁle from the NCSNv2 repo, except we set sigma_begin = 232, and sigma_end = 0.0066.
This is because of the smaller number of channels in MRI when compared to
FFHQ.
Dynamic range of the data.
MRI data exhibits a lot of variation in the dynamic range.
For example, the fastMRI dataset has max pixel value on the order of 10 − 4 , while the abdomen and Stanford knee data has max pixels on the order of 10 5 .
In order to deal with this variation, during training , we normalize each image by the 99 percentile pixel value.
During inference time, when we do not have access to the ground-truth image, we normalize the reconstruction using the 99 percentile pixel value of the pseudo-inverse .
We observe that this heuristic is suﬃcient to get good results.
Invariance to image shapes.
Due to the convolutional nature of NCSNv2, although we trained on 384 × 384 images, we can still apply them to knees, T1-weighted & FLAIR brains, and abdomens, although all of these have diﬀerent dimension shapes.
[Figure]
Hyperparameters for inference We run annealed Langevin dynamics given in Eqn (3) for T =steps.
The γ 0 and γ T hyperparameters in Eqn (3) were set to γ 0 = 232 , γ T = 0 . 00066.
The γ t values vary geometrically after every three steps, i.e.,
[Equation]
where r := (cid:16) γ T γ 0 (cid:17) 1 / ( T/ 3 − 1) =
(cid:0) 0 .
00066232 (cid:1) 1 / 2310 = 0 .
The learning rate η t in Eqn (3) also varies geometrically, as
[Equation]
T [Equation]
We use these hyperparameters for all datasets, all accelerations, and all sampling patterns .
G. E2E-VarNet Baseline
The backbone for the image reconstruction network is a U-Net with a depth of four stages, and 18 hidden channels in the ﬁrst stage, for a total of 29 million learnable parameters.
This model also include a smaller deep neural network that is used to estimate the sensitivity maps.
This is also a U-Net, with four stages, but only eight hidden channels after the ﬁrst stage, for an additional 0 .
7 million parameters.
The model is trained for a number of 12 unrolls, and separate image networks are used at each unroll.
We train this model for 50 epochs using an Adam optimizer with default PyTorch parameters and a learning rate of 2e − 4, as well as gradient clipping to a maximum magnitude of 1, on the fully-sampled MVUE reconstructions from the brain T2 contrast used to train all methods.
We use a batch size of 1 for training and use a supervised SSIM loss between the absolute values of ground truth and reconstructed MVUE images at accelerations R = { 4 , 8 } , using a vertical, equispaced sampling pattern, same as all other baselines.
Finally, it is worth mentioning that the network used to estimate the sensitivity maps explicitly uses the fully-sampled, vertical ACS region, as shown in Figure 8, both during training and inference.
This makes testing with other mask patterns non-trivial for this baseline.
To alleviate this, we always feed the image obtained from the vertical ACS region (for example, in the case of horizontal masks, we intentionally zero out other sampled lines that would fall in this region), to not introduce incoherent aliasing in this image.
[Figure]
L
[Figure]
MoDL
ConvDecoder
E2E [Figure] Varnet
Ours
G. MoDL Baseline
We train a MoDL model that uses a backbone residual network with a depth of six layers, three equispaced residual connections (that feed from the ﬁrst three layers to the last three layers) and 64 hidden channels, with a total of 220000 trainable parameters.
Unlike E2E-VarNet, the same backbone network is used across all unrolls, and the data consistency term is given by a Conjugate Gradient (CG) operator, truncated to six steps.
We train MoDL for a number of six unrolls, leading to a total of 36 CG steps and six network applications in the unroll.
We use the Adam optimizer with default PyTorch parameters and learning rate 2e − 4, as well as gradient clipping to a maximum magnitude of 1.
We train for ﬁve epochs, using a batch size of 1 on exactly the same T2 brain scans as all methods and a supervised SSIM loss at accelerations R = { 4 , 8 } using a vertical, equispaced sampling pattern and ﬁnd that this is suﬃcient to achieve excellent in-distribution reconstruction, due to the large dataset and small network size.
Since MoDL and all other methods, except E2E-VarNet, require external sensitivity map estimates to be provided to them, we use the ESPIRiT algorithm without any eigenvalue cropping to estimate a single set of sensitivity maps, one for each coil.
H Additional Results
Figure 5 shows the test PSNR evaluated in the same conditions as Figure 4 in the main text.
This highlights that our model is also robust in this metric, but also the shortcomings of PSNR in evaluating reconstruction performance, since for some methods the metric does not improve when more samples are available ( R is lower).