Source-linked AI summary

Solving Inverse Problems in Medical Imaging with Score-Based Generative Models

Yang Song, Liyue Shen, Lei Xing, Stefano Ermon

arXiv:2111.08005v2eess.IVcs.CVcs.LGstat.ML

TL;DR

Partial-measurement CT and MRI reconstruction is challenging because supervised methods require paired data tied to a fixed measurement process. The paper learns a medical-image prior with score-based generative models and samples reconstructions conditioned on measurements, achieving comparable or better performance while generalizing more effectively to changed processes.

  • Problem

    Supervised reconstruction methods require paired images and measurements synthesized for a fixed physical process, limiting adaptation when CT or MRI measurements change.

  • Method

    The method trains a score-based generative model on medical images and uses the physical measurement model during sampling to reconstruct images consistent with the prior and observations.

  • Results

    The method matches or outperforms supervised counterparts across sparse-view CT and undersampled MRI and generalizes to changed measurement processes, including using one CT model for reconstruction and metal artifact removal.

  • Takeaways & Limitations

    A single unsupervised score model can address multiple linear inverse problems within the same image domain without retraining paired data for each measurement process.

  • Takeaways & Limitations

    The method assumes the measurement mapping is linear and, in its stated formulation, assumes the operator A has full rank.

Abstract

from arXiv · show

Reconstructing medical images from partial measurements is an important inverse problem in Computed Tomography (CT) and Magnetic Resonance Imaging (MRI). Existing solutions based on machine learning typically train a model to directly map measurements to medical images, leveraging a training dataset of paired images and measurements. These measurements are typically synthesized from images using a fixed physical model of the measurement process, which hinders the generalization capability of models to unknown measurement processes. To address this issue, we propose a fully unsupervised technique for inverse problem solving, leveraging the recently introduced score-based generative models. Specifically, we first train a score-based generative model on medical images to capture their prior distribution. Given measurements and a physical model of the measurement process at test time, we introduce a sampling method to reconstruct an image consistent with both the prior and the observed measurements. Our method does not assume a fixed measurement process during training, and can thus be flexibly adapted to different measurement processes at test time. Empirically, we observe comparable or better performance to supervised learning techniques in several medical imaging tasks in CT and MRI, while demonstrating significantly better generalization to unknown measurement processes.

1 INTRODUCTION

CT and MRI reconstruction from partial measurements is difficult because supervised methods depend on paired data synthesized for a fixed measurement process. The paper instead learns an image prior with score-based generative models and samples reconstructions consistent with measurements, achieving competitive performance and stronger adaptation to changed processes.

  • Partial CT and MRI measurements reduce acquisition burden but make image reconstruction an important inverse problem.Full CT sinograms expose patients to excessive ionizing radiation, while full MRI k-space measurement is time-consuming.
  • Supervised reconstruction methods directly map measurements to images using paired datasets synthesized under a fixed physical measurement model.Changing CT projection counts or MRI downsampling ratios requires recollecting paired data.
  • The proposed unsupervised approach learns a medical-image prior with a score-based generative model and samples images consistent with observed measurements and the physical process.After training, the model can solve linear inverse problems within the same image domain.
  • Across CT and MRI tasks, the method achieves comparable or better performance than supervised counterparts and uniformly surpasses baselines when measurement counts change.It also uses one model for sparse-view CT reconstruction and metal artifact removal, indicating adaptation across measurement processes.

2 BACKGROUND

The background formulates medical reconstruction as posterior sampling for linear inverse problems and introduces score-based generation through forward diffusion and reverse-time denoising. Neural score estimates make reverse sampling practical, while related discretization and sampling methods support generation.

  • 2.1 LINEAR INVERSE PROBLEMS: A linear inverse problem recovers an unknown signal x from noisy measurements y = Ax + ϵ, requiring a prior when measurements are fewer than signal dimensions.The posterior p(x | y) combines the measurement distribution with the prior p(x).
  • 2.1 LINEAR INVERSE PROBLEMS: CT measurements are X-ray sinograms, whereas MRI measurements are Fourier-space spatial frequencies known as k-space.Both are examples of linear inverse problems in which the unknown signal is a medical image.
  • 2.1 LINEAR INVERSE PROBLEMS: Generative models estimate the unknown image prior from unpaired medical images, after which Bayes’ rule determines the posterior given the measurement model.The posterior is the target distribution for inverse-problem reconstruction.
  • 2.2 SCORE-BASED GENERATIVE MODELS: Score-based models progressively perturb images to noise with an SDE and reverse this process by estimating the time-dependent score function.The forward process uses a linear SDE whose coefficients can make the terminal distribution close to a predefined noise distribution.
  • 2.2 SCORE-BASED GENERATIVE MODELS: A neural score model trained with denoising score matching approximates ∇x log pt(x), enabling reverse-time SDE simulation from noise to data samples.The reverse process starts from an approximate terminal noise sample and removes noise while stepping backward in time.
  • 2.2 SCORE-BASED GENERATIVE MODELS: Reverse-time sampling can use Euler–Maruyama discretization, annealed Langevin dynamics, probability-flow ODE solvers, or Predictor–Corrector samplers.These alternatives provide multiple numerical routes for generating samples from the learned distribution.
  • 2.3 INVERSE PROBLEM SOLVING: The inverse-problem sampler modifies the score-based generation process to incorporate observations while sampling from the posterior.The corresponding procedure is presented as an inverse-problem sampling algorithm.

3 SOLVING INVERSE PROBLEMS WITH SCORE-BASED GENERATIVE MODELS

The method solves inverse problems by conditioning an unconditional score-based sampler on noisy measurements, using a measurement-operator decomposition and proximal data-consistency steps. It avoids paired training data and supports multiple iterative samplers without expensive SVD computation.

  • The method approximates posterior sampling by conditioning an unconditional score model on an observation through a conditional reverse-time SDE.A directly conditioned score model would require paired data, so the paper instead uses an unconditional score model and the measurement distribution.
  • It trains one unconditional score model on medical images, then incorporates test-time measurements by adding noise to the observation and applying proximal consistency steps during reverse sampling.The resulting intermediate samples are made consistent with the noisy measurement process while existing unconditional samplers reduce image noise.
  • 3.1 A CONVENIENT FORM OF THE LINEAR MEASUREMENT PROCESS: The measurement operator is decomposed as A = P(Λ)T, where T captures a shared transform and P(Λ) selects m coordinates through a subsampling mask.For sparse-view CT and undersampled MRI, T corresponds respectively to the Radon and Fourier transforms, while Λ represents measurement selection.
  • 3.2 INCORPORATING A GIVEN OBSERVATION INTO AN UNCONDITIONAL SAMPLING PROCESS: Noisy measurement-process samples are generated as ŷ_t = α(t)y + β(t)Az, enabling tractable conditioning while bypassing direct sampling from the conditional stochastic process.The construction draws z from a standard normal distribution and uses the known measurement operator A.
  • 3.2 INCORPORATING A GIVEN OBSERVATION INTO AN UNCONDITIONAL SAMPLING PROCESS: The approach modifies existing iterative score-based samplers by inserting measurement-consistency operations before their iteration updates.The paper describes the resulting procedure as compatible with Euler-Maruyama, annealed Langevin dynamics, probability-flow ODEs, and Predictor-Corrector samplers.

4 EXPERIMENTS

Experiments evaluate the method across CT and MRI reconstruction tasks, comparing it with learning-free, supervised, and score-based baselines. Results show competitive performance, improved adaptability to changed measurement counts, and reuse of one CT score model across different measurement processes.

  • Experimental setup: Experiments cover sparse-view CT, metal artifact removal, and undersampled MRI using datasets and baselines spanning learning-free, supervised, and unsupervised methods.The CT datasets are LIDC and LDCT, while MRI experiments use BraTS.
  • Competing with supervised learning approaches: The method outperforms SIN-4c-PRN on sparse-view CT for both LIDC and LDCT, including with 20 measurements versus supervised baselines trained with 23.It also outperforms SNMAR on metal artifact removal and produces reconstructions with fewer artifacts and better preserved structure.
  • Competing with supervised learning approaches: For undersampled MRI at 8× acceleration, the method ranks second with performance comparable to the top supervised method DuDoRNet.The MRI comparison includes supervised methods trained with 8× acceleration and unsupervised techniques.
  • Generalizing to different number of measurements: The same score model achieves the best MRI performance at both 24× and 4× acceleration, while supervised methods struggle or fail to generalize when measurement counts change.For sparse-view CT, supervised methods likewise struggle with different projection angles.
  • Generalizing to different measurement processes in CT: A single score model trained on CT images performs both sparse-view CT reconstruction and metal artifact removal, outperforming MAR-specific supervised methods.The two tasks use different CT measurement processes but share the same transformation component in the stated decomposition.
  • Comparing against existing score-based methods: Among score-based MRI methods, ALD + Ours uniformly outperforms Langevin and Score SDE across all tested measurement counts, while PC + Ours improves further.The comparison is performed on undersampled MRI reconstruction using BraTS.

5 CONCLUSION

The paper proposes a fully unsupervised score-based approach for linear inverse problems that avoids paired training data and adapts to measurement processes at test time. Experiments show competitive or superior reconstruction performance and broader generalization across CT and MRI settings.

  • Conclusion: The method solves linear inverse problems with score-based generative models without paired training data.It learns from medical images and uses the score model as a prior for reconstruction.
  • Conclusion: The approach adapts at test time to different projection counts, downsampling ratios, and CT measurement processes.A single model can handle sparse-view CT reconstruction and metal artifact removal.
  • Conclusion: Empirically, the method matches or outperforms supervised counterparts on sparse-view CT and undersampled MRI while generalizing better to new measurement processes.The reported scope includes several CT and MRI reconstruction tasks.

A PROOFS

The proofs characterize a full-rank linear measurement operator through an invertible coordinate transform and derive a closed-form solution for the optimization problem. The resulting expression separates measured and unmeasured components in transformed coordinates.

  • Proposition 1: If A has rank m, it can be written as A = P(Λ)T, where T is invertible and Λ is diagonal with m ones.P(Λ) reduces an n-dimensional vector to m dimensions by removing coordinates whose diagonal entries in Λ are zero.
  • Proposition 1: The proof constructs T by extending the linearly independent rows of A to a full basis of R^n, making T invertible.The first m rows are the rows of A, and the remaining rows complete the basis.
  • Lemma 1: Lemma 1 states that P(Λ)a = P(Λ)b exactly when Λa = Λb, and introduces any right inverse P^-1(Λ) for the projection operator.This equivalence is used to express the measurement constraint in transformed coordinates.
  • Theorem 1: Theorem 1 gives the solution in closed form by combining the transformed estimate, the measurement-consistent component, and the weighting parameter λ.The displayed update uses T^-1, Λ, P^-1(Λ), the measurement ŷ_t, and the estimate x̂_t.
  • Theorem 1: The optimization objective becomes a quadratic function of z after decomposing transformed coordinates into measured and unmeasured components.The measurement constraint fixes the Λ component while leaving the (I − Λ) component unconstrained by the observations.

B.1 ADDITIONAL RESULTS

Additional results report SSIM trends across measurement counts and provide visual examples for metal artifact removal. The SSIM curves generally follow the PSNR trends shown elsewhere.

  • Additional results: SSIM results are provided versus the number of measurements for multiple MRI and CT tasks, with trends generally similar to the PSNR curves.The results include MRI on BraTS and CT on LIDC.
  • Additional results: The section also provides visualizations of metal artifact removal results.Figure 7 contains example reconstructions for this task.

B.2 THE TASK OF METAL ARTIFACT REMOVAL

Metal implants create bright metal traces in CT sinograms, producing strong artifacts. Metal artifact removal is framed as reconstructing the full image from a sinogram with missing metal-trace regions.

  • Metallic implants cause strong artifacts in CT images through extremely bright sinogram regions called metal traces.
  • The measurement process removes metal traces from the sinogram, leaving the remaining sinogram data available for reconstruction.
  • Metal artifact removal is formulated as an inverse problem that reconstructs the full CT image from a partially known sinogram.
  • Perfect inpainting of the missing sinogram regions would yield an artifact-free reconstructed CT image.

B.3 DETAILS OF DATASETS

Experiments use CT datasets for reconstruction and BraTS MRI data for accelerated imaging. The implementation uses score-based models with VE SDEs and Predictor-Corrector sampling, with tuned sampling parameters.

  • Datasets: The LIDC dataset contains 1018 thoracic CT cases for lung cancer detection and diagnosis.
  • Datasets: BraTS 2021 contains 2000 cases and 8000 MRI scans spanning four MR contrasts.
  • Datasets: For MRI reconstruction, 2D slices are extracted from 3D volumes and k-space data are simulated using Fast Fourier Transform.
  • Model and sampling: The method uses NCSN++ with a Variance Exploding SDE and follows the associated training procedure.
  • Model and sampling: Predictor-Corrector sampling replaces the numerical SDE solver because it generally performs better for Variance Exploding SDEs.
  • Parameter tuning: Sampling parameters η and λ are tuned with 100 Bayesian-optimization steps on validation data before test evaluation.

B.5 TRAINING DETAILS OF BASELINE MODELS

Baseline comparisons cover standard CT reconstruction, optimization-based methods, and learned sinogram or regularizer approaches. Supervised CT baselines are trained using 23 projection angles, with additional testing across projection settings.

  • CT reconstruction baselines: Filtered back projection reconstructs CT images by mapping sinogram projections into image space using projection angles and geometry.
  • CT reconstruction baselines: FISTA-TV solves linear inverse problems with total variation regularization through iterative matrix-vector multiplication and shrinkage-threshold steps.
  • Learned baselines: cGAN completes sinogram data before CT reconstruction, avoiding costly iterative tomographic processing.
  • Experimental settings: Supervised sparse-view CT baselines are trained with 23 projection angles, while cGAN is tested on other projection settings for generalization.
  • Learned baselines: SIN-4c-PRN combines sinogram inpainting and image refinement networks connected through filtered back projection.
  • Learned baselines: Neumann networks learn nonlinear regularizers from training data for inverse problems.

B.5.2 BASELINE MODELS FOR UNDERSAMPLED MRI RECONSTRUCTION

MRI baselines include dual-domain recurrent and cascaded dense reconstruction networks, with comparisons focused on undersampled k-space recovery. Cascaded DenseNet generalizes better to more measurements than DuDoRNet.

  • Baseline models: DuDoRNet jointly recovers k-space data and images to address aliasing in both frequency and image domains.
  • Baseline models: Cascade DenseNet uses stacked dense blocks, residual connections, and data-consistency layers to reconstruct de-aliased MR images.
  • Comparison: Cascaded DenseNet generalizes better to more measurements than DuDoRNet, as shown in Figures 5 and 6.
  • Related artifact-reduction baselines: The metal-artifact baseline comparison includes linear interpolation, cGANMAR, and SNMAR, covering sinogram and image-domain artifact reduction.
Loading 2111.08005v2…