Consistency Models Improve Diffusion Inverse Solvers
Tongda Xu Ziran Zhu Dailan He Yuanyuan Wang Ming Sun Ning Li Hongwei Qin Yan Wang Jingjing Liu Ya-Qin Zhang
Diffusion inverse solvers (DIS) aim to find an image $x$ that lives on the
diffusion prior while satisfying the constraint $f(x) = y$, given an operator
$f(.)$ and measurement $y$. Most non-linear DIS use posterior mean
$\hat{x}_{0|t}=\mathbb{E}[x_0|x_t]$ to evaluate $f(.)$ and minimize the
distance $||f(\hat{x}_{0|t})-y||^2$. Previous works show that posterior
mean-based distance is biased; instead, posterior sample $x_{0|t}\sim
p_{\theta}(x_0|x_t)$ promises a better candidate. In this paper, we first
clarify when is posterior sample better: $1)$ When $f(.)$ is linear, the
distance with posterior mean is as good as single posterior sample, thus
preferable as it does not require Monte Carlo; $2)$ When $f(.)$ is non-linear,
the distance using posterior sample is better. As previous approximations to
posterior sample do not look like a real image, we propose to use consistency
model (CM) as a high quality approximation. In addition, we propose a new
family of DIS using pure CM. Empirically, we show that replacing posterior mean
by CM improves DIS performance on non-linear $f(.)$ (e.g. semantic
segmentation, image captioning). Further, our pure CM inversion works well for
both linear and non-linear $f(.)$.