MaxViT: Multi-Axis Vision Transformer

最新推荐文章于 2024-04-20 09:37:06 发布

歪歪_yy

最新推荐文章于 2024-04-20 09:37:06 发布

阅读量2.5k

点赞数 4

分类专栏：论文阅读文章标签： transformer 深度学习计算机视觉

原文链接：https://arxiv.org/abs/2204.01697

版权

论文阅读专栏收录该内容

5 篇文章 0 订阅

订阅专栏

MaxViT: Multi-Axis Vision Transformer

[论文地址]([2204.01697] MaxViT: Multi-Axis Vision Transformer (arxiv.org))

[代码地址(其中的一个实现版本)]([ChristophReich1996/MaxViT: PyTorch reimplementation of the paper “MaxViT: Multi-Axis Vision Transformer” arXiv 2022]. (github.com))

摘要

由于自注意力的机制对于图像大小方面缺乏可扩展性，限制了它们在视觉主干中的应用。本文提出了一种高效的可拓展的全局注意，该模型包括两个方面：阻塞的局部注意和拓展的全局注意。作者通过将该注意模型与卷积有效结合，并简单的将这些模块堆叠，形成了了一个分层的视觉主干网络MaxVit。值得注意的是，MaxVit能在整个网络中看到全局甚至是在早期的高分辨率的阶段。在分类任务上，该模型在ImaegNet 1K上达到86.5%的 top-1准确率，在imageNet-21K上纪进行预训练，top-1准确率可以达到88.7%。对于一些下游任务如目标检测和图像美学评估，该模型同样具有好的性能。

引言

Vit只有经过广泛的预训练才能取得良好的性能，这是因为Vit容量大并且缺少归纳偏置，这导致模型容易过拟合。为了适当调整模型容量以提高其可伸缩性，后续许多工作研究为局部注意等视觉任务定制的稀疏变换模型。这些方法通常会重新引入层级式的结构，以补偿非局部性的损失。Swin-T是一个成功的尝试，并且在ImageNet测试中，该方法首次超过了ConvNet。虽然Swin-T基于窗口的注意力具有更大的灵活性和普遍性，但由于缺乏非局部性，因此不利于在更大的数据集上扩展。与此同时，在分层网络高分辨率的早期，获取全局交互的注意力计算量巨大的，如何在计算预算下有效地整合全局和局部交互以平衡模型容量和通用性仍然是一个挑战。

本文提出了一种新的Transformer块，称为多轴自注意力（Max-SA）,其作为一个基本架构组件，可以在单个块中执行局部和全局空间交互。与完全的自我注意力相比有更大的灵活性和效率，能适应不同的输入长度，具有线性复杂度；与局部和窗口注意力不同，Max-SA通过提出一个全局感受野，获得了更强的模型容量。因为Max-SA具有线性复杂度，可以在网络中的任何曾作为一般的独立注意模块，即使是在高分辨率的早期。

为了展示其有效性和通用性，我们通过堆叠由MaxSA和卷积组成的块设计了一个简单而有效的视觉主干模型称为多轴视觉转换器（MaxVit）。虽然该模型属于视觉转换器的范畴，但MaxVit的不同之处在于，我们力求通过一个基本块来统一卷积、局部和全局注意力然后简单的重复它。实验表明，MaxVit在所有数据模式下显出改善了SOTA性能，如图所示，其性能优于所有最近基于Transformer的模型。在这里插入图片描述

本文的主要贡献有：

一个通用的Transformer架构MaxViT，可以在网络的各个阶段捕获局部和全局空间交互。
一种新颖的独立多轴注意模块，由阻塞的局部注意和扩展的全局注意组成，以线性复杂度享受全局感知
实验表明MaxVit在包括图像分类、目标检测、图像美学评估和图像生成在内的广泛任务中实现了SOTA结果。

方法

我们提出了一种新的注意模块，称为MaxSA，通过将完全幂集的注意机制分解为两种稀疏形式：窗口注意和网格注意，将普通的二次复杂度将为线性且不损失任何的局部性。优于MaxSA的灵活性和可扩展性，可以通过简单地将MAXSA的可替换曾与MBConv堆叠在一个分层体系结构中构建一个视觉主干，称为MaxVit，模型结构如图所示。MaxVit从前层到深层整个网络的全局和局部感受野中获得优化，在模型容量和泛化能力方面表现出优异的性能。

在这里插入图片描述

Attention

自注意力允许整个空间位置混合，同时也受益于基于标准化承兑相似的内容相关权重。标准的自注意力是位置不敏感的。位置敏感是ConvNet中非常重要且固有的归纳偏置。相对自我注意提出通过在注意权重中引入相对学习偏差来改善普通注意，这一辈证明在许多视觉任务中优于原始注意，这项工作中主要采用[CoAtNet]([2106.04803] CoAtNet: Marrying Convolution and Attention for All Data Sizes (arxiv.org))中定义的预规范化相对自我注意作为MaxVit的关键算子。

Multi-axis Attention

与局部卷积相比，全局交互是自我注意的关键优势之一。然而，在整个空间直接应用注意力在计算上是不可行的，因为注意力算子具有二次复杂度。为了解决这个问题，我们提出一种多轴方法，通过简单地分解空间轴，将全尺寸注意力分解为两种稀疏形式：局部和全局。让X∈R^H×W×C作为输入特征图。我们没有将注意力放在展平的空间维度HW上，而实及那个恶整分割乘一个形状向量（h/P×h/P,P×P,C），表示划分为大小P×P地非重叠窗口。将子注意力应用于局部空间维度即P×P，相当于在一个小窗口内参与。我们将利用这一块注意（block attention）进行局部互动。

尽管解决了繁重计算的问题，但人们观察到局部注意模型在大规模数据集上拟合不足，受块注意力的启发，我们提出了一种简单有效的方法来获得稀疏的全局注意力，我们将其称为网格注意（grid attention）。我们没有使用固定的窗口大小来划分特征图，而实使用固定的G×G均匀网格将张量网格化为形状（G×G，H/G×W/G,C），在分解的网格轴上应用子注意力即G×G，对应于tokens扩展的全局空间混合。通过使用相同的固定窗口和网格大小（P=G=7），可以平衡局部和全局之间的运算，两者在空间大小和序列长度方面都只有线性复杂度。值得注意的是，我们提出的MaxSA可以替换Swin注意模块，其参数和触发器数量完全相同。然而，它具有全局交互能力，不需要mask、填充、或循环移位。这使得它比以为窗口方案更易于实现。

我们顺序地将两种类型的注意叠加在一个块中，以获得局部和全局交互，如图所示。

在这里插入图片描述

我们还在多轴注意之前添加了一个MBConv块。因为将MBConv于注意一起使用进一步提高了网络地泛化和可训练性。在注意之前使用MBConv层提供了另一个优势，深度卷积可视为条件位置编码，这使得模型没有明确的位置编码层，MaxVit块的伪代码如Algo1所示。在这里插入图片描述

请注意，我们提出的独立多轴注意可以一起使用的，也可以单独用于不同的目的，块注意（block attention）用于局部交互，网格注意（grid attention）用于全局交互。这些元素可以很容易地插入许多视觉体系结构中，尤其是在高分辨率任务中，这些任务可以通过可负担的计算进行全局交互。

结构变体

网络的变体结构如图所示。

在这里插入图片描述

实验

Image Classification on ImageNet-1K

在这里插入图片描述

在大规模数据集上预训练的模型性能

在这里插入图片描述

Object Detection and Instance Segmentation

在这里插入图片描述

Image Aesthetic Assessment

在这里插入图片描述

Image Generation

在这里插入图片描述

消融实验

在这里插入图片描述

模型细节

1.骨干细节

MBConv

MaxVit 将[MBConv block]([1801.04381] MobileNetV2: Inverted Residuals and Linear Bottlenecks (arxiv.org))作为主要的卷积算子。为了提高MBConv和transformer的同质性，我们采取了[预激活结构]([2106.04803] CoAtNet: Marrying Convolution and Attention for All Data Sizes (arxiv.org))。具体的，没有下采样的MBConv计算如下：

x ← x + Proj(SE(DWConv(Conv(Norm(x)))))

其中Norm指batch norm，之后Conv是扩展的1×1卷积和GELU激活函数，DWConv是3×3的深度可分离卷积，然后经过SE注意力层，最后是一个proj1×1的前向传播实现压缩与恒等映射相加。需要注意的是，对于每个stage的第一个MBConv块，通过应用3×3和步长为2的深度可分离卷积进行下采样，短路连接也要应用池化和通道映射，表示如下:

x ← Proj(Pool2D(x)) + Proj(SE(DWConv ↓(Conv(Norm(x)))))

相对注意力

仅用单头表示的相对注意力表示如下：

RelAttention(Q, K, V ) = softmax(QKT /√d + B)V

其中Q,K,V分别表示query，key，value矩阵。B是可学习的位置感知矩阵且B∈R^{(2H-1)×(2W-1)}

多轴注意力

对于一维的输入序列，始终将输入的倒数第二个维度（…,L,C）作为空间轴，其中L，c表示序列长度和通道。

首先我们定义了block（·）操作将输入特征图转化为P×P大小的非重叠窗口，然后定义unblock（·）执行与上述相反的操作

Block : (H, W, C) → (H/P × P, W/P × P, C) → (HW/P ² , P ², C)

同样地，我们定义Grid(·）将输入特征转化为G×G的网格。与Block操作不同的是，我们需要额外的转置将网格维度放置在假定的空间轴上。同时定义Ungrid（·）执行逆操作将其返回为2D空间。

在这里插入图片描述

总结：对于给定X∈R^{（H×W×C）}，block attention可表示为：

x ← x + Unblock(RelAttention(Block(LN(x))))
x ← x + MLP(LN(x))

grid attention可表示为：

x ← x + Ungrid(RelAttention(Grid(LN(x))))
x ← x + MLP(LN(x))

与轴注意力的比较

在这里插入图片描述

轴注意力先计算列向注意再计算行向注意，计算复杂度O（N/sqrt(N)），N表示元素数。而提出的轴注意力首先使用局部注意，然后使用稀疏的全局注意,拥有O(N)线性复杂度的全局感受野。

2.结构详情

在这里插入图片描述

歪歪_yy

关注

4
点赞
踩
16

收藏

觉得还不错? 一键收藏
0
评论
MaxViT: Multi-Axis Vision Transformer

MaxViT: Multi-Axis Vision Transformer[论文地址]([2204.01697] MaxViT: Multi-Axis Vision Transformer (arxiv.org))[代码地址(其中的一个实现版本)]([ChristophReich1996/MaxViT: PyTorch reimplementation of the paper “MaxViT: Multi-Axis Vision Transformer” arXiv 2022]. (github.co
复制链接

扫一扫