CSWin Transformer: A General Vision Transformer Backbone with Cross-Shaped Windows

链接:
论文地址
代码

一、引言

本文介绍了CSWin Transformer，这是一种高效、高效的基于Transformer的通用视觉任务主干。Transformer设计中的一个具有挑战性的问题使全局自注意力的计算非常昂贵，而自注意力通常会限制每个令牌的交互领域。为了解决这个问题，本文开发了十字形窗口自注意力机制，用于在形成十字形窗口的平行水平和垂直条纹中计算自注意力，每个条纹通过将输入特征分割为相等宽度的条纹而获得。本文对条带宽度的影响进行了数学分析，并为Transformer网络的不同层改变了条带宽度，从而实现了强大的建模能力，同时限制了计算成本。还引入了局部增强位置编码（LePE），它比现有的编码方案更好地处理局部位置信息。LePE自然支持任意输入分辨率，因此对下游任务特别有效和友好。结合这些设计和分层结构，CSWin Transformer在常见视觉任务上展示了具有竞争力的性能。如下图所示，该条纹宽度是十字形窗口的一个重要参数，因为它允许我们在限制计算成本的同时实现强大的建模能力。具体来说，我们根据网络的深度调整条纹宽度：浅层的宽度较小，深层的宽度较大。更大的条纹宽度有助于长距离元素且以较小的计算成本增加实现更好的网络容量。

值得注意的是，使用CSWin自注意力机制，水平和垂直条纹中的自注意力是并行计算的。将多头分成平行组，并将不同的自注意力操作应用于不同的组。这种并行策略在扩大每个Transformer块中计算自注意力的区域的同时，不引入额外的计算成本。这种策略与现有的自注意力机制有根本不同，它在多个头部应用相同的注意操作（图上 b，c，d，e），并依次执行不同的操作。
基于CSWin自注意力机制，本文遵循分层设计，并为通用视觉任务提出了一种新的视觉转换器架构“CSWin Transformer”。该体系结构提供了显著更强的建模能力，同时限制了计算成本。为了进一步增强这种视觉Transformer，本文引入了一种有效的位置编码，即局部增强位置编码（LePE），它对于输入变化的下游任务特别有效和友好。与先前的位置编码方法相比，本文的LePE在每个Transformer块中施加位置信息，并直接对注意力结果进行操作，而不是进行注意力计算。

二、方法实现

(一)、CSWin Transformer

CSWin Transformer的整体架构如下图所示。对于大小为H×W×3的输入图像，利用重叠卷积令牌嵌入（7×7卷积层，步长为4）来获得H/4×W/4补丁令牌，每个令牌的维度为C。为了生成分层表示，整个网络由四个阶段组成。在两个相邻的级之间使用卷积层（3×3，步长2），以减少令牌的数量并使通道维度加倍。因此，构建的特征图在第i阶段具有$\frac{H}{2^{i+1}}\times \frac{W}{2^{i+1}}$标记，这与传统的CNN主干（如VGG/ResNet）类似。每个阶段由$N^i$个顺序CSWin Transformer块组成，并保持令牌的数量。CSWin Transformer Block具有与vanilla多头自注意力Transformer Block总体相似的拓扑结构，但有两个区别：1）它用本文提出的十字形窗口自注意力取代了自注意力；2） 为了引入局部感应偏置，将LePE作为并行模块添加到自注意力分支。

Cross-Shaped Window Self-Attention：尽管具有强大的长程上下文建模能力，但原始的自注意力机制的计算复杂度与特征图大小成二次关系。因此，对于以高分辨率特征图为输入的视觉任务，例如目标检测和分割，它将有巨大的计算成本。为了缓解这一问题，现有的工作建议在局部注意窗口中执行自注意力，并应用光环或移位窗口来扩大接受域。然而，每个Transformer块中的令牌仍然具有有限的关注区域，需要堆叠更多块以实现全局接受野。为了扩大关注区域并更有效地实现全局自我关注，本文提出十字形窗口自注意力机制，该机制通过在水平和垂直条纹中行执行自注意力来实现。

Horizontal and V ertical Stripes：根据多头自注意力机制，输入特征$X\in R（H\times W）\times C$将首先线性投影到K个头部，然后每个头部将在水平或垂直条纹内执行局部自注意力。 对于水平自注意力，X被均匀地划分为等宽sw的非重叠水平条带$[X^1，..，X^M]$，并且每个条带都包含$sw\times W$标记。这里，sw是条带宽度，可以调整以平衡学习能力和计算复杂性。形式上，假设第k个头的投影查询、键和值都具有维度$d_k$，那么第k个头水平条纹自注意力的输出定义为：

$X^i\in R^{(sw\times W)\times C}$和$M=H/sw$，$i=1，\dots ，M$。$W_k^Q\in R^{C\times d_k}，$d_k设置为$C/K$。可以类似地导出垂直条纹自注意力，其第k个头的输出表示为$V-Attentio{n}_k（X）$。
假设自然图像没有方向性偏差，我们将K个头平均分成两个平行的组（每个组有K/2个头，K通常是偶数）。第一组头执行水平条纹自注意力，而第二组头执行垂直条纹自注意力。最后，这两个并行组的输出将连接在一起。

其中$W^O\in R^{C\times C}$是常用的投影矩阵，将自注意力结果投影到目标输出维度（默认设置为C）。如上所述，自注意力机制设计中的一个关键见解是将多头分成不同的组，并相应地应用不同的自我关注操作。换言之，一个Transformer块中每个令牌的关注区域通过多头分组来扩大。相比之下，现有的自注意力机制在不同的多头应用相同的自注意操作。在实验部分，将表明这种设计将带来更好的性能。
Locally-Enhanced Positional Encoding：由于自注意力操作是排序不变的，因此它将忽略2D图像内的重要位置信息。为了将这些信息添加回来，在现有的视觉变换器中使用了不同的位置编码机制。在下图中，我们展示了一些典型的位置编码机制，并将它们与我们提出的局部增强位置编码进行了比较。具体而言，APE和CPE在馈送到Transformer块之前将位置信息添加到输入令牌中，而RPE和本文的LePE将位置信息合并到每个Transformer块中。但与在注意力计算中添加位置信息的RPE（即Softmax（QKT））不同，我们考虑了更直接的方式，并将位置信息施加到线性投影值上。同时，我们注意到RPE以每头的方式引入偏置，而我们的LePE是每通道偏置，这可能显示出更大的潜力来充当位置嵌入。

在数学上，我们将输入序列表示为n个元素的$x=（x_1，\dots ，x_n）$，以及相同长度的注意力$z=（z_1，\dots ，z_n）$的输出，其中$x_i,z_i\in R^C$。自注意计算可以公式化为：

其中$q_i、k_i、v_i$是输入$x_i$的线性变换得到的队列、键和值，d是特征维度，我们的局部增强位置编码作为可学习的每元素偏差执行，上式可以公式化为：

其中$z_i^k$表示向量$z_i$的第k个元素。为了使LePE适合不同的输入大小，我们为LePE设置了一个距离阈值，如果令牌i和j的切比雪夫距离大于阈值τ（默认设置中τ=3），则将其设置为0。

（二）、体系结构变体

为了与类似设置下的其他Transformer进行公平比较，本文构建了四种不同的CSWinTransformer变体，如下表所示：CSWin-T（微型）、CSWin-S（小型）、CSWin-B（基础）、CSWin-L（大型）。它们是通过改变基本通道C和每个级的块数来设计的。在所有这些变体中，每个MLP的膨胀比设置为4。前三个变体中的四个级的头数分别设置为2、4、8、16，最后一个变体中分别设置为6、12、24、48。

三、实验

为了公平比较，本文遵循DeiT中的训练策略，与其他基线Transformer架构一样，所有模型都训练了300个轮次，输入大小为224×224。使用AdamW优化器，CSWin-T/S的权重衰减为0.05，CSWin-B的权重衰减0.1。默认的批量大小和初始学习率设置为1024和0.001，并使用具有20个周期线性预热的余弦学习率调度器。对CSWin-T、CSWin-S和CSWin-B应用递增随机深度增强，最大速率分别为0.1、0.3和0.5。当报告384×384输入的结果时，使用权重对模型进行了30个时期的微调、1e-8的衰减，1e-5的学习率，512的批量大小。

(一)、分类

将CSWin Transformer与最先进的CNN和Transformer架构进行了比较。

（二）、目标检测

(三)、语义分割

(四)、消融实验

Parallel Multi-Head Grouping：在下表，我们发现“并行多头分组”是有效的，尤其是对于下游任务。当我们将并行方式替换为顺序方式时，CSWin在所有任务上的性能都会下降。当在类似的注意力区域约束下与之前的方法相比时，我们的sw=1 CSWin在ImageNet上的表现略好于轴向，同时在下游任务上以较大的优势优于轴向。我们的sw=2 CSWin表现略好于Criss Cross Attention，而CSWin的速度为2× 这进一步证明了我们的“并行”设计更高效。

Dynamic Stripe Width： 在下图中，我们研究了条纹宽度和精度之间的权衡。我们发现，随着条带宽度的增加，计算成本（FLOPS）增加，Top-1分类精度在开始时大大提高，当宽度足够大时会降低。我们的默认设置[1，2，7，7]在精度和FLOP之间实现了很好的平衡。
注意机制比较:在对CSWin自注意力的每个组成部分进行上述分析之后，我们进一步比较了现有的自注意力机制。由于某些方法在每个阶段都需要偶数层，为了公平比较，我们使用Swin-T作为主干，只改变自注意机制。具体而言，使用96个基本通道、非重叠令牌嵌入和RPE对四个阶段使用2、2、6、2个块。结果如下表所示。显然，我们的CSWin自我注意机制在所有任务中的表现都优于现有的自我注意机制。

位置编码比较:LePE专门设计用于增强局部位置信息在各种输入分辨率的下游任务。在这里，我们使用CSWin-T作为主干，只使用位置编码。在下表中，将我们的LePE与其他最近的位置编码机制（APE、CPE和RPE）进行了比较，用于图像分类、对象检测和图像分割。此外，还测试了无位置编码的变体（无PE）和CPE*，这是通过在每个Transformer块之前应用CPE获得的。根据比较结果，我们发现：
1）位置编码通过引入局部感应偏置可以带来性能增益；
2） 尽管RPE在具有固定输入分辨率的分类任务上实现了类似的性能，但我们的LePE在输入分辨率变化的下游任务上表现更好（COCO上的+1.2AP和+0.9AP，ADE20K上的+0.9 mIoU）；
3） 与APE和CPE相比，我们的LePE也取得了更好的性能。