Source-linked AI summary
Score-based diffusion models for accelerated MRI
Hyungjin Chung, Jong Chul Ye
TL;DR
Accelerated MRI reconstruction needs priors that model complex image distributions without relying on large curated k-space databases. The paper learns a continuous score function from magnitude images and alternates reverse-SDE sampling with data consistency to reconstruct measurements. It reports high-quality single- and multi-coil reconstructions, supervised-model comparisons, sampling-pattern flexibility, and uncertainty estimation, while identifying inference time and extreme or out-of-distribution cases as limitations.
Problem
Existing MRI reconstruction methods either use nontrivial hand-crafted priors or depend heavily on large, well-curated k-space databases.
Method
The method trains a continuous score function with denoising score matching and samples the conditional distribution by alternating reverse-SDE updates with data-consistency projections.
Results
The method produces high-accuracy single- and multi-coil reconstructions, outperforms supervised reconstruction models for complex images, and supports uncertainty quantification and varied sampling patterns.
Takeaways & Limitations
Magnitude-only score learning can serve as a reconstruction prior across complex-valued and parallel-imaging settings while producing multiple reconstructions for uncertainty assessment.
Takeaways & Limitations
Inference is slow at full discretization, and extreme acceleration or out-of-distribution one-dimensional undersampling can produce artifacts, unsatisfactory samples, or high posterior-detail variance.
Abstract
from arXiv · showhide
Score-based diffusion models provide a powerful way to model images using the gradient of the data distribution. Leveraging the learned score function as a prior, here we introduce a way to sample data from a conditional distribution given the measurements, such that the model can be readily used for solving inverse problems in imaging, especially for accelerated MRI. In short, we train a continuous time-dependent score function with denoising score matching. Then, at the inference stage, we iterate between numerical SDE solver and data consistency projection step to achieve reconstruction. Our model requires magnitude images only for training, and yet is able to reconstruct complex-valued data, and even extends to parallel imaging. The proposed method is agnostic to sub-sampling patterns, and can be used with any sampling schemes. Also, due to its generative nature, our approach can quantify uncertainty, which is not possible with standard regression settings. On top of all the advantages, our method also has very strong performance, even beating the models trained with full supervision. With extensive experiments, we verify the superiority of our method in terms of quality and practicality.
1. Introduction
The paper proposes a score-based diffusion framework for MRI inverse problems that combines learned image priors with measurement consistency. It uses magnitude-only training while supporting conditional reconstruction, complex-valued and parallel imaging, uncertainty quantification, and varied sampling patterns.
- Sub-sampled MRI reconstruction remains difficult because hand-crafted priors poorly model MRI data distributions, while deep learning methods require large curated k-space databases.
- The framework trains a continuous score function with denoising score matching and samples p(x|y) by alternating reverse-SDE updates with data-consistency projections.The method uses magnitude images only for training.
- The reconstruction method achieves state-of-the-art performance and outperforms supervised reconstruction models for complex images.The score function was not trained specifically for image reconstruction.
- The approach is agnostic to sub-sampling patterns and extends to complex-valued acquisitions and practical multi-coil reconstruction using the same score function.
- Its stochastic sampling procedure generates multiple reconstructions from the same measurement and quantifies uncertainty without modifying the neural network.Posterior samples vary little at low acceleration and show increasing variance as acceleration becomes more aggressive.
2. Background
Score-based diffusion models define a continuous stochastic process that transforms data toward a tractable prior and use a learned score to reverse that process. The framework supports multiple SDE choices, with denoising score matching and numerical predictor-corrector sampling providing the operational machinery.
- A continuous diffusion process connects samples from the data distribution at t=0 to tractable prior samples at t=T through an SDE.The drift and diffusion coefficients determine the stochastic process.
- Different choices of drift and diffusion functions produce distinct SDEs, including variance-preserving and variance-exploding formulations.
- The variance-exploding SDE is used here because it empirically provides higher sample quality, while variance-preserving and DDPM formulations remain compatible with the framework.
- The reverse-time SDE uses the gradient of the log data density, estimated by a time-conditional neural network score function.
- Denoising score matching trains the network by replacing the unknown data score with the score of a Gaussian perturbation kernel.For Gaussian kernels, the perturbation score is explicitly related to the noisy and clean samples through σ(t)^2.
- Numerical reverse-SDE sampling uses discretized predictor steps and optional Langevin corrector steps, yielding the predictor-corrector algorithm.The process starts from noise and iteratively updates the sample across discretized time intervals.
3. Main Contributions
The framework uses a learned score prior with data-consistency updates to sample conditional MRI reconstructions, extending from real-valued to complex-valued and parallel imaging.
- Measurement model: The accelerated MRI measurement model represents measurements y as a forward operator A applied to latent image x.For parallel imaging, A incorporates coil sensitivity maps; for single-coil acquisition, A reduces to PΩF.
- Conditional sampling: Score-based reconstruction alternates predictor-corrector reverse-SDE updates with data consistency through A∗(y − Axi).The inference algorithm applies data consistency after predictor and corrector steps.
- Data consistency: The framework relies on sensitivity normalization so the data-consistency mapping (I − λA∗A) is non-expansive for λ ∈ [0, 1].This property is stated in Proposition 1 and underpins the controlled data-consistency update.
- Complex-valued reconstruction: The method uses a magnitude-trained score function to reconstruct complex-valued images by applying predictor-corrector updates separately to real and imaginary parts.This avoids training directly on complex-valued data and can use magnitude-only DICOM images.
- Parallel imaging: The hybrid parallel-imaging approach injects coil-image dependency every m iterations and uses a linearly decreasing data-consistency weight λ.The authors report sharper reconstructions and improved performance with these intermediate updates.
- Parallel imaging: Parallel imaging is supported through SSOS-type reconstruction, which processes coil images separately before merging them with an SSOS operation.The score function is applied independently to each coil image rather than estimating GRAPPA kernels.
- Limitations: A limitation is that SSOS-type and hybrid-type parallel-imaging methods are slower because naive implementations require c times more computation than single-coil reconstruction.Parallelizing coil computations can reduce the time toward that of single-coil reconstruction when sufficient GPU resources are available.
4. Methods
The experiments train and evaluate the score model on fastMRI knee data, compare against classical and supervised baselines, and assess reconstruction quality and downstream pathology detection.
- Data: The score function is trained on approximately 25k 320 × 320 magnitude-image slices from the fastMRI knee training data.The first and last five slices of each volume were removed to avoid noise-only training examples.
- Score-model training: Training uses a VE-SDE denoising score-matching objective with λ(t) = σ^2(t), which corresponds theoretically to likelihood weighting.The noise schedule uses σmin = 0.01 and σmax = 378 in a geometric series.
- Implementation: The network is optimized with Adam, batch size 1, gradient clipping at 1.0, exponential moving averaging, and 100 epochs of training.A linear warm-up reaches 2e-4 at the 5000th optimization step, and training takes about three weeks on a single system.
- Inference: Inference uses N = 2000 iterations and M = 1 corrector step by default, requiring about 10 minutes for real-valued and 20 minutes for complex-valued reconstructions.A single score-network forward pass takes about 150 ms with a commodity GPU.
- Baselines: Comparisons include total-variation reconstruction, supervised U-Net reconstruction, DuDoRNet, and E2E-varnet.The study reports the best total-variation result after grid-searching its λ hyper-parameter.
- Evaluation: Reconstruction quality is evaluated with PSNR and SSIM, while repeated-measures ANOVA tests statistical significance.The experiments also use fastMRI+ pathology bounding boxes and a YOLOv5 detector to compare downstream pathology detection.
- Generalization: The authors investigate robustness to distribution shifts because training and inference data distributions differ in the reconstruction algorithms.This analysis examines how far the discrepancy can be increased while maintaining satisfactory results.
5. Results
Across real-valued, complex-valued, and multi-coil MRI experiments, the proposed score-based method reconstructed images with high fidelity across sampling patterns and acceleration factors. It also preserved diagnostic capability in pathology detection while exposing a multi-coil averaging limitation.
- Real-valued simulation study: The proposed method produced nearly perfect real-valued reconstructions under 2D ×8 Gaussian sampling, while TV, U-Net, and DuDoRNet retained artifacts or blur.The proposed reconstruction showed virtually no structural difference from ground truth in the zoomed region.
- Real-valued simulation study: At ×4 and ×8 acceleration, the proposed method captured high-frequency details more accurately than comparison methods across uniform and Gaussian sampling.Supervised methods smoothed details, while TV retained aliasing and cartoon-like artifacts.
- Real-valued simulation study: Under ×15 variable-density Poisson disk sampling, only the proposed method produced high-fidelity reconstructions; other methods showed severe artifacts and blurred details.The proposed method preserved anatomical structure despite the most aggressive undersampling condition.
- Quantitative evaluation: Across sampling patterns, quantitative metrics significantly favored the proposed method, except that it was on par with DuDoRNet for Gaussian 1D sampling.RM-ANOVA indicated statistically significant superiority against other methods in most cases.
- Method behavior under acceleration: Unlike supervised reconstructions that became blurrier at higher acceleration, the proposed stochastic method preserved high-frequency details and sampled high-probability modes.The method thereby supports uncertainty quantification through multiple reconstructions from the same measurement.
- Complex-valued single-coil reconstruction: For complex-valued single-coil data, the proposed method remained accurate despite training only on magnitude images and outperformed TV, U-Net, and DuDoRNet visually and quantitatively.It preserved sharp edges and texture across the tested sampling patterns, including ×15 variable-density Poisson sampling.
- Complex-valued multi-coil reconstruction: In multi-coil parallel imaging, the proposed method was the only method producing high-fidelity reconstructions regardless of sampling pattern and outperformed all methods for 2D patterns.It was on par with E2E-varnet on 1D sampling patterns, while E2E-varnet showed severe artifacts for 2D sampling.
- Pathology detection: Pathology detection from proposed reconstructions performed on par with, or sometimes better than, fully sampled images, whereas comparison reconstructions limited diagnostic capability.The fully sampled reference reached mAP 0.754; the authors caution that the small test set and imperfect detector prevent interpreting higher values as superiority.
6. Discussions
The discussion examines inference acceleration, broader applicability, related approaches, and limitations of score-based MRI reconstruction. It highlights shorter diffusion trajectories, flexible imaging use, and failure modes under distribution shift or extreme acceleration.
- Speeding up inference: CCDF shortens the reverse-diffusion horizon and can achieve similar or better performance with as few as 40 iterations.CCDF forward-diffuses an initialization and begins reverse diffusion from a shorter timestep, reducing reconstruction time.
- Related work: Unlike an earlier MRI score-matching approach, this method reportedly outperforms supervised networks while requiring only magnitude images for training.The earlier approach used separate channel complex-valued data and reportedly lagged supervised methods, particularly for single samples.
- Broader applicability: The framework is positioned for other computational-imaging problems because a single trained score function can support modified inference procedures across modalities and tasks.The paper cites applications including super-resolution and image reconstruction.
- Limitations: For out-of-distribution 1D undersampling, local aliasing-like edge artifacts can occur, whereas the paper reports no such artifact for 2D sampling.The authors caution that robustness is compromised when extending the framework to OOD data.
7. Conclusion
The paper proposes score-based reconstruction for accelerated MRI using magnitude-only training and conditional sampling through data consistency. It reports accurate single- and multi-coil reconstructions, while identifying reconstruction speed as an unresolved gap with feedforward methods.
- The method trains the gradient of the log data distribution with continuous-time denoising score matching using magnitude data.
- Data consistency projection at every step enables sampling from the conditional distribution given measurements.
- The approach produces high-accuracy reconstructions for both single-coil and multi-coil MRI.
- The method’s reconstruction speed remains behind feedforward neural-network approaches.
Details of network architecture
The NCSN++ architecture constructs the score function from image state and continuous-time inputs. Time conditioning is embedded and incorporated into image features across spatial dimensions through residual blocks.
- The NCSN++ network takes the noisy image x_t and continuous time t as inputs to construct the score function.Time conditioning is important because the model solves a continuous SDE.
- Time t is processed through a time-conditioning module before being integrated with image features.
- Random Fourier features use non-trainable weights initialized with torch.randn(d) * scale, with scale set to 16.
- An MLP matches the embedding to the image-feature channel dimension, and residual blocks broadcast temporal features across H and W.
Statistical analysis of pathological detection
The Bland–Altman analysis evaluates agreement between reconstruction methods and fully sampled data for pathological detection. The proposed method shows the least variance in the differences.
- The proposed method agrees most closely with fully sampled data, showing the least variance in Bland–Altman differences.DuDoRNet and UNet have higher variance, while TV differs radically.
- The analysis uses each class’s mAP50 score as the individual data-point measure.
Additional results
The paper provides additional reconstruction results across different anatomies in two supplementary figures.
- Additional reconstruction results for different anatomies are presented in Fig. 14 and Fig. 15.