CogView: Mastering Text-to-Image Generation via Transformers
公众号:EDPJ(添加 VX:CV_EDPJ 或直接进 Q 交流群:922230617 获取资料)
目录
0. 摘要
1. 简介
2. 方法
2.1 理论
2.2 标记化
2.3 自回归 Transformer
2.4 训练的稳定性
3. 微调
3.1 超分辨率
3.2 图像标题和自我重新排名
3.3 风格学习
3.4 工业时尚设计
4. 实验结果
4.1 机器评估
4.2 人类评价
5. 结论与讨论
参考
附录
A. 数据收集和标记器的详细信息
B. 稀疏注意力
C. 注意力分析
C.1 位置偏差
C.2 语义分割
C.3 注意力随深度而变化
C.4 注意力数值尺度
D. CogView 中的公平性:情况和解决方案
S. 总结
S.1 主要贡献
S.2 架构和方法
通用领域中的文本到图像生成长期以来一直是一个悬而未决的问题,这需要强大的生成模型和跨模态理解。 我们提出了 CogView,一个具有 VQ-VAE 标记器的 40 亿参数 Transformer 来解决这个问题。 我们还演示了各种下游任务(例如,风格学习、超分辨率、文本图像排名和时装设计)的微调策略,以及稳定预训练的方法(例如,消除 NaN 损失)。 CogView 在模糊的 MS COCO 数据集上实现了最先进的 FID,优于之前基于 GAN 的模型和最近的类似工作 DALL-E。
代码和模型位于 https://github.com/THUDM/CogView
最新模型的演示网站:https://wudao.aminer.cn/CogView/index.html(无 post-selection)
“对于画家来说,有两件事:眼睛和心灵……眼睛,我们通过它观察自然;眼睛,我们通过它观察自然; 大脑,我们通过逻辑组织感觉以进行有意义的表达。” (Paul Cézanne [17])
由于对比自监督预训练彻底改变了计算机视觉 (CV) [24, 21, 8, 32],为图像带来高级语义的视觉语言预训练正在成为视觉理解的下一个前沿 [38, 30, 39]。 在各种前置任务(pretext tasks)中,文本到图像生成期望模型(1)从像素中解耦形状、颜色、手势和其他特征,(2)理解输入文本,(3)将对象和特征与相应的单词及其特征对齐 (4)学习复杂的分布来生成不同物体和特征的重叠和复合,这就像绘画一样,超出了基本的视觉功能(与眼睛和大脑中的 V1-V4 有关[22]),需要更高水平的认知能力(与大脑中的角回(angular gyrus)更相关 [3])。
教授机器文本到图像生成的尝试可以追溯到深度生成模型的早期,当时 Mansimov 等人 [35] 向DRAW [20] 添加了文本信息。然后生成对抗网络(GAN) [19] 开始主导这项任务。
近年来,自回归生成模型兴起。
矢量量化变分自动编码器(Vector Quantized Variational AutoEncoders,VQ-VAE)[46] 的框架缓解了这个问题。
CogView 的想法很自然:对文本和图像(来自 VQ-VAE)标记进行大规模生成联合预训练。
最近的一篇工作 DALL-E [39] 独立提出了同样的想法,并且发布时间早于 CogView。 与 DALL-E 相比,CogView 在以下四个方面取得了进步:
在本节中,我们将从 VAE [26] 推导出 CogView 的理论(本文中,粗体表示随机变量,常规字体表示具体值。 有关 VAE 的基础知识,请参阅此综合教程 [12]):CogView 优化图像和文本联合似然的证据下界(Evidence Lower BOund,ELBO)。 如果没有文本 t,下面的推导将变成对 VQ-VAE 的清晰的重新解释。
假设数据集
由 N 个 i.i.d 图像变量 x 及其描述文本变量 t 的样本组成。 我们假设图像 x 可以通过涉及潜在变量 z 的随机过程生成: (1) t_i 首先从先验 p(t; θ) 生成。 (2) 然后从条件分布 p(z | t = t_i; θ) 生成 z_i。 (3) x_i 最终由 p(x | z = z_i; Ψ) 生成。 在接下来的部分中,我们将使用像 p(x_i) 这样的简写形式来指示 p(x = x_i)。
令 q(z | x_i; Φ) 为变分分布(variational distribution),它是 VAE 编码器的输出。 对数似然和证据下界 (ELBO) 可写为:
VQ-VAE 的框架与传统 VAE 的不同之处主要在于 KL 项。
目前,最强大的生成模型 Transformer (GPT) 可处理离散码本上的标记序列。 为了使用它,令
其中 |V| 是码本的大小,h*w 是 z 的维数。 序列 z_i 可以从 q(z | x_i; Φ) 中采样,也可以直接
为了简单起见,我们选择后者,以便 q(z | x_i; Φ) 成为 z_i 上的单点分布。 方程(2)可以改写为:
然后学习过程分为两个阶段:
因此,第一阶段退化为纯粹的离散自动编码器,充当图像标记器,将图像转换为标记序列; 第二阶段的 GPT 承担了大部分建模任务。 图 3 展示了 CogView 的框架。
在本节中,我们将介绍 CogView 中标记器的详细信息,并比较图像标记器的不同训练策略(VQVAE 第 1 阶段)。
文本的标记化已经得到了充分的研究,例如 BPE [16] 和 SentencePiece [28]。 在 CogView 中,我们在大型中文语料库上运行 SentencePiece 以提取 50,000 个文本标记。
图像标记器是一个离散自动编码器,类似于 VQ-VAE [46] 或 d-VAE [39] 的第一阶段。更具体地说,编码器 Φ 将形状为 H*W*3 的图像 x 映射为形状为 h*w*d 的 Enc_Φ(x),然后将每个 d 维向量量化为可学习码本
中的临近嵌入。量化结果可以用 h*w 嵌入索引表示,然后我们得到潜在变量
解码器 Ψ 将量化的矢量映射回(模糊)图像以重建输入。 在我们的 4B 参数 CogView 中,|V| = 8192,d = 256,H = W = 256,h = w = 32。
由于离散选择的存在,图像标记器的训练并不简单。 这里我们介绍四种训练图像标记器的方法:
1) 最近邻映射,直通估计器 [2],由原始 VQ-VAE 提出。 这种方法的一个常见问题 [39] 是,当码本很大并且初始化不仔细时,由于维数灾难,只会使用少量嵌入。 我们在实验中没有观察到这种现象。
2) Gumbel 采样,直通估计器。 如果我们按照原始 VAE 来根据向量之间的距离重新参数化潜在变量 z 的分类分布,即
无偏采样策略为
其中,温度 τ 逐渐降至 0。我们可以进一步使用可微分的 softmax 来近似 argmax 的 one-hot 分布。 DALL-E采用这种方法并结合许多其他技巧来稳定训练。
3) 最近邻映射,移动平均值,其中码本中的每个嵌入在训练期间定期更新为最近映射到它的向量的平均值 [46]。
4) 最近邻映射,固定码本,初始化后码本是固定的。
对比。 为了比较这些方法,我们在相同的数据集和随机种子上训练了四个具有相同架构的图像标记器,并展示了图 2 中的损失曲线。我们发现所有方法基本上是均匀匹配的,这意味着嵌入的学习 如果初始化正确的话,码本并不是很重要。 在预训练中,我们使用移动平均法的标记器。
数据的介绍和更多关于标记化的细节参见附录A。
CogView 的主干是单向 Transformer (GPT)。
预训练的前置任务(pretext task)是从左到右的标记预测,也称为语言建模。
训练细节。
目前,预训练大型模型(>2B 参数)通常依赖于 16 位精度来节省 GPU 内存并加快计算速度。 许多框架,例如 DeepSpeed ZeRO [40],甚至只支持 FP16 参数。 然而,文本到图像的预训练在 16 位精度下非常不稳定。 训练 4B 普通 pre-LN Transformer 将在 1,000 次迭代内快速导致 NaN 损失。 稳定训练是 CogView 中最具挑战性的部分,它与 DALL-E 非常一致。
我们将 DALL-E 的解决方案总结为容忍训练的数值问题。
相反,CogView 会正则化这些值。 我们发现存在两种不稳定性:溢出(以 NaN 损失为特征)和下溢(以发散损失为特征)。 提出以下技术来解决这些问题。
精度瓶颈松弛(Precision Bottleneck Relaxation,PB-Relax)。 在分析训练动态后,我们发现溢出总是发生在两个瓶颈操作处,即最终的 LayerNorm 或注意力。
三明治层规范(Sandwich LayerNorm,Sandwich-LN)。 Transformers 中的 LayerNorms [1] 对于稳定训练至关重要。 Pre-LN [50] 被证明比原始的 Post-LN 收敛更快、更稳定,并成为最近作品中 Transformer 层的默认结构。 然而,这对于文本到图像的预训练来说还不够。LayerNorm
的输出基本上与 x 的隐层大小的平方根成正比,在 CogView 中为 √d = √2560 ≈ 50。 如果某些维度的输入值明显大于其他维度(对于 Transformer 来说就是这样),这些维度的输出值也会很大 (10 ~ 100)。 在残差分支中,这些大值被放大并加回主支,从而在下一层加剧这种现象,最终导致深层数值爆炸。
这个数值爆炸背后的原因激励我们去限制层层加剧。
玩具(Toy)实验。
缩小嵌入梯度。尽管我们在使用 Sandwich-LN 后没有观察到任何下溢的迹象,但我们发现标记嵌入的梯度比其他参数的梯度大得多,因此只需将其规模缩小 α = 0.1 即可增加动态损失规模,以进一步防止下溢,可以通过 Pytorch 中的 emb = emb*alpha + emb.detach()*(1 - alpha) 来实现。 它似乎减慢了标记嵌入的更新,但实际上并没有损害我们实验中的性能,这也对应于最近的一项工作MoCo v3 [9]。
讨论。
CogView 在微调方面比 DALL-E 更进一步。 特别是,我们可以通过微调 CogView 的超分辨率和自重新排名来改进文本到图像的生成。 所有微调任务都可以在一台 DGX-2 上在一天内完成。
由于图像标记器在训练之前将 256*256 像素图像压缩为 32*32 标记序列,因此由于有损压缩,生成的图像比真实图像更模糊。
我们超分辨率微调解决方案的动机是相信 CogView 是在一般域中最复杂的分布上进行训练的,并且已经涵盖了不同分辨率的对象。
具体来说,
为了准备数据,
微调用于图像标题的 CogView 非常简单:交换输入序列中文本和图像标记的顺序。
我们提出了标题损失(Caption Loss,CapLoss)来评估图像和文本之间的对应关系。 更具体地说,
其中 t 是文本标记序列,x 是图像。 CapLoss(x, t) 是文本标记的交叉熵损失,该方法可以看作是文本到图像生成的逆提示(inverse prompting) [56] 的适应。 最后,选择 CapLosses 最低的图像。
与额外训练另一个用于重新排名的对比自监督模型(例如 CLIP [38])相比,我们的方法消耗更少的计算资源,因为我们只需要微调。 图 9 中的结果显示,通过我们的方法选择的图像在 FID 中比通过 CLIP 选择的图像表现更好。 图 6 显示了重新排名的示例。
尽管 CogView 经过预训练以尽可能覆盖多样化的图像,但无法很好地满足生成特定风格或主题图像的愿望。 我们对模型进行了四种风格的微调:国画、油画、素描、卡通。 这些风格的图片是从Google、百度、Bing 等搜索引擎页面自动提取的,关键词为 “An image of {style} style”,其中{style} 是风格名称。 我们分别针对不同风格对模型进行微调,每个模型有 1,000 张图像。
在微调过程中,图像对应的文本也是 “An image of {style} style”。 生成时,文本为 “A {object} of {style} style”,其中 {object} 是要生成的对象。 通过这种方式,CogView 可以将从预训练中学到的对象形状知识转移到微调的样式中。 图 7 显示了样式的示例。
当生成针对单个域时,纹理的复杂性大大降低。 在这些场景中,我们可以 (1) 训练 VQGAN [15] 而不是 VQVAE 作为潜在变量,以获得更真实的纹理,(2) 减少参数数量并增加序列长度以获得更高分辨率。 在这种情况下,我们的三区域稀疏注意力(three-region sparse attention)(附录 B)可以加速高分辨率图像的生成。
我们使用 50*50 个 VQGAN 图像标记在大约 1000 万个时尚-标题对上训练 3B 参数模型,并将它们解码为 800*800 像素。 图 8 显示了用于时装设计的 CogView 示例,该示例已成功部署到阿里巴巴 Rhino 时装生产中。
目前,通用领域文本到图像生成最权威的机器评估指标是 MS COCO 上的 FID,该指标未包含在我们的训练集中。
标题损失作为衡量标准。
与 CLIP 对比自我重排。 我们评估了由 CLIP 选择的 CogView 生成图像的 FID-0 和 IS,并在 MS COCO 上进行自我重新排名。 图 9 显示了不同候选数量的曲线。 自我重新排名获得更好的 FID,并随着候选数量的增加而稳步细化 FID。 CLIP 在增加 IS 方面表现更好,但如上所述,它不是此任务的合适指标。
关于 CogView 和 DALL-E 之间性能差异的讨论。
在文本到图像的生成方面,人类评估比机器评估更有说服力。 我们的人类评估包括 AttnGAN、DM-GAN、DF-GAN、CogView 生成的图像和恢复的地面实况(即由我们的图像标记器模糊的地面实况,理论上是 CogView 的上边界)之间的 2,950 组比较。 模型之间的详细信息和基于示例的比较参见附录 E。
图 10 中的结果表明,CogView 的性能大幅优于基于 GAN 的基线。 CogView 以 37.02% 的概率被选为最好的,与恢复的地面实况 (59.53%) 的性能相媲美。 图 10(b)(c) 还表明我们的超分辨率模型持续提高了图像质量,尤其是清晰度,甚至优于恢复的地面实况。
局限性。 CogView 的缺点是生成速度慢,这对于自回归模型很常见,因为每个图像都是逐标记生成的。 VQVAE带来的模糊也是一个重要的限制。 这些问题将在今后的工作中得到解决。
道德问题。
我们系统地研究了结合 VQVAE 和 Transformers 进行文本到图像生成的框架。 CogView 展示了可扩展跨模态生成预训练的有希望的结果,并且还揭示并解决了可能源自数据异构性的精度问题。 我们还介绍了针对不同下游任务微调 CogView 的方法。 我们希望 CogView 能够推进可控图像生成和跨模态知识理解的研究和应用,但需要防止它被用来创建错误信息的图像。
Ding M, Yang Z, Hong W, et al. Cogview: Mastering text-to-image generation via transformers[J]. Advances in Neural Information Processing Systems, 2021, 34: 19822-19835.
我们从多个渠道收集了约 3000 万个文本图像对,并构建了 2.5TB 的新数据集(标记化后,大小变为约 250GB)。
数据来源基本分为以下几类:
我们已经在 2.2 节中介绍了标记器,这里有一些细节。
如图 11 所示,我们设计了三区域稀疏注意力(three-region sparse attention),这是一种易于实现的文本到图像生成的稀疏注意力。 每个标记关注所有文本标记、所有枢轴标记(pivots tokens)以及其之前的相邻窗口中的块中的标记。
枢轴标记(pivots tokens)是随机选择的图像标记,类似于大鸟(big bird)[53]。 每次我们进入新层时都会对它们重新采样。 我们认为他们可以提供有关图像的全局信息。
分块窗口注意力提供了局部信息,这是最重要的区域。 一维窗口注意力的前向计算可以通过仔细填充和改变张量的步幅来有效地就地实现,因为要关注的位置在内存中已经是连续的。 然而,如果没有定制的 CUDA 内核,我们仍然需要额外的内存来进行向后计算。 我们通过将相邻标记分组到块中来缓解这个问题,其中所有标记都注意相同的标记(在因果掩蔽之前)。 更多详细信息包含在我们发布的代码中。
在我们对 4096 个标记序列的基准测试中,三区域稀疏注意力(768 个文本和枢轴标记、768 个分块窗口标记)比普通注意力快 2.5倍,并节省 40% 的 GPU 内存。 整个训练比普通注意力训练快 1.5倍,并节省 20% 的 GPU 内存。 在相同的超参数、数据和随机种子的情况下,它们的损失曲线几乎相同,这意味着稀疏注意力不会影响收敛。
然而,我们在训练 40 亿参数 CogView 时没有使用三区域稀疏注意力,因为担心它可能与 3.1 节中的超分辨率微调不兼容。 但它成功地加速了 CogView-fashion 的训练,而且没有副作用。
为了探索 CogView 的注意力机制,我们通过绘制热图并标记最受关注的标记来可视化推理过程中的注意力分布。 我们发现我们的模型的注意力头在捕获位置和语义信息方面表现出很强的能力,并且不同层之间的注意力分布有所不同。 关于注意力评分量表的分析在 C.4 节中。
注意力分布与图像的位置结构高度相关。
CogView 中的注意力还表明它也执行隐式语义分割。 有些标题突出了文本中提到的主要对象。 我们用 “桌子上有一个苹果,旁边有一个花瓶,里面有紫色的花” 作为我们实验的输入。 在图 13 中,我们用红点标记了与最受关注的标记相对应的像素,并发现注意力头成功捕获了苹果和紫色花等项目。
不同层之间的注意力模式有所不同。
作为 2.4 节的补充,我们可视化第 38 层的注意力数值尺度,该层在 CogView 中具有最大的注意力得分尺度 (Q^T)K / √ d 。 不同头部的尺度差异很大,但每个头部的方差很小(这就是为什么即使分数很大,注意力也不会退化)。 我们认为原因是模型希望不同的头部有不同的敏感度,以便它学习乘以不同的常数来得到 Q 和 K。作为副作用,这些值可能有很大的偏差。 注意力的 PB- relax 是为了消除计算过程中的偏差。
CogView 中公平情况的评估。 我们研究了不同种族和性别的比例偏差。 首先,如果在文本中给出详细的描述,例如,“黑人男性” 或 “亚洲女性”,CogView 几乎可以正确生成所有样本。 我们还通过文本 “一张脸,照片” 来测量没有具体描述的生成样本的比例。 不同种族和性别的比例如图 15 所示。(无条件)生成的人脸在种族和年龄上相对平衡,但由于数据分布,男性多于女性。
如果不指定性别,CogView 还会因刻板印象而受到性别偏见的困扰。 但是如果我们指定性别,几乎所有的性别和职业都是正确的。 我们测试了 [6] 中介绍的示例,并为文本 {男性,女性} - {“科学”,数学”,“艺术”,“文学”} 生成图像。 结果如下图所示。
单词替换解决方案。 与之前的无条件生成模型不同,我们有一个非常简单有效的种族和性别公平解决方案。
在文本中,我们可以直接在人类的词(例如 “people” 或 “person”)前面添加一些从 “白人”、“黑人”、“亚洲人”……以及“男性”、“女性”(如果没有指定)中抽取的形容词。 抽样是按照在总体中的真实比例进行的。 我们可以训练一个额外的 NER 模型来查找有关人类的单词。
由于 CogView 会根据上面的结果进行正确的预测,如果给出描述,该方法将极大地帮助解决生成模型中的公平性问题。
提出了 CogView(cognitive ability),一个具有 VQ-VAE 标记器的 40 亿参数 Transformer,来解决通用领域中的文本到图像生成问题。CogView 可适用于各种下游任务,例如风格学习(特定领域的文本到图像)、超分辨率(图像到图像)、图像标题(图像到文本),以及文本到图像重排名
提出精度瓶颈松弛和三明治层归一化(Precision Bottleneck Relaxation and Sandwich Layernorm)来解决可能由数据的异构性引起的大规模文本到图像的生成预训练的不稳定。
提出了标题损失(Caption Loss,CapLoss)来评估图像和文本之间的对应关系,以比 FID 和 Inception Score (IS) 更细的粒度来衡量文本图像生成的质量和准确性。CapLoss 是文本标记的交叉熵损失,该方法可以看作是文本到图像生成的逆提示(inverse prompting)的适应。相比于对比自监督模型(例如 CLIP),CapLoss 消耗更少的计算资源,并且有更好的 FID 性能。
提出三区域稀疏注意力(three-region sparse attention)来加速训练并节省内存而不损害性能。
CogView 的架构如图 3 所示。
精度瓶颈松弛(Precision Bottleneck Relaxation,PB-Relax):通过适度的缩放,避免数值爆炸导致的 LayerNorm 中的溢出。
三明治层归一化(Sandwich LayerNorm,Sandwich-LN):Transformer 中不同的 LayerNorm 结构如图 4 所示,Sandwich-LN 保证了各层输入值的规模在合理范围内,从而避免数值的逐层放大而最终导致深层数值爆炸。