ViTGAN: Training GANs with Vision Transformers阅读笔记

Task

图像的生成

Key idea

利用transformer作为基本的计算单元，替代传统的CNN，同时基于GAN的基础架构，去完成图像的生成任务。

 面对的问题：

(1) GAN的正则化方法与注意力机制的交互很差，导致训练中稳定性不足

(2) 在鉴别器训练的后期，对抗训练经常受到高方差梯度（或尖峰梯度）的阻碍。高方差梯度会在训练的时候导致震荡，非常影响训练的效果。

解决的方法：

(1) 在判别器中，研究者重新审视了自注意力的 Lipschitz 性质，在此基础上他们设计了一种加强了 Lipschitz 连续性的谱归一化。

(2)对于基于 ViT 的生成器，研究者尝试了多种不同的架构设计并发现了对层归一化和输出映射层的两项关键性修改。

方法

上图是整个方法的Pipeline，也是ViTGAN的架构，包括一个ViT鉴别器和一个基于ViT的生成器。在直接使用原始的ViT会使得训练变得非常的不稳定，作者对生成器和鉴别器都引入了新的技术，用来稳定训练并且促进收敛。

【注】正则化，参考：机器学习中常常提到的正则化到底是什么意思？ - 知乎

【注】Lipschitz continuous 即：存在一个实数 ，使得对于函数上的每对点，连接它们的线的斜率的绝对值不大于这个实数  。最小的  称为该函数的 Lipschitz 常数，意味着函数被一次函数上下夹逼。参考：详细解析深度学习中的 Lipschitz 条件 - 知乎

判别器的改进

在GAN的训练过程中，为了避免GAN发生模式崩溃，通常会采用梯度裁剪的方式来稳定GAN的训练过程，事实上，其他归一化手段也有类似的效果。梯度裁剪中采用的是梯度L1归一化，起到的效果就是使得由梯度塑造的生成器模型(Generator)，在各点的局部具有Lipshitz连续性，参考：深度学习中的Lipschitz约束：泛化与生成模型_qq_16568205的博客-CSDN博客_lipschitz损失

标准点积自注意力层的 Lipschitz 常数可以是无界的，这就会破坏 ViT 中的 Lipschitz 连续性。为了加强 ViT 判别器的 Lipschitz 性质，研究者采用了上述论文中提出的 L2 注意力。参考：The lipschitz constant of self-attention - 搜索结果 - 知乎

因此，为了增强ViT的Lipschitz 性质，点积相似度被替换成了欧几里得距离，并且还关联了投影矩阵的权重，以用于自注意力中的查询和键（key）。公式如下所示：

这个方法很大程度上提升了训练的稳定性。

为了进一步强化 Lipschitz 连续性，研究者还在判别器训练中使用了谱归一化。标准谱归一化是使用幂迭代来估计每层神经网络的投影矩阵的谱范数，然后再使用估计得到的谱范数来除权重矩阵，这样所得到的投影矩阵的 Lipschitz 常量就等于 1。研究者发现，Transformer 模块对 Lipschitz 常数的大小很敏感，当使用了谱归一化时，训练速度会非常慢。类似地，研究者还发现当使用了基于 ViT 的判别器时，R1 梯度惩罚项会有损 GAN 训练。另有研究发现，如果 MLP 模块的 Lipschitz 常数较小，则可能导致 Transformer 的输出坍缩为秩为 1 的矩阵。为了解决这个问题，研究者提出增大投影矩阵的谱范数。他们发现，只需在初始化时将谱范数与每一层的归一化权重矩阵相乘，便足以解决这个问题。参考Transformer也能生成图像，新型ViTGAN性能比肩基于CNN的GAN - 知乎

 谱归一化如下所示：

 生成器的改进

生成器的使用如上图所示，作者做了不同的方式的改进，A是对每一个PE增加了一个W，B是在一开始的时候增加了一个W，类似于Bert的增加一个向量的方式学习，而C是作者最终实验效果比较好的方式，将正常的transformer模块中的Norm网络替换成self-modulated layernorm (SLN)模块。

自调制型层范数（SLN）。新的做法不是将噪声向量 z 作为输入发送给 ViT，而是使用 z 来调制层范数运算。之所以称这样的操作为自调制，是因为该过程无需外部信息。

用于图块生成的隐式神经表征。为了学习从 patch 嵌入到 patch 像素值的连续映射，研究者使用了隐式神经表征。当结合傅里叶特征或正弦激活函数一起使用时，隐式表征可将所生成的样本空间约束到平滑变化的自然信号空间。研究发现，在使用基于 ViT 的生成器训练 GAN 时，隐式表征的作用尤其大。

【注】我在之前的论文阅读中，也读到了这种隐式的表达方式，之前的研究也表明，这种方法也越来越受到重视 。Adversarial Generation of Continuous Images 阅读笔记_baidu_40582876的博客-CSDN博客

实验结果

和之前的一系列工作相比，实验效果还是不错的。更多的视觉的效果 如下：

总的来说，作者的实验表明，在这种结构设计下的Transformer GAN网络已经能够和之前的GAN网络达到持平的水准。

消融实验

作者的这个消融实验让我觉得使用styleGAN2的鉴别器是不是比它自己的效果更好:-)。消融实验还是很有意思的。我个人认为，还差一个Vanilla-ViT的生成器+StyleGAN2的判别器的实验，来验证作者生成器的效果。

 

思考

我个人觉得这篇论文写得还是很好的，作者研究了基于transformer的GAN网络生成的时候，出现的训练不稳定，判别器容易过拟合，生成器不够强等问题，从理论上提出了原因和解决方案。对于训练不稳定，作者将点积相似度被替换成了欧几里得距离，并且还关联了投影矩阵的权重。同时为了进一步增强Lipschitz连续性，作者在判别器训练中使用了谱归一化。对于判别器过拟合，作者使用了重叠图像块的方法。生成器中，作者做了两点改良，用SLN和隐式神经表征 增强 了学习能力。我目前所在做的工作使用的是生成器中的方法A，效果不太好，读完论文之后考虑使用作者的方法去做微调改进。

参考文献 

[1] Kwonjoon Lee, Huiwen Chang, Lu Jiang, Han Zhang, Zhuowen Tu, Ce LiuViTGAN: Training GANs with Vision Transformers. CoRR abs/2107.04589 (2021)