Consistency Model is an Effective Posterior Sample Approximation for Diffusion Inverse Solvers

## Abstract

Abstract Diffusion Inverse Solvers (DIS) are designed to sample from the conditional distribution p θ ⁢ ( X 0 | y ) subscript 𝑝 𝜃 conditional subscript 𝑋 0 𝑦 p_{\theta}(X_{0}|y) italic_p start_POSTSUBSCRIPT italic_θ end_POSTSUBSCRIPT ( italic_X start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT | italic_y ) , with a predefined diffusion model p θ ⁢ ( X 0 ) subscript 𝑝 𝜃 subscript 𝑋 0 p_{\theta}(X_{0}) italic_p start_POSTSUBSCRIPT italic_θ end_POSTSUBSCRIPT ( italic_X start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT ) , an operator f ⁢ ( ⋅ ) 𝑓 ⋅ f(\cdot) italic_f ( ⋅ ) , and a measurement y = f ⁢ ( x 0 ′ ) 𝑦 𝑓 subscript superscript 𝑥 ′ 0 y=f(x^{\prime}_{0}) italic_y = italic_f ( italic_x start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT ) derived from an unknown image x 0 ′ subscript superscript 𝑥 ′ 0 x^{\prime}_{0} italic_x start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT . Existing DIS estimate the conditional score function by evaluating f ⁢ ( ⋅ ) 𝑓 ⋅ f(\cdot) italic_f ( ⋅ ) with an approximated posterior sample drawn from p θ ⁢ ( X 0 | X t ) subscript 𝑝 𝜃 conditional subscript 𝑋 0 subscript 𝑋 𝑡 p_{\theta}(X_{0}|X_{t}) italic_p start_POSTSUBSCRIPT italic_θ end_POSTSUBSCRIPT ( italic_X start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT | italic_X start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) . However, most prior approximations rely on the posterior means, which may not lie in the support of the image distribution, thereby potentially diverge from the appearance of genuine images. Such out-of-support samples may significantly degrade the performance of the operator f ⁢ ( ⋅ ) 𝑓 ⋅ f(\cdot) italic_f ( ⋅ ) , particularly when it is a neural network. In this paper, we introduces a novel approach for posterior approximation that guarantees to generate valid samples within the support of the image distribution, and also enhances the compatibility with neural network-based operators f ⁢ ( ⋅ ) 𝑓 ⋅ f(\cdot) italic_f ( ⋅ ) . We first demonstrate that the solution of the Probability Flow Ordinary Differential Equation (PF-ODE) with an initial value x t subscript 𝑥 𝑡 x_{t} italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT yields an effective posterior sample p θ ⁢ ( X 0 | X t = x t ) subscript 𝑝 𝜃 conditional subscript 𝑋 0 subscript 𝑋 𝑡 subscript 𝑥 𝑡 p_{\theta}(X_{0}|X_{t}=x_{t}) italic_p start_POSTSUBSCRIPT italic_θ end_POSTSUBSCRIPT ( italic_X start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT | italic_X start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT = italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) . Based on this observation, we adopt the Consistency Model (CM), which is distilled from PF-ODE, for posterior sampling. Furthermore, we design a novel family of DIS using only CM. Through extensive experiments, we show that our proposed method for posterior sample approximation substantially enhance the effectiveness of DIS for neural network operators f ⁢ ( ⋅ ) 𝑓 ⋅ f(\cdot) italic_f ( ⋅ ) (e.g., in semantic segmentation). Additionally, our experiments demonstrate the effectiveness of the new CM-based inversion techniques. The source code is provided in the supplementary material.

## 1 Introduction

Diffusion inverse solvers (DIS) are a family of algorithms that solve the inverse problem using diffusion prior [Li et al., , Moser et al., ]. More specifically, given an operator $f(.)$, a measurement $y=f(x^{\prime}_{0})$ from some unknown image $x^{\prime}_{0}$, and a diffusion model $p_{\theta}(X_{0})$, DIS aim to sample from the conditional distribution $x_{0}\sim p_{\theta}(X_{0}|y)$. For example, when $f(.)$ is a down-sampling operator, DIS is a perceptual super-resolution algorithm [Menon et al., ]. However, we cannot sample from $p_{\theta}(X_{0}|y)$ directly as it is intractable. To tackle this challenge, previous works adopt a variety of techniques such as linear projection [Wang et al., , Kawar et al., , Chung et al., , Lugmayr et al., , Song et al., , Pokle et al., , Cardoso et al., ], variational inference [Feng et al., , Mardani et al., , Janati et al., ], Bayesian filter [Dou and Song, ], sequential Monte Carlo [Wu et al., , Phillips et al., ], proximal gradient [Xu and Chi, ] and conditional score estimation [Chung et al., , Yu et al., , Zhu et al., , He et al., , Song et al., , Boys et al., , Rout et al., , ] for approximate or exact sampling.

Among those DIS techniques, the conditional score estimation methods [Chung et al., , Song et al., ] are most widely adopted, as they are suitable for general non-linear, noisy operator $f(.)$ and quite efficient in practice. During the inverse diffusion process from $X_{T}$ to $X_{0}$, they estimate the conditional score $\nabla_{x}\log p_{\theta}(X_{t}|y)$ by evaluating operator $f(.)$ with posterior sample from $p_{\theta}(X_{0}|X_{t})$. As the posterior $p_{\theta}(X_{0}|X_{t})$ is generally intractable, various approximations to posterior sample are proposed based on posterior mean: they either directly adopt posterior mean [Chung et al., , Yu et al., , Zhu et al., , He et al., ] or construct an uni-modal distribution centered at posterior mean [Song et al., , Boys et al., , Rout et al., , ].

However, those posterior-mean based approximate posterior samples are far from real images, as it is well-known that the mean of noisy images may not lie in the support of the image distribution[Ledig et al., , Blau and Michaeli, ]. Although those out-of-distribution approximations are shown to be successful for simple $f(.)$ such as down-sampling and motion blurring, they might fail for more complex $f(.)$, especially when $f(.)$ are neural networks such as segmentation or classification, which are sensitive to out-of-distribution inputs.

In this paper, we propose a novel approach to approximate posterior sample for DIS. Our approximations are guarantee to be valid images and generally perform well for neural network $f(.)$. More specifically, we first show that given initial condition $X_{t}=x_{t}$, the solution of probability flow ordinary differential equation (PF-ODE) [Song et al., ] is a valid posterior sample of true posterior $p_{\theta}(X_{0}|X_{t}=x_{t})$. Inspired by this, we propose to use consistency model (CM), a distillation of PF-ODE as a posterior sample approximation for DIS. Furthermore, we propose a new family of DIS by iteratively inverting CM in a generative adversarial network (GAN) inversion fashion [Creswell and Bharath, , Menon et al., ]. Empirically, by using CM as posterior sample approximation, we improve the DIS’s performance when operators are neural networks, such as semantic segmentation and image captioning. Furthermore,
our experiments demonstrate the CM inversion also performs well for both neural network and non-neural network operators.

## 2 Preliminaries

### 2.1 Diffusion Model

Diffusion model is a type of generative model with $T$ steps Gaussian Markovian chain in continuous space [Sohl-Dickstein et al., ]. Two widely adopted diffusion models are variance preserving (VP) and variance exploding (VE) diffusion. We follow the formulation of VE diffusion [Song et al., ], and refer VP diffusion to Ho et al. [], Kingma et al. []. We denote the source image as $X_{0}$, and the forward process of VE diffusion is a Markov chain:

$$ $\displaystyle q(X_{T},...,X_{1}|X_{0})=\prod_{t=1}^{T}q(X_{t}|X_{t-1})\textrm{ , where }q(X_{t}|X_{t-1})=\mathcal{N}(X_{t-1},(\sigma_{t}^{2}-\sigma_{t-1}^{2} )I),$ (1) $$

where $\sigma^{2}_{t}$ are hyper-parameters (typically called variance schedule). The reverse diffusion process is also a Markov chain, with transition kernel $p(X_{t-1}|X_{t})$ depending on the score function $\nabla_{X_{t}}\log p(X_{t})$. To learn such a diffusion model, one can use a neural network $s_{\theta}(t,X_{t})$
(parametrized by $\theta$)
to match the score function $\nabla_{X_{t}}\log p(X_{t})$, and the resulting reverse diffusion process is given by

$$ $\displaystyle p_{\theta}(X_{0},...,X_{T})=p(X_{T})\prod_{t=1}^{T}p_{\theta}(X_ {t-1}|X_{t}),$ $\displaystyle\textrm{where }p_{\theta}(X_{t-1}|X_{t})=\mathcal{N}(X_{t}+( \sigma_{t}^{2}-\sigma_{t-1}^{2})s_{\theta}(t,X_{t}),(\sigma_{t}^{2}-\sigma_{t- 1}^{2})I).$ (2) $$

Song et al. [] show that the reverse diffusion can be seen as a discretization of reverse stochastic differential equation (SDE) [Anderson, ]. Further, there exists a probability flow ordinary differential equation (PF-ODE) that has the same marginal distribution $p_{\theta}(X_{t})$ as the reverse SDE:

$$ $\displaystyle\textrm{reverse SDE: }dX_{t}=-\frac{d\sigma^{2}_{t}}{d_{t}}s_{ \theta}(t,X_{t})dt+\sqrt{\frac{d\sigma^{2}_{t}}{d_{t}}}dB_{t}\quad\overset{ \textrm{same }p_{\theta}(X_{t})}{\Longleftrightarrow}$ $\displaystyle\textrm{PF-ODE: }dX_{t}=-\frac{1}{2}\frac{d\sigma^{2}_{t}}{dt}s_{ \theta}(t,X_{t})dt,$ (3) $$

where $B_{t}$ is the standard Brownian motion. Therefore, solving either of them is equivalent to sampling from the reverse diffusion process. Another useful result is Tweedie’s formula [Efron, ], which provides an efficient estimation to the mean of posterior $p_{\theta}(X_{0}|X_{t})$:

$$ $\displaystyle\mathbb{E}[X_{0}|X_{t}]=X_{t}+\sigma^{2}_{t}s_{\theta}(t,X_{t}).$ (4) $$

### 2.2 Diffusion Inverse Solvers with Conditional Score Estimation

Given an operator $f(.)$, a target measurement $y=f(x^{\prime}_{0})$ from an unknown $x^{\prime}_{0}$ and a diffusion model $p_{\theta}(X_{0})$, the diffusion inverse solvers (DIS) attempt to sample from the conditional distribution $p_{\theta}(X_{0}|y)$. In this paper, we focus on DIS with conditional score estimation [Chung et al., , Song et al., ]. More specifically, this paradigm of DIS attempts to estimate the conditional score $\nabla_{X_{t}}\log p_{\theta}(X_{t}|y)$. With this conditional score at hand, sampling from $p_{\theta}(X_{0}|y)$ is as easy as solving the reverse SDE or PF-ODE in Eq. 3 with $s_{\theta}(t,X_{t})$ replaced by the conditional score.

More specifically, Chung et al. [], Song et al. [] propose to expand the conditional score into unconditional score and a term that is related to a distance $d(f(x_{0|t}),y)$, where $x_{0|t}$ is the sample from posterior $p_{\theta}(X_{0}|X_{t})$ and $d(.,.)$ is a distance:

$$ $\displaystyle\nabla_{X_{t}}\log p_{\theta}(X_{t}|y)=\nabla_{X_{t}}\log p_{ \theta}(y|X_{t})+\nabla_{X_{t}}\log p_{\theta}(X_{t}),$ $\displaystyle\nabla_{X_{t}}\log p_{\theta}(y|X_{t})=\nabla_{X_{t}}\log\mathbb{ E}_{p_{\theta}(X_{0}|X_{t})}[p_{\theta}(y|X_{0})]\approx\nabla_{X_{t}}\log\sum _{x_{0|t}^{(i)}\sim p_{\theta}(X_{0}|X_{t})}^{i=1,...,K}p_{\theta}(y|X_{0}=x_{ 0|t}^{(i)}),$ $\displaystyle p_{\theta}(y|X_{0}=x_{0|t}^{(i)})\propto\exp{-d(f(x_{0|t}^{(i)}) ,y)}.$ (5) $$

Under this formulation, an important issue is how to effectively draw differentiable samples from posterior $p_{\theta}(X_{0}|X_{t})$. Obviously, direct ancestral sampling from reverse diffusion is computationally expensive. Chung et al. [] propose to use the posterior mean computed by Tweedie’s formula in Eq. 4 as
the posterior sample. Song et al. [, ] propose to model the posterior as a Gaussian distribution with mean being the posterior mean as mean and and the covariance chosen as a hyper-parameter. Rout et al. [], Boys et al. [] improve Song et al. [, ] by using the posterior covariance computed by second order Tweedie’s formula as the covariance of Gaussian. There are several other approaches that follow conditional score estimation paradigm [Yu et al., , Chung et al., , Song et al., , He et al., , Rout et al., , Meng and Kabashima, , Dou and Song, , Chung et al., , Song et al., , He et al., ] and rely on those approximations.

## 3 Consistency Model is an Effective Posterior Sample Approximation for DIS

### 3.1 Previous Approximations are Out-of-Distribution

Most previous approximations to $p_{\theta}(X_{0}|X_{t})$ either directly use the posterior mean or construct a uni-modal distribution entered around the posterior mean. However, the posterior mean $\mathbb{E}[X_{0}|X_{t}]$ is a mean-square error (MSE) minimizer for images perturbed by Gaussian noise, which does not necessarily correspond to a valid image [Ledig et al., , Blau and Michaeli, ]. In other words, the posterior mean may not lie in the support of natural image distribution, leading to its density in both marginal and posterior distributions approaching zero.:

$$ $\displaystyle p_{\theta}(X_{0}=\mathbb{E}[X_{0}|X_{t}=x_{t}])\approx 0,p_{ \theta}(X_{0}=\mathbb{E}[X_{0}|X_{t}=x_{t}]|X_{t}=x_{t})\approx 0$ (6) $$

Figure: Figure 2: Different approximations of posterior sample, and their output after a segmentation $f(.)$.
Refer to caption: x2.png

When the operator
is a neural network, such out-of-distribution approximations can
significantly degrade the sample quality. A visual example is shown in Fig. 2. The approximation using posterior mean, and the approximation using Gaussian centered at the posterior mean are not valid image samples. Consequently, when these approximations are processed through a semantic segmentation operator $f(.)$,
the outputs are nonsensical.

### 3.2 PF-ODE provides an Effective Posterior Sample Approximation

It is known that the PF-ODE and reverse SDE in Eq. 3 have the same marginal distribution, which is the image distribution $p_{\theta}(X_{0})$ if the score is learned perfectly [Song et al., ]. Denote the solution of PF-ODE given initial condition $X_{t}=x_{t}$ as $\Phi_{0}(x_{t})$, then this solution is in support of natural image distribution, and the density of the solution is non-zero:

$$ $\displaystyle p_{\theta}(X_{0}=\Phi_{0}(x_{t}))>0.$ (7) $$

Back to the previous example in Fig. 2, when using PF-ODE as posterior sample approximation, the semantic segmentation operator $f(.)$ produces reasonable result. However, knowing $p_{\theta}(X_{0}=\Phi_{0}(x_{t}))>0$ is not enough. As we are seeking a posterior sample approximation, we need to ensure the solution $\Phi_{0}(x_{t})$ has positive density on the posterior distribution, i.e., $p_{\theta}(X_{0}=\Phi_{0}(x_{t})|X_{t}=x_{t})>0$.

To the best of our knowledge, the relationship between PF-ODE’s solution $\Phi_{0}(x_{t})$ given initial value $X_{t}=x_{t}$, and the posterior $p_{\theta}(X_{0}|X_{t}=x_{t})$ is not well understood. In this section, we show that the solution of PF-ODE has non-zero density in true posterior for any time $t$, i.e.,

$$ $\displaystyle\forall x_{t},p_{\theta}(X_{0}=\Phi_{0}(x_{t})|X_{t}=x_{t})>0.$ (8) $$

###### Assumption 3.1 .

We assume the following conditions hold:

- •
The distribution $p_{\theta}(X_{0})$ can be approximated by a $d$-dimension Gaussian Mixture Model (GMM) composed of $N$ Gaussians with same small diagonal covariance $\sigma^{2}I$ and mean $\mu^{i}$:
$\displaystyle p_{\theta}(X_{0})\approx\frac{1}{N}\sum_{i=1}^{N}\mathcal{N}(X_{
0}|\mu^{i},\sigma^{2}I).$
(9)
- •
The solution $\Phi_{0}(x_{t})$ and initial value $x_{t}$ are bounded, i.e., $||\Phi_{0}(x_{t})||<c,||x_{t}||\leq c$.
- •
As PF-ODE is margin preserving and the dimension of data is high, the solution of PF-ODE is close to the surface of a sphere centered at $\mu^{i}$ with radius $\sqrt{d}\sigma$ [Vershynin, ], i.e., $\exists k,\textrm{s.t.}||\Phi_{0}(x_{t})-\mu^{k}||^{2}\leq\sigma^{2}+d\sigma^{2}$ with probability $1-p$ for small $p$.

###### Proposition 3.2 .

The solution of PF-ODE has a positive likelihood in true posterior with high probability, i.e.,

$$ $\displaystyle p_{\theta}(X_{0}=\Phi(x_{t})|X_{t}=x_{t})\geq\frac{1}{N}\frac{1} {\sqrt{(4\pi\sigma^{2})^{d}}}\exp{(-\frac{2c^{2}}{\sigma_{t}^{2}}-\frac{d+1}{2 })}\textrm{, with probability }1-p.$ (10) $$

Despite the solution of PF-ODE has a non-zero density in true posterior with high probability, we can not tell whether the solution falls into the highest density mode of true posterior. However, when $\sigma_{t}^{2}$ is small, we can show that the solution of PF-ODE is approximately the highest density mode of true posterior.

###### Lemma 3.3 .

The PF-ODE can be written as:

$$ $\displaystyle\frac{dX_{t}}{dt}=\underbrace{\sum_{i=1}^{N}\frac{w_{i}}{2}\frac{ d\sigma_{t}^{2}}{dt}\frac{(X_{t}-\mu^{i})}{\sigma^{2}+\sigma_{t}^{2}}}_{ \textrm{velocity field }v_{t}},w_{i}=(\exp{-\frac{||X_{t}-\mu^{i}||^{2}}{2( \sigma^{2}+\sigma_{t}^{2})}})/(\sum_{j=1}^{N}\exp{-\frac{||X_{t}-\mu^{j}||^{2} }{2(\sigma^{2}+\sigma_{t}^{2})}}).$ (11) $$

Figure: Figure 3: The PF-ODE’s velocity field of a five GMM example.

We can consider the velocity field $v_{t}$ of PF-ODE in Eq. 11. The velocity field is a sum of vectors pointing to centers $\mu^{i}$ weighted by soft-max function $w_{i}$. When $t$ is small, $w_{i}$ becomes a "hard"-max which selects the closest center $\mu^{*}$ to initial point $x_{t}$, and the velocity field always points to $\mu^{*}$. On the other hand, this closest center $\mu^{*}$ is also the highest density mode in true posterior (See Appendix. A). With this coincidence, the solution of PF-ODE with $X_{t}=x_{t}$ is approximately the mode with highest density in true posterior.

When $\sigma_{t}^{2}$ is not small enough, we can still consider other conditions when PF-ODE will converge to the closest center $\mu^{*}$. Consider the five GMM example in Figure. 3 with initial point $X_{t}=x_{t}$. An obvious intuition is that the normal plane of velocity $v_{t}$ divide the space into two parts. In one part, the velocity has a negative projection on center’s direction. In the other part, the velocity has a positive projection. Then, $X_{t}$ will move away from the centers with negative projection, to the centers with positive projection. Among the centers with positive projection, when the closest center $\mu^{*}$ also has the smallest angle with velocity, it is very likely that the PF-ODE eventually converges into the closest center.

###### Assumption 3.1 .

###### Proposition 3.2 .

###### Lemma 3.3 .

### 3.3 A Toy Example

To better understand the results above, we provide a toy example in $\mathbb{R}^{2}$. As shown in Fig. 4, the source distribution $p(X_{0})$ is a five Gaussian mixture model (GMM). Each Gaussian is diagonal with standard deviation $\sigma_{0}=0.1$. The centers of Gaussians are $(-1,-1),(-1,1),(1,1),(1,-1),(0,0)$. We adopt VE diffusion with $\sigma_{T}=4$, $T=100$ and $\sigma$ schedule in Karras et al. []. The score function $\nabla_{X_{t}}\log p(X_{t})$ is computed analytically.

Figure: Figure 4: A toy example with five GMM.
Refer to caption: extracted/5636637/fig_toy1.png

As shown in Fig. 4.upper, starting from $x_{t}=(1,0.4)$, the posterior $p(X_{0}|X_{t}=x_{t})$ has most density on two right-hand side Gaussians. However, the approximation with posterior mean is close to $(1,0.4)$, where true posterior has almost no density. Furthermore, the approximation with posterior mean and covariance also concentrates on a region where true posterior has almost no density. However, the approximation with PF-ODE falls into a high density region of true posterior.

We also visualize the highest density mode’s decision boundary of true posterior and PF-ODE starting at different $\sigma^{2}_{t}$. As shown in Fig. 4.lower, the decision boundary of true posterior is always a Voronoi cell centered at $\mu^{i}$. When $\sigma_{t}^{2}$ is small, the solution of PF-ODE is similar to the true posterior. As $\sigma_{t}^{2}$ increase, the solution of PF-ODE becomes less similar to true posterior. However, at that time, the density scatters more evenly in the true posterior. The solution of PF-ODE still has non-zero density.

### 3.4 Implementation of PF-ODE Approximation with Consistency Model

Directly solving PF-ODE is also intractable for DIS. Fortunately, PF-ODE can be distilled by Consistency Model (CM) [Song et al., ]. More specifically, CM trains a one-step neural function $g_{\theta}(t,x_{t})$ to approximate the solution of PF-ODE $\Phi_{0}(x_{t})$. Its gradient is cheap to evaluate. Therefore, we can directly replace the $x_{0|t}\sim p_{\theta}(X_{0}|X_{t}=x_{t})$ step in Eq. 5 by $x_{0|t}=g_{\theta}(t,x_{t})$.

In practice, we find CM often over-fits the operator $f(.)$. More specifically, the CM approximated sample $x_{0|t}$ is numerically close to $y$ after passing the operator $f(.)$. However, by inspecting $x_{0|t}$ visually, one often concludes that it is not aligned with $y$. (See an example in Fig. 6).

In fact, the resulting overfitted sample is an adversial example [Szegedy et al., ], which is aligned with label $y$ according to neural network but not aligned with $y$ according to human eye. To make $f(.)$ robust, we propose to add a small Gaussian noise to the output of CM as literature in adverisal robustness [Li et al., ]

$$ $\displaystyle x_{0|t}=g_{\theta}(t,x_{t})+\mathcal{N}(0,\tau^{2}).$ (12) $$

### 3.5 Consistency Model Inversion

We have shown that CM can be used as posterior sample approximator for DIS. As CM can be seen as a conditional GAN conditioned on timestep $t$ and state $x_{t}$, a natural question to ask is: can CM be inverted in a GAN inversion fashion [Creswell and Bharath, , Menon et al., ]?

We first review GAN inversion briefly. GAN inversion optimizes the noise $z$ by penalizing the distance between generated image and the target for $K$ steps. And step by step, the optimized $z$ will generate an image that satisfies the constraint. Denote $h_{\theta}(.)$ as GAN, we have

$$ $\displaystyle z^{0}\sim\mathcal{N}(0,I),x^{i}=h_{\theta}(z^{i}),z^{i+1}=z^{i}- \zeta\frac{d}{dz^{i}}d(f(x^{i}),y),i=0,...,K-1.$ (13) $$

Recall that CM first initializes an $x_{T}$ from Gaussian distribution, then transforms it into target image $x_{0|t}$ by a neural network. Then iteratively CM adds a small noise back to $x_{0|t}$ and denoises again. For each iteration, CM is the same as GAN. To CM, for each iteration, we can invert CM as GAN inversion. In other words, we can stack several GAN inversions together, according to the iterative CM sampling algorithm in Algorithm 7, to obtain our final CM inversion algorithm in Algorithm 8. For each iteration, a GAN inversion is performed. Similarly, we also find that adding a small Gaussian noise is beneficial to neural-network $f(.)$.

## 4 Experiments

### 4.1 Experiment Setup

Base Diffusion Models
For diffusion model, we use a pretrained VE diffusion - EDM [Karras et al., ] provided by Song et al. []. For EDM-related methods, we adopt ancestral sampler with 1000 Euler steps. For CM [Song et al., ], we employ the official pre-trained model by Song et al. []. The details are shown in Appendix. B.

Operators We evaluate all the methods with four neural network operators: semantic segmentation, room layout estimation, image captioning and image classification. For layout estimation, we adopt the neural network by Lin et al. []. For semantic segmentation, we use the neural network by Zhou et al. []. For image captioning, we employ BLIP [Li et al., ]. For image classification, we use ResNet [He et al., ]. In addition, we also evaluate a simple non-neural network operator: down-sampling (x4).

Datasets & Metrics
Following Song et al. [] and Chung et al. [], we use the first 1000 image from LSUN Bedroom and LSUN Cat dataset [Yu et al., ] as test set. All images are resized into $256^{2}$. To evaluate sample quality, we use Fréchet Inception Distance [Heusel et al., ] and Kernel Inception Distance (KID) [Binkowski et al., ]. To evaluate consistency with the constraint, we use mIOU for segmentation and layout, CLIP score for captioning, and Accuracy for classification. For neural network $f(.)$, we use different models for DIS and testing (See Appendix. B). For down-sampling, we use image restoration metrics such as LPIPS [Zhang et al., ] and peak signal-to-noise ratio (PSNR).

Previous State-of-the-Art DIS
We compare our approach with previously published DIS that are able to solve neural network operator $f(.)$. For methods that directly use posterior mean as posterior sample, we include DPS [Chung et al., ], FreeDOM Yu et al. [] and MPGD [He et al., ]. For methods that construct an approximated posterior distribution with posterior mean as mode, we include LGD [Song et al., ] and STSL [Rout et al., ] (See Tab. 4.1). We implement all those methods with EDM and Euler ancestral sampler (See details in Appendix. B). We acknowledge that there are other very competitive works designed for latent diffusion [Chung et al., , Song et al., , He et al., , Rout et al., ] or linear operator $f(.)$ [Meng and Kabashima, , Dou and Song, , Chung et al., , Song et al., , Boys et al., , He et al., ]. However, for now we focus on pixel domain diffusion with neural network $f(.)$, and have not included them for comparison.

**Table 1: Results on neural network operators, i.e., layout estimation, segmentation, caption and classification. Bold: best in diffusion-based DIS. Underline: second best in diffusion-based DIS.**
|  | LSUN Bedroom | LSUN Cat |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | Segmentation | Layout | Caption | Classification |  |  |  |  |  |  |  |  |
|  | mIOU | FID | KID | mIOU | FID | KID | CLIP | FID | KID | Acc | FID | KID |
| Diffusion Model based |  |  |  |  |  |  |  |  |  |  |  |  |
| EDM (Base) | 0.17 | 6.35 | 1.4e-3 | 0.38 | 6.35 | 1.4e-3 | 21.27 | 6.35 | 1.4e-3 | 0.14 | 10.26 | 3.1e-3 |
| DPS | 0.27 | 22.84 | 1.0e-2 | 0.54 | 7.59 | 1.9e-3 | 22.57 | 9.49 | 2.6-e3 | 0.79 | 15.73 | 7.1e-3 |
| FreeDOM | 0.27 | 21.90 | 9.1-e3 | 0.46 | 15.27 | 8.1e-3 | 22.61 | 28.30 | 1.8e-2 | 0.84 | 32.32 | 1.7-e2 |
| MPGD | 0.24 | 82.66 | 6.9e-2 | 0.73 | 15.38 | 8.5-e3 | 21.49 | 21.14 | 1.3e-2 | 0.37 | 15.40 | 7.1e-3 |
| LGD | 0.22 | 35.69 | 2.4e-2 | 0.70 | 8.07 | 2.3e-3 | 22.58 | 8.38 | 2.6-e3 | 0.64 | 13.35 | 4.5e-3 |
| STSL | 0.27 | 19.48 | 7.4e-3 | 0.52 | 7.74 | 2.2e-3 | 22.39 | 9.70 | 2.8e-3 | 0.78 | 15.74 | 6.9e-3 |
| Proposed I | 0.34 | 18.06 | 8.2e-3 | 0.78 | 7.50 | 2.2e-3 | 22.63 | 8.16 | 2.5-e3 | 0.90 | 13.45 | 3.6e-3 |
| Consistency Model based |  |  |  |  |  |  |  |  |  |  |  |  |
| CM (Base) | 0.18 | 20.45 | 1.0e-2 | 0.37 | 20.45 | 1.0e-2 | 21.40 | 20.45 | 1.0e-2 | 0.12 | 27.15 | 1.3e-2 |
| Proposed II | 0.32 | 32.60 | 2.2e-2 | 0.82 | 15.43 | 8.1e-3 | 22.56 | 14.86 | 6.7e-3 | 0.92 | 27.35 | 1.4e-2 |

### 4.2 Main Results

Results on Neural Network Operators
We test our proposed approaches on four neural network operators: segmentation, layout estimation, caption and classification. As shown in Tab. 4.1, Fig. 1, Fig. 5, Fig. 10 and Fig. 11, both quantitatively and visually, our Proposed I (Sec. 3.4) has significant improvement on both consistency (e.g., mIOU) and sample quality (e.g., FID) over the baseline DPS [Chung et al., ]. The advantage of our approximation over other posterior mean based approximations is clearly demonstrated. This is because neural network $f(.)$ are sensitive to out-of-distribution input. On the other hand, our Proposed II (Sec. 3.5) is also quite effective compared with unconditional CM.

Figure: Figure 5: Visual results on neural network operators such as segmentation, caption and classification.
Refer to caption: x4.png

Results on Non-neural Network Operators
In additional to neural network operators, we also verify that our approaches work well for non-neural network operators such as down-sampling. Results are summarized in Tab. 6 and Fig. 6.lower. Our Proposed I is only comparable to DPS for simple operators. This is because non-neural network $f(.)$ are not that sensitive to out-of-distribution approximations. Besides, our Proposed II also works well for linear operator.

### 4.3 Ablation Study

We evaluate the effect of using CM for posterior approximation in Proposed I (Sec. 3.4) and adding randomness to CM in Proposed I & II (Sec. 3.4, Sec. 3.5) in Tab. 4. For Proposed I, we show that using CM to replace posterior mean reduces the distance with measurement $y$ in mIOU and improves the sample quality in FID. Similarly, for both Proposed I & II, adding randomness improves mIOU and reduces FID.

Table: Table 4: Ablation study of Proposed I and Proposed II with Bedroom segmentation. Bold: Method with best performance.

|  | CM | Rand | mIOU | FID |
| --- | --- | --- | --- | --- |
| Proposed I | ✗ | ✗ | 0.27 | 22.84 |
| ✓ | ✗ | 0.31 | 19.29 |  |
| ✓ | ✓ | 0.34 | 18.06 |  |
| Proposed II | - | ✗ | 0.29 | 32.93 |
| - | ✓ | 0.32 | 32.60 |  |

Table: Table 5: Ablation study of randomness and data augmentation with segmentation.

| Rand | mIOU |  |  |
| --- | --- | --- | --- |
| Model A | Model A + DA | Model B |  |
| ✗ | 0.57 | 0.43 | 0.31 |
| ✓ | 0.51 | 0.54 | 0.34 |

**Figure 6: upper. An example of over-fitting an operator $f(.)$. lower. Visual results on linear operators such as down-sampling.**
|  | Bedroom Down-sampling (x4) |  |  |  |
| --- | --- | --- | --- | --- |
|  | KID | FID | LPIPS | PSNR |
| Diffusion Model based |  |  |  |  |
| EDM (Base) | 1.4e-3 | 6.35 | 0.73 | 9.10 |
| DPS | 1.4e-3 | 4.82 | 0.11 | 26.69 |
| FreeDOM | 1.1e-3 | 4.84 | 0.11 | 26.70 |
| MPGD | 2.0e-3 | 43.36 | 0.38 | 22.84 |
| LGD | 2.8e-3 | 7.10 | 0.14 | 26.51 |
| STSL | 1.2e-3 | 4.80 | 0.12 | 26.65 |
| Proposed I | 1.9e-3 | 5.67 | 0.12 | 26.91 |
| Consistency Model based |  |  |  |  |
| CM (Base) | 1.0e-2 | 20.45 | 0.73 | 9.49 |
| Proposed II | 2.9e-3 | 6.66 | 0.14 | 26.45 |

We assume that CM benefits from randomness as it avoids overfitting the operator $f(.)$, or it makes $f(.)$ robust to adversial examples [Li et al., ]. To verify this assumption, we use Model A for $f(.)$ during DIS. During testing $f(.)$, we compare the results of Model A, Model A w/ data augmentation (DA), and a separate Model B. In Tab. 5, we show that when tested with Model A, the DIS w/o randomness outperform DIS w/ randomness. While on Model A w/ DA and Model B, the DIS w/ randomness outperforms DIS w/o randomness. This indicates that DIS w/o randomness overfits Model A. An example of such overfitting is presented in Fig. 6.upper.

### 4.4 Complexity

In Tab. 3, we show that Proposed I is $30\%$ slower than DPS while Proposed II is $50\%$ faster. Both of our approaches cost $1$ GB more VRAM.

## 5 Related Work

Diffusion Inverse Solvers An important branch of DIS focus on linear operator $f(.)$ with projection or pseudo-inverse [Wang et al., , Kawar et al., , Chung et al., , Lugmayr et al., , Song et al., , Dou and Song, , Pokle et al., , Cardoso et al., ]. For general, non-linear $f(.)$, various approaches are proposed such as Monte carlo [Wu et al., , Phillips et al., ], proximal gradient [Xu and Chi, ] and variational inference [Feng et al., , Mardani et al., , Janati et al., ]. Among those paradigms, the conditional score estimation methods are mostly adopted as they are scalable to practically large images with reasonable run-time [Chung et al., , Yu et al., , Zhu et al., , He et al., , Song et al., , Boys et al., , Rout et al., , ]. Following this paradigm, we propose to approximate posterior sample with PF-ODE, which improves the results for neural network $f(.)$.

GAN Inverse Solvers Similar to diffusion inversion, given the measurement and operator, GAN inversion finds an image $x$ on the prior of GAN by minimizing the distance of generated image and a target [Creswell and Bharath, ]. It is widely adopted for image editing and image restoration [Menon et al., , Daras et al., ]. We propose to treat CM as several stacked GANs and invert CM as a series of GAN inversions.

## 6 Discussion & Conclusion

One major limitation of our work is that all the experiments are conducted with $256^{2}$ images and diffusion models in pixel domain. Recent works already show the potential of CM on $512^{2}$ for latent diffusion models [Luo et al., , Chung et al., , Song et al., , He et al., , Rout et al., ].
It would be interesting to see how our approach works empirically for large images using latent diffusion models.

To conclude, we show that the solution of PF-ODE is an effective posterior sample. Built upon this, we propose to use CM as a high-quality approximation to posterior sample. Further, we propose a new family of DIS using only CM. Experimental results show that our proposed approaches perform well for neural network operators.

## Appendix A Proof of Main Results

We first derive some basic properties of the GMM model. More specifically, we have

$$ $\displaystyle p_{\theta}(X_{0})=\frac{1}{N}\sum_{i=1}^{N}\mathcal{N}(X_{0}|\mu ^{i},\sigma^{2}I),$ (14) $\displaystyle p_{\theta}(X_{t}|X_{0})=\mathcal{N}(X_{t}|X_{0},\sigma_{t}^{2}I),$ (15) $\displaystyle p_{\theta}(X_{t})=\frac{1}{N}\sum_{i=1}^{N}\mathcal{N}(X_{t}|\mu ^{i},(\sigma^{2}+\sigma_{t}^{2})I),$ (16) $\displaystyle p_{\theta}(X_{0}|X_{t})=\frac{p_{\theta}(X_{t}|X_{0})p_{\theta}( X_{0})}{p_{\theta}(X_{t})}$ $\displaystyle=\sum_{i=1}^{N}\frac{\mathcal{N}(X_{t}|X_{0})}{\sum_{j=1}^{N} \mathcal{N}(X_{t}|\mu^{j},(\sigma^{2}+\sigma_{t}^{2})I)}\mathcal{N}(X_{0}|\mu^ {i},\sigma^{2}I)$ $\displaystyle=\sum_{i=1}^{N}u_{i}\mathcal{N}(X_{0}|\mu^{i},\sigma^{2}I),$ $\displaystyle\textrm{where }u_{i}=(\exp-\frac{||X_{0}-X_{t}||^{2}}{2\sigma_{t} ^{2}})/(\sum_{j=1}^{N}\frac{1}{\sqrt{(1+\sigma^{2}/\sigma_{t}^{2})^{d}}}\exp{- \frac{||X_{t}-\mu^{i}||^{2}}{2(\sigma^{2}+\sigma_{t}^{2})}}).$ (17) $$

For true posterior $p_{\theta}(X_{0}|X_{t}=x_{t})$, we know that $X_{0}$ eventually converges to one of the $\mu^{i}$. Therefore, the $\mu^{*}$ closest to $X_{t}$ will have highest weighting $u_{i}$. And thus no matter what is the value of $\sigma^{2}_{t}$, the highest density mode of true posterior is always the mode $\mu^{*}$ that is closest to the initial point $x_{t}$. And the decision boundary is always the vonoroi centered at $\mu^{i}$.

Lemma 3.3
The PF-ODE can be written as:

$$ $\displaystyle\frac{dX_{t}}{dt}=\underbrace{\sum_{i=1}^{N}\frac{w_{i}}{2}\frac{ d\sigma_{t}^{2}}{dt}\frac{(X_{t}-\mu^{i})}{\sigma^{2}+\sigma_{t}^{2}}}_{ \textrm{velocity field }v_{t}},w_{i}=(\exp{-\frac{||X_{t}-\mu^{i}||^{2}}{2( \sigma^{2}+\sigma_{t}^{2})}})/(\sum_{j=1}^{N}\exp{-\frac{||X_{t}-\mu^{j}||^{2} }{2(\sigma^{2}+\sigma_{t}^{2})}}).$ $$

###### Proof.

We need to compute the score function first:

$$ $\displaystyle\nabla\log p_{\theta}(X_{t}=x_{t})=\frac{\nabla p_{\theta}(X_{t}= x_{t})}{p_{\theta}(X_{t}=x_{t})}$ $\displaystyle=\frac{1}{p_{\theta}(X_{t}=x_{t})}\nabla(\sum_{i=1}^{N}\frac{1}{N }(\mathcal{N}(X_{t}=x_{t}|\mu^{i},(\sigma^{2}+\sigma_{t}^{2})I)))$ $\displaystyle=\frac{1}{p_{\theta}(X_{t}=x_{t})}\sum_{i=1}^{N}\frac{1}{N} \mathcal{N}(X_{t}=x_{t}|\mu^{i},(\sigma^{2}+\sigma_{t}^{2}))(-\frac{(x-\mu^{i} )}{\sigma^{2}+\sigma_{t}^{2}})$ $\displaystyle=\sum_{i=1}^{N}((\exp{-\frac{||x_{t}-\mu^{i}||^{2}}{2(\sigma^{2}+ \sigma_{t}^{2})}})/(\sum_{j=1}^{N}\exp{-\frac{||x_{t}-\mu^{j}||^{2}}{2(\sigma^ {2}+\sigma_{t}^{2})}}))(-\frac{x-\mu^{i}}{\sigma^{2}+\sigma_{t}^{2}}).$ (18) $$

Take the score function into the PF-ODE definition in Eq. 3, we can obtain the result.
∎

With those basic properties, we can show that the solution of PF-ODE with initial value $X_{t}=x_{t}$ has non-zero density in true posterior $p_{\theta}(X_{0}|X_{t})$.

Proposition 3.2 The solution of PF-ODE has a positive likelihood in true posterior with high probability, i.e.,

$$ $\displaystyle p_{\theta}(X_{0}=\Phi(x_{t})|X_{t}=x_{t})\geq\frac{1}{N}\frac{1} {\sqrt{(4\pi\sigma^{2})^{d}}}\exp{(-\frac{2c^{2}}{\sigma_{t}^{2}}-\frac{d+1}{2 })}\textrm{, with probability }1-p.$ (19) $$

###### Proof.

$$ $\displaystyle p_{\theta}(X_{0}=\Phi(x_{t})|X_{t}=x_{t})=\sum_{i=1}^{N}u_{i} \mathcal{N}(X_{0}=\Phi(x_{t})|\mu^{i},\sigma^{2}I)$ (20) $\displaystyle\geq u_{j}\mathcal{N}(X_{0}=\Phi(x_{t})|\mu^{j},\sigma^{2}I),\forall j$ (21) $$

We let $k=\min_{i}\{||\mu^{i}-\Phi(x_{t})||\}$, by assumption we have $||\Phi_{0}(x_{t})-\mu^{k}||\leq\sigma^{2}+d\sigma^{2}$ with probability $1-p$.

$$ $\displaystyle p_{\theta}(X_{0}=\Phi(x_{t})|X_{t}=x_{t})\geq u_{k}\frac{1}{ \sqrt{(2\pi\sigma^{2})^{d}}}\exp{-\frac{||\mu^{k}-\Phi(x_{t})||^{2}}{2\sigma^{ 2}}}$ (22) $\displaystyle\overset{(a)}{\geq}u_{k}\frac{1}{\sqrt{(2\pi\sigma^{2})^{d}}}\exp {-\frac{(d+1)\sigma^{2}}{2\sigma^{2}}}$ (23) $\displaystyle=\frac{\exp{-\frac{||\Phi(x_{t})-x_{t}||^{2}}{2\sigma_{t}^{2}}}}{ \sum_{j=1}^{N}\sqrt{(1+\sigma^{2}/\sigma_{t}^{2})^{d}}\exp{-\frac{||x_{t}-\mu^ {i}||^{2}}{2(\sigma^{2}+\sigma_{t}^{2})}}}\frac{1}{\sqrt{(2\pi\sigma^{2})^{d}} }\exp{-\frac{d+1}{2}}$ (24) $\displaystyle\overset{(b)}{\geq}\frac{\exp{-\frac{4c^{2}}{2\sigma_{t}^{2}}}}{ \sqrt{2^{d}}\sum_{j=1}^{N}\exp{0}}\frac{1}{\sqrt{(2\pi\sigma^{2})^{d}}}\exp{- \frac{d+1}{2}}$ (25) $\displaystyle=\frac{1}{N}\frac{1}{\sqrt{(4\pi\sigma^{2})^{d}}}\exp{(-\frac{2c^ {2}}{\sigma_{t}^{2}}-\frac{d+1}{2})}$ (26) $$

(a) is due to the assumption that $||\Phi(x_{t})-\mu^{k}||^{2}\leq\sigma^{2}+d\sigma^{2}$. (b) is due to $||\Phi_{0}(x_{t})-x_{t}||\leq||\Phi_{0}(x_{t})||+||x_{t}||$. As they are both bounded by $c$, $||\Phi_{0}(x_{t})-x_{t}||^{2}$ is bounded by $4c^{2}$. And $1+\frac{\sigma^{2}}{\sigma_{t}^{2}}\leq 2$.
∎

We can study the PF-ODE in Eq. 11 informally when $\sigma_{t}^{2}$ is rather small or large. When $\sigma_{t}^{2}$ is small, the soft-max $w_{i}$ becomes a "hard"-max. Denote $k=\min_{i}\{||\mu^{i}-\Phi(x_{t})||\}$ and the PF-ODE at that time can be written as

$$ $\displaystyle\frac{dX_{t}}{dt}=-\frac{1}{2}\frac{d\sigma_{t}^{2}}{dt}(-\frac{( X_{t}-\mu^{k})}{\sigma^{2}+\sigma_{t}^{2}}).$ (27) $$

At that time, the PF-ODE is first order separable. And we can solve it with initial value $X_{t}=x_{t}$ as

$$ $\displaystyle\frac{\Phi(x_{t})-\mu^{k}}{x_{t}-\mu^{k}}=e^{h(t)},$ (28) $$

where $h(.)$ is some function of $t$ related to the $d\sigma_{t}^{2}/dt$.

Let’s assume a simple $\sigma_{t}^{2}$ schedule such as $\sigma_{t}^{2}=t^{2}$. In that case, we have

$$ $\displaystyle\Phi(x_{t})=\mu^{k}+(x_{t}-\mu^{k})e^{\frac{1}{2}\ln\frac{\sigma^ {2}}{\sigma^{2}+t^{2}}}.$ (29) $$

The solution has the form of $\mu^{k}$ with an offset term weighted by an exponential term. When $\sigma^{2}$ is small, the exponential term goes to 0 0 very fast. And therefore $\Phi(x_{t})\approx\mu^{k}$ at that time.

###### Proof.

###### Proof.

## Appendix B Additional Experiment Setup

### B.1 Implementation Details

All the experiments are implemented in Pytorch, and run in a computer with AMD EPYC 7742 CPU and Nvidia A100 GPU.

**Table 7: The model specification used for different non-linear operators.**
|  | Model A | Model B |
| --- | --- | --- |
| Segmentation | MobileNet + C1 | ResNet50 + PPM |
| Layout | Lin et al. [] | Lin et al. [] + DA |
| Caption | BLIP | CLIP |
| Classification | ResNet50 | VITB16 |

As we have shown, using the same model for $f(.)$ causes overfitting for neural network based $f(.)$. Therefore, we adopt different model for $f(.)$ in DIS and testing, and the details are shown in Tab. 7.

**Table 8: Metrics for DIS loss and evaluation.**
|  | $d(.,.)$ for DIS | metric for Test |
| --- | --- | --- |
| Segmentation | Cross Entropy | mIOU |
| Layout | Cross Entropy | mIOU |
| Caption | Cross Entropy | CLIP score |
| Classification | Cross Entropy | Accuracy |
| Downsample | MSE | MSE |

For different operators, we also have different $d(.,.)$ to evaluate the distance $d(f(x_{0|t}),y)$ during the DIS process. For all four non-linear operators, the cross entropy are used for $d(.,.)$. While for down-sample, we adopt MSE. To evaluate how consistent the generated samples are to $y$, we use y-metrics. Or to say, the metrics computed with input measurement $y$ and $f(x_{0|t})$. More specifically, for Segmentation and Layout, we evaluate consistency by y-mIOU. For image caption, we evaluate consistency by CLIP score [Hessel et al., ]. For classification, we evaluate consistency by accuracy. And for down-sample, we evaluate consistency by MSE. Note that the $d(.,.)$ used during DIS follows the convention of training corresponding $f(.)$, and the y-metric used for testing also follows the convention of testing corresponding $f(.)$.

### B.2 Details of DIS Algorithm

Below we provide detailed algorithm of different DIS methods including the ones we compare to and our own.
1
procedure DPS($p_{\theta}(.|.),T,f(.),y,d(.,.),\zeta_{t}$)
2
$x_{T}=\mathcal{N}(0,T^{2}I)$
3
for $t=T$ to $1$ do
4
$x_{t-1}\sim p_{\theta}(X_{t-1}|X_{t}=x_{t})$
5
$x_{0|t}=\mathbb{E}[X_{0}|X_{t}=x_{t}]$
6
$x_{t-1}\leftarrow x_{t-1}-\zeta_{t}d(f(x_{0|t}),y)$
7
return $x_{0}$
Algorithm 1 DPS
1
procedure FreeDOM($p_{\theta}(.|.),q(.|.)$,$T,f(.),y,d(.,.),\zeta_{t},r,K$)
2
$x_{T}=\mathcal{N}(0,T^{2}I)$
3
for $t=T$ to $1$ do
4
for $t^{\prime}=K$ to $1$ do
5
$x_{t-1}\sim p_{\theta}(X_{t-1}|X_{t}=x_{t})$
6
$x_{0|t}=\mathbb{E}[X_{0}|X_{t}=x_{t}]$
7
$x_{t-1}\leftarrow x_{t-1}-\zeta_{t}d(f(x_{0|t}),y)$
8
if $t^{\prime}\neq 1,t\in r$ then
9
$x_{t}=q(X_{t}|X_{t-1}=x_{t-1})$
10
else
11
break
12
return $x_{0}$
Algorithm 2 FreeDOM
1
procedure MPGD($q_{\theta}(.|.,.),T,f(.),y,d(.,.),\zeta_{t}$)
2
$x_{T}=\mathcal{N}(0,T^{2}I)$
3
for $t=T$ to $1$ do
4
$x_{0|t}=\mathbb{E}[X_{0}|X_{t}=x_{t}]$
5
$x_{0|t}\leftarrow x_{0|t}-\zeta_{t}d(f(x_{0|t}),y)$
6
$x_{t-1}\leftarrow q(X_{t-1}|X_{t}=x_{t},X_{t}=x_{0|t})$
7
return $x_{0}$
Algorithm 3 MPGD
1
procedure LGD($p_{\theta}(.|.),T,f(.),y,d(.,.),\zeta_{t},r_{t}$)
2
$x_{T}=\mathcal{N}(0,T^{2}I)$
3
for $t=T$ to $1$ do
4
$x_{t-1}\sim p_{\theta}(X_{t-1}|X_{t}=x_{t})$
5
$x_{0|t}=\mathbb{E}[X_{0}|X_{t}=x_{t}]+\mathcal{N}(0,r_{t}^{2}I)$
6
$x_{t-1}\leftarrow x_{t-1}-\zeta_{t}d(f(x_{0|t}),y)$
7
return $x_{0}$
Algorithm 4 LGD (Single Sample)
1
procedure STSL($p_{\theta}(.|.),T,f(.),y,d(.,.),\zeta_{t},\eta_{t}$)
2
$x_{T}=\mathcal{N}(0,T^{2}I)$
3
for $t=T$ to $1$ do
4
$x_{0|t}=\mathbb{E}[X_{0}|X_{t}=x_{t}]$
5
$x_{t}\leftarrow x_{t}-\zeta_{t}d(f(x_{0|t}),y)$
6
$\epsilon\sim\mathcal{N}(0,I)$
7
$x_{t}\leftarrow x_{t}-\eta_{t}\nabla_{x_{t}}(\epsilon^{T}(s_{\theta}(t,x_{t}+
\epsilon)-s_{\theta}(t,x_{t})))$
8
$x_{t-1}\sim p_{\theta}(X_{t-1}|X_{t}=x_{t})$
9
return $x_{0}$
Algorithm 5 STSL (Single Sample)
1
procedure Proposed-I($p_{\theta}(.|.),T,f(.),y,d(.,.),\zeta_{t},g_{\theta}(.,.),\tau$)
2
$x_{T}=\mathcal{N}(0,T^{2}I)$
3
for $t=T$ to $1$ do
4
$x_{t-1}\sim p_{\theta}(X_{t}|X_{t-1}=x_{t-1})$
5
$x_{0|t}=g_{\theta}(t,x_{t})+\mathcal{N}(0,\tau^{2}I)$
6
$x_{t-1}\leftarrow x_{t-1}-\zeta_{t}d(f(x_{0|t}),y)$
7
return $x_{0}$
Algorithm 6 Proposed I
1
procedure CM($g_{\theta}(.,.),t_{1:N}$)
2
for $n=1$ to $N-1$ do
3
$z\sim\mathcal{N}(0,I)$
4
$x_{t_{n}}\leftarrow x_{0}+t_{n}z$
5
$x_{0}\leftarrow g_{\theta}(t_{n},x_{t_{n}})$
6
return $x_{0}$
Algorithm 7 Consistency Model
1
procedure CMInversion($g_{\theta}(.,.)$, $t_{1:N}$, $f(.)$, $y$, $\tau$)
2
for $n=1$ to $N-1$ do
3
$z\sim\mathcal{N}(0,I)$
4
for $k=1$ to K do
5
$x_{t_{n}}\leftarrow x_{0}+t_{n}z$
6
$x_{0}\leftarrow g_{\theta}(t_{n},x_{t_{n}})$
7
$x_{\tau}\leftarrow x_{0}+\mathcal{N}(0,\tau^{2}I)$
8
$z\leftarrow z-\zeta\frac{d}{dz}d(f(x_{\tau}),y)$
return $x_{0}$
Algorithm 8 Proposed II

FreeDOM Yu et al. [] propose to adopt the time-travel that is designed specifically for in-painting [Lugmayr et al., ] to general operator $f(.)$. (See Algorithm. 2) More specifically, it proposes an inner loop that goes forward after a backward step with forward kernel $q(.|.)$. The new hyper-parameters are time-travel steps $K$ and time-travel range $r$.

MPGD He et al. [] propose to perform the gradient ascent directly on posterior mean instead of on $x_{t}$. And the posterior mean after gradient ascent is used to correct the score function (See Algorithm. 3). They claim that their approach is able to converge faster and outperform DPS when $T=20,100$. However, as we use $T=1000$, the advantage of their approach is not clearly shown in our experiments.

LGD Song et al. [] adopt a Gaussian approximation to posterior sample (See Algorithm. 4). More specifically, they use an additive of Gaussian noise and posterior mean as an approximation of Gaussian sample. And the mean of approximated sample is the same as real posterior sample. This approach is later improved by Boys et al. [], Rout et al. [] to second order. Or to say, they estimate the covariance of Gaussian using second order Tweedie’s formula. And the mean and covariance of approximated sample is the same as real posterior sample. The authors of LGD further propose a multi-sample approach to reduce gradient variance. However, we only use LGD with sample size $n=1$ for fair comparison. The new hyper-parameters is variance $\tau$.

STSL Rout et al. [] propose to improve Song et al. [] by estimating the posterior as Gaussian distribution with posterior mean and posterior covariance. As directly estimating the posterior covariance using second order Tweedie’s formula is expensive, they propose a Monte Carlo estimation and the resulting algorithm is shown in Algorithm. 5. However, we only use STSL with sample size $n=1$ for fair comparison.

### B.3 Hyper-parameters

**Table 9: The hyper-parameters of other DIS and proposed approaches.**
|  | Segmentation | Layout | Caption | Classification | Down-sampling |
| --- | --- | --- | --- | --- | --- |
| DPS | $\zeta=256.0$ | $\zeta=7.2$ | $\zeta=24.0$ | $\zeta=8.0$ | $\zeta=14.4$ |
| LGD | $\zeta=256.0$ | $\zeta=7.2$ | $\zeta=24.0$ | $\zeta=8.0$ | $\zeta=14.4$ |
| $K=1,\tau=0.2$ |  |  |  |  |  |
| FreeDOM | $\zeta=256.0$ | $\zeta=7.2$ | $\zeta=24.0$ | $\zeta=8.0$ | $\zeta=14.4$ |
| $K=2,r=[100,200]$ |  |  |  |  |  |
| MPGD | $\zeta=2560.0$ | $\zeta=72.0$ | $\zeta=240.0$ | $\zeta=80.0$ | $\zeta=144.0$ |
| Proposed I | $\zeta=256.0$ | $\zeta=7.2$ | $\zeta=24.0$ | $\zeta=8.0$ | $\zeta=14.4$ |
| $\tau=0.2$ |  |  |  |  |  |
| Proposed II | $\zeta_{1}=0.1,\zeta_{2}=0.005$ | $\zeta_{1}=0.1,\ \zeta_{2}=0.001$ | $\zeta_{1}=0.1,\ \zeta_{2}=0.001$ | $\zeta_{1}=0.1,\zeta_{2}=0.005$ | $\zeta_{1}=0.3,\zeta_{2}=0.03$ |
| $ts=3,6,...,18,30,39$ | $ts=75,100,125,150$ | $ts=75,100,125,150$ | $ts=25,50,75,125,150$ | $ts=75,100,125,150$ |  |
| $K=40$ | $K=151$ | $K=151$ | $K=151$ | $K=151$ |  |

Figure: Figure 7: Comparison of different inner-loop range for FreeDOM.
Refer to caption: x6.png

We list the detailed hyper-parameters of DPS, LGD, FreeDOM, MPGD [Chung et al., , Song et al., , Yu et al., , He et al., ] and two of our proposed approaches in Tab. 9. For all the methods, one common hyper-parameter is the step size $\zeta$ used in gradient descent. For LGD and our Proposed I, an additional hyper-parameter is the additional additive noise $\tau=0.2$. We do not use the multi-sample LGD as it is significantly slower than all other approaches. For FreeDOM, the additional parameters are time-travel steps $K$, and time-travel range $r$. We set $K=2$ for fair comparison, as a large $K$ make FreeDOM significantly slower than all other approaches. We set $r=[100,200]$ instead of $r=[200,500]$ in original paper [Yu et al., ]. This is because we find that setting $r=[200,500]$ in VE-diffusion has significant negative effect on sample quality. The visual comparison is in Fig. 7. For Proposed II, in addition to learning rate $\zeta$, we also control the timestep schedule $ts$ and optimization step $K$.

## Appendix C Additional Experiment Results

### C.1 Additional Visual Results

We present more visual results of non-linear operators in Fig. 10 and 11. It is shown that our Proposed I has best consistency with measurement $y$ and best sample quality for most of the times. Our Proposed II also has a good consistency and sample quality.

Furthermore, we present more visual results of down-sampling in Fig. 8. It is shown that both our Proposed II and Proposed II work as good as DPS.

Figure: Figure 8: Additional visual results on image down-sampling.
Refer to caption: x7.png

### C.2 Failure Cases

When the measurement $y$ is too far from diffusion prior, our approaches and other DIS approaches fail. An example of such failure is shown in Fig. 9. The input measurement $y$ describes a woman. However, human is not a part of LSUN bedroom dataset. And none-of the DIS approaches is able to generate a woman. And the samples generated by DIS look like unconditional sample.

Figure: Figure 9: Visual results of a failure case.
Refer to caption: x8.png

Figure: Figure 10: Additional visual results on image segmentation and layout estimation.
Refer to caption: x9.png

Figure: Figure 11: Additional visual results on image captioning and image classification.
Refer to caption: x11.png

## Appendix D Additional Discussion

### D.1 Reproducibility Statement

The proof of all theoretical results are shown in Appendix. A. For experiments, all two datasets are publicly available. In Appendix. B, we provide additional implementation details of all other DIS that we compare to. Further, detailed hyper-parameters of all baselines and our proposed approach are presented. Besides, we provide source code for reproducing empirical results as supplementary material.

### D.2 Broader Impact

The approach proposed in this paper allows conditional generation without training a new model. This saves the energy of training conditional generative diffusion model and reduces the carbon emission. Potential negative impact is the same as other conditional generative model, such as trustworthiness brought by generating fake image.

## NeurIPS Paper Checklist

- 1.
Claims
- Question: Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
- Answer: [Yes]
- Justification: Yes.
- Guidelines:
•
The answer NA means that the abstract and introduction do not include the claims made in the paper.
•
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
•
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
•
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
- 2.
Limitations
- Question: Does the paper discuss the limitations of the work performed by the authors?
- Answer: [Yes]
- Justification: In Discussion.
- Guidelines:
•
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
•
The authors are encouraged to create a separate "Limitations" section in their paper.
•
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
•
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
•
The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
•
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
•
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
•
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren’t acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
- 3.
Theory Assumptions and Proofs
- Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
- Answer: [Yes]
- Justification: In appendix.
- Guidelines:
•
The answer NA means that the paper does not include theoretical results.
•
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
•
All assumptions should be clearly stated or referenced in the statement of any theorems.
•
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
•
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
•
Theorems and Lemmas that the proof relies upon should be properly referenced.
- 4.
Experimental Result Reproducibility
- Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
- Answer: [Yes]
- Justification: Code is available in supplementary materials.
- Guidelines:
•
The answer NA means that the paper does not include experiments.
•
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
•
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
•
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
•
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a)
If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b)
If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c)
If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
- 5.
Open access to data and code
- Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
- Answer: [Yes]
- Justification: In supplementary materials.
- Guidelines:
•
The answer NA means that paper does not include experiments requiring code.
•
Please see the NeurIPS code and data submission guidelines ([https://nips.cc/public/guides/CodeSubmissionPolicy](https://nips.cc/public/guides/CodeSubmissionPolicy)) for more details.
•
While we encourage the release of code and data, we understand that this might not be possible, so “No” is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
•
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ([https://nips.cc/public/guides/CodeSubmissionPolicy](https://nips.cc/public/guides/CodeSubmissionPolicy)) for more details.
•
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
•
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
•
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
•
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
- 6.
Experimental Setting/Details
- Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
- Answer: [Yes]
- Justification: In Experimental setup and appendix.
- Guidelines:
•
The answer NA means that the paper does not include experiments.
•
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
•
The full details can be provided either with the code, in appendix, or as supplemental material.
- 7.
Experiment Statistical Significance
- Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
- Answer: [No]
- Justification: Error bar is expensive and rarely reported in our community.
- Guidelines:
•
The answer NA means that the paper does not include experiments.
•
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
•
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
•
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
•
The assumptions made should be given (e.g., Normally distributed errors).
•
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
•
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.
•
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
•
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
- 8.
Experiments Compute Resources
- Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
- Answer: [Yes]
- Justification: In Experimental setup and appendix.
- Guidelines:
•
The answer NA means that the paper does not include experiments.
•
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
•
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
•
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn’t make it into the paper).
- 9.
Code Of Ethics
- Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics [https://neurips.cc/public/EthicsGuidelines](https://neurips.cc/public/EthicsGuidelines)?
- Answer: [Yes]
- Justification: Yes.
- Guidelines:
•
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
•
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
•
The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).
- 10.
Broader Impacts
- Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
- Answer: [Yes]
- Justification: In appendix.
- Guidelines:
•
The answer NA means that there is no societal impact of the work performed.
•
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
•
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
•
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
•
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
•
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
- 11.
Safeguards
- Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
- Answer: [N/A]
- Justification: No new model is trained.
- Guidelines:
•
The answer NA means that the paper poses no such risks.
•
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
•
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
•
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
- 12.
Licenses for existing assets
- Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
- Answer: [Yes]
- Justification: In Experimental setup and appendix.
- Guidelines:
•
The answer NA means that the paper does not use existing assets.
•
The authors should cite the original paper that produced the code package or dataset.
•
The authors should state which version of the asset is used and, if possible, include a URL.
•
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
•
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
•
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, [paperswithcode.com/datasets](paperswithcode.com/datasets) has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
•
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
•
If this information is not available online, the authors are encouraged to reach out to the asset’s creators.
- 13.
New Assets
- Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
- Answer: [Yes]
- Justification: In supplementary material
- Guidelines:
•
The answer NA means that the paper does not release new assets.
•
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
•
The paper should discuss whether and how consent was obtained from people whose asset is used.
•
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
- 14.
Crowdsourcing and Research with Human Subjects
- Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
- Answer: [N/A]
- Justification: [N/A]
- Guidelines:
•
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
•
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
•
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
- 15.
Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
- Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
- Answer: [N/A]
- Justification: [N/A]
- Guidelines:
•
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
•
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
•
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
•
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: see above
soundness: 2
strengths: The paper considers an important problem: how diffusion models can be used to solve complex inverse problems, such as giving a segmented image to reproduce the underlying image. The paper acknowledges a known limitation on one assumption that prior work makes on the distribution of the data (which allows to use the Tweedie's formula), and proposes to address it by consistency models.
summary: The paper proposes using consistency models (CMs) to solve inverse problems with diffusion models. In diffusion inverse problems, common approaches need to go from xt to x0 at every iteration to be able to compute a measurement-guided gradient with respect to the measurement y from x0. The majority uses the expectation from Tweedie's formula to compute x0 from tx, which may not result in a good example for complex multi-modal distribution. The paper proposes to replace this step with a CM. They show improvement upon prior work.
weaknesses: I find the contribution of the contribution incremental, and the presentation can be improved. The main weaknesses are as follows: the mathematical formulation is confusing and the approach to training the CM is not fully fair compared to the baselines.

1. The paper trains a CM to be used for diffusion-based inverse solvers. It's not clear why such mathematical details (some are not fully proper) are included, which I do not think is the main contribution of the paper. It's unclear why in (1) the authors formulate the Markov chain following a VE formulation. Could the author explain this choice? The conventional Markov chain for diffusion models is the popular one based on VP, which has a different mean and variance such that the structure is destroyed, but the energy of the process remains the same.

2. Solving inverse problems with diffusion mostly occurs in the regime where the diffusion model is trained unconditionally without the knowledge of the measurement operator f(.) (see the DPS used for vision problems such as deblurring). However, the paper in Section 3.4 discusses that the CM is overfitting to f(.) and proposes an approach to make the framework robust. So this brings the following: the CM going from xt to x0 seems to be trained with the knowledge of the measurement. If f(.) is involved during the training, then the trained framework is not general anymore (it's problem specific). Hence, the comparison of the proposed framework to models such as DPS, where the model is not trained based on the measurement operator, is not fair. Please provide more information and clarity if this is not the case. The fair comparison would be a scenario where both methods are trained a similar condition (e.g., not having the knowledge of the forward operator).

Here are some questions concerning this:

- Is CM trained with knowledge of f(.), or if this is a misunderstanding?
- If the CM is trained with f(.), please explain and justify comparing it to methods like DPS that don't use this information?

More comments are below:


Lack of thorough literature

- The limitations of DIS are not fully explained in the intro. Indeed, the mean-based approximation is one challenge. A few others are related to whether methods such as DPS are doing posterior sampling or using the measurement to guide the process onto likely solutions (see [1]).


The paper needs improvement in presentation. Here are a few examples

- While the notations such as Xt and X0 are known to the reader with knowledge of diffusion models, these are used in the abstract and intro without introducing them. Hence, I suggest re-writing the abstract without these notations and introducing the diffusion in the introduction before using x0, xt, etc.

- Consistency models are not defined and introduced in the intro, but the authors explain that they are used to improve performance. I suggest the authors to provide a brief definition or explanation of consistency models in the introduction.

- How the results are generated for Table 1; this appears abruptly without proper explanation. I suggest to include a brief explanation of the methodology used to generate the results in Table 1.

Some terms within the manuscript are not precise and clear. Please provide clarifications.

- Section 1: with "neural network operators"? Does this refer to the measurement operator or the score function of the diffusion? I suggest to say "measurement operators" instead of "operators". Please clarify "neural network operators"?


[1] Wu, Z., Sun, Y., Chen, Y., Zhang, B., Yue, Y., & Bouman, K. L. (2024). Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors. arXiv preprint arXiv:2405.18782.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. More DIS baselines are desired such as CoSIGN [2], DDNM [3]

I am open to change my rating if authors could address my concerns (how clearly a CM model could benefits with inverse problem solving comparing to strong baselines like DDNM)
soundness: 3
strengths: 1. Easy to follow
2. The idea of using CM to approximate the PF-ODE solution is interesting
summary: This paper presents an interesting approach for posterior sampling using $p(x_0|x_t)$ being approximated via consistency model. Results show improvement over baselines such as DPS
weaknesses: 1. Using CM over the existing inverse problem solving adds a lot of computational burden, while the benefit is not clearly visible. Even though the $x_0|x_t$ may be off the manifold of ground truth image distribution, this does not imply a sacrifice in reconstruction quality as demonstrated in this work [1]. The quantitative results do not show a significant improvement over DPS.
2. CM itself can be used as a good prior for solving inverse problems: see this work [2]. This paper needs to compare with more recent works in inverse problem solving
3. More DIS baselines are desired such as DDNM [3]




[1]. Wang, Hengkang, et al. "DMPlug: A Plug-in Method for Solving Inverse Problems with Diffusion Models." arXiv preprint arXiv:2405.16749 (2024). NeurIPS 2024

[2]. Zhao, Jiankun, Bowen Song, and Liyue Shen. "CoSIGN: Few-Step Guidance of ConSIstency Model to Solve General INverse Problems." ECCV 2024

[3]. Wang, Yinhuai, Jiwen Yu, and Jian Zhang. "Zero-Shot Image Restoration Using Denoising Diffusion Null-Space Model." The Eleventh International Conference on Learning Representations.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: The presentation of Proposition 3.3 is a bit misleading. The authors should state in the assumption the condition on $\sigma$ that they use in the proof in order to get a lowerbound independent of the dimension, or remove this claim after the proposition.
soundness: 3
strengths: The idea is quite original wrt to the literature. The paper is quite illustrated and clear. The experiments are interesting and extensive; the authors compare to many existing methods, on both pixel space and latent space diffusion.
summary: This paper is about solving inverse problems using a pre-trained denoising diffusion prior. Posterior sampling with diffusion models requires an estimate of the score $\nabla _{x_t} \log p_t(x_t) + \nabla _{x_t} \log \int p(y | x_0) p _{0|t}(x_0 | x_t) d x_0$. While the first term is  estimated using the pre-trained score, the second term is usually very difficult to estimate accurately. A common approximation used in the literature involves using $\nabla _{x_t} \log p(y | E[X_0 | X_t = x_t])$ where  $E[X_0 | X_t = x_t]$ can also be estimated using the pre-trained score via Tweedie's formula. This approximation results in many efficiencies well documented in the literature and this paper tries to circumvent them using by using a sample from the PF-ODE as a replacement. Specifically, the authors use consistency models to speed up the process of sampling from the PF-ODE and ensuring that the differentiation is not costly.
weaknesses: - The most obvious weakness of the method is the use of consistency models since these are quite difficult to train and pre-trained CMs are not widely available. 

- In my opinion the justification for the method is rather weak. The paper argues that using a sample from the PF-ODE is valid because the sample has zero density and that furthermore, for the Gaussian mixture example considered, the PF-ODE sample has non-zero density under the posterior $p(x_0 | x_t)$ with high probability. At the same time it is also easy to find examples of a likelihood function $p(y|x_0)$ such $\int p(y|x_0) p(x_0 | x_t) d x_0 > 0$ for all $x_t$ but $p(y|\Phi(t, x_t)) = 0$ for $x_t$ in a set of positive Lebesgue measure. As a result, I'm not totally convinced that the argument is strong.

---

confidence: 5
contribution: 2
initial_rating: 6
presentation: 2
questions: The presented manuscript often states likelihood terms as $p_\theta(y|x)$. Please elaborate on why it contains theta, as the likelihood in general is not a learned function with the same parameters as the learned data distribution $p_\theta(x)$.

Furthermore, I would be interested as to which extend the PF-ODE solution differs from the MAP solution?
soundness: 3
strengths: The idea of using CMs in the process of inverse problems/posterior approximation seems novel and interesting.

The proposed method outperforms baseline approaches, particularly when compared to straightforward extension
summary: The authors propose a new approach using CMs to generate realistic image samples in Diffusion Inverse Solvers, improving the application of complex, non-linear neural network operators like those in  semantic segmentation, room layout estimation, image captioning, and image classification. Unlike traditional methods that produce low probability images, the incorporation of CMs is expected to maintain sample realism, resulting in more accurate posterior approximations, particularly when neural network-based operators are involved.
weaknesses: Although being interesting, the novelty of the presented work is marginally above acceptation threshold since the only contribution seems to be the use of CMs in order to compute $x_0$ from $x_t$.

The manuscript could benefit from improved clarity and organization, as certain sections are challenging to follow. See further remarks.

Further remarks: 

In the presented algorithm, the updates are stated as $\zeta_t \Delta\left(f(x_0 \mid t), y\right)$, where $\Delta$ is defined as some distance. The update of $x_t$ however should be the gradient of that distance.

In order to enhance the readability of the work, the authors should think about introducing a clear distinction between the posterior $p(x_0|x_t)$ and the posterior $p(x|y)$, given y is the observation.

The abbreviation DPS (Diffusion posterior sampling) is used but never introduced.

Meta-review: The paper explores the use of consistency models (CMs) to address inverse problems in diffusion models by replacing the traditional expectation computation step with a CM-based approach. This allows for faster sampling from the probability flow ODE (PF-ODE) while avoiding costly differentiation steps. The proposed method demonstrates improvements over prior approaches, highlighting an interesting application of CMs in solving inverse problems.

Strength: The paper discusses an interesting question, which tries to address the limitations of existing methods. Although it has been shown in previous works, the idea of leveraging CMs for solving inverse problems is interesting and hasn’t been explored too much in the literature. It may open some interesting directions for diffusion-based inverse problem-solving.

Weakness: Reviewers raise several common concerns, mainly on the justification why the additional CM is necessary in solving the inverse problems. Specifically, The approach introduces computational overhead, as it requires training both a diffusion model and a CM, which could be a significant constraint given the difficulty in training CMs. Additionally, the quantitative results show limited improvement over existing methods like DPS, and the comparison with more recent baselines would be more appreciated. Besides, the mathematical formulation and paper presentation could be further improved. 

Overall, although the target problem is interesting, the paper may need a major revision to improve clarity and method justification. Thus, I recommend reject.

Decision: Reject