PaperlayerView exact revision on arXiv

Source-linked AI summary

Robust Compressed Sensing MRI with Deep Generative Priors

Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros G. Dimakis, Jonathan I. Tamir

arXiv:2108.01368v2cs.LGcs.CVcs.ITstat.ML

TL;DR

Existing compressed-sensing MRI methods have limited acceleration and CSGM had not been demonstrated successfully on clinical MRI. The paper trains a score-based generative prior on brain MR images and uses posterior sampling with Langevin dynamics, finding competitive reconstruction and robustness to sampling and anatomy shifts. These results support posterior sampling as a flexible approach for multi-coil MRI, while clinical validation and inference efficiency remain open limitations.

  • Problem

    Sparsity-based MRI methods are limited by achievable acceleration rates, while CSGM had not been shown successful on clinical MRI or robust to distribution shifts.

  • Method

    The paper trains a score-based generative model for complex-valued brain MR images and combines it with posterior sampling for multi-coil MRI reconstruction.

  • Results

    The approach achieves competitive in-distribution reconstruction and is robust to shifts in sampling patterns and imaging anatomy, including knee and abdomen scans trained from brain MRI.

  • Takeaways & Limitations

    Posterior sampling provides reconstruction flexibility across realistic measurement and anatomy changes without training the generative model on a specific sampling scheme.

  • Takeaways & Limitations

    Results are limited to fast spin-echo imaging and retrospectively undersampled data, while inference requires high computation and a pretrained generative model.

Abstract

from arXiv · show

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 first 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. Our code and models are available at: \url{https://github.com/utcsilab/csgm-mri-langevin}.

1 Introduction

The paper applies score-based generative priors and posterior sampling to multi-coil MRI, addressing limitations of sparsity-based and end-to-end methods under distribution shifts. It reports competitive in-distribution reconstruction, robustness to sampling and anatomy shifts, theoretical support for posterior sampling, and voxel-level uncertainty estimation.

  • Sparsity-based MRI reconstruction reduces measurements but remains limited by achievable acceleration rates and hand-crafted or simple learned sparsity assumptions.
  • The paper trains a score-based generative model for complex-valued T2-weighted brain MR images without assumptions about the measurement scheme.
  • Posterior sampling with the correct prior is within constant factors of optimal recovery for arbitrary measurements, while a mismatched prior incurs an additive O(log(1/α)) loss for Gaussian measurements.
  • The method is empirically robust to shifts in sampling patterns and imaging anatomy, whereas end-to-end methods can severely degrade out of distribution.
  • Langevin dynamics with different random initializations produces multiple posterior reconstructions for voxel-level confidence intervals and uncertainty visualization without architectural changes.
  • The CSGM approach achieves competitive performance with end-to-end deep learning methods on in-distribution multi-coil MRI data using a reduced training set.

2 System Model and Algorithm

The MRI system models undersampled multicoil k-space measurements as coil-sensitive Fourier observations and reconstructs the underlying image using a learned prior. The algorithm combines posterior sampling with score estimation and Langevin dynamics, while accounting for estimator choice and sampling acceleration.

  • Multi-coil Magnetic Resonance Imaging: Each coil measures an undersampled Fourier transform of the sensitivity-weighted image plus complex Gaussian noise, using the same sampling operator across coils.
  • Multi-coil Magnetic Resonance Imaging: The acceleration factor R = N/L measures k-space undersampling, with coil redundancy making measurements compressive for moderate R.
  • Multi-coil Magnetic Resonance Imaging: Reconstruction estimates x∗ from measurements y, sensitivity maps S, and sampling operator P through a regularized optimization problem with functional prior Q.
  • Multi-coil Magnetic Resonance Imaging: MVUE references can support reconstruction-error evaluation and end-to-end training, whereas RSS is biased and cannot be directly compared numerically with MVUE.
  • Posterior Sampling: Posterior sampling draws reconstructions from µ(·|y) for observations y = Ax∗ + w, using the assumed prior distribution µ.
  • Posterior Sampling: Langevin dynamics samples the posterior through noisy gradient ascent, combining the generative model's score f(xt; βt) with the likelihood gradient AH(y − Axt).
  • Posterior Sampling: The final sampling algorithm initializes x0 from a complex Gaussian distribution and iterates updates for t = 0, ..., T−1, with inference hyperparameters including the step sizes ηt.

3 Theoretical Results

The theoretical results develop a weaker distribution-divergence framework for analyzing posterior sampling under distribution mismatch, while establishing robustness under Gaussian measurements and near-optimality under arbitrary measurements without shift. An abdominal-versus-brain example illustrates how recovery can remain possible despite very small generator mass on abdominal scans, while Fourier-measurement robustness under general shifts is not guaranteed.

  • Distributional robustness under Gaussian measurements: Prior Wasserstein-based guarantees require M ≥ O(log Covε,δ(µ)) Gaussian measurements and recover with error ε/δ^1/q under Wq closeness.The paper notes that Wasserstein closeness is unsuitable when abdominal and brain MR distributions look entirely different.
  • Distributional robustness under Gaussian measurements: The (δ, α)-W∞ divergence decomposes each distribution into components whose sufficiently large portions are close in W∞ distance.It is a strict generalization of Wasserstein distance, because Wasserstein closeness implies closeness in this divergence.
  • Distributional robustness under Gaussian measurements: Theorem 3.3 analyzes posterior sampling from ν when the ground truth follows µ and the measurements use Gaussian A and Gaussian noise with σ ≳ ε.The theorem assumes (δ, α)-W∞(µ, ν) ≤ ε and considers posterior sampling with respect to ν.
  • Distributional robustness under Gaussian measurements: For a brain-trained generator and abdominal ground-truth distribution, posterior sampling can recover abdominal scans with a polynomial additive increase in measurements despite exponentially small generator mass on abdominal scans.The argument restricts the generator and ground-truth distributions to abdominal scans and their inliers, respectively.
  • Near-optimality under arbitrary measurement processes: Under an arbitrary measurement process, posterior sampling is almost optimal for any metric when the generator and ground-truth distributions match.The result considers measurements y = A(x*) from an arbitrary forward operator and compares posterior sampling with an algorithm using y.
  • Near-optimality under arbitrary measurement processes: The paper cautions that Gaussian-measurement robustness and arbitrary-measurement accuracy cannot generally be combined into distribution-shift robustness for Fourier measurements.Fourier measurements may provide zero information for some random Fourier-sparse distributions, preventing the sampler from distinguishing them.

4 Experimental Results

Experiments evaluate the score-based generative prior against compressed-sensing and end-to-end baselines under matched and shifted sampling or anatomy conditions. The method remains competitive under distribution shifts, recovers fine details, and provides uncertainty estimates, although inference is computationally expensive.

  • In-distribution and sampling shifts: Training without a fixed sampling scheme gives the proposed method robustness across a wide distribution of inference sampling masks.The supervised baselines were trained with equispaced vertical masks, whereas the proposed approach was not tied to that sampling choice.
  • Experimental setup: The model is trained on 14,539 T2-weighted 2D brain slices using NCSNv2 and ESPIRiT sensitivity maps, while baselines include ℓ1-Wavelet, MoDL, E2E-VarNet, and ConvDecoder.All experiments use retrospective under-sampling from fastMRI and Stanford MRI data.
  • In-distribution and sampling shifts: Under matched sampling, baselines achieve better SSIM at their training accelerations, while the proposed method achieves better PSNR and surpasses baselines at higher accelerations.The baselines were trained to maximize SSIM at R = 3 and 6, and aliasing is visible at R = 6.
  • Computational cost: Reconstructing a high-resolution brain scan takes 16 minutes and 0.95 GB of memory, compared with more than 80 minutes and 6.6 GB for ConvDecoder.The method remains slower than end-to-end models, but its lower memory footprint allows multiple scans to be reconstructed in parallel.
  • In-distribution and sampling shifts: For a horizontal mask shift at R = 3, E2E-VarNet shows aliasing artifacts while the proposed method retains its performance; MoDL is more robust than E2E-VarNet.The authors attribute the difference partly to sensitivity-map estimation and note that ESPIRiT supports different sampling patterns.
  • Out-of-distribution performance: Under anatomy shifts, the method remains competitive on abdominal scans and achieves the best performance at R = 2, 4 on challenging fastMRI knee reconstructions with lower variance.ConvDecoder and E2E-VarNet suffer severe abdominal performance losses, while the knee shift is severe for all methods.
  • Uncertainty estimation: With 48 posterior samples at R = 4, pixel-wise standard deviation estimates reconstruction error while fine details such as a meniscus tear are recovered with low uncertainty.The knee pathology was absent from the T2-weighted brain-scan training data.

5 Limitations

The evaluation is preliminary and narrow in clinical scope, while practical deployment is constrained by retrospective data, inference cost, and potential variation across protected attributes.

  • Clinical evaluation: PSNR and SSIM correlate with radiologist evaluation only to an extent, so a more detailed study is required before clinical adoption.The preliminary radiologist study suggests feasibility but does not capture all real-world radiologist needs.
  • Scope and validation: Results are limited to fast spin-echo imaging and retrospectively under-sampled data, requiring prospective evaluation on larger heterogeneous MRI datasets.The authors identify prospective performance and broader data heterogeneity as necessary further study.
  • Deployment constraints: Inference currently has high compute cost and requires a pre-trained generative model, while clinical use also requires fast reconstruction.The paper notes that reconstruction speed remains slower than end-to-end models.
  • Fairness: Reconstruction quality may vary across protected attributes such as gender or race, creating potential discrimination concerns.The paper identifies this as a potential issue requiring attention.

6 Conclusions

The paper demonstrates robust multi-coil MR reconstruction under realistic sampling conditions and severe distributional shifts, with implications for clinical variability.

  • The score-based model achieved state-of-the-art performance under severe distributional shifts in multi-coil MR reconstruction.It was trained on a small subset of brain MRI scans without explicit information about the sampling scheme.
  • The method generalized to abdomen and knee MRI despite training exclusively on brain MRI.These scans came from different vendors, pulse-sequence parameters, and institutions.
  • The results address clinical variation in sampling, image orientation, receive coils, scanner hardware, and anatomy.

A Appendix: Additional Metrics

The appendix reports masked image-quality metrics and explains why MVUE rather than RSS images were used for fair evaluation, while also analyzing Wasserstein-distance relationships.

  • Additional Metrics: Masked SSIM and PSNR were reported because the method produced significant background noise.The mask zeros coordinates below 0.05 times the maximum absolute value in the fully sampled MVUE.
  • MVUE vs. RSS: All methods were trained and evaluated on MVUE images to enable fair comparison across reconstruction approaches.Numerical values therefore should not be compared directly with fastMRI leaderboard results or papers using RSS images.
  • MVUE vs. RSS: MVUE and RSS can receive substantially different SSIM scores despite virtually indistinguishable, clinically equivalent reconstructions.The fully sampled comparison found SSIM slightly above 0.8.
  • MVUE vs. RSS: The classical reconstruction baseline had SSIM 0.6275 but produced qualitatively superior reconstructions, illustrating metric limitations.Zero-filling had SSIM 0.804 and considerable artifacts.
  • Wasserstein Relationships: If Wq(µ, ν) ≤ ε, then (δ, δ)-W∞(µ, ν) ≤ ε/δ^(1/q), while the converse need not hold.The appendix constructs distributions close in truncated W∞ distance but with infinite Wq distance.

B.1 Proof of Theorem 3.3

The proof establishes posterior-sampling robustness under distribution mismatch by decomposing distributions, transferring bounds through couplings, and controlling Gaussian-measurement separation.

  • Theorem 3.3: Theorem 3.3 assumes (δ, α)-W∞(µ, ν) ≤ ε and Gaussian measurements with noise scale σ ≳ ε.Posterior sampling from ν is analyzed when the ground truth is drawn from µ.
  • Distribution Decomposition: The proof decomposes µ and ν into high-mass matched components and δ- or α-weighted residual components.This isolates the portion controlled by the Wasserstein coupling from the mismatch mass.
  • Error Bound: The resulting error probability is bounded by the mismatch mass plus a posterior-sampling error term.One intermediate bound places δ outside the conditional probability term.
  • Coupling Argument: A W∞ coupling confines paired samples within an σ-radius ball, enabling separation of near and far mixture components.The triangle inequality forces sufficiently distant posterior samples into the far component.
  • Theorem 3.4: Theorem 3.4 extends the posterior-sampling argument to arbitrary forward operators when an algorithm can produce an appropriately distributed conditional sample.

C.2 More Exemplar Reconstructions

Qualitative reconstructions show strong performance across in-distribution and shifted MRI settings, including higher acceleration, altered masks and contrasts, and anatomy changes.

  • Contrast Shifts: The method produced excellent qualitative reconstructions for T1- and FLAIR-weighted brains despite training on T2-weighted brains.
  • In-Distribution Reconstructions: At R = 3 with an equispaced vertical mask, the method was competitive with E2E-VarNet on in-distribution brain scans.
  • Higher Acceleration: At R = 6, the method retained performance on in-distribution brain reconstructions with an equispaced vertical k-space mask.
  • Higher Acceleration: At R = 12, the method showed significantly fewer artifacts than baselines under an equispaced vertical mask.
  • Mask Shifts: The method was robust to a shift from an equispaced vertical training mask to an equispaced horizontal test mask.Untrained ConvDecoder and L1-Wavelets were also robust to this mask shift.
  • Anatomy Shifts: On fastMRI knees, the method was more robust than other methods to anatomy shift despite all methods being trained on fastMRI brains.

D Appendix: fastMRI Knee

The appendix evaluates reconstructions across knee, abdominal, and fat-suppressed knee scans, including anatomy, pathology, field-of-view, and signal-to-noise shifts. The method preserves fine details qualitatively, although its numerical performance is not optimal on fat-suppressed knees.

  • Fat-suppressed knee reconstructions are qualitatively competitive and recover fine details better than deep-learning baselines, despite non-optimal numerical performance.The likely explanation given is the much lower signal-to-noise ratio of fat-suppressed knees than the brain training data.
  • R = 8 knee reconstructions show greater robustness to anatomy shift than other methods when all models were trained on fastMRI brains.
  • 48 posterior samples at R = 4 provide pixel-wise means, standard deviations, and reconstruction-error estimates for a knee with an annotated meniscus tear.The samples support voxel-level uncertainty visualization while retaining the tear in the reconstruction examples.
  • At R = 4, the method preserves meniscus-tear details better than baselines despite none of the methods seeing this pathology during training.
  • At R = 4, the brain-trained score-based model obtains good abdominal reconstructions, although kidney details can be missing or artifact-affected.
  • On Stanford knees, no statistical difference is found between L1-Wavelet, MoDL, and the proposed method under Poisson masks.The scans are 3D and differ considerably from the other datasets; training and hyperparameter selection used fastMRI brains.

F Appendix: Stanford Knee

The Stanford Knee appendix tests anatomy and dimensionality shift using axial slices extracted from a 3D scan. Under Poisson-mask sampling, the proposed method is statistically indistinguishable from L1-Wavelet and MoDL.

  • The evaluation converts a publicly available 3D-FSE knee scan into 24 equally spaced 2D axial slices of 320 × 256 pixels.
  • At Stanford Knees, Poisson masks improve incoherence, and no statistical difference appears between L1-Wavelet, MoDL, and the proposed method.All training and hyperparameter selection were performed on fastMRI brain scans.

G Appendix: Implementation

The appendix describes score-model, baseline, normalization, and evaluation procedures used across MRI anatomies and sampling settings. The convolutional score model accepts differing image shapes, while a baseline depends explicitly on vertical ACS data.

  • The score-based generator uses two channels for the real and imaginary components of complex-valued MRI, without other architectural changes.
  • Training noise scales are set to sigma_begin = 232 and sigma_end = 0.0066 because MRI uses fewer channels than FFHQ.
  • Images are normalized by their 99th-percentile pixel value to handle dynamic ranges spanning approximately 10^-4 to 10^5.At inference, the reconstruction is normalized using its 99th-percentile pixel value.
  • The convolutional NCSNv2 model trained on 384 × 384 images can be applied to knees, T1-weighted and FLAIR brains, and abdomens with different shapes.
  • Hyperparameters tuned on two validation brain scans at R = 4 were reused across anatomies and acceleration factors.
  • The fastMRI E2E-VarNet baseline has a 29-million-parameter image network and a 0.7-million-parameter sensitivity-map network, trained over 12 unrolls.
  • The baseline training uses fully sampled MVUE brain T2 reconstructions and supervised SSIM loss.
  • E2E-VarNet explicitly uses the fully sampled vertical ACS region during training and inference, making other mask patterns non-trivial to test.

H Appendix: Radiologist Study

A preliminary blinded study asked three external experts to rank reconstructions across shifted anatomy and scan settings. Experts consistently preferred the proposed method for knees, while abdominal and brain rankings showed no significant differences, and the study remains non-clinical.

  • Three external experts blindly assessed 30 questions covering abdominal, knee, and horizontally readout brain scans, all involving at least one test-time distribution shift.
  • Each question showed a fully sampled reference and three shuffled reconstructions from MoDL, ConvDecoder, and the proposed method at R = 3.
  • Experts ranked the three reconstructions from best to worst using the reference image as a perceptual guideline.
  • Participants consistently ranked the proposed method best on knee scans, while abdominal and horizontally encoded brain scans showed no significant perceived differences.
  • The knee rankings showed a significant difference and very strong inter-observer consensus, whereas abdomen and brain rankings showed weaker consensus and equivalent images.
  • The study adds evidence for robustness to test-time distribution shifts but did not evaluate every algorithm, shift, or acceleration level and is not a substitute for clinical evaluation.
Loading 2108.01368v2…