【论文阅读】Latent Consistency Models (LDMs)、LCM-LoRa-Toy模板网

这篇具有很好参考价值的文章主要介绍了【论文阅读】Latent Consistency Models (LDMs)、LCM-LoRa。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

Introduction

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

Preliminaries

Diffusion Models

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

$\large \frac{\mathrm{d}x_t}{\mathrm{d}t}=f(t)x_t+\frac{g^2(t)}{2\sigma_t}\epsilon_{\theta}(x_t,t)$

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

$\large\hat{\epsilon_{\theta}}(z_t,\omega,c,t)=(1+\omega)\epsilon_{\theta}(z_t,c,t)-\omega\epsilon_{\theta}(z_t,\varnothing,t)$

Consistency Models

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

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

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

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

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

Latent Consistency Models

Consistency Distillation in the Latent Space

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

$\large \frac{\mathrm{d}z_t}{\mathrm{d}t}=f(t)z_t+\frac{g^2(t)}{2\sigma_t}\epsilon_{\theta}(z_t,c,t)$

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

$\large f_{\theta}(z,c,t)=c_{skip}(t)z+c_{out}(t)\left(\frac{z-\sigma_t\hat{\epsilon}_{\theta}(z,c,t)}{\alpha_{t}} \right)$

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

$\large \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}}$ 的估计，计算方法如下：

$\large \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}\epsilon_{\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，损失函数可以表示为：

$\large \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}$ 的计算方法更新为

$\large \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训练的损失函数为

$\large \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(·,·,·,·)$ 的计算方法对应跨 $k$ 步，作者分别使用了DDIM、DPM-Solver、DPM-Solver++ 作为 PF-ODE solver，以DDIM为例，其对应的 $\Psi(·,·,·,·)$ 计算方法为

$\large \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)\hat{\epsilon}_{\theta}(z_{t_{n+k}},c,t_{n+k})-z_{t_{n+k}}$

再加入CFG和skipping-step之后，LCM的训练过程用如下算法所示：
【论文阅读】Latent Consistency Models (LDMs)、LCM-LoRa,论文阅读,AIGC

多步采样算法如下：
【论文阅读】Latent Consistency Models (LDMs)、LCM-LoRa,论文阅读,AIGC

Latent Consistency Fine-tuing for Customized Dataset

全量微调算法：
【论文阅读】Latent Consistency Models (LDMs)、LCM-LoRa,论文阅读,AIGC

Experiment

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

【论文阅读】Latent Consistency Models (LDMs)、LCM-LoRa,论文阅读,AIGC

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

训练时间 32 A100 GPU Hours

LCM-LoRA

【论文阅读】Latent Consistency Models (LDMs)、LCM-LoRa,论文阅读,AIGC

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