论文：https://arxiv.org/abs/2204.12451

代码：https://github.com/NVlabs/FAN

 

导读

近期诸多研究表明，Vision Transformers对输入中的各种扰动表现出很强的鲁棒性。虽然大家普遍认为这一特性的部分原因在于Self-Attention机制，但对此目前仍缺乏系统的探索。基于Vision Transformer中新出现的Visual Grouping的特性，本文研究了Self-Attention在学习鲁棒表征中的作用，表明Self-Attention可能是通过改进的中层表征来促进鲁棒性。

 

本文进一步提出了一系列的Fully Attentional Networks(FANs)，通过结合一种注意力通道来加强鲁棒表征的学习。本文的模型在具有76.8M参数的ImageNet-1k和ImageNet-C参数上达到了最先进的87.1%的精确度和35.8%的mCE，在2个下游任务语义分割和目标检测中同样实现的最先进的准确性和鲁棒性。

 

贡献

与CNN不同的是，Vision Transformer利用Self-Attention对输入的全局关系进行建模，使其在诸多方面具有不可忽略的优势，其中最为重要的优势就是对各种输入扰动的鲁棒性。与对干净图像的标准识别任务不同，几项研究表明，Vision Transformer在扰动鲁棒性方面始终优于ConvNets。

Vision Transformer的强鲁棒性可以归因于其Self-Attention的设计，但这一假设最近受到了ConvNeXt的挑战，ConvNeXt是由标准的卷积模块构建而成的网络，并在泛化和鲁棒性方面都优于Vision Transformer。

这也引出了一个关于Self-Attention在鲁棒泛化中的实际作用问题。

本文旨在探索上述问题的答案。本文的研究始于一个有趣的观察：

在图像分类过程中，Vision Transformer中会出现有意义的对象分割，这表明Self-Attention可能是通过视觉分组促进了中级表征的改善，从而增强了鲁棒性。

在进一步的研究中，作者使用频谱聚类（Spectral Clustering）分析了每个Transformer层的输出Token，其中相似矩阵的显著特征值对应于主要的Clustering成分。结果表明，显著特征值的数量和输入扰动之间存在相关性:

显著特征值的数量和扰动在中间层显著下降，这表明在这些层上，Grouping和鲁棒性是共生关系。

为了理解Grouping现象的潜在原因，作者从Information Bottleneck(IB)的角度对进行了Self-Attention解释。Information Bottleneck(IB)可以看成一种通过最小化潜在特征表示和目标类标签之间的互信息来过滤不重要信息，同时最大化潜在特征与输入原始数据之间的互信息的压缩过程。作者证明了在一定条件下，Self-Attention可以被重写为IB目标的迭代优化步骤，一定程度上解释了Grouping现象。

如上图(a)所示，之前的Vision Transformer采用了一个多头注意力设计，然后使用一个MLP Block来聚合来自多个Head的信息。由于不同的Head倾向于关注物体的不同部分，多头注意力设计本质上形成了Information Bottleneck(IB)。因此，如何聚集来自不同Head的信息尤为重要。而本文作者的目的是提出一种聚合设计，以加强Grouping的共生关系和鲁棒性。因此，本文提出了（b）中的新注意力通道处理设计，通过重加权促进通道选择。与MLP Block中的静态卷积操作不同，注意力设计是动态的，与内容相关，从而可以带来更丰富和更鲁棒的表征。

 

方法

如上图（a）所示，一个标准的ViT首先通过固定的块尺寸将输入的图像均匀地划分为n个输入块，并将每个块编码为一个Token Embedding。然后，所有这些Token被输入Transformer Blocks。每个Transformer Block利用Self-Attention进行Token Mixing，然后使用MLP进行通道特征转换。

而从下图可以观察到，通过输入高斯噪声x~N(0,1)所产生的扰动(由归一化特征范数测量)与显著特征值的数量一起迅速减少。这种观察结果表明了Grouping的共生关系，并提高了对中间块的鲁棒性。其中，(b)中类似的趋势甚至更加明显。(c)进一步比较了ViT和FAN与ResNet-50的噪声衰减，可以看出ResNet50对于噪声的鲁棒性明显低于ViTs。

下图中作者对不同Block上获得的Grouped Tokens进行了可视化，该过程也是通过逐步squeezing out不重要的组件进行可视化分组的过程。

但作者发现，在某种意义上，Visual Grouping也可以被视为某种形式的有损压缩。因此，可以从Information Bottleneck的角度提出以下解释性框架。

给定一个分布X~N(X’,)，其中X是观察到的噪声输入，X'是target clean code，IB寻求一个映射，使包含X中用于预测X'的相关信息。这个目标被表述为以下信息理论优化问题：

上式中，第1项压缩信息，第2项维护相关信息。因此，每个self-attention机制的操作可以解释如下：

即Self-Attention将相似的输入xi聚集到具有cluster结构的表示Z中。ViT中的堆叠SA模块可以广泛地看作是这种优化的迭代重复，一定程度上促进了Grouping和噪声滤波。

因此，在某种意义上，MHSA可以被解释为Information Bottleneck的mixture。

而作者更感兴趣的是在一个固定的通道总数下，Head的数量与鲁棒性之间的关系。如图5所示，拥有更多的Head可以提高表征和鲁棒性。但通道数的减少也会导致精度的下降。最好的权衡是32个通道。

通过以上IBs解释的混合，本文设计了一个通道处理模块，通过不同Head的聚合增强鲁棒表示。设计主要由2个方面驱动:

  • 为了促进更多的组合表征，引入通道权重是可取的，因为有些Head或通道确实可以捕捉到更重要的信息。
  • 重加权机制应该包括对每个通道进行更全面的考虑，以利用分组信息，而不是做出“very local”的通道道聚合决策。

实现上述目标的一个起点是引入一种类似于XCiT的Channel Self-Attention设计。如图6(a)所示，Channel Self-Attention(CA)模块采用Self-Attention设计，将MLP块移动到Self-Attention块中，然后与通道注意力分支的通道注意力矩阵进行矩阵乘法：

与SA不同的是，CA沿着通道维度而不是Token维度计算注意力矩阵，利用特征协方差进行特征变换。聚集相关值较大的强相关特征通道，分离相关值较低的异常特征。这有助于模型过滤掉不相关的信息。在CA的帮助下，该模型可以过滤不相关的特征，从而对前景和背景Token形成更精确的Token聚类。

此外，作者探索了一种具有高计算效率和参数效率的新型自注意力机制：

  • 首先，不计算Token特征之间的相关性矩阵，而是通过在通道维上平均生成Token原型Z。直观地说，Z聚合了Token所表示的每个空间位置的所有通道信息。因此，计算Token特征与Token原型Z之间的相关矩阵是有意义的。
  • 其次，不使用Softmax函数，而是使用Sigmoid函数来规范化注意力权重，然后将其与Token特征点乘，而不是使用MatMul来聚合通道信息。直观地说，并不强迫通道只选择少数重要的Token特征，而是根据空间相关性对每个通道重新加权。

定义如下：

 

实验

如表7所示，原始的ConvNeXt比Swin-Transformer具有更强的鲁棒性，但不如FAN-ViT和FAN-Swin模型。同时，FAN-hybrid具有与FAN-vit和FAN-swin相当的鲁棒性，并且对于干净数据集和已损坏数据集都具有更高的精度，这意味着FAN也可以有效地增强基于CNN的模型的鲁棒性。与FAN-Swin类似，FAN-Hybrid对于大分辨率的输入和密集的预测任务具有高效性，有利于下游任务。

作者根据模型大小将所有模型分成三组，以便进行公平比较。结果如表8所示，从结果中可以清楚地观察到，所有基于变压器的模型都比基于cnn的模型表现出更强的鲁棒性。在所有的模型尺寸下，我们提出的FAN模型都显著地超过了所有其他模型。

作者进一步评估了提出的FAN模型对分割任务的鲁棒性，使用城市景观-c进行评估，它扩展了城市景观验证集与16种自然扰动类型。

还评估了FAN在COCO-C数据集上检测任务的鲁棒性，这是与citcop-C类似的COCO的扩展。结果汇总见表10。

FAN鼓励Token特征形成grouping，并隐式地选择信息特征，这将有利于模型的泛化性能。为了验证这一点，作者直接测试了ImageNet-1K训练模型，以评估其鲁棒性，特别是对于分布外的样本。实验结果汇总见表11。在这些模型中，ResNet-50的泛化能力最弱，而最近的ConvNeXt显著提高了cnn的泛化性能。基于Transformer的模型Swin和RVT的性能与ConvNeXt相当好，并且比ResNet-50要好得多。本文提出的FANs显著优于所有这些模型，这意味着由于无关特征的有效处理，全注意架构有利于学习表示的泛化能力。