引言
Diffusion model大获成功,但是它的短板也很明显,需要大量的计算资源,并且推理速度比较慢。如何才能提升Diffusion model的计算效率。业界有各种各样的改进,无疑Latent Diffusion Models(潜在扩散模型,LDMs)是比较成功的一篇,那就来学习一下LDMS是怎么做的吧
论文贡献
1,与基于变换的方法相比,论文的方法在处理更高维度数据,可以高效地应用于高分辨率图像的合成,具体措施如下
- )使用潜在空间进行训练:作者在隐空间而不是像素空间上训练扩散模型。这使得模型可以在更高分辨率的图像上实现高效的图像合成,同时降低计算复杂性。
- )训练自动编码器:首先,作者训练了一个自动编码器,它提供了一个低维(因此高效)的表示空间,该空间在感知上等价于数据空间。这种表示空间使得模型可以更好地处理高维数据,同时保留了图像的细节。
- )在潜在空间中训练扩散模型:在自动编码器的潜在空间上训练扩散模型,使得模型可以在更高分辨率的图像上实现高效的图像合成和条件生成。这种方法在降低计算复杂性的同时,保留了扩散模型的高质量和灵活性。
2,论文在多个任务(无条件图像合成、修复、随机超分辨率)和数据集上实现了竞争性的性能,同时显著降低了计算成本。与基于像素的扩散方法相比,论文还显著降低了推断成本。
3,论文发现,与先前同时学习编码器/解码器架构的方法相比,论文的方法不需要精细权衡重建和生成能力。这确保了极其准确的重建,并且对潜在空间的正则化要求非常低。
4,此外,论文设计了一种基于交叉注意力的通用条件机制,实现了多模态训练。
- )引入交叉注意力层:作者在模型架构中引入了交叉注意力层,使得潜在扩散模型(LDMs)能够处理各种条件输入,例如文本或边界框。
- )条件潜在扩散模型:作者将交叉注意力层与 UNet 结合,以实现条件图像生成。这使得模型可以在各种条件图像合成任务中实现高效的图像生成。
- )文本到图像合成:作者在 LAION-400M 数据集上训练了一个基于文本提示的条件 LDM。他们使用 BERT-tokenizer 对文本进行编码,并实现 τθ 作为 Transformer,以从文本中推断出潜在编码,然后通过(多头)交叉注意力将其映射到 UNet。
- )多模态训练:作者在多个数据集上训练了模型,以实现类别条件、文本到图像和布局到图像的生成。这种方法可以在不同的任务上获得高度竞争性的性能,而无需针对特定任务的架构。
5,最后,论文在https://github.com/CompVis/latent-diffusion上发布了预训练的潜在扩散和自编码模型。
具体方法
1,结构整体理解
Latent Diffusion Models整体框架如图,分为左中右三个部分,左边是一个训练好的自编码模型(AutoEncoder,包括一个编码器 E 和一个解码器 D ),中间是个Diffusion Models,右边是个condition模块。
首先为什么要加入左边红色这个模块,我们可以这么想,Diffusion Models的问题在于耗时和耗费资源,那么最耗资源的部分是哪里,肯定是从清晰图像到相对模糊的这部分,因为这时候要从图像的每个像素点开始计算,占用了大量的资源和时间,
解决的方法也很直观,我们先把像素点压缩合并一下,大图像变小图像,然后从小图像开始做diffusion,因为压缩和解压缩的这个过程远比diffusion过程来的快和节省资源,所以整个过程也就一下子压缩节省了大量资源和时间。
同理我们把像素点压缩的方法替换成编码器和解码器,小图片理解成隐空间,对这就是Latent Diffusion Models的做法了。
其实到上一步Latent Diffusion Models主体工作已经可以完成,不过作者觉得可以无条件扩散生成,也可以有条件引导模型进行更好的扩散。如图1右边部分。通过使用交叉注意力机制将DM的底层UNet骨干增强为更灵活的条件图像生成器,整体是不是很自然的一个思路。
好那让我们再具体看看Latent Diffusion Models是怎么做的。对应图左中右三部分,分别为图片感知压缩(Perceptual Image Compression),潜在扩散模型(Latent Diffusion Models),条件机制(Conditioning Mechanisms)
2,图片感知压缩(Perceptual Image Compression)对应图最左边AutoEncoder部分
感知压缩的扩散模型的训练本质上是一个两阶段训练的过程的第一阶段,第一阶段训练一个自编码器
在第一阶段训练自编码器时,为了避免潜在表示空间出现高度的异化,作者使用了两种正则化方法,一种是KL-reg,另一种是VQ-reg。
具体地说,给定一个在RGB空间中的图像x∈R H×W×3,编码器E将x编码为一个潜在表示z = E(x) 编码图像,解码器 D 从潜在空间重构图像,得到 x˜ = D(z) = D(E(x)),其中 z ∈ R^(h×w×c)
3,潜在扩散模型(Latent Diffusion Models)对应图1中间部分
潜在扩散模型对应的是两阶段训练的过程的第二阶段,针对第一阶段得到的隐空间进行diffusion操作
回忆一下Diffusion Models的公式
LDM = Ex,ε∼N(0,1),th [ (ε - εθ(xt, t))^2 / 2 ]
我们把编码器 E 和一个解码器 D代入公式得到
LDM = E[x, ε ∼ N(0,1), t]((ε - ε_θ(z_t, t))^2 / 2)
好我们就得到了LDM的训练公式了
4,条件机制(Conditioning Mechanisms)对应图1右边部分
具体来说,引入一个特定领域的编码器τθ,该编码器将y投影到中间表示τθ(y)∈RM×dτ,然后通过交叉注意力层将其映射到UNet的中间层,该层算法实现如下
Attention(Q, K, V) = softmax QKT √ d · V,
其中
Q = W (i) Q · ϕi(zt),
K = W (i) K · τθ(y),
V = W (i) V · τθ(y)
好,到这里LDM的算法部分就结束了,是不是还挺直观的。