【论文阅读】Latent Consistency Models (LCMs)、LCM-LoRa

Introduction

• 提出 Latent Consistency Models (LCMs)，图像生成速度更快、质量更好.
• 提出一种简单高效的 one-stage guided consistency distillation 方法，用极少的采样步数蒸馏 Stable Diffusion，进一步提出 skipping-step 技术加快收敛过程.
• 介绍针对 LCMs 的微调方法.

Preliminaries

Diffusion Models

使用 empirical PF-ODE 表示模型的逆扩散过程：

$$\frac{\mathrm{d}{x}_t}{\mathrm{d}t}=f(t)x_t+\frac{g^2(t)}{2{\sigma }_t}{ϵ}_{\theta }(x_t,t)$$

对于 class-conditioned 扩散模型，Classifier-Free Guidance (CFG) 有效地提高了生成样本的质量，用 $\omega$表示 CFG 系数，原始的噪声预测模型可以被替换为：

$$\widehat{{ϵ}_{\theta }}(z_t,\omega ,c,t)=(1+\omega ){ϵ}_{\theta }(z_t,c,t)-\omega {ϵ}_{\theta }(z_t,\varnothing ,t)$$

Consistency Models

令 $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$.
损失函数为：

$${\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]$$

${\theta }^{-}$使用 EMA 更新，计算公式如下：

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

${\hat{\mathrm{x}}}_{t_n}^{\varphi }$是从${\mathrm{x}}_{t_{n+1}}$到${\mathrm{x}}_{t_n}$的估计，计算公式如下：

$${\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 )$$

Latent Consistency Models

Consistency Distillation in the Latent Space

针对类似 Stable Diffusion的隐空间上的条件扩散模型，其 PF- ODE 逆过程可以表示为：

$$\frac{\mathrm{d}{z}_t}{\mathrm{d}t}=f(t)z_t+\frac{g^2(t)}{2{\sigma }_t}{ϵ}_{\theta }(z_t,c,t)$$

其中$z_t$是图像隐向量，$c$是给定的条件. 类似CM中的做法，引入$f_{\theta }:(z_t,c,t)\mapsto z_0$，将其参数化为：

$$f_{\theta }(z,c,t)=c_{skip}(t)z+c_{out}(t)\left(\frac{z-{\sigma }_t{\widehat{ϵ}}_{\theta }(z,c,t)}{{\alpha }_t}\right)$$

具体的参数化形式由被蒸馏的扩散模型决定.
损失函数表示为：

$${\mathcal{L}}_{CD}(\theta ,{\theta }^{-};\Psi )={\mathbb{E}}_{z,c,n}\left[d(f_{\theta }(z_{t_{n+1}},c,t_{n+1}),f_{{\theta }^{-}}({\hat{z}}_{t_n}^{\Psi },c,t_n)\right]$$

${\hat{z}}_{t_n}^{\Psi }$为$z_{t_{n+1}}$到$z_{t_n}$的估计，计算方法如下：

$${\hat{z}}_{t_n}^{\Psi }-z_{t_{n+1}}={\int }_{t_{n+1}}^{t_n}\left(f(t)z_t+\frac{g^2(t)}{2{\sigma }_t}{ϵ}_{\theta }(z_t,c,t)\right)\mathrm{d}t\approx \Psi (z_{t_{n+1}},t_{n+1},t_n,c)$$

One-Stage Guided Distillation by Solving Augmented PF-ODE

使用CFG，损失函数可以表示为：

$${\mathcal{L}}_{CD}(\theta ,{\theta }^{-};\Psi )={\mathbb{E}}_{z,c,n}\left[d(f_{\theta }(z_{t_{n+1}},\omega ,c,t_{n+1}),f_{{\theta }^{-}}({\hat{z}}_{t_n}^{\Psi },\omega ,c,t_n)\right]$$

${\hat{z}}_{t_n}^{\Psi }$的计算方法更新为

$${\hat{z}}_{t_n}^{\Psi }-z_{t_{n+1}}\approx (1+\omega )\Psi (z_{t_{n+1}},t_{n+1},t_n,c)-\Psi (z_{t_{n+1}},t_{n+1},t_n,\varnothing )$$

Accelerating Distillation with Skipping Time Steps

扩散模型例如Stable Diffusion的总时间步长有$1000$步，LCM在训练的采样需要覆盖这$1000$步，既然相邻时间步之间的差值小，那么$f_{\theta }(z_{t_{n+1}},c,t_{n+1})$和$f_{\theta }(z_{t_n},c,t_n)$之间的差距也小，这导致计算出来的损失小、收敛慢.

作者介绍了skipping-step 方法，原来度量时间步$t_{n+1}$和$t_n$间的差距，改为度量$t_{n+k}$和$t_n$间的差距. 至此，LCM训练的损失函数为

$${\mathcal{L}}_{CD}(\theta ,{\theta }^{-};\Psi )={\mathbb{E}}_{z,c,n}\left[d(f_{\theta }(z_{t_{n+k}},\omega ,c,t_{n+k}),f_{{\theta }^{-}}({\hat{z}}_{t_n}^{\Psi },\omega ,c,t_n)\right]$$

${\hat{z}}_{t_n}^{\Psi }$中$\Psi (\cdot ,\cdot ,\cdot ,\cdot )$的计算方法对应跨$k$步，作者分别使用了DDIM、DPM-Solver、DPM-Solver++ 作为 PF-ODE solver，以DDIM为例，其对应的$\Psi (\cdot ,\cdot ,\cdot ,\cdot )$计算方法为

$$\Psi (z_{t_{n+k}},t_{n+k},t_n,c)=\frac{{\alpha }_{t_n}}{{\alpha }_{t_{n+k}}}{z}_{t_{n+k}}-{\sigma }_{t_n}\left(\frac{{\sigma }_{t_{n+k}}{\alpha }_{t_n}}{{\alpha }_{t_{n+k}}{\sigma }_{t_n}}-1\right){\widehat{ϵ}}_{\theta }(z_{t_{n+k}},c,t_{n+k})-z_{t_{n+k}}$$

再加入CFG和skipping-step之后，LCM的训练过程用如下算法所示：

多步采样算法如下：

Latent Consistency Fine-tuing for Customized Dataset

全量微调算法：

Experiment

测试数据集使用 LAION-Aesthetic-6+ 和 LAION-Aesthetic-6.5+，teacher model 是 Stable Diffusion-v2.1.

LCM的推理步数在$1$到$4$步的时候效果会比其他 baseline 方法好. 因为DPM和DPM++算实践中很常用的 ODE Solver，正常使用时推理步数在$20$以上. 所以综合速度和质量，LCM表现不错.

训练时间 32 A100 GPU Hours

LCM-LoRA

原理：在原本的 Latent Diffusion Model (LDM) 中，可以使用 LoRa 训练一个额外结构附加到模型的 TextEncoder 和 Unet 中，做到模型的风格迁移. 即图中所示的 ${\tau }^{\prime }$，它是原模型微调后额外结构的参数向量. LCM的 backbone 和被它蒸馏模型的 backbone 结构是一致的，所以LCD过程也可以视作对原模型的微调，所以也可以利用 LoRa，在初始化 student Unet 之后，整个蒸馏过程只训练 LoRa 引入的额外结构，也就是获得 ${\tau }_{\mathrm{LCM}}$. 理论上可以结合${\tau }^{\prime }$，最终做到既能加速生成，又能风格迁移.

LCD过程仅微调 LoRa，收敛更快，训练消耗显著降低.