StyleSwin: Transformer-based GAN for High-resolution Image Generation

摘要

 尽管Transformer在广泛的视觉任务中取得了诱人的成功，但在高分辨率图像生成建模方面，Transformer还没有表现出与ConvNets同等的能力。在本文中，我们试图探索使用Transformer来构建用于高分辨率图像合成的生成对抗性网络。为此，我们认为局部注意力对于在计算效率和建模能力之间取得平衡至关重要。因此，所提出的生成器在基于风格的架构中采用了Swin Transformer。为了实现更大的感受野，我们提出了双重关注，它同时利用了局部窗口和偏移窗口的上下文，从而提高了生成质量。此外，我们表明，提供基于窗口的Transformer中丢失的绝对位置的知识极大地有利于生成质量。所提出的StyleSwan可扩展到高分辨率，粗糙的几何结构和精细的结构都得益于Transformer的强大表现力。然而，在高分辨率合成期间会出现块伪影，因为以块方式执行局部关注可能会破坏空间相干性。为了解决这个问题，我们实证研究了各种解决方案，其中我们发现使用小波鉴别器来检查频谱差异可以有效地抑制伪影。大量实验表明，它优于现有的基于Transformer的GANs，尤其是在高分辨率（例如1024×1024）方面。StyleWin在没有复杂训练策略的情况下，在CelebA HQ 1024上优于StyleGAN，在FFHQ-1024上实现了同等性能，证明了使用Transformer生成高分辨率图像的前景。

1、介绍

 近年来，图像生成模型取得了巨大的进步，其中生成对抗性网络（GANs）可以说在合成高分辨率图像方面提供了最令人信服的质量。虽然早期的尝试侧重于通过适当的正则化或对抗性损失设计来稳定训练动态，但最近著名作品中显著的性能飞跃主要归因于旨在增强建模能力的架构修改，例如采用自我关注、积极的模型缩放或基于风格的生成器。最近，由于Transformer在判别模型中取得了广泛的成功，一些工作试图使用纯Transformer来构建生成网络，希望增加的表现力和对长程依赖关系建模的能力能够有利于生成复杂图像，但高质量的图像生成，尤其是在高分辨率下，仍然具有挑战性。

本文旨在探索在使用Transformer构成用于高分辨率图像生成的竞争性GAN时的关键成分。第一个障碍是抑制二次计算成本，使网络可扩展到高分辨率，例如1024×1024。我们建议利用Swin Transformer作为基本构建块，因为基于窗口的局部注意力在计算效率和建模能力之间取得了平衡。因此，我们可以利用增加的表现力来表征所有图像尺度，而不是针对更高尺度减少到逐点多层感知器（MLP），并且合成可以扩展到高分辨率，例如1024×1024，具有精细的细节。此外，局部注意力引入了局部感应偏差，因此生成器不需要从头开始重新学习图像的规律性。这些优点使得简单的Transformer网络显著优于卷积基线。

 为了与最佳的方法竞争，我们进一步提出了三种结构性建筑改进方案。首先，我们通过在基于风格的架构中使用局部注意力来增强生成模型的能力，在此期间，我们对我们的Transformer GAN的各种风格注入方法进行了实证比较。其次，我们提出了双重注意力，以扩大局部关注带来的有限感受野，其中每一层都关注局部窗口和移位窗口，有效地提高了生成器的容量，而没有太多的计算开销。此外，我们注意到，基于Conv的GANs隐式地利用零填充来推断绝对位置，这是生成的关键线索，但在基于窗口的Transformer中缺少这样的特征。我们建议通过向每一层引入正弦位置编码来解决这一问题，从而可以利用绝对位置进行图像合成。配备了上述技术，被称为StyleSwan的网络开始在256×256分辨率上显示出有利的生成质量。

尽管如此，我们在合成高分辨率图像时会观察到块伪影。我们推测，这些令人不安的伪影的出现是因为以块方式独立计算注意力会破坏空间连贯性。也就是说，虽然在判别任务中被证明是成功的，但在合成网络中应用时，分块注意力需要特殊处理。为了解决这些阻塞伪影，我们实证研究了各种解决方案，其中我们发现，在频域中检查伪影的小波鉴别器可以有效地抑制伪影，使我们基于Transformer的GAN产生视觉上令人满意的输出。

 提出的StyleSwan在多个既定基准上实现了最先进的质量，例如FFHQ、CelebAHQ和LSUN Church。特别是，我们的方法在高分辨率图像合成方面显示出令人信服的质量（图1），与没有复杂训练策略的领先的基于ConvNet的方法相比，实现了有竞争力的定量性能。在CelebA HQ 1024上，我们的方法实现了4.43的FID，优于包括StyleGAN在内的所有先前工作；而在FFHQ-1024上，我们获得了5.07的FID，接近StyleGAN2的性能。

2、相关工作

 高分辨率图像生成

在过去的十年中，图像生成模型得到了快速的改进。在各种解决方案中，生成对抗性网络（GANs）提供了具有竞争力的生成质量。虽然早期的方法侧重于稳定对抗性训练，但最近的突出工作依赖于设计具有增强能力的架构，这大大提高了生成质量。然而，当代基于GAN的方法采用了卷积主干，这被认为在建模能力方面不如Transformer。在本文中，我们有兴趣将新兴的视觉Transformer应用于高分辨率图像生成的GANs。

视觉Transformer

最近Transformer在NLP任务中的成功启发了对视觉Transformer的研究。开创性的工作ViT提出了一种用于图像分类的纯基于Transformer的架构，并展示了Transformer在视觉任务中的巨大潜力。后来，Transformer在一系列具有鉴别力的任务中占据了基准测试的主导地位。然而，Transformer块中的自注意带来了二次计算的复杂性，这限制了它在高分辨率输入中的应用。最近的一些工作通过提出在局部窗口中计算自注意来解决这个问题，从而可以实现线性计算复杂性。此外，分层结构使它们适合作为通用骨干。

 基于Transformer的GANs

最近，研究界开始探索将Transformer用于生成任务，希望增加的表现力有利于生成复杂图像。一种自然的方法是使用Transformer以自回归的方式合成像素，但较慢的推理速度限制了它们的实际应用。最近，一些工作试图提出基于Transformer的GANs，但这些方法中的大多数仅支持高达256×256分辨率的合成。值得注意的是，HiT以在其高分辨率阶段以减少MLP为代价，成功地生成了1024×1024幅图像，因此无法像基于Conv的对应物那样合成高保真度细节。相比之下，我们的StyleSwan可以使用合成精细结构，从而在高分辨率合成方面达到与领先的ConvNets相当的质量。

3、方法

 我们从一个简单的生成器架构开始，如图2（a）所示，该架构接收潜在变量

作为输入，并通过级联的Transformer块逐渐对特征图进行上采样。

 由于二次计算的复杂性，在高分辨率特征图上计算全部注意力是负担不起的。我们认为，局部关注是实现计算效率和建模能力之间权衡的好方法。我们采用Swin transformer作为基本构建块，在非重叠窗口中局部计算多头自注意（MSA）。为了支持跨相邻窗口的信息交互，Swin transformer在备选块中使用了移位窗口分区。具体而言，给定层l的输入特征图

，连续的Swin块操作如下：

其中W-MSA和SW-MSA分别表示规则和移位窗口划分下基于窗口的多头自注意，LN表示层归一化。由于这种逐块关注导致了相对于图像大小的线性计算复杂性，因此该网络可扩展到高分辨率生成，其中精细结构也可以由这些有能力的Transformer建模。

由于鉴别器严重影响对抗性训练的稳定性，我们选择直接使用基于Conv的鉴别器。在我们的实验中，我们发现，在这种基线架构下，简单地用Transformer块代替卷积，由于模型容量的提高，可以产生更稳定的训练。然而，这种天真的架构并不能使我们基于Transformer的GAN与现有技术相竞争，因此我们进行了进一步的研究，介绍如下。

 风格注入

我们首先通过使生成器适应基于样式的架构来增强模型能力，如图2（b）所示。我们学习一个非线性映射

来将潜在编码

z

从

Z

空间映射到

W

空间，这指定了注入到主合成网络中的样式。我们研究了以下风格的注入方法：

AdaNorm在归一化后调节特征图的统计数据（即均值和方差）。我们研究了多种归一化变体，包括实例归一化（IN）、批处理归一化（BN）、层归一化（LN）和最近提出的RMSnorm。由于RMSNorm去除了LN的均值中心，我们只能从W代码中预测方差。

调制MLP：不调制特征图，也可以调制线性层的权重。具体而言，我们重新缩放Transformer块内前馈网络（FFN）内各个通道的权重大小。根据[30]，这种类型的加入比AdaNorm更快。

•交叉关注：受解码器Transformer的启发，我们探索Transformer特定的风格注入，其中Transformer额外关注从W空间派生的style token。这种交叉关注的有效性也在中得到了验证。

 表1显示，所有上述类型的注入方法都显著提高了生成建模能力，除了AdaBN的训练不收敛之外，因为高分辨率合成的批量大小受到了影响。相比之下，AdaNorm带来了更充分的风格注入，这可能是因为网络可以在注意力块和FFN中两次利用风格信息，而调制的MLP和交叉注意力使用了一次风格信息。出于效率考虑，我们没有进一步研究调制MLP和交叉注意力的混合。此外，与AdaBN和AdaLN相比，AdaIN提供了更精细、更充分的特征调制，因为特征图是独立归一化和调制的，所以我们在以下实验中默认选择AdaIN。

 双重注意力机制

然而，使用局部注意力会牺牲对长程依赖关系建模的能力，这对捕捉几何图形至关重要。设Swin块使用的窗口大小为

\kappa \times \kappa

，则由于移位窗口策略，使用一个以上的Swin块，感受野在每个维度上增加

\kappa

。假设我们使用Swin块来处理64×64的特征图，并且我们默认选择

\kappa = 8

，那么需要

64/\kappa = 8

个Transformer块来跨越整个特征图。

为了实现扩大的感受野，我们提出了双重关注，允许单个Transformer块同时关注局部窗口和移位窗口的上下文。如图2（c）所示，我们将

h

个注意力头分为两组：前半部分头执行规则窗口注意力，而后半部分计算偏移窗口注意力，这两个结果被进一步连接以形成输出。具体来说，我们用

x_w

和

x_{sw}

分别表示正则和移位窗口划分下的非重叠补片，即

，则双重注意公式为：

  其中

是用于混合头以输出的投影矩阵。方程2中的注意力头部可以计算为：

其中

分别是第

i

个头的查询、密钥和值投影矩阵。 我们可以得出，每个维度的感受野增加了2.5

\kappa

，增加了一个额外的双重注意力块，这可以更有效地捕捉更大的上下文。尽管如此，对于64×64的输入，现在需要4个Transformer块来覆盖整个特征图。

局部-全局位置编码

 默认Swin块采用的相对位置编码（RPE）对像素的相对位置进行编码，并已被证明对判别任务至关重要。理论上，具有RPE的多头局部注意力层可以表示窗口大小内核的任何卷积层。然而，当用使用RPE的Transformer替换卷积层时，有一件事很少被注意到：ConvNets可以通过利用零填充的线索来推断绝对位置，但使用RPE在Swin块中缺少这种特征。另一方面，让Transformer知道绝对位置是至关重要的，因为特定组件（例如嘴）的合成在很大程度上取决于其空间坐标。

 有鉴于此，我们建议在每个尺度上引入正弦位置编码，如图2（b）所示。具体地，在比例上采样之后，特征图被添加有以下编码：

  其中

\omega_k = 1/10000^{2k}

，

(i,j)

表示2D位置。我们使用SPE而不是可学习的绝对位置编码，因为SPE允许平移不变性。在实践中，我们通过完全使用RPE和SPE来充分利用它们：在每个Transformer块中应用的RPE提供了本地上下文中的相对位置，而在每个规模上引入的SPE则告知了全局位置。

 3.2、 高分辨率合成中的块伪影

 在使用上述架构合成256×256图像时，虽然实现了最先进的质量，但直接将其用于更高分辨率的合成，例如1024×1024，会带来如图3所示的阻塞伪影，严重影响感知质量。请注意，这些绝不是由转置卷积引起的复选板伪影，因为我们使用双线性上采样，然后使用抗混叠滤波器。

我们推测阻塞伪影是由转换器引起的。为了验证这一点，我们删除了从64×64开始的注意力算子，并仅使用MLP来表征高频细节。这一次我们获得了无人工制品的结果。更进一步地，我们发现这些伪像表现出与局部注意力的窗口大小具有强相关性的周期性模式。因此，我们确信是窗口处理破坏了空间相关性并导致了阻塞伪影。为了简化，可以考虑图4中的1D情况，其中注意力是在跨步窗口中局部计算的。对于连续信号，逐窗口局部注意力可能会产生不连续的输出，因为在softmax操作之后，同一窗口内的值往往变得一致，因此相邻窗口的输出看起来相当不同。2D情况类似于由逐块编码引起的JPEG压缩伪影。

3.3、人工压缩

接下来，我们将讨论一些抑制阻塞伪影的解决方案。

无伪影生成器

我们首先尝试通过改进生成器来减少伪影。

•token共享。阻塞伪像的出现是因为注意力计算在不同的窗口中使用的键和值发生了突然的变化，所以我们建议以类似HaloNet的方式使窗口具有共享令牌。然而，工件仍然是显而易见的，因为总是存在特定窗口专用的令牌。 •从理论上讲，滑动窗口关注应该导致无伪影的结果。请注意，使用滑动注意力训练生成器的成本太高，因此我们只采用滑动窗口进行推理。 •在精细尺度上减少到MLP。正如[67]一样，人们可以消除自我关注，纯粹依靠逐点MLP进行精细结构合成，代价是牺牲对高频细节建模的能力。

伪影抑制鉴别器

事实上，我们在256×256分辨率的早期训练阶段观察到了阻塞伪影，但随着训练的进行，它们逐渐消失。换句话说，尽管基于窗口的注意力容易产生工件，生成器确实有能力提供无工件的解决方案。伪影困扰着高分辨率合成，因为鉴别器无法检查高频细节。这启发我们使用更强的鉴别器来抑制伪影。 •贴片鉴别器

具有有限的感受野，可用于专门惩罚局部结构。实验表明，使用补丁鉴别器可以部分抑制阻塞伪影。 •总变化退火

为了提倡平滑输出，我们在训练开始时应用大的总变化损失，旨在抑制网络产生伪影的趋势。然后，在接近训练结束时，体重的损失线性衰减为零。尽管可以完全去除伪影，但这种手工制作的约束优于平滑的结果，并不可避免地影响高频细节的分布匹配。

•小波鉴别器

如图5所示，在谱域中可以很容易地区分周期性伪影模式。受此启发，我们使用小波鉴别器来补充我们的空间鉴别器，并在图6中说明了其架构。鉴别器对输入图像进行分层下采样，并在离散小波分解后在每个尺度上检查相对于真实图像的频率差异。这样的小波鉴别器在对抗块伪影方面工作得非常好。同时，它不会给分布匹配带来任何副作用，有效地引导生成器生成丰富的细节。 表2比较了上述伪影抑制方法，表明有四种方法可以完全去除视觉伪影。然而，滑动窗口推理受到训练-测试差距的影响，而MLP无法在高分辨率阶段合成精细细节，这两者都会导致更高的FID分数。另一方面，随着退火的总变化仍然使FID劣化。相比之下，小波鉴别器实现了最低的FID分数，并产生了最令人视觉愉悦的结果。

4、实验

数据集

我们在以下数据集上验证了我们的StyleSwan：CelebA HQ、LSUN Church和FFHQ。CelebA HQ是CelebA数据集的高质量版本，其中包含30000张1024×1024分辨率的人脸图像。FFHQ是用于高分辨率图像生成的常用数据集。它包含70000张高质量的人脸图像，年龄、种族和背景的变化更大，眼镜、太阳镜、帽子等配件的覆盖范围更广。我们在FFHQ和CelebA HQ上合成了256×256和1024×1024分辨率的图像。LSUN Church包含了大约126000幅不同建筑风格的教堂图像，我们在这些图像上进行了256×256分辨率的实验。

评估协议

我们采用Frechet Inception Distance（FID）作为定量度量，它测量生成的图像与真实图像之间的分布差异。FID分数越低，表示发电质量越好。对于FFHQ和LSUN Church数据集，我们从原始数据集中随机抽取50000张图像作为验证集，并计算验证集和50000张生成图像之间的FID。对于CelebA HQ，我们计算了30000个生成的图像和所有训练样本之间的FID。

4.2、实现细节

 在训练过程中，我们使用Adam微分器，其中

\beta_1 = 0.0

，

\beta_2 = 0.99

。根据TTUR，我们分别为生成器和鉴别器设置了不平衡学习率5e−5和2e−4。我们使用具有

R_1

梯度惩罚的标准非饱和GAN损失来训练我们的模型，并通过在鉴别器上应用频谱归一化来稳定对抗性训练。默认情况下，我们使用小波鉴别器报告所有结果，如图6所示。使用8个32GB V100 GPU，我们能够将32幅图像拟合为一个训练批次，用于256×256分辨率的训练，并且批次大小在1024×1024分辨率上减少到16幅。为了与先前的工作进行公平的比较，我们在FFHQ-256和CelebA HQ 256数据集上报告了具有平衡一致性正则化（bCR）的FID，损失权重

\lambda_{real}=\lambda_{fake}=10

。我们没有观察到在更高分辨率下使用bCR的性能增益。请注意，我们不采用复杂的训练策略，如路径长度和混合正则化，因为我们希望对整洁的网络架构进行研究。

4.3、主要的结果

定量结果

 我们与最先进的基于Conv的GANs以及最近的基于Transformer的方法进行了比较。如表3所示，我们的StyleSwan在所有256×256合成中都获得了最先进的FID分数。特别是，在FFHQ和LSUN Church数据集上，StyleWin的表现优于StyleGAN2。除了在分辨率256×256上取得令人印象深刻的结果外，所提出的StyleSwan在高分辨率图像生成方面表现出强大的能力。如表4所示，我们以1024×1024的分辨率评估了FFHQ和CelebA HQ上的模型，其中所提出的StyleSwan也展示了最先进的性能。值得注意的是，我们在CelebA HQ 1024数据集上获得了创纪录的4.43的FID分数，同时在不涉及复杂训练策略或额外正则化的情况下，显著缩小了与领先方法StyleGAN2的差距。此外，StyleWin在1024×1024分辨率上以较大的优势优于基于Transformer的方法HiT，证明了高分辨率阶段上的自注意力有利于高保真细节合成。

定性结果

图7显示了StyleWin在FFHQ和CelebA HQ上生成的1024×1024分辨率的图像样本。我们的StyleWin在以1024×1024的分辨率合成不同年龄、背景和视点的不同图像时显示出令人信服的质量。在人脸建模之上，我们在图8中显示了LSUN Church的生成结果，表明StyleSwan能够对复杂的场景结构进行建模。全局几何的相干性和高保真度细节都证明了在所有分辨率中使用Transformer的优势。 

消融研究

为了验证拟议组件的有效性，我们在表5中进行了消融研究。与基线架构相比，由于样式注入带来的模型容量增强，我们观察到FID有了显著改进。双重注意力使每一层同时利用更大的上下文，并进一步降低FID分数。小波鉴别器有效地抑制了阻塞伪影，同时对高频进行了更强的监督，从而大大提高了FID的性能。在我们的实验中，当采用小波鉴别器时，我们观察到更快的对抗性训练。此外，在每一代尺度上引入正弦位置编码（SPE）有效地减少了FID。采用更大的模型带来了轻微的改进，并且目前Transformer的模型容量似乎不是瓶颈。从表5中我们可以看出，bCR将FID显著提高了2.69，这与最近的研究结果相吻合，即数据增强在基于Transformer的GAN中仍然至关重要，因为Transformer需要数据并且容易过拟合。然而，我们没有观察到它对更高分辨率的影响，例如1024×1024，并且我们将高分辨率合成的正则化方案留给未来的工作。

 参数和吞吐量

在表6中，我们与StyleGAN2比较了1024×1024合成的模型参数和FLOPs的数量。虽然我们的方法有更大的模型尺寸，但它实现了比StyleGAN2更低的FLOPs，这意味着该方法实现了有竞争力的生成 这意味着该方法以较少的理论计算成本实现了具有竞争力的生成质量。

5、结论

 我们提出了StyleWin，一种用于高分辨率图像生成的基于Transformer的GAN。使用局部注意力可以有效地进行计算，同时获得最大的建模能力，因为感受野在很大程度上被双重注意力补偿。此外，我们发现基于Transformer的GANs缺少一个关键特征——生成器不知道合成下的补丁的位置，因此我们引入SPE进行全局定位。由于提高了表现力，所提出的StyleSwan在256×256数据集上始终优于领先的基于Conv的方法。为了解决高分辨率合成中的块伪影，我们建议使用小波鉴别器来惩罚频谱差异。最终，所提出的StyleSwan在1024×1024的分辨率上提供了令人信服的质量，这是第一次接近性能最好的ConvNets。我们的工作有望激励更多关于在生成建模中使用有能力的Transformer的研究。