FAN(Understanding The Robustness in Vision Transformers)论文解读，鲁棒性和高效性超越ConvNeXt、Swin

在这里插入图片描述
< center >
　　　　　　
< center >

最近的研究表明，Vision Transformers对各种Corruptions表现出很强的鲁棒性。虽然这一特性部分归因于Self-Attention机制，但目前仍缺乏系统的理解。

在本文中研究了Self-Attention在学习鲁棒表征中的作用。本研究是基于Vision Transformer中新出现的Visual Grouping的特性进行深入研究的，Visual Grouping也表明Self-Attention可能是通过改进的中层表征来促进鲁棒性。

本文进一步提出了一系列的Fully Attentional Networks(FANs)，通过结合注意力通道来加强这种能力。作者在各种层次的Backbone上全面验证了FANs设计。本文的模型在具有76.8M参数的ImageNet-1k和ImageNet-C参数上达到了最先进的87.1%的精度和35.8%的mCE。作者还演示了在2个下游任务语义分割和目标检测中实现的最先进的准确性和鲁棒性。

1简介

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

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

在本文中，作者旨在找到上述问题的答案。本文的研究始于一个有趣的观察，即在图像分类过程中，Vision Transformer中会出现有意义的对象分割。这促使作者想知道Self-Attention是否通过视觉分组促进了中级表征的改善(从而增强了鲁棒性)。

作为进一步的研究，作者使用Spectral Clustering分析每个Vision Transformer层的输出Token，其中相似矩阵的显著特征值对应于主要的Clustering成分。研究表明，显著特征值的数量和输入Corruptions之间存在相关性:

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

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

如图2(a)所示，之前的Vision Transformer采用了一个多头注意力设计，然后使用一个MLP Block来聚合来自多个Head的信息。由于不同的Head倾向于关注物体的不同部分，多头注意力设计本质上形成了Information Bottleneck(IB)。因此，如何聚集来自不同Head的信息尤为重要。而本文作者的目的是提出一种聚合设计，以加强Grouping的共生关系和鲁棒性。

如图2(b)所示，提出了一种新的注意力通道处理设计，通过重加权促进通道选择。与MLP Block中的静态卷积操作不同，注意力设计是动态的，与内容相关，从而可以带来更丰富和更鲁棒的表征。

2 Fully Attentional Networks

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

在这里插入图片描述

如图可以看出，可以把Self-Attention分为Token mixing和Channel processing2个部分来看。

在这里插入图片描述

2.1 Channel processing

大多数ViTs采用MLP块将输入Token转换为特征Z：
在这里插入图片描述

该块包含2个线性层和一个GELU层。

2.2 Self-Attention有哪些有趣的特性？

在这里插入图片描述

图3(a)可以观察到，通过输入高斯噪声x~N(0,1)所产生的扰动(由归一化特征范数测量)与显著特征值的数量一起迅速减少。这种观察结果表明了Grouping的共生关系，并提高了对中间块的鲁棒性。

在图3(b)中，类似的趋势甚至更加明显。图3©进一步比较了ViT和FAN与ResNet-50的噪声衰减。可以看出ResNet50对于噪声的鲁棒性明显低于ViTs。

在这里插入图片描述

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

2.3 Information Bottleneck视角看问题？

在这里插入图片描述

Proposition

Under mild assumptions, the iterative step to optimize the objective in Eqn. (3) can be written as:

or in matrix form:

Remark
上述命题通过在普通的Self-Attention和IB之间建立了一个有趣的联系，表明Self-Attention将相似的输入聚集到具有cluster结构的表示Z中。

Self-Attention根据IB原则更新Token特征，其中矩阵K存储临时cluster中心特征，输入特征通过softmax聚类到Token特征中。新的cluster中心特征z为更新后的Token特征输出。

ViTs中的堆叠SA模块可以广泛地看作是这种优化的迭代重复，一定程度上促进了Grouping和噪声滤波。

Multi-head Self-attention

目前许多Vision Transformer架构采用MHSA设计，其中每个Head倾向于关注不同的部分。在某种意义上，MHSA可以被解释为Information Bottleneck的mixture。
在这里插入图片描述

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

2.4 Fully Attentional Networks

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

为了促进更多的组合表征，引入通道权重是可取的，因为有些Head或通道确实可以捕捉到更重要的信息。
重加权机制应该包括对每个通道进行更全面的考虑，以利用分组信息，而不是做出“very local”的通道道聚合决策。
在这里插入图片描述
实现上述目标的一个起点是引入一种类似于XCiT的Channel Self-Attention设计。如图6(a)所示，Channel Self-Attention(CA)模块采用Self-Attention设计，将MLP块移动到Self-Attention块中，然后与通道注意力分支的通道注意力矩阵进行矩阵乘法。

注意力特征转换

FAN块引入以下Channel Self-Attention(CA)来执行特征转换，其表示为：

在这里插入图片描述

2.5 高效的通道Self-attention

传统的自注意计算机制沿通道维度的应用存在2个局限性。

计算开销，引入的CA的计算复杂度与D方成二次比例，其中D为通道维数。对于金字塔模型设计，通道尺寸在顶部阶段变得越来越大。因此，直接应用CA会导致巨大的计算开销。

参数效率低，在传统的SA模块中，通过Softmax操作来增强注意力权值的注意力分布。因此，只有部分通道有助于表征学习，因为大多数通道通过与较小的注意力权值相乘而被削弱。

为了克服这些问题，作者探索了一种具有高计算效率和参数效率的新型自注意力机制。具体来说，提出了2项主要修改：

首先，不计算Token特征之间的相关性矩阵，而是通过在通道维上平均生成Token原型Z。直观地说，Z聚合了Token所表示的每个空间位置的所有通道信息。因此，计算Token特征与Token原型Z之间的相关矩阵是有意义的。

其次，不使用Softmax函数，而是使用Sigmoid函数来规范化注意力权重，然后将其与Token特征点乘，而不是使用MatMul来聚合通道信息。直观地说，并不强迫通道只选择少数重要的Token特征，而是根据空间相关性对每个通道重新加权。

在这里插入图片描述

实际上，通道特性通常被认为是独立的。价值较大的通道不应限制其他通道的重要性。结合这2种设计概念，提出了一种新的通道自注意力方法：

在这里插入图片描述
在这里，σ表示沿着Token维度的Softmax操作，Z^表示Token原型。详细的块架构设计也显示在图6中。实验证明，这种注意力机制在减少计算量的同时，显著提高了性能。

实验

3.1 消融实验

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