“稳定扩散”是如何工作的?直观解释
稳定扩散是如何工作的?直观解释
“稳定扩散”模型,通常称为“潜在扩散模型”,已经席卷全球,像Midjourney这样的工具吸引了数百万人的关注。在本文中,我将尝试解开关于这些模型的一些神秘之谜,并希望在你的脑海中描绘出它们工作的概念图像。和往常一样,本文不会深入细节,但我会在末尾提供一些相关链接。本文的主要信息来源是一篇论文,将作为第一个链接提供。

动机
图像合成有几种方法(从头开始创建新图像),包括GAN(生成对抗网络),在多样化数据上表现不佳;自回归Transformer,训练和执行速度较慢(类似于LLM Transformer逐个标记生成文本,它们逐个补丁生成图像);以及扩散模型,它们本身可以解决其中一些问题,但计算成本仍然很高。训练这些模型仍然需要数百个CPU天,并且实际利用模型需要逐步运行以生成最终图像,这也需要相当长的时间和计算资源。
扩散模型简介
为了理解扩散模型的工作原理,让我们首先看一下它们是如何训练的,这种方法有点非直观。我们首先反复对图像应用噪声,从而创建了一系列图像的“马尔可夫链”。通过这种方式,我们能够从一个原始图像中获得一些数量为T的越来越嘈杂的图像。
然后,我们的模型学习预测在某个时间步骤应用的确切噪声,并且我们可以使用其输出来“去噪”该时间步骤的图像。这实际上使我们能够从图像T到图像T-1。再强调一遍,模型是通过给它带有在某个时间T应用噪声的图像以及时间T本身来训练的,输出是将它从时间T-1带到T的噪声!

一旦我们训练了这样一个模型,我们就可以反复将其应用于随机噪声,以生成一个新的、独特的图像。上面的图像来自Steins的一篇精彩文章,可以在那里找到更详细的解释。
UNet模型
常用于预测每个时间步骤噪声的模型是“UNet”架构模型。这是一种类型的架构,它重复应用卷积层、池化层和跳跃连接,首先对图像进行降采样但增加深度(特征图),然后使用反卷积将特征图上采样回原始图像尺寸。这里有一篇由Maurício Cordeiro撰写的很好的文章,对这个模型进行了更深入的解释。
问题
这就是传统扩散模型的问题所在。第一个问题是训练时间。假设我们有N个图像,并且对图像应用噪声T次,那就有N*T个可能的输入到我们的模型中。而且通常这些是高分辨率图像,每个输入都会包含大尺寸的图像。毕竟,我们不希望Midjourney生成像素艺术品…
然后,假设我们已经训练好了我们的模型,要从随机噪声恢复到图像,我们必须反复应用权重T次!记住,我们的模型只能给出前一个时间步骤的图像,我们必须从时间步骤T回到时间步骤0。
另一个问题在于这种模型的实用性。如果你注意到,输入文本还没有被提到;然而,我们都习惯于将文本转换为图像,而不是将随机噪声转换为图像。那么,这个特性到底是如何发挥作用的,它在哪里发挥作用呢?
潜在扩散模型
解决这些问题的一个关键思想是将我们的模型分成两个独立的模型。
第一个模型被训练成将图像编码和解码为“潜在空间”,该空间保留了图像背后的许多“知觉”细节,但减少了数据维度。为了直观地理解这一点,我们可以认为在图像中有一些噪声,这些噪声我们实际上不需要学习(如像素略微改变色调的蓝天)。
第二个模型是实际的扩散器,它可以将这个潜在空间表示转换为图像。然而,这个扩散器有一个特殊的修改,使其能够“理解”并受其他领域(如文本)的输入指导。
感知图像压缩
让我们从第一个模型开始——图像编码器/解码器。记住,这个想法是以这样一种方式准备图像,以便实际的扩散器保留重要信息的同时减少维度。
这个压缩模型实际上被称为“自编码器”,同样这个模型学会了将数据编码成压缩的潜在空间,然后再解码回其原始形式。这样的模型的目标是最小化输入和重构输出之间的差异。对于我们的损失函数,我们使用两个组成部分。
- 感知损失——我们的目标是最小化图像特征(如边缘或纹理)之间的差异,而不是原始图像和解码版本之间的像素差异。有现有的工具可以从图像中推断这些特征,我们可以使用这些工具。
- 基于补丁的对抗目标——一个GAN(详见我之前关于不同类型模型的文章),其目标是通过逐个补丁地分析图像并分类判断某个补丁是真实还是伪造,以强制执行局部现实主义。
这样做的额外好处是避免模糊,因为我们优化的不是像素差异,而是特征差异。这意味着如果我们的模型生成的图像在像素着色上略有不同,但仍保留相同的“特征”,它不会受到太多惩罚,这正是我们想要的。然而,如果像素着色彼此接近,但特征不对,模型将受到相当大的惩罚(因此出现模糊)。有效地防止了解码器(以及最终的扩散器)“作弊”,创建具有接近原始数据集的像素值但实际特征“不对”的图像。
文本
那么文本在这里的作用是什么以及如何发挥作用?
在训练我们的模型时,我们不仅训练它扩散图像(或更准确地说是由我们的第一个模型产生的这些图像的潜在空间表示),我们还训练它理解文本并学会使用文本的某些部分来生成这些图像。这是因为“交叉注意力”机制,允许模型有选择地关注某些文本特征或方面,已经内置在模型本身中。下面是一个展示这一点的图表。

通过一个独立的编码器,文本被转换为“中间表示”(一种表示方法,在这种表示方法中,文本的重要部分具有更高的权重),而这些权重会影响UNet的层,从而指导生成最终的图像。这是因为模型内置了“交叉注意力”组件,类似于transformers,允许中间表示的不同部分对图像的某些部分产生较小的影响,对其他部分产生较大的影响。当然,这是通过标准的前向传播和反向传播方案进行训练的。我们的去噪UNet接收文本和一些随机噪声(都被编码),并将其最终输出与原始图像进行比较。该模型既学会了有效地使用文本,也学会了有效地去噪。如果您对前向传播和反向传播不熟悉,我强烈建议您深入了解这些主题,因为它们对机器学习至关重要。
结论
LDM在各种任务中表现出色,包括修复、超分辨率和图像生成。这个强大的模型已经成为AI图形领域的一个重要组成部分,我们仍然在发现新的用例。虽然本文应该给您一个概念性的理解,但其中还有更多的细节。如果您想了解更多信息,请查阅名为“High-Resolution Image Synthesis with Latent Diffusion Models”的论文,其中描述了我提到的许多内容的数学基础。
来源
- https://openaccess.thecvf.com/content/CVPR2022/papers/Rombach_High-Resolution_Image_Synthesis_With_Latent_Diffusion_Models_CVPR_2022_paper.pdf
- https://medium.com/@steinsfu/diffusion-model-clearly-explained-cd331bd41166
- https://medium.com/analytics-vidhya/creating-a-very-simple-u-net-model-with-pytorch-for-semantic-segmentation-of-satellite-images-223aa216e705