相关资源来自集智书童
1. 摘要
最近的研究表明,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个下游任务语义分割和目标检测中实现的最先进的准确性和鲁棒性。
2. 引言
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中的静态卷积操作不同,注意力设计是动态的,与内容相关,从而可以带来更丰富和更鲁棒的表征。
3. Fully Attentional Networks
3.1 Self-Attention的原理
一个标准的 ViT 首先通过固定的Patch Size将输入的图像均匀地划分为 n Patches, 并将每个 Patch 编码为一个 Token Embedding x i ∈ R d , i = 1 , … … , n 。然后, 所有这些 Token 被输入 Transformer Blocks 。每个 Transformer Block 利用 Self-Attention 进行 Token Mixing, 然后使用 MLP 进行通道特征转换。Transformer Block 的体系结构如图所示。
如图可以看出, 可以把 Self-Attention 分为 Token mixing 和 Channel processing 2 个部分来
看。
Token mixing
Vision Transformer 利用 Self-Attention 来聚合全局信息。假设输入 Token Embedding 张量 为 X = [ x 1 , … , x n ] ∈ R d × n ,Self-Attention 应用权重参数为 W K 、 W Q 、 W V的线性变换将 它们分别嵌入到 key K = W K X ∈ R d × n、query Q = W Q X ∈ R d × n 和value V = W V ∈ R d × n 之中。
然后, Self-Attention 模块计算注意力矩阵, 并将 Token 特征聚合如下:
其中 W L ∈ R d × n 是一个线性变换, Z = [ z 1 , … … , z n ] 是聚合的 Token 特征, 是一个缩放因 子。然后, Self-Attention 的输出到 LN 并输入到 MLP 中, 以生成下一个块的输入。
Channel processing
大多数ViTs采用MLP块将输入 Token 转换为特征Z:
Z ′ = MLP ( Z ) .
该块包含 2 个线性层和一个GELU层。
3.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不重要的组件进行可视化分组的过程。
3.3 Information Bottleneck视角看问题?
在某种意义上,Visual Grouping也可以被视为某种形式的有损压缩。因此,可以从Information Bottleneck的角度提出以下解释性框架。
给定一个分布 X ∼ N ( X ′ , ϵ ), 其中 x 是观察到的噪声输入, x ⋅ 是 target clean code, IB寻求 一个映射 f ( Z ∣ X ) , 使 Z 包含 X 中用于预测 x ⋅的相关信息。这个目标被表述为以下信息理论优 化问题:
上式中, 第 1 项压缩信息, 第 2 项维护相关信息。
对于SA Block, Z = [ z 1 , … … , z n ] ∈ R d × n 表示输出特征, X = [ x 1 , … , x n ] ∈ R d × n 表示输入。 假设 i是数据点索引, 有:
Proposition
Under mild assumptions, the iterative step to optimize the objective in Eqn. (3) can be written as:
or in matrix form:
Z = Softmax ( Q ⊤ K / d ) V ⊤ ,
with V = [ x 1 , … , x N ] ,K=[μ 1 ,…,μ N]= W K X , Q = Σ − 1 [ x 1 , … , x N ] and d = 1 / 2. Here n c , Σ and W K are learnable variables.
Remark
上述命题通过在普通的 Self-Attention 和IB之间建立了一个有趣的联系, 表明 Self-Att ention 将相似的输入 x i 聚集到具有 cluster 结构的表示 Z中。
Self-Attention 根据IB原则更新 Token 特征, 其中矩阵K存储临时 cluster 中心特征 μ c ,输入特征 x通过 softmax聚类到 Token 特征中。新的 cluster 中心特征 z 为更新后的 Toke n 特征输出。
ViTs中的堆叠SA模块可以广泛地看作是这种优化的迭代重复, 一定程度上促进了 Groupin g和噪声滤波。
Multi-head Self-attention
目前许多Vision Transformer架构采用MHSA设计,其中每个Head倾向于关注不同的部分。在某种意义上,MHSA可以被解释为Information Bottleneck的mixture。
而作者更感兴趣的是在一个固定的通道总数下,Head的数量与鲁棒性之间的关系。如图5所示,拥有更多的Head可以提高表征和鲁棒性。但通道数的减少也会导致精度的下降。最好的权衡是32个通道。
3.4 Fully Attentional Networks.
通过以上IBs解释的混合,设计了一个通道处理模块,通过不同Head的聚合增强鲁棒表示。设计主要由2个方面驱动:
1.为了促进更多的组合表征,引入通道权重是可取的,因为有些Head或通道确实可以捕捉到更重要的信息。
2.重加权机制应该包括对每个通道进行更全面的考虑,以利用分组信息,而不是做出“very local”的通道道聚合决策。
实现上述目标的一个起点是引入一种类似于XCiT的Channel Self-Attention设计。如图6(a)所示,Channel Self-Attention(CA)模块采用Self-Attention设计,将MLP块移动到Self-Attention块中,然后与通道注意力分支的通道注意力矩阵进行矩阵乘法。
注意力特征转换
FAN块引入以下Channel Self-Attention(CA)来执行特征转换,其表示为:
这里的 W 'Q ∈ R d × d 和 W' K ∈ R d × d 是线性变换参数。
与SA不同的是, CA沿着通道维度而不是 Token 维度计算注意力矩阵, 利用特征协方差 ( (( 经过线 性变换 W Q , W K )进行特征变换。聚集相关值较大的强相关特征通道, 分离相关值较低的异常 特征。这有助于模型过滤掉不相关的信息。在CA的帮助下, 该模型可以过滤不相关的特征, 从而对前景和背景 Token 形成更精确的 Token 聚类。
2.5 高效的通道Self-attention
传统的自注意计算机制沿通道维度的应用存在2个局限性。
1.计算开销,引入的CA的计算复杂度与成二次比例,其中D为通道维数。对于金字塔模型设计,通道尺寸在顶部阶段变得越来越大。因此,直接应用CA会导致巨大的计算开销。
2.参数效率低,在传统的SA模块中,通过Softmax操作来增强注意力权值的注意力分布。因此,只有部分通道有助于表征学习,因为大多数通道通过与较小的注意力权值相乘而被削弱。
为了克服这些问题,作者探索了一种具有高计算效率和参数效率的新型自注意力机制。具体来说,提出了2项主要修改:
1.首先,不计算Token特征之间的相关性矩阵,而是通过在通道维上平均生成Token原型Z。直观地说,Z聚合了Token所表示的每个空间位置的所有通道信息。因此,计算Token特征与Token原型Z之间的相关矩阵是有意义的。
2.其次,不使用Softmax函数,而是使用Sigmoid函数来规范化注意力权重,然后将其与Token特征点乘,而不是使用MatMul来聚合通道信息。直观地说,并不强迫通道只选择少数重要的Token特征,而是根据空间相关性对每个通道重新加权。
实际上,通道特性通常被认为是独立的。价值较大的通道不应限制其他通道的重要性。结合这2种设计概念,提出了一种新的通道自注意力方法:
在这里,σ表示沿着Token维度的Softmax操作,Z表示Token原型。详细的块架构设计也显示在图6中。实验证明,这种注意力机制在减少计算量的同时,显著提高了性能。
3实验
4. 实验
4.1 消融实验
如表7所示,原始的ConvNeXt比Swin-Transformer具有更强的鲁棒性,但不如FAN-ViT和FAN-Swin模型。同时,FAN-hybrid具有与FAN-vit和FAN-swin相当的鲁棒性,并且对于干净数据集和已损坏数据集都具有更高的精度,这意味着FAN也可以有效地增强基于CNN的模型的鲁棒性。与FAN-Swin类似,FAN-Hybrid对于大分辨率的输入和密集的预测任务具有高效性,有利于下游任务。
4.2 SOTA对比