Improving Diffusion Models for Inverse Problems using Manifold Constraints
Authors: Hyungjin Chung
Byeongsu Sim
Dohoon Ryu
Jong Chul Ye
Venue: NeurIPS 2022 Accept (2022)
Keywords: Diffusion model
Inverse problem
Manifold constraint

Abstract: Recently, diffusion models have been used to solve various inverse problems in an unsupervised manner with appropriate modifications to the sampling process. However, the current solvers, which recursively apply a reverse diffusion step followed by a projection-based measurement consistency step, often produce sub-optimal results. By studying the generative sampling path, here we show that current solvers throw the sample path off the data manifold, and hence the error accumulates. To address this, we propose an additional correction term inspired by the manifold constraint, which can be used synergistically with the previous solvers to make the iterations close to the manifold. The proposed manifold constraint is straightforward to implement within a few lines of code, yet boosts the performance by a surprisingly large margin. With extensive experiments, we show that our method is superior to the previous methods both theoretically and empirically, producing promising results in many applications such as image inpainting, colorization, and sparse-view computed tomography. Code available https://github.com/HJ-harry/MCG_diffusion

Improving Diffusion Models for Inverse Problems using Manifold Constraints
Hyungjin Chung".!
Byeongsu Sim +.2
Dohoon Ryu' Dept. of Bio and Brain Engineering
Jong Chul 3e3.1,2
2 Dept. of Mathematical Sciences
Kimm Jaechul Graduate School of "Equal contribution
Korea Advanced Institute of Science and Technology (KAIST)
(hj. chung, byeongsu. dy, jong. ye @ekaist..... ac.kr
Abstract
Recently, diffusion models have been used to solve various inverse problems unsupervised manner with appropriate modifications to the sampling process. However, the current solvers, which recursively apply a reverse diffusion step followed by a projection-based measurement consistency step, often produce sub- optimal results. By studying the generative sampling path, here we show that current solvers throw the sample path off the data manifold, and hence the accumulates. To address this, we propose an additional correction term inspired by the manifold constraint, which can be used synergistically with the previous solvers to make the iterations close to the manifold. The proposed manifold constraint is straightforward to implement within a few lines of code, yet boosts the performance by surprisingly large margin. With extensive experiments, we show that our method superior the previous methods both theoretically and empirically, producing promising results in many applications such as image inpainting, colorization, and sparse-view computed tomography. Code available here
Introduction
Diffusion models have shown impressive performance both as generative models themselves 41 13
and also as unsupervised inverse problem solvers [4]] 25 that do not require problem-speeific training, Specifically, given a pre-trained unconditional score function (i.e. denoiser), solving the reverse stochastic differential equation (SDE) numerically would amount sampling from the data generating distribution AT. For many different inverse problems (e.g. super-resolution [8] 9] inpainting [4] [9, compressed-sente -sensing MRI (CS-MRI) 440 9 sparse view CT (SV-CT [40], etc.).
it was shown that simple incorporation of the measurement process produces satisfactory conditional samples, even when the model was not trained for the specific problem.
Nevertheless, for certain problems (e.g. inpainting). currently used algorithms often produce unsat- isfactory results when implemented naively (e.g. boundary artifacts, shown in Fig. (b)). The authors in 32) showed that in order to produce high quality reconstructions, one needs to iterate back and forth between the noising and the denoising step least 10 times per iteration. These iterations are computationally demanding and should be avoided, considering that diffusion models slow to sample from even without such iterations. On the other hand, classic result of Tweedie's formula 37 42 shows that one can perform Bayes optimal denoising in one step, once we know the gradient of the log density. Extending such result, was recently shown that one can indeed
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Proposed method
(6) Results
Figure 1: Visual schematic of the MCG correction step. (a) ! Unconditional reverse diffusion generates æ;; ( Q; maps the noisy æ; generate doo 3 Manifold Constrained Gradient (MCG) a (W ( l1 is applied to fix the iteration manifold; ④ Takes the orthogonal complement; 8m.
② Samples from pyy,y ly), then combines Axi-1 and y,- (b) Representative results of inpainting, compared with score-SDr SDE 41. Reconstructions with score-SDE produce incoherent results, while our method produces high fidelity solutions.
perform a single-step denoising with learned score functions for denoising problems from the general exponential family [28]
In this work, we leverage the denoising result through Tweedie's formula and show that such denoised samples can be the key significantly improving the performance of reconstruction using diffusion models across arbitrary linear inverse problems, despite the simplicity in the implementation. Moreover, theoretically prove that if the score function estimation is globally optimal, the correction term from the manifold constraint enforces the sample path to stay on the plane tangent the data manifold so by combining with the reverse diffusion step, the solution becomes more stable and accurate.
Related Works
2.1 Diffusion Models
Continuous Form For a continuous diffusion process a=()) E R", 1 E [0, we set x(0) ~
Po(x) Pdotas where Pdata represents the data distribution of interest, and x(1) (I()), with Pi (α) approximating spherical Gaussian distribution, containing no information of data. Here, the forward noising process defined with the following Itō stochastic differential equation (SDE) 41
dæ f(æ t) g(t)dw,
(1)
with f Rd + Rdd defining the linear drift function, g(t) R + defining a scalar diffusion coefficient, and R denoting the standard dimensional Wiener process. The forward SDE in ( is coupled with the following reverse SDE by the Anderson's 's theorem 41
dx (x, "v log (x)) + g())dw,
with denoting the infinitesimal negative time step, and w defining the standard Wiener process running backward in time. Note that the reverse SDE defines the generative process through the score function V. log pa (æ), which in practice, is typically replaced with V. log Por por (a)())(())) to minimize the following denoising score-matching objective
min (t.-U(e.))...m
((())a=00)
| (x((), - V.. æ. Pog (2))=(())
(3)
Once the parameter 0 for the score function is estimated, one can replace the score function in with sg- (x=)), to solve the reverse SDE
Discrete Form Due to the linearity of f and g. the forward diffusion step can be implemented with simple reparameteriizaioon trick 29] Namely, the general form of the forward diffusion is
C; 0,æ0 + b;2, z NO I),
(4)
We coin method Manifold Constrained Gradient (MCG).

---

where we have replaced the continuous index t E [0, 1] with the discrete index i E N. On the other hand, the discrete reverse diffusion step can in general be represented as
- (a 80-) (x;)
No I).
(5)
where we have replaced the ground truth score function with the trained one. We detail the choice of as bis Jo in Appendix.
2.2 Conditional Generative models for Inverse problems
The main problem of our interest in this paper is the inverse problem, retrieving the unknown æ E IR" from a measurement y:
y Ha te, EER", E
Ha
(6)
where e E Rn is the noise in the measurement. Accordingly, for the case of the inverse problems, our goal to generate samples from a conditional distribution with respect the measurement y. i.e. p(zly). Accordingly. the score function V log (z(x) in (2) should be replaced by the conditional score Va log Pr (az (g). Unfortunately. this strictly restricts the generalization capability of the neural network since the conditional score should be retrained whenever the conditions change. To address this, recent conditional diffusion models [2] g utilize the unconditional score
function V_ (mr) but rely on a projection-based measurement constraint to impose the conditions. Specifically, one can apply the following:
a/-1 f(Vi, 80) + g(;)z ~~ N(0,
- AA b;.
(8)
where A, b, are functions of H y. and zo- Note that 7 is identical to the unconditional reverse diffusion step in (5) whereas
( effectively imposes condition. It was shown in [9] that any
general contraction mapping (e.g. projection onto convex sets, gradient step) may be utilized as (8) to impose the constraint.
Another recent work [25] advancing 26 establishes the state-of-the-art (SOTA) in solving noisy inverse problems with unconditional diffusion models, by running the conditional reverse diffusion process in the spectral domain achieved by performing singular value decomposition (SVD), and leveraging approximate gradient of the log likelihood term the spectral space. The authors show that feasible solutions can be obtained with as small as 20 diffusion steps.
used in a similar fashion by utilizing a ceneral-purpone unconditional denoiser in the place of proximal mappings in model-based iterative reconstruction methods [5] 3 Similarly, outside the context of diffusion models, iterative denoising followed by projection- data consistency was proposed in (4).. In such view, diffusion models can be understood generative variant of PnPs trained with multiple scales of noise.
Prior the development of diffusion models, Plug-and-Play (PnP) models 47
44
GAN-based solvers are also widly explored 4 10 20 the test time by optimizing over the latent, the parameters, or jointly.
where the pre-traine pre-trained generators tuned at
2.3 Tweedic's formula for denoising
In the case of Gaussian noise, classic result of Tweedie's 's formula 337 tells us that one can achieve the denoised result by computing the posterior expectation:
E| 5. + of log)
where the noise is modeled by & ~ (a, I))). If we consider diffusion model in which the forward step modeled as æ, ~ Na,,oo, I) (discrete form), the Tweedie's formula can be rewritten as:
E[ [m.] = (x; + V.. log p(xx))(/aa;
(10)
Tweedie's formula is in fact not only relevant to Gaussian denoising the Bayesian framework, but have also been extended to be in close relation with kernel regression K4. Moreover, it was shown
that it can be applied to arbitrary exponential noise distributions beyond Gaussian 14 28 In the following, use this key property to develop our algorithm.

---

3 Conditional Diffusion using Manifold Constraints
Although our original motivation of using the measurement constraint step ( was utilize the unconditionally trained score function in the reverse diffusion step in ( there is room for imposing additional constraints while still using the unconditionally trained score function.
Specifically, the Bayes rule p(z y) =
leads to
Ve ( = V (x) + V log pyyla).
(11)
Hence, the score function in the reverse SDE 7 can be replaced by leading to
2--1 = f(x;, 86) a dx:
W (y H + g(x:)=z
~ 1)
(12)
where a and W depend on the noise covariance, if the noise c in (6) is Gaussian.
Now, one of the important contributions of this paper is to reveal that the Bayes optimal denoising step in (TP from the Tweedie's s formula leads to a preferred condition both empirically and theoretically.
Specifically, we define the set constraint for æ, called the manifold constrained gradient (MCG), so that the gradient the measurement term stays on the manifold (see Theorem D
a E X, where X = {x E I
(æ + s0 /))/
(13)
To deal with the potential deviation from the measurement consistency, we again impose the data consistency step (8). Putting them together, the discrete reverse diffusion under the additional manifold constraint and the data consistency can be represented by
a x/-1 = f(x,, s6) - a Ox;
W (x;)) g(x,)2, 9(zz))) 2 ~~,
(14)
Ø1-1 01-1=At_111 A
(15)
We illustrate our scheme visually in Fig. (a), specifically for the task of image inpainting. The additional step leads to a dramatic performance boost, can be seen in Fig. (b). Note that while the mapping (10) does not rely on the measurement, our gradient term in (14) incorporates the information of y so that the gradient of the measurement terms stays the manifold. In the following, we study the theoretical properties of the method. Further algorithmic details and adaptations to each problem that we tackle are presented in Section
We note that the authors of [99 proposed similar gradient method for the application of temporal imputation and super-resolution. When combining (14) with (5). one can arrive at a similar gradient method proposed in 19], and hence our method can be seen inverse problems. Furthermore, there are vast literature in the context of PnP models that utilize pre- trained denoisers together with gradient of the log- aigliielioood to solve inverse problems [30 48 [[. Among them, (30) is especially relevant to this work since their method relies modified Langevin diffusion, together with Tweedie's denoising and projections to the measurement subspace.
generalization to arbitrary linear
Geometry of Diffusion Models and Manifold Constrained Gradient
In this section, we theoretically support the effectiveness of the proposed algorithm by showing the problematic behavior of the earlier algorithm and how the proposed algorithm resolves the problem. We defer all proofs in the supplementary section. To begin with, we borrow geometrical viewpoint of the data manifold.
Notation For scalar points æ. y and a set A, we use the following notations. a.A := {ax æ E
A}: d(xr, A) := inf || y |2: (A) == {ar : d(xr, A) < r): TaMM the tangent space to manifold M at æ: y: the Jacobian matrix of vector valued function f. We define Po Pdata-
To develop the theory, need assumption on the data distribution.
Assumption (Strong manifold assumption: linear structure). Suppose M C. R" is the set of all data points, here we call the data manifold. Then, the manifold coincides with the tangent space with dimension 1 < n. Mn Br(Lo)) Mr,M n Br(Z0) and TarM RR.
Moreover, the data distribution po is the uniform distribution on the data manifold M.

---

(am b
Mao
Mi
Mi
b,vn
M
M
M
(i(x;)
y Hx
(a) Geometry of diffusion model
MCG correction
Figure 2: In both (a) and (b). the central manifolds represent the data manifold M, encireled by manifolds of noisy data M, The concentration on the manifold of noisy data and the distance from the clean data manifold are prescribed by Proposition TT In (a), the backward (resp. forward) step depicted by blue (resp, red) arrows can be considered as transitions from M; to M;-1 (resp. M;-1 M,). In (b), arrows refer the directions of conventional projection onto convex sets (POCS) step (green arrow) and MCG step (red arrow) which can be predicted by Theorem
We need recall that the conventional manifold assumption is about the intrinsic geometry of data points having a low dimensional nature. However, we assume more in this work: manifold is locally linear. Although this stronger assumption might narrow the practice of the theory, the geometric approach may provide new insights on diffusion models. Under this assumption, the following proposition shows how the data perturbed by noise lies in the ambient space, illustrated pictorially in Fig. 2a
Proposition (Concentration of noisy data). Consider the distribution of noisy data p.(x;) p((x, (t))oo((aa dd, ((x, ~ (a,xx f).). Then (x;;) is concentrated on (n 1)- manifold M, = {y E R" (y,a.M) r; == b, b/ - I}. Rigorously, (M.)). 1 - 8, for some small e, O.
Remark (Geometric interpretation of the diffusion process). Considering Proposition 1 the mani- folds of noisy data can be interpreted as interpolating manifolds between the two: the hypersphere, where pure noise N(a Eo, b2x)) is concentrated, and the clean data manifold. In this regard, the diffusion steps are mere transitions from one manifold another and the diffusion process transport from the data manifold to the hypersphere through interpolating manifolds. See Fig. 22
Remark can infer from the proposition that the score functions are trained only with the data points concentrated on the noisy data manifolds. Therefore, inaecurate inference might be caused by application of a score function on points away from the noisy data manifold.
Proposition 2 (score function). Suppose is the minimizer of the denoising score matching loss in Let Q₂ be the function that maps D; to Do for each
1
Q:
Rd. x. I
- (x. + (3
aj
Then, Qi(x;) M and J2 Jo. = Jo. Rd To Intuitively, Q2 is locally an orthogonal projection onto M.
According the proposition, the score function only concerns the normal direction of the data manifold. In other words, the score function cannot discriminate two data points whose difference is tangent to the manifold. In solving inverse problems, however, we desire discriminate data points to reconstruct the original signal, and the discrimination is achievable by measurement fidelity. In order to achieve the original signal, the measurement plays a role in correcting the tangent component near the data manifold. Furthermore, with regard to remark diffusion model-based inverse problem solvers should follow the tangent component. The following theorem shows how existing algorithms and the proposed method are different in this regard.
5

---

IAGAN
DDRM
Score-SDE
(e)
(b)
(c)
(c)
(e)
(g)
Figure 3: Inpainting results on FFHQ (Ist, 2nd row) and ImageNet (3rd, 4th row). (a) Measurement, (b) Ground truth, (c) IAGAN (20] for FFHQ, LaMa 443 for ImageNet, (d) DDRM 25], (e) Score- SDE 4I]. (f) RePAINT 332, (g) MCG (Ours). Out of 256 256 image, the 1st and the 3rd row is masked with size 128 x 128 box. 92% pixels (all RGB channels) from the images in the 2nd and 4th row are blocked.
Theorem (Manifold constrained gradient). correction by the manifold constrained gradient does not leave the data manifold. Formally,
i
W H - H H E T,
Ox;
the gradient is the projection of the data fidelity term onto Tæ
This theorem suggests that in diffusion models, the naive measurement fidelity step (without consid- ering the data manifold) pushes the inference path of the manifolds and might lead to inaccurate reconstruction. (To see this pictorially, see section. [ and Fig. ( On the other hand, our correction term from the manifold constraint guides the diffusion to lie on the data manifold, leading to better reconstruction. Such geometric views are illustrated in Fig. 2b
One may concern that the suboptimality of the denoising score matching loss optimization may lead to inaccurate inference of the MCG steps. In practice, however, most of the error denoising score matching is concentrated on t ~ 1/2), and in such region, the Tweedie inference cannot make meaningful images. That is, the score function cannot detect the data manifold. Nonetheless, this regime, the magnitudes of the MCGs small when the denoising score is inaccurate, and hence the matters arising from suboptimality minimal. As + O, the estimation becomes exact, and subsequently leads to accurate implementation of the MCG.
Remark 3.
5 Experiments
For all tasks, aim to verify the superiority of our method against other diffusion model-based approaches, and also against strong supervised learning- based baselines. Further details can be found in Section.
Datasets and Implementation For inpainting, we use FFHQ 256x256 256x256 [2 to validate our method. We utilize pre-trained models from the open sourced reposi- tory based on the implementation of ADM (VP-SDE) [3]. We validate the performance on 1000 held-out validation set images for both FFHQ and ImageNet dataset. the colorization task, we FFHQ 256x256. and
24 and ImageNet
256 256 [51). We pre-trained score functions from

---

FFHQ (256 x 256)
ImageNet (256 x 256)
Random
Extreme
Wide masks
Box
Random
Wide masks
Method
FID I.PISS
FID
I.PIS
FID.J 4 I.PIPS 4 FID
I.PPS
FID
I.PIS
FID
I.PIS
FID
I.PIPS
MCG (ours)
23.7
0.089
21.4
0.186
30.6
0.366
22.1
0.099
25.4
0.157
34,8
0.308
21.2
0.148
Score-SDe
41
30.3
0.135
109.3
0.674
48.6
0.488
29.8
0.132
43.5
0.199
143.5
0.758
25.9
0.150
RePAINT
32
25.7
0.093
38.1
0,240
35.9
0.398
24.2
26.1
0.156
0.387
37.0
0.205
DDRM 25
28.4
0.109
111.6
0.774
48.1
0.532
27.5
0.113
88.8
0.386
99.6
0.767
80.6
0.398
LaMa 43
27.7
0,086
188.7
0.648
61.7
0.492
23.2
0.096
26.8
0.139
134,1
0.567
20.4
0.140
AOT-GAN 52
29.2
0,108
97.2
0.514
69.5
0.452
28.3
0.106
35.3
0.163
119.6 0.583
29.8
0.161
ICT 49
27.3
0.103
91.3
0.445
56.7
0.425
26.9
0.104
31,9
0.148
131.4 0.584
25.4
0.148
DSI
35
27.9
0.096
126.4
0.601
77.5
0.463
28.3
0.102
34.5
0.155
132.9
0.549
24.3
0.154
IAGAN
20
26.3
0.098
41.5
0.279
56.1
0.417
23.8
0.110
Table 1: Quantitative evaluation (FID, LPIPS) of inpainting task on FFHQ and ImageNet. ": Re- implemented with our score function. MCG, Score-SDE, RePAINT, and DDRM all share the same score function and differ only in the inference method. Bold: Best, under: second best.
score-SDE [4] based on VE-SDE. We use 300 validation images for testing the performance with re- spect to the LSUNN-bedroo dataset. For experiments with CT, we train model based ncsnpp as a VE-SDE from score-SDE [41, on the 2016 American Association of Physicists in Medicine (AAPM) grand challenge dataset, and process the data as in [2). Specifically, the dataset contains 3839 training images resized to x56 resolution. We simulate the CT measurement process with parallel beam geometry with evenly-spaced 180 degrees. Evaluation is performed on 421 held-out validation images from the AAPM challenge.
Inpainting Score-SDE REPAINT 332 DDRM [25 were chosen as baseline diffu- sion models to compare against the proposed method. For fair comparison, we use the same score function for all methods includ- ing MCG, and only differentiate the inference method that is used. Another class of generative models: GAN-based inverse problem solver, IA- GAN 20 is considered as a comparison method for FFHQ specifically. We also include compar- isons against supervised learning based base-
Data
FFHQ 256)
LSUN(256x 256)
Method
SSIM 1 LPIPS
SSIM
LPIPS
MCG (ours)
0.951
0.146
0.959
0.160
Score: Score-SDE 41
0.936
0.180
0.945
0.199
DDRM 25
0.948
0.154
0.957
0.182
cINN 2
0.952
0.166
0.952
0.180
pix2pix 21
0.935
0.184
0.947
0.174
Table 2: Quantitative evaluation (SSIM, LPIPS) of colorization task. Bold: best, under: second best.
lines: LaMa [43]. AOT-GAN (52). ICT [49], and DSI [35]. We use various forms of inpainting masks: box (128 x 128 sized square region is missing") 1, extreme (only the box region is existent), random (90-95% of pixels are missing), and LaMa-wide. Quantitative evaluation is performed with two
metrics- Frechet Inception Distance (FID)-Ik 17 and Learned Perceptual Image Patch Similarity
(LPIPS) 54
Our method outperforms the diffusion model baselines 32 25 by a large margin. More- over, our method is also competitive with, or even better than the best-in-class fully super- vised methods, as can be seen in Table [ In Fig. we depict representative results that show the superiority of the method, where we see in both the box-type and random dropping that MCG performs very well all experiments.
AAPM (256 x 256)
Views
18
Method
PSNR
SSIM
PSNR 1 SSIM 1
MCG (ours)
33.57
0.956
36.09
0.971
Score-CT H0
29.85
0.897
31.97
0.913
SIN-4c-PRN 50
26.96
0.850
30.23
0.917
eGAN D5
24.38
0.823
27.45
0.927
FV
3
21.57
0.791
23.92
0.861
Table 3: Quantitative evaluation (PSNR, SSIM) of CT reconstruction task. Bold: best.
Colorization We choose score-SDE [1], and DDRM (25) as diffusion-modee based compar- ison methods, and also compare against cINN (2], and pix2pix (1]] Two metrics were used for evaluation: structural similarity index (SSIM), and LPIPS. Consistent with the findings from in- painting, we achieve much improved performance than score-SDE, and also is favorable against state-of-the-an (SOTA) superivsed learning based methods. In Table 2 we see that the proposed
'The location of the box is sampled uniformly within 16 pixel margin of each side.
7

---

Colorization
Sparse-view
Measurement
Score-SDE
(Ours)
SIN-4c-PRN
Score-CT
10.76
25.04
26,83
32.94
Figure 4: Colorization results on FFHQ/ LSUN-bedroom, Sparse view CT reconstruction results on AAPM.
method outperforms all other methods in terms of both PSNR/LPIPS in LSUN-bedroom, and also achieves strong performance in the colorization of FFHQ dataset.
CT reconstruction To the best of our knowledge, 40 is the only method that tackles CT recon- struction directly with diffusion models. We compare our method against [40], which we refer to as score-CT henceforth. We also compare with the best-in-class supervised learning methods, cGAN [15] and SIN-4c-PRE PRN [50]. As a compressed sensing baseline, FISTA-TV [3] was included, along with the analytical reconstruction method, FBP. We use two standard metrics peak-signal- to-noise-ratio (PSNR), and SSIM for quantitative evaluation. From Table we see that the newly proposed MCG method outperforms the previous score-CT Core-CTT [40] by large margin. We can observe the superiority of MCG over other methods more clearly in Fig. where MCG reconstructs the measurement with high fidelity and detail. All other methods including the fully supervised baselines fall behind the proposed method.
Ablation studies We perform three ablation studies: 1) As both the MCG term and the projection term contain in- formation about the measurement we observe the contri- bution of each term to the fixed solution. To further clarify the efficacy of the gradient step combined with Tweedie's 's denoising, we also consider the case where the gradient of the log likelihood is computed not in the noiseless regime, but in the noise level matching the current iteration. Specifically, we define a/ == f(x;, 80) + g(x;)z, 1/ and implement the gra- dient step as 2., 1e.11 H 2) As the perfor- mance of diffusion models depend heavily on the number of NFES, we observe the trade-off each diffusion model when varying the NFE from 20 to 1000. Moreover, for completeness, we measure the runtime of each algorithmss including the non-diffusion based methods in wall-clock time computed with commodity GPU in Table, 3) Setting a = 0.0 reduces our method to [9] We show the difference in the performance by varying the values of a.
Method
Wall-clock time [s]
Score-SDD
38.68
RePAINT
247.6
DDRM 25
2.117
LaMa 48
0.629
AOT-GAN 52
0.082
ICT
144.6
DSI
36.64
IAGAN 20
518.47
(O, I) -1-1
Ours
81.59
Table 4: Runtime for each algorithm in Wall-clock time: Computed with single GTX 1080Ti GPU.
First, we see in Table.55thaa using only the MCG step leads to improved performance in terms of LPIPS, but introduces error the measurement consistency (measured with MSE). Combining both the projection and MCG leads to perfect data consistency along with further improved
Method
LPIPS(J) MSE(MC)
Proj.
0.138
Vz. Ta- Ha'
0.271
12.99
Proj.
V
Ha
0.124
10.7
Vz.
ly H
Proj. (Ours)
0.089
Table 5: LPIPS & Measurement consistency (MC) vs. method

---

DDRM RePAINT
0.23
PS box inseirting
box ingairting
MCG
0.2
5C0
107
12f
104
Neural Fundtoon Evalualion (NFE)
(a) LPIPS NFE
(b) LPIPS vs. a'
Figure 5: Ablation studies performed with box inpainting task FFHQ 256 data.
reconstruction. When considering gradient steps without Tweedie's denoising (i.e. keeping the noise level at the hh step), the performance heavily degrades, especially when implemented without the projection steps. Here, we see that the proposed denoising step utilize igg indeed the key the superior performance.
we immediately see that the graph of MCG stays in the lowest (best) LPIPS regime across all NFEs by large margin, except for when the NFE drops below 100. Here, DDRM 25 takes over the 1st place- allegedly due the DDIM sampling strategy they take. The performance of RePAINT deteriorates rapidly as we decrease NFE. Furthermore, we observe that the LPIPS of score-SDE 41 actually increases (i.e. worsen), as increase the number of NFEs from few hundred to one thousand. This suggests that the inference process that score-SDE takes (i.e. projection only) is inherently flawed, and cannot be corrected by taking small enough steps. In Table. we list the runtime of all the methods that were used for comparison in the task of inpainting. Note that the proposed method takes longer for compute than score-SDE albeit having the same NFE. The gap is due to the backpropagation steps that are required for the MCG step, where the gap can potentially ameliorated by switching to JAX implementation from the current Py PyTorch implementation.
Second, looking at Fig. 5a
Lastly, we observe the difference in the performance as we vary the values a. Implementation wise, we find that we yield superior results when normalizing the squared norm with the norm of itself (e.g. a a' /I (y-
o ||. where a' some constant), In order avoid cluttered notation, we instead experiment with changing the values of a' in Fig. 5b) Inspecting Fig. [bb] we see that values within the range 1, 1.0) produce satisfactory results. a values that are too low do not fully enjoy the advantages of MCG and collapses to the projection-only method, while using too high values of a results exploding gradients, and the reconstruction saturates.
Properties of our method Our proposed method is fully unsupervised and is not trained on solving specific inverse problem. For example, box masks and random masks have very different forms of erasing the pixel values. Nevertheless, our method generalizes perfectly well to such different measurement conditions, while other methods have large performance gap between the different mask shapes. We further note two appealing properties of our method as an inverse problem solver: 1) the ability generate multiple solutions given condition, and 2) the ability maintain perfect measurement consistency. The former ability often lacks in supervised learning- based methods 4 50], and the latter is often not satisfied for unsupervised GAN-based
solutions 10
Conclusion
In this work, we proposed a general framework that can greatly enhance the performance of diffusion model-based solvers for solving inverse problems. We showed several promising appli- cations - inpainting, colorization, sparse view CT reconstruction, and showed that our method can outperform the current state-of-the-as methods. We analyzed our method theoretically and show that

---

MCG prevents the data generation process from falling off the manifold, thereby reducing the errors that might accumulate at every step. Further, we showed that MCG controls the direction tangent to the data manifold, whereas the score function controls the direction that is normal, such that the two components complement each other.
Limitations and Broader Impact The proposed method is inherently stochastic since the diffusion model the main workhorse of the algorithm. When the dimension m is pushed to low values, at times, our method fails to produce high quality reconstructions, albeit being better than the other methods overall. For extreme cases of inpainting (e.g- Half masks) with the ImageNet model, we often observe artifacts in our reconstruction (e.g- generating perfectly symmetric images), which we discuss in further detail in Sec. E We note that our method is slow to sample from, inheriting the existing limitations of diffusion models. This would likely benefit from leveraging recent solvers aimed at accelerating the inference speed of diffusion models. In line with the arguments other generative model-based inverse problem solvers, our method is solver that relies heavily on the underlying diffusion model, and can thus potentially create malicious content such as deepfakes. Further, the reconstructions could intensify the social bias that is already existent in the training dataset.
Acknowledgments and Disclosure of Funding
This research supported by Field-oriented Technology Development Project for Customs Admin- istration through National Research Foundation of Korea(NRF) funded by the Ministry of Science & ICT and Korea Customs Service (NRF-2021 A 1097938, NRF-2021 122021M311A A 1097910), by the Korea Health Technology R&D Project through the Korea Health Industry Development Institute (KHIIII) which is funded by the Ministry of Health & Welfare, Republic of Korea (grant number: HU2IC0222), and by the KAIST Key Research Institute (Interdisciplinary Research Group) Project.
References
[1] Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-3 1-3113333,111- 1982.
[2] Lynton Ardizzone, Carsten Lüth, Jakob Kruse, Carsten Rother, and Ullrich Köthe. Guided image generation with conditional invertible neural networks. arXiv preprint arXiv:1907.02:3 2019.
[3] Amir Beck and Marc Teboulle. A fast iterative shrinkage-thressholdnng algorithm for linear inverse problems. SIAM journal on imaging sciences, 2(1):183-202, 2009.
[4] Ashish Bora, Ajil Jalal, Eric Price, and Alexandros Dimakis. Compressed sensing using generative models. In International Conference on Machine Learning, pages 537-546. PMLR, 2017.
[5] Stephen Boyd, Neal Parikh, Eric Chu, Borja Peleato, Jonathan Eckstein, et al. Distributed opti- mization and statistical learning via the alternating direction method of multipliers. Foundations and Trends@ in Machine learning, 3(1):1-122, 2011.
[6] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milhe and Qiao Zhang. JAX: composable transformations of Python+NumPP programs, 2018
[7] Thorsten M Buzug. Computed tomography. In Springer handbook of medical technology, pages 311-342. Springer, 2011.
[8] Jooyoung Choi, Sungwon Kim, Yonghyun Jeong, Youngjune Gwon, and Sungroh Yoon. ILVR: Conditioning method for denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2021.
[9] Hyungjin Chung, Byeongsu Sim, and Jong Chul Ye. Come-Closer Beee-Coooo---ffuse--eeee Diffuse-Faster: Acceler- ating Conditional Diffusion Models for Inverse Problems through Stochastic Contraction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
10

---

[10] Giannis Daras, Joseph Dean, Ajil Jalal, and Alexandros G Dimakis. Intermediate layer opti- mization for inverse problems using deep generative models. In International Conference on Machine Learning, 2021.
[11] Valentin De Bortoli, Alain Durmus, Marcelo Pereyra, and Ana Fernandez Vidal. Maximum likelihood estimation regularization parameters in high-dimensional inverse problems: an empirical bayesian approach. part ii: Theoretical analysis. SIAM Journal on Imaging Sciences, 13(4):1990-2 20288, 2020.
[12] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pages
leee, 2009.
[13] Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat GANs on image synthesis. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[14] Bradley Efron. Tweedic's formula and selection bias. Journal of the American Statistical Association, 106(496):1602-16) 2011.
[15] Muhammad Usman Ghani and W Clem Karl. Deep learning-basee sinogram completion for low-dose ct. In 2018 IEEE 13th Image, Video, and Multidimensional Signal Processing Workshop (/VMSP), pages 1-5. IEEE, 2018.
[16] Richard Gordon, Robert Bender, and Gabor Herman. Algebraic reconstruction techniques (art) for three-dimensional electron microscopy and x-ray photography. Journal of theoretical Biology, 29(3):471-48)) 1970.
[17] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge local nash equilibrium. In Guyon, U. Von Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[18] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33, pages 6840-6851, 2020.
[19] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. arXiv preprint 2.Xi:22244003458, 2022.
[20] Shady Abu Hussein, Tom Tirer, and Raja Giryes. Image-adaptive gan based reconstruction. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 3121-3129, 2020.
[21] Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei A Efros. Image-to-image translation with conditional adversarial networks. In Proceedings of the IEFE conference on computer vision and pattern recognition, pages 1125-1134, 2017.
[22] Zahra Kadkhodaie and Eero Simoncelli. Stochastic solutions for linear inverse problems using the prior implicit in a denoiser. In Advances in Neural Information Processing Systems, volume 34, pages 13242-13254 Curran Associates, Inc., 2021.
[23] Eunhee Kang, Junhong Min, and Jong Chul Ye. A deep convolutional neural network using directional wavelets for low-dose x-ray reconstruction. Medical physics, 44610):330--235. 2017.
[24] Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 4401-4410, 2019.
[25] Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. In ICLR Workshop on Deep Generative Models for Highly Structured Data, 2022.
[26] Bahjat Kawar, Gregory Vaksman, and Michael Elad. Snips: Solving noisy inverse problems stochastically. Advances in Neural Information Processing Systems, 34:21757-21769, 2021.

---

[27] Daniil Kazantsev, Edoardo Pasca, Martin J Turner, and Philip J Withers. Copi-regularisation toolkit for computed tomographic image reconstruction with proximal splitting algorithms. SoftwareX, 9:317-323, 2019.
[28] Kwanyoung Kim and Jong Chul Ye. Noise2score: Tweedie's approach to self-supervised image denoising without clean images. In A. Beygekzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[29] Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In 2nd International Conference on Learning Representations, ICLR, 2014.
[30] Rémi Laumont, Valentin De Bortoli, Andrés Almansa, Julie Delon, Alain Durmus, and Marcelo Pereyra. Bayesian imaging using plug & play priors: when langevin meets tweedie. SIAM Journal on Imaging Sciences, 15(2):701-737, 2022.
[31] Beatrice Laurent and Pascal Massart. Adaptive estimation of a quadratic functional by model selection. Annals of Statistics, pages 1302-1338, 2000.
[32] Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. RePaint: Inpainting using Denoising Diffusion Probabilistic Models. arXiv preprint arXiv:2201.09865 2022.
[33] Xudong Mao, Qing Li, Haoran Xie, Raymond YK Lau, Zhen Wang, and Stephen Paul Smolley, Least squares generative adversarial networks. In Proceedings of the IEEE international conference on computer vision, pages 2794-2802, 2017.
[34] Frank Ong, Peyman Milanfar, and Pascal Getreuer. Local kernels that approximate bayesian regularization and proximal operators. IEEE Transactions on Image Processing, 28(6):3007- 3019, 2019.
[35] Jialun Peng, Dong Liu, Songcen Xu, and Houqiang Li. Generating diverse structure for image inpainting with hierarchical VQ-VAE. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattenn Recognition, pages 10775-10784, 2021.
[36] Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. Advances in neural information processing systems, 32, 2019.
[37] Herbert E Robbins An empirical bayes approach to statistics. In Breakthroughs in statistics, pages 388-394, Springer, 1992.
[38] Simo Särkkä and Amno Solin. Applied stochastic differential equations, volume 10. Cambridge University Press, 2019.
[39] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In 3rd International Conference on Learning Representations, ICLR, 2015.
[40] Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medi- cal imaging with score-based generative models. In International Conference on Learning Representations, 2022.
[41] Yang Song, Jascha Sohl-Dickstein Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In 9th International Conference on Learning Representations, ICLR, 2021.
[42] Charles M Stein. Estimation of the mean of multivariate normal distribution. The annals of Statistics, pages 1135-1151, 1981.
[43] Roman Suvorov, Elizaveta Logacheva, Anton Mashikhin, Anastasia Remizova, Arsenii Ashukha, Aleksei Silvestrov, Naejin Kong, Harshith Goka. Kiwoong Park, and Victor Lempitsky. Resolution-robust large mask inpainting with fourier convolutions. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pages 2149 2149-2159, 2022.
[44] Tom Tirer and Raja Giryes. Image restoration by iterative denoising and backward projections.
IEEE Transactions Image Processing, 28(3): (2(3):1220-223 2018.
12

---

[45] Github PK Tool, Nov Sun Mon Tue Wed Thu, and Fri Sat. dkazanc/tomobar.
[46] Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
[47] Singanallur V Venkatakrishnan, Charles A Bouman, and Brendt Wohlberg. Plug-and-pel priors for model based reconstruction. In 2013 IEEE Global Conference Signal and Information Processing, pages 945-948, IEEE, 2013.
[48] Ana Fernandez Vidal, Valentin De Bortoli, Marcelo Pereyra, and Alain Durmus. Maximum likelihood estimation of regularization parameters in nigh-dimensional inverse problems: An empirical bayesian approach part i: Methodology and experiments. SIAM Journal on Imaging Sciences, 13(4):1445-198) 2020.
[49] Ziyu Wan, Jingho Zhang. Dongdong Chen, and Jing Liao. High-fidelity pluralistie image completion with transformers. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pages
2021.
[50] Haoyu Wei, Florian Schiffers, Tobias Würfl, Daming Shen, Daniel Kim, Aggelos K Katsaggelos, and Oliver Cossairt. 2-step sparse- ct reconstruction with domain-speciftr perceptual network. arXiv preprint arXiv:2012.04743, 2020.
[51] Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao. Lsun: Construction of scrge-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv: 533663335, 2015.
[52] Yanhong Zeng, Jianlong Fu, Hongyang Chao, and Baining Guo. Aggregated contextual transformations for resolutsoo image inpainting. IEEE Transactions on Visualization and Computer Graphics, 2022.
[53] Kai Zhang, Wangmeng Zuo. Shuhang Gu, and Lei Zhang. Learning deep cnn denoiser prior for image restoration. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages
2017.
[54] Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unrea- sonable effectiveness of deep features as perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 586-595, 2018.
[55] Jun-Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to-iman translation using cycle-consistent adversarial networks. In Proceedings of the IEEE international conference on computer vision, pages 2223 2223,2232, 2017.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] We discuss the limitations in (6)
(c) Did you discuss any potential negative societal impacts of your work? [Yes] We discuss potential negative impacts in (6).
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] We provide all proofs of results in supplementary material.
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or as a URL)? [Yes] We include all code for our experiments in the supplementary material. We will release the code once the paper is published.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [No] Due our limited resources we do not have time to run multiple sets of experiments.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes]
4. If you using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] We have cited the original works that released the datasets.
(b) Did you mention the license of the assets? [No] Licenses are standard and can be found
online.
(c) Did you include any new assets either in the supplemental material URL? [Yes] We include our implementation as the supplementary material. We will release the code upon publication.
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [No] All datasets used in our work are publicly available. (e) Did you discuss whether the data you using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "skWzzsMjMGl", "reviewer": "Reviewer_vg3b", "summary": "The paper proposes a simple augmentation to existing score-based diffusion solvers for inverse problems by applying an extra correction term in the diffusion step. This term, in theory, captures the on-manifold component of the current estimation error.\n\n-------------\n\nAfter rebuttal: The authors have convinced me that the empirical experiments are sound, and have also tempered their theoretical claims. I am happy with the paper in its present state, as I am convinced that it presents a novel technique for inverse imaging, and achieves a solid improvement on the prior state of the art.", "questions": "See (Overview) in Strengths and Weaknesses for questions / points to address.", "limitations": "No.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The proposed step is elegant and computationally simple.\n- The results, if correct (see below), suggest massive improvements in signal recovery.\n\nWeaknesses:\n- The main theorem does not state its assumptions clearly. Moreover, I don't believe they hold, and this is quite problematic for the proposed algorithm. (See overview.)\n- The empirical experiments are somewhat problematic. (See overview.)\n- (Minor) The use of Bayes rule in Section 3.1 to obtain Eqs. 11 and 12 is somewhat misleading; this is a general technique for conditional sampling in score-based diffusion models, see for e.g. [1, 2, 3].\n\nOverview:\nOverall, I am very excited by the direction proposed here, but I think there are currently major issues in the reasoning and execution of the paper. I believe the paper will be much stronger after addressing these problems. Ultimately, my concerns are two-fold:\n\nTheoretically, I am not convinced that the main result of the paper, upon which the proposed algorithm rests, holds in practice. Theorem 1 requires Proposition 1 to hold, and Proposition 1 assumes that the score function is a *global minimizer* of the denoising loss (Eq. 9). This is nearly impossible in a diffusion model, since the denoising function at t=T (or t=0 in the reverse process) must \"recover\" an image from pure noise --- this is a highly ill-defined problem. Another way to see that this cannot hold: If the score function was a global minimizer, the generative process for unconditional diffusion models would take 1 inference step (in practice, it takes thousands of steps in the vanilla diffusion model). The suboptimality of the score function has consequences that extend past theory, as the central manifold constraint gradient (MCG) term also relies on the global optimality of the score function to define manifold projection (Eq. 13). Thus, this key assumption is quite central to the entire paper.\n\nEmpirically, the visual quality of competing work seems like it could be somewhat misrepresented, especially in image inpainting. For example, in Figure 3, the visual quality of RePaint [4] and LaMa [5] both differ greatly from that reported in their respective papers (see Figure 1 for each paper). Moreover, looking more closely at the experimental setup of the competing works RePAINT and LaMa, the two papers share the same experiments for inpainting, and establish a baseline. Namely, their inpainting task involves recovering an original image corrupted by \"Wide\" and \"Narrow\" masks. This paper uses a different set of masks, and reports significantly different LPIP scores for RePAINT and LaMa for this task. To be clear, I think it is fair to change the experimental setup. However, I wonder if the authors did not spend enough time tuning the competing models in this new setup, and consequently misrepresented the competing works.\n\n[1] Dhariwal, P., & Nichol, A. (2021). Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34, 8780-8794.\n[2] Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., & Ganguli, S. (2015, June). Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning (pp. 2256-2265). PMLR.\n[3] Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., & Poole, B. (2020). Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456.\n[4] Lugmayr, A., Danelljan, M., Romero, A., Yu, F., Timofte, R., & Van Gool, L. (2022). Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 11461-11471).\n[5] Suvorov, R., Logacheva, E., Mashikhin, A., Remizova, A., Ashukha, A., Silvestrov, A., ... & Lempitsky, V. (2022). Resolution-robust large mask inpainting with fourier convolutions. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (pp. 2149-2159).\n", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "fukaUzCMtD4", "reviewer": "Reviewer_ze1d", "summary": "In this paper, the authors propose a new methodology to perform conditional sampling using denoising diffusion models. Namely, they combine existing conditional score-matching techniques using a stochastic contraction with a slight change of the score. Indeed, instead of considering the unconditional score functional they consider the score function associated with the conditional distribution given by some manifold constraint. The authors show that this term is in fact complementary to the existing unconditional score term as the score term only (locally) project on the data manifold, whereas the additional term moves locally on the tangent space of this manifold. They evaluate their methodologies against some baselines and show promising results for inpainting, colorization and CT-scan reconstruction.", "questions": "* The authors state that \"Accordingly, the score function ∇x log pt (x) in (2) should be replaced by the conditional score $\\nabla_x \\log p_t(x|y)$. Unfortunately, this strictly restricts the generalization capability of the neural network since the conditional score should be retrained whenever the conditions change.\" This is not necessarily true as the score could be trained using some amortization w.r.t. y. In this case $y$ can be simply passed as another input to the diffusion model.\n\n* I don't really understand how the authors derive (14) from (13) and (6). From what I see, the authors are simply plugging $\\nabla \\log p(y|\\hat{x}_0)$ in the drift of the backward Markov chain. This makes a lot of sense, but I find the explanation provided by the authors to be quite convoluted.\n\n* Assumption 1 is not the manifold assumption as understood in the literature. Usually the manifold assumption ensures that the data lives on a lower dimensional manifold but do not impose any local linear structure.  It is quite hard to believe that the data of any reasonable image distribution is linear locally. I understand that the authors need this assumption to make their results rigorous but this should be clearly stated.\n\n* I really enjoyed reading about Proposition 2 and Theorem 1. I think that together with Figure 2 they bring a lot of insight about the method. However, the authors do not discuss at all the effect of the projection step $A x_i' + b$. I would have really appreciated an ablation study comparing the effects of these two terms which both introduce some information about the measurement.\n\n* I am not an expert in this domain but I have to admit that I am a bit surprised by the FID scores obtained by the authors which seem to be dramatically better than the rest of the methods. Could the authors provide details about the experimental process and in particular how they train the different models? For example do you perform hyperparameters tuning.\n\n* I think that the comparison with DDRM is a bit biased because DDRM only use 20 steps to recover the data whereas in the appendix the authors indicate that they use between 1000 and 2000 steps depending on the task. The trade-off between sampling quality and NFE is not investigated. Similarly, when comparing with RePaint the authors chose N = 200.", "limitations": "The limitations are discussed in the conclusion.\nI found this discussion to be appropriate.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "STRENGTHS:\n\n* The paper introduces a new methodology for inpainting which makes sense from a theoretical point of view, is easy to implement and seems to yield good results.\n\nWEAKNESSES:\n\n* I think that a deeper discussion with recent works such as [1] is missing. It is not clear from reading the main document what is the state of the art conditional sampling method using diffusion models (at least previous to this work).\n\n* In terms of novelty it seems that the proposed correction was already introduced in [2]. As a consequence, the method proposed by the authors can be see as a generalization to arbitrary linear problems (as stated by the authors) and therefore has limited novely.\n\n* With the stochastic contraction term and the gradient term proposed by the authors it is not clear at all if we are sampling from the right posterior distribution. I understand that the method yields good visual results but in fields like CT-reconstruction the visual quality should not be the only metrics as we also need a models which provide good uncertainty quantification. Such considerations are missing here. As a first question: can you characterize what is the invariant measure of the proposed diffusion model?\n\n[1] Kawar et al. (2022) -- Denoising diffusion restoration models\n[2] Ho et al. (2022) -- Video diffusion models", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "UkAPQgFA-jg", "reviewer": "Reviewer_mc79", "summary": "This paper deals with diffusion models for inverse problems. An important drawback of this approach is that to produce satisfactory results, it is necessary to iterate the forward and the reverse diffusions several times. Based on the manifold hypothesis, the authors show that this occurs because the score function only acts on the normal direction of the manifold, and therefore a drift occurs pushing the inference path out of the manifold. To solve this, besides the measurement constraint equation, they propose to include a constraint on the reverse diffusion equation given by the gradient of the $L^2$ residual of the inverse problem for the Bayes optimal denoising step from Noise2Noise (Kim and Ye, Neurips 2021). The authors prove that this constraint forces the diffusion to lie on the data manifold. These claims are supported by theoretical proofs and a clear geometric interpretation. \n", "questions": "None", "limitations": "Yes", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\n(+) The paper is very well written  \n(+) The idea of combining Noise2Score and diffusion models is novel, interesting, and completely sound.   \n(+) The proposed strategy is fully backed up by the theortical arguments and proofs provided by the authors.   \n(+) Thorough comparison with state-of-the-art diffusion model-based approaches and supervised learning-based baselines confirm the superiority of the proposed approach.   \n\n## Weaknesses\n\n(-) It would be nice to show and discuss failure cases, or situations when the proposed approach does not outperform the others.  \n\nMinor comments:\n\n- table X, figure Y, section Z, etc. --> Table X, Figure Y, Section Z, etc.\n- Eq. 3: $x_t$ --> $x(t)$\n- Fix punctuation at the end of Eqs. 6 and 9\n- L71: $\\mathbb{R}^n$ --> $\\mathbb{R}^m$\n- L76: utilize\n- L77: rely\n- Eqs. 13 and 14: for consistency, write them in the discrete setting\n- L144: uses\n- L154: recall the definition of $p_0$ (it was only defined in Section 2)\n- SuppMat, L465-466: check the sentence\n- SuppMat, Fig. 6: there are two lines in red that should be in green\n - SuppMat, L502: $\\epsilon_\\theta$ --> $z_\\theta$\n - SuppMat,L507: (4) --> Table 4\n - SuppMat, L509: (1) --> Algorithm 1\n", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Q8zlZ1JNU5_", "reviewer": "Reviewer_xYeQ", "summary": "The paper suggests to improve the performance a diffusion-based reconstruction scheme by adding to its updates the gradient of a log-likelihood term that uses the link between $x_i$ to $x_0$.\nSome mathematical motivation is provided.", "questions": "My questions and comments are listed above.", "limitations": "My questions and comments are listed above.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "(1). The introduction part is lacking. You should mention other methods that use the same pre-trained priors for solving many (almost arbitrary) reconstruction tasks (i.e., in an unsupervised fashion).\nFor example, GAN-based reconstruction methods [a],[b] and the powerful plug-and-play (P&P) denoisers approach [c],[d],[e].\nThe P&P denoisers approach [c], which is much older than diffusion-based reconstruction methods, commonly uses general-porpose CNN denoisers to impose the prior within iterative schemes [d], [e]. \nSo, in fact, diffusion-based reconstruction methods can be understood as P&P variant with class-specific denoisers (determined by the training data) that can handle extreme noise levels (which makes them generative models).\nI expect these topics to be discussed in the introduction.\n\n(2). The continuous formulation of the diffusion process in the introduction part seems redundant.\n\n(3). Fix the notation: in Eq. 15 you mean y rather than y0, right? \nAlso, the \"argument\" of $\\hat{x}_0(\\cdot)$ is $x_i$, right? this should be clearly written.\n\n(4). The presentation of the MCG ingredient needs to be improved. Essentially, the \"manifold constraint\" is the information on the prior of the signal: the transformation from $x_i$ to $x_0$ (e.g.,the constraint in Eq. 13 is independent of $y$).\nYet, in the term that is added to Eq. 15 (compared to Eq. 7) you already connect it with the log-likelihood term (data fidelity term).\nThus, it should be clearly stated throughout the paper that this new term combines data fidelity term and the prior information of the transformation from $x_i$ to $x_0$.\nFurthermore, I tend to believe that there exist other works that also use (perhaps in a similar way) the log-likelihood function in their iterative diffusion reconstruction schemes. Please make sure that you cover the exiting literature. \n\n(5). Note that for $\\alpha=0$ your additional term is canceled. Thus, your experiments should include ablation study and also some discussion and examination of the effect of different values of alpha.\n\n(6). Please verify that there is no mismatch between Eq. 15 and the algorithms that you actually use in the experiments.\nFor example, in Algorithm 1 in Appendix C, line 10 is not aligned with Eq. 15 (alpha is multiplied with a matrix).\nMoreover, observe that your MCG component in line 10 of Algorithm 1 vanishes(!) when you compute (I-P'P) * d/dxi ||P'*(y0-P*x0)||:\n\n(I-P'P) * d/dxi ||P'*(y0-P*x0)|| = (I-P'P) * d/dx0 ||P'*(y0-P*x0)|| * dx0/dxi = (I-P'P) * P'P*(P*x0-y0) * dx0/dxi = 0,\n\nbecause (I-P'P) * P'P=0, since for inpainting P is m rows of the identity matrix I_n so P'P is a projection matrix.\n\n(7). In your algorithm, *all* the choices of $A$ and $b$ for the different tasks (including CT) makes Eq. 16 (and Eq. 8) coincide with the \"back-projection step\" in [e].\nThe iterative usage of such a data consistency step with pre-trained denoisers (which, as mentioned above, is very similar to diffusion-based reconstruction) has been proposed in [e] (and theoretically analyzed in follow-up works). Such very related works should be mentioned.\n\n(8). The competitors in the experiments seem somewhat weak. For example, the recovery that is presented to some of them changes even the known pixels in inpainting. Recoveries that do not utilize such known information are obviously weak. Hence, it would have been interesting and informative to see comparison with strong (non-naive) GAN-based reconstruction methods, such as [b] (obviously, with GANs that are trained on the same training data as the other methods, e.g., by using common datasets and GANs).\nThe method in [b] is a clear example of a method that does not fall into the authors' statement that \"the ability to maintain perfect measurement consistency ... is often not satisfied for unsupervised GAN-based solutions\".\n\n(9). In the experiments section, you must include a discussion on the number of iterations and run-time of the different methods that are examined.\n\n\n[a] Bora, A., Jalal, A., Price, E. and Dimakis, A.G., 2017, July. Compressed sensing using generative models. In International Conference on Machine Learning (pp. 537-546). PMLR.\n\n[b] Hussein, S.A., Tirer, T. and Giryes, R., 2020, April. Image-adaptive GAN based reconstruction. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 34, No. 04, pp. 3121-3129).\n\n[c] Venkatakrishnan, S.V., Bouman, C.A. and Wohlberg, B., 2013, December. Plug-and-play priors for model based reconstruction. In 2013 IEEE Global Conference on Signal and Information Processing (pp. 945-948). IEEE.\n\n[d] Zhang, K., Zuo, W., Gu, S. and Zhang, L., 2017. Learning deep CNN denoiser prior for image restoration. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 3929-3938).\n\n[e] Tirer, T. and Giryes, R., 2018. Image restoration by iterative denoising and backward projections. IEEE Transactions on Image Processing, 28(3), pp.1220-1234.\n", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "JsAGFFw38Oh", "reviewer": "Reviewer_qC4w", "summary": "This paper proposes a framework for solving inverse problems using diffusion models as priors. The idea is to introduce another constraint in the formulation of the posterior distribution that enforces the generated sample to be on the manifold of natural images. The introduced constraint is motivated by the Tweedie formula, that relates the score function (of the noisy data distribution) and the best MMSE estimator (denoising).\nSeveral experimental results on different imaging inverse problems show that the proposed method outperforms existing methods in FID, LPIPS, PSNR and SSIM metrics.", "questions": "For completeness, I re-list here some of the questions raised above:\n\n1. Tweedie formula vs Noise2score.The key motivation seems to be coming from the Tweedie formula. With that in mind it would be better to introduce the Tweedie formula and relevant work better. See for example Section IIB in [Ong2019]:\n\n[Ong2019] Ong, F., Milanfar, P. and Getreuer, P., 2019. Local kernels that approximate Bayesian regularization and proximal operators. IEEE Transactions on Image Processing, 28(6), pp.3007-3019.\n\n2. Motivation for the additional constraint is not clear (Section 3.1). In particular Eq (14) deduction is unclear (\\mathcal{C} is a set, so how should one interpret p(\\mathcal{C} | x) ? Could you please elaborate more on how Eq (14) is introduced (motivation/deduction). \n\n3. How is obtained W for each application, it is unclear where this W is coming from. Explaining better this term in (14) will help a lot to the understanding of the algorithm.\n\nOther minor comments:\n- l59 -  p_t(x) should be log p_t(x)\n- Eq (3) - \\epsilon is undefined maybe just add \"where \\epsilon <<1\".\n- Eq (18), Eq (20),  Eq (22). In all these equations it says \"We chose W…\" This implies that W is not deducted from the model but it's something that we need to choose. It is unclear how to do this and how it's related to the introduced motivation where everything seem to be coming naturally from the Tweedie formula. \n\nI'm more than happy to increase my final rating if during the rebuttal we manage to clarify some of the aspects regarding this work.\n", "limitations": "Limitations are correctly discussed.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper introduces an interesting adaptation to traditional sampling methods motivated by the Tweedie formula. \n- The paper is generally well written and the idea is well motivated.\n- Several experimental results on different imaging inverse problems show that the proposed method outperforms existing methods in FID, LPIPS, PSNR and SSIM metrics.\n\nWeakness:\n- The paper exposition could be better:\n- The paper relates the idea a lot to noise2score, but the key motivation is the Tweedie formula. With that in mind it would be better to introduce the Tweedie formula and relevant work better (see below)\n- Motivation for the additional constraint is not clear (Section 3.1). In particular Eq (14) deduction is unclear (\\mathcal{C} is a set, so how should one interpret p(\\mathcal{C} | x) ?\n- How is obtained W for each application.\n- Theoretical findings should be better connected to the proposed algorithm. \n\nIn general I like the paper, the idea is interesting, and the experimental results show a clear win in performance. The main current issue to me is that the motivation and exposition is not clear enough. Most of the pieces are already there but a reorganization and better discussion is needed. \n", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept