StyleGAN-T: Unlocking the Power of GANs for Fast Large-Scale Text-to-Image Synthesis
公众号:EDPJ(添加 VX:CV_EDPJ 或直接进 Q 交流群:922230617 获取资料)
目录
0. 摘要
1. 简介
2. StyleGAN-XL
3. StyleGAN-T
3.1. 重新设计生成器
3.2. 重新设计鉴别器
3.3. 变化与文本对齐的权衡
4. 实验
4.1. 与最先进技术的定量比较
4.2. 评估变化与文本对齐
4.3. 定性结果
5. 局限性和未来的工作
参考
S. 总结
S.1 主要思想
S.2 架构
由于大型预训练语言模型、大规模训练数据以及可扩展模型系列(例如扩散模型和自回归模型)的引入,文本到图像的合成最近取得了重大进展。 然而,性能最好的模型需要迭代评估来生成单个样本。 相比之下,生成对抗网络(GAN)只需要一次前向传递。 因此,它们的速度要快得多,但目前在大规模文本到图像合成方面仍远远落后于最先进的技术。 本文旨在确定恢复竞争力的必要步骤。 我们提出的模型 StyleGAN-T 解决了大规模文本到图像合成的具体要求,例如大容量、在不同数据集上的稳定训练、强大的文本对齐以及可控变化与文本对齐的权衡。 StyleGAN-T 比以前的 GAN 有了显着的改进,并且在样本质量和速度方面优于蒸馏扩散模型(以前的快速文本到图像合成的最先进技术)。
在文本到图像的合成中,根据文本提示生成新颖的图像。 由于两个关键想法,这项任务的最新技术最近取得了巨大的飞跃。 首先,使用大型预训练语言模型作为提示编码器可以根据一般语言理解来调节合成(Ramesh 等人,2022;Saharia 等人,2022)。 其次,使用由数亿个图像-标题对组成的大规模训练数据(Schuhmann 等人,2022)使模型能够合成几乎任何可以想象到的东西。
训练数据集的规模和覆盖范围继续快速增加。 因此,文本到图像模型必须能够扩展到大容量以吸收训练数据。 最近大规模文本到图像生成的成功是由扩散模型(DM)(Ramesh 等人,2022;Saharia 等人,2022;Rombach 等人,2022)和自回归模型(ARM)模型推动的,(张等人,2021;Yu 等人,2022;Gafni 等人,2022)这些模型似乎内置了这种特性,并具有处理高度多模态数据的能力。
有趣的是,生成对抗网络(GAN)(Goodfellow 等人,2014 年)——规模较小且多样性较差的数据集中的生成模型的主导家族——在这项任务中并没有特别成功(Zhou 等人,2022 年)。 我们的目标是证明他们可以重新获得竞争力。
GAN 提供的主要好处是推理速度以及通过潜在空间操作对合成结果的控制。特别地,StyleGAN(Karras 等人,2019;2020;2021)具有经过彻底研究的潜在空间,可以对生成的图像进行原则性控制(Bermano 等人,2022;Hárköonen 等人,2020;Shen 等人,2020;Abdal 等人,2021;Kafri 等人,2022)。 虽然在加速 DM 方面取得了显着进展(Salimans & Ho,2022;Karras 等人,2022;Lu 等人,2022),但它们仍然远远落后于仅需要一次前向传递的 GAN。
我们观察到 GAN 在 ImageNet (Deng et al., 2009; Dhariwal & Nichol, 2021) 合成中同样落后于扩散模型,直到在 StyleGAN-XL (Sauer et al., 2021; 2022) 中重新设计了鉴别器架构,使得 GAN 缩小了差距。我们从这个观察得到启发。在第 3 节中,我们从 StyleGAN-XL 开始,重新审视生成器和判别器架构,考虑大规模文本到图像任务的特定要求:大容量、极其多样化的数据集、强大的文本对齐,以及文本可控变化与文本对齐的权衡。
我们在 64 个 NVIDIA A100 上有 4 周的固定训练预算,可用于大规模训练我们的最终模型。 这种限制迫使我们设定优先级,因为预算可能不足以获得最先进的高分辨率结果(CompVis,2022)。 虽然 GAN 扩展到高分辨率的能力众所周知(Wang 等人,2018;Karras 等人,2020),但成功扩展到大规模文本到图像任务仍然没有记录。 因此,我们主要专注于以较低分辨率解决此任务,仅将有限的预算用于超分辨率阶段。
我们的 StyleGAN-T 在 64*64 的分辨率下实现了比当前最先进的扩散模型更好的零样本 MS COCO FID(Lin 等人,2014;Heusel 等人,2017)。 在 256*256 处,StyleGAN-T 将之前 GAN 实现的零样本 FID 减半,但仍落后于 SOTA 扩散模型。 StyleGAN-T 的主要优点包括在文本到图像合成的背景下快速的推理速度和平滑的潜在空间插值,分别如图 1 和图 2 所示。 我们将在 https://github.com/autonomousvision/stylegan-t 上提供我们的实现。
我们的架构设计基于 StyleGAN-XL(Sauer 等人,2022),与原始 StyleGAN(Karras 等人,2019)类似,首先通过映射网络处理正态分布的输入潜在代码 z,以生成中间潜在代码 w。然后,使用 StyleGAN2(Karras 等人,2020)中引入的权重解调技术,使用该中间潜在变量来调制合成网络中的卷积层。 StyleGAN-XL 的合成网络使用 StyleGAN3(Karras 等人,2021)的无失真原始操作来实现迁移等变,即强制合成网络对生成的特征没有偏好位置。
StyleGAN-XL 具有独特的鉴别器设计,其中多个鉴别器头对来自两个冻结的预训练特征提取网络的特征投影(Sauer 等人,2021)进行操作:DeiT-M(Touvron 等人,2021a)和 EfficientNet(Tan & Le) ,2019)。 它们的输出通过随机的跨通道和跨规模混合模块馈送。 这会产生两个特征金字塔,每个特征金字塔具有四个分辨率级别,然后由八个鉴别器头进行处理。 额外的预训练分类器网络用于在训练期间提供指导。
StyleGAN-XL 的合成网络是逐步训练的,一旦当前分辨率停止提高,就会通过引入新的合成层来随着时间的推移提高输出分辨率。 与之前的渐进式增长方法(Karras et al., 2018)相比,判别器结构在训练期间不会改变。 相反,早期的低分辨率图像会根据需要进行上采样以适应鉴别器。 此外,随着添加更多层,已经训练的合成层将被冻结。
对于类条件合成,StyleGAN-XL 将 one-hot 类标签的嵌入串联到 z 并使用投影鉴别器(Miyato & Koyama,2018)。
我们选择 StyleGAN-XL 作为我们的基准架构,因为它在类条件 ImageNet 合成中具有强大的性能(Sauer 等人,2022)。 在本节中,我们逐段修改此基线,依次关注生成器(第 3.1 节)、鉴别器(第 3.2 节)以及变化与文本对齐的权衡机制(第 3.3 节)。
在整个重新设计过程中,我们使用零样本 MS COCO 来衡量更改的效果。 出于实际原因,与第 4 节中的大规模实验相比,测试使用有限的计算预算、更小的模型和更小的数据集; 详细信息请参见附录 A。 我们使用 FID(Heusel 等人,2017)量化样本质量,并使用 CLIP 评分(Hessel 等人,2021)量化文本对齐。 根据现有技术(Balaji 等人,2022),我们使用在 LAION-2B 上训练的 ViT-g-14 模型计算 CLIP 分数(Schuhmann 等人,2022)。
为了将基线模型中的类别条件更改为文本条件,我们使用预训练的 CLIP ViT-L/14 文本编码器(Radford 等人,2021)嵌入文本提示,并使用它们代替类别嵌入。 因此,我们还删除了训练时分类器指导。 这种简单的调节机制与早期的文本到图像模型相匹配(Reed et al., 2016a;b)。 如表 1 所示,在我们的轻量级训练配置中,该基线达到了 51.88 的零样本 FID 和 5.58 的 CLIP 分数。 请注意,我们使用不同的 CLIP 模型来调节生成器和计算 CLIP 分数,这降低了人为增大结果的风险。
StyleGAN-XL 使用 StyleGAN3 层来实现迁移等变。 虽然等变对于各种应用来说都是可取的,但我们并不认为它对于文本到图像的合成是必需的,因为没有一个成功的基于 DM/ARM 的方法是等变的。 此外,等变约束增加了计算成本,并对大规模图像数据集通常使用的训练数据构成了某些限制(Karras 等人,2021)。
出于这些原因,我们放弃等变并切换到 StyleGAN2 主干用于合成层,包括输出跳跃连接和空间噪声输入,以促进低级细节的随机变化。 经过这些更改后,我们的生成器的高级架构如图 3a 所示。 我们还对生成器架构的细节提出了两项更改(图 3b)。
残差卷积。 由于我们的目标是显着增加模型容量,因此生成器必须能够在宽度和深度上进行扩展。 然而,在基本配置中,生成器深度的显着增加会导致训练中的早期模式崩溃。 现代 CNN 架构中的一个重要构建块(Liu et al., 2022;Dhariwal & Nichol, 2021)是一个易于优化的残差块,它可以标准化输入并缩放输出。 根据这些见解,我们对一半的卷积层进行残差,并通过 GroupNorm(Wu & He,2018)包装它们以进行归一化,并通过 Layer Scale(Touvron 等人,2021b)来缩放其贡献。 初始值较低的 10^(-5) 层缩放系数允许卷积层的贡献逐渐减弱,从而显着稳定早期训练迭代。 这种设计使我们能够显着增加总层数,轻量级配置中增加了约 2.3 倍的层,最终模型中增加了 4.5 倍的层。 为了公平起见,我们匹配 StyleGAN-XL 基线的参数计数。
更强的调节。 文本到图像的设置具有挑战性,因为每个提示的变化因素可能有很大不同。 考虑提示 “脸部特写” 和 “美丽的风景”。 第一个提示应生成具有不同眼睛颜色、肤色和比例的面孔,而第二个提示应生成不同地区、季节和白天的风景。 在基于样式的架构中,所有这些变化都必须由每层样式来实现。 因此,文本调节可能需要比简单设置更强烈地影响样式。
在早期测试中,我们观察到输入潜在 z 明显倾向于主导基线架构中的文本嵌入 c_text ,从而导致文本对齐不佳。 为了解决这个问题,我们引入了两项旨在放大 c_text 作用的更改。 首先,根据 Háarköonen 等人(2022)的观察,我们让文本嵌入绕过映射网络。 LAFITE(Zhou et al., 2022)也使用了类似的设计,假设 CLIP 文本编码器为文本调节定义了适当的中间潜在空间。 因此,我们将 c_text 直接拼接到 w 并使用一组仿射变换来生成每层样式 ~s。 其次,我们不是使用结果 ~s 按原样调制卷积,而是将其进一步分成三个相等维度的向量 ~s_(1,2,3) 并计算最终的样式向量为
此操作的关键是逐元素乘法 ⊙,它有效地将仿射变换转变为二阶多项式网络(Chrysos et al., 2020;Chrysos & Panagakis, 2021),增强其表达能力。 DF-GAN 中堆叠的基于 MLP 的条件层(Tao 等人,2022)隐式包含类似的二阶项。
总之,我们对生成器的更改将 FID 和 CLIP 分数提高了 10%,如表 1 所示。
我们从头开始重新设计了鉴别器,但保留了 StyleGAN-XL 的关键思想,即依赖于冻结的、预训练的特征网络并使用多个鉴别器头。
特征网络。 对于特征网络,我们选择使用自监督 DINO 目标(Caron 等人,2021)训练的 ViT-S(Dosovitskiy 等人,2021)。 该网络轻量,评估速度快,并以高空间分辨率编码语义信息(Amir et al., 2021)。 使用自监督特征网络的另一个好处是,它避免了可能损害 FID 的担忧(Kynkàanniemi 等人,2022)。
架构。 我们的鉴别器架构如图 3c 所示。 ViT 是各向同性的,即表示大小(tokens x 通道)和感受野(全局)在整个网络中是相同的。 这种各向同性允许我们对所有鉴别器头使用相同的架构,我们在 transformer 层之间均匀地间隔。 众所周知,多个头是有益的(Sauer 等人,2021),我们在设计中使用了五个头。
我们的鉴别器头是简约的,如图 3c 底部所示。 残差卷积的内核宽度控制标记序列中头部的感受野。 我们发现,应用于 token 序列的一维卷积与应用于空间重塑 token 的二维卷积一样执行,这表明判别任务不会从 token 中保留的任何二维结构中受益。 我们独立评估每个头部中每个 token 的铰链损失(Lim & Ye,2017)。
Sauer 等人 (2021) 使用同步 BatchNorm (Ioffe & Szegedy, 2015) 为判别器提供批量统计数据。 当扩展到多节点设置时,BatchNorm 会出现问题,因为它需要节点和 GPU 之间的通信。 我们使用一种变体来计算小型虚拟批量的批量统计信息(Hoffer 等人,2017)。批量统计信息在设备之间不同步,而是按本地小批量计算。 此外,我们不使用运行统计数据,因此不会引入 GPU 之间的额外通信开销。
数据增强。 我们在判别器中的特征网络之前应用具有默认参数的可鉴别数据增强(differentiable data augmentation)(Zhao et al., 2020)。 当分辨率大于 224*224 像素(ViT-S 训练分辨率)进行训练时,我们使用随机裁剪。
如表 1 所示,这些变化将 FID 和 CLIP 分数进一步显着提高了 40%。 这一显着的改进表明,在处理高度多样化的数据集时,设计良好的判别器至关重要。 与 StyleGAN-XL 判别器相比,我们简化的重新设计速度提高了 2.5 倍,训练速度提高了 1.5 倍。
引导(Dhariwal & Nichol,2021;Ho & Salimans,2022)是当前文本到图像扩散模型的重要组成部分。 它以原则性的方式权衡变化与感知图像质量,更喜欢与文本条件强烈一致的图像。 在实践中,引导极大地改善了结果; 因此,我们想要在 GAN 的背景下模拟它的行为。
引导生成器。StyleGAN-XL 使用预训练的 ImageNet 分类器在训练期间提供额外的梯度,引导生成器生成易于分类的图像。 该方法显着改善了结果。 在文本到图像的背景下,“分类” 涉及为图像添加标题。 因此,这种方法的自然扩展是使用 CLIP 图像编码器而不是分类器。 继 Crowson et al. (2022) 之后,在每次生成器更新时,我们将生成的图像通过 CLIP 图像编码器以获得标题 c_image,并最小化其与归一化文本嵌入 c_text 的平方球面距离:
这个额外的损失项将生成的分布引导至与输入文本编码 c_text 类似的标题图像。 因此,其效果类似于扩散模型中的指导。 图 3d 说明了我们的方法。
CLIP 已在之前的工作中用于指导合成过程中的预训练生成器(Nichol 等人,2022;Crowson 等人,2022;Liu 等人,2021)。 相反,我们在训练期间将其用作损失函数的一部分。 值得注意的是,训练期间过于强大的 CLIP 指导会损害 FID,因为它限制了分布多样性并最终开始引入图像伪影。 因此,L_CLIP在整体损失中的权重需要在图像质量、文本调节和分布多样性之间取得平衡; 我们将其设置为 0.2。 我们进一步观察到,指导仅在 64*64 像素分辨率下才有帮助。 在更高分辨率下,我们将 L_CLIP 应用于随机 64*64 像素裁剪。
如表 1 所示,CLIP 指导将 FID 和 CLIP 分数进一步提高了 20%。
指导文本编码器。 有趣的是,上面列出的使用预训练生成器的早期方法没有报告遇到低级图像伪影。 我们假设冻结的生成器充当抑制它们的先验。 我们在此基础上进一步改进文本对齐。 在我们的第一训练阶段,生成器是可训练的,文本编码器是冻结的。 然后我们引入第二阶段,其中生成器被冻结,文本编码器变得可训练。 我们仅就生成器调节而言训练文本编码器; 鉴别器和指导项(等式 2)仍然从原始冻结编码器接收 c_text。 第二阶段允许非常高的 CLIP 指导权重 50,而不会引入伪影,并且在不影响 FID 的情况下显着改善文本对齐(第 4.2 节)。 与第一阶段相比,第二阶段可以短得多。 收敛后,我们继续第一阶段。
显式截断。 通常,GAN 中使用截断技巧把变化转为更高的保真度(Marchesi,2017;Brock 等人,2019;Karras 等人,2019),其中采样的潜在 w 向着均值插值,该均值对应于给定调节输入。 这样,截断将 w 推向模型性能更好的更高密度区域。 在我们的实现中,w = [f(z); c_text],其中 f(·) 表示映射网络,因此每个提示的均值由下式给出:
因此,我们通过在训练期间跟踪 ~f,并在推理时根据缩放参数 Ψ ∈ [0,1] 在 ~w 和 w 之间插值来实现截断。
我们在图 4 中说明了截断的影响。在实践中,我们依赖于 CLIP 指导和截断的结合。 指导可以改善模型的整体文本对齐方式,而截断可以进一步提高给定样本的质量和对齐方式,从而消除一些变化。
使用第 3 节中开发的最终配置,我们扩展模型大小、数据集和训练时间。 我们的最终模型由 10 亿个参数组成; 增加模型大小时我们没有观察到任何不稳定。 我们在总共 2.5 亿个文本图像对的多个数据集的联合上进行训练。 我们使用类似于 StyleGAN-XL 的渐进式增长,但所有层仍然可训练。 超参数和数据集详细信息列于附录 A。
总训练时间为 4 周,在 64 个 A100 GPU 上使用,批量大小为 2048。我们首先训练第一阶段 3 周(分辨率高达 64*64),然后训练第二阶段 2 天(文本嵌入),最后再次训练第一阶段 5 天(分辨率提升至 512*512)。 相比之下,我们的总计算预算约为稳定扩散的四分之一(CompVis,2022)。
我们使用零样本 MS COCO 将我们的模型的性能与最先进模型进行定量比较:表 2 中的 64*64 像素输出分辨率和表 3 中的 256*256 像素输出分辨率。在低分辨率下,StyleGAN-T 在输出质量方面比所有其他方法的表现都出色,同时评估速度非常快。 在此测试中,我们在最终训练阶段之前使用该模型,即原始生成 64*64 图像的模型。 在高分辨率下,StyleGAN-T 的性能仍然明显优于 LAFITE,但在 FID 方面落后于 DM 和 ARM。
这些结果使我们得出两个结论。 首先,GAN 在低分辨率下的大规模文本到图像合成中可以匹配甚至击败当前的 DM。 其次,强大的超分辨率模型至关重要。 虽然 从 64*64 变为 256*256时,eDiff-I 的 FID 略有下降 (7.60 → 6.95) ,但目前在 StyleGAN-T 中几乎翻倍。 因此,很明显,StyleGAN-T 的超分辨率阶段表现不佳,与当前最先进的高分辨率结果存在差距。 是否可以通过额外的能力或更长时间的培训来弥补这一差距仍然是一个悬而未决的问题。
我们在图 5 中报告了 FID-CLIP 评分曲线。我们将 StyleGAN-T 与强 DM 基线(eDiff-I 的 CLIP 条件变体)和快速、蒸馏的 DM 基线 (SD-distilled) 进行比较(Meng 等人,2022) 。
使用截断,StyleGAN-T 可以将 CLIP 分数提高到 0.305,成功改善文本对齐。 StyleGAN-T 在 FID 和 CLIP 分数上均优于 SD-distilled,但仍落后于 eDiff-I。 关于速度,eDiff-I 需要 32.0 秒来生成样本。 SD-distilled 的速度明显更快,在八个采样步骤中仅需要 0.6 秒即可达到最佳性能。 StyleGAN-T 击败了两个基线,在 0.1 秒内生成了一个样本。
为了隔离文本编码器(text encoder,TE)训练的影响,我们评估了图 6 中的 FID-CLIP 分数曲线。在本实验中,我们使用相同的生成器网络,仅改变文本编码器。 由于生成器已冻结在第二阶段,因此它可以处理原始的和微调的 CLIP 文本嵌入,它们具有相同的 FID 证明了这一点。 微调文本编码器可显着提高 CLIP 分数,而不会影响 FID。
图 2 显示了 StyleGAN-T 生成的示例图像,以及它们之间的插值。 随附的视频以动画形式展示了这一点,并将其与扩散模型进行了比较,表明 GAN 的插值特性仍然相当平滑。
在不同的文本提示之间进行插值非常简单。 对于由中间潜在变量生成的图像 w0 = [f(z), c_text0],我们用新的文本条件 c_text1 替换文本条件 c_text0。 然后我们将 w0 插值到新的潜在 w1 = [f(z), c_text1],如图 7 所示。这种方法类似于 DALL·E2 的在 CLIP 嵌入之间插值的文本差异操作。 之前操纵 GAN 生成图像的工作(Patashnik 等人,2021)通常通过训练过程发现这些潜在方向,该过程需要根据提示重复,因此成本高昂。 有意义的潜在方向是我们模型的内置属性,不需要额外的训练。
通过将不同的样式附加到提示中,StyleGAN-T 可以生成多种样式,如图 8 所示。主题往往会针对固定的潜在 z 进行对齐,我们在随附的视频中展示了这一点。
与也使用 CLIP 作为底层语言模型的 DALL E 2 类似,StyleGAN-T 有时在将属性绑定到对象以及在图像中生成连贯文本方面遇到困难(图 9)。 使用更大的语言模型可能会以运行时速度变慢为代价解决这个问题(Saharia 等人,2022 年;Balaji 等人,2022 年)。
通过 CLIP 损失进行指导对于良好的文本对齐至关重要,但高指导强度会导致图像伪影。 一种可能的解决方案是在不受失真或其他图像质量问题影响的更高分辨率数据上重新训练 CLIP。 在这种情况下,鉴别器中的调节机制也可能值得重新审视。
截断改善了文本对齐,但在两个重要方面与扩散模型中的指导不同。 虽然截断总是朝向单一模式,但指导至少在理论上可以是任意多模式的。 此外,截断会在合成网络之前锐化分布,这可以以任意方式重塑分布,因此可能会撤销任何先前的锐化。 因此,截断的替代方法可能会进一步改善结果。
通过更高的容量和更长的训练来改进超分辨率阶段(即高分辨率层)是未来工作的明显途径。
“个性化”扩散模型的方法已经变得流行(Ruiz 等人,2022;Gal 等人,2022)。 他们对预训练模型进行微调,将唯一标识符与给定主题关联起来,使其能够在新颖的背景下合成同一主题的新颖图像。 这种方法可以类似地应用于 GAN。
Sauer A, Karras T, Laine S, et al. Stylegan-t: Unlocking the power of gans for fast large-scale text-to-image synthesis[J]. arXiv preprint arXiv:2301.09515, 2023.
文本到图像的合成最近取得了重大进展。 然而,性能最好的模型需要迭代评估来生成单个样本。 相比之下,生成对抗网络(GAN)只需要一次前向传递。 因此,它们的速度要快得多,但目前在大规模文本到图像合成方面仍远远落后于最先进的技术。
本文提出 StyleGAN-T 解决了大规模文本到图像合成的具体要求,例如大容量、在不同数据集上的稳定训练、强大的文本对齐以及可控变化与文本对齐的权衡
StyleGAN-T 的主要优点包括在文本到图像合成的背景下快速的推理速度和平滑的潜在空间插值。
StyleGAN-T 在 StyleGAN-XL 架构的基础上重新设计了生成器和鉴别器,如图 3 所示。StyleGAN2 主干用于合成层,包括输出跳跃连接和空间噪声输入,以促进低级细节的随机变化
新的生成器。
新的鉴别器。特征网络使用 ViT-S。ViT 是各向同性的,即表示大小和感受野在整个网络中是相同的。 这种各向同性允许对所有鉴别器头使用相同的架构。如图 3c 顶部所示。
引导生成器。在文本到图像的背景下,“分类” 涉及为图像添加标题。 因此,这种方法的自然扩展是使用 CLIP 图像编码器而不是分类器。每次生成器更新,将生成的图像通过 CLIP 图像编码器以获得标题,并最小化其与归一化文本嵌入的平方球面距离。如图 3d 所示。
指导文本编码器。