Solving Inverse Problems with Latent Diffusion Models via Hard Data Consistency
Authors: Bowen Song
Soo Min Kwon
Zecheng Zhang
Xinyu Hu
Qing Qu
Liyue Shen
Venue: ICLR 2024 spotlight (2024)
Keywords: Diffusion models
inverse problems
Primary area: generative models
TL;DR: We show how to effectively leverage latent (or stable) diffusion as generative priors for solving general inverse problems.

Abstract: Latent diffusion models have been demonstrated to generate high-quality images, while offering efficiency in model training compared to diffusion models operating in the pixel space. However, incorporating latent diffusion models to solve inverse problems remains a challenging problem due to the nonlinearity of the encoder and decoder. To address these issues, we propose ReSample, an algorithm that can solve general inverse problems with pre-trained latent diffusion models. Our algorithm incorporates data consistency by solving an optimization problem during the reverse sampling process, a concept that we term as hard data consistency. Upon solving this optimization problem, we propose a novel resampling scheme to map the measurement-consistent sample back onto the noisy data manifold and theoretically demonstrate its benefits. Lastly, we apply our algorithm to solve a wide range of linear and nonlinear inverse problems in both natural and medical images, demonstrating that our approach outperforms existing state-of-the-art approaches, including those based on pixel-space diffusion models.

Published as a conference paper at ICLR 2024
SOLVING INVERSE PROBLEMS WITH LATENT DIFFU- SION MODELS VIA HARD DATA CONSISTENCY
Bowen Song'; Soo Min Kwon'r, Zecheng Zhang?, Xinyu Hu², Qing Qu' Liyue Shen Universitty of Michigan, AKumo.AI, 'Microsoft
ABSTRACT
Latent diffusion models have been demonstrated to generate high-quality images, while offering efficiency in model training compared to diffusion models operat- ing in the pixel space. However, incorporating latent diffusion models to solve inverse problems remains a challenging problem due to the nonlinearity of the en- coder and decoder. To address these issues, we propose ReSample, algorithm that can solve general inverse problems with pre-trained latent diffusion models. Our algorithm incorporates data consistency by solving an optimization problem during the reverse sampling process, a concept that we term as hard data consis- tency. Upon solving this optimization problem, we propose novel resampling scheme to map the measurement conssisent sample back onto the noisy data man- ifold and theoretically demonstrate benefits. Lastly, we apply our algorithm solve wide range of linear and nonlinear inverse problems in both natural and medical images, demonstrating that our approach outperforms existing state-of- the-art approaches, including those based on pixel-space diffusion models.
INTRODUCTION
Inverse problems arise from a wide range of applications across many domains, including compu- tational imaging (Beck & Teboulle| 2009 Afonso al., 2011). medical imaging (Suetens) 2017 Ravishankar et al. 2019)). and remote sensing (Liu et al. 2021 2022), to name a few. When solv- ing these inverse problems, the goal reconstruct an unknown signal x. R given observed measurements E R of the form
y = A(x.) 7)
where A(-) R + R. denotes some forward measurement operator (can be linear or nonlinear) and 17 E 12m is additive noise. Usually, we are interested in the case when m < which follows many real-world scenarios. When m < n, the problem is ill-posed and some kind of regularizer (or prior) is necessary to obtain meaningful solution.
In the literature, the traditional approach of using hand-crafted priors (e.g. sparsity) is slowly being replaced by rich, learned priors such as deep generative models. Recently. there has been lot of interests using diffusion models as structural priors due to their state-of-the-ae performance in image generation (Dhariwal & Nichol, 2021a Karras et al. 2022 Song al.. 2023b Lou & Ermon 2023ah. Compared to generative adversarial networks (GANs), diffusion models are generally casier and more stable to train, making them generative prior that is more readily accessible (Dhariwal & Nichol 2021b). The most common approach for using diffusion models as priors is to resort to posterior sampling. which has been extensively explored in the literature (Song et al.) 2022 Chung et al. 2023a] |023a|2022] (022]Kawar et al.] [2022] [Song et al., [223a CChung et 2003b]] MMng & Kabashima
2022 Zhang & Zhoul [h23) However, despite their remarkable success, these techniques exhibit several limitations. The primary challenge is that the majority of existing works train these models
directly the pixel space, which requires substantial computational resources and large volume of training data (Rombach et al. (022)..
Latent diffusion models (LDMs), which embed data in order operate in lower-dimensional space, present a potential solution to this challenge, along with considerable improvements in com- putational efficiency (Rombach al. 2022 Vahdat et 202T) by training diffusion models
"Equal Contribution; Corresponding authors (bowenhw, kwonsm] [wonsm]@@micc....o @umich.edu

---

Published as a conference paper at ICLR 2024
Nonlinear Deblurring
CT Reconstruction
Measurement
Reference
ReSample (Durs)
Measurement
Reference
ReSample (Ours)
Figure 1:
Example reconstructions of our algorithm (ReSample) on two noisy inverse problems, nonlinear deblurring and CT reconstruction, on natural and medical images, respectively.
compressed latent space. They can also provide a great amount flexibility, as they can enable transfer and generalize these models to different domains by fine-tuning small amounts of training data (Ruiz et al., 2023). Nevertheless, using LDMs to solve inverse problems poses a significant challenge. The main difficulty arises from the inherent nonlinearity and nonconvexity of the decoder, making it challenging to directly apply existing solvers designed for pixel space. To address this issue, concurrent work by Rout et (2023) recently introduced a posterior sam- pling algorithm operating in the latent space (PSLD), designed to solve linear inverse problems with provable guarantees. However, observe that PSLD may reconstruct images with artifacts in the presence of measurement noise. Therefore, developing an efficient algorithm capable of addressing these challenges remains an open research question.
In this work, we introduce a novel algorithm named ReSample, which effectively employs LDMs as priors for solving general inverse problems. Our algorithm can be viewed as a two-stage process that incorporates data consistency by (1) solving hard-constrained optimization problem, ensuring we obtain the correct latent variable that is consistent with the observed measurements, and (2) employing carefully designed resampling scheme to map the preasurement-consisse onto the correct noisy data manifold. As a result, we show that our algorithm can achieve state-of- the-art performance algorithms. Notably, owing to using the latent diffusion models as generative priors, our algorithm achieves reduction in memory complexity. Below, we highlight some of our key contributions.
sample back
various inverse problem tasks and different datasets, compared to existing
We propose a novel algorithm that enables general inverse problems (linear and nonlinear) through hard data consistency.
leverage latent diffusion models for solving
Particularly, carefully design stochastic resampling scheme that can reliably map the measurement-consistern samples back onto pling process. We provide a theoretical analysis further demonstrate the superiority of the proposed stochastic resampling technique.
the noisy data manifold to continue the reverse sam-
With extensive experiments on multiple tasks and various datasets, encompassing both natural and medical images, proposed algorithm achieves state-of-the-art performance variety of linear and nonlinear inverse problems.
BACKGROUND
Denoising Diffusion Probabilistic Models. We first briefly review the basic fundamentals of dif- fusion models, namely the denoising diffusion probabilistic model (DDPM) formulation (Ho al. 2020 Let Zy (data(x) denote samples from the data distribution. Diffusion models start by progressively perturbing data noise via Gaussian kernels, which can be written as the variance- preserving stochastic differential equation (VP-SDE) (Song et al. 2021) of the form
B1 da arl Vidwa,
where 3, E (0, is the noise schedule that is monotonically increasing sequence of and is the standard Wiener process. This is generally defined such that we obtain the data distribution when i 0 and obtain Gaussian distribution when T, i.e. Ør NO, I). The objective of diffusion
2

---

Published as a conference paper at ICLR 2024
unconditioal sample
XT
ZT
Z++11
Z
Zo Ztt
reverse process
prodict &
Latent Space
encoder
optimization
Sample
Xn
D
Zo
Zt
Zg (y
process
data consistent sample
decoder
Image
y
Our Aigutthin
Sptent
messurement
Figure 2: Overview of our ReSample algorithm during the reverse sampling process condi- tioned the data constraints from measurement. The entire sampling process is conducted in the latent space upon passing the sample through the encoder. The proposed algorithm performs hard data consistency at some time steps via skipped-step mechanism.
models is to learn the corresponding reverse SDE of Equation ① which is of the form
3t
dxc
B.V log logp(z;) dt + Bidw.
where dw is the standard Wiener process running backward in time and V a. log pl (x) is the (Stein) score function. In practice, approximate the score function using neural network 8o parameter- ized by 0, which can be trained via denoising score matching (Vincent 2011
arg min I [[] (x
()) t)
Va, log p( (20) [2]
where t is uniformly sampled from [0, T] and the expectation taken over t. æ, ~ p( (o). and æo (x). Once we have access to the parameterized score function Se. approximate the reverse-time SDE and simulate it using numerical solvers (e.g. Fuler-Maruyama).
it to
Denoising Diffusion Implicit Models. As the DDPM formulation is known to have slow sam- pling process, Song et al. (2020) proposed denoising diffusion implicit models (DDIMs) that defines the diffusion process non-Markovian process remedy this (Ho et al., 2020; Song al. 2020 Lu et al. 2022). This enables a faster sampling process with the sampling steps given by
20236
84 = V 10
ar-1 - !957 so (r. 1) + noee, T=
(4)
where 04 = 3r, ā. I ai No0, I), is the temperature parameter, & controls the stochasticity of the update step, and 2o denotes the predicted zn from æ, which takes the form
(x.)))
(x. VI 0, (x,, t)):
(5)
Va
which is an application of Tweedie's formula. Here, 8p usually trained using the epsilon-matehhee score objective (Song et al.) (020). We use DDIM as the backbone of our algorithm and show how we can leverage these update steps for solving inverse problems.
Solving Inverse Problems with Diffusion Models. Given measurements y R from some forward measurement operator A(-), we can use diffusion models to solve inverse problems by replacing the score function in Equation (2) with the conditional score function Vz, p(x.l)) Then by Bayes rule, notice that we can write the conditional score as
Va, logp(x; = Va. logp(æ;) + Var logp(y(<x.)
This results in the reverse SDE the form
3r
dx
B.(V,, logp(x:) V.n. logp(y)(:))) dt v Bidw.
2

---

Published as a conference paper at ICLR 2024
In the literature, solving this reverse SDE is referred as posterior sampling. However, the issue with posterior sampling is that there does not exist an analytical formulation for the likelihood term V., log p(y (y (z)) To resolve this, there exists two lines of work: (1) to resort to alternating pro- jections onto the measurement subspace to avoid using the likelihood directly (Chung et al.) 2022. Kawar et al. 2022 Wang al.] (2022) and (2) estimate the likelihood under some mild assump- tions (Chung et al. |023a]|Snna et al. 2023a) For example, Chung et al. (2023a) proposed diffusion
posterior sampling (DPS) that uses Laplacian approximation of the likelihood, which results in the
discrete update steps
1t-1 = 100 V 0:-1 00 + 17ōE6
(6)
T;-1 = 2/-11~1 <V (()0(2)))1
(7)
where Ç E IR can be viewed tunable step-size. However, as previously mentioned, these tech- niques have limited applicability for real-world problems as they are all built on the pixel space.
Solving Inverse Problems with Latent Diffusion Models. The limited applicability of pixel- based diffusion models can be tackled by alternatively utilizing more efficient LDMs as generative priors. The setup for LDMs the following: given an image x: E IR" we have an encoder & R > R and decoder D : R + R" where k < n. Let z ((x) E I denote the embedded samples the latent space, One way of incorporating LDMs to solve inverse problems would be to replace
the update steps in Equation ([) with
z = 20 V 1- ū, (0880(Z2;0) 1) no,e,
(8)
21-1 = -
((D(Eo(2:)))1.
(9)
After incorporating LDMs, this can be viewed as a non-linear inverse problem due to the non- linearity of the decoder D(-). As this builds upon the idea behind DPS, we refer to this algorithm as Latent-DPS. While this formulation seems to work, we empirically observe that Latent-DPS often produces reconstructions that are often noisy blurry and inconsistent with the measurements. We conjecture that since the forward operator involving the decoder is highly nonconvex, the gradient update may lead towards local minimum. We provide more insights in Appendix Section D.
RESAMPLE:: INVERSE PROBLEMS USING LATENT DIFFUSION MODELS
3.1 PROPOSED METHOО
Hard Data Consistency. Similar to Latent-DPS, our algorithm involves incorporating data con- sistency into the reverse sampling process of LDMs. However, rather than gradient update as shown in Equation ( we propose to solve an optimization problem on some time steps l:
1 * 2 A(D(=)))
Ey( E arg min
(10)
where we denote εn(y) as the sample consistent with the measurements y E R" This optimization problem has been previously explored in other works that use GANs for solving inverse problems, and can be efficiently solved using iterative solvers such as gradient descent (Bora et al., 2017] Jalal et al. 20211 Shah al. 1021 Lempitsky et al.) 2018). However, it is well known that solv- ing this problem starting from a random initial point may lead to unfavorable local minima (Bora et al. 2017). To address this, solve Equation (10 starting from initial point 20(2+41), where 20 (1(+1) is the estimate of ground-truth latent vector at time 0 based on the sample at time t + 1. The intuition behind this initialization that we want to start the optimization process within local proximity of the global solution of Equation (TO) to prevent resulting in a local minimum. We term this overall concept as hard data consistency, as we strictly enforce the measurements via optimiza- tion, rather than a "soft" approach through gradient update like Latent-DPS. To obtain &0 (2+41), we use Tweedie's ss formula (Efron) (011)) that gives us an approximation of the posterior mean which takes the following formula:
20(24) E E[zo|z,]
(z+ + (1 ax)) Vlogp(z)))
(11)
Vov
4

---

Published as a conference paper at ICLR 2024
Algorithm 1 ReSample: Solving Inverse Problems with Latent Diffusion Models
Require: Measurements y. A(-). Encoder E(),, Decoder D();, Score function so(-, t), Pretrained LDM Parameters Br, Oez V/, 6, Hyperparameter to control of Time steps to perform resample C Initial noise vector
2r (O, I)
for 1 = 7 1,...,,0 do
E1 (0
ε1+1 50( t +)
Compute the score D. Predict En using Tweedie's formula
20(2(+1) IillI
(2141 /1 - 0(416(42)
z/ = / 20,20(2 (2++1) vI -1-a; āt - n012 + n/E1
> Unconditional DDIM step > ReSammle time step b Solve with initial point &n (z(+1)
if E C then
(y( E arg - (((--(D(z))Il)
2,- StochasticResample (oo Eo)))
Map back to t
else
2t z/{ 2n = D(zo)
>> Unconditional sampling not resampling > Output reconstructed image
However, we would like to note that performing hard data consistency on every reverse sampling iteration t may be very costly. To address this, we first observe that as we approach f T, the estimated 20(z++1) can deviate significantly from the ground truth. In this regime, we find that hard data consistency provides only marginal benefits. Additionally, in the literature, existing works point out the existence of a three-stage phenomenon (Yu et al. (023), where they demonstrate that data consistency primarily beneficial for the semantic and refinement stages (the latter two stages when I. is closer to 0) of the sampling process. Following this reasoning, we divide T into three sub-intervals and only apply the optimization in the latter two intervals. This approach provides both computational efficiency and accurate estimates of (0(+1))).
Furthermore, even during these two intervals, we observe that we do not need to solve the opti- mization problem every iteration f. Because of the continuity of the sampling process, after each cata-consistency optimization step, the samples in the following steps can retain similar semantic or structural information to some extent. Thus, we "reinforee" the data consistency constraint during the sampling process via a skipped-step mechanism. Empirically, we see that it is sufficient to per- form this on every 10 (or so) iterations of f. One can think of hard data consistency as guiding the sampling process towards the ground truth signal x (or respectively x²) such that it is consistent with the given measurements. Lastly, in the presence of measurement noise, minimizing Equa- tion (10) to zero loss can lead to overfitting the noise. To remedy this, we perform early stopping, where we only minimize up threshold based on the noise level. We will discuss the details of the optimization process in the Appendix. We also observe that additional Latent-DPS step after unconditional sampling can (sometimes) marginally increase the overall performance. We perform an ablation study on the performance of including Latent-DPS in the Appendix.
Remapping Back to 21. Following the flowchart in Figure the next step is to map the measuremennt-onnssseenn sample εo(y) back onto the data manifold defined by the noisy samples at time to continue the reverse sampling process. Doing so would be equivalent to computing the posterior distribution p(z p(z))y). To incorporate Eo(y) into the posterior, we propose to construct an auxiliary distribution p(&, (3)) y) to replace p(z: Here, & denotes the remapped sam- ple and 2 denotes the unconditional sample before remapping. One simple way of computing this distribution to obtain &, is shown in Proposition 1
Proposition 1 (Stochastic Encoding). Since the sample ân given 2o(y) and measurement y is con- ditionally independent of y. we have that
= p(εc)á((y) y) = = (1- T.)
(12)
We defer all of the proofs to the Appendix. P provides us a way of computing ε.. which refer to as stochastic encoding. However, observe that using stochastic encoding can incur a high variance when is farther away from t = 0, where the ground truth signal exists. This large variance can often lead to noisy image reconstructions. To address this issue, propose posterior
5

---

Published as a conference paper at ICLR 2024
Super Resolution 1x
Inpainting (Kandom 70%)
Method
LPIPS
PSNR
SSIMT
LPIPS!
PSNR
SSIMT
DPS hung et
20220
0.113
26.41 2.20
0.782
0.00
0.102
0.02
2.30
11:99
CingChung ef 20z2
0.193 + 0.03
25.92
2.35
0.740 + 0.05
0.134 + 0.03
29.53
11.3++22 2.70
0.847 + 0.05
ADMM--PP Almmal ara
2019
0.304 = 0.04
21.08 3.13
0.631 0.11
0.627 = 0.07
15.40
1::00:2000
0.312 + 0.09
DDRM Kawar al 2222
0.151 1: 0.03
29.49 1 1.93
0.817
1: 0.05
0.166 0,03
27.69 1 1.34
0.796 1: 0.04
DMPS Meng & Kabashima 2022
0.147 1 0.03
28.45
1.92
0.811 1: 0.05
0.175
0.03
28.84
1.65
0.826 1 0.03
Latent-DPSS
0.272 1 0.05
2.00
0.690
0.07
0.226
1.84
0.703
1 0.07
PSLDD LDM' Rout et al.
2025
0.209 1 0.10
27.61
1 2,95
0.704
1 0.17
0.200 1 0.08
27.07
1 2.45
0,089
0.11
ReSample (Ours)
0.144 = 0.029
30.45
: 2.09
0.832 3: 0.05
0.082 = 0.02
32.77
2 2.23
0.903 0.03
Table I: Quantitative results of super resolution and inpainting on the CelebA-HQ dataset.
Input images have additive Gaussian noise with dy = 0.01. Best results are in bold and second best results are underlined.
sampling technique that reduces the variance by additionally conditioning on z/. the unconditional sample at time L. Here, the intuition is that by using information of z/, we can get closer to the ground truth 2,, which effectively reduces the variance, Lemma under some mild assumptions, show that this new distribution p(2 (o(y), 3) is tractable Gaussian distribution.
Proposition (Stochastic Resampling). Suppose that p(z)) & 2o( y) is normally distributed such that 2 (V)),, y) = N (He, ot). If we let 2(21(2((1) y) be prior for phg then the posterior distribution p(ž, ( Eo(y), 1) is given by
oo Vivo, 2, Eoly)) +1 az)z of ax))
of (1-())))
p( Eolg) y) = N
(13)
of + (1 (1- xt)
We refer to this new mapping technique as stochastic resampling. Since we do not have access to of. it serves as a hyperparameter that we tune our algorithm. The choice of of plays a role of controlling the tradeoff between prior consistency and data consistency. If o² - 0, then we recover unconditional sampling, and if of + oo, we recover stochastic encoding. We observe that this new technique also has several desirable properties, for which we rigorously prove in the next section.
3.2 THEORETICAL RESULTS
In Section 3.1 we discussed that stochastic resampling induces less variance than stochastic encod- ing. Here, we aim to rigorously prove the validity of this statement.
Lemma 1. Let in and ê, denote the stochastically encoded and resampled image of Lo(y), respec- fively. If (TVR((2)) z >> 0, then we have that VAR( < &VA(22) E.). Theorem 1. If (0y) Ey(y), then stochastic resample is unbiased such that B[ y = E[z(].
meeasurement-consistee such that y A(D(Eo())) i.e. 20 - 20(
These two results, Lemma D and Theorem prove the benefits of stochastic resampling. At high-level, these proofs rely on the fact the posterior distributions of both stochastic encoding and resampling Gaussian and compare their respective means and variances. In the following result, we characterize the variance induced by stochastic resampling, and show that as t - O, the variance decreases, giving reconstructed image that is of better quality.
Theorem 2. Let zn denote a sample from the data distribution and 2, be a sample from the noisy perturbed distribution at time t. Then,
Cov(zulzı) (1= V2 (z) + 1 σt I.
āt
By Theorem notice that since as a₁ is an increasing sequence that converges to I as t decreases, the variance between the ground truth 2g and the estimated εσ decreases to as -> O, assuming theory, empirically show that stochastic resampling can reconstruct signals that are less noisy than stochastic encoding, as shown in the next section.
that 22 log Pz, (21) < oc. Following
EXPERIMENTS
We conduct experiments to solve both linear and nonlinear inverse problems natural and medical images. We compare our algorithm to several of-re--f-th-ann methods that directly apply the diffu- sion models that are trained in the pixel space: DPS (Chung et al. 2023a), Manifold Constrained

---

Published as a conference paper at ICLR 2024
Nonlinear Deblurring
Gaussian Dehluring
Method
LPIPS|
PSNR1
SSIM
LPIPSI
PSNR
SSIMT
Chung et
7023a
0.230 0.065
26.81 + 2.84
0.720 0.077
0.175 + 0.03
28.36
2.12
0.772 + 0.07
MCG of 2022
0.517 + 0.06
15.85 = 1.08
0.536 + 0.08
ADMM-
2019
0.199
0.073
16.17 + 1.01
0.359 0.140
0,289 = 0.04
20,98
4.51
0.602 0:15
DDRM Kawar el al 20:2
0.193 1 0.04
26.88
1.96
0.747 1: 0.07
DMPS Meng Kabashima
2022
0.206 0.04
26.45
1.83
0.726 0.07
Latent-DD
0.225 1 0.04
26.18 1: 1.73
0.703 0.07
0.205 + 0.04
27.42
1: 1.84
0.729 1 0.07
PSI.D- Rout al.
2023
0.323 4 0.09
24.21
2.79
0.548
1 0.12
ReSample (Ours)
0.153
30.18
2:21
1
0.04
30.69
2.14
0.832 0.05
Table 2: Quantitative results of Gaussian and nonlinear deblurring on the CelebA-HQ dataset. Input images have an additive Gaussian noise with dy = 0.01. Best results are in bold and second best results are underlined. For nonlinear deblurring, some basclines are omitted, as they can only
solve linear inverse problems.
Ahdominal
Head
Chest
Method
PSNRT
SSIM
PSNRT
SSIM
PSNR
SSIMI
Latent-DIS
26,80 1100
0,870 #0.026
28.64 #5.38
0.893
+0,058
25.67#1.14
0.822 #0.033
MCG Chung et al. 2022
29.41 #3.14
0.857
10.041
28.28 +3.08
0.795
10.116
27.92 -2448
0.842 100033
DPS Chung al 20238
27.33 +2.68
0.033
24.51 +2.77
0.665
10.058
24.73
-884
0.682 10.113
UnP-UNNe (iilonn al. 2021
52.84 1.22
0.942
+0.008
33,45 +3.25
0.945
+0.023
29.67 =1.14
0.891
#0.011
FBP
26.29
11.21
0.727
10.036
26.71
#5.02
0.725
10.106
24.12 1114
0.655
10.033
FBP-UNet (Jin et al.
2017
32.77
#1.21
0.937
1:0.013
31.95 +3.32
0.917 +0.048
29.78
1112
0.885 +0.016
ReSample (Ours)
35,91
11.22
0,965 100000
37.82 13.31
0,978 10.014
31.72 00912
0,922 10:011
Table 3: Quantitative results of CT reconstruction on the LDCT dataset. Best results are in bold and second best results are underlined.
Gradients (MCG) (Chung et al. 2022). Denoising Diffusion Destoration Models (DDRM) (Kawar ett al. 2022), Diffusion Model Posterior Sampling (DMPS) (Meng & Kabashimal 2022). Then we compare algorithm that uses plug-and-play approach that apply a pretrained deep denoiser for inverse problems (ADMM-PnP) (Ahmad et al., 2019). We also compare our algorithm to LaP and Posterior Sampling with Latent Diffusion (PSLD) (Rout al. 2023). a concurrent work we recently notice also tackling latent diffusion models. Various quantitative metrics are used for evaluation including Learned Perceptual Image Patch Similarity (LPIPS) distance, peak signal-to-noise rotte---raatio (PSNR), and structural similarity index (SSIM). Lastly, we conduct ablation study to compare the performance of stochastic encoding and our proposed stochastic resampling technique as mentioned in Section 3.1 and also demonstrate the memory efficiency gained by leveraging LDMs.
Experiments on Natural Images. For the experiments on natural images, we use datasets FFHQ (Kazemi& & Sullivan 2014). CelebA-HQ (Liu et al.] (015). and LSUN-Bedroom (Yu et al. 2016) with the image resolution of 256 x 256 x We take trained latent diffusion models LDM- VQ4 trained on FFHQ and CelebA-HQ provided by (Rombach et al. 2022) with autoencoders that yield images of size 64 x 64 x 3, and DDPMs (Ho et al. (020) also trained on FFHQ and CelebA- HQ training sets. Then, we sample 100 images from both the FFHQ and CelebA-HQ validation sets for testing evaluation. For computing quantitative results, all images are normalized to the range [0, 1]. All experiments had Gaussian measurement noise with standard deviation By = 0.01. Due to limited space, we put the results on FFHQ and details of the hyperparameters to the Appendix.
For linear inverse problems, we consider the following tasks: (1) Gaussian deblurring, (2) inpainting (with a random mask), and (3) super resolution. For Gaussian deblurring, we use kernel with size 61 x 61 with standard deviation 3.0. For super resolution, we use bicubic downsampling and a random mask with varying levels of missing pixels for inpainting. For nonlinear inverse problems, we consider nonlinear deblurring as proposed by (Chung et al.) (023a). The quantitative results are displayed in Tables [ and [ ReSample significantly outperforms all of the basclines across all three metrics on the CelebA-HO dataset. We also observe that ReSample performs better than or comparable to all baselines on the FFHQ dataset as shown in the Appendix. Remarkably, our method excels in handling nonlinear inverse problems, further demonstrating the flexibility of our algorithm. We further demonstrate the superiority of ReSample for handling nonlinear inverse problems in Figure ba where we show that can consistently outperform DPS.
with qualitative results in Figure 3 In Tables and 2 we see that
We have updated the baseline results for PSL.D. More details are provided in the Appendix (Section A.4
7

---

Published as a conference paper at ICLR 2024
Mesuumment
Referance
DPS
Lifent ors
Rehamplk
Mesummmett
Rereerence
DPS
Latent DPS
Recamolic
INIE Aooeett
Super
Noditesrr Dabkaring
Gaavia Callariyy
Figure 3: Qualitative results of multiple tasks on the LSUN-Bedroom and CelebA-HQ datasets.
All inverse problems have Gaussian measurement noise with variance By 0.01.
FBP
FBP-UNet
ADMM-PnP
Latent-DPS
MCG
Ours
Ground Truth
Figure 4: Qualitative results of CT reconstruction the LDCT dataset. We annotate the critical image structures in a red box, and zoom in below the image.
Experiments on Medical Images. For experiments on medical images, we fine-tune LDMs on 2000 2D computed tomography (CT) images with image resolution of 256 x 256, randomly sampled from the AAPM LDCT dataset (Moen et al. (02T) I) of 40 patients, and test on the 300 2D CT images from the remaining patients. We take the FFHQ LDMM-VQ model provided by Roobach et al. (202)) as the pre-trained LDM followed by 100K fine-tuning iterations. Then, simulate CT measurements (sinograms) with a parallel-beam geometry using 25 projection angles equally distributed across 180 degrees using the torch-radon package (Ronchetti) 2020). Following the natural images, the sinograms were perturbed with additive Gaussian measurement noise with Fy 0.01. Along with the baselines used in the natural image experiments, we also compare to the following methods: (1) Filtered Backprojection (FBP), which is a standard CT reconstruction technique, and (2) FBP-UNet, a supervised method that uses UNet model as its backbone. For FBP-Unet and PnP-UNet. we trained model on 3480 2D CT from the training set. For MCG and DPS, we used the pre-trained checkpoints provided by Chung al. (2022). We present visual reconstructions with highlighted critical anatomic structures in Figure 4 and quantitative results in T Here, we observe that our algorithm outperforms all of the baselines in terms of PSNR and SSIM. Visually, we also observe that our algorithm is able to reconstruct smoother images that have more accurate and sharper details.
Effectiveness of the Resampling Technique. Here, we validate our theoretical results by con- ducting ablation studies on stochastic resampling. Specifically, we perform experiments on the LSUN-Bedroom and CelebA-HQ datasets with tasks of Gaussian deblurring and super-resolution. As shown in Figure [ we observe that stochastic resampling reconstructs smoother images with higher PSNRs compared stochastic encoding, corroborating our theory.
8

---

Published as a conference paper at ICLR 2024
Reference
ReSample
Encoding
Reference
Rešample
Encoding
rSWe: 55
PSWR: 38
PSM:: 23.77
PSNR: 28.05
21
PSNE: 20-32
PSNR: 30.35
PSM8:
Figure 5: Effectiveness of our resampling technique compared to stochastic encoding, Results were demonstrated on the By = 0.05 to highlight the effectiveness of stochastic resample.
and CelebA-HQ datasets with measurement noise of
Measurement
Reference
DPS
Latent-DPS
ReSample
29
28
10
30
40
Fumberrs of tim stops to perform optimizatioon and resample
(b) Performance v.s. number of time steps to perform resample
(a) More examples of nonlinear deblurring
Figure 6: Left: Additional results on nonlinear deblurring highlighting the performance of ReSam-
plc. Right: Ablation study on the ReSample frequency on the performance.
Effectiveness Hard Data Consistency. In Figure 6b Sample frequency on CT reconstruction. This observation is in line with what we expect intuitively, as more ReSample time steps (i.e., more data consistency) lead to more accurate reconstructions.
we perform an ablation study on the Re-
Memory Efficiency. To demon- strate memory efficiency, we use the command nvidia-smi to moni- tor the memory consumption dur- ing solving an inverse problem. We present the memory usage for Gaus- sian deblurring on the FFHQ dataset in Table Although the entire LDM models occupy more memory due to algorithm it- self exhibits memory efficiency, resulting in lower overall memory usage. This highlights its po- tential domains like medical imaging, where memory plays crucial role in feasibility.
Table 4: Memory Usage of Different Methods for Gaussian
Deblurring on the FFHQ Dataset
Model DDPM
Algorithm
Model Only
Memory Increment
Total
DPS
1953MB
+3416MB (175%)
5369MB
MCG
MB421MB (175%)
5374MB
DMPS
+5215MB (267 %)
7168MB
DDRM
+18833MB (964 %-)
20786MB
LDM
PSLD
3969MB
+5516MB (140%)
9485MB
ReSample
+1040MB (26.2%)
5009MB
the autoencoders,
CONCLUSION
In this paper, we propose ReSample, algorithm that can effectively leverage LDMs to solve gen- eral inverse problems. We demonstrated that our algorithm can reconstruct high-quality images compared to many baselines, including those in the pixel space. One limitation of our method lies in the computational overhead of hard data consistency, which we leave as a significant challenge for future work to address and improve upon.

---

Published as a conference paper at ICLR 2024
6 REPRODUCIBILITY STATEMENT
To ensure the reproducibility of our results, we thoroughly detail the hyperparameters employed in our algorithm the Appendix. Additionally, we provide a comprehensive explanation of the configuration of all baselines used in our experiments. As we use pre-trained diffusion models throughout our experiments, they are readily accessible online. Lastly, our code is available at https://g.thien comp com/ www.moooowoooo... www.swwww.........ee.......
7 ACKNOWLEDGEMENTS
BS and I.S acknowledges support from U-M MIDAS PODS Grant and U-M MICDE Catalyst Grant, and computing resource support from NSF ACCESS Program and Google Cloud Research Credits Program. This work used NCSA Delta GPU through allocation CIS230133 13 and ELE230011 from the Advanced Cyberinfrastructure Coordination Ecosystem: Services & Support (ACCESS) pro- gram, which supported by National Science Foundation grants #2138259, #2138288 #21 #2138307, #2137603, and #2138296. SMK and QQ acknowledge support from U-M START & PODS grants, NSF CAREER CCF-2143904, NSF CCF-221 NSF CCF-2212326, ONR N00014-22- AWS AI Award, and gift grant from KLA..
REFERENCES
Manya V. Afonso, José M. Bioucas-Dias, and Mário A. T. Figueiredo. An augmented lagrangian approach to the constrained optimization formulation of imaging inverse problems. IEEE Trans- actions on Image Processing, 20(3):681-695, 2011. doi: 10.1 10.100//I12210022766244
Rizwan Ahmad, Charles A Bouman, Gregery T Buzzard, Stanley H Chan, Edward T' Reehorst, and Philip Schniter. Plug and play methods for magnetic resonance imaging. 2019.
Amir Beck and Mare Teboulle. fast iterative thiesoolggg algorithm for linear inverse problems. SIAM Journal Imaging Sciences, 2(1):183-200) 202, 2009. doi: 10. 10.1137/080716542. URL http::////oii.o11111111111111111
Ashish Bora, Ajil Jalal, Eric Price, and Alexandros Dimakis. Compressed sensing using genera- tive models. In International Conference on Machine Learning, pp. 537-546. PMLR, 2017.
Paul A. Bromiley. Products and convolutions of gaussian probability density functions. 2013.
Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. arXiv preprint arXiv:2206.0094) 2022.
Hyungjin Chung. Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In The Fleventh Interna- fional Conference on Learning Representations, 2023a. URL httttss://ppeeeeeeeeeeee
net/
Hyungjin Chung, Suhyeon Lee, and Jong Chul Ye. Fast diffusion sampler for inverse problems by geometric decomposition. arXiv preprint arXiv:2303.05754, 2023b.
Hyungjin Chung. Jong Chul Ye, Peyman Milanfar, and Mauricio Delbracio. Prompt-tuning latent diffusion models for inverse problems. arXiv preprint arXiv:2310.01110. 2023c.
Prafulla Dhariwal and Alex Nichol. Diffusion models beat gans on image synthesis. arXiv preprint arXiv:2105.05233, 202
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. Advances in Neural Information Processing Systems, 34:8780-8794 2021b.
Bradley Efron. Tweedie's 's formula and selection bias. Journal of the American Statistical Asso- ciation, 106(496) 2011. doi: 10. 198811999....11111 20 1111 URL https: :/doi. aTq/10.11190 2011
2011 Em11181
10

---

Published as a conference paper at ICLR 2024
Nic Fishman, Leo Klamer, Valentin De Bortoli, Emile Mathieu, and Michael Hutchinson. Diffusion models for constrained domains. arXiv preprint arXiv:2304.05364, 2023.
Davis Gilton, Gregory Ongie, and Rebecca Willett. Model adaptation for inverse problems in imag- ing. IEEE Transactions Computational Imaging, 7:661-674, 2021.
Harshit Gupta, Kyong Hwan Jin, Ha Q Nguyen, Michael McCann, and Michael Unser. Cnn-based projected gradient descent for consistent ct image reconstruction. IEEE transactions on medical imaging. (7(6):1440- 1453, 2018
Yoseob Han and Jong Chul Ye. Framing u-net via deep convolutional framelets: Application to sparse-view ct. IEEE transactions on medical imaging, 37(6):1418-1429, 2018.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
Shady Abu Hussein, Tom Tirer, and Raja Giryes. Image-adaptive gan based reconstruction. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 3121-3129, 2020.
Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros Dimakis, and Jon Tamir. Robust compressed sensing mri with deep generative priors. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances Neu- Curran Associates, Inc., wwt.ffeeepeeepe22111111 ral Information Processing Systems, volume 34, pp. 14938-14954,
2021. URL htttps:/ htw.://weeeddnnn...............e neurips File/ http://6644ee9111111111111111113
Kyong Hwan Jin, Michael T McCann, Emmanuel Froustey, and Michael Unser. Deep convolutional neural network for inverse problems in imaging. IEEE Transactions on Image Processing, 26(9): 4509-4522, 2017.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. arXiv preprint 2wXiv:2206.00364, 2022.
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. arXiv preprint arXiv:2201.1179) 2022.
Vahid Kazemi and Josephine Sullivan. One millisccond face alignment with an ensemble of regres- sion trees. In 2014 IEEE Conference Computer Vision and Pattern Recognition, pp. 1867- 1874, 2014. doi: 10. 202/CCVRR.20142222
Victor Lempitsky, Andrea Vedaldi, and Dmitry Ulyanov. Deep image prior. In 2018 TEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 9446-9454, 2018. doi: 10.1109/
Jingyun Liang, Jiezhang Cao, Guolei Sun, Kai Zhang, Lue Van Gool, and Radu Timofte. Swinir: Image restoration using swin transformer. In Proceedings of the IEEE/CVF international confer- ence on computer vision, pp. 1833-1844, 2021.
Guan- Hoan-HHongul Liu, Tianrong Chen, Evangelos A Theodorou, and Molei Tao. Mirror diffusion models for constrained and watermarked generation. arXiv preprint arXiv:2310.01236, 2023.
Wei Liu, Xin Xia, Lu Xiong, Yishi Lu, Letian Gao, and Zhuoping Yu. Automated vehicle sideslip angle estimation considering signal measurement characteristic. IEEE Sensors Journal, 21(19): 21675-21687, 2021.
Wei Liu, Karoll Quijano, and Melba M Crawford. Yolov5-tassel: detecting tassels in rgb uav im- agery with improved yolov5 based on transfer learning. IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing, 15:8085-8094 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), 2015.
Aaron Lou and Stefano Ermon. Reflected diffusion models. arXiv preprint arX1w:2304.04740, 2023a.
11

---

Published as a conference paper at ICLR 2024
Aaron Lou and Stefano Ermon. Reflected diffusion models. arXiv preprint arXiv:2304.04440, 2023b.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. arXiv preprint arXiv:2206.0099 2022.
Xiangming Meng and Yoshiyuki Kabashima. Diffusson model based posterior sampling for noisy linear inverse problems. arXiv preprint arXiv:2211.12343 2022.
Taylor R Moen, Baiyu Chen, David R Holmes III, Xinhui Duan, Zhicong Yu, Lifeng Yu, Shuai Leng, Joel Fletcher, and Cynthia H McCollough. Low-dose ct image and projection dataset. Medical physics, 48(2):902-911, 2021.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pp. 8162-8171. PMLR, 2021.
Saiprasad Ravishankar, Jong Chul Ye, and Jeffrey Fessler. Image reconstruction: From sparsity to data-adaptive methods and machine learning. Proceedings of the IEEE, 108(1):86-109 2019.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. Proceedings of the IEEE/CVF Con- ference on Computer Vision and Pattern Recognition, pp. 10684-10695, 2022.
Matteo Ronchetti, Torchradon: Fast differentiable routines for computed tomography. arXiv preprint (iXiv:2009.1478: 2020.
Litu Rout, Negin Raoof. Giannis Daras, Constantine Caramanis, Alexandros G. Dimakis, and Sanjay Shakkottai. Solving linear inverse problems provably via posterior sampling with latent diffusion models. arXiv preprint arXiv:2307.0065 1,069,,,,,9,,,,, 2023.
Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman. Dreambooth: Fine tuning text-lo-image diffusion models for subject-driven generation. In Pro- ceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 22500- 22510, 2023.
Viraj Shah, Rakib Hyder, M. Salman Asif, and Chinmay Hegde. Provably convergent algorithms for solving inverse problems using generative models. arXiv preprint arXiv:2105.06371, 2021.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02500 2020.
Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse-guiee diffusion models for inverse problems. In International Conference on Learning Representations, 2023a.
Yang Song. Jascha
Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011 13456. 2021.
Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imag- ing with score-based generative models. International Conference on Learning Representations, 2022. URL http://pppppneeeee.eeeeeeeerrtra
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01466, 2023b.
Paul Suetens. Fundamentals of medical imaging. Cambridge university press, 2017.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 11287-11302. Curran Associates, Inc., 2021. URL htttss///pppceeeenn...ee..oooooo hptp://007111333ccc..22222000030
fitpsperrr_e_e_______eeeeeeessst files/
12

---

Published as a conference paper at ICLR 2024
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Compu- tation, 23(7):1611-167, 2011. doi: 0.1162/NECO.a.00111
Yinhuai Wang, Jiwen Yu, and Jian Zhang. Zero-shot image restoration using denoising diffusion 2022.
null-space model. arXiv preprint
Haoyu Wei, Florian Schiffers, Tobias Würfl, Daming Shen, Daniel Kim, Aggelos Katsaggelos, and Oliver Cossairt. 2-step sparse-viey view ct reconstruction with a domain-specific perceptual net- work. arXiv preprint MiXXi:20022047433 2020.
Fisher Yu, Ari Seff. Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao. Lsun: Construction of large-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv: 0336,, 2016.
Jiwen Yu, Yinhuai Wang, Chen Zhao, Bernard Ghanem, and Jian Zhang. Freedom: Training- free energy-guided conditional diffusion model. arXiv preprint arXhv:2303.09833 2023.
Dan Zhang and Fangfang Zhou Self-supervised image denoising for real-world images with context-aware transformer. IEEE Access, 11:14340-14349 2023.
Kai Zhang, Wangmeng Zuo, Yunjin Chen, Deyu Meng, and Lei Zhang. Beyond gaussian denoiser: Residual learning of deep cnn for image denoising. IEEE transactions on image processing, 26 (7):3142-313) 3155, 2017.
Bo Zhu, Jeremiah Z Liu, Stephen F Cauley, Bruce R Rosen, and Matthew S Rosen. Image recon- struction by domain-transform manifold learning, Nature, 55577697744877-492 2018.
13

---

Published as a conference paper at ICLR 2024
Appendix
In this section, we present additional experimental results to supplement those presented in the main paper. In Section A we present additional qualitative and quantitative results to highlight the performance of ReSample. In Section B we briefly discuss some of the implementation details regarding hard data consistency. In Section C we outline all of the hyperparameters used to produce our results. In Section D we discuss some
reasons to why Latent-DPS fails to consistently accurately recover the underlying image. Lastly, in Section we present our deferred proofs.
A ADDITIONAL RESULTS
A.1 RESULTS ON FFHQ
Here, we provide additional quantitative results on the FFHQ dataset. Similar to the results in the main text the CelebA-HQ dataset, ReSample outperforms the baselines across many different tasks.
Super Resolution 4x
Inpainting (Random 70%(
Method
LPIPS
PSNR
SSIM
LPIPS
PSNR1
SSIM
DPS
0.175 1 0.04
28.47
2.09
0.793
0.03
32.32 1 2:21
0.897
MCG
0.223 0.04
2.174+3.18
0.673 + 0.07
0.178
0.03
24.89 + 2.37
0.731
0.06
ADMM-PMP
0.303 + 0.03
21.30 3.98
0.760 0.04
0.308
0.07
15.87
2.09
0.608 0.09
DDRM
0.257 = 0.03
1.79
0.05
0.287 0.03
2
0.680 0.05
DMPS
0.181 + 0.03
1.92
0.766
0.150 0.03
28.17 1.58
0.814 0.04
Latent-DPS
0.344 + 0.05
1.77
0.00
0.270
27.08 + 1.97
0.727
0.06
PSLD
0.580 + 0.05
1.58
0.07
0.270 0.06
25.6
+ 2.02
0.630 0.09
ReSample (Ours)
0.164 0.03
28.90
0.804
0.099 + 0.02
2.11
0.890 0.03
Table Comparison of quantitative results on inverse problems on the FFHQ dataset. Input images have additive Gaussian noise with By 0.01. Best results are in bold and second best results underlined.
Nonlinear Deblurring PSNR
Gaussian Deblurring PSNR
Method
LPIPS (L197 + 0.06
SSIM
LPIPS
SSIM
DPS
27.13 + 3.11
0.762 - 0.08
0.169 + 0.04
27.70 + 2.04
0.774 +
0.005 0.05
MCG
0.371 + 0.04
25.33 1.74
0.668 + 0.06
ADMM-PaP
0.424 + 0.03
21.80 + 1.42
0.497 = 0.05
0.399 + 0.03
21.23 3.01
0.675 + 0.06
DDRM
0.299 1 0.039
26.31
1.67
0.702 1 0.06
DMPS
0.227 0.036
26.04
1.75
0.699
0.066
Latent-DPS
0.319 +: 0.05
24.52
1.75
0.637 = 0.07
0.258
0.05
25.98
1.66
0.704
0.06
PSLD
0.422
20.08
2.97
0.400
0.16
ReSample (Ours)
0.171 1 0.04 30.03
1.78 0.838
1 0.04
0.201 1 0.04
28.73 1 1.87
0.794 1 0.05
Table 6: Comparison of quantitative results on inverse problems on the FFHQ dataset. Input images have an additive Gaussian noise with By 0.01. Best results are in bold and second best results are underlined.
We observe that ReSample outperforms all baselines in nonlinear deblurring and super resolution, while demonstrating comparable performance to DPS (Chung et al., 2023a) in Gaussian deblurring and inpainting. An interesting observation is that ReSample exhibits the largest performance gap from DPS in nonlinear deblurring and the smallest performance gap in random inpainting, mirroring the pattern observed in the results of the CelebA-H -HQ dataset. This may suggest that ReSample performs better than baselines when the forward operator is more complex.
A.2 ADDITIONAL RESULTS ON BOX INPAINTING
In this section, we present both qualitative and quantitative results on a more challenging image inpainting setting. More specifically, we consider the "box" inpainting setting, where the goal is to
14

---

Published as a conference paper at ICLR 2024
recover ground truth image in which large regions of pixels missing. We present results in Figure and Table 7 where we compare the performance our algorithm to PSLD (Rout et al, 2023 (latent-space algorithm) and DPS (Chung et al. 2023a) (pixel-space algorithm). Even in this challenging setting, we observe that ReSample can outperform the baselines, highlighting the effectiveness of method.
Box Inpainting
Method
LPIPS.
PSNRT
SSIM
DPS
0.127
22.85
0.861
DDRM
0.120
24.33
0.860
DMPS
0.233
22.01
0.803
Latent-DPS
0.199
23.14
0.784
PSLD
0.201
23.99
0.787
ReSample (Ours)
0.093
24.67
0.892
Table 7: Comparison of quantitative results for box inpainting CelebA-HQ dataset. Input images have an additive Gaussian noise with By 0.01. Best results are in bold and second best results are underlined.
A.3 ADDITIONAL PRELIMINARY RESULTS ON HIGH-RESOLUTION IMAGES
In this section, we present techniques for extending our algorithm to address inverse problems with high-resolution images. One straightforward method involves using arbitrary-resllutic ran- dom noise the initial input for the LDM. The convolutional architecture shared by the UNet in both the LDM and the autoencoder enables the generation of images at arbitrary resolutions, as demonstrated by Rombach et al. (2022). To illustrate the validity of this approach, randon-inpainting experiment with images of dimensions (512 x 512 x 3) and report the results in Figure 8 In Figure 8. we observe that method can achieve excellent performance high-resolution setting. Other possible methods include obtaining an accurate pair of encoders and decoders that can effectively map higher-resolution images to the low-dimensional space, and then running our algorithm using the latent diffusion model. We believe that these preliminary results represent an important step towards using LDMs as generative priors for solving inverse problems for high-resolution images.
conducted
in this
Measurement
Reference
ReSample
PSLD
DPS
PSNR:
PSNR:
LPIPS: 0.132
LPIP5: D.316
LPIPS: 0.151
PSNR: 25.37
PSNR: 24.79
PSNR: 19.60 LPIPS: 0.144
LPIPS: 0.128
LPIPS: 0.100
Figure 7: Qualitative results on box inpainting with measurement noise o, 0.01. These results highlight the effectiveness of ReSample more difficult inverse problem tasks.
A.4 DISCUSSION ON PSLD BASELINE
To better reflect the performance of PSLD (Rout et al. 2023), we re-ran several experiments with PSLD while collaborating with the original authors fine-tune hyperparameters. Furthermore,
15

---

Published as a conference paper at ICLR 2024
Measurement
Reference
ReSample
Figure Additional results on 70% random inpainting with high-resolution images (512 x 512 3) with measurement noise By 0.01.
as the work by Rout et al. (203) initially employed the stable diffusion (SD) model rather than the LDM model, have relabeled the baselines as "PSLD-LDM". We report the hyperparameters used to generate the new results, as well as the ones used to generate the previous results in Table The previous hyperparameters were chosen based on the implementation details provided by Root al.) (2023) in Section B.1, where they stated that they used available in the codebase. However, it has been brought to our attention that this was not the optimal hyperparameter. Therefore, the best results. We have also conducted additional experiments directly comparing to PSLD-LDM for inpainting tasks on the FFHQ IK validation set and present the results in Table 9 Throughout these results, we still observe that our algorithm largely outperforms the baselines, including PSLD.
0.1 and that the hyperparameters were
additionally tuned the parameter so that each bascline could obtain
Hyperparameter (Rout et al. 2023)
SR (FFHQ)
SR (CelebA)
Gaussian Deblur (CelebA)
Previous (7)
0.1
New (7)
1.0
1.1
0.15
Table 8: Different hyperparameters used to generate PSLD results. SR refers to super resolution "previous" and "new" refer to the hyperparameter used to generate the old and new results, respectively.
Inpainting (Random)
Inpainting (Box)
Method
LPIPS
PSNR1
SSIM
LPIPS]
PSNR
SSIM
DPS
0.212
29.49
0.844
0.214
23.39
0.798
PSLD-LDM
0.221
30.31
0.851
0.158
24.22
0.819
ReSample (Ours)
0.135
31.64
0.906
0.128
23.81
0.868
Table 9: Comparison of quantitative results for different inpainting tasks on the FFHQ 1K validation Best results bold and second best results underlined.
16

---

Published as a conference paper at ICLR 2024
A.5
ADDITIONAL QUALITATIVE RESULTS
Measurement
Reference
DPS
ReSample
Measurement
Reference
DPS
ReSample
VW
Figure 9: Additional results on super resolution (4x) on the CelebA-HQ dataset with Gaussian measurement noise of variance By 0.01.
Measurement
Reference
ReSample
Maasuremant
ReSample
Figure 10: Additional results Gaussian deblurring the CelebA-HQ dataset with Gaussian measurement noise of variance Bg = 0.01.
17

---

Published as a conference paper at ICLR 2024
Measurement Reference
ReSample
Measurement
ReSample
Figure 11: Additional results on inpainting with a random mask (70%) on the LSUN-Bedroom dataset with Gaussian measurement noise of variance σy 0.01.
Measurement
Reference
DDRM
DMPS
DPS
Latent-DPS
ReSample
Figure 12: Comparison of algorithms Gaussian deblurring the FFHQ dataset with Gaussian measurement noise variance dy = 0.01.
18

---

Published as a conference paper at ICLR 2024
Measurement
Reference
Latent-DPS
ReSample
Measurement
Reference
Latent-DPS
ReSample
Figure 13: Additional results Gaussian deblurring with additive Gaussian noise By 0.05.
Measurement
Reference
Latent-DPS
ReSample
Measurement
Reference
Latent-DPS
ReSample
I
Figure 14: Additional results super resolution 4x with additive Gaussian noise By 0.05.
19

---

Published as a conference paper at ICLR 2024
FBP-UNet
ADMM-PnP
Latent-DPS
Ours
Ground Truth
Figure 15: Additional results on CT reconstruction with additive Gaussian noise dy = 0.01.
20

---

Published as a conference paper at ICLR 2024
A.6 ABLATION STUDIES
Here, we present some ablation studies regarding computational efficiency and stochastic resam- pling, amongst others.
31
PSNR
PSNR
Strength of Data Consistency
Timesteps for Optimization and ReSample
Figure 16: An ablation study on few of the hyperparameters associated with ReSample. Left: Performance of different values of (hyperparameter in stochastic resampling) on the CelebA-HQ dataset. Right: CT reconstruction performance as a function of the number timesteps to perform optimization.
Effectiveness of Stochastic Resampling. Recall that in stochastic resampling, we have one hy- perparameter o7. In Section we discuss that our choice for this hyperparameter is
= (- ā
)
at
where ā is an adaptive parameter associated to the diffusion model process. Thus, the only param- eter that need to choose here is yy. To this end, perform study how y affects the image reconstruction quality. The yy term can be interpreted as a parameter that balances the prior consis- tency with the measurement consistency. In Figure [6 (left) observe that performance increases lot when increases from a small value, but plateaus afterwards. We also observe that the larger gives more fine details on the images, but may introduce additional noise. In Figure 17 visual representations corresponding to the choices of in order.
we provide
PSNR 30.21
PENR:
PSNR:
PSNR:
5
= 20
40
7 = 80
2 240
Figure 17: Visual representation of varying 7 on the CelebA-HQ dataset.
Effeet of the Skip Step Size. In this section, we conduct an ablation study to investigate the im- pact of the skip step size our algorithm. The skip step size denotes the frequency which we 'skip' before applying the next hard data consistency. For instance, skip step size of 10 implies that we apply hard data consistency every 10 iterations the reverse sampling process. Intuitively, one might expect that fewer skip steps would lead to better reconstructions, indicating more frequent hard data consistency steps. To validate this intuition, we conducted an experiment on the effect of the skip step size CT reconstruction, and the results ing inference times in Table [6] We observe that for skip step sizes ranging from to 10, the results very similar. However, considering the significantly reduced time required for skip step size
displayed in Figure 18, with correspond-
21

---

Published as a conference paper at ICLR 2024
chose skip step size of 10 in experimental setting in order to balance the trade-off off between reconstruction quality and inference time. Lastly, we would like to note that in Figure I8 skip step size of 4 exhibits (very) slight improvement over a skip step of Due to the reverse sampling process initiating with a random noise vector, there is a minor variability in the results, with the average outcomes for skip step sizes of 1 and 4 being very similar.
of 10 shown in Table 16
Algorithm ReSample (Ours)
Skip Step Size
Inference Time
PSNR
SSIM
31:20
31.74
0.922
4
8:24
31.74
0.923
10
3:52
31.72
0,922
20
2:16
31.23
0,918
50
1:20
30.86
0.915
100
1:08
30.09
0.906
200
0:54
28.74
0.869
Table 10: Inference times and corresponding performance for performing hard data consistency on with varying skip step sizes for chest CT reconstruction.
Reference
200
100
20
10
1
PSNR: 27.63
PSNR: 29.95
PSNR: 31.30
PSNR: 31.73
PSNR: 31.80
PSNR: 31.79
SSIM: 0.869
SSIM: 0.897
SSIM: 0.918
SSIM: 0.922
SSIM: 0.923
SSIM: 0.922
Figure 18: Observing the effect of the skip step size of hard data consistency optimization CT reconstruction. The numbers above the image correspond to the number of hard data consistency steps per iteration of the reverse sampling process (e.g., refers to hard data consistency on every step), and the red boxes outline the differences in reconstructions.
Computational Efficiency. Here, we present the memory usage and inference times of ReSample compared to our baselines in Ta- bles I1 and 12 We observe that the memory gap between ReSample and other algorithms widens when the pre-trained diffusion model increases in size, as the hard data consistency step of ReSample quires minimal memory. Conversely, since we only apply hard data consistency in subset of all sampling steps (which differs from isting methods that perform consistency updates at every time step), note slight increase in inference time compared algorithms such as DPS, DMPS, and MCG. ReSample outperforms PSLD in terms of inference time.
Table 11: Average inference of different algorithms for Gaussian Deblurring
Algonihm
Inference Time
seconds
ReSample
Discussion on Latent-DPS. In the main text, we briefly dis- cussed how adding the Latent- DPS gradient term into Re- Sample algorithm can improve the overall reconstruction qual- ity. Generally, we observe that adding the Latent-DPS gradient will cause very marginal boost in the PSNR, but only when the learning rate scale is chosen "correctly". Interestingly, we ob- serve that choosing âe is criti- cal in the performance Latent- DPS. To this end, we perform
Table 12: Memory Usage of different algorithms with differ- ent pretrained models for Gaussian Deblurring on FFHQ256 and ImageNet511 12 datasets
Pietrained Model
Algorihi DPS
Model Onily Memory
Total Memory
DDPM(FFH))
198MB
+3416MB (175%)
5369MB
MCG
+3421MB (175%))
5374MB
DDRM
18833MB (964 (5)
20786MB
LDMITIQ)
FSLD
3969M1
+55T6MB (140%)
SUSMM
ReSample (Ours)
1040MB
DDPM(mageNee
DES
+6637KIB (15196)
TRETMB
MCG
6637MR (151%))
1133MBB
DMPS
+873MBB (199
13125MB
DDRM
+4530MB (103 ())
5924MB
TGGTMB
ReSample (Ours)
+1322MB (30.1%)
7002MB
22

---

Published as a conference paper at ICLR 2024
brief ablation study the learning rate for Latent-DPS, where choose the learning rate to be kā for some k > 0. We vary &t and test the performance on 50 images of chest CT images and display the results in Table [3] In Table [3] we observe that k 0.5 returns the best results, and should be chosen one were to adopt the method of adding I Latent- DPS into ReSample.
k PSNR
0
0.5
1.0
1.5
2.0
2.5
31.64
31.73
31.58
31.51
31.27
31.04
Measurement Loss 9.82e
5 1.21e
4
5.03e
1.35e
3 1.96e
3 2.63e
Table 13: The effect of the Latent-DPS gradient scale (learning rate) as a function of k: > 0. We study the PSNR and measurement loss (objective function) changes for varying k for chest CT reconstruction on 50 test samples. The best results are in bold. Note that k = 0 here refers to no Latent-DPS and only using ReSample.
Training Efficiency of LDMs. To underscore the significance of our work, conducted abla- tion study focusing on the training efficiency of LDMs. This study demonstrates that LDMs require significantly fewer computational resources for training, trait particularly beneficial for down- stream applications like medical imaging. To support this assertion, we present the training time and performance of LDMs, comparing them to DDPMs (Ho al., (020) in the context of medical image synthesis. For LDMs, we utilized the pretrained autoencoder and LDM architecture provided by Rombach et al. (2022. J training them on 9000 2D CT slices from three organs across 40 patients sourced from the LDCT training set Moen et al. (02T). In the case of DDPMs, we employed the
codebase provided by Nichol & Dhariwal 2222)) to tram pixel-based diffusion models on the same
set of 2D CT training slices. The results are presented in comparison with 300 2D slices from 10 patients the validation set, detailed in Table Observing the results in Table 14 we note that by utilizing LDMs, generation quality, as measured by the FID score.
significantly reduce training time and memory usage while achieving better
Model type
Training Time Training Memory
FID
LDM (Ours)
133 mins
7772MB
100.48
DDPM
229 mins
10417MB
119.50
Table 14: Comparison of training efficiency and performance of LDMs and DDPMs on CT images, computed on a V100 GPU.
B DISCUSSION ON HARD DATA CONSISTENCY
Recall that the hard data consistency step involves solving the following optimization problem:
1 } (OD((())))
in E arg min
(14)
where we initialize using 2n(2) (z)). Instead of solving for the latent variable directly, notice that one possible technique would be to instead solve for vanilla least squares:
To(y) arg min ll A(x) x
(15)
where we instead initialize using D(2o (zz)), )), where D(.) denotes the decoder. Similarly, our goal here is to find the ( that is close D( (z)((z)) )thtt satisfies the measurement consistency: y A(x) < o², where o² is estimated noise level. Throughout the rest of the Appendix, refer to the former optimization process as latent optimization and the latter as pixel optimization.
In our experiments, we actually found that these two different formulations yield different results, in the sense that performing latent optimization gives reconstructions that are "noisy", yet much sharper with fine details, whereas pixel optimization gives results that "smoother" yet blurry with high-level semantic information. Here, the intuition is that pixel optimization does not directly change the latent variable where latent optimization directly optimizes over the latent variable.
23

---

refine the images Published as a conference paper at ICLR 2024
Moreover, the encoder E() can add additional errors the estimated Zy(y), perhaps throwing the sample off the data manifold, yielding images that are blurry as result.
There is also a significant difference in time complexity between these methods. Since latent optimization needs to backpropagate through the whole network of the decoder ())) on every gradi- step, it takes much longer obtain local minimum (or converge). Empirically, balance the trade-off between reconstruction speed and image quality, we see that using both these formula- tions for hard data consistency can only yield the best results, but also speed up the optimization process. Since pixel optimization is easy to get a global optimum, we use it first during the reverse sampling process and then use latent optimization when we are closer with the finer details.
Lastly, we would like remark that for pixel optimization, there a closed-form solution that could be leveraged under specific settings Wang et al. (2022). If the forward operator A is linear and can take the matrix form A and the measurements are noiseless (i.e., y (( then we can pose the optimization problem as
1 _ (ID( x s.t. Ar = y-
( E arg min
(16)
Then, the solution to this optimization problem is given by
( = (z))(z)) (A* AD( - A+ Aty),
(17)
where by employing the encoder, we obtain
(0 (z) & E(D( (A* (z)( A "g)).
(18)
This optimization technique does not require iterative solvers and offers great computational effi- ciency. In the following section, we provide ways in which we can compute a closed-form solution in the case in which the measurements may be noisy.
B.1 ACCELERATED PIXEL OPTIMIZATION CONJUGATE GRADIENT
Notice that since pixel optimization directly operates in the pixel space, we can use solvers such as conjugate gradients least squares for linear inverse problems. Let be the matrix form of the linear operator A(-). Then, as discussed in the previous subsection, the solution to the optimization problem in the noiseless setting is given by
4 To A A Aty),
(19)
where A = A (AAT)- and (AAT)- can be implemented by conjugate gradients. In the presence of measurement noise, we can relax this solution
a = x(A+ Aæo A A+y). y).
(20)
E 10 is a hyperparameter that can reduce the impact between the noisy component of the measurements and wy (the initial image before optimization, for which we use (z.))..))
We use this technique for CT reconstruction, where the forward operator A is the radon transform and AT is the non-filtered back projection. However, we can use this conjugate gradient technique for any linear inverse problem where the matrix A is available.
C IMPLEMENTATION DETAILS
In this section, we discuss the choices of the hyperparameters used in all of our experiments for our algorithm. All experiments are implemented in PyTorch on NVIDIA GPUs (A 100 and A40).
C.I LINEAR AND NONLINEAR INVERSE PROBLEMS ON NATURAL IMAGES
For organizational purposes, tabulate all of the hyperparameters associated to ReSample in Ta- ble [5] The parameter for the number of times to perform hard data consistency is not included in the table, did not have any explicit notation for it. For experiments across all natural images
24

---

Published as a conference paper at ICLR 2024
Notation
Definition
Hyperparameter for early stopping in hard data consistency
Variance scheduling for the resampling technique
7
Number DDIM or DDPM steps
Table 15: Summary of the hyperparameters with their respective notations for ReSample.
datasets (LSUN-Bedroom, FFHQ, CelebA-HQ), we used the same hyperparameters they seemed to all empirically give the best results.
For T, we used T 500 DDIM steps. For hard data consistency, first split T into three even sub-intervals. The first stage refers to the sub-interval closest to ' and the third stage refers to the interval closest to = O. During the second stage, performed pixel optimization, whereas in the third stage we performed latent optimization for hard data consistency as described in Section We performed this optimization every 10 iterations of t.
We set r 10 which seemed to give us the best results for noisy inverse problems, with a maxi- number of iterations of 2000 for pixel optimization and 500 for latent optimization (whichever convergence criteria came first). For the variance hyperparameter o, in the stochastic resample step, chose adaptive schedule of
11-0-11
- 0i--
(1
ae
of =
0i--
as discussed in Section natural images.
Generally, we see that 7 40 returns the best results for experiments on
C.2 LINEAR INVERSE PROBLEMS ON MEDICAL IMAGES
Since the LDMs for medical images is not readily available compared to natural images, largely had to fine-tune existing models. We discuss these in more detail this section.
Backhone Models. For the backbone latent diffusion model, use the pre-trained model from latent diffusion (Rombach et al. 2022). We select the VQ-4 autoencoder and the FFHQ-LDM with CelebA-LDM our backbone LDM.For inferencing. upon taking pre- pre-trained checkpoints provided by RRmbacc et al. (2022). we fine-tuned the models on 2000 CT images with 100K iterations and a learning rate of 10-0.
Inferencing. For T, we used a total of 7 = 1000 DDIM steps. For hard data consistency, we split T' into three sub-intervals: > 750, 300 < < 750, and < 300. During the second stage, we performed pixel optimization by using conjugate gradients as discussed previously, with 50 iterations with и 0.9. In the third stage, we performed latent optimization with as the estimated noise level r = 10-4 We set skip step size to be 10 and ? 40, with as the same as the experiments for the natural images.
C.3 IMPLEMENTATION OF BASELINES
Latent-DPS. For the Latent-DPS baseline, use T 1000 DDIM steps for CT reconstruction and 500 DDIM steps for natural image experiments. Let &. denote the learning rate. For medical images, we use G, 2.5ā, and 0.5ā, for natural images. Empirically, observe that our proposed G. step size schedule gives the best performance and robust to scale change, as previously discussed.
DPS and MCG. For DPS, we use the original DPS codebase provided by Chung et al. (2023a)
and pre-trained models trained on CelebA and FFHQ training sets for natural images. For medical images, we use the pretrained checkpoint from Chung et al. (2022) on the dataset provided by Moen et al. (2021). For MCG (Chung et al., 2022). we modified the MCG codebase by deleting the projection term tuning the gradient term for running DPS experiments
CT reconstruction.
Otherwise, we directly used the codes provided by CCung et al. (2022) for both natural and medical image experiments.
25

---

Published as a conference paper at ICLR 2024
DDRM. For DDRM, follow the original code provided by Kawar et al. (2022 WithDDPM
models trained on FFHQ and CelebA training sets adopted from the repository provided by Dhariwal We use the default parameters as displayed by Kawar et al. (2022).
& Nichol (20216).
DMPS. We follow the original code from the repository of Meng & Kabashima (2022) with the DDPM models trained on FFHQ and CelebA training sets adopted from the repository of Dhariwal & Nichol (2021b). We use the default parameters as displaye by |Meng & Kabashima (2022).
PSLD. We follow the original code from the repo Shah et al. (0221) with the pretrained LDMs on CelebA and FFHQ datasets provided by Rombach al. (2022). We use the default hyperparameters as implied in Shah et al. 2021).
ADMM-PuP and Other (Supervised) Baselines. For ADMM-PnP
the 10 iterations with
tuned for different inverse problems. We use T = 5 for CT reconstruction, = 0.1 for linear in- verse problems, and 0.075 for nonlinear deblurring. We use the pre- pre-trained model from original DrCNN repository provided by ZZang et al. (017)) We observe that ADMM-PnP tends to over- smooth the images with more iterations, which causes performance degradation. For FBP-UNet, we trained a UNet that maps FBP images to ground truth images. The UNet network architecture is the same as the one explained by Jin al. (2017).
26

---

Published as a conference paper at ICLR 2024
D MORE DISCUSSION ON THE FAILURE CASES OF LATENT-DPS
In this section, we provide a further explanation to which why Latent-DPS often fails to give accurate reconstructions.
D.1 FAILURE OF MEASUREMENT CONSISTENCY
Previously, claimed that by using Latent-DPS, it is likely that we converge a local minimum of the function lyy A A and hence cannot achieve accurate measurement consistency. Here, we validate this claim by comparing the measurement consistency loss between Latent-DPS and ReSample. We observe that ReSample is able to achieve better measurement consistency than Latent-DPS. This observation validates our motivation of using hard data consistency to improve reconstruction quality.
Anatomical Site ReSample (Ours)
Latent-DPS
Chest 5.36e-5
2.99e-3
Abdominal 6.28e-5
3,92e-3
Head 1.73e-5
2.03e-3
Table 16: Comparison of average measurement consistency loss for CT reconstruction between ReSample and Latent-DPS
D.2 FAILURE OF LINEAR MANIFOLD ASSUMPTION
We hypothesize that one reason that Latent-DPS fails to give accurate reconstructions could due to the nonlinearity of the decoder D(-). More specifically, the derivation of DPS provided by Chung al. (2023a) relies on a linear manifold assumption. For our case, since our forward model can be viewed as the form ((D(-)), where the decoder D(-) a highly nonlinear neural network, this linear manifold assumption fails hold. For example, if two images z(1) and =(2) lie on the clean data distribution manifold M at 0, a linear combination of them az(1) + bz (2) for some constants and b, may not belong to .M since ((az(1) (2)) may not give us an image that is realistic. Thus, in practice we observe that DPS reconstructions tend to be more blurry, which implies that the reverse sampling path falls out of this data manifold. We demonstrate this in Figure 19 where we show that the average of two latent vectors gives a blurry and unrealistic image.
We would like to point out that this reasoning may also explain why our algorithm outperforms DPS on nonlinear inverse tasks.
D(zo))
(1)
D(zg)
(2)
D
(20 + 2 12)
(1)
Figure Depiction of violation of the linear manifold assumption.
D.3 INACCURATE ESTIMATES OF THE POSTERIOR MEAN
In this section, we discuss how inaccurate estimates (2(2;) (i.e., the posterior mean via Tweedie's formula) can be of the reasons why Latent-DPS returns image reconstructions that noisy. This was mainly because for values of closer to = T. the estimate of (x,) may be inaccurate, leading to images that noisy at 0. Generally, find that the estimation of 6o(2.) is
27

---

Published as a conference paper at ICLR 2024
inaccurate in the early timesteps (e.g. >> 0.5T) and vary lot when t decreases. This implies that the gradient update may not point to a consistent direction when t is large. We further demonstrate this observation in Figure 20
Figure 20: Comparison of the prediction of the ground truth signal En (2)) ) for different values of t. Left: (z(2)) when 0.5T. Right: Eo (zz) when = 0. This serves to show the estimation error of the posterior mean for large values of / (i.e. when / is closer to pure noise).
E RELATED WORKS
Deep neural networks have been extensively employed as priors for solving inverse problems (Han &
Ye 2018 Bora al., 2017 Zhu al.) 2018: Gupta et al.] 2018). Numerous works focus on learning
the mapping between measurements and clean images, which we term as supervised methods (Han
Ye 2018: Zhu et al. 2018 Wei et al. 2020 Liang et al.) 2021). Supervised methods necessitate training on pairs of measurements and clean mages, requiring model retraining for each new task. On the other hand, another line of research aims at learning the prior distribution of ground truth images, solving inverse problems inference time without retraining by using the pre-trained prior
distributions (Bora et al. [017] Jalal et al. 2021 Hussein et al.) 2020 Lempitsky et al. 2018). We categorize this unsupervised methods.
Until the advent diffusion models, unsupervised methods struggled to achieve satisfactory re- construction quality compared supervised methods (Jalal et al. 2021). These methods rely on optimization within a constrained space (Bora et al.] (2017). while constructing a space that accu- rately encodes the ground truth data distribution solve the optimization problem challenging. However, with the accurate approximation of the prior distribution now provided by diffusion mod- els, unsupervised methods can outperform supervised methods for solving inverse problems more efficiently without retraining the model (Song et al.. 2021 Jalal et al. 2021).
For unsupervised approaches using diffusion models as priors, the plug-and ana-play approach has been widely applied for solving linear inverse problems (Song et 2022 Kawar et al. 2022 Wang et al.) 2022]. These methods inject the measurement on the noisy manifold or incorporate an estimate of the ground truth image into the reverse sampling procedure. We refer to these methods as hard data consistency approaches, as they directly inject the measurements into the reverse sampling pro- cess. However, to the best of our knowledge, only a few of these methods extend to nonlinear inverse problems or even latent diffusion models (Rombach et al. 2022). Other approaches focus on approx- imating the conditional score (Dhariwal & Nichol, 2021b Chung et al. [023a] (022) under some mild assumptions using gradient methods. While these methods can achieve excellent performance and be extended nonlinear inverse problems, measurement consistency may be compromised, as demonstrated in our paper. Since these methods apply data consistency only through a gradient in the reverse sampling process, we refer to these methods as soft data consistency approaches.
Recently, there has been a growing interest in modeling the data distribution by diffusion models in latent space constraint domain Liu et al. (002)); Rombach al. (2022;; Lou & Ermon (2023b): Fishman et al. (023). There has also been growing interest in solving inverse problems using latent diffusion models. In particular, Rout et al. 2023) proposed PSLD, an unsupervised soft approach that adds a gradient term at each reverse sampling step to solve linear inverse problems with LDMs. However, there are works that observe that this approach may suffer from instability due to estimating the gradient term Chung et al. (2023c). Our work proposes unsupervised hard approach that involves "resampling" and enforcing hard data consistency for solving general inverse problems (both linear and nonlinear) using latent diffusion models. By using hard data consistency approach, we can obtain much better reconstructions, highlighting the effectiveness of our algorithm.
28

---

Published as a conference paper at ICLR 2024
F DEFERRED PROOFS FOR RESAMPLE
Here, we provide the proof's regarding the theory behind our resampling technique. To make this sec- tion self-contained, we first restate all of results and discuss notation that will be used throughout this section.
Notation. We denote scalars with under-case letters (c.g. a) and vectors with bold under-case letters (e.g. x)) Recall that in the main body of the paper, E TR denotes sample in the latent space, z denotes an unconditional sample at time step t, Zo(ze) denotes a prediction of the ground truth signal 2 at time step I, Eo(y) denotes the meassreeeent-oonsiseent sample of (0(2)) using hard data consistency, and 2, denotes the re-mapped sample from Eo(y) onto the data manifold time step We use & as the next sample resume the reverse diffusion process.
Proposition (Stochastic Encoding). Since the sample &n given 2n (y) and measurement y is con- ditionally independent y, we have that
p( 9) p( NZ (.)) T))
(21)
Proof. By Tweedie's s formula, we have that Eo(y) the estimated mean of the ground truth signal Zg. By the forward process of the DDPM formulation Ho et al. (2020), we also have that
2 (y)) = N ( (v), (1 T))
Then, since y is a measurement of 2g at I = 0, we have p(yl (y) 2,) = p(yl2o(()) Finally, we get
p(E:)Eoly). p(yl 2u (3) )p()) (()
(22)
ply, 20(
ply. =0(y))))))))) 2o(1)))(((1
(23)
ply, 20(3))
= =
(24)
Proposition 2 (Stochastic Resampling). Suppose that p(z/|2,, 2o(y). y) is normally distributed such that p(z)) & (0()). y) N(p, of). If we let p( y) be prior for μhs then the posterior distribution pl 2 Zo(H), y) is given by
p( v) N of 2o (1- ax)
oo + (1 -aa) 1)
of (1 āz.) )
(25)
of + (1 -l)))
Proof. We have that
p(2, 2/,, 20, 4) 2i 2o. v)p(2z 2, 20, y y)vl (0(2,,))
(26)
pl ( 20. y)
where both p(2/, 20 y) and (Eo, y) normalizing constants. Then by Lemma 1 we have p(2: y) = - aa 2r 20, (1 - oz I). Now, we can compute the posterior distribution:
p( = k1/2/, 2o: 8) x p(zi12 = k, 20. y)p(E,
k/20 y)
(27)
kk
z\)²
(k
/ /a
x exp
+
(28)
of
1 Ct.
k
ae + varioor +11
2
x exp
(29)
This as a Gaussian distribution, which can easily shown using moment moment-generating func- tions Bromiley 2013).
29

---

Published as a conference paper at ICLR 2024
Theorem 1. If (y) is maassuuremen--oonsistee such that y A(D L.. &n = 20(2(+1)=
Eo(y), then stochastic resample is unbiased such that E| 3/ E] [z/]-
Proof. We have that [ 1 (1))(()) [15,, (2, 2//1 20 (3)) yill.
Let y of+1-a- By using Proposition 2 we have E E2, (rvā (1
of
20, 2/6 vll.
Since &n measurement--onnisseen such that y - A(D(()(())) let 1,,, we have
Zo(y) 20
Var- (z/_ + (1 ā. -
(30)
Then, we have that
E E Es [R [[7 20 + (1 - 7)=() E0, z vl
(31)
V Ōz-
12 16|2 + a V v
+ (1 >)
(32)
as both z/ and - are unconditional samples and independent of y. Now, we have
Ex [ا|ا + (1 - V logr(z/-x)1=1]) = 1ا=;_ 12 + āz k)) V logp(z((<))
(33)
Since z is the unconditional reverse sample of 2 have E[z/]] as-a El c)
and then
/ Vlogg(z/=)D(z/
Exy_, Vlog
(34)
x(z(_a)dz'
(35)
p(z
011)
(36)
0z1=
Finally, we have [[ = >B[=]] + (1 (1-7) = [[(]]-
Lemma Let é, and &, denote the stochastically encoded and resampled image of 20 (y respec- tively. If VAR(z > 0, then we have that VAR(2,) < VAR(ž,..
Proof. Recall that &, and &, are both normally distributed, with
VAR(2,) = 1 ā 1
(37)
VAR(ž,)
(38)
I- + o7
of(1 - 5.)
(39)
o + (1 ã.))
For all of 2 0, we have
0,(1 ā²)
< <1 ā.
(40)
of + (1 - āt)
VAR(2:) VAR(ž)).
(41)
Theorem Let Zn denote a sample from the data distribution and z, be sample from the noisy perturbed distribution at time l. Given that the score function Vz. log,2, Pz, (zx) is bounded,
Cov(zulzz) (1 (1-6r) )
V2 log Pz, (z))
ae I.
ā
where or 10 is decreasing sequence in
30

---

Published as a conference paper at ICLR 2024
Proof. By the LDM forward process, have
p(zz = -N(va,,1 - a.).
Consider the following variable: consider the distribution
2t =
2t, which is simply a scaled version of 2z. Then, 1-
yoa
z0)
1
d/2
exp
Et
l-a,
zul2
2
T-
T-ax
Now, want to separate the term only with &n and the term only with 20 for applying Tweedie's 's formula. Hence, let Po (z+)
(2x
4/2
exp
2
12
we achieve
α 2(1
r ( = (2() exp 21 20
(4.)
which separates the distribution interaction term &, 2o from | || term
Let X(z)) log pl 2x) log (oz).. Then by Tweedie's formula, we have 1R[2 2.]
VA(2,) and
Cov(zo) E {=)( Since Po (Et) is Gaussian distribution with mean at 0 and variance equal ai_ to I-a, we obtain
1
at αe
V = V log)(zz)
+
zz-
We observe that since é is scaled version of Zt, we can obtain the distribution of &
1 1 āg
āt
p(z))
v 0l:
Then, we can apply chain rule to first compute then gradient on and then account for &, As result we get
1- āi Vz logp(z))
v., log
V
and then
āi Vor
VA(=;) = -[z0
logp(zt) + St.
Vor
which is consistent with the score function of Chung (2023a). Afterwards, we take the gradient again with respective to the Σ,, and apply the chain rule again, then we have
1 Ot I. V²A(E) Cov(zolzz) (1 (1- 2 logp(zz) +
This gives the desired result.
31


Reviews:
[{"review_id": "CMX6woz0XW", "reviewer": "Reviewer_6mow", "summary": "The paper proposes a latent-diffusion model based inverse problem solver which is more efficient than previous solves on pixel space and a novel problem.\nThe main contribution of the paper is an unbiased and non-expansive sampling method after applying data consistency, namely stochastic resample. \nFor various linear inverse problems, the proposed method has demonstrated comparable or better performance compared to baseline, while reducing memory and time cost.", "questions": "- In the algorithm, at time point $t$, the data consistency update is done with $z_{t-1}$, rather than $z_t$ which is differ from prior works of diffusion-based inverse problem solver including DPS, DDNM and so on. In fact, to compute as in the algorithm, we need to feed-forward the diffusion model twice for each time step, which lead to two times longer sampling time. I think this is the reason that the paper emphasize the *partial* hard data-consistency update works well in the section 3.1. What if we use $z_t$ from the previous step without additional unconditional DDIM step? It seems reasonable because the mean of previous $z_t$ is also follows the unconditional sampling path. If I'm overlooking anything, please pointing it out.\n\n- How long it takes if we apply the hard data-consistency for every time-steps? Also, can authors provide the quantitative result on it? I wonder how marginal the performance gap between full hard data-consistency and partial hard data-consistency (which is a proposed method) is, as the performance would be further improved if we use ReSample more frequently according to the figure 6 and figure 15.\n\n- There is no description of $\\gamma$ in figure 15", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is clearly written and easy to follow\n- The proposed stochastic resampling is novel and theoretically solid.\n- The performance of the proposed method outperforms baseline methods, especially simple extension of DPs.", "weaknesses": "- Complex usage of multiple engineering to improve the efficiency of the method (e.g. partial data consistency, early stopping...)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "6ji64oZXKu", "reviewer": "Reviewer_1zP2", "summary": "The paper titled: Solving Inverse Problems with Latent Diffusion Models Via Hard Data Consistency presents an novel algorithm that use pre-trained latent diffusion models to solve general inverse problems. The authors demonstrated the effectiveness on multiple tasks and datasets.", "questions": "Please see above on my questions, overall a solid paper.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors propose a novel hard data consistency module, that strictly enforce the measurements via optimization, theoretical and empirical results demonstrate the superiority of ReSample over existing SOTA methods.\n\n2. The authors prove theoretical proof with analysis on the variance of the resampled images.\n\n3. The authors demonstrate the effectiveness of ReSample on multiple tasks on various benchmarks - deblurring, denoising, super resolution, inpainting CT Reconstruction.\n\n4. This paper is well written and the core idea is clearly delivered.", "weaknesses": "I didn't find clear flaw of this paper, I have a few questions instead.\n\n1. How do you scale up the algorithm to high-resolution images?\n\n2. To solve equation 10 - the L2 least square, I assume you are using gradient descent or conjugate gradient descent, one possible way to accelerate it is to use close form expression on D(z) if A is a linear operator, Do you have any thoughts on this path?\n\n3. Now, I assume the linear operator A is known, but for deblurring task, the system operator can be unknown, can you elaborate on how you could adopt your algorithms to this scenario?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "bk84Vdh3FQ", "reviewer": "Reviewer_azNZ", "summary": "In this paper, the authors propose a novel latent diffusion algorithm that ensures fidelity to the measurements for solving inverse problems. Instead of sampling in the image domain as often done in diffusion based methods, the sampling is performed in the latent domain which allows speeding up the algorithm. This approach is shown to perform well on several linear and non-linear inverse problems.", "questions": "**Major comments**\n\n1. In the DPS paper, the authors do indeed compute $\\nabla_z ||y-\\mathcal{A}(\\hat{x}_0(x_t)||_2^2$ in (7). However, their code does not use a squared norm, but a plain norm $\\nabla_z ||y-\\mathcal{A}(\\hat{x}_0(x_t)||$ [(see here)](https://github.com/DPS2022/diffusion-posterior-sampling/blob/main/guided_diffusion/condition_methods.py#L32). In practice, changing to a squared norm strongly decreases the quality of the results. Is this also the case in this work? Would the theoretical analysis in the paper still hold?\n2. It does not appear clearly in the paper that working in the latent domain improves the computational efficiency despite being a main motivation of this work by the authors. Table 7 is disappointing in that respect. Maybe the authors could emphasize other strengths of their method instead of this one (for instance, the successful results despite the difficulty to sample in constrained settings [1, 2, 3]).\n3. An important issue in my opinion is the lack of reference to methods performing similar procedure as the proposed Hard Data Consistency. While diffusion models are overall well referenced, other methods are less well documented. In particular, optimizing latent variables of a decoder as in equation (10) is at the heart of well known methods such as [4, 5]. Discussing these references (and related ones) would help relate the author's work to methods not relying on diffusion. Other relevant works [1, 2, 3] could also be added to the references.\n4. While the choice of solving (10) makes sense in practice, checking that the proposed sampling scheme works (or not) when (10) is solved at each iteration would be very insightful. So far, it is only backed by \"Because of the continuity of the sampling process, after each data-consistency optimization step, the samples in the following steps can still remain similar semantic or structural information to some extent.\" How does your method perform if a data fidelity step was performed at each iteration?\n5. The authors mention that (Rout et al., 2023) proposes a similar method; adding few sentences stressing the similarities and differences with the proposed method may be insightful.\n6. The theoretical derivations in Appendix appear in a more general form in [6]. Adding the reference when introducing Tweedie's formula may be of interest.\n\n\n**Minor comments**\n\n6. Adding PSNR metrics when displaying images would be informative (e.g. in Fig. 3 and 4, but maybe more importantly in the appendix where there is no page limit).\n7. Figure 6 is not very clear since the plot on the right does not correspond to the problem displayed on the left. Furthermore, the legend states \"ReSample frequency\" while the plot states \"# of iterations\", maybe the graph should either choose frequency in x-axis or the label should be updated. \n8. Few typos are remaining, e.g. \"in replace\" (p. 5), \"still remain\" (instead of retain, p. 5), \"noice\" (p. 6).\n9. If I understand Algorithm 1 correctly, the \"if t \\in C\" condition is not activated before Stage 2 of Fig. 16. In this context, I do not understand how information about the measurement can be visible in Stage 1. Isn't Fig. 16 misleading with that regards?\n10. More generally, is Stage 0 really necessary? I would naively that for solving eq. (10) and sampling from that may provide a good initial step, allowing to reduce the number of iterations.\n\n**References**\n\n[1] Guan-Horng Liu, Tianrong Chen, Evangelos A. Theodorou and Molei Tao, Mirror Diffusion Models for Constrained and Watermarked Generation arxiv:2310.01236, 2023\n\n[2] Nic Fishman, Leo Klarner, Valentin De Bortoli, Emile Mathieu, and Michael Hutchinson, Diffusion models for constrained domains. arXiv preprint arXiv:2304.05364, 2023.\n\n[3] Aaron Lou and Stefano Ermon. Reflected diffusion models. arXiv preprint arXiv:2304.04740, 2023\n\n[4] Ulyanov, Dmitry, Andrea Vedaldi, and Victor Lempitsky. \"Deep image prior.\" CVPR, 2018.\n\n[5] Jalal, Ajil, Marius Arvinte, Giannis Daras, Eric Price, Alexandros G. Dimakis, and Jon Tamir. \"Robust compressed sensing mri with deep generative priors.\" NeurIPS, 2021\n\n[6] Efron, Bradley. \"Tweedie’s formula and selection bias.\" Journal of the American Statistical Association 106.496 (2011): 1602-1614.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is overall well written and easy to follow. \n- The background on denoising diffusion models is clear, giving the right balance between technicalities and intuitions.\n- Experimental results are impressive. In particular, two points are to be underlined that make the method of interest for the imaging community:\n  - The authors chose to work with (relatively) difficult measurement operators that may not be trivially SVD decomposable;\n  - The authors considered mainly low measurement noise regime, which is a challenging setting for diffusion models but common setting in imaging.", "weaknesses": "- While the general background on diffusion models is well documented, key references are missing regarding diffusion models in constrained settings as well as latent space optimization methods (see references below).\n- The comparison with some other methods seems rather unfair. In particular, ADMM-PnP is ran for only 10 iterations, while it often requires at least 50 to 100 iterations to yield good results; similarly, the FBP-UNet seems to strongly under-perform compared to the (Jin et al) reference.\n- The motivations of the paper are not very consistent with the results (e.g. regarding the advantage of latent sampling on inference speed).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "gwej8WvSOY", "reviewer": "Reviewer_4YEX", "summary": "In this paper, the ReSample algorithm is introduced, leveraging the capabilities of diffusion models to address inverse problems. Unlike other approaches using diffusion models for inverse problems, ReSample trains the model in a lower-dimensional latent space, offering a computational advantage over conventional diffusion models trained in pixel space. Although a similar method was introduced very recently by Rout et al. in 2023, ReSample enhances its performance by introducing innovative sampling techniques supported by theoretical justification and validated through numerous experiments involving both linear and non-linear operators.", "questions": "-\tThe primary incentive for training the model in a lower-dimensional latent space appears to be a reduction in computational complexity. Nevertheless, the experimental outcomes presented in Table 1 and 2 demonstrate that the proposed algorithm surpasses conventional techniques in terms of reconstruction quality. Could you provide insights into the potential factors contributing to this outcome?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "-\tA design of a novel sampling method supported by theory. It is proven that the proposed sampling method is less noisy than stochastic sampling.\n\n-\tThe paper is well organized.", "weaknesses": "-\tThe idea of applying a latent space to solve inverse problems using diffusion models is not groundbreaking. Nonetheless, the paper introduces significant advancements to this approach.\n\n-\tIt is not clear why training the model in pixel space (Table 1 and 2) achieves lower quality results compared with the proposed ReSample algorithm. The rationale behinds this result remains unclear.\n\n-\tThe experimental results primarily involve random inpainting, a task considered relatively straightforward. It would be valuable to explore how the proposed algorithm performs in more complex inpainting scenarios, such as the removal of large pixel regions, as demonstrated in the approach outlined by Rout et al. in 2023.\n\n-\tA minor issue to address: There is a typographical error in the term \"noice\" after Theorem 2.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)