MOAT: ALTERNATING MOBILE CONVOLUTION AND ATTENTION BRINGS STRONG VISION MODELS

paper: https://arxiv.org/pdf/2210.01820.pdf
code: https://github.com/RooKichenn/pytorch-MOAT

一、引言

本文介绍了MOA-T，这是一个建立在MObile卷积（即反向残差块）和Attension之上的神经网络家族。与当前堆叠分离的移动卷积和transformer的工作不同，本文有效地将它们合并为MOA-T。从标准Transformer块开始，将其多层感知器替换为MObile卷积，并在自注意力操作之前对其进行重新排序。MObile卷积不仅增强了网络表示能力，而且产生了更好的下采样特征。简单的MOA-T网络非常有效，通过ImageNet-22K预训练，在ImageNet-1K上实现了89.1%的top1精度。此外，通过简单地将全局注意力转换为窗口注意力，MOA-T可以无缝地应用于需要大分辨率输入的下游任务。由于MObile卷积可以有效地在像素之间交换局部信息（从而跨窗口），MOA-T不需要额外的窗口移动机制。
在这项工作中，通过深入研究MBConv和Transformer的组合来研究微观层次的构建块设计。在仔细检查这些块后，进行了两项关键观察。首先，Transformer模块中的MLP模块与MBConv相似，因为它们都采用了反向瓶颈设计。然而，MBConv是一种更强大的运算，它采用了额外的3×3深度卷积（编码像素之间的局部交互），并且在卷积之间采用了更多的激活和归一化。其次，为了使用Transformer提取多尺度特征，可以将平均池化（步长为2）应用于自注意力层之前的输入特征。然而，汇集操作降低了自注意力的表现能力。ATtenten（MOA-T）有效地结合了MBConv和Transformer。所提出的MOA-T通过首先用MBConv块替换其MLP，然后反转注意力顺序和MBConv来修改Transformer。用MBConv替换MLP为网络带来了更多的表示能力，并且颠倒顺序（MBConv先于自注意）将下采样任务委派给MBConv内的跨深度卷积，从而学习更好的下采样内核。通过堆叠和增加网络块的通道，进一步开发了一系列MOA-T模型。
此外，通过简单地将全局注意力转换为非重叠的局部窗口注意力，MOA-T可以无缝地部署到需要大分辨率输入的下游任务。由于MBConv能够有效地在像素之间交换局部信息（支持跨窗口传播），MOA-T不需要额外的窗口移动机制。
为了探索MOA T模型的可扩展性，只需通过减少通道大小（无任何其他变化）来缩小模型，从而形成了tiny MOA-T系列，其性能也出人意料地优于基于mob Transformer的模型，如mobile Former和MobileViT。

二、实现细节

(一)、MBConv block

移动卷积（MBConv）（下图（a））也称为倒置残差块，是一种有效的构造块，已广泛用于移动模型。与ResNet中的瓶颈块不同，MBConv块采用了“反向瓶颈”的设计，以及有效的深度卷积。具体而言，首先应用1×1卷积将输入通道扩展4倍。然后，使用3×3深度卷积有效地捕获像素之间的局部空间交互。最后，通过1×1卷积将特征投影回原始通道大小，从而实现残差连接。可选的挤压和激励（SE）模块（其使用全局信息来重新加权通道激活）也可以在深度卷积之后使用。可以调整通道扩展比和深度卷积内核大小以获得更好的性能。

形式上，给定输入张量$x\in R^{H\times W\times C}$（H、W、C是其高度、宽度和通道），MBConv块表示如下：

(二)、Transformer block

Transformer块是一个强大的构建块，通过依赖数据的自注意力操作有效地捕获全局信息。它由两个主要操作组成：自注意和MLP。自注意操作基于输入张量中每对像素之间的成对相似性来计算注意图，从而使模型的感受域涵盖整个空间域。此外， 注意力图动态地依赖于输入，从而扩大了模型的表示能力（与卷积核不同，卷积核是独立于数据的）。MLP运算包含两个1×1卷积，其中第一个卷积扩展通道（4×），第二个卷积收缩通道，并且在其间使用GeLU非线性。

形式上，给定输入张量$x\in R^{H\times W\times C}$，Transformer表示如下：

(三)、Comparing MBConv and Transformer blocks

首先，MBConv和Transformer都提倡“反向瓶颈”设计，其中输入张量的通道被扩展，然后被1×1卷积投影。然而，MBConv在这两个1×1卷积之间另外采用了3×3深度卷积，并且卷积之间既有批量归一化，也有GeLU激活。
第二，为了捕获全局信息，MBConv块可以采用挤压和激励（SE）模块，而Transformer采用自注意力操作。注意，SE模块通过全局平均池化压缩空间信息，而自注意力模块保持张量的空间分辨率。
第三，在块内的不同位置执行下采样操作。为了对特征进行降采样，标准MBConv块使用跨深度卷积，而部署在现代混合模型CoAtNet中的Transformer块在自注意力之前采用平均池化操作。

(四)、MOAT block

本文尝试设计一个新的块，有效地合并MBConv和Transformer块中的最佳块。从强大的Transformer块开始，并逐步完善它。
根据第一次观察，MBConv和Transformer模块均采用“反向瓶颈”设计。由于深度卷积可以有效地编码像素之间的局部交互，这对于在ConvNets中建模平移不变至关重要，因此开始将深度卷积添加到Transformer的MLP模块中。
对于第二个观察，不将SE模块添加到MBConv块中。保持自注意力操作以捕获全局信息。
第三个观察结果至关重要。就在Transformer块中的自注意力操作之前的下采样操作（平均池化）略微降低了其表示能力。另一方面，MBConv块被设计用于具有跨步深度卷积的下采样操作，这有效地学习了每个输入通道的下采样卷积核。因此，在自注意力操作之前进一步重新排序“反向瓶颈”（包含深度卷积），将下采样操作委托给深度卷积。这样，就不需要额外的下采样层，比如CoAtNet中的平均池化，或者Swin和ConvNeXt中的补丁嵌入层。如下图所示。

形式上，给定输入张量$x\in R^{H\times W\times C}$，MOA-T块表示如下：

其中Attn表示自注意操作。MOA-T块简单地由MBConv和自注意操作组成。

(五)、Macro-level network design

在开发MOA-T块之后，采用与现有工程相同的策略：
1、在早期阶段使用卷积提高了Transformer模型的性能和训练收敛性；
2、移动卷积（MBConv）块也是混合卷积Transformer模型中的有效构建块；
3、提取多尺度主干特征有利于下游任务，如检测和分割
因此，基础架构由卷积主干、MBConv块和MOA-T块组成。此外，通过消融研究，我们发现CoAtNet-1提出的层布局有效。因此，遵循它们的层布局，形成了基础模型MOA-T-1。为了形成MOA-T模型族，在块数和通道数的维度上缩小或放大MOA-T-1，如下表所示。只在第三和第四阶段（五个阶段中）缩放块数。在每个级的第一块中执行下采样操作。注意，基础模型MOAT-1和CoAtNet-1共享相同的层布局和通道大小。然而，本文采取了不同的缩放策略：MOA-T通过交替增加深度和扩展变体之间的宽度来放大（或缩小）。

三、实验部分

(一)、ImageNet Image Classification.

ImageNet-1K上的性能。仅1K：仅使用ImageNet-1K。22K+1K:ImageNet-22K预训练和ImageNet-1K微调。

（二）、COCO Detection

COCO 2017验证集上的对象检测和实例分割。采用级联Mask-RCNN和单尺度推断（硬NMS）。†：使用ImageNet-22K预训练权重。

（三）、ADE20K Semantic Segmentation

四、消融实验

MOAT设计：从Transformer块开始，由Attn（自注意力）和MLP组成，它已经达到了很强的精度（82.6%）。直接在MLP中插入3×3深度卷积会使性能降低0.9%。如果在卷积之间额外插入批量归一化和GeLU（即，用MBConv替换MLP，但不使用挤压和激励），性能将提高到82.9%。最后，将MBConv置于Attn之前达到83.3%的性能。此外，MOA-T区块在微小模型机制中带来了更多的改进（从1.2%提高到2.6%）。

Order of MBConv and Attn in MOAT block：MOA-T块设计颠倒了注意力（Attn）和移动卷积（MBConv）的顺序，将下采样任务委托给MBConv内的跨步深度卷积。然而，下采样仍然可以以原始顺序（即Attn+MBConv）在MBConv中执行。由于操作Attn和MBConv是交错的，因此关键区别来自每个阶段的第一个块，其中Attn是在（1）空间下采样（2）通道扩展特征上操作的。为了进行研究，在MOA-T变体中使用了不同的块，使用“Attn+MLP”、“Attn+MBConv”或“MBConv+Attn”。对于“Attn+MBConv”块，我们进一步研究了该位置（Attn vs.MBConv），在那里应用了空间下采样和通道扩展操作。
在下表中，用MBConv替换MLP将MOA T-0和tiny MOA T-2的性能提高了0.3%和0.7%。其次，如果在MBConv块上同时执行空间下采样和通道道扩展，则MOA T-1和微小MOA-T-2的的性能进一步提高了0.5%和0.9%，这表明MBConv学习到了更好的下采样特征。然而，这种设计相当于将第一个Attn层移到其前一阶段，从而降低了当前阶段的表示能力。更具体地说，只有最后一个阶段会受到影响，因为一个层被移动了。第三，为了增强表示能力，颠倒Attn和MBConv的顺序允许我们将第一个Attn层保持在同一阶段。该设计进一步将MOA T-0和微小MOA T-2的性能提高了0.7%和1.2%。第四，为了补偿移位效应，还可以使用另一个1×1卷积来扩展第一Attn层的通道（然后，MBConv只执行空间下采样）。然而，这种设计与MOA-T块设计类似，但使用了更多的参数和FLOP。
Downsampling layer：对于MOA-T块设计，讨论了MBConv和Attention的重新排序。因此，不需要CoAtNet中的下采样层，也不需要Swin和ConvNeXt中的补丁嵌入层（即2×2带步长2的卷积）中的平均池化。如下表所示，使用补丁嵌入层确实比平均池化方案提高了0.2%的精度，但它需要更多的模型参数成本。另一方面，MOA-T设计（即，将下采样委托给MBConv块）以最小的参数成本和可比的FLOP显示了最佳性能。

五、附录

在MOTA网络中，对卷积和深度卷积都采用核大小3。使用多头自注意力，其中每个注意力头具有通道32。对于MBConv和MOA-T块，使用扩展比率4。MBConv块（即，第二和第三阶段）中的SE模块采用缩减比率4（相对于输入通道）。MOA-T块包括ImageNet的相对位置嵌入。然而，下游任务通常采用比ImageNet更大的输入分辨率，需要特殊的适应（例如，预训练位置嵌入的双线性插值）。为了简单起见，当在下游任务上运行MOAT时，删除了位置嵌入。

图像分类的实验设计：在ImageNet-1K上训练MOA-T模型，分辨率为224，持续300个轮次。如果在较大的ImageNet-22K上进行预训练，使用分辨率224和90轮次。之后，在ImageNet-1K上对模型进行了30个时期的微调。在微调期间，还尝试使用更大的分辨率（例如384和512）。在训练期间使用了典型的正则化方法，例如标签平滑、RandAugment、MixUp、随机深度和具有解耦权重衰减的Adam（即AdamW）。详细的超参数见如下两个表。所有实验都在TPV4上进行。MOA T-{0,1,2}用16个核训练，而MOA T-3用32个核训练。在表12中报告了MOA T模型变体的训练时间。


目标检测以及语义分割的实验设置：使用MOA T架构在COCO 2017数据集上训练Cascade Mask R-CNN。数据集包含118K个训练样本和5K个验证样本为了使MOA-T模型适应高分辨率输入，将特征划分为非重叠窗口，用于自注意力计算，其中第二个最后阶段的窗口大小设置为14，并在最后阶段利用全局注意力。由于此窗口分区，输入大小必须可被14整除。TF Vision Model Garden代码库还要求输入大小为正方形（带填充）并可被64整除。因此，选择1344作为输入大小，与基线方法中使用的大小类似（即，最长边不超过1333）。使用特征金字塔网络来整合不同层次的特征。

消融实验附录：