Robust Compressed Sensing MRI with Deep Generative Priors

1 Introduction . Compressed sensing [ 23 , 15 ] 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 ) [ 62 , 90 ] , and most MRI vendors have released products leveraging this framework to accelerate clinical workflows . Despite their successes , sparsitybased methods are limited by the achievable acceleration rates , as the sparsity assumptions are either hand-crafted or are limited to simple learned sparse codes [ 72 , 73 ] . More recently , deep learning techniques have been used as powerful data-driven reconstruction methods for inverse problems [ 49 , 68 ] . There are two broad families of deep learning inversion techniques [ 68 ] : 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 . [ 32 , 3 , 64 ] . 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 [ 4 , 19 ] 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 ∗Ajil Jalal and Marius Arvinte contributed equally to this work . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . process . The most common family of such techniques , known also as Compressed Sensing with Generative Models ( CSGM ) [ 13 ] 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 [ 13 , 34 ] including non-linear phase retrieval [ 35 ] , and improved with invertible models [ 6 ] , sparsity based deviations [ 21 ] , image adaptivity [ 42 ] , and posterior sampling [ 79 , 45 ] . 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 [ 53 ] 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 [ 77 ] . Another work [ 54 ] 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 [ 37 ] have also been recently explored [ 19 ] , showing promising results in some cases . Further , [ 17 ] studies the harder problem of learning a generative model for a class of images using only partial observations , as first proposed in AmbientGAN [ 14 ] . In this paper we train the first score-based generative model [ 80 ] 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 . 1.1 Contributions . • 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 α mass to the true distribution , we show that posterior sampling for Gaussian measurements is nearly optimal with just an additive O ( log ( 1/α ) ) 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 [ 24 ] , 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 . 1.2 Related Work . Generative priors have shown great utility to improving compressed sensing and other inverse problems , starting with [ 13 ] , who generalized the theoretical framework of compressed sensing and restricted eigenvalue conditions [ 85 , 23 , 12 , 15 , 40 , 11 , 10 , 25 ] for signals lying on the range of a deep generative model [ 29 , 55 , 81 ] . Lower bounds in [ 51 , 61 , 48 ] established that the sample complexities in [ 13 ] are order optimal . The approach in [ 13 ] has been generalized to tackle different inverse problems [ 47 , 35 , 7 , 71 , 60 , 63 , 74 , 9 ] , and different reconstruction algorithms [ 21 , 50 , 69 , 27 , 26 , 64 , 37 , 38 , 18 ] . The complexity of optimization algorithms using generative models have been analyzed in [ 28 , 39 , 58 , 36 ] . Our prior work shows that posterior sampling is instance-optimal for compressed sensing [ 45 ] , and satisfies certain fairness guarantees without explicit information about protected sensitive groups [ 46 ] . Using compressed sensing for multi-coil MRI reconstruction has led to a rich body of work in the past two decades [ 62 , 20 , 87 , 75 ] . See [ 22 ] and the recent special issue [ 44 ] for an overview of these methods . Classical approaches impose sparsity in a well-chosen basis , such as the wavelet domain [ 62 ] , or apply shallow learning that leverages low-level redundancy in the images [ 72 , 73 , 93 ] . Recent research has demonstrated the superior performance of deep neural networks for MR image reconstruction [ 76 , 32 , 3 , 82 , 83 ] . A broad class of approaches is represented by end-to-end unrolled methods , which use deep networks as learned data priors in the image [ 3 , 32 , 82 ] or k-space domain [ 84 ] . Recent work has also investigated the performance of untrained methods [ 89 , 38 ] 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 [ 67 ] proposes a CSGM-like algorithm that finetunes an entire pre-trained generator that requires a carefully tuned optimization algorithm during inference . 2 System Model and Algorithm . 2.1 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 [ 78 , 70 , 30 ] . Formally , the vector of measurements yi ∈ CL acquired by the ith coil can be characterized by the forward model [ 70 ] : yi = PFSix ∗ + wi , i = 1 , ... , Nc , ( 1 ) where x∗ ∈ CN is the image containing N pixels , Si is an operator representing the point-wise multiplication of the ith coil sensitivity map , F is the spatial Fourier transform operator , P represents the k-space sampling operator , and we assume wi ∼ Nc ( 0 , σ2I ) for simplicity . Importantly , note that the same under-sampling operator is applied to all Nc coils . The acceleration factorR 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 Nc · L > N ) [ 41 ] . Also note that we use the true acceleration factor R , and this does not match the values in fastMRI [ 56 ] 2 on certain sampling patterns . 2https : //github.com/facebookresearch/fastMRI/blob/main/fastmri/data/subsample.py , line 247 has the fastMRI definition of equispaced acceleration factors . 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 : arg min x ‖y −Ax‖22 + λQ ( x ) , ( 2 ) where we use the operator A ∈ CM×N ( with M = Nc ·L ) to subsume the discrete approximation to all linear effects , and Q is a suitably chosen functional prior for the image variable x . For example , to enforce a sparsity prior , one can penalize the ` 1 norm in the wavelet representation of x [ 62 ] . More recent approaches involve learned regularization terms parameterized by deep neural networks [ 76 , 32 , 3 ] . These models are typically trained end-to-end using a fixed training set and certain assumptions about the sampling operator . In the sequel , we present how score-based generative models can be combined with the posterior sampling [ 45 ] mechanism to reformulate ( 2 ) and achieve good quality reconstructions without any a priori assumptions about the sampling scheme . When k-space is fully sampled at the Nyquist rate and no regularization is applied , the solution to ( 2 ) corresponds to the minimum-variance unbiased estimator ( MVUE ) of x∗ , denoted by x̂MVUE [ 70 ] . Given fully sampled k-space data , this estimate can act as a reference image for evaluating reconstruction error as well as for end-to-end training . Alternatively , a reference image called the root-sum-of-squares ( RSS ) estimate can be formed by taking the inverse Fourier transform of each coil and subsequently applying the ` 2 norm for each pixel across the coil dimension , i.e . x̂RSS = √∑Nc i=1 | ( FHyi ) | 2 , where FH is the Hermitian transpose of F ( here the inverse DFT ) . Although the RSS estimate is a biased estimator , it is often used as it does not make any assumptions about the sensitivity maps , which are not explicitly measured by the MRI system . However , even if solving ( 2 ) results in perfect recovery of x∗ , there will be a bias when comparing the result to x̂RSS and thus the RSS and MVUE can not be directly compared numerically .

Summary: This paper considers the problem of compressed sensing (CS) using generative priors, which can also specifically be characterized on a multi-coil MRI setup. The algorithm itself is a more general version of the CS setup which typically considers Gaussian measurements.  The authors consider a posteriori sampling procedure to invert the measurements $y = Ax^* + w$ by estimating $\hat{x} \sim \mu(\cdot|y)$ (modeled via deep generative prior) assuming $x^* \sim \mu$, given $A$ (MRI setup or general), $y$ and zero-mean $\sigma^2$ variance Gaussian noise $w$ . Authors use an annealed Langevin Dynamics approach to sample from the posterior, which has been previously used in [73] for sampling from generative priors. Here the score function is modeled via a deep generative neural network.  Authors also present two theoretical results: one on the closeness of reconstruction $\hat{x}$ and $x^*$ under distribution shift of the deep generative prior and true distributions (assuming Gaussian A), and another on the quality of posteriori sampling as compared to any other kind of reconstruction procedure under aribitrary (not neccessarily Gaussian) measurements assuming no distributional shift.  Authors present several experiments on the fastMRI and Stanford MRI datasets and compare to baselines such as L1/wavelet, MoDL, E2EVarNet and ConvDecoder and show  improvement reconstructions on both in-distribution and out-of distribution samplings in terms of SSIM.