GroupMixFormer:Advancing Vision Transformers with Group-Mix Attention论文学习笔记

发布时间:2024年01月20日

论文地址:https://arxiv.org/pdf/2311.15157.pdf

代码地址:https://github.com/AILab-CVC/GroupMixFormer

摘要:ViT?已被证明可以通过使用多头自注意力 (MHSA) 对远程依赖关系进行建模来增强视觉识别,这通常被表述为Query-Key-Value?计算。但是,从“Query”和“Key”生成的注意力图仅捕获单个粒度的token-to-token的相关性。在本文中,作者认为自注意力应该有一个更全面的机制来捕捉?tokens和groups(即多个相邻tokens)之间的相关性,以获得更高的表征能力。因此,作者提出Group-Mix Attention(GMA)作为传统自注意力的高级替代品,它可以同时捕获具有各种group规模的token-to-token, token-to-group, and group-to-group的相关性。为此,GMA 将 Query、Key 和 Value 统一拆分为多个段,并执行不同的分组聚合以生成分组代理(group proxies)。注意力图是根据tokens和group proxies的混合计算的,用于在 Value 中重新组合tokens和groups。

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


简单总结:本文通过将QKV统一划分为5段,pre-attention branche(四段)其中三段分别做核大小为3/5/7(即滑动窗大小)的深度可分离卷积来实现聚合器,将滑动窗周围的信息聚合在一个token中,其中一个不做聚合,以保持网络在建模单个token相关性方面的能力,最后一段被建模为非注意力分支,该段有聚合器但没有注意力,如图三所示。利用了组聚合器同时捕捉token-to-token、token-to-group和group-to-group的关联。


1. Introduction

????????多头自注意力 (MHSA) 是显著提高ViT性能的一个关键模块,它支持使用长程依赖建模 、全局感受野、更高的灵活性进行网络设计和更强的鲁棒性。通常的注意力计算,关注的是单个token对之间的相关性计算。然而,经验发现,Q-K-V自注意力存在一个主要局限性,如图1所示:注意力图仅在一个粒度上描述了每个单独的token对之间的相关性(图1(a)),并且将注意力图与Value相乘会线性重组单个token。该框架显然没有考虑不同维度上不同token groups(即邻域)之间的相关性。举一个具体的例子,自注意并没有将左上角的九个tokens作为一个整体与右下角的那些组相关联。这个限制虽然很明显,但无意中被忽略了,因为Q-K-V计算似乎能够对从输入到输出的映射进行建模,因为输出中的任何entry都与输入中的每个entry有关。

????????可以观察到,生成的注意力图仅捕捉了单一粒度上的 token 到 token 的关联,从而忽略了不同 token 组之间的关联。为解决这一局限性,本文提出了 Group-Mix Attention(GMA)。GMA 将 token 分割为片段,并通过组聚合器生成组代理替代一些个别 token。GMA旨在高效计算 token-to-token、token-to-group 以及 group-to-group的关联,提供更全面的建模方法。

????????在自注意力中,关注的焦点是token对之间的线性关系,而GMA引入了组代理的概念,使其能够在更高的层次上捕捉token组之间的关联。通过组代理,GMA能够在图像中的不同区域之间建立关联,而不仅仅是在个别tokens之间。这种设计使得GMA能够更全面地理解图像的结构信息,为视觉识别任务提供更强大的建模能力。图1(c)和(d)进一步说明了GMA在计算组到组关联时的具体步骤,突显了通过组代理和聚合实现的高效计算。这种设计使得GMA不仅能够更全面地捕捉tokens之间的关联,而且在计算效率上也具有优势。

GMA具有一些吸引人的优点:

(1)GMA不仅能够模拟单个token之间的相关性,还能够模拟token group之间的相关性。将不同种类的注意力混合在一起,以便从综合方面更好地了解token。 token-to-token、token-to-group 以及 group-to-group的相关性在每个单层中同时建模,以实现更高的表示能力。

(2)GMA高效且易于实施。group-to-group的相关性是通过将组聚合到?proxy tokens中,然后计算代理之间的相关性来计算的(如图 3 所示)。这样的过程可以通过基于滑动窗口的操作(例如池化和卷积)有效地实现。

2. Related Works

省略。。。

3. GroupMix Attention and GroupMixFormer

3.1. Motivation: from Individual to Groups

作者从最初的注意力机制讨论自注意力的局限性,如下式:

作者认为上式的局限在于,不能将某些特定的token建模成具有不同粒度的组的形式,上式仅考虑了单个token之间的相关性。在本文中,作者试图利用单个模式和组模式进行综合建模。与以前跨多个阶段(通常四个阶段Transformer主干网络)的不同模式建模的方法不同,作者的方法引入了一种新的方法,在每个阶段的每个单独层中对这个建模过程进行编码。具体来说,对于组模式,试图将一些token周围的token与其他token相关联。本文提出通过在 Query、Key 和 Value 中生成组代理来实现这一点(其实就是通过卷积和池化将token周围的信息聚合到一个token中来实现),并使用代理执行 Q-K-V 计算。

3.2. GMA: Mixing Groups for Better Attention

  1. 生成组代理:?首先,GMA引入了组代理的概念,通过替换Query、Key、和Value中的一些条目为整个组的聚合结果,这是通过滑动窗口操作(例如maxpooling、convolution等)高效实现的,其中不同的组使用不同大小的滑动窗口尺寸。这种替换将注意力从仅关注个别token扩展到了关注整个组。

  2. 关联不同大小组和个别token:?GMA通过使用不同大小的聚合器,对不同大小的组进行混合,同时通过深度卷积实现聚合,使得模型可以有效地关联不同大小的组和个别token。这在视觉任务中特别重要,因为图像中的信息可能以不同大小的组织存在。

  3. 维持特征分辨率:?在聚合的过程中,GMA保持了特征的分辨率,确保在注意力计算中不降低空间分辨率。这意味着GMA为注意力计算提供了更细粒度的特征,相较于降低特征大小的方法具有更好的性能。

  4. 同时关联多个token:?由于输入是组代理,而不是单个token,GMA能够同时关联多个token,这在建模关联方面更为全面和充分。这是通过在计算中使用不同的核大小来实现的。

  5. 结合个别token和组:?最后,通过将注意力计算得到的映射与Value相乘,实现了将关联的组和个别token重新组合,从而更好地捕捉了图像中的结构信息。

总体来说,GMA通过引入组代理和巧妙的聚合操作,实现了对不同大小组和个别token关联的全面建模,从而提升了模型在视觉识别任务中的性能。这种机制使得模型能够更全面地捕捉图像中的结构信息,提高了对不同尺度和层次的视觉模式的敏感性,使得模型更加适应各种复杂的视觉场景。

code

以下是本文的核心代码

class Aggregator(nn.Module):
    def __init__(self, dim, seg=4):
        super().__init__()
        self.dim = dim
        self.seg = seg

        seg_dim = self.dim // self.seg

        self.norm0 = nn.SyncBatchNorm(seg_dim)
        self.act0 = nn.Hardswish()

        self.agg1 = SeparableConv2d(seg_dim, seg_dim, 3, 1, 1)
        self.norm1 = nn.SyncBatchNorm(seg_dim)
        self.act1 = nn.Hardswish()

        self.agg2 = SeparableConv2d(seg_dim, seg_dim, 5, 1, 2)
        self.norm2 = nn.SyncBatchNorm(seg_dim)
        self.act2 = nn.Hardswish()

        self.agg3 = SeparableConv2d(seg_dim, seg_dim, 7, 1, 3)
        self.norm3 = nn.SyncBatchNorm(seg_dim)
        self.act3 = nn.Hardswish()

        self.agg0 = Agg_0(seg_dim)


    def forward(self, x, size, num_head):
        B, N, C = x.shape
        H, W = size
        assert N == H * W

        x = x.transpose(1, 2).view(B, C, H, W)
        seg_dim = self.dim // self.seg  # 按维度进行分组,相当于分了seg个head,每个head的维度为seg_dim

        x = x.split([seg_dim]*self.seg, dim=1)

        x_local = x[4].reshape(3, B//3, seg_dim, H, W).permute(1,0,2,3,4).reshape(B//3, 3*seg_dim, H, W)
        x_local = self.agg0(x_local)  # 把qkv的信息通过深度可分离卷积进行聚合

        # 相当于在每个head上使用不同大小的核做深度可分离卷积,得到不同尺度的编码信息
        # 就是将不同的核大小视为滑动窗的大小,并将滑动窗中的信息通过卷积编入滑动窗的中心,并以该中心作为该组的proxy
        x0 = self.act0(self.norm0(x[0]))
        x1 = self.act1(self.norm1(self.agg1(x[1])))  # 使用3*3的深度可分离卷积
        x2 = self.act2(self.norm2(self.agg2(x[2])))  # 使用5*5的深度可分离卷积
        x3 = self.act3(self.norm3(self.agg3(x[3])))  # 使用7*7的深度可分离卷积

        x = torch.cat([x0, x1, x2, x3], dim=1)

        C = C // 5 * 4
        x = x.reshape(3, B//3, num_head, C//num_head, H*W).permute(0, 1, 2, 4, 3)

        return x, x_local

3.3. Architectural Configurations

1. 层级拓扑:基于提出的Group-Mix Attention(GMA),引入了一系列名为GroupMixFormer的视觉Transformer模型,模型采用了具有四个阶段的分层拓扑结构,与主流的ViT模型保持一致。

2. 图像嵌入层:模型的第一层是一个4×的图像嵌入层,将图像嵌入为tokens。此过程通过两个连续的3×3卷积层实现,每个卷积层的步幅为2,然后是另外两个3×3的卷积层,步幅为1。在每个最后的三个阶段的开始,使用了一个2×的图像嵌入,同样通过3×3卷积实现。

3. 编码器块:在每个阶段内,构建了多个编码器块。每个编码器块除了包含上一子节中介绍的GMA块之外,还包含了一个Feed-Forward Network(FFN)、Layer Normalization和identity快捷连接,这符合大多数网络架构的通用处理方法。

4. 输出层:对于图像分类任务,最终的输出tokens通过全局平均池化(GAP)后输入分类器;对于密集预测任务(例如目标检测和语义分割),各自任务的头部可以利用四个阶段输出的金字塔特征。

5. 位置编码:与一些Transformer模型不同,该模型没有采用位置编码,因为通过GMA聚合器,模型自然地破坏了排列不变性。

6. 不同配置的模型:模型实例化了四个不同配置的模型,这些配置的架构超参数包括每个阶段的编码器块数目L、嵌入维度D和MLP比例R,具体见下表。

4. Experiments

作者提出的GroupMixFormer在多个任务中都做到了SOTA,主要是这里没有和TransNext作比较,将组代理的思想和TransNext的聚焦注意力相结合应该能达到更好的效果

分类模型的结果对比

?检测的结果对比

?分割模型结果对比

在附录中做了可视化的分析

在Attention Visualization部分的可视化结果如上图所示,当中展示了输入图像(a)以及来自集成层的注意力响应图(b)。此外,还展示了来自前注意力分支和非注意力分支的输出的响应图(c)到(g)。

????????可以观察到,在应用自注意力于个别tokens时,有时会无法关注到对象,如(c)所示。在这种情况下,计算由聚合器生成的组代理之间的关联可能会有所帮助。例如,如第三行所示,通过处理由核大小为3和7的聚合器处理的组之间的关联,成功地集中注意力于狗,而在(c)中建模token-to-token的关联更多地关注背景。这些结果表明存在某些模式,因此应将一些tokens视为整体以捕捉对象特征。在GMA中,不同聚合器捕获的表示被组合在一起。这验证了全面建模token-to-token、token-to-group和group-to-group关联会导致更好的视觉识别。

文章来源:https://blog.csdn.net/athrunsunny/article/details/135714705
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。