文生图鼻祖（stable diffusion）——LDM论文解析

1、综述

今天来看一篇非常有名的论文，名字叫做《High-Resolution Image Synthesis with Latent Diffusion Models》，这篇论文被当年的CVPR收录，是文生图领域开源的最有名的算法stable diffusion的原理。

在这篇论文之前，所有的DM（diffusion model）模型都是在图片像素级别上进行操作，这使得计算开销非常大，很难大规模商用。这篇文章解决的问题就是把计算放在了网络隐层中，使得计算开销大幅度减少，使得最终的商用成为可能。

并且作者引入了条件控制机制，这就使得用文字控制图像生成或者其他类似的任务成为可能。后面著名的开源项目stable diffusion就是基于此构建的。

2、Related Work

老规矩，这部分就不详细说了，感兴趣的朋友自己看一下原文。

3、Method

整体算法的原理其实非常清晰，原理图如下所示：

具体由以下三个部分组成：

3.1 Perceptual Image Compression

上图最左边红色的框，作者称为perceptual compression model 感知压缩模型，是基于《Taming transformers for high-resolution image synthesis》这篇论文，也就是VQ-GAN。VQ-GAN在VQVAE的基础上做了一些改进，而VQVAE最大的贡献是提出了codebook这个概念，也就是把特征向量变为离散的，这里就不具体展开说了。

回到本论文，这里的自编码模型是用感知损失和patch-based的对抗损失训练出来的。具体看图里的红框，给定一张RGB图片x（形状为H*W*3），输入encoder E，得到特征向量z，使用decoder D还原回图片，能使得：

x与还原的 $\widetilde{x}$ 我们希望是一致的。

与此同时，作者还使用了KL-reg和VQ-reg这两种正则化方式来避免潜空间的方差过大。其中KL-reg相当于在潜空间上施加了KL惩罚，使得得到的特征的分布接近正态分布。VQ-reg则是使用了一个向量量化层（vector quantization layer），向量量化层将特征归一化为codebook中最近的那个特征。

3.2 Latent Diffusion Models

接下来我们来看上图中间的这个框，这里就是一个diffusion model（DM，扩散模型），只不过是在latent也就是潜空间中。

扩散模型的原理说来也简单，就是不断地在一个图像中加入高斯噪声，而模型就是学习这个去噪的过程。也就是说扩散模型可以理解为一个时序去噪自编码器。

对于一个正常的DM来说，假设时间步t = 1 . . . T，并且t在{1, . . . , T}上均匀采样，则损失函数为：

但是在这篇论文中，由于是在潜空间中训练，也就是预测zT上添加的噪声，则损失函数变为：

其中的神经网络是一个基于UNet的结构，这个也就是扩散模型中我们需要训练的网络。

上图的中间绿色部分就是在潜变量空间的扩散模型，其中上半部分是加噪过程，用于将特征 z 加噪为 zT 。下半部分是去噪过程，去噪的核心结构就是这个由交叉注意力（Cross Attention）组成的U-Net，用于将 zT 还原为 z 。由于左边红框里的自编码器是训练好的（fixed），所以只需要前向执行一次就能得到z或是 $\widetilde{x}$ ，这使得训练LDM的时候相当高效。

3.3 Conditioning Mechanisms

作者在这里引入了条件控制机制，这也是为什么文字能控制图像的生成，当然不止文字可以作为条件，作者在这里自己举了一些例子：such as text [66], semantic maps [32, 59] or other image-to-image translation tasks。

看上图最右边这个框，文字等条件信息y预先编码成特征向量，并将其通过cross-attention送入到扩散模型的去噪过程中。定义τθ为一个特征提取器，用来提取y的信息成为一个特征向量，例如可以用预训练的BERT转换文本，使用CLIP转换图像等。接着我们使用交叉注意力将条件融入到U-Net的中间层：

这里大概说一下cross-attention，这个注意力机制由QKV矩阵（Query、Key、Value）控制，由于是使用文字控制图片，所以K和V矩阵与y相关，而Q矩阵与z相关。具体细节建议大家详细去了解一下。cross-attention模块相当于在U-Net的降采样部分和上采样部分的残差块之后都添加了一个交叉注意力。

最终加入了条件的LDM损失为：