GroupMixFormer：基于Group-Mix注意力的视觉Transformer

摘要

https://arxiv.org/pdf/2311.15157.pdf
视觉转换器（ViTs）已被证明可以通过建模长程依赖关系来增强视觉识别，这种建模使用多头自注意力（MHSA），通常将其表述为查询-键-值计算。然而，从查询和键生成的注意力图只能捕获单个粒度上的令牌到令牌的相关性。在本文中，我们认为自注意力应该有一个更全面的机制来捕获令牌和令牌组（即多个相邻的令牌）之间的相关性，以获得更高的表示能力。因此，我们提出了一种先进的替代传统自注意力的方法，即Group-Mix Attention（GMA），它可以同时捕获令牌到令牌、令牌到令牌组以及令牌组到令牌组的相关性，并具有各种不同的组大小。为此，GMA将查询、键和值均匀拆分成多个片段，并执行不同的组聚合来生成组代理。注意力图是基于令牌和组代理的混合计算出来的，并用于重新组合值中的令牌和组。基于GMA，我们引入了一个强大的主干网络，即GroupMixFormer，它在图像分类、目标检测和语义分割方面取得了最先进的性能，同时参数数量比现有模型更少。例如，GroupMixFormer-L（具有70.3M个参数和$384^2$个输入）在没有外部数据的情况下，在ImageNet-1K上达到了86.2%的Top1准确率，而GroupMixFormer-B（具有45.8M个参数）在ADE20K上达到了51.2%的mIoU。相关代码和训练模型可以在以下链接中找到：https://github.com/AILab-CVC/GroupMixFormer。

1、简介

视觉转换器（ViTs）显着改善了视觉识别任务，包括图像分类[18，63]，自监督学习[3，5，10，60]，目标检测[15，37]，以及语义分割[53，54，58]。一个对性能提升有重大贡献的关键模块是多头自注意力（MHSA），它使网络设计具有长程依赖建模[43，51]，全局感受野，更高的灵活性和更强的鲁棒性[42，58]。通常，“注意力”（即Q-K-V注意力）是指将Value与Query和Key之间的相关性进行线性重组，这些相关性通常是在单个令牌对之间计算的。

然而，经验上发现Q-K-V自注意力存在一个主要的局限性，如图1所示：注意力图只描述了单个粒度上每个令牌对之间的相关性（图1(a)），并且将注意力图与Value相乘只是线性地重新组合了各个令牌。这个框架显然没有考虑到不同粒度上不同令牌组（即邻域）之间的相关性。举一个具体的例子，自注意力不能将左上角的九个令牌作为一个整体与右下角的令牌组相关联。这个局限性虽然很明显，但因为Q-K-V计算似乎足够建模输入到输出的映射，所以被无意忽略了，因为输出的任何元素都会关注输入的每个单独元素。

在本研究中，我们提出了一种更全面的建模方法，称为Group-Mix Attention（GMA），以缓解广泛使用的Q-K-V自注意力机制的上述局限性。GMA将令牌拆分为均匀且不同的片段，并用通过组聚合器生成的组代理替换一些单个令牌，如图1（b）所示。之后，我们使用查询和键（其中一些令牌已被组代理替换）来计算注意力图，并使用它重新组合Value中的组代理和单个令牌。提出的GMA具有以下优点：（1）GMA不仅能够建模单个令牌之间的相关性，还能够建模令牌组之间的相关性。不同类型的注意力被混合在一起，以从更全面的角度更好地理解令牌。在每个单层内同时建模令牌到令牌、令牌到组和组到组的相关性，以获得更高的表示能力。（2）GMA是高效且易于实现的。组到组的相关性是通过将组聚合为代理令牌然后计算代理之间的相关性来计算的（如图3所示）。这样的过程可以通过基于滑动窗口的操作（例如池化和卷积）高效地实现。

在GMA的基础上，我们开发了分层视觉转换器GroupMixFormer，它可以作为各种任务的可视化主干网。我们在标准视觉识别任务上评估了GroupMixFormers，包括图像分类、目标检测和语义分割，并与先进模型进行了比较，如图2所示。结果表明我们的设计是有效的。例如，一个小的GroupMixFormer实例（22.4M参数）在ImageNet-1K上实现了83.4%的Top-1准确率，与更大的SwinB[37]（88M参数）相当。此外，GroupMixFormer在目标检测和语义分割方面也优于先进的ViTs和CNN。在ADE20K数据集上，GroupMixFormer-B实现了51.2%的mIoU，主干大小为46M。大量实验还表明，有效地建模令牌和不同组之间的相关性对于GMA的成功至关重要。这种设计范式可以很容易地应用于其他ViT架构，作为传统自注意力的高级替代品。

2、相关工作

2.1、视觉转换器

视觉转换器（ViT） [18]首次将转换器引入计算机视觉领域。与基于CNN的架构不同，ViT在视觉令牌序列上使用顺序连接的Transformer编码器 [51]。ViT中使用的多头自注意力（MHSA）机制有效地捕捉全局依赖关系，使它们在监督 [22, 49]和自监督场景 [5, 10]中比CNN神经网络具有优势。为了提高ViT的通用性能，已经进行了一系列研究，包括数据高效训练 [49]、令牌重新设计和选择 [33, 44]、金字塔结构 [37, 53]、对自注意力机制进行调制 [6, 8, 66]，等等。大多数这些工作采用原始的Q-K-V计算，这在处理视觉信息方面是有效的。本文旨在通过引入Group-Mix Attention（GMA）来进一步推进ViT的通用性能。与现有技术不同，GMA能够建模每个单一Transformer编码器层中不仅单个令牌之间而且令牌组之间的相关性，从而提高全面的表示能力。

2.2、全面的自注意力建模

为了增强自注意力的表示能力，已经从不同的角度探索了几种方法，如下所示。(1)引入局部性已被证明是有效的，例如Swin Transformer [36, 37]和Focal Transformer [62]，它们在局部窗口内进行注意力计算。(2)使用预定义模式计算相关性可以增强自注意力的能力，如CSWin Transformer [17]和Pixelfly-Mixer [7]，它们都尝试使用预定义和精心设计的模式来计算注意力以实现更全面的建模。(3)其他网络架构[30, 31, 45, 48, 52, 54, 56, 61]也被研究用于更全面的视觉模式建模。本文关注单个粒度上令牌到令牌的相关性所造成的限制，并提出了一个先进的注意力机制（即GMA），它构建了一个更全面的自注意力原型，将我们的方法与之前的做法明确区分开来。

3、组混合注意力和GroupMixFormer

我们在本节中介绍GroupMix Attention和GroupMixFormer的动机和结构设计。

3.1. 动机：从个体到群体

我们讨论了自注意力的限制，从其原始公式开始。设 $\mathrm{X}\in {\mathbb{R}}^{N\times d}$ 是输入令牌，其中 N 是令牌数量，d 是维度。普通自注意力的输出为：
$$Y=\mathrm{Softmax}\left(X{X}^T\right)X\text{\hspace{1em}(1)}$$

注意，为了简洁起见，我们省略了归一化因子 $\frac{1}{\sqrt{d}}$。直观地说，通过矩阵乘法的定义， ${\mathrm{XX}}^{\mathrm{T}}$ 计算了每个令牌对之间的相似度/相关性。Softmax函数 $\mathrm{A}\in {\mathbb{R}}^{N\times N}$ 的输出称为注意力图。$AX$的乘法意味着根据每个位置的注意力图线性地重新组合令牌。

我们注意到这种形式的局限性。可能存在某些模式（即组模式），这些模式要求将某些特定的令牌作为具有不同粒度的组进行处理。然而，自注意力缺乏对这种模式的明确建模，因为它只考虑单个粒度（即个人模式）下成对令牌之间的相关性。本文旨在同时利用个人模式和组模式进行全面建模。与以前的方法不同，它们在多个阶段（通常在Transformer主干中有四个阶段）中分别建模不同的模式，我们的方法在每个阶段的每个单独层中引入了一种编码此建模过程的新方法。具体来说，对于组模式，我们寻求将一些令牌的邻域与其他邻域相关联。本文建议通过在查询、键和值中生成组代理，并与代理进行Q-K-V计算来实现这一点，这在第3.2节中描述。我们实验发现，通过明确建模具有不同大小和个体令牌的组之间的相关性，不仅提高了所提出的GroupMixFormer的性能，而且提高了具有不同注意力模块的其他ViT的性能（如Swin Transformer [37]和PVT [53]，如表9所示），证明升级基本组件可以使多个ViT受益。

3.2. GMA: 混合组以获得更好的注意力

我们引入GMA来模拟上述的组模式。在GMA中，我们通过替换Query、Key和Value中的一些条目来生成组代理，这些条目是通过聚合一些整个组来获得的，这可以通过滑动窗口操作 $\mathrm{Agg}(?)$ 高效地实现，例如maxpooling、卷积等。具体来说，Q/K/V条目被均匀地分为n个段，并对一些段进行聚合。为了方便起见，我们使用 ${\mathrm{X}}_i(i\in [1,?,n])$ 来表示一个段（ $\mathrm{X}$ 可能代表 $\mathrm{Q}$ 、 $\mathrm{K}$ 或 $\mathrm{V}$ ），并使用 ${\mathrm{Agg}}^i\left({\mathrm{X}}_i\right)$ 来表示聚合。请注意，聚合器可能因每个段而异。为了进行注意力计算，我们将聚合 ${\mathrm{Agg}}^i\left({\mathrm{X}}_i\right),i\in [1,?,n]$ 连接起来以产生 $X^{\prime }$ 。这样，我们获得组代理 $Q^{\prime }$ 、 $K^{\prime }$ 和 $V^{\prime }$ 。之后，我们像 [1,47,61] 中介绍的那样，在组代理上进行注意力计算，以生成输出。

在聚合过程中，我们保持特征分辨率。因此，在不降低空间分辨率的情况下，GMA为注意力计算带来了精细的特征，这优于那些减小特征大小的方法[19,55]。在本文中，我们使用具有各种内核大小的深度卷积来实现聚合器 $\mathrm{Agg}(?)$ ，尽管我们发现其他实现也有效（如表6所示）。由于注意力输入现在是组代理，因此我们同时关联 $\mathrm{K}\times \mathrm{K}$ 个令牌（ $\mathrm{K}$ 表示 $\mathrm{Agg}(?)$ 的内核大小，对于每个段可能有所不同），而不是单个令牌，这对于建模相关性更加充分和全面。

使用基于滑动窗口的操作来聚合组以生成代理的想法虽然简单，但它是将不同大小和各种粒度的个体令牌混合在一起的关键。这是因为我们为每个段使用了不同的聚合器的内核大小。这种过程可以通过将段分成不同部分，将它们分别输入具有不同内核大小的聚合器中，然后将输出连接起来来实现。此外，受[17]的启发，我们还对其中一个段使用了恒等映射，而不是聚合器，以保持网络在建模单个令牌相关性方面的能力。因此，我们可以在计算注意力图的同时建模组和令牌之间的相关性。将注意力图与值相乘可以看作是将相应的组与单个令牌重新组合在一起。

具体来说，在实现自注意力[1,47,61]之后，我们还使用三个可学习的线性投影来生成Q、K和V。之后，我们将Q/K/V均匀地分成五个段，每个段参与不同的计算。如图3（左半部分）所示，一个分支对应于前面提到的段，四个分支的输出被送入注意力计算中，被称为预注意力分支。在三个预注意力分支中，我们使用各种不同的实现（例如，最小池化、平均池化、最大池化、深度卷积）作为聚合器 $\mathrm{Agg}(?)$ ，具有不同的内核大小，分别设置为3、5、7。表6中的结果表明，每种实现都取得了良好的性能，这表明聚合是注意力进步的关键步骤，而其实施可以是灵活的。在我们的论文中，我们采用了深度卷积。我们进一步通过在最后一个预注意力分支中使用恒等映射而不是聚合器来多样化结构。除了这样一个具有注意力但没有聚合器的分支外，我们还构建了一个具有聚合器但没有注意力的分支，被称为非注意力分支。最后，输出由一个令牌集成层混合，该层由一个带有归一化[2]和激活的线性投影实现。

3.3. 架构配置

基于提出的Group-Mix Attention，我们引入了一系列名为GroupMixFormer的视觉Transformer，如图3所示。我们采用了四个阶段的分层[37,53]拓扑。第一个4 × patch嵌入层将图像嵌入到令牌中，这是通过两个连续的3 × 3卷积层实现的，每个卷积层的步长为2，另外两个3 × 3层的步长为1。在最后三个阶段开始时，我们使用2 × patch嵌入，这也是通过3 × 3卷积实现的。在每个阶段内，我们构建了几个编码器块。除了上一小节中介绍的GMA块外，编码器块还包含一个前馈网络（FFN）、层归一化[2]和恒等快捷方式，遵循[18, 37, 49, 53, 62]中的通用做法。对于图像分类，全局平均池化（GAP）后的最终输出令牌被送入分类器；对于密集预测任务（例如，目标检测和语义分割），特定任务的头部可以利用四个阶段输出的金字塔特征。由于我们已经通过GMA聚合器自然地打破了置换不变性，因此我们的模型中没有采用位置编码。

我们使用不同的架构配置实例化了四个模型。架构超参数包括每个阶段中的编码器块数量L、嵌入维度D和MLP比率R，如表1所示。根据之前的研究[37,49,53]，我们的模型从移动规模的GroupMixFormer-M（5.7M）扩展到大型的GroupMixFormer-L（70.3M）。

4、实验

在本节中，我们在标准视觉识别基准测试中评估了我们的GroupMixFormer，包括ImageNet1K [46]、MS-COCO [34]和ADE20k [68]。我们为每个场景提供了实现细节、与最先进的视觉主干网的定量比较以及消融研究，如下所示。

4.1、实现细节

我们在ImageNet-1K数据集上评估了GroupMixFormer的图像分类性能。我们遵循[49, 64,65]来增强数据并使用[37]中的训练方案。我们使用初始学习率为$10^{?3}$次方，进行300个训练周期，其中20个周期线性预热。我们使用AdamW优化器[39]，权重衰减为0.05，余弦学习率策略。对于GroupMixFormer-M/T/S/B/L，我们将随机深度丢弃率[25]设置为0.0/0.1/0.2/0.4/0.5。对于更高的分辨率（例如$384^2$或$448^2$），我们再用初始学习率为$2\times 10^{?6}$且进行5个周期线性预热的AdamW [39]进行30个周期的微调，优化过程中的权重衰减为$10^{?8}$次方。

对于目标检测和实例分割，我们采用了COCO 2017数据集。具体来说，我们将GroupMixFormer用作Mask R-CNN [24]的骨干网络进行目标检测和分割，仅用于检测的RetinaNet [35]。所有的骨干网络都通过相应的ImageNet预训练模型进行初始化。我们遵循[9]中的训练计划：初始学习率设置为$10^{?4}$，线性预热500次迭代，然后分别在第24个和第33个训练周期逐步降低到$10^{?5}$和$10^{?6}$。我们使用AdamW [39]进行Mask R-CNN和RetinaNet的训练，但前者权重衰减为0.05，后者权重衰减为$10^{?4}$。除了COCO之外，我们还使用UperNet [57]和Semantic FPN [29]在ADE20k上进行语义分割性能评估。我们遵循[37, 53]使用公共工具包[13]进行训练和评估。Semantic FPN训练80k次迭代，UperNet训练160k次迭代，两者都使用AdamW优化器。

4.2. 与最先进模型的比较

图像分类。在表2中，我们将提出的GroupMixFormer与文献中的最先进模型进行比较，其中所有报告的结果都只使用ImageNet1k进行训练。为了公平比较，我们不使用任何额外的增强，如标记token [28]、知识蒸馏、SAM [20]等。我们观察到，在类似的模型大小和计算复杂度约束下，GroupMixFormer始终实现了比ViT和CNN模型更高的Top-1准确率。具体来说，当测试分辨率为$224^2$时，GroupMixFormer-S在只有22.4M参数的情况下实现了83.4%的top-1准确率，比第二好的ViT（CSWin-T [17]）高出0.7%，比最好的CNN（ConvNext-T [38]）高出1.3%。同时，当使用224 × 244图像训练时，GroupMixFormer-B甚至实现了与Swin-B [37]相似的准确率，尽管GroupMixFormer-B的大小仅为Swin-B的一半。此外，GroupMixFormer在更高分辨率下表现出令人满意的扩展性。例如，使用分辨率为$384^2$进行微调后，GroupMixFormer-S的性能进一步提高到85.0%；大约70M参数的GroupMixFormer-L在分辨率为$224^2$时达到85.0%，在分辨率为$384^2$时达到86.2%。这些结果表明，在建模视觉模式时，全面整合token-to-token和group-to-group相关性具有优势。此外，附录中提供了来自不同聚合器的注意力响应，以支持某些token应该作为一个整体进行处理的概念。

此外，我们根据经验观察到，在GMA中实现深度卷积作为聚合器确实会导致推理速度减慢。吞吐量报告在附录中。然而，这可以通过更有效的聚合器(例如，avg-pooling)和实现工程优化(例如“torch.compile”)来改进。我们将在未来的研究中探索模型的实际速度优化。

目标检测。表3显示了使用Mask R-CNN和RetinaNet检测器在COCO上的目标检测结果。使用Mask R-CNN，GroupMixFormer在相似模型参数下实现了更高的平均精度。具体来说，GroupMixFormer-T比第二好的模型CoaT Mini高出1.0%（即47.5% vs 46.5%），同时保持较小的模型大小为30.8M。此外，我们的GroupMixFormer-B实现了APb为51.5%，超过了所有可比较的模型。使用RetinaNet，GroupMixFormer也表现出优势：GroupMixFormer-T比Swin-B高出0.5%（即46.3% vs 45.8%），尽管我们的模型小得多（即20.2M vs 98.0M）；GroupMixFormer-B比第二好的模型Focal-small高出2.9%（即50.2% vs 47.3%）。这些结果表明，GroupMixFormer在使用这两种检测器时都取得了良好的性能。一致且显著的提高证明了GroupMix机制的有效性，该机制应该能够捕捉到精细的特征以促进密集预测。

语义分割。表3还显示了使用Mask-RCNN在COCO上的语义分割结果。我们的GroupMixFormer-T令人印象深刻地实现了42.4%的APm，比Coat Mini高出0.6%，比PVTLarge高出1.7%。此外，GroupMixFormer-B的表现比Uniformer-B高出1.1%（即45.9% vs 44.8%）。在ADE20K上，我们使用UperNet和Semantic FPN并报告了结果（如表4所示）。同样，我们观察到GroupMixFormers与现有骨干网络相比，实现了持续且显著的性能提升。例如，尽管GroupMixFormer-T的模型大小小得多，但其在Semantic FPN上的表现比XCiT-S12/8高出2.0%（即46.2% vs 44.2%，14.1M vs 30.4M）。值得注意的是，尽管XCiT-M24/16的模型大小是GroupMixFormer-T的6.4倍（90.8M vs 14.1M），但GroupMixFormer-T的表现仍比后者高出0.3%（即46.2% vs 45.9%）。同样，使用UperNet时，GroupMixFormers的表现明显优于其他更大的模型，表明其在性能和效率之间达到了更好的权衡。如此显著的提升表明，Group-Mix机制能够为像素级预测产生高质量的特征。

4.3. 消融实验

在本小节中，我们进行消融实验以分析GroupMixFormer的关键设计。

(1) 我们首先通过改变GMA的结构设计来分析聚合器的必要性。

(2) 我们尝试使用各种聚合器的实现，以查看除卷积之外的其他基于滑动窗口的操作是否也有效。

(3) 我们验证GroupMixFormer的性能提升不是来自于宏观结构。

(4) 我们探索了核大小的优化配置。

(5) 我们进行实验以验证GMA不仅仅是卷积和自注意力机制的简单组合。

(6) 我们将GMA块插入到其他流行的ViT架构中，以验证GroupMixFormer的优越性能是否仅仅是由于架构设计（例如重叠的嵌入层和每个阶段中的块数）。

对于图像分类，我们在ImageNet-1k（$224^2$）上对GroupMixFormer-T进行300个epoch的训练，并在验证集上进行测试。对于目标检测和语义分割，我们在COCO上使用1×计划[9]训练Mask R-CNN。

组聚合器是必要的。表5显示了去除聚合器后的结果。我们首先通过将GMA块中的五个分支替换为恒等映射来构造GroupMixFormer-T基线，这样块就会退化为常规的自注意力模块。在第一组实验中，我们恢复了非注意力分支（Agg0）或三个预注意力分支（Agg1、Agg2和Agg3）中的聚合器。每个模型都使用与第4.1节中描述的相同配置从头开始进行训练。可以看出，聚合器都是至关重要的，因为它们将top-1准确率提高了0.4%和1.0%。

此外，表5中的第二组实验表明，在三个预注意力分支中使用聚合器的效果优于使用任何一个聚合器。在目标检测和语义分割中也有类似的实验结果。使用所有聚合器可以在一定程度上提高基线性能（例如，+0.7％APb和+0.5％APm）。这些结果表明，以更全面的方式建模相关性能够提供精细的视觉表示，从而有利于密集预测场景。

然后，我们分析了预注意力聚合器各种核大小对性能的影响。在不改变非注意力分支的情况下，我们用Agg1（3×3卷积）、Agg2（5×5）或Agg3（7×7）替换了所有预注意力聚合器。表5中的第二组结果表明，使用任何一组聚合器都可以提高分类和密集预测的性能，而3×3、5×5和7×7的组合产生了最理想的结果。具体来说，配备不同聚合器的GroupMixFormer-T在分类准确率上比基线高出+1.6％，在目标检测上高出+1.5％APb，在语义分割上高出+1.0％APm，这表明建模不同大小的组之间的相关性是提高性能的关键。

Depthwise convolution是一种有效的聚合器。请注意，聚合器的实现可以有多种。表6显示了我们关于不同聚合器实现的效果（例如，depthwise convolution[12]、max-pooling或average-pooling）。经验上观察到，由depthwise convolution实现的聚合器取得了略好的性能（分类Top-1准确率为82.5％，检测APb为42.5％，使用Mask R-CNN进行实例分割的APm为39.7％）。与最大池化和最小池化操作相比，卷积聚合器可以利用更多的可学习参数来计算相关性，从而实现更好的性能。

性能提升并非来源于宏观结构。与代表性作品[37、49、53]相比，我们的GroupMixFormer更深，补丁嵌入的实现也不同。为了证明性能提升不仅仅是由于更好的建筑超参数组合（包括Tab. 1中介绍的令牌维度、扩展比和层深度），我们将GroupMixFormer-T中的GMA块替换为Swin-attention或PVT-attention。Tab. 7中的结果显示，仅仅替换GMA会导致显著的性能下降，这证明了性能提升是由于先进的注意力机制而不是架构。

对于聚合器的核大小的最佳配置。为了找到最佳配置，我们采取了两种方法：(1)增大核大小，(2)以不同的顺序改变核配置。第一种方法涉及将核大小从(3,5,7)增加到(5,7,9)。对于第二种方法，我们在浅层部署具有较大核的聚合器，在深层部署较小核的聚合器，以及采用相反的配置。然而，正如Tab. 8所示，这些修改都没有我们最终采用的配置有效。

GMA不仅仅是卷积和自注意力的简单组合。我们进行了进一步的实验，以验证我们提出的GroupMixFormer本质上不同于卷积和自注意力的简单组合。具体来说，我们从GroupMixFormer-T中移除所有的组聚合器，并在整个自注意力模块之前插入一组以相同方式组织的卷积层（即，并行恒等映射、3×3、5×5和7×7层的组合）。Top-1准确率下降了1.0％（81.5％对比82.5％）。

聚合器是一种先进的通用构建块，可以应用于其他ViT。我们还可以通过简单地将它们插入原始注意力模块来将聚合器纳入代表性ViT（例如，Swin [37]和PVT [53]）以处理它们的Query、Key和Value。Tab. 9中的结果显示，这种策略通常可以显著提高ViT的性能。例如，带有聚合器的PVT-Small达到了80.6％的Top-1准确率，比其原始结果高出0.8％。这表明提出的聚合器通过建模组相关性来推进ViT，从而实现对令牌的全面理解。

5、结论

在本篇论文中，我们提出了一种先进的注意力机制，名为Group-Mix Attention（GMA）。与流行的仅用于建模单独令牌之间相关性的多头自注意力（MHSA）不同，提出的GMA利用组聚合器同时捕获令牌到令牌、令牌到组和组到组的相关性。我们基于GMA提出了GroupMixFormer，并实现了一系列不同大小的实用视觉主干。在标准视觉识别基准上的大量实验（包括图像分类、目标检测和语义分割）已经验证了提出的GMA和GroupMixFormer的有效性。