Variational Bayesian Imaging with an Efficient Surrogate Score-based Prior
Authors: Berthy Feng
Katherine Bouman
Venue: Accepted by TMLR ()

Abstract: We propose a surrogate function for efficient yet principled use of score-based priors in Bayesian imaging. We consider ill-posed inverse imaging problems in which one aims for a clean image posterior given incomplete or noisy measurements. Since the measurements do not uniquely determine a true image, a prior is needed to constrain the solution space. Recent work turned score-based diffusion models into principled priors for solving ill-posed imaging problems by appealing to an ODE-based log-probability function. However, evaluating the ODE is computationally inefficient and inhibits posterior estimation of high-dimensional images. Our proposed surrogate prior is based on the evidence lower bound of a score-based diffusion model. We demonstrate the surrogate prior on variational inference for efficient approximate posterior sampling of large images. Compared to the exact prior in previous work, our surrogate accelerates optimization of the variational image distribution by at least two orders of magnitude. We also find that our principled approach gives more accurate posterior estimation than non-variational diffusion-based approaches that involve hyperparameter-tuning at inference. Our work establishes a practical path forward for using score-based diffusion models as general-purpose image priors.

Published in Transactions on Machine Learning Research (07/2024)
Variational Bayesian Imaging with an Efficient Surrogate Score-based Prior
Berthy T. Feng Computing and Mathematical Sciences California Institute of Technology
bfeng@caltech. edu
Katherinc L. Bouman
Computing and Mathemalical Sciences California Institute of Technology
Reviewed on OpenReview: openreveee... net/ por//f id= db2pFKVcml
Abstract
We propose a surrogate function for efficient yet principled use of score-based priors in Bayesian imaging. We consider ill-posed inverse imaging problems which one aims for a clean image posterior given incomplete or noisy measurements. Since the measurements do not uniquely determine a true image, a prior is needed constrain the solution space. Recent work turned score-based diffusion models into principled priors for solving ill-posed imaging problems by appealing to ODF-based log- log-probability function. However, eval- uating the ODE is computationally inefficient and inhibits posterior estimation of high- dimensional images. Our proposed surrogate prior is based on the evidence lower bound of a score-based diffusion model. We demonstrate the surrogate prior on variational inference for efficient approximate posterior sampling of large images. Compared to the exact prior in previous work, our surrogate accelerates optimization of the variational image distribution by at least two orders of magnitude. We also find that our principled approach gives more accurate posterior estimation than non-variational diffusion- based approaches that involve syperparameter-tuning at inference. Our work establishes practical path forward for using score-based diffusion models as general-purpose image priors.
1 Introduction
Ill-posed image reconstruction arises when measurements of an object of interest are incomplete or noisy, making it impossible to uniquely determine the true image. Imaging in this setting requires prior to constrain images according to desired statistics. From Bayesian perspective, the prior inftuences both the uncertainty and the richness of the estimated image. Although diffusion-based generative models represent rich image priors, leveraging these priors for Bayesian image reconstruction remains challenge. True posterior sampling with an unconditional diffusion model intractable, so most previous mothods either
heavily approximate the posterior (Chung al. 2023 Jalal et al. 2021 Kawar al. 2022 Song et al.
2023a or perform non-Baycsian conditional sampling (Choi al. 2021 Chung et al. 2022a Chung Yo
2022
Graikos al. 2022 Song et al., 2022 Adam al.] 2022).
Recent work demonstrated how to turn score-based diffusion models into probabilistic priors (score-based priors) for variational Bayesian imaging (Feng et al. 2023). This method requires the exact probability of a proposed image to be evaluated with a sunputaationlll--xxeeer
ordinary differential equation (ODE), requiring days to a week to reconstruct even a 32 x 32 image. Even so, this approach is appealing because offers the same thcoretical guarantees as traditional variational inference. This can sometimes be worth the computational cost for imaging applications in which measurements are expensive or difficult obtain
(e.g., black-hole imaging (EHTC. 2019), cryo-electron microscopy (Egelman) 2016 Levy et al. 2022 and

---

Published in Transactions on Machine Learning Research (07/2024)
X-ray crystallography (Woolfson 1997 Miao et al. 2008)) Still, it would be beneficial to have efficient "surrogate" for this approach that can be used in the development and testing stages simply to reduce computational costs. We present a method for variational Bayesian inference that is both principled and computationally efficient thanks to a surrogate score-based prior.
Computing exact probabilities under a diffusion model inefficient or even intractable, but computing the evidence lower bound (ELBO) (Song et al. 2021a Ho et al., 2020) is computationally efficient and feasible for high-ddimensional images. We propose to use this lower bound In particular, use the ELBO function be plugged into any inference algorithm that requires the value or gradient of the posterior log-density. For variational inference of an image posterior, we find at least two orders of magnitude in speedup of optimizing the variational distribution. Furthermore, our approach reduces GPU memory requirements, as there is no need to evaluate and backpropagate through an ODE. These efficiency improvements make practical perform principled inference with score-based priors.
surrogate for the exact acore-based prior. substitute for the exact log- oog-robabiilty function, and it
In this paper, describe based prior | We provide experimental results to validate the proposed surrogate, including high-dimensional posterior samples of sizes up to 256 256. In the sctting of accelerated MRI, we quantify time- and memory- efficiency improvements of the surrogate over the exact prior. We also demonstrate that our approach achieves more accurate posterior estimation and higher-fidelity image reconstructions than diffusion- based methods that deviate from true Bayesian inference. Finally, we demonstrate how our approach can be used for black-hole imaging, in which images must be carefully reconstructed from sparse telescope measurements.
variational-inference approach to estimate posterior with a surrogate score-
2 Related work
2.1 Bayesian inverse imaging
Imaging can be framed an inverse problem in which hidden image x' E RR must be recovered from measurements y E RM, where
y = f(x) + n.
Usually the forward model f F + RM and statistics of the noise n E R. are assumed to be known.
Bayesian imaging accounts for the uncertainty added by the mcasurement process by formulating a postcrior distribution p(x y), whose log-densily can be decomposed into likelihood term and prior term:
logp(x y) = log x) + log + const.
(1)
Given a log-likelihood function log p(y x) and prior log-probability function logp(x), we can use estab- lished techniques for sampling from the posterior, such as Markov chain Monte Carlo (MCMC) (Brooks) et al., 2011) or variational inference (VI) (Blei et al. 2017). MCMC algorithms generate a Markov chain whose stationary distribution is the posterior, but they are generally slow to converge for high-dimensional data like images. VI instead approximates the posterior with tractable distribution. The variational dis- tribution usually parameterized and thus can be efficiently optimized represent high-dimensional data distributions. Deep Probabilistic Imaging (DPI) (Sun & Bouman, 2021 Sun et al. 2022) proposed cf- ficient VI approach specifically for computational imaging with image regularizers. In DPI, the variational distribution is a discrete normalizing flow (Kobyzev et al. 2020), which is an invertible generative model capable of representing complex distributions.
The log-prior term in Eq. [ is often defined through a regularizer. Traditional imaging approaches use hand-crafted regularizers including total variation (TV) (Vogel & Oman 1996), maximum entropy (Gull & Skilling 1984 Narayan & Nityananda] 1986), and the LI norm (Canddes & Romberg 2007] or data-driven
regularizers such as through Gaussian mixture models (Zoran & Weiss] 2011 2012). Various deep-learned regularizers have been proposed, including adversarial regularizers (Lunz al. 2018 and input-convex
neural networks (Mukherjec et al. 2020 Shumaylov et al. 2023). but they impose training or architecture restrictions that may limit generalizability and expressiveness.
Code is available at http::///tttha..cmm//tttttt///at
2

---

Published in Transactions on Machine Learning Research (07/2024)
We note that this perspective of Bayesian imaging with standalone priors differs from that of other learning- based approaches to inverse problems. Many approaches train end-to-end neural network on a. paired
dataset of measurements and ground-truth images (Pathak et al. 2016; lizuka et al. 2017 Yu et al. 2019
Zhang et al. 2022 2020; Yin et al. 2021 Saharia et al., 2022 Delbracio et al.) 2021 Delbracio & Milanfar
2023). However, for every new measurement distribution, a new paired dataset would have to be oblained, and the network would have to be re-trained. This can become cumbersome when one wishes to analyze different measurement settings with the same prior (e.g., when exploring different MRI acquisition strategies). It also prechudes analyzing the effects of different priors (e.g., by imposing different assumptions in imaging black hole (EHTC) 2019)) since the implicit prior of the network changes every time it is re-trained. Our approach, on the other hand, only requires a dataset of clean images for the prior, and the same learned prior can be used across different inverse problems by being paired with different measurennnn-likkllioow functions in a modular way, Large-scale datasets of clean images exist for many applications, such as the fastMRI (Zbontar et al.] (2018) datasct for accelerated MRI. It also usually possible to create a clean dataset through cost.yy acquisition of ideal images or simulation.
2.2 Diffusion models for inverse problems
Diffusion models (Sohl-Dicksterr Dickstein et al. 2015 Ho et al. 2020 Song & Ermon 2019 Song et al. 2019 2021b)
learn to model a rich image distribution that could be useful as a prior for image reconstruction. A diffusion model generates an image by starting from an image of noise and gradually denoising until it becomes a clean image. We discuss this process, known as reverse diffussin, in more detail in Sec. 3.1
Given inverse problem, simply adapting trained diffusion model to sample from the posterior instead of the learned prior is intractable (Feng et al.) 2023). Therefore, most diffusion- approaches do not infer
a true posterior. Some project images onto a mecsusremenn-conniitteee subspace (Song et al. 2022 Chung
ict al. 2022b Chung & Ye 2022 Choi ct al. 2021 Chung ct al. 2022a), but the projection does not account
for measurement noise and might pull images away from the true posterior. Others follow a gradient loward
higher measurement- likelihood -iikelio throughout reverse diffusion (Chung et al. 2023 Jalal et al, 2021 Graikos
ct al. 2022 Kawar et al.) 2022 Adam et al. 2022 Song al. 2023a Mardani et al. 2023). but they heavily
approximate the posterior. Overall, these dilfusion-based mods require typerparrameler-uuninng of the measurement weight. As soon as hyperparameters are introduced, there is no guarantee of sampling from a posterior that represents the true uncertainty. These ciitssoon-bssed methods are not principled enough for scientific and medical applications that require accurate uncertainty quantification.
Scoro-based priors Recent work proposed an alternative perspective: turning score-based diffusion model into a standalone, probabilistic prior (score-based prior) that can be paired with any measurement- likelihood function and plugged into established variational-inferrnce approaches. Rigorous variational in- ference is usually more computationally intensive than the aforementioned diffusion-basse approaches, so this type of approach is best-suited for applications in which finding the best posterior for a given set of measurements is worth the computational overhead.
However, the approach of Feng et al. (2023) is too memory-intensive feasibly handle large images. Even for small images, it is too time-intensive to enable fast development. This is due to a log-density function based an ODE (see Sec. 3.2 that is computationally expensive. Our work proposes efficient surrogate that makes this type of principled approach practical for high-dimensional inference. While other efforts have been made to speed up diflusion-model sampling, such as by reducing the number of function evaluations Karras et al. 2022 Zhang & Chen 2022 Song et al. 2023b), performing latent diffusion (Rombach et al., 2022 Avrahami et al. 2023), or learning more efficient reverse processes (Lipman et al.] [2022] De Bortoli et al. 2021). these approaches either break the tractability of image probabilities or only incrementally improve speed. We aim to fundamentally improve efficiency without altering the diffusion model.
3

---

Published in Transactions on Machine Learning Research (07 (2742)
3 Background
In this section, we review background on score-based diffusion models with an emphasis on evaluating probabilities of images under trained diffusion model. We then describe how a diffusion process gives rise to efficient denoising-based lower bound on these image probabilities.
3.1 Score-based diffusion models
The core idea a diffusion model is that transforms a simple distribution я to a complex image distribution through a gradual process. We follow the popular framework denoising diffusion models, which transform noise samples from π (o, to clean samples from the data distribution Paata through gradual denoising. With knowledge of the noise distribution and denoising process, we can assess the probability of novel image under this generative model.
The transformation from simple distribution to complex one occurs over many steps. To determine how the data distribution should look at each step of denoising, we turn to a stochastic differential equation (SDE) that describes a diffusion procoss from clean images to noise (Song et al. 2021b). The diffusion SDE is defined on the time interval & E [0, T] as
dx = f(x, t)dt + g())d,
(2)
where w E IRD denotes Brownian motion; g() E IR is the diffusson coefficient, which controls the rate of noise increase; and f(:,t) RD + RR is the drift coefficient, which controls the deterministic evolution of x(t). By defining stochastic trajectory (Ejo)) this SDE gives rise to a time-dependent probability distribution Pt, which the marginal distribution of x(t). We construct f(-,t) and g(t) so that po Pdatas then pr & Image generation amounts to reversing the diffusion, which requires the gradient of the data log-density (score) at every noise level in order to nudge images toward high probability under Pdata- A convolutional neural network So known a score model lcarns to approximate the truc score: s0 t) z Vx (x)..
Sampling with a reverse-time SDE Once trained, sy (x, t) is used to reverse diffusion and generate clean images from noise. This results in a distribution pa SDE denoted as such because is determined by a reverse-time SDE (Song et al. 2021b):
dx f(xx,t) so t)] dt + g(t)dw
(3)
where w RD denotes Brownian motion, and f(., t) and gt) are the same as in Eq. To generate an image, we first sample x(T) I) and then numerically solve the reverso-time SDE for x(0). The marginal distribution of x(0) is denoted by pe SDE,, which is close Pdata when the score model is well- trained. Intuitively, Eq. 3 explains how to denoise
image x(I) using the score model Sg in order to get closer to the clean image distribution. Indeed, Eq. is exactly the reverse of the forward-time SDE (Eq, 2 [Anderson 1982) when replacing Vx log (x) with (x, t), which means that removing noise corresponds to reaching higher probability under the data distribution.
3.2
Image probabilities
The generated image distribution p SD theoretically assigns density to any image. However, reverse diffussio does not lead to image distribution with tractable probabilities. In this subsection, we describe two workarounds: one based on an ODE and the other based on an ELBO related to denoising score-matching.
To compute the probability of an image x under y we need to invert it from x(0) = x to x(T). However, this not tractable through the SDE: just as it intractable to reverse a random walk, it intractable account for all the possible starting points x(T) that could have resulted x(0) through the stochastic process. Probability computation calls for an invertible process that lets us map any point from Pdata to a point from (O, and vice versa.
4

---

Published in Transactions on Machine Learning Research (07/2024)
Probability flow ODE The probability flow ODE [Sog et al, 2021b defines invertible sampling function for a distribution Pg po theoretically the same as p Tt is given by
1 dx - -9 50(())s((x,)))))))))))))))) ()
(4)
The absence of Brownian motion makes it possible to solve this ODE in both directions of time. To compute the tog-probabilite of an image x, we map x(0) = x to its corresponding noise image x(T). Under the framework of neural ODEs (Chen al.) 20188. the log-probability given by the log-probability of x(T) under (0, I) plus a normalizing factor accounting for the change in density through time:
log ODE (x(O)) log (((()) + L/ V (((((,))
for x(0) = x. Although tractable to evaluate with an ODE solver, this log-probability function is compu- tationally expensive, requiring hundreds to thousands of discrete ODE time steps to accurately evaluate. Additional time and memory costs are incurred by backpropagation through the ODE and Hutchinson- Skilling trace estimation of the divergence.
DDE Song et al. (2021a) proved that if sp(x,t) = Vx (x, t) for all Equivalence of pg SDE and p8 E [0, [] and pr = n, then almost all x E R and Pa ODE
= Pa SDE Pdata In our work, we assume that sx(x, x Vx Po t) for [0, and that pr (O,,,), so that pg 2ODE pe SDD Pdata. This assumption empirically performed well in previous work that appealed to pg ODE the exact probability distribution of the diffusion model [Feng et al. 2023 Song al. 2021b).
3.2.1 Evidence lower bound
In lieu of exact log-probability function, Song et al. (2021a) derived evidence lower bound bşDE for p p such that (x) < log (x) for any proposed image x. Essentially, this lower bound corresponds to how well the diffusion model is able to denoise given image: an image with high probability under the diffusion model is easy to denoise, whereas a low-probability image is difficult.
The lower bound, or the negative "denoising score-matching loss" Song et al, 2021a), is defined as
/ ((1))(())d SDD (x) = Exor(x xx) [log क (x'))(x))))))))))))) - 2 by
1
(6)
where
2
h(t)=
Epoc (x)) []so(x(,)) xx logoopoo xx) ||V log x)11
g(t)2 Vx f(x(,t)
(7)
(x x) denotes the transition distribution from x(0) = to x(() x'.
x'. For drift coefficient that is linear in x, this transition distribution is Gaussian: Por x) = N(x';a (I) This means that the Poe log Eq. is efficient to compute since we can evaluate by adding Gaussian noise x without having to solve an initial-value problem as with the ODE. In fact, Eq. 0 closcly related to the denoising scoro-matching objective used to train diffusion models (Song et al., 2021b).
gradient Vxr log
x) is directly proportional to the Gaussian noise that is subtracted from to get x.
Intuitively, can interpret Eq. Gas associating an image's probability with how well the score model Sy could denoise that image if it underwent diffusion. This is represented by the first term in h(tt) (Eq. 7 To assess the probability of an image x, we perturb it with Gaussian noise lo get x' and then ask the score model to estimate (x t) accurately estimates the noise, then ||sa(x', t) V log Pox |
the noise that was added. is small, and the value of 1/DE (x) becomes larger. The remaining terms h(t) are normalizing factors independent of f and (xpr(x)))) [log xxx') accounts for the probabilities of the noise images x(T) that could result from being entirely diffused.
5

---

Published in Transactions on Machine Learning Research (07 (07/2024)
4 Method
Inspired by previous theoretical work (Song et al. 2021a), we apply 6i as an efficient surrogate for the exact score-based prior in Bayesian imaging. In this section, we describe our approach for efficient inference of an approximate posterior given a score-based prior.
4.1 Approximating the posterior with VI
Given measurements y E R M (with a known log- likeliidd function) and a score-based diffusion model with parameters ! as the prior, our goal is to sample from the image posterior pe (x y). Following VI, we optimize the paramcters of variational distribution approximate po(x y
Let go denote the variational distribution with parameters ø. We would like to optimize & to minimize the KL divergence from is to the target posterior:
Φ- arg min DEL (4o llpo(- y))
(8)
arg min Ex~ae [- ply x) log (x) + log go
(9)
do can be any parameterized tractable distribution. could be Gaussian distribution with a diagonal covariance that o lu ,oT]T, where μ IR and RD (a > 0) are the mean and pixel-wise standard
deviation. de could also be RealNVP normalizing flow il is in DPI (Sun & Bouman 2021 h.
To circumvent the computational challenges of evaluating the prior term log Pg surrogate (x). This results in the following objective:
DD (X)()) replace it with the
arg min Ex plg x) -
-) (x) + log gx)))
(10)
We can also think of bODE as replacing the intractable log SDE in Eq. 8 Since - log
SDE < -bSDE, our
surrogate objective minimizes the upper bound of the KL divergence involving p SDD
4.2 Implementation details
Evaluating (x) The formula for (x) (Eq. 6 contains a time integral and expectation over Par (x' (2021a), we use importance sampling
x) that can be estimated with numerical methods. Following Song al.
with time samples & ~ p(t) for the time integral and Carlo approximation with noisy images x' ((((t),, II) for the expectation. The proposal distribution p(t)
B0FZ was empirically verified
git)²
result in lower variance in the estimation of (x)F(()) (Song et al. 2021a We provide the following formula
used in our implementation, which estimates the time integral with importance sampling and the expectation with Monto-Carle approximation, for reference:
Nx
N
N₂
SDE (x) AJ be
Σ
loga (x))
23(0)
Σ
so(x)))
Zij
g(c,) Vxy f(xg)
(11)
No
N
B(t.)
((t.)
1m
s.t. 4 ~ p(t,, Bay ~ NOO, I), x1, = o 8(4)) ((a 2 Vi= L,..,, N:.j
No.
N₁ time samples and N₂ noise samples are taken to approximate the time integral and expectation over Por x), respectively. In our experiments, N₂ N₂ 1. We find that increasing the number of time and noise samples does not efficiently decrease variance in the estimated value of byDE (x).
Optimization We use stochastic gradient descent to optimize o, Monte-Carlo-an Carlo-approximatinn the expecta- tion in Eq. 10 with a batch of x ~ 4a- Estimating (x) has higher variance than estimattig log (ODE(x) For instance, in Fig. 4 (5P (x) with N₁ = 2048, N₂ = shows higher variance than log ODE (x) with 16 trace estimators. A lower learning rate can help mitigate training instabilities caused by variance. For example, in Fig. 3b the learning rate with the exact prior was 0.0002, while the learning rate with the surrogate prior 0.00001. Please refer to Appendix B for more optimization details.

---

Published in Transactions on Machine Learning Research (07/2024)
Optimization Time
Optimization Memory
Image-Restoration Quality
126 hrs
142 GB
Time [hr]
100
Exact Surrogate
Memory [GB]
100
SSIM SSX00
Surrogate
Exact
50 hrs
50
31
GD
44 GB
0.95
19 mins.
19 mins.
45 mins.
hrr.
16 GB
0.90
G2:BB
2 2 2
2 2
2
2
2
2
2 Dimensionality
Dimensionality
Dimensionality
Figure 1: Computational efficiency of proposed surrogate prior vs. exact prior. For each image size, we estimated posterior of images conditioned on 4x-accelerated MRI measurements of knee image, using a 4x NVIDIA
Gaussian distribution with diagonal covariance as the variational distribution. The hardware RTX A6000. The surrogate prior allows for variational inference of image sizes that are prohibitively large for the exact prior. For image sizes supported by the exact prior, the surrogate improved total optimization time by over 120x while using less memory and scaling better with image size. "Image-Restoration Quality" verifies that optimization with the surrogate was done fairly, as the PSNR and SSIM of the converged postcrior (averaged over 128 samples) are at least as high as with the exact prior.
5 Experiments
We validate our approach the tasks of accolerated MRI, denoising, reconstruction from low spatial fre- quencies, and black-hole interferometric imaging. See Appendix A for details about the forward models.
5.1 Efficiency improvements
In Tab. [ and Fig. we quantify the efficiency improvements of the surrogale prior for an accelerated MRT task at different image resolu- tions. We drew a test image from the fastMRI knee dataset (Zbontar iet al. 2018) and resized it to 16 x 16, 32 32, 64 x 64, 128 x 128, and 256 x 256. For each image size, we trained a score model on images of the corresponding size from the fastMRT training set of single-coil knee scans. We then optimized Gaussian distribution with diagonal covariance lo approximate the posterior. The batch size was 64 for the surrogate and 32 for the exact prior (a smaller batch size was needed to fit 64 x 64 optimization into GPU memory). Convergence was de- fined via a minimun acceptable change in the estimated posterior mean between optimization steps.
Image size
Surrogate
Exact
16 x 16
0.029
19.5
32 x 32
0.038
41.9
64 x 64
0.090
123
128 x 128
0.294
N/A
256 x 256
1.115
N/A
Table 1: Iteration time [sec/step]. A step of gradient-based optimiza- tion of the variational distribution is two to three orders of magnitude faster with the surrogate prior.
We find at least two orders of magnitude in time improvement with the surrogate prior. Tab. I compares the iteration time between the two priors. Fig. compares the total time it takes to optimize the variational distribution. The surrogate also improves memory consumption, which in turn enables optimizing digner-dimensional posteriors. Following standard practice, we just-in-tiime (JIT) compile the optimization step to reduce time/step at the cost of GPU memory. Fig. [ shows how the surrogate prior significantly reduces memory requirements and scales better with image size. The exact prior could only handle up x before exceeding GPU memory (we tested on 4x 48GB GPUs). While memory could be reduced with smaller batch size, this would make optimization more time-consuning, On the other hand, our surrogate prior supports much larger images, MRI with Causian-approximate posterior. This type of principled inference of high-dimensional image posteriors was not possible before with the exact score-based prior.
demonstrate in Fig, 2 for 256 x 25d
LLarger images may be feasible, but their larger memory footprint might restrict the possible batch and complexity of the variational distribution.
7

---

Published in Transactions on Machine Learning Research (07/2024)
Observed y
Posterior Samples p(x y)
Original x'
16x-accel. MRI 256x256
4x-accel.
MRI
256x256
High-dimensional Bayesian inference with surrogate boseebbaedd prior. Here show posterior samples for accelerated MRI of x 256 knee images, approximated via variational inference with surrogate score- based prior. The first row shows reconstruction from 16x-reduced MRI measurements. The second row shows reconstruction given more R-space measurements, i.e., 4x-reduced MRI. Bayesian imaging at this image resolution computationally infeasible with the previous ODE-baser approach (Feng et al. 2023). Our proposed surrogate enables efficient yet principled inference with diffusion-model priors, resulting inferred posteriors where the true image is within three standard deviations of the posterior mean for 96% and 99% of the pixcls for 16x- and 4x-acccleration, respectivcly.
Figure 2:
Original
Observed
Original
Observed
Mean
Std. Dev.
True
Mean
Std. Dev.
Mean
Std. Dev.
(Exact) (Surrogate)
Surogate
Est. Est.
0.01
0.12
0.01
Exact
CelebA denoising
CIFAR-10 denoising
(a) Ground-truth (Gaussian) posterior.
(b) Complex posteriors.
Figure 3: Estimated posteriors under surrogate vs. exact prior. For each task, the variational distribution is RealNVP, and the score model is the same between both prior functions. (a) Both prior functions help recover the correct (Gaussian) posterior. The score-based prior was trained on samples from a known Gaussian distribution (originally fit to 16 x 16 face images), and the measurements spatial frequencies of test image from the prior. Since the prior and likelihood are both Gaussian, we know the ground-truth Gaussian posterior. (b) Estimated posteriors for (i) denoising a CelebA image and (ii) denoising a CIFAR-10 image. Std. dev. is averaged across the three color channels. The scoro-based prior was trained on CelebA in (i) and CIFAR-10 in (ii). Both prior functions result in comparable image quality; visual differences appear mostly in the image background.
the lowest 6.25%
5.2 Posterior estimation under surrogate exact
The surrogate prior 6 5DDE may not be an identical substitute for the exact prior log ODE, Importantly, Pg
though, we verify in Fig. Ba that both the surrogate and exact prior recover a ground-truth Gaussian posterior derived from Gaussian likelihood and prior. The variational distribution is a ReaINVP. The score model (used by both the surrogate and exact prior) trained samples from the known prior.
8

---

Published in Transactions on Machine Learning Research (07/2024)
Nonetheless, the surrogate could result in different locally-optimal variational posterior, particularly if the posterior leads to many local minima. in the variational objective. Fig. 3b compares posteriors (with unknown true distributions) approximated by a RealNVP under the surrogate versus exact prior. For each task (CelebA denoising and CIFAR-10 denoising), both prior functions used the same trained score model. We observe in these comparisons that most. of the differences appear in the image background and that both priors result in plausible mean reconstruction and uncertainty.
Visualizing the bound gap throughout optimization helps shed light on why the two priors converge to different solutions even with the same underlying score model. Fig. | shows probabilities of samples generated by de (in this case a RealNVP) as optimization progresses. At cach checkpoint of qo: we plot (x) versus log logrt (x) for samples x ~ do coming from both the exact and surrogate optimization of gs. We find that the surrogate is a valid bound for the ODE log-density: (x) S log (x) for all ~ go: except for some outliers due to variance of (x)) However, optimization follows a different trajectory depending on the prior. With the surrogate, samples x ~ 9e tend toward a region where the bound gap is small. Meanwhile, the exact prior follows loss landscape whose struttuee appears to be independent of the lower bound. Note that samples from qe optimized under the exact prior obtain higher values of (x) than samples obtained under the surrogate. Fig. 4 suggests that gradients under the surrogate tend to push the distribution along the boundary of equality between PyDE and log po path taken through gradient descent and subsequently the converged solution.
ODE This constrains the
Step 100
Step 500
Step 20000
7500
5930
Stop 20K
Exact (log p onR)
5533
Step 500
2000e
6366
Step 100
-Tececo -200383 -190003
-10000
-5000 log OnF
6993
1000
-5000
s020
O0 (X)
logpi COI
lag 10
200 (x)
Step 5000
Step 25000
Step 1000000
Plotting trajectories same scale
1500
0000
Step 1M
Surrggate
(0 }
3000
1g00 (x)
X 0000
5939
stop
Step 25K
2200
3000
1990
-3033
2000
1g (x)
2033
00E
3332
3338
10g,9 (x)
GDR
1000
log,s
2000
[x]
6000
log,5 (x (xc
Figure 4: bg (x) vs. logpoD (x) for samples x ~ 9o as optimization of á progresses. The task is from
Fig. 3b( For each plot, we took 128 samples x ~ and performed 20 estimates each of log log (x) and 03 D (x) (approximated with N₂- = 2048 for reduced variance). The density map is a KDE plot of all 128. 20 2560 values; the 128 scalter points represent the mean estimate for each x. The black line indicates perfect agreement between 1.DE (x) and log pg (x).
(x). We expect all points to lie below this black line for IŞDE to be lower bound. We find that (x) S log (xE (up to variance error), but the optimization progresses differently depending on the prior. Gradients under the surrogate push 4o (x) along the black line to increase (x) without exceeding ( Optimization under the exact prior proceeds more freely, although eventually achieves higher (x) at convergence. This visualization may help explain differences in the posterior estimated with the surrogate vs. exact prior.
5.3 Quality of posterior estimation with Bayesian approach vs. diffusion- approaches
A popular class of methods is diffusion-based approaches (discussed in Soc. 2.2 which attempt to sample from the posterior by incorporating measurements throughout the reverse diffusion process of the diffusion model that has been trained on images from the prior. Such approaches provide fast conditional sampling,
9

---

Published in Transactions on Machine Learning Research (07 (024)
meas. weight
KL weight
lowest
oracle
highest
SDE+Proj
SDE Pro
2.5
True
prior
03
DPI + exact
lowest
oracle
highest
DPS Score-ALD
Score: ALD
lowest
oracle
highest
DPI + surr.
DP5
Our (DPI SUIT.
6.00
18.3
Figure 5: Comparing our VI approach with a surrogate score-based prior to basclines on a bimodal poste- rior. In this example, the prior bimodal mixture-of-Gaussinn and the likelihood Gaussian, making the posterior a bimodal mixture-of- Gaassssass (shown in "True"). Assuming access the true prior score function, tested how well each method recovers the true posterior. Diffusion-based methods depend on hand-tuned meas. weights. Even the meas. weight giving the best KL divergence ("oracle") does not rival using our trperparameteer-ree VT approach (*DPT + surr."). Note that this "oracle" weight would not be accessible in practice, as it is determined by comparing to the ground-truth posterior. Diffusion-based base- lines either (1) incorrectly place equal weight on both posterior modes or (2) miss one the modes. DPI with either the surrogate or the exact score-based prior recovers the relative weights of both modes. (KL vs. meas. weight) Regardless of hyperparameters, diffusion- based methods do not reach our KL divergence.
KL (J)
time/opptimization step (1)
DPI + exact
0.030
130 ms
Ours: DPT + surr.
0.037
22 ms
DPS (oracle) 0.064
Scoro-ALD (oracle) 0.10
SDo+Proo (oracle) 0.12
two-component Gaussian mixture model (GMM) was fit estimated samples to obtain PDF. "Ours" achicves much lower KL div. (i.c., reverse KL from estimated posterior to true posterior) than diffusion-based baselines at their best. Time/step for DPI oplimization is lower with our surrogate than with the exact score-based prior without sacrificing much accuracy.
Table 2: Quantitative results for Fig.
but they may severely mischaracterize the posterior. Although our approach also approximates the postcrior (due to using both surrogate prior and variational distribution), we find that being grounded in variational inference helps us obtain more accurate posterior and images that more accurately reflect the true image than diffusion-baase methods that make ad-hoe approximations of the posterior.
In the following experiments, we compare to three diffusion--assd baselines: SDE+Proi (Sng el al. 2022 Score-ALD (Jalal et al. 2021), and Diffusion Posterior Sampling (DPS) CChung et al. 2023). SDE Proj projects images onto a mcasurement subspace. Score-ALD and DPS strongly approximate the postcrior throughout reverse diffusion. All baselines involve a mcasurement-wette weight byperparameter. Our approach is DPI with the surrogate prior (i.e., using RealNVP as the variational distribution).
10

---

Published in Transactions on Machine Learning Research (07 (07/202) 2024)
Curs Sarogate)
P5NK
0.8
DP5
TV Regular Iegulaaiatiion
20
0.5
Acceleration factor
Acceleration factor
(a) Image-restoration metrics.
O-Filled Recon, (16x-accel.)
Ours SDE+Proj (DPI Surr)
TV Reg.
Scare-ALD
DPS
Original
(b) Example image reconstructions for 16x acceleration.
Figure 6: Accelerated MRI (a) For each accel. factor (4x, 8x, 16x), we estimated posteriors for ten knee images measured at that accel. rate. For each method, we computed the average PSNR and SSIM of 128 estimated posterior samples. The line plot shows the average result across the ten tasks; the shaded region shows one std. dev. above and below the average. (b) An example of 1cx-accel. MRI. The cropped region exemplifies how diffusion-based baselines hallucinate more features than necessary. (a) and (b) are evidence that principled Bayesian approach can get closer to the true image than previous unsupervised methods.
5.3.1 Accuracy of posterior
A simple 2D example illustrates the accuracy gap between the diffusion-based baselines and our VI approach. We consider a bimodal posterior: the prior is a bimodal ceatuureof-GGausiian and the forward model linear projection with Gaussian noise, making the posterior a bimodal mixture-of-Gaussians This setup lets us evaluate with a true posterior and over a reasonable space of hyperparameters for baselines.
We tested how well each of the methods could recover the ground-truth posterior when given the true score function of the bimodal prior (thus avoiding potential error caused by learning the score funetion). Fig. 5 shows estimated probability density functions (PDFs) for the evaluated methods. None of the diffusion-based baselines correctly recover the bimodal posterior for any hyperparamcter value. In particular, they struggle to find the correct balance between the two posterior modes - in the best casc, they incorrectly place equal weight on each mode; in the worst case, they only recover one mode. Our VI approach with the exact or surrogate score- coore-based prior recovers both modes in correct proportion. As shown in Fig, and Tab. 2 even the best KI. divergence obtained by the diffusion-basee baselines does not rival that of VI with scoro-basod prior. We emphasize that the hyperrarameter values resulting in the "best" KL for diflusion-based methods can only be found with knowledge of the ground-truth, which inaccessible in real-world scenarios. Tn contrast, method automatically finds better KL divergence by following the Bayesian formula.
5.3.2 Image-reconstruutton quality
It would be reasonable to assume that the diffusion- based baselines, though less principled, may lead to better visual quality than Bayesian approach. However, we find that in addition to providing more reliable uncertainty, our approach achieves higher-fidelity reconstructions. We note that similarity to a ground-truth image does not. indicate a correct posterior. Still, for good prior, might be desirable for posterior samples to accurately recover the original image.
We performed multiple MRI tasks different acceleration rates and compared our approach the diffusion- based baselines, well total variation (TV) baseline. Implementations and hyperparameter settings for
11

---

Published in Transactions on Machine Learning Research (07/2024)
SDE+Proj and Score-ALD were provided by Song al. (2022) For DPS, we followed the implementation of Chung et al. (2023) and performed a hyperperameter search on an 8x-acceleration test image to find the optimal PSNR. For the TV baseline, we performed DPI with TV regularization with a regularization weight of 10 instead of the surrogate score-based prior.
We simulated MRT at three acceleration factors for ten test images, resulting thirty posteriors to be estimated. As baseline implementations do not assume measurement noise, we gave the baselines noiseless measurements and set near-zero mcasurement noise for our method. The test images were randomly drawn from the fastMRI dataset and resized to 64 x 64. The score model Se was trained on 64 x 64 fastMRI knee images and stayed fixed across all methods.
Our method achieves a marked improvement. in PSNR and SSIM (Fig. 6 Across all acceleration factors and diffusion- based baselines, our method improves PSNR by between 2.7 and 8.5 dB. Even though our method and the diffusion-based methods all use the same score model, restoration quality depends on how the prior is used for inference; whereas baselines loosely approximate the posterior and involve hyperparameters, our approach treats the diffusion model as standalone prior in Bayesian inference. Furthermore, Fig, 6 confirms that diffusion-model prior far outperforms traditional regularizer like TV..
5.4 Application case study: black-hole interferometric imaging
Computational imaging distant black holes is possible with very-long-baseee interferometry (VLBI) (Thompson et al.) 2017), by which a network of telescopes measures spatial froquencies of the sky's image. The Event Horizon Telescope (EHT) Collaboration used this technique image the black hole al the center
of the galaxy M87 EHTC 2019 2024 and the black hole at the center of our galaxy, SgrA* (EHTC 2022).
Imaging from EHT measurements requires prior since telescope observations are corrupted and sparsely sample the low spatial-frequencc space. Previously, the EHT Collaboration imposed hand-crafted regu- larizers to obtain the M87* and SgrA* images. Here we demonstrate the applicability of our method to black-hole imaging from EHT measurements. In contrast to the regularized seeiiiummlikeelhooe (RML) approaches used by the EHT Collaboration (EHTC 2019 2024 2022). our approach hyperparameter- and provides rich data-driven posterior.
Black-hole imaging exemplifies an application that calls for theoretically sound imaging approach. A set of telescope measurements must be carefully analyzed not only because it is expensive to collect but also because the inferred confidence intervals inform downstream scientific analysis. Therefore, an imaging approach like ours that estimates the posterior in a principled way is often desirable, even if it requires more computational resources.
5.4.1
Interferometry background
In VLBI, each pair of telescopes i,j provides a Fourier measurement called a visibility Viy (van Cittert, 1934 Zernike 1938). To overcome ude and phase errors, we form robust closure quantities out of the measured visibilities
[Blackburn et al., 2020). A closure phose is formed from a triplet of telescopes i, k and robust to phase errors. A log closure amplitude is formed from a set of four telescopes i, k, and robust to amplitude errors. They are given respectively as
(ptjlluke [uax [uae|lugel
Vijk 2 (UiyUyUUk) and
logea Higke
log
(12)
For the inverse problem, measurements all the non-redundant closure phases and log closure am- plitudes, which have a nonlincar forward model and Gaussian thermal noisc. Fig. 7a illustrates the VLBI measurements obtained by an EHT telescope array. These particular measurements were generated from an image of a black-hole simulation. The (u, v) coverage is sparse and constrained to low spatial frequencies. The "dirty" image, which is a naive reconstruction from the measured visibilities, is completely un-interpretable. The target image the original image blurred to the maximum resolution achievable by the EHT. Recon- structing an image beyond this intrinsic resolution is considered super-resolution.
12

---

Published in Transactions on Machine Learning Research (07 (07/202) /024)
(u, M/ Coverage
Dirty
Target
Original
Sample
Mean
Std. Dev.
Brghtnee Imperater (10*
Seighnws
(a) Simulated VLBI measurements.
(b) Estimated image posterior.
Figure 7: Black-hole imaging. (a) We simulated EHT VLBI measurements with realistic noise of synthetic black-hole image. "(u, v) Coverage" shows which points in the complex 2D Fourier plane are measured by the EHT array of radio telescopes. "Dirty" is naïve reconstruction from the sparse measured visibilitios. "Target" represents the best- best-possible image reconstruction if the low spatial frequencics (up to the intrinsic resolution of the EHT) were to be fully and perfectly measured. "Original" is the actual underlying image that generated the simulated measurements, (b) Using our pariational-inferens approach with a score- based prior trained on simulations of black holes, we approximated the image posterior. The postcrior samples recover the ring and brightness asymmetry of the original image. The pixel-wise mean and std. dev. also shown. The std. dev. indicates regions of uncertainty, such possible wisps inside and outside the thin ring.
5.4.2 Simulated-data example
We used our method to estimate an image posterior given the simulated measurements visualized in Fig. 7 The variational distribution is RealNVP, and the score-based prior trained on 64 64 images of fluid- flow simulations of black holes (Wong et al. 2022 Fig. 7b shows sample from the estimated posterior, well as the mcan and standard deviation. By using scorc-based prior trained on black-hole simulations, we are able to reconstruct images that resemble these detailed simulations while still fitting the measurements.
Our proposed surrogate prior made it possible to approximate this posterior in a reasonable amount of time. Optimization with the exact prior (Feng et al. 2023) would have taken 151 GB of GPU memory and 1.8 minutes per iteration. In contrast, optimization with the surrogale prior used 46 GB of GPU memory and 177 milliseconds per iteration on the same hardware (4x 48GB NVIDIA RIX A6000).
6 Conclusion
We have presented surrogate function that provides efficient access to score-based priors for variational Bayesian inference. Specifically, the evidence lower bound (SDE(() < log logr (X) serves as proxy for the log-prior of an image in the Bayesian log-posterior. Our experiments with variational inference show at least two orders of magnitude in runtime improvement and significant memory improvement over the ODE- based prior. We also establish that a principled approach like ours outperforms diffusion- based methods on posterior approximation and image restoration, evidence that following a traditional Bayesian approach results in more reliable image reconstructions. This advantage crucial for applications that call for a reliable image posterior while still leveraging the expressiveness of score-based prior.
Broader impacts There are many applications, such as in the medical and astronomical domains demon- strated in this paper, in which there is incomplete information in measurements that makes priors necessary to produce interpretable images. Our work proposes a way to efficiently incorporate rich diffusion-modee priors in a rigorous imaging framework. In general, however, generative models should be used with caution. They have been used to create deepfakes (Ricker et al. 2022) and may induce spurious hallucinations. Dif- fusion models have also been shown to memorize training data (Carlini al.) 2023 Somepalli al. 2023). in which one should be careful to preserve privacy when deploying diffusion model image prior.
13

---

Published in Transactions on Machine Learning Research (07 (07/2024)
Acknowledgments
The authors would like to thank Yang Song for his many technical insights and helpful feedback on the paper. BTF and KLB acknowledge funding from NSF Awards 2048237 and 1935980 and the Amazon AT4Science Partnership Discovery Grant. BTF supported by the NSF GRFP.
References
Alexandre Adam, Adam Coogan, Nikolay Malkin, Ronan Legin, Laurence Perreault-Levasseus Yashar Heza- veh, and Yoshua Bengio. Posterior samples of source galaxies in strong gravitational lenses with peore-based priors. arXiv preprint arXiv:2211.0381 2022.
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-34) -2,,6,,, 1982.
Omri Avrahami, Ohad Fried, and Dani Lischinski. Blended latent diffusion. ACM Transactions on Grophics (TOG), 42(4):1-11, 2023.
Lindy Blackburn, Dominic W. Pesce, Michael D. Johnson, Maciek Wielgus, Andrew A. Chael, Pierre Chris- tian, and Sheperd S. Doeleman. Closure statistics in interferometric data. The Astrophysical Journal, (94(1):31, may 2020. doi: 10.3547/1338-4357/bbbbb8 URL https://dx.dd: doi. oroi.org/100.. 11..3387//533-43577/ ab8469
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the Americon statistical Association, 112(518):859-88777 2017.
Steve Brooks, Andrew Gelman, Galin Jones, and Xiao-Li Meng. Handbook of markov chain monte carlo. CRC press, 2011.
Emmanuel Candes and Justin Romberg. Sparsity and incoherence in compressive sampling. Inverse problems, 23(3);969, 2007.
Nicolas Carlini, Jamie Hayes, Milad Nasr, Matthew Jagielski, Vikash Sehwag, Florian Tramer, Borja Balle, Daphne Ippolito, and Eric Wallace. Extracting training data from diffusion models. In 32nd USENIX Security Symposium (USENIX Security 23), pp. 5253-5270, 2023.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. NewrIPS, 31, 2018.
Jooyoung Choi, Sungwon Kim, Yonghyun Jeong, Youngjune Gwon, and Sungroh Yoon. llv: Conditioning method for denoising diffusion probabilistic models. In ICCV. IEEE, 2021.
Hyungjin Chung and Jong Chul Ye. Score-based diffusion models for accelerated mri. Medical Image Analysis, 80:102479, 2022.
Hyungjin Chung, Bycongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. arXiv preprint arXiv:2206.0094/ 2022a.
Hyungjin Chung, Bycongsu Sim, and Jong Chul Ye. Conne-laerrddifisse--ste Accelerating conditional
diffuusio models for inverse problems through stochastic contraction. In Proceedings of the IEEE/CVE Conference Computer Vision and Pattern Recognition, pp. DPA13-1242,, 2022b.
Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul Ye. Diffu- sion posterior sampling for general noisy inverse problems. In The Fleventh International Conference Learning Representations, 2023. URL httpp://
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score- score-based generative modeling. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (ods.), Advances in Newral Information Processing Systems, volume 34, pp. 17695-17709, Curran Associates, Inc., 2021. URL D/tp. www.see.eeiiiioee. cttpp/ssccppppprraeeeeeeee files/
14

---

Published in Transactions on Machine Learning Research (07/2024)
M. Delbracio, H. Talchei, and P. Milanfar. Projected distribution loss for image enhancement. In 2021 IEFE International Conference Compulational Pholography (TCCP), pp- 1-12, Los Alamitos, CA. USA, may 2021. IEEE Computer Society. doi: 10.1 /tt/:///CCP11111221..11220 1109/ICCP51581 Lececonputersocietyyyoo 110. 2021
URL https://doi.
9466271
Mauricio Delbracio and Peyman Milanfar. Inversion by direct iteration: An alternative to denoising diffusion for image restoration. arXiv preprinl Xrx:223303. 11435, 2023.
Edward H Egelman. The current revolution in cryo-cm. Biophysical journal, 110(5):1008-1012 2016.
Event Horizon Telescope Collaboration EHTC. First sagittarius a* event horizon telescope results. iii. imaging of the galactic center supermassive black hole. The Astrophysical Journal Letlers, 930(2):L14, may 2022, doi: 10.3847 /2884772441-221////00 ac6429. URL htttps://dd doi. wttp:///....88/21111211111/////1
The Event Horizon Telescope Collaboration EHTC. First m87 event horizon telescope results. iv. imaging the central supermassive black hole. The Astrophysical Jourmal Letters, 875(1):L4, apr 2019. doi: 10.3847/ 2041--213/ 1al1-8823/abbe85. ab0085. URL https://dx doi org/ 000..033877/00111133340660000000
The Event Horizon Telescope Collaboration EHTC. The persistent shadow of the supermassive black hole of m 87-i. observations, calibration, imaging, and analysis. Astronomy & Astrophysies, 681:A79, 2024.
Berthy T Feng, Jamie Smith, Michael Rubinstein, Huiwen Chang, Katherine L Bouman, and William T Freeman. Score- based diffusion models as principled priors for inverse imaging. In International Conference on Computer Vision (ICCV). IEEE, 2023.
Alexandros Graíkos, Nikolay Malkin, Nebojsa Jojie, and Dimitris Samaras. Diffusion models as plug-and- play priors. In Siairty--iath Conference on Neural Information Processing Systems, 2022. URL https:
//axxiv. perxivv.oggpdtt. org/påf/ http://ff/222.....e/ 220//2226660 090...099012..00 pdf
Stephen F Gull and John Skilling. Maximum entropy method in image processing. In Ice proceedings f (commumications, vador and signal processing), volume 131, pp. 646-646 659. IET, 1984.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information. Processing Systems, 33:6840-666 (0006685,444 2020.
Satoshi Tizuka, Edgar Simo-Serra, and Hiroshi Tshikawa. Globally and locally consistent image completion. ACM Transactions on Graphics (ToG), 36(4):1 14, 2017.
Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros G Dimakis, and Jonathan I Taamir. Robust compressod sensing mri with deep generative priors. NewrIPS, 2021.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- generative models. arXiv preprint XrXiv:2206. 000064, 2022.
Bahjat Kawar, Michael Flad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. In Advances in Newral Information Processing Systems, 2022.
Ivan Kobyzov, Simon JD Prince, and Marcus A Brubaker. Normalizing flows: An introduction and review current methods. IEEE transactions on pattern analysis and machine intelligence, 43(11):3964-399 39743 2020.
Axel Levy, Gordon Wetzstein, Julien NP Martel, Frederic Poitevin, and Ellen Zhong, Amortized inference for heterogencous reconstruction in cryo-em. Advances in Neural Information Processing Systems, 35: 13038-1304B 2022.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210, 1..:2220.002777, 2022.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
15

---

Published in Transactions on Machine Learning Research (07/2024)
Sebastian Lunz, Ozan Öktem, and Carola-Bibian Schönlieb. Adversarial regularizers in inverse problems. Aduances in neural information processing systems, 31, 2018.
Morteza Mardani, Jiaming Song, Jan Kautz, and Arash Vahdat. A variational perspective on solving inverse problems with diffusion models. arXiv preprint (XXx:2305044311 2023.
Jianwei Miao, Tetsuya Ishikawa, Qun Shen, and Thomas Earnest. Extending x-ray crystallography to allow the imaging of noncrystalline materials, cells, and single protein complexes. Annu. Rev. Phys. Chem 59: 387 410, 2008.
Subhadip Mukherjee, Sören Dittmer, Zakhar Shumaylov, Scbastian Lunz, Ozan Öktem, and Carola-Bibian Schöulieb. Learned convex regularizers for inverse problems. urXiv preprint arXiv:2008. 02833, 2020.
Ramesh Narayan and Rajaram Nityananda. Maximum entropy image restoration in astronomy. Annual review of astronomy and astrophysics, 24(1):127-11 170, 1986.
Deepak Pathak, Philipp Krahenbuhl, Jeff Donahue, Trevor Darrell, and Alexei A Efros. Context encoders: Feature learning by inpainting. In Proceedings of the IEEE conference on computer vision and pallern recognition, pp. 2536 2544, 2016.
Jonas Ricker, Simon Damm, Thorsten Holz, and Asja Fischer. Towards the detection of diffusion model deepfakes. arXiv preprint wwXiv2220011571,, 2022.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 10684-106999 10695, 2022.
Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022.
Zakhar Shumaylov, Jeremy Budd, Subhadip Mukherjee, and data-driven convex-nonconvex regularization. arXiv preprint orXiw:2310.05812 2023.
Schönlieb. Provably convergent
Jascha Sohl-Dickstein Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning PMLR, 2015.
using noncquilibrium thermodynamics. In Int. Conf. Machine Learning, pp.
Gowthami Somepalli, Vasu Singla, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Understanding and mitigating copying in diffusion models. In A. Oh, T. Naumann, A. Globerson, K. Sacnko, M. Hardt, and S. Levine (eds.), Adoonces in Neural Information Processing Systems, volume 36, pp. 47783 47803. Cur- ran Associates, Inc., 2023. URL https: /ttroeee//rrocediingg neurips wwwrips.cccpapprr.... wwtppcpp..tlleeeeeeeeeeeeeeeeeee fitles/papee//22///Ei 35226611188000011222222337715511
Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoiinerse-guided diffusion models for inverse problems. In International Conference on Learning Representations, 2023a. URL https:
//openreview. www.:eeeoomddddggeeeeaaeaaa
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In NewrIPS, pp. 11895-11990 2019.
Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. In Proceedings of the Thirty-Fifth Conference on Uncertainty én Artificial Intelligence, UAT, pp. 204, 2019. URL
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score based diffusion models. In Thirty-Fifth Conference on Neural Information Processing Systems, 2021a.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poolc. Score- generative modeling through stochastic differential equations. In ICLR, 2021b. URL https: httppparereie..tt.ffoorrraaaaaip
16

---

Published in Transactions on Machine Learning Research (07/2024)
Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with http://wwww......eooooee///////a
score-based generative models. In ICLR, 2022. URL
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303014409 2023b.
He Sun and Katherine L Bouman. Deep probabilistic imaging: Uncertainty quantification and multi-modal solution characterization for computational imaging. In AAAT, pp. 2628 2238, 2021.
He Sun, Katherine L Bouman, Paul Tiede, Jason J Wang, Sarah Blunt, and Dimitri Mawet. alpha-deep probabilistic inference (alpha-dpi): efficient uncertainty quantification from exoplanet astrometry to black hole feature extraction. arXi preprint or arXiv:2211 10.22200008066, 2022.
A Richard Thompson, James M Moran, and George W Swenson. Interferometry and synthesis in radio astronomy- Springer Nature, 2017.
Muhammad Usman and Philipp G Batchelor. Optimized sampling patterns for practical compressed mri. In SAMPTA 09, pp. Poster-session 2009.
Pieter Hendrík van Cittert. Die wahrscheinliche schwingungsverteillng in einer von einer lichtquelle direkt oder mittels einer linse beleuchteten ebene. Physica, ((1-6):201 210, 1934.
Curtis R Vogel and Mary E Oman. Iterative methods for total variation denoising. STAM Journal Scientific Computing, 17(1):227 238, 1996.
George N Wong, Ben S Prather, Vedant Dhruv, Benjamin R Ryan, Monika Mošcibrodzka, Chi-kwan Chan, Abhishek V Joshi, Ricardo Yarza, Angelo Ricarte, Hotaka Shiokawa, et al. Patoka: Simulating electro- magnetic observables of black hole accretion. The Astrophysical Journal Supplement Series, 259(2):64, 2022.
Michael M Woolfson. An introduction to X-ray crystallography, Cambridge University Press, 1997.
Tianwei Yin, Zihui Wu, He Sun, Adrian V Dalca, Yisong Yuc, and Katherine L Bouman. End-to-end sequential sampling and reconstruction for mri. ar preprint or orXiv:1005006600 2021.
Jiahui Yu, Zhe Lin, Jimei Yang. Xiaohui Shen, Xin Lu, and Thomas S Huang. Free-form image inpainting with gated convolution. In Proceedings of the IEEE/CVF international conference on compuler vision, pp. 4471-44800 14400 2019.
Jure Zbontar, Florian Knoll, Amuroop Sriram, Tullic Murrell, Zhengnan Huang, Matthew J Muckley, Aaron Delazio, Ruben Stern, Patricia Johnson, Mary Bruno, et al. fastmri: An open dataset and benchmarks for accelerated mri. arXiv preprint arXiv:1811.08821 2018.
Frederik Zernike. The concept of degree of coherence and its application to optical problems. Physica, 5(8): 335-795, 1938.
Jinwei Zhang, Hang Zhang, Alan Wang, Qihao Zhang, Mert Sabuncu, Pascal Spincemaille, Thanh D Nguyen, and Yi Wang. Extending loupe for k-space under-sampling pattern optimization in multi-coil mri. In International Workshop on Machine Learning for Medical Image Reconstruction, pp. 91 101. Springer, 2020.
Kaihao Zhang, Wenqi Ren, Wenhan Luo, Wei-Sheng Lai, Björn Stenger, Ming-Hsuan Yang, and Hongdong Li. Deep image deblurring: A survey. International Journal of Computer Vision, (30(9))210322133 2022.
Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. arXio preprinl ar Xrx::22004. 13902, 2022.
Daniel Zoran and Yair Weiss. From learning models of natural image patches to whole image restoration. In ICCV, pp. 479-40 486. IEEF, 2011.
Daniel Zoran and Yair Weiss. Natural images, gaussian mixtures and dead leaves. NewrIPS, 25, 2012.
17

---

Published in Transactions on Machine Learning Research (07 (224))
A Forward models
In this appendix, we describe the forward models of the inverse problems explored in the main text: ac- celerated MRI, denoising, and reconstruction from low spatial frequencies ("deblurring"), These tasks have forward models of the form
== Ax e~ E~N(0,0
(13)
with the corresponding log- oog-likel function
logply x) x
20 lly Ax
(14)
A.1 Accelerated MRI
Accelerated MRI collects sparse spatial-frequency measurements k-space an underlying anatomical image. As the acceleration rate increases, the number of measurements decreases. The forward model can be written as
y Mox e~ oz)
(15)
where x E CP and y E CM, F denotes the 2D Fourier transform, and M {0, 13 is a binary sampling mask that reduces the number of non-zero measurements M << D. Often oy assumed be small (c.g.. corresponding to an SNR of at least 30 dB). We use Poisson-disc sampling (Usman & Batchelor 2009) for the sampling mask. 16x a6cecccelr,, for example, corresponds to sampling mask with only 1/16 nonzero elements.
Experimental setup In our experiments, we assumed that oyy is 0.05% of the DC (zero-frequency) amplitude. This corresponds to a maximum SNR of 40 dB. The only exception for comparison to baselines (Fig. 6), since baseline methods do not account for measurement noise. In this case, we let |oyl = 0.1% of the DC amplitude along the horizontal direction of the true image, which amounts to very low level of noise.
A.2 Denoising
The denoising forward model is simply
-=x+6, C~,
(16)
where x E R,, and oy determines the level i..d. Gaussian noise added to the clean image to get E R
Experimental setup In our presented experiments on denoising, By = 0.2, which is 20% of the dynamic range.
A.3 Deblurring
We refer to the lask of reconstruction from the lowest spatial frequencies as deblurring. The forward model is given by
y Dx Dx + E e~ N=
(17)
where x E CD, E CM, and D E MXD performs a 2D discrete Fourier transform (DFT) with only the first. M DFT components.
Experimental setup In our presented experiments on deblurring, the measurements are the lowest 6.25% of the DFT components, and loyl
18

---

Published in Transactions on Machine Learning Research (07 (07/2024)
B
Experiment details
For the sake of reproducibility, we detail the experimental setup behind each figure. Our code will be made publicly available. Some common implementation details are that the exact prior (log (g was always estimated with trace estimators. The RealNVP had 32 afine-coupling layers unless stated otherwise.
B.1 Variational distributions
We first describe the two types of variational distributions considered in our experiments: a RealNVP normalizing flow and a multivariate Gaussian with diagonal covariance.
RealNVP The architecture of the RealNVP is determined by the number of affine-coupling layers and the width of each layer. For images up to 64 x 64, we use 32 affine-cupling layers and set the number of hidden neurons in the first layer to 1/8 of the image dimensionality (o.g., 32 -22322222 3/8 for 32 x 32 RGB images). We use batch normalization in the network. Please refer to the original DPI (Sn & Bouman, 2021 Py PoTorch implementation for details on the architecture. Our implementation is an adaptation of this codebase in
JAX.
Gaussian Other experiments the variational family. In this case, the parameters the mean image and the pixcl-wise standard deviation. We initialize the mean at 0.5 and the standard deviation at 0.1 for all pixels. To sample, we take the absolute value of the standard deviation and construct the diagonal covariance matrix.
multivariate Gaussian distribution with diagonal covariance matrix
B.2 MRI efficiency experiment (Fig. Tab. 1
Score model For each image size, the score model was an NCSN++ architecture with 64 filters in the first layer and trained with the VP SDE with Bmin 0.1, 9max 10.
Variational optimization For each task (i.e., each image size and prior), the variational distribution was multivariate Gaussian with diagonal covariance. The batch size was 64, learning rate 0.0002, and gradient clip 1. A convergence criterion based on the loss value is difficult to define due to high variance of the loss (we used 1 time sample to estimate bxxx).. We defined convergence criterion based the change in the mean of the variational distribution. Specifically, every 10000 steps, we evaluated snapshot of the variational Gaussian and computed 5 = || Pprovll l| /prelll where Peurr and (prev are the current and previous snapshot means, respectively. was considered converged. Since convergence rate depends on the image size and the prior used, we set a different & for each task:
< E for some threshold two snapshots in a row, then the optimization
16 16 (surrogate): 0.002
32 32 (surrogate): 0.003
64 64 (surrogate): 0.005
128 x 128 (surrogate): E = 0.007
256 x 256 (surrogate): E = 0.009
16 x 16 (exact): s 0.0025
32 x 32 (exact): = 0.0027
64 64 (exact): 0.005
We were conservative in defining the convergence and checked that optimization under the surrogate actually achieved better sample quality than optimization under the exact prior (see Fig.
19

---

Published in Transactions on Machine Learning Research (07/2024)
Data The test image is from the fastMRI Zbontar et al. 2018 single-coil knee test dataset and resized to 64 64 with antialiasing.
B.3 256x256 MRI examples (Fig. 2
The 4x-accoleration result is from the efficiency experiment (Fig. and Tab. 1 on the 256 x 256 test image. The 16x-acceleration result came from a similar setup, where the variational distribution was Gaussian with diagonal covariance. Optimization was done with batch size of 64, learning rate of 0.00001, and gradient clip of 0.0002. We ran optimization for 270K steps (optimization for 4x-acceleration was done in 100K steps with the convergence criterion).
In the figure caption, we report that the true image is within three standard deviations of the inferred posterior mean for 96% and 99% of the pixels for 16x- and 4x-acceleration, respectively. This was computed based on the mean and standard deviation of 128 samples from the inferred posterior. We find the same result when using the exact mean and standard deviation of the inferred posterior: with respect to the inferred posterior, the truc image within three standard deviations of the mcan for 96.7% and 99.0% of the pixels for 16x- and 4x-acceleration, respectively.
B.4 Ground-truth posterior (Fig. 3a
Data
The mean and covariance of the ground-truth Gaussian prior were fit. with PCA (with 256 principal components) to training data from the CelebA dataset (Liu et al. 2015 The CelebA images were resized to 16 x 16 with antialiasing.
Score model The score model was based on the DDPM++ deep continuous archietecture of [onn et al. (2021b) with 128 filters in the first layer. It was trained with the VP SDE with 8min = 0.1 and Bmax 20 for 100K steps.
Variational optimization The variational distribution was RealNVP. Under the surrogate prior, opti- mization was done with learning rate of 0.00005 and gradient clip of 1. Under the exact prior, the learning rale was 0.0002 and gradient clip 1. Both priors used a batch size of 64.
B.5 32x32 image denoising (Fig. 3b
Variational optimization For both CelebA denoising (i) and CIFAR-10 denoising (ii), the variational distribution was RealNVP. Optimization under the exact prior was done with learning rate of 0.0002 and gradient clip of 1 for 20K steps. Optimization under the surrogate prior done with learning rate of 0.00001 and gradient clip of 1. For CelebA, the batch size was 64 and training was done for 1.72M steps (convergence was probably achieved carlier, but we continued training to be conservative). For CIFAR-10, the batch size was 128 and training was done for 550K steps.
Score model For both (i) and (ii), the score model had NCSN++ architecture with 64 filters in the first layer. For the CelebA prior, it was trained with the VP SDE with 8min 0.1 and 9max = 20 and with images that were resized without antialiasing. For the CIFAR-10 prior, it was trained with the VP SDE with 8min = 0.1 and 9max = 10.
Data The CelebA and CIFAR-10 images are both 32 x 32. The CelebA image was resized without an- tialiasing.
B.6 Bound gap (Fig.) 4
Visualization of the bound gap shown for optimization of the RealNVP from Fig 3b ii) (i.c., 32 x 32 CelebA denoising). For the plots comparing the lower-bound to the ODE loo-broby,,lity, we used 2048 time samples to estimate bx(x).
20

---

Published in Transactions on Machine Learning Research (07/2024)
B.7 Accuracy of posterior (Fig. 5 Tab. 2)
Variational optimization For both the exact score-based prior and surrogale score-based prior, the variational distribution was a RealNVP with 16 affine-cupllig layers, and it was optimized for 12K iter- ations with a batch size of 2560 and learning rate 10- For the surrogate score-based prior, bg(x) was approximated with N₂ = NN
Basclines For this 2D experiment, we implemented the diffusion-based basclines exactly according their proposed algorithms. For SDE+Proj,, we tested the following values for the measurement weight A: linspace (0.001, 0.5, num=100). For Score-ALD, hyperparameter 1/rr and tested the following values for yr: linspace (100, 0.8, num=100), For DPS, we tested the following values for the scale parameter Ç: exp (linspaee (1og(0 (o.o(..000))) 001), log (0. 15), num=100)).
distilled all hyperparameters into one global
Evaluation Since the diffusion-based approaches only provide samples (not probability densities), we ap- proximated the probability density function (PDF) from the estimated posterior samples. For each method, we a two-component Gaussian mixture model (GMM) to 10000 samples. The reverse KT. divergence was approximated with the log-density function of the fitted GMM and the log-density function of the true posterior, evaluated on these 10000 samples.
B.8
Image-restoration metrics (Fig. 6
Score model The score model was the same the one used for the 64 64 image in the MRI efficiency experiment (Fig.
Variational optimization The variational distribution was RealNVP. Optimization was done with learning rate of 0.00001 and gradient clip of 0.0002. We used the same convergence criterion as the one used in the MRI efficiency experiment with 0.005. The same convergence criterion was also used for the TV results but with a maximum number of steps of 50000. The TV regularization weight was 105.
Baseline hyperparameters For PDo+Proj, we used the projection CS solver provided by SSng et al. (2022) with the hyperparameters snr=0.517, coeff=1. For AcDr we used the langevin CS solver with the hyperparameters n_ each=3, snr=0.212, projection sigma_rate=0. used scale=0.5. This was the best scale out of [10, 1, 0.9, 0.5, 0.3, 0.1, 0.1,0.001] for test image in terms of PSNR with respect to the true image.
rate=0.713.
For DPS, we
B.9 Black-hole interferometric imaging (Fig.
Score model The score model was trained on images of general relativistic magnetoo-yydrodyyanii (GRMHD) simulations (Wong et al. 2022) resized to 64 x 64. During training, the images were randomly flipped horizontally, and they were randomly zoomed so that the ring diameter would vary between 35 and 48 pas. The score model had an NCSN+ architecture with 64 filters in the first layer; it was trained with the VP SDE with 8min = 0.1 and 9max = 20 for 100K steps.
Variational optimization The variational distribution with a RealNVP. Optimization was done with a learning rate of 0.00001 and gradient clip of for 100K steps.
Interferometric imaging assumptions Imaging was done for field of view of 160 pas. A flux-constraint loss added to the DPI optimization objective that all posterior images would have total flux around 173 (computed as the median total flux of images sampled from the score-based GRMHD prior).
21


Reviews:
[{"review_id": "w3YOlM6Jwu", "reviewer": "Reviewer_7FR5", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths.\nThe use of diffusion models as strong priors is interesting and innovative. The argumentation for modeling choices is clear and  coherent.\nThe paper provides a fairly clear and concise technical summary of the score based diffusion framework.\nWriting is clear and concise.\nThe practical relevance and potential impact is strong.\n\nWeaknesses.\nBeacause I am not so familiar with the field, it is unclear to me if the comparison with other approaches is fair and comprehensive.\nI did not see a link to a software implementation. I would strongly suggest providing code.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper studies the use of a score based diffusion model as a prior in Bayesian imaging. The exact computation of probability densities under the diffusion model is expensive, and the paper uses an approximation that is orders of magnitude faster. This is demonstrated on several examples (simulated as well as real) and validated in terms of efficiency, posterior estimation, and reconstruction quality.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please revise the abstract so that the problem setting is more clearly defined.\n\nCan you more clearly define from the onset what is technically meant by \"ill-posed image reconstruction\". Perhaps just a short description along the lines of the first few sentences in sec. 2.1, so the setting is clear from the onset.\n\nIs there a missing \"dt\" in equation 2?\n\nCould you provide references for eq. 2 and 3. (perhaps Song 2021b or Anderson 1982). Would it be possible to also provide some intuition about eq. 3 or some hint to how it is derived? \n\nIn section 4: \"Inspired by previous ... we propose bSDE ...\" Consider rephrasing 'propose' to 'apply' or similar, to clarify that bSDE is not a novel contribution.\n\nI am not so familiar with the imaging application, so I am unsure if the proposed method is evaluated fairly against appropriate baselines.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "vT26FRPNqu", "reviewer": "Reviewer_qKCT", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\n- The paper is clear, well-written, and logically presented. \n\n- Experiments are clearly defined and answer specific questions. \n\n- The empirical results show dramatic improvements in computational efficiency. \n\n- The code will be made publically available (according to the appendix).\n\n**Weaknesses:**\n\n- While the theory is clearly presented, it is mostly a recap of the methods developed in Song 2021a. A theoretical analysis of, for instance, how the approximation error in the learned score function affects performance would have strengthened the paper.\n\n**Comments and suggestions:**\n\nThe setting considered is where an unconditional diffusion model is trained a priori. For this to make sense, a dataset of sufficiently related images needs to exist. A discussion of where this situation arises could be helpful, which could be contrasted to blind image restoration where paired (high quality-low quality) images are available. Furthermore, it would have been interesting to understand the effect of domain shifts between the data the prior is trained on and that of the actual measurements. That said, I understand that the focus of this paper is on the quality of the approximate inference and not on comparing different (learned) priors. \n\nIn section 5.2, does the prior have a significant influence? My suggestion would be to add results from a maximum likelihood approach or a simple conventional prior like TV or L1.\n\n\n**Nitpicking:**\n - Rephrase \"In an ill-posed inverse problem, there is inherent uncertainty in image reconstruction.\"\n - more-accurate -> more accurate\n - Avoid starting sentences with mathematical expressions.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes an approximate inference technique for Bayesian imaging where the prior is defined by a pretrained score-based diffusion model.\nSpecifically, the authors use the ELBO as a substitute for the exact log-probability function and show, empirically, that compared to the exact approach of Feng et al. (2023) this yields large computational savings while resulting in a close approximation of the posterior.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I don't have requests I consider critical, but please consider my suggestions above to strengthen the work.", "broader_impact_concerns": "No.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Rc0n3TPIgQ", "reviewer": "Reviewer_jZTo", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n•\tThe proposed approach is useful and provides an accurate surrogate function with orders of magnitude improved efficiency providing a useful computational framework using score-based diffusion models as priors.\n\n•\tThe approach is sound.\n\n•\tThe results are compelling.\n\nWeaknesses\n\n•\tThe approach technically combines insights from (Feng et al., 2023) with estimation from (Song et al., 2021a). As such, the novelty lies in combining score-based diffusion priors with efficient ELBO based surrogate functions for the problem of ill-posed image reconstruction. The novelty from a technical perspective is thus somewhat incremental.\n\n•\tThe approach could be contrasted also conventional image reconstruction procedures based on standard priors as reviewed at the end of section 2.1. It is unclear how the approach compares to non-diffusion based procedures.\n\n•\tThe paper includes no broader impact discussion which should be included.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a variational inference based surrogate modeling procedure for diffusion models that can characterize complicated distributions by use of score based priors. The approach is found to accelerate inference by orders of magnitudes when compared to exact ODE based prior evaluations. The approach is considered in the context of ill-posed image reconstruction in which the validity of the prior imposed can have substantial influence on the quality of image reconstruction. On synthetic and real datasets the merits of the proposed approach is established when compared to exact modeling and some existing diffusion based baselines.\n\nThe approach is leveraging the work of (Feng et al., 2023) establishing how score based diffusion models can be used as priors relying on expensive ODE solvers thus the need for an efficient surrogate modeling approach scaling this procedure. The approach here explores that the evidence lower bound of the ODE based prior can be efficiently computed leveraging the work of (Song et al., 2021a) to define such efficient ELBO based surrogate function optimized using importance sampling and MC approximations.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Why have traditional priors as reviewed end section 2.1 not been compared against in the experimentation and assessment of results?\n\nIt is well known that ELBOs in general can be tightened by importance weighting. Could such importance weighted ELBOs improve upon the surrogate function accuracy and would this be useful?\n\nThere are many strategies for speeding up the computations associated with diffusion models and the paper has not positioned the current work in this larger literature. It would improve the paper to do this and also explain why some of these other efforts are not relevant/useful in the considered context, i.e. latent diffusion, reducing steps of the diffusion, reverse processes learned to be reduced in diffusion steps etc. Would these existing approaches be relevant also in the considered context?\n\nPlease include a broader impact discussion (see also next section)", "broader_impact_concerns": "Inverse models are influenced by the inductive biases imposed through the prior and it would therefore be good to further elaborate on the trustworthiness of such reconstructions and if reconstructions can be hallucinated resulting in critical erroneous interpretations. Furthermore, diffusion models can negative influence society in terms of deep fake technologies and disinformation – could the considered technology contribute to such negative aspects?\n\nIt would be good to include discussions of the broader impact of the considered research, which is presently lacking from the current manuscript.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: In this paper, the authors focus on inverse problem with deep priors given by diffusion models. Here they train a model using the forward KL to match the posterior in a setting where the likelihood is known and the log probability of the prior is lower bounded with a continuous time ELBO. The authors successfully apply their methods to accelerated MRI, denoising, reconstruction from low spatial frequencies, and black-hole interferometric imaging. The methodology is sound. The experiments are quite compelling. Based on this assessment and the feedback of the reviewers I recommend the acceptance of the manuscript as is.
Decision: Accept (poster)