隐空间扩散模型（LDM）学习笔记

前言

这段时间在项目中接触到了LDM，很早之前了解过GAN，感觉二者还是有很大的不同，因此阅读了论文原文以及网上的公开介绍，结合自己的理解总结了一下，方便后续查看。如涉及到其他网络著作权，可联系删除。
这篇笔记主要围绕隐空间扩散模型优势、稳定扩散模型原理、隐空间机制、生成模型部署训练以及图像生成等方面进行阐述说明。

论文原文： High-Resolution Image Synthesis with Latent Diffusion Models
GITHUB代码链接：https://github.com/CompVis/latent-diffusion

1. 隐空间扩散模型的优势

图像合成是最近发展最引人注目的计算机视觉领域之一，也是计算需求最大的领域之一。特别是复杂自然场景的高分辨率合成，目前主要是基于似然的模型，自回归（autoregressive, AR）transformers 中可能包含数十亿个参数。相比之下基于GAN的图像合成主要局限在了可变性相对有限的数据，其对抗学习过程难以扩展到建模复杂的多模态分布。最近，根据去噪自编码器（denoising autoencoders）层次结构构建的扩散模型已在图像合成及其他方面取得了可观的成果，同时取得了类条件图像合成（class-conditional image synthesis）和超分辨率（super-resolution）中的最先进成果。除此之外，作为基于似然的模型，它们不会像GAN那样存在模式崩溃和训练不稳定性问题；并且，通过大量利用参数共享，它们可以对自然图像的高复杂分布进行建模，而无需像AR模型那样需要数以亿计的参数。
然而，普通的基于似然的模型，如DM，为了在训练和评估过程中对微小细节进行建模，其需要在RGB空间反复进行函数估计和梯度计算，因此消耗了极高的存储空间和计算资源。例如，训练出一个强大的DM通常需要数百个 GPU days，而生成50000个样本则需要在单个A100 GPU上训练大约5天。
隐空间扩散模型则是通过“感知压缩”和“语义压缩”达到感知上等效、但计算上更合适的生成效果。其并没有使用传统的过度空间压缩减少计算消耗，而是利用自编码器的设计生成一个“隐空间”，在这个低维度的空间进行学习训练，减少了计算复杂度。

2. LDM的原理和流程

隐空间扩散模型（Latent Diffusion Model，LDM）分离了压缩学习和生成学习阶段，其结构如下图所示。在设计方面，通过设计使用自编码器，模型的正向扩散和反向生成都在从感知上与图像空间等效的“隐”空间中进行学习，抛弃掉不必要的图像细节同时降低了计算复杂度；利用从UNet 架构继承的DM的归纳偏置，其对于具有空间结构的数据特别有效（隐空间同样有效），因此减轻了对于激进的、降低质量的压缩级别的需求；通过引入交叉注意力机制，使其支持图像、文本等条件监督的图像生成，更加灵活自由。下面简要从感知压缩模型、隐式扩散模型和条件学习机制几个方面介绍LDM的原理和流程。

2.1 感知压缩模型

感知压缩模型由一个感知损失和一个基于补丁（patch）的对抗损失训练的自编码器组成。感知损失可以帮助模型关注高级语义特征而不仅仅限制在像素级特征，而对抗性损失的结合则强调了局部的真实性，通过将重建强制限定在图像流内避免了因依赖像素空间损失（如L2损失）造成的模糊现象。
进一步的，给定输入的RGB图像，编码器将图像编码为隐空间表示，解码器在隐空间重建图像，即
，。更重要的是，编码器按照比例对图像进行下采样，从而减小了图像的尺寸，降低数据维度。同时，使用两类正则化方式避免高方差，一种是KL正则化，另一种是向量量化正则化。KL正则化类似于变分自编码器（VAE），对隐空间施加KL散度惩罚使其接近标准正态分布，从而帮助控制隐空间的变异性，使模型更加稳定。向量量化正则化则是在解码器中使用向量量化层，类似于VQGAN中的量化操作，这有助于模型学习更离散、更有结构的潜在表示。
与其他依赖一维排序进行学习的模型不同，隐空间模型保持了隐空间的二维结构，更好地利用了图像数据的空间结构，从而能够在较小的压缩率下实现高质量的图像重构。在细节方面，相较于以前的工作，这种方法更好地保留了输入图像的细节，提高了重构图像的质量。

2.2 隐式扩散模型

（1）扩散模型（Diffusion Model，DM）

扩散模型是一种概率模型，设计用来学习数据分布。这是通过逐渐去噪一个正态分布变量完成的，相当于学习一个固定马尔可夫链（Markov Chain）的逆过程，其中链的长度为T。在图像合成领域，最成功的扩散模型依赖于变分下界（variational lower bound）的重新加权变体，对应于去噪得分匹配（denoising score-matching）。这些模型可以被解释为一系列等权重的去噪自编码器，通过训练这些自编码器将加噪的输入逐渐去噪，从而得到一张重构生成的全新图像。其中，为输入的噪声版本。
DM对应的目标函数可以被简化为：

其中，从中均匀采样。

（2）隐式扩散模型（Latent Diffusion Model，LDM）

在上述训练好的感知压缩模型（即自编码器）的基础上，可以得到一个高效、低维的隐空间，该空间抛弃掉高频次、不易观察且无用的细节信息，能够良好地专注数据的语义部分，且在低维的空间中训练更加高效。
与依赖于自回归、基于注意力的变换器模型在高度压缩的离散隐空间中的工作不同，在隐式扩散模型LDM中使用到了图像特有的归纳偏置，即利用了图像数据的一些自然特性，比如空间上的连续性和局部区域的相关性。通过由2D卷积层构成的UNet关注图像的空间连续性、使用 “重新加权的边界”关注图像的局部相关性。UNet是一种特别适合图像任务的网络结构，它通过在网络中反复应用卷积操作来捕捉图像的空间层次结构。这样的设计使得模型不仅能够处理图像的高层语义信息，也能精细地处理图像的细节。“重新加权边界”则是一种训练策略，在模型训练时会特别关注那些对人的感知更为重要的图像部分，使得重建的图像在视觉上更加真实和自然。
则LDM的目标函数可以表示为：

神经网络的主干部分为时间条件下的。由于前向扩散过程是固定的，因此在训练期间可以高效地获取，通过解码器D将来自的样本单次传递解码到图像像素空间。

3. 条件隐式扩散模型

在上述章节，扩散模型学习了数据分布，则原则上也可以学习模拟条件分布，其中 y 可以是文本、语义地图或其他图像到图像的转换任务的输入。这种条件分布的建模能力使得扩散模型能够依据这些外部输入控制图像的合成过程。
（1）条件自编码器
实现条件分布的一种方式是使用条件去噪自编码器。这种自编码器根据输入 y 和时间步 t 以及当前的潜在状态来预测去噪后的数据。
（2）引入交叉注意力机制
为了使LDM更加灵活，对主干去噪网络引入交叉注意力机制。这种机制适用于学习基于多种输入模态的注意力模型，提高了模型处理不同类型输入的能力。
同时，为了处理多模态的输入，引入了一个域特定的编码器，它将输入y投影到一个中间表示上。这个中间表示随后通过交叉注意力层Attention被映射到UNet的中间层。

其中，。表示
的中间层，是投影矩阵。交叉注意力层使用查询Q，键K，值V的机制，其中Q是基于的变换，而K和V是基于y的域特定编码的变换。这样可以有效地将输入y的信息集成到图像生成过程中。
基于“图像-条件”对，条件隐式扩散模型可以表示为：

通过最小化预测的去噪结果和实际的去噪结果，训练上述目标函数，其中编码器和域编码器是联合进行优化的。

4. 一些抱怨

在实现过程中，至少需要12G的显存才能运行，反正我的2080和TITAN是不成了。