Self-Attention Generative Adversarial Networks（SAGAN）理解

介绍

由Han Zhang及其合作伙伴去年提出的是一种新型的GAN架构。该网络架构通过引入注意力机制解决了卷积结构带来的感受野大小限制，并且使得生成器能够自动识别并关注图像的不同区域。实验表明，在Inception score方面，其表现从36.8提升至52.52；同时，在Frechet Inception distance（FID）指标上取得了显著改进。

GAN之前存在的问题： 对于类别中结构约束较少的情况（如海洋、天空等），GAN能够生成较好的结果；但在面对包含较多几何或结构约束的情况（如合成图像中的狗——四足动物）时，则容易出现失败现象。这是因为这些复杂的形状和细节需要建立远距离依赖(distant dependencies)，而卷积操作由于其局部性的特点，在感受野大小的限制下难以捕捉到这些远距离依赖关系。尽管可以通过增加网络深度或扩大卷积核尺寸来一定程度地解决这一问题，但这种做法却会削弱卷积网络在参数量和计算效率方面的优势。

论文的主要贡献：

将自注意力机制整合到GAN框架中，并对卷积结构进行补充设计，在模型中实现了对图像区域内长距离依赖关系以及多层次关联关系的建模能力，并对其进行了可视化验证；通过在判别器和生成器中采用谱归一化方法显著提升了生成器性能表现；在训练过程中采用Two Timescale Update Rule (TTUR)策略，并对判别器设定较高的学习率以实现生成器与判别器的同步更新比例为1:1的优化目标，在此策略下显著加快了模型收敛速度并降低了训练所需时间。

实现原理

自注意力机制生成对抗网络

该方法的主要实现灵感源自之前一篇论文的研究成果：Non-local Neural Networks，如果有兴趣进一步了解相关内容的话可以通过链接直接访问相关文献进行阅读。为了有效连接相距较远的空间区域该模型采用了自注意力模块其架构如图所示

在某个卷积层输出后的特征图（如上图所示为x），若忽略Batch维度的影响，则本模块的具体流程如下

各分支通过不同的一维卷积操作对输入进行了处理。各分支处理后均未改变特征图的空间尺寸，并且各分支处理后的深度信息有所变化：其中f分支与g分支改变了深度信息而h分支则维持了原有的深度信息。
通过融合高度与宽度信息生成了一个[H×W, H×W]的空间注意力矩阵。
将注意力机制与基底空间特征进行矩阵运算生成新的基底空间特征。随后通过一次深度可变层对其进行转换并重构二维空间分布。
整合后的结果即为y = γo + x。

清晰地展示了TensorFlow版本的实现代码如下：( https://github.com/taki0112/Self-Attention-GAN-Tensorflow )

    def attention(self, x, ch):
      f = conv(x, ch // 8, kernel=1, stride=1, sn=self.sn, scope='f_conv') # [bs, h, w, c']
      g = conv(x, ch // 8, kernel=1, stride=1, sn=self.sn, scope='g_conv') # [bs, h, w, c']
      h = conv(x, ch, kernel=1, stride=1, sn=self.sn, scope='h_conv') # [bs, h, w, c]
    
      # N = h * w
      s = tf.matmul(hw_flatten(g), hw_flatten(f), transpose_b=True) # # [bs, N, N]
    
      beta = tf.nn.softmax(s)  # attention map
    
      o = tf.matmul(beta, hw_flatten(h)) # [bs, N, C]
      gamma = tf.get_variable("gamma", [1], initializer=tf.constant_initializer(0.0))
    
      o = tf.reshape(o, shape=x.shape) # [bs, h, w, C]
      x = gamma * o + x
    
      return x

上文简要阐述了自注意力机制的过程

基于f(x)和g(x)的输出生成Attention Map：需要注意的是，在传统的Gram矩阵计算方法（如[1]所示），由于其对称性特点所导致的不同维度之间的关系会被直接反映出来。然而，在本方法中由于其对称性特点所导致的不同维度之间的关系会被反向反映出来。传统的Gram矩阵会得到一个大小为[C,C]的输出结果（如图2所示），这反映了原始特征图中不同通道之间的关联程度；而在此处的方法中则会得到一个大小为[HxW,HxW]的结果（如图3所示），这则反映了原始特征图中各个像素位置之间相互关联的关系程度。在对每个行向量应用了softmax函数后（注意，在这一过程中我们仅对单个行向量进行了归一化处理），每一行就对应了一种特定的空间注意力机制。实际上这反映了一个像素位置与其他所有像素位置之间的影响关系模式；HxW个行也就对应了原始图像中的全部像素位置数量。

h(x)的结果与Attention Map矩阵相乘： 如前所述，在Attention Map中每一行对应一种注意力机制。这里h(x)具有H×W行和C列的数量维度。因此，在进行两个矩阵相乘操作时，则是将h(x)每一个通道的所有像素通过全部注意力机制进行处理得到最终结果。这一过程使得每个像素能够与其他所有像素产生联系，并从而实现了对自身区域的关注。（通过选取简单数值绘制输出矩阵的过程来直观展示其工作原理会更加清晰易懂）

最终特征图的表示： 在模型中所生成的结果表达式可以表示为y等于γ乘以o加上x，在这里γ被定义为一个可学习参数，在初始状态下被设置为零值。值得注意的是，在模型刚开始进入训练阶段时，默认情况下我们希望网络能够先集中于学习局部特征的相关信息；然而随着训练过程逐步推进，在后续的学习过程中网络会逐渐开始利用self-attention机制来捕捉更长距离的空间关系特征。

最后，该网络的生成对抗损失函数采用了合页(Hinge)损失的版本：

稳定训练过程的技巧

主要有两点：
1、在生成器和判别器中均应用了Spectral normalization方法：该方法能够约束参数分布（Lipschitz条件），从而有效解决梯度爆炸等问题，并提高网络的稳定性。最初该方法仅用于判别器，在SAGAN中生成器与判别器同时引入了Spectral normalization。
2、在训练过程中采用了Two Timescale Update Rule（TTUR）策略：通常情况下，在GAN模型中生成器（G）与判别器（D）采用交替训练方式，在一次G的更新后需要多次更新D（常见的比例为一次G更新后五次D更新）。SAGAN模型采用TTUR策略具体实现为：G与D分别使用不同的学习率（其中G的学习率为1\times1e^{-4}、D的学习率为4\times1e^{-4}），这样在每次生成器更新后判别器只需较少的迭代即可收敛。

下图是上面两个技巧的消融实验结果：

实验部分

生成效果

作者不仅对生成效果进行了与残差网络的比较分析之外，在不同层级分别采用了自注意力机制与残差连接的组合策略

可以看出，在上层特征中引入注意力机制能够显著提升性能。从经验来看，生成网络和分类网络相互之间存在反向关系。越靠近输出层的位置上层结构越能够清晰地表现出细节特征。由此可知，在上层引入注意力机制能够取得更好的效果。

Attention机制的可视化

作者旨在阐述Attention机制的工作原理，在最接近输出层的self-attention模块上实现了Attention Map的可视化呈现，并得出了令人瞩目的效果。

如上图所示，在我们之前的分析基础上

此图展示了取五个点对注意力机制的可视化表示。通过观察该图可以看到网络不仅能够区分前景与背景 还能识别出物体中的一些细节结构 在生成图像的过程中也会更加合理

总结

SAGAN网络基于引入自注意力机制的优势，在图像处理领域取得了显著进展。该方法通过构建自适应感受野机制，在提升模型识别能力的同时减少了参数数量以达到更好的平衡效果。

参考博客：<>