视觉变换器和 MLP 模型的多轴方法

自 2012 年推出AlexNet以来，卷积神经网络一直是计算机视觉领域的主要机器学习架构。最近，受到自然语言处理中Transformers演变的启发，注意力机制已被广泛纳入视觉模型中。这些注意力机制会增强输入数据的某些部分，同时最小化其他部分，以便网络可以专注于数据中虽小但很重要的部分。Vision Transformer (ViT) 为完全摆脱卷积的计算机视觉模型设计开辟了新局面。ViT 将图像块视为一系列单词，并在其上应用 Transformer 编码器。在足够大的数据集上进行训练后，ViT 在图像识别方面表现出色。

虽然卷积和注意力机制都足以实现良好的性能，但它们都不是必需的。例如，MLP-Mixer采用简单的多层感知器(MLP) 在所有空间位置混合图像块，从而形成全 MLP 架构。在训练和推理所需的准确度和计算量之间进行权衡方面，它是现有最先进视觉模型的有竞争力的替代方案。然而，ViT 和 MLP 模型都难以扩展到更高的输入分辨率，因为计算复杂度会随着图像大小的增加而呈二次方增加。

今天，我们介绍了一种简单有效的新多轴方法，改进了原有的 ViT 和 MLP 模型，可以更好地适应高分辨率、密集预测任务，并且可以自然地适应不同的输入大小，具有很高的灵活性和较低的复杂度。基于这种方法，我们为高级和低级视觉任务构建了两个主干模型。我们在ECCV 2022中展示的“ MaxViT：多轴视觉变换器”中描述了第一个模型，并表明它显著提高了图像分类、对象检测、分割、质量评估和生成等高级任务的最新水平。第二个模型在CVPR 2022的“ MAXIM：用于图像处理的多轴 MLP ”中提出，它基于类似UNet的架构，在去噪、去模糊、去雾、去雨和低光增强等低级成像任务上取得了具有竞争力的性能。为了促进对高效 Transformer 和 MLP 模型的进一步研究，我们开源了MaxViT和MAXIM的代码和模型。

使用 MAXIM 逐帧进行图像去模糊的演示。

概述

我们的新方法基于多轴注意力机制，它将 ViT 中使用的全尺寸注意力机制（每个像素关注所有像素）分解为两种稀疏形式——局部和（稀疏）全局。如下图所示，多轴注意力机制包含块注意力机制和网格注意力机制的顺序堆栈。块注意力机制在非重叠窗口（中间特征图中的小块）内工作以捕获局部模式，而网格注意力机制在稀疏采样的均匀网格上工作以进行长距离（全局）交互。网格注意力机制和块注意力机制的窗口大小可以完全作为超参数进行控制，以确保计算复杂度与输入大小呈线性关系。

所提出的多轴注意力机制依次引导局部阻塞注意力和扩张全局注意力，然后是FFN，复杂度仅为线性。相同颜色的像素一起被关注。

这种低复杂度的注意力机制可以显著提高其对许多视觉任务的广泛适用性，尤其是对于高分辨率视觉预测，比 ViT 中使用的原始注意力机制具有更高的通用性。我们利用这种多轴注意力机制构建了两个骨干实例——MaxViT 和 MAXIM，分别用于高级任务和低级任务。

最大ViT

在 MaxViT 中，我们首先通过将 MBConv（由EfficientNet，V2提出）与多轴注意力 连接起来，构建单个 MaxViT 块（如下所示） 。无论输入分辨率如何，这个单个块都可以编码局部和全局视觉信息。然后，我们简单地将由注意力和卷积组成的重复块堆叠在分层架构中（类似于ResNet、CoAtNet），从而产生我们的同质 MaxViT 架构。值得注意的是，MaxViT 与以前的分层方法不同，因为它可以在整个网络中全局“查看”，即使在早期的高分辨率阶段也是如此，从而在各种任务上展示出更强大的模型容量。

MaxViT 的元架构。

格言

我们的第二个主干网络MAXIM是一种通用的类似UNet的架构，专为低级图像到图像预测任务而量身定制。MAXIM 探索使用门控多层感知器(gMLP) 网络（带有门控机制的修补混合 MLP ）对局部和全局方法进行并行设计。MAXIM 的另一个贡献是交叉门控块，可用于应用两个不同输入信号之间的相互作用。该块可以作为交叉注意模块的有效替代方案，因为它仅使用廉价的门控 MLP 运算符与各种输入进行交互，而不依赖于计算繁重的交叉注意。此外，MAXIM 中所有提出的组件（包括门控 MLP 和交叉门控块）都具有与图像大小成线性关系的复杂度，这使得它在处理高分辨率图片时更加高效。

结果

我们展示了 MaxViT 在各种视觉任务上的有效性。在图像分类方面，MaxViT 在各种设置下都取得了最佳效果：仅使用ImageNet-1K训练，MaxViT 就达到了 86.5% 的 top-1 准确率；使用ImageNet-21K（14M 图像，21k 类）预训练，MaxViT 实现了 88.7% 的 top-1 准确率；使用JFT（300M 图像，18k 类）预训练，我们最大的模型 MaxViT-XL 以 475M 参数实现了 89.5% 的高准确率。

MaxViT 与 ImageNet-1K 上最先进模型的性能比较。顶部：224x224 图像分辨率下的准确度与 FLOP 性能缩放曲线。底部：ImageNet-1K 微调设置下的准确度与参数缩放曲线。

对于下游任务，MaxViT 作为主干在广泛的任务中提供了良好的性能。对于COCO数据集上的对象检测和分割，MaxViT 主干实现了 53.4 AP ，优于其他基础级模型，同时仅需要约 60% 的计算成本。对于图像美学评估，MaxViT 模型在与人类意见分数的线性相关性方面比最先进的MUSIQ模型提高了 3.5%。独立的 MaxViT 构建块在图像生成方面也表现出有效的性能，在 ImageNet-1K 无条件生成任务上获得了更好的FID和IS分数，并且参数数量明显少于最先进的模型HiT。

与 UNet 类似的 MAXIM 主干是为图像处理任务定制的，在 20 个测试数据集中的 15 个数据集上也展示了最先进的结果，包括去噪、去模糊、去雨、去雾和低光增强，同时所需的参数和 FLOP 数量比竞争模型更少或相当。MAXIM 恢复的图像显示更多恢复的细节，视觉伪影更少。

MAXIM 对图像去模糊、去雨和低光增强的视觉结果。

概括

近两年来的研究显示，ConvNets 和 Vision Transformers 可以实现类似的性能。我们的工作提出了一种统一的设计，充分利用了两者的优点——高效卷积和稀疏注意力，并证明了在此基础上构建的模型 MaxViT 可以在各种视觉任务上实现最先进的性能。更重要的是，MaxViT 可以很好地扩展到非常大的数据量。我们还展示了使用 MLP 运算符的替代多轴设计 MAXIM 在广泛的低级视觉任务上实现了最先进的性能。

尽管我们在视觉任务的背景下展示了我们的模型，但所提出的多轴方法可以轻松扩展到语言建模，以在线性时间内捕获局部和全局依赖关系。受此工作的启发，我们预计值得研究高维或多模态信号（如视频、点云和视觉语言模型）中的其他形式的稀疏注意力。

我们已经开源了MAXIM和MaxViT的代码和模型，以方便未来对有效注意力和 MLP 模型的研究。

致谢

我们要感谢我们的合著者：Hossein Talebi、Han Zhang、Feng Yang、Peyman Milanfar 和 Alan Bovik。我们还要感谢 Xianzhi Du、Long Zhao、Wuyang Chen、Hanxiao Liu、Zihang Dai、Anurag Arnab、Sungjoon Choi、Junjie Ke、Mauricio Delbracio、Irene Zhu、Innfarn Yoo、Huiwen Chang 和 Ce Liu 的宝贵讨论和支持。