The choice of prior is central to solving ill-posed imaging inverse problems, making it essential to select one consistent with the measurements y to avoid severe bias. In Bayesian inverse problems, this could be achieved by evaluating the model evidence p(y∣M) under different models M that specify the prior and then selecting the one with the highest value. Diffusion models are the state-of-the-art approach to solving inverse problems with a data-driven prior; however, directly computing the model evidence with respect to a diffusion prior is intractable. Furthermore, most existing model evidence estimators require either many pointwise evaluations of the unnormalized prior density or an accurate clean prior score. We propose DiME, an estimator of the model evidence of a diffusion prior by integrating over the time-marginals of posterior sampling methods. Our method leverages the large amount of intermediate samples naturally obtained during the reverse diffusion sampling process to obtain an accurate estimation of the model evidence using only a handful of posterior samples (e.g., 20). We also demonstrate how to implement our estimator in tandem with recent diffusion posterior sampling methods. Empirically, our estimator matches the model evidence when it can be computed analytically, and it is able to both select the correct diffusion model prior and diagnose prior misfit under different highly ill-conditioned, non-linear inverse problems, including a real-world black hole imaging problem.
Figures & tables
Figure 1: Visualization of the unbiased estimators Θhigh,Θlow of the likelihood score ∇xtlogp(y∣xt) described in Lemma 2 for both in-distribution (left) and out-of-distribution (right) measurements y . At large diffusion time steps, the high noise estimator (top) uses the distance between x~0 and E[x0∣xt] . At small diffusion time steps, the low noise estimator (bottom) uses the likelihood score at x~0 . The estimator is larger (longer arrows) when y is out-of-distribution, as there greater KL divergence from the posterior to the prior; this results in DiME computing a lower evidence. Note the scaling factors σt2at and Σx0∣xt of the estimator are not shown.
In-distribution x∗
Out-of-distribution x∗
Saddle point x∗
Estimate
Rel. Error ↓
Estimate
Rel. Error ↓
Estimate
Rel. Error ↓
Ground truth logp(y)
−128.7±0.0
0.0
−1680.2±0
0.0
−251.3±0.0
0.0
Naive MC (1000 samples)
−3283±130
2451%
−41289±577
2357%
−6029±170
2299%
Naive MC (10000 samples)
−3041±113
2263%
−40123±487
2288%
−5666±144
2155%
Original DAPS heuristic
−316.2±79.5
146%
−1735.4±4.7
3.3%
−269.6±3.3
7.3%
TI
−132.8±1.3
3.2%
−1773.8±11.8
5.6%
−254.2±3.2
1.2%
Table 1: Evidence estimates on Mixture of Gaussians study with analytic logp(y) and ground truth x∗ in-distribution (top), out-of-distribution (middle), and at a saddle point (bottom). Mean and standard deviation of 50 trials displayed each with 20 sample paths. When the analytic prior score is available, DiME performs comparably to the baselines for all cases, while DiME-PnPDM performs well when the ground truth is roughly in-distribution. *Important: as DiME-PnPDM outputs an unnormalized estimate, we normalize via the analytic constant factor for ease of comparison.
Figure 2: Model evidence confusion matrix for Gaussian phase retrieval ( left ) and Fourier phase retrieval ( right ) for each (ground truth measurement, model) pair of MNIST digits. Our method selects the correct model for all cases. Posterior samples shown for dotted matrix entries below. For Gaussian phase retrieval, DiME estimates higher model likelihood for visually similar digits, such as 4 and 9. For Fourier phase retrieval, both translations and reflections are invariances as seen in the posterior samples, so DiME estimates a high likelihood of model 9 given a measurement of a 6.
Figure 3: Model evidence estimates on real M87* observations across 5 different priors using exact DAPS ( left ) and Gaussian approximation DAPS ( right ). Our method concludes that, of these prior models, GRMHD is the most likely model. The violin plots show the distribution of evidences from 20 different posterior-sample paths from the DAPS sampling process, and the mean is the overall evidence estimate. Gaussian approximation DAPS gives highly accurate evidence estimates with 7x less compute than exact DAPS, but with slightly higher variance.
Figure 4: M87* reconstructions using the 5 priors using exact DAPS ( left ) and Gaussian approximation DAPS ( middle ). While all posterior samples have a data fit of around reduced χ2≈1 , they have varying path evidence estimates (shown in the top left of each image). Right: Example unconditional samples from each prior.
Figure 5: Left: GRMHD model validation results on M87* observations by comparing to evidence of in-distribution measurements y . Our method shows that the evidence of M87* observations have a z -score of about -0.81 compared to the evidence distribution of GRMHD measurements, indicating that M87* is statistically in-distribution of GRMHD. The evidence of simulated measurements of out-of-distribution images are also shown, demonstrating that measurements from out-of-distribution (OOD) images have OOD evidence. Right: Mean reconstruction and posterior samples of the OOD images with the highest and lowest distance to the GRMHD prior. The blurred ground truth image, or the maximum resolution that can be obtained from measurements, is also displayed.
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Task
lr
Max steps
Mixture of Gaussians (in-distribution)
5×10−4
2,000
Mixture of Gaussians (out-of-distribution)
2×10−5
10,000
Mixture of Gaussians (saddle point)
2×10−5
10,000
Fourier phase retrieval
1×10−3
1,000
Gaussian phase retrieval
5×10−4
1,000
M87* black hole imaging
1×10−5
1,000
Appendix
Table 7
Figure 6: Model evidence confusion matrix using sequential Monte Carlo (SMC) for Gaussian phase retrieval ( left ) and Fourier phase retrieval ( right ) for each (ground truth measurement, model) pair of MNIST digits. When SMC uses the learned diffusion score ∇xlogp(xt) at t=0.1 as the prior score, it often fails to select the correct model.
σmin=0.1
σmin=0.05
SMC schedule
GT 8 Model 6
GT 8 Model 7
GT 8 Model 8
GT 8 Model 6
GT 8 Model 7
GT 8 Model 8
Linear
−182±9
−212±9
−210±11
−279±26
−352±3
−426±40
Exponential
−183±7
−224±5
−239±6
−292±20
−383±9
−468±12
Sigmoidal
−175±5
−219±13
−191±2
−273±15
−374±9
−480±17
Adaptive 0.5
−195±31
−202±23
−261±1
−296±52
−346±16
−489±33
Adaptive 0.6
−184±6
−189±7
−259±4
−279±45
−373±17
−485±26
Appendix
Table 2: SMC model evidence estimates on Gaussian phase retrieval.
σmin=0.1
σmin=0.05
SMC schedule
GT 8 Model 6
GT 8 Model 7
GT 8 Model 8
GT 8 Model 6
GT 8 Model 7
GT 8 Model 8
Linear
−416±11
−431±4
−409±1
−485±25
−577±2
−563±2
Exponential
−422±20
−427±2
−467±11
−527±10
−582±12
−585±9
Sigmoidal
−415±9
−451±3
−411±7
−520±13
−580±7
−575±5
Adaptive 0.5
−409±4
−409±6
−460±8
−524±33
−563±18
−567±10
Adaptive 0.6
−431±20
−437±11
−458±8
−510±9
−566±14
−562±13
Appendix
Table 3: SMC model evidence estimates on Fourier phase retrieval.
While diffusion priors generate high-quality posterior samples across many inverse problems, they are often trained on limited training sets or purely simulated data, thus inheriting the errors and biases of these underlying sources. Current approaches to finetuning diffusion models rely on a large number of observations with varying forward operators, which can be difficult to collect for many applications, and thus lead to overfitting when the measurement set is small. We propose a method for tuning a prior from only a single observation by combining existing diffusion priors into a single product-of-experts prior and identifying the exponents that maximize the Bayesian evidence. We validate our method on real-world inverse problems, including black hole imaging, where the true prior is unknown a priori, and image deblurring with text-conditioned priors. We find that the evidence is often maximized by priors that extend beyond those trained on a single dataset. By generalizing the prior through exponent weighting, our approach enables posterior sampling from both tempered and combined diffusion models, yielding more flexible priors that improve the trustworthiness of the resulting posterior image distribution.
This article addresses the issue of estimating observation parameters (response and error parameters) in inverse problems. The focus is on cases where regularization is introduced in a Bayesian framework and the prior is modeled by a diffusion process. In this context, the issue of posterior sampling is known to be thorny, and a recent paper proposes a notably simple and effective solution. Additionally, it opens an remarkable flexibility when it comes to estimating observation parameters. The proposed strategy enables to define an optimal estimator for both observation parameters and image of interest. Furthermore, the strategy provides a means for uncertainty quantification. In addition, MCMC algorithms allow for the computation of estimates and properties of posteriors, while offering some guarantees. The paper presents several numerical experiments that clearly confirm the computational efficiency and the quality of both estimates and uncertainty quantification.
Jean-François Giovannelli
Groupe Signal-Image, IMS (Univ. Bordeaux, CNRS, BINP), Talence, France
Diffusion models (DMs) have proven to be effective in modeling high-dimensional distributions, leading to their widespread adoption for representing complex priors in Bayesian inverse problems (BIPs). However, current DM-based posterior sampling methods proposed for solving common BIPs rely on heuristic approximations to the generative process. To exploit the generative capability of DMs and avoid the usage of such approximations, we propose an ensemble-based algorithm that performs posterior sampling without the use of heuristic approximations. Our algorithm is motivated by existing work that combines DM-based methods with the sequential Monte Carlo (SMC) method. By examining how the prior evolves through the diffusion process encoded by the pre-trained score function, we derive a modified partial differential equation (PDE) governing the evolution of the corresponding posterior distribution. This PDE includes a modified diffusion term and a reweighting term, which can be simulated via stochastic weighted particle methods. Theoretically, we prove that the error between the true posterior and the empirical distribution of the generated samples can be bounded in terms of the training error of the pre-trained score function and the number of particles in the ensemble. Empirically, we validate our algorithm on several inverse problems in imaging to show that our method gives more accurate reconstructions compared to existing DM-based methods. Our code is available at the following Github repository~https://github.com/HaoxuanSteveC00/AFDPS-TMLR.
Haoxuan Chen, Yinuo Ren, Martin Renqiang Min +2
ICME, Stanford University, Stanford, CA 94305 · Machine Learning Department, NEC Labs America, Princeton, NJ 08540 · Department of Mathematics and ICME, Stanford University, Stanford, CA 94305