Pale Transformer: A General Vision Transformer Backbone with Pale-Shaped Attention

摘要

为了降低全局自注意力引起的二次计算复杂度，各种方法将注意力范围限制在局部区域内以提高其效率。因此，它们在单个注意力层中的感受野不够大，导致上下文建模不足。为了解决这个问题，文章提出了一种Pale-Shaped的自注意力（PS-Attention），它在pale-shaped的区域内执行自注意力。与全局自注意力相比，PS-Attention 可以显著降低计算和内存成本。同时，它可以在与以前的局部自注意力机制相似的计算复杂度下捕获更丰富的上下文信息。基于 PS-Attention，文章提出了一个具有层次结构的通用 Vision Transformer 主干，命名为 Pale Transformer。

方法

Pale-Shaped Attention

为了捕获从短期到长期的依赖关系，提出了Pale-Shaped Attention(PS-Attention)，它在一个Pale-Shaped区域(简称pale)中计算自注意力。如图1(e)的粉色阴影所示，一个pale包含$s_r$个交错的行和$s_c$个交错的列，它覆盖了包含$(s_{r}w+s_{c}h-s_r{s}_c)$个Token的区域。

图1

定义$(s_r,s_c)$为pale size。给定一个输入特征图，首先将其分割成多个相同大小(s_r,s_c)的pale 。Pale的个数等于$N_r=\frac{h}{s_r},N_c=\frac{w}{s_c}$，可以通过填充或插补操作来保证。对于所有pale，相邻行或列之间的间隔是相同的。然后在每个pale中分别进行自注意力计算。如图1所示，PS-Attention的感受野比之前所有的局部自注意力机制都要广泛和丰富得多，能够实现更强大的上下文建模能力。

Efficient Vision Transformers

为了进一步提高效率，将上面提到的普通PS-Attention分解为行注意和列注意，它们分别在行Token组和列Token组内执行自注意力。

图2

复杂度分析

给定尺寸为h×w×c的输入特征，Pale-size为$(s_r,s_c)$，标准全局自注意力的计算复杂度为：

本文提出的PS-Attention并行实现下的计算复杂度为：
由于$2hw>>(s_{c}h+s_{r}w+27)$始终保持不变，因此与全局算法相比，可以明显减轻计算量和内存负担。

Pale Transformer Block

图3

Pale Transformer Block由3个顺序部分组成，用于动态生成位置嵌入的条件位置编码(CPE)，用于捕获上下文信息的PS-Attention模块，以及用于特征投影的MLP模块

整体网络

图4

Pale Transformer由四个层次级阶段组成，用于捕获多尺度特征。每个阶段包含一个Patch合并层和多个Pale Transformer Block。patch merge层的目标是对输入特征进行一定比例的空间下采样，并将通道维数扩展2倍以获得更好的表示能力。

实验

消融实验

Pale-size的影响：增加Pale-size(从1到7)可以持续提高所有任务的性能，而进一步增加到9不会带来明显的、持续的改进，但会带来更多的FLOPs。

不同注意力对比：为了直接比较PS-Attention与最相关的Axial-based Attention，分别用轴向自注意力和cross-shaped window self-attention替代了Pale-T的PS-Attention。PS-Attention明显优于这两种机制。

不同的位置嵌入：

ImageNet-1K分类

在相同的计算复杂度下，与先进的CNN相比，Pale模型分别比RegNet模型好+3.4%、+2.6%和+2.0%。同时Pale Transformer的性能优于最先进的基于Transformer的Backbone，在类似型号尺寸和FLOPs下的所有变种中，其性能比最相关的CSWin Transformer高0.7%。

COCO目标检测