Title: Direct Diffusion Bridge using Data Consistency for Inverse Problems
Venue: NeurIPS 2023
Accepted: yes
Keywords: Diffusion models

Inverse problems

Diffusion bridge
Abstract: Diffusion model-based inverse problem solvers have shown impressive performance, but are limited in speed, mostly as they require reverse diffusion sampling starting from noise. Several recent works have tried to alleviate this problem by building a diffusion process, directly bridging the clean and the corrupted for specific inverse problems. In this paper, we first unify these existing works under the name Direct Diffusion Bridges (DDB), showing that while motivated by different theories, the resulting algorithms only differ in the choice of parameters. Then, we highlight a critical limitation of the current DDB framework, namely that it does not ensure data consistency. To address this problem, we propose a modified inference procedure that imposes data consistency without the need for fine-tuning. We term the resulting method data Consistent DDB (CDDB), which outperforms its inconsistent counterpart in terms of both perception and distortion metrics, thereby effectively pushing the Pareto-frontier toward the optimum. Our proposed method achieves state-of-the-art results on both evaluation criteria, showcasing its superiority over existing methods. Code is open-sourced [here](https://github.com/HJ-harry/CDDB).

Direct Diffusion Bridge using Data Consistency for Inverse Problems

Abstract

Diffusion model-based inverse problem solvers have shown impressive performance, but are limited in speed, mostly as they require reverse diffusion sampling starting from noise. Several recent works have tried to alleviate this problem by building a diffusion process, directly bridging the clean and the corrupted for specific inverse problems. In this paper, we first unify these existing works under the name Direct Diffusion Bridges (DDB), showing that while motivated by different theories, the resulting algorithms only differ in the choice of parameters. Then, we highlight a critical limitation of the current DDB framework, namely that it does not ensure data consistency. To address this problem, we propose a modified inference procedure that imposes data consistency without the need for fine-tuning. We term the resulting method data Consistent DDB (CDDB), which outperforms its inconsistent counterpart in terms of both perception and distortion metrics, thereby effectively pushing the Pareto-frontier toward the optimum. Our proposed method achieves state-of-the-art results on both evaluation criteria, showcasing its superiority over existing methods. Code is open-sourced at https://github.com/HJ-harry/CDDB

1 Introduction

Diffusion models [15, 38] have become the de facto standard of recent vision foundation models [32, 31, 33]. Among their capabilities is the use of diffusion models as generative priors that can serve as plug-and-play building blocks for solving inverse problems in imaging [18, 38, 21, 5]. D iffusion model-based i nverse problem s olvers (DIS) have shown remarkable performance and versatility, as one can leverage the powerful generative prior regardless of the given problem at hand, scaling to linear [18, 38, 21], non-linear [5, 36], and noisy problems [21, 5].

Although there are many advantages of DIS, one natural limitation is its slow inference. Namely, the overall process of inference-starting from Gaussian noise and being repeatedly denoised to form a clean image-is kept the same, although there are marginal changes made to keep the sampling process consistent with respect to the given measurement. In such cases, the distance between the reference Gaussian distribution and the data distribution remains large, requiring inevitably a large number of sampling steps to achieve superior sample quality. On the other hand, the distribution of the measurements is much more closely related to the distribution of the clean images. Thus, intuitively, it would cost us much less compute if we were allowed to start the sampling process directly from the measurement, as in the usual method of direct inversion in supervised learning schemes.

Interestingly, several recent works aimed to tackle this problem under several different theoretical motivations: 1) Schrödinger bridge with paired data [26], 2) a new formulation of the diffusion process

28.5

DDB

CDDB

40

DDB

CDDB

37.5

40.0

★

DDB

CDDB

28.0

CDDB(deep)

35

CDDB(deep)

35.0

worse

10

27.5

32.5

PSNR(

27.0

FID(

一

→

30

FID(

→

30.0

26.5

27.5

50

26.0

25

25.0

1000

100

20,

betrer

22.5

25.5

20

20.0

1000

100

50

Z

5

10

20

NFE

50

100

1000

2

5

10

20

NFE

50

100

1000

25.5

26.0

26.5

PSNR( ↑ )

27.0

27.5

28.0

28.5

via constant-speed continual degradation [9], and 3) Ornstein-Uhlenbeck stochastic differential equation (OU-SDE) [28]. While developed from distinct motivations, the resulting algorithms can be understood in a unifying framework with minor variations: we define this new class of methods as D irect D iffusion B ridges (DDB; Section 3.1). In essence, DDB defines the diffusion process from the clean image distribution in t = 0 to the measurement distribution in t = 1 as the convex combination between the paired data, such that the samples x t goes through continual degradation as t = 0 → 1 . In training, one trains a time-conditional neural network G θ that learns a mapping to x 0 for all timesteps, resulting in an iterative sampling procedure that reverts the measurement process.

Using such an iterative sampling process, one can flexibly choose the number of neural function evaluations (NFE) to generate reconstructions that meet the desiderata: with low NFE, less distortion can be achieved as the reconstruction regresses towards the mean [9]; with high NFE, one can opt for high perceptual quality at the expense of some distortion from the ground truth. This intriguing property of DDB creates a Pareto-frontier of reconstruction quality, where our desire would be to maximally pull the plot towards high perception and low distortion (bottom right corner of Fig. 1c).

In this work, we assert that DDB is missing a crucial component of data consistency , and devise methods to make the models consistent with respect to the given measurement by only modifying the sampling algorithm, without any fine-tuning of the pre-trained model. We refer to this new class of models as data C onsistent D irect D iffusion B ridge (CDDB; Section 3.2), and show that CDDB is capable of pushing the Pareto-frontier further towards the optima (lower distortion: Fig. 1a, higher perception: Fig. 1b, overall trend: Fig. 1c) across a variety of tasks. Theoretically, we show that CDDB is a generalization of DDS (Decomposed Diffusion Sampling) [6], a recently proposed method tailored for DIS with Gaussian diffusion, which guarantees stable and fast sampling. We then propose another variation, CDDB-deep, which can be derived as the DDB analogue of the DPS [5] by considering deeper gradients, which even further boosts the performance for certain tasks and enables the application to nonlinear problems where one cannot compute gradients in the usual manner (e.g. JPEG restoration). In the experiments, we showcase the strengths of each algorithm and show how one can flexibly construct and leverage the algorithms depending on the circumstances.

2 Background

2.1 Diffusion models

Diffusion models [15, 38, 22, 19] defines the forward data noising process p ( x t | x 0 ) as

where α t , σ t controls the signal component and the noise component, respectively, and are usually designed such that α 2 t + σ 2 t = 1 [15, 22]. Starting from the data distribution p data := p ( x 0 ) , the noising process in (1) gradually maps p ( x t ) towards isotropic Gaussian distribution as t → 1 , i.e. p ( x 1 ) ≃ N (0 , I ) . Training a neural network to reverse the process amounts to training a residual denoiser

such that ϵ ( t ) θ ∗ ( x t ) ≃ x t -α t x 0 σ t . Furthermore, it can be shown that epsilon matching is equivalent to the denoising score matching (DSM) [16, 37] objective up to a constant with different parameterization

such that s ( t ) θ ∗ ( x t ) ≃ -x t -α t x 0 σ 2 t = -ϵ ( t ) θ ∗ ( x t ) /σ t . Moreover, for optimal θ ∗ and under regularity conditions, s θ ∗ ( x t ) = ∇ x t log p ( x t ) . Then, sampling from the distribution can be performed by solving the reverse-time generative SDE/ODE [38, 19] governed by the score function. It is also worth mentioning that the posterior mean, or the so-called denoised estimate can be computed via Tweedie's formula [13]

In practice, DDPM/DDIM solvers [15, 35] work by iteratively refining these denoised estimates.

2.2 Diffusion model-based inverse problem solving with gradient guidance

Suppose now that we are given a measurement y obtained through some Gaussian linear measurement process A , where our goal is to sample from the posterior distribution p ( x | y ) . Starting from the sampling process of running the reverse SDE/ODE to sample from the prior distribution, one can modify the score function to adapt it for posterior sampling [38, 5]. By Bayes rule, ∇ x t log p ( x t | y ) = ∇ x t log p ( x t ) + ∇ x t log p ( y | x t ) , where ∇ x t log p ( x t ) ≃ s θ ∗ ( x t ) . However, ∇ x t log p ( y | x t ) is intractable. Several methods have been proposed to approximate this time-dependent likelihood, two of the most widely used being DPS [5] and Π GDM[36]. DPS proposes the following Jensen approximation 1

of which the chain rule is based on the denominator layout notation [41]. Here, we see that the gradient term can be represented as the Jacobian ( J ) vector ( V ) product (JVP). In the original implementation of DPS, the two terms are not computed separately, but computed directly as ∇ x t ∥ y -A ˆ x 0 | t ∥ 2 2 , where the whole term can be handled with backpropagation. By this choice, DPS can also handle non-linear operators when the gradients can be computed, e.g. phase retrieval, forward model given as a neural network. On the other hand, Π GDMproposes

where A † := A ⊤ ( AA ⊤ ) -1 is the Moore-Penrose pseudo-inverse. Using the JVP for implementation, it is no longer required that the whole term is differentiable. For this reason, Π GDMcan be applied to cases where we have non-differentiable, non-linear measurements given that an operation analogous to pseudo-inverse can be derived, e.g. JPEG restoration. Notably, the update step of DPS can be achieved by simply pre-conditioning Π GDM with A ⊤ A . Implementing DIS with DPS (5) or Π GDM (6) amounts to augmenting the gradient descent steps in between the ancestral sampling iterations.

While these methods are effective and outperforms the prior projection-based approaches [38, 21], they also have several drawbacks. Namely, the incorporation of the U-Net Jacobian is slow, computeheavy, and often unstable [12, 34]. For example, when applied to MRI reconstruction in medical imaging, DPS results in noisy reconstructions [6] possibly due to unstable incorporation of the Wirtinger derivatives [23], and Π GDMis hard to use as it is non-trivial to compute A † . In order to circumvent these issues, DDS [6] proposed to use numerical optimization (i.e. conjugate gradients; CG) in the clean image domain, bypassing the need to compute J . Consequently, DDS achieves fast and stable reconstructions for inverse problems in medical imaging.

3 Main Contributions

3.1 Direct Diffusion Bridge

We consider the case where we can sample x 0 := x ∼ p ( x ) , and x 1 := y ∼ p ( y | x ) 2 , i.e. paired data for training. Adopting the formulation of I 2 SB [26] we define the posterior of x t to be the product of Gaussians N ( x t ; x 0 , γ 2 t ) and N ( x t ; x 1 , ¯ γ 2 t ) , such that

Note that the sampling of x t from (7) can be done by the reparametrization trick

where α t := γ 2 t γ 2 t +¯ γ 2 t , σ 2 t := γ 2 t ¯ γ 2 t γ 2 t +¯ γ 2 t 3 . This diffusion bridge introduces a continual degradation process by taking a convex combination of ( x 0 , x 1 ) , starting from the clean image at t = 0 to maximal degradation at t = 1 , with additional stochasticity induced by the noise component σ t . Our goal is to train a time-dependent neural network that maps any x t to x 0 that recovers the clean image. The training objective for I 2 SB [26] analogous to denoising score matching (DSM) [16] reads

which is also equivalent to training a residual network G θ with min θ E [ ∥ G θ ( x t ) -x 0 ∥ 2 2 ] . For brevity, we simply denote the trained networks as G θ ∗ even if it is parametrized otherwise. Once the network is trained, we can reconstruct x 0 starting from x 1 by, for example, using DDPM ancestral sampling [15], where the posterior for s < t reads

with α 2 s | t := γ 2 s γ 2 t , σ 2 s | t := ( γ 2 t -γ 2 s ) γ 2 s γ 2 t . At inference, x 0 is replaced with a neural network-estimated ˆ x 0 | t to yield x s ∼ p ( x s | ˆ x 0 | t , x t ) .

However, when the motivation is to introduce 1) a tractable training objective that learns to recover the clean image along the degradation trajectory, and 2) devise a sampling method to gradually revert the degradation process, we find that the choices made for the parameters in (8),(7) can be arbitrary, as long as the marginal can be retrieved in the sampling process (10). In Table 1, we summarize the

choices made in [9, 26] to emphasize that the difference stems mostly from the parameter choices and not something fundamental. Concretely, sampling x t from paired data can always be represented as (8): a convex combination of x 0 and x 1 with some additional noise. Reverse diffusion at inference can be represented as (10): a convex combination of x 0 and x t with some stochasticity. We define the methods that belong to this category as Direct Diffusion Bridge (DDB) henceforth. Below, we formally state the equivalence between the algorithms, with proofs given in Appendix A.

Theorem 1. Let the parameters of InDI [9] in Table 1 be t := γ 2 t γ 2 t +¯ γ 2 t , ϵ 2 t := ¯ γ 2 t γ 2 t ( γ 2 t + ¯ γ 2 t ) . Then, InDI and I 2 SB are equivalent.

The equivalence relation will be useful when we derive our CDDB algorithm in Section. 3.2. As a final note, IR-SDE [28] does not strictly fall into this category as the sampling process is derived from running the reverse SDE. However, the diffusion process can still be represented as (8) by setting α t = 1 -e -¯ θ t , σ 2 t = λ 2 (1 -e -2 ¯ θ t ) , and the only difference comes from the sampling procedure.

3.2 Data Consistent Direct Diffusion Bridge

Motivation Regardless of the choice in constructing DDB, there is a crucial component that is missing from the framework. While the sampling process (10) starts directly from the measurement (or equivalent), as the predictions ˆ x 0 | t = G θ ( x t ) are imperfect and are never guaranteed to preserve the measurement condition y = Ax , the trajectory can easily deviate from the desired path, while the residual blows up. Consequently, this may result in inferior sample quality, especially in terms of distortion. In order to mitigate this downside, our strategy is to keep the DDB sampling strategy (10) intact and augment the steps to constantly guide the trajectory to satisfy the data consistency, similar in spirit to gradient guidance in DIS. Here, we focus on the fact that the clean image estimates ˆ x 0 | t is produced at every iteration, which can be used to compute the residual with respect to the measurement y . Taking a gradient step that minimizes this residual after every sampling step results in Algorithm 1, which we name data Consistent DDB (CDDB). In the following, we elaborate on how the proposed method generalizes DDS which was developed for DIS.

CDDB as a generalization of DDS [6] Rewriting (10) with reparameterization trick

we see that the iteration decomposes into three terms: the denoised component, the deterministic noise, and the stochastic noise. The key observation of DDIM [35] is that if the score network is fully expressive, then the deterministic noise term x t -ˆ x 0 | t becomes Gaussian such that it satisfies the total variance condition

allowing (11) to restore the correct marginal N ( x s ; x 0 , σ 2 s ) . Under this condition, DDS showed that using a few step of numerical optimization ensure the updates from the denoised image ˆ x 0 | t remain on the clean manifold. Furthermore, subsequent noising process using deterministic and stochastic noises can then be used to ensure the transition to the correct noisy manifold [6].

Under this view, our algorithm can be written concisely as

where we make the update step only to the clean denoised component, and leave the other components as is. In order to achieve proper sampling that obeys the marginals, it is important to show that the remaining components constitute the correct noise variance and the condition assuming Gaussianity should be (12). In the following, we show that this is indeed satisfied for the two cases of direction diffusion bridge (DDB):

Theorem 2. The total variance condition (12) is satisfied for both I 2 SB and InDI.

Proof. For InDI, considering the noise variance σ t = tϵ t in Table 1,

Due to the equivalence in Theorem 1, the condition is automatically satisfied in I 2 SB. We show that this is indeed the case in Appendix A.

In other words, given that the gradient descent update step in CDenoise( x t ) does not leave the clean data manifold, it is guaranteed that the intermediate samples generated by (13) will stay on the correct noisy manifold [6]. In this regard, CDDB can be thought of as the DDB-generalized version of DDS. Similar to DDS, CDDB does not require the computation of heavy U-Net Jacobians and hence introduces negligible computation cost to the inference procedure, while being robust in the choice of step size.

CDDB-deep As shown in DPS and Π GDM,taking deeper gradients by considering U-Net Jacobians is often beneficial for reconstruction performance. Moreover, it even provides way to impose data consistency for non-linear inverse problems, where standard gradient methods are not feasible. In order to devise an analogous method, we take inspiration from DPS, and propose to augment the solver with a gradient step that maximizes the time-dependent likelihood (w.r.t. the measurement) p ( y | x t ) . Specifically, we use the Jensen approximation from [5]

where the last equality is naturally satisfied from the training objective (9). Using the approximation used in (15), the correcting step under the Gaussian measurement model yields

Implementing (16) in the place of the shallow gradient update step of Algorithm 1, we achieve CDDB-deep (see Algorithm 2). From our initial experiments, we find that preconditioning with A † as in Π GDMimproves performance by a small margin, and hence use this setting as default.

4 Experiments

4.1 Setup

Model, Dataset For a representative DDB, we choose I 2 SB [26] along with the pre-trained model weights for the following reasons: 1) it is open-sourced 4 , 2) it stands as the current state-of-the-art, 3) the model architecture is based on ADM [11], which induces fair comparison against other DIS methods. All experiments are based on ImageNet 256 × 256 [10], a benchmark that is considered to be much more challenging for inverse problem solving based on generative models [5], compared to more focused datasets such as FFHQ [20]. We follow the standards of [26] and test our method on the following degradations: sr4x-{bicubic, pool}, deblur-{uniform, gauss}, and JPEG restoration with 1k validation images.

Degraded

Ground Truth

PnP-ADMM

DPS

I2SB

CDDB (ours)

CDDB-deep (ours)

SR

(

×

4)

Deblur

bicubic

pool

gauss

uniform

Method

PSNR

↑

SSIM

↑

LPIPS

↓

FID

↓

PSNR

↑

SSIM

↑

LPIPS

↓

FID

↓

PSNR

↑

SSIM

↑

LPIPS

↓

FID

↓

PSNR

↑

SSIM

↑

LPIPS

↓

FID

↓

CDDB (ours)

26.41

0.860

0.198

19.88

26.36

0.855

0.184

17.79

37.02

0.978

0.059

5.007

31.26

0.927

0.193

23.15

I

2

SB [26]

25.22

0.802

0.260

24.13

25.08

0.800

0.258

23.53

36.01

0.973

0.067

5.800

30.75

0.919

0.198

23.01

DPS [5]

19.89

0.498

0.384

63.37

21.01

0.562

0.326

49.34

27.21

0.766

0.244

34.58

22.51

0.565

0.357

60.00

Π

GDM[36]

26.20

0.850

0.252

29.36

26.07

0.849

0.256

26.97

-

-

-

-

-

-

-

-

DDRM [21]

26.05

0.838

0.270

46.49

25.54

0.848

0.257

40.40

36.73

0.975

0.071

4.346

29.21

0.901

0.210

19.97

DDNM[40]

26.41

0.801

0.230

38.63

26.04

0.792

0.218

33.15

-

-

-

-

-

-

-

-

DDS [6]

26.41

0.801

0.230

38.64

26.04

0.792

0.218

33.15

33.27

0.945

0.057

6.442

27.88

0.829

0.193

26.07

PnP-ADMM [4] 26.16

0.788

0.350

74.06

25.85

0.733

0.372

72.63

28.18

0.800

0.325

60.27

25.47

0.701

0.416

83.76

ADMM-TV

22.55

0.595

0.493

122.7

22.31

0.574

0.512

119.4

24.67

0.773

0.324

50.74

21.72

0.600

0.491

98.15

Baselines, Evaluation Along with the most important comparison against I 2 SB, we also include comparisons with state-of-the-art DIS methods including DDRM [21], DPS [5], Π GDM [36], DDNM [40], and DDS [6]. For choosing the NFE and the hyper-parameters for each method, we closely abide to the original advised implementation: DDRM (20 NFE), DPS (1000 NFE), Π GDM(100 NFE), DDNM (100 NFE), DDS (100 NFE). We find that increasing the NFE for methods other than DPS does not induce performance gain. We also note that we exclude Π GDM and DDNM for the baseline comparison in the deblurring problem, as directly leveraging the pseudoinverse matrix may result in unfair boost in performance [29]. We additionally perform comparisons against PnP-ADMM [4] and ADMM-TV, which are methods tailored towards higher SNR. For I 2 SB along with the proposed method, we choose 100 NFE for JPEG restoration as we found it to be the most stable, and choose 1000 NFE for all other tasks. Further details on the experimental setup can be found in Appendix C.

4.2 Results

Comparison against baselines Throughout the experiments, we thoroughly analyze both distortion and perception of the reconstructions obtained through our method against other DDB, DIS, and iterative optimization methods. Note that for the recent diffusion-based methods, analysis has been

Degraded

I2SB

CDDB (ours)

Ground Truth

Degraded

I2SB

CDDB (ours)

Ground Truth

56宠物网

56宠物风

www.S6pel.s

www.S6pel.

mainly focused on perceptual metrics [5, 36, 26], mainly because these methods excel on these metrics, but often compromising distortion metrics. DDB methods often take this to the extreme, where one can achieve the best PSNR with very little NFE, and the PSNR consistently degrades as one increases the NFE [9, 26] (See Fig. 1c). Despite this fact, the standard setting in DDB methods is to set a high NFE as one can achieve much improved perceptual quality. On the other hand, conventional iterative methods are often highly optimized for less distortion, albeit with low perceptual quality. While this trade-off may seem imperative, we show that CDDB can improve both aspects, putting it in the place of the state-of-the-art on most experiments (See Tab. 2, 3). A similar trend can be observed in Fig. 2, where we see that CDDB greatly improves the performance of DDB, while also outperforming DIS and iterative optimization methods.

CDDBpushes forward the Pareto-frontier It is widely known that there exists an inevitable trade-off of distortion when aiming for higher perceptual quality [3, 24]. This phenomenon has been reconfirmed consistently in the recent DIS [5, 36] and DDS [9, 26] methods. For DDB, one can flexibly control this trade-off by simply choosing different NFE values, creating a Pareto-frontier with higher NFE tailored towards perceptual quality. While this property is intriguing, the gain we achieve when increasing the

NFE decreases exponentially , and eventually reaches a bound when NFE > 1000. In contrast, we show in Fig. 1 that CDDB pushes the bound further towards the optima. Specifically, 20 NFE CDDB outperforms 1000 NFE DDB in PSNR by > 2 db, while having lower FID (i.e. better perceptual quality). To this point, CDDB induces dramatic acceleration (> 50 × ) to DDB.

CDDB vs. CDDB-deep The two algorithms presented in this work share the same spirit but have different advantages. CDDB generally has higher speed and stability, possibly due to guaranteed convergence. As a result, it robustly increases the performance of SR and deblurring. In contrast, considering the case of inpainting and JPEG restoration, CDDB cannot improve the performance of DDB. For inpainting, the default setting of I 2 SB ensures consistency by iteratively applying replacement, as implemented in [40]. As the measurement stays in the pixel space, the gradients cannot impose any constraint on the missing pixel values. CDDB-deep is useful in such a situation, as the U-Net Jacobian has a global effect on all the pixels, improv-

Degraded

I2SB

CDDB (ours)

Ground Truth

ing the performance by inducing coherence. CDDB-deep also enables the extension to nonlinear inverse problems where one cannot take standard gradient steps. This is illustrated for the case of JPEG restoration in Tab. 3 and Fig. 4, where we see overall improvement in performance compared to I 2 SB.

Noise robustness DIS methods are often designed such that they are robust to measurement noise [21, 5]. In contrast, this is not the case for DDB as they are trained in a supervised fashion: If not explicitly trained with synthetic adding of noise, the method does not generalize well to noisy measurements, as can be seen in Fig. 5. On the other hand, note that with CDDB, we are essentially incorporating a Gaussian likelihood model, which naturally enhances the robustness to noise. As a result, while I 2 SB tends to propagate noise (best seen in the background), we do not observe such artifacts when using CDDB.

5 Discussion

Extension to other related works Going beyond the paired inverse problem setting and considering the Schrödinger Bridge (SB) problem [25, 8], or more generally transport mapping problems [27] between the two unmatched distributions, it is often desirable to control the deviation from the start of sampling. A concrete example would be the case of image-to-image translation [43] where one does not want to alter the content of the image. As CDDB can be thought of as a regularization method that penalizes the deviation from the starting point, the application is general and can be extended to such SB problems at inference time by using the gradients that minimize the distance from the start point. We leave this direction for future work.

Data consistency in supervised learning frameworks The first applications of supervised deep learning to solve inverse problems in medical imaging (e.g. CT [17], MRI reconstruction [39]) mostly involved directly inverting the measurement signal without considering the measurement constraints. The works that followed [1, 14] naturally extended the algorithms by incorporating measurement consistency steps in between the forward passes through the neural network. Analogously, CDDB is a natural extension of DDB but with high flexibility, as we do not have to pre-determine the number of forward passes [1] or modify the training algorithm [14].

6 Conclusion

In this work, we unify the seemingly different algorithms under the class of direct diffusion bridges (DDB) and identify the crucial missing part of the current methods: data consistency. Our train-free modified inference procedure named consistent DDB (CDDB) fixes this problem by incorporating consistency-imposing gradient steps in between the reverse diffusion steps, analogous to the recent DIS methods. We show that CDDB can be seen as a generalization of representative DIS methods (DDS, DPS) in the DDB framework. We validate the superiority of our method with extensive experiments on diverse inverse problems, achieving state-of-the-art sample quality in both distortion and perception. Consequently, we show that CDDB can push the Pareto-frontier of the reconstruction toward the desired optimum.

Limitations and societal impact The proposed method assumes prior knowledge of the forward operator. While we limit our scope to non-blind inverse problems, the extension of CDDB to blind inverse problems [7, 30] will be a possible direction of research. Moreover, for certain inverse problems (e.g. inpainting), even when do observe improvements in qualitative results, the quantitative metrics tend to slightly decrease overall. Finally, inheriting from DDS/DIS methods, our method relies on strong priors that are learned from the training data distribution. This may potentially lead to reconstructions that intensify social bias and should be considered in practice.

𝜎 = 0.015

I2SB

CDDB (ours)

Ground Truth

Acknowledgments and Disclosure of Funding

This research was supported by the KAIST Key Research Institute (Interdisciplinary Research Group) Project, by the National Research Foundation of Korea under Grant NRF-2020R1A2B5B03001980, by the Korea Medical Device Development Fund grant funded by the Korea government (the Ministry of Science and ICT, the Ministry of Trade, Industry and Energy, the Ministry of Health & Welfare, the Ministry of Food and Drug Safety) (Project Number: 1711137899, KMDF_PR_20200901_0015), by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government(MSIT) (No.2019-0-00075, Artificial Intelligence Graduate School Program(KAIST)), and by the Field-oriented Technology Development Project for Customs Administration through National Research Foundation of Korea(NRF) funded by the Ministry of Science & ICT and Korea Customs Service(NRF-2021M3I1A1097938).

References

Hemant K Aggarwal, Merry P Mani, and Mathews Jacob. MoDL: Model-based deep learning architecture for inverse problems. IEEE transactions on medical imaging , 38(2):394-405, 2018.

Thilo Balke, Fernando Davis, Cristina Garcia-Cardona, Michael McCann, Luke Pfister, and Brendt Wohlberg. Scientific Computational Imaging COde (SCICO). Software library available from https://github.com/lanl/scico , 2022.

Yochai Blau and Tomer Michaeli. The perception-distortion tradeoff. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 6228-6237, 2018.

Stanley H Chan, Xiran Wang, and Omar A Elgendy. Plug-and-play ADMM for image restoration: Fixed-point convergence and applications. IEEE Transactions on Computational Imaging , 3(1):84-98, 2016.

Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Marc Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In International Conference on Learning Representations , 2023.

Hyungjin Chung, Suhyeon Lee, and Jong Chul Ye. Fast diffusion sampler for inverse problems by geometric decomposition. arXiv preprint arXiv:2303.05754 , 2023.

Hyungjin Chung, Dohoon Ryu, Michael T Mccann, Marc L Klasky, and Jong Chul Ye. Solving 3d inverse problems using pre-trained 2d diffusion models. IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2023.

Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems , 34:17695-17709, 2021.

Mauricio Delbracio and Peyman Milanfar. Inversion by direct iteration: An alternative to denoising diffusion for image restoration. arXiv preprint arXiv:2303.11435 , 2023.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A largescale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee, 2009.

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.

Yilun Du, Conor Durkan, Robin Strudel, Joshua B Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Grathwohl. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and mcmc. In International conference on machine learning . PMLR, 2023.

Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association , 106(496):1602-1614, 2011.

Harshit Gupta, Kyong Hwan Jin, Ha Q Nguyen, Michael T McCann, and Michael Unser. Cnn-based projected gradient descent for consistent ct image reconstruction. IEEE transactions on medical imaging , 37(6):1440-1453, 2018.

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

Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research , 6(4), 2005.

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.

Zahra Kadkhodaie and Eero P Simoncelli. Stochastic solutions for linear inverse problems using the prior implicit in a denoiser. In A. Beygelzimer, Y . Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , 2021.

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. In Proc. NeurIPS , 2022.

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.

Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

Diederik P Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. arXiv preprint arXiv:2107.00630 , 2021.

Ken Kreutz-Delgado. The complex gradient operator and the cr-calculus. arXiv preprint arXiv:0906.4835 , 2009.

Christian Ledig, Lucas Theis, Ferenc Huszár, Jose Caballero, Andrew Cunningham, Alejandro Acosta, Andrew Aitken, Alykhan Tejani, Johannes Totz, Zehan Wang, et al. Photo-realistic single image super-resolution using a generative adversarial network. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 4681-4690, 2017.

Christian Léonard. A survey of the schrödinger problem and some of its connections with optimal transport. arXiv preprint arXiv:1308.0215 , 2013.

Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A Theodorou, Weili Nie, and Anima Anandkumar. I 2 SB: Image-to-Image Schrödinger Bridge. In International conference on machine learning . PMLR, 2023.

Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations , 2023.

Ziwei Luo, Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B Schön. Image restoration with mean-reverting stochastic differential equations. In International conference on machine learning . PMLR, 2023.

Sladjana Miljkovi´ c, Marko Miladinovi´ c, Predrag Stanimirovi´ c, and Igor Stojanovi´ c. Application of the pseudoinverse computation in reconstruction of blurred images. Filomat , 26(3):453-465, 2012.

Naoki Murata, Koichi Saito, Chieh-Hsin Lai, Yuhta Takida, Toshimitsu Uesaka, Yuki Mitsufuji, and Stefano Ermon. Gibbsddrm: A partially collapsed gibbs sampler for solving blind inverse problems with denoising diffusion restoration. In International conference on machine learning . PMLR, 2023.

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 , 2022.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10684-10695, 2022.

Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems , 35:36479-36494, 2022.

Tim Salimans and Jonathan Ho. Should ebms model the energy or the score? In Energy Based Models Workshop-ICLR 2021 , 2021.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In 9th International Conference on Learning Representations, ICLR , 2021.

Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse-guided diffusion models for inverse problems. In International Conference on Learning Representations , 2023.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems , volume 32, 2019.

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.

Shanshan Wang, Zhenghang Su, Leslie Ying, Xi Peng, Shun Zhu, Feng Liang, Dagan Feng, and Dong Liang. Accelerating magnetic resonance imaging via deep learning. In 2016 IEEE 13th international symposium on biomedical imaging (ISBI) , pages 514-517. IEEE, 2016.

Yinhuai Wang, Jiwen Yu, and Jian Zhang. Zero-shot image restoration using denoising diffusion null-space model. In The Eleventh International Conference on Learning Representations , 2023.

Jong Chul Ye. Geometry of Deep Learning . Springer, 2022.

Kai Zhang, Wangmeng Zuo, Yunjin Chen, Deyu Meng, and Lei Zhang. Beyond a gaussian denoiser: Residual learning of deep CNN for image denoising. IEEE transactions on image processing , 26(7):3142-3155, 2017.

Jun-Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. In Proceedings of the IEEE international conference on computer vision , pages 2223-2232, 2017.

Peer reviews:
comments: SUMMARY:
This paper focuses on the diffusion model-based inversion problem. The paper first analyzes the current works and unifies them with the Direct Diffusion Bridges. Then, the authors point out that the data consistency is ignored by current works and propose CDDB. Experiments show that the proposed CDDB works well on various inverse tasks like inpainting, super-resolution, and deblurring.

STRENGTHS:
1. This paper summarized the current work in a systematic way.

2. The differences between this work and previous works are clear.

3. The paper is supported well by theorems and experiments.

WEAKNESSES:
I am not so familiar with the related works and it is hard for me to judge the novelty and contribution of this paper, but the logic of this paper is reasonable, so I can only give borderline accept and I will increase my score if more experts support this paper. I am also wondering whether the work can be applied to the current popular textual inversion problem, since its current form only supports from one noisy image to a clean image.

QUESTIONS:
Please refer to the weakness part.

LIMITATIONS:
Some limitations are discusses in the Section 5.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-16

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors propose Consistent Direct Diffusion Bridge (CDDB) a modification of the Direct Diffusion Bridge (DDB) procedure [1,2] that includes a data consistency term similar to [3,4]. The authors propose two ways to introduce the consistency: the first one forgets about the Jacobian term which is hard to compute in the guidance approximation. The second one (called CDDB-deep) compute the gradient through the observation term. It is more costly but can lead to improved results. The authors briefly justify their method in the spirit of [5] and then showcase the efficiency of their method on several inverse problems (inpainting/deblurring) on ImageNet 256x256.\n\n[1] Liu et al. (2023) -- SB: Image-to-Image Schrödinger Bridge\n\n[2] Delbracio, Milanfar (2023) -- Inversion by Direct Iteration: An Alternative to Denoising Diffusion for Image Restoration\n\n[3] Song et al. (2023) -- Pseudoinverse-Guided Diffusion Models for Inverse Problems\n\n[4] Chung et al. (2023) -- Improving Diffusion Models for Inverse Problems using Manifold Constraints\n\nSTRENGTHS:\n* The paper is very well-written and the explanations are very clear. The presentation of the method is sound and the authors did a good job at recalling the basics on diffusion models and I2SB.\n\n* I really appreciated the discussion on the difference (and ultimately the equivalence) between I2SB [1] and INDI [2]. This discussion is very well presented.\n\n* The experiments are strong. The authors have done a very thorough work investigating several metrics on challenging problems in ImageNet 256x256. \n\n* While one could complain that the novelty of the method is relatively low (basically combining I2SB with a guidance term), I think the experiments are extensive enough to warrant publication.\n\nWEAKNESSES:\n* I think this is overall a good paper and the next point is not the central problem tackled by the paper but I think it should be clarified (or at least the related remarks should be tamed). The authors spend a few lines deriving the variance preserving condition (Equation (12)) but not a lot is said afterwards. The authors basically refer to [1] to justify their method \"Furthermore, subsequent noising process using deterministic and stochastic noises can then be used to ensure the transition to the correct noisy manifold\". It is not clear to me that the results of [1] apply directly in the context of I2SB. Also, the current work is not self-contained as it relies heavily on the work of [1] and in particular the conjugate gradient method is never explained in the manuscript. I think that in the related work or in the theoretical section the authors should put more effort into giving details about what is [1] and how to adapt it to the current context (even though this is quite clear from reading the paper). In the experiments, the authors claim that that \"CDDB generally has higher speed and stability, possibly due to guaranteed convergence\". In the paper, no convergence results are provided. I don't see anything that would suggest that CDDB has guaranteed convergence. Overall, I think that the theoretical analysis is the weak point of the paper.\n\n* The authors \"exclude PiGDM for the baseline comparison in the deblurring problem, as directly leveraging the pseudo-inverse matrix may result in unfair boost in performances\". If this is the case then why not include PiGDM and use CDDB-deep which is also preconditioned with the pseudo-inverse? The fact that PiGDM performs (better?) than the proposed method because it makes a smart use of the pseudo-inverse does not seem like a good reason to remove it from the baseline.\n\n[1] Chung et al. (2023) -- Fast Diffusion Sampler for Inverse Problems by Geometric Decomposition\n\nQUESTIONS:\n* Another method that performs transfer tasks is Rectified Flow (albeit in a different context) [1] (or other transport based methods like [2]). Since these methods are not optimised to promote the fidelity to the term $Ax=y$ they should have a higher FID score but they might give better LPIPS, especially Rectified Flow which is deterministic. Can the authors comment on that?\n\n* I find the comment on the link with Schrodinger Bridge to be quite interesting but also quite confusing. In SB problems the coupling is not known beforehand. Hence, I don't really see how one could leverage a relation of the form $Ax=y$.\n\n* is it also possible to use some preconditioning with the pseudo-inverse in CDDB? (since in CDDB-deep \"we find that preconditioning with the pseudo inverse as in PiGDM improves performance\" it is a natural question.\n\n*Line 87: there is no footnote associated with the superscript 1.\n\n[1] Liu et al. (2023) -- Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow\n\n[2] Shi et al. (2023) -- Diffusion Schrödinger Bridge Matching\n\nLIMITATIONS:\nAdressed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the existing works about Direct Diffusion Bridges (DDB) with a unified scheme and limitations, and proposes a modified inference procedure that imposes data consistency without the need for fine-tuning, called data Consistent DDB (CDDB), as a new diffusion model-based inverse problem solvers.\n\nSTRENGTHS:\n+ The paper is well written and organized. Especially, the writing of the background section clearly describes and sorts out several important relevant works in a unified perspective.\n\n+ This paper solves an important problem to study the data consistency problem with diffusion models, which can be used to solve inverse problems generally.\n\n+ The experiments performed on natural images are comprehensive, and results show satisfying reconstruction quality, with comparison methods in both DIS and DDB methods.\n\nWEAKNESSES:\n- The paper studies a series of DDB methods which match the clean data/image distribution with measurements distribution, with the key contribution to add the data consistency term into the inference sampling steps. Actually, this kind of data consistency term has been introduced in quite a few previous works about diffusion model-based inverse problem solving, so called DIS/DDS methods in this paper. Although this paper claims the proposed approach is a generalization of the DDS method, after reading the whole manuscript, one essential question is still not very clear why DDB methods should outperform DIS/DDS methods, if they share a similar scheme from some perspective and both consider data consistency constraints during the sampling steps.\n\n- As mentioned above, the idea of updating the clean part for data consistency from DDIM formulation has been introduced in some previous works, such as DDS paper [5] and DDNM paper [39]. Although the paper claims that the proposed approach is a generalization of method in [5], the comparison with these two similar papers is missing in the experiments.\n\n- Besides, for the proposed DDB methods, it seems to require the paired data for training the distribution matching, different from the series of DIS methods which generally are trained without paired data. Although the paper discuss the relevance with supervised learning frameworks, the proposed method may need to be also compared with other supervised methods or conditional diffusion methods, which are missing in the current experiments and results. Besides, in the Discussion section, the paper claims the CDDB is flexible and does not have to pre-determine the number of forward passes or modify the training algorithm. This flexibility is also carried with all the DIS methods generally which are used to solve inverse problems based on trained diffusion models, but not the unique characteristics from DDB methods, which may need to be clarified clearly in the paper.\n\n- In the proposed CDDB (deep) method, it relies on estimating a pseudo-inverse to preserve data consistency in each sampling step. But for solving non-linear inverse problems, it rarely uses a pseudo-inverse formulation, which may not be straightforward to obtain pseudo-inverse for most nonlinear problems.\n\n- As mentioned in the Discussion, in order to show the advantage of DDB methods to match two distributions with data consistency, a great example is image translation problem, which is kind of surprising not included in the scope of this paper. But meanwhile, for image translation problem, since the forward matrix A is not explicit, is it still possible to use the proposed CDDB method to preserve the data consistency in the image translation problem?\n\nQUESTIONS:\nPlease see the weaknesses for specific questions. \n\nMinor comments:\n- In Table 2, the best results are not bolded correctly in the PSNR of “pool” images as mentioned in the caption?\n\nLIMITATIONS:\nThe authors discuss the limitations and societal impact of the proposed method at the end of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a unified view of several existing diffusion-model-based methods for solving the inverse problem (i.e., DPS and $\\Pi$GDM), and proposes a new approach for this problem called Direct Diffusion Bridge (DDB) which is derived from the DDPM ancestral sampling, and to some extent, links to the Image-to-Image Schrodinger bridge (I$^2$SB). The main idea behind DDB is using a trained neural network $G_{\\theta^*}$ to compute a denoised image $\\hat{x}\\_{0|t}$, which is an approximation of $x_0$, from $x_t$. The paper also proposes an improved version of DDB called Consistent DDB (CDDB), which according to the authors, is more beneficial for reconstruction (of $x_0$).\n\nSTRENGTHS:\n- The unified view between DPS and $\\Pi$GDM is interesting though I think it’s not hard to figure out as both methods share the same motivation.\n- Experimental results seem to support the proposed method there are improvements in both perception and signal-to-noise ratio compared to previous works. However, the fairness in settings of different methods should be clarified in the paper.\n\nWEAKNESSES:\n1) The presentation in the paper makes the motivation for the proposed method “Direct Diffusion Bridge” (DDB) unclear. From my view, neither the unified view between DPS and $\\Pi$GDM (Section 2.2) nor the I$^2$SB model (Section 3.1) motivates the design of the denoised image $\\hat{x}_{0|t}$ in the paper. In fact, it seems like in Section 3.1 the authors try to link the posterior of $x_t$ in the I$^2$SB model (Eq. 7) with the posterior of $x_s$ (s < t) in the DDPM (Eq. 10), which, I guess, leads to the name “Direct Diffusion Bridge”.\n\n2) I think the main contribution *in terms of technique* in this paper is a new way to compute $\\hat{x}\\_0$ (line 2 in Algorithm). Other parts of Algorithm 1 (lines 3 -> 7) were already proposed in previous works (e.g., DPS, DDS). However, the trained network $G\\_{\\theta^*}$, which plays an important part in the algorithm, is not described in detail in the paper. It is only mentioned briefly at lines 120 -> 121. Besides, no ablation study was conducted to validate the use of $G\\_{\\theta^*}$ with other methods for computing $\\hat{x}\\_0$ from $x\\_t$ (e.g., Tweedie’s formula).\n\n3) The authors don’t mention how $G\\_{\\theta^*}$ can handle $x\\_t$ with different time steps $t$ in the paper. Does it take both $x_t$ and $t$ as input? In fact, from my view, $G\\_{\\theta^*}$ is very similar to the network in the paper Consistency Models [1]. However, I don’t see the authors discuss Consistency Models in their paper. Since Consistency Models can be directly used for DIS (a lot of experiments on inverse sampling can be found in [1]), I am keen to see the empirical comparison between this method and Consistency Models.\n\n4) Please correct me if I am wrong but I cannot find the NFE of the proposed method in the main paper or the supplementary material. Therefore, it is hard to fairly compare the performance of the proposed method with other baselines.\n\n[1] Consistency Models, Song et al., ICML 2023\n\nQUESTIONS:\n1) What are the architecture and configurations of the residual network $G\\_{\\theta^*}$? I couldn’t find any detail about $G\\_{\\theta^*}$ in the main paper or the supplementary material. Does it depend on time $t$ besides $x_t$\n\n2) What are the technical differences between the method in this paper and the use of Consistency Models for the inverse sampling problem?\n\n3) In Algorithm 1, is $x\\_1$ $y$ or Gaussian random noise? Since I$^2$SB can take the corrupted image $y$ as input, do the authors think their method can do so?\n\n4) In Eq 15, $\\mathbb{E}[x_t]$ should be $\\mathbb{E}[x_0|x_t]$.\n\n5) The main difference between Algorithm 1 and Algorithm 2 is the computation of $g$. In Alg. 1, $g$ is computed as $\\nabla_{\\hat{x}\\_{0|t}} \\log p(y|\\hat{x}\\_{0|t})$ while in Alg. 2, $g$ is computed as $\\nabla_{x_t} \\log p(y|\\hat{x}\\_{0|t})$. In my view, Alg. 2 is *more theoretically correct* than Alg. 1 rather than just being more beneficial for reconstruction as mentioned in the paper in lines (181 -> 183). Thus, I hope the authors could make this point clearer in their paper rather than just mentioning CDDB as an alternative for DDB when the inverse problem is non-linear. Even when the inverse problem is linear, we still need the term $\\frac{\\partial \\hat{x}\\_{0|t}}{\\partial x_t}$ which is the Jacobian in Eqs. 5, 6.\n\n6) At line 181, the authors should be clear about the “U-Net Jacobians”. Is the U-Net here the noise network $\\epsilon$ of the diffusion model or the network $G\\_{\\theta}$? I guess it should be $G\\_{\\theta}$ since we are considering $\\frac{\\partial \\hat{x}\\_{0|t}}{\\partial x_{t}}$. The authors give me the impression that they are trying to hide details about $G_{\\theta}$, which in fact plays a critical role in the paper.\n\nLIMITATIONS:\nPlease refer to my questions above", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-08-01"}