Dalle2 技术小结

背景

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

Dalle2的总体架构

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

其一是DiffusionPrior \(P(z_i|y)\)，其作用是获得输入文本\(y\)对应的CLIP图片向量\(z_i\)。

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

DiffusionPrior模块

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

前置知识：

DIffusion model的原理，可参考博客³。
CLIP的基本原理⁴。

训练阶段

训练数据：图片文本对\(\{(x_i,y_i)| i=1,2,...,N\}\)

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

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

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

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

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⁶，Prior模块的很大的意义是给定文本embedding，准确预测出其对应CLIP的image embedding，更有益于重建的效果。如下图将CLIP的image embedding送入decode生成的图像（第三行）从相关性与质量上明显优于直接送入text embedding的效果（第二行）。

DiffusionDecode 模块

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

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

1. 如何实现classifier-free guidance⁷

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

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

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

最后来看一下Dalle2的pipeline

Dalle2的生成质量和目前的局限性

Dalle2的生成效果

基于MS-COCO caption生成效果

Other

Dalle2的局限性

1. CLIP的embedding没有将属性精确的绑定到对象，导致进行解码重建时会混淆属性和对象。

2. CLIP的embedding可能没有精确编码渲染文本的拼写信息，导致连续文本的生成会出现问题。

3. 由于解码器是在\(64 \times 64\)分辨率生成图像，虽有进行上采样，这可能导致难以在复杂场景生成细节。

莫叶何竹