这里提出了一种先进的注意力机制 Group-Mix Attention(GMA)。与流行的多头自注意力相对,文本方法仅对个别token之间的关联进行建模不同,所提出的GMA利用了组聚合器同时捕捉token-to-token、token-to-group和group-to-group的关联。最后,基于GMA提出了GroupMixFormer,并实例化了一系列具有不同规模的实用视觉骨干网络。最新开源注意力机制,高效神经网络架构助力各大视觉任务显著涨点

标题:Advancing Vision Transformers with Group-Mix Attention

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

源码:https://github.com/AILab-CVC/GroupMixFormer

本文致力于解决 Vision Transformers(ViTs) 在视觉识别任务中采用的传统 Query-Key-Value(Q-K-V) 自注意机制存在的局限性。文章主张现有的自注意框架仅在单一粒度上捕捉了 token 之间的关联,忽略了 token 组之间的关联。

为解决这一问题,作者提出了 Group-Mix Attention(GMA) 作为传统自注意的先进替代方案。GMA 引入了一种机制,能够同时捕捉各种组大小的 token-to-token、token-to-group 以及 group-to-group 的关联。为实现这一目标,GMA 将 Query、Key 和 Value 进行划分,执行不同的组聚合并生成组代理。基于 token 和 group 代理的混合,计算了注意力图,有助于全面捕捉关联(全局注意力)。进一步地,基于 GMA 搭建了 GroupMixFormer,构成一个强大的骨干网络,在图像分类、目标检测和语义分割方面展现出比现有模型更好的性能,同时参数更少。

GroupMixFormer_人工智能

动机

老生常谈,ViTs 成功的一大关键要素得益于多头自注意力(MHSA)模块。MHSA 使得网络设计具备长距离依赖建模、全局感受野、灵活性和鲁棒性等优势。然而,如下图所示:

GroupMixFormer_建模_02

可以观察到,生成的注意力图仅捕捉了单一粒度上的 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(Group-Mix Attention)的作用机制:

GroupMixFormer_卷积_03

  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关联的全面建模,从而提升了模型在视觉识别任务中的性能。这种机制使得模型能够更全面地捕捉图像中的结构信息,提高了对不同尺度和层次的视觉模式的敏感性,使得模型更加适应各种复杂的视觉场景。

网络架构

最后,基于 GMA,我们便可以构建出一个强大的骨干网络——GroupMixFormer!

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,具体见下表。

GroupMixFormer_Group_04

伪代码实现

GroupMixFormer_卷积_05

实验

分类/检测/分割

GroupMixFormer_Group_06

GroupMixFormer_人工智能_07

可视化分析

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

GroupMixFormer_Group_08

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

总结

本文提出了一种先进的注意力机制 Group-Mix Attention(GMA)。与流行的多头自注意力相对,文本方法仅对个别token之间的关联进行建模不同,所提出的GMA利用了组聚合器同时捕捉token-to-token、token-to-group和group-to-group的关联。最后,基于GMA提出了GroupMixFormer,并实例化了一系列具有不同规模的实用视觉骨干网络。在标准的视觉识别基准上进行的广泛实验证明了所提出的GMA和GroupMixFormer的有效性,包括图像分类、目标检测和语义分割等任务。