(2021|PMLR,DALL·E,自回归,Transformer)零样本文本到图像生成

Zero-Shot Text-to-Image Generation

公众号:EDPJ(添加 VX:CV_EDPJ 或直接进 Q 交流群:922230617 获取资料)

目录

0. 摘要

1. 简介

2. 方法 

2.1. 第一阶段:学习视觉码本 

2.2. 第二阶段:学习先验知识 

2.3. 数据收集 

2.4. 混合精度训练 

2.5. 分布优化 

2.6. 样本生成

3. 实验 

3.1. 定量结果 

3.2. 数据重叠分析 

3.3. 定性研究结果 

4. 结论 

参考

附录

A. 离散 VAE (dVAE)的详细信息

A.1. 架构

A.3. 逻辑拉普拉斯分布

B. Transformer 的详细信息 

B.1. 架构 

D. 混合精度训练指南 

S. 总结

S.1 主要思想

S.2 方法


0. 摘要

文本到图像生成传统上专注于寻找更好的建模假设以在固定数据集上进行训练。 这些假设可能涉及复杂的架构、辅助损失(auxiliary losses)或辅助信息(side information)(例如,训练期间提供的对象部分标签或分割掩模)。 我们描述了一种基于 transformer 的简单方法来完成此任务,该 transformer 将文本和图像 token 自回归地建模为单个数据流。 凭借足够的数据和规模,当以零样本方式评估时,我们的方法与以前的特定领域模型具有竞争力。

1. 简介

文本到图像合成的现代机器学习方法始于 Mansimov 等人 (2015) 的工作,他表明 DRAW(Gregor 等人,2015)生成模型,当扩展到以图像标题为条件时,也可以生成新颖的视觉场景。 Reed 等人 (2016b) 后来证明,使用生成对抗网络 (Goodfellow et al., 2014),而不是循环变分自动编码器,可以提高图像保真度。Reed 等人(2016b)表明该系统不仅可以生成具有可识别属性的对象,而且还可以零样本泛化到其他类别。

在接下来的几年里,通过多种方法的结合,研究不断取得进展。 其中包括通过多尺度生成器等修改来改进生成模型架构(Zhang et al., 2017; 2018)、整合注意力和辅助(auxiliary)损失(Xu et al., 2018),以及利用文本之外的其他调节信息源( Reed 等人,2016a;Li 等人,2019;Koh 等人,2021)。 

另外,Nguyen 等人(2017)提出了一种基于能量的条件图像生成框架,相对于当代方法,该框架在样本质量方面取得了很大的进步。 他们的方法可以结合预训练的判别模型,并且表明,当应用于在 MS-COCO 上预训练的获取标题(captioning)模型时,该方法能够执行文本到图像的生成。 最近,Cho 等人(2020)还提出了一种方法,涉及优化预训练的跨模态掩码语言模型的输入。虽然自 Mansimov 等人 (2015) 以来,视觉保真度显着提高,但样本仍然可能遭受严重的伪影,例如对象扭曲、不合逻辑的对象放置或前景和背景元素的不自然混合。 

大规模生成模型推动的最新进展提出了进一步改进的可能途径。 具体来说,当计算、模型大小和数据仔细扩展时,自回归 transformers(Vaswani 等人,2017)在文本(Radford 等人,2019)、图像(Chen 等人,2019)和音频(Dhariwal 等人,2020)等多个领域取得了令人印象深刻的结果。 

相比之下,文本到图像的生成通常在相对较小的数据集(例如 MS-COCO 和 CUB-200)上进行评估(Welinder 等人,2010)。 数据集大小和模型大小会成为当前方法的限制因素吗? 在这项工作中,我们证明,对从互联网收集的 2.5 亿个图像文本对训练 120 亿个参数的自回归 transformers,可以生成灵活、高保真度的可通过自然语言控制的图像生成模型。 

由此产生的系统在流行的 MS-COCO 数据集上零样本地实现了高质量图像生成,而无需使用任何训练标签。 90% 的人类评估者认为它优于该在数据集上训练的先前工作。 我们还发现它能够执行复杂的任务,例如初级水平的图像到图像的转换。 这以前需要定制方法(Isola et al., 2017),而不是作为单一大型生成模型的能力而出现。 

代码:https://github.com/openai/DALL-E

2. 方法 

我们的目标是训练一个 Transformer(Vaswani 等人,2017)将文本和图像标记(token )自回归地建模为单个数据流。 然而,直接使用像素作为图像 token 将需要大量的内存来存储高分辨率图像。 似然目标倾向于优先考虑对像素之间的短程依赖性进行建模(Salimans et al., 2017),因此大部分建模能力将用于捕获高频细节,而不是使我们在视觉上可识别对象的低频结构。 

我们通过使用两阶段训练程序来解决这些问题,类似于(Oord 等人,2017 年;Razavi 等人,2019 年): 

  • 第 1 阶段。我们训练一个离散变分自动编码器 (discrete variational autoencoder,dVAE),将每个 256*256 RGB 图像压缩为 32*32 图像标记(token)网格,其中每个元素可以假设有 8192 个可能值。 这将 transformer 的上下文大小减少了 192 倍,而视觉质量没有大幅下降(见图 1)。 
  • 第 2 阶段。我们将最多 256 个 BPE 编码的文本标记与 32*32 = 1024 个图像标记连接起来,并训练一个自回归转换器来对文本和图像标记上的联合分布进行建模。 

整个过程可以被视为最大化编码的 RGB 图像的图像 x、文本标题 y 和标记 z 上模型分布的联合似然的证据下界 (evidence lower bound,ELB)(Kingma & Welling,2013;Rezende 等人,2014)。 我们使用如下因式分解对该分布进行建模

得出下界 

其中,

  • q_φ 表示在给定 RGB 图像 x 的情况下 dVAE 编码器生成的 32*32 个图像标记上的分布。我们假设给定 z,y 条件独立于 x。
  • p_θ 表示给定图像标记时 dVAE 解码器生成的 RGB 图像上的分布 
  • p_Ψ 表示由 transformer 建模的文本标题 y 和图像标记 z 的联合分布。 

请注意,该界限仅适用于 β = 1,而在实践中我们发现使用更大的值很有帮助(Higgins 等人,2016)。 以下小节更详细地描述了这两个阶段。 (在 ImageNet(Deng et al., 2009)的初步实验中,我们尝试联合最大化对应于 φ、θ  和 Ψ 的 ELB,但无法改进两阶段训练)

2.1. 第一阶段:学习视觉码本 

在训练的第一阶段,我们最大化关于 φ 和 θ 的 ELB,这对应于仅在图像上训练 dVAE。 我们将初始先验 p_Ψ 设置为 K = 8192 个码本向量上的均匀分类分布,将 q_φ 设置为由编码器输出的 32*32 网格中相同空间位置处的 8192 个逻辑值(logits)参数化的分类分布。 

ELB 现在变得难以优化:因为 q_φ 是离散分布,我们无法使用重新参数化梯度来最大化它。 Oord 等人(2017);Razavi 等人 (2019) 使用在线聚类分配程序与直通估计器(straight-through estimator)相结合来解决这个问题 (Bengio et al., 2013)。 相反,我们使用 gumbel-softmax 松弛(relaxation)(Jang et al., 2016; Maddison et al., 2016),用对 

的期望替换对 q_φ 的期望,其中随着温度 T→0,松弛变得紧密。p_θ 的似然使用逻辑拉普拉斯分布进行评估(推导参见附录 A.3)。 

使用带有指数加权迭代平均的 Adam (Kingma & Ba, 2014) 来最大化松弛的 ELB。 附录 A.2 给出了超参数的完整描述,但我们发现以下内容对于稳定训练尤其重要: 

  • 松弛温度和步长的具体退火计划。 我们发现,将 T 退火至 1/16 足以缩小松弛验证 ELB 与真实验证 ELB (使用 q_φ 而不是 q^T_φ)之间的差距。 
  • 在编码器末端和解码器开始处使用 1*1 卷积。 我们发现,减小松弛周围的卷积的感受野可以使其更好地泛化到真正的 ELB。 
  • 将编码器和解码器的 resblock 的传出激活乘以一个小常数,以确保初始化时的稳定训练。 

我们还发现,将 KL 权重增加到 β = 6.6 可以促进更好的码本使用,并最终在训练结束时导致更小的重建误差。(这与这两项之间通常的权衡相反。 我们推测,对于较小的 β 值,松弛产生的噪声会导致优化器在训练开始时减少码本的使用,从而导致收敛时的 ELB 较差。) 

2.2. 第二阶段:学习先验知识 

在第二阶段,我们固定 φ 和 θ,并通过最大化关于 Ψ 的 ELB 来学习文本和图像标记的先验分布。 这里,p_Ψ 由 120 亿参数的稀疏 transformer 表示(Child et al., 2019)。 

给定文本-图像对,我们使用最多 256 个词汇量为 16,384 的标记对小写(lowercased)标题进行 BPE 编码(在训练过程中,我们应用 10% BPE dropout (Provilkov et al., 2019),其使用在神经机器转换文献中很常见)(Sennrich et al., 2015),并使用词汇量为 8192 的 32*32 = 1024 个标记对图像进行编码。 图像标记是使用 dVAE 编码器 logits 的 argmax 采样获得的,不添加任何 gumbel 噪声。最后,文本和图像标记被连接并自回归建模为单个数据流。 

严格来说,等式 1 要求从 dVAE 编码器 logits 指定的分类分布中进行采样,而不是采用 argmax。 在 ImageNet 的初步实验中,我们发现这是超参数化(overparameterized)机制中有用的正则化器,并且允许使用软目标来训练 transformer 以实现交叉熵损失。 我们在这里决定不这样做,因为所考虑的模型处于参数化不足(underparameterized)的状态。

Transformer 是一种仅解码器模型,其中每个图像标记都可以注意其 64 个自注意力层中任意一层中的所有文本标记。 完整的架构在附录 B.1 中描述。 模型中使用了三种不同类型的自注意力掩模。 与文本到文本注意相对应的注意掩码部分是标准因果掩码,而图像到图像注意部分使用行、列或卷积注意掩码。 

我们发现对所有三种交互(“文本注意文本”、“图像注意文本” 和 “图像注意图像”)使用单一注意(single attention)操作比使用独立标准化的单独关注(separate attention)操作表现得更好。 

我们将文本标题的长度限制为 256 个标记,但并不完全清楚如何处理最后一个文本标记和图像开始标记之间的 “填充” 位置。 一种选择是在自注意力操作中将这些标记的 logits 设置为 1。 相反,我们选择为 256 个文本位置中的每个位置单独学习一个特殊的填充标记。 仅当没有可用的文本标记时才使用此标记。 在概念标题(Conceptual Captions)(Sharma 等人,2018)的初步实验中,我们发现这会导致更高的验证损失,但在分布外标题上表现更好。 

我们通过一批数据中每种类型的总数来标准化文本和图像标记的交叉熵损失。 由于我们主要对图像建模感兴趣,因此我们将文本的交叉熵损失乘以 1/8,将图像的交叉熵损失乘以 7/8。 使用带有指数加权迭代平均的 Adam 来优化目标; 附录 B.2 更详细地描述了训练程序。 我们保留了大约 606,000 张图像进行验证,没有发现收敛时过度拟合的迹象。 

2.3. 数据收集 

我们对高达 1.2 亿参数的模型的初步实验是在 Conceptual Captions 上进行的,Conceptual Captions 是一个包含 330 万个文本图像对的数据集,是作为 MS-COCO 的扩展而开发的(Lin 等人,2014 年)。 

为了扩展到 120 亿个参数,我们通过从互联网收集 2.5 亿个文本图像对,创建了一个与 JFT-300M (Sun et al., 2017) 类似规模的数据集。 该数据集不包括 MS-COCO,但包括 Conceptual Captions 和 YFCC100M (Thomee 等人,2016)的过滤子集。 由于 MS-COCO 是从后者创建的,因此我们的训练数据包括 MS-COCO 验证图像的一小部分(但没有任何标题)。 我们在第 3 节中给出的定量结果中对此进行了控制,发现它对结果没有明显的影响。 我们在附录 C 中提供了有关数据收集过程的更多详细信息。 

2.4. 混合精度训练 

为了节省 GPU 内存并提高吞吐量,大多数参数、Adam 矩(moments)和激活都以 16 位精度存储。 我们还使用激活检查点并在向后传递期间重新计算 resblock 内的激活。 让模型以 16 位精度训练超过 10 亿个参数而不发散,是该项目中最具挑战性的部分。 

我们认为这种不稳定的根本原因是 16 位梯度的下溢(underflow)。 附录 D 介绍了我们制定的一套指南,以避免在训练大规模生成模型时出现下溢。 在这里,我们描述了这些准则之一:每 resblock 梯度缩放。 

与之前的工作类似(Liu et al., 2020),我们发现,当我们从较早的 resblock 移动到较晚的 resblock 时,resblock 的激活梯度范数单调减小(更好的初始化方案(Liu et al., 2020)可能能够避免这种情况,但我们在实验中没有成功使用替代方案)。随着模型变得更深更宽, 后续 resblock 的激活梯度的真实指数可能会低于 16 位格式的最小指数。 因此,它们四舍五入为零,这种现象称为下溢(underflow)。 我们发现消除下溢可以实现稳定的训练收敛。 

当最小和最大激活梯度(绝对值)跨越的范围符合 16 位格式的指数范围时,标准损失缩放(Standard loss scaling)(Micikevicius 等人,2017)能够避免下溢。 在 NVIDIA V100 GPU 上,该指数范围由 5 位指定。 虽然这足以训练相同大小的普通(vanilla)语言模型,但我们发现对于文本到图像模型来说该范围太小。 

我们的修复如图 4 所示,涉及为模型中的每个 resblock 使用单独的 “梯度比例”。 这可以被视为一个称为 Flexpoint(Köster 等人,2017)的更通用混合精度训练框架的实用替代方案,其优点是不需要专门的 GPU 内核。 我们发现 Sun 等人 (2020) 独立开发了类似的程序来训练 4 位精度的卷积网络。 

2.5. 分布优化 

当以 16 位精度存储时,我们的 120 亿参数模型消耗约 24 GB 内存,这超过了 16 GB NVIDIA V100 GPU 的内存。 我们使用参数分片(parameter sharding)(Rajbhandari 等人,2019)来解决这个问题。 如图 5 所示,参数分片使我们能够通过与计算密集型操作重叠来几乎完全隐藏机器内通信的延迟。 

在用于训练模型的集群上,机器之间的带宽远低于同一机器上 GPU 之间的带宽。 这使得用于平均机器间梯度(all-reduce)的操作成本成为训练期间的主要瓶颈。 通过使用 PowerSGD 压缩梯度,我们能够大幅降低这一成本(Vogels 等人,2019)。 

在我们的实现中,机器中的每个 GPU 独立于其相邻 GPU 计算参数分片梯度的低秩因子(其他操作仍存在机内通信; 我们的意思是,跨分片的低秩因子在连接时不被视为集体逼近全参数矩阵的梯度)。计算出低秩因子后,每台机器将其误差缓冲区设置为在其八个 GPU 上未压缩梯度平均 (从reduce-scatter 获得)与从低秩因子获得的解压缩梯度之间的残差。 

PowerSGD 将针对未压缩参数梯度的大型通信操作替换为针对其低秩因子的两个小得多的通信操作。 对于给定的压缩秩 r 和 transformer 激活大小 d_model,压缩率由 1 - 5r / (8d_model) 给出(参见附录 E.1)。 表 1 显示我们可以实现约 85% 的压缩率,与模型大小无关。 

在附录 E.2 中,我们描述了使 PowerSGD 大规模运行所需的各种细节。 这些包括: 

  • 通过在反向传播期间将梯度累积到误差缓冲区中来节省内存,而不是分配单独的缓冲区。 
  • 最大限度地减少错误缓冲区归零的情况(例如,由于在混合精度反向传播期间或从检查点恢复训练时遇到非有限值)。 
  • 通过使用 Householder 正交化而不是 Gram-Schmidt 以及向输入添加一小部分单位矩阵来提高数值稳定性。 
  • 通过对错误缓冲区、其低秩因子以及涉及它们的 all-reduce 通信操作使用自定义 16 位浮点格式来避免下溢。 

我们还发现了 Vogels 等人(2019)描述的 Q 矩阵的热启动程序是不必要的:通过在训练开始时将 Q 固定为随机高斯矩阵并且从不更新它,我们能够获得等效的结果。我们验证了当 Q 固定时重建真实梯度的误差比使用热启动更新时更高,因此有趣的是这不会影响损失。 相比之下,每次更新时重新采样 Q 都会导致性能大幅下降。

2.6. 样本生成

与 Razavi 等人 (2019) 类似,我们使用预训练的对比模型 (Radford et al., 2021) 对从 Transformer 中抽取的样本进行重新排序。 给定标题和候选图像,对比模型根据图像与标题的匹配程度分配分数。 图 6 显示了增加样本数量 N 的效果,我们从中选择前 k 个图像。 这个过程可以看作是一种语言引导的搜索(Andreas et al., 2017),也类似于 Xu 等人(2018)提出的辅助文本图像匹配损失。除非另有说明,用于定性和定量结果的所有样本均在未降低温度的情况下获得(即使用 t = 1)(图 2 除外),并使用 N = 512 进行重新排序。 

3. 实验 

3.1. 定量结果 

我们通过将模型与之前的三种方法进行比较来评估我们的零样本模型:AttnGAN(Xu 等人,2018)、DMGAN(Zhu 等人,2019)和 DF-GAN(Tao 等人,2020),其中,最后一种方法报告了 MS-COCO 上的最佳 Inception Score (Salimans et al., 2016) 和 Fréchet Inception Distance (Heusel et al., 2017)。 图 3 对我们模型中的样本与之前工作中的样本进行了定性比较。 

我们还进行了类似于 Koh 等人 (2021) 使用的人工评估,来比较我们的方法与 DF-GAN,其结果如图 7 所示。给定一个标题,在更好地匹配标题方面,我们模型中的样本获得 93% 的多数票。在更现实方面,获得了 90% 的多数票。 

图 9(a) 显示我们的模型在 MS-COCO 上获得的 FID 分数与先前最佳方法相比也相差 2 分,尽管从未接受过标题训练。 我们的训练数据包含 YFCC100M 的过滤子集,我们发现它包含来自下一节中描述的重复数据删除过程的 MS-COCO 验证集中大约 21% 的图像。 为了隔离这种影响,我们计算了包含这些图像(实线)和不包含这些图像(虚线)的验证集的 FID 统计数据,发现结果没有显着变化。 

使用 dVAE 编码器的标记训练 transformer 使我们能够将其建模容量分配给使我们能够在视觉上识别图像的低频信息。 然而,它也对模型不利,因为严重的压缩使其无法产生高频细节。 为了测试其对定量评估的影响,我们在对验证图像和模型样本应用不同半径的高斯滤波器后计算图 9(a) 中的 FID 和 IS。对于半径为 1 的轻微模糊,我们的方法以大约 6 个点的优势实现了最佳 FID。随着模糊半径的增加,我们的方法与其他方法之间的差距往往会扩大。当模糊半径大于或等于 2 时,我们也获得最高的 IS。

我们的模型在 CUB 数据集上的表现明显较差,我们的模型与领先的先前方法之间的 FID 存在近 40 点的差距(图 9(b))。 我们发现该数据集的重叠率为 12%,并且再次观察到删除这些图像后结果没有显着差异。 我们推测,我们的零样本方法不太可能与 CUB 等专门发行版相比。 我们认为微调是一个有希望的改进方向,并将这项研究留到未来的工作中。 该数据集中的标题模型的样本如图 8 所示。 

最后,图 9(c) 显示了随着用于对比模型重新排序的样本量的增加,MS-COCO 的 FID 和 IS 有了明显的改进。 这种趋势一直持续到样本量达到 32,之后我们观察到收益递减。 

3.2. 数据重叠分析 

我们使用了 Radford 等人 (2021) 描述的重复数据删除程序,来确定要删除哪些图像。 对于每个验证图像,我们使用专门为此任务训练的对比模型在训练数据中找到最接近的图像。 然后,我们根据与训练数据中最接近的匹配项的接近程度对图像进行降序排序。 手动检查结果后,我们通过手动选择旨在最小化假负样本率的保留阈值来确定要删除的图像。 

3.3. 定性研究结果 

我们发现我们的模型能够以我们最初没有预料到的方式进行概括。 当给出标题 “由手风琴制成的貘...”(图 2a)时,模型似乎画出了一只身体有手风琴的貘,或者键盘或贝斯呈现貘的躯干或腿形状的手风琴。 这表明它已经具备了在高抽象层次上构建不寻常概念的基本能力。 

我们的模型似乎还能够进行组合泛化,例如在渲染文本时(图 2b)或在探究诸如“穿着圣诞毛衣的小刺猬遛狗的插图”(图 2c)之类的句子时。 像后者这样的提示要求模型执行变量绑定(Smolensky,1990;Greff 等人,2020)——圣诞毛衣里的是刺猬,而不是狗。 然而,我们注意到该模型在任务上的表现不一致,有时会画出两只都穿着圣诞毛衣的动物,或者画出一只刺猬正在遛一只较小的刺猬。 

在有限的可靠性范围内,我们还发现我们的模型能够通过自然语言控制进行零样本图像到图像的转换(图 2d)。 当为模型提供标题“顶部的猫与底部的草图完全相同”以及猫照片的图像标记网格的顶部 15*32 部分时,它能够在底部绘制一个看起来相似的猫的草图。 

这适用于其他几种类型的转换,包括图像操作(例如,更改图像的颜色、将其转换为灰度或将其上下翻转)和风格转换(例如,在贺卡、邮票或手机壳上画猫)。 一些转换(例如仅涉及改变动物颜色的转换)表明该模型能够执行基本的目标分割。 我们在附录 G 中提供了零样本图像到图像转换的其他示例。 

4. 结论 

我们研究了一种基于自回归 transformer 的文本到图像生成的简单方法,当它大规模执行时。 我们发现,大规模可以提高泛化能力,无论是相对于之前的特定领域方法的零样本性能,还是单个生成模型产生的能力范围。 我们的研究结果表明,提高泛化能力作为规模函数可能是推动这项任务取得进展的有用驱动力。 

参考

Ramesh A, Pavlov M, Goh G, et al. Zero-shot text-to-image generation[C]//International Conference on Machine Learning. PMLR, 2021: 8821-8831.

如何评价DALL-E模型的实现? - 金雪锋的回答 - 知乎 https://www.zhihu.com/question/447757686/answer/1764970196

附录

A. 离散 VAE (dVAE)的详细信息

A.1. 架构

dVAE 编码器和解码器是具有瓶颈式 resblock 的卷积(LeCun 等人,1998)ResNets(He 等人,2016)。 该模型主要使用 3*3 卷积,以及沿跳跃连接的 1*1 卷积,其中特征图的数量在 resblock 的输入和输出之间变化。 编码器的第一个卷积是 7*7,编码器的最后一个卷积(产生 32*32*8192 输出,用作图像标记的分类分布的 logits)是 1*1。解码器的第一个和最后一个卷积是 1*1。编码器使用最大池化(我们发现它比平均池化产生更好的 ELB)来对特征图进行下采样,解码器使用最近邻上采样。 架构的精确细节在代码发布的 dvae/encoder.py 和 dvae/decoder.py 文件中给出。 

A.3. 逻辑拉普拉斯分布

训练 VAE 时通常使用 “L1” 和 “L2” 重建目标。 这些目标分别对应于等式 1 中 ln p (x | y, z) 的拉普拉斯分布和高斯分布。 这种建模选择存在奇怪的不匹配:像素值位于有界区间内,但是这两种分布都受到整条实线的支持(supported by the entire real line)。 因此,一定量的可能性将被置于像素值的允许范围之外。

我们提出了拉普拉斯分布的一种变体,它也受到有界区间的支持。 这解决了正在建模的像素值的范围与用于建模它们的分布的支持之间的差异。 我们考虑将 sigmoid 函数应用于拉普拉斯分布随机变量而获得的随机变量的 pdf。 该 pdf 定义在 (0,1) 上,由下式给出 

我们称之为 Logit-Laplace 分布。 我们使用等式 2 的 RHS 的对数作为 dVAE 训练目标的重建项。 

dVAE 的解码器生成六个特征图,表示正在重建的图像的 RGB 通道的 Logit-Laplace 分布的充分统计数据。 前三个特征图代表 RGB 通道的 μ 参数,后三个代表 ln b。 在将图像输入 dVAE 编码器之前,我们使用 φ: [0, 255] → (ε, 1-ε) 转换其值,由下式给出

这将 dVAE 解码器建模的像素值的范围限制为 (ε, 1-ε),从而避免了等式 2 中的 x(1 - x) 引起的数值问题。我们使用 ε = 0.1。 为了重建图像以进行手动检查或计算指标,我们忽略 ln b 并计算

其中,μ 是由 dVAE 解码器输出的前三个特征图给出的。  

B. Transformer 的详细信息 

B.1. 架构 

我们的模型是与 Child 等人 (2019) 描述的同类的仅解码器稀疏变换器,对图像标记的上下文部分行和列嵌入进行广播。 图 10 显示了我们模型中使用的嵌入方案的完整描述。我们使用 64 个注意力层,每个注意力层使用 62 个注意力头,每个头状态大小为 64。 

该模型使用三种稀疏注意掩码,如图 11 所示。卷积注意掩码(图 11(d))仅在最后一个自注意层中使用。其它层,给定自注意力层的索引 i(i ∈ [1,63]),如果 i - 2 mod 4 = 0,我们使用列注意力掩码(图 11(c)),否则使用行注意力。 例如,前四个自注意力层分别使用 “行,列,行,行”。 除了卷积注意掩码之外,我们发现在最终的自注意层中使用卷积注意掩码时,其性能比行和密集因果注意掩码小幅提升,这与 Child 等人(2019)使用的配置相同。 

D. 混合精度训练指南 

该项目最具挑战性的部分是让模型以 16 位精度训练超过 10 亿个参数。 在检测训练各个部分的下溢并修改代码以消除它之后,我们能够做到这一点。 我们在此介绍的这一流程的结果是制定了一套指导方针。(对于像 TPU 这样具有 bfloat16 格式本机支持的硬件,可能需要较少的指导原则,因为较大的 8 位指数范围使得下溢发生的可能性较小。)

  • 使用每 resblock 梯度缩放(图 4)而不是标准损失缩放。 我们的模型使用 128 个梯度尺度,每个梯度尺度对应一个 resblock。 所有梯度尺度都初始化为 M·2^13,其中 M 是数据并行副本的数量(即GPU的数量)。 在我们的设置中,当该 resblock 中的任何参数梯度没有非有限值时,每次参数更新时每个梯度比例都会乘以 2^(1/1000)。 否则,我们将梯度除以 √2 并跳过更新。 我们还不允许在 125 次更新的窗口内连续划分相同的等级范围。 更新后,所有等级范围均被限制在 [M·27,M·224] 范围内。 图 12 显示了 28 亿参数模型训练早期阶段的梯度尺度。 
  • 仅在性能确实需要时才使用 16 位精度。 特别地,以 32 位精度存储所有增益、偏差、嵌入和非嵌入,以及 32 位梯度(包括远程通信)和 32 位 Adam moments.。 我们禁用这些参数的梯度压缩(尽管 PowerSGD 对于增益和偏差等一维参数没有意义)。 文本和图像标记的 logits 以 32 位精度计算和存储。 我们发现,以 16 位精度存储嵌入有时会在优化初期导致发散,而使用 16 位逻辑会导致训练曲线出现小幅变化,因此出于谨慎考虑,我们转而使用 32 位精度。 
  • 划分梯度时避免下溢。 对于数据并行训练,我们需要将梯度除以数据并行工作线程总数 M。一种方法是将损失除以每台机器的批量大小,然后在机器上对它们求和(使用 all-reduce)之前将参数梯度除以 M。 为了节省时间和空间,梯度通常以 16 位精度计算和存储。 当 M 很大时,这种除法可能会导致梯度求和之前发生下溢。 另一方面,如果我们尝试先对梯度求和,然后再除以它们,则可能会在 all-reduce 中遇到溢出。我们针对这个问题的解决方案试图在 all-reduce 之前最大程度地减少除法中的信息丢失,而不会有溢出的危险。 为此,我们将损失除以总体批量大小(其中包括 M 作为因子)而不是每台机器的批量大小,并将梯度尺度乘以 M 进行补偿,如第一点所述。 然后,在 all-reduce 操作之前,我们将梯度除以一个手动调整的常数,以避免下溢和溢出。 这是通过检查每个参数梯度的标量分量的绝对值的指数(即以 2 为底的对数)的直方图来完成的。 由于梯度缩放使梯度接近 16 位格式指数范围的右端,我们发现相同的常数对于 16 位梯度模型中的所有参数都适用。 使用 PowerSGD 时,我们为 P 和 Q 矩阵选择不同的常数。

S. 总结

S.1 主要思想

本文使用从互联网收集的 2.5 亿个图像文本对训练 120 亿个参数的自回归 transformers(该 transformer 将文本和图像 token 自回归地建模为单个数据流),获得灵活、高保真度的可通过自然语言控制的图像生成模型。该模型在流行的 MS-COCO 数据集上实现了零样本高质量图像生成,而无需使用任何训练标签。

S.2 方法

通过使用两阶段训练程序,来训练一个 Transformer,将文本和图像标记(token )自回归地建模为单个数据流。

  • 第 1 阶段。我们训练一个离散变分自动编码器 (discrete variational autoencoder,dVAE),将每个 256*256 RGB 图像压缩为 32*32 图像标记(token)网格,其中每个元素可以假设有 8192 个可能值。 这将 transformer 的上下文大小减少了 192 倍,而视觉质量没有大幅下降。 
  • 第 2 阶段。我们将最多 256 个 BPE 编码的文本标记与 32*32 = 1024 个图像标记连接起来,并训练一个自回归转换器来对文本和图像标记上的联合分布进行建模。

整个过程可以被视为最大化编码的 RGB 图像的图像 x、文本标题 y 和标记 z 上模型分布的联合似然的证据下界 (evidence lower bound,ELB) 

其中,

  • q_φ 表示在给定 RGB 图像 x 的情况下 dVAE 编码器生成的 32*32 个图像标记上的分布。我们假设给定 z,y 条件独立于 x。
  • p_θ 表示给定图像标记时 dVAE 解码器生成的 RGB 图像上的分布 
  • p_Ψ 表示由 transformer 建模的文本标题 y 和图像标记 z 的联合分布。 

优化

  • 因为 q_φ 是离散分布,我们无法使用重新参数化梯度来最大化它。本文使用 gumbel-softmax 松弛(relaxation) 对 q_φ 进行优化。
  • p_θ 的似然使用逻辑拉普拉斯分布进行评估。 
  • 使用带有指数加权迭代平均的 Adam (Kingma & Ba, 2014) 来最大化松弛的 ELB。

第 2 阶段,固定 φ 和 θ,并通过最大化关于 Ψ 的 ELB 来学习文本和图像标记的先验分布。其中,p_Ψ 由 120 亿参数的稀疏 transformer 表示。

样本生成。使用预训练的对比模型对从 Transformer 中抽取的样本进行重新排序。 给定标题和候选图像,对比模型根据图像与标题的匹配程度分配分数,从中选择分数最高的前 k 个图像。 这个过程可以看作是一种语言引导的搜索,也类似于辅助文本图像匹配损失。

  • 0
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值