diffusion model（五）stable diffusion底层原理（latent diffusion model， LDM）

LDM: 在隐空间用diffusion model合成高质量的图片！

[论文地址] High-Resolution Image Synthesis with Latent Diffusion Models

[github] https://github.com/compvis/latent-diffusion

系列阅读

• diffusion model（一）DDPM技术小结 (denoising diffusion probabilistic)
• diffusion model（二）—— DDIM技术小结
• diffusion model（三）—— classifier guided diffusion model
• diffusion model（四）文生图diffusion model（classifier-free guided）
• diffusion model（五）stable diffusion底层原理（latent diffusion model， LDM

1 背景

近期扩散模型（diffusion model，DM）在图像生成取得了最先进的结果。但是传统的扩散模型是在像素空间(pixel space)进行优化的，高分辨率图像的训练往往需要更大的显存，更多的梯度回传，导致DM模型训练成本高，推理时延长。当下很多工作从优化采样策略¹²、分层方法³来解决这两个缺点，但效果有限。latent diffusion model提供了一个新的思路：它将扩散过程从传统的像素空间转到了隐空间（DM模型输入参数量大大降低），极大提升了DM的训练效率和推理效率，使得DM模型得以在单张消费级显卡应用，降低了AI图片生成的上手成本。目前火出圈的Dalle2⁴, Stable Diffusion⁵都用到了LDM技术。下面我们来看LDM是如何做的。

2 方法

2.1 整体架构

LDM主要参考了VQGAN⁶的思路，其整体架构如下图所示。与传统Diffusion model在像素空间(pixel-based diffusion model)重建不同的是，LDM是在隐空间进行重建。为了得到图片的隐空间表征，LDM需要预选训练一个VAE模型。记原始图像为$x,x\in {\mathbb{R}}^{H\times W\times 3}$通过VAE的编码器$\mathcal{E}$获得图片的隐空间表示$z,z\in {\mathbb{R}}^{\frac{H}{f}\times \frac{W}{f}\times c}$， $f$为下采样率。在重建阶段，首先通过DM预测隐空间的表征$\tilde{z}$，再将其送入到VAE的解码器中$\mathcal{D}$重建像素空间的表征$\tilde{x}$。本质上LDM是一种二阶段的图片生成方法。

通过在隐空间重建的操作，相较原本的像素空间重建，一个时间步的计算量近似降低了近$f^2$倍。有同学会问通道数不是从$3\to c$了吗，为什么整体计算量没有扩到$\frac{c}{3}$倍，其实这只会增加第一层卷积的计算量，相较整体模型而言较低。

原本DM模型的优化目标
$$L_{DM}={\mathbb{E}}_{x,ϵ\sim \mathcal{N}(0,1),t}[\parallel ϵ-{ϵ}_{\theta }(x_t,t){\parallel }_2^2]$$
LDM的优化目标
$$L_{LDM}={\mathbb{E}}_{\mathcal{E}(x),ϵ\sim \mathcal{N}(0,1),t}[\parallel ϵ-{ϵ}_{\theta }(z_t,t){\parallel }_2^2]$$
$t$在 { 1 , . . . , T } \{1, ..., T\} {1,...,T}以服从均匀分布进行采样

通过上面的描述不难看出LDM就是套了一层VAE的DM！利用VAE搭建起隐空间和像素空间的桥梁。

2.2 更多细节

2.2.1 感知压缩的权衡

上文提到LDM先通过VAE的encode将图片从像素空间$x,x\in {\mathbb{R}}^{H\times W\times 3}$变换到隐空间$z,z\in {\mathbb{R}}^{\frac{H}{f}\times \frac{W}{f}\times c}$。当下采样率$f$越大，DM的过程越快，但f越大可能导致丢失过多的高频信号，影响重建结果。因此需要对不同的下采样率进行权衡。论文中分别对 f ∈ { 1 , 2 , 4 , 8 , 16 , 32 } f \in \{1,2,4,8,16,32\} f∈{1,2,4,8,16,32}进行了试验。基于试验结果，作者建议采样率控制在$4-16$能够取得效率和质量的均衡。

从下图可以看出，当采样率为4-16时，LDM的训练收益比最大。

2.2.2 LDM的训练策略与预测

LDM采用了二阶段的训练策略(two-stage)。需要先训练VAE模型再训练DM模型。DM训练时会利用已训练好的VAE的encode将数据从像素空间切换到隐空间。

预测时，在隐空间随机采样一个噪声，用DM进行去噪。最后将DM的预测结果传给VAE的Decode进行解码。

2.2.3 给生成过程引入控制信号

如果说不带控制信号的LDM是对$p(z)$建模的话，那么加上控制信号的建模可以表述成$p(z|y)$。控制信号$y$可以是文本、图片布局、轮廓图等。当有有一些更为细粒度的生成存在多个控制信号，即对$p(z|y_1,\cdots ,y_N)$。为了使得生成过程考虑控制信息，作者在原有的Unet backbone上引入了一个交叉注意力机制，来融入控制信号。首先通过一个$\text{domain specific encoder}{\tau }_{\theta }$将$y$投影到为一个中间表征${\tau }_{\theta }(y)\in {\mathbb{R}}^{M\times d_{\tau }}$,随后和DM某层的输出进行融合。

假定${\phi }_i(z_t)$是隐变量$z$在Unet在层$i$时间步为$t$时的输出，下面需要将${\phi }_i(z_t)$与${\tau }_{\theta }(y)$用交叉注意力机制融合
$$\begin{array}{rl}\mathrm{Attention}(Q,K,V)& =\mathrm{softmax}(\frac{Q{K}^T}{\sqrt{d}})\cdot V\\ Q=W_Q^{(i)}\cdot {\phi }_i(z_t);K& =W_K^{(i)}\cdot {\tau }_{\theta }(y);V=W_V^{(i)}\cdot {\tau }_{\theta }(y)\end{array}$$
此时模型的优化目标为
$$L_{LD{M}_{cond}}={\mathbb{E}}_{\mathcal{E}(x),ϵ\sim \mathcal{N}(0,1),t}[\parallel ϵ-{ϵ}_{\theta }(z_t,t,{\tau }_{\theta }(y)){\parallel }_2^2]$$

参考文献