Robust Compressed Sensing MRI with Deep Generative Priors

Introduction

Compressed sensing 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) , and most MRI vendors have released products leveraging this framework to accelerate clinical workflows. Despite their successes, sparsity-based methods are limited by the achievable acceleration rates, as the sparsity assumptions are either hand-crafted or are limited to simple learned sparse codes .

More recently, deep learning techniques have been used as powerful data-driven reconstruction methods for inverse problems . There are two broad families of deep learning inversion techniques : 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. Network architectures that include both CNN blocks and the imaging forward model have grown in popularity, as they combine deep learning with the compressed sensing optimization framework, see e.g. . End-to-end methods are trained for specific imaging anatomy and measurement models and show excellent performance in these tasks. However, reconstruction quality is known to suffer when applied out of distribution, and recently has been shown to severely degrade under certain types of natural measurement and anatomy perturbations.

In this paper we study deep learning inversion techniques based on distribution learning. These models are trained without reference to measurements, and so easily adapt to changes in the measurement process. The most common family of such techniques, known also as Compressed Sensing with Generative Models (CSGM) uses pre-trained generative models as priors. Generative models are extremely powerful at representing image statistics and CSGM has been successfully applied to numerous inverse problems including non-linear phase retrieval , and improved with invertible models , sparsity based deviations , image adaptivity , and posterior sampling . 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 trains a StyleGAN for magnitude-only DICOM images but requires the presence of side-information and studies Gaussian, real-valued 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 . Another work trained an Invertible Neural Network on complex-valued single-coil MR images and showed very good performance in comparison to sparsity and GAN priors. Untrained and unamortized generators have also been recently explored , showing promising results in some cases. Further, studies the harder problem of learning a generative model for a class of images using only partial observations, as first proposed in AmbientGAN .

In this paper we train the first score-based generative model for MR images. We show that we can faithfully represent MR images without any assumptions on the measurement system. As a consequence, we are able to reconstruct retrospectively under-sampled MRI data under a variety of realistic sampling schemes. We show that our reconstruction algorithm is competitive with end-to-end supervised training when the test-data are matched to the training data and that it is robust to various out-of-distribution shifts, while in some cases end-to-end methods significantly degrade.

We successfully train a score-based deep generative model 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.

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 $\alpha$ mass to the true distribution, we show that posterior sampling for Gaussian measurements is nearly optimal with just an additive $O(\log(1/\alpha))$ loss.

We empirically show that our approach is robust to test-time distribution shifts including different 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 flexibility in choosing scan parameters – a common situation in routine clinical imaging. Perhaps surprisingly, the latter indicates that a specialized training set may offer sufficient regularization for a larger class of images. In contrast, 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.

Our method can be used to obtain multiple samples from the posterior by running Langevin dynamics with different random initializations. This allows us to get multiple reconstructions which can be used to obtain confidence intervals for each reconstructed voxel and visualize our reconstruction uncertainty on a voxel-by-voxel resolution. Uncertainty quantification can be incorporated into end-to-end methods, e.g., using variational auto-encoders , but this requires changes to the architecture. Our method does not require any modification 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 are robust when evaluated out-of-distribution.

2 Related Work

Generative priors have shown great utility to improving compressed sensing and other inverse problems, starting with , who generalized the theoretical framework of compressed sensing and restricted eigenvalue conditions for signals lying on the range of a deep generative model . Lower bounds in established that the sample complexities in are order optimal. The approach in has been generalized to tackle different inverse problems , and different reconstruction algorithms . The complexity of optimization algorithms using generative models have been analyzed in . Our prior work shows that posterior sampling is instance-optimal for compressed sensing , and satisfies certain fairness guarantees without explicit information about protected sensitive groups .

Using compressed sensing for multi-coil MRI reconstruction has led to a rich body of work in the past two decades . See and the recent special issue for an overview of these methods. Classical approaches impose sparsity in a well-chosen basis, such as the wavelet domain , or apply shallow learning that leverages low-level redundancy in the images . Recent research has demonstrated the superior performance of deep neural networks for MR image reconstruction . A broad class of approaches is represented by end-to-end unrolled methods, which use deep networks as learned data priors in the image or k-space domain . Recent work has also investigated the performance of untrained methods for MR reconstruction and has shown competitive results. A much less explored line of research is MR image reconstruction with generative priors. The work in proposes a CSGM-like algorithm that finetunes an entire pre-trained generator that requires a carefully tuned optimization algorithm during inference.

System Model and Algorithm

The acceleration factor $R$ denotes the degree of under-sampling in the $k$-space domain, i.e., $R=N/L$. Due to the multiple coils, the measurements may not be compressive for small $R$. However, due to redundancy between the coils, the measurements are compressive for moderate values of $R$ (even if $N_{c}\cdot L>N$) . Also note that we use the true acceleration factor $R$, and this does not match the values in fastMRI https://github.com/facebookresearch/fastMRI/blob/main/fastmri/data/subsample.py, line 247 has the fastMRI definition of equispaced acceleration factors. on certain sampling patterns.

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^{*}$. Prior work formulates this as a regularized optimization problem:

2 Posterior Sampling

In order to sample from the posterior, we use Langevin Dynamics . Assuming we have access to $\nabla_{x}\log\mu(x|y)$, we can sample from $\mu(x|y)$ by running noisy gradient ascent:

Prior work has shown that as $t\to\infty$ and $\eta_{t}\to 0$, Langevin dynamics will correctly sample from $\mu(x|y)$. In practice, vanilla Langevin Dynamics are slow to converge. Hence, the work in proposes annealed Langevin Dynamics, where the marginal distribution of $x$ at iteration $t$ is modelled as $\mu_{t}=\mu*\mathcal{N}(0,\beta_{t}^{2})$ and the generative model is trained to estimate the score function $f(x_{t};\beta_{t}):=\nabla_{x_{t}}\log((\mu*\mathcal{N}(0,\beta_{t}^{2})(x_{t}))$.

Since the distribution of $y|x^{*}$ is Gaussian in Eqn (2), we obtain $\nabla_{x_{t}}\log\mu(y|x_{t})=\frac{A^{H}(y-Ax_{t})}{\sigma^{2}}$. We find that it is also helpful to anneal this term, and we set it to $\frac{A^{H}(y-Ax_{t})}{\sigma^{2}+\gamma_{t}^{2}}$, where $\gamma_{t}\to 0$ is a decreasing sequence. An application of Bayes’ rule gives: $\nabla_{x_{t}}\log\mu(x_{t}|y)=f(x_{t};\beta_{t})+\frac{A^{H}(y-Ax_{t})}{\sigma^{2}+\gamma_{t}^{2}}$.

Putting everything together, our final algorithm is: for $x_{0}\sim\mathcal{N}_{c}(0,I)$ and for all $t=0,\cdots,T-1$,

Note that the parameters $T,\{\beta_{t}\}_{t=0}^{T-1}$ were fixed during training of the generative model, and hence the only hyperparameters during inference are $\{\eta_{t}\}_{t=0}^{T-1},\sigma$ and $\{\gamma_{t}\}_{t=0}^{T-1}$. Scripts in our codebase describe hyperparameter values used in our experiments.

Theoretical Results

For two probability distributions $\mu,\nu$ on some normed space $\Omega$, and for any $q\geq 1$, the Wasserstein-$q$ and Wasserstein-$\infty$ distances are defined as:

where $\Pi(\mu,\nu)$ denotes the set of joint distributions whose marginals are $\mu,\nu$. The above definition says that if $\mathcal{W}_{\infty}(\mu,\nu)\leq\varepsilon$, and $(u,v)\sim\gamma$, then $\|u-v\|\leq\varepsilon$ almost surely.

The $(\varepsilon,\delta)-$approximate covering number , is defined as the smallest number of $\varepsilon$-radius balls required to cover $1-\delta$ mass under a distribution.

Distributional robustness under Gaussian measurements.

First, we consider mismatch between the ground-truth distribution, denoted by $\mu$, and the generator distribution, denoted by $\nu$. Prior work has shown that if (i) $\mathcal{W}_{q}(\mu,\nu)\leq\varepsilon$ for some $q\geq 1$ and (ii) we are given $M\geq O(\log\operatorname{Cov}_{\varepsilon,\delta}(\mu))$ Gaussian measurements, then posterior sampling with respect to $\nu$ will recover $x^{*}\sim\mu$ up to an error of $\varepsilon/\delta^{1/q}$ with probability $1-\delta$. Closeness in Wasserstein distance is a reasonable assumption in certain examples, such as when $\mu$ is the distribution of celebrity faces and $\nu$ is the distribution of a generator trained on FlickrFaces . However, this assumption is unsatisfactory when we consider distributions of abdominal and brain MR scans, for example, since images of these anatomies look entirely different.

We define the following weaker notion of divergence between distributions. Informally, this new definition tells us that $\nu$ and $\mu$ are “close” if they can each be split into components which are close in $\mathcal{W}_{\infty}$ distance, such that the close components contain a sufficiently large fraction under $\nu$ and $\mu$. Formally, this is defined as:

For two probability distributions $\nu$ and $\mu$, and parameters $\delta,\alpha\in$, the $(\delta,\alpha)\text{-}\mathcal{W}_{\infty}$ divergence is defined as

Lemma B.1 highlights that this is a strict generalization of Wasserstein distances, in the sense that closeness in Wasserstein distance implies closeness in this new divergence.

Since the $(\delta,\alpha)\text{-}\mathcal{W}_{\infty}$ divergence is a generalization of Wasserstein distances, it is not clear that the main Theorem in holds for distributions that are close in this new divergence. The following result shows a rather surprising fact: if $(\delta,\alpha)\text{-}\mathcal{W}_{\infty}(\mu,\nu)\leq\varepsilon$ then posterior sampling with $M=O\left(\log\left(\frac{1}{1-\alpha}\right)+\log\operatorname{Cov}_{\varepsilon,\delta}(\mu)\right)$ measurements will still succeed with probability $\geq 1-O(\delta)$.

Then for $M\geq O\left(\log\left(\frac{1}{1-\alpha}\right)+\min(\log\operatorname{Cov}_{\sigma,\delta}(\mu),\log\operatorname{Cov}_{\sigma,\delta}(\nu))\right)$, there exists a universal constant $c>0$ such that with probability at least $1-e^{-\Omega(M)}$ over $A,w$,

For our running example of $\nu$ being a generator trained on brain scans, and $\mu$ the distribution of abdominal scans, we can set $\nu^{\prime}$ to be the distribution of our generator restricted to abdominal scans, and we can let $\mu^{\prime}$ be the distribution restricted to “inliers” in $\mu$. This shows that even if our generator places an exponentially small probability mass(i.e., $1-\alpha\ll 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-\alpha))$).

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 fixed measurement process. This result also shows that posterior sampling is good with respect to any metric.

then posterior sampling $\widehat{x}\sim\mu(\cdot|y)$ will satisfy

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 $\mu$ and $\nu$ 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 $\mu$ rather than $\nu$.

Experimental Results

We perform retrospective under-sampling in all experiments, i.e., given fully-sampled k-space measurements from the NYU fastMRI and Stanford MRI 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 .

We first 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 . As our results show, this affords us a significant degree of robustness across a wide distribution of sampling masks during inference.

We train a score-based model, NCSNv2 , on a small subset of scans from the NYU fastMRI brain dataset. Specifically, we train using T2-weighted images at a field strength of 3 Tesla for a total of 14,539 2D training slices. We calculate the MVUE from the fully sampled data and use the ESPIRiT algorithm applied to the fully-sampled central portion of k-space to estimate the sensitivity maps. The backbone network for our model is a RefineNet . 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 train the MoDL and E2E-VarNet baselines from scratch on the same training dataset as our method, at acceleration factors $R=\{3,6\}$ and equispaced under-sampling, with a supervised SSIM loss on the magnitude MVUE image, for $40$ and $15$ epochs, respectively, using a batch size of $1$. For the ConvDecoder baseline, we use the architecture for brain data in that outputs a complex image estimate and optimize the number of fitting iterations on a subset of samples from the training data. We find that $10000$ iterations are sufficient to reach a stable average performance at $R=3$. Put together, all of our baselines are tailored to estimate the complex image $x$, thus all comparisons are fair. We evaluate reconstruction performance using the complex MVUE of the fully sampled data as a reference image and measure the peak signal-to-noise ratio (PSNR) and structural similarity index (SSIM) between the absolute values of the reconstruction and ground-truth MVUE images.

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 three rows) shows qualitative results and Figures 2a & 5a respectively show PSNR & SSIM values, for the case where there is no mismatch between the training and inference sampling patterns. As the baselines were trained to maximize SSIM at $R=3\;\&\;6$, we see that they achieve better SSIM scores than us at these accelerations, although there is clear aliasing in the baselines at $R=6$. We achieve better PSNR values at these accelerations, which supports the claim that our method does not overfit to a particular metric (Theorem 3.4). This also highlights the importance of qualitative evaluations in medical image reconstruction and the limitations of existing image quality metrics . From the third row of Figure 1, and Figures 2a & 5a, we notice that our method surpasses baselines at higher accelerations.

2 Out-of-Distribution Performance

Here we consider shifts in the forward sampling operator at test-time, while still evaluating on the same anatomy as the training conditions. We measure robustness by evaluating the average incurred performance loss when the sampling pattern changes. Recall that our proposed approach does not use any explicit information about the sampling pattern $P$ during training, hence we anticipate the highest degree of robustness.

Figure 1 (fourth row) shows qualitative reconstructions when the measurements are obtained from an equispaced, horizontal sampling mask, with an acceleration factor $R=3$. It can be observed that the reconstructions output by E2E-VarNet show aliasing artifacts. Based on the statistical results in Figure 2b & 5b, our method retains its performance.

Furthermore, this experiment reveals that MoDL is more robust to this type of mask shift when compared to E2E-VarNet, even though it uses a smaller network. This is explained by the fact that E2E-VarNet does not use external sensitivity map estimates, but uses a deep neural network for end-to-end map estimation. While this improves performance on in-distribution samples, the performance drop is strong evidence that accurate sensitivity map estimation is vital for robust generalization, and both our proposed approach and MoDL benefit from the external ESPIRiT algorithm, which is compatible with different sampling patterns.

We do note that retrospectively flipping the horizontal and vertical sampling direction is not necessarily representative of prospective sampling in the horizontal direction due to the discrete nature of the phase encoding direction in MRI, and this may contribute to the higher scores compared to the vertical mask experiments.

Test-time anatomy shifts.

We now consider the more difficult problem of reconstructing different anatomies than the ones seen during. This was previously investigated in , which concluded that all methods suffer a drastic shift due to the various changes in scan parameters between body parts. In contrast to prior work, our main finding is that the proposed score-based model retains a significant degree of robustness under these shifts, and outputs excellent qualitative reconstructions. In some cases, some end-to-end methods retain robustness as well.

Finally, Figures 2d & 5d show PSNR and 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 significantly lower variance. Appendix D shows more examples of knee reconstructions with and without fat suppression, and Figure 20 shows metrics on fat suppressed knees.

3 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 $\widehat{x}_{1},\cdots,\widehat{x}_{K}\sim\mu(\cdot|y)$, for $K$ sufficiently large. Now, using the conditional mean estimate $\bar{x}=\sum_{i=1}^{K}\widehat{x}_{i}/K$, we can compute the pixel-wise standard deviation $\sqrt{\sum_{i=1}^{K}|\widehat{x}_{i}-\bar{x}|^{2}/K}$, and this gives an estimate of the error in each pixel. As shown in Fig 4, the pixel-wise standard deviation is a good estimate of the ground truth error $|x^{*}-\bar{x}|$. Additionally, notice that the reconstructions are able to recover fine details such as the annotated meniscus tearhttps://discuss.fastmri.org/t/219 in Fig 4 and predict low uncertainty for these features.

Figure 17 in Appendix D shows another example of an annotated meniscus tear. Figures 18 and 19 show comparisons with baselines on the same examples.

4 Radiologist Study

We have conducted a preliminary blind assessment of overall image quality with two board-certified radiologists and one faculty member who uses neuroimaging for their research. These experts were not involved in our research. We have found that our algorithm was ranked best for knee scans, and tied with the baselines for abdominal and brain scans, supporting our robustness claims in the paper. For more details, please see Appendix H.

Limitations

We reported PSNR and SSIM values as they are correlated with radiologist evaluation upto an extent, and our preliminary radiologist study in Section 4.4 suggests the feasibility of clinical adoption. These metrics do not capture the needs of real-world radiologists, and a more detailed study is required before the proposed techniques can be clinically adopted.

Though promising, our initial results were 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.

Finally, there are potential issues related to discrimination. Specifically, it is possible that the quality of the reconstructed images varies across protected attributes, such as gender or race .

Conclusions

This paper reports the first successful application of the CSGM framework for robust 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 small subset of brain MRI scans without any explicit information about the sampling scheme. This shows state-of-the-art performance under severe distributional shifts, making our model applicable in a wide range of clinical settings.

Our method shows a considerable degree of generalization to out-of-distribution samples such as abdomen and knee MRI, even when trained exclusively on brain MRI. Notably, these scans were acquired using different MRI vendors with different pulse sequence parameters and at different institutions. We postulate that adding a small set of diverse training samples to our generative model could further improve robustness, and we hypothesize that these samples may not necessarily be restricted to MR images.

The results presented in this work represent an important step to applying deep learning models in the clinic, as there is a natural variation in sampling, image orientation, receive coils, scanner hardware, and anatomy in clinical practice.

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 N00014-19-1-2590, NIH Grant U24EB029240, and an AWS Machine Learning Research Award.

We thank the anonymous NeurIPS reviewers for their helpful and considerate feedback.

Finally, we would like to thank the experts who graciously helped with our image assessment study.

References

Appendix A Appendix: Additional Metrics

Figure 5 shows the test SSIM evaluated in the same conditions as Figure 2 in the main text. This highlights that our model is also robust in this metric.

We observe that our method has significant noise in the background. Hence, we also report the masked SSIM and PSNR values in Figures 6 and 7. The mask zeros out all coordinates whose absolute value is smaller than 0.05 times the maximum absolute value in the fully-sampled MVUE.

The difference in numerical values between our results and the publicly available fastMRI leaderboard, as well as original results in the published baseline papers baselines comes from training and evaluating all methods on MVUE instead of RSS images. This is a design choice that we have made for all baselines, since our goal is to compare with a wide range of previous methods in a fair way.

Algorithms that output a complex-valued image (such as ours and L1-Wavelet) as a solution to the optimization in Eqn (2) will artificially perform worse (w.r.t. E2E methods) when compared to the RSS ground truth, even when the output is of similar or higher quality, due to the bias in the RSS. Since there is no way to obtain a good RSS score with these algorithms, this justifies our choice to train and evaluate all methods on MVUE.

To the best of our knowledge, a rigorous, reproducible comparison between end-to-end models trained on RSS or MVUE images has not been made in prior work. The recent work of has also discussed this point. To illustrate our claim of incompatibility between the two estimates, as well as the importance of qualitative inspection, we provide two simple, easy-to-verify examples.

We compare the fully sampled MVUE reconstruction (with ESPiRIT estimated maps) with the fully sampled RSS reconstruction, on T2 brain scans: we find that the SSIM is slightly larger than $0.8$. This is a large penalty (as per Fig. 1), even though the two images are virtually indistinguishable and known to be clinically equivalent (see discussions of SENSE vs. GRAPPA in ). This would unfairly penalize the family of methods that explicitly solve the inverse problem. Since E2E methods can be trained to target the MVUE directly, this justifies our choice for using the MVUE as the reference image.

We point to the public knee fastMRI leaderboard at https://fastmri.org/leaderboards. Selecting "Multi-coil Knee" and "4x" acceleration, we inspect the two following submissions:

"zero-filling", which does zero-filling RSS reconstruction, has an SSIM of $0.804$ and considerable artifacts.

"Baseline Classical Reconstruction Model", which applies compressed sensing with the ESPiRIT algorithm, has a much poorer SSIM score of $0.6275$, but produces qualitatively superior reconstructions.

Appendix B Appendix: Theory

If two distributions $\mu$ and $\nu$ satisfy $\mathcal{W}_{q}(\mu,\nu)\leq\varepsilon$ for some $q\geq 1$, then they satisfy $(\delta,\delta)\text{-}\mathcal{W}_{\infty}(\mu,\nu)\leq\varepsilon/\delta^{1/q}$. Futhermore, there exist distributions that satisfy $(\delta,\delta)\text{-}\mathcal{W}_{\infty}(\mu,\nu)\leq\varepsilon$, but $\mathcal{W}_{q}(\mu,\nu)=\infty$ for all $q\geq 1$.

Now, we can split the distribution $\Gamma$ into two unnormalized components $\Gamma^{\prime},\Gamma^{\prime\prime}$ defined as

Using $\Gamma^{\prime},\Gamma^{\prime\prime}$, we can define measures $\mu^{\prime},\mu^{\prime\prime},\nu^{\prime},\nu^{\prime\prime}$, via

where $B$ is any measurable set and $\Omega$ is the state-space.

Since $\Gamma$ is a valid coupling between $\mu,\nu$, and $\Gamma^{\prime},\Gamma^{\prime\prime}$ are disjoint distributions, for any measurable $B\subseteq\Omega$, we have:

Using Eqn (5), we can conclude that $\mu^{\prime}(\Omega)\geq 1-\delta,\mu^{\prime\prime}(\Omega)\leq\delta$. Setting $\mu^{\prime}\leftarrow\mu^{\prime}/\mu^{\prime}(\Omega)$ and $\mu^{\prime\prime}\leftarrow\mu^{\prime\prime}/\mu^{\prime\prime}(\Omega)$, we can now rewrite $\mu$ as $\mu=(1-\delta)\mu^{\prime}+\delta\mu^{\prime\prime}$. A similar argument for $\nu$ gives $\nu=(1-\delta)\nu^{\prime}+\delta\nu^{\prime\prime}$.

By construction, $\mu^{\prime},\nu^{\prime}$ can be $\mathcal{W}_{\infty}$ coupled via $\Gamma^{\prime}$ to within a distance of $\varepsilon/\delta^{1/q}$. This shows that $(\delta,\delta)\text{-}\mathcal{W}_{\infty}(\mu,\nu)\leq\varepsilon/\delta^{1/q}$.

Now we need to show that two distributions can be close in $(\delta,\delta)\text{-}\mathcal{W}_{\infty}$, but $\mathcal{W}_{q}=\infty$ for all $q$. Consider two scalar distributions $\mu,\nu$ defined as

Clearly, these distributions satisfy $(\delta,\delta)\text{-}\mathcal{W}_{\infty}(\mu,\nu)\leq\varepsilon$, but $\mathcal{W}_{q}(\mu,\nu)\approx r$ for all $q$. As $r\to\infty$, we get $\mathcal{W}_{q}(\mu,\nu)\to\infty$ for all $q\geq 1$.

In order to prove the Theorem, we make use of the following three Lemmas from .

For $c\in$, 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^{*}\sim H_{z^{*}}$ where $z^{*}\sim Bernoulli(c)$.

Define $\widehat{c}_{y}=\frac{ch_{1}(y)}{(1-c)h_{0}(y)+ch_{1}(y)}$, and let $\widehat{z}|y\sim Bernoulli(\widehat{c}_{y})$ be the posterior sampling of $z^{*}$ given $y$. Then we have

Let $x^{*}\sim R$ and $z^{*}\sim P$ and let $y$ and $u$ be generated from $x^{*}$ and $z^{*}$ via a Gaussian measurement process with $m$ measurements and noise rate $\sigma$. Let $\widehat{x}\sim P(\cdot|y,A)$ and $\widehat{z}\sim P(\cdot|u,A)$. For any $d>0$, we have

Then for $M\geq O\left(\log\left(\frac{1}{1-\alpha}\right)+\min(\log\operatorname{Cov}_{\sigma,\delta}(\mu),\log\operatorname{Cov}_{\sigma,\delta}(\nu))\right)$, there exists a universal constant $c>0$ such that with probability at least $1-e^{-\Omega(M)}$ over $A,w$,

We know from $(\delta,\alpha)\text{-}\mathcal{W}_{\infty}(\mu,\nu)\leq\varepsilon$ that there exist $\mu^{\prime},\nu^{\prime},\mu^{\prime\prime},\nu^{\prime\prime}$ and a finite distribution $Q$ supported on a set $S$ such that

$\mathcal{W}_{\infty}(\mu^{\prime},\nu^{\prime})\leq\varepsilon$,

$\min\{\mathcal{W}_{\infty}(\nu^{\prime},Q),\mathcal{W}_{\infty}(\mu^{\prime},Q)\}\leq\sigma$,

$\mu=(1-\delta)\mu^{\prime}+\delta\mu^{\prime\prime}$ and $\nu=(1-\alpha)\nu^{\prime}+\alpha\nu^{\prime\prime}$.

Suppose $\mathcal{W}_{\infty}(\nu^{\prime},Q)\leq\sigma$. If not, then $\mathcal{W}_{\infty}(\mu^{\prime},Q)\leq\sigma$, and by (1), we see that $\mathcal{W}_{\infty}(\nu^{\prime},Q)\leq\sigma+\varepsilon$, and we will use this in the proof instead. By decomposing $\mu=(1-\delta)\mu^{\prime}+\delta\mu^{\prime\prime}$, we have

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,\widehat{x}$. By Lemma B.4, we can replace $x^{*}\sim\mu^{\prime}$ with $z^{*}\sim\nu^{\prime}$, replace $y=Ax^{*}+w$ with $u=Az^{*}+w,$ and replace $\widehat{x}\sim\nu(\cdot|A,y)$ with $\widehat{z}\sim\nu(\cdot|A,u)$ to get the following bound

We now bound the second term in the right hand side of the above inequality. Let $\Gamma$ denote an optimal $\mathcal{W}_{\infty}-$coupling between $\nu^{\prime}$ and $Q$.

where the last inequality is satisfied if $m\geq 4\log\left(\frac{1}{1-\alpha}\right)+4\log\left(|S|\right).$

Substituting in Eqn (7), if $c>4\exp{\left(2+\frac{8\varepsilon\left(\varepsilon+2\sigma\right)}{\sigma^{2}}\right)},$ we have

This implies that there exists a set $S_{A,w}$ over $A,w$ satisfying $\Pr_{A,w}[S_{A,w}]\geq 1-e^{-\Omega(m)},$ such that for all $A,w\in S_{A,w},$ we have

At the beginning of the proof, we had assumed that $\mathcal{W}_{\infty}(\nu^{\prime},Q)\leq\sigma$. If instead $\mathcal{W}_{\infty}(\mu^{\prime},Q)\leq\sigma$, then we need to replace $\sigma$ in the above bound by $\sigma+\varepsilon$. Rescaling $c$ in the above bound gives us the Theorem statement.

B.2 Proof of Theorem 3.4

then posterior sampling $\widehat{x}\sim\mu(\cdot|y)$ will satisfy

By the statement of the Lemma, and conditioning on the measurements $y$, we have

Using a similar conditioning for the event $d(x^{*},\widehat{x})\leq 2\varepsilon$, we get

where the second line follows from a triangle inequality, the third line follows since $x^{*},\widehat{x}$ are independent conditioned on $y$, the fourth line follows since $\widehat{x}|y$ is distributed according to $x^{*}|y$, and the fifth line follows from Jensen’s inequality.

Appendix C Appendix: fastMRI Brain

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-filled k-space. This shows that different types of masks lead to different types of aliasing patterns in the image domain, motivating the need for robust image reconstruction algorithms.

C.2 More Exemplar Reconstructions

Figures 9 throughout 14 show detailed qualitative reconstructions on different brain scans from the fastMRI dataset. We highlight Figures 13 and 14, which represent a contrast shift from the in-distribution data (T1 and FLAIR vs. T2, respectively). Our method still produces excellent qualitative reconstructions.

Appendix D Appendix: fastMRI Knee

Figure 15 and Figure 16 show further examples of proton density knee reconstructions.

Figure 18 and Figure 19 show comparisons of our method and baselines on knees with meniscus tears. Figure 17 shows uncertainty estimates from our algorithm on a knee with a meniscus tear.

Figure 20 shows PSNR and SSIM on fat-suppressed(FS) knees. Our approach is not optimal numerically, likely due to a much lower signal-to-noise ratio in FS knees than the brain training data. However, Figures 18, 19, 21, 22 show that our qualitative reconstructions are competitive, and recovers fine details (like meniscus tears) better than the deep learning baselines.

Appendix E Appendix: Abdomen

Figure 23 shows an additional example of a reconstructed abdominal scan. This is obtained from the same volume as the figure in the main text, and has a resolution of $158\times 320$ voxels, but a much larger field of view, leading to a resolution shift for all models.

Appendix F Appendix: Stanford Knee

Figures 24 and 25 show quantitative and qualitative reconstruction under an anatomy shift induced by testing axial knee scans. In this case, we first 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\times 256$ pixels.

Appendix G Appendix: Implementation

We use the implementation from https://github.com/ermongroup/ncsnv2. As raw MRI scans 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) configs file 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 complex image. We observe that this heuristic is sufficient to get good results.

Invariance to image shapes.

Due to the convolutional nature of NCSNv2, although we trained on $384\times 384$ images, we can still apply them to knees, T1-weighted & FLAIR brains, and abdomens, although all of these have different dimension shapes.

Hyperparameters

We tuned our hyperparameters on two validation brain scans, at an acceleration of $R=4$. We then reused these hyperparameters on all anatomies, all accelerations. Please see our GitHub link: https://github.com/utcsilab/csgm-mri-langevin for the hyperparameter values.

G.2 E2E-VarNet Baseline

We use the architecture publicly available in the fastMRI official repository. The backbone for the image reconstruction network is a U-Net with a depth of four stages, and $18$ hidden channels in the first 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 first 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.

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.

G.3 MoDL Baseline

We use the PyTorch MoDL implementation publicly available at https://github.com/utcsilab/deep-jsense and train a MoDL model that uses a backbone residual network with a depth of six layers, three equispaced residual connections (that feed hidden signals from the first 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.

Since MoDL and all other methods (including ours) except E2E-VarNet, require external sensitivity map estimates to be provided to them, we use the ESPIRiT algorithm from the BART toolbox without any eigenvalue cropping to estimate a single set of sensitivity maps, one for each coil.

Appendix H Appendix: Radiologist Study

We performed a preliminary image quality assessment experiment with two board-certified radiologists and a faculty member that uses neuro-imaging in their research.

The three external experts were not involved with our research and have performed the image quality assessment blindly. Each of them was presented with ten scans from the following anatomies and scan parameters: abdominal scans, knee scans and brain scans with a horizontal readout direction, leading to a total of 30 quality assessment questions. Note that all anatomies represent test-time distributional shifts in at least one aspect.

In each question, the experts were shown four images:

The fully-sampled reference image, explicitly marked as "Reference".

The results of three reconstruction algorithms at acceleration factor R=3: MoDL, ConvDecoder and our method. The order of the reconstructions was shuffled for each question, and the reconstructions were labeled as "1", "2" and "3".

We chose to compare with MoDL and ConvDecoder since these method had the best overall quantitative and qualitative (according to our own pre-assessment) robust performance. The participants were instructed to rank the three reconstructions from best to worst quality, while using the "Reference" image as a perceptual guideline. Table 1 shows the average and standard deviation (in parentheses) of the ranking for each anatomy, obtained using a total of 30 data points (3 participants x 10 scans per anatomy).

In Table 1, a lower ranking is better, the best possible ranking is 1, and the worst 3. We draw the following conclusions:

Participants consistently ranked our method as best on the knee scans, which supports the distributional shift robustness claimed in the main paper, and detailed in Appendices D, E and A.

Participants did not perceive a significant difference between all methods when applied to abdominal or brain scans with a horizontal phase encode direction. In the brain case, this supports the qualitative results shown in Appendix C, Figure 9.

In the abdominal case, this partially correlates with Figure 2c, regarding the quantitative tie between our approach and MoDL.

To quantify the statistical significance of the above results, we perform a Wilcoxson Rank Sum test to determine if the rankings of different algorithms are drawn from different populations. We evaluate if our proposed method leads to different rankings than MoDL and the ConvDecoder, and show the p-values in Table 2.

The results show a significant difference in the case of knees, while no significant difference is present for abdomen and brain. Finally, to evaluate inter-observer agreement between the three reviewers, we calculated the intra-class correlation (ICC) coefficient separately for each anatomy by aggregating the ten questions related to that anatomy and evaluating the ICC2 coefficient in a pairwise manner at a $5\%$ significance level.

The results are shown in Table 3, where we also include the p-value and the $95\%$ confidence interval for the ICC2 estimate. This indicates that there exists a very strong consensus regarding the ranking on the knee anatomy, while for abdomen and brain this consensus is much weaker, which together with Table 2 indicates that the images were considered equivalent.

This preliminary image quality assessment gives additional evidence (in addition to the quantitative metrics of SSIM and PSNR) that our method maintains robustness to distribution shifts at test time. As our quantitative results show, other methods maintain robustness in some but not all cases. Due to time limitations, we were not able to ask the reviewers to evaluate every algorithm and every distribution shift including different levels of acceleration. We stress that this preliminary study is not a substitute for a rigorous clinical evaluation which is necessary before considering using our proposed method in a clinical setting.