SG-Former: Self-guided Transformer with Evolving Token Reallocation

系列博客目录

---

---

摘要

Vision Transformer（ViT）在各种视觉任务中展现了令人印象深刻的成功。然而，其计算成本较高，随着token序列长度的增加，计算量呈平方增长，这在处理大特征图时极大地限制了其能力。为了减轻计算成本，之前的工作要么依赖于局部小区域的精细自注意力，要么使用全局自注意力，但通过缩短序列长度来实现，从而导致粗粒度的表示。在本文中，我们提出了一种新颖的模型，称为Self-guided Transformer（SG-Former），旨在通过自适应的细粒度实现有效的全局自注意力。我们方法的核心是利用一个重要性图，该图通过混合尺度自注意力进行估算，并在训练过程中自我演化，用来基于每个区域的重要性重新分配token。直观地，我们将更多的token分配给显著区域，以实现精细的注意力，同时将较少的token分配给不重要的区域，从而在保证效率的同时保持全局感受野。

所提出的SG-Former在性能上超越了现有的最先进方法：我们的基准模型在ImageNet-1K上达到了84.7%的Top-1准确率，在COCO上达到了51.2的mAP bbAP，在ADE20K上达到了52.7的mIoU，比Swin Transformer提高了+1.3% / +2.7 mAP / +3 mIoU，同时计算成本更低，参数更少。代码可在以下网址获取：https://github.com/OliverRensu/SG-Former

引言

Transformer[49]起源于自然语言处理（NLP），最近在视觉学习中展示了最先进的性能。Vision Transformer（ViT）[10]的开创性工作引入了自注意力模块，明确地对图像块之间的长距离依赖关系进行了建模，这克服了卷积中局部感受野的固有局限，从而提升了各种任务的性能[9, 27, 56, 70, 74, 1]。

尽管取得了巨大的成功，但自注意力的计算成本随着序列长度的增加呈平方增长，这极大地限制了其在大规模输入上的应用。为了减少计算成本，ViT采用了较大步幅的patch嵌入来缩短序列长度。然而，这种操作不可避免地导致自注意力只能应用于小尺寸的特征图，且具有较粗的粒度。为了在高分辨率特征上计算自注意力，一些方法[27, 9, 48]提出将自注意力的计算区域限制在局部窗口内，而不是整个特征图（即精细粒度的局部自注意力）。例如，Swin Transformer设计了窗口注意力，CSWin设计了跨形状注意力。因此，这些方法[27, 9, 48]牺牲了在每个自注意力层中建模全局信息的能力。另一类方法[53, 59, 52]旨在聚合整个key-value特征图上的token，以减少全局序列长度（即粗粒度的全局注意力）。例如，Pyramid Vision Transformer（PVT）[53]通过使用大步幅的大卷积核均匀地聚合整个特征图上的token，从而导致整个特征图上的信息呈现均匀的粗粒度。

在本文中，我们介绍了一种新颖的Transformer模型，称为自指导Transformer（SG-Former），旨在通过一种不断演化的自注意力设计实现具有自适应细粒度的全局注意力。SG-Former的核心思想是，在保留整个特征图的长距离依赖的同时，根据图像区域的重要性重新分配token。直观地，我们倾向于将更多的token分配给显著区域，以便每个token可以在细粒度上与显著区域进行交互，同时在不重要的区域分配较少的token，以提高效率。SG-Former通过全局感受野高效地估算自注意力，并自适应地关注显著区域的细粒度信息。如图2所示，我们的SG-Former根据自身获得的重要性图，在显著区域（如狗）分配更多的token，而在不重要区域（如墙壁）分配较少的token。相比之下，PVT采用预定义策略均匀地聚合token。

具体来说，我们保留查询token，但为了实现高效的全局自注意力，重新分配键（key）和值（value）token。图像区域的重要性以得分图（score map）的形式通过混合尺度自注意力进行估算，并进一步用于引导token的重新分配。换句话说，给定一张输入图像，token的重新分配是通过自我引导来完成的，这意味着每张图像都会经历一个仅为其自身定制的token重新分配。因此，重新分配的token较少受到人为因素的影响。此外，这种自我引导会随着训练过程中逐步准确的得分图预测而发展。得分图在重新分配的有效性上起着重要作用，因此我们提出了一种混合尺度自注意力，它在一个层中考虑了不同粒度的信息，并且计算成本与Swin Transformer相同。混合尺度自注意力中的不同粒度信息是通过将头（heads）分组并多样化每组的方式来实现的，以便针对不同的注意力粒度。混合尺度自注意力还为整个Transformer提供了混合尺度的信息。

因此，我们的贡献总结如下：

• 我们通过统一混合尺度的信息提取，提出了一种新颖的Transformer模型——SG-Former，涵盖了一个自注意力层内的精细粒度局部信息和全局粗粒度信息。通过统一的局部和全局混合尺度信息，我们预测得分图以识别区域的重要性。

• 通过得分图，我们构建了自我引导的注意力机制，自动定位显著区域，并保持显著区域的精细粒度，以便准确提取信息，而对不重要区域保持粗粒度，以降低计算成本。

• 我们在多种下游任务（包括分类、目标检测和分割）上评估了所提出的Transformer骨干网络。实验结果表明，SG-Former在相似模型规模下始终优于之前的视觉Transformer模型。

3 方法

SG-Former的整体流程如图3所示。SG-Former与之前的CNN和Transformer模型[17, 53, 27, 9]共享相同的patch嵌入层和四阶段金字塔结构。输入图像 $X\in {\mathbb{R}}^{H\times W\times 3}$ 首先通过patch嵌入层在输入层进行4倍下采样。每两个阶段之间都有一个下采样层，采样率为2倍。因此，在第i阶段，特征图的大小为$H/2^{i+1}\times W/2^{i+1}$。除了最后一个阶段外，每个阶段都有$N_i$ 个Transformer模块，这些模块由两种类型的模块交替组成：（i）混合尺度Transformer模块和（ii）自我引导Transformer模块。混合尺度自注意力提取混合尺度对象和多粒度信息，并指导区域的重要性；其细节将在第3.3节中介绍。自我引导自注意力建模全局信息，同时根据来自混合尺度Transformer模块的重要性信息，在显著区域保持精细粒度；其细节将在第3.2节中介绍。

3.2 自我引导注意力

尽管传统的自注意力机制能够建模扩展的范围信息，但其计算成本和内存消耗与序列长度的平方成正比，这限制了其在大尺寸特征图上的应用，尤其是在诸如分割和检测等计算机视觉任务中。最近的研究[53, 52]建议通过合并几个token为一个来减少序列长度，即进行token聚合。然而，这种聚合方法对每个token平等对待，忽略了不同token之间固有的重要性差异。这样的聚合面临两个问题：（i）在显著区域，信息可能会丢失或与不相关的信息混合；（ii）在不重要的区域或背景区域，对于简单语义而言，许多token（占据序列的大比例）是冗余的，而它们还会消耗大量计算资源。基于这个观察，我们提出了自我引导注意力，它利用重要性作为指导来进行token聚合。具体来说，在显著区域，保留更多的token以获取精细粒度信息；而在不重要区域，保留较少的token，以保持自注意力的全局视野并同时降低计算成本。如图4所示，“自我引导”表示Transformer在训练过程中自我决定我们的计算成本降低策略，而非依赖人工引入的先验知识，例如Swin中的窗口注意力[27]、CSWin中的跨形状注意力[9]、PVT中的静态空间缩减[53]。输入特征图 ( X \in \mathbb{R}^{h \times w \times c} ) 首先被投影到查询（Query, Q）、键（Key, K）和值（Value, V）。接下来，H个独立的自注意力头并行计算自注意力。为了减少计算成本，同时保持注意力后的特征图大小不变，我们固定Q的长度，但通过重要性引导的聚合模块（IAM）对K和V的token进行聚合。
$Q=X{W}^Q,K=IAM(X,S,r)W^K,V=IAM(X,S,r)W^V.$
IAM的目标是在显著区域将更少token聚合为一个（即保留更多信息），而在背景区域将更多token聚合为一个（即保留较少的信息）。在公式（1）中，得分图 $S\in {\mathbb{R}}^{h\times w}$ 包含具有多种粒度的区域重要性信息（见第3.3节）。

我们将得分图的值按升序排序，并将 $S$ 均匀地分成 $n$ 个子区域 $S_1,\cdots ,S_n$。因此，$S_n$ 和 $S_1$ 分别代表最重要和最不重要的区域。同时，所有的token在 $X$ 中根据 $S_1,\cdots ,S_n$ 被分组为 $X_1,\cdots ,X_n$。在公式（1）中，$r$ 表示聚合率，即每 $r$ 个token被聚合为一个token。我们在不同重要性的区域设置不同的聚合率 $r_1,\cdots ,r_n$，使得每个子区域有一个聚合率，并且子区域的重要性越高，聚合率越小。不同阶段的 $r$ 的具体值列在表1中。因此，IAM通过不同的聚合率按组聚合重新分配输入特征 $X_1,\cdots ,X_n$，每组token通过连接聚合得到：
${\hat{X}}_i=F(X_i,r_i),\hat{X}=Cat({\hat{X}}_1,\cdots ,{\hat{X}}_n)$
其中，$F(X,r)$ 是聚合函数，我们通过具有输入维度为 $r$，输出维度为1的全连接层来实现。${\hat{X}}_i$ 中的token数量等于 $X_i$ 中token数量除以 $r_i$。

3.3 混合尺度注意力

混合尺度注意力有两个目的：（i）在计算成本不超过Swin Transformer中窗口注意力的情况下，提取混合尺度的全局和精细粒度信息；（ii）为自我引导注意力提供重要性信息。如图5所示，输入特征 $X$ 被投影到查询（Query, Q）、键（Key, K）和值（Value, V）。然后，多头自注意力采用 $H$ 个独立的头。通常，这些 $H$ 个独立的头在相同的局部区域内执行，因此缺乏头的多样性。相比之下，我们将 $H$ 个头均匀划分为 $h$ 组，并将混合尺度和多感受野的注意力注入这些$h$ 组，每组中有 $H/h$ 个头。在属于第$j$ 组的第 $i$ 个头中，具有尺度$s_j$（其中 $j=1,\cdots ,h$），每 $s_j\times s_j$ 个token在 {K, V} 中被合并为一个token。接下来，我们将 {Q, K, V} 切分成窗口。{K, V} 的窗口大小设置为 $M$，并在所有组中保持不变。为了使 {Q} 和 {K, V} 的窗口大小与 {K, V} 中的token合并对齐，{Q} 的窗口大小被选择为$s_{j}M\times s_{j}M$，是 {K, V} 窗口大小的 $s_j$倍。注意力的感受野通过 $s_j$ 倍显著扩大：

$$Q_i=X{W}_i^Q,$$
$$K_i=\text{Merge}(X,s_j)W_i^K,\text{\hspace{1em}(3)}$$
$$V_i=\text{Merge}(X,s_j)W_i^V,$$

“$Merge(X,s_j)$” 表示将 $X$ 中每个 $s_j\times s_j$ 的 token 合并为一个 token，这通过步幅为 $s_j$ 的卷积来实现。特殊情况是当 $s_j$ 等于 1 时，这时不进行任何 token 合并，{Q, K, V} 的窗口大小保持不变。

$${\text{Atten}}_i=\text{Softmax}\left(\frac{P(Q_i,s_{j}M)P(K_i^{\top },M)}{\sqrt{d_h}}\right),h_i={\text{Atten}}_{i}P(V_i,M)\text{\hspace{1em}(4)}$$

“P(X, $s_j$M)” 表示窗口划分，窗口大小为 $s_j$M× $s_j$M。$Atte{n}_j$ 是注意力图。这里有一个特殊情况：当$s_j$M×$s_j$M 等于 h×w 时，无需进行窗口划分，{K, V} 中的所有 token 都会被 {Q} 所关注，从而实现全局信息提取。

$$S_i=\frac{1}{hw}\sum _{m=1}^{hw}{\text{Atten}}_{m,n}^i,S=\sum _{i=1}^h{S}_i,\text{\hspace{1em}(5)}$$

其中 $S$ 是最终的重要性图，通过对 $S_i$ 求和得到，结合了全局信息和细粒度信息的混合尺度引导。