【论文阅读】Consistency Models

Introduction

• 相比于单步生成的模型（例如 GANs, VAEs, normalizing flows)，扩散模型的迭代式生成过程需要 10 到 2000 步计算来采样，导致推理速度低，实时性应用受限.

• 本文的目的是创造高效、单步的生成，同时不牺牲迭代采样的优势。在数据到噪声的 PF-ODE 轨迹上，学习轨迹上任意点到轨迹起点的映射，对这些映射的建模成为 consistency model.
  

• 两种训练 consistency model的方法

  1. 使用 numerical ODE solver 和预训练的扩散模型在 PF-ODE 轨迹上生成若干相邻点对，通过最小化模型输出点对间的距离（相似度），蒸馏出 consistency model.
  2. 不依赖预训练扩散模型，独立训练一个 consistency model.

• 在一些数据集上测试.

Diffusion Models

使用$p_{data}(\mathrm{x})$表示数据分布，扩散模型使用如下随机微分公式对服从原分布的数据进行扩散：

$${\mathrm{dx}}_t=\mu (,x_t,t)+\sigma (t){\mathrm{dw}}_t$$

其中 $t$为时间步，范围是 $0$ 到 $T$，$\mu (\cdot ,\cdot )$ 和 $\sigma (\cdot )$分别是布朗运动中的漂移系数和扩散系数，${\mathbf{x}}_t$服从分布 $p_t(\mathrm{x})$，${\mathrm{x}}_0$服从分布 $p_{data}(\mathrm{x})$. 该方程的一个重要属性是，其存在一个 PF-ODE 方程：

$${\mathrm{dx}}_t=\left[\mu ({\mathrm{x}}_t,t)-\frac{1}{2}\sigma (t{)}^2\nabla \log p_t({\mathrm{x}}_t)\right]\mathrm{d}t$$

其中$\nabla \log p_t(\mathrm{x})$是 $p_t(\mathrm{x})$的 score function.
在 SDE 中，令漂移系数 $\mu (\mathrm{x},t)=0$， 扩散系数 $\sigma (t)=\sqrt{2t}$. 使用得分匹配的方式训练模型 $s_{\varphi }(\mathrm{x},t)\approx \nabla \log p_t(\mathrm{x})$，代入 PF-ODE 方程，得到 empirical PF-ODE：

$$\frac{{\mathrm{dx}}_t}{\mathrm{d}t}=-t{s}_{\varphi }({\mathrm{x}}_t,t）$$

采样时，使用 ${\hat{\mathrm{x}}}_T\sim \mathcal{N}(0,T^{2}I)$初始化，再使用 numerical ODE solver（例如 Euler, Heun）按时间步倒推出 ${\hat{x}}_0$. 为了防止数值不稳定，会在 $t=ϵ$是提前终止，$ϵ$为一个正小数，同时将 ${\hat{\mathrm{x}}}_{ϵ}$作为结果.

扩散模型的瓶颈在于采样速度慢， ODE solver 利用得分模型$s_{\varphi }(\mathrm{x},t)$迭代求解，消耗算力多. 目前存在一些更快的 ODE solver，但是仍然需要大于 $10$ 步的采样. 也存在一些蒸馏方法，但是大多数方法需要从扩散模型中采集巨大的数据集，同样消耗算力多.

Consistency Models

Definition

根据 PF-ODE 得到一条解路径 { x t } t ∈ [ ϵ , T ] \{\mathrm{x}_t\}_{t\in[\epsilon, T]} {xt​}t∈[ϵ,T]​，将 consistency function 定义为：

$$f:({\mathrm{x}}_t,t)\mapsto {\mathrm{x}}_{ϵ}$$

对于该路径上的任意点 $({\mathrm{x}}_t,t)$，其输出是一致的. 对于任意的 $t,t^{\prime }\in [ϵ,T]$，有 $f({\mathrm{x}}_t,t)=f({\mathrm{x}}_{t^{\prime }},t^{\prime })$恒成立.

Parameterization

令 $F_{\theta }(\mathrm{x},t)$表示任意形式的神经网络，使用 sikp connection 可以将模型表示为：

$$f_{\theta }(\mathrm{x},t)=c_{skip}(t)\mathrm{x}+c_{out}(t)F_{\theta }(\mathrm{x},t)$$

其中边界条件为 $c_{skip}(ϵ)=1$，$c_{out}(ϵ)=0$.
具体为：

$$c_{skip}(t)=\frac{{\sigma }_{data}^2}{(t-ϵ{)}^2+{\sigma }_{data}^2}$$

$$c_{out}(t)=\frac{{\sigma }_{data}(t-ϵ)}{\sqrt{{\sigma }_{data}^2+t^2}}$$

${\sigma }_{data}$取值 $0.5$.

Sampling

有了一个训练好的 consistency model $f_{\theta }(\cdot ,\cdot )$之后，从高斯噪声 $\mathcal{N}(0,T^{2}I)$采样${\hat{\mathrm{x}}}_T$，再代入模型一步推出 ${\hat{\mathrm{x}}}_{ϵ}=f_{\theta }(\widehat{{\mathrm{x}}_T},T)$.为了提高质量，也可以进行多步采样，算法如下：

Training Consistency Models via Distillation

作者的第一个方法是在预训练的得分模型 $s_{\varphi }(\mathrm{x},t)$上蒸馏.

首先考虑将 $ϵ$到 $T$的时间离散化成 $N-1$ 个间隔，也即 $t_1=ϵ<t_2<t_3<...<t_N=T$. 在实践中，使用如下公式：

$$t_i={\left({ϵ}^{1/\rho }+\frac{i-1}{N-1}\left(T^{1/\rho }-{ϵ}^{1/\rho }\right)\right)}^{\rho }$$

其中 $\rho =7$. 当 $N$充分大时，可以获得 ${\mathrm{x}}_{t_n}$到 ${\mathrm{x}}_{t_{n+1}}$的准确估计，于是 ${\hat{\mathrm{x}}}_{t_n}^{\varphi }$可以定义为：

$${\hat{\mathrm{x}}}_{t_n}^{\varphi }={\mathrm{x}}_{t_{n+1}}+(t_n-t_{n+1})\Phi ({\mathrm{x}}_{t_{n+1}},t_{n+1};\varphi )$$

$\Phi (...;\varphi )$为 one-step ODE solver（比如Euler）.

从数据集中采样 $\mathrm{x}$,通过 SDE 加噪$\mathcal{N}(\mathrm{x},t_{n+1}^{2}I)$得到 ${\mathrm{x}}_{t_{n+1}}$, 然后使用 ODE solver 求解出${\hat{\mathrm{x}}}_{t_n}^{\varphi }$，通过最小化在 ${\hat{\mathrm{x}}}_{t_n}^{\varphi }$ 和 ${\mathrm{x}}_{t_{n+1}}$计算结果的差距训练模型.

Definition 1
consistency distillation loss (CD)表示为：

$${\mathcal{L}}_{CD}^N(\theta ,{\theta }^{-};\varphi )=\mathbb{E}\left[\lambda (t_n)d(f_{\theta }({\mathrm{x}}_{t_{n+1}},t_{n+1}),f_{{\theta }^{-}}({\hat{\mathrm{x}}}_{t_n}^{\varphi },t_n)\right]$$

其中，$\lambda (\cdot )\in {\mathbb{R}}^{+}$是正权重函数，${\theta }^{-}$是$\theta$在优化过程中历史值的均值. $d(\cdot ,\cdot )$是一个度量函数，满足当且仅当两个输入相等时为 $0$，其余情况大于 $0$.

作者考虑 $d(\cdot ,\cdot )$ 使用 $l_1$ 以及 $l_2$，在实验中$\lambda (t_n)\equiv 1$表现较好. ${\theta }^{-}$使用 EMA 更新，计算公式如下：

$${\theta }^{-}\leftarrow \mathrm{stopgard}(\mu {\theta }^{-}+(1-\mu )\theta )$$

其中 $0\le \mu <1$. 使用 EMA 可以使训练更稳定，同时能提高模型的表现.
模型训练算法如下：

Training Consistency Models in Isolation

consistency model 可以不依赖预训练扩散模型训练，使用如下无偏估计替换$\nabla \log p_t(\mathrm{x})$：

$$\nabla \log p_t(\mathrm{x})=-\mathbb{E}\left[\frac{{\mathrm{x}}_t-\mathrm{x}}{t^2}|{\mathrm{x}}_t\right]$$

consistency training loss (CT)表示为：

$${\mathcal{L}}_{CD}^N(\theta ,{\theta }^{-})=\mathbb{E}\left[\lambda (t_n)d(f_{\theta }(\mathrm{x}+t_{n+1}\mathrm{z},t_{n+1}),f_{{\theta }^{-}}(\mathrm{x}+t_n\mathrm{z},t_n)\right]$$

其中 $\mathrm{z}\sim \mathcal{N}(0,I)$. 损失函数的计算依赖于$f_{\theta }$和 $f_{{\theta }^{-}}$，且与扩散模型的无关.

为了提升模型效果，使用 schedule function $N(\cdot )$控制 $N$ 增长. 直觉上，当 $N$ 小的时候，使用 consistency distillation loss 模型在一开始收敛更快，同时方差小、偏差大. 反之，在训练结束时，应当使 $N$ 大，这样方差大、偏差小。同时，使用 schedule function $\mu (\cdot )$替换 $\mu$，让它随着 $N$ 增长而变化.
$N(\cdot )$和$\mu (\cdot )$具体为

$$N(k)=\lceil \sqrt{\frac{k}{K}((s_1+1{)}^2-s_0^2)+s_0^2}-1\rceil +1$$

$$\mu (k)=\exp \left(\frac{s_0\log {\mu }_0}{N(k)}\right)$$

$K$表示整体训练步数，$s_0$表示开始的离散化步数.

训练算法如下：

Experiment

关于 CD ，作者分别使用 $l_1$,$l_2$,$\mathrm{LPIPS}$作为度量函数，使用一阶Euler和二阶Heun座位 ODE solver，$N$取 { 9 , 12 , 18 , 36 , 50 , 60 , 80 , 120 } \{9,12,18,36,50,60,80,120\} {9,12,18,36,50,60,80,120}，使用相应的预训练扩散模型做初始化. 使用 CT 训练的模型则随机初始化.

(a) 对比不同的度量函数在 CD 上的表现，其中 LPIPS 的效果最好.
(b, c) 对不不同 ODE solver 和$N$在 CD 上的表现，使用 Heun 且$N$取$18$时效果最好.在取相同的 $N$时，二阶Heun的表现优于一阶Euler，因为高阶的 ODE solver 的估计误差更小. 当$N$充分大时，模型对$N$变得不敏感.
(d) 根据之前的结论，关于 CT 的实验使用 LPIPS 作为度量函数. 更小的$N$收敛更快，但是采样结构更差；使用自适应的$N(\cdot )$和$\mu (\cdot )$效果最好.

对比 CD 和 progressive disillation(PD) 在不同数据集上的效果，CD 的表现普遍比 PD 好.

对比 CT 和其它生成模型，仅使用一步或两步生成.

Zero-Shot Image Editing