Advancing Vision Transformers with Group-Mix Attention 论文解析

论文地址：

https://arxiv.org/pdf/2311.15157v1

代码地址：

GitHub - AILab-CVC/GroupMixFormer: GroupMixAttention and GroupMixFormer

---

目录

摘要

创新点

群体混合注意力 GMA

        自注意力

        GMA

分层视觉转换器 GroupMixFormer

---

Advancing Vision Transformers with Group-Mix Attention

摘要

视觉变形(ViTs) 通过对 多头自注意(MHSA) 的远程依赖关系建模来增强视觉识别，而多头自注意通常被表述为查询键值计算。然而，由 Query 和 Key 生成的注意图仅以单一粒度捕获 token 到 token 的相关性。本文认为自注意应该有一个更全面的机制来捕捉 token 和 group (即多个相邻token)之间的相关性，以获得更高的表征能力。

因此，我们提出 群体混合注意(GMA) 作为传统自我注意的高级替代品，它可以同时捕获不同群体规模的 token-token、 token-组 和 组-组 的相关性。为此，GMA 将 Query、 Key 和 Value 统一划分为不同的段，并进行不同的组聚合，生成组代理。注意图是基于 token 和 组代理的混合计算的，并用于重新组合值中的 token 和组。

基于GMA，我们引入了一个强大的骨干，即 GroupMixFormer，它在图像分类、目标检测和语义分割方面达到了最先进的性能，比现有模型的参数更少。例如，GroupMixFormer-L(参数为70.3M，输入3842个)在没有外部数据的情况下，在ImageNet-1K上达到了86.2%的Top1精度，而GroupMixFormer-B(参数为45.8M)在ADE20K上达到了51.2%的mIoU。

创新点

由上述可得本篇论文的创新点有以下两点：

1. 群体混合注意力（GMA）
2. 基于GMA提出的骨干网络 —— GroupMixFormer

用以解决以下问题：

• 在自注意力机制未能考虑不同粒度的不同令牌组（即邻域）之间的相关性。 

接下来我们来详细解释这两点创新点

群体混合注意力 GMA

GMA将 token 分割成统一且不同的部分，并且通过组聚合器生成的组代理替换一些单独的 token，如下图的（b）。之后使用查询和密钥（其中一些 token 已经被组代理替换）计算注意力图，并使用它来将组代理与值中的各个标记重新组合在一起

下图为文中给出的自注意力和组混合注意力GMA之间的概念比较

 在图中（a）和（b）中，为7 * 7的单维 token。

注意力图仅描述单个粒度上每个单独 token 对之间的相关性（a），并且乘以 V，仅能线性地重新组合各个 token 。

与自注意力计算单个 token 对的相关性不同，GMA 通过组聚合器创建 tokens 组（比如九个相邻 token）的代理，然后通过代理计算组与组之间的相关性。

在（c）和（d）中，展示了具有七个四维 token 的 GMA 具体计算。为了计算两个突出显示的组（每个组由三个 token 组成）之间的相关性，我们将他们聚合成两个代理以进一步相乘。组聚合可以通过滑动窗口的算子有效地实现。

自注意力

首先讨论一下自注意力的局限性，令为输入标记，其中 N 为标记数量，d 为维度，则普通自注意力的输出是：

 式中为了简洁起见忽略了归一化因子

 直观上根据矩阵乘法的定义，计算的是每两个标记之间的相关性。softmax 函数的输出（以下被称为A）被称为注意力图。AX相乘则意味着根据每个位置的注意力图线性的重新组合 token。

若我们现在存在某种模式（即组模式），需要将某种特定标记视为具有不同粒度的组。然而，自注意力明显缺乏对此类模式进行建模的明确机制，因为它仅仅思考单个粒度（即单个模式）的单个 token 对之间的相关性。

GMA

文章寻求利用个体模式和组模式进行综合建模。与先前跨多个阶段（通常是Transformer主干的四个阶段）建模不同模式的方法不同，我们的方法引入了一种在每个阶段的每个单独层内编码此建模过程的方法。

具体来说：

将 Q / K / V 条目统一分为 n 段，并对某些段进行聚合（每个分段的聚合器可能不同）。为了执行注意力计算，我们将聚合段和 n 段拼接（concat）在一起这样我们就得到了组代理Q '、K '、V '。之后对组代理执行注意力计算以生成输出。文中使用不同内核大小的深度卷积来实现聚合器。

此外，文章还在一个片段上采用身份映射而不是聚合器来维持网络对单个令牌相关性进行建模的能力。因此我们可以在计算注意力图时对组和 token 之间的相关性进行建模。将注意力图与 V 相乘可以被视为将响应的组与各个 token 相应地重新组合在一起。

 上图为GMA结构设计以及GroupMixFormer的架构，图中逻辑为：

        使用三个可学习的线性投影来生成Q、K 和 V。然后将其分为五个部分，并且使用不同内核大小（分别为3，5，7）的聚合器在其中四个上生成组代理，以便对不同粒度的单个令牌和组代理的混合进行注意力计算。这四个分支被称为预注意力分支，在其中三个注意力分支中使用不同的实现（例如：最小池化、平均池化、最大池化、深度卷积）作为具有不同内核大小的聚合器，最后一个预注意力分支不采用聚合器是为了进一步使结构多样化。使其成为恒等映射。除了这个带有注意力但是没有聚合器的分支之外，还构造了一个带有聚合器但没有注意力的分支，称为非注意力分支。最后输出由 token 集成层混合，该层具有归一化和激活的线性投影来简单实现。

文章认为 GMA 的有点有如下两点：

1. GMA不仅能够对单个 token 之间的相关性进行建模，而且还能够对 tokens 组之间的相关性进行建模。不同种类的注意力混合在一起，从全面的角度更好的理解 token。token 到 token、token 到 组、组 到 组的相关性在每个单层内同时建模，以获得更高的表示能力
2. GMA高效且易于实施，组与组之间的相关性是通过将组聚合为代理 token，然后计算代理之间的相关性来计算的，这样的过程可以通过滑动窗口的操作（例如池化和卷积）有效实现

分层视觉转换器 GroupMixFormer

在GMA基础上开发了一个分层视觉转换器，可以作为各种视觉任务的主干网。 

如图所示，文章采用具有四级的分层拓扑。前四个Patch embedding将图像嵌入到 token 中，这是通过两个连续的3 * 3卷积层实现的，每个卷积层的步长为2，另外两个3 * 3层的步长为1。在后三层开始之前，也就是第二个Patch Embedding采用两个3 * 3卷积实现。每一阶段的编码块中，除了包含我们上面介绍过的GMA之外，还包括前馈网络（FFN）、层归一化（LayerNorm）和恒等快捷方式（跳跃连接）。

对于图像分类，最终输出标记根据全局平均池化（GAP）后输入分类器；

对于密集预测任务（例如：对象检测和语义分割），特定任务头可以利用四个阶段输出的金字塔特征来进行预测。

模型中没有采用位置编码器，因为文章自然地打破了GMA聚合器的排列不变性。