Robust Compressed Sensing MRI with Deep Generative Priors

In plain words

An MRI machine measures only part of the data it could collect, because a shorter scan is kinder to the patient. The image then has to be rebuilt from partial measurements. A common approach trains a network on many scans with one fixed measurement pattern. Such networks can do poorly when the scan looks different from their training data.

This paper trains a generative model on brain scans only and uses it as a prior, meaning a learned sense of what a plausible brain image looks like. To rebuild an image, the method starts from noise and repeatedly nudges it in two directions: toward images that the model considers plausible, and toward images that agree with the measurements actually taken. The size of these nudges is gradually reduced, so the process settles on a final image. The model never sees the measurement pattern during training, so the same model can be used with other patterns and with other body parts.

The authors report that the method is competitive with supervised networks when the test scan resembles the training data, and that it holds up better when it does not. The price is speed: one reconstruction takes minutes rather than the seconds that supervised networks need. Because each run starts from different noise, repeated runs give a spread of images, and that spread can show which pixels the method is unsure about.

TL;DR

• Posterior sampling with a score-based prior for brain MRI (NCSNv2 trained on T2 fastMRI brain slices), run as annealed Langevin dynamics on the multi-coil k-space forward model §2.1 §2.2 §4 eq. 4.
• The prior is trained without the sampling pattern, so the same weights serve other patterns and anatomies §4.2 §6.
• Author-reported: competitive with supervised baselines in-distribution (better PSNR, lower SSIM at R = 3 and 6), and more robust under sampling and anatomy shifts §4.1 §4.2.
• Cost: 16 minutes and 0.95 GB for one high-resolution brain on an RTX 2080Ti §4.1.
• The released code departs from the equations in several places. See Paper vs code before implementing.

Problem

Compressed sensing reduces the number of k-space measurements needed for MRI, which shortens scans §1. Sparsity-based methods are limited in the acceleration they reach §1. Supervised deep networks reconstruct well when test data match training data, but their quality drops under out-of-distribution shifts in the sampling pattern or the anatomy §1 §4.2.

Compressed Sensing with Generative Models (CSGM) instead trains a generative model on the images and reconstructs by searching the model's range. It does not depend on the measurement pattern at training time §1. Before this paper, CSGM had worked mainly on faces and digits, and it is known to perform poorly on out-of-distribution samples §1.

The paper's theory covers the case where the generator's distribution differs from the true one. Its guarantees use Gaussian measurement matrices, and it notes that they do not carry over to Fourier measurements in general §3.

Method

Forward model. A coil with sensitivity map S_i measures the image x through the Fourier transform F and a sampling operator P, plus noise eq. 1 §2.1. The same sampling pattern applies to every coil. The reference image is the minimum-variance unbiased estimate (MVUE) from the fully sampled data, built with ESPIRiT sensitivity maps §2.1 §4.

Posterior sampling. For a measurement y = A x + w with Gaussian noise and a prior mu, the output is a draw from the posterior mu(x | y) §2.2. Langevin dynamics produces such draws from the score of the posterior eq. 3 §2.2. Annealed Langevin dynamics uses a family of noise-perturbed priors, each with its own score, going from large to small noise §2.2.

Score model. The prior is a score network trained with annealed denoising on T2-weighted brain slices from the NYU fastMRI dataset, 3 T, 14,539 two-dimensional training slices §4. The network is NCSNv2 §4 app. G.1. Its input and output have two channels, the real and imaginary parts, and it is trained on 384 by 384 images app. G.1.

Algorithm. The paper's final update is an annealed Langevin step with a data term eq. 4 §2.2. The gradient of the measurement likelihood is A^H(y - A x) divided by sigma^2 + gamma_t^2, where gamma_t decreases over time. The paper states this but gives no values for sigma, gamma_t, the step sizes eta_t or the number of noise levels T §2.2.

x_{t+1} = x_t + \eta_t \left( f(x_t; \beta_t) + \frac{A^H (y - A x_t)}{\gamma_t^2 + \sigma^2} \right) + \sqrt{2\eta_t}\, \zeta_t, \qquad \zeta_t \sim \mathcal{N}(0, I) \qquad [\text{eq. 4}]

Here f is the score of the prior at noise level beta_t, and the initial state is drawn from a standard complex Gaussian §2.2. The paper states that T and the beta_t were fixed during training of the score model, so inference only sets eta_t, sigma and gamma_t §2.2.

# Posterior sampling by annealed Langevin  [eq. 4] [§2.2]
# Input:  y (masked k-space), A (coil maps and mask), score f
# Output: one posterior sample x, complex (H, W)

def posterior_sample(y, A, score, eta, gamma, sigma):
    # Initial state from a standard complex Gaussian  [§2.2]
    x = complex_normal(0, 1, shape=(H, W))
    for t in range(T):
        # Prior term at the current noise level  [§2.2] [eq. 4]
        prior_grad = score(x, t)
        # Data term with its annealed scale  [§2.2] [eq. 4]
        residual = y - forward(A, x)
        data_grad = adjoint(A, residual) / (gamma[t] ** 2 + sigma ** 2)
        # Noise with variance 2 * eta_t  [§2.2] [eq. 4]
        noise = normal(0, 1, shape=(H, W)) * sqrt(2 * eta[t])
        x = x + eta[t] * (prior_grad + data_grad) + noise
    return x

Normalisation. During training, each image is divided by its 99th-percentile pixel value. At inference, the reconstruction is scaled using the 99th percentile of the pseudo-inverse image, since the ground truth is unavailable app. G.1.

Uncertainty. Several independent posterior samples can be averaged, and their pixel-wise standard deviation estimates the error at each pixel §4.3. For a knee at R = 4, the paper uses 48 reconstructions fig. 4.

Hyperparameters

The paper's appendix refers readers to the GitHub repository for the sampler values app. G.1. The values below come from the repository's brain_T2 configuration, unless marked as stated in the paper. The other anatomy configurations differ in some values, listed under Paper vs code.

Results

Author-reported. Setup: retrospective under-sampling of fastMRI brain, knee and abdomen scans, and Stanford knees, with baselines trained on MVUE targets §4 app. A.1. The score model is trained on T2 brains only §4. Metrics are PSNR and SSIM on magnitude images §4. The paper gives the key comparisons in figures. Their numbers are not stated in the text, so only the figure references are given here.

Figure-only results. PSNR across acceleration factors and anatomy shifts is in fig. 2. SSIM is in fig. 5, masked SSIM in fig. 6, masked PSNR in fig. 7, and fat-suppressed knees in fig. 20. Stanford knees under Poisson masks are in fig. 24. The authors report no statistical difference between l1-Wavelet, MoDL and their method on that dataset fig. 24.

Qualitative claims §4.1 §4.2:

• In-distribution (vertical equispaced masks, R = 3 and 6): baselines have higher SSIM, while this method has higher PSNR. At higher accelerations it surpasses the baselines fig. 1.
• Sampling shift (horizontal masks, R = 3): this method keeps its performance. E2E-VarNet shows aliasing artefacts there fig. 1.
• Anatomy shift: this method, MoDL and l1-Wavelet keep their advantage on abdomen scans, while ConvDecoder and E2E-VarNet degrade severely §4.2. On knees it gives the best results at R = 2 and 4, with lower variance §4.2.

The radiologist study used three experts, with 30 scores per anatomy app. H. The authors call it preliminary, and they state that it is not a clinical evaluation app. H.

Theory §3:

• Theorem 3.3 (Gaussian measurements). If the true and generator distributions are close in the (delta, alpha)-W_infinity divergence with parameter epsilon, and the number of measurements M grows like log(1/(1 - alpha)) plus a covering-number term, then posterior sampling recovers x within c(epsilon + sigma) with probability at least 1 - delta - exp(-Omega(M)). The constant c is universal but not given §3.
• Theorem 3.4 (arbitrary measurements). If some algorithm is within epsilon with probability 1 - delta, posterior sampling is within 2 epsilon with probability at least 1 - 2 delta §3 app. B.
• Lemma (appendix B): W_q-closeness implies (delta, delta)-W_infinity closeness with epsilon divided by delta^(1/q) app. B.

Paper vs code

Code: official repository https://github.com/utcsilab/csgm-mri-langevin at commit 91e0a5ac658feaf6460045706ef16dd140894374, MIT licence LICENSE. The paper's abstract links it. The owner is the Computational Sensing and Imaging Lab at UT Austin, which is the lab of author Jonathan I. Tamir. Code citations below refer to that commit. Score-model training code and the NCSNv2 sources are not in the snapshot, because the ncsnv2 submodule is empty (.gitmodules). The baseline code is not in the repository either.

• Data-consistency term. The paper's update divides the measurement gradient by sigma^2 + gamma_t^2 eq. 4 §2.2. The code does not use sigma or gamma. It normalises the measurement gradient to unit norm, rescales it to the norm of the prior score, and multiplies by 5.0 main.py brain_T2.yaml. The paper states no value or rule for sigma or gamma_t. Treat the code's rule as what produced the published results, and note that it is not the stated equation.
• Initial state. The paper draws the start from a standard complex Gaussian §2.2. The code draws each of the two channels uniformly from [0, 1) main.py.
• Starting noise level. The paper runs the whole annealing schedule. The code skips every level up to start_iter = 1155 main.py brain_T2.yaml. The config's comment calls this a different starting noise level.
• Inner steps. The paper does not state a step count per level. The code uses 3 steps for levels up to 1800 and the configured count above that, with 1800 hard-coded main.py.
• Step size. The paper calls the step size eta_t and gives no value eq. 4. The code uses step_lr times (sigma_t / sigma_end)^2 main.py.
• Parameter averaging. The code loads an exponential moving average of the score network's weights, with rate 0.999 main.py. The paper does not mention it.
• Normalisation. The paper uses the 99th percentile of the pseudo-inverse image app. G.1. The code takes the 99th percentile of a coil-combined image from the masked k-space main.py utils.py. The coil combination is a SENSE-style weighted sum, which is not the same as a pseudo-inverse.
• Sampling pattern. The paper says R is the true acceleration N / L §2.1. For R = 4 at 384 lines, the code keeps 30 centre lines and every fourth outer line, which gives 119 of 384 lines. That is about 3.2 times, by our count from the code. The paper's footnote says the code's definition differs from fastMRI's §2.1.
• Sensitivity maps. The paper computes ESPIRiT from the fully sampled central portion of k-space §4. The code runs ecalib on the full k-space of each slice estimate_maps.py. It does not set the calibration-region size.
• Only one batch is processed. The main loop breaks after the first batch of the loader main.py. The released script therefore reconstructs one slice per run, which is marked as a leftover in the code.
• Other configurations differ. The paper says one set of hyperparameters was tuned on two brain scans and reused for all anatomies and accelerations app. G.1. The configurations differ: the inner step count is 3 for brain_T1, brain_FLAIR, knees and stanford_knees, and 4 for brain_T2 and abdomen brain_T1.yaml abdomen.yaml. The brain_T1 and brain_FLAIR configurations do not set start_iter, which main.py reads main.py. We did not run the code, so we did not confirm whether this fails.
• Hard-coded CUDA. The gradient is cast to torch.cuda.FloatTensor main.py, so the sampler needs a GPU.
• Matches. The image size of 384 by 384 for brains, the geometric noise schedule from 232 to 0.0066, the two-channel representation and the NCSNv2 backbone match app. G.1 brain_T2.yaml brain_T2.yaml.

Pitfalls

• The equation is not the algorithm. Eq. 4 has sigma and gamma_t, which the paper never gives values for. The released sampler uses a different data term and extra rules (see Paper vs code). An implementer who follows only the paper will not get the published numbers, and one who follows only the code will get rules the paper does not describe.
• Epoch counts conflict in the paper. The main text says MoDL and E2E-VarNet were trained for 40 and 15 epochs, respectively §4. The appendix says E2E-VarNet for 40 epochs app. G.2 and MoDL for 15 app. G.3. The two readings disagree on which baseline gets which count.
• Theorems do not cover the experiments. Theorem 3.3 assumes Gaussian measurements with specific normalisations and a noise level sigma at least of order epsilon §3. The experiments use Fourier masks, and the paper says the result does not carry over to Fourier measurements in general §3.
• Effective acceleration. Because the equispaced rule keeps every R-th outer line plus the centre, the sampled fraction is above 1 / R dataloaders.py dataloaders.py. Compare results with the paper's nominal R and the code's sampled fraction.
• Metrics are not comparable with RSS references. The paper's references are MVUE images, which give different numbers from RSS-trained leaderboard results. It reports an SSIM of slightly above 0.8 between MVUE and RSS on T2 brains, and quotes fastMRI leaderboard figures for RSS-based methods app. A.1. Do not compare the numbers directly.
• Output varies from run to run. The reconstruction is a random draw from the posterior. Results and uncertainty maps depend on the random seed and on the number of samples taken §4.3.
• Cost scales with the schedule. By our count from the brain_T2 configuration, reconstructing one image runs about 3,975 score evaluations: 645 levels at 3 steps and 510 levels at 4 steps. This is derived from the code, not stated in the paper. The paper reports 16 minutes on an RTX 2080Ti §4.1.
• Dynamic range. Training normalises each image by its 99th-percentile pixel value, which the paper says is needed because MRI values span orders of magnitude app. G.1. Skipping the scale breaks the step size.
• Scope. All data are retrospectively under-sampled, and the study covers fast spin-echo imaging only §5. The authors list fairness across protected attributes as an open issue §5.
• Training code is missing. The score model is trained with NCSNv2 changes described only in the appendix app. G.1. The repository does not include the training loop, so the training normalisation and noise-level function cannot be checked from it.

Lineage

• Builds on: CSGM, the generative-prior framework for compressed sensing of Bora et al. §1; posterior sampling for compressed sensing §2.2; NCSNv2 for the score model §4 app. G.1 (the repository's README gives arXiv:2006.09011 for the earlier score-based work, README.md).
• Compared against: l1-Wavelet with BART, MoDL, E2E-VarNet, and ConvDecoder §4.
• Superseded by: not stated in this wiki.