【论文阅读】Consistency Models-Toy模板网

这篇具有很好参考价值的文章主要介绍了【论文阅读】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})$ 表示数据分布，扩散模型使用如下随机微分公式对服从原分布的数据进行扩散：

$\large \mathrm{dx}_t = \mu(\mathrm,{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 方程：

$\large\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_{\phi}(\mathrm{x},t) \approx \nabla\log{p_t(\mathrm{x})}$ ，代入 PF-ODE 方程，得到 empirical PF-ODE：

$\large \frac{\mathrm{dx}_t}{\mathrm{d}t}=-ts_{\phi}(\mathrm{x}_t,t）$

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

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

Consistency Models

Definition

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

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

对于该路径上的任意点 $(\mathrm{x}_t, t)$ ，其输出是一致的. 对于任意的 $\in [\epsilon, T]$ ，有 $f(\mathrm{x}_t, t) =f(\mathrm{x}_{t'}, t')$ 恒成立.
【论文阅读】Consistency Models,论文阅读,AIGC

Parameterization

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

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

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

$\large c_{skip}(t)=\frac{\sigma_{data}^2}{(t-\epsilon)^2+\sigma_{data}^2}$

$\large c_{out}(t)=\frac{\sigma_{data}(t-\epsilon)}{\sqrt{\sigma_{data}^2+t^2}}$

$\sigma_{data}$ 取值 $0.5$ .

Sampling

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

【论文阅读】Consistency Models,论文阅读,AIGC

Training Consistency Models via Distillation

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

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

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

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

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

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

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

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

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

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

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

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

其中 $0\le\mu<1$ . 使用 EMA 可以使训练更稳定，同时能提高模型的表现.
模型训练算法如下：
【论文阅读】Consistency Models,论文阅读,AIGC

Training Consistency Models in Isolation

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

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

consistency training loss (CT)表示为：

$\large \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(·)$ 替换 $\mu$ ，让它随着 $N$ 增长而变化.
$N (\cdot)$ 和 $\mu(·)$ 具体为

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

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

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

训练算法如下：
【论文阅读】Consistency Models,论文阅读,AIGC

Experiment

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

(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(·)$ 效果最好.