分享自:

MaxViT:多轴视觉变换器

期刊:ECCVDOI:10.1007/978-3-031-20053-3_27

本文档为Zhengzhong Tu等多位研究者发表于2022年欧洲计算机视觉会议(ECCV)的研究论文,标题为“MaxViT: Multi-Axis Vision Transformer”。该工作提出了一种新型的视觉Transformer骨干网络,旨在高效地融合局部与全局视觉信息。

第一作者及研究机构 本项研究的主要第一作者为Zhengzhong Tu,其隶属于美国德克萨斯大学奥斯汀分校,并同时与Google Research团队合作。研究团队其他成员包括来自Google Research的Hossein Talebi, Han Zhang, Feng Yang, Peyman Milanfar, Yinxiao Li,以及来自德克萨斯大学奥斯汀分校的Alan Bovik。论文发表于ECCV 2022会议,并被收录于其会议论文集LNCS系列中。

学术背景与研究动机 本研究属于计算机视觉领域,具体聚焦于图像识别骨干网络的架构设计。近年来,Vision Transformer(ViT)模型在多个视觉任务上展现出强大潜力,但其核心的自注意力(Self-Attention)机制存在计算复杂度随图像尺寸呈二次方增长的问题,这限制了其在处理高分辨率图像时的可扩展性。虽然已有工作(如Swin Transformer)通过引入局部窗口注意力来降低计算量,但这种做法牺牲了模型在浅层(即高分辨率阶段)捕获全局上下文信息的能力,可能导致模型容量受限,特别是在大规模数据上训练时。因此,如何在保持线性计算复杂度的前提下,设计一个既能捕获局部细节又能理解全局结构的视觉骨干网络,成为了一个关键挑战。本研究的目标正是为了解决这一矛盾,提出一种高效的、可扩展的多轴自注意力机制,并以此为基础构建一个强大的通用视觉骨干网络MaxViT。

详细研究流程与方法 本研究主要包含核心模块设计、架构构建、实验验证与消融研究四大环节。

  1. 核心模块:多轴自注意力(MAX-SA)设计 这是本研究的核心创新点。研究团队提出将标准的全局密集自注意力分解为两种稀疏形式:块注意力(Block Attention)和网格注意力(Grid Attention)。

    • 块注意力(Block Attention):将输入特征图分割成一系列非重叠的、固定大小(如7x7)的窗口(Windows)。自注意力计算仅在每个窗口内部进行,从而实现了局部空间交互。这与Swin Transformer的窗口注意力类似,计算复杂度相对于图像大小是线性的。
    • 网格注意力(Grid Attention):为了实现全局交互,研究者提出了一个新颖的网格划分方法。不同于固定大小的窗口,网格注意力将特征图划分为一个固定数量的、均匀的网格(如7x7个网格单元)。每个网格单元覆盖的特征区域大小会自适应于输入图像的分辨率。随后,自注意力计算在网格维度上进行,即关注于分布在图像不同位置的、稀疏的网格单元。这一操作等效于一种空洞的、全局的空间混合。关键在于,通过将网格大小设置为与块注意力窗口大小相同(例如均为7),可以均衡局部与全局操作的计算负担,两者均保持线性复杂度。 研究者将这两种注意力机制顺序堆叠,构成了多轴自注意力(MAX-SA)模块。这种设计使得单个模块内部即可完成从局部到全局的信息融合,且无需像Swin那样进行复杂的窗口移位(Shifted Windows)操作,实现更为简便。
  2. 架构构建:MaxViT块与整体网络 研究团队进一步将MAX-SA模块与卷积操作有效融合,提出了统一的MaxViT基本构建块(MaxViT Block)。一个MaxViT块通常包含以下几个顺序执行的子层:

    • MBConv层:一个移动倒置瓶颈卷积(Mobile Inverted Bottleneck Convolution)块,包含深度可分离卷积和挤压激励(Squeeze-and-Excitation)模块。该层放置在注意力层之前,研究者发现这能提升模型的泛化能力和可训练性。深度卷积可被视为一种条件位置编码(Conditional Position Encoding, CPE),因此模型无需显式的位置编码层。
    • 块注意力层:执行局部窗口内的自注意力。
    • 网格注意力层:执行跨全局网格的自注意力。
    • 前馈网络:标准的Transformer前馈网络,用于非线性变换。 整个MaxViT网络遵循经典的卷积网络分层设计(类似ResNet),包含一个初始的卷积干(Stem)和四个主要阶段(Stage)。每个阶段由多个重复的MaxViT块堆叠而成,随着网络加深,空间分辨率逐步减半,通道数逐步增加。研究者构建了从Tiny(MaxViT-T)到Extra Large(MaxViT-XL)一系列不同规模与深度的模型变体,以评估其可扩展性。整个设计追求简洁性:仅通过重复相同的基本块来构建强大的分层视觉骨干。
  3. 实验验证与性能评估 研究团队在广泛的视觉任务上验证了MaxViT的有效性,主要实验流程如下:

    • 图像分类(ImageNet-1K/21K, JFT-300M):这是核心评估任务。首先在ImageNet-1K数据集上从头训练不同规模的MaxViT模型,输入分辨率包括224x224,并在更高分辨率(384,512)上进行微调。同时,也在更大规模的数据集(ImageNet-21K和谷歌内部的JFT-300M)上进行预训练,然后在ImageNet-1K上微调,以检验其在大数据下的可扩展性。评估指标为Top-1准确率,并与当前最先进的卷积网络(如ConvNeXt, EfficientNetV2)、纯Transformer(如ViT, Swin)以及混合模型(如CoAtNet)进行对比。
    • 下游任务迁移
      • 目标检测与实例分割:在COCO 2017数据集上,将MaxViT作为骨干网络嵌入到两阶段检测框架(如Mask R-CNN with FPN)中。评估指标包括平均精度(AP, AP50, AP75)。
      • 图像美学评估:在AVA数据集上,训练MaxViT模型预测图像的美学分数,评估指标为皮尔逊线性相关系数(PLCC)和斯皮尔曼秩相关系数(SRCC)。
      • 图像生成:在ImageNet-1K上进行无条件图像生成任务,使用生成对抗网络(GAN)框架,其中生成器由MaxViT块构建。评估指标为弗雷歇初始距离(FID)和初始分数(IS)。
  4. 消融研究与分析 为了验证设计选择的有效性,研究者进行了详细的消融实验(主要基于MaxViT-T模型在ImageNet-1K上的表现):

    • 全局网格注意力的作用:通过移除或替换(用块注意力替代)不同阶段的网格注意力,证明了全局交互在网络的各个阶段(尤其是早期高分辨率阶段)都能带来性能提升。
    • MBConv层的作用:移除MBConv层会导致性能显著下降,证实了卷积层对于提供局部归纳偏置和提升模型训练稳定性的重要性。
    • 块内顺序研究:比较了MBConv(C)、块注意力(BA)、网格注意力(GA)六种排列组合。发现C-BA-GA(局部到全局)的顺序在分类任务上表现最佳,而GA-BA-C(全局到局部)的顺序在图像生成任务上更优,表明任务特性影响最优的信息处理流程。
    • 顺序vs并行设计:将MAX-SA的顺序堆叠设计与并行执行块和网格注意力的方案对比,结果表明顺序设计在性能和计算效率上均显著优于并行方案。
    • 垂直布局:研究了各阶段块数量的分配策略,发现MaxViT采用的布局在大规模模型上比Swin的布局缩放性更好。

主要研究结果 1. 图像分类SOTA性能:MaxViT在ImageNet-1K上取得了当时领先的准确率。例如,MaxViT-L在224x224分辨率下达到85.17%的Top-1准确率;在512x512分辨率下,MaxViT-L达到86.7%的准确率。更重要的是,在各种计算量(FLOPs)和参数量水平上,MaxViT均优于或媲美同时期的先进模型,如图1所示的缩放曲线所示,显示了其优异的效率。在更大规模数据预训练(ImageNet-21K和JFT)后,MaxViT继续表现出强大的可扩展性,例如MaxViT-XL在JFT预训练后达到89.53%的准确率。 2. 下游任务卓越表现: * 目标检测与分割:在COCO数据集上,MaxViT骨干网络带来了显著的性能提升。例如,MaxViT-S在AP指标上超越了参数量更大的Swin-B和ConvNeXt-B等模型,且计算成本更低。 * 图像美学评估:MaxViT-T在AVA数据集上取得了优于或接近当时最佳方法(如MUSIQ)的相关性分数,展现了其在感知任务上的潜力。 * 图像生成:使用MaxViT块构建的生成器在ImageNet-1K的128x128图像生成任务上,取得了优于专门为生成任务设计的Transformer模型(如HIT)的FID和IS分数,且参数更少,证明了MaxViT模块在生成建模方面的通用性。 3. 消融实验结果:所有消融研究结果均有力地支持了MaxViT的核心设计决策:全局网格注意力、MBConv的整合、C-BA-GA的顺序以及分层堆叠架构的有效性。这些结果为理解模型为何有效提供了实证依据。

研究结论与价值 本研究的结论是,MaxViT通过其创新的多轴自注意力机制,成功地统一了卷积的局部归纳偏置与自注意力的全局建模能力,创造了一个既高效又强大的通用视觉骨干网络。其科学价值在于: 1. 方法论创新:提出了一种新颖的、计算高效的全局-局部注意力分解方法(块注意力与网格注意力),为解决Transformer在视觉任务中计算复杂度与模型容量之间的矛盾提供了新思路。 2. 架构设计贡献:展示了通过简单地重复一个精心设计的、融合了卷积与多轴注意力的基本块,即可构建出性能卓越的分层视觉骨干,这为未来的架构设计提供了简洁而有效的范例。 3. 实证性能领先:在图像分类、目标检测、图像美学评估和图像生成等多个核心视觉任务上,MaxViT均取得了当时最先进或极具竞争力的性能,证明了其作为“通用视觉模块”的潜力。 其应用价值在于,MaxViT模型可以作为一种强大的现成骨干网络,直接应用于各种需要高性能视觉特征提取的下游任务中,如自动驾驶、内容理解、图像编辑等。

研究亮点 1. 核心创新突出:多轴自注意力(MAX-SA)是本文最关键的创新点,其“网格注意力”的设计巧妙地在线性复杂度下实现了全局感知,是区别于以往局部窗口注意力和轴向注意力的核心所在。 2. 简洁统一的设计哲学:整个MaxViT架构追求极致的简洁性,仅通过重复相同的MaxViT块来构建,避免了复杂的定制化模块,体现了“大道至简”的设计理念。 3. 广泛的实验验证:研究不仅在标准图像分类上证明其优越性,更在检测、分割、美学评估、生成等多个差异巨大的任务上进行了全面验证,充分证明了模型的通用性和鲁棒性。 4. 深入的消融分析:论文提供了详尽的消融实验,系统地剖析了每个设计组件的作用、模块内的最佳顺序以及与其他设计选择(如并行)的对比,使得研究结论坚实可靠。

其他有价值的补充 论文最后还讨论了MaxViT在语言建模、视频、点云等多模态信号处理上的潜在扩展性,并简要提及了大规模模型训练可能带来的环境(碳排放)和社会(偏见、虚假信息生成)影响,体现了研究者的社会责任意识。源代码和预训练模型已在GitHub上公开,促进了研究的可复现性和后续发展。

上述解读依据用户上传的学术文献,如有不准确或可能侵权之处请联系本站站长:admin@fmread.com