consistency model

paper: Consistency Models

看这一篇笔记之前，最好先把 "score-based SDE" 这篇笔记看了。score-based SDE 这篇笔记就是说，现有的生成模型，都可以用 SDE 去描述它们的前向和反向过程。而且发现，反向过程还可以用确定性的 ODE 过程来描述。

这篇 paper 受到反向过程可以用确定性的 ODE 过程来描述这个启发，提出了 consistency model ，具体来说就是训练一个模型，对于任意时间刻的采样 $(x_t,t)$ ，都能直接吐出 $x_0$ 。这是怎么做到的，为什么以前不能做到？

1 引入

以前的反向过程是用 SDE 描述的，这就导致从高级别噪声的分布 $x_{t+1}$ 走到下一级别噪声的分布 $x_t$ 的过程中，路径是不一样的。（还记得 NCSN 的分层式采样吗？每一个噪声级别都会进行 T 次迭代，终点作为下一个噪声级别的起点）

所以，如果让我们发明的模型对于这过程中的每一个步都去映射到 $x_0$ ，这个工作量未免太大了点，因为随机过程的过程千千万万啊。

但是如果用 ODE 描述，过程就确定下来了，此时再对这轨迹上的每一点做映射，难度就小多了。

‍

2 Method

通过 score-based SDE 这篇笔记，我们知道生成模型可以用 SDE 和 ODE 描述，consistency model 这篇 paper 也是一样的，其中，规定了漂流项 $f_t(x)=0$ ，扩散项系数 $g(t)=\sqrt{2t}$ 。那么，该模型的前向过程和反向过程我分别用 SDE 和 ODE 写出来：

嗯？${\nabla }_x\log p(x)$ 是不是很熟悉？是的，它就是 score。我们可以直接先训练一个 score-based model $s_{\theta }(x_t,t)$ 来得到它，所以不妨写为以下形式：

我们的目标是，能拥有一个模型 $f$，使得：$f:(x(t),t)\to x(ϵ),t\in [ϵ,T]$ 。

（不用 0 而是 $ϵ$ 是为了数值稳定性）

那么如何训练这个 $f$ 呢？有两种方式，从已经学好的 scored model 蒸馏，或者从头开始训练。

2.1 Consistency Distillation (CD)

第一种方法很好理解，就是前面说的，用一个训好的 score-based model 去估计 ${\nabla }_x\log p(x)$ 。

那么最直接的想法就是，采样 $(x_t,t)$ ，然后顺着 ODE 概率流求解出 $x_0$ ，然后直接优化 $\Vert f_{\theta }(x_t,t)-x_0{\Vert }^2$。但是这样做的缺点是，每一个 $(x_t,t)$ ，都要跑一遍 ODE Solver 去得到 $x_0$ 后才能算一次 loss，效率太低了。

所以论文的做法，是监督一条轨迹上相邻两个点的输出一致，具体如下：

首先离散化为 N 步，$ϵ=t_1<t_2<\cdots <t_N=T$。然后先均匀采样一个 $t$ ，正向跑出 $(x_{t+1},t+1)$，然后再顺着 ODE 概率流往回跑到 $(x_t,t)$ ，然后优化 $f_{\theta }(x_{t+1},t+1)$ 与 $f_{\theta }(x_t,t)$ 之间的差距。

loss 函数如下：

这里注意，${\theta }^{-}$ 是 $\theta$ 参数的指数滑动平均(EMA)，这种训练方式比起两边都用 $\theta$ 梯度更加稳定，效果显著。

2.2 Consistency Training (CT)

也就是你没有先验模型来估计 score，那么此时你只能从数据中估计 score 。

我们在 score-based model 这篇笔记里推过如下式子：

即可以用 $-\frac{\tilde{x}-x}{{\sigma }^2}$ 来估计 score。

那么接下来同理去优化 loss 就好了。

3 实验

效果上，diffusion > CD > diffusion distillation > CT

但是推理速度 CD, CT 比传统 diffusion 快得多的多。

‍