diffusion model(六)：Dalle2 技术小结

type

status

date

slug

summary

背景

随着AIGC概念及diffusion model在视觉领域的发光发热。OpenAI结合CLIP与diffusion model提出了一个二阶段的文本生成图片范式Dalle21，能够更精细的从语义层面控制图片的生成效果。得益于其训练范式，Dalle2的图片生成的质量不弱于GLIDE2,但生成器的多样性有明显提升。下面我们来看它是怎么做的吧！

Dalle2的总体架构

Dalle2作为一个文到图的生成模型，其输入是文本，其输出是基于该文本对应的图片即。Dalle2将该过程分为两个阶段

其一是DiffusionPrior ，其作用是获得输入文本对应的CLIP图片向量。

其二是DiffusionDecode ，作用是根据输入的文本和估计的图片向量生成图片。下面我们对这两个模块进行详细介绍。

上述两个阶段可用公式(1)进行表述

可能会有读者疑问为什么。这是因为zi是x的CLIP的image embedding，CLIP是一个确定的函数，故相等。

DiffusionPrior模块

Prior模块的作用是给定文本，生成该文本对应的CLIP模型的图片向量。（generate possible CLIP image embedding from a given text caption）

前置知识：

DIffusion model的原理，可参考博客[5]。

CLIP的基本原理[4]。

训练阶段

训练数据：图片文本对

对DuffusionPrior的训练主要包含以下几个步骤：

Step1：将图片文本对输入到CLIP[4]模型中获得对应的文本向量编码（text embedding）文本中每个token的编码（text encoding）及图像向量编码（image embedding）。这里值得注意的是：CLIP的模型权重冻结，不参加训练。

Step2: 对image embedding（此处的image embedding就是diffusion model中的）。随机采样时间步，对进行加噪,得到第时刻的noise image embedding 。

Step3: 将文本向量编码、文本token编码、时间步向量、时刻的图像向量编码输入到模型中（其本质是一个decode-only transformer with a causal attention mask）。该模型的作用是预测分布分布的均值。

Step4: 计算预测分布与实际分布的KL散度实现训练。

PS：作者训练中用了很多trick，由于本文主要关注在理解dalle2的工作原理，故未详细描述，若想深入了解可参考代码。

推理阶段

推理阶段主要包含以下步骤：

Step1：将文本内容y送入到CLIP模型中得到文本向量文本向量编码（text embedding）文本中每个token的编码（text encoding）。

Step2: 从高斯分布中随机采样一组向量作为T时刻的图片向量编码。将文本向量编码、文本中每个token的编码、时间步编码、T时刻的图片向量编码送入到DiffusionPriorNetwork中，预测T−1时刻的图片向量编码。

Step3: 不断重复步骤2，直至时间步到0得到所需的image embedding。

可能会有读者有疑问：CLIP模型不是已经在文本特征与图片特征对齐了吗？为什么要多此一举设计DiffusionPrior模块来根据文本embedding来生成图片embedding而不是直接用文本embedding来作为图片embedding。

这是因为虽然CLIP在训练过程中通过对比损失使得图片特征与文本特征进行对齐，但是模态间还是有存在Gap[6]，Prior模块的很大的意义是给定文本embedding，准确预测出其对应CLIP的image embedding，更有益于重建的效果。如下图将CLIP的image embedding送入decode生成的图像（第三行）从相关性与质量上明显优于直接送入text embedding的效果（第二行）。

DiffusionDecode 模块

Decode部分也是个diffusion model。其主要作用是根据image embedding和文本embedding进行图像重建。这里值得强调的是，之前基于文本生成图片的任务只有文本向量作为条件（condition）。与之不同的是，Dalle2生成图片的条件有两个：其一是文本向量；其二是基于Prior模块估计的图片向量，这个操作使得它的生成效果更逼真、稳定。

decode模块有两个比较关键训练trick

如何实现classifier-free guidance7

dalle2的实现是每次随机mask掉10%的CLIP embedding，并且每次随机丢弃50%的text caption。

如何获得更高分辨率的生成图像

Dalle2采用了渐进式的生成方法。Dalle2训练了两个diffusion upsampler模型，第一个模型从64×64采样到256×256,第二个进一步采样到1024×1024。为了提升upsamper的稳健性，第一个上采样阶段用了高斯模糊；第二个上采样阶段用了多种BSR degradation操作[8]

最后来看一下Dalle2的pipeline