本文提出了一种高效、可扩展的多轴注意力模型,该模型包括两个方面:局部注意力被阻塞和全局注意力被扩张。这些设计选择允许在任意输入分辨率下的全局-局部空间交互,只有线性复杂度。还通过有效地将注意力模型与卷积混合在一起,提出了一个新的架构元素,并相应地提出了一个简单的分层视觉主干,称为MaxViT,通过简单地在多个阶段重复基本构建块。值得注意的是,即使在早期的高分辨率阶段,MaxViT也能够“看到”整个网络的全局。
受自然语言处理中Transformer[85]等自注意模型进化的启发,许多研究人员开始将注意机制引入视觉。Vision Transformer (ViT)可能是第一个完全基于Transformer的视觉架构,其中图像补丁被简单地视为单词序列,并且在这些视觉token上应用了一个Transformer编码器。当在大规模数据集上进行预训练时,ViT可以在图像识别上取得令人信服的结果。
如果没有广泛的预训练,ViT在图像识别方面表现不佳。这是由于Transformer具有较强的模型能力,具有较少的感应偏置,从而导致过拟合。为了适当地正则化模型容量并提高其可扩展性,许多后续工作研究了为局部注意力等视觉任务量身定制的稀疏Transformer模型。这些方法通常重新引入层次结构来弥补非局部性的损失。
Swin Transformer通过在移位的非重叠窗口上应用自注意力来修改Transformer。该方法首次在纯视觉Transformer的ImageNet基准测试中优于ConvNets。尽管与ViT中使用的全注意力相比,基于窗口的注意力具有更大的灵活性和泛化性,但由于非局域性的丧失,基于窗口的注意力的模型容量有限,因此在ImageNet-21K和JFT等更大的数据体系上扩展不利。然而,在层次网络的早期或高分辨率阶段通过全注意力获取全局交互是计算量很大的,因为注意力算子需要二次复杂度。在一定的计算预算下,如何有效地结合全局和局部相互作用来平衡模型的能力和泛化性是一个挑战。
1. 引言
提出了一种新型的Transformer模块,称为多轴自关注(Max-SA),它能够作为一个基本的架构组件,在单个块中执行局部和全局空间交互。与完全自注意力相比,Max-SA具有更大的灵活性和效率,即对线性复杂度的不同输入长度具有自然适应性; 与(移位的)窗口/局部注意相比,Max-SA通过提出一个全局接受域来允许更强的模型容量。此外,Max-SA仅具有线性复杂性,可以在网络的任何层中用作通用的独立注意力模块,即使在早期的高分辨率阶段也是如此。
进一步设计了一种简单而有效的视觉骨干,称为多轴视觉变压器(MaxViT),通过分层堆叠由Max-SA和卷积组成的重复块。虽然提出的模型属于混合Vision Transformer的范畴,但MaxViT与以前的方法不同,因为力求简单,通过设计一个统一卷积、局部和全局注意力的基本块,然后简单地重复它。实验表明,在所有数据体制下,MaxViT显著提高了最先进的(SOTA)性能,适用于广泛的视觉任务,包括分类、目标检测和分割、图像美学评估和图像生成。
2. 相关工作
2.3 混合模型
纯基于Transformer的视觉模型由于相对较少的归纳偏差而泛化较差。Vision Transformer 也表现出不合格的可优化性。一个有趣的简单改进是采用Transformer层和卷积层的混合设计,例如使用几个卷积来替换粗糙的patchify stem(将图像分割成一系列小块的步骤)。
2.4 GANs Transformer
Transformer 在生成对抗网络(GANs)中也被证明是有效的[26]。TransGAN构建了一个纯Transformer GAN,精心设计了局部注意力层和上采样层,在小规模数据集上显示了有效性。GANformer则探索了有效的全局注意力机制来改进StyleGAN generator。
HiT提出了一种基于局部-全局注意力的高效Transformer生成器,可以扩展到1K高分辨率图像生成。
3. 方法
Blocked Multi-Axis Self-Attention (Max-SA)
受到稀疏方法的启发,研究者们提出了一种新的注意力模块,称为“Blocked Multi-Axis Self-Attention”(Max-SA)。这个模块通过将全密集(fully dense)的注意力机制分解为两种稀疏形式(窗口注意力和网格注意力)来降低计算复杂度。
稀疏形式的优点
- 降低复杂度:传统的注意力机制(如Transformer中的自注意力)具有二次复杂度(即O(n^2)),其中n是序列长度。通过将注意力机制分解为稀疏形式,Max-SA能够将复杂度降低到线性级别,而不会损失非局部性(non-locality)的特性。
- 非局部性:非局部性意味着每个位置都可以与序列中的任何其他位置进行交互。Max-SA通过其稀疏形式的设计,仍然能够保持这种非局部性。
顺序设计的优势
- 简单性:Max-SA的设计是顺序的,这意味着每个模块可以按照特定的顺序依次应用。
- 灵活性:每个模块既可以单独使用,也可以以任何顺序组合,这为模型设计提供了很大的灵活性。
- 性能:尽管设计简单且灵活,但Max-SA的性能甚至超过了以前的方法。
MaxViT的构建
- 层次结构:MaxViT通过将Max-SA模块与MBConv(MobileNetV2中的深度可分离卷积变体)交替堆叠在一个层次结构中,形成了一种新颖的模型结构。
- 全局和局部感受野:在整个网络中,从浅层到深层,MaxViT都能够利用全局和局部的感受野,这使得模型在模型容量和泛化能力方面均表现出色。
引入了一种新型的注意模块,称为阻塞多轴自注意(Max-SA),通过将全密集的注意机制分解为两种稀疏形式{窗口注意力和网格注意力},将普通注意的二次复杂度降低到线性,而不会损失任何非局域性。
全密集(fully dense)的注意力机制指的是在注意力计算过程中,对于输入序列中的每一个元素(例如,在自然语言处理中的每个单词或在计算机视觉中的每个像素),都会与其他所有元素计算注意力权重。这种机制不区分哪些元素是更重要的,而是对所有元素都给予相同的关注。
与传统的全连接层类似,