Skip to content
Shunted Self-Attention via Multi-Scale Token Aggregation 阅读笔记

  • Sucheng Ren,Daquan Zhou,Shengfeng He,Jiashi Feng,Xinchao Wang
  • National University of Singapore
  • South China University of Technology
  • ByteDance Inc.
  • CVPR2022 Oral
  • Code

背景

传统ViT通过将图片切分为Patch后产生对应的Token然后渐进式的通过全局注意力提取Token特征,这种方法虽然能处理特征的长程依赖关系但是过多的Token会导致在处理复杂的全局关系时丢失细小物体的特征,并且当Token较多时,网络的参数量与计算量会进一步上升。例如上图中天花板的灯在基于ViT的PVT中并没有得到很好的注意,而文中提出的网络结构能够更好地关注细小物体的特征。


Method

本文总体架构与ViT类似,均采用4阶段模式,辅以Patch Embedding、Linear Embedding操作,每个阶段进行一次下采样减小特征图分辨率并将通道数翻倍。

Shunted Transformer Block

Shunted Self-Attention

与ViT一样,先将输入序列FRh×w×c映射为Q,K,V,然后再经过MHSA。但是与其不同的是本文的结构将K,V的长度通过下采样的方式进行缩减以减少计算量并且以不同的长度捕获多尺度信息。其通过MTA(Multi-scale Token Aggregation)实现,公式如下:

Qi=XWIQKi,Vi=MTA(X,ri)WiK,MTA(X,ri)WiVVi=Vi+LE(Vi)

其中i为网络的第i个阶段,MTA(,ri)为第i阶段的采样率为ri的下采样MTA层(MTA的具体实现为一个Stride=ri的卷积层),WiQ,WiK,WiV为第i阶段的Linear Projection参数,LE()是一个对MTA进行局部增强的深度卷积层。最终的Self-Attention公式与ViT类似,即:

hi=Softmax(QiKiTdh)Vi

通过上述公式可知,当r增大时,K,V中更多的Token得以被融合并且K,V得以缩短,因此能在降低计算量的同时增强对大目标的捕获能力。相反地,当r减小时,更多的细节可以被关注但也会增加计算成本。

Detail-specific Feedforward Layers

为了进一步补充局部信息,本文的网络在前馈层添加一个Detail Specific模块,因此该网络的前馈层的表述公式为:

x=FC(x;θ1)x=FC(σ(x+DS(x;θ));θ2)

其中,DS(;θ)为使用参数θ的深度卷积层,θ1,θ2分别为两个FC层的参数。

Patch Embedding

论文Scaled ReLU Matters for Training Vision Transformers表明在Patch Embedding阶段使用卷积操作可以得到高质量Token,其效果要比使用单个大Stride的无重叠卷积效果更优秀。

本文中使用两个卷积层和一个Projection层作为Patch Embedding,第一个卷积层为Stride=2的7x7卷积,第二个卷积层为Stride=1的3x3卷积,最终通过一个Stride=2的无重叠Projection层来产生长度为H4×W4的输入序列。

Architecture Details and Variants

网络的输入为W×H×3的原始图片,经过Patch Embedding后产生C个长度为H4×W4的Patch序列。每个阶段的Linear Projection部分由一个Stride=2的卷积层构成。每个阶段的最后特征图分辨率减半并且通道数翻倍。网络的几种配置规格如下表:


实验与结果

ImageNet-1K 图像分类

COCO 2017 目标检测

1xSchedule(12 epochs)的fine-tuning阶段将输入图的短边resize为800px,长边不大于1333px;而3xSchedule(36 epochs)的fine-tuning阶段采取多尺度训练策略,将较短的尺寸调整到480至800之间。

ADK20K 语义分割

下表为与使用MiT Backbone的SegFormer框架的对比:

消融实验

Patch Embedding

其中Non-Overlap指ViT中使用的策略,而Overlap为Swin和PVT的Embedding策略。

Token Aggregation Function

Detail-specific Feed-Forward

上次更新于: