Large scale GAN training for high fidelity natural image synthesis

  这篇论文是Heriot-Watt大学和Google DeepMind向ICLR 2019提交的一篇论文，最近的讨论热度非常高，公众号各种推送，这里博主也来细读学习一下。同门粗看了一下，据说文中有很多晦涩的地方，还是老老实实先翻译一遍吧=-=。

原文链接：https://arxiv.org/abs/1809.11096?context=cs.LG

---

摘要

  尽管最近在生成图像建模方面取得了进展，但是想要从ImageNet这样复杂的数据集中成功生成高分辨率和多样化的图像仍然是一个难以实现的目标。为此，我们在迄今为止最大的数据集上训练了GAN，并针对这种大规模研究了它的不稳定性。我们发现，在生成器中使用正交正则化（Orthogonal Regularization）是一个简单的“截断技巧”，通过截断潜在的解空间来很好的控制样本真实性和多样性之间的平衡。 我们的这个改动使模型达到了类条件（class-conditional）图像合成的新高度。在$128\times 128$的ImageNet上训练我们的模型（BigGANs）时，Inception Socre（IS）达到了166.3，Fréchet Inception Distance（FID）达到了9.6，之前最好的结果IS是52.52，FID是18.65。

1. 前言

  近年来，生成图像建模取得了显著的进展，通过直接从数据中学习模型，GAN处于生成高保真、多样化图像研究的前沿。GAN训练是动态的，并且几乎对它设置的每个方面都很敏感（从优化参数到模型结构），但是大量研究产生了经验上和理论上的见解，使得在各种设置下稳定训练成为可能。尽管取得了这些进展，当前条件ImageNet建模最高水平（《Self-attention generative adversarial networks》）的Inception Score仅为52.5，而真实图像的得分为233。
　　在本文中，我们致力于缩小GANs生成图片和ImageNet数据集中真实图片在真实性和多样性之间的差距。我们为了实现这个目标做了以下三方面的贡献：

• 我们证明了GANs极大的受益于scaling。我们模型训练的参数是现有方法的2~4倍，batch size是现有方法的8倍。我们介绍了两种简单通用的能提升稳定性的结构改变，并修改了正则化方案来提升训练，显著地提升了性能。
• 作为我们改动的一个副作用，我们的模型变得适合于“截断技巧”，这是一种允许显示、细粒度控制样本多样性和保真性之间平衡的简单取样技术。
• 我们发现了针对大规模数据集GANs的不稳定性，并从经验上对其进行了特征化。利用得到的见解，我们证明了现有技术和新技术的结合能够降低这种不稳定性，但是训练的稳定性只能通过极大的性能代价来实现。

  我们的改动极大地提升了类条件GANs的性能。当在$128\times 128$图片大小的ImageNet上训练时，先比现有方法，我们的模型（BigGANs）将IS和FID分数分别从52.52和18.65提升到了166.3和9.6。我们同样在$256\times 256$和$512\times 512$的ImageNet上成功训练了BigGANs，IS和FID分数分别为233.0，9.3和241.4和10.9。最后，我们在一个更大的数据集JFT-300M上训练了我们的模型，并证明了我们的改动很好的从ImageNet迁移到了新数据集上。

2. 背景

  GAN包括生成器$G$和判别器$D$，生成器用于将随机噪声映射到样本，判别器用于区分真实和生成的样本。形式上，在GAN的损失函数的初始形式中涉及在双人玩家min-max问题中寻找一个纳什均衡：$$\underset{G}{\min }\underset{D}{\max }{\mathbb{E}}_{x\sim q_{data(x)}}[\log D(x)]+{\mathbb{E}}_{z\sim p(z)}[\log (1-D(G(z)))]\text{\hspace{1em}(1)}$$其中$z\in {\mathbb{R}}^{d_z}$是例如$\mathcal{N}(0,I)$或$\mathcal{U}[-1,1]$这样的分布$p(z)$的潜变量。当把GAN应用到图像时，$G$和$D$通常是卷积神经网络。如果没有辅助的稳定技巧，GAN的训练过程是出了名的脆弱，这需要精确的超参数调整和架构选择才能使其正常工作。
　　因此，最近许多研究都关注与如何对普通GAN进行改进来获得稳定性，产生了大量经验和理论的见解。其中一类方法是对目标函数进行改进来鼓励模型的收敛。另一种方法是通过梯度惩罚来限制生成器$D$，这样既可以抵消无边界损失函数的作用，又可以保证判别器$D$处处为生成器$G$提供梯度。
　　谱归一化和我们的工作是相关的，它通过估计参数的第一个奇异值对参数进行归一化，从而强制实现Lipschitz连续，诱导自适应地调整最大奇异值的方向。另一项相关的工作分析了生成器Jacobian矩阵的条件数，并发现$G$的性能确实和此相关。Zhang等人则发现在生成器中使用谱归一化可以提升稳定性，并允许每次迭代减少D的训练步数。我们在这些分析的基础上进一步深入探索了GAN训练失败的原因。
　　还有一些其他工作关注GAN架构的选择，比如SA-GAN增加了一个自注意力模块来提升$G$和$D$对全局建模的能力。ProGAN通过队单个模型进行一系列分辨率递增的训练，得到在单类情况下的高分辨率GANs。
　　条件GANs可以通过各种方式将类信息送给模型。Odena等人将one-hot类向量和噪声向量concat之后送给$G$，并修改目标函数来鼓励辅助分类器能以最大概率预测样本类别。de Vires和Dumoulin通过在BatchNorm层中提供类条件增益和偏置，改进了给$G$传递类条件的方式。在Miyato和Koyama的论文里，$D$使用特征和一组学习到的类嵌入之间的余弦相似性来区分真实和生成的样本，有效地鼓励了生成样本的特征和学习的类原型之间的匹配。
　　客观评价生成模型是很困难的。许多工作都提出测量样本质量的启发式方法，其中Inception Score（IS）和Frechet Inception Distance（FIS）用的较多。我们使用它们来近似测量样本的质量，并和前人工作做比较。

3. 扩大GANs训练规模

  本节中，我们探讨如何扩大GANs的训练规模，从而获得更大模型和更大batch的性能收益。我们使用Zhang的SA-GAN结构作为基准方法（使用hinge loss作为GAN目标函数）。我们分别通过类条件BatchNorm方式和投影向$G$和$D$提供类信息。优化的设置和Zhang保持一致（特别是在$G$中应用了谱归一化），不过我们减半了学习率，并每个$G$步执行两个$D$步。在测试阶段，我们使用$G$权重的移动均值，衰减率为0.9999。我们使用正交初始化（Orthogonal Initialization），虽然前人工作中用的是$\mathcal{N}(0,0.02I)$或Xavier初始化。每个模型在一张128到512核心的Google TPU v3 Pod上训练，和标准实现不同的是，我们计算所有设备上的BatchNorm统计信息，而不是每个设备单独计算。我们发现就算对于我们最大的$512\times 512$模型，分辨率渐进增长也是不必要的。

  我们以增加基准模型（SA-GAN）的batch size开始，然后立马发现了这样做的极大好处。如Table 1的1~4行所示，简单的将batch size扩大8倍后，基准方法的IS分数提高了46%。我们猜测这是因为大batch覆盖了更多的模式，为两个子网络提供了更好的梯度。增加batch size的一个显著的副作用是我们模型可以在更少的迭代次数后达到更好的最终效果，不过训练过程也变得不稳定，还遇到过训练失败的情况。我们将在第四节讨论这个的原因和分支。对于这些失败的实验，我们在训练崩溃后马上停止训练，并记录崩溃之前检查点记录的得分。
　　我们然后将每一层的宽度（通道数量）增加了50%，相当于每个模型的参数将近增加一倍。这导致模型的IS得分提高了21%，我们猜测这是因为增加后的模型容量变得和数据集复杂度相匹配了。但是对ImageNet分类模型进行参数翻倍，并没有相似的性能提升，反而性能下降了（也就是过拟合）。
　　我们注意到$G$中条件BatchNorm层使用的类嵌入$c$包含大量的权重参数。相较于为每个嵌入使用单独的层，我们选择使用共享嵌入，即线性投影每一层的增益和偏差。这减少了计算和内存开销，提升了37%的训练速度（达到指定性能的迭代次数）。接下来，我们使用一种分层潜在空间的变体，在$G$的多层中都加入噪声向量$z$，而不只是在输入层加入。这个设计背后的直觉是允许$G$使用潜在空间直接影响不同分别率和层次级别的特征。（没懂，隐藏层加入噪声不应该就只是让模型对噪声更加鲁棒吗？）对于我们的结构，这很容易实现，只要将将$z$分解为每个分辨率的一个块，然后和投影到BatchNorm增益和偏置的条件向量$c$相连。已经有人考虑过这个概念的变体了，我们只是对它进行了小部分的修改。这种层次降低了内存和计算的开销（主要是通过减少第一个线性层的参数预算），并提供了约4%的性能提升，18%的训练速度提升。

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

  不像那些需要通过潜在空间反向传播的模型，GANs可以使用任意先验$p(z)$，目前的前人工作主要选择从$\mathcal{N}(0,I)$或$\mathcal{U}[-1,1]$种提取$z$。我们对这个最优选择表示了质疑，并在附录E中探索了其他选择。
　　值得注意的是，我们最好的结果在训练和测试阶段使用了不同的潜在分布。模型使用$z\sim \mathcal{N}(0,I)$训练，并从truncated normal分布（范围以外的值会在范围内重新取样）中提取$z$，这样马上获得了IS和FID分数的上升。我们将之称之为“截断技巧”：通过对大小超过指定阈值的值进行重采样来截断$z$向量，可以以降低全部样本多样性为代价来提高单个样本的质量。 图2(a)说明了这一点：当阈值降低时，$z$的元素将截断至接近零（潜在分布的模式），单个样本接近$G$的输出分布的模式。

  这个技术允许对给定的$G$的样本质量和多样性之间的平衡进行细粒度和事后（post-hoc）选择。尤其是，我们可以对某个范围内的阈值计算FID和IS，得到和PR曲线相似的“多样性-保真性”曲线（图16）。因为IS不会惩罚类条件模型的多样性，所以降低截断阈值会直接导致IS的上升。FID会惩罚多样性，但是也奖励精度，所以在一开始可以看到FID有些许上升，但是随着截断到零和多样性的减少，FID迅速下降。由采样和训练潜在分布不同引起的分布偏移在许多模型中是存在问题的。我们的一些较大的模型不适合截断，当输入截断噪声时会产生饱和的人工纹理（如图2(b)所示）。为了克服这个问题，我们尝试通过使$G$变得平滑来增强截断的适用性，这样$z$的整个空间就会映射到好的输出样本。基于这个目的，我们使用正交正则化，直接强制满足正交条件：$$R_{\beta }(W)=\beta \Vert W^{\top }W-I{\Vert }_F^2\text{\hspace{1em}(2)}$$其中$W$是权重矩阵，$\beta$是超参数。这种正则化通常被认为过于苛刻的，所以我们探索了几种旨在放松约束又同时为我们模型提供期望的平滑性的正交正则化变体形式。我们找到效果最好的方法是从正则化中 移除对角项，目的是最小化滤波器之间的余弦相似度，但又不限制它们的范数：$$R_{\beta }(W)=\beta \Vert W\top W\odot (\text{1}-I){\Vert }_F^2\text{\hspace{1em}(3)}$$其中$\text{1}$表示所有元素为1的矩阵。我们选择$\beta =10^{-4）}$，发现这个小的额外正则化足以提高我们模型适用于截断的可能性。浏览表1，我们注意到如果没有正交正则化，只有16%的模型适用于截断，而有正则化的情况下，这一比例是60%。

3.2 小结

  我们发现当前的GAN技术已经足以支持扩展到大模型和分布式的、大批次的训练。我们可以显著提高当前技术水平，并且不需要像Karras等人那样明确的多尺度方法就能训练出$512\times 512$分辨率的模型。尽管有这些改进，我们的模型还是遇到了训练崩溃，所以在实践中还是需要早停。在接下来的两个小节，我们将研究在前人工作中稳定的设置应用到大规模时会变得不稳定。

4. 分析

4.1 表征不稳定性：生成器

  以往许多工作从各种分析角度对GAN的稳定性进行了研究，但是我们注意到产生不稳定性的设置是小规模的，有必要直接对大规模的设置进行直接分析。和Odena等人一样，我们在训练过程中监视了权重、梯度和损失统计数据，期望能搜寻一种能够预示训练开始崩溃的指标。我们发现每个特征矩阵Top 3的奇异值${\sigma }_0,{\sigma }_1,{\sigma }_2$包含的信息量最大。它们可以通过Alrnoldi iteration方法有效地计算，这种算法是power iteration的扩展，用于估计额外的奇异值和奇异向量。如图3(a)和附录F中所示，出现了一种明显的模式：大多数的$G$层具有良好的谱线模式，但是有些层（通常是$G$的第一层，over-complete and not convolutional）表现不佳，谱线在整个训练过程中不断增长，在崩溃时爆发。

  为了弄清楚这种现象是训练崩溃的原因还是说仅仅是一种症状，我们研究了附加条件对$G$的影响，借此抵消谱线爆炸。首先我们直接对每个权重最大的奇异值${\sigma }_0$进行正则化，朝一个固定值${\sigma }_{reg}$或者朝第二个奇异值的某个比率$r\cdot sg({\sigma }_1)$。或者，我们使用一个局部奇异值分解来代替钳位${\sigma }_0$。给定一个权重$W$，按照入下方式进行限制：$$W=W-\max (0,{\sigma }_0-{\sigma }_{clamp})v_0{u}_0\text{\hspace{1em}(4)}$$其中${\sigma }_{clamp}$设为${\sigma }_{reg}$或$r\cdot sg({\sigma }_1)$。我们注意到无论使不使用谱归一化，这些技术都有防止${\sigma }_0$或$\frac{{\sigma }_0}{{\sigma }_1}$逐渐增长并爆炸，虽然某些情况下，它们只能轻微提升性能，并不能防止训练崩溃。这些证据显示对$G$进行限制可以提升稳定性，但并不能足以保证稳定性。我们接下来来看看判别器$D$。

4.2 表征不稳定性：判别器

  和$G$一样，我们分析了$D$权重的谱线来了解它的行为，然后通过施加额外约束来稳定训练。图3(b)和附录F呈现了$D$中${\sigma }_0$的典型曲线。和$G$不一样的是，我们发现$D$的谱线是有噪声的，$\frac{{\sigma }_0}{{\sigma }_1}$表现良好，奇异值随着训练上升，在训练崩溃时跳变而不是爆炸。
　　$D$中的尖峰可能表明它周期性地接收非常大的梯度，但我们观察到Frobenius范数是平滑的（附录F），这表明这种效应主要集中在前几个奇异值的方向上。我们假设这种噪声是通过对抗训练过程进行优化的结果，在这个过程中，$G$周期性产生对$D$产生强烈干扰的批次。如果这种谱噪声和不稳定有因果关系，一个自然的对策是使用梯度惩罚，显示地对$D$的Jacobian改变进行正则化。我们尝试了Mecheder等人的$R_1$零中心梯度惩罚：$$R_1:=\frac{\gamma }{2}{\mathbb{E}}_{p\mathcal{D}(x)}[\Vert \nabla D(x){\Vert }_F^2]\text{\hspace{1em}(5)}$$$\gamma$默认建议值为10的情况下，训练变得稳定，并且提升了$G$和$D$谱线的平滑度和有界性，但是性能严重下降了，IS分数下降了45%。减小惩罚力度能够部分减轻这种性能下降，但是会导致不良谱线的增加；就算将惩罚强度降到1（不会突然发生崩溃的最低强度），IS也降低了20%。使用正交正则化、DropOut、L2重复进行了这个实验，出现了相似的结果：对$D$进行足够高的惩罚可以稳定训练，但是性能会大幅下降。
　　我们同样注意到在训练过程中$D$的loss接近于0，但是在崩溃处会经历一个急剧的向上突变。这种现象的一个可能的解释是$D$对训练集过拟合了，它记住了训练样本而不是学习到了一些真实图像和生成图像之间有意义的边界。对$D$这种记忆作一个简单的测试，我们在ImageNet训练集和验证集上对未发生崩溃的判别器进行评估，并测试样本分类的正确率。结果训练集准确率一直保持在98%以上，而验证集只有50~55%，和随机猜测差不多。这证明了$D$确实是记住了训练集；我们认为和$D$的角色是相符的，$D$不是明确的泛化，而是提取训练数据为$G$提供一个有效的学习信号。

4.3 小结

  我们发现在对抗训练过程中的稳定性来自$G$和$D$之间的相互作用，而不是只取决于其中某一个。虽然不良条件的症状可以用来跟踪和识别不稳定性，而且保证合理的条件对训练来说是必须的，但是这并不足以防止最终训练崩溃。通过严格限制$D$来增强稳定性是可行的，但是这样做会带来性能上的巨大损失。使用当前的技术，可以通过放松这种条件限制并允许在训练后期发生崩溃可以达到更好的性能，此时模型已经经过充分训练能够获得良好的结果。

5. 实验

5.1 ImageNet评估

  我们使用表1第8行的设置对$128\times 128$、$256\times 256$和$512\times 512$大小的ImageNet ILSVRC 2012数据集上对我们的模型进行了评估。每个分辨率下的结构细节在附录B中提供，结果样本如图4所示，其它的样本在附录A中，IS和FID结果在表2中。由于我们的模型能够平衡样本多样性和质量，但是还不清楚如何和现有技术进行合理的比较；相应地，我们报告了三种设置下的结果，详细的曲线在附录D中可以找到。首先，我们报告了在截断设置下获得的最佳FID分数时FID/IS的值。然后，我们报告了在截断设置下，我们模型的IS分数和真实验证数据获得的IS分数一样的情况下的FID分数，我们推断这是一个可行的测量指标，用来衡量在保持良好“对象”等级下能达到的最大样本多样性。最后，我们报告了每个模型最大IS分数下的FID值，阐述了为了达到最佳质量需要牺牲多少多样性。在上面三种情况中，我们的模型的IS和FID分数大大超出现有方法。
　　我们观察到$D$对训练集发生了过拟合，结合我们模型的样本质量，提出了一个明显的问题：$G$是否也是简单记住了训练点。为了验证这一点，我们在预训练好的分类器的像素空间和特征空间中进行了类-近邻分析（附录A）。另外，我们在图8和图9中给出了样本插值和类插值（其中$z$为常量）。我们的模型在不同样本间进行了有说服力的插值，且样本最近邻的视觉效果也是不同的，这表明我们的模型不是简单记住了训练数据。

我们注意到部分训练好的模型的失效模式和之前发现的不同。之前大多数失效的模式都涉及到局部的人工痕迹、由纹理块而不是对象组成的图像、或者典型的模式崩溃。我们注意到类溢出现象：属于一个类的图像包含另外一个类的特征，如图4(d)所示。我们同样发现对于我们的模型，ImageNet中许多类比其他类更难生成，我们的模型相比人群（数据集的一小部分，拥有更大尺度的结构）能更成功地生成狗（组成数据集的很大一部分，主要靠纹理区分）。更多的讨论见附录A。

5.2 JFT-300M数据集上的额外实验

  为了确认我们设计的选择是否对更大更复杂更多样的数据集也有效，我们同样报告了我们系统在JFT-300M的子集上的结果。完整的JFT-300M数据集包括300M张真实世界的图像，报刊18K个类别。因为类别分布是长尾的，我们对数据集进行了采样，只保留了8.5K最常见的分类。最后数据集包含292M张图像——比ImageNet大了两个数量级。对于多标签的图像，我们每次取样到该图片时随机选择一个标签。为了计算在该数据集下训练好的GAN的IS和FID分数，我们使用在该数据集上训练的Inception V2分类器。定量的结果图表3所示。所有模型训练的额batch size为2048。我们将模型的简化版本（SA-GAN，但是用了更大的batch size）和完整版本进行了比较，后者使用了应用到ImageNet上所有使结果更好的技术（共享嵌入，层次潜在空间，正交正则化）。我们的结果表明这些技术在模型容量一致（64个基础通道）但数据集大得多的情况下也明显提升了性能。我们进一步表明，如果将模型的基础通道数量提升到128,可以获得明显的性能提升，但对于ImageNet GANs额外的模型容量并没有带来益处。
　　在图18中（附录D），我们给出了在这个数据集上训练的模型的截断图。不像ImageNet，在$\sigma \approx 0$出产生最高的保真性，对于我们的JFT-300M模型，当$\sigma$在0.5到1的范围时，IS通常是最大的。我们怀疑这是因为JFT-300M标签的类间多样性，以及图像分布的相对复杂性（在多个尺度包含多个对象）。有趣的是，ImageNet训练的模型在没有强正则化时会发生崩溃，但JFT-300M训练的模型在成千上万次迭代之后还是保持稳定。这意味着移动到超过ImageNet更大的数据集上能部分缓解GAN稳定性的问题。
　　在不改变基础模型或训练和正则化技术（超出扩展容量）的情况下，我们在该数据集上获得的提升意味着我们从ImageNet扩展到这样大尺度和复杂的数据集上的发现对于图像的生成模型是前所未有的。

6. 结论

  我们证明了用于对多类自然图像建模的GAN，不论从生成样本的保真性或多样性的角度，都很大程度上受益于尺度提升。因此，我们的模型在ImageNet GAn中设置了一个新水准，大大提高了技术水平。我们同样分析了大尺度GANs的训练行为，用网络权重的奇异值表征了它们额稳定性，并讨论了稳定性和性能的相互作用。