Simplestory's Blog

LDM

Latent Diffusion Models

Word count: 2.3kReading time: 9 min
2024/08/17

以往的diffusion model(DM)一般会在全图范围内进行操作,这样在运行DM模型所需的显存会变得非常大。而latent diffusion model(LDM)则是将输入通过预训练好的autoencoder encoder映射到潜在空间,在潜在空间上进行DM操作,最后得到的输出再通过autoencoder decoder部分映射回原图空间。作者还通过cross attention的方式将限制条件引入DM中,从而实现可控的图像生成。

Autoencoder

Autoencoder是一类神经网络,采用无监督的方式进行学习。训练方式是生成的图片与原图做比较,即关注的是模型重建的能力,AE模型也不具备生成能力。该网络包含两部分:

  • Encoder:将原始high-dimension输入数据转换为latent low- dimension code。输入尺寸大于输出尺寸。
  • Decoder:将上面Encoder得到的latent low-dimension code恢复为原始的输入数据。

在LDM中,作者为了避免latent空间很可能出现的大方差使用了两种可选的正则化方式\(\text{KL-reg}\)\(\text{VQ-reg}\)。损失函数使用的是感知损失函数(perceptual loss)以及包含对抗训练过程的Discriminator loss。

  • KL-reg:该惩罚项下的自编码模型称为AutoencoderKL,可以说是标准的VAE(带KL损失)。该模型是通过对学习的隐变量添加一个较小的KL惩罚,以使隐变量接近标准正态分布。
  • VQ-reg:这个自编码模型可以视为VQ-GAN,但是向量量化层(vector quantized layer)被放进decoder中。

Loss function

LDM中的感知损失是利用可学习的感知图像块相似度(LPIPS,Learned Perceptual Image Patch Similarity)来衡量两张图片之间的差别。而Discriminator loss则是借鉴了taming transformer的训练过程。

LPIPS

在计算图像相似度中,传统方式如L2、SSIM和PSNR与人类认为的会有较大的出入,特别是当图片平滑时,传统方式大概率会失效。作者使用了深度学习网络特征来度量图像相似度,并通过大量的实验验证了该设计的有效性。这里作者引入了两个新的评价指标:

  • 2AFC(two alternative forced choice):给实验员一张原始图像与两张经过变化的图像,然后判断哪一张更接近原始图像。这种评价指标有个问题就是实验员可能会有意识地选择自己关注哪些方面的相似度。
  • JND(just noticeable difference):先后给实验员两张图片,每张展示1秒,中间有250毫秒的间隔,然后让实验员判断两张图片是否相同。只让实验员判断两张图是否相似一定程度上避免了两个图像都部分相似的时候实验员选择的主观性。

LPIPS的结构如下图所示:

  • 将两个patch \(\mathbf{x}\)\(\mathbf{x}_0\)分别输入到特征提取模型\(\mathcal{F}\)中,\(\mathcal{F}\)可以是VGG、AlexNet等网络(在LDM中是VGG-16)
  • 提取\(\mathcal{F}\)中不同层(准确说应该是不同阶段)的输出并沿channel维度做归一化得到\(\hat{y}^l, \hat{y}^l \in \mathbb{R}^{H_l\times W_l\times C_l}\),其中\(l\)表示提取的是第\(l\)
  • 使用\(w^l\in\mathbb{R}^{C_l}\)\(\hat{y}^l\)\(\hat{y}^l\)按channel进行缩放,之后计算它们的L2距离。最后取平均值即可得到结果

公示描述如下:

\[ d(x,x_0)=\sum_l\frac{1}{H_lW_l}\sum_{h,w}\vert\vert w_l\odot(\hat{y}^l_{hw} - \hat{y}^l_{0hw})\vert\vert_2^2 \]

当对\(\forall l\),有\(w_l=1\),上式可以看为余弦距离。

训练时如上图,主要是把两个距离作为小型网络\(\mathcal{G}\)的输入计算得分\(\hat{h}\in(0,1)\),然后用交叉熵损失进行训练,图中的\(h\)即是作者引入的新指标,这里用的是2AFC。

在LPIPS论文中,作者给出了三种模型训练方式:

  • lin: 固定\(\mathcal{F}\)预训练权重,只训练\(w\)权重
  • tune: 加载\(\mathcal{F}\)预训练权重,并和\(w\)一起进行训练
  • scratch: 使用随机高斯分布来初始化模型权重,重新开始训练

Discriminator loss

使用PatchGAN的判别器获取图片patch级别的真实性评价并利用对抗损失函数进行优化。参考VQ-GAN,其损失函数大致可以表示如下:

\[ L_D = [\log(D(x)) + \log(1-D(\hat{x}))] \]

其中\(D\)是判别器,\(x\)是真实样本,\(\hat{x}\)则为生成样本。这部分的详细内容可以参考VQ-GAN。

Conditoinal Denoiser

DM部分与DDPM类似,也是由UNet来预测噪声。Latent Diffusion Model和Stable Diffusion都采用了带time-conditional的UNet结构来预测每一步的噪声,但官方文件并没有给出详细的UNet结构图。这里参考了另外一篇论文,获得time-conditional UNet的大致结构如下:

注意:该结构是基于SD2.1版本的,且使用的条件约束是text condition。跟SD1.5相比,主要更改的点是将text encoder从CLIP ViT-L/14替换为OpenCLIP-ViT/H,其余基本一致。

从上图可看出,与原始UNet相比,主要是将原本的卷积结构替换为了Cross Attention Down BlockDown BlockCross Attention Mid BlockUp BlockCross Attention Up Block模块以及添加了time embedding和condition embedding(在上图中为text embedding)作为额外输入。

额外输入

首先是time embedding。依据diffusion model的前向过程,我们会有一个时间步(timestep),而UNet的目标是预测前向过程中每个timestep所添加的随机噪声,所以timestep对UNet来说是一个必要的信息。这里参考了transformer对timestep做了正弦位置编码(Sinusoidal Positional Embedding),之后经过两层线性层得到time embedding。

condition embedding则需要依据不同类型的条件约束使用不同的网络提取特征得到condition embedding。如图中是文本类型,对应使用的模型是OpenCLIP。

额外模块

对于Cross Attention Down BlockCross Attention Up Block模块。其中Cross Attention Down Block的结构如下:

Cross Attention Up Block的结构如下:

上面两个模块中,time embedding都是只与ResNet Block做操作,而condition embedding则是通过交叉注意力参与到Transformer Model中。

ResNet Block详细结构如下。time embedding会通过线性层映射到合适的尺寸以便与卷积层的输出进行逐元素相加。最后的\(1\times 1\)卷积是可选项,目的是用于调整输出特征图的通道数。

Transformer Model详细结构如下。该结构可以分为三部分。第一部分将输入的特征图转为一维向量输出到transformer中。第二部分则是一个常规的transformer,其中第二个Multi-Head Attention中的K,V值由condition embedding产生。第三部分则是将transformer输出的一维向量转换为与输入相同尺寸的特征图。

其余模块(Up Block、Down Block和Cross Attention Mid Block)的结构如下:

Appendix

Latent diffusion model与Stable diffusion model的区别(两个模型结构上基本是一样的):

  • 相比Latent diffusion的训练数据(LAION-400M),stable diffusion使用了更大量的数据(LAION-2B-en)进行训练,同时数据也有经过初步筛选,质量会高一些
  • Latent diffusion model中使用的text encoder是随机初始化的transformer(即从零开始训练),而stable diffusion model则是通过预训练的CLIP模型
  • 训练的尺寸也有所不同。例如Latent diffusion model只是在\(256\times 256\)的分辨率上训练,而stable diffusion model会先在\(256\times 256\)上预训练,之后在\(512\times 512\)上微调

更详细的回答可以参考这里

Stable diffusion各版本的区别。自stable diffusion发布以来,官方已经推出了许多版本的SD模型,但这些模型原理结构上的改变其实并不大,更多的可能是生成风格上的差别。

  • SD 1.x 版本的text encoder采用的是OpenAI的CLIP ViT-L/14模型,其模型参数量为123.65M;而SD 2.0采用了更大的text encoder:基于OpenCLIP在laion-2b数据集上训练的CLIP ViT-H/14模型,其参数量为354.03M
  • 在text encoder中,SD 1.x 提取的是倒数第一层的特征。由于倒数第一层的特征之后就是CLIP的对比学习任务,所以倒数第一层的特征可能部分丢失细粒度语义信息。SD 2.0提取的是text encoder倒数第二层的特征
  • SD 2.0不同stage的attention模块是固定attention head dim为64,而SD 1.x 则是不同stage的attention模块采用固定attention head数量,明显SD 2.0的这种设定更常用,但是这个变动不会影响模型参数
  • 训练数据上,SD 1.x 版本其实最后主要采用LAION-2B中美学评分为5以上的子集来训练,而SD 2.0版本采用评分在4.5以上的子集,相当于扩大了训练数据集

更详细内容可以参考这里

参考

High-Resolution Image Synthesis with Latent Diffusion Models

The Unreasonable Effectiveness of Deep Features as a Perceptual Metric

The Big Data Myth: Using Diffusion Models for Dataset Generation to Train Deep Detection Models

Latent Diffusion(Stable-Diffusion) Model (video)

Latent Diffusion(Stable-Diffusion) Model (pdf)

The Illustrated Stable Diffusion

From Autoencoder to Beta-VAE

A Beginner's Guide to Variational Methods: Mean-Field Approximation

CATALOG
  1. 1. Autoencoder
    1. 1.1. Loss function
      1. 1.1.1. LPIPS
      2. 1.1.2. Discriminator loss
  2. 2. Conditoinal Denoiser
    1. 2.1. 额外输入
    2. 2.2. 额外模块
  3. 3. Appendix
  4. 4. 参考