Consistency Models 阅读笔记

简介

Diffusion models需要多步迭代采样才能生成一张图片，这导致生成速度很慢。一致性模型（Consistency models）的提出是为了加速生成过程。
Consistency models可以直接一步采样就生成图片，但是也允许进行多步采样来提高生成的质量。
Consistency models可以从预训练的扩散模型蒸馏得到，也可以作为独立的生成模型从头训练得到。

PF ODE

论文中考虑的PF ODE(Probability Flow Ordinary Differential Equation)形式如下：
$$\frac{d{\mathbf{x}}_t}{dt}=-t{s}_{\varphi }({\mathbf{x}}_t,t)$$其中$s_{\varphi }({\mathbf{x}}_t,t)\approx \nabla \log p_t(\mathbf{x})$是分数函数，$t\in [0,T]$。
从初始分布${\hat{\mathbf{x}}}_{\mathbf{T}}\sim \mathcal{N}(0,T^2\mathbf{I})$中采样，然后逆向求解ODE，得到的${\hat{\mathbf{x}}}_0$是近似服从数据分布的样本。值得注意的是，为了保证数值稳定，在本文中用${\hat{\mathbf{x}}}_{ϵ}$当做最后的近似样本，$ϵ$是一个接近0的小正数。

一致性模型（Consistency models）

给定一个PF ODE(Probability Flow Ordinary Differential Equation) { x t } t ∈ [ ϵ , T ] \{\mathbf x_t\}_{t\in[\epsilon, T]} {xt​}t∈[ϵ,T]​，一致性函数（consistency function）被定义为$f:({\mathbf{x}}_t,t)\to {\mathbf{x}}_{ϵ}$，其中$ϵ$是一个接近0的小正数，是ODE求解器停止的位置。一致性函数具有self-consistency性质，即对于PF ODE轨迹上的任意点输出都是一样的。一致性模型$f_{\theta }$的是从数据中估计的一致性函数。

给一个训练好的一致性模型$f_{\theta }(\cdot ,\cdot )$，可以通过一致性模型一步生成了结果：首先从初始分布中采样${\hat{\mathbf{x}}}_{\mathbf{T}}\sim \mathcal{N}(0,T^2\mathbf{I})$，然后用一致性模型计算${\hat{\mathbf{x}}}_{ϵ}=f_{\theta }({\hat{\mathbf{x}}}_{\mathbf{T}},T)$。也可以调用一致性模型多次生成更准确的结果，如Algorithm 1所示，迭代的去噪和添加噪声。

Consistency models的训练算法有两种，一种是从预训练的扩散模型蒸馏（Algorithm 2），一种是作为独立的生成模型从头训练（Algorithm 3）。

感觉一致性模型和EDM¹有共同之处，在每一步中都想恢复出${\mathbf{x}}_0$，但是一致性模型训练时优化的目标是self-consistency性质，通过self-consistency性质来在保证每一步中都能直接恢复出${\mathbf{x}}_0$。