Score-Based Diffusion Prior

• Score-Based Diffusion Prior is defined via a time-dependent score function that models the full data distribution, enabling an implicit and data-driven probabilistic prior.
• It integrates reverse-time SDE integration with measurement consistency projections to solve inverse problems, notably enhancing MRI reconstruction with uncertainty quantification.
• Empirical benchmarks show that the method improves quantitative metrics such as PSNR and SSIM compared to traditional deterministic and handcrafted priors.

A score-based diffusion prior refers to a probabilistic prior over images (or other signals) defined implicitly by the score function—i.e., the gradient of the logarithmic probability density—learned by a score-based diffusion model. In this context, the prior governs a continuous distribution over possible images, encapsulating statistical properties of the training data, and is realized through the time-dependent score function estimated by the denoising score matching objective. This modeling framework is central to several recent advances in Bayesian inference for inverse imaging, especially in domains such as accelerated MRI, where the goal is to reconstruct images from highly incomplete or noisy measurements by sampling from the conditional posterior distribution under an expressive, data-driven prior.

1. Foundational Principles of Score-Based Diffusion Priors

The key idea underlying a score-based diffusion prior is to model the prior distribution $p(x)$ on clean data via the time-dependent score function $\nabla_x \log p_t(x)$, which tracks the evolution of the data's distribution under a forward diffusion process. Typically, a stochastic differential equation (SDE) of the form

$dx = f(x, t) dt + g(t) dW$

is selected, where $f$ is the drift, $g$ encodes the diffusion coefficient, and $W$ denotes Brownian motion. The diffusion process non-linearly "noises" the data until, at a large $t=T$, the distribution $p_T(x)$ is tractable (e.g., standard normal).

A neural network $s_\theta(x, t)$ is trained to approximate the score function $\nabla_x \log p_t(x)$ using denoising score matching: $$\min_\theta\, \mathbb{E}_{t\sim\mathrm{Uniform}[0,T],\, x_0\sim p_0,\, \mathrm{noise}}\big[\lambda(t) \|s_\theta(x_t, t) - \nabla_x \log p_{0t}(x_t|x_0)\|^2 \big],$$ where $x_t$ is formed by adding Gaussian noise with variance $\sigma(t)^2$ to $x_0$, and $\lambda(t)$ is a weighting function.

This approach eliminates the need for explicitly parameterizing or sampling from the prior. The score network learns the shape of the data distribution at all noise levels, thereby encoding a powerful implicit prior that captures complex, non-Gaussian data statistics.

2. Integration with Inverse Problem Solving

In typical inverse problems, only noisy and/or incomplete measurements $y = A x$ are available, where $A$ may involve subsampling (e.g., in MRI, a Fourier mask). Conventional strategies use handcrafted regularizers (such as total variation), which are either too simple or require manual tuning. In contrast, the score-based diffusion prior allows leveraging the learned score function as the regularization mechanism.

During inference, an initial sample is evolved backwards via the numerically-solved reverse-time SDE: $$dx = [f(x, t) - g(t)^2 s_\theta(x, t)] dt + g(t) d\bar{W},$$ where $d\bar{W}$ denotes time-reversed Brownian increments. This "denoising" operation is interleaved with measurement consistency projections. For instance, in MRI reconstruction, projections onto the data manifold occur via

$x_i \leftarrow x_i + A^* (y - A x_i),$

where $A^*$ is the adjoint operator. This iterative algorithm ensures that each sample both conforms to the learned score-based prior and satisfies measurement constraints.

The framework readily extends to settings with complex-valued data (by separate treatment of real and imaginary channels) and to parallel imaging with multiple receiver coils, through either coil-wise or hybrid updates.

3. Generative Nature and Uncertainty Quantification

Unlike deterministic regression networks, score-based diffusion priors enable sampling from the full posterior. The stochasticity of the reverse-time diffusion process (through random seed initialization and injection of noise in each update step) permits the generation of multiple distinct, data-consistent reconstructions.

By producing an ensemble of samples, pixelwise uncertainty can be quantified via empirical statistics (e.g., mean, standard deviation over reconstructions), directly enabling rigorous uncertainty analysis for operator decision support. This is not feasible with conventional deterministic models—highlighting a significant methodological advancement offered by diffusion priors.

4. Empirical Performance and Practical Considerations

Experimental benchmarks on both simulated and real data (e.g., for accelerated MRI) consistently show that score-based diffusion priors lead to marked improvements in standard quantitative metrics such as peak signal-to-noise ratio (PSNR) and structural similarity (SSIM), compared to classical compressed sensing (e.g., total variation), supervised deep learning (e.g., U-Net), and advanced variational architectures (e.g., E2E-varnet, DuDoRNet).

The approach generalizes robustly to a wide range of measurement operators (subsampling masks), anatomy, and contrasts not present in the training set—confirming the flexibility of the learned prior. Models trained on magnitude images alone suffice for complex-valued multi-coil MRI, a practical advantage given the limited availability of calibrated raw data.

Furthermore, the generative prior is invariant to subsampling patterns: a single pre-trained model can reconstruct under diverse acquisition protocols with no retraining.

5. Implementation Workflow

A typical pipeline using a score-based diffusion prior in inverse problems encompasses:

1. Training: Learn the time-dependent score function $s_\theta(x, t)$ using magnitude images with denoising score matching, requiring only oracle (fully sampled) data, usually available in DICOM formats.
2. Inference (Reconstruction): Alternate between predictor steps (reverse-SDE integration, e.g., Euler–Maruyama) and POCS-like data consistency projections, iteratively updating $x_i$.
3. Extension to Parallel Imaging: Apply the steps to each coil image and merge (SSOS), or inject coil sensitivity information in a hybrid data-consistency step.

Computational costs are primarily dictated by the score network's evaluation per SDE step and the number of iterations. The process is scalable and can be parallelized.

6. Limitations, Generalization, and Future Directions

Current limitations include increased inference time relative to purely feed-forward networks (due to iterative SDE integration and repeated projections), and the risk of over-regularization if the prior excessively dominates poorly constrained regions. Careful balancing between score prior and measurement projections is required to avoid artifacts or loss of relevant "noisy" signal.

Nevertheless, due to the parameterization of the score function, the prior generalizes well to out-of-distribution (OOD) data. The ability to reconstruct both real and complex-valued data and to handle arbitrary subsampling patterns without retraining establishes broad practical relevance. Further generalizations to other imaging modalities and operators, as well as improved uncertainty calibration, remain active directions of research.

7. Summary Table: Score-Based Diffusion Priors vs. Other Priors

Feature	Score-Based Diffusion Prior	Total Variation / Handcrafted Prior	Supervised DL Prior (e.g., U-Net)
Expressiveness	High (learns full distribution)	Low (simple statistics)	Moderate (data-driven, but deterministic)
Handles Arbitrary Sampling	Yes	No	No
Quantifies Uncertainty	Yes (multi-sample)	No	No
Generalization to OOD	Strong	Weak	Weak/moderate
Plug-and-play Flexibility	Yes	Limited	Limited
Supervision Required	Unsupervised (score matching)	No	Fully supervised

The adoption of score-based diffusion models as priors constitutes a transition from deterministic, handcrafted, or point-estimate-focused regularizers to probabilistic, data-driven, and generative image models, offering superior flexibility, robust uncertainty quantification, and empirical performance across a range of inverse problem settings (Chung et al., 2021).