Large scale gan training for high fidelity natural image synthesis

目录

0. 摘要

1. 简介

2. 背景

3. 扩展 GAN

3.1 使用截断技巧权衡多样性和保真度

3.2 总结 

4. 分析

4.1 表征不稳定性：生成器

4.2 表征不稳定性：判别器 

4.3 总结

5. 实验

5.1 在 IMAGENET 上的评估

5.2 在 JFT-300M 上的附加评估

6. 结论

参考

附录

B. 结构细节

S. 总结

S.1 核心思想

S.2 解决办法（谱归一化）

S.3 其他贡献

0. 摘要

尽管最近在生成图像建模方面取得了进展，但从 ImageNet 等复杂数据集成功生成高分辨率、多样化的样本仍然是一个难以实现的目标。为此，我们以迄今为止尝试过的最大规模，训练生成对抗网络，并研究这种规模特有的不稳定性。我们发现，将正交正则化应用于生成器可以使其适用于简单的“截断技巧”，从而可以通过减少生成器输入的方差来精细控制样本保真度和多样性之间的权衡。我们的修改导致模型在类条件图像合成中设定了最新的技术水平。当在 ImageNet 上以 128X128 分辨率进行训练时，我们的模型 (BigGAN) 获得了 IS 166.5 和 FID 7.4，比之前的最佳 IS 52.52 和 FID 18.65 有所提高。

1. 简介

在这项工作中，我们着手缩小 GAN 生成的图像与来自 ImageNet 数据集的真实世界图像在保真度和多样性方面的差距。我们为实现这一目标做出以下三项贡献：

• 我们证明 GAN 从尺度中获益匪浅，训练模型的参数是现有技术的两到四倍，batch 大小是现有技术的八倍。我们介绍了两个简单的、通用的架构更改，以提高可扩展性，并修改正则化方案以改进调节，从而显着提高性能。

• 作为我们修改的副产物，我们的模型变得可以使用“截断技巧”，这是一种简单的采样技术，可以明确、细粒度地控制样本多样性和保真度之间的权衡。

• 我们发现了大规模 GAN 特有的不稳定性，并根据经验对其进行了表征。利用从该分析中得出的见解，我们证明了将新颖技术和现有技术相结合可以减少这些不稳定性，但只有以极大的性能成本才能实现完整的训练稳定性。

2. 背景

生成对抗网络 (GAN) 涉及生成器 (G) 和鉴别器 (D) 网络，其目的分别是将随机噪声映射到样本并区分真实样本和生成样本。形式上，GAN 的目标，在其原始形式中表示为，为如下 min-max 问题找到纳什均衡：

其中 z 是从简单分布 p(z)，例如 N(0,1) 或 U[1,1]，中提取的隐变量。当应用于图像时，G 和 D 通常是卷积神经网络。如果没有辅助稳定技术，这个训练过程是出了名的脆弱，需要微调超参数和架构选择才能工作。 

因此，最近的许多研究都集中在，利用越来越多的经验和理论见解修改普通（vanilla）GAN 以获得稳定性。

• 一项工作侧重于改变目标函数以促使收敛。

• 另一条线侧重于通过梯度惩罚或归一化来约束 D

• 两者都是为了消除使用无界的损失函数，并确保 D 为 G 提供梯度。

与我们的工作特别相关的是谱归一化，它通过使用第一个奇异值的动态估计对其参数进行归一化，在 D 上强制执行 Lipschitz 连续性，从而引入自适应正则化最大奇异方向的后向动态（backwards dynamics）。类似的， Odena 等人分析了 G 的雅可比行列式的条件数，发现性能取决于 G 的调节。Zhang 等人发现在 G 中使用谱归一化可以提高稳定性，从而减少每次迭代的 D 步数。我们扩展这些分析以进一步了解 GAN 训练的表现。

其他作品侧重于架构的选择，例如 SA-GAN 添加了自注意力块以提高 G 和 D 对全局结构建模的能力。ProGAN 通过在一系列分辨率递增的序列中训练单个模型，在单类别设置中训练高分辨率 GAN。

在条件 GAN 中，类别信息可以通过多种方式输入到模型中。在 Odena 等人的作品中，它通过将 1-hot 类别向量连接到噪声向量来提供给 G，目标被修改为，鼓励条件样本最大化相应的由辅助分类器预测的类别概率。de Vries 等人，通过在 BatchNorm 层中为其提供类别条件增益和偏差，来修改把类别条件传递给 G 的方式。在 Miyato & Koyama 的作品中，使用特征与一组学到的类别 embedding 之间的余弦相似性，作为区分真实样本和生成样本的额外证据，从而调节 D，有效地鼓励生成特征与学到的类别原型匹配的样本。

3. 扩展 GAN

在本节中，我们将探索扩大 GAN 训练规模的方法，以获得更大模型和更大批量的性能优势。作为基线，我们采用 Zhang 等人的 SA-GAN 架构，它使用铰链损失（hinge loss）GAN 目标。我们通过类别条件 BatchNorm 向 G 提供类信息，并通过投影向 D 提供类信息。优化设置遵循 Zhang 等人（特别是在 G 中使用 Spectral Norm）并进行了修改，我们将学习率减半并在每个 G 步执行两次 D 步。为了进行评估，我们采用了 Karras 等人之后的 G 权重的移动平均值，衰减为 0.9999。我们使用正交初始化，而之前的工作使用 N(0,0.02I) 或 Xavier 初始化。每个模型都在 Google TPUv3 Pod 的 128 到 512 个内核上进行训练，并在所有设备上计算 G 中的 BatchNorm 统计数据，而不是像通常那样按设备计算。我们发现渐进式增长即使对于我们的 512X512 模型也是不必要的。其他详细信息在附录 C 中。

我们首先增加基线模型的 batch 大小，并立即发现这样做的巨大好处。表 1 的第 1-4 行表明，只需将 batch 大小增加 8 倍即可将最先进的 IS 提高 46%。我们推测这是每 batch 覆盖更多模式的结果，为两个网络提供更好的梯度。这种扩展的一个显着副作用是，我们的模型在更少的迭代中达到更好的最终性能，但变得不稳定并经历完全训练崩溃。我们在第 4 节中讨论了这种情况的原因和后果。对于这些实验，我们报告了崩溃前保存的检查点的分数。

然后我们将每层的宽度（通道数）增加 50%，大约使两个模型中的参数数量增加一倍。这导致 IS 进一步提高 21%，我们认为这是由于模型的容量相对于数据集的复杂性有所增加。将深度加倍最初并没有带来改进——我们稍后在 BigGAN-deep 模型中解决了这个问题，它使用了不同的残差（residual）块结构。

我们注意到，用于 G 中条件 BatchNorm 层的类别 embedding c 包含大量权重。我们没有为每个embedding 使用单独的层，而是选择使用共享的 embedding，它线性投影到每一层的增益和偏差。这减少了计算和内存成本，并将训练速度（达到给定性能所需的迭代次数）提高了 37%。接下来，我们将噪声向量 z 的直跳连接 (skip-z) 添加到 G 的多个层，而不仅仅是初始层。这种设计背后的直觉是允许 G 使用隐空间直接影响不同分辨率和层次结构级别的特征。在 BigGAN 中，这是通过将 z 分成每个分辨率的一个块，并将每个块连接到条件向量 c 来实现的，该向量 c 被投影到 BatchNorm 增益和偏差。在 BigGAN-deep 中，我们使用更简单的设计，将整个 z 与条件向量连接起来，而不将其分成块。以前的工作已经考虑了这个概念的变体；我们的实现是对该设计的一个小修改。Skip-z 提供了大约 4% 的适度性能改进，并将训练速度进一步提高了 18%。

3.1 使用截断技巧权衡多样性和保真度

与需要通过隐层反向传播的模型不同，GAN 可以使用任意先验 p(z)，但之前的绝大多数工作都选择 N(0,I) 或 U[-1,1]。我们质疑这种选择的最优性，并在附录 E 中探索替代方案。

值得注意的是，我们最好的结果来自于使用与训练中不同的隐分布进行采样。使用 z~N(0,I) 训练的模型，并从截断的正态分布中对 z 进行采样（其中超出范围的值被重新采样以落在该范围内），立即提供了 IS 和 FID 的提升。我们称之为截断技巧：通过对幅度高于所选阈值的值进行重采样来截断 z 向量，以减少整体样本多样性为代价提高单个样本质量。图 2(a) 证明了这一点：随着阈值降低，z 的元素被截断为零（隐分布的模式），单个样本接近 G 的输出分布模式。

这种技术允许在给定 G 的样本质量和多样性之间进行细粒度的 post-hoc 选择权衡。值得注意的是，我们可以计算一系列阈值的 FID 和 IS，获得多样性保真度曲线，让人联想到精度-重建曲线。由于 IS 不会惩罚类条件模型中缺乏多样性，因此降低截断阈值会导致 IS 的直接增加（类似于精度）。FID 惩罚缺乏多样性（类似于重建）但也奖励精确度，因此我们最初看到 FID 有适度的改善，但随着截断接近零和多样性减少，FID 急剧恶化。我们的一些较大模型不适合截断，在输入截断噪声时会产生饱和伪影（图 2(b)）。为了抵消这一点，我们试图通过将 G 调节为平滑来强制截断的适用性，以便将 z 的整个空间映射到良好的输出样本。为此，我们求助于正交正则化，它直接强制执行正交条件：

其中 W 是权重矩阵，β 是超参数。 众所周知，这种正则化通常过于局限，因此我们探索了几种旨在放松约束的变体，同时仍然赋予我们的模型所需的平滑度。 我们发现最有效的版本从正则化中删除了对角线项，旨在最小化过滤器之间的成对余弦相似性但不限制它们的范数：

其中， 1 表示所有元素都设置为 1 的矩阵。我们遍历 β 值并选择 10^(-4)，发现这个小的附加惩罚足以提高我们的模型将被截断的可能性。 在表 1 中的运行中，我们观察到在没有正交正则化的情况下，只有 16% 的模型可以截断，而使用正交正则化进行训练时，这一比例为 60%。

3.2 总结 

我们发现当前的 GAN 技术足以扩展到大型模型和分布式、大批量训练。 我们发现我们可以显着改进现有技术并训练高达 512X512 分辨率的模型，而无需像 Karras 等人那样的显式多尺度方法。尽管有这些改进，我们的模型还是经历了训练崩溃，需要在实践中提前停止。 在接下来的两节中，我们将研究为什么在以前的工作中稳定的设置在大规模应用时会变得不稳定。

4. 分析

4.1 表征不稳定性：生成器

许多以前的工作已经从各种分析角度研究了 GAN 的稳定性，但是我们观察到的不稳定性发生在小规模稳定的设置中，因此需要在大规模内进行直接分析。 我们在训练期间监控一系列权重、梯度和损失统计数据，以寻找可能预示训练崩溃开始的指标。 我们发现每个权重矩阵中前三个奇异值 σ0，σ1，σ2 最能提供信息。可以使用 Alrnoldi 迭代法有效地计算它们，该方法扩展了 Miyato 等人使用的幂迭代法，估计额外的奇异向量和值。 一个清晰的模式出现了，如图 3(a) 和附录 F 所示：大多数 G 层都具有良好的谱范数，但有些层（通常是 G 中的第一层，它是过完备的并且不是卷积的）是病态的，具有在整个训练过程中增长并在崩溃时爆炸的谱范数。

为了确定这种病态是崩溃的原因还是仅仅是一种症状，我们研究了对 G 施加额外条件以明确抵消谱爆炸的影响。 首先，我们直接将每个权重的最大奇异值 σ0 正则化为固定值 σ_reg 或第二大奇异值 σ1 的某个比率 r，r·sg(σ1) （使用 sg 停止梯度操作以防止正则化增加 σ1)。 或者，我们采用部分奇异值分解来代替钳位（clamp） σ0。 给定一个权重 W，它的第一个奇异向量 u0 和 v0，以及 σ_clamp（clamp σ0 的值），我们的权重变为： 

其中 σ_clamp 设置为 σ_reg 或 r·sg(σ1)。 我们观察到，无论是否使用谱归一化，这些技术都具有防止 σ0 或 σ0 / σ1 逐渐增加和爆炸的效果，即使在某些情况下它们略微提高了性能，但没有任何组合可以防止训练崩溃。 这一证据表明，虽然调节 G 可能会提高稳定性，但不足以确保稳定性。 因此，我们将注意力转向 D。

4.2 表征不稳定性：判别器 

与 G 一样，我们分析 D 的权重谱以深入了解其行为，然后通过施加额外的约束来寻求稳定训练。 图 3(b) 显示了 D 为 0 的典型图（附录 F 中有更多图）。 与 G 不同，我们看到谱有噪声，σ0 / σ1 表现良好，并且奇异值在整个训练过程中增长但仅在崩溃时跳跃，而不是爆炸。

D 光谱中的尖峰可能表明它周期性地接收非常大的梯度，但我们观察到 Frobenius 范数是平滑的（附录 F），表明这种效应主要集中在前几个奇异方向上。 我们假设这种噪声是通过对抗训练过程进行优化的结果，其中 G 周期性地产生强烈扰动 D 的 batches。 如果这种频谱噪声与不稳定性有因果关系，一个对策是采用梯度惩罚，显式正则化 D 的雅可比矩阵的变化。 我们探索 R1 零中心梯度惩罚：

γ 使用默认的建议强度 10，训练变得稳定并提高了 G 和 D 中谱的平滑度和有界性，但性能严重下降，导致 IS 减少 45%。 减少惩罚可以部分缓解这种退化，但会导致越来越多的不良谱； 即使惩罚强度降低到 1（不会发生突然坍塌的最低强度），IS 也会降低 20%。 用正交正则化、DropOut和 L2 的各种强度重复这个实验，揭示了这些正则化策略的类似行为：对 D 有足够高的惩罚，可以实现训练稳定性 ，但性能代价很高。

我们还观察到 D 的损失在训练期间接近于零，但在崩溃时经历了急剧的向上跳跃（附录 F）。 对这种行为的一种可能解释是 D 过度拟合训练集，记忆训练示例而不是学习真实图像和生成图像之间的一些有意义的边界。 作为 D 记忆的简单测试，我们在 ImageNet 训练集和验证集上评估未崩溃的鉴别器，并测量样本被分类为真实或生成的百分比。 虽然训练准确率始终高于 98%，但验证准确率在 50-55% 的范围内，并不比随机猜测好（无论正则化策略如何）。 这证实了 D 确实在记忆训练集； 我们认为这符合 D 的作用，D 的作用不是明确地泛化，而是提取训练数据并为 G 提供有用的学习信号。附录 G 中提供了额外的实验和讨论。

4.3 总结

识别不稳定性，确保合理的调节证明是训练所必需的，但不足以防止最终的训练崩溃。 可以通过强烈约束 D 来加强稳定性，但这样做会导致性能上的巨大成本。 使用当前的技术，可以通过放松这种约束并允许在训练的后期阶段发生崩溃来实现更好的最终性能，届时模型将得到充分训练以获得良好的结果。

5. 实验

5.1 在 IMAGENET 上的评估

我们在 ImageNet ILSVRC 2012 上以 128X128、256X256 和 512X512 分辨率评估我们的模型，采用表 1 第 8 行中的设置。我们的模型生成的样本显示在图 4，附录 A 中包含其他示例。 我们在表 2 中报告了 IS 和 FID。由于我们的模型能够用样本多样性来换取质量，因此不清楚如何最好地与现有技术进行比较；我们相应地报告了三种设置下的值，附录 D 中有完整的曲线。

• 首先，我们报告了在获得最佳 FID 的截断设置下的 FID/IS 值。

• 其次，我们报告了截断设置下的 FID，我们的模型的 IS 与真实验证数据所获得的相同，因此推断这是对实现的最大样本多样性的可行衡量，同时仍能达到良好的“客观性”水平。

• 第三，我们报告每个模型达到的最大 IS 的 FID，以证明必须权衡多少多样性才能最大限度地提高质量。 在所有这三种情况下，我们的模型都优于 Miya 等人和 Zhang 等人以前取得的最先进的 IS 和 FID 分数。

除了在第一版论文中介绍并用于大多数实验（除非另有说明）的 BigGAN 模型之外，我们还提出了一个使用不同配置的残差块的 4 倍深度模型（BigGAN-deep）。 从表 2 可以看出，BigGAN-deep 在所有分辨率和指标上都大大优于 BigGAN。 这证实了我们的发现扩展到其他架构，并且深度的增加导致样本质量的提高。 BigGAN 和 BigGAN-deep 架构都在附录 B 中描述。

我们观察到 D 过度拟合（记忆）训练集，再加上我们模型的样本质量，提出了一个明显的问题，即 G 是否简单地记住了训练数据。 为了测试这一点，我们在像素空间和预训练分类器网络的特征空间（附录 A）中执行类最近邻分析。 此外，我们在图 8 和图 9 中展示了样本之间的插值和类别插值（其中 z 保持不变）。我们的模型令人信服地在不同的样本之间进行插值，并且其样本的最近邻在视觉上是不同的，这表明我们的模型不只是简单地记住训练数据。

我们注意到，我们的部分训练模型的某些失败模式与之前观察到的不同。 大多数以前的失败都涉及局部伪影（artifacts）、由纹理斑点而不是对象组成的图像、或规范的模式崩溃。 我们观察到类泄漏，其中一个类的图像包含另一个类的属性，如图 4(d) 所示。 我们还发现 ImageNet 上的许多类对于我们的模型来说比其他类更难； 我们的模型在生成狗（占数据集的很大一部分，并且主要通过它们的纹理区分）方面比人群（占数据集的一小部分并且具有更大规模的结构）更成功。 附录 A 中提供了进一步的讨论。

5.2 在 JFT-300M 上的附加评估

为了确认我们的设计选择对更大、更复杂和多样化的数据集有效，我们还在 JFT-300M 的子集上展示了我们系统的结果。 完整的 JFT-300M 数据集包含 300M 张带有 18K 个类别标签的真实世界图像。 由于类别分布严重的拖尾（long-tailed），我们对数据集进行二次采样，仅保留具有 8.5K 个最常见标签的图像。 生成的数据集包含 292M 图像——比 ImageNet 大两个数量级。 对于具有多个标签的图像，每当对图像进行采样时，我们都会随机且独立地对单个标签进行采样。 为了计算在此数据集上训练的 GAN 的 IS 和 FID，我们使用在此数据集上训练的 Inception v2 分类器。 定量结果如表 3 所示。所有模型均以 2048 的 batch 大小进行训练。我们将模型的消融版本（与 SA-GAN 相当，但 batch 大小更大）与“完整的” BigGAN 模型进行了比较，BigGAN 模型使用了应用在 ImageNet 上获得最佳结果的所有的技术（共享 embedding、skip-z 和正交正则化）。 我们的结果表明，即使在相同模型容量（64 个基本通道）下设置这个更大的数据集，这些技术也能显着提高性能。 我们进一步表明，对于这种规模的数据集，我们看到将模型容量扩展到 128 个基本通道的显着额外改进，而对于 ImageNet GAN，增加容量并没有好处。

在图 19（附录 D）中，我们展示了在该数据集上训练的模型的截断图。 与 ImageNet 不同，截断极限 σ ≈ 0 往往会产生最高的保真度分数，当截断值范围为 0.5 到 1 时，我们的 JFT-300M 模型通常会最大化 IS。我们怀疑这至少部分是由于 JFT-300M 标签的类内可变性，以及图像分布的相对复杂性，其中包括具有各种尺度的多个对象的图像。 有趣的是，与在 ImageNet 上训练的模型不同，在没有大量正则化的情况下训练往往会崩溃（第 4 节），在 JFT-300M 上训练的模型在数十万次迭代中保持稳定。 这表明从 ImageNet 转移到更大的数据集可能会部分缓解 GAN 的稳定性问题。

我们在这个数据集上实现的基线 GAN 模型的改进，没有改变底层模型或训练和正则化技术（超出扩展容量），表明我们的发现从 ImageNet 扩展到具有规模和复杂性的数据集，这是迄今为止图像生成模型前所未有的。

6. 结论

我们已经证明，在生成样本的保真度和多样性方面，经过训练可以对多个类别的自然图像进行建模的生成对抗网络，从规模扩大中受益匪浅。 因此，我们的模型在 ImageNet GAN 模型中设定了新的性能水平，大大提高了现有技术水平。 我们还分析了大规模 GAN 的训练行为，根据权重的奇异值表征了它们的稳定性，并讨论了稳定性与性能之间的相互作用。

参考

Brock, Andrew, Jeff Donahue, and Karen Simonyan. "Large scale GAN training for high fidelity natural image synthesis." arXiv preprint arXiv:1809.11096 (2018).

附录

B. 结构细节

BigGAN 模型的结构如上图所示。

S. 总结

S.1 核心思想

BigGAN 的架构如图 15 所示。本文的核心思想是，更大的 batch 以及模型规模（更多的参数），可以带来更好的性能。然而，随着模型规模的扩大，训练的不稳定性也越发明显。

S.2 解决办法（谱归一化）

谱归一化，它通过使用第一个奇异值的动态估计对其参数进行归一化，在 D 上强制执行 Lipschitz 连续性，从而引入自适应正则化最大奇异方向的 backwards dynamics。

对生成器的权重谱研究发现，虽然调节 G 可能会提高稳定性，但不足以确保稳定性。

对鉴别器的权重谱研究发现，虽然可以通过强烈约束 D 来加强稳定性，但这样做会导致性能上的巨大代价。可行的办法是，放松这种约束并允许在训练的后期阶段发生崩溃，然后使用早停来实现更好的最终性能。

此外，作者认为，D 在记忆训练集符合 D 的作用，因为 D 的作用不是明确地泛化，而是提取训练数据并为 G 提供有用的学习信号。

S.3 其他贡献

截断。作者使用截断实现生成多样性和保真度的权衡。

具体做法是，在训练时，从截断的分布（例如，正态分布）中对 z 进行采样（其中超出范围的值被重新采样以落在该范围内）。

然而，一些较大模型不适合截断，在输入截断噪声时会产生饱和伪影。为解决问题，作者使用正交正则化，通过将 G 调节为平滑来增加截断的适用性。

BigGAN-deep。对于 BigGAN，增加每层的宽度（通道数）可以大幅度提升性能。而增加深度并没有带来改进。为解决这个问题，作者提出了 BigGAN-deep，它使用了不同的 residual 块结构。