LDM：High-Resolution Image Synthesis with Latent Diffusion Models

LDM: High-Resolution Image Synthesis with Latent Diffusion Models

TL; DR：关键的技术点有两个。一是通过感知压缩将 Diffusion 从像素空间的转换到隐空间来做，降低训练/采样的算力成本，提高生成质量。二是通过交叉注意力，将各种条件（如文本、图片、表征等）加入到 UNet 的生成过程中，实现高效的可控生成。

本篇提出了大名鼎鼎的 LDM（以及 Stable Diffusion）。在这之前的扩散模型，都是做在像素空间上的，现如今图片的像素数动辄上百万，而且扩散模型的训练需要多步去噪，在像素空间进行训练成本太高。借鉴 AE、VAE、VQGAN 这一类图像压缩表征的方法，LDM 提出先将图片压缩为隐空间的特征图，然后在隐空间中进行去噪训练，之后再将去噪结果解码到像素空间，完成图像生成。详细介绍的文章已经很多了，这里简要讲笔者认为的两个关键技术点。分别是 VAE 感知压缩和交叉注意力条件机制。

VAE感知压缩

感知压缩是 LDM 最关键的一个点，其名称中的 latent，正是隐空间的意思。在已有的方法中，AE、VAE、VQVAE、VQGAN 都可以作为图像压缩表征与重建的方法。这些方法在之前的文章中都已经介绍过了。简单来说，AE 就是经典的自编码器，输入图片，压缩后重构出原图；VAE 将中间的隐层特征约束为高斯分布，方便直接采样生成；VQVAE 引入了离散的向量表征 codebook，并通过额外训练一个自回归模型在隐空间采样后解码为真实图片；VQGAN 在 VQVAE 的基础上添加了感知损失和对抗损失来避免仅使用 L2 重构损失造成的图片模糊。

LDM 中的压缩感知模型采用了 VQGAN 中的感知损失和对抗损失来优化解码质量。为了避免隐层特征图的方差太大，LDM 试验了两种正则化方法：KL-reg 和 VQ-reg。KL-reg 将隐层特征约束为一个高斯分布，这就完全等价于 vae 了，其编码器的输出会被视作高斯分布的均值和方差，然后从该分布中随机一个 latent 出来，送到解码器中，重构原始图像。可以在 diffusers相关实现 得到验证。这也是目前 Stable Diffusion 最常用的正则化方式。VQ-reg 则在解码器中融合进了一个向量化层，相当于是 VQGAN。将从接近无限的真实 code 中采样，约束为了从有限的 codebook 中采样，也能达到约束隐层特征分布方差的正则作用。不过目前看使用 VQ-reg 的模型远不如 KL-reg 多。

交叉注意力条件机制

在各种生成模型成熟之后，可控生成都是必须要思考的一个问题，因为只有可控的生成才有现实意义。LDM 提出使用交叉注意力机制，将条件注入到 UNet 的去噪过程中，来实现可控生成。

对于每一种条件 $y$，先使用不同的编码器 ${\tau }_{\theta }$ 进行预处理，将其转换为一个隐层表示 ${\tau }_{\theta }(y)\in {\mathbb{R}}^{M\times d_{\tau }}$ ，然后通过交叉注意力层 (cross-attention layer) 映射到 UNet 的中间层中：
$$\text{Attention}(Q,K,V)=\text{softmax}(\frac{Q{K}^T}{\sqrt{d}})\cdot V$$
其中：
$$Q=W_Q\cdot \varphi (z_t),K=W_K\cdot {\tau }_{\theta }(y),V=W_V\cdot {\tau }_{\theta }(y)$$
$\varphi (z_t)$ 是 UNet 的中间层表征，这里要看清楚谁和谁乘，如何交叉注意。

总结

Stable Diffusion 的影响力不用多说，本文主要提了 LDM 技术上的两个关键点。其中 VAE 感知压缩要首先深刻理解一系列自编码压缩方法（之前的文章都有介绍），然后搞懂 KL-reg 和 VQ-reg 两种正则化方式实际上相当于什么方法，以及为什么要加正则。交叉注意力的条件生成机制也是目前比较常用的，理解好交叉注意力机制本身是如何操作的即可，别的就是如何将不同的条件形式进行高效编码。