Transformer
已成为深度学习中的主要架构之一,尤其是作为计算机视觉中卷积神经网络 (CNN) 的强大替代品。然而,由于Self-Attention
在长序列表示上的二次复杂性,特别是对于高分辨率密集预测任务,先前工作中的Transformer
训练和推理可能非常昂贵。为此,我们提出了一种新颖的少注意力视觉Transformer
(LIT),基于Transformers
中早期的Self-Attention
仍然专注于局部模式,并在最近的分层视觉Transformers
中带来较小的好处。具体来说,提出了一个分层
Transformer
,使用纯多层感知器(MLP)在早期阶段对丰富的局部模式进行编码,同时应用Self-Attention
模块在更深层捕获更长的依赖关系。此外,进一步提出了一种Learned Deformable Token Merging Module
,以非均匀方式自适应地融合信息Patch
。所提出的
LIT
在图像识别任务(包括图像分类、对象检测和实例分割)上取得了可观的性能,可作为许多视觉任务的强大支柱。
1简介
Transformers
在自然语言处理(NLP)和最近在计算机视觉(CV)领域取得了长足的进步。受CNN中金字塔设计的启发,最近的分层视觉转换器(HVT)将Transformer Block
划分为多个阶段,并随着网络的深入逐渐缩小特征图。然而,早期的高分辨率特征图导致了很长的token
序列,由于Self-Attention
的二次复杂度,带来了巨大的计算成本和内存消耗。例如,大小为56×56×96
的特征图在一个Multi-Head Self-Attention
(MSA)中需要2.0G FLOPs,而ResNet-18
的整个模型只需要1.8G FLOPs。如此巨大的计算成本使得将Transformer
应用到广泛的计算机视觉任务中变得非常困难。
在HVT
的早期阶段,已经做出了一些努力来降低计算成本。例如,一些工作减少了MSA
层中Self-Attention
Head的数量或进一步减少了Transformer Block
的数量。另一行工作建议通过启发式近似来权衡MSA
的准确性和效率,例如空间缩减注意力(SRA
)和基于Shift Window Multi-Head Self-Attention
(SW-MSA
)。也有研究当特征图的分辨率相当大时使用卷积层。然而,早期采用Self-Attention
层对最终性能有多大贡献仍不清楚。
在本文中提出了一种Less attention Vision Transformer
(LIT)来解决上述HVT
问题。具体来说,建议在早期Stage使用MLP
层来捕获局部信息,同时引入具有足够数量的Head来处理的MSA
层后期的长期依赖关系。
作者的动机来自2个方面。首先,先前的CNN
和Transformer
相关研究表明,模型浅层会专注于局部信息,而深层倾向于捕获高级语义或全局关系,这就产生了在早期Stage是否有必要使用Self-Attention
的问题。其次,从理论的角度来看,一个具有足够Head的Self-Attention
应用于图像可以表达任何卷积层。
然而,MSA
层中较少的Head
理论上阻碍了逼近具有大kernel-size
的卷积层的能力,其中极端情况与1×1卷积
一样具有表现力,可以被视为独立应用于每个像素的标准FC
层。虽然最近的HVT
在早期阶段采用很少的Head
来提供金字塔表示,但作者认为这不是最优的,因为这样的配置会引入高计算和内存成本,但带来的收益却很小。
需要强调的是,通过在前期Stage利用MLP Block
可以避免Self-Attention
在高分辨率特征图上产生的巨大计算成本和内存占用。此外,在后期Stage应用Self-Attention
来捕获远程依赖关系。综合实验结果表明,经过如此简单的架构设计后可以带来模型性能和效率之间的最佳平衡点。
此外,最近的HVT
要么采用标准卷积层,要么采用线性投影层来合并邻近的Token
,旨在控制特征图的规模。然而,考虑到并非每个像素对输出单元的贡献都相同,这种方法阻碍了Vision Transformer
对几何变换建模的表示能力。
为此,作者提出了一个受Deformable Convolutions
启发的Deformable Token Merging
(DTM
) 模块,在该模块中,学习了一个偏移网格以自适应地增加空间采样位置,进而合并来自特征图中子窗口的相邻Patch
。通过这种方式可以获得更多信息的下采样Token
用于后续处理。
主要贡献:
- 首先,确定了早期
MSA
层在最近的HVT
中的微小贡献,并在早期Stage提出了一个简单的HVT
结构,其中包含纯MLP Block
; - 其次,提出了一个
Deformable Token Merging
模块,以自适应地合并更多信息Patch
以提供分层表示,并具有增强建模能力; - 最后,进行了广泛的实验,以表明所提出的
LIT
在计算复杂度和内存消耗方面的有效性。
2本文方法
2.1 整体架构
图1
LIT
的整体架构如图1所示。设为输入的RGB图像,其中H和W分别表示高度和宽度。首先将I拆分为不重叠的Patch
,Patch
大小为4×4,因此每个Patch
的初始特征维度为4×4×3=48。
接下来,利用线性嵌入层将每个Patch
投影到维度,用作以下过程的初始输入。整个模型分为4个Stage。让s∈[1,2,3,4] 作为Stage的索引,在每个Stage使用个Block,其中前2个Stage仅使用MLP Block
来编码局部信息,后两个Stage使用标准Transformer Block
来编码处理更长的依赖关系。
在每个Stage,将输入特征图缩放为,其中和分别表示第s个Stage的块大小和隐藏维度。对于最后2个Stage,在每个Transformer Block
中设置个Self-Attention
heads。
2.2 LIT中的Block设计
如图1所示,LIT
采用2种类型的模块:MLP
模块和Transformer
模块。在前期Stage应用MLP Block
。具体来说,MLP Block
建立在MLP
之上,该MLP
由2个FC
层组成,中间有GELU
。对于第s阶段的每个MLP
,使用的扩展比。
具体来说,第1个FC
层将token
的维度从扩展到,另一个FC
层将维度缩减回。形式上,令为第s阶段的输入,l为块的索引,MLP Block
可以表示为:
其中LN
表示层归一化。在最后阶段,ViT
中描述的Transformer Block
包含一个MSA
层和一个MLP
,可以表示为:
使用这种架构有2个主要优点:
- 首先,避免了在早期Stage由长序列引入的巨大计算成本和内存占用;
- 其次,与最近使用子窗口缩小注意力图或减少注意力图的工作不同在key和value矩阵的空间维度上,在最后2个Stage保留标准
MSA
层,以保持LIT处理远程依赖关系的能力。
备注
在这里通过考虑卷积层
、FC层
和MSA层
之间的关系来证明在前期Stage应用纯MLP Block
的合理性。
这里建议参考一下字节跳动近期新出的TRT-ViT,有对应的结论给出,同时给出了非常详细的模型设计准则,对应推文的链接如下:
建议背诵 | 字节用4大准则教你设计一个拥有CNN的速度,Transformer精度的模型!
首先,从回顾标准卷积层开始。令为输入特征图,令为卷积权重,其中K为Kernel-size
,和分别为输入和输出通道维度。为简单起见,省略了偏置项,并用表示,其中(i,j)表示像素索引,。给定K×K
个采样位置的卷积核,像素p的输出可以表示为:
其中是采样索引到预先指定的偏移量ΔK的双射映射。
例如,令ΔK={(−1,−1),(−1,0),...,(0,1),(1,1)}为3×3的kernel,膨胀率为1,则g(0)=(-1,-1)表示第1个采样偏移量。当K=1时,权重张量W等价于一个矩阵,使得$。在这种情况下,Eq(4)可以表示FC层,像素p的输出定义为:
最后,让是MSA
层中的Head数,$是第h个Head的可学习参数。在特定的相对位置编码方案下,Cordonnier等人证明了像素p处的MSA层
的输出可以表示为:
其中是Head到像素位移的双射映射。在这种情况下,Eq(6)可以看作是Kernel-size
为 的卷积层近似。
从Eqs(4)-(6)观察到,虽然具有足够数量的Head的MSA
层能够逼近任何卷积层,但理论上更少的Head限制了这种逼近的能力。作为极端情况,具有一个Head的MSA
层是只能逼近FC
层。
请注意,MSA
层在实践中肯定不等同于卷积层。然而,d'Ascoli 等人观察到,早期的MSA
层可以在训练时学会卷积的类似表征。考虑最近的HVT
在早期Stage采用很少的Head,这种卷积行为可能会限制在小的感受野内。
图 3
在图3中,在可视化中显示PVT-S
中的早期MSA
层确实只关注query像素周围的一小块区域,而删除它们会导致性能轻微下降,但会显着降低模型复杂性。这证明了在前2个stage应用纯MLP Block
是合理的。
2.3 Deformable Token Merging
以前关于 HVT
的工作依赖于Patch Merge
来实现金字塔特征表示。然而,从规则网格合并Patch
,并忽略了并非每个Patch
对输出单元的贡献相同的事实。受可变形卷积的启发,提出了一个Deformable Token Merging
模块来学习偏移网格,以自适应地采样更多信息块。形式上,可变形卷积被表述为:
与Eq(4)中的标准卷积运算相比,DC为每个预先指定的偏移量g(k)学习一个偏移量Δg(k)。学习Δg(k) 需要一个单独的卷积层,它也应用于输入特征图X。为了以自适应方式合并Patch
,在DTM模块中采用一个DC层,可以表示为:
其中 BN
表示批量归一化,并使用了 GELU
。
在消融研究中,当物体的尺度和形状发生变化时,DTM
中的采样位置会进行自适应调整,从而受益于学习到的偏移量。另请注意,与Baseline
中的常规网格采样相比,轻量级DTM
引入了可忽略的FLOP和参数,因此使其成为最近HVT
的即插即用模块。