【论文阅读】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 (?,?)$ 和 $\sigma (?)$分别是布朗运动中的漂移系数和扩散系数，${\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?\log p_t({\mathrm{x}}_t)\right]\mathrm{d}t$$

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

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

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

扩散模型的瓶颈在于采样速度慢， ODE solver 利用得分模型$$s_{?}(\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)?{\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 }(?,?)$之后，从高斯噪声 $\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_{?}(\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}^{?}$可以定义为：

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

$\Phi (...;?)$为 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}^{?}$，通过最小化在 ${\hat{\mathrm{x}}}_{t_n}^{?}$ 和 ${\mathrm{x}}_{t_{n+1}}$计算结果的差距训练模型.

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

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

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

作者考虑 $d(?,?)$ 使用 $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 可以不依赖预训练扩散模型训练，使用如下无偏估计替换$?\log p_t(\mathrm{x})$：

$$?\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}～\mathcal{N}(0,I)$. 损失函数的计算依赖于$f_{\theta }$和 $f_{{\theta }^{?}}$，且与扩散模型的无关.

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

$$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(?)$和$\mu (?)$效果最好.

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

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

Zero-Shot Image Editing