Shunted Self-Attention via Multi-Scale Token Aggregation

近期提出的ViT模型在各种计算机视觉任务上展现了令人鼓舞的结果，这要归功于能够通过自注意力对补丁或令牌的长期依赖性进行建模。然而这些方法通常指定每个令牌相似感受野。这种约束不可避免地限制了每个自注意力层在捕捉多尺度特征的能力。本文提出一种新的注意力，称为分流自注意力（Shunted Self-attention，SAA），允许ViT在每个注意力层的混合尺度上对注意力进行建模。SAA关键思想是将异构感受野大小注入到令牌中：在计算自注意力矩阵之前，选择性合并令牌以表示更大的目标特征，同时保留某些令牌以保留细粒度特征。这种新颖的合并方案使得自注意力模块能够学习目标之间的关系，同时减少令牌数量和计算成本。

分流Transformer块

为了降低处理高分辨率特征图时的计算成本，PVT引入了空间约简注意力（Spatial Reduction Attention，SRA）取代原来的多头自注意力（MSA）。然而SRA倾向于在一个自注意力层中存储太多令牌，并且仅在单个尺度上提供令牌特征。这些限制阻碍了模型捕捉多尺度目标的能力，尤其是小尺寸目标。

分流自注意力

输入序列$F$首先投影到查询$Q$，键$K$和值$V$向量。之后多头自注意力采用$H$独立注意力头并行计算自注意力。为了减少计算成本，遵循PVT并减少K,V长度而不是Swin中按照区域划分 { Q , K , V } \{Q,K,V\} {Q,K,V}。本文的SSA与PVT的SRA不同点之处在于，在同一自注意力层的注意力头之间，K、V长度不相同。相反，长度在不同头中变化，以捕捉不同粒度的信息。这提供了多尺度令牌聚合。键和值$K,V$对于不同头下采样到不同尺寸：
$$\begin{array}{rl}{Q}_i& =X{W}_i^Q\\ K_i,V_i& =MTA(X,r_i)W_i^K,MTA(X,r_i)W_i^V\\ V_i& =V_i+LE(V_i)\end{array}$$
其中$MTA(\cdot .r_i)$是第i层的多尺度令牌聚合层，其中下采样率为$r_i$。实践中采用和卷积和步幅为$r_i$的卷积实现下采样。$W_i^Q,W_i^K,W_i^V$是第i头线性投影参数。在注意力头之间有变量KaTeX parse error: Expected '}', got 'EOF' at end of input: r_{i]。因此自注意力中键和值可以捕捉不同尺度。$LE(\cdot )$是MTA通过深度卷积对值$V$的局部增强分量。
分量自注意力可以描述为：$h_i=\text{Softmax}(\frac{Q_i{K}_i^T}{\sqrt{d_k}})V_i$。
得益于多尺度键核值，本文的分流自注意力在捕捉多尺度目标时更强大。计算成本的降低可以取决于$r$的值。当$r$变大，$K,V$中更多令牌被合并，且$K,V$长度更短。因此，计算成本较低，但它仍然保留了捕捉大目标的能力。当$r$变小，更多细节被保留但带入更多计算成本。融合各种$r$ 到一个自注意力层使得捕捉多细粒度特征。

具体细节前馈层

在传统的前馈层中，全连接层是逐点的且无交叉令牌信息可以学习。这里本文目标是通过指定前馈层中的细节来补充局部信息。通过在前馈层的两个全连接层之间添加数据特定层来补充前馈层中局部细节。
$$\begin{array}{rl}{x}^{\prime }& =FC(x;{\theta }_1)\\ x^{\prime \prime }& =FC(\sigma (x^{\prime }+DS(x^{\prime };theta));{\theta }_2)\end{array}$$
其中$DS(\cdot ;\theta )$是具体细节层，参数为$\theta$，实践中由深度卷积实现。

补丁嵌入层

如何映射图像到序列对于模型性能是重要的。ViT直接划分输入图像为16x16无重叠补丁。 Scaled relu
matters for training vision transformers发现在补丁嵌入使用卷积提供了高质量补丁序列，它帮助Transformer比传统的大跨度非重叠补丁嵌入看得更好。在本文模型中，根据模型尺寸采用具有重叠得不同卷积层。将步幅为2和零填充得7x7卷积层作为补丁嵌入的第一层，并根据模型大小添加步幅为1的额外3x3卷积层。最后使用步幅为2的无重叠投影层生成大小为$\frac{H}{4}\times \frac{W}{4}$。