|
- 论文地址：https:///abs/2204.01697
Transformer最近在计算机视觉界引起了极大的关注。然而，自注意力机制在图像大小方面缺乏可扩展性，限制了其在最先进的视觉主干中的广泛采用。在本文中，作者提出了一种高效且可扩展的注意力模型，称之为多轴注意力，它包括两个方面：块局部（blocked local）注意力和扩展全局（dilated global）注意力。这些设计选择允许在任意输入分辨率上进行全局-局部空间交互，仅具有线性复杂度。作者还提出了一个新的架构元素，有效地将本文提出的注意力模型与卷积相结合，并通过简单地在多个阶段重复基本构建块提出了一个简单的分层视觉主干，称为MaxViT。值得注意的是，MaxViT能够在整个网络中“看到”全局，甚至在早期的高分辨率阶段。作者证明了本文的模型在广泛的视觉任务中的有效性。在图像分类方面，MaxViT在各种设置下都达到了最先进的性能：在没有额外数据的情况下，MaxViT达到了86.5%的ImageNet-1K top-1精度；通过ImageNet-21K预训练，本文的模型达到了88.7%的top-1精度。对于下游任务，MaxViT作为主干在目标检测和视觉美感评估方面提供了良好的性能。作者还表明，本文提出的模型在ImageNet上表现出强大的生成建模能力，证明了MaxViT块作为通用视觉模块的优越潜力。自AlexNet以来，卷积神经网络（ConvNet）一直是计算机视觉的主导架构设计选择。ConvNet通过深入、拓宽、添加密集连接、高效可分离卷积、空洞卷积、使用编码器-解码器框架甚至引入现代微设计组件，继续在众多视觉问题上表现出色。同时，受自然语言处理中Transformers等自注意力模型演变的启发，许多研究人员已经开始将注意力机制引入视觉。视觉Transformer（ViT）可能是第一个完全基于Transformer的视觉架构，其中图像patch仅被视为文字序列，Transformer编码器应用于这些视觉token。当对大规模数据集进行预训练时，ViT可以在图像识别方面取得令人信服的结果。然而，已经观察到，如果没有广泛的预训练，ViT在图像识别方面表现不佳。这是因为Transformer的模型容量很强，充满了较少的归纳偏置，这会导致过配合。为了适当地调整模型容量并提高其可扩展性，许多后续工作研究了为视觉任务（如局部注意）定制的稀疏Transformer模型。这些方法通常重新引入层次结构以补偿非局部性的损失。Swin Transformer就是这样一种成功的尝试，通过在移位的非重叠窗口上应用自注意力来修改Transformer。这种方法在ImageNet基准上优于纯视觉Transformer和ConvNet。尽管与ViT中使用的完全注意力相比，基于窗口的注意力具有更大的灵活性和通用性，但由于非局部性的损失，已观察到基于窗口的注意力的模型容量有限，因此在更大的数据区域（如ImageNet-21K和JFT）上的扩展不利。然而，由于注意力算子需要二次复杂度，因此在层次网络的早期或高分辨率阶段通过完全注意力获取全局交互需要大量计算。如何有效地结合全局和局部交互，在计算预算下平衡模型容量和可推广性仍然是一个挑战。在本文中，作者提出了一种新型的Transformer模块，称为多轴自注意力（Max SA），它可以作为基本的架构组件，在单个块中执行局部和全局空间交互。与完全自注意力相比，Max SA具有更大的灵活性和效率，即自然适应线性复杂度的不同输入长度；与窗口/局部注意力不同，Max SA通过提出一个全局接受域，允许更强的模型容量。此外，MaxSA仅具有线性复杂度，即使在早期的高分辨率阶段，也可以在网络的任何层中用作一般的独立注意力模块。
为了证明其有效性和通用性，作者进一步设计了一种简单但有效的视觉主干，称为多轴视觉Transformer（MaxViT），通过分层堆叠由Max-SA和卷积组成的重复块。虽然作者提出的模型属于混合视觉Transformer的范畴，但MaxViT与以前的方法不同之处在于，我们力求简单，通过设计一个基本块来统一卷积、局部和全局注意力，然后简单地重复它。本文的实验表明，对于广泛的视觉任务，MaxViT在所有数据模式下显著提高了最先进的（SOTA）性能。具体而言，如上图所示，MaxViT在精度与浮点运算以及精度与参数曲线方面均优于所有最近基于Transformer的模型。提出一种通用的Transformer架构：MaxViT，可以在网络的每个阶段捕获局部和全局空间交互。 提出一种新颖的独立多轴注意力模块，由块局部注意力和扩展全局注意力组成，以线性复杂度享受全局感知。 大量实验表明，MaxViT在各种数据模式下实现了SOTA结果，用于广泛的任务，包括图像分类、目标检测、图像美学评估和图像生成。
方法
本文引入了一种新型的注意力模块，称为多轴自注意力（MaxSA），通过将完全密集的注意力机制分解为两种稀疏形式（窗口注意力和网格注意力）将普通注意力的二次复杂度降低到线性，而不失任何非局部性。由于Max SA的灵活性和可扩展性，可以通过在分层架构中简单地将Max SA的替代层和MBConv堆叠，来构建一个视觉主干，称之为MaxViT，如上图所示。MaxViT在整个网络中，从浅到深，从全局和局部接受域中受益，在模型容量和泛化能力方面表现出卓越的性能。3.1 Attention自注意力允许整个空间位置的空间混合，同时也受益于基于归一化成对相似性的内容相关权重。标准自注意力是位置不可感知的，即非平移等变性，这是一种重要的归纳偏见。相对自注意力被提出，通过在注意力权重中引入相对学习偏置来改善普通注意力，已证明在许多视觉任务中，相对学习偏置始终优于原始注意力。在这项工作中，作者主要采用预归一化相对自注意力作为MaxViT中的关键算子。3.2 Multi-axis Attention与局部卷积相比，全局交互是自注意力的关键优势之一。然而，直接沿整个空间应用注意力在计算上是不可行的，因为注意力算子需要二次复杂度。为了解决这个问题，作者提出了一种多轴方法，通过简单地分解空间轴，将全尺寸注意力分解为两种稀疏形式——局部和全局。设为输入特征图。作者没有将注意力放在铺平的空间维度HW上，而是将特征分为形状张量，表示划分为大小为P×P的非重叠窗口。将自注意力应用于局部空间维度，即P×P，相当于在一个小窗口内进行注意力。作者将使用这种块注意力（block attention）来进行局部交互。尽管绕过了繁重的完全自注意力计算，但在大规模数据集上观察到局部注意力模型不足。受块注意力的启发，作者提出了一种简单但有效的方法来获得稀疏的全局注意力，作者称之为网格注意力（grid attention）。作者使用固定的G×G均匀网格将张量网格化为形状，而不是使用固定的窗口大小对特征图进行分区，从而使窗口具有自适应大小。在分解的网格轴上使用自注意力，即G×G，对应于token的扩展、全局空间混合。通过使用相同的固定窗口和网格大小，可以充分平衡局部和全局运算之间的计算，两者在空间大小或序列长度方面仅具有线性复杂度。注意，本文提出的Max SA模块可以用相同数量的参数和FLOPs来替代Swin注意力模块。然而，它具有全局交互能力，不需要mask、填充或循环移位，使其更易于实现，优于移位窗口方案。
作者连续交替的采用了这两种注意力机制，以获取局部和全局建模。作者还在多轴注意力之前采用了MBConv，因为之前的工作表明这个模块增加模型的泛化性能。其中，MBConv中的深度卷积也可以看做是位置编码，MaxViT的伪代码如算法1所示。
本文提出的多轴注意力可以一起使用，也可以单独使用，用于不同的目的——块注意力用于局部交互，网格注意力用于全局混合。这些元素可以很容易地插入到许多视觉体系结构中，尤其是在高分辨率任务中。3.3 Architecture Variants作者设计了一系列极其简单的架构变体，以探索本文提出的MaxViT块的有效性。作者使用了类似于常见ConvNet实践的层次主干，其中输入首先在stem阶段（S0）使用Conv3x3层进行降采样。网络主体包含四个阶段（S1-S4），每个阶段的分辨率为前一阶段的一半，通道数增加了一倍（隐藏维度）。在本文的网络中，作者在整个主干中使用相同的MaxViT块。在每个阶段的第一个MBConv块的深度Conv3x3层中应用下采样。默认情况下，反向瓶颈和SE模块的膨胀率和收缩率分别为4和0.25。作者将所有注意力块的注意力头大小设置为32。通过增加每个阶段B的块数和通道维度C来扩大模型。下表总结了MaxViT变体的架构配置。
实验
作者在上表中展示了ImageNet-1K分类的性能比较。在基本224×224设置下，MaxViT在整个FLOPs中大大优于最近的强混合模型CoAtNet，此外在性能上也有明显的优势。
上表显示了在ImageNet-21K上预训练的模型的结果。值得注意的是，MaxViT-B模型实现了88.38%的精度，仅使用43%的参数计数和38%的浮点运算，就比之前的最佳模型CoAtNet-4高出0.28%。
上图全面显示了模型大小比较MaxViT量表明显优于先前具有类似复杂性的基于注意力的模型。此外，MaxViT XL模型实现了新的SOTA性能，在分辨率512×512下进行微调时，精度达到88.70%。
本文的模型也可扩展到大规模训练数据–MaxViT XL在4.75亿个参数下实现了89.53%的高精度，在可比模型大小下优于以前的模型。
上表展示了在目标检测和分割任务中使用的MaxViT主干模型在精度和效率方面大大优于所有其他主干，包括各种模型大小的Swin、ConvNeXt和UViT。
为了将本文的模型与现有方法进行评估和比较，作者在上表中总结了结果。对于相似的输入分辨率，提出的MaxViT-T模型优于现有的图像美学评估方法。
在生成任务上，本文的MaxViT实现了更好的FID，并且参数数量明显减少。这些结果证明了MaxViT块对生成任务的有效性。
本文的主要贡献之一是网格注意力模块，它允许在线性时间内进行稀疏的全局交互，使模型能够捕获所有阶段的全局信息。作者进行了两次消融以了解其收益：1）在每个阶段完全消除了全局注意力；2） 将网格注意力替换为块注意力，以保持相同的参数计数和FLOPs。如上表所示，与仅使用局部注意力或卷积相比，在早期阶段启用全局注意力可以进一步提高性能。
作者还通过删除每个阶段中的所有MBConv来消除MaxViT中MBConv层的使用。如上表所示，在MaxViT中使用MBConv层显著提高了性能。
本文提出了三个不同的模块来构建MaxViT块——MBConv、块和网格注意力——它捕捉从局部到全局的空间交互。为了研究最有效的组合方式，作者使用所有6种排列评估了MaxViT-T模型，结果如上表所示。
作者还探究了顺序和并行堆叠模块的影响，从上表中可以看出顺序方法在参数和计算量较少的情况下显著优于并行方法。原因可能是并行设计学习互补线索，它们之间的交互更少，而本文的顺序堆栈能够学习局部和全局层之间更强大的融合。
上图展示了本文方法和Swin在不同模型大小下的实验结果。
总结虽然这几年的最新研究表明，ConvNet和视觉Transformer可以在图像识别方面实现类似的性能，但本文的工作提出了一种统一的设计，它利用了两个模型的优点（高效卷积和稀疏注意力）并证明了一个建立在上面的模型，即MaxViT，可以在各种视觉任务上实现最先进的性能，更重要的是，可以很好地扩展到大规模数据。尽管作者在视觉任务的背景下提出了本文的模型，但所提出的多轴方法可以很容易地扩展到语言建模，以在线性时间内捕获局部和全局依存关系。此外，还可以在MLP-Mixer或gMLP模型上采用所提出的多轴方法，以提高可扩展性和效率。参考资料[1]https:///abs/2204.01697研究领域：FightingCV公众号运营者，研究方向为多模态内容理解，专注于解决视觉模态和语言模态相结合的任务，促进Vision-Language模型的实地应用。
|
