(2023|ICML,LLM,标记掩蔽,并行解码)Muse:使用掩蔽生成 Transformer 的文本到图像生成

Muse: Text-To-Image Generation via Masked Generative Transformers

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

目录

0. 摘要

1. 简介

2. 模型

2.1. 预训练文本编码器

2.2. 使用 VQGAN 进行语义标记化

2.3. 基础模型 

2.4. 超分辨模型

2.5. 解码器微调 

2.6. 可变掩码率

2.7. 无分类器引导

2.8. 推理时的迭代并行解码 

3. 结果

3.1. 定性性能 

3.2. 量化性能

3.2.1. 人工评估

3.2.2. 推断速度 

3.3. 图像编辑

3.3.1. 文本引导的修复/扩展

3.3.2. 零样本无掩蔽编辑

4. 相关工作

4.1. 图像生成模型

4.2. 图像标记化

4.3. 大型语言模型

4.4. 文本-图像模型

4.5. 利用生成模型进行图像编辑

5. 讨论和社会影响 

附录

VQGAN Architecture

S. 总结

S.1 主要贡献

S.2 架构和方法


0. 摘要

我们提出了 Muse,这是一个文本到图像的 Transformer 模型,它在图像生成性能上实现了最先进的水平,同时比扩散或自回归模型高效得多。Muse 在离散标记空间中进行了一个带掩码的建模任务:给定从预训练的大型语言模型(LLM)中提取的文本嵌入,Muse 经过训练来预测随机掩码的图像标记。与像 Imagen 和 DALL-E2 这样的像素空间扩散模型相比,Muse 由于使用了离散标记并且需要较少的采样迭代,因此效率显著更高。与像 Parti 这样的自回归模型相比,Muse 由于使用并行解码更加高效。预训练 LLM 的使用使其具备了精细的语言理解能力,从而实现了高保真度的图像生成和对视觉概念(如对象、它们的空间关系、姿势、基数等)的理解。我们的 900M 参数模型在 CC3M 上取得了新的 SOTA,FID 得分为 6.06。Muse 的 3B 参数模型在零样本 COCO 评估上取得了 FID 为 7.88 的成绩,CLIP 分数为 0.32。Muse 还直接支持许多图像编辑应用,无需微调或反演模型:修复、扩展和免掩模式编辑。更多结果可在http://muse-model.github.io上找到。

1. 简介

以文本提示为条件的图像生成模型在过去几年中在质量和灵活性方面取得了巨大的飞跃 (Ramesh et al., 2022; Nichol et al., 2021; Saharia et al., 2022; Yu et al., 2022; Rombach et al., 2022; Midjourney, 2022)。这得益于深度学习架构的创新 (Van Den Oord et al., 2017; Vaswani et al., 2017),以及新颖的训练范式,如用于语言 (Devlin et al., 2018; Raffel et al., 2020) 和视觉任务 (He et al., 2022; Chang et al., 2022) 的掩蔽建模;新的生成模型家族,如扩散模型 (Ho et al., 2020; Rombach et al., 2022; Saharia et al., 2022) 和基于掩蔽的生成 (Chang et al., 2022);以及大规模图像-文本对数据集的可用性 (Schuhmann et al., 2021)。

在这项工作中,我们提出了一种新的文本到图像合成模型,采用了掩蔽图像建模方法 (Chang et al., 2022)。我们的图像解码器架构以来自预训练的并且冻结的 T5-XXL (Raffel et al., 2020) 大型语言模型 (LLM) 编码器的嵌入为条件。与 Imagen (Saharia et al., 2022) 的结果一致,我们发现在预训练的 LLM 的条件下对于逼真和高质量的图像生成是至关重要的。我们的模型 (除了 VQGAN 量化器) 基于Transformer (Vaswani et al., 2017) 架构构建。

我们已经训练了一系列的 Muse 模型,其规模从 632M 参数到 3B 参数不等(对于图像解码器;T5-XXL 模型还有额外的 4.6B 参数)。每个模型由几个子模型组成(图 3):首先,我们有一对VQGAN “ 标记器” 模型 (Esser et al., 2021b),它们可以将输入图像编码成一系列离散标记,也可以将标记序列解码回图像。我们使用了两个 VQGAN,一个用于 256x256 的分辨率(“低分辨率”),另一个用于 512x512 的分辨率(“高分辨率”)。其次,我们有一个基本的掩蔽图像模型,其中包含了我们大部分的参数。这个模型接受一系列部分掩蔽的低分辨率标记,并预测每个掩蔽标记的边际分布,条件是未掩蔽的标记和一个T5-XXL文本嵌入。第三,我们有一个 “超分辨率” transformer 模型,它将(未掩蔽的)低分辨率标记转化为高分辨率标记,同样是在 T5-XXL 文本嵌入的条件下。我们在第 2 节中详细解释了我们的流程。 

与建立在级联像素空间扩散模型上的 Imagen(Saharia等人,2022)或 Dall-E2(Ramesh等人,2022)相比,Muse 由于使用离散标记而更加高效;它可以被视为一个带有吸收状态([MASK])的离散扩散过程(Austin等人,2021)。与最新的自回归模型 Parti(Yu等人,2022)相比,Muse 由于使用并行解码而更加高效。基于在类似硬件(TPU-v4芯片)上的比较,我们估计 Muse在推理时比 Imagen-3B 或 Parti-3B 模型快 10 倍以上,并比 Stable Diffusion v1.4(Rombach等人,2022)快 3 倍(请参见第 3.2.2 节)。所有这些比较都是在相同大小的图像上进行的,分辨率要么是 256 x 256,要么是 512 x 512。尽管 Muse 和 Stable Diffusion(Rombach等人,2022)都在 VQGAN 的潜在空间中工作,但 Muse 更快,我们认为这是因为 Stable Diffusion v1.4 中使用了扩散模型,需要在推理时执行更多迭代。 

然而,Muse 的效率改进并不会降低生成的图像质量或对输入文本提示的语义理解。我们根据多个标准评估我们的输出,包括 CLIP 分数(Radford等人,2021)和 FID(Heusel等人,2017)。前者是图像与文本的对应度量,后者是图像质量和多样性的度量。我们的 3B 参数模型在 COCO(Lin等人,2014)零样本验证基准上获得了 0.32 的 CLIP 分数和 7.88 的 FID 分数,与其他大规模文本到图像模型的表现相媲美(见表 2)。我们的 632M(基本)+ 268M(超分辨率)参数模型在 CC3M(Sharma等人,2018)数据集上进行训练和评估时,实现了 6.06 的 FID 分数,明显低于文献中报告的所有其他结果(见表 1)。此外,我们还使用人工评分员在 PartiPrompts(Yu等人,2022)评估套件上评估我们的生成结果,评分员发现 Muse 生成的图像与其文本提示的对齐度要比 Stable Diffusion v1.4(Rombach等人,2022)高 2.7 倍。 

Muse 生成的图像反映了输入标题中的不同词性,包括名词、动词和形容词。此外,我们提供了多对象属性理解的证据,如组合性和基数,以及图像风格理解。请参见图 1,其中有许多这些示例,以及我们网站 http://muse-model.github.io 上的更多示例。Muse 的基于掩码的训练使其具备许多零样本图像编辑功能,一部分如图 2 所示,包括零样本、文本引导的 inpainting 和 outpainting,以及无掩码编辑。更多细节请参见第 3 节。

我们的贡献包括:

  • 我们提出了一种用于文本到图像生成的最先进模型,实现了出色的 FID 和 CLIP 分数(这是图像生成质量、多样性以及与文本提示的一致性的定量度量)。
  • 由于使用了量化图像标记和并行解码,我们的模型明显比类似模型更快。
  • 我们的架构实现了开箱即用的零样本编辑功能,包括 inpainting,outpainting 和无掩码编辑。

(2023|ICML,LLM,标记掩蔽,并行解码)Muse:使用掩蔽生成 Transformer 的文本到图像生成_第1张图片

2. 模型

我们的模型由许多组件构建而成。在这里,我们按照它们的训练顺序为每个组件提供一个概述,同时将许多架构和参数的细节放在附录中。图3提供了模型架构的概述。

2.1. 预训练文本编码器

与 Saharia 等人(2022)的发现类似,我们发现利用预训练的大型语言模型(LLM)有助于高质量的图像生成。从 LLM(例如 T5-XXL)提取的嵌入包含有关对象(名词)、动作(动词)、视觉属性(形容词)、空间关系(介词)以及基数和组成等丰富信息。我们的假设是,Muse 模型学会将 LLM 嵌入中的这些丰富的视觉和语义概念映射到生成的图像上;最近的研究表明(Merullo等人,2022),LLM 学到的概念表示与在视觉任务上训练的模型学到的表示大致是线性可映射的。给定一个文本标题输入,我们将其通过冻结的 T5-XXL 编码器,生成一个 4096 维的语言嵌入向量序列。这些嵌入向量经过线性投影,投影到我们的 Transformer 模型(基和超分辨率)的隐藏大小。 

2.2. 使用 VQGAN 进行语义标记化

我们模型的核心组件之一是从 VQGAN (Esser 等人,2021b) 模型获得的语义标记的使用。该模型由一个编码器和一个解码器组成,具有一个量化层,将输入图像映射为从学习的码本中获得的标记序列。我们完全使用卷积层构建我们的编码器和解码器,以支持对不同分辨率的图像进行编码。编码器具有多个降采样块,以减小输入的空间维度,而解码器具有相应数量的上采样块,将潜变量映射回原始图像尺寸。对于尺寸为 H x W 的图像,编码的标记大小为 H/f x W/f,其中 f 是降采样比例。我们训练了两个 VQGAN 模型:一个降采样比率 f = 16,另一个降采样比率 f = 8。我们使用尺寸为 256 x 256 像素的图像在基础模型中使用 f = 16 VQGAN 模型获得标记,从而获得了空间尺寸为 16 x 16 的标记。我们使用尺寸为 512 x 512 像素的图像在超分辨模型中使用 f = 8 VQGAN 模型获得标记,相应的标记具有空间尺寸为 64 x 64。如在先前的工作中所提到的 (Esser 等人,2021b),编码后的离散标记捕捉了图像的高层语义,同时忽略了低层噪音。此外,这些标记的离散性质允许我们在输出阶段使用交叉熵损失来预测下一阶段中的掩蔽标记。 

2.3. 基础模型 

我们的基础模型是一个掩码 transformer(Vaswani 等人,2017;Devlin 等人,2018),其中输入是投影的 T5 嵌入和图像标记。我们保留所有文本嵌入不加掩码,并随机掩蔽其中的一部分图像标记(参见第 2.6 节),并将它们替换为特殊的 [MASK] 标记(Chang 等人,2022)。然后,我们将图像标记线性映射到所需的 transformer 输入/隐藏大小的图像输入嵌入,以及学习的 2D 位置嵌入。按照以前的 transformer 架构(Vaswani 等人,2017),我们使用多个 transformer 层,包括自注意块、交叉注意块和 MLP 块来提取特征。在输出层,我们使用 MLP 将每个被掩码的图像嵌入转换为一组与 VQGAN 码本大小相对应的标记,并使用与真实标记的交叉熵损失作为目标。在训练时,基础模型被训练以在每个步骤预测所有被掩码的标记。然而,在推断中,掩码预测是以迭代方式执行的,这显著提高了质量。有关详细信息,请参见第 2.8 节。

2.4. 超分辨模型

我们发现直接预测 512 x 512 分辨率会导致模型关注低级细节而忽略大尺度语义。因此,我们发现使用级联模型是有益的:首先是一个基础模型,生成一个 16 x 16 的潜变量图(对应于 256 x 256 的图像),然后是一个超分辨模型,将基础潜变量图上采样到 64 x 64 的潜变量图(对应于 512 x 512 的图像)。超分辨模型是在基础模型训练之后进行训练的。

如第 2.2 节所述,我们训练了两个 VQGAN 模型,一个在 16 x 16 的潜变量分辨率和 256 x 256 的空间分辨率,第二个在 64 x 64 的潜变量分辨率和 512 x 512 的空间分辨率。由于我们的基础模型输出与 16 x 16 潜变量映射对应的标记,我们的超分辨过程学会了将低分辨率潜变量图 “翻译” 为高分辨率潜变量图,然后通过高分辨率 VQGAN 进行解码,以生成最终的高分辨率图像。这个潜变量图翻译模型也是在文本条件和交叉关注的情况下进行训练,类似于基础模型,如图 4 所示。

(2023|ICML,LLM,标记掩蔽,并行解码)Muse:使用掩蔽生成 Transformer 的文本到图像生成_第2张图片

2.5. 解码器微调 

为了进一步提高我们模型生成精细细节的能力,我们增加了 VQGAN 解码器的容量,增加了更多的残差层和通道,同时保持编码器的容量不变。然后,我们微调新的解码器层,同时保持 VQGAN编码器的权重、码本和 transformer (即基础模型和超分辨模型)冻结。这使我们能够改善我们的视觉质量,而无需重新训练任何其他模型组件(因为视觉标记的 “语言” 保持不变)。这在附录的图 13 中有所显示,我们可以看到,微调后的解码器可以重建更锐利的细节。我们还在附录中提供了微调后的解码器架构的详细信息。

2.6. 可变掩码率

与(Chang等人,2022)一样,我们使用基于余弦调度的可变掩码率来训练我们的模型:对于每个训练示例,我们从一个带有密度函数

的截断反余弦分布中采样一个掩码率 r ∈ [0, 1]。这具有预期的掩码率为 0.64,具有对更高掩码率的强烈偏向。对更高掩码率的偏好使预测问题变得更加困难。与自回归方法相反,自回归方法学习了一些固定顺序的标记的条件分布 P(x_i | x_

2.7. 无分类器引导

我们采用无分类器引导(classifier-free guidance,CFG)(Ho&Salimans,2022)来提高我们的生成质量和文本-图像对齐。在训练时,我们随机选择 10% 的样本,去除文本条件(因此注意力减少到图像标记的自注意力)。在推断时,我们为每个被掩蔽的标记计算一个条件 logit ℓ_c 和一个无条件 logit​​​​​​​ ℓ_u。然后,我们通过从无条件 logit 中移动一定数量 t,即引导尺度,形成最终的 logit ℓ_g:

直观地,CFG 使用多样性置换忠实度。不同于以往的方法,我们通过采样过程线性增加了指导尺度 t,以减少对多样性的打击。这允许早期的标记以更自由地进行采样,具有较低或无指导,但增加了后期标记的条件提示的影响。

我们还利用这一机制来实现负向提示(NegPrompt,2022),通过使用以 “负向提示” 为条件的 logit 来替换无条件的 logit ℓ_u。这鼓励生成的图像具有与正向提示 c 关联的特征,并去除与负向提示 ℓ_u 关联的特征。

2.8. 推理时的迭代并行解码 

我们模型推理时间效率的关键组件是使用并行解码,在单个前向传递中预测多个输出标记。支持并行解码效果的关键假设是马尔可夫属性,即许多标记在给定其他标记的情况下是条件独立的。解码是基于余弦进度表(Chang 等人,2022)执行的,该表选择要在该步骤预测的具有最高置信度的掩蔽标记的一部分。然后,这些标记在剩下的步骤中被设置为未掩蔽,并且掩蔽标记集合适当减小。使用这一过程,我们能够在我们的基础模型中仅使用 24 个解码步骤执行 256 个标记的推理,以及在我们的超分辨率模型中使用 8 个解码步骤执行 4096 个标记的推理,与自回归模型(例如 (Yu 等人,2022))需要的 256 或 4096 步,以及扩散模型(例如 (Rombach 等人,2022; Saharia 等人,2022))的数百步相比。我们注意到,最近的方法,包括渐进蒸馏(Salimans&Ho,2022)和更好的 ODE 求解器(Lu 等人,2022),已经大大减少了扩散模型的采样步骤,但它们尚未在大规模文本到图像生成中得到广泛验证。我们将在将来的工作中与这些更快的方法进行比较,同时注意到相似的蒸馏方法对于我们的模型也是可能的。

(2023|ICML,LLM,标记掩蔽,并行解码)Muse:使用掩蔽生成 Transformer 的文本到图像生成_第3张图片

3. 结果

我们训练了多个基本的 Transformer 模型,其参数大小不同,从 600M 到 3B 参数不等。这些模型中的每一个都接收来自 T5-XXL 模型的输出嵌入,该模型是经过预训练并冻结的,由 4.6B 参数组成。我们最大的基本模型有 3B 参数,包括 48 个 Transformer 层,其中包括文本到图像的交叉注意力和图像标记之间的自注意力。所有基本模型共享相同的图像标记器。我们使用一个包含 19 个 ResNet 块和大小为 8192 的量化码本的 CNN 模型进行标记化。更大的码本大小并没有带来性能改进。超分辨率模型由 32 个多轴 Transformer 层(Zhao 等人,2021)组成,其中包括从连接的文本和图像嵌入到高分辨率图像的交叉注意力和高分辨率图像标记之间的自注意力。这个模型将一系列标记从一个潜在空间转换到另一个:第一个潜在空间是基本模型标记器的潜在空间,包含 16x16 的标记,然后转换到更高分辨率标记器的 64x64 的标记的潜在空间。在标记转换后,高分辨率标记器的解码器用于将其转换为高分辨率图像空间。附录中提供了更多配置的详细信息。 我们在包含 460M 个文本-图像对的 Imagen 数据集上进行训练(Saharia 等人,2022)。训练步骤数为 1M,批量大小为 512,在 512 核 TPU-v4 芯片上进行(Jouppi 等人,2020)。这需要大约 1 周的训练时间。 我们使用 Adafactor 优化器(Shazeer&Stern,2018)以节省内存消耗,这使我们能够在没有模型并行化的情况下容纳 3B 参数模型。在训练期间,我们还避免执行模型权重的指数移动平均(EMA),再次以节省 TPU 内存。为了获得 EMA 的好处,我们每 5000 步进行一次检查点,然后使用衰减因子为 0.7 在检查点权重上离线执行 EMA。这些平均权重形成了最终的基本模型权重。

3.1. 定性性能 

图 6 在定性展示了 Muse 针对具有不同属性的文本提示的能力。图 6 的左上角显示了示例,展示了对基本基数的理解。对于具有非一致基数的对象,Muse 不会多次生成相同对象像素,而是添加上下文变化,以使整体图像更加真实,例如,大象的大小和方向,酒瓶包装颜色,以及网球旋转。图 6 的右上角演示了对多对象构图和相对性的理解。Muse 不会将对象随机放置在不同位置,而是生成保留文本中介词对象关系的图像,例如,on vs under,left vs right 等。 图 6 的中间左侧展示了 Muse 生成涵盖多种风格的图像的能力,既特定于著名艺术家(例如,伦勃朗),又通用于整个风格(例如,波普艺术和中国水墨画)。图 6 的中间右侧展示了 Muse 渲染单词和短语的能力。文本生成与生成大多数其他对象的方式根本不同。模型不是学习对象名称与其特征之间的映射(例如,“大象” 映射到 “大”、“灰色” 和 “吃花生”),而是必须学习不同的方式。它必须学习短语、单词和字母之间的层次理解。 图 6 的左下方演示了 Muse 渲染图像时使用整个文本提示,而不仅仅关注少数显著的词。最后,图 7 显示了 Muse、Dall-E 2(Ramesh 等人,2022)和 Imagen(Saharia 等人,2022)在一些选择性提示下的比较,显示 Muse 与 Imagen 相当,对于许多提示在定性上更好。

然而,正如图 6 的右下角所示,Muse 在生成与某些类型的提示完全对齐的图像方面存在局限性。对于指示应直接呈现长的、多词短语的提示,Muse 倾向于不正确地呈现这些短语,通常导致(不希望的)重复呈现的词语或仅呈现短语的一部分。此外,提示指示高基数对象的倾向通常会导致生成的图像不能正确反映所需的基数(例如,在提示中指定为 10 时仅呈现 7 个酒瓶)。总的来说,Muse 以正确的对象基数呈现的能力随着基数的增加而减小。Muse 的另一种难以应对的提示类型是具有多个基数的提示(例如,“四只猫和三只狗的队伍”)。对于这些情况,Muse 倾向于在呈现中至少有一个基数不正确。

(2023|ICML,LLM,标记掩蔽,并行解码)Muse:使用掩蔽生成 Transformer 的文本到图像生成_第4张图片

3.2. 量化性能

在表 1 和表 2 中,我们展示了我们在 CC3M(Sharma等人,2018)和 COCO(Lin等人,2014)数据集上与其他方法的性能,使用 Frechet Inception Distance(FID)(Heusel等人,2017)来衡量质量和样本多样性,以及 CLIP(Radford等人,2021)分数,用于衡量图像/文本对齐度。对于CC3M 的结果,两个 Muse 模型均在 CC3M上 训练。COCO 的结果是零样本学习,使用了与Imagen(Saharia等人,2022)相同数据集上训练的模型。

我们的 632M 模型在 CC3M 上取得了 SOTA 结果,显著改善了 FID 分数的技术水平,并且也达到了 CLIP 分数的技术水平。我们的 3B 模型获得了 7.88 的 FID 分数,略优于具有相似参数数量的Parti-3B 模型的 8.1 分。我们的 CLIP 分数为 0.32,高于 Imagen 的 0.29 分(当 FID 显著高于 20 时获得)。对于 FID 为 7.27 的情况,Imagen 获得了约 0.27 的CLIP分数(请参见(Saharia等人,2022)中的图 4)。

我们的采样算法(第 2.8节)具有许多超参数,例如指导尺度、采样温度、是否在线性增加采样期间的指导等。我们对这些参数进行了评估扫描。我们发现一些采样参数的子集在帕累托(Pareto)效率的意义上是高效的,即我们不能改善 FID 而不影响 CLIP。这使我们能够研究多样性和图像/文本对齐之间的权衡关系,如图 8 所示。

(2023|ICML,LLM,标记掩蔽,并行解码)Muse:使用掩蔽生成 Transformer 的文本到图像生成_第5张图片

3.2.1. 人工评估

与之前的研究类似(Yu等人,2022年;Saharia等人,2022年),我们进行了并排评估,其中人工评分员被呈现一段文本提示和两张图像,每张图像由不同的文本到图像模型使用该提示生成。评分员被要求通过问题来评估提示-图像对齐性,问题是:“哪张图像更符合描述?”每个图像对都是匿名的,随机排序(左侧与右侧)。评分员可以选择任何一张图像,或表示他们不关心。每个(提示,图像对)三元组由五名独立的评分员评估;评分员通过 Google 内部众包计算团队提供,对 Muse 团队完全匿名。对于呈现给评分员的提示集,我们使用了 PartiPrompts(Yu等人,2022年),这是一个包含 1650 个文本提示的收藏,旨在测量模型在各种类别下的能力。对于两个文本到图像模型,我们将 Muse(3B 参数)与 Stable Diffusion v1.4(Rombach等人,2022年)进行了比较,后者在推理速度方面与 Muse 最相似。对于每个提示,生成了 16 个图像实例,选择了具有最高CLIP 分数(Radford等人,2021年)的那个。稳定扩散图像是通过 CompVis Stable Diffusion v1.4笔记本(CompVis,2022年)生成的。我们要求至少有 3 名评分员一致同意的结果才能计入某个模型的支持。通过这个分析,我们发现,Muse 70.6% 被选择为比 Stable Diffusion 更好的对齐,Stable Diffusion 25.4% 被选择为比 Muse 更好的对齐,而 4% 没有评分员达成共识。这些结果与Muse 具有显着更好的描述匹配能力(2.7 倍)一致。图 9 显示了评分员一致性 3、4 和所有 5 种可能投票的结果细分。所有 5 名评分员认为 Muse 的对齐性比 Stable Diffusion 更好的提示是较大的贡献者。 除了测量对齐性,其他研究作品(Yu等人,2022年;Saharia等人,2022年)还测量了图像的逼真度,通常通过与评分员问题类似的问题,如“哪个图像更逼真?”。然而,我们注意到在检查此类结果时必须小心。尽管问题的目的不是如此,但模型完全模式崩溃,以便生成与提示无关的相同足够逼真的图像,几乎总是会在此问题上表现得比在图像生成过程中考虑提示的模型更好。我们建议这种类型的问题只适用于具有类似对齐性的模型之间。由于 Muse 的对齐性明显优于Stable Diffusion,我们没有通过人工评分员评估逼真度。我们认为这是一个开放研究领域。 

(2023|ICML,LLM,标记掩蔽,并行解码)Muse:使用掩蔽生成 Transformer 的文本到图像生成_第6张图片

3.2.2. 推断速度 

在表 3 中,我们将 Muse 的推理时间与其他几个流行的模型进行了比较。我们在 TPUv4 加速器上对 Parti-3B、Imagen 和 Muse-3B 进行了内部基准测试。对于 Stable Diffusion/LDM,我们使用了最快的报告基准测试(Lambda Labs,2022年),该基准测试是在 A100 GPU 上进行的。对于 Stable Diffusion,我们测试的 TPU 实现并不比 A100 实现更快。我们还报告了 LDM 的推理时间,使用了 250 个迭代,这是在表 2 中实现 FID 所使用的配置。尽管 Muse 的参数数量相当(比Stable Diffusion/LDM 多约 3 倍),但 Muse 明显比竞争性扩散或自回归模型更快。Muse 比Imagen 更快的速度是由于使用离散标记并需要较少的采样迭代。Muse 比 Parti 更快的速度是由于使用并行解码。Muse 比 Stable Diffusion 更快的速度主要归因于需要较少的采样迭代。 

3.3. 图像编辑

通过利用我们的模型可以根据图像标记的任意子集进行条件处理的事实,我们可以直接使用该模型进行各种图像编辑应用,无需额外的训练或模型微调。

3.3.1. 文本引导的修复/扩展

我们的采样过程(第2.8节)为我们提供了文本引导修复和扩展:我们将输入图像转换为一组标记,掩蔽掉与局部区域相对应的标记,然后在未掩蔽的标记和文本提示的条件下对掩蔽的标记进行采样。我们通过多尺度方法集成了超分辨率:对于尺寸为 512x512 的图像,我们首先将其减少到256x256 并将两个图像转换为高分辨率和低分辨率标记。然后,我们掩蔽了每组标记的相应区域。接下来,我们使用并行采样算法修复低分辨率标记。最后,我们在这些低分辨率标记的条件下使用相同的采样算法修复高分辨率标记。我们在图 2 和图 10 中展示了这方面的示例。

(2023|ICML,LLM,标记掩蔽,并行解码)Muse:使用掩蔽生成 Transformer 的文本到图像生成_第7张图片

3.3.2. 零样本无掩蔽编辑

我们在零样本情况下使用 Muse 进行真实输入图像的无掩蔽图像编辑。该方法直接针对(标记化的)图像进行操作,不需要 “反演” 完整的生成过程,与最近的零样本图像编辑技术(利用生成模型的方法)形成对比(Gal等人,2022b;Patashnik等人,2021;Kim等人,2022;Mokady等人,2022)。 首先,我们将输入图像转换为视觉标记。接下来,我们迭代地掩蔽和重新采样一组随机标记,以文本提示为条件。我们可以将这看作类似于吉布斯(Gibbs)采样过程,其中我们固定一些标记并在其条件下重新采样其他标记。这会将标记化的图像移动到给定文本提示的条件分布的典型集中。 

我们使用低分辨率基础模型进行编辑,然后在最终输出上执行超分辨率(在编辑提示的条件下)。在示例中(图2、图11),我们每次迭代重新采样 8% 的标记,共进行 100 次迭代,引导尺度为4。我们还对标记的逻辑值执行前 k(k = 3)采样,以防止与输入相差太远。迭代的特性允许对最终输出进行控制。图 12 显示了一些中间编辑(不包括超分辨率);在这个示例中,用户可能会更喜欢第 50 次或第 75 次迭代而不是最终输出。

(2023|ICML,LLM,标记掩蔽,并行解码)Muse:使用掩蔽生成 Transformer 的文本到图像生成_第8张图片

4. 相关工作

4.1. 图像生成模型

变分自动编码器(Van Den Oord等人,2017年)和生成对抗模型(GANs)已经在图像生成性能方面表现出色,提出了许多变种,包括卷积和 Transformer 架构(例如(Goodfellow等人,2020年;Esser等人,2021年;Karras等人,2019年;Brock等人,2018年;Donahue和Simonyan,2019年))。直到最近,GANs 被认为是最先进的技术。基于渐进去噪原理的扩散模型现在能够以相等或更高的保真度合成图像和视频(Ho等人,2020年;Kingma等人,2021年;Ho等人,2022年)。融合了多种方法原理的混合方法也表现出色(Chang等人,2022年;Lezama等人,2022年),表明各种方法之间存在更多的互补性可以利用。 

4.2. 图像标记化

图像标记化因将大部分计算从输入(像素)空间转移到潜在空间(Rombach等人,2022年),或者实现更有效的损失函数(如分类而不是回归)而变得有用,多种生成模型都能够受益于它(Chang等人,2022年;Lezama等人,2022年;Li等人,2022年)。已经开发了许多标记化方法,例如离散 VAE(Rolfe,2016年)、VQVAE(Van Den Oord等人,2017年)和 VQGAN(Esser等人,2021年),其中后者是性能最好的,因为它结合了感知和对抗性损失以实现出色的重建。ViT-VQGAN(Yu等人,2021年)将 VQGAN 扩展到 Transformer 架构。我们使用 VQGAN而不是 ViT-VQGAN,因为我们发现它对我们的模型表现更好,但需要注意,性能更好的标记化模型并不总是转化为性能更好的文本到图像模型。

4.3. 大型语言模型

我们的工作利用了 T5,这是一个在多个文本到文本任务上进行预训练的大型语言模型(LLM)(Raffel等人,2020年)。已经证明 LLMs(包括 T5、BERT(Devlin等人,2018年)和 GPT(Brown等人,2020年;Radford等人,2019年))学习了强大的嵌入,使其能够进行少样本迁移学习。我们在模型中利用了这一能力。 现代 LLMs 都是基于标记预测任务进行训练(自回归或非自回归)。在这项工作中,我们利用了关于标记预测的能力,使用 transformer 来预测视觉标记。

4.4. 文本-图像模型

利用配对的文本-图像数据被证明是一种强大的表示学习和生成模型学习范式。CLIP(Radford等人,2021年)和 ALIGN(Jia等人,2021年)训练模型来对齐文本和图像嵌入的配对,展现出出色的迁移和少样本能力。Imagen(Saharia等人,2022年)和 Parti(Yu等人,2022年)使用类似的大规模文本-图像数据集(Schuhmann等人,2021年;2022年)来学习如何从文本输入预测图像,取得了在 FID 和人工评估上的出色结果。一个关键技巧是使用无分类器引导(Ho&Salimans,2022年;Dhariwal&Nichol,2021年),它在多样性和质量之间进行权衡。

4.5. 利用生成模型进行图像编辑

GANs 已经被广泛研究,用于图像编辑和操纵能力(参见(Xia等人,2022年)进行调查)。许多技术已经开发在扩散模型上,以实现编辑、个性化和转换到标记空间(Gal等人,2022年;Meng等人,2021年;Ruiz等人,2022年;Kawar等人,2022年;Brooks等人,2022年;Hertz等人,2022年;Mokady等人,2022年)。Dreambooth(Ruiz等人,2022年)和 Imagic(Kawar等人,2022年)涉及到生成模型的微调。ImagenEditor(Wang等人,2022年)将编辑任务构建为文本引导的图像修复,涉及用户指定的掩蔽。 

5. 讨论和社会影响 

Muse 模型证实了(Saharia等人,2022年)的发现,即冻结的大型预训练语言模型在文本到图像生成中作为强大的文本编码器。在我们的初步实验中,我们还尝试从头开始学习语言模型,但发现性能明显比使用预训练 LLM 差,特别是在长提示和罕见词汇方面。我们还展示了基于 Transformer 架构的非扩散、非自回归模型(non-diffusion, non-autoregressive)可以在推理时间上与扩散模型相匹敌,同时更加高效。我们取得了最先进的 CLIP 分数,显示了图像和文本之间的卓越对齐。我们还展示了我们方法的灵活性,具有许多图像编辑应用。

我们认识到生成模型具有多种应用,对人类社会产生多样潜在影响。生成模型(Saharia等人,2022年;Yu等人,2022年;Rombach等人,2022年;Midjourney,2022年)具有显著的潜力来增强人类创造力(Hughes等人,2021年)。然而,众所周知,它们也可以被用来传播虚假信息、骚扰以及各种社会和文化偏见(Franks&Waldman,2018年;Whittaker等人,2020年;Srinivasan&Uchino,2021年;Steed&Caliskan,2021年)。由于这些重要考虑因素,我们选择目前不发布代码或公共演示。

数据集偏见是另一个重要的伦理考虑因素,因为需要大多数是自动策划的大型数据集。这些数据集存在各种潜在问题,如同意和主体意识(Paullada等人,2021年;Dulhanty,2020年;Scheuerman等人,2021年)。许多常用数据集往往反映了负面的社会刻板印象和观点(Prabhu&Birhane,2020年)。因此,训练这些数据集可能简单地放大这些偏见,需要进行大量额外的研究,以了解如何减轻这些偏见并生成不包含这些偏见的数据集:这是一个超出本文范围的非常重要的课题(Buolamwini&Gebru,2018年;Hendricks等人,2018年)。

鉴于上述考虑因素,我们不建议在不考虑各种用例和潜在危害的情况下使用文本到图像生成模型。我们特别警告不要使用这些模型来生成人、人类和面孔。 

附录

VQGAN Architecture

我们的 VQGAN 架构与之前的工作(Esser等人,2021b年)类似。它由多个残差块、下采样(编码器)块和上采样(解码器)块组成。主要的区别在于,我们移除了非局部块,使编码器和解码器完全卷积,以支持不同的图像尺寸。在基本的 VQGAN 模型中,我们在每个分辨率中应用 2 个残差块,基本通道维度为 128。对于微调的解码器,我们在每个分辨率中应用 4 个残差块,还将基本通道维度设置为 256。 

S. 总结

S.1 主要贡献

本文提出了一个文本到图像的 Transformer 模型 Muse,它在离散标记空间中进行了一个带掩码的建模任务:给定从预训练的大型语言模型(LLM)中提取的文本嵌入,Muse 经过训练来并行地预测随机掩蔽的图像标记。

  • 与像素空间扩散模型(例如,Imagen 和 DALL-E2)相比,Muse 使用 VQ-GAN 获得了离散标记并且需要较少的采样迭代,因此效率显著更高。 
  • 与自回归模型(例如,Parti)相比,Muse 由于使用并行解码而更加高效。

(2023|ICML,LLM,标记掩蔽,并行解码)Muse:使用掩蔽生成 Transformer 的文本到图像生成_第9张图片

S.2 架构和方法

本文使用的框架如图 3 所示。

文本标记。本文使用冻结的预训练 T5-XXL 编码器获得文本标记。

图像标记。训练两个具有不同降采样率的 VQ-GAN 来分别获得用于基础模型和超分辨率模型的图像标记。

基础模型

  • 其输入是文本标记和图像标记。文本标记不加掩码,图像标记被随机掩蔽一部分。
  • Transformer 具有多个 transformer 层,包括自注意块、交叉注意块和 MLP 块来提取特征。
  • 在输出层,使用 MLP 将每个被掩码的图像标记转换为一组与 VQGAN 码本大小相对应的标记,并使用与真实标记的交叉熵损失作为目标。
  • 在训练时,基础模型被训练以在每个步骤预测所有被掩码的标记。
  • 在推断时,掩码预测是以并行迭代的方式执行的,这显著提高了质量。

推理时的迭代并行解码。 

  • 支持并行解码(在单个前向传递中预测多个输出标记)的关键假设是马尔可夫属性,即许多标记在给定其他标记的情况下是条件独立的。
  • 解码是基于余弦进度表迭代执行的,该表选择要在该步骤预测的具有最高置信度的掩蔽标记的一部分。然后,这些标记在剩下的步骤中被设置为未掩蔽,并且掩蔽标记集合适当减小。 

你可能感兴趣的:(论文笔记,人工智能)