Title: Efficient Bayesian Computational Imaging with a Surrogate Score-Based Prior
Venue: NeurIPS 2023
Accepted: no
Keywords: computational imaging

inverse problems

generative priors
Abstract: We propose a surrogate function for efficient use of score-based priors for Bayesian inverse imaging. Recent work turned score-based diffusion models into probabilistic priors for solving ill-posed imaging problems by appealing to an ODE-based log-probability function. However, evaluating this function 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 posterior sampling of large images. Compared to the exact prior used in previous work, our surrogate prior accelerates optimization of the variational distribution by at least two orders of magnitude. We also find that our principled approach achieves higher-fidelity image-reconstruction than non-Bayesian baselines that involve hyperparameter-tuning at inference. Our work establishes a practical path forward for using score-based diffusion models as general-purpose priors for computational imaging.

1

Abstract

We propose a surrogate function for efficient use of score-based priors for Bayesian inverse imaging. Recent work turned score-based diffusion models into probabilistic priors for solving ill-posed imaging problems by appealing to an ODE-based log-probability function. However, evaluating this function 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 posterior sampling of large images. Compared to the exact prior used in previous work, our surrogate prior accelerates optimization of the variational distribution by at least two orders of magnitude. We also find that our principled approach achieves higher-fidelity image-reconstruction than non-Bayesian baselines that involve hyperparameter-tuning at inference. Our work establishes a practical path forward for using score-based diffusion models as general-purpose priors for computational imaging.

1 Introduction

Ill-posed image reconstruction requires a prior to constrain the reconstruction according to desired image statistics. From a Bayesian perspective, the prior influences 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 a challenge. True posterior sampling with an unconditional diffusion model is intractable, so most previous methods heavily approximate the posterior [9; 13; 14; 19] or disregard measurement noise [5; 7; 8; 6; 11; 24; 1]. Recent work demonstrated how to turn score-based diffusion models into probabilistic priors ( scorebased priors ) for Bayesian imaging [10]. However, this method requires the exact probability of a proposed image to be evaluated with a computationally-expensive ordinary differential equation (ODE), requiring days to a week to reconstruct even a 32 × 32 image [10]. We present a method for Bayesian inference with a score-based prior that is both principled and computationally efficient.

Although computing exact probabilities under a diffusion model is inefficient or even intractable, 27 computing the evidence lower-bound [22; 12] is computationally efficient and feasible for high28 dimensional images. Thus we propose to use this evidence lower-bound as a surrogate for the 29 exact score-based prior. In particular, we use the evidence lower-bound of a score-based diffusion 30 model [22] as a substitute for the exact log-probability function. This function can be plugged into 31 any inference algorithm that requires the value or gradient of the posterior log-density. When it is 32 used in variational inference, we find at least two orders of magnitude in speedup of optimizing the 33 variational distribution. Furthermore, our approach reduces GPU memory requirements, as there 34 is no need to evaluate and backpropagate through an ODE. These efficiency improvements make it 35 practical to perform inference with score-based priors. 36

<latexi sh

1_b

64="7fz

n

FWpDNg3QS2

RMrBU0Y

>A

c

V

Hy

+

T

C

5KIo

m

Z

Eq/

G

u

Pv

w

L

d

9

JO

k

j

<latexi sh

1_b

64="KHEYI2S

p

8

7L

PQ

O

9

C

A

>

3

c

VD

N

FJ

U

k

M

u

m

T

/B

R

df+

j

wy

r

n

G5

Z

X

z

<latexi sh

1_b

64="

8H

v

XWTPZ

K/

7

B5R0

I

>A

3

c

VDLSgM

F

j

9

Y

E

k

d

Ny r2

o

O

G

q

f

+

m

u

n

z

w

C

U

Q

Observed

Posterior Samples

Original

16x-accel.

MRI

256x256

4x-accel.

MRI

256x256

y

x

∗

x

∼

p

(

x

|

y

)

In this paper, we describe our variational-inference approach to efficiently estimate a posterior with 37 a surrogate score-based prior. We provide experimental results to validate the proposed surrogate 38 prior, including high-dimensional posterior samples of sizes up to 256 × 256 , a resolution infeasible 39 with the exact prior. In the setting of accelerated MRI, we quantify time- and memory-efficiency 40 improvements of the surrogate over the exact prior. We also demonstrate how our proposed approach 41 achieves higher-quality image reconstructions than methods that deviate from true Bayesian inference. 42

2 Related work 43

2.1 Bayesian inverse imaging 44

Image reconstruction can be framed as an inverse problem: a hidden image x ∗ ∈ R D must be 45 recovered from measurements y ∈ R M , where 46

It is usually assumed that the forward model f : R D → R M is known and that the measurement 47 noise ϵ ∈ R M is a random variable with a known distribution. With an ill-posed inverse problem, 48 there is inherent uncertainty in image reconstruction. 49

Bayesian imaging accounts for the uncertainty by formulating a posterior distribution p ( x | y ) . The 50 posterior can be decomposed into a likelihood term and a prior term: 51

Given a log-likelihood function log p ( y | x ) and a prior log-probability function log p ( x ) , we can 52 use established techniques for sampling from the posterior, such as Markov chain Monte Carlo 53 (MCMC) [3] or variational inference [2]. MCMC algorithms generate a Markov chain whose 54 stationary distribution is the posterior, but they are generally slow to converge for high-dimensional 55 data like images. Variational inference instead approximates the posterior with a tractable distribution 56 (e.g., Gaussian). The variational distribution is usually parameterized and thus can be efficiently 57 optimized to represent high-dimensional data distributions. Deep Probabilistic Imaging (DPI) [25; 26] 58 proposed an efficient variational-inference approach specifically for computationtal imaging with 59 traditional regularizers; in DPI, the variational distribution is a discrete normalizing flow [15], which 60 is an invertible generative model capable of representing complex distributions. 61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

2.2 Diffusion models for inverse problems

Primarily developed for image generation, diffusion models [18; 12; 20; 21; 23] 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 it until it becomes a clean image. We discuss this process, known as reverse diffusion , in more detail in Sec. 3.1.

Given an inverse problem, simply adapting a pretrained diffusion model to sample from the posterior instead of the learned prior is intractable [10]. Therefore, most diffusion-based approaches do not infer a true Bayesian posterior. Some methods project images onto a measurement-consistent subspace [24; 8; 6; 5; 7], but the projection does not account for measurement noise and might pull images away from a true posterior. Other methods follow a gradient toward higher likelihood throughout reverse diffusion [9; 13; 11; 14; 1; 19; 17], but these methods heavily approximate the posterior. Overall, these diffusion-based methods require hyperparameter-tuning to balance the measurements and the prior. As soon as hyperparameters are introduced, there is no guarantee of sampling from a posterior that represents the true uncertainty.

Score-based priors. Alternatively, a score-based diffusion model can be turned into a standalone, probabilistic prior ( score-based prior ) that can be paired with any measurement-likelihood function and plugged into established Bayesian-inference approaches. Feng et al. [10] proposed to do this with a log-density function based on the ODE associated with reverse diffusion (see Sec. 3.2). This function provides the log-probability of any image under the diffusion model's generative distribution, but it is computationally expensive to evaluate. When used in iterative optimization algorithms, it incurs prohibitively high time and memory costs.

3 Background

In this section, we review background on score-based diffusion models with an emphasis on evaluating probabilities of images with a pretrained diffusion model. We then describe how a diffusion process gives rise to an efficient denoising-based lower-bound on these image probabilities.

3.1 Score-based diffusion models

The core idea of a diffusion model is that it transforms a simple distribution π to a complex image distribution through a gradual process. In this work, we follow the popular framework of denoising diffusion models, which transform noise samples from π = N ( 0 , I ) to clean samples from the data distribution p data through gradual denoising. With knowledge of the noise distribution and the denoising process, we can assess the probability of a novel image under this generative model.

The transformation from a simple distribution to a complex one occurs over many steps. To determine how the data distribution should look at each step of the denoising process, we turn to a stochastic differential equation (SDE) that describes a diffusion process from clean images to noise. The diffusion SDE is defined on the time interval t ∈ [0 , T ] and has the form

where w ∈ R D denotes Brownian motion. g ( t ) ∈ R is the diffusion coefficient, which controls the rate of noise increase. f ( · , t ) : R D → R D is the drift coefficient, which controls the deterministic evolution of x ( t ) . By defining a stochastic trajectory { x ( t ) } t ∈ [0 ,T ] , this SDE gives rise to a timedependent probability distribution p t , which is the marginal distribution of x ( t ) . We construct f ( · , t ) and g ( t ) so that if p 0 = p data, then p T ≈ π . 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 p data. A convolutional neural network s θ known as a score model is trained to approximate the true score: s θ ( x , t ) ≈ ∇ x log p t ( x ) .

3.2 Image probabilities under a score-based diffusion model

Once trained, s θ ( x , t ) is used in a reverse-diffusion process to generate clean images from noise. 106 The generated image distribution theoretically assigns a probability density to every possible image. 107 However, reverse diffusion does not lead to an image distribution with tractable probabilities. In this 108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

where

p 0 t ( x ′ | x ) denotes the transition distribution from x (0) = x to x ( t ) = x ′ . For a drift coefficient 143 that is linear in x , this transition distribution is Gaussian: p 0 t ( x ′ | x ) = N ( x ′ ; α ( t ) x , β ( t ) 2 I ) . This 144 means that the gradient ∇ x ′ log p 0 t ( x ′ | x ) is directly proportional to the Gaussian noise that is 145 subtracted from x ′ to get x . Eq. 7 is efficient to compute since we can evaluate it by adding Gaussian 146

subsection, we describe two workarounds: one based on an ordinary differential equation (ODE) and the other based on a denoising score-matching objective.

Sampling with a reverse-time SDE. Reversing diffusion (Eq. 3) with a score model s θ ( x , t ) results in a distribution p SDE θ , denoted as such because it is determined by a reverse-time SDE:

¯ w ∈ R D denotes Brownian motion, and f ( · , t ) and g ( t ) are the same as in Eq. 3. To generate an image, we first sample x ( T ) ∼ N ( 0 , I ) and then numerically solve the reverse-time SDE for x (0) . p SDE θ is the marginal distribution of x (0) , which for a well-trained score model is close to p data .

To compute the probability of an image x under p SDE θ , we need to invert this image from x (0) = x to x ( T ) . However, this is not tractable through the SDE: just as it is intractable to reverse a random walk, it is intractable to account for all the possible starting points x ( T ) that could have resulted in x (0) through the stochastic process. Probability computation calls for an invertible process that lets us map any point from p data to N ( 0 , I ) and vice versa.

Computing probabilities with an ODE. The probability flow ODE [23] defines an invertible sampling function for an image distribution p ODE θ theoretically the same as p SDE θ . It is given by

The absence of Brownian motion makes it possible to solve this ODE in both directions of time. To compute the log-probability of an image x , we map x (0) = x to its corresponding noise image x ( T ) . Under the framework of neural ODEs [4], the log-probability is given by the log-probability of x ( T ) under N ( 0 , I ) plus a normalization factor accounting for the change in density through time:

Although tractable to evaluate with an ODE solver, this log-probability function is computationally 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 HutchinsonSkilling trace estimation of the divergence.

Equivalence of p SDE θ and p ODE θ . Song et al. [22] proved that if s θ ( x , t ) ≡ ∇ x log p t ( x , t ) for all t ∈ [0 , T ] and p T = π , then p ODE θ = p SDE θ = p data. In our work, we assume that s θ ( x , t ) ≈ ∇ x log p t ( x , t ) for almost all x ∈ R D and t ∈ [0 , T ] and that p T ≈ N ( 0 , I ) , so that p ODE θ ≈ p SDE θ ≈ p data. This assumption empirically performed well in previous work that appealed to p ODE θ as the exact probability distribution of the diffusion model [10; 23].

3.3 Evidence lower bound of a score-based diffusion model

In lieu of an exact log-probability function, Song et al. [22] derived an evidence lower-bound for p SDE θ such that b SDE θ ( x ) ≤ log p SDE θ ( x ) for any proposed image x . Essentially, this lower-bound corresponds to how well the diffusion model is able to denoise a 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' [22], is defined as

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

noise to x without having to solve an initial-value problem as with the ODE. In fact, Eq. 7 is closely related to the denoising score-matching objective used to efficiently train diffusion models [23].

Intuitively, we can interpret Eq. 7 as associating an image's probability with how well the score model s θ could denoise that image if it underwent diffusion. This is represented by the first term in h ( t ) (Eq. 8). To assess the probability of an image x , we perturb it with Gaussian noise to get x ′ and then ask the score model to estimate the noise that was added. If s θ ( x , t ) accurately estimates the noise, then ∥ s θ ( x ′ , t ) -∇ x ′ log p 0 t ( x ′ | x ) ∥ 2 2 is small, and the value of b SDE θ ( x ) becomes larger.

The remaining terms in h ( t ) are normalizing factors independent of θ . The term E p 0 T ( x ′ | x ) [log π ( x ′ )] accounts for the probabilities of the noise images x ( T ) that could result from x being entirely diffused.

4 Method

Inspired by previous theoretical work [22], we propose b SDE θ as an efficient surrogate prior for the exact score-based prior in Bayesian imaging. In this section, we describe our approach for efficient posterior inference with a score-based prior.

4.1 Variational inference with a surrogate score-based prior

Given measurements y ∈ R M (with a known log-likelihood function) and a score-based diffusion model (parameterized by θ ) as the prior, our goal is to sample from the image posterior p θ ( x | y ) . We follow a variational-inference approach by optimizing the parameters of a variational distribution to closely approximate the target posterior.

Let q ϕ denote the variational distribution with parameters ϕ , and we assume q ϕ to have tractable log-probabilities. We optimize ϕ to minimize the KL divergence from q ϕ to the target posterior:

q ϕ can be various types of distributions. For example, it could be a Gaussian distribution with a diagonal covariance matrix so that ϕ := [ µ ⊤ , σ ⊤ ] ⊤ , where µ ∈ R D and σ ∈ R D ( σ > 0 ) are the mean and pixel-wise standard deviation. As DPI showed [25], q ϕ could also be a RealNVP normalizing flow with network parameters ϕ .

To circumvent the computational challenges of evaluating the prior term log p ODE θ ( x ) , we replace it with the surrogate b SDE θ ( x ) . This results in the following objective:

We can also think of b SDE θ as replacing the intractable log p SDE θ in Eq. 9. Since -log p SDE θ ≤ -b SDE θ , our surrogate objective minimizes the upper-bound of a valid KL divergence involving p SDE θ .

4.2 Implementation details

Evaluating b SDE θ ( x ) . The formula for b SDE θ ( x ) (Eq. 7) contains a time integral and expectation over p 0 t ( x ′ | x ) that can be estimated with numerical methods. Following Song et al. [22], we use importance sampling with time samples t ∼ p ( t ) for the time integral and Monte-Carlo approximation with noisy images x ( α ( t ) x , β ( t ) 2 I ) for the expectation. The proposal distribution p ( t ) :=

β ( t ) 2 Z was empirically verified to result in lower variance in the estimation of b SDE θ ( x ) [22]. We provide the following formula used in our implementation, which estimates the time integral with importance sampling and the expectation with Monte-Carlo approximation, for reference:

′ ∼ N g ( t ) 2

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

N t is the number of time samples used to approximate the time integral, and N z is the number of noise samples taken to approximate the expectation over p 0 t ( x ′ | x ) . In our experiments, we set N t = N z = 1 . Increasing the number of time and noise samples does not efficiently decrease variance in the estimated value of b SDE θ ( x ) . We use the Variance Preserving (VP) SDE.

Optimization. We use stochastic gradient descent to optimize ϕ , Monte-Carlo approximating the expectation in Eq. 10 with a batch of x ∼ q ϕ . We find that estimating b SDE θ ( x ) has higher variance than estimating log p ODE θ ( x ) . For example, in Fig. 4, b SDE θ ( x ) with N t = 2048 , N z = 1 shows higher variance than log p ODE θ ( x ) with 16 trace estimators. When optimizing a complex distribution like RealNVP, a lower learning-rate helps 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 was 0 . 00001 . Please refer to the supplemental text for more optimization details.

5 Experiments

We validate our proposed approach on the tasks of accelerated MRI, image denoising, and reconstruction from low spatial frequencies. We highlight accelerated (or compressed sensing) MRI because in addition to being a real-world imaging problem that calls for accurate posterior estimation, it is the focus of much related work [24; 13]. In MRI, measurements in a spatial-frequency space ( κ -space) are obtained to help reveal a hidden anatomical image. Accelerated MRI reduces the number of κ -space measurements, thus reducing the scan time but also making the image reconstruction ill-posed. The supplemental text provides details on how measurements were generated for all tasks.

5.1 Efficiency improvements

In Tab. 1 and Fig. 2, we quantify the efficiency improvements of the surrogate prior for an accelerated MRI task at different image resolutions. We drew a test image from the fastMRI knee dataset [27] and resized it to 16 × 16 , 32 × 32 , 64 × 64 , 128 × 128 , and 256 × 256 . For each image size, we trained a score model on training images of the corresponding size from the fastMRI dataset of single-coil knee scans. We then optimized a Gaussian distribution with diagonal covariance to 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 × 64 optimization into GPU memory). Convergence was defined by setting a minimum acceptable change in the mean of the estimated posterior between optimization steps.

We find at least two orders of magnitude in time improvement with the surrogate prior. Tab. 1 compares the iteration time between the two priors. Fig. 2 compares the total time it takes to optimize the variational distribution. The surrogate also significantly improves memory consumption, which in turn enables optimizing higher-dimensional posteriors. Following standard practice, we just-in-time (JIT) compile the optimization step to reduce time/step at the cost of GPU memory. Fig. 2 shows how the surrogate prior significantly reduces memory requirements and scales better with image size. The exact prior could only handle up to 32 × 32 before exceeding GPU memory (we tested on 4x 48GB GPUs). While memory could be reduced with a smaller batch size, this would make optimization more timeconsuming. On the other hand, our surrogate prior supports much larger images, as we demonstrate in Fig. 1 for 256 × 256 1 MRI with a Gaussian-approximated posterior. This type of principled inference of high-dimensional image posteriors was not possible before with the exact score-based prior.

5.2 Posterior estimation under the surrogate vs. exact prior

We cannot expect the surrogate prior b SDE θ to be an identical substitute for the exact prior log p ODE θ . Importantly, though, we verify in Fig. 3a that both the surrogate and the exact prior recover a groundtruth Gaussian posterior derived from a Gaussian likelihood and prior. The variational distribution

×

232

233

234

235

236

237

238

239

240

241

242

243

244

245

Dimensionality

Memory [GB]

Dimensionality

Optimization Time

Optimization Memory

PSNR

SSIM

Image-Restoration Quality

Dimensionality

39 hrs.

43 hrs.

126 hrs.

19 mins.

19 mins.

45 mins.

2 hrs.

9 hrs.

16 GB

31 GB

89 GB

8 GB

12 GB

17 GB

44 GB

142 GB

Time [sec]

= Exact

= Surrogate

= Exact

= Surrogate

-0.35

Mean

True

Est.

(Surrogate)

Est.

(Exact)

0.89

0.10

0.16

(a) Ground-truth (Gaussian) posterior.

Surrogate

Exact

Mean

Original

Observed

Original

Observed

Mean

0.01

0.12

0.01

0.1

(i) CelebA denoising

(ii) CIFAR-10 denoising

(b) Complex posteriors.

used for inference is a RealNVP, and the score model (used by both the surrogate and exact prior) was trained on samples from the known Gaussian prior.

Nonetheless, the surrogate could result in a different locally-optimal variational posterior, particularly if the posterior is complex with various local minima in the variational objective. Fig. 3b compares posteriors (with unknown true distribution) 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 pretrained score model. We observe in these comparisons that most of the differences appear in the image background and that both priors result in a plausible mean reconstruction and uncertainty.

Visualizing the bound bap throughout optimization helps shed light on why the two priors converge to different solutions even if the underlying score model is the same. Fig. 4 shows probabilities of samples generated by q ϕ (in this case, a RealNVP) as optimization progresses. At each checkpoint of q ϕ , we plot log p ODE θ ( x ) versus b SDE θ ( x ) (approximated with N t = 2048 for reduced variance) for samples x ∼ q ϕ coming from both the exact and surrogate optimization of q ϕ . Importantly, we find that the surrogate is a valid bound for the ODE log-density: b SDE θ ( x ) ≤ log p ODE θ ( x ) for all

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

x ∼ q ϕ ( x ) , except for some outliers due to variance of b SDE θ ( x ) . However, we find that optimization follows a different trajectory depending on the prior. With the surrogate, samples x ∼ q ϕ tend toward a region where the bound gap is small (i.e., b SDE θ ( x ) is close to log p ODE θ ( x ) ). Meanwhile, the exact prior follows a loss landscape whose structure appears to be independent of the lower-bound. Note that samples from q ϕ optimized under the exact prior obtain higher values of b SDE θ ( x ) than samples obtained under the surrogate. The observations in Fig. 4 suggest that gradients under the surrogate tend to push the q ϕ distribution along the boundary of equality between b SDE θ and log p ODE θ . This constrains the path taken through gradient descent and subsequently the converged solution.

Exact

Surrogate

Step 20000

Step 1000000

Step 500

Step 100

Step 5000

Step 25000

Step 100

Step 500

Step 20K

Step 5K

Step 25K

Step 1M

Plotting

trajectories on

same scale

log

p

ODE

θ

(               )

b

SDE

θ

(         )

log

p

ODE

θ

(

x

)

b

SDE

θ

(

x

)

log

p

ODE

θ

(

x

)

log

p

ODE

θ

(

x

)

b

SDE

θ

(

x

)

b

SDE

θ

(

x

)

log

p

ODE

θ

(

x

)

log

p

ODE

θ

(

x

)

log

p

ODE

θ

(

x

)

b

SDE

θ

(

x

)

b

SDE

θ

(

x

)

b

SDE

θ

(

x

)

b

SDE

θ

(

x

)

b

SDE

θ

(

x

)

log

p

ODE

θ

(

x

)

log

p

ODE

θ

(

x

)

00

00

0

00

00

00

00

00

00

00

00

00

5.3 Image-reconstruction quality

It would be reasonable to assume that diffusion-based approaches discussed in Sec. 2, although less principled, may lead to better visual quality than a 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 a good prior, it might be desirable for posterior samples to accurately reflect the true underlying image.

We performed multiple MRI tasks at different acceleration rates and compared our approach to three baselines: SDE+Proj [24], Score-ALD [13], and Diffusion Posterior Sampling ( DPS ) [9]. SDE+Proj projects images onto a measurement subspace. Score-ALD and DPS approximate the posterior throughout reverse diffusion. All baselines involve at least one measurement-weight hyperparameter. The implementations and hyperparameter settings for SDE+Proj and Score-ALD were provided by Song et al. [24]. For DPS, we followed the implementation of Chung et al. [9] and performed a hyperparameter search on an 8 × -acceleration test image to find the optimal PSNR.

We simulated MRI at three different acceleration factors for ten test images, resulting in thirty 267 posterior distributions to be estimated. As baseline implementations do not account for measurement 268 noise, we gave the baselines noiseless measurements and set a near-zero measurement noise for our 269 method. The test images were randomly sampled from the fastMRI dataset and resized to 64 × 64 . 270

<latexi sh

1_b

64="zVurDO9TRSko

A7q3G

LQ

0U

>

C

c

J

gNBE

2

W

5

Y/

wZ

y

n

m

Hj

M

P+

v

F

d

K

p

I

f

8

X

<latexi sh

1_b

64="

W93/YyrfI

A8

F

qM

P

nc

E

>

B

H

VDJSgN

O2

z

5

o

K

jm

U

T

C

0

L

Q

p

7

G

Z

R

+

d

k

<latexi sh

1_b

64="

UWrp2

HXC

Y

nw g0mE

>A

c

VDJS

N

O

j

9

5

R

/

T

G

k

u

y

Lv

F

Q

o

d

7

3Pz

Z

8q

M

I

K

f

<latexi sh

1_b

64="

UWrp2

HXC

Y

nw g0mE

>A

c

VDJS

N

O

j

9

5

R

/

T

G

k

u

y

Lv

F

Q

o

d

7

3Pz

Z

8q

M

I

K

f

<latexi sh

1_b

64="oPUr

A

Q

2O3

R

5Fm

H

>

CDX

c

V

JSgNBE

j

0

w

u

M

I

y

f

p

vd

W

/

T

n

q

Z

L

+

Gz

9

K

k

7

<latexi sh

1_b

64="oPUr

A

Q

2O3

R

5Fm

H

>

CDX

c

V

JSgNBE

j

0

w

u

M

I

y

f

p

vd

W

/

T

n

q

Z

L

+

Gz

9

K

k

7

<latexi sh

1_b

64="

UWrp2

HXC

Y

nw g0mE

>A

c

VDJS

N

O

j

9

5

R

/

T

G

k

u

y

Lv

F

Q

o

d

7

3Pz

Z

8q

M

I

K

f

<latexi sh

1_b

64="

UWrp2

HXC

Y

nw g0mE

>A

c

VDJS

N

O

j

9

5

R

/

T

G

k

u

y

Lv

F

Q

o

d

7

3Pz

Z

8q

M

I

K

f

<latexi sh

1_b

64="oPUr

A

Q

2O3

R

5Fm

H

>

CDX

c

V

JSgNBE

j

0

w

u

M

I

y

f

p

vd

W

/

T

n

q

Z

L

+

Gz

9

K

k

7

<latexi sh

1_b

64="oPUr

A

Q

2O3

R

5Fm

H

>

CDX

c

V

JSgNBE

j

0

w

u

M

I

y

f

p

vd

W

/

T

n

q

Z

L

+

Gz

9

K

k

7

<latexi sh

1_b

64="

UWrp2

HXC

Y

nw g0mE

>A

c

VDJS

N

O

j

9

5

R

/

T

G

k

u

y

Lv

F

Q

o

d

7

3Pz

Z

8q

M

I

K

f

<latexi sh

1_b

64="

UWrp2

HXC

Y

nw g0mE

>A

c

VDJS

N

O

j

9

5

R

/

T

G

k

u

y

Lv

F

Q

o

d

7

3Pz

Z

8q

M

I

K

f

<latexi sh

1_b

64="oPUr

A

Q

2O3

R

5Fm

H

>

CDX

c

V

JSgNBE

j

0

w

u

M

I

y

f

p

vd

W

/

T

n

q

Z

L

+

Gz

9

K

k

7

<latexi sh

1_b

64="oPUr

A

Q

2O3

R

5Fm

H

>

CDX

c

V

JSgNBE

j

0

w

u

M

I

y

f

p

vd

W

/

T

n

q

Z

L

+

Gz

9

K

k

7

<latexi sh

1_b

64="

UWrp2

HXC

Y

nw g0mE

>A

c

VDJS

N

O

j

9

5

R

/

T

G

k

u

y

Lv

F

Q

o

d

7

3Pz

Z

8q

M

I

K

f

<latexi sh

1_b

64="

UWrp2

HXC

Y

nw g0mE

>A

c

VDJS

N

O

j

9

5

R

/

T

G

k

u

y

Lv

F

Q

o

d

7

3Pz

Z

8q

M

I

K

f

<latexi sh

1_b

64="oPUr

A

Q

2O3

R

5Fm

H

>

CDX

c

V

JSgNBE

j

0

w

u

M

I

y

f

p

vd

W

/

T

n

q

Z

L

+

Gz

9

K

k

7

<latexi sh

1_b

64="oPUr

A

Q

2O3

R

5Fm

H

>

CDX

c

V

JSgNBE

j

0

w

u

M

I

y

f

p

vd

W

/

T

n

q

Z

L

+

Gz

9

K

k

7

Avg. PSNR

= Ours

(DPI + Surrogate)

= SDE+Proj

= DPS

= Score-ALD

Avg. SSIM

Acceleration factor

Acceleration factor

(a) Image-restoration metrics.

Original

Ours

(DPI + Surr.)

SDE+Proj

Score-ALD

0-Filled Recon.

(16x-accel.)

DPS

Our approach was DPI with the surrogate prior, meaning we optimized a RealNVP to approximate 271 each posterior and used the lower-bound function b SDE θ as the prior log-density. The score model s θ 272 was trained on 64 × 64 images of knee scans from fastMRI and stayed fixed across all methods. 273

274

275

276

277

278

Our method achieves a marked improvement in PSNR and SSIM over the three baselines (Fig. 5). Across all acceleration factors and baselines, our method improves PSNR by between 2 . 7 and 8 . 5 dB. Even though each method uses 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 a standalone prior in Bayesian inference.

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

6 Conclusion

We have presented a surrogate function that provides efficient access to score-based priors for Bayesian inference. We empirically verify that the evidence lower-bound b SDE θ ( x ) ≤ log p SDE θ ( x ) can serve as a proxy for evaluating the log-prior of an image under a trained diffusion model. Paired with any log-likelihood function, b SDE θ ( x ) can be plugged into a Bayesian-inference algorithm. Our experiments with variational inference show at least two orders of magnitude in runtime improvement and significant memory improvement over the ODE-based prior. This enables inference of images previously too large for a strictly Bayesian approach, such as 256 × 256 pixels. We also establish that a principled approach like ours outperforms baselines on image-restoration metrics, evidence that following a Bayesian approach results in more-reliable image reconstructions.

Limitations. A variational approach like ours depends on the expressiveness of the variational distribution. Improvements may be possible by using a diffusion model instead of a discrete normalizing flow as the variational distribution. We also note that there are open theoretical questions about b SDE θ as it relates to p ODE θ [16]. Broader impact. Our proposed framework for efficient estimation of high-dimensional, sophisticated posteriors has broad potential impact for computational imaging. Many imaging tasks, especially in science and medicine, would benefit from accurate uncertainty quantification with principled, data-driven priors.

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

References

Alexandre Adam, Adam Coogan, Nikolay Malkin, Ronan Legin, Laurence Perreault-Levasseur, Yashar Hezaveh, and Yoshua Bengio. Posterior samples of source galaxies in strong gravitational lenses with score-based priors. arXiv preprint arXiv:2211.03812 , 2022.

David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association , 112(518):859-877, 2017.

Steve Brooks, Andrew Gelman, Galin Jones, and Xiao-Li Meng. Handbook of markov chain monte carlo . CRC press, 2011.

Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. NeurIPS , 31, 2018.

Jooyoung Choi, Sungwon Kim, Yonghyun Jeong, Youngjune Gwon, and Sungroh Yoon. Ilvr: 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, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. arXiv preprint arXiv:2206.00941 , 2022.

Hyungjin Chung, Byeongsu Sim, and Jong Chul Ye. Come-closer-diffuse-faster: Accelerating conditional diffusion models for inverse problems through stochastic contraction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 12413-12422, 2022.

Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Marc Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=OnD9zGAGT0k .

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. arXiv preprint arXiv:2304.11751 , 2023.

Alexandros Graikos, Nikolay Malkin, Nebojsa Jojic, and Dimitris Samaras. Diffusion models as plugand-play priors. In Thirty-Sixth Conference on Neural Information Processing Systems , 2022. URL https://arxiv.org/pdf/2206.09012.pdf .

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros G Dimakis, and Jonathan I Tamir. Robust compressed sensing mri with deep generative priors. NeurIPS , 2021.

Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. In Advances in Neural Information Processing Systems , 2022.

Ivan Kobyzev, Simon JD Prince, and Marcus A Brubaker. Normalizing flows: An introduction and review of current methods. IEEE transactions on pattern analysis and machine intelligence , 43(11):3964-3979, 2020.

Cheng Lu, Kaiwen Zheng, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Maximum likelihood training for score-based diffusion ODEs by high order denoising score matching. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 14429-14460. PMLR, 17-23 Jul 2022. URL https://proceedings.mlr.press/ v162/lu22f.html .

Morteza Mardani, Jiaming Song, Jan Kautz, and Arash Vahdat. A variational perspective on solving inverse problems with diffusion models. arXiv preprint arXiv:2305.04391 , 2023.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Int. Conf. Machine Learning , pages 2256-2265. PMLR, 2015.

Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse-guided diffusion models for inverse problems. In International Conference on Learning Representations , 2023. URL https: //openreview.net/forum?id=9_gsMA8MRKQ .

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In NeurIPS , pages 11895-11907, 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 in Artificial Intelligence, UAI , page 204, 2019. URL http://auai.org/uai2019/proceedings/papers/204. pdf .

Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. In Thirty-Fifth Conference on Neural Information Processing Systems , 2021.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In ICLR , 2021. URL https://openreview.net/forum?id=PxTIG12RRHS .

Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. In ICLR , 2022. URL https://openreview.net/forum?id= vaRCHVj0uGI .

He Sun and Katherine L Bouman. Deep probabilistic imaging: Uncertainty quantification and multi-modal solution characterization for computational imaging. In AAAI , pages 2628-2637, 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. arXiv preprint arXiv:2201.08506 , 2022.

Jure Zbontar, Florian Knoll, Anuroop Sriram, Tullie Murrell, Zhengnan Huang, Matthew J Muckley, Aaron Defazio, Ruben Stern, Patricia Johnson, Mary Bruno, et al. fastmri: An open dataset and benchmarks for accelerated mri. arXiv preprint arXiv:1811.08839 , 2018.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper deals with variational inference (VI) of the posterior distribution when using a diffusion based generative model as the prior distribution. The paper proposes using a lower bound on the log probability of the prior distribution (instead of calculating it through the  ODE in standard diffusion based generative models as in [1]) for the optimization of the Kullback-Leibler (KL) divergence between in the VI framework. This results in a more efficient training procedure for the VI than using the true log probability as well as a more efficient memory footprint. \n\nThe authors then validate their method using an MRI dataset as well as the most common CelebA and Cifar datasets.\n\n[1] Berthy T Feng, Jamie Smith, Michael Rubinstein, Huiwen Chang, Katherine L Bouman, and William T318\nFreeman. Score-based diffusion models as principled priors for inverse imaging. arXiv preprint319\narXiv:2304.11751, 2023\n\nSTRENGTHS:\n* The paper proposes a solution to estimating the KL that is indeed way faster than the previous known methods and this allows for **effective gains in training time and memory footprints**. They provide some visual evidence that those gains do not alter considerably the quality of the variational approximation of the posterior distribution.\n\n* The presentation of the problem is good and they motivate the need for an efficient algorithm with a real word application (Accelerated MRI).\n\n* The explanation of the method is clear and reproducible. I expect practitioners would be able to clearly implement the algorithm with the description given in the main part of the text.\n\nWEAKNESSES:\n* The paper for me gives a slightly overstated presentation of the variational inference as a sampling from a true Bayesian posterior. The Bayesian posterior is clearly defined once one states that the prior is the diffusion based generative model. Therefore, when doing Variational Inference, we are only approximating (to an unknown degree) this posterior distribution, so I no way the samples generated by the outcome distribution of the VI procedure are not an approximation of the Bayesian problem.  Of course, VI is still a useful approach, but I would say that I agree that the proposed algorithm is closer to the \"true Bayesian inference\" (line 42) than other such as DPS [2] or SMC-DIFF [3].\n\n* The numerical evaluation and comparison with other algorithms is insufficient.\n 1) **Evaluating the distance to the true posterior**: The numeric compare only visually the posterior with the posterior obtained from [1]. When comparing with [2], the paper compare SSIM and PSNR as well as visual assessment of the reconstructions. As stated in line 258 of the paper, for ill posed inverse problems comparison to the \"true image\" can not be considered an adequate metric. I'd suggest considering an example where the posterior is analytically available (for example, when considering a Gaussian likelihood with a diffusion model over a mixture of Gaussians). In such case, several metrics can be used to compare the different methods ([1], [2]) such as the sliced wasserstein or even the KL.\n\n2) **Complexity**: When comparing with other methods such as [2], we should keep in mind that the proposed method needs an optimization problem for each measurement $y$ (minimization of the KL). This is not the case for some of the \"posterior diffusion samplers\" such as [2] and [3]. Therefore, the actual computing time needed once we receive a measurement $y$ is smaller for [2] and [3] quite considerably.\n\nI'd be inclined to augment my grade if the quantification of the error to the true posterior is better understood, specially in comparison with [1] and [2]. \n\n\n\n[2] Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Marc Louis Klasky, & Jong Chul Ye (2023). Diffusion Posterior Sampling for General Noisy Inverse Problems. In The Eleventh International Conference on Learning Representations .\n\n[3] Brian L. Trippe, Jason Yim, Doug Tischer, David Baker, Tamara Broderick, Regina Barzilay, & Tommi S. Jaakkola (2023). Diffusion Probabilistic Modeling of Protein Backbones in 3D for the motif-scaffolding problem. In The Eleventh International Conference on Learning Representations .\n\nQUESTIONS:\n1. I do not understand figure 3(a). How can the prior be Gaussian and at the same time come from a diffusion generative model on CelebA? I understand that somehow the score was trained to sample from a Gaussian (if the target distribution is a Gaussian the score is analytically available), but it's not clear how. If the posterior is Gaussian, why not using a Gaussian family for VI instead of RealNVP?  Also, if everything is Gaussian, one could in principle calculate several metrics between each algorithm and the true distribution, which would be a non subjective metric for the posterior.\n\n2. Figure 4 suggests that the minimum found by both the proposed algorithm and [1] are quite different. What is the KL gap between the two?\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper focuses on solving inverse problems using diffusion based probabilistic models. The approach considered consists in minimizing the KL divergence between a variational posterior and the true posterior of the diffusion model. Computing this KL involves approximating the log probability of the diffusion model's marginal (which is assumed to approximate the true data generating distribution). A previous paper [1] used the very expensive ODE approach to approximate this log probability, making the whole method prohibitively expensive and inefficient when factoring in the computational cost. The present paper suggests instead minimizing an upper bound on the KL divergence, using a lower bound on the log probability of interest that was derived in [2]. \n\n[1] *Feng, B.T., Smith, J., Rubinstein, M., Chang, H., Bouman, K.L. and Freeman, W.T., 2023. Score-Based diffusion models as principled priors for inverse imaging. arXiv preprint arXiv:2304.11751.* \n\n[2] *Song, Yang, Conor Durkan, Iain Murray, and Stefano Ermon. \"Maximum likelihood training of score-based diffusion models.\" Advances in Neural Information Processing Systems 34 (2021): 1415-1428.*\n\nSTRENGTHS:\nThe proposed method is more efficient than that proposed in [1]. It results in a drastic reduction of computational time and memory cost while maintaining having very similar performance.\n\nWEAKNESSES:\n- I believe that the KL approach is sound and reliable since the resulting variational approximation will likely not sample outside the support of the posterior. It is however extremely costly in terms of computational time and memory cost. It takes 9 hours to obtain a variational approximation over a **single image** of dimension 2^16. On the other hand, classical posterior sampling methods such as DPS [3] take less than one minute for larger images. One might then argue that such methods do not target the exact posterior and it is true! However, the approach in the present paper is also not guaranteed to sample the correct posterior; the forward KL suffers from mode collapse and posteriors over high dimensional images are highly multimodal. The only advantage of this approach is that it will likely not \"hallucinate\" but I do not think that it is in anyway competitive with other existing methods when one factors in the computational cost. \n\n- The authors should have at least illustrated their method on simple toy examples where the posterior is available and multimodal, so that we can see if it indeed recovers the posterior completely. \n\n[3] Chung, Hyungjin, Jeongsol Kim, Michael T. Mccann, Marc L. Klasky, and Jong Chul Ye. \"Diffusion posterior sampling for general noisy inverse problems.\" ICLR (2023)\n\nQUESTIONS:\n- Figure 2, page 7. How come the present method scales better with image size? Shouldn't it be the opposite? \n\n- The main reason for the huge computational cost is the fact this method requires differentiating the score network with respect to its input (by the chain rule). This is because the authors seek an approximation of the log probability. However, it seems to me that what is actually need is $\\nabla_x \\log p_{\\theta}(x)$ (again by the chain rule). While this score is not available in practice, can't it be approximated by taking for example the score at a time $t$ close to $0$?\n\nLIMITATIONS:\nsee above.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAuthors propose a non-amortized variational inference approach to solve large-scale Bayesian inference problems where the prior is based on an cheap-to-evaluate approximation to a pretrained diffusion model.\n\nSTRENGTHS:\n**Originality.** This paper hits the nail on the head when it comes to large-scale Bayesian inference in the context of inverse problems, especially when diffusion models are used as priors.\n\n**Quality and clarity.** The paper is well-written and easy to follow. The authors have done a great job in explaining the proposed approach and the experiments are well-designed to demonstrate the effectiveness of the approach.\n\n**Significance.** This approach allows to leverage the power of diffusion models in approximating complex distributions in solving large-scale Bayesian inference problems. This is a significant contribution to the field and I am interested to apply this approach to another inverse problem domain.\n\nWEAKNESSES:\n* My primary concern revolves around the justification for selecting diffusion models as prior distributions over alternative generative models. It would greatly enhance the paper to thoroughly examine the advantages and disadvantages of diffusion models compared to other generative models, specifically within the framework of large-scale Bayesian inference. It would be ideal to include a comparison with amortized normalizing flows and/or injective flows. This raises the question: why not initially employ a normalizing flow to learn the the prior or full posterior distribution (amortized VI)?\n\nQUESTIONS:\n* I would like to know the authors' stance on amortized vs. non-amortized variational inference. What would be required to amortize their objective function, enabling the use of $q_{\\phi}$ to approximate the posterior distribution for any new observation after training? This approach has the potential to make training costs offline, enabling fast posterior sampling during test time.\n\nLIMITATIONS:\n* It would be beneficial to provide additional comments on the limitations that arise when dealing with \"out-of-distribution\" data (unknown being out of diffusion model distribution).", "TITLE": "Official Review", "RECOMMENDATION": 9, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}